Skip to content

Commit de0e3ef

Browse files
smolasmola
committed
wip
1 parent b69714d commit de0e3ef

File tree

4 files changed

+245
-6
lines changed

4 files changed

+245
-6
lines changed

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

+6-2
Original file line numberDiff line numberDiff line change
@@ -28,8 +28,7 @@ public class ApiSecurityRequestSampler {
2828
private final long expirationTimeInMs;
2929
private final int capacity;
3030
private final TimeSource timeSource;
31-
32-
final NonBlockingSemaphore counter =
31+
private final NonBlockingSemaphore counter =
3332
NonBlockingSemaphore.withPermitCount(MAX_POST_PROCESSING_TASKS);
3433

3534
public ApiSecurityRequestSampler() {
@@ -86,6 +85,11 @@ public boolean sampleRequest(AppSecRequestContext ctx) {
8685
return updateApiAccessIfExpired(hash);
8786
}
8887

88+
/** Release one permit for the sampler. This must be called after processing a span. */
89+
void releaseOne() {
90+
counter.release();
91+
}
92+
8993
private boolean updateApiAccessIfExpired(final long hash) {
9094
final long currentTime = timeSource.getCurrentTimeMillis();
9195

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

+3-2
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import datadog.trace.bootstrap.instrumentation.api.SpanPostProcessor;
1414
import java.util.Collections;
1515
import java.util.function.BooleanSupplier;
16+
import javax.annotation.Nonnull;
1617
import org.slf4j.Logger;
1718
import org.slf4j.LoggerFactory;
1819

@@ -29,7 +30,7 @@ public AppSecSpanPostProcessor(
2930
}
3031

3132
@Override
32-
public void process(AgentSpan span, BooleanSupplier timeoutCheck) {
33+
public void process(@Nonnull AgentSpan span, @Nonnull BooleanSupplier timeoutCheck) {
3334
final RequestContext ctx_ = span.getRequestContext();
3435
if (ctx_ == null) {
3536
return;
@@ -58,7 +59,7 @@ public void process(AgentSpan span, BooleanSupplier timeoutCheck) {
5859
} catch (Exception e) {
5960
log.debug("Error closing AppSecRequestContext", e);
6061
}
61-
sampler.counter.release();
62+
sampler.releaseOne();
6263
}
6364
}
6465

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

+234
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,234 @@
1+
package com.datadog.appsec.api.security
2+
3+
import com.datadog.appsec.event.EventProducerService
4+
import com.datadog.appsec.event.ExpiredSubscriberInfoException
5+
import com.datadog.appsec.gateway.AppSecRequestContext
6+
import datadog.trace.api.gateway.RequestContext
7+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan
8+
import datadog.trace.test.util.DDSpecification
9+
10+
class AppSecSpanPostProcessorTest extends DDSpecification {
11+
12+
void 'schema extracted on happy path'() {
13+
given:
14+
def sampler = Mock(ApiSecurityRequestSampler)
15+
def producer = Mock(EventProducerService)
16+
def subInfo = Mock(EventProducerService.DataSubscriberInfo)
17+
def span = Mock(AgentSpan)
18+
def reqCtx = Mock(RequestContext)
19+
def ctx = Mock(AppSecRequestContext)
20+
def processor = new AppSecSpanPostProcessor(sampler, producer)
21+
22+
when:
23+
processor.process(span, { false })
24+
25+
then:
26+
noExceptionThrown()
27+
1 * span.getRequestContext() >> reqCtx
28+
1 * reqCtx.getData(_) >> ctx
29+
1 * ctx.isKeepOpenForApiSecurityPostProcessing() >> true
30+
1 * sampler.sampleRequest(_) >> true
31+
1 * producer.getDataSubscribers(_) >> subInfo
32+
1 * subInfo.isEmpty() >> false
33+
1 * producer.publishDataEvent(_, ctx, _, _)
34+
1 * ctx.setKeepOpenForApiSecurityPostProcessing(false)
35+
1 * ctx.close()
36+
1 * sampler.releaseOne()
37+
0 * _
38+
}
39+
40+
void 'no schema extracted if sampling is false'() {
41+
given:
42+
def sampler = Mock(ApiSecurityRequestSampler)
43+
def producer = Mock(EventProducerService)
44+
def span = Mock(AgentSpan)
45+
def reqCtx = Mock(RequestContext)
46+
def ctx = Mock(AppSecRequestContext)
47+
def processor = new AppSecSpanPostProcessor(sampler, producer)
48+
49+
when:
50+
processor.process(span, { false })
51+
52+
then:
53+
noExceptionThrown()
54+
1 * span.getRequestContext() >> reqCtx
55+
1 * reqCtx.getData(_) >> ctx
56+
1 * ctx.isKeepOpenForApiSecurityPostProcessing() >> true
57+
1 * sampler.sampleRequest(_) >> false
58+
1 * ctx.setKeepOpenForApiSecurityPostProcessing(false)
59+
1 * ctx.close()
60+
1 * sampler.releaseOne()
61+
0 * _
62+
}
63+
64+
void 'permit is released even if request context close throws'() {
65+
given:
66+
def sampler = Mock(ApiSecurityRequestSampler)
67+
def producer = Mock(EventProducerService)
68+
def span = Mock(AgentSpan)
69+
def reqCtx = Mock(RequestContext)
70+
def ctx = Mock(AppSecRequestContext)
71+
def processor = new AppSecSpanPostProcessor(sampler, producer)
72+
73+
when:
74+
processor.process(span, { false })
75+
76+
then:
77+
noExceptionThrown()
78+
1 * span.getRequestContext() >> reqCtx
79+
1 * reqCtx.getData(_) >> ctx
80+
1 * ctx.isKeepOpenForApiSecurityPostProcessing() >> true
81+
1 * sampler.sampleRequest(_) >> true
82+
1 * producer.getDataSubscribers(_) >> null
83+
1 * ctx.setKeepOpenForApiSecurityPostProcessing(false)
84+
1 * ctx.close() >> { throw new RuntimeException() }
85+
1 * sampler.releaseOne()
86+
0 * _
87+
}
88+
89+
void 'context is cleaned up on timeout'() {
90+
given:
91+
def sampler = Mock(ApiSecurityRequestSampler)
92+
def producer = Mock(EventProducerService)
93+
def span = Mock(AgentSpan)
94+
def reqCtx = Mock(RequestContext)
95+
def ctx = Mock(AppSecRequestContext)
96+
def processor = new AppSecSpanPostProcessor(sampler, producer)
97+
98+
when:
99+
processor.process(span, { true })
100+
101+
then:
102+
noExceptionThrown()
103+
1 * span.getRequestContext() >> reqCtx
104+
1 * reqCtx.getData(_) >> ctx
105+
1 * ctx.isKeepOpenForApiSecurityPostProcessing() >> true
106+
1 * ctx.setKeepOpenForApiSecurityPostProcessing(false)
107+
1 * ctx.close()
108+
1 * sampler.releaseOne()
109+
0 * _
110+
}
111+
112+
void 'process null request context does nothing'() {
113+
given:
114+
def sampler = Mock(ApiSecurityRequestSampler)
115+
def producer = Mock(EventProducerService)
116+
def span = Mock(AgentSpan)
117+
def processor = new AppSecSpanPostProcessor(sampler, producer)
118+
119+
when:
120+
processor.process(span, { false })
121+
122+
then:
123+
noExceptionThrown()
124+
1 * span.getRequestContext() >> null
125+
0 * _
126+
}
127+
128+
void 'process null appsec request context does nothing'() {
129+
given:
130+
def sampler = Mock(ApiSecurityRequestSampler)
131+
def producer = Mock(EventProducerService)
132+
def span = Mock(AgentSpan)
133+
def reqCtx = Mock(RequestContext)
134+
def processor = new AppSecSpanPostProcessor(sampler, producer)
135+
136+
when:
137+
processor.process(span, { false })
138+
139+
then:
140+
noExceptionThrown()
141+
1 * span.getRequestContext() >> reqCtx
142+
1 * reqCtx.getData(_) >> null
143+
0 * _
144+
}
145+
146+
void 'process already closed context does nothing'() {
147+
given:
148+
def sampler = Mock(ApiSecurityRequestSampler)
149+
def producer = Mock(EventProducerService)
150+
def span = Mock(AgentSpan)
151+
def reqCtx = Mock(RequestContext)
152+
def ctx = Mock(AppSecRequestContext)
153+
def processor = new AppSecSpanPostProcessor(sampler, producer)
154+
155+
when:
156+
processor.process(span, { false })
157+
158+
then:
159+
noExceptionThrown()
160+
1 * span.getRequestContext() >> reqCtx
161+
1 * reqCtx.getData(_) >> ctx
162+
1 * ctx.isKeepOpenForApiSecurityPostProcessing() >> false
163+
0 * _
164+
}
165+
166+
void 'process throws on null span'() {
167+
given:
168+
def sampler = Mock(ApiSecurityRequestSampler)
169+
def producer = Mock(EventProducerService)
170+
def processor = new AppSecSpanPostProcessor(sampler, producer)
171+
172+
when:
173+
processor.process(null, { false })
174+
175+
then:
176+
thrown(NullPointerException)
177+
0 * _
178+
}
179+
180+
void 'empty event subscription does not break the process'() {
181+
given:
182+
def sampler = Mock(ApiSecurityRequestSampler)
183+
def producer = Mock(EventProducerService)
184+
def subInfo = Mock(EventProducerService.DataSubscriberInfo)
185+
def span = Mock(AgentSpan)
186+
def reqCtx = Mock(RequestContext)
187+
def ctx = Mock(AppSecRequestContext)
188+
def processor = new AppSecSpanPostProcessor(sampler, producer)
189+
190+
when:
191+
processor.process(span, { false })
192+
193+
then:
194+
noExceptionThrown()
195+
1 * span.getRequestContext() >> reqCtx
196+
1 * reqCtx.getData(_) >> ctx
197+
1 * ctx.isKeepOpenForApiSecurityPostProcessing() >> true
198+
1 * sampler.sampleRequest(_) >> true
199+
1 * producer.getDataSubscribers(_) >> subInfo
200+
1 * subInfo.isEmpty() >> true
201+
1 * ctx.setKeepOpenForApiSecurityPostProcessing(false)
202+
1 * ctx.close()
203+
1 * sampler.releaseOne()
204+
0 * _
205+
}
206+
207+
void 'expired event subscription does not break the process'() {
208+
given:
209+
def sampler = Mock(ApiSecurityRequestSampler)
210+
def producer = Mock(EventProducerService)
211+
def subInfo = Mock(EventProducerService.DataSubscriberInfo)
212+
def span = Mock(AgentSpan)
213+
def reqCtx = Mock(RequestContext)
214+
def ctx = Mock(AppSecRequestContext)
215+
def processor = new AppSecSpanPostProcessor(sampler, producer)
216+
217+
when:
218+
processor.process(span, { false })
219+
220+
then:
221+
noExceptionThrown()
222+
1 * span.getRequestContext() >> reqCtx
223+
1 * reqCtx.getData(_) >> ctx
224+
1 * ctx.isKeepOpenForApiSecurityPostProcessing() >> true
225+
1 * sampler.sampleRequest(_) >> true
226+
1 * producer.getDataSubscribers(_) >> subInfo
227+
1 * subInfo.isEmpty() >> false
228+
1 * producer.publishDataEvent(_, ctx, _, _) >> { throw new ExpiredSubscriberInfoException() }
229+
1 * ctx.setKeepOpenForApiSecurityPostProcessing(false)
230+
1 * ctx.close()
231+
1 * sampler.releaseOne()
232+
0 * _
233+
}
234+
}

‎internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/SpanPostProcessor.java

+2-2
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ public interface SpanPostProcessor {
2020
* much as possible. This method is guaranteed to be called even if post-processing of previous
2121
* spans have timed out.
2222
*/
23-
void process(AgentSpan span, BooleanSupplier timeoutCheck);
23+
void process(@Nonnull AgentSpan span, @Nonnull BooleanSupplier timeoutCheck);
2424

2525
class Holder {
2626
private static final SpanPostProcessor NOOP = new NoOpSpanPostProcessor();
@@ -32,6 +32,6 @@ class Holder {
3232

3333
class NoOpSpanPostProcessor implements SpanPostProcessor {
3434
@Override
35-
public void process(AgentSpan span, BooleanSupplier timeoutCheck) {}
35+
public void process(@Nonnull AgentSpan span, @Nonnull BooleanSupplier timeoutCheck) {}
3636
}
3737
}

0 commit comments

Comments
 (0)