egress-proxies-operator/internal/provider/kvm/image.go

package kvm

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	libvirt "github.com/digitalocean/go-libvirt"
	"github.com/go-logr/logr"

	"gitea.home.hrajfrisbee.cz/kacerr/egress-proxies-operator/internal/provider"
)

// imageSource is a resolved placement.image: either an entry from the
// provider's images map (Key set) or a URL given directly.
type imageSource struct {
	Key    string
	URL    string
	SHA256 string
	Format string
}

// resolveImage maps placement.image to a source. Anything that is neither a
// configured key nor an http(s) URL is a spec error the reconciler should
// surface, not retry.
func resolveImage(cfg provider.KVMConfig, image string) (imageSource, error) {
	if img, ok := cfg.Images[image]; ok {
		return imageSource{Key: image, URL: img.URL, SHA256: img.SHA256, Format: img.Format}, nil
	}
	u, err := url.Parse(image)
	if err == nil && (u.Scheme == "http" || u.Scheme == "https") && u.Host != "" {
		return imageSource{URL: image, Format: formatQcow2}, nil
	}
	return imageSource{}, fmt.Errorf("%w: placement.image %q is neither a key in kvm.images nor an http(s) URL",
		provider.ErrPermanent, image)
}

// volumeName is the cached base image's name in the pool: readable for
// configured keys, a URL hash for ad-hoc URLs.
func (s imageSource) volumeName() string {
	if s.Key != "" {
		return "base-" + s.Key + "." + s.Format
	}
	sum := sha256.Sum256([]byte(s.URL))
	return "base-" + hex.EncodeToString(sum[:8]) + "." + s.Format
}

// markerName is the sidecar volume holding the sha256 of the verified
// upload. Its presence is what says "the base volume is complete".
func (s imageSource) markerName() string {
	return s.volumeName() + ".sha256"
}

// imageCache fetches base images into a host's pool on first use. It keeps
// no state about what is cached — the pool is the source of truth, checked
// on every call — only a per-(host, image) lock so concurrent Creates for
// the same image on the same host result in one download.
type imageCache struct {
	http *http.Client

	mu    sync.Mutex
	locks map[string]*sync.Mutex
}

func newImageCache(client *http.Client) *imageCache {
	if client == nil {
		client = http.DefaultClient
	}
	return &imageCache{http: client, locks: make(map[string]*sync.Mutex)}
}

func (c *imageCache) lockFor(key string) *sync.Mutex {
	c.mu.Lock()
	defer c.mu.Unlock()
	m, ok := c.locks[key]
	if !ok {
		m = &sync.Mutex{}
		c.locks[key] = m
	}
	return m
}

// ensure makes src present in pool and returns the base volume's path on
// the host. The image is downloaded synchronously through the libvirt
// stream when the pool does not already hold a verified copy.
func (c *imageCache) ensure(ctx context.Context, log logr.Logger, api libvirtAPI, hostName string, pool libvirt.StoragePool, src imageSource) (string, error) {
	m := c.lockFor(hostName + "/" + src.volumeName())
	m.Lock()
	defer m.Unlock()

	base, ready, err := c.lookup(api, pool, src)
	if err != nil {
		return "", err
	}
	if ready {
		log.V(1).Info("base image already cached", "volume", src.volumeName())
		return api.StorageVolGetPath(base)
	}
	if err := deleteIfExists(api, pool, src.markerName()); err != nil {
		return "", err
	}
	if err := deleteIfExists(api, pool, src.volumeName()); err != nil {
		return "", err
	}
	base, err = c.fetch(ctx, log, api, hostName, pool, src)
	if err != nil {
		return "", err
	}
	return api.StorageVolGetPath(base)
}

// lookup reports whether the pool holds a complete, verified base volume:
// the volume and its marker both exist and, when a checksum is configured,
// the marker records that checksum.
func (c *imageCache) lookup(api libvirtAPI, pool libvirt.StoragePool, src imageSource) (libvirt.StorageVol, bool, error) {
	base, err := api.StorageVolLookupByName(pool, src.volumeName())
	if hasCode(err, libvirt.ErrNoStorageVol) {
		return libvirt.StorageVol{}, false, nil
	}
	if err != nil {
		return libvirt.StorageVol{}, false, fmt.Errorf("looking up base volume: %w", err)
	}
	marker, err := api.StorageVolLookupByName(pool, src.markerName())
	if hasCode(err, libvirt.ErrNoStorageVol) {
		return base, false, nil
	}
	if err != nil {
		return libvirt.StorageVol{}, false, fmt.Errorf("looking up image marker: %w", err)
	}
	if src.SHA256 == "" {
		return base, true, nil
	}
	var buf bytes.Buffer
	if err := api.StorageVolDownload(marker, &buf, 0, 0, 0); err != nil {
		return libvirt.StorageVol{}, false, fmt.Errorf("reading image marker: %w", err)
	}
	return base, strings.EqualFold(strings.TrimSpace(buf.String()), src.SHA256), nil
}

