diff --git a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/baseTest/groovy/AkkaHttpServerInstrumentationTest.groovy b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/baseTest/groovy/AkkaHttpServerInstrumentationTest.groovy index 87752bd275c..65ed26909e7 100644 --- a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.0/src/baseTest/groovy/AkkaHttpServerInstrumentationTest.groovy +++ b/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> callback = + BiFunction> bodyCallback = cbp.getCallback(EVENTS.requestBodyProcessed()); - if (callback == null) { + BiFunction, Flow> filenamesCallback = + cbp.getCallback(EVENTS.requestFilesFilenames()); + if (bodyCallback == null && filenamesCallback == null) { return; } - // conversion to map string -> list of string java.lang.Iterable strictParts = st.getStrictParts(); Map> conv = new HashMap<>(); + List filenames = new ArrayList<>(); for (akka.http.javadsl.model.Multipart.FormData.BodyPart.Strict part : strictParts) { + Optional 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 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 transformStringUnmarshaller( @@ -389,6 +420,7 @@ public static Unmarshaller transformStrictFormUnmarshall private static void handleStrictFormData(StrictForm sf) { Iterator> iterator = sf.fields().iterator(); Map> conv = new HashMap<>(); + List filenames = new ArrayList<>(); while (iterator.hasNext()) { Tuple2 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 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, Flow> filenamesCb = + cbp.getCallback(EVENTS.requestFilesFilenames()); + if (filenamesCb != null) { + Flow 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) { diff --git a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.6/src/test/groovy/AkkaHttp102ServerInstrumentationTests.groovy b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.6/src/test/groovy/AkkaHttp102ServerInstrumentationTests.groovy index a8c4134f9f5..6e517914b17 100644 --- a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.6/src/test/groovy/AkkaHttp102ServerInstrumentationTests.groovy +++ b/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 diff --git a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.6/src/test/groovy/AkkaHttpServerInstrumentationTest.groovy b/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.6/src/test/groovy/AkkaHttpServerInstrumentationTest.groovy index 9917d24108d..1fe0b3a3be0 100644 --- a/dd-java-agent/instrumentation/akka/akka-http/akka-http-10.6/src/test/groovy/AkkaHttpServerInstrumentationTest.groovy +++ b/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 true } + @Override + boolean testBodyFilenames() { + true + } + @Override boolean testBodyJson() { true diff --git a/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/test/groovy/test/boot/SpringBootBasedTest.groovy b/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/test/groovy/test/boot/SpringBootBasedTest.groovy index e98faaa09ae..7d38e56fe65 100644 --- a/dd-java-agent/instrumentation/spring/spring-webmvc/spring-webmvc-3.1/src/test/groovy/test/boot/SpringBootBasedTest.groovy +++ b/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 true } + @Override + boolean testBodyFilenames() { + true + } + @Override boolean testBodyJson() { true