Skip to content

Commit 94715ff

Browse files
Added preSampler for short-circuit Api Security
1 parent e92a04f commit 94715ff

File tree

5 files changed

+39
-20
lines changed

5 files changed

+39
-20
lines changed

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,13 @@ public boolean updateApiAccessIfExpired(String route, String method, int statusC
7777
return isNewOrUpdated;
7878
}
7979

80+
public boolean isApiAccessExpired(String route, String method, int statusCode) {
81+
long currentTime = System.currentTimeMillis();
82+
long hash = computeApiHash(route, method, statusCode);
83+
return !apiAccessMap.containsKey(hash)
84+
|| currentTime - apiAccessMap.get(hash) > expirationTimeInMs;
85+
}
86+
8087
private void cleanupExpiredEntries(long currentTime) {
8188
while (!apiAccessQueue.isEmpty()) {
8289
Long oldestHash = apiAccessQueue.peekFirst();

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

Lines changed: 15 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -14,25 +14,28 @@ public ApiSecurityRequestSampler(final Config config) {
1414
}
1515

1616
public boolean sampleRequest(AppSecRequestContext ctx) {
17-
if (!config.isApiSecurityEnabled() || ctx == null) {
17+
if (!isValid(ctx)) {
1818
return false;
1919
}
2020

21-
String route = ctx.getRoute();
22-
if (route == null) {
23-
return false;
24-
}
21+
return apiAccessTracker.updateApiAccessIfExpired(
22+
ctx.getRoute(), ctx.getMethod(), ctx.getResponseStatus());
23+
}
2524

26-
String method = ctx.getMethod();
27-
if (method == null) {
25+
public boolean preSampleRequest(AppSecRequestContext ctx) {
26+
if (!isValid(ctx)) {
2827
return false;
2928
}
3029

31-
int statusCode = ctx.getResponseStatus();
32-
if (statusCode == 0) {
33-
return false;
34-
}
30+
return apiAccessTracker.isApiAccessExpired(
31+
ctx.getRoute(), ctx.getMethod(), ctx.getResponseStatus());
32+
}
3533

36-
return apiAccessTracker.updateApiAccessIfExpired(route, method, statusCode);
34+
private boolean isValid(AppSecRequestContext ctx) {
35+
return config.isApiSecurityEnabled()
36+
&& ctx != null
37+
&& ctx.getRoute() != null
38+
&& ctx.getMethod() != null
39+
&& ctx.getResponseStatus() != 0;
3740
}
3841
}

dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@
1010
import static datadog.trace.api.UserIdCollectionMode.DISABLED;
1111
import static datadog.trace.api.UserIdCollectionMode.SDK;
1212
import static datadog.trace.api.telemetry.LogCollector.SEND_TELEMETRY;
13-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
1413
import static datadog.trace.util.Strings.toHexString;
1514

1615
import com.datadog.appsec.AppSecSystem;
@@ -767,6 +766,7 @@ private NoopFlow onRequestEnded(RequestContext ctx_, IGSpanInfo spanInfo) {
767766
}
768767

769768
TraceSegment traceSeg = ctx_.getTraceSegment();
769+
Map<String, Object> tags = spanInfo.getTags();
770770

771771
// AppSec report metric and events for web span only
772772
if (traceSeg != null) {
@@ -788,7 +788,7 @@ private NoopFlow onRequestEnded(RequestContext ctx_, IGSpanInfo spanInfo) {
788788
traceSeg.setTagTop("network.client.ip", ctx.getPeerAddress());
789789

790790
// Reflect client_ip as actor.ip for backward compatibility
791-
Object clientIp = spanInfo.getTags().get(Tags.HTTP_CLIENT_IP);
791+
Object clientIp = tags.get(Tags.HTTP_CLIENT_IP);
792792
if (clientIp != null) {
793793
traceSeg.setTagTop("actor.ip", clientIp);
794794
}
@@ -829,10 +829,15 @@ private NoopFlow onRequestEnded(RequestContext ctx_, IGSpanInfo spanInfo) {
829829
}
830830

831831
// Route used in post-processing
832-
Object route = spanInfo.getTags().get(Tags.HTTP_ROUTE);
832+
Object route = tags.get(Tags.HTTP_ROUTE);
833833
if (route instanceof String) {
834834
ctx.setRoute((String) route);
835835
}
836+
if (requestSampler.preSampleRequest(ctx)) {
837+
// The request is pre-sampled - we need to post-process it
838+
spanInfo.setRequiresPostProcessing(true);
839+
}
840+
836841
return NoopFlow.INSTANCE;
837842
}
838843

@@ -1027,7 +1032,6 @@ private Flow<Void> maybePublishRequestData(AppSecRequestContext ctx) {
10271032

10281033
try {
10291034
GatewayContext gwCtx = new GatewayContext(false);
1030-
activeSpan().getLocalRootSpan().setRequiresPostProcessing(true);
10311035
return producerService.publishDataEvent(subInfo, ctx, bundle, gwCtx);
10321036
} catch (ExpiredSubscriberInfoException e) {
10331037
this.initialReqDataSubInfo = null;

dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/api/security/ApiSecurityRequestSamplerTest.groovy

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ class ApiSecurityRequestSamplerTest extends DDSpecification {
1515
void 'Api Security Sample Request'() {
1616
when:
1717
def span = Mock(AppSecRequestContext) {
18-
getSavedRawURI() >> route
18+
getRoute() >> route
1919
getMethod() >> method
2020
getResponseStatus() >> statusCode
2121
}

dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy

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

33
import com.datadog.appsec.AppSecSystem
4+
import com.datadog.appsec.api.security.ApiSecurityRequestSampler
45
import com.datadog.appsec.config.TraceSegmentPostProcessor
56
import com.datadog.appsec.event.EventDispatcher
67
import com.datadog.appsec.event.EventProducerService
@@ -79,7 +80,10 @@ class GatewayBridgeSpecification extends DDSpecification {
7980
}
8081

8182
TraceSegmentPostProcessor pp = Mock()
82-
GatewayBridge bridge = new GatewayBridge(ig, eventDispatcher, null, [pp])
83+
ApiSecurityRequestSampler requestSampler = Mock(ApiSecurityRequestSampler) {
84+
preSampleRequest(_ as AppSecRequestContext) >> false
85+
}
86+
GatewayBridge bridge = new GatewayBridge(ig, eventDispatcher, requestSampler, [pp])
8387

8488
Supplier<Flow<AppSecRequestContext>> requestStartedCB
8589
BiFunction<RequestContext, AgentSpan, Flow<Void>> requestEndedCB
@@ -164,7 +168,6 @@ class GatewayBridgeSpecification extends DDSpecification {
164168
1 * spanInfo.getTags() >> ['http.client_ip': '1.1.1.1']
165169
1 * mockAppSecCtx.transferCollectedEvents() >> [event]
166170
1 * mockAppSecCtx.peerAddress >> '2001::1'
167-
1 * mockAppSecCtx.close()
168171
1 * traceSegment.setTagTop("_dd.appsec.enabled", 1)
169172
1 * traceSegment.setTagTop("_dd.runtime_family", "jvm")
170173
1 * traceSegment.setTagTop('appsec.event', true)
@@ -974,7 +977,9 @@ class GatewayBridgeSpecification extends DDSpecification {
974977
getData(RequestContextSlot.APPSEC) >> mockAppSecCtx
975978
getTraceSegment() >> traceSegment
976979
}
977-
final spanInfo = Mock(AgentSpan)
980+
final spanInfo = Mock(AgentSpan) {
981+
getTags() >> ['http.route':'/']
982+
}
978983

979984
when:
980985
requestEndedCB.apply(mockCtx, spanInfo)

0 commit comments

Comments
 (0)