Properly capture lambda payloads for all handler types.

Author: purple4reina

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

@@ -66,20 +66,14 @@ protected boolean defaultEnabled() {
 
   @Override
   public void methodAdvice(MethodTransformer transformer) {
-    // two args
-    transformer.applyAdvice(
-        isMethod()
-            .and(named("handleRequest"))
-            .and(takesArgument(1, named("com.amazonaws.services.lambda.runtime.Context"))),
-        getClass().getName() + "$ExtensionCommunicationAdvice");
-    // three args (streaming)
+    // lambda under the hood converts all handlers to streaming handlers via
+    // lambdainternal.EventHandlerLoader$PojoHandlerAsStreamHandler.handleRequest
+    // full spec here : https://docs.aws.amazon.com/lambda/latest/dg/java-handler.html
     transformer.applyAdvice(
         isMethod()
             .and(named("handleRequest"))
             .and(takesArgument(2, named("com.amazonaws.services.lambda.runtime.Context"))),
         getClass().getName() + "$ExtensionCommunicationAdvice");
-    // full spec here : https://docs.aws.amazon.com/lambda/latest/dg/java-handler.html
-
   }
 
   public static class ExtensionCommunicationAdvice {
@@ -108,7 +102,7 @@ static AgentScope enter(
     static void exit(
         @Origin String method,
         @Enter final AgentScope scope,
-        @Advice.Return(typing = DYNAMIC) final Object result,
+        @Advice.Argument(1) final Object result,
         @Advice.Thrown final Throwable throwable) {
 
       if (scope == null) {

dd-trace-core/src/main/java/datadog/trace/lambda/LambdaHandler.java

@@ -10,6 +10,7 @@
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanContext;
 import datadog.trace.core.CoreTracer;
 import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
 import java.nio.charset.StandardCharsets;
 import java.util.Base64;
 import okhttp3.ConnectionPool;
@@ -63,6 +64,7 @@ public class LambdaHandler {
   private static final JsonAdapter<Object> adapter =
       new Moshi.Builder()
           .add(ByteArrayInputStream.class, new ReadFromInputStreamJsonAdapter())
+          .add(ByteArrayOutputStream.class, new ReadFromOutputStreamJsonAdapter())
           .add(SkipUnsupportedTypeJsonAdapter.newFactory())
           .build()
           .adapter(Object.class);

dd-trace-core/src/main/java/datadog/trace/lambda/ReadFromOutputStreamJsonAdapter.java

@@ -0,0 +1,26 @@
+package datadog.trace.lambda;
+
+import com.squareup.moshi.JsonAdapter;
+import com.squareup.moshi.JsonReader;
+import com.squareup.moshi.JsonWriter;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import okio.BufferedSink;
+
+public final class ReadFromOutputStreamJsonAdapter extends JsonAdapter<ByteArrayOutputStream> {
+
+  @Override
+  public ByteArrayOutputStream fromJson(JsonReader reader) throws IOException {
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public void toJson(JsonWriter writer, ByteArrayOutputStream outputStream) throws IOException {
+    if (outputStream != null) {
+      BufferedSink sink = writer.valueSink();
+      byte[] bytes = outputStream.toByteArray();
+      sink.write(bytes);
+      sink.flush();
+    }
+  }
+}

dd-trace-core/src/test/groovy/datadog/trace/lambda/LambdaHandlerTest.groovy

@@ -308,4 +308,16 @@ class LambdaHandlerTest extends DDCoreSpecification {
     then:
     result == body
   }
+
+  def "test moshi toJson OutputStream"() {
+    given:
+    def body = "{\"body\":\"bababango\",\"statusCode\":\"200\"}"
+    def myEvent = new ByteArrayOutputStream(body.getBytes())
+
+    when:
+    def result = LambdaHandler.writeValueAsString(myEvent)
+
+    then:
+    result == body
+  }
 }