a73x

internal/agent/syncclient/client_test.go

Ref:   Size: 18.2 KiB   History

package syncclient

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"sync/atomic"
	"testing"
	"time"

	"github.com/a73x/eitri/internal/agent/reconcile"
	"github.com/a73x/eitri/internal/agent/seed"
	"github.com/a73x/eitri/internal/agent/selfupdate"
	"github.com/a73x/eitri/internal/agent/state"
	"github.com/a73x/eitri/internal/pb"
	"github.com/a73x/eitri/internal/server/hosttoken"
	"github.com/a73x/eitri/internal/server/hub"
	"github.com/a73x/eitri/internal/server/registry"
	"github.com/a73x/eitri/internal/server/store"
	"github.com/a73x/eitri/internal/server/syncsvc"
	"github.com/a73x/eitri/internal/transport"
	"github.com/a73x/eitri/internal/version"
	"github.com/quic-go/quic-go"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// noopProv satisfies the reconcile provisioner seam with no side effects so we
// can drive a real Client.Run against a real syncsvc server over QUIC loopback.
type noopProv struct{}

func (noopProv) Preflight(context.Context) error                             { return nil }
func (noopProv) PrepareRootDisk(context.Context, state.VMSpec, string) error { return nil }
func (noopProv) Boot(context.Context, string, state.VMSpec) error            { return nil }
func (noopProv) Shutdown(context.Context, string) error                      { return nil }
func (noopProv) Destroy(context.Context, string) error                       { return nil }
func (noopProv) Running(string) bool                                         { return false }
func (noopProv) Address(string) string                                       { return "10.77.1.2" }
func (noopProv) NetworkAddress(string) string                                { return "" }
func (noopProv) FailureReason(string) string                                 { return "" }

// testQUICIdle is the deliberately-short idle timeout the test listeners use so
// reconnect/drop cases don't wait out the production SyncMaxIdleTimeout.
const testQUICIdle = 5 * time.Second

// TestHelloFactsCarryAgentVersion proves helloFacts stamps the running
// binary's version onto the host facts it sends in Hello — the one fact
// the agent knows about itself rather than the host.
func TestHelloFactsCarryAgentVersion(t *testing.T) {
	f := helloFacts(context.Background(), nil)
	if got, want := f.GetAgentVersion(), version.Version; got != want {
		t.Fatalf("agent_version = %q, want %q", got, want)
	}
}

// TestAdvertisedCapacityComputesOnce proves the Statfs+Sysinfo capacity probe
// runs a single time per client (host totals are fixed for the session) while
// the cheap cap clamp still applies on every call.
func TestAdvertisedCapacityComputesOnce(t *testing.T) {
	orig := computeCapacity
	defer func() { computeCapacity = orig }()
	var calls int
	computeCapacity = func(string) *pb.Capacity {
		calls++
		return &pb.Capacity{Vcpus: 8, MemMb: 16384, DiskGb: 500}
	}
	c := &Client{MaxVCPUs: 2} // cap vCPUs to prove the clamp runs each call
	for range 5 {
		got := c.advertisedCapacity("/state")
		require.Equal(t, int64(2), got.GetVcpus())
		require.Equal(t, int64(16384), got.GetMemMb())
	}
	require.Equal(t, 1, calls, "raw capacity syscalls must run exactly once, not per report")
}

// TestAdvertisedCapacityReprobesUntilSuccess proves a transiently-failed probe
// (capacity() silently zeroes a dimension whose syscall failed) is NOT memoized:
// the agent re-probes until it gets a full reading, then caches that — otherwise
// a boot-time blip would freeze a 0-capacity advertisement for the whole session
// where the old per-report probe self-healed within one tick.
func TestAdvertisedCapacityReprobesUntilSuccess(t *testing.T) {
	orig := computeCapacity
	defer func() { computeCapacity = orig }()
	var calls int
	computeCapacity = func(string) *pb.Capacity {
		calls++
		if calls == 1 {
			return &pb.Capacity{Vcpus: 8, MemMb: 0, DiskGb: 500} // transient mem probe failure
		}
		return &pb.Capacity{Vcpus: 8, MemMb: 16384, DiskGb: 500}
	}
	c := &Client{}

	got1 := c.advertisedCapacity("/state")
	require.Equal(t, int64(0), got1.GetMemMb(), "first (failed) probe surfaces the zero, not cached")

	got2 := c.advertisedCapacity("/state")
	require.Equal(t, int64(16384), got2.GetMemMb(), "must re-probe after a zeroed dimension")
	require.Equal(t, 2, calls)

	got3 := c.advertisedCapacity("/state")
	require.Equal(t, int64(16384), got3.GetMemMb())
	require.Equal(t, 2, calls, "a full reading is memoized: no further syscalls")
}

