fleet-manager: checkpoint 2 — gRPC agent plane (bidi mTLS stream)

The steady-state host<->manager transport, resolving the design's OPEN FORK to bespoke gRPC (not OTLP+Wazuh). One long-lived, host-dialed, mTLS bidi stream: health/metrics up, desired-state/commands down — no polling, real-time control, still no inbound host port. - proto/fleet_agent.proto: FleetAgent.Session(stream AgentUp) -> stream ManagerDown. AgentUp = Hello | HealthSnapshot | CommandResult; ManagerDown = DesiredState | Command(reboot/drain/resync). int64 unix times (no well-known-type dep). Generated stubs committed (CI needs no protoc). - agent.go: gRPC server with mTLS verified by the FLEET CA (the host's reused TPM identity cert — distinct from the admin plane's admin CA). Identity is the peer cert SERIAL, taken from the verified cert = inventory key, never self-asserted; an unenrolled serial is refused (enrolment gates the stream). Per-conn send channel decouples admin pushes from stream.Send. Stream registry (serial->conn) with reconnect eviction. On connect + on change, desired-state is pushed live; inventory stays authoritative so offline hosts reconcile on reconnect. Metrics go through a MetricsSink seam (no-op now; InfluxDB in ckpt 4). - main.go: agent plane starts only when its cert/key/CA are set; admin can run standalone. Lifecycle PUT now live-pushes and reports pushed_live. Full mTLS integration test (in-process fleet CA + host client cert): initial desired-state on connect, Hello+snapshot -> inventory + online presence, live desired-state push, and unenrolled-serial rejection. go build/test/vet/gofmt clean. Deps (server-side only): google.golang.org/grpc + protobuf.
2026-07-04 03:40:55 +00:00
parent 0572d21c07
commit 80cbb833cc
11 changed files with 1700 additions and 13 deletions
@@ -0,0 +1,243 @@
+package main
+
+import (
+	"context"
+	"crypto/ecdsa"
+	"crypto/elliptic"
+	"crypto/rand"
+	"crypto/tls"
+	"crypto/x509"
+	"crypto/x509/pkix"
+	"encoding/pem"
+	"math/big"
+	"net"
+	"os"
+	"path/filepath"
+	"testing"
+	"time"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials"
+
+	fleetpb "fleet-manager/proto"
+)
+
+// TestAgentSession drives the full agent plane over real mTLS: a host dials in
+// with a fleet-CA client cert, the manager derives its serial from that cert,
+// pushes the current desired-state on connect, ingests a Hello + HealthSnapshot
+// into inventory, and delivers a live desired-state change over the same stream.
+func TestAgentSession(t *testing.T) {
+	dir := t.TempDir()
+	caCert, caKey := makeCA(t)
+	caPath := writeFile(t, dir, "ca.crt", pemCert(caCert.Raw))
+
+	srvCertPEM, srvKeyPEM := makeLeaf(t, caCert, caKey, "fleet-manager", big.NewInt(1), true)
+	srvCrt := writeFile(t, dir, "srv.crt", srvCertPEM)
+	srvKey := writeFile(t, dir, "srv.key", srvKeyPEM)
+
+	const serialHexStr = "1234abcd"
+	cliCertPEM, cliKeyPEM := makeLeaf(t, caCert, caKey, "fleet-host", big.NewInt(0x1234abcd), false)
+	cliCert, err := tls.X509KeyPair(cliCertPEM, cliKeyPEM)
+	if err != nil {
+		t.Fatalf("client keypair: %v", err)
+	}
+
+	store, err := newJSONStore("") // in-memory
+	if err != nil {
+		t.Fatalf("store: %v", err)
+	}
+	if err := store.Upsert(&Host{
+		Identity: Identity{CertSerial: serialHexStr, Hostnqn: "nqn.test:host-" + serialHexStr, StarshipName: "test-ship"},
+		Desired:  Desired{TargetVersion: "0.5.0", KioskURL: "http://k/"},
+	}); err != nil {
+		t.Fatalf("seed: %v", err)
+	}
+
+	hub := newAgentHub(store, nil)
+	gs, err := agentServer(config{agentCert: srvCrt, agentKey: srvKey, agentClientCA: caPath}, hub)
+	if err != nil {
+		t.Fatalf("agentServer: %v", err)
+	}
+	lis, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("listen: %v", err)
+	}
+	go gs.Serve(lis)
+	defer gs.Stop()
+
+	// ---- host dials in ----
+	caPool := x509.NewCertPool()
+	caPool.AddCert(caCert)
+	creds := credentials.NewTLS(&tls.Config{
+		Certificates: []tls.Certificate{cliCert},
+		RootCAs:      caPool,
+		ServerName:   "localhost",
+	})
+	conn, err := grpc.NewClient(lis.Addr().String(), grpc.WithTransportCredentials(creds))
+	if err != nil {
+		t.Fatalf("dial: %v", err)
+	}
+	defer conn.Close()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	stream, err := fleetpb.NewFleetAgentClient(conn).Session(ctx)
+	if err != nil {
+		t.Fatalf("session: %v", err)
+	}
+
+	// 1. current desired-state arrives on connect.
+	down, err := stream.Recv()
+	if err != nil {
+		t.Fatalf("recv desired: %v", err)
+	}
+	if got := down.GetDesired().GetTargetVersion(); got != "0.5.0" {
+		t.Fatalf("initial desired target_version = %q, want 0.5.0", got)
+	}
+
+	// 2. host reports Hello (running an older version) + a health snapshot.
+	if err := stream.Send(&fleetpb.AgentUp{Msg: &fleetpb.AgentUp_Hello{Hello: &fleetpb.Hello{ImageVersion: "0.4.2", Kernel: "5.14.0"}}}); err != nil {
+		t.Fatalf("send hello: %v", err)
+	}
+	if err := stream.Send(&fleetpb.AgentUp{Msg: &fleetpb.AgentUp_Snapshot{Snapshot: &fleetpb.HealthSnapshot{
+		TmpfsUsedPct: 63, NvmeRoot: "ok", Auth: "up", Selinux: "enforcing", KioskActive: true,
+	}}}); err != nil {
+		t.Fatalf("send snapshot: %v", err)
+	}
+
+	// 3. inventory reflects the report + online presence.
+	waitFor(t, 3*time.Second, func() bool {
+		h, ok := store.Get(serialHexStr)
+		return ok && h.Online && h.Version.ImageVersion == "0.4.2" && h.Health.TmpfsUsedPct == 63 && h.Health.NVMeRoot == "ok"
+	}, "inventory did not reflect the agent report")
+
+	// 4. an admin changes desired-state -> live push over the same stream.
+	if _, _, err := store.SetDesired(serialHexStr, Desired{TargetVersion: "0.6.0"}); err != nil {
+		t.Fatalf("set desired: %v", err)
+	}
+	if !hub.pushDesired(serialHexStr, Desired{TargetVersion: "0.6.0"}) {
+		t.Fatal("pushDesired reported the host as not live, but it is connected")
+	}
+	down2, err := stream.Recv()
+	if err != nil {
+		t.Fatalf("recv desired push: %v", err)
+	}
+	if got := down2.GetDesired().GetTargetVersion(); got != "0.6.0" {
+		t.Fatalf("pushed desired target_version = %q, want 0.6.0", got)
+	}
+}
+
+// TestAgentRejectsUnenrolled: a valid fleet-CA cert whose serial has no inventory
+// record must be refused (enrolment gates the stream).
+func TestAgentRejectsUnenrolled(t *testing.T) {
+	dir := t.TempDir()
+	caCert, caKey := makeCA(t)
+	caPath := writeFile(t, dir, "ca.crt", pemCert(caCert.Raw))
+	srvCertPEM, srvKeyPEM := makeLeaf(t, caCert, caKey, "fleet-manager", big.NewInt(1), true)
+	srvCrt := writeFile(t, dir, "srv.crt", srvCertPEM)
+	srvKey := writeFile(t, dir, "srv.key", srvKeyPEM)
+	cliCertPEM, cliKeyPEM := makeLeaf(t, caCert, caKey, "fleet-host", big.NewInt(0xdead), false)
+	cliCert, _ := tls.X509KeyPair(cliCertPEM, cliKeyPEM)
+
+	store, _ := newJSONStore("") // empty — the host is NOT enrolled
+	hub := newAgentHub(store, nil)
+	gs, err := agentServer(config{agentCert: srvCrt, agentKey: srvKey, agentClientCA: caPath}, hub)
+	if err != nil {
+		t.Fatalf("agentServer: %v", err)
+	}
+	lis, _ := net.Listen("tcp", "127.0.0.1:0")
+	go gs.Serve(lis)
+	defer gs.Stop()
+
+	caPool := x509.NewCertPool()
+	caPool.AddCert(caCert)
+	creds := credentials.NewTLS(&tls.Config{Certificates: []tls.Certificate{cliCert}, RootCAs: caPool, ServerName: "localhost"})
+	conn, _ := grpc.NewClient(lis.Addr().String(), grpc.WithTransportCredentials(creds))
+	defer conn.Close()
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	stream, err := fleetpb.NewFleetAgentClient(conn).Session(ctx)
+	if err != nil {
+		t.Fatalf("session open: %v", err)
+	}
+	if _, err := stream.Recv(); err == nil {
+		t.Fatal("expected the stream to be rejected for an unenrolled host, got nil error")
+	}
+}
+
+// ---- test helpers ---------------------------------------------------------
+
+func makeCA(t *testing.T) (*x509.Certificate, *ecdsa.PrivateKey) {
+	t.Helper()
+	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
+	if err != nil {
+		t.Fatalf("ca key: %v", err)
+	}
+	tmpl := &x509.Certificate{
+		SerialNumber:          big.NewInt(1),
+		Subject:               pkix.Name{CommonName: "test-fleet-ca"},
+		NotBefore:             time.Now().Add(-time.Hour),
+		NotAfter:              time.Now().Add(time.Hour),
+		IsCA:                  true,
+		KeyUsage:              x509.KeyUsageCertSign | x509.KeyUsageDigitalSignature,
+		BasicConstraintsValid: true,
+	}
+	der, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &key.PublicKey, key)
+	if err != nil {
+		t.Fatalf("ca cert: %v", err)
+	}
+	cert, _ := x509.ParseCertificate(der)
+	return cert, key
+}
+
+func makeLeaf(t *testing.T, ca *x509.Certificate, caKey *ecdsa.PrivateKey, cn string, serial *big.Int, server bool) (certPEM, keyPEM []byte) {
+	t.Helper()
+	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
+	if err != nil {
+		t.Fatalf("leaf key: %v", err)
+	}
+	tmpl := &x509.Certificate{
+		SerialNumber: serial,
+		Subject:      pkix.Name{CommonName: cn},
+		NotBefore:    time.Now().Add(-time.Hour),
+		NotAfter:     time.Now().Add(time.Hour),
+		KeyUsage:     x509.KeyUsageDigitalSignature,
+		ExtKeyUsage:  []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth},
+	}
+	if server {
+		tmpl.ExtKeyUsage = append(tmpl.ExtKeyUsage, x509.ExtKeyUsageServerAuth)
+		tmpl.DNSNames = []string{"localhost"}
+		tmpl.IPAddresses = []net.IP{net.ParseIP("127.0.0.1")}
+	}
+	der, err := x509.CreateCertificate(rand.Reader, tmpl, ca, &key.PublicKey, caKey)
+	if err != nil {
+		t.Fatalf("leaf cert: %v", err)
+	}
+	kb, _ := x509.MarshalECPrivateKey(key)
+	return pemCert(der), pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: kb})
+}
+
+func pemCert(der []byte) []byte {
+	return pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: der})
+}
+
+func writeFile(t *testing.T, dir, name string, b []byte) string {
+	t.Helper()
+	p := filepath.Join(dir, name)
+	if err := os.WriteFile(p, b, 0o600); err != nil {
+		t.Fatalf("write %s: %v", name, err)
+	}
+	return p
+}
+
+func waitFor(t *testing.T, d time.Duration, cond func() bool, msg string) {
+	t.Helper()
+	deadline := time.Now().Add(d)
+	for time.Now().Before(deadline) {
+		if cond() {
+			return
+		}
+		time.Sleep(20 * time.Millisecond)
+	}
+	t.Fatal(msg)
+}