Add server.request.body.filenames support for Akka HTTP

Author: jandro996

dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/baseTest/groovy/AkkaHttpServerInstrumentationTest.groovy

@@ -74,6 +74,11 @@ abstract class AkkaHttpServerInstrumentationTest extends HttpServerTest<AkkaHttp
     true
   }
 
+  @Override
+  boolean testBodyFilenames() {
+    true
+  }
+
   @Override
   boolean testBodyJson() {
     true
@@ -223,6 +228,11 @@ abstract class AkkaHttpServerInstrumentationSyncTest extends AkkaHttpServerInstr
     false
   }
 
+  @Override
+  boolean testBodyFilenames() {
+    false
+  }
+
   @Override
   boolean testBodyJson() {
     false
@@ -283,6 +293,11 @@ class AkkaHttpServerInstrumentationBindAndHandleTest extends AkkaHttpServerInstr
     akkaHttpVersion != '10.0.10'
   }
 
+  @Override
+  boolean testBodyFilenames() {
+    akkaHttpVersion != '10.0.10'
+  }
+
   @Override
   boolean testBodyUrlencoded() {
     akkaHttpVersion != '10.0.10'
@@ -309,6 +324,11 @@ class AkkaHttpServerInstrumentationBindAndHandleAsyncWithRouteAsyncHandlerTest e
     akkaHttpVersion != '10.0.10'
   }
 
+  @Override
+  boolean testBodyFilenames() {
+    akkaHttpVersion != '10.0.10'
+  }
+
   @Override
   boolean testBodyUrlencoded() {
     akkaHttpVersion != '10.0.10'

dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/latestDepTest/groovy/AkkaHttp102ServerInstrumentationTests.groovy

@@ -36,6 +36,11 @@ class AkkaHttp102ServerInstrumentationBindSyncTest extends AkkaHttpServerInstrum
     false
   }
 
+  @Override
+  boolean testBodyFilenames() {
+    false
+  }
+
   @Override
   boolean testBodyJson() {
     false

dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/main/java/datadog/trace/instrumentation/akkahttp/appsec/UnmarshallerHelpers.java

@@ -26,6 +26,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.WeakHashMap;
 import java.util.function.BiFunction;
 import org.slf4j.Logger;
@@ -199,17 +200,28 @@ private static void handleMultipartStrictFormData(
     }
 
     CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC);
-    BiFunction<RequestContext, Object, Flow<Void>> callback =
+    BiFunction<RequestContext, Object, Flow<Void>> bodyCallback =
         cbp.getCallback(EVENTS.requestBodyProcessed());
