Add support for host methods

DataDog · Mariovido · Dec 9, 2024 · Nov 8, 2024 · Nov 8, 2024 · Nov 8, 2024
commit af7c2aa7eacc8840977d977eb6efe5310e4660a8
@@ -1,5 +1,6 @@
 package datadog.trace.instrumentation.apachehttpclient5;
 
+import datadog.trace.api.iast.util.PropagationUtils;
 import java.net.URI;
 import java.net.URISyntaxException;
 import org.apache.hc.core5.http.Header;
@@ -15,6 +16,7 @@ public class HostAndRequestAsHttpUriRequest extends BasicClassicHttpRequest {
   public HostAndRequestAsHttpUriRequest(final HttpHost httpHost, final HttpRequest httpRequest) {
     super(httpRequest.getMethod(), httpHost, httpRequest.getPath());
     actualRequest = httpRequest;
+    PropagationUtils.taintObjectIfTainted(this, httpHost.toURI());
  }
 
   @Override

@@ -19,7 +19,7 @@ public class IastHttpUriRequestBaseInstrumentation extends InstrumenterModule.Ia
     implements Instrumenter.ForSingleType, Instrumenter.HasTypeAdvice {
 
   public IastHttpUriRequestBaseInstrumentation() {
-    super("testApache");
+    super("apache-httpclient", "httpclient5");
   }
 
   @Override

@@ -0,0 +1,47 @@
+package datadog.trace.instrumentation.apachehttpcore5;
+
+import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
+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.net.InetAddress;
+import net.bytebuddy.asm.Advice;
+
+@AutoService(InstrumenterModule.class)
+public class IastHttpHostInstrumentation extends InstrumenterModule.Iast
+    implements Instrumenter.ForSingleType {
+
+  public IastHttpHostInstrumentation() {
+    super("httpcore", "apache-httpcore", "apache-http-core");
+  }
+
+  @Override
+  public String instrumentedType() {
+    return "org.apache.hc.core5.http.HttpHost";
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isConstructor()
+            .and(takesArguments(String.class, InetAddress.class, String.class, int.class)),
+        IastHttpHostInstrumentation.class.getName() + "$CtorAdvice");
+  }
+
+  public static class CtorAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    @Propagation
+    public static void afterCtor(
+        @Advice.This final Object self, @Advice.Argument(2) final Object host) {
+      final PropagationModule module = InstrumentationBridge.PROPAGATION;
+      if (module != null) {
+        module.taintObjectIfTainted(self, host);
+      }
+    }
+  }
+}
@@ -0,0 +1,31 @@
+package datadog.trace.instrumentation.apachehttpcore5
+
+import datadog.trace.agent.test.AgentTestRunner
+import datadog.trace.api.iast.InstrumentationBridge
+import datadog.trace.api.iast.propagation.PropagationModule
+import org.apache.hc.core5.http.HttpHost
+
+class IastHttpHostInstrumentationTest extends AgentTestRunner {
+
+  @Override
+  protected void configurePreAgent() {
+    injectSysConfig('dd.iast.enabled', 'true')
+  }
+
+  void 'test constructor'(){
+    given:
+    final module = Mock(PropagationModule)
+    InstrumentationBridge.registerIastModule(module)
+
+    when:
+    HttpHost.newInstance(*args)
+
+    then:
+    1 * module.taintObjectIfTainted( _ as HttpHost, 'localhost')
+
+    where:
+    args | _
+    ['localhost'] | _
+    ['localhost', 8080] | _
+  }
+}
@@ -95,17 +95,24 @@ public String okHttp3(@RequestParam(value = "url") final String url) {
   @PostMapping("/apache-httpclient5")
   public String apacheHttpClient5(
       @RequestParam(value = "url", required = false) final String url,
+      @RequestParam(value = "urlHandler", required = false) final String urlHandler,
       @RequestParam(value = "host", required = false) final String host) {
     CloseableHttpClient client = HttpClients.createDefault();
-    org.apache.hc.client5.http.classic.methods.HttpGet request =
-        new org.apache.hc.client5.http.classic.methods.HttpGet(url);
     try {
       if (host != null) {
-        //        final HttpHost httpHost = new HttpHost(host);
-        //        final BasicHttpRequest request = new BasicHttpRequest("GET", "/");
-        //        client.execute(httpHost, request);
+        final org.apache.hc.core5.http.HttpHost httpHost =
+            new org.apache.hc.core5.http.HttpHost(host);
+        final org.apache.hc.client5.http.classic.methods.HttpGet request =
+            new org.apache.hc.client5.http.classic.methods.HttpGet("/");
+        client.execute(httpHost, request);
       } else if (url != null) {
+        final org.apache.hc.client5.http.classic.methods.HttpGet request =
+            new org.apache.hc.client5.http.classic.methods.HttpGet(url);
         client.execute(request);
+      } else if (urlHandler != null) {
+        final org.apache.hc.client5.http.classic.methods.HttpGet request =
+            new org.apache.hc.client5.http.classic.methods.HttpGet(urlHandler);
+        client.execute(request, response -> null);
       }
       client.close();
     } catch (Exception e) {

@@ -724,7 +724,7 @@ abstract class AbstractIastSpringBootTest extends AbstractIastServerSmokeTest {
     'host'    | 'dd.datad0g.com'
   }
 
-  void 'ssrf is present (#path)'() {
+  void 'ssrf is present (#path) (#parameter)'() {
     setup:
     final url = "http://localhost:${httpPort}/ssrf/${path}"
     final body = new FormBody.Builder().add(parameter, value).build()
@@ -739,27 +739,28 @@ abstract class AbstractIastSpringBootTest extends AbstractIastServerSmokeTest {
         return false
       }
       final parts = vul.evidence.valueParts
-      if (parameter == 'url') {
+      if (parameter == 'url' || parameter == 'urlHandler') {
         return parts.size() == 1
         && parts[0].value == value && parts[0].source.origin == 'http.request.parameter' && parts[0].source.name == parameter
       } else if (parameter == 'host') {
         String protocol = protocolSecure ? 'https://' : 'http://'
-        return parts.size() == 2
-        && parts[0].value == protocol + value && parts[0].source.origin == 'http.request.parameter' && parts[0].source.name == parameter
-        && parts[1].value == '/' && parts[1].source == null
+        String finalValue = protocol + value + (endSlash ? '/' : '')
+        return parts[0].value.endsWith(finalValue) && parts[0].source.origin == 'http.request.parameter' && parts[0].source.name == parameter
       } else {
         throw new IllegalArgumentException("Parameter $parameter not supported")
       }
     }
 
     where:
-    path                  | parameter | value                     | protocolSecure
-    "apache-httpclient4"  | "url"     | "https://dd.datad0g.com/" | true
-    "apache-httpclient4"  | "host"    | "dd.datad0g.com"          | false
-    "apache-httpclient5"  | "url"     | "https://dd.datad0g.com/" | false
-    "commons-httpclient2" | "url"     | "https://dd.datad0g.com/" | true
-    "okHttp2"             | "url"     | "https://dd.datad0g.com/" | true
-    "okHttp3"             | "url"     | "https://dd.datad0g.com/" | true
+    path                  | parameter    | value                     | protocolSecure | endSlash
+    "apache-httpclient4"  | "url"        | "https://dd.datad0g.com/" | true           | true
+    "apache-httpclient4"  | "host"       | "dd.datad0g.com"          | false          | false
+    "apache-httpclient5"  | "url"        | "https://dd.datad0g.com/" | true           | true
+    "apache-httpclient5"  | "urlHandler" | "https://dd.datad0g.com/" | true           | true
+    "apache-httpclient5"  | "host"       | "dd.datad0g.com"          | false          | true
+    "commons-httpclient2" | "url"        | "https://dd.datad0g.com/" | true           | true
+    "okHttp2"             | "url"        | "https://dd.datad0g.com/" | true           | true
+    "okHttp3"             | "url"        | "https://dd.datad0g.com/" | true           | true
   }
 
   void 'test iast metrics stored in spans'() {

@@ -2,6 +2,7 @@
 
 import datadog.trace.api.iast.InstrumentationBridge;
 import datadog.trace.api.iast.propagation.CodecModule;
+import datadog.trace.api.iast.propagation.PropagationModule;
 import datadog.trace.api.iast.propagation.StringModule;
 import java.net.URI;
 
@@ -44,4 +45,15 @@ public static StringBuilder onStringBuilderAppend(final String path, final Strin
     }
     return sb;
   }
+
+  public static void taintObjectIfTainted(final Object target, final Object input) {
+    final PropagationModule module = InstrumentationBridge.PROPAGATION;
+    if (module != null) {
+      try {
+        module.taintObjectIfTainted(target, input);
+      } catch (final Throwable e) {
+        module.onUnexpectedException("taintObjectIfTainted threw", e);
+      }
+    }
+  }
 }