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

package kvm2

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

	"github.com/digitalocean/go-libvirt"
	"github.com/go-logr/logr/funcr"
	"go.opentelemetry.io/otel/attribute"
	sdktrace "go.opentelemetry.io/otel/sdk/trace"
	"go.opentelemetry.io/otel/sdk/trace/tracetest"
	logf "sigs.k8s.io/controller-runtime/pkg/log"

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

const imageKey = "debian-13"

var testPool = libvirt.StoragePool{Name: fakePoolName, UUID: fakeUUID(fakePoolName)}

// imageFixture is one host with a pooled fake, a dialer that hands out
// twins of it for dialFresh, an image server, a span recorder and a log
// capture. Tests that need a single ordered call log point the dialer at
// the pooled fake instead (sameConnection).
type imageFixture struct {
	t      *testing.T
	p      *Provider
	pooled *fakeBackend
	dialer *stubDialer
	srv    *httptest.Server
	hits   atomic.Int32
	sr     *tracetest.SpanRecorder
	logs   *logCapture
	ctx    context.Context
	host   provider.KVMHost

	mu    sync.Mutex
	fresh []*fakeBackend
	// onFresh configures each twin before it is handed out.
	onFresh func(*fakeBackend)
}

func newImageFixture(t *testing.T, handler http.HandlerFunc) *imageFixture {
	t.Helper()
	fx := &imageFixture{t: t, pooled: newFakeBackend(), logs: &logCapture{}}
	fx.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fx.hits.Add(1)
		handler(w, r)
	}))
	t.Cleanup(fx.srv.Close)
	fx.dialer = &stubDialer{fresh: func() backend {
		tw := fx.pooled.twin()
		fx.mu.Lock()
		if fx.onFresh != nil {
			fx.onFresh(tw)
		}
		fx.fresh = append(fx.fresh, tw)
		fx.mu.Unlock()
		return tw
	}}
	fx.sr = tracetest.NewSpanRecorder()
	tp := sdktrace.NewTracerProvider(sdktrace.WithSpanProcessor(fx.sr))
	t.Cleanup(func() { _ = tp.Shutdown(context.Background()) })
	cfg := provider.KVMConfig{
		Images:       map[string]provider.KVMImage{imageKey: {URL: fx.srv.URL + "/images/debian-13-generic-amd64.qcow2"}},
		MachineTypes: map[string]provider.KVMMachineType{"small": {VCPUs: 1, MemoryMiB: 1024}},
		Hosts:        []provider.KVMHost{{Name: "beelink", Address: "beelink:22", User: "virt", Bridge: "br0"}},
	}
	fx.p = newProvider("kvm2-home", cfg, fx.dialer, withHTTPClient(fx.srv.Client()), withTracing(tracing.WithTracerProvider(tp)))
	fx.host = fx.p.cfg.Hosts[0]
	log := funcr.New(func(prefix, args string) {
		fx.logs.mu.Lock()
		defer fx.logs.mu.Unlock()
		fx.logs.lines = append(fx.logs.lines, prefix+" "+args)
	}, funcr.Options{Verbosity: 1})
	fx.ctx = logf.IntoContext(context.Background(), log)
	return fx
}

// sameConnection makes dialFresh return the pooled fake, so one call log
// carries every call in order.
func (fx *imageFixture) sameConnection() {
	fx.dialer.fresh = nil
	fx.dialer.backend = fx.pooled
}

func (fx *imageFixture) src() imageSource {
	fx.t.Helper()
	src, err := resolveImage(fx.p.name, fx.p.cfg, imageKey)
	if err != nil {
		fx.t.Fatalf("resolveImage(%q) error = %v", imageKey, err)
	}
	return src
}

func (fx *imageFixture) ensure(ctx context.Context, src imageSource) (string, error) {
	return fx.p.ensureImage(ctx, fx.host, testPool, fx.pooled, src)
}

func (fx *imageFixture) freshBackends() []*fakeBackend {
	fx.mu.Lock()
	defer fx.mu.Unlock()
	return append([]*fakeBackend(nil), fx.fresh...)
}

// putVol places a volume in the pool directly, as a previous run would have.
func (fx *imageFixture) putVol(name string, content []byte) {
	fx.pooled.mu.Lock()
	defer fx.pooled.mu.Unlock()
	fx.pooled.pools[fakePoolName].vols[name] = &fakeVol{content: content, capacity: uint64(len(content))}
}

func randomBody(t *testing.T, n int) []byte {
	t.Helper()
	b := make([]byte, n)
	if _, err := rand.Read(b); err != nil {
		t.Fatal(err)
	}
	return b
}

func digestOf(b []byte) string {
	sum := sha256.Sum256(b)
	return hex.EncodeToString(sum[:])
}

func serveBytes(body []byte) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Length", strconv.Itoa(len(body)))
		w.WriteHeader(http.StatusOK)
		if r.Method != http.MethodHead {
			_, _ = w.Write(body)
		}
	}
}

func serveStatus(code int) http.HandlerFunc {
	return func(w http.ResponseWriter, _ *http.Request) { http.Error(w, http.StatusText(code), code) }
}

