iop/apps/control-plane/internal/wire/edge_server_test.go

package wire

import (
	"context"
	"crypto/ed25519"
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"errors"
	"fmt"
	"math/big"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"go.uber.org/zap/zaptest"

	proto_socket "git.toki-labs.com/toki/proto-socket/go"
	"google.golang.org/protobuf/proto"
	"iop/packages/go/auth"
	iop "iop/proto/gen/iop"
)

type wireCertFiles struct {
	cert string
	key  string
}

func writeWireCertificate(t *testing.T, dir, name string, template, parent *x509.Certificate, parentKey ed25519.PrivateKey) (wireCertFiles, ed25519.PrivateKey) {
	t.Helper()
	publicKey, privateKey, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	if parentKey == nil {
		parentKey = privateKey
	}
	der, err := x509.CreateCertificate(rand.Reader, template, parent, publicKey, parentKey)
	if err != nil {
		t.Fatal(err)
	}
	certPath := filepath.Join(dir, name+".crt")
	keyPath := filepath.Join(dir, name+".key")
	keyDER, err := x509.MarshalPKCS8PrivateKey(privateKey)
	if err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(certPath, pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: der}), 0o600); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(keyPath, pem.EncodeToMemory(&pem.Block{Type: "PRIVATE KEY", Bytes: keyDER}), 0o600); err != nil {
		t.Fatal(err)
	}
	return wireCertFiles{cert: certPath, key: keyPath}, privateKey
}

func edgeWireTLSConfigs(t *testing.T, edgeName string) (*tls.Config, *tls.Config) {
	t.Helper()
	dir := t.TempDir()
	now := time.Now()
	caTemplate := &x509.Certificate{SerialNumber: big.NewInt(1), Subject: pkix.Name{CommonName: "wire-ca"}, NotBefore: now.Add(-time.Hour), NotAfter: now.Add(time.Hour), IsCA: true, BasicConstraintsValid: true, KeyUsage: x509.KeyUsageCertSign}
	caFiles, caKey := writeWireCertificate(t, dir, "ca", caTemplate, caTemplate, nil)
	caPEM, err := os.ReadFile(caFiles.cert)
	if err != nil {
		t.Fatal(err)
	}
	caBlock, _ := pem.Decode(caPEM)
	caCert, err := x509.ParseCertificate(caBlock.Bytes)
	if err != nil {
		t.Fatal(err)
	}
	serverURI, _ := url.Parse("spiffe://iop/control-plane/cp-1")
	serverTemplate := &x509.Certificate{SerialNumber: big.NewInt(2), DNSNames: []string{"control-plane.internal"}, URIs: []*url.URL{serverURI}, NotBefore: now.Add(-time.Hour), NotAfter: now.Add(time.Hour), KeyUsage: x509.KeyUsageDigitalSignature, ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth}}
	serverFiles, _ := writeWireCertificate(t, dir, "server", serverTemplate, caCert, caKey)
	edgeURI, _ := url.Parse("spiffe://iop/edge/" + edgeName)
	edgeTemplate := &x509.Certificate{SerialNumber: big.NewInt(3), URIs: []*url.URL{edgeURI}, NotBefore: now.Add(-time.Hour), NotAfter: now.Add(time.Hour), KeyUsage: x509.KeyUsageDigitalSignature, ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth}}
	edgeFiles, _ := writeWireCertificate(t, dir, "edge", edgeTemplate, caCert, caKey)
	serverConfig, err := auth.LoadServerTLSWithIdentity(serverFiles.cert, serverFiles.key, caFiles.cert, "edge", "")
	if err != nil {
		t.Fatal(err)
	}
	clientConfig, err := auth.LoadClientTLSWithIdentity(edgeFiles.cert, edgeFiles.key, caFiles.cert, "control-plane.internal", "control-plane", "cp-1")
	if err != nil {
		t.Fatal(err)
	}
	return serverConfig, clientConfig
}

func startEdgeServer(t *testing.T) (*EdgeServer, int) {
	t.Helper()
	logger := zaptest.NewLogger(t)
	port := getFreePort(t)
	listenAddr := fmt.Sprintf("127.0.0.1:%d", port)

	server, err := NewEdgeServer(listenAddr, logger)
	if err != nil {
		t.Fatalf("NewEdgeServer failed: %v", err)
	}

	ctx, cancel := context.WithCancel(context.Background())
	t.Cleanup(cancel)

	if err := server.Start(ctx); err != nil {
		t.Fatalf("edge server Start failed: %v", err)
	}
	t.Cleanup(func() { _ = server.Stop() })

	// Give a tiny moment for the listener to bind.
	time.Sleep(50 * time.Millisecond)
	return server, port
}

func dialEdge(t *testing.T, ctx context.Context, port int) *proto_socket.TcpClient {
	t.Helper()
	client, err := proto_socket.DialTcp(
		ctx,
		"127.0.0.1",
		port,
		proto_socket.DefaultHeartbeatIntervalSec,
		proto_socket.DefaultHeartbeatWaitSec,
		EdgeParserMap(),
	)
	if err != nil {
		t.Fatalf("failed to DialTcp: %v", err)
	}
	return client
}

func startTLSEdgeServer(t *testing.T, tlsConfig *tls.Config) (*EdgeServer, int) {
	t.Helper()
	port := getFreePort(t)
	server, err := NewEdgeServerTLS(fmt.Sprintf("127.0.0.1:%d", port), tlsConfig, zaptest.NewLogger(t))
	if err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	t.Cleanup(cancel)
	if err := server.Start(ctx); err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = server.Stop() })
	return server, port
}