// TestHeartbeatMarginInvariant guards the cross-package coupling documented on
// DefaultTickInterval and registry.OnlineWindow: the server must not declare a
// host offline before it has had room to miss ~3 reports, or healthy hosts flap.
func TestHeartbeatMarginInvariant(t *testing.T) {
	require.GreaterOrEqual(t, registry.OnlineWindow, 3*DefaultTickInterval,
		"OnlineWindow (%s) must stay >= 3x the agent report tick (%s) to avoid Online/Offline flapping",
		registry.OnlineWindow, DefaultTickInterval)
}

// genTestCert returns a fresh self-signed server cert+key and its pinned
// fingerprint — the bootstrap every test QUIC server shares.
func genTestCert(t *testing.T) (certPEM, keyPEM []byte, fp string) {
	t.Helper()
	var err error
	certPEM, keyPEM, err = transport.GenerateServerCert()
	require.NoError(t, err)
	fp, err = transport.CertFingerprint(certPEM)
	require.NoError(t, err)
	return
}

// listenTestQUIC opens a loopback QUIC listener with the shared short test idle
// timeout, from an already-generated cert+key (so a harness can re-listen on the
// same port with the same identity).
func listenTestQUIC(t *testing.T, addr string, certPEM, keyPEM []byte) *quic.Listener {
	t.Helper()
	tlsConf, err := transport.ServerTLS(certPEM, keyPEM)
	require.NoError(t, err)
	lis, err := quic.ListenAddr(addr, tlsConf, &quic.Config{MaxIdleTimeout: testQUICIdle})
	require.NoError(t, err)
	return lis
}

// serverHarness owns a real syncsvc.Service on a fixed UDP loopback port so it
// can be stopped and restarted (TestReconnectAfterDrop) on the same address.
type serverHarness struct {
	t       *testing.T
	st      *store.Store
	reg     *registry.Registry
	hub     *hub.Hub
	secret  []byte
	certPEM []byte
	keyPEM  []byte
	fp      string
	addr    string
	cancel  context.CancelFunc
	lis     *quic.Listener
	svc     *syncsvc.Service
}

// testTenant is the tenant the harness provisions — through the real JIT
// path, the only way tenants are born. The name has no significance.
const testTenant = "default"

// seedTestTenant JIT-provisions testTenant on a fresh store.
func seedTestTenant(t *testing.T, st *store.Store) {
	t.Helper()
	_, err := st.CreateTenantForIdentity("https://test-issuer", "test-subject", testTenant+"@test.local")
	require.NoError(t, err)
}

func newServerHarness(t *testing.T) *serverHarness {
	t.Helper()
	st, err := store.Open(t.TempDir()+"/db", "10.77.0.0/16")
	require.NoError(t, err)
	t.Cleanup(func() { st.Close() })
	seedTestTenant(t, st)
	certPEM, keyPEM, fp := genTestCert(t)
	h := &serverHarness{
		t: t, st: st, reg: registry.New(time.Now), hub: hub.New(),
		secret: []byte("s3cret"), certPEM: certPEM, keyPEM: keyPEM, fp: fp,
	}
	h.start("127.0.0.1:0")
	t.Cleanup(h.stop)
	return h
}

func (h *serverHarness) start(addr string) {
	h.lis = listenTestQUIC(h.t, addr, h.certPEM, h.keyPEM)
	h.addr = h.lis.Addr().String()
	ctx, cancel := context.WithCancel(context.Background())
	h.cancel = cancel
	h.svc = syncsvc.New(h.st, h.reg, h.hub, h.secret, 0)
	go h.svc.Serve(ctx, h.lis) //nolint:errcheck
}

func (h *serverHarness) stop() {
	if h.cancel != nil {
		h.cancel()
	}
	if h.lis != nil {
		h.lis.Close()
	}
}

