feat: add support for aws sdk v2 (Rd 12691) (#65)

Author: nadav3396

.circleci/config.yml

@@ -33,53 +33,47 @@ workflows:
 
     - lumigo-orb/be-deploy:
         context: common
+        save_project_folder: false
         requires:
           - lumigo-orb/is_environment_available
 
-    - lumigo-orb/pre_build_artifacts:
+    - lumigo-orb/prep-it-resources:
         context: common
         requires:
           - lumigo-orb/is_environment_available
 
-    - lumigo-orb/integration-test-prep:
-        context:
-          - common
-          - java
-        install_maven_dependencies: true
-        pre_builds_available: true
-        run_test_cleanup: false
-        requires:
-          - lumigo-orb/be-deploy
-          - lumigo-orb/pre_build_artifacts
-
-    - lumigo-orb/integration-test-cleanup:
-        name: pre-test-cleanup
+    - lumigo-orb/prep-k8s-and-operator:
         context: common
         requires:
-          - lumigo-orb/integration-test-prep
+          - lumigo-orb/is_environment_available
 
     - lumigo-orb/integration-test-parallel:
         context: common
         run_test_cleanup: false
         requires:
-          - pre-test-cleanup
+          - lumigo-orb/be-deploy
+          - lumigo-orb/prep-it-resources
+          - lumigo-orb/prep-k8s-and-operator
 
-    - lumigo-orb/integration-test-cleanup:
-        name: post-test-cleanup
+    - lumigo-orb/e2e-test:
         context: common
         requires:
-          - lumigo-orb/integration-test-parallel
+          - lumigo-orb/be-deploy
+          - lumigo-orb/prep-it-resources
+          - lumigo-orb/prep-k8s-and-operator
 
-    - lumigo-orb/e2e-test:
+    - lumigo-orb/integration-test-cleanup:
+        name: post-test-cleanup
         context: common
         requires:
-          - pre-test-cleanup
+          - lumigo-orb/integration-test-parallel
+          - lumigo-orb/e2e-test
 
     - lumigo-orb/workflow-completed-successfully:
         context: common
         requires:
-          - test
           - lumigo-orb/integration-test-parallel
+          - test
           - lumigo-orb/e2e-test
 
     - deploy:

README.md

@@ -102,3 +102,9 @@ class MyFunction implements RequestHandler<String, String> {
 
 Add the environment variable `JAVA_TOOL_OPTIONS` to your Lambda functions and set it to
 `-Djdk.attach.allowAttachSelf=true` in addition to the manual code mentioned above.
+
+### Supported Instrumentation Libraries
+
+- Aws SDK V1 
+- Aws SDK V2
+- Apache HTTP Client

pom.xml

@@ -62,7 +62,7 @@
         <jackson.version>2.9.8</jackson.version>
         <httpcomponents.httpclient.version>4.5.6</httpcomponents.httpclient.version>
         <httpcomponents.httpcore.version>4.4.10</httpcomponents.httpcore.version>
-        <junit.version>5.4.2</junit.version>
+        <junit.version>5.10.2</junit.version>
     </properties>
     <repositories>
         <repository>
@@ -96,7 +96,7 @@
         <dependency>
             <groupId>com.amazonaws</groupId>
             <artifactId>aws-java-sdk-s3</artifactId>
-            <version>1.11.505</version>
+            <version>1.12.261</version>
             <scope>provided</scope>
         </dependency>
         <dependency>
@@ -117,7 +117,33 @@
             <version>3.14.2</version>
         </dependency>
 
-
+        <!-- AWS sdk V2 dependencies -->
+        <dependency>
+            <groupId>software.amazon.awssdk</groupId>
+            <artifactId>core</artifactId>
+            <version>2.25.45</version>
+            <type>pom</type>
+        </dependency>
+        <dependency>
+            <groupId>software.amazon.awssdk</groupId>
+            <artifactId>dynamodb</artifactId>
+            <version>2.25.45</version>
+        </dependency>
+        <dependency>
+            <groupId>software.amazon.awssdk</groupId>
+            <artifactId>sqs</artifactId>
+            <version>2.25.45</version>
+        </dependency>
+        <dependency>
+            <groupId>software.amazon.awssdk</groupId>
+            <artifactId>sns</artifactId>
+            <version>2.25.45</version>
+        </dependency>
+        <dependency>
+            <groupId>software.amazon.awssdk</groupId>
+            <artifactId>kinesis</artifactId>
+            <version>2.25.45</version>
+        </dependency>
 
         <!-- Tracer dependencies -->
         <dependency>

src/main/java/io/lumigo/core/SpansContainer.java

@@ -5,8 +5,9 @@
 import com.amazonaws.services.lambda.runtime.Context;
 import io.lumigo.core.configuration.Configuration;
 import io.lumigo.core.network.Reporter;
-import io.lumigo.core.parsers.AwsParserFactory;
 import io.lumigo.core.parsers.event.EventParserFactory;
+import io.lumigo.core.parsers.v1.AwsSdkV1ParserFactory;
+import io.lumigo.core.parsers.v2.AwsSdkV2ParserFactory;
 import io.lumigo.core.utils.AwsUtils;
 import io.lumigo.core.utils.JsonUtils;
 import io.lumigo.core.utils.StringUtils;
@@ -21,6 +22,11 @@
 import org.apache.http.client.methods.HttpEntityEnclosingRequestBase;
 import org.apache.http.client.methods.HttpUriRequest;
 import org.pmw.tinylog.Logger;
+import software.amazon.awssdk.awscore.AwsResponse;
+import software.amazon.awssdk.core.SdkResponse;
+import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
+import software.amazon.awssdk.core.interceptor.SdkExecutionAttribute;
+import software.amazon.awssdk.core.sync.RequestBody;
 
 public class SpansContainer {
 
@@ -354,14 +360,88 @@ public void addHttpSpan(Long startTime, Request<?> request, Response<?> response
                                                         response.getHttpResponse().getStatusCode())
                                                 .build())
                                 .build());
-        AwsParserFactory.getParser(request.getServiceName()).parse(httpSpan, request, response);
+        AwsSdkV1ParserFactory.getParser(request.getServiceName())
+                .safeParse(httpSpan, request, response);
+        httpSpans.add(httpSpan);
+    }
+
+    public void addHttpSpan(
+            Long startTime,
+            final software.amazon.awssdk.core.interceptor.Context.AfterExecution context,
+            final ExecutionAttributes executionAttributes) {
+        HttpSpan httpSpan = createBaseHttpSpan(startTime);
+        String spanId = null;
+        for (Map.Entry<String, List<String>> header : context.httpResponse().headers().entrySet()) {
+            if ("x-amzn-requestid".equalsIgnoreCase(header.getKey())
+                    || "x-amz-requestid".equalsIgnoreCase(header.getKey())) {
+                spanId = header.getValue().get(0);
+            }
+        }
+        if (spanId != null) {
+            httpSpan.setId(spanId);
+        }
+        httpSpan.getInfo()
+                .setHttpInfo(
+                        HttpSpan.HttpInfo.builder()
+                                .host(context.httpRequest().getUri().getHost())
+                                .request(
+                                        HttpSpan.HttpData.builder()
+                                                .headers(
+                                                        callIfVerbose(
+                                                                () ->
+                                                                        extractHeadersV2(
+                                                                                context.httpRequest()
+                                                                                        .headers())))
+                                                .uri(
+                                                        callIfVerbose(
+                                                                () ->
+                                                                        context.httpRequest()
+                                                                                .getUri()
+                                                                                .toString()))
+                                                .method(context.httpRequest().method().name())
+                                                .body(
+                                                        callIfVerbose(
+                                                                () ->
+                                                                        extractBodyFromRequest(
+                                                                                context
+                                                                                        .requestBody())))
+                                                .build())
+                                .response(
+                                        HttpSpan.HttpData.builder()
+                                                .headers(
+                                                        callIfVerbose(
+                                                                () ->
+                                                                        extractHeadersV2(
+                                                                                context.httpResponse()
+                                                                                        .headers())))
+                                                .body(
+                                                        callIfVerbose(
+                                                                () ->
+                                                                        extractBodyFromResponse(
+                                                                                context
+                                                                                        .response())))
+                                                .statusCode(context.httpResponse().statusCode())
+                                                .build())
+                                .build());
+
+        Logger.debug(
+                "Trying to extract aws custom properties for service: "
+                        + executionAttributes.getAttribute(SdkExecutionAttribute.SERVICE_NAME));
+        AwsSdkV2ParserFactory.getParser(
+                        executionAttributes.getAttribute(SdkExecutionAttribute.SERVICE_NAME))
+                .safeParse(httpSpan, context);
+
         httpSpans.add(httpSpan);
     }
 
     private static String extractHeaders(Map<String, String> headers) {
         return JsonUtils.getObjectAsJsonString(headers);
     }
 
