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

package kvm

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"sync"
	"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
}

// 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)

// 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
}

type hostEntry struct {
	mu     sync.Mutex
	client *hostClient
}

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{}
		p.entries[name] = e
	}
	return e
}

func (p *hostPool) get(ctx context.Context, log logr.Logger, host provider.KVMHost) (*hostClient, error) {
	e := p.entry(host.Name)
	e.mu.Lock()
	defer e.mu.Unlock()
	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.mu.Lock()
	defer e.mu.Unlock()
	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 (anything that is not a libvirt.Error) also drops the
// cached connection so the next call redials instead of reusing a dead
// 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
	}
	api := c.api
	if log.V(1).Enabled() {
		api = &tracingAPI{inner: c.api, log: log}
	}
	err = awaitOrCancel(ctx, func() error { return fn(api) }, func() {
		log.V(1).Info("libvirt connection dropped", "reason", "context cancelled", "error", ctx.Err().Error())
		p.drop(host.Name, c)
	})
	if err != nil && !isLibvirtError(err) && !errors.Is(err, ctx.Err()) {
		log.V(1).Info("libvirt connection dropped", "reason", "transport error", "error", err.Error())
		p.drop(host.Name, c)
	}
	return err
}

// 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 loaded once at startup so a
// missing or unreadable file fails the operator fast rather than on the
// first Create.
type sshAuth struct {
	signer  ssh.Signer
	hostKey ssh.HostKeyCallback
}

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("parsing kvm.sshKeyFile: %w", err)
	}
	hostKey := ssh.InsecureIgnoreHostKey() //nolint:gosec // opt-in via config, documented as insecure
	if !cfg.InsecureIgnoreHostKey {
		hostKey, err = knownhosts.New(cfg.KnownHostsFile)
		if err != nil {
			return nil, fmt.Errorf("loading kvm.knownHostsFile: %w", err)
		}
	}
	return &sshAuth{signer: signer, hostKey: hostKey}, nil
}

// dialer returns the production hostDialer: SSH to the host, forward its
// libvirt unix socket through the connection, and open a libvirt session
// over that.
func (a *sshAuth) dialer() hostDialer {
	return func(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, 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, 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, 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
		}
	}
}