func dialTLSEdge(t *testing.T, ctx context.Context, port int, tlsConfig *tls.Config, parsers proto_socket.ParserMap) *proto_socket.TcpClient {
	t.Helper()
	client, err := proto_socket.DialTcpTLS(ctx, "127.0.0.1", port, tlsConfig, EdgeHeartbeatIntervalSec, EdgeHeartbeatWaitSec, parsers)
	if err != nil {
		t.Fatal(err)
	}
	return client
}

func projectionClientParserMap() proto_socket.ParserMap {
	parsers := EdgeParserMap()
	parsers[proto_socket.TypeNameOf(&iop.PrincipalProjectionApplyRequest{})] = func(data []byte) (proto.Message, error) {
		request := &iop.PrincipalProjectionApplyRequest{}
		return request, proto.Unmarshal(data, request)
	}
	return parsers
}

func TestEdgeServerHelloRegistersEdge(t *testing.T) {
	server, port := startEdgeServer(t)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	client := dialEdge(t, ctx, port)
	defer client.Close()

	req := &iop.EdgeHelloRequest{
		EdgeId:       "edge-dgx-group",
		EdgeName:     "DGX Group",
		Version:      "1.2.3",
		Capabilities: []string{"node-registry", "run-dispatch"},
	}

	res, err := proto_socket.SendRequestTyped[*iop.EdgeHelloRequest, *iop.EdgeHelloResponse](
		&client.Communicator,
		req,
		2*time.Second,
	)
	if err != nil {
		t.Fatalf("SendRequestTyped failed: %v", err)
	}

	if !res.Accepted {
		t.Errorf("expected response.Accepted to be true, got false (reason=%q)", res.Reason)
	}
	if res.Protocol != Protocol {
		t.Errorf("expected response.Protocol to be %q, got %q", Protocol, res.Protocol)
	}
	if res.ServerTimeUnixNano <= 0 {
		t.Errorf("expected server time unix nano to be positive, got %d", res.ServerTimeUnixNano)
	}

	state, ok := server.Registry().Snapshot("edge-dgx-group")
	if !ok {
		t.Fatalf("edge not registered after hello")
	}
	if !state.Connected {
		t.Errorf("expected registry Connected to be true")
	}
	if state.EdgeName != "DGX Group" {
		t.Errorf("edge_name: got %q want %q", state.EdgeName, "DGX Group")
	}
	if len(state.Capabilities) != 2 {
		t.Errorf("capabilities: got %v", state.Capabilities)
	}
	if state.LastSeen.IsZero() {
		t.Errorf("expected LastSeen to be set")
	}
}

func TestEdgeServerRejectsMissingEdgeID(t *testing.T) {
	server, port := startEdgeServer(t)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	client := dialEdge(t, ctx, port)
	defer client.Close()

	req := &iop.EdgeHelloRequest{
		EdgeName: "no id edge",
	}

	res, err := proto_socket.SendRequestTyped[*iop.EdgeHelloRequest, *iop.EdgeHelloResponse](
		&client.Communicator,
		req,
		2*time.Second,
	)
	if err != nil {
		t.Fatalf("SendRequestTyped failed: %v", err)
	}

	if res.Accepted {
		t.Errorf("expected response.Accepted to be false for missing edge_id")
	}
	if res.Reason == "" {
		t.Errorf("expected a rejection reason for missing edge_id")
	}
	if server.Registry().Len() != 0 {
		t.Errorf("expected registry to stay empty, got %d entries", server.Registry().Len())
	}
}

func TestEdgeServerAcceptsMatchingAuthenticatedEnrollmentName(t *testing.T) {
	serverTLS, clientTLS := edgeWireTLSConfigs(t, "edge-a")
	server, port := startTLSEdgeServer(t, serverTLS)
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()
	client := dialTLSEdge(t, ctx, port, clientTLS, EdgeParserMap())
	defer client.Close()

	response, err := proto_socket.SendRequestTyped[*iop.EdgeHelloRequest, *iop.EdgeHelloResponse](
		&client.Communicator,
		&iop.EdgeHelloRequest{EdgeId: "edge-a", EdgeName: "Edge A"},
		2*time.Second,
	)
	if err != nil {
		t.Fatal(err)
	}
	if !response.GetAccepted() {
		t.Fatalf("matching authenticated identity rejected: %q", response.GetReason())
	}
	if state, ok := server.Registry().Snapshot("edge-a"); !ok || !state.Connected {
		t.Fatalf("matching authenticated edge was not registered: state=%+v ok=%v", state, ok)
	}
}

func TestEdgeServerRejectsSameRoleWrongEnrollmentName(t *testing.T) {
	serverTLS, clientTLS := edgeWireTLSConfigs(t, "edge-a")
	server, port := startTLSEdgeServer(t, serverTLS)
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()
	client := dialTLSEdge(t, ctx, port, clientTLS, EdgeParserMap())
	defer client.Close()

	response, err := proto_socket.SendRequestTyped[*iop.EdgeHelloRequest, *iop.EdgeHelloResponse](
		&client.Communicator,
		&iop.EdgeHelloRequest{EdgeId: "edge-b", EdgeName: "Edge B"},
		2*time.Second,
	)
	if err != nil {
		t.Fatal(err)
	}
	if response.GetAccepted() {
		t.Fatal("same-role certificate enrolled a different edge name")
	}
	if server.Registry().Len() != 0 {
		t.Fatalf("rejected authenticated identity mutated registry: %d entries", server.Registry().Len())
	}
	if _, ok := server.activeClient("edge-b"); ok {
		t.Fatal("rejected authenticated identity became an active client")
	}
}

