Skip to content

Commit 9ad3778

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

File tree

11 files changed

+1124
-38
lines changed

11 files changed

+1124
-38
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.1'
18+
implementation group: 'io.sqreen', name: 'libsqreen', version: '16.0.0'
1919
implementation libs.moshi
2020

2121
testImplementation libs.bytebuddy

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

Lines changed: 75 additions & 13 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;
@@ -37,8 +39,8 @@
3739
import com.datadog.ddwaf.exception.InvalidRuleSetException;
3840
import com.datadog.ddwaf.exception.UnclassifiedWafException;
3941
import com.squareup.moshi.JsonAdapter;
40-
import com.squareup.moshi.Moshi;
41-
import com.squareup.moshi.Types;
42+
import com.squareup.moshi.JsonReader;
43+
import com.squareup.moshi.JsonWriter;
4244
import datadog.remoteconfig.ConfigurationEndListener;
4345
import datadog.remoteconfig.ConfigurationPoller;
4446
import datadog.remoteconfig.PollingRateHinter;
@@ -60,6 +62,7 @@
6062
import java.util.ArrayList;
6163
import java.util.Collections;
6264
import java.util.HashMap;
65+
import java.util.LinkedHashMap;
6366
import java.util.List;
6467
import java.util.Map;
6568
import java.util.Set;
@@ -92,10 +95,7 @@ public class AppSecConfigServiceImpl implements AppSecConfigService {
9295
new WAFInitializationResultReporter();
9396
private final WAFStatsReporter statsReporter = new WAFStatsReporter();
9497

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

100100
private boolean hasUserWafConfig;
101101
private boolean defaultConfigActivated;
@@ -104,7 +104,7 @@ public class AppSecConfigServiceImpl implements AppSecConfigService {
104104
Collections.newSetFromMap(new ConcurrentHashMap<>());
105105
private final Set<String> ignoredConfigKeys =
106106
Collections.newSetFromMap(new ConcurrentHashMap<>());
107-
private final String DEFAULT_WAF_CONFIG_RULE = "DEFAULT_WAF_CONFIG";
107+
private final String DEFAULT_WAF_CONFIG_RULE = "ASM_DD/default";
108108
private String currentRuleVersion;
109109
private List<AppSecModule> modulesToUpdateVersionIn;
110110

@@ -137,6 +137,7 @@ private void subscribeConfigurationPoller() {
137137
private long getRulesAndDataCapabilities() {
138138
long capabilities =
139139
CAPABILITY_ASM_DD_RULES
140+
| CAPABILITY_ASM_DD_MULTICONFIG
140141
| CAPABILITY_ASM_IP_BLOCKING
141142
| CAPABILITY_ASM_EXCLUSIONS
142143
| CAPABILITY_ASM_EXCLUSION_DATA
@@ -148,7 +149,8 @@ private long getRulesAndDataCapabilities() {
148149
| CAPABILITY_ENDPOINT_FINGERPRINT
149150
| CAPABILITY_ASM_SESSION_FINGERPRINT
150151
| CAPABILITY_ASM_NETWORK_FINGERPRINT
151-
| CAPABILITY_ASM_HEADER_FINGERPRINT;
152+
| CAPABILITY_ASM_HEADER_FINGERPRINT
153+
| CAPABILITY_ASM_TRACE_TAGGING_RULES;
152154
if (tracerConfig.isAppSecRaspEnabled()) {
153155
capabilities |= CAPABILITY_ASM_RASP_SQLI;
154156
capabilities |= CAPABILITY_ASM_RASP_SSRF;
@@ -210,7 +212,8 @@ public void accept(ConfigKey configKey, byte[] content, PollingRateHinter pollin
210212
}
211213
final String key = configKey.toString();
212214
Map<String, Object> contentMap =
213-
ADAPTER.fromJson(Okio.buffer(Okio.source(new ByteArrayInputStream(content))));
215+
(Map<String, Object>)
216+
ADAPTER.fromJson(Okio.buffer(Okio.source(new ByteArrayInputStream(content))));
214217
if (contentMap == null || contentMap.isEmpty()) {
215218
ignoredConfigKeys.add(key);
216219
} else {
@@ -255,7 +258,7 @@ private class AppSecConfigChangesDDListener extends AppSecConfigChangesListener
255258
@Override
256259
protected void beforeApply(final String key, final Map<String, Object> config) {
257260
if (defaultConfigActivated) { // if we get any config, remove the default one
258-
log.debug("Removing default config");
261+
log.debug("Removing default config ASM_DD/default");
259262
try {
260263
wafBuilder.removeConfig(DEFAULT_WAF_CONFIG_RULE);
261264
} catch (UnclassifiedWafException e) {
@@ -466,7 +469,8 @@ private static Map<String, Object> loadDefaultWafConfig() throws IOException {
466469
throw new IOException("Resource " + DEFAULT_CONFIG_LOCATION + " not found");
467470
}
468471

469-
Map<String, Object> ret = ADAPTER.fromJson(Okio.buffer(Okio.source(is)));
472+
Map<String, Object> ret =
473+
(Map<String, Object>) ADAPTER.fromJson(Okio.buffer(Okio.source(is)));
470474

471475
StandardizedLogging._initialConfigSourceAndLibddwafVersion(log, "<bundled config>");
472476
if (log.isInfoEnabled()) {
@@ -483,7 +487,8 @@ private static Map<String, Object> loadUserWafConfig(Config tracerConfig) throws
483487
return null;
484488
}
485489
try (InputStream is = new FileInputStream(filename)) {
486-
Map<String, Object> ret = ADAPTER.fromJson(Okio.buffer(Okio.source(is)));
490+
Map<String, Object> ret =
491+
(Map<String, Object>) ADAPTER.fromJson(Okio.buffer(Okio.source(is)));
487492

488493
StandardizedLogging._initialConfigSourceAndLibddwafVersion(log, filename);
489494
if (log.isInfoEnabled()) {
@@ -512,6 +517,7 @@ public void close() {
512517
this.configurationPoller.removeCapabilities(
513518
CAPABILITY_ASM_ACTIVATION
514519
| CAPABILITY_ASM_DD_RULES
520+
| CAPABILITY_ASM_DD_MULTICONFIG
515521
| CAPABILITY_ASM_IP_BLOCKING
516522
| CAPABILITY_ASM_EXCLUSIONS
517523
| CAPABILITY_ASM_EXCLUSION_DATA
@@ -529,7 +535,8 @@ public void close() {
529535
| CAPABILITY_ENDPOINT_FINGERPRINT
530536
| CAPABILITY_ASM_SESSION_FINGERPRINT
531537
| CAPABILITY_ASM_NETWORK_FINGERPRINT
532-
| CAPABILITY_ASM_HEADER_FINGERPRINT);
538+
| CAPABILITY_ASM_HEADER_FINGERPRINT
539+
| CAPABILITY_ASM_TRACE_TAGGING_RULES);
533540
this.configurationPoller.removeListeners(Product.ASM_DD);
534541
this.configurationPoller.removeListeners(Product.ASM_DATA);
535542
this.configurationPoller.removeListeners(Product.ASM);
@@ -601,4 +608,59 @@ private static WafConfig createWafConfig(Config config) {
601608
}
602609
return wafConfig;
603610
}
611+
612+
private static class SafeMapAdapter extends JsonAdapter<Object> {
613+
@Override
614+
public Object fromJson(JsonReader reader) throws IOException {
615+
switch (reader.peek()) {
616+
case BEGIN_OBJECT:
617+
Map<String, Object> map = new LinkedHashMap<>();
618+
reader.beginObject();
619+
while (reader.hasNext()) {
620+
map.put(reader.nextName(), fromJson(reader));
621+
}
622+
reader.endObject();
623+
return map;
624+
625+
case BEGIN_ARRAY:
626+
List<Object> list = new ArrayList<>();
627+
reader.beginArray();
628+
while (reader.hasNext()) {
629+
list.add(fromJson(reader));
630+
}
631+
reader.endArray();
632+
return list;
633+
634+
case STRING:
635+
return reader.nextString();
636+
case NUMBER:
637+
String numberStr = reader.nextString();
638+
try {
639+
if (numberStr.contains(".")) {
640+
return Double.parseDouble(numberStr);
641+
} else {
642+
return Long.parseLong(numberStr);
643+
}
644+
} catch (NumberFormatException e) {
645+
// Fallback to string if parsing fails
646+
return numberStr;
647+
}
648+
649+
case BOOLEAN:
650+
return reader.nextBoolean();
651+
652+
case NULL:
653+
reader.nextNull();
654+
return null;
655+
656+
default:
657+
throw new IllegalStateException("Unexpected token: " + reader.peek());
658+
}
659+
}
660+
661+
@Override
662+
public void toJson(JsonWriter writer, Object value) throws IOException {
663+
throw new UnsupportedOperationException("Serialization not supported");
664+
}
665+
}
604666
}

‎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
@@ -224,6 +224,11 @@ private void initOrUpdateWafHandle(AppSecModuleConfigurer.Reconfiguration reconf
224224
reconf.reloadSubscriptions();
225225
}
226226

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

406-
if (!events.isEmpty()) {
407-
if (!reqCtx.isThrottled(rateLimiter)) {
412+
if (resultWithData.keep) {
413+
if (!isThrottled) {
408414
AgentSpan activeSpan = AgentTracer.get().activeSpan();
409415
if (activeSpan != null) {
410-
log.debug("Setting force-keep tag on the current span");
416+
log.debug("Setting force-keep tag and manual keep tag on the current span");
411417
// Keep event related span, because it could be ignored in case of
412418
// reduced datadog sampling rate.
413419
activeSpan.getLocalRootSpan().setTag(Tags.ASM_KEEP, true);
@@ -418,18 +424,19 @@ public void onDataAvailable(
418424
.getLocalRootSpan()
419425
.setTag(Tags.PROPAGATED_TRACE_SOURCE, ProductTraceSource.ASM);
420426
} else {
421-
// If active span is not available the ASM_KEEP tag will be set in the GatewayBridge
422-
// when the request ends
427+
// If active span is not available then we need to set manual keep in GatewayBridge
423428
log.debug("There is no active span available");
424429
}
425-
reqCtx.reportEvents(events);
426430
} else {
427431
log.debug("Rate limited WAF events");
428432
if (!gwCtx.isRasp) {
429433
reqCtx.setWafRateLimited();
430434
}
431435
}
432436
}
437+
if (resultWithData.events && !events.isEmpty() && !isThrottled) {
438+
reqCtx.reportEvents(events);
439+
}
433440

434441
if (flow.isBlocking()) {
435442
if (!gwCtx.isRasp) {
@@ -438,8 +445,8 @@ public void onDataAvailable(
438445
}
439446
}
440447

441-
if (resultWithData.derivatives != null) {
442-
reqCtx.reportDerivatives(resultWithData.derivatives);
448+
if (resultWithData.attributes != null && !resultWithData.attributes.isEmpty()) {
449+
reqCtx.reportDerivatives(resultWithData.attributes);
443450
}
444451
}
445452

@@ -564,6 +571,10 @@ private Collection<AppSecEvent> buildEvents(Waf.ResultWithData actionWithData) {
564571
}
565572
Collection<WAFResultData> listResults;
566573
try {
574+
if (actionWithData.data == null || actionWithData.data.isEmpty()) {
575+
log.debug("WAF returned no data");
576+
return emptyList();
577+
}
567578
listResults = RES_JSON_ADAPTER.fromJson(actionWithData.data);
568579
} catch (IOException e) {
569580
throw new UndeclaredThrowableException(e);

0 commit comments

Comments
 (0)