From 0f41388bd78668ceae6d5c12b05868bd0ca8fd1f Mon Sep 17 00:00:00 2001
From: Jason Cameron <git@jasoncameron.dev>
Date: Sat, 29 Mar 2025 23:24:06 -0400
Subject: [PATCH] Add periodic cleanup job for DecayMap (#8) (#158)

* Add periodic cleanup job for DecayMap

see https://github.com/TecharoHQ/anubis/issues/8

* Refactor: Improve DecayMap cleanup tests and add Len method

- Refactored DecayMap cleanup tests to use the new Len method
  for more precise assertions.
- Added a Len method to DecayMap to retrieve the number of
  entries.
- Simplified conditional checks in Get method.

* chore(changelog): add entry

* fix(tests): Use Impl.expire for decaymap cleanup

Signed-off-by: Jason Cameron <git@jasoncameron.dev>

---------

Signed-off-by: Jason Cameron <git@jasoncameron.dev>
---
 cmd/anubis/main.go        | 16 ++++++++++++++++
 decaymap/decaymap.go      | 20 ++++++++++++++++++++
 decaymap/decaymap_test.go | 29 +++++++++++++++++++++++++++++
 docs/docs/CHANGELOG.md    |  2 +-
 lib/anubis.go             |  4 ++++
 5 files changed, 70 insertions(+), 1 deletion(-)

diff --git a/cmd/anubis/main.go b/cmd/anubis/main.go
index a7e7350..4cee20c 100644
--- a/cmd/anubis/main.go
+++ b/cmd/anubis/main.go
@@ -135,6 +135,20 @@ func makeReverseProxy(target string) (http.Handler, error) {
 	return rp, nil
 }
 
+func startDecayMapCleanup(ctx context.Context, s *libanubis.Server) {
+	ticker := time.NewTicker(1 * time.Hour)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ticker.C:
+			s.CleanupDecayMap()
+		case <-ctx.Done():
+			return
+		}
+	}
+}
+
 func main() {
 	flagenv.Parse()
 	flag.Parse()
@@ -210,6 +224,8 @@ func main() {
 		go metricsServer(ctx, wg.Done)
 	}
 
+	go startDecayMapCleanup(ctx, s)
+
 	var h http.Handler
 	h = s
 	h = internal.RemoteXRealIP(*useRemoteAddress, *bindNetwork, h)
diff --git a/decaymap/decaymap.go b/decaymap/decaymap.go
index 7498bb5..57ee6c2 100644
--- a/decaymap/decaymap.go
+++ b/decaymap/decaymap.go
@@ -85,3 +85,23 @@ func (m *Impl[K, V]) Set(key K, value V, ttl time.Duration) {
 		expiry: time.Now().Add(ttl),
 	}
 }
+
+// Cleanup removes all expired entries from the DecayMap.
+func (m *Impl[K, V]) Cleanup() {
+	m.lock.Lock()
+	defer m.lock.Unlock()
+
+	now := time.Now()
+	for key, entry := range m.data {
+		if now.After(entry.expiry) {
+			delete(m.data, key)
+		}
+	}
+}
+
+// Len returns the number of entries in the DecayMap.
+func (m *Impl[K, V]) Len() int {
+	m.lock.RLock()
+	defer m.lock.RUnlock()
+	return len(m.data)
+}
diff --git a/decaymap/decaymap_test.go b/decaymap/decaymap_test.go
index c930e08..c1830ed 100644
--- a/decaymap/decaymap_test.go
+++ b/decaymap/decaymap_test.go
@@ -29,3 +29,32 @@ func TestImpl(t *testing.T) {
 		t.Error("got value even though it was supposed to be expired")
 	}
 }
+
+func TestCleanup(t *testing.T) {
+	dm := New[string, string]()
+
+	dm.Set("test1", "hi1", 1*time.Second)
+	dm.Set("test2", "hi2", 2*time.Second)
+	dm.Set("test3", "hi3", 3*time.Second)
+
+	dm.expire("test1") // Force expire test1
+	dm.expire("test2") // Force expire test2
+
+	dm.Cleanup()
+
+	finalLen := dm.Len() // Get the length after cleanup
+
+	if finalLen != 1 { // "test3" should be the only one left
+		t.Errorf("Cleanup failed to remove expired entries. Expected length 1, got %d", finalLen)
+	}
+
+	if _, ok := dm.Get("test1"); ok { // Verify Get still behaves correctly after Cleanup
+		t.Error("test1 should not be found after cleanup")
+	}
+	if _, ok := dm.Get("test2"); ok {
+		t.Error("test2 should not be found after cleanup")
+	}
+	if val, ok := dm.Get("test3"); !ok || val != "hi3" {
+		t.Error("test3 should still be found after cleanup")
+	}
+}
diff --git a/docs/docs/CHANGELOG.md b/docs/docs/CHANGELOG.md
index 2ddda01..4189e57 100644
--- a/docs/docs/CHANGELOG.md
+++ b/docs/docs/CHANGELOG.md
@@ -10,8 +10,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
+- Added a periodic cleanup routine for the decaymap that removes expired entries, ensuring stale data is properly pruned.
 - Added a no-store Cache-Control header to the challenge page
-
 - Hide the directory listings for Anubis' internal static content
 - Changed `--debug-x-real-ip-default` to `--use-remote-address`, getting the IP address from the request's socket address instead.
 - DroneBL lookups have been disabled by default
diff --git a/lib/anubis.go b/lib/anubis.go
index 114356c..c61b110 100644
--- a/lib/anubis.go
+++ b/lib/anubis.go
@@ -529,3 +529,7 @@ func (s *Server) checkRemoteAddress(b policy.Bot, addr net.IP) bool {
 
 	return ok
 }
+
+func (s *Server) CleanupDecayMap() {
+	s.DNSBLCache.Cleanup()
+}