Properly capture lambda payloads for all handler types. (#8264)

* Properly capture lambda payloads for all handler types.

* Remove unused import.

* Remove non-stream handler test.

* Properly create OutputStream for testing.

Author: purple4reina

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

@@ -10,7 +10,6 @@
 import static net.bytebuddy.asm.Advice.OnMethodExit;
 import static net.bytebuddy.asm.Advice.Origin;
 import static net.bytebuddy.asm.Advice.This;
-import static net.bytebuddy.implementation.bytecode.assign.Assigner.Typing.DYNAMIC;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 
@@ -66,20 +65,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 +101,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-java-agent/instrumentation/aws-lambda-handler/src/test/groovy/Handler.java

@@ -8,21 +8,6 @@ abstract class LambdaHandlerInstrumentationTest extends VersionedNamingTestBase
     null
   }
 
-  def "test lambda handler"() {
-    when:
-    new Handler().handleRequest(null, null)
-
-    then:
-    assertTraces(1) {
-      trace(1) {
-        span {
-          operationName operation()
-          errored false
-        }
-      }
-    }
-  }
-
   def "test lambda streaming handler"() {
     when:
     def input = new ByteArrayInputStream(StandardCharsets.UTF_8.encode("Hello").array())

dd-java-agent/instrumentation/aws-lambda-handler/src/test/groovy/LambdaHandlerInstrumentationTest.groovy

@@ -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/LambdaHandler.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/main/java/datadog/trace/lambda/ReadFromOutputStreamJsonAdapter.java

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