// enroll creates a host and returns a valid credential for it.
func (h *serverHarness) enroll() (hostID, cred string) {
	tok, _ := h.st.CreateEnrollmentToken(testTenant)
	host, err := h.st.RedeemEnrollmentToken(tok, store.EnrollFacts{Name: "h", OS: "linux", Arch: "amd64", Provisioner: "cloudhv", Remote: ""})
	require.NoError(h.t, err)
	return host.ID, hosttoken.Mint(h.secret, host.ID, host.CredGeneration, time.Now())
}

func newClient(t *testing.T, addr, fp, hostID, cred string) *Client {
	t.Helper()
	agentSt, err := state.Open(t.TempDir())
	require.NoError(t, err)
	id := state.Identity{
		HostID: hostID, Credential: cred,
		ServerQUICAddr: addr, ServerCertSHA256: fp,
	}
	require.NoError(t, agentSt.SaveIdentity(id))
	engine := &reconcile.Engine{
		St: agentSt, Prov: noopProv{},
		Images:  func(context.Context, string, string, func(int64, int64)) (string, error) { return "/x.raw", nil },
		Seed:    func(string, seed.Params) error { return nil },
		HostKey: state.LoadOrCreateHostKey,
		// CIDR/grace not exercised by these tests.
		BootID: func() string { return "boot-test" }, Now: time.Now,
		TombstoneGrace: time.Hour, VanishGrace: time.Hour, MaxCreateAttempts: 3,
	}
	return &Client{
		Engine: engine, St: agentSt, Identity: id,
		StateDir: t.TempDir(), TickInterval: 100 * time.Millisecond,
		ReconnectBackoff: 100 * time.Millisecond,
	}
}

// TestHelloAdvertisesHostNetworks pins the advertisement: the named networks
// this agent is configured to serve ride in the Hello, which is where the
// control plane's admission reads them. In the handshake and not in the report
// because the set is this agent's command line — it changes only with a
// restart, and a restart is a new session.
func TestHelloAdvertisesHostNetworks(t *testing.T) {
	certPEM, keyPEM, fp := genTestCert(t)
	lis := listenTestQUIC(t, "127.0.0.1:0", certPEM, keyPEM)
	defer lis.Close()

	hellos := make(chan *pb.Hello, 1)
	acceptCtx, stopAccept := context.WithCancel(context.Background())
	defer stopAccept()
	go func() {
		conn, err := lis.Accept(acceptCtx)
		if err != nil {
			return
		}
		up, err := conn.AcceptStream(acceptCtx)
		if err != nil {
			return
		}
		var first pb.AgentMessage
		if err := transport.ReadMsg(up, &first, transport.DefaultMaxFrame); err != nil {
			return
		}
		hellos <- first.GetHello()
		_ = conn.CloseWithError(0, "captured")
	}()

	c := newClient(t, lis.Addr().String(), fp, "host-nets", "host-nets.cred")
	c.HostNetworks = []string{"lab", "lan"}
	runCtx, stopRun := context.WithCancel(context.Background())
	done := make(chan struct{})
	go func() { defer close(done); c.Run(runCtx) }()
	defer func() { stopRun(); <-done }()

	select {
	case h := <-hellos:
		assert.Equal(t, []string{"lab", "lan"}, h.GetHostNetworks())
	case <-time.After(10 * time.Second):
		t.Fatal("no Hello arrived")
	}
}

func TestReconnectAfterDrop(t *testing.T) {
	h := newServerHarness(t)
	hostID, cred := h.enroll()
	c := newClient(t, h.addr, h.fp, hostID, cred)

	// Seed a VM so there's content; reports flow on each tick.
	require.NoError(t, h.st.CreateVM(store.VM{ID: "vm1", HostID: hostID, Name: "a",
		ImageURL: "u", ImageSHA256: "s", VCPUs: 1, MemMB: 512, DiskGB: 5, PowerState: "stopped"}))

	ctx := t.Context()
	go c.Run(ctx)

	// First connection: wait for the host to register a report.
	require.Eventually(t, func() bool {
		_, ok := h.reg.Get(hostID)
		return ok
	}, 5*time.Second, 50*time.Millisecond, "client should connect and report initially")

	addr := h.addr
	// Drop the server, then bring it back on the SAME UDP port.
	h.stop()
	time.Sleep(200 * time.Millisecond)
	h.reg = registry.New(time.Now) // fresh registry: a re-report proves reconnect
	h.start(addr)

	require.Eventually(t, func() bool {
		_, ok := h.reg.Get(hostID)
		return ok
	}, 15*time.Second, 100*time.Millisecond, "client should reconnect after the server restarts on the same port")
}

