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