// headFails answers HEAD with code and hands every other method to next.
func headFails(code int, next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if r.Method == http.MethodHead {
			http.Error(w, http.StatusText(code), code)
			return
		}
		next(w, r)
	}
}

func TestResolveImage_table(t *testing.T) {
	t.Parallel()
	cfg := withDefaults(provider.KVMConfig{Images: map[string]provider.KVMImage{
		"debian-13": {URL: "https://cloud.debian.org/debian-13.qcow2", SHA256: "ABCD"},
		"alpine":    {URL: "https://alpine.example/alpine.img", Format: provider.KVMFormatRaw, Firmware: provider.KVMFirmwareBIOS},
	}})
	const rawURL = "http://mirror.example/x/y.qcow2"
	tests := []struct {
		name       string
		image      string
		want       imageSource
		wantVolume string
		wantErr    string
	}{
		{
			name:       "catalogue key with defaults applied",
			image:      "debian-13",
			want:       imageSource{Entry: testProvider, Key: "debian-13", KVMImage: provider.KVMImage{URL: "https://cloud.debian.org/debian-13.qcow2", SHA256: "ABCD", Format: "qcow2", Firmware: "uefi"}},
			wantVolume: "base-" + testProvider + "-debian-13.qcow2",
		},
		{
			name:       "catalogue key with explicit format and firmware",
			image:      "alpine",
			want:       imageSource{Entry: testProvider, Key: "alpine", KVMImage: provider.KVMImage{URL: "https://alpine.example/alpine.img", Format: "raw", Firmware: "bios"}},
			wantVolume: "base-" + testProvider + "-alpine.raw",
		},
		{
			name:       "http URL used directly with defaults",
			image:      rawURL,
			want:       imageSource{Entry: testProvider, Key: urlImageKey(rawURL), KVMImage: provider.KVMImage{URL: rawURL, Format: "qcow2", Firmware: "uefi"}},
			wantVolume: "base-" + testProvider + "-" + urlImageKey(rawURL) + ".qcow2",
		},
		{
			name:    "unknown key is permanent and names the value",
			image:   "ubuntu-24",
			wantErr: `placement.image "ubuntu-24" is neither a key in kvm2.images nor an http(s) URL`,
		},
		{
			name:    "non-http scheme is permanent",
			image:   "file:///var/lib/images/x.qcow2",
			wantErr: `placement.image "file:///var/lib/images/x.qcow2" is neither`,
		},
		{
			name:    "empty is permanent",
			image:   "",
			wantErr: `placement.image "" is neither`,
		},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()
			got, err := resolveImage(testProvider, cfg, tc.image)
			if tc.wantErr != "" {
				if err == nil || !errors.Is(err, provider.ErrPermanent) || !strings.Contains(err.Error(), tc.wantErr) {
					t.Fatalf("resolveImage(%q) = %+v, %v; want ErrPermanent containing %q", tc.image, got, err, tc.wantErr)
				}
				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)
			}
			if got.volumeName() != tc.wantVolume || got.markerName() != tc.wantVolume+".sha256" {
				t.Errorf("volume names = %q, %q; want %q and its .sha256", got.volumeName(), got.markerName(), tc.wantVolume)
			}
		})
	}
}

