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

package kvm

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"io"
	"net/http"
	"net/http/httptest"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/go-logr/logr"

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

var testImage = bytes.Repeat([]byte("QFI\xfb fake qcow2 payload "), 1024)

func testImageSHA() string {
	sum := sha256.Sum256(testImage)
	return hex.EncodeToString(sum[:])
}

// imageServer serves testImage and counts hits. Status overrides the
// response, noLength forces chunked encoding (no Content-Length), delay
// slows each response so concurrent callers overlap.
type imageServer struct {
	*httptest.Server
	hits     atomic.Int32
	status   int
	noLength bool
	delay    time.Duration
}

func newImageServer(t *testing.T) *imageServer {
	t.Helper()
	s := &imageServer{status: http.StatusOK}
	s.Server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		s.hits.Add(1)
		time.Sleep(s.delay)
		if s.status != http.StatusOK {
			w.WriteHeader(s.status)
			return
		}
		// Bodies larger than the server's write buffer would otherwise go
		// out chunked without a Content-Length, unlike real image mirrors.
		if s.noLength {
			w.(http.Flusher).Flush()
		} else {
			w.Header().Set("Content-Length", strconv.Itoa(len(testImage)))
		}
		_, _ = w.Write(testImage)
	}))
	t.Cleanup(s.Close)
	return s
}

func testSource(srv *imageServer, sha string) imageSource {
	return imageSource{Key: "debian-13", KVMImage: provider.KVMImage{URL: srv.URL + "/debian.qcow2", SHA256: sha, Format: formatQcow2}}
}

func ensureTest(t *testing.T, api *fakeAPI, src imageSource) (string, error) {
	t.Helper()
	c := newImageCache(nil)
	pool, err := api.StoragePoolLookupByName("default")
	if err != nil {
		t.Fatal(err)
	}
	return c.ensure(context.Background(), logr.Discard(), api, "nas1", pool, src)
}

func TestResolveImage(t *testing.T) {
	t.Parallel()
	cfg := withDefaults(provider.KVMConfig{Images: map[string]provider.KVMImage{
		"debian-13": {URL: "https://cloud.debian.org/d.qcow2", SHA256: "abc"},
	}})
	tests := []struct {
		name    string
		image   string
		want    imageSource
		wantErr bool
	}{
		{"configured key", "debian-13", imageSource{Key: "debian-13", KVMImage: provider.KVMImage{URL: "https://cloud.debian.org/d.qcow2", SHA256: "abc", Format: "qcow2", Firmware: "uefi"}}, false},
		{"direct https url", "https://x.example/img.qcow2", imageSource{KVMImage: provider.KVMImage{URL: "https://x.example/img.qcow2", Format: "qcow2", Firmware: "uefi"}}, false},
		{"direct http url", "http://mirror/img.qcow2", imageSource{KVMImage: provider.KVMImage{URL: "http://mirror/img.qcow2", Format: "qcow2", Firmware: "uefi"}}, false},
		{"unknown key", "ubuntu", imageSource{}, true},
		{"file url", "file:///tmp/x.qcow2", imageSource{}, true},
		{"empty", "", imageSource{}, true},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()
			got, err := resolveImage(cfg, tc.image)
			if tc.wantErr {
				wantClass(t, err, provider.ErrPermanent)
				return
			}
			if err != nil {
				t.Fatalf("resolveImage(%q) error = %v", tc.image, err)
			}
			if got != tc.want {
				t.Errorf("resolveImage(%q) = %+v, want %+v", tc.image, got, tc.want)
			}
		})
	}
}

func TestImageSource_volumeNames(t *testing.T) {
	t.Parallel()
	keyed := imageSource{Key: "debian-13", KVMImage: provider.KVMImage{URL: "https://a/b.qcow2", Format: "qcow2"}}
	if keyed.volumeName() != "base-debian-13.qcow2" || keyed.markerName() != "base-debian-13.qcow2.sha256" {
		t.Errorf("keyed names = %s / %s", keyed.volumeName(), keyed.markerName())
	}
	direct := imageSource{KVMImage: provider.KVMImage{URL: "https://a/b.qcow2", Format: "qcow2"}}
	name := direct.volumeName()
	if !strings.HasPrefix(name, "base-") || !strings.HasSuffix(name, ".qcow2") || len(name) != len("base-")+16+len(".qcow2") {
		t.Errorf("direct-url name = %q, want base-<16 hex>.qcow2", name)
	}
	other := imageSource{KVMImage: provider.KVMImage{URL: "https://a/c.qcow2", Format: "qcow2"}}
	if other.volumeName() == name {
		t.Error("different URLs produced the same volume name")
	}
}

