software.amazon.awssdk.auth.credentials.AwsSessionCredentials Java Examples

@Override
public AwsCredentials resolveCredentials() {
    String accessKey = environment.getProperty(ACCESS_KEY_ENV_VAR, String.class, environment.getProperty(ALTERNATE_ACCESS_KEY_ENV_VAR, String.class, (String) null));

    String secretKey = environment.getProperty(SECRET_KEY_ENV_VAR, String.class, environment.getProperty(ALTERNATE_SECRET_KEY_ENV_VAR, String.class, (String) null));
    accessKey = StringUtils.trim(accessKey);
    secretKey = StringUtils.trim(secretKey);
    String sessionToken = StringUtils.trim(environment.getProperty(AWS_SESSION_TOKEN_ENV_VAR, String.class, (String) null));

    if (StringUtils.isBlank(accessKey) || StringUtils.isBlank(secretKey)) {
        throw SdkClientException.create(
                "Unable to load AWS credentials from environment "
                        + "(" + ACCESS_KEY_ENV_VAR + " (or " + ALTERNATE_ACCESS_KEY_ENV_VAR + ") and "
                        + SECRET_KEY_ENV_VAR + " (or " + ALTERNATE_SECRET_KEY_ENV_VAR + "))");
    }

    return sessionToken == null
            ? AwsBasicCredentials.create(accessKey, secretKey)
            : AwsSessionCredentials.create(accessKey, secretKey, sessionToken);
}
 

@Override
public AwsCredentials resolveCredentials() {
    AWSCredentials current = oldCredentialsProvider.getCredentials();
    if (current instanceof AWSSessionCredentials) {
        return AwsSessionCredentials.create(current.getAWSAccessKeyId(), current.getAWSSecretKey(), ((AWSSessionCredentials) current).getSessionToken());
    }
    return new AwsCredentials() {
        @Override
        public String accessKeyId() {
            return current.getAWSAccessKeyId();
        }

        @Override
        public String secretAccessKey() {
            return current.getAWSSecretKey();
        }
    };
}
 

private static DynamoDbClient buildClient() {
  final AwsSessionCredentials awsCreds = AwsSessionCredentials.create("access_key_id", "secret_key_id", "session_token");
  return DynamoDbClient
    .builder()
    .endpointOverride(URI.create("http://localhost:8000"))
    .region(Region.US_WEST_2)
    .credentialsProvider(StaticCredentialsProvider.create(awsCreds))
    .overrideConfiguration(ClientOverrideConfiguration.builder().apiCallTimeout(Duration.ofSeconds(1)).build())
    .build();

  // final AwsClientBuilder.EndpointConfiguration endpointConfiguration = new
  // AwsClientBuilder.EndpointConfiguration(
  // "http://localhost:8000", "us-west-2");
  // final BasicAWSCredentials awsCreds = new
  // BasicAWSCredentials("access_key_id", "secret_key_id");
  // final AmazonDynamoDBClientBuilder builder = AmazonDynamoDBClientBuilder
  // .standard()
  // .withEndpointConfiguration(endpointConfiguration)
  // .withCredentials(new AWSStaticCredentialsProvider(awsCreds))
  // .withClientConfiguration(new
  // ClientConfiguration().withConnectionTimeout(1));
  //
  // return builder.build();
}
 

public AmazonWebServicesClientProxy(final boolean inHandshakeMode,
                                    final LoggerProxy loggerProxy,
                                    final Credentials credentials,
                                    final Supplier<Long> remainingTimeToExecute,
                                    final DelayFactory override) {
    this.inHandshakeMode = inHandshakeMode;
    this.loggerProxy = loggerProxy;
    this.remainingTimeInMillis = remainingTimeToExecute;

    BasicSessionCredentials basicSessionCredentials = new BasicSessionCredentials(credentials.getAccessKeyId(),
                                                                                  credentials.getSecretAccessKey(),
                                                                                  credentials.getSessionToken());
    this.v1CredentialsProvider = new AWSStaticCredentialsProvider(basicSessionCredentials);

    AwsSessionCredentials awsSessionCredentials = AwsSessionCredentials.create(credentials.getAccessKeyId(),
        credentials.getSecretAccessKey(), credentials.getSessionToken());
    this.v2CredentialsProvider = StaticCredentialsProvider.create(awsSessionCredentials);
    this.override = Objects.requireNonNull(override);
}
 