func TestEnsureImage_downloadsAndWritesMarkerLast(t *testing.T) {
	t.Parallel()
	body := randomBody(t, 1<<20)
	fx := newImageFixture(t, serveBytes(body))
	fx.sameConnection()
	src := fx.src()

	path, err := fx.ensure(fx.ctx, src)
	if err != nil {
		t.Fatalf("ensureImage() error = %v", err)
	}
	if want := fakeVolPath(fakePoolName, src.volumeName()); path != want {
		t.Errorf("path = %q, want %q", path, want)
	}
	base := fx.pooled.vol(src.volumeName())
	if base == nil || !bytes.Equal(base.content, body) {
		t.Fatalf("base volume missing or content differs (%d bytes)", len(body))
	}
	if base.capacity != uint64(len(body)) {
		t.Errorf("base capacity = %d, want the Content-Length %d", base.capacity, len(body))
	}
	marker := fx.pooled.vol(src.markerName())
	if marker == nil || string(marker.content) != digestOf(body) {
		t.Fatalf("marker = %+v, want the body's digest %s", marker, digestOf(body))
	}
	if marker.capacity != markerCapacityBytes {
		t.Errorf("marker capacity = %d, want %d", marker.capacity, markerCapacityBytes)
	}

	log := fx.pooled.callLog()
	wantTail := []string{"StorageVolCreateXML", "StorageVolUpload", "StoragePoolRefresh", "StorageVolCreateXML", "StorageVolUpload", "StorageVolGetPath"}
	if len(log) < len(wantTail) || strings.Join(log[len(log)-len(wantTail):], ",") != strings.Join(wantTail, ",") {
		t.Errorf("call log = %v, want it to end with %v (base create+upload, refresh, marker create+upload, path)", log, wantTail)
	}
	if fx.hits.Load() != 1 {
		t.Errorf("server hits = %d, want 1", fx.hits.Load())
	}

	for _, want := range []string{`"msg"="fetching base image into host pool"`, `"msg"="base image cached"`} {
		lines := fx.logs.matching(want)
		if len(lines) != 1 {
			t.Errorf("got %d lines with %s, want 1:\n%s", len(lines), want, strings.Join(fx.logs.all(), "\n"))
			continue
		}
		for _, kv := range []string{`"provider"="kvm2-home"`, `"host"="beelink"`, `"volume"="base-` + testProvider + `-debian-13.qcow2"`, `"bytes"=` + strconv.Itoa(len(body)), `"traceID"=`} {
			if !strings.Contains(lines[0], kv) {
				t.Errorf("line lacks %s: %s", kv, lines[0])
			}
		}
	}
	if cached := fx.logs.matching(`"msg"="base image cached"`); len(cached) == 1 && !strings.Contains(cached[0], `"sha256"="`+digestOf(body)+`"`) {
		t.Errorf("cached line lacks the digest: %s", cached[0])
	}

	spans := fx.sr.Ended()
	if len(spans) != 1 || spans[0].Name() != "kvm2.download" {
		t.Fatalf("spans = %v, want exactly one kvm2.download", spanNames(spans))
	}
	attrs := spans[0].Attributes()
	for _, want := range []attribute.KeyValue{
		attribute.String("kvm2.host", "beelink"),
		attribute.String("kvm2.volume", "base-"+testProvider+"-debian-13.qcow2"),
		attribute.String("kvm2.url_host", strings.TrimPrefix(fx.srv.URL, "http://")),
		attribute.Int64("kvm2.bytes", int64(len(body))),
		attribute.Bool("kvm2.sha256_configured", false),
	} {
		if !hasAttr(attrs, want) {
			t.Errorf("span attributes %v lack %v", attrs, want)
		}
	}
	for _, a := range attrs {
		if strings.Contains(a.Value.Emit(), src.URL) || strings.Contains(a.Value.Emit(), "/images/") {
			t.Errorf("span attribute %s carries the full URL: %s", a.Key, a.Value.Emit())
		}
	}
}

func spanNames(spans []sdktrace.ReadOnlySpan) []string {
	out := make([]string, len(spans))
	for i, s := range spans {
		out[i] = s.Name()
	}
	return out
}

func TestEnsureImage_cachedWithMarkerSkipsDownload(t *testing.T) {
	t.Parallel()
	body := randomBody(t, 4096)
	fx := newImageFixture(t, serveBytes(body))
	src := fx.src()
	fx.putVol(src.volumeName(), body)
	fx.putVol(src.markerName(), []byte(digestOf(body)))

	path, err := fx.ensure(fx.ctx, src)
	if err != nil {
		t.Fatalf("ensureImage() error = %v", err)
	}
	if want := fakeVolPath(fakePoolName, src.volumeName()); path != want {
		t.Errorf("path = %q, want %q", path, want)
	}
	if fx.hits.Load() != 0 {
		t.Errorf("server hits = %d, want 0", fx.hits.Load())
	}
	if fx.dialer.dials.Load() != 0 {
		t.Errorf("dials = %d, want 0: a cached image needs no dedicated connection", fx.dialer.dials.Load())
	}
	if n := fx.pooled.callCount("StorageVolDelete"); n != 0 {
		t.Errorf("StorageVolDelete called %d times, want 0", n)
	}
	if len(fx.sr.Ended()) != 0 {
		t.Errorf("spans = %v, want none without a download", spanNames(fx.sr.Ended()))
	}
}

func TestEnsureImage_cachedMarkerMatchesConfiguredSHA256CaseInsensitively(t *testing.T) {
	t.Parallel()
	body := randomBody(t, 4096)
	fx := newImageFixture(t, serveBytes(body))
	src := fx.src()
	src.SHA256 = strings.ToUpper(digestOf(body))
	fx.putVol(src.volumeName(), body)
	fx.putVol(src.markerName(), []byte(digestOf(body)))

	if _, err := fx.ensure(fx.ctx, src); err != nil {
		t.Fatalf("ensureImage() error = %v", err)
	}
	if fx.hits.Load() != 0 {
		t.Errorf("server hits = %d, want 0", fx.hits.Load())
	}
}

func TestEnsureImage_missingMarkerRedownloads(t *testing.T) {
	t.Parallel()
	body := randomBody(t, 8192)
	fx := newImageFixture(t, serveBytes(body))
	src := fx.src()
	fx.putVol(src.volumeName(), []byte("cut off half way"))

	if _, err := fx.ensure(fx.ctx, src); err != nil {
		t.Fatalf("ensureImage() error = %v", err)
	}
	if fx.hits.Load() != 1 {
		t.Errorf("server hits = %d, want 1", fx.hits.Load())
	}
	if n := fx.pooled.callCount("StorageVolDelete"); n != 1 {
		t.Errorf("StorageVolDelete called %d times on the pooled connection, want 1 (the leftover base)", n)
	}
	base := fx.pooled.vol(src.volumeName())
	if base == nil || !bytes.Equal(base.content, body) {
		t.Errorf("base volume was not replaced by the download")
	}
	marker := fx.pooled.vol(src.markerName())
	if marker == nil || string(marker.content) != digestOf(body) {
		t.Errorf("marker = %+v, want %s", marker, digestOf(body))
	}
}

