Skip to content

Commit 5978655

Browse files
smolasmola
committed
wip
1 parent bf3ebfe commit 5978655

10 files changed

Lines changed: 284 additions & 459 deletions

File tree

‎dd-java-agent/appsec/src/main/java/com/datadog/appsec/AppSecSystem.java‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,7 @@
11
package com.datadog.appsec;
22

3+
import com.datadog.appsec.api.security.ApiSecurityProcessor;
34
import com.datadog.appsec.api.security.ApiSecuritySampler;
4-
import com.datadog.appsec.api.security.ApiSecuritySamplerImpl;
5-
import com.datadog.appsec.api.security.AppSecSpanPostProcessor;
65
import com.datadog.appsec.blocking.BlockingServiceImpl;
76
import com.datadog.appsec.config.AppSecConfigService;
87
import com.datadog.appsec.config.AppSecConfigServiceImpl;
@@ -76,7 +75,7 @@ private static void doStart(SubscriptionService gw, SharedCommunicationObjects s
7675
// This should be low overhead since the post-processor exits early if there's no AppSec
7776
// context.
7877
SpanPostProcessor.Holder.INSTANCE =
79-
new AppSecSpanPostProcessor(requestSampler, REPLACEABLE_EVENT_PRODUCER);
78+
new ApiSecurityProcessor(requestSampler, REPLACEABLE_EVENT_PRODUCER);
8079
} else {
8180
requestSampler = new ApiSecuritySampler.NoOp();
8281
}

dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/AppSecSpanPostProcessor.java renamed to dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiSecurityProcessor.java

Lines changed: 6 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -11,26 +11,24 @@
1111
import datadog.trace.api.gateway.RequestContextSlot;
1212
import datadog.trace.api.internal.TraceSegment;
1313
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
14-
import datadog.trace.bootstrap.instrumentation.api.SpanPostProcessor;
14+
1515
import java.util.Collections;
16-
import java.util.function.BooleanSupplier;
1716
import javax.annotation.Nonnull;
1817
import org.slf4j.Logger;
1918
import org.slf4j.LoggerFactory;
2019

