Forget idle scanner IPs in the brute-force detector

internal/sshwatch/brute.go

@@ -9,13 +9,18 @@ import (
 // Window. Each IP alerts once per quiet period: after firing, it must fall
 // silent for a full Window before it can fire again, preventing alert spam
 // during an ongoing attack.
+//
+// State is kept only for IPs seen within the last Window or so: an
+// internet-facing sshd sees thousands of distinct scanner IPs a day, and the
+// agent runs for months under a small memory ceiling.
 type BruteDetector struct {
 	Window    time.Duration
 	Threshold int
 
-	mu       sync.Mutex
-	attempts map[string][]time.Time
-	fired    map[string]time.Time
+	mu        sync.Mutex
+	attempts  map[string][]time.Time
+	fired     map[string]time.Time
+	lastSweep time.Time
 }
 
 // NewBruteDetector returns a detector with the given window and threshold.
@@ -33,6 +38,11 @@ func (b *BruteDetector) Record(ip string, at time.Time) bool {
 	b.mu.Lock()
 	defer b.mu.Unlock()
 
+	if at.Sub(b.lastSweep) >= b.Window {
+		b.prune(at)
+		b.lastSweep = at
+	}
+
 	cutoff := at.Add(-b.Window)
 	kept := b.attempts[ip][:0]
 	for _, t := range b.attempts[ip] {
@@ -67,3 +77,34 @@ func (b *BruteDetector) Count(ip string, now time.Time) int {
 	}
 	return n
 }
+
+// prune forgets every IP whose newest attempt has left the window. Callers
+// hold b.mu. Newest is the max timestamp in the slice, not the last element,
+// so out-of-order Record times cannot drop an IP that still has in-window
+// attempts.
+//
+// Forgetting such an IP cannot change a later decision. Its attempts would
+// be filtered out by the next Record anyway, and its fired time is never
+// later than its newest attempt, so it is outside the quiet period too:
+// the IP is free to fire again, exactly as if it had never been seen.
+func (b *BruteDetector) prune(now time.Time) {
+	cutoff := now.Add(-b.Window)
+	for ip, times := range b.attempts {
+		if len(times) == 0 || !newestAttempt(times).After(cutoff) {
+			delete(b.attempts, ip)
+			delete(b.fired, ip)
+		}
+	}
+}
+
+// newestAttempt returns the latest timestamp in times. Record does not require
+// chronological order, so prune must not assume the last element is newest.
+func newestAttempt(times []time.Time) time.Time {
+	max := times[0]
+	for _, t := range times[1:] {
+		if t.After(max) {
+			max = t
+		}
+	}
+	return max
+}

internal/sshwatch/geoip.go

@@ -32,6 +32,11 @@ func (g GeoInfo) String() string {
 	return s
 }
 
+// geoCacheMax bounds the lookup cache. Lookups happen for successful logins
+// and live sessions, so this is many times what a real host needs, while
+// keeping a long-running agent's memory flat.
+const geoCacheMax = 1024
+
 // GeoResolver looks up IPs against ip-api.com with an in-memory cache.
 // Lookups are best-effort: on any failure it returns a zero GeoInfo, never
 // an error — alerts must not depend on a third-party service being up.
@@ -91,6 +96,10 @@ func (g *GeoResolver) Lookup(ctx context.Context, ip string) GeoInfo {
 	}
 	info = GeoInfo{Country: body.Country, City: body.City, ISP: body.ISP}
 	g.mu.Lock()
+	if len(g.cache) >= geoCacheMax {
+		// Start over rather than track recency: a miss costs one lookup.
+		clear(g.cache)
+	}
 	g.cache[ip] = info
 	g.mu.Unlock()
 	return info

internal/sshwatch/sshwatch_test.go

@@ -151,6 +151,86 @@ func TestBruteDetector(t *testing.T) {
 	}
 }
 
+func TestBruteDetectorForgetsIdleIPs(t *testing.T) {
+	b := NewBruteDetector(time.Minute, 5)
+	// A scan: thousands of IPs, each failing once or twice, one of which
+	// crosses the threshold.
+	for i := 0; i < 5000; i++ {
+		ip := fmt.Sprintf("203.0.%d.%d", i/256, i%256)
+		at := t0.Add(time.Duration(i) * time.Millisecond)
+		b.Record(ip, at)
+		b.Record(ip, at)
+	}
+	attacker := "198.51.100.66"
+	for i := 0; i < 5; i++ {
+		b.Record(attacker, t0.Add(5*time.Second))
+	}
+
+	// Long after the scan, a single new attempt sweeps everything idle.
+	later := t0.Add(time.Hour)
+	b.Record("192.0.2.1", later)
+	b.mu.Lock()
+	attempts, fired := len(b.attempts), len(b.fired)
+	b.mu.Unlock()
+	if attempts != 1 || fired != 0 {
+		t.Fatalf("after the window: %d IPs with attempts, %d with fired state; want 1 and 0", attempts, fired)
+	}
+
+	// Forgetting must not change behaviour: the attacker fires afresh.
+	for i := 0; i < 4; i++ {
+		if b.Record(attacker, later.Add(time.Duration(i)*time.Second)) {
+			t.Fatalf("fired at attempt %d after being forgotten", i+1)
+		}
+	}
+	if !b.Record(attacker, later.Add(4*time.Second)) {
+		t.Error("a new burst from a forgotten IP must fire")
+	}
+}
+
+func TestBruteDetectorPruneOutOfOrder(t *testing.T) {
+	b := NewBruteDetector(time.Minute, 3)
+	ip := "198.51.100.42"
+	sec := func(n int) time.Time { return t0.Add(time.Duration(n) * time.Second) }
+
+	// Seed another IP so the first Record establishes lastSweep.
+	b.Record("192.0.2.1", sec(0))
+	// Fresh attempt, then an older one appended out of order so the slice's last
+	// element is outside the upcoming sweep window while the max is still inside.
+	b.Record(ip, sec(50))
+	b.Record(ip, sec(5))
+
+	// Sweep at 70s: cutoff is 10s. Last element (5s) is outside; max (50s) is not.
+	// Using last-element prune would drop the IP; max-based prune must keep it.
+	b.Record("192.0.2.2", sec(70))
+	if b.Count(ip, sec(70)) != 1 {
+		t.Fatalf("Count after out-of-order prune = %d, want 1 (the in-window attempt)", b.Count(ip, sec(70)))
+	}
+	b.Record(ip, sec(72))
+	if !b.Record(ip, sec(74)) {
+		t.Fatal("third in-window attempt must fire; prune must not have dropped the IP")
+	}
+}
+
+func TestBruteDetectorKeepsActiveIPs(t *testing.T) {
+	b := NewBruteDetector(time.Minute, 3)
+	ip, other := "198.51.100.7", "192.0.2.1"
+	sec := func(n int) time.Time { return t0.Add(time.Duration(n) * time.Second) }
+	b.Record(other, sec(0)) // first sweep
+	b.Record(ip, sec(40))
+	b.Record(ip, sec(50))
+	b.Record(other, sec(60)) // a full window later: sweeps, ip is still inside the window
+	if !b.Record(ip, sec(70)) {
+		t.Fatal("third attempt inside the window must fire across a sweep")
+	}
+	// Still attacking, inside the quiet period. The Record at 120s sweeps
+	// again; it must not drop the fired time and allow a re-fire.
+	for _, n := range []int{100, 110, 120, 130} {
+		if b.Record(ip, sec(n)) {
+			t.Fatalf("ongoing attack re-fired at %ds", n)
+		}
+	}
+}
+
 func TestParseUtmp(t *testing.T) {
 	var buf bytes.Buffer
 	buf.Write(EncodeUtmpRecord(2, "", "~", "reboot", t0.Add(-time.Hour))) // BOOT_TIME, skipped
@@ -243,6 +323,26 @@ func TestGeoResolver(t *testing.T) {
 	}
 }
 
+func TestGeoResolverCacheIsBounded(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		fmt.Fprint(w, `{"status":"success","country":"Germany"}`)
+	}))
+	defer srv.Close()
+
+	g := NewGeoResolver()
+	g.BaseURL = srv.URL
+	g.Client = srv.Client()
+	for i := 0; i < geoCacheMax+10; i++ {
+		g.Lookup(context.Background(), fmt.Sprintf("203.0.%d.%d", i/256, i%256))
+	}
+	g.mu.Lock()
+	n := len(g.cache)
+	g.mu.Unlock()
+	if n > geoCacheMax {
+		t.Errorf("cache holds %d entries, cap is %d", n, geoCacheMax)
+	}
+}
+
 func TestGeoResolverFailureIsQuiet(t *testing.T) {
 	g := NewGeoResolver()
 	g.BaseURL = "http://127.0.0.1:1"