Refactor UntrustedDeserializationModule call to onObject

DataDog · Aug 21, 2024 · e79cb51 · e79cb51
1 parent 99cfdba
commit e79cb51
Show file tree

Hide file tree

Showing 5 changed files with 15 additions and 17 deletions.
diff --git a/...nt/agent-iast/src/main/java/com/datadog/iast/sink/UntrustedDeserializationModuleImpl.java b/...nt/agent-iast/src/main/java/com/datadog/iast/sink/UntrustedDeserializationModuleImpl.java
@@ -3,7 +3,6 @@
 import com.datadog.iast.Dependencies;
 import com.datadog.iast.model.VulnerabilityType;
 import datadog.trace.api.iast.sink.UntrustedDeserializationModule;
-import java.io.InputStream;
 import javax.annotation.Nullable;
 
 public class UntrustedDeserializationModuleImpl extends SinkModuleBase
@@ -14,10 +13,10 @@ public UntrustedDeserializationModuleImpl(final Dependencies dependencies) {
   }
 
   @Override
-  public void onInputStream(@Nullable InputStream is) {
-    if (is == null) {
+  public void onObject(@Nullable Object object) {
+    if (object == null) {
       return;
     }
-    checkInjection(VulnerabilityType.UNTRUSTED_DESERIALIZATION, is);
+    checkInjection(VulnerabilityType.UNTRUSTED_DESERIALIZATION, object);
   }
 }
diff --git a/...gent-iast/src/test/groovy/com/datadog/iast/sink/UntrustedDeserializationModuleTest.groovy b/...gent-iast/src/test/groovy/com/datadog/iast/sink/UntrustedDeserializationModuleTest.groovy
@@ -22,29 +22,29 @@ class UntrustedDeserializationModuleTest extends IastModuleImplTestBase {
     return Mock(Reporter)
   }
 
-  void 'test null value'() {
+  void 'test null value with object null'() {
     when:
-    module.onInputStream(null)
+    module.onObject(null)
 
     then:
     0 * _
   }
 
   void 'test untrusted deserialization detection' () {
     setup:
-    def inputStream = Mock(InputStream)
+    def object = Mock(Object)
 
     when:
-    module.onInputStream(inputStream)
+    module.onObject(object)
 
-    then: 'without tainted input stream'
+    then: 'without tainted object'
     0 * reporter.report(_, _)
 
     when:
-    taint(inputStream)
-    module.onInputStream(inputStream)
+    taint(object)
+    module.onObject(object)
 
-    then: 'with tainted input stream'
+    then: 'with tainted object'
     1 * reporter.report(_, { Vulnerability vul -> vul.type == VulnerabilityType.UNTRUSTED_DESERIALIZATION})
   }
 

diff --git a/...a-io/src/main/java/datadog/trace/instrumentation/java/lang/ObjectInputStreamCallSite.java b/...a-io/src/main/java/datadog/trace/instrumentation/java/lang/ObjectInputStreamCallSite.java
@@ -18,7 +18,7 @@ public static void beforeConstructorUntrusted(@CallSite.Argument(0) final InputS
 
     if (module != null) {
       try {
-        module.onInputStream(is);
+        module.onObject(is);
       } catch (Throwable e) {
         module.onUnexpectedException("before constructor untrusted threw", e);
       }

diff --git a/...rc/test/groovy/datadog/trace/instrumentation/java/io/ObjectInputStreamCallSiteTest.groovy b/...rc/test/groovy/datadog/trace/instrumentation/java/io/ObjectInputStreamCallSiteTest.groovy
@@ -12,7 +12,7 @@ class ObjectInputStreamCallSiteTest extends AgentTestRunner {
     injectSysConfig('dd.iast.enabled', 'true')
   }
 
-  void 'test onInputStream'() {
+  void 'test onObject'() {
     setup:
     final module = Mock(UntrustedDeserializationModule)
     InstrumentationBridge.registerIastModule(module)
@@ -23,6 +23,6 @@ class ObjectInputStreamCallSiteTest extends AgentTestRunner {
     TestObjectInputStreamSuite.init(inputStream)
 
     then:
-    1 * module.onInputStream(_)
+    1 * module.onObject(_)
   }
 }
diff --git a/internal-api/src/main/java/datadog/trace/api/iast/sink/UntrustedDeserializationModule.java b/internal-api/src/main/java/datadog/trace/api/iast/sink/UntrustedDeserializationModule.java
@@ -1,10 +1,9 @@
 package datadog.trace.api.iast.sink;
 
 import datadog.trace.api.iast.IastModule;
-import java.io.InputStream;
 import javax.annotation.Nullable;
 
 public interface UntrustedDeserializationModule extends IastModule {
 
-  void onInputStream(@Nullable InputStream is);
+  void onObject(@Nullable Object object);
 }