diff --git a/README.md b/README.md index ff89b0252..84f4e9053 100644 --- a/README.md +++ b/README.md @@ -73,7 +73,7 @@ To activate Zen you then just have to add the following `-javaagent` to your Jav ``` java -javaagent:/opt/zen/agent.jar -jar build/myapp.jar ``` -To use user-blocking and/or rate-limiting features, you will have to include the following Jarfile into your repository +To use user blocking, rate limiting, or custom event tracking, include `agent_api.jar` in your application. ### Gradle Add the following code to your `build.gradle` file. ```gradle @@ -130,7 +130,8 @@ See [Reporting to Aikido](#reporting-to-your-aikido-security-dashboard) to learn ## Additional configuration -[Configure Zen using environment variables for authentication, mode settings, debugging, and more.](https://help.aikido.dev/doc/configuration-via-env-vars/docrSItUkeR9) +* [Track custom events](./docs/track.md): trigger events that Playbooks can act on +* [Configure Zen using environment variables for authentication, mode settings, debugging, and more](https://help.aikido.dev/doc/configuration-via-env-vars/docrSItUkeR9) ## License diff --git a/agent_api/src/main/java/dev/aikido/agent_api/Track.java b/agent_api/src/main/java/dev/aikido/agent_api/Track.java new file mode 100644 index 000000000..d3be8ac73 --- /dev/null +++ b/agent_api/src/main/java/dev/aikido/agent_api/Track.java @@ -0,0 +1,86 @@ +package dev.aikido.agent_api; + +import dev.aikido.agent_api.background.cloud.api.events.CustomEvent; +import dev.aikido.agent_api.context.Context; +import dev.aikido.agent_api.context.ContextObject; +import dev.aikido.agent_api.context.SpringWebfluxContextObject; +import dev.aikido.agent_api.helpers.env.Token; +import dev.aikido.agent_api.helpers.logging.LogManager; +import dev.aikido.agent_api.helpers.logging.Logger; +import dev.aikido.agent_api.storage.AttackQueue; +import java.util.concurrent.atomic.AtomicBoolean; + +public final class Track { + private Track() {} + + private static final Logger logger = LogManager.getLogger(Track.class); + private static final int MAX_EVENTS_PER_REQUEST = 25; + private static final AtomicBoolean loggedWarningTrackCalledWithoutContext = new AtomicBoolean(false); + private static final AtomicBoolean loggedWarningWebfluxNotSupported = new AtomicBoolean(false); + + /** + * External function for applications to track a custom event, e.g. a + * failed login or a signup. Only works inside an HTTP request. + */ + public static void track(String eventName) { + if (eventName == null || eventName.isEmpty()) { + logger.info("track(...) expects a non-empty string as event name."); + return; + } + + ContextObject currentContext = Context.get(); + if (currentContext == null) { + if (!Context.isBypassed()) { + logWarningTrackCalledWithoutContext(); + } + return; + } + + // WebFlux's shared threads can hold another request's context, so we'd send the wrong IP and user. + if (currentContext instanceof SpringWebfluxContextObject) { + logWarningWebfluxNotSupported(); + return; + } + + // No reporting worker runs without a token, so queued events would remain in memory. + if (Token.fromEnv() == null) { + return; + } + + int trackedEventCount = currentContext.incrementTrackedEventCount(); + if (trackedEventCount > MAX_EVENTS_PER_REQUEST) { + if (trackedEventCount == MAX_EVENTS_PER_REQUEST + 1) { + logger.warn( + "track(...) was called more than %d times during this request. " + + "Only the first %d events were tracked.", + MAX_EVENTS_PER_REQUEST, MAX_EVENTS_PER_REQUEST); + } + return; + } + + AttackQueue.add(CustomEvent.createAPIEvent(eventName, currentContext)); + } + + private static void logWarningTrackCalledWithoutContext() { + if (!loggedWarningTrackCalledWithoutContext.compareAndSet(false, true)) { + return; + } + logger.warn("track(...) was called without a context. The event will not be tracked. " + + "Make sure to call track(...) within an HTTP request."); + } + + private static void logWarningWebfluxNotSupported() { + if (!loggedWarningWebfluxNotSupported.compareAndSet(false, true)) { + return; + } + logger.warn("track(...) is not supported for Spring WebFlux yet. The event will not be tracked."); + } + + /** + * Resets internal warning state. Only intended for use in tests. + */ + static void reset() { + loggedWarningTrackCalledWithoutContext.set(false); + loggedWarningWebfluxNotSupported.set(false); + } +} diff --git a/agent_api/src/main/java/dev/aikido/agent_api/background/cloud/api/events/CustomEvent.java b/agent_api/src/main/java/dev/aikido/agent_api/background/cloud/api/events/CustomEvent.java new file mode 100644 index 000000000..53725a0ec --- /dev/null +++ b/agent_api/src/main/java/dev/aikido/agent_api/background/cloud/api/events/CustomEvent.java @@ -0,0 +1,49 @@ +package dev.aikido.agent_api.background.cloud.api.events; + +import static dev.aikido.agent_api.background.cloud.GetManagerInfo.getManagerInfo; +import static dev.aikido.agent_api.helpers.UnixTimeMS.getUnixTimeMS; + +import dev.aikido.agent_api.background.cloud.GetManagerInfo; +import dev.aikido.agent_api.context.ContextObject; +import dev.aikido.agent_api.context.User; + +public final class CustomEvent { + private CustomEvent() {} + + public record RequestData( + // note that URL is not included in the request metadata + String method, String ipAddress, String userAgent, String source, String route) {} + + public record UserData(String id, String name) {} + + public record CustomEventEvent( + String type, String name, RequestData request, GetManagerInfo.ManagerInfo agent, UserData user, long time) + implements APIEvent {} + + public static CustomEventEvent createAPIEvent(String eventName, ContextObject context) { + return new CustomEventEvent( + "custom", + eventName, + buildRequestData(context), + getManagerInfo(), + buildUserData(context), + getUnixTimeMS()); + } + + private static UserData buildUserData(ContextObject context) { + if (context.getUser() == null) { + return null; + } + User user = context.getUser(); + return new UserData(user.id(), user.name()); + } + + private static RequestData buildRequestData(ContextObject context) { + return new RequestData( + context.getMethod(), + context.getRemoteAddress(), + context.getHeader("user-agent"), + context.getSource(), + context.getRoute()); + } +} diff --git a/agent_api/src/main/java/dev/aikido/agent_api/context/ContextObject.java b/agent_api/src/main/java/dev/aikido/agent_api/context/ContextObject.java index 339358384..787432230 100644 --- a/agent_api/src/main/java/dev/aikido/agent_api/context/ContextObject.java +++ b/agent_api/src/main/java/dev/aikido/agent_api/context/ContextObject.java @@ -24,6 +24,7 @@ public class ContextObject { protected transient ArrayList redirectStartNodes; protected transient Map> cache = new HashMap<>(); protected transient Optional forcedProtectionOff = Optional.empty(); + protected transient int trackedEventCount = 0; public boolean middlewareExecuted() {return executedMiddleware; } public void setExecutedMiddleware(boolean value) { executedMiddleware = value; } @@ -104,6 +105,14 @@ public RouteMetadata getRouteMetadata() { return new RouteMetadata(route, url, method); } + /** + * Increments the number of custom events tracked during this request and + * returns the new count. Used to enforce a per-request limit on track(...) calls. + */ + public int incrementTrackedEventCount() { + return ++trackedEventCount; + } + @Override public String toString() { return "ContextObject{" + diff --git a/agent_api/src/main/java/dev/aikido/agent_api/storage/AttackQueue.java b/agent_api/src/main/java/dev/aikido/agent_api/storage/AttackQueue.java index 8bde5a36e..94dfbad04 100644 --- a/agent_api/src/main/java/dev/aikido/agent_api/storage/AttackQueue.java +++ b/agent_api/src/main/java/dev/aikido/agent_api/storage/AttackQueue.java @@ -14,7 +14,9 @@ public final class AttackQueue { private static final Logger logger = LogManager.getLogger(AttackQueue.class); - private static final BlockingQueue queue = new LinkedBlockingQueue<>(); + // Keeps memory bounded when our API is slow or down. add() drops new events when the queue is full. + private static final int MAX_QUEUE_SIZE = 100; + private static final BlockingQueue queue = new LinkedBlockingQueue<>(MAX_QUEUE_SIZE); private AttackQueue() { diff --git a/agent_api/src/test/java/background/cloud/api/CustomEventTest.java b/agent_api/src/test/java/background/cloud/api/CustomEventTest.java new file mode 100644 index 000000000..872838280 --- /dev/null +++ b/agent_api/src/test/java/background/cloud/api/CustomEventTest.java @@ -0,0 +1,52 @@ +package background.cloud.api; + +import static org.junit.jupiter.api.Assertions.*; + +import com.google.gson.Gson; +import com.google.gson.JsonObject; +import dev.aikido.agent_api.background.cloud.api.events.CustomEvent; +import dev.aikido.agent_api.context.ContextObject; +import dev.aikido.agent_api.context.User; +import java.util.Set; +import org.junit.jupiter.api.Test; +import utils.EmptySampleContextObject; + +class CustomEventTest { + + @Test + void createAPIEvent_WithValidContext_ReturnsCustomEventEvent() { + ContextObject context = new EmptySampleContextObject("test", "/api/resource", "POST"); + context.setUser(new User("user-1", "Jane Doe", "192.168.1.1", 1000L)); + + CustomEvent.CustomEventEvent event = CustomEvent.createAPIEvent("my-custom-event", context); + + assertNotNull(event); + assertEquals("custom", event.type()); + assertEquals("my-custom-event", event.name()); + assertNotNull(event.request()); + assertEquals("POST", event.request().method()); + assertEquals("web", event.request().source()); + assertEquals("/api/resource", event.request().route()); + assertEquals("192.168.1.1", event.request().ipAddress()); + assertNotNull(event.agent()); + assertNotNull(event.user()); + assertEquals("user-1", event.user().id()); + assertEquals("Jane Doe", event.user().name()); + assertTrue(event.time() > 0); + + JsonObject serializedUser = + new Gson().toJsonTree(event).getAsJsonObject().getAsJsonObject("user"); + assertEquals(Set.of("id", "name"), serializedUser.keySet()); + assertEquals("user-1", serializedUser.get("id").getAsString()); + assertEquals("Jane Doe", serializedUser.get("name").getAsString()); + } + + @Test + void createAPIEvent_WithContextButNoUser_ReturnsCustomEventEventWithNullUser() { + ContextObject context = new EmptySampleContextObject("test", "/api/resource", "GET"); + + CustomEvent.CustomEventEvent event = CustomEvent.createAPIEvent("my-custom-event", context); + + assertNull(event.user()); + } +} diff --git a/agent_api/src/test/java/dev/aikido/agent_api/TrackTest.java b/agent_api/src/test/java/dev/aikido/agent_api/TrackTest.java new file mode 100644 index 000000000..3fa196688 --- /dev/null +++ b/agent_api/src/test/java/dev/aikido/agent_api/TrackTest.java @@ -0,0 +1,155 @@ +package dev.aikido.agent_api; + +import static dev.aikido.agent_api.helpers.UnixTimeMS.getUnixTimeMS; +import static org.junit.jupiter.api.Assertions.*; +import static utils.EmptyAPIResponses.emptyAPIResponse; + +import dev.aikido.agent_api.background.cloud.api.APIResponse; +import dev.aikido.agent_api.background.cloud.api.events.APIEvent; +import dev.aikido.agent_api.background.cloud.api.events.CustomEvent; +import dev.aikido.agent_api.collectors.WebRequestCollector; +import dev.aikido.agent_api.context.Context; +import dev.aikido.agent_api.context.ContextObject; +import dev.aikido.agent_api.context.SpringWebfluxContextObject; +import dev.aikido.agent_api.storage.AttackQueue; +import dev.aikido.agent_api.storage.ServiceConfigStore; +import java.net.InetSocketAddress; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.junit.jupiter.api.*; +import org.junitpioneer.jupiter.ClearEnvironmentVariable; +import org.junitpioneer.jupiter.SetEnvironmentVariable; +import org.junitpioneer.jupiter.StdIo; +import org.junitpioneer.jupiter.StdOut; +import utils.EmptySampleContextObject; + +@SetEnvironmentVariable(key = "AIKIDO_LOG_LEVEL", value = "trace") +@SetEnvironmentVariable(key = "AIKIDO_TOKEN", value = "invalid-token-2") +public class TrackTest { + @BeforeEach + @AfterEach + public void resetState() { + Context.reset(); + AttackQueue.clear(); + ServiceConfigStore.updateFromAPIResponse(emptyAPIResponse); + Track.reset(); + } + + @Test + @StdIo + public void testTrackWithInvalidEventName(StdOut out) throws Exception { + Track.track(""); + assertTrue(out.capturedString().contains("expects a non-empty string as event name.")); + assertEquals(0, AttackQueue.getSize()); + } + + @Test + @StdIo + public void testTrackWithNullEventName(StdOut out) throws Exception { + Track.track(null); + assertTrue(out.capturedString().contains("expects a non-empty string as event name.")); + assertEquals(0, AttackQueue.getSize()); + } + + @Test + @StdIo + public void testTrackWithoutContext(StdOut out) throws Exception { + Track.track("my-event"); + assertTrue(out.capturedString().contains("track(...) was called without a context.")); + assertEquals(0, AttackQueue.getSize()); + } + + @Test + @StdIo + public void testTrackWithoutContextOnlyLogsOnce(StdOut out) throws Exception { + Track.track("my-event"); + Track.track("my-event"); + int occurrences = + out.capturedString().split("track\\(\\.\\.\\.\\) was called without a context\\.", -1).length - 1; + assertEquals(1, occurrences); + } + + @Test + @StdIo + public void testTrackForBypassedIpDoesNotQueueOrWarn(StdOut out) { + 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(new EmptySampleContextObject("test", "/track-me", "POST")); + + Track.track("my-custom-event"); + + assertEquals(0, AttackQueue.getSize()); + assertFalse(out.capturedString().contains("track(...) was called without a context.")); + } + + @Test + public void testTrackDoesNotQueueEventForWebfluxRequest() { + Context.set(new SpringWebfluxContextObject( + "POST", "http://localhost/login", new InetSocketAddress("1.2.3.4", 443), + new HashMap<>(), Map.of(), Map.of())); + + Track.track("my-custom-event"); + + assertEquals(0, AttackQueue.getSize()); + } + + @Test + @ClearEnvironmentVariable(key = "AIKIDO_TOKEN") + public void testTrackWithoutTokenDoesNotQueueEvent() { + ContextObject context = new EmptySampleContextObject("test", "/track-me", "POST"); + Context.set(context); + + Track.track("my-custom-event"); + + assertEquals(0, AttackQueue.getSize()); + } + + @Test + public void testTrackSendsEventToQueue() throws InterruptedException { + ContextObject context = new EmptySampleContextObject("test", "/track-me", "POST"); + Context.set(context); + + Track.track("my-custom-event"); + + assertEquals(1, AttackQueue.getSize()); + APIEvent event = AttackQueue.get(); + assertInstanceOf(CustomEvent.CustomEventEvent.class, event); + CustomEvent.CustomEventEvent customEvent = (CustomEvent.CustomEventEvent) event; + assertEquals("custom", customEvent.type()); + assertEquals("my-custom-event", customEvent.name()); + assertEquals("POST", customEvent.request().method()); + assertEquals("/track-me", customEvent.request().route()); + } + + @Test + public void testTrackAllowsUpToLimitPerRequest() { + ContextObject context = new EmptySampleContextObject("test", "/track-me", "POST"); + Context.set(context); + + for (int i = 0; i < 25; i++) { + Track.track("event-" + i); + } + + assertEquals(25, AttackQueue.getSize()); + } + + @Test + @StdIo + public void testTrackDropsEventsOverLimitAndLogsWarningOnlyOnce(StdOut out) { + ContextObject context = new EmptySampleContextObject("test", "/track-me", "POST"); + Context.set(context); + + for (int i = 0; i < 30; i++) { + Track.track("event-" + i); + } + + assertEquals(25, AttackQueue.getSize()); + long warningCount = out.capturedString() + .lines() + .filter(line -> line.contains("Only the first 25 events were tracked.")) + .count(); + assertEquals(1, warningCount); + } +} diff --git a/agent_api/src/test/java/storage/AttackQueueTest.java b/agent_api/src/test/java/storage/AttackQueueTest.java new file mode 100644 index 000000000..6d6258d98 --- /dev/null +++ b/agent_api/src/test/java/storage/AttackQueueTest.java @@ -0,0 +1,25 @@ +package storage; + +import dev.aikido.agent_api.background.cloud.api.events.APIEvent; +import dev.aikido.agent_api.storage.AttackQueue; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +public class AttackQueueTest { + @BeforeEach + @AfterEach + public void clearQueue() { + AttackQueue.clear(); + } + + @Test + public void testDropsEventsWhenFull() { + for (int i = 0; i < 150; i++) { + AttackQueue.add(new APIEvent() {}); + } + assertEquals(100, AttackQueue.getSize()); + } +} diff --git a/docs/track.md b/docs/track.md new file mode 100644 index 000000000..09bf173eb --- /dev/null +++ b/docs/track.md @@ -0,0 +1,18 @@ +# Track custom events + +Use `Track.track()` to report events that only your application knows about, such as failed logins. [Playbooks](https://help.aikido.dev/zen-firewall/zen-features/playbooks) can act when an event occurs repeatedly, for example by blocking an IP after three failed logins in five minutes. + +```java +import dev.aikido.agent_api.Track; + +if (!authenticate(request)) { + Track.track("user.login_failed"); + throw new UnauthorizedException(); +} +``` + +After adding `Track.track()`, trigger the event at least once. It will then appear on the Playbooks page in the Aikido dashboard. From there, you can create a playbook and choose what should happen when the event occurs. Calling `Track.track()` by itself does not create a playbook or block anything. + +Call `Track.track()` while handling an HTTP request. Zen associates the event with the request's IP address. Playbook counts are per IP, not across your whole app. If you call [`SetUser.setUser()`](./user.md) before `Track.track()`, Zen also includes the current user. `SetUser.setUser()` is optional. Events without a user are still tracked. + +Event names can use any format. We recommend lowercase, dot-separated names such as `user.login_failed`. diff --git a/end2end/spring_boot_postgres.py b/end2end/spring_boot_postgres.py index a5ce1100a..d201e0a99 100644 --- a/end2end/spring_boot_postgres.py +++ b/end2end/spring_boot_postgres.py @@ -1,4 +1,5 @@ from utils import App, Request +from utils.test_custom_events import test_custom_events spring_boot_postgres_app = App(8080) @@ -21,3 +22,6 @@ ) spring_boot_postgres_app.test_all_payloads() + +test_custom_events(spring_boot_postgres_app.urls["enabled"], spring_boot_postgres_app.event_handler) +print("✅ Tested custom events") diff --git a/end2end/utils/test_custom_events.py b/end2end/utils/test_custom_events.py new file mode 100644 index 000000000..04e3cb0e7 --- /dev/null +++ b/end2end/utils/test_custom_events.py @@ -0,0 +1,33 @@ +import time +import requests +from utils.assert_equals import assert_eq +from utils.EventHandler import filter_on_event_type + +def wait_for_custom_events(event_handler, timeout=10): + deadline = time.time() + timeout + while time.time() < deadline: + custom_events = filter_on_event_type(event_handler.fetch_events_from_mock(), "custom") + if custom_events: + return custom_events + time.sleep(0.5) + return [] + +def test_custom_events(url, event_handler): + event_handler.reset() + + res = requests.post(url + "/api/track", headers={ + "X-Forwarded-For": "203.0.113.7", + "User-Agent": "zen-e2e", + }) + assert_eq(res.status_code, equals=200) + + custom_events = wait_for_custom_events(event_handler) + assert_eq(len(custom_events), equals=1) + event = custom_events[0] + assert_eq(event["name"], equals="user.login_failed") + assert_eq(event["request"]["ipAddress"], equals="203.0.113.7") + assert_eq(event["request"]["userAgent"], equals="zen-e2e") + assert_eq(event["request"]["method"], equals="POST") + assert_eq(event["request"]["route"], equals="/api/track") + assert_eq(event["agent"]["library"], equals="firewall-java") + assert "user" not in event diff --git a/sample-apps/SpringBootPostgres/src/main/java/com/example/demo/TrackController.java b/sample-apps/SpringBootPostgres/src/main/java/com/example/demo/TrackController.java new file mode 100644 index 000000000..117b55917 --- /dev/null +++ b/sample-apps/SpringBootPostgres/src/main/java/com/example/demo/TrackController.java @@ -0,0 +1,14 @@ +package com.example.demo; + +import dev.aikido.agent_api.Track; +import org.springframework.web.bind.annotation.*; + +@RestController +@RequestMapping("/api/track") +public class TrackController { + @PostMapping + public String track() { + Track.track("user.login_failed"); + return "OK"; + } +}