diff --git a/agent_api/src/main/java/dev/aikido/agent_api/SetRateLimitGroup.java b/agent_api/src/main/java/dev/aikido/agent_api/SetRateLimitGroup.java index 7ef31cac4..c1f65edef 100644 --- a/agent_api/src/main/java/dev/aikido/agent_api/SetRateLimitGroup.java +++ b/agent_api/src/main/java/dev/aikido/agent_api/SetRateLimitGroup.java @@ -17,6 +17,9 @@ public static void setRateLimitGroup(String groupId) { ContextObject currentContext = Context.get(); if (currentContext == null) { + if (Context.isBypassed()) { + return; + } logger.warn("setRateLimitGroup(...) was called without a context. Make sure to call setRateLimitGroup(...) within an HTTP request."); return; } diff --git a/agent_api/src/main/java/dev/aikido/agent_api/collectors/WebRequestCollector.java b/agent_api/src/main/java/dev/aikido/agent_api/collectors/WebRequestCollector.java index 1d66a2126..705635775 100644 --- a/agent_api/src/main/java/dev/aikido/agent_api/collectors/WebRequestCollector.java +++ b/agent_api/src/main/java/dev/aikido/agent_api/collectors/WebRequestCollector.java @@ -46,6 +46,7 @@ public static Res report(ContextObject newContext) { PendingHostnamesStore.clear(); if (config.isIpBypassed(newContext.getRemoteAddress())) { + Context.markBypassed(); return null; // do not set context when the IP address is bypassed (zen = off) } diff --git a/agent_api/src/main/java/dev/aikido/agent_api/context/Context.java b/agent_api/src/main/java/dev/aikido/agent_api/context/Context.java index 0389cc903..fc2521401 100644 --- a/agent_api/src/main/java/dev/aikido/agent_api/context/Context.java +++ b/agent_api/src/main/java/dev/aikido/agent_api/context/Context.java @@ -6,6 +6,8 @@ public final class Context { private Context() {} static final ThreadLocal threadLocalContext = new ThreadLocal<>(); + // Bypassed requests have no context, so this flag tells them apart from code running outside a request. + static final ThreadLocal threadLocalBypassed = ThreadLocal.withInitial(() -> false); public static ContextObject get() { return threadLocalContext.get(); } @@ -14,5 +16,12 @@ public static void set(ContextObject contextObject) { } public static void reset() { threadLocalContext.remove(); + threadLocalBypassed.remove(); + } + public static void markBypassed() { + threadLocalBypassed.set(true); + } + public static boolean isBypassed() { + return threadLocalBypassed.get(); } } diff --git a/agent_api/src/test/java/SetRateLimitGroupTest.java b/agent_api/src/test/java/SetRateLimitGroupTest.java index d82bbd13f..624676770 100644 --- a/agent_api/src/test/java/SetRateLimitGroupTest.java +++ b/agent_api/src/test/java/SetRateLimitGroupTest.java @@ -39,7 +39,7 @@ public static void clean() { @AfterEach public void tearDown() throws SQLException { - Context.set(null); + Context.reset(); } @Test @@ -58,6 +58,16 @@ public void testWithNoContext(StdOut out) throws SQLException, IOException { assertTrue(out.capturedString().contains("setRateLimitGroup(...) was called without a context. Make sure to call setRateLimitGroup(...) within an HTTP request.")); } + @Test + @StdIo + public void testWithBypassedRequestDoesNotWarn(StdOut out) { + Context.markBypassed(); + + SetRateLimitGroup.setRateLimitGroup("id"); + + assertFalse(out.capturedString().contains("was called without a context")); + } + @Test @StdIo public void testEmptyGroupId(StdOut out) throws SQLException { diff --git a/agent_api/src/test/java/collectors/WebRequestCollectorTest.java b/agent_api/src/test/java/collectors/WebRequestCollectorTest.java index d47e639a9..6ff97791a 100644 --- a/agent_api/src/test/java/collectors/WebRequestCollectorTest.java +++ b/agent_api/src/test/java/collectors/WebRequestCollectorTest.java @@ -49,6 +49,23 @@ void testReport_contextSet() { } + @Test + void testReport_bypassedIpIsMarkedUntilNextRequest() { + List bypassedIps = List.of("192.168.1.1"); + ServiceConfigStore.updateFromAPIResponse(new APIResponse( + true, "", getUnixTimeMS(), List.of(), List.of(), bypassedIps, false, null, true, false, List.of(), 0 + )); + + WebRequestCollector.report(contextObject); + assertNull(Context.get()); + assertTrue(Context.isBypassed()); + + contextObject.setIp("10.0.0.1"); + WebRequestCollector.report(contextObject); + assertNotNull(Context.get()); + assertFalse(Context.isBypassed()); + } + @Test void testReport_ipNotAllowed() { setEmptyConfigWithEndpointList(List.of(new Endpoint(