// TestEnsureImage_sha256ChangeRefusesInsteadOfReplacing pins the rule the
// contract already stated and the code contradicted. Every overlay on the
// host names the base as its backing file, and libvirt unlinks it under
// running domains without complaint, so a digest that no longer matches is
// for the operator to resolve, not for the next Create to resolve by
// deleting.
func TestEnsureImage_sha256ChangeRefusesInsteadOfReplacing(t *testing.T) {
	t.Parallel()
	body := randomBody(t, 8192)
	fx := newImageFixture(t, serveBytes(body))
	src := fx.src()
	src.SHA256 = digestOf(body)
	digestA := strings.Repeat("a", 64)
	fx.putVol(src.volumeName(), []byte("the old image"))
	fx.putVol(src.markerName(), []byte(digestA))

	_, err := fx.ensure(fx.ctx, src)
	if !errors.Is(err, provider.ErrPermanent) || !strings.Contains(err.Error(), "add the image under a new key") {
		t.Fatalf("ensureImage() error = %v, want ErrPermanent telling the operator what to do", err)
	}
	if fx.hits.Load() != 0 {
		t.Errorf("server hits = %d, want 0: nothing should be refetched", fx.hits.Load())
	}
	if n := fx.pooled.callCount("StorageVolDelete"); n != 0 {
		t.Errorf("StorageVolDelete called %d times, want 0: the base backs live overlays", n)
	}
	if base := fx.pooled.vol(src.volumeName()); base == nil || !bytes.Equal(base.content, []byte("the old image")) {
		t.Errorf("base volume = %+v, want the cached one untouched", base)
	}
	if marker := fx.pooled.vol(src.markerName()); marker == nil || string(marker.content) != digestA {
		t.Errorf("marker = %+v, want the cached digest untouched", marker)
	}
}

func TestEnsureImage_markerWithoutDigestRedownloads(t *testing.T) {
	t.Parallel()
	body := randomBody(t, 4096)
	fx := newImageFixture(t, serveBytes(body))
	src := fx.src()
	fx.putVol(src.volumeName(), []byte("whatever"))
	fx.putVol(src.markerName(), []byte("not a digest"))

	if _, err := fx.ensure(fx.ctx, src); err != nil {
		t.Fatalf("ensureImage() error = %v", err)
	}
	if fx.hits.Load() != 1 {
		t.Errorf("server hits = %d, want 1", fx.hits.Load())
	}
	if marker := fx.pooled.vol(src.markerName()); marker == nil || string(marker.content) != digestOf(body) {
		t.Errorf("marker = %+v, want %s", marker, digestOf(body))
	}
}

func TestEnsureImage_sha256MismatchIsPermanentAndRemovesVolume(t *testing.T) {
	t.Parallel()
	body := randomBody(t, 8192)
	fx := newImageFixture(t, serveBytes(body))
	src := fx.src()
	src.SHA256 = strings.Repeat("f", 64)

	_, err := fx.ensure(fx.ctx, src)
	if !errors.Is(err, provider.ErrPermanent) {
		t.Fatalf("ensureImage() error = %v, want ErrPermanent", err)
	}
	if !strings.Contains(err.Error(), digestOf(body)) || !strings.Contains(err.Error(), src.SHA256) {
		t.Errorf("error %q should name both digests", err)
	}
	if fx.pooled.vol(src.volumeName()) != nil {
		t.Error("base volume left behind after a digest mismatch")
	}
	if fx.pooled.vol(src.markerName()) != nil {
		t.Error("marker written although the digest did not match")
	}
	if n := fx.pooled.callCount("StoragePoolRefresh"); n != 0 {
		t.Errorf("StoragePoolRefresh called %d times, want 0", n)
	}
}

func TestEnsureImage_noContentLengthIsPermanent(t *testing.T) {
	t.Parallel()
	fx := newImageFixture(t, func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusOK)
		w.(http.Flusher).Flush()
		_, _ = w.Write([]byte("chunked body"))
	})
	src := fx.src()

	_, err := fx.ensure(fx.ctx, src)
	if !errors.Is(err, provider.ErrPermanent) {
		t.Fatalf("ensureImage() error = %v, want ErrPermanent", err)
	}
	if !strings.Contains(err.Error(), "no Content-Length") || !strings.Contains(err.Error(), src.URL) {
		t.Errorf("error %q should say Content-Length is missing and name the URL", err)
	}
	if fx.dialer.dials.Load() != 0 || fx.pooled.callCount("StorageVolCreateXML") != 0 {
		t.Error("a rejected response must not dial or create a volume")
	}
}

func TestEnsureImage_http404IsPermanent(t *testing.T) {
	t.Parallel()
	fx := newImageFixture(t, serveStatus(http.StatusNotFound))
	_, err := fx.ensure(fx.ctx, fx.src())
	if !errors.Is(err, provider.ErrPermanent) || !strings.Contains(err.Error(), "404") {
		t.Fatalf("ensureImage() error = %v, want ErrPermanent naming 404", err)
	}
	if fx.pooled.vol(fx.src().volumeName()) != nil {
		t.Error("volume created for a 404")
	}
}

