Fix org.json iast instrumentation test for latest dependency

dd-java-agent/instrumentation/iast-instrumenter/src/main/resources/datadog/trace/instrumentation/iastinstrumenter/iast_exclusion.trie

@@ -256,6 +256,8 @@
 1 org.jooq.*
 1 org.jruby.*
 1 org.json.*
+#Needed for JSON propagation
+2 org.json.JSONTokener
 1 org.jsoup.*
 1 org.junit.*
 1 org.jvnet.hk2.*

dd-java-agent/instrumentation/java-io/src/main/java/datadog/trace/instrumentation/java/lang/InputStreamReaderCallSite.java

@@ -12,6 +12,7 @@
 @CallSite(spi = IastCallSites.class)
 public class InputStreamReaderCallSite {
 
+  @CallSite.After("void java.io.InputStreamReader.<init>(java.io.InputStream)")
   @CallSite.After(
       "void java.io.InputStreamReader.<init>(java.io.InputStream, java.nio.charset.Charset)")
   public static InputStreamReader afterInit(

dd-java-agent/instrumentation/java-io/src/test/groovy/datadog/trace/instrumentation/java/io/InputStreamReaderCallSiteTest.groovy

@@ -14,10 +14,17 @@ class InputStreamReaderCallSiteTest extends  BaseIoCallSiteTest{
     InstrumentationBridge.registerIastModule(iastModule)
 
     when:
-    TestInputStreamReaderSuite.init(new ByteArrayInputStream("test".getBytes()), Charset.defaultCharset())
+    TestInputStreamReaderSuite.init(*args)
 
     then:
     1 * iastModule.taintObjectIfTainted(_ as InputStreamReader, _ as InputStream)
     0 * _
+
+    where:
+    args << [
+      [new ByteArrayInputStream("test".getBytes()), Charset.defaultCharset()],
+      // InputStream input
+      [new ByteArrayInputStream("test".getBytes())]// Reader input
+    ]
   }
 }

dd-java-agent/instrumentation/java-io/src/test/java/foo/bar/TestInputStreamReaderSuite.java

@@ -9,4 +9,8 @@ public class TestInputStreamReaderSuite {
   public static InputStreamReader init(final InputStream in, Charset charset) {
     return new InputStreamReader(in, charset);
   }
+
+  public static InputStreamReader init(final InputStream in) {
+    return new InputStreamReader(in);
+  }
 }

dd-java-agent/instrumentation/org-json/build.gradle

@@ -1,10 +1,25 @@
 muzzle {
   pass {
+    name = 'all'
     group = 'org.json'
     module = 'json'
     versions = '[20070829, ]'
     assertInverse = true
   }
+  pass {
+    name = 'before_20241224'
+    group = 'org.json'
+    module = 'json'
+    versions = '[20070829, 20241224)'
+    assertInverse = true
+  }
+  pass {
+    name = 'after_20241224'
+    group = 'org.json'
+    module = 'json'
+    versions = '[20241224, ]'
+    assertInverse = true
+  }
 }
 
 apply from: "$rootDir/gradle/java.gradle"
@@ -16,7 +31,7 @@ dependencies {
   testImplementation group: 'org.json', name: 'json', version: '20230227'
 
   testRuntimeOnly project(':dd-java-agent:instrumentation:iast-instrumenter')
+  testRuntimeOnly project(':dd-java-agent:instrumentation:java-io') //Needed for Reader
 
-  //FIXME: ASM
-  latestDepTestImplementation group: 'org.json', name: 'json', version: '20240303'
+  latestDepTestImplementation group: 'org.json', name: 'json', version: '+'
 }

dd-java-agent/instrumentation/org-json/gradle.lockfile

@@ -116,7 +116,7 @@ org.hamcrest:hamcrest-core:1.3=latestDepTestCompileClasspath,latestDepTestRuntim
 org.hamcrest:hamcrest:2.2=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath
 org.jctools:jctools-core:3.3.0=instrumentPluginClasspath,latestDepTestRuntimeClasspath,muzzleTooling,runtimeClasspath,testRuntimeClasspath
 org.json:json:20230227=compileClasspath,testCompileClasspath,testRuntimeClasspath
-org.json:json:20240303=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath
+org.json:json:20250107=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath
 org.junit.jupiter:junit-jupiter-api:5.9.2=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath
 org.junit.jupiter:junit-jupiter-engine:5.9.2=latestDepTestRuntimeClasspath,testRuntimeClasspath
 org.junit.platform:junit-platform-commons:1.9.2=latestDepTestCompileClasspath,latestDepTestRuntimeClasspath,testCompileClasspath,testRuntimeClasspath

dd-java-agent/instrumentation/org-json/src/main/java/datadog/trace/instrumentation/json/JSONArrayInstrumentation.java

@@ -5,6 +5,7 @@
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
 import static net.bytebuddy.matcher.ElementMatchers.returns;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
 import com.google.auto.service.AutoService;
@@ -14,8 +15,6 @@
 import datadog.trace.api.iast.Propagation;
 import datadog.trace.api.iast.propagation.PropagationModule;
 import net.bytebuddy.asm.Advice;
-import org.json.JSONArray;
-import org.json.JSONObject;
 
 @AutoService(InstrumenterModule.class)
 public class JSONArrayInstrumentation extends InstrumenterModule.Iast
@@ -25,6 +24,11 @@ public JSONArrayInstrumentation() {
     super("org-json");
   }
 
+  @Override
+  public String muzzleDirective() {
+    return "all";
+  }
+
   @Override
   public String instrumentedType() {
     return "org.json.JSONArray";
@@ -33,18 +37,11 @@ public String instrumentedType() {
   @Override
   public void methodAdvice(MethodTransformer transformer) {
     transformer.applyAdvice(
-        isConstructor().and(takesArguments(String.class)),
+        isConstructor().and(takesArguments(0)).and(takesArgument(0, named("org.json.JSONTokener"))),
         getClass().getName() + "$ConstructorAdvice");
-    transformer.applyAdvice(
-        isMethod()
-            .and(isPublic())
-            .and(returns(Object.class))
-            .and(named("get"))
-            .and(takesArguments(1)),
-        getClass().getName() + "$GetAdvice");
     transformer.applyAdvice(
         isMethod().and(isPublic()).and(returns(Object.class)).and(named("opt")),
-        getClass().getName() + "$OptAdvice");
+        packageName + ".OptAdvice");
   }
 
   public static class ConstructorAdvice {
@@ -57,44 +54,4 @@ public static void afterInit(@Advice.This Object self, @Advice.Argument(0) final
       }
     }
   }
-
-  public static class GetAdvice {
-    @Advice.OnMethodExit(suppress = Throwable.class)
-    @Propagation
-    public static void afterMethod(@Advice.This Object self, @Advice.Return final Object result) {
-      boolean isString = result instanceof String;
-      boolean isJson = !isString && (result instanceof JSONObject || result instanceof JSONArray);
-      if (!isString && !isJson) {
-        return;
-      }
-      final PropagationModule iastModule = InstrumentationBridge.PROPAGATION;
-      if (iastModule != null) {
-        if (isString) {
-          iastModule.taintStringIfTainted((String) result, self);
-        } else {
-          iastModule.taintObjectIfTainted(result, self);
-        }
-      }
-    }
-  }
-
-  public static class OptAdvice {
-    @Advice.OnMethodExit(suppress = Throwable.class)
-    @Propagation
-    public static void afterMethod(@Advice.This Object self, @Advice.Return final Object result) {
-      boolean isString = result instanceof String;
-      boolean isJson = !isString && (result instanceof JSONObject || result instanceof JSONArray);
-      if (!isString && !isJson) {
-        return;
-      }
-      final PropagationModule iastModule = InstrumentationBridge.PROPAGATION;
-      if (iastModule != null) {
-        if (isString) {
-          iastModule.taintStringIfTainted((String) result, self);
-        } else {
-          iastModule.taintObjectIfTainted(result, self);
-        }
-      }
-    }
-  }
 }

