llmgw/internal/agent/mtls_test.go

package agent

import (
	"context"
	"log/slog"
	"net"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"golang.org/x/net/http2"
	"golang.org/x/net/http2/h2c"

	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/agentsvc"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
	"gitea.mynodes.xyz/ilker/llmgw/internal/tlsconf"
)

// pki issues a cluster's certificates into a temporary directory and returns
// the configurations the control plane and one agent would load.
func pki(t *testing.T, nodeID string) (server, client tlsconf.Config) {
	t.Helper()
	dir := t.TempDir()

	ca, err := tlsconf.NewCA("test CA")
	if err != nil {
		t.Fatal(err)
	}
	if err := ca.Write(dir); err != nil {
		t.Fatal(err)
	}
	srv, err := ca.IssueServer("control-plane", []string{"127.0.0.1"})
	if err != nil {
		t.Fatal(err)
	}
	if err := srv.Write(dir, "server"); err != nil {
		t.Fatal(err)
	}
	cli, err := ca.IssueClient(nodeID)
	if err != nil {
		t.Fatal(err)
	}
	if err := cli.Write(dir, "agent"); err != nil {
		t.Fatal(err)
	}

	caPath := filepath.Join(dir, "ca.pem")
	return tlsconf.Config{CA: caPath, Cert: filepath.Join(dir, "server.pem"), Key: filepath.Join(dir, "server-key.pem")},
		tlsconf.Config{CA: caPath, Cert: filepath.Join(dir, "agent.pem"), Key: filepath.Join(dir, "agent-key.pem")}
}

// startTLS runs a real service behind a real TLS listener, with identity
// required — the deployed configuration rather than an approximation of it.
func startTLS(t *testing.T, serverTLS tlsconf.Config) (*agentsvc.Service, *inventory.Store, string) {
	t.Helper()

	inv := inventory.NewStore()
	svc := agentsvc.New(inv, newRecorder(), true, slog.New(slog.DiscardHandler))
	path, handler := svc.Handler()
	mux := http.NewServeMux()
	mux.Handle(path, handler)

	tc, err := serverTLS.ServerConfig()
	if err != nil {
		t.Fatal(err)
	}
	raw, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	srv := httptest.NewUnstartedServer(mux)
	srv.Listener = raw
	srv.TLS = tc
	srv.StartTLS()
	t.Cleanup(srv.Close)

	return svc, inv, "https://" + raw.Addr().String()
}

func runAgent(t *testing.T, cfg Config) *Agent {
	t.Helper()

	cfg.MgmtAddr = "127.0.0.1"
	cfg.HeartbeatInterval = 50 * time.Millisecond
	cfg.ReconnectBase = 20 * time.Millisecond
	cfg.ReconnectMax = 100 * time.Millisecond
	cfg.Discovery.SkipGPUs = true
	cfg.Discovery.IBRoot = filepath.Join(t.TempDir(), "absent")
	cfg.Superv = supervisor.Config{
		WorkDir: t.TempDir(), LogLines: 64,
		StopTimeout: 2 * time.Second, HealthInterval: 20 * time.Millisecond,
	}

	client, err := Client(cfg.TLS, 5*time.Second)
	if err != nil {
		t.Fatal(err)
	}
	ag, err := New(cfg, client, slog.New(slog.DiscardHandler))
	if err != nil {
		t.Fatal(err)
	}

	ctx, cancel := context.WithCancel(context.Background())
	go func() { _ = ag.Run(ctx) }()
	t.Cleanup(func() {
		cancel()
		ag.Supervisor().StopAll(context.Background())
	})
	return ag
}

func registered(svc *agentsvc.Service, inv *inventory.Store, nodeID string) func() bool {
	return func() bool {
		_, ok := inv.Get(nodeID)
		return ok && len(svc.ConnectedNodes()) == 1
	}
}

func eventually(t *testing.T, limit time.Duration, pred func() bool) bool {
	t.Helper()
	deadline := time.Now().Add(limit)
	for time.Now().Before(deadline) {
		if pred() {
			return true
		}
		time.Sleep(10 * time.Millisecond)
	}
	return false
}

func TestAgentRegistersOverMutualTLS(t *testing.T) {
	serverTLS, clientTLS := pki(t, "llmgw-node01")
	svc, inv, url := startTLS(t, serverTLS)

	cfg := DefaultConfig()
	cfg.NodeID = "llmgw-node01"
	cfg.Server = url
	cfg.TLS = clientTLS
	runAgent(t, cfg)

	if !eventually(t, 10*time.Second, registered(svc, inv, "llmgw-node01")) {
		t.Fatal("the agent never registered over mutual TLS")
	}
}