public void callClientWithCredentialsProvider(Instant credentialsExpirationDate, int numTimesInvokeCredentialsProvider) {
    Credentials credentials = Credentials.builder().accessKeyId("a").secretAccessKey("b").sessionToken("c").expiration(credentialsExpirationDate).build();
    RequestT request = getRequest();
    ResponseT response = getResponse(credentials);

    when(callClient(stsClient, request)).thenReturn(response);

    try (StsCredentialsProvider credentialsProvider = createCredentialsProviderBuilder(request).stsClient(stsClient).build()) {
        for (int i = 0; i < numTimesInvokeCredentialsProvider; ++i) {
            AwsSessionCredentials providedCredentials = (AwsSessionCredentials) credentialsProvider.resolveCredentials();
            assertThat(providedCredentials.accessKeyId()).isEqualTo("a");
            assertThat(providedCredentials.secretAccessKey()).isEqualTo("b");
            assertThat(providedCredentials.sessionToken()).isEqualTo("c");
        }
    }
}
 

/** Tests that we can call GetFederatedSession to start a federated session. */
@Test
public void testGetFederatedSessionToken() throws Exception {
    if (CREDENTIALS_PROVIDER_CHAIN.resolveCredentials() instanceof AwsSessionCredentials) {
        log.warn(() -> "testGetFederatedSessionToken() skipped due to the current credentials being session credentials. " +
                       "Session credentials cannot be used to get federation tokens.");
        return;
    }

    GetFederationTokenRequest request = GetFederationTokenRequest.builder()
                                                                 .durationSeconds(SESSION_DURATION)
                                                                 .name("Name").build();
    GetFederationTokenResponse result = sts.getFederationToken(request);

    assertNotNull(result.credentials().accessKeyId());
    assertNotNull(result.credentials().expiration());
    assertNotNull(result.credentials().secretAccessKey());
    assertNotNull(result.credentials().sessionToken());

    assertNotNull(result.federatedUser().arn());
    assertNotNull(result.federatedUser().federatedUserId());
}
 

/** Tests that we can call GetSession to start a session. */
@Test
public void testGetSessionToken() throws Exception {
    if (CREDENTIALS_PROVIDER_CHAIN.resolveCredentials() instanceof AwsSessionCredentials) {
        log.warn(() -> "testGetSessionToken() skipped due to the current credentials being session credentials. " +
                       "Session credentials cannot be used to get other session tokens.");
        return;
    }

    GetSessionTokenRequest request = GetSessionTokenRequest.builder().durationSeconds(SESSION_DURATION).build();
    GetSessionTokenResponse result = sts.getSessionToken(request);

    assertNotNull(result.credentials().accessKeyId());
    assertNotNull(result.credentials().expiration());
    assertNotNull(result.credentials().secretAccessKey());
    assertNotNull(result.credentials().sessionToken());
}
 

private static DynamoDbClient buildClient() {
  final AwsSessionCredentials awsCreds = AwsSessionCredentials.create("access_key_id", "secret_key_id", "session_token");
  return DynamoDbClient.builder()
    .endpointOverride(URI.create("http://localhost:8000"))
    .region(Region.US_WEST_2)
    .credentialsProvider(StaticCredentialsProvider.create(awsCreds))
    .overrideConfiguration(ClientOverrideConfiguration.builder()
      .apiCallTimeout(Duration.ofSeconds(1)).build())
    .build();
}
 

@Test
public void profileFileWithSessionCredentialsLoadsCorrectly() {
    ProfileFile profileFile = allTypesProfile();
    assertThat(profileFile.profile("profile-with-session-token")).hasValueSatisfying(profile -> {
        assertThat(profile.property(ProfileProperty.REGION)).isNotPresent();
        assertThat(new ProfileCredentialsUtils(profile, profileFile::profile).credentialsProvider()).hasValueSatisfying(credentialsProvider -> {
            assertThat(credentialsProvider.resolveCredentials()).satisfies(credentials -> {
                assertThat(credentials).isInstanceOf(AwsSessionCredentials.class);
                assertThat(credentials.accessKeyId()).isEqualTo("defaultAccessKey");
                assertThat(credentials.secretAccessKey()).isEqualTo("defaultSecretAccessKey");
                Assertions.assertThat(((AwsSessionCredentials) credentials).sessionToken()).isEqualTo("awsSessionToken");
            });
        });
    });
}
 

@Test
public void equalsHashCode() {
    AwsSessionCredentials credentials =
        AwsSessionCredentials.create("test", "key", "sessionToken");

    AwsSessionCredentials anotherCredentials =
        AwsSessionCredentials.create("test", "key", "sessionToken");
    assertThat(credentials).isEqualTo(anotherCredentials);
    assertThat(credentials.hashCode()).isEqualTo(anotherCredentials.hashCode());
}
 