func TestEdgeServerDisconnectMarksEdgeDisconnected(t *testing.T) {
	server, port := startEdgeServer(t)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	client := dialEdge(t, ctx, port)

	req := &iop.EdgeHelloRequest{
		EdgeId:   "edge-dgx-group",
		EdgeName: "DGX Group",
	}
	if _, err := proto_socket.SendRequestTyped[*iop.EdgeHelloRequest, *iop.EdgeHelloResponse](
		&client.Communicator,
		req,
		2*time.Second,
	); err != nil {
		t.Fatalf("SendRequestTyped failed: %v", err)
	}

	if err := client.Close(); err != nil {
		t.Fatalf("client Close failed: %v", err)
	}

	// The server detects the closed connection asynchronously; poll until the
	// disconnect listener flips the registry state.
	deadline := time.Now().Add(2 * time.Second)
	for {
		state, ok := server.Registry().Snapshot("edge-dgx-group")
		if ok && !state.Connected {
			if state.DisconnectInfo.Reason == "" {
				t.Errorf("expected a disconnect reason to be recorded")
			}
			return
		}
		if time.Now().After(deadline) {
			t.Fatalf("edge was not marked disconnected in time (ok=%v)", ok)
		}
		time.Sleep(20 * time.Millisecond)
	}
}

func sendEdgeHello(t *testing.T, client *proto_socket.TcpClient, edgeID string) {
	t.Helper()
	if _, err := proto_socket.SendRequestTyped[*iop.EdgeHelloRequest, *iop.EdgeHelloResponse](
		&client.Communicator,
		&iop.EdgeHelloRequest{EdgeId: edgeID, EdgeName: "DGX Group"},
		2*time.Second,
	); err != nil {
		t.Fatalf("SendRequestTyped failed: %v", err)
	}
}

func TestEdgeServerStaleDisconnectDoesNotClearReconnect(t *testing.T) {
	server, port := startEdgeServer(t)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	const edgeID = "edge-dgx-group"

	// First connection enrolls the edge.
	conn1 := dialEdge(t, ctx, port)
	sendEdgeHello(t, conn1, edgeID)

	// Second connection with the same edge_id reconnects and becomes current.
	conn2 := dialEdge(t, ctx, port)
	defer conn2.Close()
	sendEdgeHello(t, conn2, edgeID)

	// The stale first connection closing must not clear the live reconnection.
	if err := conn1.Close(); err != nil {
		t.Fatalf("conn1 Close failed: %v", err)
	}

	// Give the server time to process the stale disconnect, then assert the edge
	// stays connected under the second connection's token.
	time.Sleep(300 * time.Millisecond)
	state, ok := server.Registry().Snapshot(edgeID)
	if !ok {
		t.Fatalf("edge missing from registry")
	}
	if !state.Connected {
		t.Fatalf("stale disconnect cleared the live reconnection; Connected=false")
	}
}

func TestEdgeServerStaleConnectionNodeEventDoesNotPolluteReconnect(t *testing.T) {
	server, port := startEdgeServer(t)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	const edgeID = "edge-dgx-group"

	conn1 := dialEdge(t, ctx, port)
	defer conn1.Close()
	sendEdgeHello(t, conn1, edgeID)

	conn2 := dialEdge(t, ctx, port)
	defer conn2.Close()
	sendEdgeHello(t, conn2, edgeID)

	if err := conn1.Send(&iop.EdgeNodeEvent{
		EventId: "evt-stale",
		Type:    "node.connected",
		NodeId:  "node-stale",
	}); err != nil {
		t.Fatalf("send stale EdgeNodeEvent: %v", err)
	}
	time.Sleep(300 * time.Millisecond)
	if got := server.Registry().NodeEvents(edgeID, ""); len(got) != 0 {
		t.Fatalf("stale connection event polluted registry: %+v", got)
	}

	if err := conn2.Send(&iop.EdgeNodeEvent{
		EventId: "evt-current",
		Type:    "node.connected",
		NodeId:  "node-current",
	}); err != nil {
		t.Fatalf("send current EdgeNodeEvent: %v", err)
	}

	var records []EdgeNodeEventRecord
	deadline := time.Now().Add(2 * time.Second)
	for {
		records = server.Registry().NodeEvents(edgeID, "")
		if len(records) == 1 {
			break
		}
		if time.Now().After(deadline) {
			t.Fatalf("current node event was not recorded in time; got %d records", len(records))
		}
		time.Sleep(20 * time.Millisecond)
	}
	if records[0].Event.GetEventId() != "evt-current" {
		t.Fatalf("recorded event_id: got %q want evt-current", records[0].Event.GetEventId())
	}
	if records[0].Event.GetNodeId() != "node-current" {
		t.Fatalf("recorded node_id: got %q want node-current", records[0].Event.GetNodeId())
	}
}

func TestEdgeServerRequestsStatusFromConnectedEdge(t *testing.T) {
	server, port := startEdgeServer(t)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	client := dialEdge(t, ctx, port)
	defer client.Close()

	const edgeID = "edge-dgx-group"

	// The test acts as the Edge: it answers a Control Plane status request with a
	// service-backed node snapshot.
	proto_socket.AddRequestListenerTyped(&client.Communicator, func(req *iop.EdgeStatusRequest) (*iop.EdgeStatusResponse, error) {
		return &iop.EdgeStatusResponse{
			RequestId:            req.GetRequestId(),
			EdgeId:               edgeID,
			EdgeName:             "DGX Group",
			ObservedTimeUnixNano: time.Now().UnixNano(),
			Nodes: []*iop.EdgeNodeSnapshot{
				{NodeId: "node-1", Alias: "alpha", Label: "node0", Connected: true},
			},
		}, nil
	})

	sendEdgeHello(t, client, edgeID)

	resp, err := server.RequestStatus(edgeID, 2*time.Second)
	if err != nil {
		t.Fatalf("RequestStatus: %v", err)
	}
	if resp.GetRequestId() == "" {
		t.Errorf("expected a non-empty correlation request_id")
	}
	if resp.GetEdgeId() != edgeID {
		t.Errorf("edge_id: got %q want %q", resp.GetEdgeId(), edgeID)
	}
	if len(resp.GetNodes()) != 1 {
		t.Fatalf("nodes: got %d want 1", len(resp.GetNodes()))
	}
	if node := resp.GetNodes()[0]; node.GetNodeId() != "node-1" || node.GetLabel() != "node0" {
		t.Errorf("node snapshot: %+v", node)
	}
}

