From f3c4c8c7ebc765830f86e109b7bb994905f062ea Mon Sep 17 00:00:00 2001
From: Lauri Tulmin <ltulmin@splunk.com>
Date: Wed, 28 Jun 2023 11:05:10 +0300
Subject: [PATCH] Fix context leak on call to AmazonS3.generatePresignedUrl

---
 .../aws-sdk-1.11/javaagent/build.gradle.kts   |  3 +++
 .../awssdk/v1_11/Aws1ClientTest.groovy        | 19 +++++++++++++++++++
 .../groovy/Aws0ClientTest.groovy              | 16 ++++++++++++++++
 .../awssdk/v1_11/TracingRequestHandler.java   |  8 ++++++++
 4 files changed, 46 insertions(+)

diff --git a/instrumentation/aws-sdk/aws-sdk-1.11/javaagent/build.gradle.kts b/instrumentation/aws-sdk/aws-sdk-1.11/javaagent/build.gradle.kts
index c226dad0fc1e..b654433cefe6 100644
--- a/instrumentation/aws-sdk/aws-sdk-1.11/javaagent/build.gradle.kts
+++ b/instrumentation/aws-sdk/aws-sdk-1.11/javaagent/build.gradle.kts
@@ -65,6 +65,9 @@ testing {
         implementation("com.amazonaws:aws-java-sdk-kinesis:1.11.0")
         implementation("com.amazonaws:aws-java-sdk-dynamodb:1.11.0")
         implementation("com.amazonaws:aws-java-sdk-sns:1.11.0")
+
+        // needed by S3
+        implementation("javax.xml.bind:jaxb-api:2.3.1")
       }
     }
 
diff --git a/instrumentation/aws-sdk/aws-sdk-1.11/javaagent/src/test/groovy/io/opentelemetry/javaagent/instrumentation/awssdk/v1_11/Aws1ClientTest.groovy b/instrumentation/aws-sdk/aws-sdk-1.11/javaagent/src/test/groovy/io/opentelemetry/javaagent/instrumentation/awssdk/v1_11/Aws1ClientTest.groovy
index c4a4593ef0c1..b8dc9f841c72 100644
--- a/instrumentation/aws-sdk/aws-sdk-1.11/javaagent/src/test/groovy/io/opentelemetry/javaagent/instrumentation/awssdk/v1_11/Aws1ClientTest.groovy
+++ b/instrumentation/aws-sdk/aws-sdk-1.11/javaagent/src/test/groovy/io/opentelemetry/javaagent/instrumentation/awssdk/v1_11/Aws1ClientTest.groovy
@@ -6,8 +6,11 @@
 package io.opentelemetry.javaagent.instrumentation.awssdk.v1_11
 
 import com.amazonaws.AmazonWebServiceClient
+import com.amazonaws.ClientConfiguration
 import com.amazonaws.Request
 import com.amazonaws.auth.BasicAWSCredentials
+import com.amazonaws.auth.NoOpSigner
+import com.amazonaws.auth.SignerFactory
 import com.amazonaws.handlers.RequestHandler2
 import com.amazonaws.regions.Regions
 import com.amazonaws.services.s3.AmazonS3Client
@@ -110,4 +113,20 @@ class Aws1ClientTest extends AbstractAws1ClientTest implements AgentTestTrait {
       }
     }
   }
+
+  def "calling generatePresignedUrl does not leak context"() {
+    setup:
+    SignerFactory.registerSigner("noop", NoOpSigner)
+    def client = AmazonS3ClientBuilder.standard()
+      .withRegion(Regions.US_EAST_1)
+      .withClientConfiguration(new ClientConfiguration().withSignerOverride("noop"))
+      .build()
+
+    when:
+    client.generatePresignedUrl("someBucket", "someKey", new Date())
+
+    then:
+    // expecting no active span after call to generatePresignedUrl
+    !Span.current().getSpanContext().isValid()
+  }
 }
diff --git a/instrumentation/aws-sdk/aws-sdk-1.11/javaagent/src/test_before_1_11_106/groovy/Aws0ClientTest.groovy b/instrumentation/aws-sdk/aws-sdk-1.11/javaagent/src/test_before_1_11_106/groovy/Aws0ClientTest.groovy
index f21ab34c2073..ceb6e4258bae 100644
--- a/instrumentation/aws-sdk/aws-sdk-1.11/javaagent/src/test_before_1_11_106/groovy/Aws0ClientTest.groovy
+++ b/instrumentation/aws-sdk/aws-sdk-1.11/javaagent/src/test_before_1_11_106/groovy/Aws0ClientTest.groovy
@@ -11,6 +11,8 @@ import com.amazonaws.auth.AWSCredentialsProviderChain
 import com.amazonaws.auth.BasicAWSCredentials
 import com.amazonaws.auth.EnvironmentVariableCredentialsProvider
 import com.amazonaws.auth.InstanceProfileCredentialsProvider
+import com.amazonaws.auth.NoOpSigner
+import com.amazonaws.auth.SignerFactory
 import com.amazonaws.auth.SystemPropertiesCredentialsProvider
 import com.amazonaws.auth.profile.ProfileCredentialsProvider
 import com.amazonaws.handlers.RequestHandler2
@@ -272,4 +274,18 @@ class Aws0ClientTest extends AgentInstrumentationSpecification {
       }
     }
   }
+
+  def "calling generatePresignedUrl does not leak context"() {
+    setup:
+    SignerFactory.registerSigner("noop", NoOpSigner)
+    def client = new AmazonS3Client(new ClientConfiguration().withSignerOverride("noop"))
+      .withEndpoint("${server.httpUri()}")
+
+    when:
+    client.generatePresignedUrl("someBucket", "someKey", new Date())
+
+    then:
+    // expecting no active span after call to generatePresignedUrl
+    !Span.current().getSpanContext().isValid()
+  }
 }
diff --git a/instrumentation/aws-sdk/aws-sdk-1.11/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/TracingRequestHandler.java b/instrumentation/aws-sdk/aws-sdk-1.11/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/TracingRequestHandler.java
index c427c9f729c5..a6755552197c 100644
--- a/instrumentation/aws-sdk/aws-sdk-1.11/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/TracingRequestHandler.java
+++ b/instrumentation/aws-sdk/aws-sdk-1.11/library/src/main/java/io/opentelemetry/instrumentation/awssdk/v1_11/TracingRequestHandler.java
@@ -36,6 +36,14 @@ final class TracingRequestHandler extends RequestHandler2 {
   @Override
   @SuppressWarnings("deprecation") // deprecated class to be updated once published in new location
   public void beforeRequest(Request<?> request) {
+    // GeneratePresignedUrlRequest doesn't result in actual request, beforeRequest is the only
+    // method called for it. Span created here would never be ended and scope would be leaked when
+    // running with java agent.
+    if ("com.amazonaws.services.s3.model.GeneratePresignedUrlRequest"
+        .equals(request.getOriginalRequest().getClass().getName())) {
+      return;
+    }
+
     Context parentContext = Context.current();
     if (!requestInstrumenter.shouldStart(parentContext, request)) {
       return;