Skip to content

Commit 5cf596e

Browse files
millemsmillems
authored
Pass _X_AMZN_TRACE_ID environment variable through to X-Amzn-Trace-Id header for AWS services running on Lambda. (#3043)
1 parent 76aef12 commit 5cf596e

File tree

6 files changed

+269
-9
lines changed

6 files changed

+269
-9
lines changed

‎core/aws-core/src/main/java/software/amazon/awssdk/awscore/client/builder/AwsDefaultClientBuilder.java

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
import software.amazon.awssdk.awscore.endpoint.FipsEnabledProvider;
3434
import software.amazon.awssdk.awscore.eventstream.EventStreamInitialRequestInterceptor;
3535
import software.amazon.awssdk.awscore.interceptor.HelpfulUnknownHostExceptionInterceptor;
36+
import software.amazon.awssdk.awscore.interceptor.TraceIdExecutionInterceptor;
3637
import software.amazon.awssdk.awscore.internal.defaultsmode.AutoDefaultsModeDiscovery;
3738
import software.amazon.awssdk.awscore.internal.defaultsmode.DefaultsModeConfiguration;
3839
import software.amazon.awssdk.awscore.internal.defaultsmode.DefaultsModeResolver;
@@ -419,7 +420,8 @@ private List<ExecutionInterceptor> addAwsInterceptors(SdkClientConfiguration con
419420

420421
private List<ExecutionInterceptor> awsInterceptors() {
421422
return Arrays.asList(new HelpfulUnknownHostExceptionInterceptor(),
422-
new EventStreamInitialRequestInterceptor());
423+
new EventStreamInitialRequestInterceptor(),
424+
new TraceIdExecutionInterceptor());
423425
}
424426

425427
@Override

‎core/aws-core/src/main/java/software/amazon/awssdk/awscore/interceptor/TraceIdExecutionInterceptor.java

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License").
5+
* You may not use this file except in compliance with the License.
6+
* A copy of the License is located at
7+
*
8+
* http://aws.amazon.com/apache2.0
9+
*
10+
* or in the "license" file accompanying this file. This file is distributed
11+
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12+
* express or implied. See the License for the specific language governing
13+
* permissions and limitations under the License.
14+
*/
15+
16+
package software.amazon.awssdk.awscore.interceptor;
17+
18+
import java.util.Optional;
19+
import software.amazon.awssdk.annotations.SdkInternalApi;
20+
import software.amazon.awssdk.core.interceptor.Context;
21+
import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
22+
import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
23+
import software.amazon.awssdk.http.SdkHttpRequest;
24+
import software.amazon.awssdk.utils.SystemSetting;
25+
26+
/**
27+
* The {@code TraceIdExecutionInterceptor} copies the {@link #TRACE_ID_ENVIRONMENT_VARIABLE} value to the
28+
* {@link #TRACE_ID_HEADER} header, assuming we seem to be running in a lambda environment.
29+
*/
30+
@SdkInternalApi
31+
public class TraceIdExecutionInterceptor implements ExecutionInterceptor {
32+
private static final String TRACE_ID_HEADER = "X-Amzn-Trace-Id";
33+
private static final String TRACE_ID_ENVIRONMENT_VARIABLE = "_X_AMZN_TRACE_ID";
34+
private static final String LAMBDA_FUNCTION_NAME_ENVIRONMENT_VARIABLE = "AWS_LAMBDA_FUNCTION_NAME";
35+
36+
@Override
37+
public SdkHttpRequest modifyHttpRequest(Context.ModifyHttpRequest context, ExecutionAttributes executionAttributes) {
38+
Optional<String> traceIdHeader = traceIdHeader(context);
39+
if (!traceIdHeader.isPresent()) {
40+
Optional<String> lambdafunctionName = lambdaFunctionNameEnvironmentVariable();
41+
Optional<String> traceId = traceIdEnvironmentVariable();
42+
43+
if (lambdafunctionName.isPresent() && traceId.isPresent()) {
44+
return context.httpRequest().copy(r -> r.putHeader(TRACE_ID_HEADER, traceId.get()));
45+
}
46+
}
47+
48+
return context.httpRequest();
49+
}
50+
51+
private Optional<String> traceIdHeader(Context.ModifyHttpRequest context) {
52+
return context.httpRequest().firstMatchingHeader(TRACE_ID_HEADER);
53+
}
54+
55+
private Optional<String> traceIdEnvironmentVariable() {
56+
// CHECKSTYLE:OFF - This is not configured by the customer, so it should not be configurable by system property
57+
return SystemSetting.getStringValueFromEnvironmentVariable(TRACE_ID_ENVIRONMENT_VARIABLE);
58+
// CHECKSTYLE:ON
59+
}
60+
61+
private Optional<String> lambdaFunctionNameEnvironmentVariable() {
62+
// CHECKSTYLE:OFF - This is not configured by the customer, so it should not be configurable by system property
63+
return SystemSetting.getStringValueFromEnvironmentVariable(LAMBDA_FUNCTION_NAME_ENVIRONMENT_VARIABLE);
64+
// CHECKSTYLE:ON
65+
}
66+
}

‎core/aws-core/src/test/java/software/amazon/awssdk/awscore/interceptor/TraceIdExecutionInterceptorTest.java

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License").
5+
* You may not use this file except in compliance with the License.
6+
* A copy of the License is located at
7+
*
8+
* http://aws.amazon.com/apache2.0
9+
*
10+
* or in the "license" file accompanying this file. This file is distributed
11+
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12+
* express or implied. See the License for the specific language governing
13+
* permissions and limitations under the License.
14+
*/
15+
16+
package software.amazon.awssdk.awscore.interceptor;
17+
18+
import static org.assertj.core.api.Assertions.assertThat;
19+
20+
import java.net.URI;
21+
import org.junit.jupiter.api.Test;
22+
import org.mockito.Mockito;
23+
import software.amazon.awssdk.core.SdkRequest;
24+
import software.amazon.awssdk.core.interceptor.Context;
25+
import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
26+
import software.amazon.awssdk.core.interceptor.InterceptorContext;
27+
import software.amazon.awssdk.http.SdkHttpMethod;
28+
import software.amazon.awssdk.http.SdkHttpRequest;
29+
import software.amazon.awssdk.testutils.EnvironmentVariableHelper;
30+
31+
public class TraceIdExecutionInterceptorTest {
32+
@Test
33+
public void nothingDoneWithoutEnvSettings() {
34+
EnvironmentVariableHelper.run(env -> {
35+
resetRelevantEnvVars(env);
36+
Context.ModifyHttpRequest context = context();
37+
assertThat(modifyHttpRequest(context)).isSameAs(context.httpRequest());
38+
});
39+
}
40+
41+
@Test
42+
public void headerAddedWithEnvSettings() {
43+
EnvironmentVariableHelper.run(env -> {
44+
resetRelevantEnvVars(env);
45+
env.set("AWS_LAMBDA_FUNCTION_NAME", "foo");
46+
env.set("_X_AMZN_TRACE_ID", "bar");
47+
Context.ModifyHttpRequest context = context();
48+
assertThat(modifyHttpRequest(context).firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("bar");
49+
});
50+
}
51+
52+
@Test
53+
public void headerNotAddedIfHeaderAlreadyExists() {
54+
EnvironmentVariableHelper.run(env -> {
55+
resetRelevantEnvVars(env);
56+
env.set("AWS_LAMBDA_FUNCTION_NAME", "foo");
57+
env.set("_X_AMZN_TRACE_ID", "bar");
58+
Context.ModifyHttpRequest context = context(SdkHttpRequest.builder()
59+
.uri(URI.create("https://localhost"))
60+
.method(SdkHttpMethod.GET)
61+
.putHeader("X-Amzn-Trace-Id", "existing")
62+
.build());
63+
assertThat(modifyHttpRequest(context)).isSameAs(context.httpRequest());
64+
});
65+
}
66+
67+
@Test
68+
public void headerNotAddedIfNotInLambda() {
69+
EnvironmentVariableHelper.run(env -> {
70+
resetRelevantEnvVars(env);
71+
env.set("_X_AMZN_TRACE_ID", "bar");
72+
Context.ModifyHttpRequest context = context();
73+
assertThat(modifyHttpRequest(context)).isSameAs(context.httpRequest());
74+
});
75+
}
76+
77+
@Test
78+
public void headerNotAddedIfNoTraceIdEnvVar() {
79+
EnvironmentVariableHelper.run(env -> {
80+
resetRelevantEnvVars(env);
81+
env.set("_X_AMZN_TRACE_ID", "bar");
82+
Context.ModifyHttpRequest context = context();
83+
assertThat(modifyHttpRequest(context)).isSameAs(context.httpRequest());
84+
});
85+
}
86+
87+
private Context.ModifyHttpRequest context() {
88+
return context(SdkHttpRequest.builder()
89+
.uri(URI.create("https://localhost"))
90+
.method(SdkHttpMethod.GET)
91+
.build());
92+
}
93+
94+
95+
private Context.ModifyHttpRequest context(SdkHttpRequest request) {
96+
return InterceptorContext.builder()
97+
.request(Mockito.mock(SdkRequest.class))
98+
.httpRequest(request)
99+
.build();
100+
}
101+
102+
private SdkHttpRequest modifyHttpRequest(Context.ModifyHttpRequest context) {
103+
return new TraceIdExecutionInterceptor().modifyHttpRequest(context, new ExecutionAttributes());
104+
}
105+
106+
private void resetRelevantEnvVars(EnvironmentVariableHelper env) {
107+
env.remove("AWS_LAMBDA_FUNCTION_NAME");
108+
env.remove("_X_AMZN_TRACE_ID");
109+
}
110+
}

‎test/codegen-generated-classes-test/src/test/java/software/amazon/awssdk/services/TraceIdTest.java

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License").
5+
* You may not use this file except in compliance with the License.
6+
* A copy of the License is located at
7+
*
8+
* http://aws.amazon.com/apache2.0
9+
*
10+
* or in the "license" file accompanying this file. This file is distributed
11+
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12+
* express or implied. See the License for the specific language governing
13+
* permissions and limitations under the License.
14+
*/
15+
16+
package software.amazon.awssdk.services;
17+
18+
import static org.assertj.core.api.Assertions.assertThat;
19+
20+
import org.junit.jupiter.api.Test;
21+
import software.amazon.awssdk.auth.credentials.AnonymousCredentialsProvider;
22+
import software.amazon.awssdk.awscore.interceptor.TraceIdExecutionInterceptor;
23+
import software.amazon.awssdk.http.AbortableInputStream;
24+
import software.amazon.awssdk.http.HttpExecuteResponse;
25+
import software.amazon.awssdk.http.SdkHttpResponse;
26+
import software.amazon.awssdk.regions.Region;
27+
import software.amazon.awssdk.services.protocolrestjson.ProtocolRestJsonClient;
28+
import software.amazon.awssdk.testutils.EnvironmentVariableHelper;
29+
import software.amazon.awssdk.testutils.service.http.MockSyncHttpClient;
30+
import software.amazon.awssdk.utils.StringInputStream;
31+
32+
/**
33+
* Verifies that the {@link TraceIdExecutionInterceptor} is actually wired up for AWS services.
34+
*/
35+
public class TraceIdTest {
36+
@Test
37+
public void traceIdInterceptorIsEnabled() {
38+
EnvironmentVariableHelper.run(env -> {
39+
env.set("AWS_LAMBDA_FUNCTION_NAME", "foo");
40+
env.set("_X_AMZN_TRACE_ID", "bar");
41+
42+
try (MockSyncHttpClient mockHttpClient = new MockSyncHttpClient();
43+
ProtocolRestJsonClient client = ProtocolRestJsonClient.builder()
44+
.region(Region.US_WEST_2)
45+
.credentialsProvider(AnonymousCredentialsProvider.create())
46+
.httpClient(mockHttpClient)
47+
.build()) {
48+
mockHttpClient.stubNextResponse(HttpExecuteResponse.builder()
49+
.response(SdkHttpResponse.builder()
50+
.statusCode(200)
51+
.build())
52+
.responseBody(AbortableInputStream.create(new StringInputStream("{}")))
53+
.build());
54+
client.allTypes();
55+
assertThat(mockHttpClient.getLastRequest().firstMatchingHeader("X-Amzn-Trace-Id")).hasValue("bar");
56+
}
57+
});
58+
}
59+
}

‎test/http-client-tests/src/main/java/software/amazon/awssdk/http/HttpTestUtils.java

Lines changed: 13 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -76,8 +76,18 @@ public static CompletableFuture<byte[]> sendHeadRequest(int serverPort, SdkAsync
7676
}
7777

7878
private static CompletableFuture<byte[]> sendRequest(int serverPort,
79-
SdkAsyncHttpClient client,
80-
SdkHttpMethod httpMethod) {
79+
SdkAsyncHttpClient client,
80+
SdkHttpMethod httpMethod) {
81+
SdkHttpFullRequest request = SdkHttpFullRequest.builder()
82+
.method(httpMethod)
83+
.protocol("https")
84+
.host("127.0.0.1")
85+
.port(serverPort)
86+
.build();
87+
return sendRequest(client, request);
88+
}
89+
90+
public static CompletableFuture<byte[]> sendRequest(SdkAsyncHttpClient client, SdkHttpFullRequest request) {
8191
ByteArrayOutputStream responsePayload = new ByteArrayOutputStream();
8292
AtomicBoolean responsePayloadReceived = new AtomicBoolean(false);
8393
return client.execute(AsyncExecuteRequest.builder()
@@ -98,12 +108,7 @@ public void onStream(Publisher<ByteBuffer> stream) {
98108
public void onError(Throwable error) {
99109
}
100110
})
101-
.request(SdkHttpFullRequest.builder()
102-
.method(httpMethod)
103-
.protocol("https")
104-
.host("127.0.0.1")
105-
.port(serverPort)
106-
.build())
111+
.request(request)
107112
.requestContentPublisher(new EmptyPublisher())
108113
.build())
109114
.thenApply(v -> responsePayloadReceived.get() ? responsePayload.toByteArray() : null);

‎test/http-client-tests/src/main/java/software/amazon/awssdk/http/SdkAsyncHttpClientH1TestSuite.java

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import static io.netty.handler.codec.http.HttpHeaderValues.TEXT_PLAIN;
2323
import static io.netty.handler.codec.http.HttpResponseStatus.INTERNAL_SERVER_ERROR;
2424
import static io.netty.handler.codec.http.HttpResponseStatus.OK;
25+
import static org.assertj.core.api.AssertionsForClassTypes.assertThatThrownBy;
2526
import static org.assertj.core.api.AssertionsForInterfaceTypes.assertThat;
2627

2728
import io.netty.bootstrap.ServerBootstrap;
@@ -45,6 +46,7 @@
4546
import io.netty.handler.ssl.SslContext;
4647
import io.netty.handler.ssl.SslContextBuilder;
4748
import io.netty.handler.ssl.util.SelfSignedCertificate;
49+
import java.net.URI;
4850
import java.nio.charset.StandardCharsets;
4951
import java.util.ArrayList;
5052
import java.util.List;
@@ -132,6 +134,19 @@ public void headRequestResponsesHaveNoPayload() {
132134
assertThat(responseData).isNull();
133135
}
134136

137+
@Test
138+
public void naughtyHeaderCharactersDoNotGetToServer() {
139+
String naughtyHeader = "foo\r\nbar";
140+
assertThatThrownBy(() -> HttpTestUtils.sendRequest(client,
141+
SdkHttpFullRequest.builder()
142+
.uri(URI.create("https://localhost:" + server.port()))
143+
.method(SdkHttpMethod.POST)
144+
.appendHeader("h", naughtyHeader)
145+
.build())
146+
.join())
147+
.hasCauseInstanceOf(Exception.class);
148+
}
149+
135150
private static class Server extends ChannelInitializer<Channel> {
136151
private static final byte[] CONTENT = "helloworld".getBytes(StandardCharsets.UTF_8);
137152
private ServerBootstrap bootstrap;
@@ -141,6 +156,7 @@ private static class Server extends ChannelInitializer<Channel> {
141156
private SslContext sslCtx;
142157
private boolean return500OnFirstRequest;
143158
private boolean closeConnection;
159+
private volatile HttpRequest lastRequestReceived;
144160

145161
public void init() throws Exception {
146162
SelfSignedCertificate ssc = new SelfSignedCertificate();
@@ -178,6 +194,8 @@ private class BehaviorTestChannelHandler extends ChannelDuplexHandler {
178194
@Override
179195
public void channelRead(ChannelHandlerContext ctx, Object msg) {
180196
if (msg instanceof HttpRequest) {
197+
lastRequestReceived = (HttpRequest) msg;
198+
181199
HttpResponseStatus status;
182200
if (ctx.channel().equals(channels.get(0)) && return500OnFirstRequest) {
183201
status = INTERNAL_SERVER_ERROR;

0 commit comments

Comments
 (0)