func TestEdgeServerRecordsNodeLifecycleEventFromEdge(t *testing.T) {
	server, port := startEdgeServer(t)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	client := dialEdge(t, ctx, port)
	defer client.Close()

	if err := client.Send(&iop.EdgeNodeEvent{
		EventId: "evt-before-hello",
		Type:    "node.connected",
		NodeId:  "node-early",
	}); err != nil {
		t.Fatalf("send pre-hello EdgeNodeEvent: %v", err)
	}
	time.Sleep(100 * time.Millisecond)
	if got := server.Registry().NodeEvents("edge-dgx-group", ""); len(got) != 0 {
		t.Fatalf("expected pre-hello node event to be ignored, got %+v", got)
	}

	sendEdgeHello(t, client, "edge-dgx-group")

	event := &iop.EdgeNodeEvent{
		EventId:  "evt-node-1",
		Type:     "node.connected",
		Source:   "edge",
		NodeId:   "node-1",
		Alias:    "alpha",
		Reason:   "registered",
		Metadata: map[string]string{"rack": "r1"},
	}
	if err := client.Send(event); err != nil {
		t.Fatalf("send EdgeNodeEvent: %v", err)
	}

	var records []EdgeNodeEventRecord
	deadline := time.Now().Add(2 * time.Second)
	for {
		records = server.Registry().NodeEvents("edge-dgx-group", "")
		if len(records) == 1 {
			break
		}
		if time.Now().After(deadline) {
			t.Fatalf("node event was not recorded in time; got %d records", len(records))
		}
		time.Sleep(20 * time.Millisecond)
	}
	got := records[0]
	if got.EdgeID != "edge-dgx-group" {
		t.Errorf("edge_id: got %q want %q", got.EdgeID, "edge-dgx-group")
	}
	if got.Event.GetEventId() != "evt-node-1" || got.Event.GetType() != "node.connected" {
		t.Fatalf("unexpected recorded event: %+v", got.Event)
	}
	if got.Event.GetNodeId() != "node-1" || got.Event.GetAlias() != "alpha" {
		t.Fatalf("unexpected node identity: %+v", got.Event)
	}
	if got.Event.GetMetadata()["rack"] != "r1" {
		t.Fatalf("metadata rack: got %q", got.Event.GetMetadata()["rack"])
	}
	if got.Event.GetTimestamp() == 0 {
		t.Fatal("expected registry to fill missing event timestamp")
	}
	if event.GetTimestamp() != 0 {
		t.Fatal("registry mutated original event timestamp")
	}

	records[0].Event.Metadata["rack"] = "mutated"
	again := server.Registry().NodeEvents("edge-dgx-group", "node-1")
	if len(again) != 1 {
		t.Fatalf("filtered node events: got %d want 1", len(again))
	}
	if again[0].Event.GetMetadata()["rack"] != "r1" {
		t.Fatalf("registry event was mutated through returned snapshot: %+v", again[0].Event.GetMetadata())
	}
}

type commandEdgeRecorder struct {
	mu       sync.Mutex
	received []string
}

func (r *commandEdgeRecorder) operations() []string {
	r.mu.Lock()
	defer r.mu.Unlock()
	return append([]string(nil), r.received...)
}

// dialCommandEdge dials the server as an Edge that answers command requests with
// an "ok" response and records the operations it was asked to run.
func dialCommandEdge(t *testing.T, ctx context.Context, port int, edgeID string) *commandEdgeRecorder {
	t.Helper()
	client := dialEdge(t, ctx, port)
	t.Cleanup(func() { _ = client.Close() })
	rec := &commandEdgeRecorder{}
	proto_socket.AddRequestListenerTyped(&client.Communicator, func(req *iop.EdgeCommandRequest) (*iop.EdgeCommandResponse, error) {
		rec.mu.Lock()
		rec.received = append(rec.received, req.GetOperation())
		rec.mu.Unlock()
		return &iop.EdgeCommandResponse{
			RequestId: req.GetRequestId(),
			CommandId: req.GetCommandId(),
			EdgeId:    edgeID,
			Status:    "ok",
			Summary:   "applied " + req.GetOperation(),
		}, nil
	})
	sendEdgeHello(t, client, edgeID)
	return rec
}

func TestEdgeServerSendsCommandToTargetEdgeOnly(t *testing.T) {
	server, port := startEdgeServer(t)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	recA := dialCommandEdge(t, ctx, port, "edge-a")
	recB := dialCommandEdge(t, ctx, port, "edge-b")

	resp, err := server.SendCommand("edge-a", "restart-node", "node-1", map[string]string{"force": "true"}, 2*time.Second)
	if err != nil {
		t.Fatalf("SendCommand: %v", err)
	}
	if resp.GetStatus() != "ok" || resp.GetEdgeId() != "edge-a" {
		t.Fatalf("unexpected response: %+v", resp)
	}
	if resp.GetCommandId() == "" {
		t.Error("expected a non-empty command_id correlation")
	}

	if got := recA.operations(); len(got) != 1 || got[0] != "restart-node" {
		t.Fatalf("edge-a received: %+v", got)
	}
	if got := recB.operations(); len(got) != 0 {
		t.Fatalf("edge-b should not receive a targeted command, got: %+v", got)
	}

	cmds := server.Registry().Commands("edge-a")
	if len(cmds) != 1 {
		t.Fatalf("edge-a audit: got %d want 1", len(cmds))
	}
	if cmds[0].Operation != "restart-node" || cmds[0].Status != "ok" {
		t.Fatalf("unexpected audit record: %+v", cmds[0])
	}
	if got := server.Registry().Commands("edge-b"); len(got) != 0 {
		t.Fatalf("edge-b audit should be empty, got %+v", got)
	}
}

