diff --git a/nodes_shard.go b/nodes_shard.go index 5fec8b0..d1eb642 100644 --- a/nodes_shard.go +++ b/nodes_shard.go @@ -91,6 +91,23 @@ func (m *nodeMap) Snapshot(nodeID uint32) (addr *net.UDPAddr, ok bool) { return n.addr, true } +// claimNotify returns the node's address for a notify and records the +// time, unless one was sent to it within minGap (throttled). +func (m *nodeMap) claimNotify(nodeID uint32, now time.Time, minGap time.Duration) (addr *net.UDPAddr, ok, throttled bool) { + s := m.shardFor(nodeID) + s.mu.Lock() + defer s.mu.Unlock() + n, exists := s.nodes[nodeID] + if !exists { + return nil, false, false + } + if !n.lastNotify.IsZero() && now.Sub(n.lastNotify) < minGap { + return n.addr, true, true + } + n.lastNotify = now + return n.addr, true, false +} + // Has reports whether the node is locally registered. func (m *nodeMap) Has(nodeID uint32) bool { s := m.shardFor(nodeID) diff --git a/server.go b/server.go index 5212b67..708dd95 100644 --- a/server.go +++ b/server.go @@ -28,6 +28,8 @@ import ( type beaconNode struct { addr *net.UDPAddr lastSeen time.Time + // lastNotify is when NotifyNode last sent this node a MsgNotify. + lastNotify time.Time } // relayJob is a pre-parsed relay packet dispatched to a worker. @@ -1306,6 +1308,44 @@ func (s *Server) SendPunchCommand(nodeID uint32, targetIP net.IP, targetPort uin return err } +// MsgNotify is a beacon → node message, [0x0A][kind(1)], telling the node +// that the registry is holding something for it. It carries no node ID, +// address or payload, so it discloses nothing to the node or to anyone +// watching the path; the node reacts by polling the registry over its own +// authenticated connection. Daemons that predate it drop it as an unknown +// beacon message. (The other beacon message types live in common/protocol; +// this one should move there. Must match the daemon's beaconMsgNotify.) +const MsgNotify byte = 0x0A + +// NotifyKindHandshake is the MsgNotify kind for a relayed trust-handshake +// request or answer waiting at the registry. +const NotifyKindHandshake byte = 0x01 + +// notifyMinInterval is the least time between two notifies to one node. +// A notify only asks the node to poll, and a poll collects everything that +// is waiting, so more often than this buys nothing. +const notifyMinInterval = time.Second + +// NotifyNode tells a node that the registry holds a relayed handshake for +// it. Best effort: the node may not be registered with this beacon, may be +// on the compat (WSS) transport, or may predate MsgNotify — it then finds +// the handshake on its next poll, as before. A notify within +// notifyMinInterval of the previous one to the same node is skipped (nil). +func (s *Server) NotifyNode(nodeID uint32) error { + if s.conn == nil { + return fmt.Errorf("beacon: not listening") + } + nodeAddr, ok, throttled := s.nodes.claimNotify(nodeID, time.Now(), notifyMinInterval) + if !ok { + return fmt.Errorf("node %d: %w", nodeID, protocol.ErrNodeNotFound) + } + if throttled { + return nil + } + _, err := s.conn.WriteToUDP([]byte{MsgNotify, NotifyKindHandshake}, nodeAddr) + return err +} + // --- Reap --- // reapLoop periodically removes stale node entries that haven't sent a diff --git a/zz_notify_node_test.go b/zz_notify_node_test.go new file mode 100644 index 0000000..593a061 --- /dev/null +++ b/zz_notify_node_test.go @@ -0,0 +1,62 @@ +// SPDX-License-Identifier: AGPL-3.0-or-later + +package beacon + +import ( + "errors" + "net" + "testing" + "time" + + "github.com/pilot-protocol/common/protocol" +) + +// TestNotifyNode: a registered node receives the two-byte notify — nothing +// else — at most once per notifyMinInterval; an unknown node is an error. +func TestNotifyNode(t *testing.T) { + t.Parallel() + s := New() + go s.ListenAndServe("127.0.0.1:0") + <-s.Ready() + defer s.Close() + + if err := s.NotifyNode(99999); !errors.Is(err, protocol.ErrNodeNotFound) { + t.Fatalf("unknown node: err = %v, want ErrNodeNotFound", err) + } + + node, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1)}) + if err != nil { + t.Fatal(err) + } + defer node.Close() + s.nodes.Upsert(7, node.LocalAddr().(*net.UDPAddr), time.Now(), maxBeaconNodes) + + for i := 0; i < 20; i++ { + if err := s.NotifyNode(7); err != nil { + t.Fatalf("NotifyNode: %v", err) + } + } + buf := make([]byte, 64) + _ = node.SetReadDeadline(time.Now().Add(2 * time.Second)) + n, _, err := node.ReadFromUDP(buf) + if err != nil { + t.Fatalf("no notify received: %v", err) + } + if n != 2 || buf[0] != MsgNotify || buf[1] != NotifyKindHandshake { + t.Fatalf("notify = %x, want [%02x %02x]", buf[:n], MsgNotify, NotifyKindHandshake) + } + _ = node.SetReadDeadline(time.Now().Add(300 * time.Millisecond)) + if n, _, err := node.ReadFromUDP(buf); err == nil { + t.Fatalf("a second notify (%x) arrived within %s of the first", buf[:n], notifyMinInterval) + } + + // After the interval the next one goes out. + s.nodes.shardFor(7).nodes[7].lastNotify = time.Now().Add(-notifyMinInterval) + if err := s.NotifyNode(7); err != nil { + t.Fatalf("NotifyNode: %v", err) + } + _ = node.SetReadDeadline(time.Now().Add(2 * time.Second)) + if _, _, err := node.ReadFromUDP(buf); err != nil { + t.Fatalf("no notify after the interval: %v", err) + } +}