+    private static String extractHeadersV2(Map<String, List<String>> headers) {
+        return JsonUtils.getObjectAsJsonString(headers);
+    }
+
     private static String extractHeaders(Header[] headers) {
         Map<String, String> headersMap = new HashMap<>();
         if (headers != null) {
@@ -373,27 +453,31 @@ private static String extractHeaders(Header[] headers) {
     }
 
     protected static String extractBodyFromRequest(Request<?> request) {
-        return extractBodyFromRequest(request.getContent());
+        return extractBodyFromStream(request.getContent());
+    }
+
+    protected static String extractBodyFromRequest(Optional<RequestBody> request) {
+        return request.map(
+                        requestBody ->
+                                extractBodyFromStream(
+                                        requestBody.contentStreamProvider().newStream()))
+                .orElse(null);
     }
 
     protected static String extractBodyFromRequest(HttpUriRequest request) throws Exception {
         if (request instanceof HttpEntityEnclosingRequestBase) {
             HttpEntity entity = ((HttpEntityEnclosingRequestBase) request).getEntity();
             if (entity != null) {
-                return extractBodyFromRequest(entity.getContent());
+                return extractBodyFromStream(entity.getContent());
             }
         }
         return null;
     }
 
-    protected static String extractBodyFromRequest(InputStream stream) {
-        return StringUtils.extractStringForStream(stream, MAX_STRING_SIZE);
-    }
-
     protected static String extractBodyFromResponse(HttpResponse response) throws IOException {
-        return StringUtils.extractStringForStream(
-                response.getEntity() != null ? response.getEntity().getContent() : null,
-                MAX_STRING_SIZE);
+        return response.getEntity() != null
+                ? extractBodyFromStream(response.getEntity().getContent())
+                : null;
     }
 
     protected static String extractBodyFromResponse(Response response) {
@@ -402,6 +486,17 @@ protected static String extractBodyFromResponse(Response response) {
                 : null;
     }
 
+    protected static String extractBodyFromResponse(SdkResponse response) {
+        if (response instanceof AwsResponse) {
+            return JsonUtils.getObjectAsJsonString(response.toBuilder());
+        }
+        return null;
+    }
+
+    protected static String extractBodyFromStream(InputStream stream) {
+        return StringUtils.extractStringForStream(stream, MAX_STRING_SIZE);
+    }
+
     public String getPatchedRoot() {
         return String.format(
                 "Root=%s-0000%s-%s%s",

src/main/java/io/lumigo/core/instrumentation/agent/Loader.java

@@ -4,6 +4,7 @@
 import static net.bytebuddy.matcher.ElementMatchers.not;
 
 import io.lumigo.core.instrumentation.impl.AmazonHttpClientInstrumentation;
+import io.lumigo.core.instrumentation.impl.AmazonHttpClientV2Instrumentation;
 import io.lumigo.core.instrumentation.impl.ApacheHttpInstrumentation;
 import net.bytebuddy.agent.builder.AgentBuilder;
 import org.pmw.tinylog.Logger;
@@ -14,17 +15,25 @@ public static void instrument(java.lang.instrument.Instrumentation inst) {
         ApacheHttpInstrumentation apacheHttpInstrumentation = new ApacheHttpInstrumentation();
         AmazonHttpClientInstrumentation amazonHttpClientInstrumentation =
                 new AmazonHttpClientInstrumentation();
+        AmazonHttpClientV2Instrumentation amazonHttpClientV2Instrumentation =
+                new AmazonHttpClientV2Instrumentation();
         AgentBuilder builder =
                 new AgentBuilder.Default()
                         .disableClassFormatChanges()
                         .with(AgentBuilder.RedefinitionStrategy.RETRANSFORMATION)
                         .ignore(
                                 not(nameStartsWith("com.amazonaws.http.AmazonHttpClient"))
-                                        .and(not(nameStartsWith("org.apache.http.impl.client"))))
+                                        .and(not(nameStartsWith("org.apache.http.impl.client")))
+                                        .and(
+                                                not(
+                                                        nameStartsWith(
+                                                                "software.amazon.awssdk.core.client.builder.SdkDefaultClientBuilder"))))
                         .type(apacheHttpInstrumentation.getTypeMatcher())
                         .transform(apacheHttpInstrumentation.getTransformer())
                         .type(amazonHttpClientInstrumentation.getTypeMatcher())
-                        .transform(amazonHttpClientInstrumentation.getTransformer());
+                        .transform(amazonHttpClientInstrumentation.getTransformer())
+                        .type(amazonHttpClientV2Instrumentation.getTypeMatcher())
+                        .transform(amazonHttpClientV2Instrumentation.getTransformer());
 
         builder.installOn(inst);
         Logger.debug("Finish Instrumentation");

src/main/java/io/lumigo/core/instrumentation/impl/AmazonHttpClientInstrumentation.java

@@ -24,7 +24,6 @@ public ElementMatcher<TypeDescription> getTypeMatcher() {
 
     @Override
     public AgentBuilder.Transformer.ForAdvice getTransformer() {
-
         return new AgentBuilder.Transformer.ForAdvice()
                 .include(Loader.class.getClassLoader())
                 .advice(isMethod().and(named("execute")), AmazonHttpClientAdvice.class.getName());