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

package kvm

import (
	"context"
	"errors"
	"fmt"
	"io"
	"io/fs"
	"net"
	"os"
	"sync"
	"sync/atomic"
	"time"

	libvirt "github.com/digitalocean/go-libvirt"
	"github.com/go-logr/logr"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/knownhosts"

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

const (
	defaultHandshakeTimeout = 30 * time.Second
	keepaliveInterval       = 15 * time.Second
)

// libvirtAPI is the slice of *libvirt.Libvirt the provider uses, and the
// test seam: the fake in kvm_test.go implements exactly this.
type libvirtAPI interface {
	DomainLookupByName(name string) (libvirt.Domain, error)
	DomainDefineXML(xml string) (libvirt.Domain, error)
	DomainCreate(d libvirt.Domain) error
	DomainDestroy(d libvirt.Domain) error
	DomainUndefineFlags(d libvirt.Domain, flags libvirt.DomainUndefineFlagsValues) error
	DomainGetState(d libvirt.Domain, flags uint32) (state int32, reason int32, err error)
	DomainGetMetadata(d libvirt.Domain, typ int32, uri libvirt.OptString, flags libvirt.DomainModificationImpact) (string, error)
	DomainInterfaceAddresses(d libvirt.Domain, source uint32, flags uint32) ([]libvirt.DomainInterface, error)
	ConnectListAllDomains(needResults int32, flags libvirt.ConnectListAllDomainsFlags) ([]libvirt.Domain, uint32, error)
	StoragePoolLookupByName(name string) (libvirt.StoragePool, error)
	StoragePoolRefresh(pool libvirt.StoragePool, flags uint32) error
	StorageVolLookupByName(pool libvirt.StoragePool, name string) (libvirt.StorageVol, error)
	StorageVolGetPath(vol libvirt.StorageVol) (string, error)
	StorageVolCreateXML(pool libvirt.StoragePool, xml string, flags libvirt.StorageVolCreateFlags) (libvirt.StorageVol, error)
	StorageVolUpload(vol libvirt.StorageVol, r io.Reader, offset, length uint64, flags libvirt.StorageVolUploadFlags) error
	StorageVolDownload(vol libvirt.StorageVol, w io.Writer, offset, length uint64, flags libvirt.StorageVolDownloadFlags) error
	StorageVolDelete(vol libvirt.StorageVol, flags libvirt.StorageVolDeleteFlags) error
	ConnectGetLibVersion() (uint64, error)
	ConnectGetHostname() (string, error)
	StoragePoolGetInfo(pool libvirt.StoragePool) (state uint8, capacity, allocation, available uint64, err error)
}

// hostClient is one live connection to a hypervisor.
type hostClient struct {
	api libvirtAPI
	// done is closed once the connection is known to be lost. A nil channel
	// never fires, which suits fakes.
	done  <-chan struct{}
	close func()
}

func (c *hostClient) alive() bool {
	select {
	case <-c.done:
		return false
	default:
		return true
	}
}

// hostDialer opens a connection to one host. The SSH implementation is the
// real one; tests inject their own.
type hostDialer func(ctx context.Context, host provider.KVMHost) (*hostClient, error)

// dialStage names the step of a connection attempt that failed, so the
// health check can say why a host is unreachable without parsing messages.
type dialStage string

const (
	stageSSHMaterial    dialStage = "ssh-material"
	stageTCP            dialStage = "tcp"
	stageSSHHandshake   dialStage = "ssh-handshake"
	stageLibvirtConnect dialStage = "libvirt-connect"
)

// dialError tags a dial failure with its stage. It keeps the message and
// unwraps to the cause, so classify sees exactly what it saw before.
type dialError struct {
	stage dialStage
	err   error
}

func (e *dialError) Error() string { return e.err.Error() }
func (e *dialError) Unwrap() error { return e.err }

// hostPool caches one connection per host, dialing lazily and redialing
// after the connection drops. Dials for different hosts don't block each
// other; concurrent callers for the same host share a single dial.
type hostPool struct {
	dial hostDialer

	mu      sync.Mutex
	entries map[string]*hostEntry
}

// hostEntry is one host's slot in the pool. lock is a one-slot semaphore
// rather than a mutex so that waiting for it can end with the caller's
// context: whoever holds it may be dialing a dead host for as long as the
// dial takes.
type hostEntry struct {
	lock   chan struct{}
	client *hostClient
	// probing is set while a probe's fn runs, including one whose caller
	// has already given up on it.
	probing atomic.Bool
}

func newHostPool(dial hostDialer) *hostPool {
	return &hostPool{dial: dial, entries: make(map[string]*hostEntry)}
}

func (p *hostPool) entry(name string) *hostEntry {
	p.mu.Lock()
	defer p.mu.Unlock()
	e, ok := p.entries[name]
	if !ok {
		e = &hostEntry{lock: make(chan struct{}, 1)}
		p.entries[name] = e
	}
	return e
}

// acquire takes e.lock. Only waiting is bounded by ctx: a free slot is
// taken even when ctx has already ended, so an uncontended caller behaves
// as it did with a mutex, and a contended one gives up with ctx's error.
func (e *hostEntry) acquire(ctx context.Context) error {
	select {
	case e.lock <- struct{}{}:
		return nil
	default:
	}
	select {
	case e.lock <- struct{}{}:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

func (e *hostEntry) release() { <-e.lock }

func (p *hostPool) get(ctx context.Context, log logr.Logger, host provider.KVMHost) (*hostClient, error) {
	e := p.entry(host.Name)
	if err := e.acquire(ctx); err != nil {
		return nil, fmt.Errorf("waiting for the connection to host %s (%s): %w", host.Name, host.Address, err)
	}
	defer e.release()
	if e.client != nil {
		if e.client.alive() {
			return e.client, nil
		}
		log.V(1).Info("libvirt connection lost, redialing", "address", host.Address)
		e.client.close()
		e.client = nil
	}
	start := time.Now()
	c, err := p.dial(ctx, host)
	if err != nil {
		log.V(1).Info("libvirt host dial failed", "address", host.Address, "duration", time.Since(start).Round(time.Millisecond), "error", err.Error())
		return nil, fmt.Errorf("connecting to host %s (%s): %w", host.Name, host.Address, err)
	}
	log.V(1).Info("libvirt host connected", "address", host.Address, "duration", time.Since(start).Round(time.Millisecond))
	e.client = c
	return c, nil
}

// drop closes and forgets c if it is still the cached client for host, so a
// connection that was already replaced by a redial is left alone.
func (p *hostPool) drop(host string, c *hostClient) {
	e := p.entry(host)
	e.lock <- struct{}{}
	defer e.release()
	if e.client == c {
		e.client = nil
	}
	c.close()
}

// call runs fn against host's connection. go-libvirt RPCs take no context
// and block until the daemon answers, so when ctx ends first the connection
// is closed — which unblocks fn with an error — and ctx.Err() is returned.
// A transport error also drops the cached connection so the next call
// redials instead of reusing a dead link. "Transport error" means anything
// that is neither a libvirt.Error nor an error fn already classified with a
// provider sentinel (a failed seed build, a bad placement): those are local
// verdicts, not evidence about the link. When log has V(1) enabled, every
// RPC fn makes is traced through it (see tracingAPI), as are connection
// events.
func (p *hostPool) call(ctx context.Context, log logr.Logger, host provider.KVMHost, fn func(api libvirtAPI) error) error {
	log = log.WithName("libvirt")
	c, err := p.get(ctx, log, host)
	if err != nil {
		return err
	}
	err = awaitOrCancel(ctx, func() error { return fn(traced(c, log)) }, func() {
		log.V(1).Info("libvirt connection dropped", "reason", "context cancelled", "error", ctx.Err().Error())
		p.drop(host.Name, c)
	})
	p.dropOnTransportError(ctx, log, host, c, err)
	return err
}

// errProbeBusy refuses a probe while the previous one on the same host has
// not come back yet.
var errProbeBusy = errors.New("the previous check of this host is still waiting for libvirt to answer")

// probe is call for the health check, which must never cost the reconciles
// their connection: when ctx ends first, fn is abandoned and the connection
// left alone. A link that is really dead is closed by the keepalive, which
// ends fn too. At most one probe per host runs at a time, an abandoned one
// included, so a daemon that stopped answering cannot pile up goroutines;
// until it answers, further probes get errProbeBusy.
func (p *hostPool) probe(ctx context.Context, log logr.Logger, host provider.KVMHost, fn func(api libvirtAPI) error) error {
	log = log.WithName("libvirt")
	e := p.entry(host.Name)
	if !e.probing.CompareAndSwap(false, true) {
		return errProbeBusy
	}
	c, err := p.get(ctx, log, host)
	if err != nil {
		e.probing.Store(false)
		return err
	}
	err = awaitOrCancel(ctx, func() error {
		defer e.probing.Store(false)
		return fn(traced(c, log))
	}, func() {
		log.V(1).Info("libvirt probe abandoned, connection kept", "error", ctx.Err().Error())
	})
	p.dropOnTransportError(ctx, log, host, c, err)
	return err
}

// traced returns c's API, logging every RPC when log has V(1) enabled.
func traced(c *hostClient, log logr.Logger) libvirtAPI {
	if log.V(1).Enabled() {
		return &tracingAPI{inner: c.api, log: log}
	}
	return c.api
}

// dropOnTransportError drops c when err is evidence against the link:
// neither a libvirt error, nor a verdict already classified, nor the
// caller's own context ending.
func (p *hostPool) dropOnTransportError(ctx context.Context, log logr.Logger, host provider.KVMHost, c *hostClient, err error) {
	if err != nil && !isLibvirtError(err) && !isClassified(err) && !errors.Is(err, ctx.Err()) {
		log.V(1).Info("libvirt connection dropped", "reason", "transport error", "error", err.Error())
		p.drop(host.Name, c)
	}
}

// awaitOrCancel runs fn in the background and returns its result, unless ctx
// ends first: then abort is invoked (it must make fn return eventually, e.g.
// by closing the connection fn is blocked on) and ctx.Err() is returned. The
// abandoned fn's result is discarded.
func awaitOrCancel(ctx context.Context, fn func() error, abort func()) error {
	errc := make(chan error, 1)
	go func() { errc <- fn() }()
	select {
	case err := <-errc:
		return err
	case <-ctx.Done():
		abort()
		return ctx.Err()
	}
}

// sshAuth holds the pieces of the SSH client config shared by every host:
// the private key and the host-key policy. Both are loaded on the first
// connection rather than at startup — see lazySSH.
type sshAuth struct {
	signer  ssh.Signer
	hostKey ssh.HostKeyCallback
}

// newSSHAuth loads the key material. A file that is not there yet stays
// unclassified, so the provider treats it as transient and heals when the
// Secret appears; content that will never parse is permanent, since
// retrying cannot fix it.
func newSSHAuth(cfg provider.KVMConfig) (*sshAuth, error) {
	key, err := os.ReadFile(cfg.SSHKeyFile)
	if err != nil {
		return nil, fmt.Errorf("reading kvm.sshKeyFile: %w", err)
	}
	signer, err := ssh.ParsePrivateKey(key)
	if err != nil {
		return nil, fmt.Errorf("%w: parsing kvm.sshKeyFile: %v", provider.ErrPermanent, err)
	}
	hostKey := ssh.InsecureIgnoreHostKey() //nolint:gosec // opt-in via config, documented as insecure
	if !cfg.InsecureIgnoreHostKey {
		hostKey, err = knownhosts.New(cfg.KnownHostsFile)
		if errors.Is(err, fs.ErrNotExist) {
			return nil, fmt.Errorf("loading kvm.knownHostsFile: %w", err)
		}
		if err != nil {
			return nil, fmt.Errorf("%w: loading kvm.knownHostsFile: %v", provider.ErrPermanent, err)
		}
	}
	return &sshAuth{signer: signer, hostKey: hostKey}, nil
}

// lazySSH defers loading the SSH material until it is first needed and
// caches it once it works. The kvm-ssh Secret is an optional mount: the
// manager has to start without it and keep every other provider running,
// and a Secret created afterwards has to heal the KVM proxies on the next
// reconcile rather than needing a restart.
type lazySSH struct {
	cfg provider.KVMConfig

	mu   sync.Mutex
	auth *sshAuth
}

// newLazySSH keeps only the SSH fields of cfg, since the loader lives as
// long as the Provider.
func newLazySSH(cfg provider.KVMConfig) *lazySSH {
	return &lazySSH{cfg: provider.KVMConfig{
		SSHKeyFile:            cfg.SSHKeyFile,
		KnownHostsFile:        cfg.KnownHostsFile,
		InsecureIgnoreHostKey: cfg.InsecureIgnoreHostKey,
	}}
}

func (l *lazySSH) load() (*sshAuth, error) {
	l.mu.Lock()
	defer l.mu.Unlock()
	if l.auth != nil {
		return l.auth, nil
	}
	auth, err := newSSHAuth(l.cfg)
	if err != nil {
		return nil, err
	}
	l.auth = auth
	return auth, nil
}

// dial is the production hostDialer.
func (l *lazySSH) dial(ctx context.Context, host provider.KVMHost) (*hostClient, error) {
	a, err := l.load()
	if err != nil {
		return nil, &dialError{stage: stageSSHMaterial, err: err}
	}
	return a.dial(ctx, host)
}

// dial SSHes to the host, forwards its libvirt unix socket through the
// connection, and opens a libvirt session over that.
func (a *sshAuth) dial(ctx context.Context, host provider.KVMHost) (*hostClient, error) {
	client, err := a.dialSSH(ctx, host)
	if err != nil {
		return nil, err
	}
	hc, err := connectLibvirt(ctx, client, host.Socket)
	if err != nil {
		client.Close()
		return nil, err
	}
	return hc, nil
}

// dialSSH establishes the SSH connection. The handshake is bounded by ctx's
// deadline, or defaultHandshakeTimeout when ctx has none, since a
// half-open hypervisor must not pin a reconcile worker indefinitely.
func (a *sshAuth) dialSSH(ctx context.Context, host provider.KVMHost) (*ssh.Client, error) {
	deadline, ok := ctx.Deadline()
	if !ok {
		deadline = time.Now().Add(defaultHandshakeTimeout)
	}
	var d net.Dialer
	nc, err := d.DialContext(ctx, "tcp", host.Address)
	if err != nil {
		return nil, &dialError{stage: stageTCP, err: fmt.Errorf("ssh dial: %w", err)}
	}
	_ = nc.SetDeadline(deadline)
	conn, chans, reqs, err := ssh.NewClientConn(nc, host.Address, &ssh.ClientConfig{
		User:            host.User,
		Auth:            []ssh.AuthMethod{ssh.PublicKeys(a.signer)},
		HostKeyCallback: a.hostKey,
	})
	if err != nil {
		nc.Close()
		return nil, &dialError{stage: stageSSHHandshake, err: fmt.Errorf("ssh handshake: %w", err)}
	}
	_ = nc.SetDeadline(time.Time{})
	return ssh.NewClient(conn, chans, reqs), nil
}

// unixForward adapts an SSH client to go-libvirt's socket.Dialer: each Dial
// opens a direct-streamlocal channel to the libvirt socket on the host.
type unixForward struct {
	client *ssh.Client
	path   string
}

func (u unixForward) Dial() (net.Conn, error) {
	return u.client.Dial("unix", u.path)
}

// connectLibvirt opens the libvirt session over the SSH client and starts a
// keepalive loop so a dead link is noticed within roughly two intervals
// rather than at the TCP stack's leisure. Closing the returned client tears
// down the SSH connection, which in turn ends the libvirt session.
func connectLibvirt(ctx context.Context, client *ssh.Client, socket string) (*hostClient, error) {
	lv := libvirt.NewWithDialer(unixForward{client: client, path: socket})
	if err := awaitOrCancel(ctx, lv.Connect, func() { client.Close() }); err != nil {
		return nil, &dialError{stage: stageLibvirtConnect, err: fmt.Errorf("libvirt connect via %s: %w", socket, err)}
	}
	stop := make(chan struct{})
	var once sync.Once
	go keepalive(client, keepaliveInterval, stop)
	return &hostClient{
		api:  lv,
		done: lv.Disconnected(),
		close: func() {
			once.Do(func() {
				close(stop)
				client.Close()
			})
		},
	}, nil
}

// keepalive sends OpenSSH keepalive requests until stop closes; a failed or
// unanswered request closes the client, which surfaces as Disconnected on
// the libvirt side and makes the pool redial.
func keepalive(client *ssh.Client, interval time.Duration, stop <-chan struct{}) {
	t := time.NewTicker(interval)
	defer t.Stop()
	for {
		select {
		case <-stop:
			return
		case <-t.C:
		}
		replied := make(chan error, 1)
		go func() {
			_, _, err := client.SendRequest("keepalive@openssh.com", true, nil)
			replied <- err
		}()
		select {
		case <-stop:
			return
		case err := <-replied:
			if err != nil {
				client.Close()
				return
			}
		case <-time.After(interval):
			client.Close()
			return
		}
	}
}