beacon: punch-grant authoritative-key verification + panic-recover hardening

recover.go

@@ -0,0 +1,53 @@
+// SPDX-License-Identifier: AGPL-3.0-or-later
+
+package beacon
+
+import (
+	"log/slog"
+	"runtime/debug"
+	"sync/atomic"
+)
+
+// recoveredPanicCount is a process-global counter incremented every time
+// recoverHandler swallows a panic. Exposed via RecoveredPanicCount so the
+// dashboard / Prom scrape can alert on a non-zero rate — a panic is still a
+// bug and must be visible; it just must not take the process down. The
+// beacon runs embedded in the registry binary, so an unrecovered panic on a
+// packet-handling goroutine takes the whole trust root with it.
+var recoveredPanicCount atomic.Uint64
+
+// RecoveredPanicCount returns the total number of panics swallowed by
+// recoverHandler since process start.
+func RecoveredPanicCount() uint64 {
+	return recoveredPanicCount.Load()
+}
+
+// recoverHandler is the standard panic-recovery shim placed at the top of
+// every packet-handling and background-loop goroutine in this package.
+// Mirrors the shim in rendezvous/accept. Usage:
+//
+//	defer recoverHandler("handlePacket")
+//
+// It must be the OUTERMOST defer in the function: defers run LIFO, so any
+// mutex unlock / buffer return registered later still executes before the
+// recover observes the panic.
+func recoverHandler(label string) {
+	r := recover()
+	if r == nil {
+		return
+	}
+	count := recoveredPanicCount.Add(1)
+	slog.Error("beacon: panic recovered",
+		"where", label,
+		"panic", r,
+		"recovered_total", count,
+		"stack", string(debug.Stack()),
+	)
+}
+
+// safely runs fn under recoverHandler. Used by the periodic background
+// loops so one bad tick cannot retire the ticker goroutine.
+func safely(label string, fn func()) {
+	defer recoverHandler(label)
+	fn()
+}

server.go

