Skip to content

Commit 9ad217d

Browse files
smolasmola
committed
Lazy API Security initialization
1 parent a7ce6e7 commit 9ad217d

File tree

3 files changed

+90
-60
lines changed

3 files changed

+90
-60
lines changed

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

Lines changed: 27 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,8 @@ public class AppSecSystem {
4343
private static ReplaceableEventProducerService REPLACEABLE_EVENT_PRODUCER; // testing
4444
private static Runnable STOP_SUBSCRIPTION_SERVICE;
4545
private static Runnable RESET_SUBSCRIPTION_SERVICE;
46+
private static final AtomicBoolean API_SECURITY_INITIALIZED = new AtomicBoolean(false);
47+
private static volatile ApiSecuritySampler API_SECURITY_SAMPLER = new ApiSecuritySampler.NoOp();
4648

4749
public static void start(SubscriptionService gw, SharedCommunicationObjects sco) {
4850
try {
@@ -69,18 +71,6 @@ private static void doStart(SubscriptionService gw, SharedCommunicationObjects s
6971
EventDispatcher eventDispatcher = new EventDispatcher();
7072
REPLACEABLE_EVENT_PRODUCER.replaceEventProducerService(eventDispatcher);
7173

72-
ApiSecuritySampler requestSampler;
73-
if (Config.get().isApiSecurityEnabled()) {
74-
requestSampler = new ApiSecuritySamplerImpl();
75-
// When DD_API_SECURITY_ENABLED=true, ths post-processor is set even when AppSec is inactive.
76-
// This should be low overhead since the post-processor exits early if there's no AppSec
77-
// context.
78-
SpanPostProcessor.Holder.INSTANCE =
79-
new AppSecSpanPostProcessor(requestSampler, REPLACEABLE_EVENT_PRODUCER);
80-
} else {
81-
requestSampler = new ApiSecuritySampler.NoOp();
82-
}
83-
8474
ConfigurationPoller configurationPoller = sco.configurationPoller(config);
8575
// may throw and abort startup
8676
APP_SEC_CONFIG_SERVICE =
@@ -94,7 +84,7 @@ private static void doStart(SubscriptionService gw, SharedCommunicationObjects s
9484
new GatewayBridge(
9585
gw,
9686
REPLACEABLE_EVENT_PRODUCER,
97-
requestSampler,
87+
() -> API_SECURITY_SAMPLER,
9888
APP_SEC_CONFIG_SERVICE.getTraceSegmentPostProcessors());
9989

10090
loadModules(eventDispatcher, sco.monitoring);
@@ -129,6 +119,9 @@ public static void setActive(boolean status) {
129119
log.debug("AppSec is now {}", status ? "active" : "inactive");
130120
ProductChangeCollector.get()
131121
.update(new ProductChange().productType(ProductChange.ProductType.APPSEC).enabled(status));
122+
if (status) {
123+
maybeInitializeApiSecurity();
124+
}
132125
}
133126

134127
public static void stop() {
@@ -196,6 +189,27 @@ private static void reloadSubscriptions(
196189
}
197190
}
198191

192+
private static void maybeInitializeApiSecurity() {
193+
if (!Config.get().isApiSecurityEnabled()) {
194+
return;
195+
}
196+
if (!ActiveSubsystems.APPSEC_ACTIVE) {
197+
return;
198+
}
199+
// We initialize API Security the first time AppSec becomes active.
200+
// We never de-initialize it, as that could lead to a leak of open WAF contexts in-flight.
201+
if (API_SECURITY_INITIALIZED.compareAndSet(false, true)) {
202+
ApiSecuritySampler requestSampler = new ApiSecuritySamplerImpl();
203+
SpanPostProcessor.Holder.INSTANCE =
204+
new AppSecSpanPostProcessor(requestSampler, REPLACEABLE_EVENT_PRODUCER);
205+
if (SpanPostProcessor.Holder.INSTANCE == SpanPostProcessor.Holder.NOOP) {
206+
SpanPostProcessor.Holder.INSTANCE =
207+
new AppSecSpanPostProcessor(requestSampler, REPLACEABLE_EVENT_PRODUCER);
208+
API_SECURITY_SAMPLER = requestSampler;
209+
}
210+
}
211+
}
212+
199213
public static boolean isStarted() {
200214
return STARTED.get();
201215
}

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

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -52,8 +52,10 @@
5252
import java.util.Map;
5353
import java.util.Set;
5454
import java.util.concurrent.ConcurrentHashMap;
55+
import java.util.function.Supplier;
5556
import java.util.regex.Pattern;
5657
import java.util.stream.Collectors;
58+
import javax.annotation.Nonnull;
5759
import org.slf4j.Logger;
5860
import org.slf4j.LoggerFactory;
5961

@@ -89,7 +91,7 @@ public class GatewayBridge {
8991

9092
private final SubscriptionService subscriptionService;
9193
private final EventProducerService producerService;
92-
private final ApiSecuritySampler requestSampler;
94+
private final Supplier<ApiSecuritySampler> requestSamplerSupplier;
9395
private final List<TraceSegmentPostProcessor> traceSegmentPostProcessors;
9496

9597
// subscriber cache
@@ -115,11 +117,11 @@ public class GatewayBridge {
115117
public GatewayBridge(
116118
SubscriptionService subscriptionService,
117119
EventProducerService producerService,
118-
ApiSecuritySampler requestSampler,
120+
@Nonnull Supplier<ApiSecuritySampler> requestSamplerSupplier,
119121
List<TraceSegmentPostProcessor> traceSegmentPostProcessors) {
120122
this.subscriptionService = subscriptionService;
121123
this.producerService = producerService;
122-
this.requestSampler = requestSampler;
124+
this.requestSamplerSupplier = requestSamplerSupplier;
123125
this.traceSegmentPostProcessors = traceSegmentPostProcessors;
124126
}
125127

@@ -778,6 +780,7 @@ private boolean maybeSampleForApiSecurity(
778780
if (route != null) {
779781
ctx.setRoute(route.toString());
780782
}
783+
ApiSecuritySampler requestSampler = requestSamplerSupplier.get();
781784
return requestSampler.preSampleRequest(ctx);
782785
}
783786

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

Lines changed: 57 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,7 @@ class GatewayBridgeSpecification extends DDSpecification {
8686

8787
TraceSegmentPostProcessor pp = Mock()
8888
ApiSecuritySamplerImpl requestSampler = Mock(ApiSecuritySamplerImpl)
89-
GatewayBridge bridge = new GatewayBridge(ig, eventDispatcher, requestSampler, [pp])
89+
GatewayBridge bridge = new GatewayBridge(ig, eventDispatcher, () -> requestSampler, [pp])
9090

9191
Supplier<Flow<AppSecRequestContext>> requestStartedCB
9292
BiFunction<RequestContext, AgentSpan, Flow<Void>> requestEndedCB
@@ -258,8 +258,9 @@ class GatewayBridgeSpecification extends DDSpecification {
258258
ctx.data.rawURI = '/'
259259
ctx.data.peerAddress = '0.0.0.0'
260260
eventDispatcher.getDataSubscribers(_) >> nonEmptyDsInfo
261-
eventDispatcher.publishDataEvent(nonEmptyDsInfo, ctx.data, _ as DataBundle, _ as GatewayContext) >>
262-
{ bundle = it[2]; NoopFlow.INSTANCE }
261+
eventDispatcher.publishDataEvent(nonEmptyDsInfo, ctx.data, _ as DataBundle, _ as GatewayContext) >> {
262+
bundle = it[2]; NoopFlow.INSTANCE
263+
}
263264

264265
and:
265266
reqHeadersDoneCB.apply(ctx)
@@ -277,8 +278,9 @@ class GatewayBridgeSpecification extends DDSpecification {
277278

278279
when:
279280
eventDispatcher.getDataSubscribers(_) >> nonEmptyDsInfo
280-
eventDispatcher.publishDataEvent(nonEmptyDsInfo, ctx.data, _ as DataBundle, _ as GatewayContext) >>
281-
{ bundle = it[2]; gatewayContext = it[3]; NoopFlow.INSTANCE }
281+
eventDispatcher.publishDataEvent(nonEmptyDsInfo, ctx.data, _ as DataBundle, _ as GatewayContext) >> {
282+
bundle = it[2]; gatewayContext = it[3]; NoopFlow.INSTANCE
283+
}
282284

283285
and:
284286
reqHeadersDoneCB.apply(ctx)
@@ -298,8 +300,9 @@ class GatewayBridgeSpecification extends DDSpecification {
298300

299301
when:
300302
eventDispatcher.getDataSubscribers(_) >> nonEmptyDsInfo
301-
eventDispatcher.publishDataEvent(nonEmptyDsInfo, ctx.data, _ as DataBundle, _ as GatewayContext) >>
302-
{ bundle = it[2]; gatewayContext = it[3]; NoopFlow.INSTANCE }
303+
eventDispatcher.publishDataEvent(nonEmptyDsInfo, ctx.data, _ as DataBundle, _ as GatewayContext) >> {
304+
bundle = it[2]; gatewayContext = it[3]; NoopFlow.INSTANCE
305+
}
303306

304307
and:
305308
reqHeadersDoneCB.apply(ctx)
@@ -319,8 +322,9 @@ class GatewayBridgeSpecification extends DDSpecification {
319322

320323
when:
321324
eventDispatcher.getDataSubscribers(_) >> nonEmptyDsInfo
322-
eventDispatcher.publishDataEvent(nonEmptyDsInfo, ctx.data, _ as DataBundle, _ as GatewayContext) >>
323-
{ bundle = it[2]; gatewayContext = it[3]; NoopFlow.INSTANCE }
325+
eventDispatcher.publishDataEvent(nonEmptyDsInfo, ctx.data, _ as DataBundle, _ as GatewayContext) >> {
326+
bundle = it[2]; gatewayContext = it[3]; NoopFlow.INSTANCE
327+
}
324328

325329
and:
326330
reqHeadersDoneCB.apply(ctx)
@@ -339,9 +343,12 @@ class GatewayBridgeSpecification extends DDSpecification {
339343
def adapter = TestURIDataAdapter.create(uri, supportsRaw)
340344

341345
when:
342-
eventDispatcher.getDataSubscribers({ KnownAddresses.REQUEST_URI_RAW in it }) >> nonEmptyDsInfo
343-
eventDispatcher.publishDataEvent(nonEmptyDsInfo, ctx.data, _ as DataBundle, _ as GatewayContext) >>
344-
{ bundle = it[2]; gatewayContext = it[3]; NoopFlow.INSTANCE }
346+
eventDispatcher.getDataSubscribers({
347+
KnownAddresses.REQUEST_URI_RAW in it
348+
}) >> nonEmptyDsInfo
349+
eventDispatcher.publishDataEvent(nonEmptyDsInfo, ctx.data, _ as DataBundle, _ as GatewayContext) >> {
350+
bundle = it[2]; gatewayContext = it[3]; NoopFlow.INSTANCE
351+
}
345352

346353
and:
347354
requestMethodURICB.apply(ctx, 'GET', adapter)
@@ -373,9 +380,12 @@ class GatewayBridgeSpecification extends DDSpecification {
373380
def adapter = TestURIDataAdapter.create(uri)
374381

375382
when:
376-
eventDispatcher.getDataSubscribers({ KnownAddresses.REQUEST_URI_RAW in it }) >> nonEmptyDsInfo
377-
eventDispatcher.publishDataEvent(nonEmptyDsInfo, ctx.data, _ as DataBundle, _ as GatewayContext) >>
378-
{ bundle = it[2]; gatewayContext = it[3]; NoopFlow.INSTANCE }
383+
eventDispatcher.getDataSubscribers({
384+
KnownAddresses.REQUEST_URI_RAW in it
385+
}) >> nonEmptyDsInfo
386+
eventDispatcher.publishDataEvent(nonEmptyDsInfo, ctx.data, _ as DataBundle, _ as GatewayContext) >> {
387+
bundle = it[2]; gatewayContext = it[3]; NoopFlow.INSTANCE
388+
}
379389

380390
and:
381391
requestMethodURICB.apply(ctx, 'GET', adapter)
@@ -406,9 +416,12 @@ class GatewayBridgeSpecification extends DDSpecification {
406416
GatewayContext gatewayContext
407417

408418
when:
409-
eventDispatcher.getDataSubscribers({ KnownAddresses.REQUEST_PATH_PARAMS in it }) >> nonEmptyDsInfo
410-
eventDispatcher.publishDataEvent(nonEmptyDsInfo, ctx.data, _ as DataBundle, _ as GatewayContext) >>
411-
{ bundle = it[2]; gatewayContext = it[3]; NoopFlow.INSTANCE }
419+
eventDispatcher.getDataSubscribers({
420+
KnownAddresses.REQUEST_PATH_PARAMS in it
421+
}) >> nonEmptyDsInfo
422+
eventDispatcher.publishDataEvent(nonEmptyDsInfo, ctx.data, _ as DataBundle, _ as GatewayContext) >> {
423+
bundle = it[2]; gatewayContext = it[3]; NoopFlow.INSTANCE
424+
}
412425

413426
and:
414427
pathParamsCB.apply(ctx, [a: 'b'])
@@ -663,9 +676,9 @@ class GatewayBridgeSpecification extends DDSpecification {
663676

664677
when:
665678
Flow<?> flow = requestBodyProcessedCB.apply(ctx, new Object() {
666-
@SuppressWarnings('UnusedPrivateField')
667-
private String foo = 'bar'
668-
})
679+
@SuppressWarnings('UnusedPrivateField')
680+
private String foo = 'bar'
681+
})
669682

670683
then:
671684
1 * eventDispatcher.publishDataEvent(nonEmptyDsInfo, ctx.data, _ as DataBundle, _ as GatewayContext) >>
@@ -762,9 +775,9 @@ class GatewayBridgeSpecification extends DDSpecification {
762775

763776
when:
764777
Flow<?> flow = grpcServerRequestMessageCB.apply(ctx, new Object() {
765-
@SuppressWarnings('UnusedPrivateField')
766-
private String foo = 'bar'
767-
})
778+
@SuppressWarnings('UnusedPrivateField')
779+
private String foo = 'bar'
780+
})
768781

769782
then:
770783
1 * eventDispatcher.publishDataEvent(nonEmptyDsInfo, ctx.data, _ as DataBundle, _ as GatewayContext) >>
@@ -919,28 +932,28 @@ class GatewayBridgeSpecification extends DDSpecification {
919932

920933
void 'no appsec events if was not created request context in request_start event'() {
921934
RequestContext emptyCtx = new RequestContext() {
922-
final Object data = null
923-
BlockResponseFunction blockResponseFunction
924-
925-
@Override
926-
Object getData(RequestContextSlot slot) {
927-
data
928-
}
929-
930-
@Override
931-
final TraceSegment getTraceSegment() {
932-
GatewayBridgeSpecification.this.traceSegment
933-
}
934-
935-
@Override
936-
def <T> T getOrCreateMetaStructTop(String key, Function<String, T> defaultValue) {
937-
return null
938-
}
939-
940-
@Override
941-
void close() throws IOException {}
935+
final Object data = null
936+
BlockResponseFunction blockResponseFunction
937+
938+
@Override
939+
Object getData(RequestContextSlot slot) {
940+
data
942941
}
943942

943+
@Override
944+
final TraceSegment getTraceSegment() {
945+
GatewayBridgeSpecification.this.traceSegment
946+
}
947+
948+
@Override
949+
def <T> T getOrCreateMetaStructTop(String key, Function<String, T> defaultValue) {
950+
return null
951+
}
952+
953+
@Override
954+
void close() throws IOException {}
955+
}
956+
944957
StoredBodySupplier supplier = Stub()
945958
IGSpanInfo spanInfo = Stub(AgentSpan)
946959
Object obj = 'obj'

0 commit comments

Comments
 (0)