Cache total stats in memory

honeypot/ssh/store/cache_test.go

@@ -11,3 +11,59 @@ func TestCacheStore(t *testing.T) {
 	s := store.NewCachingStore(backend)
 	testLoginAttemptStore(s, t)
 }
+
+func TestCacheTotalStats(t *testing.T) {
+	backend := &store.MemoryStore{}
+	// Add initial attempts, to ensure that the cache is correcly initialized with existing attempts
+	attempts := randomAttempts(50000)
+	for _, attempt := range attempts {
+		err := backend.AddAttempt(attempt)
+		if err != nil {
+			t.Fatalf("Error adding attempts: %s", err)
+		}
+	}
+	s := store.NewCachingStore(backend)
+	cacheTotals, err := s.Stats(store.LoginStatsTotals, 0)
+	if err != nil {
+		t.Fatalf("Error getting cached stats: %s", err)
+	}
+	backendTotals, err := backend.Stats(store.LoginStatsTotals, 0)
+	if err != nil {
+		t.Fatalf("Error getting cached stats: %s", err)
+	}
+	for i := range cacheTotals {
+		if cacheTotals[i].Count != backendTotals[i].Count || cacheTotals[i].Name != backendTotals[i].Name {
+			t.Fatalf("Mismatched totals: Cache: %+v Backend: %+v", cacheTotals[i], backendTotals[i])
+		}
+	}
+
+	// Add the same attempts again, to ensure that duplicates are handled correctly
+	for _, attempt := range attempts {
+		err := s.AddAttempt(attempt)
+		if err != nil {
+			t.Fatalf("Error adding attempts: %s", err)
+		}
+	}
+
+	// Add some new attempts
+	attempts = randomAttempts(10000)
+	for _, attempt := range attempts {
+		err := s.AddAttempt(attempt)
+		if err != nil {
+			t.Fatalf("Error adding attempts: %s", err)
+		}
+	}
+	cacheTotals, err = s.Stats(store.LoginStatsTotals, 0)
+	if err != nil {
+		t.Fatalf("Error getting cached stats: %s", err)
+	}
+	backendTotals, err = backend.Stats(store.LoginStatsTotals, 0)
+	if err != nil {
+		t.Fatalf("Error getting cached stats: %s", err)
+	}
+	for i := range cacheTotals {
+		if cacheTotals[i].Count != backendTotals[i].Count || cacheTotals[i].Name != backendTotals[i].Name {
+			t.Fatalf("Mismatched totals: Cache: %+v Backend: %+v", cacheTotals[i], backendTotals[i])
+		}
+	}
+}