21-
public class AppSecSpanPostProcessor implements SpanPostProcessor {
20+
public class ApiSecurityProcessor {
2221

23-
private static final Logger log = LoggerFactory.getLogger(AppSecSpanPostProcessor.class);
22+
private static final Logger log = LoggerFactory.getLogger(ApiSecurityProcessor.class);
2423
private final ApiSecuritySampler sampler;
2524
private final EventProducerService producerService;
2625

27-
public AppSecSpanPostProcessor(ApiSecuritySampler sampler, EventProducerService producerService) {
26+
public ApiSecurityProcessor(ApiSecuritySampler sampler, EventProducerService producerService) {
2827
this.sampler = sampler;
2928
this.producerService = producerService;
3029
}
3130

32-
@Override
33-
public void process(@Nonnull AgentSpan span, @Nonnull BooleanSupplier timeoutCheck) {
31+
public void process(@Nonnull AgentSpan span) {
3432
final RequestContext ctx_ = span.getRequestContext();
3533
if (ctx_ == null) {
3634
return;
@@ -40,16 +38,8 @@ public void process(@Nonnull AgentSpan span, @Nonnull BooleanSupplier timeoutChe
4038
return;
4139
}
4240

43-
if (!ctx.isKeepOpenForApiSecurityPostProcessing()) {
44-
return;
45-
}
46-
4741
try {
48-
if (timeoutCheck.getAsBoolean()) {
49-
log.debug("Timeout detected, skipping API security post-processing");
50-
return;
51-
}
52-
if (!sampler.sampleRequest(ctx)) {
42+
if (!sampler.sample(ctx)) {
5343
log.debug("Request not sampled, skipping API security post-processing");
5444
return;
5545
}

‎dd-java-agent/appsec/src/main/java/com/datadog/appsec/api/security/ApiSecuritySampler.java‎

Lines changed: 194 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1,35 +1,209 @@
11
package com.datadog.appsec.api.security;
22

33
import com.datadog.appsec.gateway.AppSecRequestContext;
4-
import javax.annotation.Nonnull;
4+
import datadog.trace.util.AgentTaskScheduler;
55

6-
public interface ApiSecuritySampler {
7-
/**
8-
* Prepare a request context for later sampling decision. This method should be called at request
9-
* end, and is thread-safe. If a request can potentially be sampled, this method will return true.
10-
* If this method returns true, the caller MUST call {@link #releaseOne()} once the context is not
11-
* needed anymore.
12-
*/
13-
boolean preSampleRequest(final @Nonnull AppSecRequestContext ctx);
6+
import java.util.Random;
7+
import java.util.concurrent.Executor;
8+
import java.util.concurrent.atomic.AtomicBoolean;
9+
import java.util.concurrent.atomic.AtomicInteger;
10+
import java.util.concurrent.atomic.AtomicLong;
11+
import java.util.concurrent.atomic.AtomicReference;
1412

15-
/** Get the final sampling decision. This method is NOT required to be thread-safe. */
16-
boolean sampleRequest(AppSecRequestContext ctx);
13+
/**
14+
* Internal map for API Security sampling.
15+
* See "[RFC-1021] API Security Sampling Algorithm for thread-based concurrency".
16+
*/
17+
final public class ApiSecuritySampler {
1718

18-
/** Release one permit for the sampler. This must be called after processing a span. */
19-
void releaseOne();
19+
private static final int DEFAULT_MAX_ITEM_COUNT = 4096;
20+
private static final int DEFAULT_INTERVAL_SECONDS = 30;
2021

21-
final class NoOp implements ApiSecuritySampler {
22-
@Override
23-
public boolean preSampleRequest(@Nonnull AppSecRequestContext ctx) {
22+
private final MonotonicClock clock;
23+
private final Executor executor;
24+
private final int intervalSeconds;
25+
private final AtomicReference<Table> table;
26+
private final AtomicBoolean rebuild = new AtomicBoolean(false);
27+
private final long zero;
28+
private final long maxItemCount;
29+
30+
public ApiSecuritySampler() {
31+
this(DEFAULT_MAX_ITEM_COUNT, DEFAULT_INTERVAL_SECONDS, new Random().nextLong(), new DefaultMonotonicClock(), AgentTaskScheduler.INSTANCE);
32+
}
33+
34+
public ApiSecuritySampler(final int maxItemCount, final int intervalSeconds, final long zero, final MonotonicClock clock, Executor executor) {
35+
table = new AtomicReference<>(new Table(maxItemCount));
36+
this.maxItemCount = maxItemCount;
37+
this.intervalSeconds = intervalSeconds;
38+
this.zero = zero;
39+
this.clock = clock != null ? clock : new DefaultMonotonicClock();
40+
this.executor = executor != null ? executor : AgentTaskScheduler.INSTANCE;
41+
}
42+
43+
public boolean sample(AppSecRequestContext ctx) {
44+
final String route = ctx.getRoute();
45+
if (route == null) {
2446
return false;
2547
}
26-
27-
@Override
28-
public boolean sampleRequest(AppSecRequestContext ctx) {
48+
final String method = ctx.getMethod();
49+
if (method == null) {
2950
return false;
3051
}
52+
final int statusCode = ctx.getResponseStatus();
53+
if (statusCode <= 0) {
54+
return false;
55+
}
56+
final long hash = computeApiHash(route, method, statusCode);
57+
return sample(hash);
58+
}
3159

60+
public boolean sample(long key) {
61+
if (key == 0L) {
62+
key = zero;
63+
}
64+
final int now = clock.now();
65+
final Table table = this.table.get();
66+
Table.FindSlotResult findSlotResult;
67+
while (true) {
68+
findSlotResult = table.findSlot(key);
69+
if (!findSlotResult.exists) {
70+
final int newCount = table.count.incrementAndGet();
71+
if (newCount > maxItemCount && rebuild.compareAndSet(false, true)) {
72+
runRebuild();
73+
}
74+
if (newCount > maxItemCount * 2) {
75+
table.count.decrementAndGet();
76+
return false;
77+
}
78+
if (!findSlotResult.entry.key.compareAndSet(0, key)) {
79+
if (findSlotResult.entry.key.get() == key) {
80+
// Another thread just added this entry
81+
return false;
82+
}
83+
// This entry was just claimed for another key, try another slot.
84+
table.count.decrementAndGet();
85+
continue;
86+
}
87+
final long newEntryData = buildDataEntry(now, now);
88+
if (findSlotResult.entry.data.compareAndSet(0, newEntryData)) {
89+
return true;
90+
} else {
91+
return false;
92+
}
93+
}
94+
break;
95+
}
96+
long curData = findSlotResult.entry.data.get();
97+
final int stime = getStime(curData);
98+
final int deadline = now - intervalSeconds;
99+
if (stime <= deadline) {
100+
final long newData = buildDataEntry(now, now);
101+
while (!findSlotResult.entry.data.compareAndSet(curData, newData)) {
102+
curData = findSlotResult.entry.data.get();
103+
if (getStime(curData) == getAtime(curData)) {
104+
// Another thread just issued a keep decision
105+
return false;
106+
}
107+
if (getStime(curData) > now) {
108+
// Another thread is in our fugure, but did not issue a keep decision.
109+
return true;
110+
}
111+
}
112+
return true;
113+
}
114+
final long newData = buildDataEntry(getStime(curData), now);
115+
while (getAtime(curData) < now) {
116+
if (!findSlotResult.entry.data.compareAndSet(curData, newData)) {
117+
curData = findSlotResult.entry.data.get();
118+
}
119+
}
120+
return false;
121+
}
122+
123+
private void runRebuild() {
124+
// TODO
125+
}
126+
127+
private static class Table {
128+
private final Entry[] table;
129+
private final AtomicInteger count = new AtomicInteger(0);
130+
private final int maxItemCount;
131+
132+
public Table(int maxItemCount) {
133+
this.maxItemCount = maxItemCount;
134+
final int size = 2 * maxItemCount + 1;
135+
table = new Entry[size];
136+
for (int i = 0; i < size; i++) {
137+
table[i] = new Entry();
138+
}
139+
}
140+
141+
public FindSlotResult findSlot(final long key) {
142+
final int startIndex = (int) (key % (2L * maxItemCount));
143+
int index = startIndex;
144+
do {
145+
final Entry slot = table[index];
146+
final long slotKey = slot.key.get();
147+
if (slotKey == key) {
148+
return new FindSlotResult(slot, true);
149+
} else if (slotKey == 0L) {
150+
return new FindSlotResult(slot, false);
151+
}
152+
index++;
153+
if (index >= table.length) {
154+
index = 0;
155+
}
156+
} while (index != startIndex);
157+
return new FindSlotResult(table[(int)(maxItemCount * 2)], false);
158+
}
159+
160+
static class FindSlotResult {
161+
public final Entry entry;
162+
public final boolean exists;
163+
164+
public FindSlotResult(final Entry entry, final boolean exists) {
165+
this.entry = entry;
166+
this.exists = exists;
167+
}
168+
}
169+
170+
static class Entry {
171+
private final AtomicLong key = new AtomicLong(0L);
172+
private final AtomicLong data = new AtomicLong(0L);
173+
}
174+
}
175+
176+
interface MonotonicClock {
177+
int now();
178+
}
179+
180+
static class DefaultMonotonicClock implements MonotonicClock {
32181
@Override
33-
public void releaseOne() {}
182+
public int now() {
183+
return (int) (System.nanoTime() / 1_000_000);
184+
}
185+
}
186+
187+
long buildDataEntry(final int stime, final int atime) {
188+
long result = stime;
189+
result <<= 32;
190+
result |= atime & 0xFFFFFFFFL;
191+
return result;
192+
}
193+
194+
int getStime(final long data) {
195+
return (int) (data >> 32);
196+
}
197+
198+
int getAtime(final long data) {
199+
return (int) (data & 0xFFFFFFFFL);
200+
}
201+
202+
private long computeApiHash(final String route, final String method, final int statusCode) {
203+
long result = 17;
204+
result = 31 * result + route.hashCode();
205+
result = 31 * result + method.hashCode();
206+
result = 31 * result + statusCode;
207+
return result;
34208
}
35209
}

0 commit comments

Comments
 (0)