From c95a54ad4812e8dda20b5e61e7ca024a7337b577 Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Fri, 20 Jun 2025 14:49:24 +0200 Subject: [PATCH 01/10] Update Statistics.java to add rateLimited --- .../storage/statistics/Statistics.java | 31 ++++++++++++------- 1 file changed, 20 insertions(+), 11 deletions(-) diff --git a/agent_api/src/main/java/dev/aikido/agent_api/storage/statistics/Statistics.java b/agent_api/src/main/java/dev/aikido/agent_api/storage/statistics/Statistics.java index ce53f76a..64e258ac 100644 --- a/agent_api/src/main/java/dev/aikido/agent_api/storage/statistics/Statistics.java +++ b/agent_api/src/main/java/dev/aikido/agent_api/storage/statistics/Statistics.java @@ -10,19 +10,19 @@ public class Statistics { private final Map ipAddressMatches = new HashMap<>(); private final Map userAgentMatches = new HashMap<>(); private int totalHits; + private final int aborted; // We don't use the "aborted" field right now + private int rateLimited; private int attacksDetected; private int attacksBlocked; private long startedAt; - public Statistics(int totalHits, int attacksDetected, int attacksBlocked) { - this.totalHits = totalHits; - this.attacksDetected = attacksDetected; - this.attacksBlocked = attacksBlocked; - this.startedAt = UnixTimeMS.getUnixTimeMS(); - } - public Statistics() { - this(0, 0, 0); + this.totalHits = 0; + this.rateLimited = 0; + this.aborted = 0; + this.attacksDetected = 0; + this.attacksBlocked = 0; + this.startedAt = UnixTimeMS.getUnixTimeMS(); } @@ -35,6 +35,14 @@ public int getTotalHits() { return totalHits; } + public void incrementRateLimited() { + rateLimited += 1; + } + + public int getRateLimited() { + return rateLimited; + } + // attack stats public void incrementAttacksDetected(String operation) { @@ -104,8 +112,7 @@ public void addMatchToUserAgents(String key) { public StatsRecord getRecord() { long endedAt = UnixTimeMS.getUnixTimeMS(); return new StatsRecord(this.startedAt, endedAt, new StatsRequestsRecord( - /* total */ totalHits, - /* aborted */ 0, // Unknown statistic, default to 0, + totalHits, aborted, rateLimited, /* attacksDetected */ Map.of( "total", attacksDetected, "blocked", attacksBlocked @@ -118,6 +125,7 @@ public StatsRecord getRecord() { public void clear() { this.totalHits = 0; + this.rateLimited = 0; this.attacksBlocked = 0; this.attacksDetected = 0; this.startedAt = UnixTimeMS.getUnixTimeMS(); @@ -127,7 +135,8 @@ public void clear() { } // Stats records for sending out the heartbeat : - public record StatsRequestsRecord(long total, long aborted, Map attacksDetected) { + public record StatsRequestsRecord(long total, long aborted, long rateLimited, + Map attacksDetected) { } public record StatsRecord(long startedAt, long endedAt, StatsRequestsRecord requests, From 082eaa20ea810239ad7507916090f9f3ceef572b Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Fri, 20 Jun 2025 15:33:32 +0200 Subject: [PATCH 02/10] Update RouteEntry.java: add rateLimitedCount --- .../aikido/agent_api/storage/routes/RouteEntry.java | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/agent_api/src/main/java/dev/aikido/agent_api/storage/routes/RouteEntry.java b/agent_api/src/main/java/dev/aikido/agent_api/storage/routes/RouteEntry.java index 9bfe2635..d17ebf60 100644 --- a/agent_api/src/main/java/dev/aikido/agent_api/storage/routes/RouteEntry.java +++ b/agent_api/src/main/java/dev/aikido/agent_api/storage/routes/RouteEntry.java @@ -1,17 +1,15 @@ package dev.aikido.agent_api.storage.routes; -import com.google.gson.*; import dev.aikido.agent_api.api_discovery.APISpec; import dev.aikido.agent_api.context.RouteMetadata; -import java.lang.reflect.Type; - import static dev.aikido.agent_api.api_discovery.APISpecMerger.mergeAPISpecs; public class RouteEntry { final String method; final String path; private int hits; + private int rateLimitedCount; private APISpec apispec; public RouteEntry(String method, String path) { @@ -32,6 +30,13 @@ public int getHits() { return hits; } + public void incrementRateLimitCount() { + rateLimitedCount++; + } + + public int getRateLimitCount() { + return rateLimitedCount; + } public void updateApiSpec(APISpec newApiSpec) { this.apispec = mergeAPISpecs(newApiSpec, this.apispec); } From 182dca53c7262e1d049321371eea3421e4a13de1 Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Fri, 20 Jun 2025 15:37:33 +0200 Subject: [PATCH 03/10] Add test cases to RouteEntry --- agent_api/src/test/java/storage/RouteEntryTest.java | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/agent_api/src/test/java/storage/RouteEntryTest.java b/agent_api/src/test/java/storage/RouteEntryTest.java index d25fea8a..b0878d8f 100644 --- a/agent_api/src/test/java/storage/RouteEntryTest.java +++ b/agent_api/src/test/java/storage/RouteEntryTest.java @@ -37,4 +37,17 @@ public void testGsonWithoutSerializer() throws IOException { ); } + @Test + public void testIncrementRateLimitedCount() { + // Initial count should be 0 + assertEquals(0, route1.getRateLimitCount()); + + // Increment the rate limited count + route1.incrementRateLimitCount(); + assertEquals(1, route1.getRateLimitCount()); + + // Increment again + route1.incrementRateLimitCount(); + assertEquals(2, route1.getRateLimitCount()); + } } From 92c659066cce78500569ab640540ddc0d1f583fa Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Fri, 20 Jun 2025 15:37:50 +0200 Subject: [PATCH 04/10] Update StatisticsTest test cases --- .../src/test/java/storage/StatisticsTest.java | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/agent_api/src/test/java/storage/StatisticsTest.java b/agent_api/src/test/java/storage/StatisticsTest.java index 6d0101c9..3ed716be 100644 --- a/agent_api/src/test/java/storage/StatisticsTest.java +++ b/agent_api/src/test/java/storage/StatisticsTest.java @@ -52,15 +52,24 @@ public void testClear() { @Test public void testConstructor() { - Statistics stats2 = new Statistics(100, 5, 1); - assertEquals(100, stats2.getTotalHits()); - assertEquals(5, stats2.getAttacksDetected()); - assertEquals(1, stats2.getAttacksBlocked()); + Statistics stats2 = new Statistics(); + assertEquals(0, stats2.getTotalHits()); + assertEquals(0, stats2.getRateLimited()); + assertEquals(0, stats2.getAttacksDetected()); + assertEquals(0, stats2.getAttacksBlocked()); } @Test public void testStatsRecord() { - Statistics stats2 = new Statistics(100, 5, 1); + Statistics stats2 = new Statistics(); + stats2.incrementTotalHits(100); + stats2.incrementAttacksDetected("op2"); + stats2.incrementAttacksDetected("op2"); + stats2.incrementAttacksDetected("op2"); + stats2.incrementAttacksDetected("op2"); + stats2.incrementAttacksDetected("op2"); + stats2.incrementAttacksBlocked("op2"); + stats2.registerCall("operation1", OperationKind.FS_OP); Statistics.StatsRecord statsRecord = stats2.getRecord(); assertEquals(5, statsRecord.requests().attacksDetected().get("total")); From a9fd605e4c023f01e27e10d8b59e502a00dc6839 Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Tue, 24 Jun 2025 13:21:26 +0200 Subject: [PATCH 05/10] Add extra test cases for getRateLimited to StatisticsTest --- agent_api/src/test/java/storage/StatisticsTest.java | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/agent_api/src/test/java/storage/StatisticsTest.java b/agent_api/src/test/java/storage/StatisticsTest.java index 3ed716be..c66dfbbe 100644 --- a/agent_api/src/test/java/storage/StatisticsTest.java +++ b/agent_api/src/test/java/storage/StatisticsTest.java @@ -34,9 +34,11 @@ public void testClear() { stats.incrementAttacksDetected("test2"); stats.incrementAttacksDetected("test1"); stats.incrementAttacksDetected("test1"); + stats.incrementRateLimited(); assertEquals(3, stats.getAttacksDetected()); assertEquals(2, stats.getAttacksBlocked()); assertEquals(20, stats.getTotalHits()); + assertEquals(1, stats.getRateLimited()); assertEquals(2, stats.getOperations().get("test1").getAttacksDetected().get("total")); assertEquals(1, stats.getOperations().get("test1").getAttacksDetected().get("blocked")); @@ -47,7 +49,7 @@ public void testClear() { assertEquals(0, stats.getAttacksBlocked()); assertEquals(0, stats.getAttacksDetected()); assertEquals(0, stats.getTotalHits()); - + assertEquals(0, stats.getRateLimited()); } @Test From ac8774ad8758ca24215cf7dc576c4773eeb70e17 Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Tue, 24 Jun 2025 13:28:39 +0200 Subject: [PATCH 06/10] Add rate limit increment to Routes & to it's store --- .../agent_api/storage/routes/Routes.java | 22 ++++++++++++------- .../agent_api/storage/routes/RoutesStore.java | 11 ++++++++++ .../src/test/java/storage/RoutesTest.java | 14 ++++++++++++ 3 files changed, 39 insertions(+), 8 deletions(-) diff --git a/agent_api/src/main/java/dev/aikido/agent_api/storage/routes/Routes.java b/agent_api/src/main/java/dev/aikido/agent_api/storage/routes/Routes.java index 64503b5f..c8f538e4 100644 --- a/agent_api/src/main/java/dev/aikido/agent_api/storage/routes/Routes.java +++ b/agent_api/src/main/java/dev/aikido/agent_api/storage/routes/Routes.java @@ -19,24 +19,30 @@ public Routes() { this(1000); // Default max size } - private void initializeRoute(RouteMetadata routeMetadata) { + private void ensureRoute(RouteMetadata routeMetadata) { manageRoutesSize(); String key = routeToKey(routeMetadata); - routes.put(key, new RouteEntry(routeMetadata)); + if(!routes.containsKey(key)) { + routes.put(key, new RouteEntry(routeMetadata)); + } } public void incrementRoute(RouteMetadata routeMetadata) { - String key = routeToKey(routeMetadata); - if (!routes.containsKey(key)) { - // if the route does not yet exist, create it. - initializeRoute(routeMetadata); - } - RouteEntry route = routes.get(key); + ensureRoute(routeMetadata); + RouteEntry route = this.get(routeMetadata); if (route != null) { route.incrementHits(); } } + public void incrementRateLimitCount(RouteMetadata routeMetadata) { + ensureRoute(routeMetadata); + RouteEntry route = this.get(routeMetadata); + if (route != null) { + route.incrementRateLimitCount(); + } + } + public RouteEntry get(RouteMetadata routeMetadata) { String key = routeToKey(routeMetadata); return routes.get(key); diff --git a/agent_api/src/main/java/dev/aikido/agent_api/storage/routes/RoutesStore.java b/agent_api/src/main/java/dev/aikido/agent_api/storage/routes/RoutesStore.java index 4fcb6b31..1af518e1 100644 --- a/agent_api/src/main/java/dev/aikido/agent_api/storage/routes/RoutesStore.java +++ b/agent_api/src/main/java/dev/aikido/agent_api/storage/routes/RoutesStore.java @@ -59,6 +59,17 @@ public static void addRouteHits(RouteMetadata routeMetadata) { } } + public static void addRouteRateLimitedCount(RouteMetadata routeMetadata) { + mutex.lock(); + try { + routes.incrementRateLimitCount(routeMetadata); + } catch (Throwable e) { + logger.debug("Error occurred incrementing route rate limit count: %s", e.getMessage()); + } finally { + mutex.unlock(); + } + } + public static void clear() { mutex.lock(); try { diff --git a/agent_api/src/test/java/storage/RoutesTest.java b/agent_api/src/test/java/storage/RoutesTest.java index 90eabb43..141d7f4c 100644 --- a/agent_api/src/test/java/storage/RoutesTest.java +++ b/agent_api/src/test/java/storage/RoutesTest.java @@ -51,6 +51,20 @@ void testIncrementNonExistentRoute() { assertEquals(1, routes.size()); } + @Test + void testIncrementRouteRateLimitCount() { + routes.incrementRateLimitCount(routeMetadata1); + RouteEntry entry = routes.get(routeMetadata1); + assertNotNull(entry); + assertEquals(1, entry.getRateLimitCount()); + } + + @Test + void testIncrementNonExistentRouteRateLimit() { + routes.incrementRateLimitCount(routeMetadata1); + assertEquals(1, routes.size()); + } + @Test void testManageRoutesSize() { routes.incrementRoute(routeMetadata1); From 4ac6b1437f15b611f4a944c8522a7627be79146e Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Tue, 24 Jun 2025 13:29:18 +0200 Subject: [PATCH 07/10] Add incrementRateLimited to StatisticsStore --- .../agent_api/storage/statistics/StatisticsStore.java | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/agent_api/src/main/java/dev/aikido/agent_api/storage/statistics/StatisticsStore.java b/agent_api/src/main/java/dev/aikido/agent_api/storage/statistics/StatisticsStore.java index 3d1fad42..e611793c 100644 --- a/agent_api/src/main/java/dev/aikido/agent_api/storage/statistics/StatisticsStore.java +++ b/agent_api/src/main/java/dev/aikido/agent_api/storage/statistics/StatisticsStore.java @@ -35,6 +35,15 @@ public static void incrementHits() { } } + public static void incrementRateLimited() { + mutex.lock(); + try { + stats.incrementRateLimited(); + } finally { + mutex.unlock(); + } + } + public static void incrementAttacksDetected(String operation) { mutex.lock(); try { From a9fc757574784189b606df6605c7c6fe76901df5 Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Tue, 24 Jun 2025 13:29:33 +0200 Subject: [PATCH 08/10] ShouldBlockRequest trigger rate limit increment --- .../main/java/dev/aikido/agent_api/ShouldBlockRequest.java | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/agent_api/src/main/java/dev/aikido/agent_api/ShouldBlockRequest.java b/agent_api/src/main/java/dev/aikido/agent_api/ShouldBlockRequest.java index 41dbb4fb..e321aacc 100644 --- a/agent_api/src/main/java/dev/aikido/agent_api/ShouldBlockRequest.java +++ b/agent_api/src/main/java/dev/aikido/agent_api/ShouldBlockRequest.java @@ -5,6 +5,8 @@ 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.routes.RoutesStore; +import dev.aikido.agent_api.storage.statistics.StatisticsStore; public final class ShouldBlockRequest { private ShouldBlockRequest() { @@ -34,6 +36,10 @@ public static ShouldBlockRequestResult shouldBlockRequest() { context.getRouteMetadata(), context.getUser(), context.getRemoteAddress() ); if (rateLimitDecision.block()) { + // increment rate-limiting stats both globally and on the route : + StatisticsStore.incrementRateLimited(); + RoutesStore.addRouteRateLimitedCount(context.getRouteMetadata()); + BlockedRequestResult blockedRequestResult = new BlockedRequestResult( "ratelimited", rateLimitDecision.trigger(), context.getRemoteAddress() ); From a65492b5f530494a6bbfdf8da6e8107194158fb8 Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Tue, 24 Jun 2025 13:37:07 +0200 Subject: [PATCH 09/10] Update test cases in ShouldBlockRequest --- .../src/test/java/ShouldBlockRequestTest.java | 31 ++++++++++++++++--- 1 file changed, 27 insertions(+), 4 deletions(-) diff --git a/agent_api/src/test/java/ShouldBlockRequestTest.java b/agent_api/src/test/java/ShouldBlockRequestTest.java index 3b731e03..eeddc6c2 100644 --- a/agent_api/src/test/java/ShouldBlockRequestTest.java +++ b/agent_api/src/test/java/ShouldBlockRequestTest.java @@ -4,7 +4,10 @@ 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.routes.RoutesStore; +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 +41,18 @@ public SampleContextObject() { public static void clean() { Context.set(null); ServiceConfigStore.updateFromAPIResponse(emptyAPIResponse); + StatisticsStore.clear(); + RoutesStore.clear(); + RateLimiterStore.clear(); }; @AfterEach public void tearDown() throws SQLException { Context.set(null); ServiceConfigStore.updateFromAPIResponse(emptyAPIResponse); + StatisticsStore.clear(); + RoutesStore.clear(); + RateLimiterStore.clear(); } @Test @@ -59,6 +68,7 @@ public void testNoContext() throws SQLException { // Test with thread cache not set : var res2 = ShouldBlockRequest.shouldBlockRequest(); assertFalse(res2.block()); + assertEquals(0, StatisticsStore.getStatsRecord().requests().rateLimited()); } @Test @@ -112,7 +122,8 @@ public void testUserSet() throws SQLException { @Test public void testEndpointsExistButNoMatch() throws SQLException { - Context.set(null); + ContextObject ctx = new SampleContextObject(); + Context.set(ctx); setEmptyConfigWithEndpointList(List.of( new Endpoint("POST", "/api2/*", 1, 1000, Collections.emptyList(), false, false, false) )); @@ -121,7 +132,6 @@ public void testEndpointsExistButNoMatch() throws SQLException { var res1 = ShouldBlockRequest.shouldBlockRequest(); assertFalse(res1.block()); - Context.set(null); setEmptyConfigWithEndpointList(List.of( new Endpoint("POST", "/api2/*", 1, 1000, Collections.emptyList(), false, false, true) )); @@ -133,7 +143,8 @@ public void testEndpointsExistButNoMatch() throws SQLException { @Test public void testEndpointsExistWithMatch() throws SQLException { - Context.set(null); + ContextObject ctx = new SampleContextObject(); + Context.set(ctx); setEmptyConfigWithEndpointList(List.of( new Endpoint("GET", "/api/*", 1, 1000, Collections.emptyList(), false, false, false) )); @@ -142,7 +153,6 @@ public void testEndpointsExistWithMatch() throws SQLException { var res1 = ShouldBlockRequest.shouldBlockRequest(); assertFalse(res1.block()); - Context.set(null); setEmptyConfigWithEndpointList(List.of( new Endpoint("GET", "/api/*", 1, 1000, Collections.emptyList(), false, false, true) )); @@ -150,6 +160,19 @@ public void testEndpointsExistWithMatch() throws SQLException { // Test with match & rate-limiting enabled : 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()); + assertEquals(2, RoutesStore.getRoutesAsList()[0].getRateLimitCount()); + } @Test From fdebc03a2e8b2434da989505e305f1a6a53a2cba Mon Sep 17 00:00:00 2001 From: BitterPanda63 Date: Tue, 24 Jun 2025 14:20:15 +0200 Subject: [PATCH 10/10] Fix RouteEntryTest case by addingg rateLimitedCount field --- agent_api/src/test/java/storage/RouteEntryTest.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/agent_api/src/test/java/storage/RouteEntryTest.java b/agent_api/src/test/java/storage/RouteEntryTest.java index b0878d8f..e77388dd 100644 --- a/agent_api/src/test/java/storage/RouteEntryTest.java +++ b/agent_api/src/test/java/storage/RouteEntryTest.java @@ -32,7 +32,7 @@ public void testGsonWithoutSerializer() throws IOException { Gson gson = new Gson(); String json = gson.toJson(route1); assertEquals( - "{\"method\":\"GET\",\"path\":\"/api/1\",\"hits\":0,\"apispec\":{\"body\":{\"schema\":{\"type\":\"object\",\"properties\":{\"oldProp\":{\"type\":\"string\",\"optional\":false}},\"optional\":false},\"type\":\"oldType\"},\"auth\":[{\"type\":\"apiKey\"}]}}", + "{\"method\":\"GET\",\"path\":\"/api/1\",\"hits\":0,\"rateLimitedCount\":0,\"apispec\":{\"body\":{\"schema\":{\"type\":\"object\",\"properties\":{\"oldProp\":{\"type\":\"string\",\"optional\":false}},\"optional\":false},\"type\":\"oldType\"},\"auth\":[{\"type\":\"apiKey\"}]}}", json ); }