func TestAuthRejectedClassified(t *testing.T) {
	h := newServerHarness(t)
	hostID, _ := h.enroll()
	// Use a junk credential so the server rejects every connection with
	// CodeAuthRejected.
	c := newClient(t, h.addr, h.fp, hostID, "host-x.deadbeef")

	// session() is the unit that classifies the server's rejection; assert it
	// returns the permanent-auth sentinel (not a transient error). This is what
	// drives Run's 60s backoff branch instead of the 5s transient retry.
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	err := c.session(ctx)
	require.ErrorIs(t, err, errPermanentAuth,
		"a server auth rejection must classify as permanent, not transient")

	// And confirm Run takes the long-backoff path and does NOT tight-loop. We
	// hook dial-counting by counting connections the server accepts: after the
	// first rejected session Run sleeps 60s, so within ~1s the server must see at
	// most one connection attempt.
	var accepts atomic.Int64
	cs := newCountingServer(t, accepts.Add)
	csHostID, _ := cs.enroll()
	dc := newClient(t, cs.addr, cs.fp, csHostID, "host-x.deadbeef")
	rctx, rcancel := context.WithCancel(context.Background())
	go dc.Run(rctx)
	time.Sleep(1 * time.Second)
	rcancel()
	require.LessOrEqual(t, accepts.Load(), int64(1),
		"Run must not tight-loop on a permanently rejected credential")
}

// echoConsole implements the Console interface: replays a fixed backlog, then
// echoes input back uppercased — enough to prove both directions + ordering.
type echoConsole struct{ backlog string }

func (e *echoConsole) Attach(ctx context.Context, vmID string, rw io.ReadWriter, onReady func() error) error {
	if vmID != "vm-ok" {
		return fmt.Errorf("unknown vm %q", vmID)
	}
	if err := onReady(); err != nil {
		return err
	}
	if _, err := io.WriteString(rw, e.backlog); err != nil {
		return err
	}
	buf := make([]byte, 64)
	for {
		n, err := rw.Read(buf)
		if n > 0 {
			if _, werr := rw.Write(bytes.ToUpper(buf[:n])); werr != nil {
				return werr
			}
		}
		if err != nil {
			return nil //nolint:nilerr // stream closed by peer = clean end
		}
	}
}

// startConsoleClient runs a real client (with the given Console handler, which
// may be nil) against a fresh harness and blocks until the host is connected —
// the precondition for OpenConsole to find a live connection.
func startConsoleClient(t *testing.T, console Console) (h *serverHarness, hostID string) {
	t.Helper()
	h = newServerHarness(t)
	hostID, cred := h.enroll()
	c := newClient(t, h.addr, h.fp, hostID, cred)
	c.Console = console
	ctx, cancel := context.WithCancel(context.Background())
	t.Cleanup(cancel)
	go c.Run(ctx)
	require.Eventually(t, func() bool {
		_, ok := h.reg.Get(hostID)
		return ok
	}, 5*time.Second, 50*time.Millisecond, "client should connect and report")
	return h, hostID
}

func TestConsoleStreamEndToEnd(t *testing.T) {
	h, hostID := startConsoleClient(t, &echoConsole{backlog: "BOOT|"})

	stream, err := h.svc.OpenConsole(context.Background(), hostID, "vm-ok")
	require.NoError(t, err)
	defer stream.Close()

	got := make([]byte, 5)
	_, err = io.ReadFull(stream, got)
	require.NoError(t, err)
	assert.Equal(t, "BOOT|", string(got), "backlog replays first")

	_, err = stream.Write([]byte("hi"))
	require.NoError(t, err)
	_, err = io.ReadFull(stream, got[:2])
	require.NoError(t, err)
	assert.Equal(t, "HI", string(got[:2]), "input reaches the console and output returns")
}