-    if (callback == null) {
+    BiFunction<RequestContext, List<String>, Flow<Void>> filenamesCallback =
+        cbp.getCallback(EVENTS.requestFilesFilenames());
+    if (bodyCallback == null && filenamesCallback == null) {
       return;
     }
 
-    // conversion to map string -> list of string
     java.lang.Iterable<akka.http.javadsl.model.Multipart.FormData.BodyPart.Strict> strictParts =
         st.getStrictParts();
     Map<String, List<String>> conv = new HashMap<>();
+    List<String> filenames = new ArrayList<>();
     for (akka.http.javadsl.model.Multipart.FormData.BodyPart.Strict part : strictParts) {
+      Optional<String> filenameOpt = part.getFilename();
+      if (filenameOpt.isPresent() && !filenameOpt.get().isEmpty()) {
+        filenames.add(filenameOpt.get());
+      }
+
+      if (bodyCallback == null) {
+        continue;
+      }
+
       akka.http.javadsl.model.HttpEntity.Strict entity = part.getEntity();
       if (!(entity instanceof HttpEntity.Strict)) {
         continue;
@@ -232,8 +244,27 @@ private static void handleMultipartStrictFormData(
       curStrings.add(s);
     }
 
-    // callback execution
-    executeCallback(reqCtx, callback, conv, "multipartFormDataUnmarshaller");
+    if (bodyCallback != null) {
+      executeCallback(reqCtx, bodyCallback, conv, "multipartFormDataUnmarshaller");
+    }
+
+    if (filenamesCallback != null && !filenames.isEmpty()) {
+      Flow<Void> filenamesFlow = filenamesCallback.apply(reqCtx, filenames);
+      Flow.Action filenamesAction = filenamesFlow.getAction();
+      if (filenamesAction instanceof Flow.Action.RequestBlockingAction) {
+        Flow.Action.RequestBlockingAction rba = (Flow.Action.RequestBlockingAction) filenamesAction;
+        BlockResponseFunction brf = reqCtx.getBlockResponseFunction();
+        if (brf != null) {
+          boolean success = brf.tryCommitBlockingResponse(reqCtx.getTraceSegment(), rba);
+          if (success) {
+            if (brf instanceof AkkaBlockResponseFunction) {
+              ((AkkaBlockResponseFunction) brf).setUnmarshallBlock(true);
+            }
+            throw new BlockingException("Blocked request (multipart file upload)");
+          }
+        }
+      }
+    }
   }
 
   public static Unmarshaller<HttpEntity, String> transformStringUnmarshaller(
@@ -389,6 +420,7 @@ public static Unmarshaller<HttpEntity, StrictForm> transformStrictFormUnmarshall
   private static void handleStrictFormData(StrictForm sf) {
     Iterator<Tuple2<String, StrictForm.Field>> iterator = sf.fields().iterator();
     Map<String, List<String>> conv = new HashMap<>();
+    List<String> filenames = new ArrayList<>();
     while (iterator.hasNext()) {
       Tuple2<String, StrictForm.Field> next = iterator.next();
       String fieldName = next._1();
@@ -413,9 +445,13 @@ private static void handleStrictFormData(StrictForm sf) {
         strings.add((String) strictFieldValue);
       } else if (strictFieldValue
           instanceof akka.http.scaladsl.model.Multipart$FormData$BodyPart$Strict) {
-        HttpEntity.Strict sentity =
-            ((akka.http.scaladsl.model.Multipart$FormData$BodyPart$Strict) strictFieldValue)
-                .entity();
+        akka.http.scaladsl.model.Multipart$FormData$BodyPart$Strict bodyPart =
+            (akka.http.scaladsl.model.Multipart$FormData$BodyPart$Strict) strictFieldValue;
+        Optional<String> filenameOpt = bodyPart.getFilename();
+        if (filenameOpt.isPresent() && !filenameOpt.get().isEmpty()) {
+          filenames.add(filenameOpt.get());
+        }
+        HttpEntity.Strict sentity = bodyPart.entity();
         String s =
             sentity
                 .getData()
@@ -426,6 +462,36 @@ private static void handleStrictFormData(StrictForm sf) {
     }
 
     handleArbitraryPostData(conv, "HttpEntity -> StrictForm unmarshaller");
+
+    if (!filenames.isEmpty()) {
+      AgentSpan span = activeSpan();
+      RequestContext reqCtx;
+      if (span != null
+          && (reqCtx = span.getRequestContext()) != null
+          && reqCtx.getData(RequestContextSlot.APPSEC) != null) {
+        CallbackProvider cbp = AgentTracer.get().getCallbackProvider(RequestContextSlot.APPSEC);
+        BiFunction<RequestContext, List<String>, Flow<Void>> filenamesCb =
+            cbp.getCallback(EVENTS.requestFilesFilenames());
+        if (filenamesCb != null) {
+          Flow<Void> filenamesFlow = filenamesCb.apply(reqCtx, filenames);
+          Flow.Action filenamesAction = filenamesFlow.getAction();
+          if (filenamesAction instanceof Flow.Action.RequestBlockingAction) {
+            Flow.Action.RequestBlockingAction rba =
+                (Flow.Action.RequestBlockingAction) filenamesAction;
+            BlockResponseFunction brf = reqCtx.getBlockResponseFunction();
+            if (brf != null) {
+              boolean success = brf.tryCommitBlockingResponse(reqCtx.getTraceSegment(), rba);
+              if (success) {
+                if (brf instanceof AkkaBlockResponseFunction) {
+                  ((AkkaBlockResponseFunction) brf).setUnmarshallBlock(true);
+                }
+                throw new BlockingException("Blocked request (multipart file upload)");
+              }
+            }
+          }
+        }
+      }
+    }
   }
 
   private static Object tryConvertingScalaContainers(Object obj, int depth) {

dd-java-agent/instrumentation/akka/akka-http/akka-http-10.6/src/test/groovy/AkkaHttp102ServerInstrumentationTests.groovy

@@ -36,6 +36,11 @@ class AkkaHttp102ServerInstrumentationBindSyncTest extends AkkaHttpServerInstrum
     false
   }
 
+  @Override
+  boolean testBodyFilenames() {
+    false
+  }
+
   @Override
   boolean testBodyJson() {
     false

dd-java-agent/instrumentation/akka/akka-http/akka-http-10.6/src/test/groovy/AkkaHttpServerInstrumentationTest.groovy

@@ -75,6 +75,11 @@ abstract class AkkaHttpServerInstrumentationTest extends HttpServerTest<AkkaHttp
     true
   }
 
+  @Override
+  boolean testBodyFilenames() {
+    true
+  }
+
   @Override
   boolean testBodyJson() {
     true
@@ -223,6 +228,11 @@ abstract class AkkaHttpServerInstrumentationSyncTest extends AkkaHttpServerInstr
     false
   }
 
+  @Override
+  boolean testBodyFilenames() {
+    false
+  }
+
   @Override
   boolean testBodyJson() {
     false
@@ -279,6 +289,11 @@ class AkkaHttpServerInstrumentationBindAndHandleTest extends AkkaHttpServerInstr
     true
   }
 
+  @Override
+  boolean testBodyFilenames() {
+    true
+  }
+
   @Override
   boolean testBodyUrlencoded() {
     true
@@ -305,6 +320,11 @@ class AkkaHttpServerInstrumentationBindAndHandleAsyncWithRouteAsyncHandlerTest e
     true
   }
 
+  @Override
+  boolean testBodyFilenames() {
+    true
+  }
+
   @Override
   boolean testBodyUrlencoded() {
     true

dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/latestDepTest/groovy/test/boot/SpringBootBasedTest.groovy

@@ -109,6 +109,11 @@ class SpringBootBasedTest extends HttpServerTest<ConfigurableApplicationContext>
     true
   }
 
+  @Override
+  boolean testBodyFilenames() {
+    true
+  }
+
   @Override
   boolean testBodyJson() {
     true

dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/test/groovy/test/boot/SpringBootBasedTest.groovy

@@ -119,6 +119,11 @@ class SpringBootBasedTest extends HttpServerTest<ConfigurableApplicationContext>
     true
   }
 
+  @Override
+  boolean testBodyFilenames() {
+    true
+  }
+
   @Override
   boolean testBodyJson() {
     true