software.amazon.awssdk.services.sts.StsClient Java Examples

/**
 * Discovers an AWS account.
 *
 * @return null, if no account was discovered. Otherwise the discovered {@link AwsAccount}.
 */
public static AwsAccount discover() {
  try {
    var account = new AwsAccount();

    // use STS to find account id using the default provider
    var stsClient = StsClient.builder().credentialsProvider(DEFAULT_PROVIDER).build();

    var identity = stsClient.getCallerIdentity();

    account.setAutoDiscovered(true);
    account.setAccountId(identity.account());
    account.setUser(identity.arn());

    return account;
  } catch (SdkClientException ex) {
    // TODO: log error, etc.
    return null;
  }
}
 

@Override
public void validate() throws IOException {
  try {
    // use STS to find account id and user

    var builder = StsClient.builder();

    if (!this.isAutoDiscovered()) {
      builder.region(Region.of(this.region));
      builder.credentialsProvider(() -> this);
    }

    var stsClient = builder.build();

    var identity = stsClient.getCallerIdentity();

    this.accountId = identity.account();
    this.user = identity.arn();

    LOGGER.info("Account {} validated with user {}.", this.accountId, this.user);
  } catch (SdkClientException | StsException ex) {
    // TODO: log error, etc.
    throw new IOException(ex.getMessage());
  }
}
 

@Test(expected = RuntimeException.class)
public void testVerifyCredentialsNoRetryOnAuthnError() {
  PowerMockito.mockStatic(StsClient.class);
  StsClient mockedClient = mock(StsClient.class);
  StsClientBuilder mockedClientBuilder = mock(StsClientBuilder.class);
  when(mockedClientBuilder.credentialsProvider(any(AwsCredentialsProvider.class))).thenReturn(mockedClientBuilder);
  when(mockedClientBuilder.region(any(Region.class))).thenReturn(mockedClientBuilder);
  when(mockedClientBuilder.build()).thenReturn(mockedClient);
  when(StsClient.builder()).thenReturn(mockedClientBuilder);

  TestExtendedS3FileSystem fs = new TestExtendedS3FileSystem();
  AtomicInteger retryAttemptNo = new AtomicInteger(0);
  when(mockedClient.getCallerIdentity(any(GetCallerIdentityRequest.class))).then(invocationOnMock -> {
    retryAttemptNo.incrementAndGet();
    throw StsException.builder().message("The security token included in the request is invalid. (Service: Sts, Status Code: 403, Request ID: a7e2e92e-5ebb-4343-87a1-21e4d64edcd4)").build();
  });
  fs.verifyCredentials(new Configuration());
  assertEquals(1, retryAttemptNo.get());
}
 

@Test
public void testVerifyCredentialsRetry() {
  PowerMockito.mockStatic(StsClient.class);
  StsClient mockedClient = mock(StsClient.class);
  StsClientBuilder mockedClientBuilder = mock(StsClientBuilder.class);
  when(mockedClientBuilder.credentialsProvider(any(AwsCredentialsProvider.class))).thenReturn(mockedClientBuilder);
  when(mockedClientBuilder.region(any(Region.class))).thenReturn(mockedClientBuilder);
  when(mockedClientBuilder.build()).thenReturn(mockedClient);
  when(StsClient.builder()).thenReturn(mockedClientBuilder);

  TestExtendedS3FileSystem fs = new TestExtendedS3FileSystem();
  AtomicInteger retryAttemptNo = new AtomicInteger(1);
  when(mockedClient.getCallerIdentity(any(GetCallerIdentityRequest.class))).then(invocationOnMock -> {
    if (retryAttemptNo.incrementAndGet() < 10) {
      throw new SdkClientException("Unable to load credentials from service endpoint.");
    }
    return null;
  });

  fs.verifyCredentials(new Configuration());
  assertEquals(10, retryAttemptNo.get());
}
 

/**
 * Checks if credentials are valid using GetCallerIdentity API call.
 */