func TestNodeCannotClaimAnotherNodesIdentity(t *testing.T) {
	// The point of the whole exercise. node01 holds a valid certificate — it is
	// a legitimate member of the cluster — and still cannot register as node02,
	// because the id it claims is checked against the one it can prove.
	serverTLS, clientTLS := pki(t, "llmgw-node01")
	svc, inv, url := startTLS(t, serverTLS)

	cfg := DefaultConfig()
	cfg.NodeID = "llmgw-node02"
	cfg.Server = url
	cfg.TLS = clientTLS
	runAgent(t, cfg)

	// The agent reconnects on failure, so this has to stay false throughout
	// rather than merely start false.
	if eventually(t, 2*time.Second, func() bool {
		_, ok := inv.Get("llmgw-node02")
		return ok || len(svc.ConnectedNodes()) > 0
	}) {
		t.Fatal("an agent registered under a node id its certificate does not authorize")
	}
}

func TestUnauthenticatedAgentIsRefused(t *testing.T) {
	// The defensive half of the check: identity is required, but the connection
	// carried none. In the shipped wiring TLS and this flag are tied together,
	// so reaching here means someone untied them — and the answer must still be
	// refusal rather than a silently trusted node id.
	inv := inventory.NewStore()
	svc := agentsvc.New(inv, newRecorder(), true, slog.New(slog.DiscardHandler))
	path, handler := svc.Handler()
	mux := http.NewServeMux()
	mux.Handle(path, handler)

	srv := httptest.NewServer(h2c.NewHandler(mux, &http2.Server{}))
	t.Cleanup(srv.Close)

	cfg := DefaultConfig()
	cfg.NodeID = "llmgw-node01"
	cfg.Server = srv.URL
	runAgent(t, cfg)

	if eventually(t, 2*time.Second, registered(svc, inv, "llmgw-node01")) {
		t.Fatal("an agent with no proven identity reached the control plane")
	}
}

func TestPlaintextAgentIsRejectedWhenTLSIsConfigured(t *testing.T) {
	// A misconfiguration worth catching at startup rather than at runtime: an
	// agent holding certificates but dialling http:// would connect in the
	// clear and never learn that it had.
	_, clientTLS := pki(t, "llmgw-node01")

	cfg := DefaultConfig()
	cfg.NodeID = "llmgw-node01"
	cfg.Server = "http://127.0.0.1:9443"
	cfg.TLS = clientTLS

	if _, err := New(cfg, nil, slog.New(slog.DiscardHandler)); err == nil {
		t.Fatal("an agent with certificates was allowed to dial a plaintext address")
	}
}

func TestBareAddressGainsTheRightScheme(t *testing.T) {
	_, clientTLS := pki(t, "llmgw-node01")

	secure := Config{NodeID: "n", Server: "10.0.0.140:9443", TLS: clientTLS}
	if err := secure.normalize(); err != nil {
		t.Fatal(err)
	}
	if secure.Server != "https://10.0.0.140:9443" {
		t.Errorf("a configured agent dials %q", secure.Server)
	}

	plain := Config{NodeID: "n", Server: "10.0.0.140:9443"}
	if err := plain.normalize(); err != nil {
		t.Fatal(err)
	}
	if plain.Server != "http://10.0.0.140:9443" {
		t.Errorf("an unconfigured agent dials %q", plain.Server)
	}
}

func TestRefusedAgentLearnsWhy(t *testing.T) {
	// An agent that cannot register must be able to say why. Reporting a
	// refusal as a clean disconnect leaves an operator with a silent reconnect
	// loop and nothing to act on — and resets the backoff, so the loop is a
	// fast one.
	serverTLS, clientTLS := pki(t, "llmgw-node01")
	_, _, url := startTLS(t, serverTLS)

	cfg := DefaultConfig()
	cfg.NodeID = "llmgw-node02"
	cfg.Server = url
	cfg.TLS = clientTLS
	cfg.Discovery.SkipGPUs = true
	cfg.Discovery.IBRoot = filepath.Join(t.TempDir(), "absent")
	cfg.Superv = supervisor.Config{WorkDir: t.TempDir(), LogLines: 8}

	client, err := Client(cfg.TLS, 5*time.Second)
	if err != nil {
		t.Fatal(err)
	}
	ag, err := New(cfg, client, slog.New(slog.DiscardHandler))
	if err != nil {
		t.Fatal(err)
	}

	err = ag.session(context.Background())
	if err == nil {
		t.Fatal("a refused session reported success")
	}
	if !strings.Contains(err.Error(), "llmgw-node01") || !strings.Contains(err.Error(), "llmgw-node02") {
		t.Errorf("the error names neither certificate nor claim: %v", err)
	}
}