/**
 * Parse the process output to retrieve the credentials.
 */
private AwsCredentials credentials(JsonNode credentialsJson) {
    String accessKeyId = getText(credentialsJson, "AccessKeyId");
    String secretAccessKey = getText(credentialsJson, "SecretAccessKey");
    String sessionToken = getText(credentialsJson, "SessionToken");

    Validate.notEmpty(accessKeyId, "AccessKeyId cannot be empty.");
    Validate.notEmpty(secretAccessKey, "SecretAccessKey cannot be empty.");

    if (sessionToken != null) {
        return AwsSessionCredentials.create(accessKeyId, secretAccessKey, sessionToken);
    } else {
        return AwsBasicCredentials.create(accessKeyId, secretAccessKey);
    }
}
 

@Override
public AwsCredentials resolveCredentials() {
    String accessKey = trim(loadSetting(SdkSystemSetting.AWS_ACCESS_KEY_ID).orElse(null));
    String secretKey = trim(loadSetting(SdkSystemSetting.AWS_SECRET_ACCESS_KEY).orElse(null));
    String sessionToken = trim(loadSetting(SdkSystemSetting.AWS_SESSION_TOKEN).orElse(null));

    if (StringUtils.isEmpty(accessKey)) {
        throw SdkClientException.builder()
                                .message(String.format("Unable to load credentials from system settings. Access key must be" +
                                         " specified either via environment variable (%s) or system property (%s).",
                                         SdkSystemSetting.AWS_ACCESS_KEY_ID.environmentVariable(),
                                         SdkSystemSetting.AWS_ACCESS_KEY_ID.property()))
                                .build();
    }

    if (StringUtils.isEmpty(secretKey)) {
        throw SdkClientException.builder()
                                .message(String.format("Unable to load credentials from system settings. Secret key must be" +
                                         " specified either via environment variable (%s) or system property (%s).",
                                         SdkSystemSetting.AWS_SECRET_ACCESS_KEY.environmentVariable(),
                                         SdkSystemSetting.AWS_SECRET_ACCESS_KEY.property()))
                                .build();
    }

    return sessionToken == null ? AwsBasicCredentials.create(accessKey, secretKey)
                                : AwsSessionCredentials.create(accessKey, secretKey, sessionToken);
}
 

/**
 * Load a set of session credentials that have been configured in this profile.
 */
private AwsCredentialsProvider sessionProfileCredentialsProvider() {
    requireProperties(ProfileProperty.AWS_ACCESS_KEY_ID,
                      ProfileProperty.AWS_SECRET_ACCESS_KEY,
                      ProfileProperty.AWS_SESSION_TOKEN);
    AwsCredentials credentials = AwsSessionCredentials.create(properties.get(ProfileProperty.AWS_ACCESS_KEY_ID),
                                                              properties.get(ProfileProperty.AWS_SECRET_ACCESS_KEY),
                                                              properties.get(ProfileProperty.AWS_SESSION_TOKEN));
    return StaticCredentialsProvider.create(credentials);
}
 

/**
 * Loads the individual access key ID and secret key from the specified credentials, trimming any extra whitespace from the
 * credentials.
 *
 * <p>Returns either a {@link AwsSessionCredentials} or a {@link AwsBasicCredentials} object, depending on the input type.
 *
 * @return A new credentials object with the sanitized credentials.
 */
protected AwsCredentials sanitizeCredentials(AwsCredentials credentials) {
    String accessKeyId = StringUtils.trim(credentials.accessKeyId());
    String secretKey = StringUtils.trim(credentials.secretAccessKey());

    if (credentials instanceof AwsSessionCredentials) {
        AwsSessionCredentials sessionCredentials = (AwsSessionCredentials) credentials;
        return AwsSessionCredentials.create(accessKeyId,
                                            secretKey,
                                            StringUtils.trim(sessionCredentials.sessionToken()));
    }

    return AwsBasicCredentials.create(accessKeyId, secretKey);
}
 

