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

package kvm2

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/xml"
	"fmt"
	"io"
	"slices"
	"sync"
	"sync/atomic"

	"github.com/digitalocean/go-libvirt"

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

// fakeBackend is the in-memory backend of spec entry 21: enough of a
// libvirt host for the provider's Create, Get, Delete and ListByTag, with
// per-method error hooks and blocking, call counters and an ordered log.
type fakeBackend struct {
	// mu, domains and pools are the host's state; a twin shares them, so
	// two connections to the same fake host see the same volumes.
	mu       *sync.Mutex
	domains  map[string]*fakeDomain
	pools    map[string]*fakePool
	hostname string
	version  uint64

	// hooks returns the error a method fails with before touching state.
	hooks map[string]func() error
	// block makes a method wait on its channel (or ctx) before running.
	block map[string]chan struct{}
	// discardUploads makes StorageVolUpload stream the reader through a
	// sha256 in fixed chunks instead of holding the bytes, for tests that
	// push more through the fake than they want in memory.
	discardUploads bool
	// startDelayPolls, when set, makes DomainCreate leave the domain in
	// nostate; it turns running on the startDelayPolls-th DomainGetState
	// after that, as a host whose QEMU is still settling would.
	startDelayPolls int
	// stateAfterStart, when not nostate, is what DomainCreate leaves the
	// domain in instead of running, as a guest that dies at once would.
	stateAfterStart libvirt.DomainState
	// undefineFlags records the flags of the last DomainUndefineFlags per
	// domain name.
	undefineFlags map[string]libvirt.DomainUndefineFlagsValues

	calls  map[string]int
	log    []string
	closed int
	deadCh chan struct{}
	deadMu sync.Once
}

type fakeDomain struct {
	state    libvirt.DomainState
	xml      string
	metadata string
	addrs    map[uint32][]libvirt.DomainInterface
	uuid     libvirt.UUID
	id       int32
	// pendingStart counts the DomainGetState calls left before a delayed
	// start (see startDelayPolls) turns the domain running.
	pendingStart int
}

type fakePool struct {
	state     libvirt.StoragePoolState
	capacity  uint64
	available uint64
	vols      map[string]*fakeVol
}

type fakeVol struct {
	content  []byte
	capacity uint64
	// uploaded and sha256 record a discarded upload (see discardUploads).
	uploaded int64
	sha256   string
}

const fakePoolName = "default"

func newFakeBackend() *fakeBackend {
	return &fakeBackend{
		mu:            &sync.Mutex{},
		domains:       make(map[string]*fakeDomain),
		pools:         map[string]*fakePool{fakePoolName: newFakePool()},
		hostname:      "fake-host",
		version:       11006000,
		hooks:         make(map[string]func() error),
		block:         make(map[string]chan struct{}),
		undefineFlags: make(map[string]libvirt.DomainUndefineFlagsValues),
		calls:         make(map[string]int),
		deadCh:        make(chan struct{}),
	}
}

// twin returns a second connection to the same fake host: it shares the
// domains and pools (and their lock) but has its own hooks, counters, log,
// close count and liveness, so a test can tell which connection did what.
func (f *fakeBackend) twin() *fakeBackend {
	return &fakeBackend{
		mu:             f.mu,
		domains:        f.domains,
		pools:          f.pools,
		hostname:       f.hostname,
		version:        f.version,
		hooks:          make(map[string]func() error),
		block:          make(map[string]chan struct{}),
		discardUploads: f.discardUploads,
		undefineFlags:  make(map[string]libvirt.DomainUndefineFlagsValues),
		calls:          make(map[string]int),
		deadCh:         make(chan struct{}),
	}
}

// vol returns the volume by name in the default pool, or nil.
func (f *fakeBackend) vol(name string) *fakeVol {
	f.mu.Lock()
	defer f.mu.Unlock()
	return f.pools[fakePoolName].vols[name]
}

// closeCount is closed read under the lock, for twins closed by another
// goroutine.
func (f *fakeBackend) closeCount() int {
	f.mu.Lock()
	defer f.mu.Unlock()
	return f.closed
}

