Add track() API for custom event tracking

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
 

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);
+    }
+}

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());
+    }
+}

agent_api/src/main/java/dev/aikido/agent_api/context/ContextObject.java

@@ -24,6 +24,7 @@ public class ContextObject {
     protected transient ArrayList<RedirectNode> redirectStartNodes;
     protected transient Map<String, Map<String, String>> cache = new HashMap<>();
     protected transient Optional<Boolean> 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{" +

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<APIEvent> 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<APIEvent> queue = new LinkedBlockingQueue<>(MAX_QUEUE_SIZE);
 
     private AttackQueue() {
 

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());
+    }
+}

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<String> 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);
+    }
+}