// fetch streams the image from its URL straight into a new base volume,
// hashing on the way, then writes the marker. Nothing is buffered in the
// operator beyond the HTTP body's read buffer.
func (c *imageCache) fetch(ctx context.Context, log logr.Logger, api libvirtAPI, hostName string, pool libvirt.StoragePool, src imageSource) (libvirt.StorageVol, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, src.URL, nil)
	if err != nil {
		return libvirt.StorageVol{}, fmt.Errorf("%w: building image request: %v", provider.ErrPermanent, err)
	}
	resp, err := c.http.Do(req)
	if err != nil {
		return libvirt.StorageVol{}, fmt.Errorf("%w: fetching image %s: %v", provider.ErrTransient, src.URL, err)
	}
	defer resp.Body.Close()
	switch {
	case resp.StatusCode >= 200 && resp.StatusCode < 300:
	case resp.StatusCode >= 400 && resp.StatusCode < 500:
		return libvirt.StorageVol{}, fmt.Errorf("%w: image %s returned HTTP %d", provider.ErrPermanent, src.URL, resp.StatusCode)
	default:
		return libvirt.StorageVol{}, fmt.Errorf("%w: image %s returned HTTP %d", provider.ErrTransient, src.URL, resp.StatusCode)
	}
	if resp.ContentLength < 0 {
		return libvirt.StorageVol{}, fmt.Errorf("%w: image %s has no Content-Length; the base volume needs a known size", provider.ErrPermanent, src.URL)
	}

	log.Info("fetching base image into host pool",
		"pool", pool.Name, "volume", src.volumeName(),
		"url", src.URL, "bytes", resp.ContentLength)
	start := time.Now()

	baseXML, err := buildVolumeXML(volumeSpec{Name: src.volumeName(), SizeBytes: resp.ContentLength, Format: src.Format})
	if err != nil {
		return libvirt.StorageVol{}, err
	}
	base, err := api.StorageVolCreateXML(pool, baseXML, 0)
	if err != nil {
		return libvirt.StorageVol{}, fmt.Errorf("creating base volume: %w", err)
	}
	hasher := sha256.New()
	if err := api.StorageVolUpload(base, io.TeeReader(resp.Body, hasher), 0, 0, 0); err != nil {
		_ = api.StorageVolDelete(base, 0)
		return libvirt.StorageVol{}, fmt.Errorf("uploading image to base volume: %w", err)
	}
	digest := hex.EncodeToString(hasher.Sum(nil))
	if src.SHA256 != "" && !strings.EqualFold(digest, src.SHA256) {
		_ = api.StorageVolDelete(base, 0)
		return libvirt.StorageVol{}, fmt.Errorf("%w: image %s sha256 is %s, config expects %s",
			provider.ErrPermanent, src.URL, digest, src.SHA256)
	}
	// libvirt sized the volume from Content-Length; a refresh makes it
	// re-read the real virtual size from the uploaded qcow2 header.
	if err := api.StoragePoolRefresh(pool, 0); err != nil {
		return libvirt.StorageVol{}, fmt.Errorf("refreshing pool after upload: %w", err)
	}

	markerXML, err := buildVolumeXML(rawVolume(src.markerName(), int64(len(digest))))
	if err != nil {
		return libvirt.StorageVol{}, err
	}
	marker, err := api.StorageVolCreateXML(pool, markerXML, 0)
	if err != nil {
		return libvirt.StorageVol{}, fmt.Errorf("creating image marker: %w", err)
	}
	if err := api.StorageVolUpload(marker, strings.NewReader(digest), 0, 0, 0); err != nil {
		_ = api.StorageVolDelete(marker, 0)
		return libvirt.StorageVol{}, fmt.Errorf("writing image marker: %w", err)
	}
	log.Info("base image cached",
		"volume", src.volumeName(), "sha256", digest,
		"bytes", resp.ContentLength, "duration", time.Since(start).Round(time.Second))
	return base, nil
}

func deleteIfExists(api libvirtAPI, pool libvirt.StoragePool, name string) error {
	vol, err := api.StorageVolLookupByName(pool, name)
	if hasCode(err, libvirt.ErrNoStorageVol) {
		return nil
	}
	if err != nil {
		return fmt.Errorf("looking up volume %s: %w", name, err)
	}
	if err := api.StorageVolDelete(vol, 0); err != nil && !hasCode(err, libvirt.ErrNoStorageVol) {
		return fmt.Errorf("deleting volume %s: %w", name, err)
	}
	return nil
}