protected void verifyCredentials(Configuration conf) throws RuntimeException {
    AwsCredentialsProvider awsCredentialsProvider = getAsync2Provider(conf);
    final StsClientBuilder stsClientBuilder = StsClient.builder()
      // Note that AWS SDKv2 client will close the credentials provider if needed when the client is closed
      .credentialsProvider(awsCredentialsProvider)
      .region(getAWSRegionFromConfigurationOrDefault(conf));
    try (StsClient stsClient = stsClientBuilder.build()) {
      retryer.call(() -> {
        GetCallerIdentityRequest request = GetCallerIdentityRequest.builder().build();
        stsClient.getCallerIdentity(request);
        return true;
      });
    } catch (Retryer.OperationFailedAfterRetriesException e) {
      throw new RuntimeException("Credential Verification failed.", e);
    }
}
 

@BeforeClass
public static void setupFixture() {
    createBucket(BUCKET);

    s3control = S3ControlClient.builder()
                               .region(Region.US_WEST_2)
                               .credentialsProvider(CREDENTIALS_PROVIDER_CHAIN)
                               .build();

    sts = StsClient.builder()
                   .region(Region.US_WEST_2)
                   .credentialsProvider(CREDENTIALS_PROVIDER_CHAIN)
                   .build();

    accountId = sts.getCallerIdentity().account();
    s3control.createAccessPoint(r -> r.accountId(accountId)
                                      .bucket(BUCKET)
                                      .name(AP_NAME));
}
 

private StsProfileCredentialsProvider(AwsCredentialsProvider parentCredentialsProvider, Profile profile) {
    String roleArn = requireProperty(profile, ProfileProperty.ROLE_ARN);
    String roleSessionName = profile.property(ProfileProperty.ROLE_SESSION_NAME)
                                    .orElseGet(() -> "aws-sdk-java-" + System.currentTimeMillis());
    String externalId = profile.property(ProfileProperty.EXTERNAL_ID).orElse(null);

    AssumeRoleRequest assumeRoleRequest = AssumeRoleRequest.builder()
                                                           .roleArn(roleArn)
                                                           .roleSessionName(roleSessionName)
                                                           .externalId(externalId)
                                                           .build();

    this.stsClient = StsClient.builder()
                              .applyMutation(client -> configureEndpoint(client, profile))
                              .credentialsProvider(parentCredentialsProvider)
                              .build();

    this.parentCredentialsProvider = parentCredentialsProvider;
    this.credentialsProvider = StsAssumeRoleCredentialsProvider.builder()
                                                               .stsClient(stsClient)
                                                               .refreshRequest(assumeRoleRequest)
                                                               .build();
}
 

@Override
public software.amazon.awssdk.auth.credentials.AwsCredentialsProvider getV2CredentialsProvider() {
    StsClient client = StsClient.create();
    return StsAssumeRoleCredentialsProvider.builder().stsClient(client).refreshRequest((req) -> {
        req.roleArn(roleArn).roleSessionName(roleSessionName).build();
    }).build();
}
 

public STSCredentialProviderV2(Configuration conf) {
  AwsCredentialsProvider awsCredentialsProvider = null;

  if (S3StoragePlugin.ACCESS_KEY_PROVIDER.equals(conf.get(Constants.ASSUMED_ROLE_CREDENTIALS_PROVIDER))) {
    awsCredentialsProvider = StaticCredentialsProvider.create(AwsBasicCredentials.create(
      conf.get(Constants.ACCESS_KEY), conf.get(Constants.SECRET_KEY)));
  } else if (S3StoragePlugin.EC2_METADATA_PROVIDER.equals(conf.get(Constants.ASSUMED_ROLE_CREDENTIALS_PROVIDER))) {
    awsCredentialsProvider = InstanceProfileCredentialsProvider.create();
  }

  final StsClientBuilder builder = StsClient.builder()
    .credentialsProvider(awsCredentialsProvider)
    .region(S3FileSystem.getAWSRegionFromConfigurationOrDefault(conf))
    .httpClientBuilder(initConnectionSettings(conf));
  S3FileSystem.getStsEndpoint(conf).ifPresent(e -> {
    try {
      builder.endpointOverride(new URI(e));
    } catch (URISyntaxException use) {
      throw UserException.sourceInBadState(use).buildSilently();
    }
  });

  initUserAgent(builder, conf);

  final AssumeRoleRequest assumeRoleRequest = AssumeRoleRequest.builder()
    .roleArn(conf.get(Constants.ASSUMED_ROLE_ARN))
    .roleSessionName(UUID.randomUUID().toString())
    .build();

  this.stsAssumeRoleCredentialsProvider = StsAssumeRoleCredentialsProvider.builder()
    .refreshRequest(assumeRoleRequest)
    .stsClient(builder.build())
    .build();
}
 

