Add standalone ASM billing support

dd-java-agent/agent-iast/src/main/java/com/datadog/iast/Reporter.java

@@ -98,6 +98,7 @@ private VulnerabilityBatch getOrCreateVulnerabilityBatch(final AgentSpan span) {
       // TODO: We need to check if we can have an API with more fine-grained semantics on why traces
       // are kept.
       segment.setTagTop(Tags.ASM_KEEP, true);
+      segment.setTagTop(Tags.PROPAGATED_APPSEC, true);
       return batch;
     }
 

dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/ReporterTest.groovy

@@ -80,6 +80,7 @@ class ReporterTest extends DDSpecification {
       ]
     }''', batch.toString(), true)
     1 * traceSegment.setTagTop('asm.keep', true)
+    1 * traceSegment.setTagTop('_dd.p.appsec', true)
     0 * _
   }
 
@@ -146,6 +147,7 @@ class ReporterTest extends DDSpecification {
       ]
     }''', batch.toString(), true)
     1 * traceSegment.setTagTop('asm.keep', true)
+    1 * traceSegment.setTagTop('_dd.p.appsec', true)
     0 * _
   }
 
@@ -269,6 +271,7 @@ class ReporterTest extends DDSpecification {
     1 * traceSegment.getDataTop('iast') >> null
     1 * traceSegment.setDataTop('iast', _ as VulnerabilityBatch)
     1 * traceSegment.setTagTop('asm.keep', true)
+    1 * traceSegment.setTagTop('_dd.p.appsec', true)
     1 * traceSegment.setTagTop('_dd.iast.enabled', 1)
     0 * _
   }

dd-java-agent/appsec/src/main/java/com/datadog/appsec/powerwaf/PowerWAFModule.java