func TestImageCache_ensure_freshFetch(t *testing.T) {
	t.Parallel()
	srv := newImageServer(t)
	api := newFakeAPI("default")
	path, err := ensureTest(t, api, testSource(srv, testImageSHA()))
	if err != nil {
		t.Fatalf("ensure() error = %v", err)
	}
	if path != "/var/lib/libvirt/images/base-debian-13.qcow2" {
		t.Errorf("path = %q", path)
	}
	base := api.vol("default", "base-debian-13.qcow2")
	if base == nil || !bytes.Equal(base.data, testImage) {
		t.Fatalf("base volume missing or content differs")
	}
	if base.format != "qcow2" || base.capacity != int64(len(testImage)) {
		t.Errorf("base volume format/capacity = %s/%d, want qcow2/%d", base.format, base.capacity, len(testImage))
	}
	marker := api.vol("default", "base-debian-13.qcow2.sha256")
	if marker == nil || string(marker.data) != testImageSHA() {
		t.Fatalf("marker missing or wrong: %+v", marker)
	}
	if !api.hasCall("StoragePoolRefresh") {
		t.Error("pool was not refreshed after the upload")
	}
	if srv.hits.Load() != 1 {
		t.Errorf("http hits = %d, want 1", srv.hits.Load())
	}
}

func TestImageCache_ensure_noChecksumConfigured(t *testing.T) {
	t.Parallel()
	srv := newImageServer(t)
	api := newFakeAPI("default")
	if _, err := ensureTest(t, api, testSource(srv, "")); err != nil {
		t.Fatalf("ensure() error = %v", err)
	}
	if marker := api.vol("default", "base-debian-13.qcow2.sha256"); marker == nil || string(marker.data) != testImageSHA() {
		t.Errorf("marker should still record the computed digest, got %+v", marker)
	}
	// Second call: ready without a checksum to compare, so no marker read.
	if _, err := ensureTest(t, api, testSource(srv, "")); err != nil {
		t.Fatalf("second ensure() error = %v", err)
	}
	if srv.hits.Load() != 1 {
		t.Errorf("http hits = %d, want 1", srv.hits.Load())
	}
}

func TestImageCache_ensure_existingVolumes(t *testing.T) {
	t.Parallel()
	tests := []struct {
		name     string
		base     []byte
		marker   []byte
		wantHits int32
	}{
		{"already cached", testImage, []byte(testImageSHA()), 0},
		{"marker missing is a partial upload", []byte("half an image"), nil, 1},
		{"configured checksum changed", []byte("an older release"), []byte(strings.Repeat("0", 64)), 1},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()
			srv := newImageServer(t)
			api := newFakeAPI("default")
			api.putVol("default", "base-debian-13.qcow2", "qcow2", tc.base)
			if tc.marker != nil {
				api.putVol("default", "base-debian-13.qcow2.sha256", "raw", tc.marker)
			}
			path, err := ensureTest(t, api, testSource(srv, testImageSHA()))
			if err != nil {
				t.Fatalf("ensure() error = %v", err)
			}
			if path != "/var/lib/libvirt/images/base-debian-13.qcow2" {
				t.Errorf("path = %q", path)
			}
			if srv.hits.Load() != tc.wantHits {
				t.Errorf("http hits = %d, want %d", srv.hits.Load(), tc.wantHits)
			}
			if touched := api.hasCall("StorageVolCreateXML") || api.hasCall("StorageVolDelete"); touched != (tc.wantHits > 0) {
				t.Errorf("volumes touched = %v, want %v: calls = %v", touched, tc.wantHits > 0, api.calls)
			}
			if base := api.vol("default", "base-debian-13.qcow2"); base == nil || !bytes.Equal(base.data, testImage) {
				t.Error("base volume does not hold the verified image")
			}
			if marker := api.vol("default", "base-debian-13.qcow2.sha256"); marker == nil || string(marker.data) != testImageSHA() {
				t.Errorf("marker = %+v, want the verified digest", marker)
			}
		})
	}
}

func TestImageCache_ensure_checksumMismatchIsPermanent(t *testing.T) {
	t.Parallel()
	srv := newImageServer(t)
	api := newFakeAPI("default")
	_, err := ensureTest(t, api, testSource(srv, strings.Repeat("a", 64)))
	wantClass(t, err, provider.ErrPermanent)
	if names := api.volNames("default"); len(names) != 0 {
		t.Errorf("volumes left behind after a checksum mismatch: %v", names)
	}
}

func TestImageCache_ensure_httpErrors(t *testing.T) {
	t.Parallel()
	tests := []struct {
		name   string
		status int
		want   error
	}{
		{"404 is permanent", http.StatusNotFound, provider.ErrPermanent},
		{"403 is permanent", http.StatusForbidden, provider.ErrPermanent},
		{"503 is transient", http.StatusServiceUnavailable, provider.ErrTransient},
		{"500 is transient", http.StatusInternalServerError, provider.ErrTransient},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()
			srv := newImageServer(t)
			srv.status = tc.status
			api := newFakeAPI("default")
			_, err := ensureTest(t, api, testSource(srv, ""))
			if !errors.Is(err, tc.want) {
				t.Fatalf("ensure() error = %v, want %v", err, tc.want)
			}
			if names := api.volNames("default"); len(names) != 0 {
				t.Errorf("volumes created despite HTTP %d: %v", tc.status, names)
			}
		})
	}
}