func TestEdgeServerSendCommandFailsWhenEdgeNotConnected(t *testing.T) {
	server, _ := startEdgeServer(t)
	if _, err := server.SendCommand("missing-edge", "restart-node", "", nil, time.Second); err == nil {
		t.Fatal("expected error sending command to a non-connected edge")
	}
	if got := server.Registry().Commands("missing-edge"); len(got) != 0 {
		t.Fatalf("expected no audit record for a never-connected edge, got %+v", got)
	}
}

func TestEdgeServerRecordsCommandLifecycleEventsFromEdge(t *testing.T) {
	server, port := startEdgeServer(t)

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	client := dialEdge(t, ctx, port)
	defer client.Close()
	const edgeID = "edge-dgx-group"

	proto_socket.AddRequestListenerTyped(&client.Communicator, func(req *iop.EdgeCommandRequest) (*iop.EdgeCommandResponse, error) {
		_ = client.Send(&iop.EdgeCommandEvent{CommandId: req.GetCommandId(), EdgeId: edgeID, Phase: "accepted", Summary: "queued"})
		_ = client.Send(&iop.EdgeCommandEvent{CommandId: req.GetCommandId(), EdgeId: edgeID, Phase: "running", Summary: "executing"})
		return &iop.EdgeCommandResponse{
			RequestId: req.GetRequestId(),
			CommandId: req.GetCommandId(),
			EdgeId:    edgeID,
			Status:    "ok",
			Summary:   "done",
		}, nil
	})

	sendEdgeHello(t, client, edgeID)

	resp, err := server.SendCommand(edgeID, "drain", "", nil, 2*time.Second)
	if err != nil {
		t.Fatalf("SendCommand: %v", err)
	}

	var rec EdgeCommandRecord
	deadline := time.Now().Add(2 * time.Second)
	for {
		got, ok := server.Registry().Command(edgeID, resp.GetCommandId())
		if ok && len(got.Events) == 2 {
			rec = got
			break
		}
		if time.Now().After(deadline) {
			t.Fatalf("command events not recorded in time; record=%+v", got)
		}
		time.Sleep(20 * time.Millisecond)
	}
	if rec.Status != "ok" || rec.Summary != "done" {
		t.Fatalf("unexpected final command record: %+v", rec)
	}
	if rec.Events[0].Phase != "accepted" || rec.Events[1].Phase != "running" {
		t.Fatalf("unexpected event phases: %+v", rec.Events)
	}
}

func TestEdgeRegistryCommandsAreBoundedAndCopied(t *testing.T) {
	registry := NewEdgeRegistry()
	at := time.Unix(1780142500, 0).UTC()

	for i := 0; i < recentCommandCapacity+2; i++ {
		cmdID := fmt.Sprintf("cmd-%03d", i)
		registry.RecordCommandRequest("edge-a", &iop.EdgeCommandRequest{CommandId: cmdID, Operation: "op"}, at.Add(time.Duration(i)*time.Second))
		registry.RecordCommandResult("edge-a", &iop.EdgeCommandResponse{CommandId: cmdID, Status: "ok", Summary: "done"}, at.Add(time.Duration(i)*time.Second).Add(time.Millisecond))
	}

	cmds := registry.Commands("edge-a")
	if len(cmds) != recentCommandCapacity {
		t.Fatalf("commands len: got %d want %d", len(cmds), recentCommandCapacity)
	}
	if cmds[0].CommandID != "cmd-002" {
		t.Fatalf("expected oldest retained cmd-002, got %q", cmds[0].CommandID)
	}
	if cmds[0].Status != "ok" {
		t.Fatalf("expected result status recorded, got %q", cmds[0].Status)
	}

	registry.RecordCommandRequest("edge-b", &iop.EdgeCommandRequest{CommandId: "cmd-b", Operation: "drain"}, at)
	for i := 0; i < commandEventCapacity+3; i++ {
		registry.RecordCommandEvent("edge-b", &iop.EdgeCommandEvent{CommandId: "cmd-b", Phase: fmt.Sprintf("phase-%03d", i)}, at.Add(time.Duration(i)*time.Nanosecond))
	}
	rec, ok := registry.Command("edge-b", "cmd-b")
	if !ok {
		t.Fatal("edge-b cmd-b missing")
	}
	if len(rec.Events) != commandEventCapacity {
		t.Fatalf("events len: got %d want %d", len(rec.Events), commandEventCapacity)
	}
	if rec.Events[0].Phase != "phase-003" {
		t.Fatalf("expected oldest retained phase-003, got %q", rec.Events[0].Phase)
	}

	rec.Events[0].Phase = "mutated"
	again, _ := registry.Command("edge-b", "cmd-b")
	if again.Events[0].Phase == "mutated" {
		t.Fatal("Command returned mutable registry storage")
	}
}

func TestEdgeServerRequestStatusFailsWhenEdgeNotConnected(t *testing.T) {
	server, _ := startEdgeServer(t)
	if _, err := server.RequestStatus("missing-edge", time.Second); err == nil {
		t.Fatal("expected error requesting status from a non-connected edge")
	}
}