@@ -109,6 +109,8 @@ type Server struct {
 	// handlePacket the same way UDP datagrams do.
 	wssServer *bwss.Server
 
+	authKeyLookup atomic.Pointer[authKeyLookupFn]
+
 	done chan struct{} // closed on shutdown
 
 	// Close idempotency. The actual teardown runs at most once
@@ -217,6 +219,16 @@ func NewWithPeers(beaconID uint32, peers []string) *Server {
 	return s
 }
 
+type authKeyLookupFn func(nodeID uint32) (ed25519.PublicKey, bool)
+
+func (s *Server) SetAuthoritativeKeyLookup(fn authKeyLookupFn) {
+	if fn == nil {
+		s.authKeyLookup.Store(nil)
+		return
+	}
+	s.authKeyLookup.Store(&fn)
+}
+
 // EnableCompatWSS attaches a WSS-bridge listener for compat-mode
 // daemons. After Start, the beacon's relay worker checks the WSS
 // peer map BEFORE the UDP tier-1/2 lookups: relay packets destined
@@ -235,6 +247,11 @@ func (s *Server) EnableCompatWSS(bindAddr string, pubKeyLookup bwss.PubKeyLookup
 	if s.wssServer != nil {
 		return fmt.Errorf("beacon: compat WSS already enabled")
 	}
+	if pubKeyLookup != nil {
+		s.SetAuthoritativeKeyLookup(func(id uint32) (ed25519.PublicKey, bool) {
+			return pubKeyLookup(id)
+		})
+	}
 	ws, err := bwss.New(bwss.Config{
 		BindAddr:     bindAddr,
 		PubKeyLookup: pubKeyLookup,
@@ -454,18 +471,40 @@ const readBatchCap = 32
 // allocation contention. On non-Linux platforms the x/net/ipv4
 // fallback degrades to per-message ReadFrom, preserving correctness.
 func (s *Server) readLoop(conn *net.UDPConn) error {
+	// A panic on the receive path must not silently retire this socket's
+	// reader: recover, log, and resume the batch loop. Only a closed
+	// socket ends it. The per-packet boundary lives in handlePacket; this
+	// one covers the batch bookkeeping around it.
+	for {
+		err, done := s.readLoopBatches(conn)
+		if done {
+			return err
+		}
+		select {
+		case <-s.done:
+			return nil
+		default:
+		}
+	}
+}
+
+// readLoopBatches runs the recvmmsg loop until the socket closes (done=true)
+// or a panic unwinds it (done=false, caller resumes).
+func (s *Server) readLoopBatches(conn *net.UDPConn) (err error, done bool) {
+	defer recoverHandler("readLoop")
+
 	pc := ipv4.NewPacketConn(conn)
 	msgs := make([]ipv4.Message, readBatchCap)
 	for i := range msgs {
 		msgs[i].Buffers = [][]byte{make([]byte, 65535)}
 	}
 	for {
-		n, err := pc.ReadBatch(msgs, 0)
-		if err != nil {
-			if opErr, ok := err.(*net.OpError); ok && opErr.Err.Error() == "use of closed network connection" {
-				return nil
+		n, rerr := pc.ReadBatch(msgs, 0)
+		if rerr != nil {
+			if opErr, ok := rerr.(*net.OpError); ok && opErr.Err.Error() == "use of closed network connection" {
+				return nil, true
 			}
-			slog.Debug("beacon read batch error", "err", err)
+			slog.Debug("beacon read batch error", "err", rerr)
 			continue
 		}
 		for i := 0; i < n; i++ {
@@ -530,6 +569,11 @@ func (s *Server) RelayDropped() uint64 { return s.relayDropped.Load() }
 func (s *Server) RelayNotFound() uint64 { return s.relayNotFound.Load() }
 
 func (s *Server) handlePacket(data []byte, remote *net.UDPAddr) {
+	// Outermost defer: a panic anywhere in the dispatch below drops this
+	// one datagram and leaves the read loop running. Every inbound packet
+	// on this path is unauthenticated remote input.
+	defer recoverHandler("handlePacket")
+
 	// Empty-frame guard. The UDP readLoop screens len(data)<1, but the
 	// WSS OnFrame path only enforces an upper bound — a 0-byte binary
 	// frame from any authenticated WSS peer used to crash the beacon
@@ -579,7 +623,7 @@ func (s *Server) handleDiscover(data []byte, remote *net.UDPAddr, wantDest bool)
 	if len(data) >= 4+ed25519.PublicKeySize {
 		pub := make(ed25519.PublicKey, ed25519.PublicKeySize)
 		copy(pub, data[4:4+ed25519.PublicKeySize])
-		s.nodePubKeys.Store(nodeID, pub)
+		s.nodePubKeys.LoadOrStore(nodeID, pub)
 	}
 
 	// Per-nodeID endpoint update rate limit (PILOT-334) — prevents a single
@@ -617,6 +661,12 @@ func (s *Server) handleDiscover(data []byte, remote *net.UDPAddr, wantDest bool)
 	copy(reply[2:2+len(ip)], ip)
 	binary.BigEndian.PutUint16(reply[2+len(ip):], uint16(remote.Port))
 
+	// The compat WSS bridge starts accepting frames in EnableCompatWSS,
+	// which runs before ListenAndServe binds the UDP sockets — a frame
+	// arriving in that window reaches here with no socket to reply on.
+	if s.conn == nil {
+		return
+	}
 	if _, err := s.conn.WriteToUDP(reply, remote); err != nil {
 		slog.Debug("beacon discover reply failed", "node_id", nodeID, "err", err)
 	}
@@ -634,18 +684,29 @@ func (s *Server) verifyPunchGrant(trailer []byte, requesterID, targetID uint32)
 		return false
 	}
 	sig := trailer[punchGrantExpirySize : punchGrantExpirySize+ed25519.SignatureSize]
-	v, ok := s.nodePubKeys.Load(targetID)
-	if !ok {
-		return false
-	}
-	pub, ok := v.(ed25519.PublicKey)
+	pub, ok := s.targetPubKey(targetID)
 	if !ok || len(pub) != ed25519.PublicKeySize {
 		return false
 	}
 	preimage := fmt.Sprintf("punch-grant:%d:%d:%d", requesterID, targetID, expiry)
 	return ed25519.Verify(pub, []byte(preimage), sig)
 }
 
+func (s *Server) targetPubKey(targetID uint32) (ed25519.PublicKey, bool) {
+	if p := s.authKeyLookup.Load(); p != nil {
+		return (*p)(targetID)
+	}
+	v, ok := s.nodePubKeys.Load(targetID)
+	if !ok {
+		return nil, false
+	}
+	pub, ok := v.(ed25519.PublicKey)
+	if !ok {
+		return nil, false
+	}
+	return pub, true
+}
+
 func (s *Server) handlePunchRequest(data []byte, remote *net.UDPAddr) {
 	if len(data) < 8 {
 		return
@@ -897,6 +958,26 @@ const relayFlushAfter = 2 * time.Millisecond
 // With per-fd workers, sends parallelise across fds and the only
 // serialisation left is per-worker (its own batch state).
 func (s *Server) relayWorker(sendConn *ipv4.PacketConn, rawConn *net.UDPConn) {
+	// A panic while shaping one relay job must not permanently retire this
+	// worker — that would cut the beacon's drain capacity by 1/N for the
+	// life of the process. Resume with fresh batch state instead.
+	for {
+		if s.relayWorkerLoop(sendConn, rawConn) {
+			return
+		}
+		select {
+		case <-s.done:
+			return
+		default:
+		}
+	}
+}
+
+// relayWorkerLoop is the relayWorker body. Returns true when the server is
+// shutting down, false when a panic unwound it and the caller should resume.
+func (s *Server) relayWorkerLoop(sendConn *ipv4.PacketConn, rawConn *net.UDPConn) (shutdown bool) {
+	defer recoverHandler("relayWorker")
+
 	msgs := make([]ipv4.Message, 0, relayBatchCap)
 	payloadsToReturn := make([][]byte, 0, relayBatchCap)
 	outBufsToReturn := make([]*[]byte, 0, relayBatchCap)
@@ -983,7 +1064,7 @@ func (s *Server) relayWorker(sendConn *ipv4.PacketConn, rawConn *net.UDPConn) {
 		select {
 		case <-s.done:
 			flush()
-			return
+			return true
 		case <-timer.C:
 			timerActive = false
 			flush()
@@ -1138,6 +1219,11 @@ func (s *Server) relayStatsLoop() {
 
 // SendPunchCommand tells a node to send UDP to a target endpoint.
 func (s *Server) SendPunchCommand(nodeID uint32, targetIP net.IP, targetPort uint16) error {
+	// See handleDiscover: reachable via the compat WSS bridge before
+	// ListenAndServe has bound the UDP sockets.
+	if s.conn == nil {
+		return fmt.Errorf("beacon: not listening")
+	}
 	// Snapshot returns the address under RLock so we don't race with
 	// a concurrent Upsert mutating beaconNode.addr — see
 	// nodes_shard.go Get() caveat (race detector flag, 2026-05-19).
@@ -1177,7 +1263,7 @@ func (s *Server) reapLoop() {
 	for {
 		select {
 		case <-ticker.C:
-			s.reapStaleNodes()
+			safely("reapStaleNodes", s.reapStaleNodes)
 		case <-s.done:
 			return
 		}
@@ -1238,7 +1324,7 @@ func (s *Server) gossipLoop() {
 	for {
 		select {
 		case <-ticker.C:
-			s.sendGossip()
+			safely("sendGossip", s.sendGossip)
 		case <-s.done:
 			return
 		}
@@ -1420,11 +1506,11 @@ func (s *Server) registryDiscoveryLoop() {
 	defer ticker.Stop()
 
 	// Run immediately, then on tick
-	s.registryDiscover()
+	safely("registryDiscover", s.registryDiscover)
 	for {
 		select {
 		case <-ticker.C:
-			s.registryDiscover()
+			safely("registryDiscover", s.registryDiscover)
 		case <-s.done:
 			return
 		}

server_test.go

@@ -30,7 +30,7 @@ func registerNode(t *testing.T, beaconAddr *net.UDPAddr, nodeID uint32) *net.UDP
 
 	// Read discover reply
 	buf := make([]byte, 64)
-	conn.SetReadDeadline(time.Now().Add(2 * time.Second))
+	conn.SetReadDeadline(time.Now().Add(10 * time.Second))
 	n, err := conn.Read(buf)
 	if err != nil {
 		t.Fatalf("read discover reply: %v", err)

wss/server.go

@@ -419,6 +419,13 @@ func (s *Server) runAuth(ctx context.Context, conn *websocket.Conn) (uint32, err
 	if !ok {
 		return 0, fmt.Errorf("unknown node_id %d", reply.NodeID)
 	}
+	// ed25519.Verify panics on a key that is not exactly PublicKeySize
+	// bytes. PubKeyLookup is supplied by the embedding process (in
+	// production, the registry's pubkey index) so its output is not
+	// guaranteed here; reject rather than pass it through.
+	if len(pubKey) != ed25519.PublicKeySize {
+		return 0, fmt.Errorf("node_id %d has malformed public key (%d bytes)", reply.NodeID, len(pubKey))
+	}
 	sig, err := base64.StdEncoding.DecodeString(reply.Sig)
 	if err != nil {
 		return 0, fmt.Errorf("bad sig encoding: %w", err)