protected SdkHttpFullRequest.Builder doPresign(SdkHttpFullRequest request,
                                               Aws4SignerRequestParams requestParams,
                                               U signingParams) {

    SdkHttpFullRequest.Builder mutableRequest = request.toBuilder();

    long expirationInSeconds = getSignatureDurationInSeconds(requestParams, signingParams);
    addHostHeader(mutableRequest);

    AwsCredentials sanitizedCredentials = sanitizeCredentials(signingParams.awsCredentials());
    if (sanitizedCredentials instanceof AwsSessionCredentials) {
        // For SigV4 pre-signing URL, we need to add "X-Amz-Security-Token"
        // as a query string parameter, before constructing the canonical
        // request.
        mutableRequest.putRawQueryParameter(SignerConstant.X_AMZ_SECURITY_TOKEN,
                                            ((AwsSessionCredentials) sanitizedCredentials).sessionToken());
    }

    // Add the important parameters for v4 signing
    Map<String, List<String>> canonicalizedHeaders = canonicalizeSigningHeaders(mutableRequest.headers());
    String signedHeadersString = getSignedHeadersString(canonicalizedHeaders);

    addPreSignInformationToRequest(mutableRequest, signedHeadersString, sanitizedCredentials,
                                   requestParams, expirationInSeconds);

    String contentSha256 = calculateContentHashPresign(mutableRequest, signingParams);

    String canonicalRequest = createCanonicalRequest(mutableRequest, canonicalizedHeaders, signedHeadersString,
                                                     contentSha256, signingParams.doubleUrlEncode());

    String stringToSign = createStringToSign(canonicalRequest, requestParams);

    byte[] signingKey = deriveSigningKey(sanitizedCredentials, requestParams);

    byte[] signature = computeSignature(stringToSign, signingKey);

    mutableRequest.putRawQueryParameter(SignerConstant.X_AMZ_SIGNATURE, BinaryUtils.toHex(signature));

    return mutableRequest;
}
 

protected SdkHttpFullRequest.Builder doSign(SdkHttpFullRequest request,
                                            Aws4SignerRequestParams requestParams,
                                            T signingParams) {

    SdkHttpFullRequest.Builder mutableRequest = request.toBuilder();
    AwsCredentials sanitizedCredentials = sanitizeCredentials(signingParams.awsCredentials());
    if (sanitizedCredentials instanceof AwsSessionCredentials) {
        addSessionCredentials(mutableRequest, (AwsSessionCredentials) sanitizedCredentials);
    }

    addHostHeader(mutableRequest);
    addDateHeader(mutableRequest, requestParams.getFormattedRequestSigningDateTime());

    String contentSha256 = calculateContentHash(mutableRequest, signingParams);
    mutableRequest.firstMatchingHeader(SignerConstant.X_AMZ_CONTENT_SHA256)
                  .filter(h -> h.equals("required"))
                  .ifPresent(h -> mutableRequest.putHeader(SignerConstant.X_AMZ_CONTENT_SHA256, contentSha256));

    Map<String, List<String>> canonicalHeaders = canonicalizeSigningHeaders(mutableRequest.headers());
    String signedHeadersString = getSignedHeadersString(canonicalHeaders);

    String canonicalRequest = createCanonicalRequest(mutableRequest,
                                                     canonicalHeaders,
                                                     signedHeadersString,
                                                     contentSha256,
                                                     signingParams.doubleUrlEncode());

    String stringToSign = createStringToSign(canonicalRequest, requestParams);

    byte[] signingKey = deriveSigningKey(sanitizedCredentials, requestParams);

    byte[] signature = computeSignature(stringToSign, signingKey);

    mutableRequest.putHeader(SignerConstant.AUTHORIZATION,
                             buildAuthorizationHeader(signature, sanitizedCredentials, requestParams, signedHeadersString));

    processRequestPayload(mutableRequest, signature, signingKey, requestParams, signingParams);

    return mutableRequest;
}
 

private static DynamoDbAsyncClient buildAsyncClient() {
  final AwsSessionCredentials awsCreds = AwsSessionCredentials.create("access_key_id", "secret_key_id", "session_token");
  final DynamoDbAsyncClient build = DynamoDbAsyncClient.builder()
    .endpointOverride(URI.create("http://localhost:8000"))
    .region(Region.US_WEST_2)
    .credentialsProvider(StaticCredentialsProvider.create(awsCreds))
    .overrideConfiguration(ClientOverrideConfiguration.builder()
      .apiCallTimeout(Duration.ofSeconds(1)).build())
    .build();
  return build;
}
 

private CredentialsProvider prepareMockProvider() {
    return new CredentialsProvider() {
        @Override
        public AwsSessionCredentials get() {
            return MockCreds;
        }

        @Override
        public void setCredentials(Credentials credentials) {

        }
    };
}
 

