handshake: trust-transition preconditions, replay accounting, key-bound allowlist gate

handshake.go

@@ -79,6 +79,7 @@ const (
 	handshakeRecvTimeout  = 10 * time.Second       // time to wait for handshake message
 	handshakeCloseDelay   = 500 * time.Millisecond // delay before closing after send to let data flush
 	maxReplaySetEntries   = 8192                   // cap replay set to prevent unbounded growth between reaps
+	maxReplayPerPeer      = 256                    // cap replay entries attributable to any single peer
 	maxPendingHandshakes  = 256                    // cap pending (unapproved) handshake requests
 	maxPendingPerSource   = 16
 	maxJustificationLen   = 1024
@@ -92,6 +93,14 @@ const (
 	pendingHandshakeTTL = 30 * 24 * time.Hour
 )
 
+// replayEntry is one recorded handshake message hash: when it was first
+// accepted, and which peer node ID it was attributed to. The peer field
+// backs the per-peer accounting in recordReplay.
+type replayEntry struct {
+	seen time.Time
+	peer uint32
+}
+
 // Manager handles the trust handshake protocol on port 444.
 type Manager struct {
 	mu              sync.RWMutex
@@ -111,8 +120,9 @@ type Manager struct {
 	done            chan struct{}        // closed in Stop to signal drain goroutine exit
 
 	// Replay protection
-	replayMu  sync.Mutex
-	replaySet map[[32]byte]time.Time // message hash → first seen
+	replayMu   sync.Mutex
+	replaySet  map[[32]byte]replayEntry // message hash → first seen + attributed peer
+	replayPeer map[uint32]int           // peer → number of entries it holds in replaySet
 
 	// WaitForTrust notification: per-node list of channels closed when trust
 	// transitions to granted. Separate mu from hm.mu (always taken AFTER hm.mu)
@@ -132,7 +142,8 @@ func NewManager(rt Runtime) *Manager {
 		pendingBySource: make(map[uint32]int),
 		outgoing:        make(map[uint32]time.Time),
 		revoked:         make(map[uint32]time.Time),
-		replaySet:       make(map[[32]byte]time.Time),
+		replaySet:       make(map[[32]byte]replayEntry),
+		replayPeer:      make(map[uint32]int),
 		trustWaiters:    make(map[uint32][]chan struct{}),
 		dirty:           make(chan struct{}, 1),
 		done:            make(chan struct{}),
@@ -570,21 +581,15 @@ func (hm *Manager) processMessage(stream coreapi.Stream, msg *HandshakeMsg) {
 		return
 	}
 
+	// The replay set is consulted here but only written after the message
+	// has cleared authentication below, so an unauthenticated sender
+	// cannot consume replay-set capacity on behalf of a peer.
 	replayChallenge := fmt.Sprintf("handshake:%d:%d", msg.NodeID, hm.rt.NodeID())
 	msgHash := sha256.Sum256([]byte(msg.Type + "\x00" + replayChallenge + "\x00" + msg.Signature))
-	hm.replayMu.Lock()
-	if _, seen := hm.replaySet[msgHash]; seen {
-		hm.replayMu.Unlock()
+	if hm.replaySeen(msgHash) {
 		slog.Warn("handshake replay detected", "peer_node_id", msg.NodeID)
 		return
 	}
-	if len(hm.replaySet) >= maxReplaySetEntries {
-		hm.replayMu.Unlock()
-		slog.Warn("handshake replay set full, rejecting", "peer_node_id", msg.NodeID)
-		return
-	}
-	hm.replaySet[msgHash] = now
-	hm.replayMu.Unlock()
 
 	// M12 fix: verify P2P signature if the sender provides a public key
 	registryBound := false
@@ -632,18 +637,29 @@ func (hm *Manager) processMessage(stream coreapi.Stream, msg *HandshakeMsg) {
 		}
 	}
 
-	if msg.Type == HandshakeAccept || msg.Type == HandshakeRevoke {
+	// accept / reject / revoke all act on existing local state (they grant
+	// trust, cancel an outgoing request, or drop trust), so each must
+	// arrive on a transport whose authenticated node ID matches the
+	// claimed sender. Only request creates fresh state and is exempt.
+	if msg.Type == HandshakeAccept || msg.Type == HandshakeRevoke || msg.Type == HandshakeReject {
 		if stream == nil || stream.RemoteAddr().Node != msg.NodeID {
 			var authenticated uint32
 			if stream != nil {
 				authenticated = stream.RemoteAddr().Node
 			}
-			slog.Warn("handshake: accept/revoke node_id does not match authenticated sender, rejecting",
+			slog.Warn("handshake: accept/reject/revoke node_id does not match authenticated sender, rejecting",
 				"claimed_node_id", msg.NodeID, "authenticated_node_id", authenticated, "type", msg.Type)
 			return
 		}
 	}
 
+	// Message is authenticated: record it so an identical copy replayed
+	// later is dropped by the replaySeen check above.
+	if !hm.recordReplay(msg.NodeID, msgHash, now) {
+		slog.Warn("handshake replay detected", "peer_node_id", msg.NodeID)
+		return
+	}
+
 	switch msg.Type {
 	case HandshakeRequest:
 		hm.handleRequest(stream, msg, registryBound)
@@ -656,14 +672,101 @@ func (hm *Manager) processMessage(stream coreapi.Stream, msg *HandshakeMsg) {
 	}
 }
 
+// replaySeen reports whether msgHash is already recorded.
+func (hm *Manager) replaySeen(msgHash [32]byte) bool {
+	hm.replayMu.Lock()
+	defer hm.replayMu.Unlock()
+	_, seen := hm.replaySet[msgHash]
+	return seen
+}
+
+// recordReplay stores msgHash against peerNodeID. Returns false when the
+// hash was already present (the message is a duplicate).
+//
+// Capacity is enforced by eviction, not refusal. A peer that already
+// holds maxReplayPerPeer entries displaces its own oldest entry, so one
+// talkative peer cannot crowd out the rest. When the whole set is at
+// maxReplaySetEntries, expired entries are dropped first and the
+// globally oldest entry after that. Recording therefore always succeeds
+// for a message that is not itself a duplicate.
+func (hm *Manager) recordReplay(peerNodeID uint32, msgHash [32]byte, now time.Time) bool {
+	hm.replayMu.Lock()
+	defer hm.replayMu.Unlock()
+
+	if _, seen := hm.replaySet[msgHash]; seen {
+		return false
+	}
+	for hm.replayPeer[peerNodeID] >= maxReplayPerPeer {
+		if !hm.evictOldestReplayLocked(peerNodeID, true) {
+			break
+		}
+	}
+	if len(hm.replaySet) >= maxReplaySetEntries {
+		hm.reapReplayLocked(now)
+	}
+	for len(hm.replaySet) >= maxReplaySetEntries {
+		if !hm.evictOldestReplayLocked(0, false) {
+			break
+		}
+	}
+
+	hm.replaySet[msgHash] = replayEntry{seen: now, peer: peerNodeID}
+	hm.replayPeer[peerNodeID]++
+	return true
+}
+
+// evictOldestReplayLocked drops the oldest replay entry, restricted to
+// entries attributed to peerNodeID when scoped is true. Reports whether
+// an entry was removed. Caller MUST hold replayMu.
+func (hm *Manager) evictOldestReplayLocked(peerNodeID uint32, scoped bool) bool {
+	var (
+		oldestHash [32]byte
+		oldestSeen time.Time
+		found      bool
+	)
+	for hash, e := range hm.replaySet {
+		if scoped && e.peer != peerNodeID {
+			continue
+		}
+		if !found || e.seen.Before(oldestSeen) {
+			oldestHash, oldestSeen, found = hash, e.seen, true
+		}
+	}
+	if !found {
+		return false
+	}
+	hm.deleteReplayLocked(oldestHash)
+	return true
+}
+
+// deleteReplayLocked removes one entry and keeps the per-peer counter in
+// step. Caller MUST hold replayMu.
+func (hm *Manager) deleteReplayLocked(hash [32]byte) {
+	e, ok := hm.replaySet[hash]
+	if !ok {
+		return
+	}
+	delete(hm.replaySet, hash)
+	if hm.replayPeer[e.peer] <= 1 {
+		delete(hm.replayPeer, e.peer)
+		return
+	}
+	hm.replayPeer[e.peer]--
+}
+
 // reapReplay removes expired entries from the replay set.
 func (hm *Manager) reapReplay() {
 	hm.replayMu.Lock()
 	defer hm.replayMu.Unlock()
-	threshold := time.Now().Add(-2 * handshakeMaxAge)
-	for hash, seen := range hm.replaySet {
-		if seen.Before(threshold) {
-			delete(hm.replaySet, hash)
+	hm.reapReplayLocked(time.Now())
+}
+
+// reapReplayLocked is reapReplay's body. Caller MUST hold replayMu.
+func (hm *Manager) reapReplayLocked(now time.Time) {
+	threshold := now.Add(-2 * handshakeMaxAge)
+	for hash, e := range hm.replaySet {
+		if e.seen.Before(threshold) {
+			hm.deleteReplayLocked(hash)
 		}
 	}
 }
@@ -833,7 +936,10 @@ func (hm *Manager) handleRequest(stream coreapi.Stream, msg *HandshakeMsg, regis
 	// into auto-accept.
 	if registryBound {
 		// Trust gate via L10 TrustChecker. Nil → trust nothing.
-		name, ok := hm.rt.IsTrusted(peerNodeID)
+		// msg.PublicKey is the key the peer authenticated with and that
+		// the registry lookup bound to peerNodeID, so it is the key the
+		// allowlist pin is checked against.
+		name, ok := hm.trustedAgentName(peerNodeID, msg.PublicKey)
 		if ok {
 			hm.markTrustedLocked(peerNodeID, &TrustRecord{
 				NodeID:     peerNodeID,
@@ -908,6 +1014,11 @@ func (hm *Manager) handleRequest(stream coreapi.Stream, msg *HandshakeMsg, regis
 }
 
 // handleAccept processes a handshake acceptance from a peer.
+//
+// An acceptance only completes a handshake this node started: it is
+// matched against, and consumes, the hm.outgoing entry created by
+// SendRequest. Accepts with no matching entry are dropped, mirroring
+// processRelayedApproval's precondition on the relay path.
 func (hm *Manager) handleAccept(msg *HandshakeMsg) {
 	peerNodeID := msg.NodeID
 	slog.Info("handshake accepted by peer", "peer_node_id", peerNodeID)
@@ -925,6 +1036,13 @@ func (hm *Manager) handleAccept(msg *HandshakeMsg) {
 		delete(hm.revoked, peerNodeID)
 	}
 
+	// Trust follows a request we actually made: an accept with no
+	// matching outgoing entry is unsolicited and is dropped.
+	if _, ok := hm.outgoing[peerNodeID]; !ok {
+		slog.Warn("dropping handshake acceptance with no matching outgoing request", "peer_node_id", peerNodeID)
+		return
+	}
+
 	// An accept that carries a different key than the record we hold
 	// comes from a different key holder behind the same node ID. Drop
 	// the record instead of rebinding it — re-trusting is an explicit
@@ -1089,7 +1207,10 @@ func (hm *Manager) processRelayedRequest(fromNodeID uint32, justification string
 	// Auto-approve when the sender is in the embedded trusted-agents list.
 	// The registry relay vouches for the sender's node_id (only authenticated
 	// nodes can relay), so no separate registryBound check is needed here.
-	if name, ok := hm.rt.IsTrusted(fromNodeID); ok {
+	//
+	// The relay carries no peer public key, so the empty key is passed:
+	// allowlist entries that carry a pin do not match on this path.
+	if name, ok := hm.trustedAgentName(fromNodeID, ""); ok {
 		hm.markTrustedLocked(fromNodeID, &TrustRecord{
 			NodeID:     fromNodeID,
 			ApprovedAt: time.Now(),
@@ -1448,6 +1569,17 @@ func (hm *Manager) handleRevokeMsg(msg *HandshakeMsg) {
 	}
 }
 
+// trustedAgentName consults the allowlist gate for peerNodeID, passing
+// the authenticated peer key when the Runtime implements the key-bound
+// form. Runtimes that only implement IsTrusted keep the node-ID-only
+// answer.
+func (hm *Manager) trustedAgentName(peerNodeID uint32, pubKeyB64 string) (string, bool) {
+	if kc, ok := hm.rt.(KeyedTrustChecker); ok {
+		return kc.IsTrustedWithKey(peerNodeID, pubKeyB64)
+	}
+	return hm.rt.IsTrusted(peerNodeID)
+}
+
 // IsTrusted returns whether a peer has been approved.
 //
 // This answers the node-ID-only question, for call sites that have no

runtime.go

@@ -89,6 +89,21 @@ type Runtime interface {
 	Registry() RegistryClient
 }
 
+// KeyedTrustChecker is the key-bound form of Runtime.IsTrusted. It is
+// an optional extension: the handshake manager type-asserts its Runtime
+// for it and falls back to IsTrusted when the adapter does not provide
+// it, so existing Runtime implementations keep compiling unchanged.
+//
+// pubKeyB64 is the base64 Ed25519 key the peer proved possession of on
+// this handshake, or "" when no authenticated key is in scope (the
+// registry-relayed path). Implementations that carry a pinned key for
+// the node MUST require an exact match and MUST treat "" as a miss;
+// entries with no pin answer the node-ID-only question, matching
+// IsTrusted.
+type KeyedTrustChecker interface {
+	IsTrustedWithKey(nodeID uint32, pubKeyB64 string) (name string, ok bool)
+}
+
 // RegistryClient is the subset of pkg/registry/client.Client the
 // handshake manager uses. Defined here as an interface so the manager
 // stays free of import cycles and tests can supply a fake.

zz_ceiling_branches_test.go

@@ -11,6 +11,8 @@ import (
 	"path/filepath"
 	"testing"
 	"time"
+
+	"github.com/pilot-protocol/common/coreapi"
 )
 
 // This file targets every coverage hole identified by the iter-baseline
@@ -522,16 +524,17 @@ func TestReapStalePending_PersistsViaSaveTrust(t *testing.T) {
 }
 
 // -----------------------------------------------------------------------
-// processMessage: replay set FULL — second message past the cap is dropped
+// processMessage: a full replay set evicts rather than dropping the
+// message, and stays at the cap
 // -----------------------------------------------------------------------
 
-func TestProcessMessage_ReplaySetFullDropsNewMessages(t *testing.T) {
+func TestProcessMessage_ReplaySetFullStillDispatches(t *testing.T) {
 	t.Parallel()
 	hm := newTestHM(t, "")
 	t.Cleanup(hm.Stop)
 
-	// Fill the replay set to exactly the cap so the next message hits the
-	// "set full" branch.
+	// Fill the replay set to exactly the cap, spread across many peers so
+	// no single peer is over its own limit.
 	hm.replayMu.Lock()
 	for i := 0; i < maxReplaySetEntries; i++ {
 		var h [32]byte
@@ -540,7 +543,9 @@ func TestProcessMessage_ReplaySetFullDropsNewMessages(t *testing.T) {
 		h[1] = byte(i >> 8)
 		h[2] = byte(i >> 16)
 		h[3] = byte(i >> 24)
-		hm.replaySet[h] = time.Now()
+		peer := uint32(1000 + i)
+		hm.replaySet[h] = replayEntry{seen: time.Now(), peer: peer}
+		hm.replayPeer[peer]++
 	}
 	hm.replayMu.Unlock()
 
@@ -550,15 +555,22 @@ func TestProcessMessage_ReplaySetFullDropsNewMessages(t *testing.T) {
 		Reason:    "test",
 		Timestamp: time.Now().Unix(),
 	}
-	// Seed an outgoing entry so a dispatch would be observable.
+	// Seed an outgoing entry so the dispatch is observable.
 	hm.outgoing[77] = time.Now()
-	hm.processMessage(nil, msg)
+	hm.processMessage(&addrStream{addr: coreapi.Addr{Node: 77}}, msg)
 
 	hm.mu.RLock()
 	_, stillOut := hm.outgoing[77]
 	hm.mu.RUnlock()
-	if !stillOut {
-		t.Fatal("replay-set-full path must early-return — outgoing[77] should be preserved")
+	if stillOut {
+		t.Fatal("a full replay set must not block an authenticated message — outgoing[77] should have been cleared")
+	}
+
+	hm.replayMu.Lock()
+	size := len(hm.replaySet)
+	hm.replayMu.Unlock()
+	if size > maxReplaySetEntries {
+		t.Fatalf("replay set grew past the cap: %d > %d", size, maxReplaySetEntries)
 	}
 }