Skip to content

Commit b6e4650

Browse files
sezen-datadogsezen-datadog
committed
Add new features for trace tagging rules
Signed-off-by: sezen.leblay <[email protected]>
1 parent ad6d6b9 commit b6e4650

File tree

11 files changed

+1038
-41
lines changed

11 files changed

+1038
-41
lines changed

‎dd-java-agent/appsec/build.gradle

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ dependencies {
1515
implementation project(':internal-api')
1616
implementation project(':communication')
1717
implementation project(':telemetry')
18-
implementation group: 'io.sqreen', name: 'libsqreen', version: '15.0.0'
18+
implementation group: 'io.sqreen', name: 'libsqreen', version: '15.0.2'
1919
implementation libs.moshi
2020

2121
testImplementation libs.bytebuddy

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

Lines changed: 79 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_AUTO_USER_INSTRUM_MODE;
66
import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_CUSTOM_BLOCKING_RESPONSE;
77
import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_CUSTOM_RULES;
8+
import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_DD_MULTICONFIG;
89
import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_DD_RULES;
910
import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_EXCLUSIONS;
1011
import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_EXCLUSION_DATA;
@@ -18,6 +19,7 @@
1819
import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_RASP_SSRF;
1920
import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_REQUEST_BLOCKING;
2021
import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_SESSION_FINGERPRINT;
22+
import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_TRACE_TAGGING_RULES;
2123
import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_TRUSTED_IPS;
2224
import static datadog.remoteconfig.Capabilities.CAPABILITY_ASM_USER_BLOCKING;
2325
import static datadog.remoteconfig.Capabilities.CAPABILITY_ENDPOINT_FINGERPRINT;
@@ -36,8 +38,8 @@
3638
import com.datadog.ddwaf.exception.InvalidRuleSetException;
3739
import com.datadog.ddwaf.exception.UnclassifiedWafException;
3840
import com.squareup.moshi.JsonAdapter;
39-
import com.squareup.moshi.Moshi;
40-
import com.squareup.moshi.Types;
41+
import com.squareup.moshi.JsonReader;
42+
import com.squareup.moshi.JsonWriter;
4143
import datadog.remoteconfig.ConfigurationEndListener;
4244
import datadog.remoteconfig.ConfigurationPoller;
4345
import datadog.remoteconfig.PollingRateHinter;
@@ -58,6 +60,7 @@
5860
import java.util.Collections;
5961
import java.util.HashMap;
6062
import java.util.HashSet;
63+
import java.util.LinkedHashMap;
6164
import java.util.List;
6265
import java.util.Map;
6366
import java.util.Set;
@@ -90,15 +93,12 @@ public class AppSecConfigServiceImpl implements AppSecConfigService {
9093
new WAFInitializationResultReporter();
9194
private final WAFStatsReporter statsReporter = new WAFStatsReporter();
9295

93-
private static final JsonAdapter<Map<String, Object>> ADAPTER =
94-
new Moshi.Builder()
95-
.build()
96-
.adapter(Types.newParameterizedType(Map.class, String.class, Object.class));
96+
private static final JsonAdapter<Object> ADAPTER = new SafeMapAdapter();
9797

9898
private boolean hasUserWafConfig;
9999
private boolean defaultConfigActivated;
100100
private final Set<String> usedDDWafConfigKeys = new HashSet<>();
101-
private final String DEFAULT_WAF_CONFIG_RULE = "DEFAULT_WAF_CONFIG";
101+
private final String DEFAULT_WAF_CONFIG_RULE = "ASM_DD/default";
102102
private String currentRuleVersion;
103103
private List<AppSecModule> modulesToUpdateVersionIn;
104104

@@ -129,6 +129,7 @@ private void subscribeConfigurationPoller() {
129129

130130
long capabilities =
131131
CAPABILITY_ASM_DD_RULES
132+
| CAPABILITY_ASM_DD_MULTICONFIG
132133
| CAPABILITY_ASM_IP_BLOCKING
133134
| CAPABILITY_ASM_EXCLUSIONS
134135
| CAPABILITY_ASM_EXCLUSION_DATA
@@ -140,7 +141,8 @@ private void subscribeConfigurationPoller() {
140141
| CAPABILITY_ENDPOINT_FINGERPRINT
141142
| CAPABILITY_ASM_SESSION_FINGERPRINT
142143
| CAPABILITY_ASM_NETWORK_FINGERPRINT
143-
| CAPABILITY_ASM_HEADER_FINGERPRINT;
144+
| CAPABILITY_ASM_HEADER_FINGERPRINT
145+
| CAPABILITY_ASM_TRACE_TAGGING_RULES;
144146
if (tracerConfig.isAppSecRaspEnabled()) {
145147
capabilities |= CAPABILITY_ASM_RASP_SQLI;
146148
capabilities |= CAPABILITY_ASM_RASP_SSRF;
@@ -185,7 +187,8 @@ public void accept(ConfigKey configKey, byte[] content, PollingRateHinter pollin
185187
}
186188
} else {
187189
Map<String, Object> contentMap =
188-
ADAPTER.fromJson(Okio.buffer(Okio.source(new ByteArrayInputStream(content))));
190+
(Map<String, Object>)
191+
ADAPTER.fromJson(Okio.buffer(Okio.source(new ByteArrayInputStream(content))));
189192
try {
190193
handleWafUpdateResultReport(configKey.toString(), contentMap);
191194
} catch (AppSecModule.AppSecModuleActivationException e) {
@@ -211,7 +214,7 @@ private class AppSecConfigChangesDDListener extends AppSecConfigChangesListener
211214
public void accept(ConfigKey configKey, byte[] content, PollingRateHinter pollingRateHinter)
212215
throws IOException {
213216
if (defaultConfigActivated) { // if we get any config, remove the default one
214-
log.debug("Removing default config");
217+
log.debug("Removing default config ASM_DD/default");
215218
try {
216219
wafBuilder.removeConfig(DEFAULT_WAF_CONFIG_RULE);
217220
} catch (UnclassifiedWafException e) {
@@ -306,7 +309,10 @@ private void subscribeAsmFeatures() {
306309
private void distributeSubConfigurations(
307310
String key, AppSecModuleConfigurer.Reconfiguration reconfiguration) {
308311
if (usedDDWafConfigKeys.isEmpty() && !defaultConfigActivated && !hasUserWafConfig) {
309-
// no config left in the WAF builder, add the default config
312+
// ASM_DD Failure Fallback: If none of the configurations obtained through ASM_DD were loaded
313+
// successfully,
314+
// libraries must revert back to the default configuration
315+
log.debug("No ASM_DD configurations loaded, falling back to default configuration");
310316
init();
311317
}
312318
for (Map.Entry<String, SubconfigListener> entry : subconfigListeners.entrySet()) {
@@ -427,7 +433,8 @@ private static Map<String, Object> loadDefaultWafConfig() throws IOException {
427433
throw new IOException("Resource " + DEFAULT_CONFIG_LOCATION + " not found");
428434
}
429435

430-
Map<String, Object> ret = ADAPTER.fromJson(Okio.buffer(Okio.source(is)));
436+
Map<String, Object> ret =
437+
(Map<String, Object>) ADAPTER.fromJson(Okio.buffer(Okio.source(is)));
431438

432439
StandardizedLogging._initialConfigSourceAndLibddwafVersion(log, "<bundled config>");
433440
if (log.isInfoEnabled()) {
@@ -444,7 +451,8 @@ private static Map<String, Object> loadUserWafConfig(Config tracerConfig) throws
444451
return null;
445452
}
446453
try (InputStream is = new FileInputStream(filename)) {
447-
Map<String, Object> ret = ADAPTER.fromJson(Okio.buffer(Okio.source(is)));
454+
Map<String, Object> ret =
455+
(Map<String, Object>) ADAPTER.fromJson(Okio.buffer(Okio.source(is)));
448456

449457
StandardizedLogging._initialConfigSourceAndLibddwafVersion(log, filename);
450458
if (log.isInfoEnabled()) {
@@ -473,6 +481,7 @@ public void close() {
473481
this.configurationPoller.removeCapabilities(
474482
CAPABILITY_ASM_ACTIVATION
475483
| CAPABILITY_ASM_DD_RULES
484+
| CAPABILITY_ASM_DD_MULTICONFIG
476485
| CAPABILITY_ASM_IP_BLOCKING
477486
| CAPABILITY_ASM_EXCLUSIONS
478487
| CAPABILITY_ASM_EXCLUSION_DATA
@@ -490,7 +499,8 @@ public void close() {
490499
| CAPABILITY_ENDPOINT_FINGERPRINT
491500
| CAPABILITY_ASM_SESSION_FINGERPRINT
492501
| CAPABILITY_ASM_NETWORK_FINGERPRINT
493-
| CAPABILITY_ASM_HEADER_FINGERPRINT);
502+
| CAPABILITY_ASM_HEADER_FINGERPRINT
503+
| CAPABILITY_ASM_TRACE_TAGGING_RULES);
494504
this.configurationPoller.removeListeners(Product.ASM_DD);
495505
this.configurationPoller.removeListeners(Product.ASM_DATA);
496506
this.configurationPoller.removeListeners(Product.ASM);
@@ -558,4 +568,59 @@ private static WafConfig createWafConfig(Config config) {
558568
}
559569
return wafConfig;
560570
}
571+
572+
private static class SafeMapAdapter extends JsonAdapter<Object> {
573+
@Override
574+
public Object fromJson(JsonReader reader) throws IOException {
575+
switch (reader.peek()) {
576+
case BEGIN_OBJECT:
577+
Map<String, Object> map = new LinkedHashMap<>();
578+
reader.beginObject();
579+
while (reader.hasNext()) {
580+
map.put(reader.nextName(), fromJson(reader));
581+
}
582+
reader.endObject();
583+
return map;
584+
585+
case BEGIN_ARRAY:
586+
List<Object> list = new ArrayList<>();
587+
reader.beginArray();
588+
while (reader.hasNext()) {
589+
list.add(fromJson(reader));
590+
}
591+
reader.endArray();
592+
return list;
593+
594+
case STRING:
595+
return reader.nextString();
596+
case NUMBER:
597+
String numberStr = reader.nextString();
598+
try {
599+
if (numberStr.contains(".")) {
600+
return Double.parseDouble(numberStr);
601+
} else {
602+
return Long.parseLong(numberStr);
603+
}
604+
} catch (NumberFormatException e) {
605+
// Fallback to string if parsing fails
606+
return numberStr;
607+
}
608+
609+
case BOOLEAN:
610+
return reader.nextBoolean();
611+
612+
case NULL:
613+
reader.nextNull();
614+
return null;
615+
616+
default:
617+
throw new IllegalStateException("Unexpected token: " + reader.peek());
618+
}
619+
}
620+
621+
@Override
622+
public void toJson(JsonWriter writer, Object value) throws IOException {
623+
throw new UnsupportedOperationException("Serialization not supported");
624+
}
625+
}
561626
}

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

Lines changed: 19 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -223,6 +223,11 @@ private void initOrUpdateWafHandle(AppSecModuleConfigurer.Reconfiguration reconf
223223
reconf.reloadSubscriptions();
224224
}
225225

226+
/**
227+
* Creates a rate limiter for AppSec events. The rate limiter accounts for when libddwaf returns
228+
* keep with a value of true, rather than when events are present, as specified in the technical
229+
* specification.
230+
*/
226231
private static RateLimiter getRateLimiter(Monitoring monitoring) {
227232
if (monitoring == null) {
228233
return null;
@@ -401,12 +406,13 @@ public void onDataAvailable(
401406
}
402407
}
403408
Collection<AppSecEvent> events = buildEvents(resultWithData);
409+
boolean isThrottled = reqCtx.isThrottled(rateLimiter);
404410

405-
if (!events.isEmpty()) {
406-
if (!reqCtx.isThrottled(rateLimiter)) {
411+
if (resultWithData.keep) {
412+
if (!isThrottled) {
407413
AgentSpan activeSpan = AgentTracer.get().activeSpan();
408414
if (activeSpan != null) {
409-
log.debug("Setting force-keep tag on the current span");
415+
log.debug("Setting force-keep tag and manual keep tag on the current span");
410416
// Keep event related span, because it could be ignored in case of
411417
// reduced datadog sampling rate.
412418
activeSpan.getLocalRootSpan().setTag(Tags.ASM_KEEP, true);
@@ -417,18 +423,19 @@ public void onDataAvailable(
417423
.getLocalRootSpan()
418424
.setTag(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.ASM);
419425
} else {
420-
// If active span is not available the ASM_KEEP tag will be set in the GatewayBridge
421-
// when the request ends
426+
// If active span is not available then we need to set manual keep in GatewayBridge
422427
log.debug("There is no active span available");
423428
}
424-
reqCtx.reportEvents(events);
425429
} else {
426430
log.debug("Rate limited WAF events");
427431
if (!gwCtx.isRasp) {
428432
reqCtx.setWafRateLimited();
429433
}
430434
}
431435
}
436+
if (resultWithData.events && !events.isEmpty() && !isThrottled) {
437+
reqCtx.reportEvents(events);
438+
}
432439

433440
if (flow.isBlocking()) {
434441
if (!gwCtx.isRasp) {
@@ -437,8 +444,8 @@ public void onDataAvailable(
437444
}
438445
}
439446

440-
if (resultWithData.derivatives != null) {
441-
reqCtx.reportDerivatives(resultWithData.derivatives);
447+
if (resultWithData.attributes != null && !resultWithData.attributes.isEmpty()) {
448+
reqCtx.reportDerivatives(resultWithData.attributes);
442449
}
443450
}
444451

@@ -559,6 +566,10 @@ private Waf.ResultWithData runWafTransient(
559566
private Collection<AppSecEvent> buildEvents(Waf.ResultWithData actionWithData) {
560567
Collection<WAFResultData> listResults;
561568
try {
569+
if (actionWithData.data == null || actionWithData.data.isEmpty()) {
570+
log.debug("WAF returned no data");
571+
return emptyList();
572+
}
562573
listResults = RES_JSON_ADAPTER.fromJson(actionWithData.data);
563574
} catch (IOException e) {
564575
throw new UndeclaredThrowableException(e);

0 commit comments

Comments
 (0)