func newFakePool() *fakePool {
	return &fakePool{state: libvirt.StoragePoolRunning, capacity: 100 << 30, available: 50 << 30, vols: make(map[string]*fakeVol)}
}

func (f *fakeBackend) dead() <-chan struct{} { return f.deadCh }

func (f *fakeBackend) markDead() { f.deadMu.Do(func() { close(f.deadCh) }) }

// addDomain defines a domain directly, bypassing DomainDefineXML.
func (f *fakeBackend) addDomain(name string, state libvirt.DomainState, metadata string) *fakeDomain {
	f.mu.Lock()
	defer f.mu.Unlock()
	d := &fakeDomain{state: state, metadata: metadata, addrs: make(map[uint32][]libvirt.DomainInterface), uuid: fakeUUID(name)}
	f.domains[name] = d
	return d
}

func (f *fakeBackend) failWith(method string, err error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.hooks[method] = func() error { return err }
}

func (f *fakeBackend) blockOn(method string) chan struct{} {
	f.mu.Lock()
	defer f.mu.Unlock()
	ch := make(chan struct{})
	f.block[method] = ch
	return ch
}

func (f *fakeBackend) callCount(method string) int {
	f.mu.Lock()
	defer f.mu.Unlock()
	return f.calls[method]
}

func (f *fakeBackend) callLog() []string {
	f.mu.Lock()
	defer f.mu.Unlock()
	return slices.Clone(f.log)
}

func fakeUUID(name string) libvirt.UUID {
	var u libvirt.UUID
	copy(u[:], name)
	return u
}

func fakeErr(code libvirt.ErrorNumber, format string, args ...any) error {
	return libvirt.Error{Code: uint32(code), Message: fmt.Sprintf(format, args...)}
}

// enter records the call, honours a block and runs the hook. It is called
// without f.mu held; state access happens afterwards under the lock.
func (f *fakeBackend) enter(ctx context.Context, method string) error {
	f.mu.Lock()
	f.calls[method]++
	f.log = append(f.log, method)
	block := f.block[method]
	hook := f.hooks[method]
	f.mu.Unlock()
	if block != nil {
		select {
		case <-block:
		case <-ctx.Done():
			return ctx.Err()
		}
	}
	if hook != nil {
		return hook()
	}
	return nil
}

func (f *fakeBackend) domain(name string) (*fakeDomain, error) {
	d, ok := f.domains[name]
	if !ok {
		return nil, fakeErr(libvirt.ErrNoDomain, "Domain not found: no domain with matching name '%s'", name)
	}
	return d, nil
}

func (f *fakeBackend) pool(name string) (*fakePool, error) {
	p, ok := f.pools[name]
	if !ok {
		return nil, fakeErr(libvirt.ErrNoStoragePool, "Storage pool not found: no storage pool with matching name '%s'", name)
	}
	return p, nil
}

func (f *fakeBackend) volume(vol libvirt.StorageVol) (*fakeVol, error) {
	p, err := f.pool(vol.Pool)
	if err != nil {
		return nil, err
	}
	v, ok := p.vols[vol.Name]
	if !ok {
		return nil, fakeErr(libvirt.ErrNoStorageVol, "Storage volume not found: no storage vol with matching name '%s'", vol.Name)
	}
	return v, nil
}

func (f *fakeBackend) libvirtDomain(name string, d *fakeDomain) libvirt.Domain {
	id := int32(-1)
	if isActive(d.state) {
		id = d.id
	}
	return libvirt.Domain{Name: name, UUID: d.uuid, ID: id}
}