func TestEdgeRegistrySnapshotsAreSortedAndCopied(t *testing.T) {
	registry := NewEdgeRegistry()
	at := time.Unix(1780142300, 0).UTC()
	registry.MarkConnected(&iop.EdgeHelloRequest{
		EdgeId:       "edge-b",
		EdgeName:     "Edge B",
		Capabilities: []string{"run"},
	}, at)
	registry.MarkConnected(&iop.EdgeHelloRequest{
		EdgeId:       "edge-a",
		EdgeName:     "Edge A",
		Capabilities: []string{"status"},
	}, at.Add(time.Second))

	snapshots := registry.Snapshots()
	if len(snapshots) != 2 {
		t.Fatalf("expected 2 snapshots, got %d", len(snapshots))
	}
	if snapshots[0].EdgeID != "edge-a" || snapshots[1].EdgeID != "edge-b" {
		t.Fatalf("expected snapshots sorted by edge_id, got %+v", snapshots)
	}

	snapshots[0].Capabilities[0] = "mutated"
	again, ok := registry.Snapshot("edge-a")
	if !ok {
		t.Fatal("edge-a missing from registry")
	}
	if again.Capabilities[0] != "status" {
		t.Fatalf("snapshot capabilities alias registry storage: %+v", again.Capabilities)
	}
}

func TestEdgeRegistryNodeEventsAreBoundedAndCopied(t *testing.T) {
	registry := NewEdgeRegistry()
	receivedAt := time.Unix(1780142400, 0).UTC()

	for i := 0; i < recentNodeEventCapacity+2; i++ {
		registry.RecordNodeEvent("edge-a", &iop.EdgeNodeEvent{
			EventId: fmt.Sprintf("evt-%03d", i),
			Type:    "node.connected",
			NodeId:  "node-1",
			Metadata: map[string]string{
				"index": fmt.Sprintf("%d", i),
			},
		}, receivedAt.Add(time.Duration(i)*time.Nanosecond))
	}

	events := registry.NodeEvents("edge-a", "")
	if len(events) != recentNodeEventCapacity {
		t.Fatalf("node events len: got %d want %d", len(events), recentNodeEventCapacity)
	}
	if events[0].Event.GetEventId() != "evt-002" {
		t.Fatalf("expected oldest retained event evt-002, got %q", events[0].Event.GetEventId())
	}
	events[0].Event.Metadata["index"] = "mutated"
	again := registry.NodeEvents("edge-a", "")
	if again[0].Event.GetMetadata()["index"] == "mutated" {
		t.Fatal("NodeEvents returned mutable registry storage")
	}
}

func testWireProjection(generation uint64) *iop.PrincipalProjection {
	now := time.Now()
	return &iop.PrincipalProjection{
		Generation:        generation,
		IssuedAtUnixNano:  now.UnixNano(),
		ExpiresAtUnixNano: now.Add(time.Minute).UnixNano(),
	}
}

func TestEdgeServerPeriodicallyRefreshesEveryActiveEdge(t *testing.T) {
	port := getFreePort(t)
	server, err := NewEdgeServer(fmt.Sprintf("127.0.0.1:%d", port), zaptest.NewLogger(t))
	if err != nil {
		t.Fatal(err)
	}
	var generation atomic.Uint64
	generation.Store(1)
	server.SetCredentialPlane(func(context.Context) (*iop.PrincipalProjection, error) {
		return testWireProjection(generation.Load()), nil
	}, nil, 20*time.Millisecond)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	if err := server.Start(ctx); err != nil {
		t.Fatal(err)
	}

	clients := make([]*proto_socket.TcpClient, 0, 2)
	var applied [2]atomic.Int32
	var seenGenerationTwo [2]atomic.Bool
	for index, edgeID := range []string{"edge-refresh-a", "edge-refresh-b"} {
		client, dialErr := proto_socket.DialTcp(ctx, "127.0.0.1", port, EdgeHeartbeatIntervalSec, EdgeHeartbeatWaitSec, projectionClientParserMap())
		if dialErr != nil {
			t.Fatal(dialErr)
		}
		clients = append(clients, client)
		index := index
		proto_socket.AddRequestListenerTyped(&client.Communicator, func(request *iop.PrincipalProjectionApplyRequest) (*iop.PrincipalProjectionApplyResponse, error) {
			applied[index].Add(1)
			if request.GetProjection().GetGeneration() == 2 {
				seenGenerationTwo[index].Store(true)
			}
			return &iop.PrincipalProjectionApplyResponse{Applied: true, AcceptedGeneration: request.GetProjection().GetGeneration()}, nil
		})
		sendEdgeHello(t, client, edgeID)
	}
	defer func() {
		for _, client := range clients {
			_ = client.Close()
		}
	}()

	deadline := time.Now().Add(2 * time.Second)
	for (applied[0].Load() == 0 || applied[1].Load() == 0) && time.Now().Before(deadline) {
		time.Sleep(10 * time.Millisecond)
	}
	if applied[0].Load() == 0 || applied[1].Load() == 0 {
		t.Fatalf("periodic refresh did not reach every edge: a=%d b=%d", applied[0].Load(), applied[1].Load())
	}

	generation.Store(2)
	if err := server.BroadcastProjection(ctx); err != nil {
		t.Fatal(err)
	}
	if !seenGenerationTwo[0].Load() || !seenGenerationTwo[1].Load() {
		t.Fatal("higher-generation mutation broadcast did not coexist with periodic refresh")
	}

	if err := server.Stop(); err != nil {
		t.Fatal(err)
	}
	counts := [2]int32{applied[0].Load(), applied[1].Load()}
	time.Sleep(80 * time.Millisecond)
	if applied[0].Load() != counts[0] || applied[1].Load() != counts[1] {
		t.Fatalf("projection refresh continued after stop: before=%v after=[%d %d]", counts, applied[0].Load(), applied[1].Load())
	}
}

