Skip to content

Commit 7b98a0f

Browse files
bitterpanda63bitterpanda63
committed
Add excluded user ids & a check in the rate limit function
1 parent 27ea0d5 commit 7b98a0f

11 files changed

Lines changed: 124 additions & 28 deletions

File tree

‎agent_api/src/main/java/dev/aikido/agent_api/background/cloud/api/APIResponse.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ public record APIResponse(
1515
boolean blockNewOutgoingRequests,
1616
List<Domain> domains,
1717
boolean receivedAnyStats,
18-
boolean block
18+
boolean block,
19+
List<String> excludedUserIdsFromRateLimiting
1920
) {
2021
}

‎agent_api/src/main/java/dev/aikido/agent_api/background/cloud/api/ReportingApiHTTP.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -152,7 +152,7 @@ private static APIResponse getUnsuccessfulAPIResponse(String error) {
152152
return new APIResponse(
153153
false, // Success
154154
error,
155-
0, null, null, null, false, null, false, false // Unimportant values.
155+
0, null, null, null, false, null, false, false, null // Unimportant values.
156156
);
157157
}
158158
}

‎agent_api/src/main/java/dev/aikido/agent_api/ratelimiting/ShouldRateLimit.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,10 @@ public static RateLimitDecision shouldRateLimit(RouteMetadata routeMetadata, Use
2323
return NO_RATE_LIMIT;
2424
}
2525

26+
if (user != null && ServiceConfigStore.getConfig().isUserExcludedFromRateLimiting(user.id())) {
27+
return NO_RATE_LIMIT;
28+
}
29+
2630
long windowSizeInMS = rateLimitedEndpoint.getRateLimiting().windowSizeInMS();
2731
long maxRequests = rateLimitedEndpoint.getRateLimiting().maxRequests();
2832

