Fix org.json iast instrumentation test for latest dependency (#8347)

What Does This Do
Rework tests to be more accurate and avoid problems with new versions
Remove redundant instrumentation as get methods always call opt methods
Be more accurate with instrumented constructors as they call each other internally
Improve java.io.InputStream propagation
Motivation
Fix latest dependency test for org.json (20250107)

Author: jandro996

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);
+      }
+    }
+  }
+}