func TestEdgeServerBroadcastProjectionAttemptsAllActiveEdges(t *testing.T) {
	port := getFreePort(t)
	server, err := NewEdgeServer(fmt.Sprintf("127.0.0.1:%d", port), zaptest.NewLogger(t))
	if err != nil {
		t.Fatal(err)
	}
	server.SetCredentialPlane(func(context.Context) (*iop.PrincipalProjection, error) {
		return testWireProjection(1), nil
	}, nil, 0)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	if err := server.Start(ctx); err != nil {
		t.Fatal(err)
	}
	defer server.Stop()

	var attempts [2]atomic.Int32
	for index, edgeID := range []string{"edge-rejects", "edge-accepts"} {
		client, dialErr := proto_socket.DialTcp(ctx, "127.0.0.1", port, EdgeHeartbeatIntervalSec, EdgeHeartbeatWaitSec, projectionClientParserMap())
		if dialErr != nil {
			t.Fatal(dialErr)
		}
		defer client.Close()
		index := index
		proto_socket.AddRequestListenerTyped(&client.Communicator, func(request *iop.PrincipalProjectionApplyRequest) (*iop.PrincipalProjectionApplyResponse, error) {
			attempts[index].Add(1)
			return &iop.PrincipalProjectionApplyResponse{Applied: index == 1, AcceptedGeneration: request.GetProjection().GetGeneration()}, nil
		})
		sendEdgeHello(t, client, edgeID)
	}

	if err := server.BroadcastProjection(ctx); err == nil {
		t.Fatal("expected aggregate error when one edge rejects the projection")
	}
	if attempts[0].Load() != 1 || attempts[1].Load() != 1 {
		t.Fatalf("broadcast did not attempt every active edge: reject=%d accept=%d", attempts[0].Load(), attempts[1].Load())
	}
}

// dialProjectionEdge dials the server as an Edge whose PrincipalProjectionApply
// handler is driven by the returned start/release channels: it signals start on
// its first push and then blocks until release is closed. It is used to hold one
// Edge's projection push while asserting the fan-out does not stall the others.
func dialProjectionEdge(t *testing.T, ctx context.Context, port int, edgeID string, applied *atomic.Int32) (start <-chan struct{}, release chan struct{}) {
	t.Helper()
	startCh := make(chan struct{})
	releaseCh := make(chan struct{})
	var startOnce sync.Once
	client, dialErr := proto_socket.DialTcp(ctx, "127.0.0.1", port, EdgeHeartbeatIntervalSec, EdgeHeartbeatWaitSec, projectionClientParserMap())
	if dialErr != nil {
		t.Fatal(dialErr)
	}
	t.Cleanup(func() { _ = client.Close() })
	proto_socket.AddRequestListenerTyped(&client.Communicator, func(request *iop.PrincipalProjectionApplyRequest) (*iop.PrincipalProjectionApplyResponse, error) {
		if applied != nil {
			applied.Add(1)
		}
		startOnce.Do(func() { close(startCh) })
		<-releaseCh
		return &iop.PrincipalProjectionApplyResponse{Applied: true, AcceptedGeneration: request.GetProjection().GetGeneration()}, nil
	})
	sendEdgeHello(t, client, edgeID)
	return startCh, releaseCh
}