@Before
public void setup() {
    StsClient sts = StsClient.create();
    accountId = sts.getCallerIdentity().account();
    client = S3ControlClient.builder()
                            .overrideConfiguration(o -> o.addExecutionInterceptor(new AssertPayloadIsSignedExecutionInterceptor()))
                            .build();
}
 

private StsWebIdentityCredentialsProvider(WebIdentityTokenCredentialProperties credentialProperties) {
    String roleSessionName = credentialProperties.roleSessionName();
    String sessionName = roleSessionName != null ? roleSessionName : "aws-sdk-java-" + System.currentTimeMillis();

    OrRetryCondition retryCondition = OrRetryCondition.create(new StsRetryCondition(),
                                                              RetryCondition.defaultRetryCondition());

    this.stsClient = StsClient.builder()
                              .applyMutation(this::configureEndpoint)
                              .credentialsProvider(AnonymousCredentialsProvider.create())
                              .overrideConfiguration(o -> o.retryPolicy(r -> r.retryCondition(retryCondition)))
                              .build();

    AssumeRoleWithWebIdentityRequest request = AssumeRoleWithWebIdentityRequest.builder()
                                                                               .roleArn(credentialProperties.roleArn())
                                                                               .roleSessionName(sessionName)
                                                                               .build();

    AssumeRoleWithWebIdentityRequestSupplier supplier =
        new AssumeRoleWithWebIdentityRequestSupplier(request,
                                                     credentialProperties.webIdentityTokenFile());

    this.credentialsProvider =
        StsAssumeRoleWithWebIdentityCredentialsProvider.builder()
                                                       .stsClient(stsClient)
                                                       .refreshRequest(supplier)
                                                       .build();
}
 

@Override
protected AssumeRoleWithWebIdentityResponse callClient(StsClient client, AssumeRoleWithWebIdentityRequest request) {
    return client.assumeRoleWithWebIdentity(request);
}
 

@Override
protected AssumeRoleWithSamlResponse callClient(StsClient client, AssumeRoleWithSamlRequest request) {
    return client.assumeRoleWithSAML(request);
}
 

@Override
protected AssumeRoleResponse callClient(StsClient client, AssumeRoleRequest request) {
    return client.assumeRole(request);
}
 

@Override
protected GetFederationTokenResponse callClient(StsClient client, GetFederationTokenRequest request) {
    return client.getFederationToken(request);
}
 

@Override
protected GetSessionTokenResponse callClient(StsClient client, GetSessionTokenRequest request) {
    return client.getSessionToken(request);
}
 

@Override
protected Credentials getUpdatedCredentials(StsClient stsClient) {
    return stsClient.getFederationToken(getFederationTokenRequest).credentials();
}
 

@Override
protected Credentials getUpdatedCredentials(StsClient stsClient) {
    return stsClient.getSessionToken(getSessionTokenRequest).credentials();
}
 

@Override
protected Credentials getUpdatedCredentials(StsClient stsClient) {
    AssumeRoleRequest assumeRoleRequest = assumeRoleRequestSupplier.get();
    Validate.notNull(assumeRoleRequest, "Assume role request must not be null.");
    return stsClient.assumeRole(assumeRoleRequest).credentials();
}
 

@Override
protected Credentials getUpdatedCredentials(StsClient stsClient) {
    AssumeRoleWithWebIdentityRequest request = assumeRoleWithWebIdentityRequest.get();
    notNull(request, "AssumeRoleWithWebIdentityRequest can't be null");
    return stsClient.assumeRoleWithWebIdentity(request).credentials();
}
 

@Override
protected Credentials getUpdatedCredentials(StsClient stsClient) {
    AssumeRoleWithSamlRequest assumeRoleWithSamlRequest = assumeRoleWithSamlRequestSupplier.get();
    Validate.notNull(assumeRoleWithSamlRequest, "Assume role with saml request must not be null.");
    return stsClient.assumeRoleWithSAML(assumeRoleWithSamlRequest).credentials();
}
 

/**
 * Implemented by a child class to call STS and get a new set of credentials to be used by this provider.
 */
protected abstract Credentials getUpdatedCredentials(StsClient stsClient);
 

protected abstract ResponseT callClient(StsClient client, RequestT request);