AppSec Lambda: Analyze HTTP response

dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java

@@ -126,7 +126,7 @@ static void exit(
         }
         String lambdaRequestId = awsContext.getAwsRequestId();
 
-        AgentTracer.get().notifyAppSecEnd(span);
+        AgentTracer.get().notifyAppSecEnd(span, result);
         // Force the resource name back to the literal placeholder marker right
         // before finish so that the Datadog Lambda Extension's filter
         // (filter_span_from_lambda_library_or_runtime in

dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/HandlerStreamingWith404Response.java

@@ -0,0 +1,24 @@
+import com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.RequestStreamHandler;
+import java.io.BufferedWriter;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.io.OutputStreamWriter;
+import java.io.PrintWriter;
+import java.nio.charset.StandardCharsets;
+
+public class HandlerStreamingWith404Response implements RequestStreamHandler {
+  @Override
+  public void handleRequest(InputStream inputStream, OutputStream outputStream, Context context)
+      throws IOException {
+    PrintWriter writer =
+        new PrintWriter(
+            new BufferedWriter(new OutputStreamWriter(outputStream, StandardCharsets.UTF_8)));
+    writer.write(
+        "{\"statusCode\": 404, "
+            + "\"headers\": {\"content-type\": \"text/html\"}, "
+            + "\"body\": \"Not Found\"}");
+    writer.close();
+  }
+}

dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/HandlerStreamingWithApiGwResponse.java

@@ -0,0 +1,24 @@
+import com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.RequestStreamHandler;
+import java.io.BufferedWriter;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.io.OutputStreamWriter;
+import java.io.PrintWriter;
+import java.nio.charset.StandardCharsets;
+
+public class HandlerStreamingWithApiGwResponse implements RequestStreamHandler {
+  @Override
+  public void handleRequest(InputStream inputStream, OutputStream outputStream, Context context)
+      throws IOException {
+    PrintWriter writer =
+        new PrintWriter(
+            new BufferedWriter(new OutputStreamWriter(outputStream, StandardCharsets.UTF_8)));
+    writer.write(
+        "{\"statusCode\": 200, "
+            + "\"headers\": {\"content-type\": \"application/json\", \"x-custom\": \"custom-val\"}, "
+            + "\"body\": \"{\\\"result\\\": \\\"ok\\\"}\"}");
+    writer.close();
+  }
+}

dd-java-agent/instrumentation/aws-java/aws-java-lambda-handler-1.2/src/test/java/HandlerStreamingWithRawJson.java

@@ -0,0 +1,15 @@
+import com.amazonaws.services.lambda.runtime.Context;
+import com.amazonaws.services.lambda.runtime.RequestStreamHandler;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.nio.charset.StandardCharsets;
+
+/** Writes valid JSON that is not in API Gateway response format (no statusCode/headers/body). */
+public class HandlerStreamingWithRawJson implements RequestStreamHandler {
+  @Override
+  public void handleRequest(InputStream inputStream, OutputStream outputStream, Context context)
+      throws IOException {
+    outputStream.write("{\"result\": \"hello\"}".getBytes(StandardCharsets.UTF_8));
+  }
+}