func TestEdgeServerBroadcastProjectionDoesNotLetStalledEdgeBlockHealthyEdge(t *testing.T) {
	port := getFreePort(t)
	server, err := NewEdgeServer(fmt.Sprintf("127.0.0.1:%d", port), zaptest.NewLogger(t))
	if err != nil {
		t.Fatal(err)
	}
	server.SetCredentialPlane(func(context.Context) (*iop.PrincipalProjection, error) {
		return testWireProjection(1), nil
	}, nil, 0)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	if err := server.Start(ctx); err != nil {
		t.Fatal(err)
	}
	defer server.Stop()

	// The stalled Edge accepts the request but never answers until released.
	stalledStart, release := dialProjectionEdge(t, ctx, port, "edge-stalled", nil)
	var releaseOnce sync.Once
	closeRelease := func() { releaseOnce.Do(func() { close(release) }) }
	defer closeRelease()

	// The healthy Edge answers immediately.
	var healthyApplied atomic.Int32
	client, dialErr := proto_socket.DialTcp(ctx, "127.0.0.1", port, EdgeHeartbeatIntervalSec, EdgeHeartbeatWaitSec, projectionClientParserMap())
	if dialErr != nil {
		t.Fatal(dialErr)
	}
	defer client.Close()
	proto_socket.AddRequestListenerTyped(&client.Communicator, func(request *iop.PrincipalProjectionApplyRequest) (*iop.PrincipalProjectionApplyResponse, error) {
		healthyApplied.Add(1)
		return &iop.PrincipalProjectionApplyResponse{Applied: true, AcceptedGeneration: request.GetProjection().GetGeneration()}, nil
	})
	sendEdgeHello(t, client, "edge-healthy")

	broadcastErr := make(chan error, 1)
	go func() { broadcastErr <- server.broadcastProjection(ctx, 300*time.Millisecond) }()

	// The stalled push must begin and the healthy push must complete while the
	// stalled handler is still blocked: the fan-out is concurrent, not serial.
	select {
	case <-stalledStart:
	case <-time.After(2 * time.Second):
		t.Fatal("stalled edge projection push never began")
	}
	deadline := time.Now().Add(2 * time.Second)
	for healthyApplied.Load() == 0 && time.Now().Before(deadline) {
		time.Sleep(5 * time.Millisecond)
	}
	if healthyApplied.Load() == 0 {
		t.Fatal("healthy edge was not attempted while the stalled edge blocked the fan-out")
	}

	// The broadcast must return within the per-client deadline bound, reporting the
	// stalled edge, without ever releasing the stalled handler.
	select {
	case err := <-broadcastErr:
		if err == nil {
			t.Fatal("expected an aggregate error for the stalled edge")
		}
		if !strings.Contains(err.Error(), "edge-stalled") {
			t.Fatalf("aggregate error did not name the stalled edge: %v", err)
		}
		if strings.Contains(err.Error(), "edge-healthy") {
			t.Fatalf("healthy edge must not appear in the aggregate error: %v", err)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("broadcast did not complete within the per-client deadline; a stalled edge blocked it")
	}
	closeRelease()
}

func TestEdgeServerStopCancelsStalledProjectionRefresh(t *testing.T) {
	port := getFreePort(t)
	server, err := NewEdgeServer(fmt.Sprintf("127.0.0.1:%d", port), zaptest.NewLogger(t))
	if err != nil {
		t.Fatal(err)
	}
	server.SetCredentialPlane(func(context.Context) (*iop.PrincipalProjection, error) {
		return testWireProjection(1), nil
	}, nil, 20*time.Millisecond)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	if err := server.Start(ctx); err != nil {
		t.Fatal(err)
	}

	// The only Edge stalls its projection push forever unless released.
	stalledStart, release := dialProjectionEdge(t, ctx, port, "edge-stalled", nil)
	var releaseOnce sync.Once
	defer releaseOnce.Do(func() { close(release) })

	// Wait until the periodic refresh has begun a push that the handler holds.
	select {
	case <-stalledStart:
	case <-time.After(2 * time.Second):
		t.Fatal("periodic projection refresh never reached the stalled edge")
	}

	// Stop must cancel the in-flight refresh and complete promptly without the
	// stalled handler ever responding.
	stopped := make(chan error, 1)
	go func() { stopped <- server.Stop() }()
	select {
	case stopErr := <-stopped:
		if stopErr != nil {
			t.Fatalf("Stop returned error while a projection push was stalled: %v", stopErr)
		}
	case <-time.After(3 * time.Second):
		t.Fatal("Stop did not complete while a periodic projection push was stalled")
	}
}

func TestEdgeServerBroadcastProjectionCancellationDoesNotWaitForPriorBatch(t *testing.T) {
	port := getFreePort(t)
	server, err := NewEdgeServer(fmt.Sprintf("127.0.0.1:%d", port), zaptest.NewLogger(t))
	if err != nil {
		t.Fatal(err)
	}
	server.SetCredentialPlane(func(context.Context) (*iop.PrincipalProjection, error) {
		return testWireProjection(1), nil
	}, nil, 0)

	if err := server.acquireProjectionBatch(context.Background()); err != nil {
		t.Fatal(err)
	}
	released := false
	defer func() {
		if !released {
			server.releaseProjectionBatch()
		}
	}()

	bCtx, bCancel := context.WithCancel(context.Background())
	broadcastErr := make(chan error, 1)
	go func() {
		broadcastErr <- server.BroadcastProjection(bCtx)
	}()

	bCancel()

	select {
	case err := <-broadcastErr:
		if !errors.Is(err, context.Canceled) {
			t.Fatalf("expected context.Canceled, got %v", err)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("BroadcastProjection did not observe cancellation while waiting for prior batch")
	}

	released = true
	server.releaseProjectionBatch()
}

type doneObservedContext struct {
	context.Context
	observed chan struct{}
	once     sync.Once
}

func (c *doneObservedContext) Done() <-chan struct{} {
	c.once.Do(func() { close(c.observed) })
	return c.Context.Done()
}

func TestEdgeServerStopCancelsProjectionRefreshWaitingForPriorBatch(t *testing.T) {
	port := getFreePort(t)
	server, err := NewEdgeServer(fmt.Sprintf("127.0.0.1:%d", port), zaptest.NewLogger(t))
	if err != nil {
		t.Fatal(err)
	}
	server.SetCredentialPlane(func(context.Context) (*iop.PrincipalProjection, error) {
		return testWireProjection(1), nil
	}, nil, 0)

	if err := server.acquireProjectionBatch(context.Background()); err != nil {
		t.Fatal(err)
	}
	released := false
	defer func() {
		if !released {
			server.releaseProjectionBatch()
		}
	}()

	refreshCtx, refreshCancel := context.WithCancel(context.Background())
	observed := make(chan struct{})
	doneCtx := &doneObservedContext{
		Context:  refreshCtx,
		observed: observed,
	}

	refreshDone := make(chan struct{})
	server.refreshMu.Lock()
	server.refreshCancel = refreshCancel
	server.refreshDone = refreshDone
	server.refreshMu.Unlock()

	broadcastErrCh := make(chan error, 1)
	go func() {
		defer close(refreshDone)
		broadcastErrCh <- server.BroadcastProjection(doneCtx)
	}()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	if err := server.Start(ctx); err != nil {
		t.Fatal(err)
	}

	select {
	case <-observed:
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for projection refresh to evaluate batch admission")
	}

	stopped := make(chan error, 1)
	go func() { stopped <- server.Stop() }()

	select {
	case stopErr := <-stopped:
		if stopErr != nil {
			t.Fatalf("Stop returned error while projection refresh was queued behind prior batch: %v", stopErr)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("Stop did not complete while projection refresh was queued behind prior batch")
	}

	select {
	case broadcastErr := <-broadcastErrCh:
		if !errors.Is(broadcastErr, context.Canceled) {
			t.Fatalf("expected context.Canceled from queued broadcast, got %v", broadcastErr)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("queued broadcast did not return after Stop")
	}

	released = true
	server.releaseProjectionBatch()
}

func TestEdgeHeartbeatWaitExceedsInterval(t *testing.T) {
	if EdgeHeartbeatWaitSec <= EdgeHeartbeatIntervalSec {
		t.Fatalf("EdgeHeartbeatWaitSec (%d) must exceed EdgeHeartbeatIntervalSec (%d)",
			EdgeHeartbeatWaitSec, EdgeHeartbeatIntervalSec)
	}
}