Merge branch 'master' into alejandro.gonzalez/APPSEC-61873-3

Author: jandro996

dd-java-agent/agent-aiguard/src/main/java/com/datadog/aiguard/AIGuardInternal.java

@@ -73,6 +73,7 @@ public BadConfigurationException(final String message) {
   static final String META_STRUCT_MESSAGES = "messages";
   static final String META_STRUCT_CATEGORIES = "attack_categories";
   static final String META_STRUCT_SDS = "sds";
+  static final String META_STRUCT_TAG_PROBS = "tag_probs";
 
   public static void install() {
     final Config config = Config.get();
@@ -258,13 +259,18 @@ public Evaluation evaluate(final List<Message> messages, final Options options)
         final List<String> tags = (List<String>) result.get("tags");
         @SuppressWarnings("unchecked")
         final List<?> sdsFindings = (List<?>) result.get("sds_findings");
+        @SuppressWarnings("unchecked")
+        final Map<String, Number> tagProbs = (Map<String, Number>) result.get("tag_probs");
         span.setTag(ACTION_TAG, action);
         if (reason != null) {
           span.setTag(REASON_TAG, reason);
         }
         if (tags != null && !tags.isEmpty()) {
           metaStruct.put(META_STRUCT_CATEGORIES, tags);
         }
+        if (tagProbs != null && !tagProbs.isEmpty()) {
+          metaStruct.put(META_STRUCT_TAG_PROBS, tagProbs);
+        }
         if (sdsFindings != null && !sdsFindings.isEmpty()) {
           metaStruct.put(META_STRUCT_SDS, sdsFindings);
         }
@@ -273,9 +279,9 @@ public Evaluation evaluate(final List<Message> messages, final Options options)
         WafMetricCollector.get().aiGuardRequest(action, shouldBlock);
         if (shouldBlock) {
           span.setTag(BLOCKED_TAG, true);
-          throw new AIGuardAbortError(action, reason, tags, sdsFindings);
+          throw new AIGuardAbortError(action, reason, tags, tagProbs, sdsFindings);
         }
-        return new Evaluation(action, reason, tags, sdsFindings);
+        return new Evaluation(action, reason, tags, tagProbs, sdsFindings);
       }
     } catch (AIGuardAbortError e) {
       span.addThrowable(e);

dd-java-agent/agent-aiguard/src/test/groovy/com/datadog/aiguard/AIGuardInternalTests.groovy

@@ -168,7 +168,7 @@ class AIGuardInternalTests extends DDSpecification {
         return mockResponse(
           request,
           200,
-          [data: [attributes: [action: suite.action, reason: suite.reason, tags: suite.tags ?: [], is_blocking_enabled: suite.blocking]]]
+          [data: [attributes: [action: suite.action, reason: suite.reason, tags: suite.tags ?: [], tag_probs: suite.tagProbabilities ?: [:], is_blocking_enabled: suite.blocking]]]
           )
       }
     }
@@ -210,12 +210,14 @@ class AIGuardInternalTests extends DDSpecification {
       error.action == suite.action
       error.reason == suite.reason
       error.tags == suite.tags
+      error.tagProbabilities == suite.tagProbabilities
       error.sds == []
     } else {
       error == null
       eval.action == suite.action
       eval.reason == suite.reason
       eval.tags == suite.tags
+      eval.tagProbabilities == suite.tagProbabilities
       eval.sds == []
     }
     assertTelemetry('ai_guard.requests', "action:$suite.action", "block:$throwAbortError", 'error:false')
