From 709fbdbc64d0ed492d52d1b02aeeae69d56cd78e Mon Sep 17 00:00:00 2001 From: Teodor Calin Date: Sun, 21 Jun 2026 01:03:18 +0300 Subject: [PATCH 1/5] beacon: optional destination-carrying relay delivery (0x08/0x09) --- server.go | 32 ++++++++++++++++++++++++++------ 1 file changed, 26 insertions(+), 6 deletions(-) diff --git a/server.go b/server.go index f4af19d..0e5abe7 100644 --- a/server.go +++ b/server.go @@ -41,6 +41,7 @@ type Server struct { conns []*net.UDPConn // SO_REUSEPORT sockets — one per reader goroutine sendBatchConns []*ipv4.PacketConn // ipv4 wrappers of conns[i] for WriteBatch (sendmmsg) — one per worker fd nodes *nodeMap // sharded node_id → observed endpoint + last-seen + dstNodes sync.Map // node_id → struct{}: endpoints that requested dest-carrying delivery (0x09) readyCh chan struct{} relayCh chan relayJob // buffered channel for relay workers pool sync.Pool // reusable payload buffers (inner relay payload) @@ -531,7 +532,9 @@ func (s *Server) handlePacket(data []byte, remote *net.UDPAddr) { switch msgType { case protocol.BeaconMsgDiscover: - s.handleDiscover(data[1:], remote) + s.handleDiscover(data[1:], remote, false) + case protocol.BeaconMsgDiscoverEx: + s.handleDiscover(data[1:], remote, true) case protocol.BeaconMsgPunchRequest: s.handlePunchRequest(data[1:], remote) case protocol.BeaconMsgRelay: @@ -543,7 +546,7 @@ func (s *Server) handlePacket(data []byte, remote *net.UDPAddr) { } } -func (s *Server) handleDiscover(data []byte, remote *net.UDPAddr) { +func (s *Server) handleDiscover(data []byte, remote *net.UDPAddr, wantDest bool) { if len(data) < 4 { return } @@ -557,6 +560,13 @@ func (s *Server) handleDiscover(data []byte, remote *net.UDPAddr) { nodeID := binary.BigEndian.Uint32(data[0:4]) + // This endpoint requested destination-carrying delivery; remember the + // node id so its final-hop relay frames include the destination (0x09), + // letting the receiver resolve which local endpoint the frame is for. + if wantDest { + s.dstNodes.Store(nodeID, struct{}{}) + } + // Per-nodeID endpoint update rate limit (PILOT-334) — prevents a single // nodeID from flapping its endpoint via rapid Discover messages. s.discoverRateMu.Lock() @@ -925,10 +935,20 @@ func (s *Server) relayWorker(sendConn *ipv4.PacketConn) { var addr *net.UDPAddr if localOk { addr = localAddr - outBuf = outBuf[:1+4+len(job.payload)] - outBuf[0] = protocol.BeaconMsgRelayDeliver - binary.BigEndian.PutUint32(outBuf[1:5], job.senderID) - copy(outBuf[5:], job.payload) + if _, wantDest := s.dstNodes.Load(job.destID); wantDest { + // Destination-carrying endpoint: include the dest node id + // so the receiver can resolve which endpoint the frame is for. + outBuf = outBuf[:1+4+4+len(job.payload)] + outBuf[0] = protocol.BeaconMsgRelayDeliverDest + binary.BigEndian.PutUint32(outBuf[1:5], job.senderID) + binary.BigEndian.PutUint32(outBuf[5:9], job.destID) + copy(outBuf[9:], job.payload) + } else { + outBuf = outBuf[:1+4+len(job.payload)] + outBuf[0] = protocol.BeaconMsgRelayDeliver + binary.BigEndian.PutUint32(outBuf[1:5], job.senderID) + copy(outBuf[5:], job.payload) + } } else { // Tier 2: peer beacon lookup. peerNodes is an // atomic.Pointer; reads are lock-free. From 9d4c6f1db54df94376ce669678823d963c65d2fb Mon Sep 17 00:00:00 2001 From: Teodor Calin Date: Sun, 21 Jun 2026 15:20:10 +0300 Subject: [PATCH 2/5] =?UTF-8?q?go:=20bump=20github.com/pilot-protocol/comm?= =?UTF-8?q?on=20v0.5.0=20=E2=86=92=20v0.5.4?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit v0.5.4 adds BeaconMsgDiscoverEx (0x08) and BeaconMsgRelayDeliverDest (0x09) which this PR references in server.go. --- go.mod | 2 +- go.sum | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/go.mod b/go.mod index 896d49b..eb545c9 100644 --- a/go.mod +++ b/go.mod @@ -4,7 +4,7 @@ go 1.25.10 require ( github.com/coder/websocket v1.8.15 - github.com/pilot-protocol/common v0.5.0 + github.com/pilot-protocol/common v0.5.4 golang.org/x/net v0.56.0 ) diff --git a/go.sum b/go.sum index 3a9a6ea..e01b756 100644 --- a/go.sum +++ b/go.sum @@ -1,7 +1,7 @@ github.com/coder/websocket v1.8.15 h1:6B2JPeOGlpff2Uz6vOEH1Vzpi0iUz20A+lPVhPHtNUA= github.com/coder/websocket v1.8.15/go.mod h1:NX3SzP+inril6yawo5CQXx8+fk145lPDC6pumgx0mVg= -github.com/pilot-protocol/common v0.5.0 h1:+rakrcpwbg8dfuQDWAgOI1N/XC3c6WSwar2SlW+cZ28= -github.com/pilot-protocol/common v0.5.0/go.mod h1:yrAwPXGVMbXU+SADvOCmbdXjK/wJ3uA0KshyLvRlej4= +github.com/pilot-protocol/common v0.5.4 h1:lz/cOXOjYCNIJvkWpb+pxXWpCqNKM/IY7D5fnVAq+8g= +github.com/pilot-protocol/common v0.5.4/go.mod h1:yrAwPXGVMbXU+SADvOCmbdXjK/wJ3uA0KshyLvRlej4= golang.org/x/net v0.56.0 h1:Rw8j/hFzGvJUZwNBXnAtf5sVDVt+65SK2C7IxCxZt5o= golang.org/x/net v0.56.0/go.mod h1:D3Ku6r+V6JROoZK144D2XfMHFcMq/0zSfLelVTCFKec= golang.org/x/sys v0.46.0 h1:noSf2Fq6F8DBgS+LysIkx7rIExoNHJsxOAtPp4rthXw= From 58981d640d5bc720e4f611230e63ba37db8304ee Mon Sep 17 00:00:00 2001 From: Teodor Calin Date: Sun, 21 Jun 2026 22:22:12 +0300 Subject: [PATCH 3/5] relay: per-packet WriteToUDP fallback for dual-stack send sockets MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ipv4.PacketConn.WriteBatch (sendmmsg) fails with EINVAL when the send socket is bound to the IPv6/dual-stack wildcard ([::], the default for a ':port' listen on Linux) and a relay destination carries an IPv4 address. Because the node fleet is IPv4, this meant every relay delivery to a node failed and was silently dropped — relay-routed nodes were unreachable. Detect a dual-stack send socket at worker start and route the flush through per-packet WriteToUDP (which maps IPv4 destinations onto the dual-stack socket correctly) instead of the always-failing WriteBatch. The sendmmsg fast path is unchanged on IPv4-bound sockets, and a per-packet fallback also covers any unexpected partial-batch error. --- server.go | 81 +++++++++++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 76 insertions(+), 5 deletions(-) diff --git a/server.go b/server.go index 0e5abe7..06544d4 100644 --- a/server.go +++ b/server.go @@ -40,6 +40,7 @@ type Server struct { conn *net.UDPConn // primary read+write socket; equal to conns[0] conns []*net.UDPConn // SO_REUSEPORT sockets — one per reader goroutine sendBatchConns []*ipv4.PacketConn // ipv4 wrappers of conns[i] for WriteBatch (sendmmsg) — one per worker fd + sendRawConns []*net.UDPConn // raw conns[i], parallel to sendBatchConns — used for the per-packet WriteToUDP fallback when WriteBatch fails (e.g. dual-stack [::] socket + IPv4 destination, which sendmmsg via ipv4.PacketConn rejects with "invalid argument") nodes *nodeMap // sharded node_id → observed endpoint + last-seen dstNodes sync.Map // node_id → struct{}: endpoints that requested dest-carrying delivery (0x09) readyCh chan struct{} @@ -371,11 +372,14 @@ func (s *Server) ListenAndServe(addr string) error { } if sharedPort { s.sendBatchConns = make([]*ipv4.PacketConn, len(s.conns)) + s.sendRawConns = make([]*net.UDPConn, len(s.conns)) for i, c := range s.conns { s.sendBatchConns[i] = ipv4.NewPacketConn(c) + s.sendRawConns[i] = c } } else { s.sendBatchConns = []*ipv4.PacketConn{ipv4.NewPacketConn(s.conn)} + s.sendRawConns = []*net.UDPConn{s.conn} } slog.Info("beacon listening", "addr", s.conn.LocalAddr(), "beacon_id", s.beaconID, "peers", len(s.peers), "readers", readers) @@ -400,7 +404,8 @@ func (s *Server) ListenAndServe(addr string) error { workers = 8 } for i := 0; i < workers; i++ { - go s.relayWorker(s.sendBatchConns[i%len(s.sendBatchConns)]) + idx := i % len(s.sendBatchConns) + go s.relayWorker(s.sendBatchConns[idx], s.sendRawConns[idx]) } // Start relay stats logger (every 60s) @@ -851,18 +856,67 @@ const relayFlushAfter = 2 * time.Millisecond // WriteBatch globally — the send-side bottleneck before fan-out. // 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) { +func (s *Server) relayWorker(sendConn *ipv4.PacketConn, rawConn *net.UDPConn) { msgs := make([]ipv4.Message, 0, relayBatchCap) payloadsToReturn := make([][]byte, 0, relayBatchCap) outBufsToReturn := make([]*[]byte, 0, relayBatchCap) + // sendmmsg via ipv4.PacketConn.WriteBatch rejects the whole batch with + // EINVAL ("invalid argument") when the send socket is bound to the + // IPv6/dual-stack wildcard ([::]) and a destination carries an IPv4 + // address — the common case, because production binds the beacon to + // ":9001" which resolves to [::]:9001 on Linux and the node fleet is + // IPv4. The result was that EVERY relay-deliver to an IPv4 node failed + // here, silently blackholing relay-only nodes. When the send socket is + // dual-stack we can't use the batch fast path for those, so route the + // flush through per-packet WriteToUDP on the raw conn (which maps IPv4 + // destinations onto the dual-stack socket correctly). On an IPv4-bound + // socket the sendmmsg fast path is unaffected. + sendDualStack := isWildcardIPv6(rawConn.LocalAddr()) + + writeToUDPFlush := func() (sent int) { + for i := range msgs { + dst, ok := msgs[i].Addr.(*net.UDPAddr) + if !ok { + continue + } + if _, werr := rawConn.WriteToUDP(msgs[i].Buffers[0], dst); werr == nil { + sent++ + } else { + slog.Debug("relay per-packet send failed", "dst", dst, "err", werr) + } + } + return sent + } + flush := func() { if len(msgs) == 0 { return } - n, err := sendConn.WriteBatch(msgs, 0) - if err != nil { - slog.Debug("relay write batch failed", "n", n, "of", len(msgs), "err", err) + var n int + if sendDualStack { + n = writeToUDPFlush() + } else { + var err error + n, err = sendConn.WriteBatch(msgs, 0) + if err != nil { + // Unexpected on an IPv4-bound socket, but fall back to + // per-packet for the messages the batch didn't send rather + // than dropping them. + slog.Debug("relay write batch failed, falling back to per-packet send", + "n", n, "of", len(msgs), "err", err) + for i := n; i < len(msgs); i++ { + dst, ok := msgs[i].Addr.(*net.UDPAddr) + if !ok { + continue + } + if _, werr := rawConn.WriteToUDP(msgs[i].Buffers[0], dst); werr == nil { + n++ + } else { + slog.Debug("relay per-packet fallback send failed", "dst", dst, "err", werr) + } + } + } } if n > 0 { s.relayForwarded.Add(uint64(n)) @@ -1001,6 +1055,23 @@ func (s *Server) relayWorker(sendConn *ipv4.PacketConn) { } } +// isWildcardIPv6 reports whether addr is an IPv6/dual-stack UDP socket (e.g. +// bound to [::] from a ":port" listen). On such a socket, sendmmsg via +// ipv4.PacketConn.WriteBatch fails with EINVAL for IPv4 destinations, so the +// relay worker must use per-packet WriteToUDP (which maps IPv4 destinations +// onto the dual-stack socket correctly). A socket bound to an IPv4 address +// (including 0.0.0.0) keeps the batch fast path. +func isWildcardIPv6(addr net.Addr) bool { + ua, ok := addr.(*net.UDPAddr) + if !ok { + return false + } + // To4()==nil means the local bind is genuinely IPv6 (not IPv4 / not an + // IPv4-mapped form). That's exactly the case where WriteBatch chokes on + // IPv4 destinations. + return ua.IP.To4() == nil +} + func (s *Server) returnPayload(buf []byte) { buf = buf[:cap(buf)] s.pool.Put(&buf) From c86af2cb2a2a3283dbc5f9fc675a04e90c8c61f3 Mon Sep 17 00:00:00 2001 From: Teodor Calin Date: Thu, 23 Jul 2026 23:33:18 +0300 Subject: [PATCH 4/5] beacon: WS3 punch-token verification gate (default-off) Adds RequirePunchToken gate on handlePunchRequest so an endpoint is only released against a valid target-issued punch token. Gate is an atomic.Bool defaulting to false (no behavior change; not yet wired to a CLI flag), so this is wire-compatible and does NOT enforce until explicitly enabled. Adds zz_punch_token_test.go. Co-Authored-By: Claude Opus 4.8 --- server.go | 52 +++++++ zz_punch_token_test.go | 299 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 351 insertions(+) create mode 100644 zz_punch_token_test.go diff --git a/server.go b/server.go index 06544d4..6b5754e 100644 --- a/server.go +++ b/server.go @@ -3,6 +3,7 @@ package beacon import ( + "crypto/ed25519" "encoding/binary" "encoding/json" "fmt" @@ -71,6 +72,9 @@ type Server struct { // (e.g. specialist boxes that serve high-volume query traffic). punchWhitelistAll atomic.Bool + requirePunchToken atomic.Bool + nodePubKeys sync.Map + // breakerAllow consults the parent process's breaker manager for // the named beacon switch (beacon.punch, beacon.relay, // beacon.discover). Returns (allow, reason). Open = drop the @@ -572,6 +576,12 @@ func (s *Server) handleDiscover(data []byte, remote *net.UDPAddr, wantDest bool) s.dstNodes.Store(nodeID, struct{}{}) } + if len(data) >= 4+ed25519.PublicKeySize { + pub := make(ed25519.PublicKey, ed25519.PublicKeySize) + copy(pub, data[4:4+ed25519.PublicKeySize]) + s.nodePubKeys.Store(nodeID, pub) + } + // Per-nodeID endpoint update rate limit (PILOT-334) — prevents a single // nodeID from flapping its endpoint via rapid Discover messages. s.discoverRateMu.Lock() @@ -612,6 +622,30 @@ func (s *Server) handleDiscover(data []byte, remote *net.UDPAddr, wantDest bool) } } +const punchGrantExpirySize = 8 +const punchGrantTrailerSize = punchGrantExpirySize + ed25519.SignatureSize + +func (s *Server) verifyPunchGrant(trailer []byte, requesterID, targetID uint32) bool { + if len(trailer) < punchGrantTrailerSize { + return false + } + expiry := int64(binary.BigEndian.Uint64(trailer[0:punchGrantExpirySize])) + if time.Now().Unix() > expiry { + return false + } + sig := trailer[punchGrantExpirySize : punchGrantExpirySize+ed25519.SignatureSize] + v, ok := s.nodePubKeys.Load(targetID) + if !ok { + return false + } + pub, ok := v.(ed25519.PublicKey) + 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) handlePunchRequest(data []byte, remote *net.UDPAddr) { if len(data) < 8 { return @@ -663,6 +697,12 @@ rateLimitBypass: requesterID := binary.BigEndian.Uint32(data[0:4]) targetID := binary.BigEndian.Uint32(data[4:8]) + if s.requirePunchToken.Load() { + if !s.verifyPunchGrant(data[8:], requesterID, targetID) { + return + } + } + // Update requester's endpoint (handles symmetric NAT port changes). s.nodes.Upsert(requesterID, remote, time.Now(), maxBeaconNodes) @@ -1177,6 +1217,14 @@ func (s *Server) reapStaleNodes() { } } s.discoverRateMu.Unlock() + + s.nodePubKeys.Range(func(k, _ interface{}) bool { + id, ok := k.(uint32) + if ok && !s.nodes.Has(id) { + s.nodePubKeys.Delete(k) + } + return true + }) } // --- Gossip --- @@ -1358,6 +1406,10 @@ func (s *Server) SetPunchWhitelist(ips []string) { s.punchWhitelistAll.Store(all) } +func (s *Server) SetRequirePunchToken(v bool) { + s.requirePunchToken.Store(v) +} + // registryDiscoveryLoop registers this beacon with the registry and discovers // peers every 30 seconds. Requires the beacon to be listening (conn bound). func (s *Server) registryDiscoveryLoop() { diff --git a/zz_punch_token_test.go b/zz_punch_token_test.go new file mode 100644 index 0000000..9836748 --- /dev/null +++ b/zz_punch_token_test.go @@ -0,0 +1,299 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +package beacon + +import ( + "crypto/ed25519" + "crypto/rand" + "encoding/binary" + "errors" + "fmt" + "net" + "testing" + "time" + + "github.com/pilot-protocol/common/protocol" +) + +func registerNodeWithPubKey(t *testing.T, beaconAddr *net.UDPAddr, nodeID uint32, pub ed25519.PublicKey) *net.UDPConn { + t.Helper() + conn, err := net.DialUDP("udp", nil, beaconAddr) + if err != nil { + t.Fatalf("dial beacon: %v", err) + } + + msg := make([]byte, 1+4+len(pub)) + msg[0] = protocol.BeaconMsgDiscover + binary.BigEndian.PutUint32(msg[1:5], nodeID) + copy(msg[5:], pub) + if _, err := conn.Write(msg); err != nil { + t.Fatalf("send discover: %v", err) + } + + buf := make([]byte, 64) + conn.SetReadDeadline(time.Now().Add(2 * time.Second)) + n, err := conn.Read(buf) + if err != nil { + t.Fatalf("read discover reply: %v", err) + } + if n < 1 || buf[0] != protocol.BeaconMsgDiscoverReply { + t.Fatalf("unexpected reply type: 0x%02x", buf[0]) + } + + return conn +} + +func mintPunchGrant(t *testing.T, priv ed25519.PrivateKey, requesterID, targetID uint32, expiry int64) []byte { + t.Helper() + trailer := make([]byte, punchGrantTrailerSize) + binary.BigEndian.PutUint64(trailer[0:punchGrantExpirySize], uint64(expiry)) + preimage := fmt.Sprintf("punch-grant:%d:%d:%d", requesterID, targetID, expiry) + sig := ed25519.Sign(priv, []byte(preimage)) + copy(trailer[punchGrantExpirySize:], sig) + return trailer +} + +func sendPunchRequest(t *testing.T, conn *net.UDPConn, requesterID, targetID uint32, trailer []byte) { + t.Helper() + msg := make([]byte, 1+8+len(trailer)) + msg[0] = protocol.BeaconMsgPunchRequest + binary.BigEndian.PutUint32(msg[1:5], requesterID) + binary.BigEndian.PutUint32(msg[5:9], targetID) + copy(msg[9:], trailer) + if _, err := conn.Write(msg); err != nil { + t.Fatalf("send punch request: %v", err) + } +} + +func expectPunchCommand(t *testing.T, conn *net.UDPConn, wantAddr *net.UDPAddr) { + t.Helper() + buf := make([]byte, 64) + conn.SetReadDeadline(time.Now().Add(3 * time.Second)) + n, err := conn.Read(buf) + if err != nil { + t.Fatalf("read punch command: %v", err) + } + if n < 2 || buf[0] != protocol.BeaconMsgPunchCommand { + t.Fatalf("unexpected reply type: 0x%02x", buf[0]) + } + ipLen := int(buf[1]) + if n < 2+ipLen+2 { + t.Fatalf("short punch command: %d bytes", n) + } + gotIP := net.IP(buf[2 : 2+ipLen]) + gotPort := binary.BigEndian.Uint16(buf[2+ipLen : 2+ipLen+2]) + if !gotIP.Equal(wantAddr.IP) { + t.Fatalf("punch command ip: got %v, want %v", gotIP, wantAddr.IP) + } + if int(gotPort) != wantAddr.Port { + t.Fatalf("punch command port: got %d, want %d", gotPort, wantAddr.Port) + } +} + +func expectNoPacket(t *testing.T, conn *net.UDPConn, timeout time.Duration) { + t.Helper() + buf := make([]byte, 64) + conn.SetReadDeadline(time.Now().Add(timeout)) + n, err := conn.Read(buf) + if err == nil { + t.Fatalf("expected no packet, got %d bytes (type 0x%02x)", n, buf[0]) + } + var netErr net.Error + if !errors.As(err, &netErr) || !netErr.Timeout() { + t.Fatalf("expected read timeout, got %v", err) + } +} + +func TestPunchToken_FlagOff_BehavesAsToday(t *testing.T) { + t.Parallel() + s := New() + go s.ListenAndServe("127.0.0.1:0") + <-s.Ready() + defer s.Close() + + beaconAddr := beaconUDPAddr(t, s) + + requesterID := uint32(9001) + targetID := uint32(9002) + + conn1 := registerNode(t, beaconAddr, requesterID) + defer conn1.Close() + conn2 := registerNode(t, beaconAddr, targetID) + defer conn2.Close() + + requesterAddr := conn1.LocalAddr().(*net.UDPAddr) + targetAddr := conn2.LocalAddr().(*net.UDPAddr) + + sendPunchRequest(t, conn1, requesterID, targetID, nil) + + expectPunchCommand(t, conn1, targetAddr) + expectPunchCommand(t, conn2, requesterAddr) +} + +func TestPunchToken_FlagOn_NoToken_Dropped(t *testing.T) { + t.Parallel() + s := New() + s.SetRequirePunchToken(true) + go s.ListenAndServe("127.0.0.1:0") + <-s.Ready() + defer s.Close() + + beaconAddr := beaconUDPAddr(t, s) + + requesterID := uint32(9101) + targetID := uint32(9102) + + _, priv, err := ed25519.GenerateKey(rand.Reader) + if err != nil { + t.Fatalf("generate key: %v", err) + } + pub := priv.Public().(ed25519.PublicKey) + + conn1 := registerNode(t, beaconAddr, requesterID) + defer conn1.Close() + conn2 := registerNodeWithPubKey(t, beaconAddr, targetID, pub) + defer conn2.Close() + + sendPunchRequest(t, conn1, requesterID, targetID, nil) + + expectNoPacket(t, conn1, 500*time.Millisecond) + expectNoPacket(t, conn2, 500*time.Millisecond) +} + +func TestPunchToken_FlagOn_ValidGrant_Succeeds(t *testing.T) { + t.Parallel() + s := New() + s.SetRequirePunchToken(true) + go s.ListenAndServe("127.0.0.1:0") + <-s.Ready() + defer s.Close() + + beaconAddr := beaconUDPAddr(t, s) + + requesterID := uint32(9201) + targetID := uint32(9202) + + _, priv, err := ed25519.GenerateKey(rand.Reader) + if err != nil { + t.Fatalf("generate key: %v", err) + } + pub := priv.Public().(ed25519.PublicKey) + + conn1 := registerNode(t, beaconAddr, requesterID) + defer conn1.Close() + conn2 := registerNodeWithPubKey(t, beaconAddr, targetID, pub) + defer conn2.Close() + + requesterAddr := conn1.LocalAddr().(*net.UDPAddr) + targetAddr := conn2.LocalAddr().(*net.UDPAddr) + + expiry := time.Now().Add(60 * time.Second).Unix() + trailer := mintPunchGrant(t, priv, requesterID, targetID, expiry) + + sendPunchRequest(t, conn1, requesterID, targetID, trailer) + + expectPunchCommand(t, conn1, targetAddr) + expectPunchCommand(t, conn2, requesterAddr) +} + +func TestPunchToken_FlagOn_ExpiredGrant_Rejected(t *testing.T) { + t.Parallel() + s := New() + s.SetRequirePunchToken(true) + go s.ListenAndServe("127.0.0.1:0") + <-s.Ready() + defer s.Close() + + beaconAddr := beaconUDPAddr(t, s) + + requesterID := uint32(9301) + targetID := uint32(9302) + + _, priv, err := ed25519.GenerateKey(rand.Reader) + if err != nil { + t.Fatalf("generate key: %v", err) + } + pub := priv.Public().(ed25519.PublicKey) + + conn1 := registerNode(t, beaconAddr, requesterID) + defer conn1.Close() + conn2 := registerNodeWithPubKey(t, beaconAddr, targetID, pub) + defer conn2.Close() + + expiry := time.Now().Add(-60 * time.Second).Unix() + trailer := mintPunchGrant(t, priv, requesterID, targetID, expiry) + + sendPunchRequest(t, conn1, requesterID, targetID, trailer) + + expectNoPacket(t, conn1, 500*time.Millisecond) + expectNoPacket(t, conn2, 500*time.Millisecond) +} + +func TestPunchToken_FlagOn_WrongPairGrant_Rejected(t *testing.T) { + t.Parallel() + s := New() + s.SetRequirePunchToken(true) + go s.ListenAndServe("127.0.0.1:0") + <-s.Ready() + defer s.Close() + + beaconAddr := beaconUDPAddr(t, s) + + legitRequesterID := uint32(9401) + attackerID := uint32(9402) + targetID := uint32(9403) + + _, priv, err := ed25519.GenerateKey(rand.Reader) + if err != nil { + t.Fatalf("generate key: %v", err) + } + pub := priv.Public().(ed25519.PublicKey) + + legitConn := registerNode(t, beaconAddr, legitRequesterID) + defer legitConn.Close() + attackerConn := registerNode(t, beaconAddr, attackerID) + defer attackerConn.Close() + targetConn := registerNodeWithPubKey(t, beaconAddr, targetID, pub) + defer targetConn.Close() + + expiry := time.Now().Add(60 * time.Second).Unix() + grantForLegitRequester := mintPunchGrant(t, priv, legitRequesterID, targetID, expiry) + + sendPunchRequest(t, attackerConn, attackerID, targetID, grantForLegitRequester) + + expectNoPacket(t, attackerConn, 500*time.Millisecond) + expectNoPacket(t, targetConn, 500*time.Millisecond) +} + +func TestPunchToken_NotFoundVsUnauthorized_Indistinguishable(t *testing.T) { + t.Parallel() + s := New() + s.SetRequirePunchToken(true) + go s.ListenAndServe("127.0.0.1:0") + <-s.Ready() + defer s.Close() + + beaconAddr := beaconUDPAddr(t, s) + + requesterID := uint32(9501) + unregisteredTargetID := uint32(9502) + registeredTargetID := uint32(9503) + + _, priv, err := ed25519.GenerateKey(rand.Reader) + if err != nil { + t.Fatalf("generate key: %v", err) + } + pub := priv.Public().(ed25519.PublicKey) + + requesterConnA := registerNode(t, beaconAddr, requesterID) + defer requesterConnA.Close() + targetConn := registerNodeWithPubKey(t, beaconAddr, registeredTargetID, pub) + defer targetConn.Close() + + sendPunchRequest(t, requesterConnA, requesterID, unregisteredTargetID, nil) + expectNoPacket(t, requesterConnA, 500*time.Millisecond) + + sendPunchRequest(t, requesterConnA, requesterID, registeredTargetID, nil) + expectNoPacket(t, requesterConnA, 500*time.Millisecond) + expectNoPacket(t, targetConn, 500*time.Millisecond) +} From decbd9413c02ed8666bdf1769c0be9dd89bade09 Mon Sep 17 00:00:00 2001 From: Teodor Calin Date: Thu, 23 Jul 2026 23:45:09 +0300 Subject: [PATCH 5/5] beacon: wire --require-punch-token / BEACON_REQUIRE_PUNCH_TOKEN (default off) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Makes the WS3 RequirePunchToken gate flippable in prod. Default false — no behavior change, wire-compatible with old agents. Enforcement stays deferred. Co-Authored-By: Claude Opus 4.8 --- cmd/beacon/main.go | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/cmd/beacon/main.go b/cmd/beacon/main.go index 005632a..f592a79 100644 --- a/cmd/beacon/main.go +++ b/cmd/beacon/main.go @@ -28,6 +28,7 @@ func main() { logLevel := flag.String("log-level", "info", "log level (debug, info, warn, error)") logFormat := flag.String("log-format", "text", "log format (text, json)") punchWhitelist := flag.String("punch-whitelist", "", "PILOT-342: comma-separated source IPs that bypass SEC-026 punch-request rate limits (use for trusted operator boxes, test rigs, paired beacons). Env: BEACON_PUNCH_WHITELIST.") + requirePunchToken := flag.Bool("require-punch-token", false, "WS3: require a valid target-issued punch token before releasing a peer endpoint. Default false (not enforcing, wire-compatible with old agents). Env: BEACON_REQUIRE_PUNCH_TOKEN=1.") flag.Parse() if *configPath != "" { @@ -71,6 +72,10 @@ func main() { s.SetPunchWhitelist(strings.Split(wlSrc, ",")) slog.Info("punch-request whitelist configured", "ips", strings.Count(wlSrc, ",")+1) } + if *requirePunchToken || os.Getenv("BEACON_REQUIRE_PUNCH_TOKEN") == "1" { + s.SetRequirePunchToken(true) + slog.Info("punch-token enforcement enabled (WS3)") + } if *healthAddr != "" { go func() {