From 027c707cb755453ed1e0bfdd78551f697cebadf0 Mon Sep 17 00:00:00 2001 From: Rey Abolofia Date: Fri, 16 Jun 2023 12:00:25 -0700 Subject: [PATCH 1/3] Deserialize json from byte streams. --- .../aws-lambda-handler/build.gradle | 8 +++ .../lambda/LambdaHandlerInstrumentation.java | 8 +++ .../datadog/trace/lambda/LambdaHandler.java | 1 + .../ReadFromInputStreamJsonAdapter.java | 51 +++++++++++++++++++ .../trace/lambda/LambdaHandlerTest.groovy | 15 +++++- 5 files changed, 82 insertions(+), 1 deletion(-) create mode 100644 dd-trace-core/src/main/java/datadog/trace/lambda/ReadFromInputStreamJsonAdapter.java diff --git a/dd-java-agent/instrumentation/aws-lambda-handler/build.gradle b/dd-java-agent/instrumentation/aws-lambda-handler/build.gradle index 2c2c7a11dac..291eb78aca0 100644 --- a/dd-java-agent/instrumentation/aws-lambda-handler/build.gradle +++ b/dd-java-agent/instrumentation/aws-lambda-handler/build.gradle @@ -1,3 +1,10 @@ +muzzle { + pass { + group = 'com.amazonaws' + module = 'aws-lambda-java-core' + versions = '[1.2.1,)' + } +} apply from: "$rootDir/gradle/java.gradle" @@ -16,5 +23,6 @@ latestDepTest { } dependencies { + compileOnly group: 'com.amazonaws', name: 'aws-lambda-java-core', version: '1.2.1' testImplementation group: 'com.amazonaws', name: 'aws-lambda-java-core', version: '1.2.1' } diff --git a/dd-java-agent/instrumentation/aws-lambda-handler/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java b/dd-java-agent/instrumentation/aws-lambda-handler/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java index 7c65bccb147..6595bc76d15 100644 --- a/dd-java-agent/instrumentation/aws-lambda-handler/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java +++ b/dd-java-agent/instrumentation/aws-lambda-handler/src/main/java/datadog/trace/instrumentation/aws/v1/lambda/LambdaHandlerInstrumentation.java @@ -14,8 +14,10 @@ import static net.bytebuddy.matcher.ElementMatchers.isMethod; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; +import com.amazonaws.services.lambda.runtime.RequestHandler; import com.google.auto.service.AutoService; import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.bootstrap.CallDepthThreadLocalMap; import datadog.trace.bootstrap.instrumentation.api.AgentScope; import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.bootstrap.instrumentation.api.AgentTracer; @@ -85,6 +87,10 @@ static AgentScope enter( @Advice.Argument(0) final Object event, @Origin("#m") final String methodName) { + if (CallDepthThreadLocalMap.incrementCallDepth(RequestHandler.class) > 0) { + return null; + } + AgentSpan.Context lambdaContext = AgentTracer.get().notifyExtensionStart(event); final AgentSpan span; if (null == lambdaContext) { @@ -107,6 +113,8 @@ static void exit( return; } + CallDepthThreadLocalMap.reset(RequestHandler.class); + try { final AgentSpan span = scope.span(); span.finish(); diff --git a/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaHandler.java b/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaHandler.java index 2874476641c..5af2c6de95f 100644 --- a/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaHandler.java +++ b/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaHandler.java @@ -56,6 +56,7 @@ public class LambdaHandler { new Moshi.Builder() // we need to bypass abstract Classes as we can't JSON serialize them .add(SkipUnsupportedTypeJsonAdapter.newFactory()) + .add(ReadFromInputStreamJsonAdapter.newFactory()) .build() .adapter(Object.class); diff --git a/dd-trace-core/src/main/java/datadog/trace/lambda/ReadFromInputStreamJsonAdapter.java b/dd-trace-core/src/main/java/datadog/trace/lambda/ReadFromInputStreamJsonAdapter.java new file mode 100644 index 00000000000..6c66f32c531 --- /dev/null +++ b/dd-trace-core/src/main/java/datadog/trace/lambda/ReadFromInputStreamJsonAdapter.java @@ -0,0 +1,51 @@ +package datadog.trace.lambda; + +import com.squareup.moshi.JsonAdapter; +import com.squareup.moshi.JsonReader; +import com.squareup.moshi.JsonWriter; +import com.squareup.moshi.Moshi; +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.lang.annotation.Annotation; +import java.lang.reflect.Type; +import java.util.Set; +import okio.BufferedSink; + +public final class ReadFromInputStreamJsonAdapter extends JsonAdapter { + + @Override + public T fromJson(JsonReader reader) throws IOException { + throw new UnsupportedOperationException(); + } + + @Override + public void toJson(JsonWriter writer, T value) throws IOException { + BufferedSink sink = writer.valueSink(); + byte[] bytes = getInputBytes(value); + sink.write(bytes); + sink.flush(); + } + + private byte[] getInputBytes(T value) throws IOException { + ByteArrayInputStream inputStream = (ByteArrayInputStream) value; + inputStream.mark(0); + byte[] bytes = new byte[inputStream.available()]; + inputStream.read(bytes); + inputStream.reset(); + return bytes; + } + + public static Factory newFactory() { + return new Factory() { + @Override + public JsonAdapter create( + Type requestedType, Set annotations, Moshi moshi) { + if (requestedType == ByteArrayInputStream.class) { + return new ReadFromInputStreamJsonAdapter<>(); + } + return moshi.nextAdapter(this, requestedType, annotations); + } + }; + } +} \ No newline at end of file diff --git a/dd-trace-core/src/test/groovy/datadog/trace/lambda/LambdaHandlerTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/lambda/LambdaHandlerTest.groovy index 7465c09e19c..c273436e8eb 100644 --- a/dd-trace-core/src/test/groovy/datadog/trace/lambda/LambdaHandlerTest.groovy +++ b/dd-trace-core/src/test/groovy/datadog/trace/lambda/LambdaHandlerTest.groovy @@ -11,6 +11,7 @@ import com.amazonaws.services.lambda.runtime.events.SNSEvent import com.amazonaws.services.lambda.runtime.events.S3Event import com.amazonaws.services.lambda.runtime.events.APIGatewayProxyRequestEvent import com.amazonaws.services.lambda.runtime.events.models.s3.S3EventNotification +import java.io.ByteArrayInputStream import static datadog.trace.agent.test.server.http.TestHttpServer.httpServer @@ -223,4 +224,16 @@ class LambdaHandlerTest extends DDCoreSpecification { then: result == "{\"body\":\"bababango\",\"httpMethod\":\"POST\"}" } -} + + def "test moshi toJson InputStream"() { + given: + def body = "{\"body\":\"bababango\",\"httpMethod\":\"POST\"}" + def myEvent = new ByteArrayInputStream(body.getBytes()) + + when: + def result = LambdaHandler.writeValueAsString(myEvent) + + then: + result == body + } +} \ No newline at end of file From 01128fcedff0ea3d848f8d5dd7334b8249fedec0 Mon Sep 17 00:00:00 2001 From: Rey Abolofia Date: Fri, 16 Jun 2023 13:20:14 -0700 Subject: [PATCH 2/3] Spotless fixes. --- .../datadog/trace/lambda/ReadFromInputStreamJsonAdapter.java | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/dd-trace-core/src/main/java/datadog/trace/lambda/ReadFromInputStreamJsonAdapter.java b/dd-trace-core/src/main/java/datadog/trace/lambda/ReadFromInputStreamJsonAdapter.java index 6c66f32c531..45454317caa 100644 --- a/dd-trace-core/src/main/java/datadog/trace/lambda/ReadFromInputStreamJsonAdapter.java +++ b/dd-trace-core/src/main/java/datadog/trace/lambda/ReadFromInputStreamJsonAdapter.java @@ -5,7 +5,6 @@ import com.squareup.moshi.JsonWriter; import com.squareup.moshi.Moshi; import java.io.ByteArrayInputStream; -import java.io.ByteArrayOutputStream; import java.io.IOException; import java.lang.annotation.Annotation; import java.lang.reflect.Type; @@ -48,4 +47,4 @@ public JsonAdapter create( } }; } -} \ No newline at end of file +} From f0880304cce95824acde3f320791bb196815ff70 Mon Sep 17 00:00:00 2001 From: Joey Zhao <5253430+joeyzhao2018@users.noreply.github.com> Date: Thu, 22 Jun 2023 10:45:32 -0400 Subject: [PATCH 3/3] change T to ByteArrayInputStream --- .../datadog/trace/lambda/LambdaHandler.java | 4 +-- .../ReadFromInputStreamJsonAdapter.java | 36 ++++++------------- .../trace/lambda/LambdaHandlerTest.groovy | 3 +- 3 files changed, 13 insertions(+), 30 deletions(-) diff --git a/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaHandler.java b/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaHandler.java index 5af2c6de95f..bcf0ffc8834 100644 --- a/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaHandler.java +++ b/dd-trace-core/src/main/java/datadog/trace/lambda/LambdaHandler.java @@ -10,6 +10,7 @@ import datadog.trace.bootstrap.instrumentation.api.AgentSpan; import datadog.trace.core.propagation.ExtractedContext; import datadog.trace.core.propagation.PropagationTags; +import java.io.ByteArrayInputStream; import okhttp3.MediaType; import okhttp3.OkHttpClient; import okhttp3.Request; @@ -54,9 +55,8 @@ public class LambdaHandler { private static final MediaType jsonMediaType = MediaType.parse("application/json"); private static final JsonAdapter adapter = new Moshi.Builder() - // we need to bypass abstract Classes as we can't JSON serialize them + .add(ByteArrayInputStream.class, new ReadFromInputStreamJsonAdapter()) .add(SkipUnsupportedTypeJsonAdapter.newFactory()) - .add(ReadFromInputStreamJsonAdapter.newFactory()) .build() .adapter(Object.class); diff --git a/dd-trace-core/src/main/java/datadog/trace/lambda/ReadFromInputStreamJsonAdapter.java b/dd-trace-core/src/main/java/datadog/trace/lambda/ReadFromInputStreamJsonAdapter.java index 45454317caa..76b6081d24f 100644 --- a/dd-trace-core/src/main/java/datadog/trace/lambda/ReadFromInputStreamJsonAdapter.java +++ b/dd-trace-core/src/main/java/datadog/trace/lambda/ReadFromInputStreamJsonAdapter.java @@ -3,48 +3,32 @@ import com.squareup.moshi.JsonAdapter; import com.squareup.moshi.JsonReader; import com.squareup.moshi.JsonWriter; -import com.squareup.moshi.Moshi; import java.io.ByteArrayInputStream; import java.io.IOException; -import java.lang.annotation.Annotation; -import java.lang.reflect.Type; -import java.util.Set; import okio.BufferedSink; -public final class ReadFromInputStreamJsonAdapter extends JsonAdapter { +public final class ReadFromInputStreamJsonAdapter extends JsonAdapter { @Override - public T fromJson(JsonReader reader) throws IOException { + public ByteArrayInputStream fromJson(JsonReader reader) throws IOException { throw new UnsupportedOperationException(); } @Override - public void toJson(JsonWriter writer, T value) throws IOException { - BufferedSink sink = writer.valueSink(); - byte[] bytes = getInputBytes(value); - sink.write(bytes); - sink.flush(); + public void toJson(JsonWriter writer, ByteArrayInputStream inputStream) throws IOException { + if (inputStream != null) { + BufferedSink sink = writer.valueSink(); + byte[] bytes = getInputBytes(inputStream); + sink.write(bytes); + sink.flush(); + } } - private byte[] getInputBytes(T value) throws IOException { - ByteArrayInputStream inputStream = (ByteArrayInputStream) value; + private byte[] getInputBytes(ByteArrayInputStream inputStream) throws IOException { inputStream.mark(0); byte[] bytes = new byte[inputStream.available()]; inputStream.read(bytes); inputStream.reset(); return bytes; } - - public static Factory newFactory() { - return new Factory() { - @Override - public JsonAdapter create( - Type requestedType, Set annotations, Moshi moshi) { - if (requestedType == ByteArrayInputStream.class) { - return new ReadFromInputStreamJsonAdapter<>(); - } - return moshi.nextAdapter(this, requestedType, annotations); - } - }; - } } diff --git a/dd-trace-core/src/test/groovy/datadog/trace/lambda/LambdaHandlerTest.groovy b/dd-trace-core/src/test/groovy/datadog/trace/lambda/LambdaHandlerTest.groovy index c273436e8eb..53676ce0d11 100644 --- a/dd-trace-core/src/test/groovy/datadog/trace/lambda/LambdaHandlerTest.groovy +++ b/dd-trace-core/src/test/groovy/datadog/trace/lambda/LambdaHandlerTest.groovy @@ -11,7 +11,6 @@ import com.amazonaws.services.lambda.runtime.events.SNSEvent import com.amazonaws.services.lambda.runtime.events.S3Event import com.amazonaws.services.lambda.runtime.events.APIGatewayProxyRequestEvent import com.amazonaws.services.lambda.runtime.events.models.s3.S3EventNotification -import java.io.ByteArrayInputStream import static datadog.trace.agent.test.server.http.TestHttpServer.httpServer @@ -236,4 +235,4 @@ class LambdaHandlerTest extends DDCoreSpecification { then: result == body } -} \ No newline at end of file +}