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() { diff --git a/server.go b/server.go index 0e5abe7..6b5754e 100644 --- a/server.go +++ b/server.go @@ -3,6 +3,7 @@ package beacon import ( + "crypto/ed25519" "encoding/binary" "encoding/json" "fmt" @@ -40,6 +41,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{} @@ -70,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 @@ -371,11 +376,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 +408,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) @@ -567,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() @@ -607,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 @@ -658,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) @@ -851,18 +896,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 +1095,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) @@ -1106,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 --- @@ -1287,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) +}