func TestEnsureImage_http403And410ArePermanent(t *testing.T) {
	t.Parallel()
	for _, code := range []int{http.StatusForbidden, http.StatusGone, http.StatusBadRequest} {
		t.Run(strconv.Itoa(code), func(t *testing.T) {
			t.Parallel()
			fx := newImageFixture(t, serveStatus(code))
			if _, err := fx.ensure(fx.ctx, fx.src()); !errors.Is(err, provider.ErrPermanent) {
				t.Fatalf("ensureImage() error = %v, want ErrPermanent", err)
			}
		})
	}
}

func TestEnsureImage_http503IsTransient(t *testing.T) {
	t.Parallel()
	for _, code := range []int{http.StatusServiceUnavailable, http.StatusTooManyRequests, http.StatusBadGateway} {
		t.Run(strconv.Itoa(code), func(t *testing.T) {
			t.Parallel()
			fx := newImageFixture(t, serveStatus(code))
			_, err := fx.ensure(fx.ctx, fx.src())
			if !errors.Is(err, provider.ErrTransient) || errors.Is(err, provider.ErrPermanent) {
				t.Fatalf("ensureImage() error = %v, want ErrTransient", err)
			}
		})
	}
}

func TestEnsureImage_connectionRefusedIsTransient(t *testing.T) {
	t.Parallel()
	fx := newImageFixture(t, serveBytes([]byte("never served")))
	src := fx.src()
	fx.srv.Close()

	_, err := fx.ensure(fx.ctx, src)
	if !errors.Is(err, provider.ErrTransient) || errors.Is(err, provider.ErrPermanent) {
		t.Fatalf("ensureImage() error = %v, want ErrTransient", err)
	}
	if !strings.Contains(err.Error(), "fetching image") {
		t.Errorf("error %q should say what it was doing", err)
	}
}

func TestEnsureImage_shortReadRemovesVolumeAndIsTransient(t *testing.T) {
	t.Parallel()
	fx := newImageFixture(t, func(w http.ResponseWriter, _ *http.Request) {
		w.Header().Set("Content-Length", "4096")
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write(make([]byte, 1024))
		w.(http.Flusher).Flush()
		conn, _, err := w.(http.Hijacker).Hijack()
		if err != nil {
			return
		}
		_ = conn.Close()
	})
	src := fx.src()

	_, err := fx.ensure(fx.ctx, src)
	if !errors.Is(err, provider.ErrTransient) || errors.Is(err, provider.ErrPermanent) {
		t.Fatalf("ensureImage() error = %v, want ErrTransient", err)
	}
	if !strings.Contains(err.Error(), "1024 of 4096") {
		t.Errorf("error %q should say how far the body got", err)
	}
	if fx.pooled.vol(src.volumeName()) != nil {
		t.Error("base volume left behind after a short read")
	}
	if fx.pooled.vol(src.markerName()) != nil {
		t.Error("marker written after a short read")
	}
	if fresh := fx.freshBackends(); len(fresh) != 1 || fresh[0].closeCount() != 1 {
		t.Errorf("fresh connections = %d, want 1 and closed", len(fresh))
	}
}

func TestEnsureImage_concurrentCallersFetchOnce(t *testing.T) {
	t.Parallel()
	body := randomBody(t, 256<<10)
	fx := newImageFixture(t, serveBytes(body))
	src := fx.src()

	const n = 20
	var wg sync.WaitGroup
	paths := make([]string, n)
	errs := make([]error, n)
	for i := range n {
		wg.Add(1)
		go func() {
			defer wg.Done()
			paths[i], errs[i] = fx.ensure(fx.ctx, src)
		}()
	}
	wg.Wait()
	want := fakeVolPath(fakePoolName, src.volumeName())
	for i := range n {
		if errs[i] != nil || paths[i] != want {
			t.Errorf("caller %d: path = %q, err = %v; want %q, nil", i, paths[i], errs[i], want)
		}
	}
	if fx.hits.Load() != 1 {
		t.Errorf("server hits = %d, want 1", fx.hits.Load())
	}
	baseUploads := 0
	for _, b := range fx.freshBackends() {
		baseUploads += b.callCount("StorageVolUpload")
	}
	if baseUploads != 1 || len(fx.freshBackends()) != 1 {
		t.Errorf("base uploads = %d over %d dedicated connections, want 1 over 1", baseUploads, len(fx.freshBackends()))
	}
	if n := fx.pooled.callCount("StorageVolUpload"); n != 1 {
		t.Errorf("pooled uploads = %d, want 1 (the marker)", n)
	}
	if got := len(fx.sr.Ended()); got != 1 {
		t.Errorf("download spans = %d, want 1", got)
	}
}