@Override
protected void addSessionCredentials(SdkHttpFullRequest.Builder mutableRequest,
                                     AwsSessionCredentials credentials) {
    mutableRequest.putHeader(SignerConstant.X_AMZ_SECURITY_TOKEN, credentials.sessionToken());
}
 

public AwsSessionCredentials get() {
    return this.awsSessionCredentials;
}
 

public void setCredentials(final Credentials credentials) {
    this.awsSessionCredentials = AwsSessionCredentials.create(credentials.getAccessKeyId(), credentials.getSecretAccessKey(),
        credentials.getSessionToken());
}
 

@Test
public void testResponseDescriptors() throws Exception {
    String responseBody = "{\"LastEvaluatedTableName\":\"baz\",\"TableNames\":[\"foo\",\"bar\",\"baz\"]}";
    SdkHttpResponse mockResponse = SdkHttpResponse.builder()
            .statusCode(200)
            .putHeader("x-amzn-requestid", "1111-2222-3333-4444")
            .putHeader("Content-Length", "84")
            .putHeader("Content-Type", "application/x-amz-json-1.0")
            .build();
    SdkHttpClient mockClient = mockSdkHttpClient(mockResponse, responseBody);

    DynamoDbClient client = DynamoDbClient.builder()
            .httpClient(mockClient)
            .endpointOverride(URI.create("http://example.com"))
            .region(Region.of("us-west-42"))
            .credentialsProvider(StaticCredentialsProvider.create(
                    AwsSessionCredentials.create("key", "secret", "session")
            ))
            .overrideConfiguration(ClientOverrideConfiguration.builder()
                    .addExecutionInterceptor(new TracingInterceptor())
                    .build()
            )
            .build();

    Segment segment = AWSXRay.getCurrentSegment();
    client.listTables(ListTablesRequest.builder()
            .limit(3)
            .build()
    );

    Assert.assertEquals(1, segment.getSubsegments().size());
    Subsegment subsegment = segment.getSubsegments().get(0);
    Map<String, Object> awsStats = subsegment.getAws();
    @SuppressWarnings("unchecked")
    Map<String, Object> httpResponseStats = (Map<String, Object>) subsegment.getHttp().get("response");

    Assert.assertEquals("ListTables", awsStats.get("operation"));
    Assert.assertEquals(3, awsStats.get("limit"));
    Assert.assertEquals("1111-2222-3333-4444", awsStats.get("request_id"));
    Assert.assertEquals(3, awsStats.get("table_count"));
    Assert.assertEquals("us-west-42", awsStats.get("region"));
    Assert.assertEquals(0, awsStats.get("retries"));
    Assert.assertEquals(84L, httpResponseStats.get("content_length"));
    Assert.assertEquals(200, httpResponseStats.get("status"));
    Assert.assertEquals(false, subsegment.isInProgress());
}
 

@Test
public void testLambdaInvokeSubsegmentContainsFunctionName() throws Exception {
    SdkHttpClient mockClient = mockSdkHttpClient(generateLambdaInvokeResponse(200));

    LambdaClient client = LambdaClient.builder()
            .httpClient(mockClient)
            .endpointOverride(URI.create("http://example.com"))
            .region(Region.of("us-west-42"))
            .credentialsProvider(StaticCredentialsProvider.create(
                    AwsSessionCredentials.create("key", "secret", "session")
            ))
            .overrideConfiguration(ClientOverrideConfiguration.builder()
                    .addExecutionInterceptor(new TracingInterceptor())
                    .build()
            )
            .build();


    Segment segment = AWSXRay.getCurrentSegment();

    client.invoke(InvokeRequest.builder()
            .functionName("testFunctionName")
            .build()
    );

    Assert.assertEquals(1, segment.getSubsegments().size());
    Subsegment subsegment = segment.getSubsegments().get(0);
    Map<String, Object> awsStats = subsegment.getAws();
    @SuppressWarnings("unchecked")
    Map<String, Object> httpResponseStats = (Map<String, Object>) subsegment.getHttp().get("response");

    Assert.assertEquals("Invoke", awsStats.get("operation"));
    Assert.assertEquals("testFunctionName", awsStats.get("function_name"));
    Assert.assertEquals("1111-2222-3333-4444", awsStats.get("request_id"));
    Assert.assertEquals("extended", awsStats.get("id_2"));
    Assert.assertEquals("Failure", awsStats.get("function_error"));
    Assert.assertEquals("us-west-42", awsStats.get("region"));
    Assert.assertEquals(0, awsStats.get("retries"));
    Assert.assertEquals(2L, httpResponseStats.get("content_length"));
    Assert.assertEquals(200, httpResponseStats.get("status"));
    Assert.assertEquals(false, subsegment.isInProgress());
}
 

