From 04cda746be50ba7f3b01423d1e4fd2346ae2e698 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Alejandro=20Gonz=C3=A1lez=20Garc=C3=ADa?= Date: Wed, 26 Jun 2024 07:55:20 +0200 Subject: [PATCH] Apply appsec rate limiter on event instead of when request end (#7221) What Does This Do check rate limiter once per request only waf events are rate limited Motivation We need to propagate sampling decision on appsec event, not on request end --- .../java/com/datadog/appsec/AppSecSystem.java | 24 +----- .../appsec/gateway/AppSecRequestContext.java | 11 +++ .../datadog/appsec/gateway/GatewayBridge.java | 74 +++++++++---------- .../datadog/appsec/gateway/RateLimiter.java | 2 +- .../appsec/powerwaf/PowerWAFModule.java | 45 ++++++++++- .../appsec/AppSecSystemSpecification.groovy | 31 -------- .../AppSecRequestContextSpecification.groovy | 20 +++++ ...ayBridgeIGRegistrationSpecification.groovy | 2 +- .../gateway/GatewayBridgeSpecification.groovy | 26 +------ .../gateway/RateLimiterSpecification.groovy | 12 +++ .../PowerWAFModuleSpecification.groovy | 58 +++++++++++---- 11 files changed, 171 insertions(+), 134 deletions(-) diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/AppSecSystem.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/AppSecSystem.java index d0eabe70bc0..a356db4efc0 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/AppSecSystem.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/AppSecSystem.java @@ -7,20 +7,17 @@ import com.datadog.appsec.event.EventDispatcher; import com.datadog.appsec.event.ReplaceableEventProducerService; import com.datadog.appsec.gateway.GatewayBridge; -import com.datadog.appsec.gateway.RateLimiter; import com.datadog.appsec.powerwaf.PowerWAFModule; import com.datadog.appsec.util.AbortStartupException; import com.datadog.appsec.util.StandardizedLogging; import datadog.appsec.api.blocking.Blocking; import datadog.appsec.api.blocking.BlockingService; import datadog.communication.ddagent.SharedCommunicationObjects; -import datadog.communication.monitor.Counter; import datadog.communication.monitor.Monitoring; import datadog.remoteconfig.ConfigurationPoller; import datadog.trace.api.Config; import datadog.trace.api.ProductActivation; import datadog.trace.api.gateway.SubscriptionService; -import datadog.trace.api.time.SystemTimeSource; import datadog.trace.bootstrap.ActiveSubsystems; import datadog.trace.util.Strings; import java.util.Collections; @@ -81,17 +78,14 @@ private static void doStart(SubscriptionService gw, SharedCommunicationObjects s sco.createRemaining(config); - RateLimiter rateLimiter = getRateLimiter(config, sco.monitoring); - GatewayBridge gatewayBridge = new GatewayBridge( gw, REPLACEABLE_EVENT_PRODUCER, - rateLimiter, requestSampler, APP_SEC_CONFIG_SERVICE.getTraceSegmentPostProcessors()); - loadModules(eventDispatcher); + loadModules(eventDispatcher, sco.monitoring); gatewayBridge.init(); RESET_SUBSCRIPTION_SERVICE = gatewayBridge::stop; @@ -112,18 +106,6 @@ private static void doStart(SubscriptionService gw, SharedCommunicationObjects s } } - private static RateLimiter getRateLimiter(Config config, Monitoring monitoring) { - RateLimiter rateLimiter = null; - int appSecTraceRateLimit = config.getAppSecTraceRateLimit(); - if (appSecTraceRateLimit > 0) { - Counter counter = monitoring.newCounter("_dd.java.appsec.rate_limit.dropped_traces"); - rateLimiter = - new RateLimiter( - appSecTraceRateLimit, SystemTimeSource.INSTANCE, () -> counter.increment(1)); - } - return rateLimiter; - } - public static boolean isActive() { return ActiveSubsystems.APPSEC_ACTIVE; } @@ -144,11 +126,11 @@ public static void stop() { APP_SEC_CONFIG_SERVICE.close(); } - private static void loadModules(EventDispatcher eventDispatcher) { + private static void loadModules(EventDispatcher eventDispatcher, Monitoring monitoring) { EventDispatcher.DataSubscriptionSet dataSubscriptionSet = new EventDispatcher.DataSubscriptionSet(); - final List modules = Collections.singletonList(new PowerWAFModule()); + final List modules = Collections.singletonList(new PowerWAFModule(monitoring)); for (AppSecModule module : modules) { log.debug("Starting appsec module {}", module.getName()); try { diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java index 6893e4845b9..f01cdf7009b 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/AppSecRequestContext.java @@ -16,6 +16,7 @@ import java.util.*; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicIntegerFieldUpdater; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -103,6 +104,9 @@ public class AppSecRequestContext implements DataBundle, Closeable { private boolean pathParamsPublished; private Map apiSchemas; + private AtomicBoolean rateLimited = new AtomicBoolean(false); + private volatile boolean throttled; + // should be guarded by this private Additive additive; // set after additive is set @@ -486,4 +490,11 @@ boolean commitApiSchemas(TraceSegment traceSegment) { apiSchemas.forEach(traceSegment::setTagTop); return true; } + + public boolean isThrottled(RateLimiter rateLimiter) { + if (rateLimiter != null && rateLimited.compareAndSet(false, true)) { + throttled = rateLimiter.isThrottled(); + } + return throttled; + } } diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java index ae86eac021e..cbe3c858aec 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java @@ -70,7 +70,6 @@ public class GatewayBridge { private final SubscriptionService subscriptionService; private final EventProducerService producerService; - private final RateLimiter rateLimiter; private final ApiSecurityRequestSampler requestSampler; private final List traceSegmentPostProcessors; @@ -90,12 +89,10 @@ public class GatewayBridge { public GatewayBridge( SubscriptionService subscriptionService, EventProducerService producerService, - RateLimiter rateLimiter, ApiSecurityRequestSampler requestSampler, List traceSegmentPostProcessors) { this.subscriptionService = subscriptionService; this.producerService = producerService; - this.rateLimiter = rateLimiter; this.requestSampler = requestSampler; this.traceSegmentPostProcessors = traceSegmentPostProcessors; } @@ -141,47 +138,42 @@ public void init() { pp.processTraceSegment(traceSeg, ctx, collectedEvents); } - if (rateLimiter == null || !rateLimiter.isThrottled()) { - // If detected any events - mark span at appsec.event - if (!collectedEvents.isEmpty()) { - // Keep event related span, because it could be ignored in case of - // reduced datadog sampling rate. - traceSeg.setTagTop(Tags.ASM_KEEP, true); - traceSeg.setTagTop("appsec.event", true); - traceSeg.setTagTop("network.client.ip", ctx.getPeerAddress()); - - // Reflect client_ip as actor.ip for backward compatibility - Object clientIp = spanInfo.getTags().get(Tags.HTTP_CLIENT_IP); - if (clientIp != null) { - traceSeg.setTagTop("actor.ip", clientIp); - } + // If detected any events - mark span at appsec.event + if (!collectedEvents.isEmpty()) { + // Set asm keep in case that root span was not available when events are detected + traceSeg.setTagTop(Tags.ASM_KEEP, true); + traceSeg.setTagTop("appsec.event", true); + traceSeg.setTagTop("network.client.ip", ctx.getPeerAddress()); + + // Reflect client_ip as actor.ip for backward compatibility + Object clientIp = spanInfo.getTags().get(Tags.HTTP_CLIENT_IP); + if (clientIp != null) { + traceSeg.setTagTop("actor.ip", clientIp); + } - // Report AppSec events via "_dd.appsec.json" tag - AppSecEventWrapper wrapper = new AppSecEventWrapper(collectedEvents); - traceSeg.setDataTop("appsec", wrapper); - - // Report collected request and response headers based on allow list - writeRequestHeaders(traceSeg, REQUEST_HEADERS_ALLOW_LIST, ctx.getRequestHeaders()); - writeResponseHeaders( - traceSeg, RESPONSE_HEADERS_ALLOW_LIST, ctx.getResponseHeaders()); - - // Report collected stack traces - StackTraceCollection stackTraceCollection = ctx.transferStackTracesCollection(); - if (stackTraceCollection != null) { - Object flatStruct = ObjectFlattener.flatten(stackTraceCollection); - if (flatStruct != null) { - traceSeg.setMetaStructTop("_dd.stack", flatStruct); - } - } + // Report AppSec events via "_dd.appsec.json" tag + AppSecEventWrapper wrapper = new AppSecEventWrapper(collectedEvents); + traceSeg.setDataTop("appsec", wrapper); - } else if (hasUserTrackingEvent(traceSeg)) { - // Report all collected request headers on user tracking event - writeRequestHeaders(traceSeg, REQUEST_HEADERS_ALLOW_LIST, ctx.getRequestHeaders()); - } else { - // Report minimum set of collected request headers - writeRequestHeaders( - traceSeg, DEFAULT_REQUEST_HEADERS_ALLOW_LIST, ctx.getRequestHeaders()); + // Report collected request and response headers based on allow list + writeRequestHeaders(traceSeg, REQUEST_HEADERS_ALLOW_LIST, ctx.getRequestHeaders()); + writeResponseHeaders(traceSeg, RESPONSE_HEADERS_ALLOW_LIST, ctx.getResponseHeaders()); + + // Report collected stack traces + StackTraceCollection stackTraceCollection = ctx.transferStackTracesCollection(); + if (stackTraceCollection != null) { + Object flatStruct = ObjectFlattener.flatten(stackTraceCollection); + if (flatStruct != null) { + traceSeg.setMetaStructTop("_dd.stack", flatStruct); + } } + } else if (hasUserTrackingEvent(traceSeg)) { + // Report all collected request headers on user tracking event + writeRequestHeaders(traceSeg, REQUEST_HEADERS_ALLOW_LIST, ctx.getRequestHeaders()); + } else { + // Report minimum set of collected request headers + writeRequestHeaders( + traceSeg, DEFAULT_REQUEST_HEADERS_ALLOW_LIST, ctx.getRequestHeaders()); } // If extracted any Api Schemas - commit them if (!ctx.commitApiSchemas(traceSeg)) { diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/RateLimiter.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/RateLimiter.java index 5aadfda0c86..a04b9c162f2 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/RateLimiter.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/RateLimiter.java @@ -47,7 +47,7 @@ public RateLimiter(int limitPerSec, TimeSource timeSource, ThrottledCallback cb) this.throttledCb = cb; } - public final boolean isThrottled() { + public boolean isThrottled() { long curSec = this.timeSource.getNanoTicks(); long storedState; long newState; diff --git a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/powerwaf/PowerWAFModule.java b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/powerwaf/PowerWAFModule.java index fc1bbf9137b..9994532dfcd 100644 --- a/dd-java-agent/appsec/src/main/java/com/datadog/appsec/powerwaf/PowerWAFModule.java +++ b/dd-java-agent/appsec/src/main/java/com/datadog/appsec/powerwaf/PowerWAFModule.java @@ -13,6 +13,7 @@ import com.datadog.appsec.event.data.DataBundle; import com.datadog.appsec.event.data.KnownAddresses; import com.datadog.appsec.gateway.AppSecRequestContext; +import com.datadog.appsec.gateway.RateLimiter; import com.datadog.appsec.report.AppSecEvent; import com.datadog.appsec.stack_trace.StackTraceEvent; import com.datadog.appsec.stack_trace.StackTraceEvent.Frame; @@ -21,13 +22,17 @@ import com.squareup.moshi.Moshi; import com.squareup.moshi.Types; import datadog.appsec.api.blocking.BlockingContentType; +import datadog.communication.monitor.Counter; +import datadog.communication.monitor.Monitoring; import datadog.trace.api.Config; import datadog.trace.api.ProductActivation; import datadog.trace.api.gateway.Flow; import datadog.trace.api.telemetry.LogCollector; import datadog.trace.api.telemetry.WafMetricCollector; +import datadog.trace.api.time.SystemTimeSource; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentTracer; +import datadog.trace.bootstrap.instrumentation.api.Tags; import datadog.trace.util.stacktrace.StackWalkerFactory; import io.sqreen.powerwaf.Additive; import io.sqreen.powerwaf.Powerwaf; @@ -146,9 +151,18 @@ static void createLimitsObject() { private final PowerWAFInitializationResultReporter initReporter = new PowerWAFInitializationResultReporter(); private final PowerWAFStatsReporter statsReporter = new PowerWAFStatsReporter(); + private final RateLimiter rateLimiter; private String currentRulesVersion; + public PowerWAFModule() { + this(null); + } + + public PowerWAFModule(Monitoring monitoring) { + this.rateLimiter = getRateLimiter(monitoring); + } + @Override public void config(AppSecModuleConfigurer appSecConfigService) throws AppSecModuleActivationException { @@ -327,6 +341,21 @@ private PowerwafConfig createPowerwafConfig() { return pwConfig; } + private static RateLimiter getRateLimiter(Monitoring monitoring) { + if (monitoring == null) { + return null; + } + RateLimiter rateLimiter = null; + int appSecTraceRateLimit = Config.get().getAppSecTraceRateLimit(); + if (appSecTraceRateLimit > 0) { + Counter counter = monitoring.newCounter("_dd.java.appsec.rate_limit.dropped_traces"); + rateLimiter = + new RateLimiter( + appSecTraceRateLimit, SystemTimeSource.INSTANCE, () -> counter.increment(1)); + } + return rateLimiter; + } + @Override public String getName() { return "powerwaf"; @@ -444,7 +473,21 @@ public void onDataAvailable( } } Collection events = buildEvents(resultWithData); - reqCtx.reportEvents(events); + + if (!events.isEmpty() && !reqCtx.isThrottled(rateLimiter)) { + AgentSpan activeSpan = AgentTracer.get().activeSpan(); + if (activeSpan != null) { + log.debug("Setting force-keep tag on the current span"); + // Keep event related span, because it could be ignored in case of + // reduced datadog sampling rate. + activeSpan.getLocalRootSpan().setTag(Tags.ASM_KEEP, true); + } else { + // If active span is not available the ASK_KEEP tag will be set in the GatewayBridge + // when the request ends + log.debug("There is no active span available"); + } + reqCtx.reportEvents(events); + } if (flow.isBlocking()) { reqCtx.setBlocked(); diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/AppSecSystemSpecification.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/AppSecSystemSpecification.groovy index 99284d237f8..a92960e0a13 100644 --- a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/AppSecSystemSpecification.groovy +++ b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/AppSecSystemSpecification.groovy @@ -7,7 +7,6 @@ import com.datadog.appsec.report.AppSecEvent import com.datadog.appsec.util.AbortStartupException import datadog.communication.ddagent.DDAgentFeaturesDiscovery import datadog.communication.ddagent.SharedCommunicationObjects -import datadog.communication.monitor.Counter import datadog.communication.monitor.Monitoring import datadog.remoteconfig.ConfigurationChangesTypedListener import datadog.remoteconfig.ConfigurationEndListener @@ -83,36 +82,6 @@ class AppSecSystemSpecification extends DDSpecification { 1 * traceSegment.setTagTop('actor.ip', '1.1.1.1') } - void 'honors appsec.trace.rate.limit'() { - BiFunction> requestEndedCB - RequestContext requestContext = Mock() - TraceSegment traceSegment = Mock() - AppSecRequestContext appSecReqCtx = Mock() - def sco = sharedCommunicationObjects() - Counter throttledCounter = Mock() - IGSpanInfo span = Mock(AgentSpan) - - setup: - injectSysConfig('dd.appsec.trace.rate.limit', '5') - - when: - AppSecSystem.start(subService, sco) - 7.times { requestEndedCB.apply(requestContext, span) } - - then: - span.getTags() >> ['http.client_ip':'1.1.1.1'] - 1 * sco.monitoring.newCounter('_dd.java.appsec.rate_limit.dropped_traces') >> throttledCounter - 1 * subService.registerCallback(EVENTS.requestEnded(), _) >> { requestEndedCB = it[1]; null } - 7 * requestContext.getData(RequestContextSlot.APPSEC) >> appSecReqCtx - 7 * requestContext.traceSegment >> traceSegment - 7 * appSecReqCtx.transferCollectedEvents() >> [Stub(AppSecEvent)] - // allow for one extra in case we move to another second and round down the prev count - (5..6) * appSecReqCtx.getRequestHeaders() >> [:] - (5..6) * appSecReqCtx.getResponseHeaders() >> [:] - (5..6) * traceSegment.setDataTop("appsec", _) - (1..2) * throttledCounter.increment(1) - } - void 'throws if the config file is not parseable'() { setup: Path path = Files.createTempFile('dd-trace-', '.json') diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/AppSecRequestContextSpecification.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/AppSecRequestContextSpecification.groovy index 701db55097b..9cd74259c45 100644 --- a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/AppSecRequestContextSpecification.groovy +++ b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/AppSecRequestContextSpecification.groovy @@ -224,4 +224,24 @@ class AppSecRequestContextSpecification extends DDSpecification { ctx.additive == null !additive.online } + + void 'test isThrottled'(){ + setup: + def rateLimiter = Mock(RateLimiter) + def appSecRequestContext = new AppSecRequestContext() + + when: 'rate limiter is called and throttled is set' + def result = appSecRequestContext.isThrottled(rateLimiter) + + then: + 1 * rateLimiter.isThrottled() >> true + assert result + + when: 'rate limiter is not called more than once per appsec context returns first result' + def result2 = appSecRequestContext.isThrottled(rateLimiter) + + then: + 0 * rateLimiter.isThrottled() + result == result2 + } } diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeIGRegistrationSpecification.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeIGRegistrationSpecification.groovy index 7457fd23e98..eee76de12e1 100644 --- a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeIGRegistrationSpecification.groovy +++ b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeIGRegistrationSpecification.groovy @@ -10,7 +10,7 @@ class GatewayBridgeIGRegistrationSpecification extends DDSpecification { SubscriptionService ig = Mock() EventDispatcher eventDispatcher = Mock() - GatewayBridge bridge = new GatewayBridge(ig, eventDispatcher, null, null, []) + GatewayBridge bridge = new GatewayBridge(ig, eventDispatcher, null, []) void 'request_body_start and request_body_done are registered'() { given: diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy index bbb6146eed8..9017a99a676 100644 --- a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy +++ b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy @@ -18,7 +18,6 @@ import datadog.trace.api.gateway.RequestContextSlot import datadog.trace.api.gateway.SubscriptionService import datadog.trace.api.http.StoredBodySupplier import datadog.trace.api.internal.TraceSegment -import datadog.trace.api.time.TimeSource import datadog.trace.bootstrap.instrumentation.api.AgentSpan import datadog.trace.bootstrap.instrumentation.api.URIDataAdapter import datadog.trace.bootstrap.instrumentation.api.URIDataAdapterBase @@ -59,9 +58,8 @@ class GatewayBridgeSpecification extends DDSpecification { i }() - RateLimiter rateLimiter = new RateLimiter(10, { -> 0L } as TimeSource, RateLimiter.ThrottledCallback.NOOP) TraceSegmentPostProcessor pp = Mock() - GatewayBridge bridge = new GatewayBridge(ig, eventDispatcher, rateLimiter, null, [pp]) + GatewayBridge bridge = new GatewayBridge(ig, eventDispatcher, null, [pp]) Supplier> requestStartedCB BiFunction> requestEndedCB @@ -139,7 +137,6 @@ class GatewayBridgeSpecification extends DDSpecification { 1 * mockAppSecCtx.transferCollectedEvents() >> [event] 1 * mockAppSecCtx.peerAddress >> '2001::1' 1 * mockAppSecCtx.close() - 1 * traceSegment.setTagTop('asm.keep', true) 1 * traceSegment.setTagTop("_dd.appsec.enabled", 1) 1 * traceSegment.setTagTop("_dd.runtime_family", "jvm") 1 * traceSegment.setTagTop('appsec.event', true) @@ -152,27 +149,6 @@ class GatewayBridgeSpecification extends DDSpecification { flow.action == Flow.Action.Noop.INSTANCE } - void 'event publishing is rate limited'() { - AppSecEvent event = Stub() - AppSecRequestContext mockAppSecCtx = Mock(AppSecRequestContext) - mockAppSecCtx.requestHeaders >> [:] - RequestContext mockCtx = Stub(RequestContext) { - getData(RequestContextSlot.APPSEC) >> mockAppSecCtx - getTraceSegment() >> traceSegment - } - IGSpanInfo spanInfo = Mock(AgentSpan) - - when: - 11.times {requestEndedCB.apply(mockCtx, spanInfo) } - - then: - 11 * mockAppSecCtx.transferCollectedEvents() >> [event] - 11 * mockAppSecCtx.close() - 11 * mockAppSecCtx.closeAdditive() - 10 * spanInfo.getTags() >> ['http.client_ip':'1.1.1.1'] - 10 * traceSegment.setDataTop("appsec", _) - } - void 'actor ip calculated from headers'() { AppSecRequestContext mockAppSecCtx = Mock(AppSecRequestContext) mockAppSecCtx.requestHeaders >> [ diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/RateLimiterSpecification.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/RateLimiterSpecification.groovy index 80fd7b14f7c..023e64c5cd0 100644 --- a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/RateLimiterSpecification.groovy +++ b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/RateLimiterSpecification.groovy @@ -163,4 +163,16 @@ class RateLimiterSpecification extends Specification { 2 * mock.nanoTicks >> initialTime - 1_000_000_000L count == 11 } + + void 'test NoOP'(){ + setup: + RateLimiter rateLimiter = new RateLimiter(10, { -> 0L } as TimeSource, RateLimiter.ThrottledCallback.NOOP) + def count = 0 + + when: + 15.times {rateLimiter.isThrottled() || count++ } + + then: + count == 10 + } } diff --git a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/powerwaf/PowerWAFModuleSpecification.groovy b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/powerwaf/PowerWAFModuleSpecification.groovy index f5326857fe1..60ac5e3544f 100644 --- a/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/powerwaf/PowerWAFModuleSpecification.groovy +++ b/dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/powerwaf/PowerWAFModuleSpecification.groovy @@ -17,6 +17,7 @@ import com.datadog.appsec.gateway.AppSecRequestContext import com.datadog.appsec.report.AppSecEvent import com.datadog.appsec.stack_trace.StackTraceEvent import com.datadog.appsec.test.StubAppSecConfigService +import datadog.communication.monitor.Monitoring import datadog.trace.api.ConfigDefaults import datadog.trace.api.internal.TraceSegment import datadog.appsec.api.blocking.BlockingContentType @@ -47,6 +48,7 @@ class PowerWAFModuleSpecification extends DDSpecification { protected AgentTracer.TracerAPI tracer = Mock(AgentTracer.TracerAPI) { activeSpan() >> Mock(AgentSpan) { getSpanId() >> 777 + getLocalRootSpan() >> Mock(AgentSpan) } getSpanId() >> 777 } @@ -203,11 +205,12 @@ class PowerWAFModuleSpecification extends DDSpecification { 1 * ctx.getOrCreateAdditive(_ as PowerwafContext, true) >> { pwafAdditive = it[0].openAdditive() } - 1 * tracer.activeSpan() + 2 * tracer.activeSpan() 1 * ctx.reportEvents(_ as Collection) 1 * ctx.getWafMetrics() 1 * ctx.closeAdditive() 1 * flow.isBlocking() + 1 * ctx.isThrottled(null) 0 * _ } @@ -235,11 +238,12 @@ class PowerWAFModuleSpecification extends DDSpecification { 1 * ctx.getOrCreateAdditive(_ as PowerwafContext, true) >> { pwafAdditive = it[0].openAdditive() } - 1 * tracer.activeSpan() + 2 * tracer.activeSpan() 1 * ctx.reportEvents(_ as Collection) 1 * ctx.getWafMetrics() 2 * ctx.closeAdditive() 1 * flow.isBlocking() + 1 * ctx.isThrottled(null) 0 * _ when: 'merges new waf data with the one in the rules config' @@ -275,11 +279,12 @@ class PowerWAFModuleSpecification extends DDSpecification { 1 * ctx.getOrCreateAdditive(_ as PowerwafContext, true) >> { pwafAdditive = it[0].openAdditive() } - 1 * tracer.activeSpan() + 2 * tracer.activeSpan() 1 * ctx.reportEvents(_ as Collection) 1 * ctx.getWafMetrics() 1 * ctx.closeAdditive() 1 * flow.isBlocking() + 1 * ctx.isThrottled(null) 0 * _ when: @@ -298,11 +303,12 @@ class PowerWAFModuleSpecification extends DDSpecification { 1 * ctx.getOrCreateAdditive(_ as PowerwafContext, true) >> { pwafAdditive = it[0].openAdditive() } - 1 * tracer.activeSpan() + 2 * tracer.activeSpan() 1 * ctx.reportEvents(_ as Collection) 1 * ctx.getWafMetrics() 1 * ctx.closeAdditive() 1 * flow.isBlocking() + 1 * ctx.isThrottled(null) 0 * _ when: 'changes the rules config' @@ -353,11 +359,12 @@ class PowerWAFModuleSpecification extends DDSpecification { 1 * ctx.getOrCreateAdditive(_ as PowerwafContext, true) >> { pwafAdditive = it[0].openAdditive() } - 1 * tracer.activeSpan() + 2 * tracer.activeSpan() 1 * ctx.reportEvents(_ as Collection) 1 * ctx.getWafMetrics() 1 * ctx.closeAdditive() 1 * flow.isBlocking() + 1 * ctx.isThrottled(null) 0 * _ when: @@ -428,11 +435,12 @@ class PowerWAFModuleSpecification extends DDSpecification { 1 * ctx.getOrCreateAdditive(_, true) >> { pwafAdditive = it[0].openAdditive() } - 1 * tracer.activeSpan() + 2 * tracer.activeSpan() 1 * ctx.reportEvents(_ as Collection) 1 * ctx.getWafMetrics() 1 * ctx.closeAdditive() >> { pwafAdditive.close() } 1 * ctx.setBlocked() + 1 * ctx.isThrottled(null) 0 * _ when: @@ -507,12 +515,13 @@ class PowerWAFModuleSpecification extends DDSpecification { 1 * ctx.getOrCreateAdditive(_, true) >> { pwafAdditive = it[0].openAdditive() } - 2 * tracer.activeSpan() + 3 * tracer.activeSpan() // we get two events: one for origin rule, and one for the custom one 1 * ctx.reportEvents(hasSize(2)) 1 * ctx.getWafMetrics() 1 * ctx.closeAdditive() 1 * ctx.setBlocked() + 1 * ctx.isThrottled(null) 0 * _ } @@ -583,11 +592,12 @@ class PowerWAFModuleSpecification extends DDSpecification { rba.blockingContentType == BlockingContentType.AUTO }) 1 * ctx.getOrCreateAdditive(_, true) >> { it[0].openAdditive() } - 1 * tracer.activeSpan() + 2 * tracer.activeSpan() 1 * ctx.reportEvents(_ as Collection) 1 * ctx.getWafMetrics() 1 * ctx.closeAdditive() 1 * flow.isBlocking() + 1 * ctx.isThrottled(null) 0 * _ } @@ -611,6 +621,7 @@ class PowerWAFModuleSpecification extends DDSpecification { 1 * ctx.closeAdditive() 1 * ctx.reportEvents(_) 1 * ctx.setBlocked() + 1 * ctx.isThrottled(null) 0 * ctx._(*_) flow.blocking == true flow.action instanceof Flow.Action.RequestBlockingAction @@ -675,6 +686,7 @@ class PowerWAFModuleSpecification extends DDSpecification { 1 * ctx.closeAdditive() 1 * ctx.reportEvents(_) 1 * ctx.setBlocked() + 1 * ctx.isThrottled(null) 0 * ctx._(*_) flow.blocking == true flow.action.statusCode == 418 @@ -700,6 +712,7 @@ class PowerWAFModuleSpecification extends DDSpecification { 1 * ctx.closeAdditive() 1 * ctx.reportEvents(_) 1 * ctx.setBlocked() + 1 * ctx.isThrottled(null) 0 * ctx._(*_) metrics == null } @@ -751,6 +764,7 @@ class PowerWAFModuleSpecification extends DDSpecification { 1 * ctx.getWafMetrics() >> metrics 1 * ctx.reportEvents(*_) 1 * ctx.setBlocked() + 1 * ctx.isThrottled(null) 0 * ctx._(*_) flow.blocking == true } @@ -1010,12 +1024,13 @@ class PowerWAFModuleSpecification extends DDSpecification { then: 1 * reconf.reloadSubscriptions() 1 * ctx.getOrCreateAdditive(_, true) >> { pwafAdditive = it[0].openAdditive() } - 1 * tracer.activeSpan() + 2 * tracer.activeSpan() 1 * ctx.reportEvents(_ as Collection) 1 * ctx.getWafMetrics() 1 * flow.setAction({ it.blocking }) 1 * ctx.closeAdditive() 1 * flow.isBlocking() + 1 * ctx.isThrottled(null) 0 * _ } @@ -1109,12 +1124,13 @@ class PowerWAFModuleSpecification extends DDSpecification { 1 * reconf.reloadSubscriptions() 1 * ctx.getOrCreateAdditive(_, true) >> { pwafAdditive = it[0].openAdditive() } - 1 * tracer.activeSpan() + 2 * tracer.activeSpan() 1 * ctx.reportEvents(_ as Collection) 1 * ctx.getWafMetrics() 1 * flow.setAction({ it.blocking }) 1 * ctx.closeAdditive() >> {pwafAdditive.close()} 1 * flow.isBlocking() + 1 * ctx.isThrottled(null) _ * ctx.increaseTimeouts() 0 * _ @@ -1212,10 +1228,11 @@ class PowerWAFModuleSpecification extends DDSpecification { 1 * ctx.getWafMetrics() 1 * flow.isBlocking() 1 * flow.setAction({ it.blocking }) - 1 * tracer.activeSpan() + 2 * tracer.activeSpan() 1 * ctx.reportEvents(_ as Collection) 1 * ctx.closeAdditive() >> {pwafAdditive.close()} _ * ctx.increaseTimeouts() + 1 * ctx.isThrottled(null) 0 * _ when: 'removing c restores the state before c was added (rule disabled)' @@ -1343,12 +1360,13 @@ class PowerWAFModuleSpecification extends DDSpecification { then: 1 * ctx.getOrCreateAdditive(_, true) >> { pwafAdditive = it[0].openAdditive() } - 1 * tracer.activeSpan() + 2 * tracer.activeSpan() 1 * ctx.reportEvents(_ as Collection) >> { it[0].iterator().next().ruleMatches[0].parameters[0].value == '/cybercop' } 1 * ctx.getWafMetrics() 1 * flow.isBlocking() + 1 * ctx.isThrottled(null) 0 * _ when: @@ -1359,12 +1377,13 @@ class PowerWAFModuleSpecification extends DDSpecification { 1 * ctx.getOrCreateAdditive(_, true) >> { pwafAdditive } 1 * flow.setAction({ it.blocking }) - 1 * tracer.activeSpan() + 2 * tracer.activeSpan() 1 * ctx.reportEvents(_ as Collection) >> { it[0].iterator().next().ruleMatches[0].parameters[0].value == 'user-to-block-1' } 1 * ctx.getWafMetrics() 1 * ctx.closeAdditive() + 1 * ctx.isThrottled(null) 1 * flow.isBlocking() 0 * _ } @@ -1403,6 +1422,19 @@ class PowerWAFModuleSpecification extends DDSpecification { n << (1..3) } + void 'honors appsec.trace.rate.limit'() { + setup: + injectSysConfig('dd.appsec.trace.rate.limit', '5') + def monitoring = Mock(Monitoring) + + when: + def waf = new PowerWAFModule(monitoring) + + then: + waf.rateLimiter.limitPerSec == 5 + + } + private Map getDefaultConfig() { def service = new StubAppSecConfigService() service.init()