func TestImageCache_ensure_noContentLengthIsPermanent(t *testing.T) {
	t.Parallel()
	srv := newImageServer(t)
	srv.noLength = true
	api := newFakeAPI("default")
	_, err := ensureTest(t, api, testSource(srv, ""))
	wantClass(t, err, provider.ErrPermanent)
	if !strings.Contains(err.Error(), "Content-Length") {
		t.Errorf("ensure() error = %v, want one about Content-Length", err)
	}
}

func TestImageCache_ensure_uploadFailureCleansUp(t *testing.T) {
	t.Parallel()
	srv := newImageServer(t)
	api := newFakeAPI("default")
	api.uploadHook = func(_ string, r io.Reader) error {
		_, _ = io.CopyN(io.Discard, r, 100)
		return io.ErrUnexpectedEOF
	}
	_, err := ensureTest(t, api, testSource(srv, testImageSHA()))
	if !errors.Is(err, provider.ErrTransient) && classify(err) != provider.ErrTransient {
		t.Fatalf("ensure() error = %v, want a transient failure", err)
	}
	if names := api.volNames("default"); len(names) != 0 {
		t.Errorf("volumes left behind after a failed upload: %v", names)
	}
	// Next attempt starts clean and succeeds.
	api.uploadHook = nil
	if _, err := ensureTest(t, api, testSource(srv, testImageSHA())); err != nil {
		t.Fatalf("retry ensure() error = %v", err)
	}
}

func TestImageCache_ensure_missingPoolIsPermanent(t *testing.T) {
	t.Parallel()
	srv := newImageServer(t)
	api := newFakeAPI("default")
	c := newImageCache(nil)
	pool, _ := api.StoragePoolLookupByName("default")
	delete(api.pools, "default")
	_, err := c.ensure(context.Background(), logr.Discard(), api, "nas1", pool, testSource(srv, ""))
	if classify(err) != provider.ErrPermanent {
		t.Fatalf("ensure() error = %v, want one classifying as Permanent", err)
	}
}

func TestImageCache_ensure_concurrentCallersDownloadOnce(t *testing.T) {
	t.Parallel()
	srv := newImageServer(t)
	srv.delay = 50 * time.Millisecond
	api := newFakeAPI("default")
	c := newImageCache(nil)
	pool, _ := api.StoragePoolLookupByName("default")
	src := testSource(srv, testImageSHA())
	var wg sync.WaitGroup
	paths := make([]string, 4)
	errs := make([]error, 4)
	for i := range 4 {
		wg.Go(func() {
			paths[i], errs[i] = c.ensure(context.Background(), logr.Discard(), api, "nas1", pool, src)
		})
	}
	wg.Wait()
	for i := range 4 {
		if errs[i] != nil || paths[i] != "/var/lib/libvirt/images/base-debian-13.qcow2" {
			t.Errorf("caller %d: path=%q err=%v", i, paths[i], errs[i])
		}
	}
	if srv.hits.Load() != 1 {
		t.Errorf("http hits = %d, want 1 for four concurrent callers", srv.hits.Load())
	}
}

func TestImageCache_ensure_differentHostsFetchIndependently(t *testing.T) {
	t.Parallel()
	srv := newImageServer(t)
	c := newImageCache(nil)
	src := testSource(srv, "")
	for _, host := range []string{"nas1", "lab2"} {
		api := newFakeAPI("default")
		pool, _ := api.StoragePoolLookupByName("default")
		if _, err := c.ensure(context.Background(), logr.Discard(), api, host, pool, src); err != nil {
			t.Fatalf("ensure(%s) error = %v", host, err)
		}
	}
	if srv.hits.Load() != 2 {
		t.Errorf("http hits = %d, want one per host", srv.hits.Load())
	}
}

func TestImageCache_ensure_cancelledContextIsTransient(t *testing.T) {
	t.Parallel()
	srv := newImageServer(t)
	api := newFakeAPI("default")
	c := newImageCache(nil)
	pool, _ := api.StoragePoolLookupByName("default")
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	_, err := c.ensure(ctx, logr.Discard(), api, "nas1", pool, testSource(srv, ""))
	if classify(err) != provider.ErrTransient {
		t.Fatalf("ensure() error = %v, want transient", err)
	}
	if names := api.volNames("default"); len(names) != 0 {
		t.Errorf("volumes created despite cancelled context: %v", names)
	}
}