@Test
public void test500Exception() throws Exception {
    SdkHttpClient mockClient = mockSdkHttpClient(generateLambdaInvokeResponse(500));

    LambdaClient client = LambdaClient.builder()
            .httpClient(mockClient)
            .endpointOverride(URI.create("http://example.com"))
            .region(Region.of("us-west-42"))
            .credentialsProvider(StaticCredentialsProvider.create(
                    AwsSessionCredentials.create("key", "secret", "session")
            ))
            .overrideConfiguration(ClientOverrideConfiguration.builder()
                    .addExecutionInterceptor(new TracingInterceptor())
                    .build()
            )
            .build();


    Segment segment = AWSXRay.getCurrentSegment();

    try {
        client.invoke(InvokeRequest.builder()
                .functionName("testFunctionName")
                .build()
        );
    } catch (Exception e) {
        // ignore SDK errors
    } finally {
        Assert.assertEquals(1, segment.getSubsegments().size());
        Subsegment subsegment = segment.getSubsegments().get(0);
        Map<String, Object> awsStats = subsegment.getAws();
        @SuppressWarnings("unchecked")
        Map<String, Object> httpResponseStats = (Map<String, Object>) subsegment.getHttp().get("response");
        Cause cause = subsegment.getCause();

        Assert.assertEquals("Invoke", awsStats.get("operation"));
        Assert.assertEquals("testFunctionName", awsStats.get("function_name"));
        Assert.assertEquals("1111-2222-3333-4444", awsStats.get("request_id"));
        Assert.assertEquals("extended", awsStats.get("id_2"));
        Assert.assertEquals("us-west-42", awsStats.get("region"));
        // 500 exceptions are retried
        Assert.assertEquals(3, awsStats.get("retries"));
        Assert.assertEquals(2L, httpResponseStats.get("content_length"));
        Assert.assertEquals(500, httpResponseStats.get("status"));
        Assert.assertEquals(false, subsegment.isError());
        Assert.assertEquals(false, subsegment.isThrottle());
        Assert.assertEquals(true, subsegment.isFault());
        Assert.assertEquals(1, cause.getExceptions().size());
        Assert.assertEquals(true, cause.getExceptions().get(0).isRemote());
    }
}
 

@Test
public void testAsyncLambdaInvokeSubsegmentContainsFunctionName() {
    SdkAsyncHttpClient mockClient = mockSdkAsyncHttpClient(generateLambdaInvokeResponse(200));

    LambdaAsyncClient client = LambdaAsyncClient.builder()
            .httpClient(mockClient)
            .endpointOverride(URI.create("http://example.com"))
            .region(Region.of("us-west-42"))
            .credentialsProvider(StaticCredentialsProvider.create(
                    AwsSessionCredentials.create("key", "secret", "session")
            ))
            .overrideConfiguration(ClientOverrideConfiguration.builder()
                    .addExecutionInterceptor(new TracingInterceptor())
                    .build()
            )
            .build();

    Segment segment = AWSXRay.getCurrentSegment();

    client.invoke(InvokeRequest.builder()
            .functionName("testFunctionName")
            .build()
    ).join();

    Assert.assertEquals(1, segment.getSubsegments().size());
    Subsegment subsegment = segment.getSubsegments().get(0);
    Map<String, Object> awsStats = subsegment.getAws();
    @SuppressWarnings("unchecked")
    Map<String, Object> httpResponseStats = (Map<String, Object>) subsegment.getHttp().get("response");

    Assert.assertEquals("Invoke", awsStats.get("operation"));
    Assert.assertEquals("testFunctionName", awsStats.get("function_name"));
    Assert.assertEquals("1111-2222-3333-4444", awsStats.get("request_id"));
    Assert.assertEquals("extended", awsStats.get("id_2"));
    Assert.assertEquals("Failure", awsStats.get("function_error"));
    Assert.assertEquals("us-west-42", awsStats.get("region"));
    Assert.assertEquals(0, awsStats.get("retries"));
    Assert.assertEquals(2L, httpResponseStats.get("content_length"));
    Assert.assertEquals(200, httpResponseStats.get("status"));
    Assert.assertEquals(false, subsegment.isInProgress());
}
 