@@ -481,6 +481,10 @@ public void onDataAvailable(
             // Keep event related span, because it could be ignored in case of
             // reduced datadog sampling rate.
             activeSpan.getLocalRootSpan().setTag(Tags.ASM_KEEP, true);
+            // If APM is disabled, inform downstream services that the current
+            // distributed trace contains at least one ASM event and must inherit
+            // the given force-keep priority
+            activeSpan.getLocalRootSpan().setTag(Tags.PROPAGATED_APPSEC, true);
           } else {
             // If active span is not available the ASK_KEEP tag will be set in the GatewayBridge
             // when the request ends

dd-smoke-tests/asm-standalone-billing/build.gradle

@@ -0,0 +1,31 @@
+plugins {
+  id 'java'
+  id 'org.springframework.boot' version '2.7.15'
+  id 'io.spring.dependency-management' version '1.0.15.RELEASE'
+  id 'java-test-fixtures'
+}
+
+apply from: "$rootDir/gradle/java.gradle"
+description = 'ASM Standalone Billing Tests.'
+
+java {
+  sourceCompatibility = '1.8'
+}
+
+repositories {
+  mavenCentral()
+}
+
+dependencies {
+  implementation 'org.springframework.boot:spring-boot-starter-web'
+  implementation group: 'io.opentracing', name: 'opentracing-api', version: '0.32.0'
+  implementation group: 'io.opentracing', name: 'opentracing-util', version: '0.32.0'
+  implementation project(':dd-trace-api')
+  testImplementation project(':dd-smoke-tests')
+  testImplementation(testFixtures(project(":dd-smoke-tests:iast-util")))
+}
+
+tasks.withType(Test).configureEach {
+  dependsOn "bootJar"
+  jvmArgs "-Ddatadog.smoketest.springboot.shadowJar.path=${tasks.bootJar.archiveFile.get()}"
+}

dd-smoke-tests/asm-standalone-billing/src/main/java/datadog/smoketest/asmstandalonebilling/AppConfig.java

@@ -0,0 +1,25 @@
+package datadog.smoketest.asmstandalonebilling;
+
+import java.util.EnumSet;
+import javax.servlet.ServletContext;
+import javax.servlet.ServletException;
+import javax.servlet.SessionTrackingMode;
+import org.springframework.boot.web.servlet.ServletContextInitializer;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+
+@Configuration
+public class AppConfig {
+  @Bean
+  public ServletContextInitializer servletContextInitializer() {
+    return new SessionTrackingConfig();
+  }
+
+  private class SessionTrackingConfig implements ServletContextInitializer {
+    @Override
+    public void onStartup(ServletContext servletContext) throws ServletException {
+      EnumSet<SessionTrackingMode> sessionTrackingModes = EnumSet.of(SessionTrackingMode.COOKIE);
+      servletContext.setSessionTrackingModes(sessionTrackingModes);
+    }
+  }
+}

dd-smoke-tests/asm-standalone-billing/src/main/java/datadog/smoketest/asmstandalonebilling/Controller.java

@@ -0,0 +1,84 @@
+package datadog.smoketest.asmstandalonebilling;
+
+import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
+import io.opentracing.Span;
+import io.opentracing.util.GlobalTracer;
+import java.io.IOException;
+import java.util.Map;
+import javax.servlet.http.HttpServletResponse;
+import org.springframework.http.MediaType;
+import org.springframework.http.ResponseEntity;
+import org.springframework.web.bind.annotation.*;
+import org.springframework.web.client.RestTemplate;
+
+@RestController
+@RequestMapping("/rest-api")
+public class Controller {
+
+  @GetMapping("/greetings")
+  public String greetings(
+      @RequestParam(name = "url", required = false) String url,
+      @RequestParam(name = "forceKeep", required = false) boolean forceKeep) {
+    if (forceKeep) {
+      forceKeepSpan();
+    }
+    if (url != null) {
+      RestTemplate restTemplate = new RestTemplate();
+      return restTemplate.getForObject(url, String.class);
+    }
+    return "Hello  I'm service " + System.getProperty("dd.service.name");
+  }
+
+  @GetMapping(value = "/returnheaders", produces = MediaType.APPLICATION_JSON_VALUE)
+  public ResponseEntity<Map<String, String>> returnheaders(
+      @RequestHeader Map<String, String> headers) {
+    return ResponseEntity.ok(headers);
+  }
+
+  @GetMapping("/appsec/{id}")
+  public String pathParam(
+      @PathVariable("id") String id,
+      @RequestParam(name = "url", required = false) String url,
+      @RequestParam(name = "forceKeep", required = false) boolean forceKeep) {
+    if (forceKeep) {
+      forceKeepSpan();
+    }
+    if (url != null) {
+      RestTemplate restTemplate = new RestTemplate();
+      return restTemplate.getForObject(url, String.class);
+    }
+    return id;
+  }
+
+  @GetMapping("/iast")
+  @SuppressFBWarnings
+  public void write(
+      @RequestParam(name = "injection", required = false) String injection,
+      @RequestParam(name = "url", required = false) String url,
+      @RequestParam(name = "forceKeep", required = false) boolean forceKeep,
+      final HttpServletResponse response) {
+    if (forceKeep) {
+      forceKeepSpan();
+    }
+    if (injection != null) {
+      try {
+        response.getWriter().write(injection);
+      } catch (IOException e) {
+        throw new RuntimeException(e);
+      }
+    }
+    if (url != null) {
+      RestTemplate restTemplate = new RestTemplate();
+      restTemplate.getForObject(url, String.class);
+    }
+  }
+
+  private String forceKeepSpan() {
+    final Span span = GlobalTracer.get().activeSpan();
+    if (span != null) {
+      span.setTag("manual.keep", true);
+      return span.context().toSpanId();
+    }
+    return null;
+  }
+}

dd-smoke-tests/asm-standalone-billing/src/main/java/datadog/smoketest/asmstandalonebilling/SpringbootApplication.java

@@ -0,0 +1,11 @@
+package datadog.smoketest.asmstandalonebilling;
+
+import org.springframework.boot.SpringApplication;
+import org.springframework.boot.autoconfigure.SpringBootApplication;
+
+@SpringBootApplication
+public class SpringbootApplication {
+  public static void main(String[] args) {
+    SpringApplication.run(SpringbootApplication.class, args);
+  }
+}

dd-smoke-tests/asm-standalone-billing/src/test/groovy/datadog/smoketest/asmstandalonebilling/AbstractAsmStandaloneBillingSmokeTest.groovy

@@ -0,0 +1,80 @@
+package datadog.smoketest.asmstandalonebilling
+
+import datadog.smoketest.AbstractServerSmokeTest
+import datadog.trace.api.sampling.PrioritySampling
+import datadog.trace.test.agent.decoder.DecodedTrace
+
+abstract class AbstractAsmStandaloneBillingSmokeTest extends AbstractServerSmokeTest {
+
+  @Override
+  File createTemporaryFile(int processIndex) {
+    return null
+  }
+
+  @Override
+  String logLevel() {
+    return 'debug'
+  }
+
+  @Override
+  Closure decodedTracesCallback() {
+    return {} // force traces decoding
+  }
+
+  protected ProcessBuilder createProcess(String[] properties){
+    createProcess(-1, properties)
+  }
+
+
+  protected ProcessBuilder createProcess(int processIndex, String[] properties){
+    def port = processIndex == -1 ? httpPort : httpPorts[processIndex]
+    String springBootShadowJar = System.getProperty("datadog.smoketest.springboot.shadowJar.path")
+    List<String> command = []
+    command.add(javaPath())
+    command.addAll(defaultJavaProperties)
+    command.addAll(properties)
+    command.addAll((String[]) ['-jar', springBootShadowJar, "--server.port=${port}"])
+    ProcessBuilder processBuilder = new ProcessBuilder(command)
+    processBuilder.directory(new File(buildDirectory))
+    // Spring will print all environment variables to the log, which may pollute it and affect log assertions.
+    processBuilder.environment().clear()
+    return processBuilder
+  }
+
+  protected DecodedTrace getServiceTrace(String serviceName) {
+    return traces.find { trace ->
+      trace.spans.find { span ->
+        span.service == serviceName
+      }
+    }
+  }
+
+  protected DecodedTrace getServiceTraceFromUrl(String url) {
+    return traces.find { trace ->
+      trace.spans.find { span ->
+        span.meta["http.url"] == url
+      }
+    }
+  }
+
+  protected checkRootSpanPrioritySampling(DecodedTrace trace, byte priority) {
+    return trace.spans[0].metrics['_sampling_priority_v1'] == priority
+  }
+
+  protected isSampledBySampler(DecodedTrace trace) {
+    def samplingPriority = trace.spans[0].metrics['_sampling_priority_v1']
+    return samplingPriority == PrioritySampling.SAMPLER_KEEP || samplingPriority == PrioritySampling.SAMPLER_DROP
+  }
+
+  protected hasAppsecPropagationTag(DecodedTrace trace) {
+    return trace.spans[0].meta['_dd.p.appsec'] == "1"
+  }
+
+  protected hasApmDisabledTag(DecodedTrace trace) {
+    return trace.spans[0].metrics['_dd.apm.enabled'] == 0
+  }
+
+  protected hasASMEvents(DecodedTrace trace){
+    return trace.spans[0].meta['_dd.iast.json'] != null || trace.spans[0].meta['_dd.appsec.json'] != null
+  }
+}