Merge pull request #224 from AikidoSec/report-total-rate-limited

Count total rate limited requests

Author: bitterpanda63

agent_api/src/main/java/dev/aikido/agent_api/ShouldBlockRequest.java

@@ -6,6 +6,7 @@
 import dev.aikido.agent_api.ratelimiting.ShouldRateLimit;
 import dev.aikido.agent_api.storage.ServiceConfigStore;
 import dev.aikido.agent_api.storage.ServiceConfiguration;
+import dev.aikido.agent_api.storage.statistics.StatisticsStore;
 
 public final class ShouldBlockRequest {
     private ShouldBlockRequest() {
@@ -54,6 +55,7 @@ public static ShouldBlockRequestResult shouldBlockRequest() {
             context.getRouteMetadata(), context.getUser(), context.getRateLimitGroup(), context.getRemoteAddress()
         );
         if (rateLimitDecision.block()) {
+            StatisticsStore.incrementRateLimitedStats();
             BlockedRequestResult blockedRequestResult = new BlockedRequestResult(
                 /* type */    "ratelimited",
                 /* trigger */ rateLimitDecision.trigger(),

agent_api/src/main/java/dev/aikido/agent_api/storage/statistics/Statistics.java

@@ -14,6 +14,7 @@ public class Statistics {
     private int attacksBlocked;
     private int attackWavesDetected;
     private int attackWavesBlocked;
+    private int rateLimited;
     private long startedAt;
 
     public Statistics(int totalHits, int attacksDetected, int attacksBlocked) {
@@ -22,6 +23,7 @@ public Statistics(int totalHits, int attacksDetected, int attacksBlocked) {
         this.attacksBlocked = attacksBlocked;
         this.attackWavesDetected = 0;
         this.attackWavesBlocked = 0;
+        this.rateLimited = 0;
         this.startedAt = UnixTimeMS.getUnixTimeMS();
     }
 
@@ -80,6 +82,13 @@ public int getAttackWavesBlocked() {
         return attackWavesBlocked;
     }
 
+    public void incrementRateLimitStats() {
+        this.rateLimited += 1;
+    }
+    public int getRateLimitedStats() {
+        return this.rateLimited;
+    }
+
     // operations
     public void registerCall(String operation, OperationKind kind) {
         if (!this.operations.containsKey(operation)) {
@@ -128,7 +137,7 @@ public StatsRecord getRecord() {
         return new StatsRecord(
             this.startedAt,
             endedAt,
-            new StatsRequestsRecord(totalHits, /* aborted: unknown */ 0, attackStats, attackWaveStats),
+            new StatsRequestsRecord(totalHits, /* aborted: unknown */ 0, rateLimited, attackStats, attackWaveStats),
             getOperations(),
             Map.of("breakdown", getIpAddresses()),
             Map.of("breakdown", getUserAgents())
@@ -141,6 +150,7 @@ public void clear() {
         this.attacksDetected = 0;
         this.attackWavesBlocked = 0;
         this.attackWavesDetected = 0;
+        this.rateLimited = 0;
         this.startedAt = UnixTimeMS.getUnixTimeMS();
         this.operations.clear();
         this.ipAddressMatches.clear();
@@ -154,6 +164,7 @@ public record StatsTotalAndBlocked(int total, int blocked) {
     public record StatsRequestsRecord(
         long total,
         long aborted,
+        long rateLimited,
         StatsTotalAndBlocked attacksDetected,
         StatsTotalAndBlocked attackWaves
     ) {

agent_api/src/main/java/dev/aikido/agent_api/storage/statistics/StatisticsStore.java

@@ -97,4 +97,13 @@ public static void incrementAttackWavesDetected() {
             mutex.unlock();
         }
     }
+
+    public static void incrementRateLimitedStats() {
+        mutex.lock();
+        try {
+            stats.incrementRateLimitStats();
+        } finally {
+            mutex.unlock();
+        }
+    }
 }

agent_api/src/test/java/ShouldBlockRequestTest.java

@@ -4,7 +4,9 @@
 import dev.aikido.agent_api.context.Context;
 import dev.aikido.agent_api.context.ContextObject;
 import dev.aikido.agent_api.context.User;
+import dev.aikido.agent_api.storage.RateLimiterStore;
 import dev.aikido.agent_api.storage.ServiceConfigStore;
+import dev.aikido.agent_api.storage.statistics.StatisticsStore;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
@@ -38,12 +40,16 @@ public SampleContextObject() {
     public static void clean() {
         Context.set(null);
         ServiceConfigStore.updateFromAPIResponse(emptyAPIResponse);
+        StatisticsStore.clear();
+        RateLimiterStore.clear();
     };
 
     @AfterEach
     public void tearDown() throws SQLException {
         Context.set(null);
         ServiceConfigStore.updateFromAPIResponse(emptyAPIResponse);
+        StatisticsStore.clear();
+        RateLimiterStore.clear();
     }
 
     @Test
@@ -135,7 +141,7 @@ public void testEndpointsExistButNoMatch() throws SQLException {
     public void testEndpointsExistWithMatch() throws SQLException {
         Context.set(null);
         setEmptyConfigWithEndpointList(List.of(
-                new Endpoint("GET", "/api/*", 1, 1000, Collections.emptyList(), false, false, false)
+            new Endpoint("GET", "/api/*", 1, 1000, Collections.emptyList(), false, false, false)
         ));
 
         // Test with match & rate-limiting disabled :
@@ -144,14 +150,40 @@ public void testEndpointsExistWithMatch() throws SQLException {
 
         Context.set(null);
         setEmptyConfigWithEndpointList(List.of(
-                new Endpoint("GET", "/api/*", 1, 1000, Collections.emptyList(), false, false, true)
+            new Endpoint("GET", "/api/*", 1, 1000, Collections.emptyList(), false, false, true)
         ));
 
         // Test with match & rate-limiting enabled :
         var res2 = ShouldBlockRequest.shouldBlockRequest();
         assertFalse(res2.block());
     }
 
+    @Test
+    public void testEndpointsExistAndGetsRateLimited() throws SQLException {
+        setEmptyConfigWithEndpointList(List.of(
+                new Endpoint("GET", "/api/*", 2, 1000, Collections.emptyList(), false, false, true)
+        ));
+
+        // Test with match
+        Context.set(new SampleContextObject());
+        var res1 = ShouldBlockRequest.shouldBlockRequest();
+        assertFalse(res1.block());
+
+        // Test with match
+        var res2 = ShouldBlockRequest.shouldBlockRequest();
+        assertFalse(res2.block());
+        assertEquals(0, StatisticsStore.getStatsRecord().requests().rateLimited());
+
+        var res3 = ShouldBlockRequest.shouldBlockRequest();
+        var res4 = ShouldBlockRequest.shouldBlockRequest();
+        assertTrue(res3.block());
+        assertTrue(res4.block());
+        assertEquals("ip", res3.data().trigger());
+        assertEquals("192.168.1.1", res3.data().ip());
+        assertEquals("ratelimited", res3.data().type());
+        assertEquals(2, StatisticsStore.getStatsRecord().requests().rateLimited());
+    }
+
     @Test
     public void testThreadClientInvalid() throws SQLException {
         Context.set(new SampleContextObject());

agent_api/src/test/java/storage/StatisticsTest.java

@@ -34,11 +34,14 @@ public void testClear() {
         stats.incrementAttacksDetected("test2");
         stats.incrementAttacksDetected("test1");
         stats.incrementAttacksDetected("test1");
+        stats.incrementRateLimitStats();
+        stats.incrementRateLimitStats();
         assertEquals(3, stats.getAttacksDetected());
         assertEquals(2, stats.getAttacksBlocked());
         assertEquals(20, stats.getTotalHits());
         assertEquals(2, stats.getOperations().get("test1").getAttacksDetected().get("total"));
         assertEquals(1, stats.getOperations().get("test1").getAttacksDetected().get("blocked"));
+        assertEquals(2, stats.getRateLimitedStats());
 
         assertFalse(stats.getOperations().containsKey("test2"));
         // Reset :
@@ -47,6 +50,7 @@ public void testClear() {
         assertEquals(0, stats.getAttacksBlocked());
         assertEquals(0, stats.getAttacksDetected());
         assertEquals(0, stats.getTotalHits());
+        assertEquals(0, stats.getRateLimitedStats());
 
     }
 
@@ -56,6 +60,7 @@ public void testConstructor() {
         assertEquals(100, stats2.getTotalHits());
         assertEquals(5, stats2.getAttacksDetected());
         assertEquals(1, stats2.getAttacksBlocked());
+        assertEquals(0, stats.getRateLimitedStats());
     }
 
     @Test