@Test
public void test400Exception() throws Exception {
    SdkHttpClient mockClient = mockSdkHttpClient(generateLambdaInvokeResponse(400));

    LambdaClient client = LambdaClient.builder()
            .httpClient(mockClient)
            .endpointOverride(URI.create("http://example.com"))
            .region(Region.of("us-west-42"))
            .credentialsProvider(StaticCredentialsProvider.create(
                    AwsSessionCredentials.create("key", "secret", "session")
            ))
            .overrideConfiguration(ClientOverrideConfiguration.builder()
                    .addExecutionInterceptor(new TracingInterceptor())
                    .build()
            )
            .build();


    Segment segment = AWSXRay.getCurrentSegment();

    try {
        client.invoke(InvokeRequest.builder()
                .functionName("testFunctionName")
                .build()
        );
    } catch (Exception e) {
        // ignore SDK errors
    } finally {
        Assert.assertEquals(1, segment.getSubsegments().size());
        Subsegment subsegment = segment.getSubsegments().get(0);
        Map<String, Object> awsStats = subsegment.getAws();
        @SuppressWarnings("unchecked")
        Map<String, Object> httpResponseStats = (Map<String, Object>) subsegment.getHttp().get("response");
        Cause cause = subsegment.getCause();

        Assert.assertEquals("Invoke", awsStats.get("operation"));
        Assert.assertEquals("testFunctionName", awsStats.get("function_name"));
        Assert.assertEquals("1111-2222-3333-4444", awsStats.get("request_id"));
        Assert.assertEquals("extended", awsStats.get("id_2"));
        Assert.assertEquals("us-west-42", awsStats.get("region"));
        Assert.assertEquals(0, awsStats.get("retries"));
        Assert.assertEquals(2L, httpResponseStats.get("content_length"));
        Assert.assertEquals(400, httpResponseStats.get("status"));
        Assert.assertEquals(false, subsegment.isInProgress());
        Assert.assertEquals(true, subsegment.isError());
        Assert.assertEquals(false, subsegment.isThrottle());
        Assert.assertEquals(false, subsegment.isFault());
        Assert.assertEquals(1, cause.getExceptions().size());
        Assert.assertEquals(true, cause.getExceptions().get(0).isRemote());
    }
}
 

@Test
public void testAsync400Exception() {
    SdkAsyncHttpClient mockClient = mockSdkAsyncHttpClient(generateLambdaInvokeResponse(400));

    LambdaAsyncClient client = LambdaAsyncClient.builder()
            .httpClient(mockClient)
            .endpointOverride(URI.create("http://example.com"))
            .region(Region.of("us-west-42"))
            .credentialsProvider(StaticCredentialsProvider.create(
                    AwsSessionCredentials.create("key", "secret", "session")
            ))
            .overrideConfiguration(ClientOverrideConfiguration.builder()
                    .addExecutionInterceptor(new TracingInterceptor())
                    .build()
            )
            .build();

    Segment segment = AWSXRay.getCurrentSegment();

    try {
        client.invoke(InvokeRequest.builder()
                .functionName("testFunctionName")
                .build()
        ).get();
    } catch (Exception e) {
        // ignore exceptions
    } finally {
        Assert.assertEquals(1, segment.getSubsegments().size());
        Subsegment subsegment = segment.getSubsegments().get(0);
        Map<String, Object> awsStats = subsegment.getAws();
        @SuppressWarnings("unchecked")
        Map<String, Object> httpResponseStats = (Map<String, Object>) subsegment.getHttp().get("response");
        Cause cause = subsegment.getCause();

        Assert.assertEquals("Invoke", awsStats.get("operation"));
        Assert.assertEquals("testFunctionName", awsStats.get("function_name"));
        Assert.assertEquals("1111-2222-3333-4444", awsStats.get("request_id"));
        Assert.assertEquals("extended", awsStats.get("id_2"));
        Assert.assertEquals("us-west-42", awsStats.get("region"));
        Assert.assertEquals(0, awsStats.get("retries"));
        Assert.assertEquals(2L, httpResponseStats.get("content_length"));
        Assert.assertEquals(400, httpResponseStats.get("status"));
        Assert.assertEquals(false, subsegment.isInProgress());
        Assert.assertEquals(true, subsegment.isError());
        Assert.assertEquals(false, subsegment.isThrottle());
        Assert.assertEquals(false, subsegment.isFault());
        Assert.assertEquals(1, cause.getExceptions().size());
        Assert.assertEquals(true, cause.getExceptions().get(0).isRemote());
    }
}
 

