diff --git a/routing/routing.go b/routing/routing.go index e8a54e0..4afe4b4 100644 --- a/routing/routing.go +++ b/routing/routing.go @@ -20,6 +20,12 @@ import ( // BeaconTTL is how long a beacon registration is valid without re-register. const BeaconTTL = 60 * time.Second +// BeaconRegisterCooldown is the minimum interval between successive +// HandleBeaconRegister calls for the same beacon ID. This prevents an +// attacker (even an admin-token-armed one) from churning through many +// distinct beacon IDs faster than one per cooldown window. +const BeaconRegisterCooldown = 10 * time.Second + type BeaconEntry struct { ID uint32 Addr string @@ -81,10 +87,18 @@ func (st *Store) HandleBeaconRegister(msg map[string]interface{}) (map[string]in } st.mu.Lock() + now := st.now() + if existing, ok := st.beacons[beaconID]; ok { + if elapsed := now.Sub(existing.LastSeen); elapsed < BeaconRegisterCooldown { + st.mu.Unlock() + remain := BeaconRegisterCooldown - elapsed + return nil, fmt.Errorf("beacon %d re-registered too soon: cooldown %v remaining", beaconID, remain.Round(time.Second)) + } + } st.beacons[beaconID] = &BeaconEntry{ ID: beaconID, Addr: addr, - LastSeen: st.now(), + LastSeen: now, } st.mu.Unlock() diff --git a/routing/zz_routing_test.go b/routing/zz_routing_test.go index d8dc3cf..667dcac 100644 --- a/routing/zz_routing_test.go +++ b/routing/zz_routing_test.go @@ -199,6 +199,60 @@ func TestHandlePunchRequesterNotParticipant(t *testing.T) { } } +func TestBeaconRegisterCooldown(t *testing.T) { + t.Parallel() + st := routing.NewStore(nil) + + base := time.Now() + st.SetClock(func() time.Time { return base }) + + // First register should succeed. + resp, err := st.HandleBeaconRegister(map[string]interface{}{ + "beacon_id": float64(1), + "addr": "1.2.3.4:9001", + }) + if err != nil { + t.Fatalf("first register: %v", err) + } + if resp["type"] != "beacon_register_ok" { + t.Fatalf("unexpected type: %v", resp["type"]) + } + + // Immediate re-register within cooldown should fail. + _, err = st.HandleBeaconRegister(map[string]interface{}{ + "beacon_id": float64(1), + "addr": "1.2.3.4:9001", + }) + if err == nil { + t.Fatal("expected cooldown error for immediate re-register") + } + + // Advance past cooldown — re-register should succeed again. + st.SetClock(func() time.Time { return base.Add(routing.BeaconRegisterCooldown + time.Second) }) + resp, err = st.HandleBeaconRegister(map[string]interface{}{ + "beacon_id": float64(1), + "addr": "1.2.3.4:9001", + }) + if err != nil { + t.Fatalf("re-register after cooldown: %v", err) + } + if resp["type"] != "beacon_register_ok" { + t.Fatalf("unexpected type after cooldown: %v", resp["type"]) + } + + // Different beacon ID should NOT be affected by the cooldown on ID 1. + resp, err = st.HandleBeaconRegister(map[string]interface{}{ + "beacon_id": float64(2), + "addr": "5.6.7.8:9001", + }) + if err != nil { + t.Fatalf("different beacon ID should bypass cooldown: %v", err) + } + if resp["type"] != "beacon_register_ok" { + t.Fatalf("unexpected type for different beacon: %v", resp["type"]) + } +} + func TestHandlePunchBackendVerifyError(t *testing.T) { t.Parallel() be := &stubBackend{