test: add transport layer tests (Phase 2) — node/ coverage 42% to 63.5%

13 tests covering all Phase 2 TODO items: reusable test pair helper,
full handshake with challenge-response verification, protocol version
rejection, allowlist rejection, SMSG encrypted message round-trip,
message deduplication, rate limiting (100 burst verified), MaxConns
enforcement, keepalive timeout cleanup, graceful disconnect, and
concurrent sends with race detection clean.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

node/transport_test.go

@@ -0,0 +1,592 @@
+package node
+
+import (
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"path/filepath"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/gorilla/websocket"
+)
+
+// --- Test Helpers ---
+
+// testNode creates a NodeManager with a generated identity in a temp directory.
+func testNode(t *testing.T, name string, role NodeRole) *NodeManager {
+	t.Helper()
+	dir := t.TempDir()
+	nm, err := NewNodeManagerWithPaths(
+		filepath.Join(dir, "private.key"),
+		filepath.Join(dir, "node.json"),
+	)
+	if err != nil {
+		t.Fatalf("create node manager %q: %v", name, err)
+	}
+	if err := nm.GenerateIdentity(name, role); err != nil {
+		t.Fatalf("generate identity %q: %v", name, err)
+	}
+	return nm
+}
+
+// testRegistry creates a PeerRegistry with open auth in a temp directory.
+func testRegistry(t *testing.T) *PeerRegistry {
+	t.Helper()
+	dir := t.TempDir()
+	reg, err := NewPeerRegistryWithPath(filepath.Join(dir, "peers.json"))
+	if err != nil {
+		t.Fatalf("create registry: %v", err)
+	}
+	t.Cleanup(func() { reg.Close() })
+	return reg
+}
+
+// testTransportPair holds everything needed for transport pair tests.
+type testTransportPair struct {
+	Server     *Transport
+	Client     *Transport
+	ServerNode *NodeManager
+	ClientNode *NodeManager
+	ServerReg  *PeerRegistry
+	ClientReg  *PeerRegistry
+	HTTPServer *httptest.Server
+	ServerAddr string // "127.0.0.1:PORT"
+}
+
+// setupTestTransportPair creates a server transport (backed by httptest) and a
+// client transport, both with generated identities and open-auth registries.
+func setupTestTransportPair(t *testing.T) *testTransportPair {
+	return setupTestTransportPairWithConfig(t, DefaultTransportConfig(), DefaultTransportConfig())
+}
+
+// setupTestTransportPairWithConfig allows custom configs for server and client.
+func setupTestTransportPairWithConfig(t *testing.T, serverCfg, clientCfg TransportConfig) *testTransportPair {
+	t.Helper()
+
+	serverNM := testNode(t, "server", RoleWorker)
+	clientNM := testNode(t, "client", RoleController)
+	serverReg := testRegistry(t)
+	clientReg := testRegistry(t)
+
+	serverTransport := NewTransport(serverNM, serverReg, serverCfg)
+	clientTransport := NewTransport(clientNM, clientReg, clientCfg)
+
+	// Use httptest.Server with the transport's WebSocket handler
+	mux := http.NewServeMux()
+	mux.HandleFunc(serverCfg.WSPath, serverTransport.handleWSUpgrade)
+	ts := httptest.NewServer(mux)
+
+	u, _ := url.Parse(ts.URL)
+
+	tp := &testTransportPair{
+		Server:     serverTransport,
+		Client:     clientTransport,
+		ServerNode: serverNM,
+		ClientNode: clientNM,
+		ServerReg:  serverReg,
+		ClientReg:  clientReg,
+		HTTPServer: ts,
+		ServerAddr: u.Host,
+	}
+
+	t.Cleanup(func() {
+		clientTransport.Stop()
+		serverTransport.Stop()
+		ts.Close()
+	})
+
+	return tp
+}
+
+// connectClient establishes a connection from client to server transport.
+func (tp *testTransportPair) connectClient(t *testing.T) *PeerConnection {
+	t.Helper()
+
+	peer := &Peer{
+		ID:      tp.ServerNode.GetIdentity().ID,
+		Name:    "server",
+		Address: tp.ServerAddr,
+		Role:    RoleWorker,
+	}
+	tp.ClientReg.AddPeer(peer)
+
+	pc, err := tp.Client.Connect(peer)
+	if err != nil {
+		t.Fatalf("client connect failed: %v", err)
+	}
+	return pc
+}
+
+// --- Unit Tests for Sub-Components ---
+
+func TestMessageDeduplicator(t *testing.T) {
+	t.Run("MarkAndCheck", func(t *testing.T) {
+		d := NewMessageDeduplicator(5 * time.Minute)
+
+		if d.IsDuplicate("msg-1") {
+			t.Error("should not be duplicate before marking")
+		}
+
+		d.Mark("msg-1")
+
+		if !d.IsDuplicate("msg-1") {
+			t.Error("should be duplicate after marking")
+		}
+
+		if d.IsDuplicate("msg-2") {
+			t.Error("different ID should not be duplicate")
+		}
+	})
+
+	t.Run("Cleanup", func(t *testing.T) {
+		d := NewMessageDeduplicator(50 * time.Millisecond)
+		d.Mark("msg-1")
+
+		if !d.IsDuplicate("msg-1") {
+			t.Error("should be duplicate immediately after marking")
+		}
+
+		time.Sleep(60 * time.Millisecond)
+		d.Cleanup()
+
+		if d.IsDuplicate("msg-1") {
+			t.Error("should not be duplicate after TTL + cleanup")
+		}
+	})
+
+	t.Run("ConcurrentAccess", func(t *testing.T) {
+		d := NewMessageDeduplicator(5 * time.Minute)
+		var wg sync.WaitGroup
+		for i := 0; i < 100; i++ {
+			wg.Add(1)
+			go func(id int) {
+				defer wg.Done()
+				msgID := "msg-" + time.Now().String()
+				d.Mark(msgID)
+				d.IsDuplicate(msgID)
+			}(i)
+		}
+		wg.Wait()
+	})
+}
+
+func TestPeerRateLimiter(t *testing.T) {
+	t.Run("AllowUpToBurst", func(t *testing.T) {
+		rl := NewPeerRateLimiter(10, 5)
+
+		for i := 0; i < 10; i++ {
+			if !rl.Allow() {
+				t.Errorf("should allow message %d (within burst)", i)
+			}
+		}
+
+		if rl.Allow() {
+			t.Error("should reject message after burst exhausted")
+		}
+	})
+
+	t.Run("RefillAfterTime", func(t *testing.T) {
+		rl := NewPeerRateLimiter(5, 10) // 5 burst, 10/sec refill
+
+		// Exhaust all tokens
+		for i := 0; i < 5; i++ {
+			rl.Allow()
+		}
+
+		if rl.Allow() {
+			t.Error("should reject after exhaustion")
+		}
+
+		// Wait for refill
+		time.Sleep(1100 * time.Millisecond)
+
+		if !rl.Allow() {
+			t.Error("should allow after refill")
+		}
+	})
+}
+
+// --- Transport Integration Tests ---
+
+func TestTransport_FullHandshake(t *testing.T) {
+	tp := setupTestTransportPair(t)
+	pc := tp.connectClient(t)
+
+	// Shared secret must be derived
+	if len(pc.SharedSecret) == 0 {
+		t.Error("shared secret should be derived after handshake")
+	}
+
+	// Allow server goroutines to register the connection
+	time.Sleep(50 * time.Millisecond)
+
+	if tp.Server.ConnectedPeers() != 1 {
+		t.Errorf("server connected peers: got %d, want 1", tp.Server.ConnectedPeers())
+	}
+	if tp.Client.ConnectedPeers() != 1 {
+		t.Errorf("client connected peers: got %d, want 1", tp.Client.ConnectedPeers())
+	}
+
+	// Verify peer identity was exchanged correctly
+	serverID := tp.ServerNode.GetIdentity().ID
+	serverConn := tp.Client.GetConnection(serverID)
+	if serverConn == nil {
+		t.Fatal("client should have connection to server by server ID")
+	}
+	if serverConn.Peer.Name != "server" {
+		t.Errorf("peer name: got %q, want %q", serverConn.Peer.Name, "server")
+	}
+}
+
+func TestTransport_HandshakeRejectWrongVersion(t *testing.T) {
+	tp := setupTestTransportPair(t)
+
+	// Dial raw WebSocket and send handshake with unsupported version
+	wsURL := "ws://" + tp.ServerAddr + "/ws"
+	conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
+	if err != nil {
+		t.Fatalf("raw dial: %v", err)
+	}
+	defer conn.Close()
+
+	clientIdentity := tp.ClientNode.GetIdentity()
+	payload := HandshakePayload{
+		Identity: *clientIdentity,
+		Version:  "99.99", // Unsupported
+	}
+	msg, _ := NewMessage(MsgHandshake, clientIdentity.ID, "", payload)
+	data, _ := MarshalJSON(msg)
+
+	if err := conn.WriteMessage(websocket.TextMessage, data); err != nil {
+		t.Fatalf("write handshake: %v", err)
+	}
+
+	_, respData, err := conn.ReadMessage()
+	if err != nil {
+		t.Fatalf("read response: %v", err)
+	}
+
+	var resp Message
+	if err := json.Unmarshal(respData, &resp); err != nil {
+		t.Fatalf("unmarshal response: %v", err)
+	}
+
+	var ack HandshakeAckPayload
+	resp.ParsePayload(&ack)
+
+	if ack.Accepted {
+		t.Error("should reject incompatible protocol version")
+	}
+	if !strings.Contains(ack.Reason, "incompatible protocol version") {
+		t.Errorf("expected version rejection reason, got: %s", ack.Reason)
+	}
+}
+
+func TestTransport_HandshakeRejectAllowlist(t *testing.T) {
+	tp := setupTestTransportPair(t)
+
+	// Switch server to allowlist mode WITHOUT adding client's key
+	tp.ServerReg.SetAuthMode(PeerAuthAllowlist)
+
+	peer := &Peer{
+		ID:      tp.ServerNode.GetIdentity().ID,
+		Name:    "server",
+		Address: tp.ServerAddr,
+		Role:    RoleWorker,
+	}
+	tp.ClientReg.AddPeer(peer)
+
+	_, err := tp.Client.Connect(peer)
+	if err == nil {
+		t.Fatal("should reject peer not in allowlist")
+	}
+	if !strings.Contains(err.Error(), "rejected") {
+		t.Errorf("expected rejection error, got: %v", err)
+	}
+}
+
+func TestTransport_EncryptedMessageRoundTrip(t *testing.T) {
+	tp := setupTestTransportPair(t)
+
+	received := make(chan *Message, 1)
+	tp.Server.OnMessage(func(conn *PeerConnection, msg *Message) {
+		received <- msg
+	})
+
+	pc := tp.connectClient(t)
+
+	// Send an encrypted message from client to server
+	clientID := tp.ClientNode.GetIdentity().ID
+	serverID := tp.ServerNode.GetIdentity().ID
+	sentMsg, _ := NewMessage(MsgPing, clientID, serverID, PingPayload{
+		SentAt: time.Now().UnixMilli(),
+	})
+
+	if err := pc.Send(sentMsg); err != nil {
+		t.Fatalf("send: %v", err)
+	}
+
+	select {
+	case msg := <-received:
+		if msg.Type != MsgPing {
+			t.Errorf("type: got %s, want %s", msg.Type, MsgPing)
+		}
+		if msg.ID != sentMsg.ID {
+			t.Error("message ID mismatch after encrypt/decrypt round-trip")
+		}
+		if msg.From != clientID {
+			t.Errorf("from: got %s, want %s", msg.From, clientID)
+		}
+
+		var payload PingPayload
+		msg.ParsePayload(&payload)
+		if payload.SentAt == 0 {
+			t.Error("payload should have SentAt timestamp")
+		}
+	case <-time.After(5 * time.Second):
+		t.Fatal("timeout waiting for message")
+	}
+}
+
+func TestTransport_MessageDedup(t *testing.T) {
+	tp := setupTestTransportPair(t)
+
+	var count atomic.Int32
+	tp.Server.OnMessage(func(conn *PeerConnection, msg *Message) {
+		count.Add(1)
+	})
+
+	pc := tp.connectClient(t)
+
+	clientID := tp.ClientNode.GetIdentity().ID
+	serverID := tp.ServerNode.GetIdentity().ID
+	msg, _ := NewMessage(MsgPing, clientID, serverID, PingPayload{SentAt: time.Now().UnixMilli()})
+
+	// Send the same message twice
+	if err := pc.Send(msg); err != nil {
+		t.Fatalf("first send: %v", err)
+	}
+	time.Sleep(100 * time.Millisecond) // Ensure first is processed and marked
+
+	if err := pc.Send(msg); err != nil {
+		t.Fatalf("second send: %v", err)
+	}
+	time.Sleep(100 * time.Millisecond) // Allow time for second to be processed (or dropped)
+
+	if got := count.Load(); got != 1 {
+		t.Errorf("expected 1 message delivered (dedup), got %d", got)
+	}
+}
+
+func TestTransport_RateLimiting(t *testing.T) {
+	tp := setupTestTransportPair(t)
+
+	var count atomic.Int32
+	tp.Server.OnMessage(func(conn *PeerConnection, msg *Message) {
+		count.Add(1)
+	})
+
+	pc := tp.connectClient(t)
+
+	clientID := tp.ClientNode.GetIdentity().ID
+	serverID := tp.ServerNode.GetIdentity().ID
+
+	// Send 150 messages rapidly (rate limiter burst = 100)
+	for i := 0; i < 150; i++ {
+		msg, _ := NewMessage(MsgPing, clientID, serverID, PingPayload{SentAt: time.Now().UnixMilli()})
+		pc.Send(msg)
+	}
+
+	time.Sleep(1 * time.Second) // Allow processing
+
+	received := int(count.Load())
+	t.Logf("rate limiting: %d/150 messages delivered", received)
+
+	if received >= 150 {
+		t.Error("rate limiting should have dropped some messages")
+	}
+	if received < 50 {
+		t.Errorf("too few messages received (%d), rate limiter may be too aggressive", received)
+	}
+}
+
+func TestTransport_MaxConnsEnforcement(t *testing.T) {
+	// Server with MaxConns=1
+	serverNM := testNode(t, "maxconns-server", RoleWorker)
+	serverReg := testRegistry(t)
+
+	serverCfg := DefaultTransportConfig()
+	serverCfg.MaxConns = 1
+	serverTransport := NewTransport(serverNM, serverReg, serverCfg)
+
+	mux := http.NewServeMux()
+	mux.HandleFunc(serverCfg.WSPath, serverTransport.handleWSUpgrade)
+	ts := httptest.NewServer(mux)
+	t.Cleanup(func() {
+		serverTransport.Stop()
+		ts.Close()
+	})
+
+	u, _ := url.Parse(ts.URL)
+	serverAddr := u.Host
+
+	// First client connects successfully
+	client1NM := testNode(t, "client1", RoleController)
+	client1Reg := testRegistry(t)
+	client1Transport := NewTransport(client1NM, client1Reg, DefaultTransportConfig())
+	t.Cleanup(func() { client1Transport.Stop() })
+
+	peer1 := &Peer{ID: serverNM.GetIdentity().ID, Name: "server", Address: serverAddr, Role: RoleWorker}
+	client1Reg.AddPeer(peer1)
+
+	_, err := client1Transport.Connect(peer1)
+	if err != nil {
+		t.Fatalf("first connection should succeed: %v", err)
+	}
+
+	// Allow server to register the connection
+	time.Sleep(50 * time.Millisecond)
+
+	// Second client should be rejected (MaxConns=1 reached)
+	client2NM := testNode(t, "client2", RoleController)
+	client2Reg := testRegistry(t)
+	client2Transport := NewTransport(client2NM, client2Reg, DefaultTransportConfig())
+	t.Cleanup(func() { client2Transport.Stop() })
+
+	peer2 := &Peer{ID: serverNM.GetIdentity().ID, Name: "server", Address: serverAddr, Role: RoleWorker}
+	client2Reg.AddPeer(peer2)
+
+	_, err = client2Transport.Connect(peer2)
+	if err == nil {
+		t.Fatal("second connection should be rejected when MaxConns=1")
+	}
+}
+
+func TestTransport_KeepaliveTimeout(t *testing.T) {
+	// Use short keepalive settings so the test is fast
+	serverCfg := DefaultTransportConfig()
+	serverCfg.PingInterval = 100 * time.Millisecond
+	serverCfg.PongTimeout = 100 * time.Millisecond
+
+	clientCfg := DefaultTransportConfig()
+	clientCfg.PingInterval = 100 * time.Millisecond
+	clientCfg.PongTimeout = 100 * time.Millisecond
+
+	tp := setupTestTransportPairWithConfig(t, serverCfg, clientCfg)
+	tp.connectClient(t)
+
+	// Verify connection is established
+	time.Sleep(50 * time.Millisecond)
+	if tp.Server.ConnectedPeers() != 1 {
+		t.Fatalf("server should have 1 peer initially, got %d", tp.Server.ConnectedPeers())
+	}
+
+	// Close the underlying WebSocket on the client side to simulate network failure.
+	// The server's readLoop will detect the broken connection and clean up.
+	clientID := tp.ClientNode.GetIdentity().ID
+	serverPeerID := tp.ServerNode.GetIdentity().ID
+	clientConn := tp.Client.GetConnection(serverPeerID)
+	if clientConn == nil {
+		t.Fatal("client should have connection to server")
+	}
+	clientConn.Conn.Close()
+
+	// Wait for server to detect and clean up
+	deadline := time.After(2 * time.Second)
+	for {
+		select {
+		case <-deadline:
+			t.Fatalf("server did not clean up connection: still has %d peers", tp.Server.ConnectedPeers())
+		default:
+			if tp.Server.ConnectedPeers() == 0 {
+				// Verify registry updated
+				peer := tp.ServerReg.GetPeer(clientID)
+				if peer != nil && peer.Connected {
+					t.Error("registry should show peer as disconnected")
+				}
+				return
+			}
+			time.Sleep(50 * time.Millisecond)
+		}
+	}
+}
+
+func TestTransport_GracefulClose(t *testing.T) {
+	tp := setupTestTransportPair(t)
+
+	received := make(chan *Message, 10)
+	tp.Server.OnMessage(func(conn *PeerConnection, msg *Message) {
+		received <- msg
+	})
+
+	pc := tp.connectClient(t)
+
+	// Allow connection to fully establish
+	time.Sleep(50 * time.Millisecond)
+
+	// Graceful close should send a MsgDisconnect before closing
+	pc.GracefulClose("test shutdown", DisconnectNormal)
+
+	// Check if disconnect message was received
+	select {
+	case msg := <-received:
+		if msg.Type != MsgDisconnect {
+			t.Errorf("expected disconnect message, got %s", msg.Type)
+		}
+		var payload DisconnectPayload
+		msg.ParsePayload(&payload)
+		if payload.Reason != "test shutdown" {
+			t.Errorf("disconnect reason: got %q, want %q", payload.Reason, "test shutdown")
+		}
+		if payload.Code != DisconnectNormal {
+			t.Errorf("disconnect code: got %d, want %d", payload.Code, DisconnectNormal)
+		}
+	case <-time.After(2 * time.Second):
+		t.Error("timeout waiting for disconnect message")
+	}
+}
+
+func TestTransport_ConcurrentSends(t *testing.T) {
+	tp := setupTestTransportPair(t)
+
+	var count atomic.Int32
+	tp.Server.OnMessage(func(conn *PeerConnection, msg *Message) {
+		count.Add(1)
+	})
+
+	pc := tp.connectClient(t)
+
+	clientID := tp.ClientNode.GetIdentity().ID
+	serverID := tp.ServerNode.GetIdentity().ID
+
+	// Spawn 10 goroutines each sending 5 messages concurrently
+	const goroutines = 10
+	const msgsPerGoroutine = 5
+	var wg sync.WaitGroup
+
+	for g := 0; g < goroutines; g++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for i := 0; i < msgsPerGoroutine; i++ {
+				msg, _ := NewMessage(MsgPing, clientID, serverID, PingPayload{SentAt: time.Now().UnixMilli()})
+				pc.Send(msg)
+			}
+		}()
+	}
+
+	wg.Wait()
+	time.Sleep(1 * time.Second) // Allow delivery
+
+	got := int(count.Load())
+	// All messages should be delivered (unique IDs, within rate limit burst of 100)
+	expected := goroutines * msgsPerGoroutine
+	if got != expected {
+		t.Errorf("concurrent sends: got %d/%d messages delivered", got, expected)
+	}
+}