@@ -555,6 +557,9 @@ class AIGuardInternalTests extends DDSpecification {
     if (suite.tags) {
       assert meta.attack_categories == suite.tags
     }
+    if (suite.tagProbabilities)  {
+      assert meta.tag_probs == suite.tagProbabilities
+    }
     final receivedMessages = snakeCaseJson(meta.messages)
     final expectedMessages = snakeCaseJson(suite.messages)
     JSONAssert.assertEquals(expectedMessages, receivedMessages, JSONCompareMode.NON_EXTENSIBLE)
@@ -774,15 +779,17 @@ class AIGuardInternalTests extends DDSpecification {
     private final AIGuard.Action action
     private final String reason
     private final List<String> tags
+    private final Map<String, Double> tagProbabilities
     private final boolean blocking
     private final String description
     private final String target
     private final List<AIGuard.Message> messages
 
-    TestSuite(AIGuard.Action action, String reason, List<String> tags, boolean blocking, String description, String target, List<AIGuard.Message> messages) {
+    TestSuite(AIGuard.Action action, String reason, Map<String, Double> tagProbabilities, boolean blocking, String description, String target, List<AIGuard.Message> messages) {
       this.action = action
       this.reason = reason
-      this.tags = tags
+      this.tags = new ArrayList<>(tagProbabilities.keySet())
+      this.tagProbabilities = tagProbabilities
       this.blocking = blocking
       this.description = description
       this.target = target
@@ -791,9 +798,9 @@ class AIGuardInternalTests extends DDSpecification {
 
     static List<TestSuite> build() {
       def actionValues = [
-        [ALLOW, 'Go ahead', []],
-        [DENY, 'Nope', ['deny_everything', 'test_deny']],
-        [ABORT, 'Kill it with fire', ['alarm_tag', 'abort_everything']]
+        [ALLOW, 'Go ahead', [:]],
+        [DENY, 'Nope', ['deny_everything': 0.2D, 'test_deny': 0.8D]],
+        [ABORT, 'Kill it with fire', ['alarm_tag': 0.1D, 'abort_everything': 0.9D]]
       ]
       def blockingValues = [true, false]
       def suiteValues = [

dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/instrumentation/decorator/BaseDecoratorTest.groovy

@@ -5,11 +5,16 @@ import datadog.trace.bootstrap.instrumentation.api.AgentSpan
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext
 import datadog.trace.bootstrap.instrumentation.api.ErrorPriorities
 import datadog.trace.bootstrap.instrumentation.api.Tags
+import datadog.trace.config.inversion.ConfigHelper
 import datadog.trace.test.util.DDSpecification
 import spock.lang.Shared
 
 class BaseDecoratorTest extends DDSpecification {
 
+  def setupSpec() {
+    ConfigHelper.get().setConfigInversionStrict(ConfigHelper.StrictnessPolicy.TEST)
+  }
+
   @Shared
   def decorator = newDecorator()
 

dd-java-agent/agent-installer/src/test/groovy/datadog/trace/agent/test/DefaultInstrumenterForkedTest.groovy

@@ -2,6 +2,7 @@ package datadog.trace.agent.test
 
 import datadog.environment.EnvironmentVariables
 import datadog.trace.agent.tooling.InstrumenterModule
+import datadog.trace.config.inversion.ConfigHelper
 import datadog.trace.agent.tooling.bytebuddy.matcher.DDElementMatchers
 import datadog.trace.agent.tooling.bytebuddy.outline.TypePoolFacade
 import datadog.trace.test.util.DDSpecification
@@ -12,6 +13,10 @@ class DefaultInstrumenterForkedTest extends DDSpecification {
     DDElementMatchers.registerAsSupplier()
   }
 
+  def setupSpec() {
+    ConfigHelper.get().setConfigInversionStrict(ConfigHelper.StrictnessPolicy.TEST)
+  }
+
   def "default enabled"() {
     setup:
     def target = new TestDefaultInstrumenter("test")

dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/InstrumenterIndexTest.groovy

@@ -1,10 +1,15 @@
 package datadog.trace.agent.tooling
 
+import datadog.trace.config.inversion.ConfigHelper
 import datadog.trace.test.util.DDSpecification
 import spock.lang.Shared
 
 class InstrumenterIndexTest extends DDSpecification {
 
+  def setupSpec() {
+    ConfigHelper.get().setConfigInversionStrict(ConfigHelper.StrictnessPolicy.TEST)
+  }
+
   @Shared
   def unknownInstrumentation = new InstrumenterModule('unknown') {}
 

dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/BaseCallSiteTest.groovy

@@ -4,6 +4,7 @@ import datadog.trace.agent.tooling.bytebuddy.csi.Advices
 import datadog.trace.agent.tooling.bytebuddy.csi.CallSiteInstrumentation
 import datadog.trace.agent.tooling.bytebuddy.csi.CallSiteSupplier
 import datadog.trace.agent.tooling.bytebuddy.csi.CallSiteTransformer
+import datadog.trace.config.inversion.ConfigHelper
 import datadog.trace.test.util.DDSpecification
 import groovy.transform.CompileDynamic
 import net.bytebuddy.agent.builder.AgentBuilder
@@ -31,6 +32,10 @@ import static net.bytebuddy.matcher.ElementMatchers.named
 @CompileDynamic
 class BaseCallSiteTest extends DDSpecification {
 
+  def setupSpec() {
+    ConfigHelper.get().setConfigInversionStrict(ConfigHelper.StrictnessPolicy.TEST)
+  }
+
   protected CallSites mockCallSites(final byte type = BEFORE, final CallSiteAdvice advice, final Pointcut target, final String... helpers) {
     return Stub(CallSites) {
       accept(_ as CallSites.Container) >> {

dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/muzzle/MuzzleVersionScanPluginTest.groovy

@@ -2,6 +2,7 @@ package datadog.trace.agent.tooling.muzzle
 
 import datadog.trace.agent.tooling.Instrumenter
 import datadog.trace.agent.tooling.InstrumenterModule
+import datadog.trace.config.inversion.ConfigHelper
 import datadog.trace.test.util.DDSpecification
 import net.bytebuddy.matcher.ElementMatcher
 
@@ -20,6 +21,10 @@ import static datadog.trace.agent.tooling.muzzle.TestInstrumentationClasses.Vali
 
 class MuzzleVersionScanPluginTest extends DDSpecification {
 
+  def setupSpec() {
+    ConfigHelper.get().setConfigInversionStrict(ConfigHelper.StrictnessPolicy.TEST)
+  }
+
   def "test assertInstrumentationMuzzled advice"() {
     setup:
     def instrumentationLoader = new ServiceEnabledClassLoader(InstrumenterModule,

dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/InstrumentationSpecification.groovy

@@ -10,6 +10,7 @@ import static datadog.trace.api.config.TraceInstrumentationConfig.CODE_ORIGIN_FO
 import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.closePrevious
 import static datadog.trace.util.AgentThreadFactory.AgentThread.TASK_SCHEDULER
 
+
 import ch.qos.logback.classic.Level
 import ch.qos.logback.classic.util.ContextInitializer
 import com.datadog.debugger.agent.ClassesToRetransformFinder

dd-java-agent/instrumentation-testing/src/test/groovy/AgentTestRunnerTest.groovy

@@ -1,5 +1,6 @@
 import com.google.common.reflect.ClassPath
 import datadog.trace.agent.test.InstrumentationSpecification
+import datadog.trace.config.inversion.ConfigHelper
 import datadog.trace.agent.test.BootstrapClasspathSetupListener
 import datadog.trace.api.GlobalTracer
 import datadog.trace.api.Platform
@@ -26,6 +27,8 @@ class AgentTestRunnerTest extends InstrumentationSpecification {
   @Override
   void configurePreAgent() {
     super.configurePreAgent()
+    // Opt out of strict config validation - test module loads test instrumentations with fake names
+    ConfigHelper.get().setConfigInversionStrict(ConfigHelper.StrictnessPolicy.TEST)
 
     injectSysConfig(TRACE_CLASSES_EXCLUDE, "config.exclude.packagename.*, config.exclude.SomeClass,config.exclude.SomeClass\$NestedClass")
   }

dd-java-agent/instrumentation-testing/src/test/groovy/ConfigResetTest.groovy

@@ -1,10 +1,18 @@
 import datadog.environment.EnvironmentVariables
 import datadog.trace.agent.test.InstrumentationSpecification
+import datadog.trace.config.inversion.ConfigHelper
 import datadog.trace.api.Config
 import spock.lang.Shared
 
 class ConfigResetTest extends InstrumentationSpecification {
 
+  @Override
+  protected void configurePreAgent() {
+    super.configurePreAgent()
+    // Opt out of strict config validation - test module loads test instrumentations with fake names
+    ConfigHelper.get().setConfigInversionStrict(ConfigHelper.StrictnessPolicy.TEST)
+  }
+
   @Shared
   def sharedInstance = checkStaticAssertions()