@Test
public void testThrottledException() throws Exception {
    SdkHttpClient mockClient = mockSdkHttpClient(generateLambdaInvokeResponse(429));

    LambdaClient client = LambdaClient.builder()
            .httpClient(mockClient)
            .endpointOverride(URI.create("http://example.com"))
            .region(Region.of("us-west-42"))
            .credentialsProvider(StaticCredentialsProvider.create(
                    AwsSessionCredentials.create("key", "secret", "session")
            ))
            .overrideConfiguration(ClientOverrideConfiguration.builder()
                    .addExecutionInterceptor(new TracingInterceptor())
                    .build()
            )
            .build();


    Segment segment = AWSXRay.getCurrentSegment();

    try {
        client.invoke(InvokeRequest.builder()
                .functionName("testFunctionName")
                .build()
        );
    } catch (Exception e) {
        // ignore SDK errors
    } finally {
        Assert.assertEquals(1, segment.getSubsegments().size());
        Subsegment subsegment = segment.getSubsegments().get(0);
        Map<String, Object> awsStats = subsegment.getAws();
        @SuppressWarnings("unchecked")
        Map<String, Object> httpResponseStats = (Map<String, Object>) subsegment.getHttp().get("response");
        Cause cause = subsegment.getCause();

        Assert.assertEquals("Invoke", awsStats.get("operation"));
        Assert.assertEquals("testFunctionName", awsStats.get("function_name"));
        Assert.assertEquals("1111-2222-3333-4444", awsStats.get("request_id"));
        Assert.assertEquals("extended", awsStats.get("id_2"));
        Assert.assertEquals("us-west-42", awsStats.get("region"));
        // throttled requests are retried
        Assert.assertEquals(3, awsStats.get("retries"));
        Assert.assertEquals(2L, httpResponseStats.get("content_length"));
        Assert.assertEquals(429, httpResponseStats.get("status"));
        Assert.assertEquals(true, subsegment.isError());
        Assert.assertEquals(true, subsegment.isThrottle());
        Assert.assertEquals(false, subsegment.isFault());
        Assert.assertEquals(1, cause.getExceptions().size());
        Assert.assertEquals(true, cause.getExceptions().get(0).isRemote());
    }
}
 

@Test
public void testAsyncThrottledException() {
    SdkAsyncHttpClient mockClient = mockSdkAsyncHttpClient(generateLambdaInvokeResponse(429));

    LambdaAsyncClient client = LambdaAsyncClient.builder()
            .httpClient(mockClient)
            .endpointOverride(URI.create("http://example.com"))
            .region(Region.of("us-west-42"))
            .credentialsProvider(StaticCredentialsProvider.create(
                    AwsSessionCredentials.create("key", "secret", "session")
            ))
            .overrideConfiguration(ClientOverrideConfiguration.builder()
                    .addExecutionInterceptor(new TracingInterceptor())
                    .build()
            )
            .build();

    Segment segment = AWSXRay.getCurrentSegment();

    try {
        client.invoke(InvokeRequest.builder()
                .functionName("testFunctionName")
                .build()
        ).get();
    } catch (Exception e) {
        // ignore exceptions
    } finally {
        Assert.assertEquals(1, segment.getSubsegments().size());
        Subsegment subsegment = segment.getSubsegments().get(0);
        Map<String, Object> awsStats = subsegment.getAws();
        @SuppressWarnings("unchecked")
        Map<String, Object> httpResponseStats = (Map<String, Object>) subsegment.getHttp().get("response");
        Cause cause = subsegment.getCause();

        Assert.assertEquals("Invoke", awsStats.get("operation"));
        Assert.assertEquals("testFunctionName", awsStats.get("function_name"));
        Assert.assertEquals("1111-2222-3333-4444", awsStats.get("request_id"));
        Assert.assertEquals("extended", awsStats.get("id_2"));
        Assert.assertEquals("us-west-42", awsStats.get("region"));
        // throttled requests are retried
        Assert.assertEquals(3, awsStats.get("retries"));
        Assert.assertEquals(2L, httpResponseStats.get("content_length"));
        Assert.assertEquals(429, httpResponseStats.get("status"));
        Assert.assertEquals(true, subsegment.isError());
        Assert.assertEquals(true, subsegment.isThrottle());
        Assert.assertEquals(false, subsegment.isFault());
        Assert.assertEquals(1, cause.getExceptions().size());
        Assert.assertEquals(true, cause.getExceptions().get(0).isRemote());
    }
}
 

public AwsSessionCredentials getSessionCredentials() {
    return sessionCredentials;
}