Skip to content

Commit b616b2b

Browse files
timokoesslerhansott
timokoessler
authored and
hansott
committed
Fix review comments
1 parent c8e5c3d commit b616b2b

6 files changed

Lines changed: 92 additions & 89 deletions

File tree

‎README.md‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,7 @@ To activate Zen you then just have to add the following `-javaagent` to your Jav
7373
```
7474
java -javaagent:/opt/zen/agent.jar -jar build/myapp.jar
7575
```
76-
To use user-blocking and/or rate-limiting features, you will have to include the following Jarfile into your repository
76+
To use user blocking, rate limiting, or custom event tracking, include `agent_api.jar` in your application.
7777
### Gradle
7878
Add the following code to your `build.gradle` file.
7979
```gradle
@@ -130,7 +130,8 @@ See [Reporting to Aikido](#reporting-to-your-aikido-security-dashboard) to learn
130130

131131
## Additional configuration
132132

133-
[Configure Zen using environment variables for authentication, mode settings, debugging, and more.](https://help.aikido.dev/doc/configuration-via-env-vars/docrSItUkeR9)
133+
* [Track custom events](./docs/track.md): trigger events that Playbooks can act on
134+
* [Configure Zen using environment variables for authentication, mode settings, debugging, and more](https://help.aikido.dev/doc/configuration-via-env-vars/docrSItUkeR9)
134135

135136
## License
136137

‎agent_api/src/main/java/dev/aikido/agent_api/Track.java‎

Lines changed: 21 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -3,15 +3,18 @@
33
import dev.aikido.agent_api.background.cloud.api.events.CustomEvent;
44
import dev.aikido.agent_api.context.Context;
55
import dev.aikido.agent_api.context.ContextObject;
6+
import dev.aikido.agent_api.helpers.env.Token;
67
import dev.aikido.agent_api.helpers.logging.LogManager;
78
import dev.aikido.agent_api.helpers.logging.Logger;
89
import dev.aikido.agent_api.storage.AttackQueue;
10+
import java.util.concurrent.atomic.AtomicBoolean;
911

1012
public final class Track {
1113
private Track() {}
14+
1215
private static final Logger logger = LogManager.getLogger(Track.class);
1316
private static final int MAX_EVENTS_PER_REQUEST = 25;
14-
private static boolean loggedWarningTrackCalledWithoutContext = false;
17+
private static final AtomicBoolean loggedWarningTrackCalledWithoutContext = new AtomicBoolean(false);
1518

1619
/**
1720
* External function for applications to track a custom event, e.g. a
@@ -29,32 +32,37 @@ public static void track(String eventName) {
2932
return;
3033
}
3134

32-
if (currentContext.incrementTrackedEventCount() > MAX_EVENTS_PER_REQUEST) {
33-
logger.warn(
34-
"track(...) was called more than %d times during this request. Dropping event: %s",
35-
MAX_EVENTS_PER_REQUEST, eventName
36-
);
35+
// No reporting worker runs without a token, so queued events would remain in memory.
36+
if (Token.fromEnv() == null) {
37+
return;
38+
}
39+
40+
int trackedEventCount = currentContext.incrementTrackedEventCount();
41+
if (trackedEventCount > MAX_EVENTS_PER_REQUEST) {
42+
if (trackedEventCount == MAX_EVENTS_PER_REQUEST + 1) {
43+
logger.warn(
44+
"track(...) was called more than %d times during this request. "
45+
+ "Only the first %d events were tracked.",
46+
MAX_EVENTS_PER_REQUEST, MAX_EVENTS_PER_REQUEST);
47+
}
3748
return;
3849
}
3950

4051
AttackQueue.add(CustomEvent.createAPIEvent(eventName, currentContext));
4152
}
4253

4354
private static void logWarningTrackCalledWithoutContext() {
44-
if (loggedWarningTrackCalledWithoutContext) {
55+
if (!loggedWarningTrackCalledWithoutContext.compareAndSet(false, true)) {
4556
return;
4657
}
47-
logger.warn(
48-
"track(...) was called without a context. The event will not be tracked. " +
49-
"Make sure to call track(...) within an HTTP request."
50-
);
51-
loggedWarningTrackCalledWithoutContext = true;
58+
logger.warn("track(...) was called without a context. The event will not be tracked. "
59+
+ "Make sure to call track(...) within an HTTP request.");
5260
}
5361

5462
/**
5563
* Resets internal warning state. Only intended for use in tests.
5664
*/
5765
public static void reset() {
58-
loggedWarningTrackCalledWithoutContext = false;
66+
loggedWarningTrackCalledWithoutContext.set(false);
5967
}
6068
}

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

Lines changed: 29 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -1,53 +1,52 @@
11
package dev.aikido.agent_api.background.cloud.api.events;
22

3+
import static dev.aikido.agent_api.background.cloud.GetManagerInfo.getManagerInfo;
4+
import static dev.aikido.agent_api.helpers.UnixTimeMS.getUnixTimeMS;
5+
36
import dev.aikido.agent_api.background.cloud.GetManagerInfo;
47
import dev.aikido.agent_api.context.ContextObject;
58
import dev.aikido.agent_api.context.User;
69

7-
import static dev.aikido.agent_api.background.cloud.GetManagerInfo.getManagerInfo;
8-
import static dev.aikido.agent_api.helpers.UnixTimeMS.getUnixTimeMS;
9-
1010
public final class CustomEvent {
1111
private CustomEvent() {}
12+
1213
public record RequestData(
13-
// note that URL is not included in the request metadata
14-
String method,
15-
String ipAddress,
16-
String userAgent,
17-
String source,
18-
String route
19-
) {}
14+
// note that URL is not included in the request metadata
15+
String method, String ipAddress, String userAgent, String source, String route) {}
16+
17+
public record UserData(String id, String name) {}
2018

2119
public record CustomEventEvent(
22-
String type,
23-
String name,
24-
RequestData request,
25-
GetManagerInfo.ManagerInfo agent,
26-
User user,
27-
long time
28-
) implements APIEvent {}
20+
String type, String name, RequestData request, GetManagerInfo.ManagerInfo agent, UserData user, long time)
21+
implements APIEvent {}
2922

3023
public static CustomEventEvent createAPIEvent(String eventName, ContextObject context) {
3124
return new CustomEventEvent(
32-
"custom", // type
33-
eventName, // name
34-
buildRequestData(context), // request
35-
getManagerInfo(), // agent
36-
context != null ? context.getUser() : null, // user
37-
getUnixTimeMS() // time
38-
);
25+
"custom",
26+
eventName,
27+
buildRequestData(context),
28+
getManagerInfo(),
29+
buildUserData(context),
30+
getUnixTimeMS());
31+
}
32+
33+
private static UserData buildUserData(ContextObject context) {
34+
if (context == null || context.getUser() == null) {
35+
return null;
36+
}
37+
User user = context.getUser();
38+
return new UserData(user.id(), user.name());
3939
}
4040

4141
private static RequestData buildRequestData(ContextObject context) {
4242
if (context == null) {
4343
return null;
4444
}
4545
return new RequestData(
46-
context.getMethod(),
47-
context.getRemoteAddress(),
48-
context.getHeader("user-agent"),
49-
context.getSource(),
50-
context.getRoute()
51-
);
46+
context.getMethod(),
47+
context.getRemoteAddress(),
48+
context.getHeader("user-agent"),
49+
context.getSource(),
50+
context.getRoute());
5251
}
5352
}

‎agent_api/src/test/java/TrackTest.java‎

Lines changed: 24 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,18 @@
1+
import static org.junit.jupiter.api.Assertions.*;
2+
13
import dev.aikido.agent_api.Track;
24
import dev.aikido.agent_api.background.cloud.api.events.APIEvent;
35
import dev.aikido.agent_api.background.cloud.api.events.CustomEvent;
46
import dev.aikido.agent_api.context.Context;
57
import dev.aikido.agent_api.context.ContextObject;
68
import dev.aikido.agent_api.storage.AttackQueue;
79
import org.junit.jupiter.api.*;
10+
import org.junitpioneer.jupiter.ClearEnvironmentVariable;
811
import org.junitpioneer.jupiter.SetEnvironmentVariable;
912
import org.junitpioneer.jupiter.StdIo;
1013
import org.junitpioneer.jupiter.StdOut;
1114
import utils.EmptySampleContextObject;
1215

13-
import static org.junit.jupiter.api.Assertions.*;
14-
1516
@SetEnvironmentVariable(key = "AIKIDO_LOG_LEVEL", value = "trace")
1617
@SetEnvironmentVariable(key = "AIKIDO_TOKEN", value = "invalid-token-2")
1718
public class TrackTest {
@@ -58,10 +59,22 @@ public void testTrackWithoutContext(StdOut out) throws Exception {
5859
public void testTrackWithoutContextOnlyLogsOnce(StdOut out) throws Exception {
5960
Track.track("my-event");
6061
Track.track("my-event");
61-
int occurrences = out.capturedString().split("track\\(\\.\\.\\.\\) was called without a context\\.", -1).length - 1;
62+
int occurrences =
63+
out.capturedString().split("track\\(\\.\\.\\.\\) was called without a context\\.", -1).length - 1;
6264
assertEquals(1, occurrences);
6365
}
6466

67+
@Test
68+
@ClearEnvironmentVariable(key = "AIKIDO_TOKEN")
69+
public void testTrackWithoutTokenDoesNotQueueEvent() {
70+
ContextObject context = new EmptySampleContextObject("test", "/track-me", "POST");
71+
Context.set(context);
72+
73+
Track.track("my-custom-event");
74+
75+
assertEquals(0, AttackQueue.getSize());
76+
}
77+
6578
@Test
6679
public void testTrackSendsEventToQueue() throws InterruptedException {
6780
ContextObject context = new EmptySampleContextObject("test", "/track-me", "POST");
@@ -93,15 +106,20 @@ public void testTrackAllowsUpToLimitPerRequest() {
93106

94107
@Test
95108
@StdIo
96-
public void testTrackDropsEventsOverLimitPerRequest(StdOut out) {
109+
public void testTrackDropsEventsOverLimitAndLogsWarningOnlyOnce(StdOut out) {
97110
ContextObject context = new EmptySampleContextObject("test", "/track-me", "POST");
98111
Context.set(context);
99112

100-
for (int i = 0; i < 26; i++) {
113+
for (int i = 0; i < 30; i++) {
101114
Track.track("event-" + i);
102115
}
103116

104117
assertEquals(25, AttackQueue.getSize());
105-
assertTrue(out.capturedString().contains("Dropping event: event-25"));
118+
long warningCount = out.capturedString()
119+
.lines()
120+
.filter(line -> line.contains("Only the first 25 events were tracked."))
121+
.count();
122+
assertEquals(1, warningCount);
123+
assertFalse(out.capturedString().contains("event-25"));
106124
}
107125
}

‎agent_api/src/test/java/background/cloud/api/CustomEventTest.java‎

Lines changed: 11 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,25 +1,25 @@
11
package background.cloud.api;
22

3+
import static org.junit.jupiter.api.Assertions.*;
4+
5+
import com.google.gson.Gson;
6+
import com.google.gson.JsonObject;
37
import dev.aikido.agent_api.background.cloud.api.events.CustomEvent;
48
import dev.aikido.agent_api.context.ContextObject;
59
import dev.aikido.agent_api.context.User;
10+
import java.util.Set;
611
import org.junit.jupiter.api.Test;
712
import utils.EmptySampleContextObject;
813

9-
import static org.junit.jupiter.api.Assertions.*;
10-
1114
class CustomEventTest {
1215

1316
@Test
1417
void createAPIEvent_WithValidContext_ReturnsCustomEventEvent() {
15-
// Arrange
1618
ContextObject context = new EmptySampleContextObject("test", "/api/resource", "POST");
1719
context.setUser(new User("user-1", "Jane Doe", "192.168.1.1", 1000L));
1820

19-
// Act
2021
CustomEvent.CustomEventEvent event = CustomEvent.createAPIEvent("my-custom-event", context);
2122

22-
// Assert
2323
assertNotNull(event);
2424
assertEquals("custom", event.type());
2525
assertEquals("my-custom-event", event.name());
@@ -33,14 +33,18 @@ void createAPIEvent_WithValidContext_ReturnsCustomEventEvent() {
3333
assertEquals("user-1", event.user().id());
3434
assertEquals("Jane Doe", event.user().name());
3535
assertTrue(event.time() > 0);
36+
37+
JsonObject serializedUser =
38+
new Gson().toJsonTree(event).getAsJsonObject().getAsJsonObject("user");
39+
assertEquals(Set.of("id", "name"), serializedUser.keySet());
40+
assertEquals("user-1", serializedUser.get("id").getAsString());
41+
assertEquals("Jane Doe", serializedUser.get("name").getAsString());
3642
}
3743

3844
@Test
3945
void createAPIEvent_WithNullContext_ReturnsCustomEventEventWithNullRequestAndUser() {
40-
// Act
4146
CustomEvent.CustomEventEvent event = CustomEvent.createAPIEvent("my-custom-event", null);
4247

43-
// Assert
4448
assertNotNull(event);
4549
assertEquals("custom", event.type());
4650
assertEquals("my-custom-event", event.name());
@@ -52,13 +56,10 @@ void createAPIEvent_WithNullContext_ReturnsCustomEventEventWithNullRequestAndUse
5256

5357
@Test
5458
void createAPIEvent_WithContextButNoUser_ReturnsCustomEventEventWithNullUser() {
55-
// Arrange
5659
ContextObject context = new EmptySampleContextObject("test", "/api/resource", "GET");
5760

58-
// Act
5961
CustomEvent.CustomEventEvent event = CustomEvent.createAPIEvent("my-custom-event", context);
6062

61-
// Assert
6263
assertNull(event.user());
6364
}
6465
}

‎docs/track.md‎

Lines changed: 4 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
1-
# Tracking events
1+
# Track custom events
22

3-
`track` lets you record things happening in your app — like failed logins, signups, or password resets. [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.
3+
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.
44

55
```java
66
import dev.aikido.agent_api.Track;
@@ -21,30 +21,6 @@ public void login(HttpServletRequest request) {
2121

2222
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.
2323

24-
Call `Track.track(...)` while handling an HTTP request. Zen associates the event with the request's IP address, and playbook counts are per IP, not across your whole app. Zen also picks up the user agent and current user (if you called [`setUser`](./user.md)) — you don't need to pass those yourself. Calling `setUser` is optional; events without a user are still tracked.
24+
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`](./user.md) before `Track.track(...)`, Zen also includes the current user. Calling `setUser` is optional — events without a user are still tracked.
2525

26-
## More examples
27-
28-
```java
29-
Track.track("user.signed_up");
30-
Track.track("user.password_reset_requested");
31-
Track.track("plan.invite_sent");
32-
Track.track("payment.failed");
33-
```
34-
35-
## Naming events
36-
37-
Event names can use any format. We recommend lowercase, dot-separated names to group related events, such as:
38-
39-
- `user.login_failed`
40-
- `user.login_succeeded`
41-
- `user.signed_up`
42-
- `user.password_reset_requested`
43-
- `payment.failed`
44-
- `plan.invite_sent`
45-
46-
## Things to know
47-
48-
`track` only works inside an HTTP request. If you call it in a background job or outside of a request, nothing gets sent and you'll see a warning in the console.
49-
50-
Each request accepts up to 25 tracked events. Any further calls within the same request are dropped and logged as a warning.
26+
Event names can use any format. We recommend lowercase, dot-separated names such as `user.login_failed`.

0 commit comments

Comments
 (0)