func TestConsoleRefusedUnknownVM(t *testing.T) {
	h, hostID := startConsoleClient(t, &echoConsole{backlog: "BOOT|"})

	// Unknown vm → agent replies ok=false, OpenConsole errors.
	_, err := h.svc.OpenConsole(context.Background(), hostID, "vm-nope")
	require.Error(t, err)
	assert.Contains(t, err.Error(), "console refused")
}

func TestConsoleWithoutHandlerRefused(t *testing.T) {
	h, hostID := startConsoleClient(t, nil) // client.Console left nil

	_, err := h.svc.OpenConsole(context.Background(), hostID, "vm-ok")
	require.Error(t, err)
	assert.Contains(t, err.Error(), "console refused")
}

// newCountingServer is a serverHarness whose accept loop calls onAccept(1) for
// every connection, so tests can observe reconnect attempts.
func newCountingServer(t *testing.T, onAccept func(int64) int64) *serverHarness {
	t.Helper()
	st, err := store.Open(t.TempDir()+"/db", "10.77.0.0/16")
	require.NoError(t, err)
	t.Cleanup(func() { st.Close() })
	seedTestTenant(t, st)
	certPEM, keyPEM, fp := genTestCert(t)
	h := &serverHarness{
		t: t, st: st, reg: registry.New(time.Now), hub: hub.New(),
		secret: []byte("s3cret"), certPEM: certPEM, keyPEM: keyPEM, fp: fp,
	}
	h.lis = listenTestQUIC(t, "127.0.0.1:0", certPEM, keyPEM)
	h.addr = h.lis.Addr().String()
	ctx, cancel := context.WithCancel(context.Background())
	h.cancel = cancel
	go func() {
		for {
			conn, err := h.lis.Accept(ctx)
			if err != nil {
				return
			}
			onAccept(1)
			// Reject everything: close immediately as auth-rejected. The client's
			// AcceptStream/session surfaces this as errPermanentAuth.
			_ = conn.CloseWithError(transport.CodeAuthRejected, "test reject")
		}
	}()
	t.Cleanup(func() { cancel(); h.lis.Close() })
	return h
}

func TestMaybeUpgradeSkipsOwnVersionAndSingleFlights(t *testing.T) {
	c := &Client{}
	applied := make(chan selfupdate.Update, 2)
	block := make(chan struct{})
	c.applyUpgrade = func(ctx context.Context, u selfupdate.Update) error {
		applied <- u
		<-block // hold the flight open
		return nil
	}

	// Same version: no-op.
	c.maybeUpgrade(context.Background(), &pb.AgentUpgrade{Version: version.Version})
	select {
	case <-applied:
		t.Fatal("must not apply own version")
	default:
	}

	// New version: applies once; a second call while in flight is dropped.
	up := &pb.AgentUpgrade{Version: "v99.0.0", Url: "u", Sha256: "s"}
	c.maybeUpgrade(context.Background(), up)
	c.maybeUpgrade(context.Background(), up)
	got := <-applied
	if got.Version != "v99.0.0" {
		t.Fatalf("applied %+v", got)
	}
	select {
	case <-applied:
		t.Fatal("second in-flight apply must be dropped")
	default:
	}
	close(block)
}

// TestLogGuestCIDROnlyOnChange pins the observability that was missing while
// two implementations of the macOS reader silently reported nothing. An empty
// answer is CORRECT on the wire — the fleet keeps its existing record — which
// is exactly why it has to be visible somewhere: a host that cannot see its own
// network otherwise looks identical to one that can.
func TestLogGuestCIDROnlyOnChange(t *testing.T) {
	last := unreported

	// The first answer always logs, including the first empty one.
	logGuestCIDR("", &last)
	if last != "" {
		t.Fatalf("last = %q, want the empty answer recorded", last)
	}
	// A repeat does not: this runs every tick, forever.
	logGuestCIDR("", &last)
	if last != "" {
		t.Fatalf("last = %q, want unchanged", last)
	}
	// An answer arriving is a change.
	logGuestCIDR("192.168.64.0/24", &last)
	if last != "192.168.64.0/24" {
		t.Fatalf("last = %q, want the new subnet", last)
	}
	// And an answer going away is a change too — the interesting one.
	logGuestCIDR("", &last)
	if last != "" {
		t.Fatalf("last = %q, want the loss recorded", last)
	}
}