func TestEnsureImage_cancelledMidStreamReturnsPromptly(t *testing.T) {
	t.Parallel()
	started := make(chan struct{})
	fx := newImageFixture(t, func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Length", strconv.Itoa(1<<20))
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write(make([]byte, 1024))
		w.(http.Flusher).Flush()
		close(started)
		<-r.Context().Done()
	})
	src := fx.src()

	ctx, cancel := context.WithCancel(fx.ctx)
	defer cancel()
	go func() {
		<-started
		cancel()
	}()
	begin := time.Now()
	_, err := fx.ensure(ctx, src)
	elapsed := time.Since(begin)
	if elapsed > time.Second {
		t.Errorf("ensureImage() took %v after cancellation, want under 1s", elapsed)
	}
	if !errors.Is(err, provider.ErrTransient) || !errors.Is(err, context.Canceled) {
		t.Fatalf("ensureImage() error = %v, want ErrTransient wrapping context.Canceled", err)
	}
	if fx.pooled.vol(src.volumeName()) != nil {
		t.Error("base volume left behind after cancellation; the best-effort delete must not use the cancelled ctx")
	}
	if fx.pooled.vol(src.markerName()) != nil {
		t.Error("marker written after cancellation")
	}
}

// TestEnsureImage_streamsWithoutBuffering proves the body is never
// materialised: the server counts what it has sent, and the fake records
// that count the moment StorageVolUpload is entered. Had ensureImage read
// the body first (io.ReadAll, a bytes.Buffer, a temp file), the server would
// have delivered all 64 MiB before the upload began; streaming means the
// upload starts while the server has sent only what the socket buffers
// hold. This holds under -race and in parallel, unlike a process-wide
// allocation measurement. The fake discards the bytes through a sha256, so
// the test itself holds no 64 MiB either.
func TestEnsureImage_streamsWithoutBuffering(t *testing.T) {
	t.Parallel()
	const size = 64 << 20
	block := randomBody(t, 64<<10)
	var served atomic.Int64
	fx := newImageFixture(t, func(w http.ResponseWriter, _ *http.Request) {
		w.Header().Set("Content-Length", strconv.Itoa(size))
		w.WriteHeader(http.StatusOK)
		for sent := 0; sent < size; sent += len(block) {
			if _, err := w.Write(block); err != nil {
				return
			}
			served.Add(int64(len(block)))
		}
	})
	fx.pooled.discardUploads = true
	var servedAtUploadStart atomic.Int64
	fx.onFresh = func(tw *fakeBackend) {
		tw.hooks["StorageVolUpload"] = func() error {
			servedAtUploadStart.Store(served.Load())
			return nil
		}
	}
	src := fx.src()
	src.SHA256 = func() string {
		h := sha256.New()
		for sent := 0; sent < size; sent += len(block) {
			h.Write(block)
		}
		return hex.EncodeToString(h.Sum(nil))
	}()

	if _, err := fx.ensure(fx.ctx, src); err != nil {
		t.Fatalf("ensureImage() error = %v", err)
	}
	if got := servedAtUploadStart.Load(); got >= 16<<20 {
		t.Errorf("server had sent %d bytes when the upload started, want under 16 MiB: the body was buffered", got)
	}
	base := fx.pooled.vol(src.volumeName())
	if base == nil || base.uploaded != size || base.sha256 != src.SHA256 {
		t.Fatalf("base = %+v, want %d bytes uploaded with digest %s", base, size, src.SHA256)
	}
	if marker := fx.pooled.vol(src.markerName()); marker == nil || marker.sha256 != digestOf([]byte(src.SHA256)) {
		t.Errorf("marker not written after the streamed upload")
	}
}

func TestEnsureImage_usesDedicatedConnectionForUpload(t *testing.T) {
	t.Parallel()
	body := randomBody(t, 64<<10)
	fx := newImageFixture(t, serveBytes(body))
	src := fx.src()

	if _, err := fx.ensure(fx.ctx, src); err != nil {
		t.Fatalf("ensureImage() error = %v", err)
	}
	fresh := fx.freshBackends()
	if len(fresh) != 1 {
		t.Fatalf("dedicated connections = %d, want 1", len(fresh))
	}
	dedicated := fresh[0]
	if dedicated.callCount("StorageVolCreateXML") != 1 || dedicated.callCount("StorageVolUpload") != 1 {
		t.Errorf("dedicated connection calls = %v, want the base create and upload", dedicated.callLog())
	}
	if dedicated.closeCount() != 1 {
		t.Errorf("dedicated connection closed %d times, want 1", dedicated.closeCount())
	}
	if fx.pooled.closeCount() != 0 {
		t.Error("pooled connection was closed")
	}
	if n := fx.pooled.callCount("StorageVolUpload"); n != 1 {
		t.Errorf("pooled uploads = %d, want 1 (the marker only)", n)
	}
	if fx.pooled.callCount("StorageVolLookupByName") == 0 || fx.pooled.callCount("StoragePoolRefresh") != 1 {
		t.Errorf("pooled connection calls = %v, want the lookups and one refresh", fx.pooled.callLog())
	}
	if fx.pooled.callCount("StorageVolCreateXML") != 1 {
		t.Errorf("pooled StorageVolCreateXML = %d, want 1 (the marker)", fx.pooled.callCount("StorageVolCreateXML"))
	}
	if fx.dialer.hosts[0] != "beelink" {
		t.Errorf("dialled %q, want beelink", fx.dialer.hosts[0])
	}
}