dd-java-agent/instrumentation/org-json/src/main/java/datadog/trace/instrumentation/json/JSONCookieInstrumentation.java

@@ -19,6 +19,11 @@ public JSONCookieInstrumentation() {
     super("org-json");
   }
 
+  @Override
+  public String muzzleDirective() {
+    return "all";
+  }
+
   @Override
   public String instrumentedType() {
     return "org.json.Cookie";

dd-java-agent/instrumentation/org-json/src/main/java/datadog/trace/instrumentation/json/JSONObject20241224Instrumentation.java

@@ -0,0 +1,84 @@
+package datadog.trace.instrumentation.json;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.ClassLoaderMatchers.hasClassNamed;
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.isPublic;
+import static net.bytebuddy.matcher.ElementMatchers.returns;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.iast.InstrumentationBridge;
+import datadog.trace.api.iast.Propagation;
+import datadog.trace.api.iast.propagation.PropagationModule;
+import java.util.Map;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@AutoService(InstrumenterModule.class)
+public class JSONObject20241224Instrumentation extends InstrumenterModule.Iast
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+  public JSONObject20241224Instrumentation() {
+    super("org-json");
+  }
+
+  static final ElementMatcher.Junction<ClassLoader> AFTER_20241224 =
+      hasClassNamed("org.json.StringBuilderWriter");
+
+  @Override
+  public String muzzleDirective() {
+    return "after_20241224";
+  }
+
+  @Override
+  public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
+    return AFTER_20241224;
+  }
+
+  @Override
+  public String instrumentedType() {
+    return "org.json.JSONObject";
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    // public JSONObject(JSONTokener x, JSONParserConfiguration jsonParserConfiguration)
+    transformer.applyAdvice(
+        isConstructor()
+            .and(takesArguments(2))
+            .and(takesArgument(0, named("org.json.JSONTokener")))
+            .and(takesArgument(1, named("org.json.JSONParserConfiguration"))),
+        getClass().getName() + "$ConstructorAdvice");
+    // private JSONObject(Map<?, ?> m, int recursionDepth, JSONParserConfiguration
+    // jsonParserConfiguration)
+    transformer.applyAdvice(
+        isConstructor()
+            .and(takesArguments(3))
+            .and(takesArgument(0, Map.class))
+            .and(takesArgument(1, int.class))
+            .and(takesArgument(2, named("org.json.JSONParserConfiguration"))),
+        getClass().getName() + "$ConstructorAdvice");
+    transformer.applyAdvice(
+        isMethod()
+            .and(isPublic())
+            .and(returns(Object.class))
+            .and(named("opt"))
+            .and(takesArguments(String.class)),
+        packageName + ".OptAdvice");
+  }
+
+  public static class ConstructorAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    @Propagation
+    public static void afterInit(@Advice.This Object self, @Advice.Argument(0) final Object input) {
+      final PropagationModule iastModule = InstrumentationBridge.PROPAGATION;
+      if (iastModule != null && input != null) {
+        iastModule.taintObjectIfTainted(self, input);
+      }
+    }
+  }
+}