‎agent_api/src/main/java/dev/aikido/agent_api/storage/ServiceConfiguration.java‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ public class ServiceConfiguration {
2424
private boolean middlewareInstalled;
2525
private IPList bypassedIPs = new IPList();
2626
private HashSet<String> blockedUserIDs = new HashSet<>();
27+
private HashSet<String> excludedUserIdsFromRateLimiting = new HashSet<>();
2728
private List<Endpoint> endpoints = new ArrayList<>();
2829
private OutboundDomains outboundDomains = new OutboundDomains();
2930

@@ -43,6 +44,9 @@ public void updateConfig(APIResponse apiResponse) {
4344
if (apiResponse.blockedUserIds() != null) {
4445
this.blockedUserIDs = new HashSet<>(apiResponse.blockedUserIds());
4546
}
47+
if (apiResponse.excludedUserIdsFromRateLimiting() != null) {
48+
this.excludedUserIdsFromRateLimiting = new HashSet<>(apiResponse.excludedUserIdsFromRateLimiting());
49+
}
4650
if (apiResponse.endpoints() != null) {
4751
this.endpoints = apiResponse.endpoints();
4852
}
@@ -78,6 +82,10 @@ public boolean isUserBlocked(String userId) {
7882
return this.blockedUserIDs.contains(userId);
7983
}
8084

85+
public boolean isUserExcludedFromRateLimiting(String userId) {
86+
return this.excludedUserIdsFromRateLimiting.contains(userId);
87+
}
88+
8189
public boolean isIpBypassed(String ip) {
8290
return this.bypassedIPs.matches(ip);
8391
}

‎agent_api/src/test/java/ShouldBlockRequestTest.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,7 @@ public void testUserSet() throws SQLException {
8787
ServiceConfigStore.updateFromAPIResponse(new APIResponse(
8888
true, "", getUnixTimeMS(), List.of(),
8989
/* blockedUserIds */ List.of("ID1", "ID2", "ID3"), List.of(),
90-
false, null, false, true
90+
false, null, false, true, List.of()
9191
));
9292
var res2 = ShouldBlockRequest.shouldBlockRequest();
9393
assertTrue(res2.block());
@@ -227,7 +227,7 @@ public void testBlockedUserWithMultipleEndpoints() throws SQLException {
227227
);
228228
List<String> blockedUserIds = List.of("ID1");
229229
ServiceConfigStore.updateFromAPIResponse(new APIResponse(
230-
true, "", getUnixTimeMS(), endpoints, blockedUserIds, List.of(), false, null,true, false
230+
true, "", getUnixTimeMS(), endpoints, blockedUserIds, List.of(), false, null,true, false, List.of()
231231
));
232232

233233
// Call the method

‎agent_api/src/test/java/collectors/DNSRecordCollectorTest.java‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ public void cleanup() {
4747
AttackQueue.clear();
4848
// Reset domain config
4949
ServiceConfigStore.updateFromAPIResponse(new APIResponse(
50-
true, null, 0L, null, null, null, false, List.of(), true, false
50+
true, null, 0L, null, null, null, false, List.of(), true, false, List.of()
5151
));
5252
}
5353

@@ -116,7 +116,7 @@ public void testHostnameSameWithContextAsAStoredSSRFAttack() {
116116
public void testBlockedDomain() {
117117
ServiceConfigStore.updateFromAPIResponse(new APIResponse(
118118
true, null, 0L, null, null, null,
119-
false, List.of(new Domain("blocked.example.com", "block")), true, true
119+
false, List.of(new Domain("blocked.example.com", "block")), true, true, List.of()
120120
));
121121
assertThrows(BlockedOutboundException.class, () ->
122122
DNSRecordCollector.report("blocked.example.com", new InetAddress[]{inetAddress1})
@@ -127,7 +127,7 @@ public void testBlockedDomain() {
127127
public void testAllowedDomainNotBlocked() {
128128
ServiceConfigStore.updateFromAPIResponse(new APIResponse(
129129
true, null, 0L, null, null, null,
130-
false, List.of(new Domain("allowed.example.com", "allow")), true, true
130+
false, List.of(new Domain("allowed.example.com", "allow")), true, true, List.of()
131131
));
132132
assertDoesNotThrow(() ->
133133
DNSRecordCollector.report("allowed.example.com", new InetAddress[]{inetAddress1})
@@ -138,7 +138,7 @@ public void testAllowedDomainNotBlocked() {
138138
public void testUnknownDomainBlockedWhenBlockNewOutgoingRequests() {
139139
ServiceConfigStore.updateFromAPIResponse(new APIResponse(
140140
true, null, 0L, null, null, null,
141-
true, List.of(), true, true
141+
true, List.of(), true, true, List.of()
142142
));
143143
assertThrows(BlockedOutboundException.class, () ->
144144
DNSRecordCollector.report("unknown.example.com", new InetAddress[]{inetAddress1})

‎agent_api/src/test/java/collectors/WebRequestCollectorTest.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -212,7 +212,7 @@ void testReport_userAgentBlocked_Ip_Bypassed() {
212212

213213
List<String> bypassedIps = List.of("192.168.1.1");
214214
ServiceConfigStore.updateFromAPIResponse(new APIResponse(
215-
true, "", getUnixTimeMS(), List.of(), List.of(), bypassedIps, false, null, true, false
215+
true, "", getUnixTimeMS(), List.of(), List.of(), bypassedIps, false, null, true, false, List.of()
216216
));
217217

218218

@@ -231,7 +231,7 @@ void testReport_ipBlockedUsingLists_Ip_Bypassed() {
231231

232232
List<String> bypassedIps = List.of("192.168.1.1");
233233
ServiceConfigStore.updateFromAPIResponse(new APIResponse(
234-
true, "", getUnixTimeMS(), List.of(), List.of(), bypassedIps, false, null, true, false
234+
true, "", getUnixTimeMS(), List.of(), List.of(), bypassedIps, false, null, true, false, List.of()
235235
));
236236

237237
WebRequestCollector.Res response = WebRequestCollector.report(contextObject);
@@ -251,7 +251,7 @@ void testReport_ipNotAllowedUsingLists_Ip_Bypassed() {
251251

252252
List<String> bypassedIps = List.of("192.168.1.1");
253253
ServiceConfigStore.updateFromAPIResponse(new APIResponse(
254-
true, "", getUnixTimeMS(), List.of(), List.of(), bypassedIps, false, null, true, false
254+
true, "", getUnixTimeMS(), List.of(), List.of(), bypassedIps, false, null, true, false, List.of()
255255
));
256256

257257

‎agent_api/src/test/java/ratelimiting/ShouldRateLimitTest.java‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -304,4 +304,40 @@ public void testRateLimitsByGroupIfUserIsNotSet() {
304304
assertEquals(new ShouldRateLimit.RateLimitDecision(true, "group"),
305305
ShouldRateLimit.shouldRateLimit(metadata, null, "group1", "4.3.2.1"));
306306
}
307+
308+
@Test
309+
public void testDoesNotRateLimitExcludedUsers() {
310+
List<Endpoint> endpoints = new ArrayList<>();
311+
endpoints.add(new Endpoint("POST", "/login",
312+
/*maxRequests*/ 3, /*windowSizeMS*/ 1000, List.of(),
313+
false, false, true));
314+
utils.EmptyAPIResponses.setEmptyConfigWithEndpointListAndExcludedUsers(
315+
endpoints, List.of("excluded-user-id"));
316+
317+
RouteMetadata routeMetadata = createRouteMetadata("POST", "/login");
318+
User excludedUser = new User("excluded-user-id", "John Doe", "1.1.1.1", 0);
319+
320+
for (int i = 0; i < 5; i++) {
321+
assertEquals(new ShouldRateLimit.RateLimitDecision(false, null),
322+
ShouldRateLimit.shouldRateLimit(routeMetadata, excludedUser, null, "1.2.3.4"));
323+
}
324+
}
325+
326+
@Test
327+
public void testDoesNotRateLimitExcludedUsersInGroup() {
328+
List<Endpoint> endpoints = new ArrayList<>();
329+
endpoints.add(new Endpoint("POST", "/login",
330+
/*maxRequests*/ 3, /*windowSizeMS*/ 1000, List.of(),
331+
false, false, true));
332+
utils.EmptyAPIResponses.setEmptyConfigWithEndpointListAndExcludedUsers(
333+
endpoints, List.of("excluded-user-id"));
334+
335+
RouteMetadata routeMetadata = createRouteMetadata("POST", "/login");
336+
User excludedUser = new User("excluded-user-id", "John Doe", "1.1.1.1", 0);
337+
338+
for (int i = 0; i < 5; i++) {
339+
assertEquals(new ShouldRateLimit.RateLimitDecision(false, null),
340+
ShouldRateLimit.shouldRateLimit(routeMetadata, excludedUser, "group1", "1.2.3.4"));
341+
}
342+
}
307343
}

‎agent_api/src/test/java/storage/ServiceConfigStoreTest.java‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ public void setUp() {
1717
// Reset to a known state: no domains, blockNewOutgoingRequests=false
1818
ServiceConfigStore.updateFromAPIResponse(new APIResponse(
1919
true, null, 0L, null, null, null,
20-
false, null, true, false
20+
false, null, true, false, null
2121
));
2222
}
2323

@@ -32,7 +32,7 @@ public void testShouldBlockOutgoingRequestBlockedDomain() {
3232
true, null, 0L, null, null, null,
3333
false,
3434
List.of(new Domain("blocked.com", "block")),
35-
true, false
35+
true, false, null
3636
));
3737
assertTrue(ServiceConfigStore.shouldBlockOutgoingRequest("blocked.com"));
3838
}
@@ -43,7 +43,7 @@ public void testShouldBlockOutgoingRequestAllowedDomain() {
4343
true, null, 0L, null, null, null,
4444
true,
4545
List.of(new Domain("allowed.com", "allow")),
46-
true, false
46+
true, false, null
4747
));
4848
assertFalse(ServiceConfigStore.shouldBlockOutgoingRequest("allowed.com"));
4949
}
@@ -54,7 +54,7 @@ public void testShouldBlockOutgoingRequestUnknownWhenBlockNewEnabled() {
5454
true, null, 0L, null, null, null,
5555
true,
5656
List.of(new Domain("allowed.com", "allow")),
57-
true, false
57+
true, false, null
5858
));
5959
assertTrue(ServiceConfigStore.shouldBlockOutgoingRequest("unknown.com"));
6060
}

‎agent_api/src/test/java/storage/ServiceConfigurationTest.java‎

Lines changed: 53 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,8 @@ public void testUpdateConfig() {
3838
false,
3939
null,
4040
true,
41-
true
41+
true,
42+
null
4243
);
4344

4445
serviceConfiguration.updateConfig(apiResponse);
@@ -68,7 +69,8 @@ public void testUpdateConfigWithUnsuccessfulResponse() {
6869
false,
6970
null,
7071
false,
71-
false
72+
false,
73+
null
7274
);
7375

7476
serviceConfiguration.updateConfig(apiResponse);
@@ -309,7 +311,8 @@ public void testReceivedAnyStats() {
309311
false,
310312
null,
311313
false,
312-
true
314+
true,
315+
null
313316
);
314317

315318
serviceConfiguration.updateConfig(apiResponse);
@@ -329,7 +332,8 @@ public void testIsIpBypassedWithEmptyBypassedList() {
329332
false,
330333
null,
331334
true,
332-
true
335+
true,
336+
null
333337
));
334338

335339
assertFalse(serviceConfiguration.isIpBypassed("192.168.1.1"));
@@ -347,7 +351,8 @@ public void testIsIpBypassedWithMultipleBypassedEntries() {
347351
false,
348352
null,
349353
true,
350-
true
354+
true,
355+
null
351356
);
352357

353358
serviceConfiguration.updateConfig(apiResponse);
@@ -369,12 +374,46 @@ public void testIsUserBlockedWithEmptyBlockedUserList() {
369374
false,
370375
null,
371376
true,
372-
true
377+
true,
378+
null
373379
));
374380

375381
assertFalse(serviceConfiguration.isUserBlocked("user1"));
376382
}
377383

384+
@Test
385+
public void testExcludedUserIdsFromRateLimiting() {
386+
// Initially empty
387+
assertFalse(serviceConfiguration.isUserExcludedFromRateLimiting("user1"));
388+
389+
// Populated via updateConfig
390+
serviceConfiguration.updateConfig(new APIResponse(
391+
true, null, 12345L, null, null, null,
392+
false, null, true, true,
393+
List.of("user1", "user2")
394+
));
395+
assertTrue(serviceConfiguration.isUserExcludedFromRateLimiting("user1"));
396+
assertTrue(serviceConfiguration.isUserExcludedFromRateLimiting("user2"));
397+
assertFalse(serviceConfiguration.isUserExcludedFromRateLimiting("user3"));
398+
399+
// Subsequent update replaces the set
400+
serviceConfiguration.updateConfig(new APIResponse(
401+
true, null, 12345L, null, null, null,
402+
false, null, true, true,
403+
List.of("user3")
404+
));
405+
assertFalse(serviceConfiguration.isUserExcludedFromRateLimiting("user1"));
406+
assertTrue(serviceConfiguration.isUserExcludedFromRateLimiting("user3"));
407+
408+
// Empty list clears all
409+
serviceConfiguration.updateConfig(new APIResponse(
410+
true, null, 12345L, null, null, null,
411+
false, null, true, true,
412+
Collections.emptyList()
413+
));
414+
assertFalse(serviceConfiguration.isUserExcludedFromRateLimiting("user3"));
415+
}
416+
378417
@Test
379418
public void testIsUserBlockedWithMultipleBlockedUsers() {
380419
APIResponse apiResponse = new APIResponse(
@@ -387,7 +426,8 @@ public void testIsUserBlockedWithMultipleBlockedUsers() {
387426
false,
388427
null,
389428
true,
390-
true
429+
true,
430+
null
391431
);
392432

393433
serviceConfiguration.updateConfig(apiResponse);
@@ -408,7 +448,8 @@ public void testGetEndpoints() {
408448
false,
409449
null,
410450
true,
411-
true
451+
true,
452+
null
412453
);
413454

414455
serviceConfiguration.updateConfig(apiResponse);
@@ -428,7 +469,8 @@ public void testGetEndpointsWithEmptyList() {
428469
false,
429470
null,
430471
true,
431-
true
472+
true,
473+
null
432474
);
433475

434476
serviceConfiguration.updateConfig(apiResponse);
@@ -573,7 +615,7 @@ public void testShouldBlockOutgoingRequest() {
573615
true, null, 0L, null, null, null,
574616
true,
575617
List.of(new Domain("example.com", "block"), new Domain("allowed.com", "allow")),
576-
true, true
618+
true, true, null
577619
);
578620
serviceConfiguration.updateConfig(apiResponse);
579621

@@ -588,7 +630,7 @@ public void testShouldBlockOutgoingRequest() {
588630
true, null, 0L, null, null, null,
589631
false,
590632
List.of(new Domain("example.com", "block"), new Domain("allowed.com", "allow")),
591-
true, true
633+
true, true, null
592634
);
593635
serviceConfiguration.updateConfig(apiResponse2);
594636

0 commit comments

Comments
 (0)