func TestEnsureImage_dialFailureIsReturnedAndCreatesNothing(t *testing.T) {
	t.Parallel()
	fx := newImageFixture(t, serveBytes([]byte("image")))
	fx.dialer.err = transientf("ssh: connect: no route to host")
	src := fx.src()

	_, err := fx.ensure(fx.ctx, src)
	if !errors.Is(err, provider.ErrTransient) || !strings.Contains(err.Error(), "no route to host") {
		t.Fatalf("ensureImage() error = %v, want the dial error", err)
	}
	if fx.pooled.vol(src.volumeName()) != nil || fx.pooled.vol(src.markerName()) != nil {
		t.Error("volumes created although the dedicated connection failed")
	}
}

func TestEnsureImage_lockWaiterGivesUpWithContext(t *testing.T) {
	t.Parallel()
	release := make(chan struct{})
	fx := newImageFixture(t, func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Length", "4096")
		w.WriteHeader(http.StatusOK)
		w.(http.Flusher).Flush()
		select {
		case <-release:
		case <-r.Context().Done():
			return
		}
		_, _ = w.Write(make([]byte, 4096))
	})
	src := fx.src()
	firstDone := make(chan error, 1)
	go func() {
		_, err := fx.ensure(fx.ctx, src)
		firstDone <- err
	}()
	for fx.hits.Load() == 0 {
		time.Sleep(time.Millisecond)
	}

	ctx, cancel := context.WithTimeout(fx.ctx, 50*time.Millisecond)
	defer cancel()
	begin := time.Now()
	_, err := fx.ensure(ctx, src)
	if time.Since(begin) > time.Second || !errors.Is(err, provider.ErrTransient) || !errors.Is(err, context.DeadlineExceeded) {
		t.Errorf("waiter returned after %v with %v, want promptly with ErrTransient wrapping DeadlineExceeded", time.Since(begin), err)
	}
	close(release)
	if err := <-firstDone; err != nil {
		t.Errorf("first caller error = %v", err)
	}
	if fx.hits.Load() != 1 {
		t.Errorf("server hits = %d, want 1: the waiter must not fetch", fx.hits.Load())
	}
}

func TestCheckPoolSpace_belowNeededIsQuota(t *testing.T) {
	t.Parallel()
	be := newFakeBackend()
	be.pools[fakePoolName].available = 5 << 30
	err := checkPoolSpace(context.Background(), be, testPool, fakePoolName, 6<<30)
	if !errors.Is(err, provider.ErrQuotaExceeded) {
		t.Fatalf("checkPoolSpace() error = %v, want ErrQuotaExceeded", err)
	}
	for _, want := range []string{"storage pool default", strconv.Itoa(5 << 30), strconv.Itoa(6 << 30)} {
		if !strings.Contains(err.Error(), want) {
			t.Errorf("error %q lacks %q", err, want)
		}
	}
}

func TestCheckPoolSpace_poolNotRunningIsPermanent(t *testing.T) {
	t.Parallel()
	for _, state := range []libvirt.StoragePoolState{libvirt.StoragePoolInactive, libvirt.StoragePoolBuilding, libvirt.StoragePoolDegraded, libvirt.StoragePoolInaccessible} {
		t.Run(state.String(), func(t *testing.T) {
			t.Parallel()
			be := newFakeBackend()
			be.pools[fakePoolName].state = state
			err := checkPoolSpace(context.Background(), be, testPool, fakePoolName, 1)
			if !errors.Is(err, provider.ErrPermanent) || !strings.Contains(err.Error(), "storage pool default is not running") {
				t.Fatalf("checkPoolSpace() error = %v, want ErrPermanent 'not running'", err)
			}
		})
	}
}

func TestCheckPoolSpace_enoughIsNil(t *testing.T) {
	t.Parallel()
	be := newFakeBackend()
	be.pools[fakePoolName].available = 50 << 30
	for _, needed := range []uint64{0, 1, 50 << 30} {
		if err := checkPoolSpace(context.Background(), be, testPool, fakePoolName, needed); err != nil {
			t.Errorf("checkPoolSpace(needed=%d) error = %v, want nil", needed, err)
		}
	}
}

func TestCheckPoolSpace_backendErrorIsReturned(t *testing.T) {
	t.Parallel()
	be := newFakeBackend()
	be.failWith("StoragePoolGetInfo", fakeErr(libvirt.ErrNoStoragePool, "no pool"))
	err := checkPoolSpace(context.Background(), be, testPool, fakePoolName, 1)
	if !hasCode(err, libvirt.ErrNoStoragePool) || errors.Is(err, provider.ErrQuotaExceeded) {
		t.Fatalf("checkPoolSpace() error = %v, want the libvirt error unclassified", err)
	}
}

func TestNeededBytes(t *testing.T) {
	t.Parallel()
	tests := []struct {
		name       string
		imageBytes int64
		diskGB     int64
		want       uint64
	}{
		{name: "cached image and default disk", imageBytes: 0, diskGB: 10, want: 10<<30 + seedReserveBytes},
		{name: "uncached image adds its size", imageBytes: 435 << 20, diskGB: 10, want: 435<<20 + 10<<30 + seedReserveBytes},
		{name: "zero disk keeps the seed reserve", imageBytes: 0, diskGB: 0, want: seedReserveBytes},
		{name: "negative sizes count as zero", imageBytes: -1, diskGB: -3, want: seedReserveBytes},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()
			if got := neededBytes(tc.imageBytes, tc.diskGB); got != tc.want {
				t.Errorf("neededBytes(%d, %d) = %d, want %d", tc.imageBytes, tc.diskGB, got, tc.want)
			}
		})
	}
}