func (f *fakeBackend) DomainLookupByName(ctx context.Context, name string) (libvirt.Domain, error) {
	if err := f.enter(ctx, "DomainLookupByName"); err != nil {
		return libvirt.Domain{}, err
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	d, err := f.domain(name)
	if err != nil {
		return libvirt.Domain{}, err
	}
	return f.libvirtDomain(name, d), nil
}

func (f *fakeBackend) DomainDefineXML(ctx context.Context, doc string) (libvirt.Domain, error) {
	if err := f.enter(ctx, "DomainDefineXML"); err != nil {
		return libvirt.Domain{}, err
	}
	var parsed struct {
		Name     string `xml:"name"`
		Metadata struct {
			Inner string `xml:",innerxml"`
		} `xml:"metadata"`
	}
	if err := xml.Unmarshal([]byte(doc), &parsed); err != nil || parsed.Name == "" {
		return libvirt.Domain{}, fakeErr(libvirt.ErrXMLError, "XML error: %v", err)
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	d, ok := f.domains[parsed.Name]
	if !ok {
		d = &fakeDomain{state: libvirt.DomainShutoff, addrs: make(map[uint32][]libvirt.DomainInterface), uuid: fakeUUID(parsed.Name)}
		f.domains[parsed.Name] = d
	}
	d.xml = doc
	d.metadata = parsed.Metadata.Inner
	return f.libvirtDomain(parsed.Name, d), nil
}

func (f *fakeBackend) DomainCreate(ctx context.Context, dom libvirt.Domain) error {
	if err := f.enter(ctx, "DomainCreate"); err != nil {
		return err
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	d, err := f.domain(dom.Name)
	if err != nil {
		return err
	}
	if isActive(d.state) {
		return fakeErr(libvirt.ErrOperationInvalid, "Requested operation is not valid: domain is already running")
	}
	switch {
	case f.startDelayPolls > 0:
		d.state = libvirt.DomainNostate
		d.pendingStart = f.startDelayPolls
		return nil
	case f.stateAfterStart != libvirt.DomainNostate:
		d.state = f.stateAfterStart
		return nil
	}
	d.state = libvirt.DomainRunning
	d.id = int32(len(f.domains))
	return nil
}

func (f *fakeBackend) DomainDestroy(ctx context.Context, dom libvirt.Domain) error {
	if err := f.enter(ctx, "DomainDestroy"); err != nil {
		return err
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	d, err := f.domain(dom.Name)
	if err != nil {
		return err
	}
	if !isActive(d.state) {
		return fakeErr(libvirt.ErrOperationInvalid, "Requested operation is not valid: domain is not running")
	}
	d.state = libvirt.DomainShutoff
	return nil
}

func (f *fakeBackend) DomainUndefineFlags(ctx context.Context, dom libvirt.Domain, flags libvirt.DomainUndefineFlagsValues) error {
	if err := f.enter(ctx, "DomainUndefineFlags"); err != nil {
		return err
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	if _, err := f.domain(dom.Name); err != nil {
		return err
	}
	f.undefineFlags[dom.Name] = flags
	delete(f.domains, dom.Name)
	return nil
}

func (f *fakeBackend) DomainGetState(ctx context.Context, dom libvirt.Domain, _ uint32) (int32, int32, error) {
	if err := f.enter(ctx, "DomainGetState"); err != nil {
		return 0, 0, err
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	d, err := f.domain(dom.Name)
	if err != nil {
		return 0, 0, err
	}
	if d.pendingStart > 0 {
		d.pendingStart--
		if d.pendingStart == 0 {
			d.state = libvirt.DomainRunning
			d.id = int32(len(f.domains))
		}
	}
	return int32(d.state), 0, nil
}

func (f *fakeBackend) DomainGetMetadata(ctx context.Context, dom libvirt.Domain, typ int32, uri libvirt.OptString, _ libvirt.DomainModificationImpact) (string, error) {
	if err := f.enter(ctx, "DomainGetMetadata"); err != nil {
		return "", err
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	d, err := f.domain(dom.Name)
	if err != nil {
		return "", err
	}
	if typ != int32(libvirt.DomainMetadataElement) || len(uri) != 1 || uri[0] != metadataURI || d.metadata == "" {
		return "", fakeErr(libvirt.ErrNoDomainMetadata, "Requested metadata element is not present")
	}
	return d.metadata, nil
}

func (f *fakeBackend) DomainInterfaceAddresses(ctx context.Context, dom libvirt.Domain, source uint32, _ uint32) ([]libvirt.DomainInterface, error) {
	if err := f.enter(ctx, "DomainInterfaceAddresses"); err != nil {
		return nil, err
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	d, err := f.domain(dom.Name)
	if err != nil {
		return nil, err
	}
	if !isActive(d.state) {
		return nil, fakeErr(libvirt.ErrOperationInvalid, "Requested operation is not valid: domain is not running")
	}
	return slices.Clone(d.addrs[source]), nil
}

func (f *fakeBackend) ConnectListAllDomains(ctx context.Context, _ int32, _ libvirt.ConnectListAllDomainsFlags) ([]libvirt.Domain, uint32, error) {
	if err := f.enter(ctx, "ConnectListAllDomains"); err != nil {
		return nil, 0, err
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	names := make([]string, 0, len(f.domains))
	for name := range f.domains {
		names = append(names, name)
	}
	slices.Sort(names)
	out := make([]libvirt.Domain, 0, len(names))
	for _, name := range names {
		out = append(out, f.libvirtDomain(name, f.domains[name]))
	}
	return out, uint32(len(out)), nil
}

func (f *fakeBackend) StoragePoolLookupByName(ctx context.Context, name string) (libvirt.StoragePool, error) {
	if err := f.enter(ctx, "StoragePoolLookupByName"); err != nil {
		return libvirt.StoragePool{}, err
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	if _, err := f.pool(name); err != nil {
		return libvirt.StoragePool{}, err
	}
	return libvirt.StoragePool{Name: name, UUID: fakeUUID(name)}, nil
}

func (f *fakeBackend) StoragePoolRefresh(ctx context.Context, pool libvirt.StoragePool, _ uint32) error {
	if err := f.enter(ctx, "StoragePoolRefresh"); err != nil {
		return err
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	_, err := f.pool(pool.Name)
	return err
}

func (f *fakeBackend) StoragePoolGetInfo(ctx context.Context, pool libvirt.StoragePool) (uint8, uint64, uint64, uint64, error) {
	if err := f.enter(ctx, "StoragePoolGetInfo"); err != nil {
		return 0, 0, 0, 0, err
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	p, err := f.pool(pool.Name)
	if err != nil {
		return 0, 0, 0, 0, err
	}
	return uint8(p.state), p.capacity, p.capacity - p.available, p.available, nil
}

func (f *fakeBackend) StorageVolLookupByName(ctx context.Context, pool libvirt.StoragePool, name string) (libvirt.StorageVol, error) {
	if err := f.enter(ctx, "StorageVolLookupByName"); err != nil {
		return libvirt.StorageVol{}, err
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	vol := libvirt.StorageVol{Pool: pool.Name, Name: name, Key: fakeVolPath(pool.Name, name)}
	if _, err := f.volume(vol); err != nil {
		return libvirt.StorageVol{}, err
	}
	return vol, nil
}

func fakeVolPath(pool, name string) string { return "/pool/" + pool + "/" + name }

func (f *fakeBackend) StorageVolGetPath(ctx context.Context, vol libvirt.StorageVol) (string, error) {
	if err := f.enter(ctx, "StorageVolGetPath"); err != nil {
		return "", err
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	if _, err := f.volume(vol); err != nil {
		return "", err
	}
	return fakeVolPath(vol.Pool, vol.Name), nil
}

func (f *fakeBackend) StorageVolCreateXML(ctx context.Context, pool libvirt.StoragePool, doc string, _ libvirt.StorageVolCreateFlags) (libvirt.StorageVol, error) {
	if err := f.enter(ctx, "StorageVolCreateXML"); err != nil {
		return libvirt.StorageVol{}, err
	}
	var parsed struct {
		Name     string `xml:"name"`
		Capacity uint64 `xml:"capacity"`
	}
	if err := xml.Unmarshal([]byte(doc), &parsed); err != nil || parsed.Name == "" {
		return libvirt.StorageVol{}, fakeErr(libvirt.ErrXMLError, "XML error: %v", err)
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	p, err := f.pool(pool.Name)
	if err != nil {
		return libvirt.StorageVol{}, err
	}
	if _, exists := p.vols[parsed.Name]; exists {
		return libvirt.StorageVol{}, fakeErr(libvirt.ErrStorageVolExist, "storage volume '%s' exists already", parsed.Name)
	}
	p.vols[parsed.Name] = &fakeVol{capacity: parsed.Capacity}
	return libvirt.StorageVol{Pool: pool.Name, Name: parsed.Name, Key: fakeVolPath(pool.Name, parsed.Name)}, nil
}

func (f *fakeBackend) StorageVolUpload(ctx context.Context, vol libvirt.StorageVol, r io.Reader, _, _ uint64, _ libvirt.StorageVolUploadFlags) error {
	if err := f.enter(ctx, "StorageVolUpload"); err != nil {
		return err
	}
	var (
		content  []byte
		uploaded int64
		digest   string
	)
	if r != nil {
		var err error
		if f.discardUploads {
			h := sha256.New()
			if uploaded, err = io.CopyBuffer(h, r, make([]byte, 64<<10)); err != nil {
				return err
			}
			digest = hex.EncodeToString(h.Sum(nil))
		} else if content, err = io.ReadAll(r); err != nil {
			return err
		}
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	v, err := f.volume(vol)
	if err != nil {
		return err
	}
	v.content = content
	v.uploaded = uploaded
	v.sha256 = digest
	return nil
}

func (f *fakeBackend) StorageVolDownload(ctx context.Context, vol libvirt.StorageVol, w io.Writer, _, _ uint64, _ libvirt.StorageVolDownloadFlags) error {
	if err := f.enter(ctx, "StorageVolDownload"); err != nil {
		return err
	}
	f.mu.Lock()
	v, err := f.volume(vol)
	var content []byte
	if v != nil {
		content = slices.Clone(v.content)
	}
	f.mu.Unlock()
	if err != nil {
		return err
	}
	if w == nil {
		return nil
	}
	_, err = w.Write(content)
	return err
}

func (f *fakeBackend) StorageVolDelete(ctx context.Context, vol libvirt.StorageVol, _ libvirt.StorageVolDeleteFlags) error {
	if err := f.enter(ctx, "StorageVolDelete"); err != nil {
		return err
	}
	f.mu.Lock()
	defer f.mu.Unlock()
	if _, err := f.volume(vol); err != nil {
		return err
	}
	delete(f.pools[vol.Pool].vols, vol.Name)
	return nil
}

func (f *fakeBackend) ConnectGetLibVersion(ctx context.Context) (uint64, error) {
	if err := f.enter(ctx, "ConnectGetLibVersion"); err != nil {
		return 0, err
	}
	return f.version, nil
}

func (f *fakeBackend) ConnectGetHostname(ctx context.Context) (string, error) {
	if err := f.enter(ctx, "ConnectGetHostname"); err != nil {
		return "", err
	}
	return f.hostname, nil
}

func (f *fakeBackend) Close() error {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.closed++
	return nil
}

// stubDialer covers every dialer a test needs: it counts dials, fails with
// err when set, waits on block (or ctx) when set, and hands out backend, or
// one from fresh per call, or a new fakeBackend when neither is given.
type stubDialer struct {
	backend backend
	fresh   func() backend
	err     error
	block   chan struct{}

	dials atomic.Int32
	mu    sync.Mutex
	hosts []string
}

func (d *stubDialer) dial(ctx context.Context, host provider.KVMHost) (backend, error) {
	d.dials.Add(1)
	d.mu.Lock()
	d.hosts = append(d.hosts, host.Name)
	d.mu.Unlock()
	if d.block != nil {
		select {
		case <-d.block:
		case <-ctx.Done():
			return nil, fmt.Errorf("%w: dialing %s: %w", provider.ErrTransient, host.Name, ctx.Err())
		}
	}
	if d.err != nil {
		return nil, d.err
	}
	switch {
	case d.fresh != nil:
		return d.fresh(), nil
	case d.backend != nil:
		return d.backend, nil
	default:
		return newFakeBackend(), nil
	}
}