From 422bf7fc4b5ed9dfc215d717514f3c6140851ef5 Mon Sep 17 00:00:00 2001 From: Teodor Calin Date: Sun, 26 Jul 2026 15:09:27 +0300 Subject: [PATCH] beacon: relay budget keyed on transport source, WSS peer-map ownership, bounded WSS writes, gossip peer-set gate MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit L6 — dispatchRelay charged the per-source relay budget against the senderID field of the frame body, which the sender writes. One origin could therefore open an unbounded number of independent budgets by varying that field per packet. The budget is now keyed on the observed datagram endpoint, or, for frames handed over by the compat WSS bridge (which have no datagram source of their own), on the peer id the bridge verified against the node's registered key at connect time. The key is a comparable value type built without allocating, so the relay hot path keeps its per-packet cost. Keying on the full source endpoint rather than the address alone keeps co-located fleets — hosts running hundreds of daemons behind one address — on per-daemon budgets instead of collapsing them onto one. H4 — wss dropPeer removed whatever connection was registered for a node id. A node that reconnects installs a fresh connection under the same key while the previous connection's read loop is still unwinding, and that loop's cleanup then deregistered the live replacement. dropPeer now takes the *wssPeer and compare-and-deletes. M11 — wss WriteFrame bounded the write with IdleTimeout (90s). It is called from the beacon's relay workers, so a peer that stopped draining its socket could hold a worker for that whole window. Added Config.WriteTimeout (default 2s) for the per-frame deadline; the idle window keeps governing connection lifetime. A peer whose write expires is dropped, so later calls return immediately instead of re-entering the same stall. M10 — handleSync rewrites the node-id → beacon routing table the relay workers read, and accepted a sync from any source. Added an opt-in peer-set gate (SetStrictGossip / -strict-gossip / BEACON_STRICT_GOSSIP=1) restricting syncs to addresses in the peer set learned from -peers and the registry. Default off, because a beacon behind DNAT egresses from an address other than the one it advertises and enabling the check there would stop the mesh converging. Route expiry (gossipPeerTTL, swept from reapStaleNodes) is unconditional: routes published by a beacon that stops gossiping used to stand forever. Co-Authored-By: Claude Opus 5 --- cmd/beacon/main.go | 5 + ratelimit_shard.go | 74 ++++++++++-- server.go | 145 ++++++++++++++++++++++- wss/server.go | 55 +++++++-- wss/zz_peer_lifecycle_test.go | 217 ++++++++++++++++++++++++++++++++++ zz_fuzz_dispatch_test.go | 2 +- zz_gossip_peer_set_test.go | 117 ++++++++++++++++++ zz_relay_source_limit_test.go | 149 +++++++++++++++++++++++ zz_server_branches_test.go | 6 +- 9 files changed, 739 insertions(+), 31 deletions(-) create mode 100644 wss/zz_peer_lifecycle_test.go create mode 100644 zz_gossip_peer_set_test.go create mode 100644 zz_relay_source_limit_test.go diff --git a/cmd/beacon/main.go b/cmd/beacon/main.go index f592a79..5537909 100644 --- a/cmd/beacon/main.go +++ b/cmd/beacon/main.go @@ -29,6 +29,7 @@ func main() { 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.") + strictGossip := flag.Bool("strict-gossip", false, "accept gossip sync messages only from addresses in the peer set (-peers plus registry-discovered beacons). Default false (any source may publish node routes into the peer mesh), because a beacon behind DNAT egresses from an address other than the one it advertises. Env: BEACON_STRICT_GOSSIP=1.") flag.Parse() if *configPath != "" { @@ -76,6 +77,10 @@ func main() { s.SetRequirePunchToken(true) slog.Info("punch-token enforcement enabled (WS3)") } + if *strictGossip || os.Getenv("BEACON_STRICT_GOSSIP") == "1" { + s.SetStrictGossip(true) + slog.Info("gossip peer-set enforcement enabled") + } if *healthAddr != "" { go func() { diff --git a/ratelimit_shard.go b/ratelimit_shard.go index 321fdf9..d6e8cab 100644 --- a/ratelimit_shard.go +++ b/ratelimit_shard.go @@ -3,6 +3,7 @@ package beacon import ( + "net" "sync" "time" ) @@ -110,31 +111,88 @@ func (rl *punchRateLimiter) sweep(cutoff time.Time) { } } -// relayRateLimiter is a sharded senderID -> sliding-window map. +// relaySourceKey identifies the transport-level origin of a relay +// frame. For a UDP datagram that is the observed source endpoint; for +// a frame handed over by the compat WSS bridge it is the peer id the +// bridge authenticated at connect time. Both are established by the +// transport, not by fields inside the frame the sender wrote. +// +// It is a comparable value type so it can be a map key with no +// per-packet allocation on the relay hot path. +type relaySourceKey struct { + ip [16]byte // observed source address in 16-byte form; zero when bridged + port uint16 // observed source port; zero when bridged + node uint32 // bridge-authenticated peer id; zero for datagrams + bridged bool +} + +// relaySourceForUDP builds the key for a datagram source. It writes the +// IPv4-mapped form by hand rather than calling net.IP.To16, which +// allocates for a 4-byte address. +func relaySourceForUDP(remote *net.UDPAddr) relaySourceKey { + var k relaySourceKey + if remote == nil { + return k + } + switch len(remote.IP) { + case net.IPv4len: + k.ip[10], k.ip[11] = 0xff, 0xff + copy(k.ip[12:], remote.IP) + case net.IPv6len: + copy(k.ip[:], remote.IP) + } + k.port = uint16(remote.Port) + return k +} + +// relaySourceForBridge builds the key for a frame delivered by the +// compat WSS bridge, which has no datagram source of its own. +func relaySourceForBridge(peerID uint32) relaySourceKey { + return relaySourceKey{node: peerID, bridged: true} +} + +func shardRelaySource(k relaySourceKey) uint32 { + var h uint32 = 2166136261 + for i := 0; i < len(k.ip); i++ { + h ^= uint32(k.ip[i]) + h *= 16777619 + } + h ^= uint32(k.port) + h *= 16777619 + h ^= k.node + h *= 16777619 + if k.bridged { + h ^= 1 + h *= 16777619 + } + return h & (rateLimitShards - 1) +} + +// relayRateLimiter is a sharded relay-source -> sliding-window map. type relayRateLimiter struct { shards [rateLimitShards]struct { mu sync.Mutex - m map[uint32]*relaySourceWindow + m map[relaySourceKey]*relaySourceWindow } } func newRelayRateLimiter() *relayRateLimiter { rl := &relayRateLimiter{} for i := range rl.shards { - rl.shards[i].m = make(map[uint32]*relaySourceWindow) + rl.shards[i].m = make(map[relaySourceKey]*relaySourceWindow) } return rl } -// allow reports whether a relay from senderID is within its per-second +// allow reports whether a relay from src is within its per-second // budget, advancing the window as needed. -func (rl *relayRateLimiter) allow(senderID uint32, nowNano int64, maxPerSecond uint32) bool { - sh := &rl.shards[shardU32(senderID)] +func (rl *relayRateLimiter) allow(src relaySourceKey, nowNano int64, maxPerSecond uint32) bool { + sh := &rl.shards[shardRelaySource(src)] sh.mu.Lock() defer sh.mu.Unlock() - w, ok := sh.m[senderID] + w, ok := sh.m[src] if !ok || nowNano-w.windowStart >= int64(time.Second) { - sh.m[senderID] = &relaySourceWindow{windowStart: nowNano, count: 1} + sh.m[src] = &relaySourceWindow{windowStart: nowNano, count: 1} return true } if w.count >= maxPerSecond { diff --git a/server.go b/server.go index 484f6f6..d041812 100644 --- a/server.go +++ b/server.go @@ -98,6 +98,19 @@ type Server struct { peerMu sync.RWMutex // protects s.peers only (peerNodes is atomic) healthOk atomic.Bool + // strictGossip gates the peer-set check in handleSync. Off by + // default so a mesh whose members reach each other from an address + // other than the one they advertise keeps converging; operators + // turn it on once every beacon's observed source address matches + // its registered address. See SetStrictGossip. + strictGossip atomic.Bool + + // gossipSeen records, per gossip source address, when that source + // last delivered a sync. reapStaleNodes uses it to drop peerNodes + // entries owned by a beacon that has stopped gossiping. Guarded by + // peerWriteMu, alongside the peerNodes copy-on-write. + gossipSeen map[string]time.Time + registryAddr string // registry address for dynamic peer discovery advertiseAddr string // address to register (overrides auto-detect from TCP local addr) registryAdminToken string // admin token sent with beacon_register (required by SEC-002) @@ -186,6 +199,7 @@ func NewWithPeers(beaconID uint32, peers []string) *Server { punchRL: newPunchRateLimiter(), relayRL: newRelayRateLimiter(), discoverRL: newDiscoverRateLimiter(), + gossipSeen: make(map[string]time.Time), } emptyPeers := make(map[uint32]*net.UDPAddr) s.peerNodes.Store(&emptyPeers) @@ -269,7 +283,13 @@ func (s *Server) EnableCompatWSS(bindAddr string, pubKeyLookup bwss.PubKeyLookup IP: net.ParseIP("192.0.2.1"), Port: int(senderID & 0xFFFF), } - s.handlePacket(frame, synth) + // senderID here is the id the WSS bridge verified against + // the node's registered Ed25519 key during the connect + // challenge, so it is a transport-established identity and + // is what per-source budgets are charged against on this + // path (the synthetic address is shared by every bridged + // peer and carries no identity). + s.handlePacketFrom(frame, synth, senderID) }, }) if err != nil { @@ -567,7 +587,16 @@ func (s *Server) RelayDropped() uint64 { return s.relayDropped.Load() } // node was not registered with the beacon. func (s *Server) RelayNotFound() uint64 { return s.relayNotFound.Load() } +// handlePacket dispatches an inbound UDP datagram. func (s *Server) handlePacket(data []byte, remote *net.UDPAddr) { + s.handlePacketFrom(data, remote, 0) +} + +// handlePacketFrom dispatches an inbound frame. bridgePeer is the peer +// id the compat WSS bridge authenticated at connect time, or 0 when the +// frame arrived as a UDP datagram; it selects which transport-level +// identifier the relay budget is charged against. +func (s *Server) handlePacketFrom(data []byte, remote *net.UDPAddr, bridgePeer uint32) { // 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. @@ -590,7 +619,11 @@ func (s *Server) handlePacket(data []byte, remote *net.UDPAddr) { case protocol.BeaconMsgPunchRequest: s.handlePunchRequest(data[1:], remote) case protocol.BeaconMsgRelay: - s.dispatchRelay(data[1:]) + if bridgePeer != 0 { + s.dispatchRelay(data[1:], relaySourceForBridge(bridgePeer)) + } else { + s.dispatchRelay(data[1:], relaySourceForUDP(remote)) + } case protocol.BeaconMsgSync: s.handleSync(data[1:], remote) default: @@ -798,7 +831,11 @@ rateLimitBypass: // dispatchRelay parses the relay header and dispatches to a worker goroutine. // The read loop stays fast — no locks (other than a sharded RLock for the // pre-check), no syscalls, no allocations on the hot path. -func (s *Server) dispatchRelay(data []byte) { +// +// src is the transport-level origin of the frame, supplied by the caller. +// The sender id carried inside the frame is routing metadata echoed back to +// the destination and is not used for accounting. +func (s *Server) dispatchRelay(data []byte, src relaySourceKey) { if len(data) < 8 { return } @@ -861,15 +898,21 @@ func (s *Server) dispatchRelay(data []byte) { } // Per-source rate limit: prevent one sender from flooding the relay - // queue and squeezing out legitimate traffic. A DoS source targeting - // a known destination can saturate the 524288-deep relayCh at rates + // queue and squeezing out legitimate traffic. A source targeting a + // known destination can saturate the 524288-deep relayCh at rates // far above normal — this cap gives each source a fixed share. + // + // The budget is keyed on src (the observed datagram endpoint, or the + // bridge-authenticated peer id for WSS-delivered frames) rather than + // on senderID: senderID is a field of the frame body, so keying on it + // would let one origin open an unbounded number of independent + // budgets by varying that field between packets. now := time.Now().UnixNano() // Source exceeded per-second budget — silently drop. The sender's // daemon retries (3-attempt path in pkg/daemon/daemon.go relay // branch), so a drop here is eventually self-healing for honest // senders. - if !s.relayRL.allow(senderID, now, maxRelaysPerSourcePerSecond) { + if !s.relayRL.allow(src, now, maxRelaysPerSourcePerSecond) { return } @@ -1258,6 +1301,9 @@ func (s *Server) reapStaleNodes() { s.relayRL.sweep(time.Now().Add(-relaySourceCleanupInterval).UnixNano()) s.discoverRL.sweep(time.Now().Add(-discoverMinInterval * 2)) + // Withdraw routes published by beacons that stopped gossiping. + s.sweepStaleGossip(time.Now().Add(-gossipPeerTTL)) + s.nodePubKeys.Range(func(k, _ interface{}) bool { id, ok := k.(uint32) if ok && !s.nodes.Has(id) { @@ -1316,6 +1362,81 @@ func (s *Server) sendGossip() { slog.Debug("gossip sent", "beacon_id", s.beaconID, "nodes", len(nodeIDs), "peers", len(peers)) } +// SetStrictGossip controls whether inbound gossip sync messages must +// come from an address in the peer set — the beacons configured at +// construction plus those learned from the registry in +// registryDiscover. With it off (the default) any source may publish +// node-id → beacon routes into the peer mesh, so it should be enabled +// wherever every mesh member's observed source address matches the +// address it registers. Off by default because a beacon behind DNAT +// egresses from an address other than the one it advertises, and +// enabling the check there would stop the mesh from converging. +func (s *Server) SetStrictGossip(v bool) { + s.strictGossip.Store(v) +} + +// StrictGossip reports whether the handleSync peer-set check is active. +func (s *Server) StrictGossip() bool { + return s.strictGossip.Load() +} + +// isKnownGossipPeer reports whether remote's address belongs to a beacon +// in the peer set. Only the address is compared, not the port: peers are +// recorded by their advertised listen endpoint, and gossip egresses from +// that same socket, but a middlebox in the path may rewrite the port. +func (s *Server) isKnownGossipPeer(remote *net.UDPAddr) bool { + if remote == nil { + return false + } + s.peerMu.RLock() + defer s.peerMu.RUnlock() + for _, p := range s.peers { + if p != nil && p.IP.Equal(remote.IP) { + return true + } + } + return false +} + +// gossipPeerTTL is how long a peerNodes entry survives without its +// owning beacon delivering another sync. gossipLoop ticks every 10s, so +// this tolerates several consecutive losses before the routes a silent +// beacon published are withdrawn. +const gossipPeerTTL = 90 * time.Second + +// sweepStaleGossip drops peerNodes entries owned by a gossip source that +// has not delivered a sync since cutoff, so routes to a beacon that went +// away stop being handed to the relay workers. +func (s *Server) sweepStaleGossip(cutoff time.Time) { + s.peerWriteMu.Lock() + defer s.peerWriteMu.Unlock() + + stale := make(map[string]struct{}) + for key, last := range s.gossipSeen { + if last.Before(cutoff) { + stale[key] = struct{}{} + delete(s.gossipSeen, key) + } + } + if len(stale) == 0 { + return + } + + cur := *s.peerNodes.Load() + next := make(map[uint32]*net.UDPAddr, len(cur)) + for id, addr := range cur { + if _, drop := stale[addr.String()]; drop { + continue + } + next[id] = addr + } + if len(next) == len(cur) { + return + } + s.peerNodes.Store(&next) + slog.Debug("gossip peers expired", "sources", len(stale), "routes_dropped", len(cur)-len(next)) +} + // handleSync processes an incoming gossip sync message from a peer beacon. func (s *Server) handleSync(data []byte, remote *net.UDPAddr) { // Need at least beaconID(4) + nodeCount(2) @@ -1323,6 +1444,14 @@ func (s *Server) handleSync(data []byte, remote *net.UDPAddr) { return } + // Peer-set gate (opt-in, see SetStrictGossip). A sync rewrites the + // node-id → beacon routing table the relay workers read, so when the + // gate is on only members of the known peer set may publish into it. + if s.strictGossip.Load() && !s.isKnownGossipPeer(remote) { + slog.Debug("gossip sync from address outside the peer set", "from", remote) + return + } + peerBeaconID := binary.BigEndian.Uint32(data[0:4]) nodeCount := binary.BigEndian.Uint16(data[4:6]) @@ -1356,6 +1485,10 @@ func (s *Server) handleSync(data []byte, remote *net.UDPAddr) { } } s.peerNodes.Store(&next) + if s.gossipSeen == nil { + s.gossipSeen = make(map[string]time.Time) + } + s.gossipSeen[remote.String()] = time.Now() s.peerWriteMu.Unlock() slog.Debug("gossip sync received", "peer_beacon_id", peerBeaconID, "nodes", nodeCount, "from", remote) diff --git a/wss/server.go b/wss/server.go index 6464f75..cafbfd1 100644 --- a/wss/server.go +++ b/wss/server.go @@ -61,6 +61,14 @@ const DefaultAuthTimeout = 10 * time.Second // the outer ceiling. const DefaultIdleTimeout = 90 * time.Second +// DefaultWriteTimeout bounds a single outbound frame write. It is +// deliberately much shorter than DefaultIdleTimeout: WriteFrame is +// called from the beacon's relay workers, so the deadline is the +// longest a worker can sit inside one write before the peer is +// dropped and the worker is handed back. Idle-connection lifetime is +// a separate concern and keeps using IdleTimeout. +const DefaultWriteTimeout = 2 * time.Second + // Config configures a Server. type Config struct { // BindAddr is the host:port to listen on. Typically @@ -84,6 +92,11 @@ type Config struct { // IdleTimeout overrides DefaultIdleTimeout. IdleTimeout time.Duration + // WriteTimeout overrides DefaultWriteTimeout. It caps how long a + // single WriteFrame call may block on a peer whose socket is not + // draining; on expiry the write fails and the peer is dropped. + WriteTimeout time.Duration + // MaxPeers caps the number of concurrent WSS peer connections. // New upgrades beyond this are rejected with 503. 0 = unlimited. MaxPeers int @@ -177,6 +190,9 @@ func New(cfg Config) (*Server, error) { if cfg.IdleTimeout == 0 { cfg.IdleTimeout = DefaultIdleTimeout } + if cfg.WriteTimeout == 0 { + cfg.WriteTimeout = DefaultWriteTimeout + } return &Server{ cfg: cfg, peers: make(map[uint32]*wssPeer), @@ -253,6 +269,11 @@ func (s *Server) Close() error { // given node ID. Returns false if destID is not currently connected // — the caller (relay router) should fall back to UDP. Returns false // + logs on write error (peer is then dropped). +// +// The write is bounded by Config.WriteTimeout so a peer that stops +// draining its socket cannot pin the calling goroutine for the whole +// idle window. On expiry the peer is dropped, which makes subsequent +// calls for the same node ID return immediately. func (s *Server) WriteFrame(destID uint32, frame []byte) bool { s.mu.RLock() p := s.peers[destID] @@ -261,12 +282,12 @@ func (s *Server) WriteFrame(destID uint32, frame []byte) bool { return false } p.writeMu.Lock() - ctx, cancel := context.WithTimeout(context.Background(), s.cfg.IdleTimeout) - defer cancel() + ctx, cancel := context.WithTimeout(context.Background(), s.cfg.WriteTimeout) err := p.conn.Write(ctx, websocket.MessageBinary, frame) + cancel() p.writeMu.Unlock() if err != nil { - s.dropPeer(destID, "write error: "+err.Error()) + s.dropPeer(p, "write error: "+err.Error()) return false } s.framesOut.Add(1) @@ -453,7 +474,7 @@ func (s *Server) runAuth(ctx context.Context, conn *websocket.Conn) (uint32, err // peerReadLoop drains binary frames from a peer's WS connection into // the OnFrame callback. Exits on Close, read error, or idle timeout. func (s *Server) peerReadLoop(p *wssPeer) { - defer s.dropPeer(p.nodeID, "read loop exit") + defer s.dropPeer(p, "read loop exit") for { if p.closed.Load() { @@ -504,18 +525,26 @@ func (s *Server) peerGenForNode(nodeID uint32) uint32 { return p.gen.Load() } -// dropPeer removes a peer from the map and closes the underlying WS -// connection. Idempotent — safe to call from both the read loop exit -// and an explicit Close path. -func (s *Server) dropPeer(nodeID uint32, reason string) { +// dropPeer closes p's underlying WS connection and removes it from the +// peer map, but only if p is still the connection installed for its +// node ID. A node that reconnects installs a fresh *wssPeer under the +// same key while the previous connection's read loop is still +// unwinding; the compare-and-delete keeps that unwinding loop from +// deregistering the live connection it was replaced by. +// +// Idempotent — safe to call from both the read loop exit and an +// explicit Close path. +func (s *Server) dropPeer(p *wssPeer, reason string) { + if p == nil { + return + } s.mu.Lock() - p := s.peers[nodeID] - if p != nil { - delete(s.peers, nodeID) + if s.peers[p.nodeID] == p { + delete(s.peers, p.nodeID) } s.mu.Unlock() - if p != nil && !p.closed.Swap(true) { + if !p.closed.Swap(true) { _ = p.conn.Close(websocket.StatusNormalClosure, reason) - slog.Info("wss peer disconnected", "node_id", nodeID, "reason", reason) + slog.Info("wss peer disconnected", "node_id", p.nodeID, "reason", reason) } } diff --git a/wss/zz_peer_lifecycle_test.go b/wss/zz_peer_lifecycle_test.go new file mode 100644 index 0000000..5360a2f --- /dev/null +++ b/wss/zz_peer_lifecycle_test.go @@ -0,0 +1,217 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +package wss_test + +import ( + "context" + "crypto/ed25519" + "encoding/base64" + "encoding/json" + "fmt" + "testing" + "time" + + cw "github.com/coder/websocket" + + "github.com/pilot-protocol/beacon/wss" + "github.com/pilot-protocol/common/crypto" +) + +// startServer starts a real wss.Server on an ephemeral loopback port +// with caller-chosen timeouts and returns its ws:// URL. +func startServer(t *testing.T, pubKeys map[uint32]ed25519.PublicKey, idle, write time.Duration) (*wss.Server, string) { + t.Helper() + s, err := wss.New(wss.Config{ + BindAddr: "127.0.0.1:0", + AuthTimeout: 2 * time.Second, + IdleTimeout: idle, + WriteTimeout: write, + PubKeyLookup: func(id uint32) (ed25519.PublicKey, bool) { + k, ok := pubKeys[id] + return k, ok + }, + OnFrame: func(uint32, []byte) {}, + }) + if err != nil { + t.Fatalf("wss.New: %v", err) + } + if err := s.Start(); err != nil { + t.Fatalf("Server.Start: %v", err) + } + t.Cleanup(func() { _ = s.Close() }) + return s, waitForServer(t, s) +} + +// rawPeer dials the compat endpoint and completes the Ed25519 auth +// handshake by hand, returning the live connection. +// +// It deliberately does not use the daemon transport: that transport +// reconnects on its own, which makes "this connection was replaced" +// scenarios impossible to observe — the replaced side immediately dials +// back in and replaces its replacement, forever. A raw connection stays +// down once the server closes it. +func rawPeer(t *testing.T, wsURL string, id *crypto.Identity, nodeID uint32) *cw.Conn { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + conn, _, err := cw.Dial(ctx, wsURL, &cw.DialOptions{Subprotocols: []string{"pilot.v1"}}) + if err != nil { + t.Fatalf("dial: %v", err) + } + conn.SetReadLimit(wss.MaxFrameSize) + t.Cleanup(func() { _ = conn.CloseNow() }) + + _, chBytes, err := conn.Read(ctx) + if err != nil { + t.Fatalf("read challenge: %v", err) + } + var challenge struct { + Type string `json:"type"` + Nonce string `json:"nonce"` + Timestamp int64 `json:"ts"` + } + if err := json.Unmarshal(chBytes, &challenge); err != nil { + t.Fatalf("parse challenge: %v", err) + } + + signed := fmt.Sprintf("compat_auth:%d:%d:%s", nodeID, challenge.Timestamp, challenge.Nonce) + sig := ed25519.Sign(ed25519.PrivateKey(id.PrivateKey), []byte(signed)) + reply, _ := json.Marshal(map[string]any{ + "type": "auth_reply", + "node_id": nodeID, + "public_key": base64.StdEncoding.EncodeToString(id.PublicKey), + "sig": base64.StdEncoding.EncodeToString(sig), + }) + if err := conn.Write(ctx, cw.MessageText, reply); err != nil { + t.Fatalf("write auth reply: %v", err) + } + if _, _, err := conn.Read(ctx); err != nil { + t.Fatalf("read auth_ok: %v", err) + } + return conn +} + +// TestServer_ReplacedPeerDoesNotEvictItsReplacement pins the peer-map +// ownership rule: when a node reconnects, the previous connection's read +// loop unwinds and runs its cleanup, and that cleanup must not remove +// the connection that replaced it. Otherwise the node is registered as +// connected for a moment and then silently deregistered, and every relay +// for a peer that only has a WSS path is dropped as unroutable. +func TestServer_ReplacedPeerDoesNotEvictItsReplacement(t *testing.T) { + t.Parallel() + id, _ := crypto.GenerateIdentity() + nodeID := uint32(31337) + + s, wsURL := startServer(t, map[uint32]ed25519.PublicKey{ + nodeID: ed25519.PublicKey(id.PublicKey), + }, 30*time.Second, 2*time.Second) + + rawPeer(t, wsURL, id, nodeID) + if !waitForCondition(2*time.Second, func() bool { return s.IsConnected(nodeID) }) { + t.Fatal("first connection never registered") + } + + // Same node reconnects. handleUpgrade installs the new connection + // and closes the old one, which unblocks the old read loop. + second := rawPeer(t, wsURL, id, nodeID) + + // The replacement must stay registered while the old read loop + // finishes unwinding. + deadline := time.Now().Add(1500 * time.Millisecond) + for time.Now().Before(deadline) { + if !s.IsConnected(nodeID) { + t.Fatal("node deregistered after reconnect: the replaced connection's cleanup removed its replacement") + } + time.Sleep(10 * time.Millisecond) + } + + // And it must actually be usable. + payload := []byte("to-the-live-connection") + if !s.WriteFrame(nodeID, payload) { + t.Fatal("WriteFrame to the reconnected node returned false") + } + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + typ, frame, err := second.Read(ctx) + if err != nil { + t.Fatalf("Read on the live connection: %v", err) + } + if typ != cw.MessageBinary || string(frame) != string(payload) { + t.Errorf("read frame = %q (type %v); want %q binary", frame, typ, payload) + } +} + +// TestServer_WriteFrameIsBoundedByWriteTimeout pins that one outbound +// frame cannot pin the calling goroutine for the idle window. The +// beacon's relay workers call WriteFrame, so the write deadline is the +// longest a worker can be held by a peer that stopped draining its +// socket; the idle window is orders of magnitude longer. +func TestServer_WriteFrameIsBoundedByWriteTimeout(t *testing.T) { + t.Parallel() + id, _ := crypto.GenerateIdentity() + nodeID := uint32(24680) + + const ( + idleTimeout = 60 * time.Second + writeTimeout = 300 * time.Millisecond + // Generous headroom over writeTimeout, but far below + // idleTimeout, so the assertion distinguishes the two. + bound = 15 * time.Second + ) + + s, wsURL := startServer(t, map[uint32]ed25519.PublicKey{ + nodeID: ed25519.PublicKey(id.PublicKey), + }, idleTimeout, writeTimeout) + + // Authenticate, then never read again, so the server's writes back + // up in the kernel buffers. + rawPeer(t, wsURL, id, nodeID) + if !waitForCondition(2*time.Second, func() bool { return s.IsConnected(nodeID) }) { + t.Fatal("stalled peer never registered") + } + + frame := make([]byte, 256*1024) + done := make(chan time.Duration, 1) + go func() { + start := time.Now() + for i := 0; i < 2000; i++ { + if !s.WriteFrame(nodeID, frame) { + done <- time.Since(start) + return + } + } + done <- -1 + }() + + select { + case elapsed := <-done: + if elapsed < 0 { + t.Fatal("2000 writes to a peer that never reads all reported success") + } + if elapsed > bound { + t.Fatalf("WriteFrame blocked for %s on a stalled peer; the write deadline should bound it near %s, not the %s idle window", + elapsed, writeTimeout, idleTimeout) + } + case <-time.After(bound): + t.Fatalf("WriteFrame still blocked after %s on a stalled peer; the write deadline should bound it near %s, not the %s idle window", + bound, writeTimeout, idleTimeout) + } + + // A peer whose write timed out is dropped, so the relay router stops + // selecting it instead of retrying into the same stall. + if !waitForCondition(2*time.Second, func() bool { return !s.IsConnected(nodeID) }) { + t.Error("peer still registered after its write deadline expired; it should have been dropped") + } +} + +// TestServer_WriteTimeoutDefaults pins the default so a caller that does +// not set WriteTimeout still gets a bounded write rather than the idle +// window. +func TestServer_WriteTimeoutDefaults(t *testing.T) { + t.Parallel() + if wss.DefaultWriteTimeout >= wss.DefaultIdleTimeout { + t.Fatalf("DefaultWriteTimeout (%s) must be well below DefaultIdleTimeout (%s)", + wss.DefaultWriteTimeout, wss.DefaultIdleTimeout) + } +} diff --git a/zz_fuzz_dispatch_test.go b/zz_fuzz_dispatch_test.go index 9cce866..9e1a412 100644 --- a/zz_fuzz_dispatch_test.go +++ b/zz_fuzz_dispatch_test.go @@ -225,7 +225,7 @@ func FuzzBeaconDispatchRelay(f *testing.F) { if len(data) > maxRelayPayload { data = data[:maxRelayPayload] } - s.dispatchRelay(data) + s.dispatchRelay(data, relaySourceForUDP(fuzzAddr)) // Drain so the buffered channel cannot fill across iterations. for { select { diff --git a/zz_gossip_peer_set_test.go b/zz_gossip_peer_set_test.go new file mode 100644 index 0000000..36dd81b --- /dev/null +++ b/zz_gossip_peer_set_test.go @@ -0,0 +1,117 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +package beacon + +import ( + "encoding/binary" + "net" + "testing" + "time" +) + +// syncBody builds a gossip sync body (everything after the message-type +// byte) advertising the given node ids on behalf of peerBeaconID. +func syncBody(peerBeaconID uint32, nodeIDs ...uint32) []byte { + b := make([]byte, 6+4*len(nodeIDs)) + binary.BigEndian.PutUint32(b[0:4], peerBeaconID) + binary.BigEndian.PutUint16(b[4:6], uint16(len(nodeIDs))) + for i, id := range nodeIDs { + binary.BigEndian.PutUint32(b[6+4*i:10+4*i], id) + } + return b +} + +func peerRoute(s *Server, nodeID uint32) *net.UDPAddr { + return (*s.peerNodes.Load())[nodeID] +} + +// TestStrictGossipRejectsSourceOutsidePeerSet pins the peer-set gate: a +// sync from an address that is not a known peer beacon must not install +// routes into the peer mesh the relay workers read. +func TestStrictGossipRejectsSourceOutsidePeerSet(t *testing.T) { + t.Parallel() + s := NewWithPeers(1, []string{"198.51.100.7:9001"}) + defer s.Close() + s.SetStrictGossip(true) + if !s.StrictGossip() { + t.Fatal("SetStrictGossip(true) did not take effect") + } + + outsider := &net.UDPAddr{IP: net.IPv4(203, 0, 113, 66), Port: 9001} + s.handleSync(syncBody(2, 500, 501), outsider) + + if got := peerRoute(s, 500); got != nil { + t.Fatalf("gossip from %s installed a route for node 500 -> %s; want no route", outsider, got) + } + + // The configured peer is accepted on the same code path, so the gate + // rejects on provenance and not by refusing all gossip. + known := &net.UDPAddr{IP: net.IPv4(198, 51, 100, 7), Port: 9001} + s.handleSync(syncBody(2, 500, 501), known) + if got := peerRoute(s, 500); got == nil { + t.Fatal("gossip from the configured peer was rejected; want the route installed") + } +} + +// TestStrictGossipDefaultsOff pins that the gate is opt-in, so existing +// deployments keep converging until an operator turns it on. +func TestStrictGossipDefaultsOff(t *testing.T) { + t.Parallel() + s := NewWithPeers(1, nil) + defer s.Close() + + if s.StrictGossip() { + t.Fatal("strict gossip is on by default; it must be opt-in") + } + outsider := &net.UDPAddr{IP: net.IPv4(203, 0, 113, 67), Port: 9001} + s.handleSync(syncBody(2, 600), outsider) + if peerRoute(s, 600) == nil { + t.Fatal("with the gate off, gossip from any source should still install routes") + } +} + +// TestStrictGossipMatchesPeerOnAddressNotPort covers a peer whose gossip +// egresses from a rewritten source port. +func TestStrictGossipMatchesPeerOnAddressNotPort(t *testing.T) { + t.Parallel() + s := NewWithPeers(1, []string{"198.51.100.8:9001"}) + defer s.Close() + s.SetStrictGossip(true) + + remapped := &net.UDPAddr{IP: net.IPv4(198, 51, 100, 8), Port: 34567} + s.handleSync(syncBody(2, 700), remapped) + if peerRoute(s, 700) == nil { + t.Fatal("gossip from a known peer address on a rewritten port was rejected") + } +} + +// TestStaleGossipRoutesExpire pins that routes published by a beacon +// that stops gossiping are withdrawn, instead of standing forever and +// pointing relay traffic at a beacon that is gone. +func TestStaleGossipRoutesExpire(t *testing.T) { + t.Parallel() + s := NewWithPeers(1, nil) + defer s.Close() + + gone := &net.UDPAddr{IP: net.IPv4(203, 0, 113, 10), Port: 9001} + live := &net.UDPAddr{IP: net.IPv4(203, 0, 113, 11), Port: 9001} + s.handleSync(syncBody(2, 800), gone) + s.handleSync(syncBody(3, 801), live) + + if peerRoute(s, 800) == nil || peerRoute(s, 801) == nil { + t.Fatal("setup: both gossip sources should have installed routes") + } + + // Age out only the first source, then sweep. + s.peerWriteMu.Lock() + s.gossipSeen[gone.String()] = time.Now().Add(-2 * gossipPeerTTL) + s.peerWriteMu.Unlock() + s.sweepStaleGossip(time.Now().Add(-gossipPeerTTL)) + + if got := peerRoute(s, 800); got != nil { + t.Errorf("route for node 800 survived its source going silent: %s", got) + } + if peerRoute(s, 801) == nil { + t.Error("route for node 801 was withdrawn even though its source is still gossiping") + } +} diff --git a/zz_relay_source_limit_test.go b/zz_relay_source_limit_test.go new file mode 100644 index 0000000..f71303e --- /dev/null +++ b/zz_relay_source_limit_test.go @@ -0,0 +1,149 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +package beacon + +import ( + "encoding/binary" + "net" + "testing" + "time" +) + +// relayFrame builds a bare relay body (the bytes dispatchRelay parses, +// i.e. everything after the message-type byte). +func relayFrame(senderID, destID uint32, payload string) []byte { + b := make([]byte, 8+len(payload)) + binary.BigEndian.PutUint32(b[0:4], senderID) + binary.BigEndian.PutUint32(b[4:8], destID) + copy(b[8:], payload) + return b +} + +func drainRelayCh(s *Server) int { + n := 0 + for { + select { + case job := <-s.relayCh: + s.returnPayload(job.payload) + n++ + default: + return n + } + } +} + +// TestRelayBudgetIsPerDatagramSource pins that the relay budget is +// charged against the observed datagram source, not the sender id +// carried in the frame body. One source that varies the sender id on +// every packet must still be held to the single per-source budget. +func TestRelayBudgetIsPerDatagramSource(t *testing.T) { + t.Parallel() + s := New() + defer s.Close() + + dest := uint32(4242) + s.nodes.Upsert(dest, &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 1}, time.Now(), maxBeaconNodes) + + src := &net.UDPAddr{IP: net.IPv4(203, 0, 113, 9), Port: 40000} + + // Three times the per-source budget, every packet claiming a + // different sender id. + const attempts = maxRelaysPerSourcePerSecond * 3 + for i := 0; i < attempts; i++ { + s.dispatchRelay(relayFrame(uint32(i+1), dest, "x"), relaySourceForUDP(src)) + } + + got := drainRelayCh(s) + if got > maxRelaysPerSourcePerSecond { + t.Fatalf("one datagram source enqueued %d relays with rotating sender ids; budget is %d per second", + got, maxRelaysPerSourcePerSecond) + } + if got == 0 { + t.Fatalf("no relays enqueued at all; the limiter rejected everything") + } +} + +// TestRelayBudgetIsIndependentPerSource is the companion property: two +// distinct datagram sources must not share one budget. +func TestRelayBudgetIsIndependentPerSource(t *testing.T) { + t.Parallel() + s := New() + defer s.Close() + + dest := uint32(4243) + s.nodes.Upsert(dest, &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 1}, time.Now(), maxBeaconNodes) + + a := &net.UDPAddr{IP: net.IPv4(198, 51, 100, 1), Port: 1111} + b := &net.UDPAddr{IP: net.IPv4(198, 51, 100, 2), Port: 2222} + + // Exhaust a's budget entirely, then send one packet from b. + for i := 0; i < maxRelaysPerSourcePerSecond+50; i++ { + s.dispatchRelay(relayFrame(1, dest, "x"), relaySourceForUDP(a)) + } + before := drainRelayCh(s) + if before == 0 { + t.Fatal("source a enqueued nothing") + } + + s.dispatchRelay(relayFrame(1, dest, "y"), relaySourceForUDP(b)) + if n := drainRelayCh(s); n != 1 { + t.Fatalf("source b enqueued %d relays after source a exhausted its budget; want 1", n) + } +} + +// TestRelaySourceKeyDistinguishesEndpoints checks the key construction +// itself: same address different port, different address same port, and +// the bridged form must all be distinct, and the IPv4 / IPv4-in-IPv6 +// spellings of one address must be equal. +func TestRelaySourceKeyDistinguishesEndpoints(t *testing.T) { + t.Parallel() + + v4 := relaySourceForUDP(&net.UDPAddr{IP: net.IPv4(192, 0, 2, 5).To4(), Port: 9001}) + v4b := relaySourceForUDP(&net.UDPAddr{IP: net.IPv4(192, 0, 2, 5).To16(), Port: 9001}) + if v4 != v4b { + t.Error("the 4-byte and 16-byte spellings of one address produced different keys") + } + + otherPort := relaySourceForUDP(&net.UDPAddr{IP: net.IPv4(192, 0, 2, 5).To4(), Port: 9002}) + if v4 == otherPort { + t.Error("two ports on one address produced the same key") + } + + otherIP := relaySourceForUDP(&net.UDPAddr{IP: net.IPv4(192, 0, 2, 6).To4(), Port: 9001}) + if v4 == otherIP { + t.Error("two addresses on one port produced the same key") + } + + bridged := relaySourceForBridge(7) + if bridged == v4 || bridged == relaySourceForBridge(8) { + t.Error("bridged keys collided") + } + if relaySourceForUDP(nil) != (relaySourceKey{}) { + t.Error("a nil address should produce the zero key, not panic or vary") + } +} + +// TestRelayBudgetForBridgedFramesUsesPeerID pins that frames handed over +// by the compat WSS bridge are budgeted per authenticated peer — every +// bridged frame shares one synthetic datagram address, so keying on that +// address would put the whole bridge on a single budget. +func TestRelayBudgetForBridgedFramesUsesPeerID(t *testing.T) { + t.Parallel() + s := New() + defer s.Close() + + dest := uint32(4244) + s.nodes.Upsert(dest, &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 1}, time.Now(), maxBeaconNodes) + + for i := 0; i < maxRelaysPerSourcePerSecond+50; i++ { + s.dispatchRelay(relayFrame(999, dest, "x"), relaySourceForBridge(1001)) + } + if drainRelayCh(s) == 0 { + t.Fatal("bridged peer 1001 enqueued nothing") + } + + s.dispatchRelay(relayFrame(999, dest, "y"), relaySourceForBridge(1002)) + if n := drainRelayCh(s); n != 1 { + t.Fatalf("bridged peer 1002 enqueued %d relays after peer 1001 exhausted its budget; want 1", n) + } +} diff --git a/zz_server_branches_test.go b/zz_server_branches_test.go index 110f6fa..cd3fa60 100644 --- a/zz_server_branches_test.go +++ b/zz_server_branches_test.go @@ -97,7 +97,7 @@ func TestDispatchRelay_OversizePayloadDropped(t *testing.T) { data := make([]byte, 8+maxRelayPayload+1) binary.BigEndian.PutUint32(data[0:4], 1) binary.BigEndian.PutUint32(data[4:8], 11) - s.dispatchRelay(data) + s.dispatchRelay(data, relaySourceForUDP(addr)) // No assertion — we just want the branch hit (and no panic). } @@ -107,7 +107,7 @@ func TestDispatchRelay_TooShortNoOp(t *testing.T) { t.Parallel() s := New() defer s.Close() - s.dispatchRelay([]byte{0x01, 0x02}) // < 8 + s.dispatchRelay([]byte{0x01, 0x02}, relaySourceKey{}) // < 8 } // TestDispatchRelay_DestInPeerMesh covers the peerMap-hit branch in @@ -129,7 +129,7 @@ func TestDispatchRelay_DestInPeerMesh(t *testing.T) { binary.BigEndian.PutUint32(data[0:4], 1) binary.BigEndian.PutUint32(data[4:8], 42) data[8], data[9], data[10], data[11] = 'X', 'Y', 'Z', '!' - s.dispatchRelay(data) + s.dispatchRelay(data, relaySourceForUDP(peer)) } // TestRelayStatsLoop_ClosesOnDone exercises the s.done branch of