func TestImageSizeIfUncached_headThenGetFallback(t *testing.T) {
	t.Parallel()
	body := randomBody(t, 12345)
	tests := []struct {
		name        string
		handler     http.HandlerFunc
		cached      bool
		wantSize    int64
		wantMethods []string
		wantClass   error
		wantErr     string
	}{
		{
			name:        "HEAD answers with the length",
			handler:     serveBytes(body),
			wantSize:    int64(len(body)),
			wantMethods: []string{http.MethodHead},
		},
		{
			name:        "HEAD not allowed falls back to GET with the body closed",
			handler:     headFails(http.StatusMethodNotAllowed, serveBytes(body)),
			wantSize:    int64(len(body)),
			wantMethods: []string{http.MethodHead, http.MethodGet},
		},
		{
			name:        "HEAD forbidden falls back to GET",
			handler:     headFails(http.StatusForbidden, serveBytes(body)),
			wantSize:    int64(len(body)),
			wantMethods: []string{http.MethodHead, http.MethodGet},
		},
		{
			name:        "HEAD not implemented falls back to GET",
			handler:     headFails(http.StatusNotImplemented, serveBytes(body)),
			wantSize:    int64(len(body)),
			wantMethods: []string{http.MethodHead, http.MethodGet},
		},
		{
			name:        "HEAD refused and GET 404 is permanent from the GET",
			handler:     headFails(http.StatusMethodNotAllowed, serveStatus(http.StatusNotFound)),
			wantMethods: []string{http.MethodHead, http.MethodGet},
			wantClass:   provider.ErrPermanent,
			wantErr:     "404",
		},
		{
			name:        "marker present means nothing to download",
			handler:     serveBytes(body),
			cached:      true,
			wantSize:    0,
			wantMethods: nil,
		},
		{
			name:        "404 on both is permanent",
			handler:     serveStatus(http.StatusNotFound),
			wantMethods: []string{http.MethodHead, http.MethodGet},
			wantClass:   provider.ErrPermanent,
			wantErr:     "404",
		},
		{
			name:        "503 on both is transient",
			handler:     serveStatus(http.StatusServiceUnavailable),
			wantMethods: []string{http.MethodHead, http.MethodGet},
			wantClass:   provider.ErrTransient,
		},
		{
			name: "no Content-Length is permanent",
			handler: func(w http.ResponseWriter, _ *http.Request) {
				w.WriteHeader(http.StatusOK)
				w.(http.Flusher).Flush()
			},
			wantMethods: []string{http.MethodHead},
			wantClass:   provider.ErrPermanent,
			wantErr:     "no Content-Length",
		},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()
			var mu sync.Mutex
			var methods []string
			fx := newImageFixture(t, func(w http.ResponseWriter, r *http.Request) {
				mu.Lock()
				methods = append(methods, r.Method)
				mu.Unlock()
				tc.handler(w, r)
			})
			src := fx.src()
			if tc.cached {
				fx.putVol(src.markerName(), []byte(strings.Repeat("0", 64)))
			}
			got, err := fx.p.imageSizeIfUncached(fx.ctx, fx.pooled, testPool, src)
			mu.Lock()
			defer mu.Unlock()
			if strings.Join(methods, ",") != strings.Join(tc.wantMethods, ",") {
				t.Errorf("methods = %v, want %v", methods, tc.wantMethods)
			}
			if tc.wantClass != nil {
				if !errors.Is(err, tc.wantClass) || !strings.Contains(err.Error(), tc.wantErr) {
					t.Fatalf("imageSizeIfUncached() error = %v, want %v containing %q", err, tc.wantClass, tc.wantErr)
				}
				return
			}
			if err != nil {
				t.Fatalf("imageSizeIfUncached() error = %v", err)
			}
			if got != tc.wantSize {
				t.Errorf("size = %d, want %d", got, tc.wantSize)
			}
		})
	}
}

func TestReadMarker_capsTheRead(t *testing.T) {
	t.Parallel()
	be := newFakeBackend()
	huge := bytes.Repeat([]byte("x"), 1<<20)
	be.pools[fakePoolName].vols["m"] = &fakeVol{content: huge}
	got, err := readMarker(context.Background(), be, libvirt.StorageVol{Pool: fakePoolName, Name: "m"})
	if err != nil {
		t.Fatal(err)
	}
	if len(got) != markerReadCap {
		t.Errorf("read %d bytes, want the cap %d", len(got), markerReadCap)
	}
	w := &cappedWriter{limit: 4}
	if _, err := io.Copy(w, strings.NewReader("abcdefgh")); err != nil || string(w.buf) != "abcd" {
		t.Errorf("cappedWriter kept %q, want abcd", w.buf)
	}
}