Move to ProcessImplInstrumentation approach

DataDog · jandro996 · Dec 19, 2024 · Sep 13, 2024 · Nov 28, 2024 · Nov 29, 2024
commit b488d8a3c7689755d2c3c7e61bc440a37264857a
@@ -260,7 +260,7 @@ private Flow<Void> onNetworkConnection(RequestContext ctx_, String url) {
     }
   }
 
-  private Flow<Void> onShellCmd(RequestContext ctx_, Object command) {
+  private Flow<Void> onShellCmd(RequestContext ctx_, String[] command) {
     AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC);
     if (ctx == null) {
       return NoopFlow.INSTANCE;

@@ -92,7 +92,7 @@ class GatewayBridgeSpecification extends DDSpecification {
   TriFunction<RequestContext, UserIdCollectionMode, String, Flow<Void>> userIdCB
   TriFunction<RequestContext, UserIdCollectionMode, String, Flow<Void>> loginSuccessCB
   TriFunction<RequestContext, UserIdCollectionMode, String, Flow<Void>> loginFailureCB
-  BiFunction<RequestContext, Object, Flow<Void>> shellCmdCB
+  BiFunction<RequestContext, String[], Flow<Void>> shellCmdCB
 
   void setup() {
     callInitAndCaptureCBs()

diff --git a/...va-lang/src/main/java/datadog/trace/instrumentation/java/lang/ProcessBuilderCallSite.java b/...va-lang/src/main/java/datadog/trace/instrumentation/java/lang/ProcessBuilderCallSite.java
@@ -1,7 +1,6 @@
 package datadog.trace.instrumentation.java.lang;
 
 import datadog.trace.agent.tooling.csi.CallSite;
-import datadog.trace.api.appsec.RaspCallSites;
 import datadog.trace.api.iast.IastCallSites;
 import datadog.trace.api.iast.InstrumentationBridge;
 import datadog.trace.api.iast.Sink;
@@ -12,9 +11,7 @@
 
 // TODO deal with the environment
 @Sink(VulnerabilityTypes.COMMAND_INJECTION)
-@CallSite(
-    spi = {IastCallSites.class, RaspCallSites.class},
-    helpers = ShellCmdRaspHelper.class)
+@CallSite(spi = IastCallSites.class)
 public class ProcessBuilderCallSite {
 
   @CallSite.Before("java.lang.Process java.lang.ProcessBuilder.start()")
@@ -25,23 +22,14 @@ public static void beforeStart(@CallSite.This @Nullable final ProcessBuilder sel
     // be careful when fetching the environment as it does mutate the instance
     final List<String> cmd = self.command();
     if (cmd != null && cmd.size() > 0) {
-      iastCallback(cmd);
-      raspCallback(cmd);
-    }
-  }
-
-  private static void iastCallback(List<String> cmd) {
-    final CommandInjectionModule module = InstrumentationBridge.COMMAND_INJECTION;
-    if (module != null) {
-      try {
-        module.onProcessBuilderStart(cmd);
-      } catch (final Throwable e) {
-        module.onUnexpectedException("beforeStart threw", e);
+      final CommandInjectionModule module = InstrumentationBridge.COMMAND_INJECTION;
+      if (module != null) {
+        try {
+          module.onProcessBuilderStart(cmd);
+        } catch (final Throwable e) {
+          module.onUnexpectedException("beforeStart threw", e);
+        }
       }
     }
   }
-
-  private static void raspCallback(List<String> list) {
-    ShellCmdRaspHelper.INSTANCE.beforeShellCmd(list.toArray(new String[0]));
-  }
 }
@@ -23,6 +23,7 @@ public static AgentSpan startSpan(@Advice.Argument(0) final String[] command) th
     span.setResourceName(ProcessImplInstrumentationHelpers.determineResource(command));
     span.setTag("component", "subprocess");
     ProcessImplInstrumentationHelpers.setTags(span, command);
+    ProcessImplInstrumentationHelpers.shiRaspCheck(command);
     return span;
   }
 

diff --git a/...tion/java-lang/src/main/java/datadog/trace/instrumentation/java/lang/RuntimeCallSite.java b/...tion/java-lang/src/main/java/datadog/trace/instrumentation/java/lang/RuntimeCallSite.java
@@ -1,7 +1,6 @@
 package datadog.trace.instrumentation.java.lang;
 
 import datadog.trace.agent.tooling.csi.CallSite;
-import datadog.trace.api.appsec.RaspCallSites;
 import datadog.trace.api.iast.IastCallSites;
 import datadog.trace.api.iast.InstrumentationBridge;
 import datadog.trace.api.iast.Sink;
@@ -11,24 +10,34 @@
 import javax.annotation.Nullable;
 
 @Sink(VulnerabilityTypes.COMMAND_INJECTION)
-@CallSite(
-    spi = {IastCallSites.class, RaspCallSites.class},
-    helpers = ShellCmdRaspHelper.class)
+@CallSite(spi = IastCallSites.class)
 public class RuntimeCallSite {
 
   @CallSite.Before("java.lang.Process java.lang.Runtime.exec(java.lang.String)")
   public static void beforeStart(@CallSite.Argument @Nullable final String command) {
     if (command != null) { // runtime fails if null
-      iastCallback(command);
-      raspCallback(command);
+      final CommandInjectionModule module = InstrumentationBridge.COMMAND_INJECTION;
+      if (module != null) {
+        try {
+          module.onRuntimeExec(command);
+        } catch (final Throwable e) {
+          module.onUnexpectedException("beforeExec threw", e);
+        }
+      }
     }
   }
 
   @CallSite.Before("java.lang.Process java.lang.Runtime.exec(java.lang.String[])")
   public static void beforeExec(@CallSite.Argument @Nullable final String[] cmdArray) {
     if (cmdArray != null && cmdArray.length > 0) { // runtime fails if null or empty
-      iastCallback(cmdArray);
-      raspCallback(cmdArray);
+      final CommandInjectionModule module = InstrumentationBridge.COMMAND_INJECTION;
+      if (module != null) {
+        try {
+          module.onRuntimeExec(cmdArray);
+        } catch (final Throwable e) {
+          module.onUnexpectedException("beforeExec threw", e);
+        }
+      }
     }
   }
 
@@ -37,8 +46,14 @@ public static void beforeExec(
       @CallSite.Argument @Nullable final String command,
       @CallSite.Argument @Nullable final String[] envp) {
     if (command != null) { // runtime fails if null
-      iastCallback(envp, command);
-      raspCallback(command);
+      final CommandInjectionModule module = InstrumentationBridge.COMMAND_INJECTION;
+      if (module != null) {
+        try {
+          module.onRuntimeExec(envp, command);
+        } catch (final Throwable e) {
+          module.onUnexpectedException("beforeExec threw", e);
+        }
+      }
     }
   }
 
@@ -48,8 +63,14 @@ public static void beforeExec(
       @CallSite.Argument @Nullable final String[] cmdArray,
       @CallSite.Argument @Nullable final String[] envp) {
     if (cmdArray != null && cmdArray.length > 0) { // runtime fails if null or empty
-      iastCallback(envp, cmdArray);
-      raspCallback(cmdArray);
+      final CommandInjectionModule module = InstrumentationBridge.COMMAND_INJECTION;
+      if (module != null) {
+        try {
+          module.onRuntimeExec(envp, cmdArray);
+        } catch (final Throwable e) {
+          module.onUnexpectedException("beforeExec threw", e);
+        }
+      }
     }
   }
 
@@ -60,8 +81,14 @@ public static void beforeExec(
       @CallSite.Argument @Nullable final String[] envp,
       @CallSite.Argument @Nullable final File dir) {
     if (command != null) { // runtime fails if null
-      iastCallback(envp, command);
-      raspCallback(command);
+      final CommandInjectionModule module = InstrumentationBridge.COMMAND_INJECTION;
+      if (module != null) {
+        try {
+          module.onRuntimeExec(envp, command);
+        } catch (final Throwable e) {
+          module.onUnexpectedException("beforeExec threw", e);
+        }
+      }
     }
   }
 
@@ -72,60 +99,14 @@ public static void beforeExec(
       @CallSite.Argument @Nullable final String[] envp,
       @CallSite.Argument @Nullable final File dir) {
     if (cmdArray != null && cmdArray.length > 0) { // runtime fails if null or empty
-      iastCallback(envp, cmdArray);
-      raspCallback(cmdArray);
-    }
-  }
-
-  private static void iastCallback(String command) {
-    final CommandInjectionModule module = InstrumentationBridge.COMMAND_INJECTION;
-    if (module != null) {
-      try {
-        module.onRuntimeExec(command);
-      } catch (final Throwable e) {
-        module.onUnexpectedException("iastCallback threw", e);
-      }
-    }
-  }
-
-  private static void iastCallback(String[] cmdArray) {
-    final CommandInjectionModule module = InstrumentationBridge.COMMAND_INJECTION;
-    if (module != null) {
-      try {
-        module.onRuntimeExec(cmdArray);
-      } catch (final Throwable e) {
-        module.onUnexpectedException("iastCallback threw", e);
+      final CommandInjectionModule module = InstrumentationBridge.COMMAND_INJECTION;
+      if (module != null) {
+        try {
+          module.onRuntimeExec(envp, cmdArray);
+        } catch (final Throwable e) {
+          module.onUnexpectedException("beforeExec threw", e);
+        }
       }
     }
   }
-
-  private static void iastCallback(String[] envp, String command) {
-    final CommandInjectionModule module = InstrumentationBridge.COMMAND_INJECTION;
-    if (module != null) {
-      try {
-        module.onRuntimeExec(envp, command);
-      } catch (final Throwable e) {
-        module.onUnexpectedException("iastCallback threw", e);
-      }
-    }
-  }
-
-  private static void iastCallback(String[] envp, String[] cmdArray) {
-    final CommandInjectionModule module = InstrumentationBridge.COMMAND_INJECTION;
-    if (module != null) {
-      try {
-        module.onRuntimeExec(envp, cmdArray);
-      } catch (final Throwable e) {
-        module.onUnexpectedException("iastCallback threw", e);
-      }
-    }
-  }
-
-  private static void raspCallback(String command) {
-    ShellCmdRaspHelper.INSTANCE.beforeShellCmd(command);
-  }
-
-  private static void raspCallback(String[] cmdArray) {
-    ShellCmdRaspHelper.INSTANCE.beforeShellCmd(cmdArray);
-  }
 }
diff --git a/...n/java-lang/src/main/java/datadog/trace/instrumentation/java/lang/ShellCmdRaspHelper.java b/...n/java-lang/src/main/java/datadog/trace/instrumentation/java/lang/ShellCmdRaspHelper.java
diff --git a/.../lang/ShellCmdRaspHelperForkedTest.groovy → ...trumentationShellCmdRaspForkedTest.groovy b/.../lang/ShellCmdRaspHelperForkedTest.groovy → ...trumentationShellCmdRaspForkedTest.groovy
@@ -9,13 +9,14 @@ import datadog.trace.api.gateway.RequestContextSlot
 import datadog.trace.api.internal.TraceSegment
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan
 import datadog.trace.bootstrap.instrumentation.api.AgentTracer
+import datadog.trace.bootstrap.instrumentation.api.java.lang.ProcessImplInstrumentationHelpers
 import spock.lang.Shared
 
 import java.util.function.BiFunction
 
 import static datadog.trace.api.gateway.Events.EVENTS
 
-class ShellCmdRaspHelperForkedTest  extends  AgentTestRunner {
+class ProcessImplInstrumentationShellCmdRaspForkedTest extends  AgentTestRunner {
 
   @Shared
   protected static final ORIGINAL_TRACER = AgentTracer.get()
@@ -49,7 +50,7 @@ class ShellCmdRaspHelperForkedTest  extends  AgentTestRunner {
     injectSysConfig(AppSecConfig.APPSEC_RASP_ENABLED, 'true')
   }
 
-  void 'test Helper'() {
+  void 'test shiRaspCheck'() {
 
     setup:
     final callbackProvider = Mock(CallbackProvider)
@@ -58,15 +59,10 @@ class ShellCmdRaspHelperForkedTest  extends  AgentTestRunner {
     tracer.getCallbackProvider(RequestContextSlot.APPSEC) >> callbackProvider
 
     when:
-    ShellCmdRaspHelper.INSTANCE.beforeShellCmd(*args)
+    ProcessImplInstrumentationHelpers.shiRaspCheck(['cat etc/password'] as String[])
 
     then:
     1 * callbackProvider.getCallback(EVENTS.shellCmd()) >> listener
-    1 * listener.apply(reqCtx, expected) >> flow
-
-    where:
-    args                 | expected
-    ['cat etc/password'] | 'cat etc/password'
-    [['cat etc/password', 'cat etc/password'] as String[]] | ['cat etc/password', 'cat etc/password']
+    1 * listener.apply(reqCtx, ['cat etc/password']) >> flow
   }
 }
@@ -573,11 +573,12 @@ class SpringBootSmokeTest extends AbstractAppSecServerSmokeTest {
 
     where:
     endpoint                    | cmd                                      | params
-    'cmd'                       | ['cat etc/password']                     | null
+    //TODO these test that receive String cmd instead of String[] cmd are not working due to internally are converted to ["cat", "etc/password"] and the WAF is not blocking them
+    //    'cmd'    | ['cat etc/password'] | null
+    //    'cmdWithParams'      | ['cat etc/password'] | ['param']
+    //    'cmdParamsAndFile'      | ['cat etc/password'] | ['param']
     'arrayCmd'                  | ['cat etc/password', 'cat etc/password'] | null
-    'cmdWithParams'             | ['cat etc/password']                     | ['param']
     'arrayCmdWithParams'        | ['cat etc/password', 'cat etc/password'] | ['param']
-    'cmdParamsAndFile'          | ['cat etc/password']                     | ['param']
     'arrayCmdWithParamsAndFile' | ['cat etc/password', 'cat etc/password'] | ['param']
     'processBuilder'            | ['cat etc/password', 'cat etc/password'] | null
   }

@@ -314,8 +314,8 @@ public EventType<TriFunction<RequestContext, UserIdCollectionMode, String, Flow<
   private static final EventType SHELL_CMD = new ET<>("shell.cmd", SHELL_CDM_ID);
 
   @SuppressWarnings("unchecked")
-  public EventType<BiFunction<RequestContext, Object, Flow<Void>>> shellCmd() {
-    return (EventType<BiFunction<RequestContext, Object, Flow<Void>>>) SHELL_CMD;
+  public EventType<BiFunction<RequestContext, String[], Flow<Void>>> shellCmd() {
+    return (EventType<BiFunction<RequestContext, String[], Flow<Void>>>) SHELL_CMD;
   }
 
   static final int MAX_EVENTS = nextId.get();