Adds accountId to STS credentials providers (#4040)

Adds accountId to STS credentials providers

Author: cenedhryn

core/auth/src/main/java/software/amazon/awssdk/auth/credentials/AwsBasicCredentials.java

@@ -44,6 +44,8 @@ public final class AwsBasicCredentials implements AwsCredentials,
                                                   ToCopyableBuilder<AwsBasicCredentials.Builder, AwsBasicCredentials> {
     /**
      * A set of AWS credentials without an access key or secret access key, indicating that anonymous access should be used.
+     * <p>
+     * This should be accessed via {@link AnonymousCredentialsProvider#resolveCredentials()}.
      */
     // TODO(sra-identity-and-auth): Check if this static member can be removed after cleanup
     @SdkInternalApi
@@ -53,12 +55,14 @@ public final class AwsBasicCredentials implements AwsCredentials,
     private final String secretAccessKey;
     private final boolean validateCredentials;
     private final String providerName;
+    private final String accountId;
 
     private AwsBasicCredentials(Builder builder) {
         this.accessKeyId = trimToNull(builder.accessKeyId);
         this.secretAccessKey = trimToNull(builder.secretAccessKey);
         this.validateCredentials = builder.validateCredentials;
         this.providerName = builder.providerName;
+        this.accountId = builder.accountId;
 
         if (builder.validateCredentials) {
             Validate.notNull(this.accessKeyId, "Access key ID cannot be blank.");
@@ -77,6 +81,7 @@ protected AwsBasicCredentials(String accessKeyId, String secretAccessKey) {
         this.secretAccessKey = trimToNull(secretAccessKey);
         this.validateCredentials = false;
         this.providerName = null;
+        this.accountId = null;
     }
 
     public static Builder builder() {
@@ -88,7 +93,7 @@ public static Builder builder() {
      *
      * @param accessKeyId The AWS access key, used to identify the user interacting with AWS.
      * @param secretAccessKey The AWS secret access key, used to authenticate the user interacting with AWS.
-     * */
+     */
     public static AwsBasicCredentials create(String accessKeyId, String secretAccessKey) {
         return builder().accessKeyId(accessKeyId)
                         .secretAccessKey(secretAccessKey)
@@ -119,11 +124,20 @@ public Optional<String> providerName() {
         return Optional.ofNullable(providerName);
     }
 
+    /**
+     * Retrieve the AWS account id associated with this credentials identity, if found.
+     */
+    @Override
+    public Optional<String> accountId() {
+        return Optional.ofNullable(accountId);
+    }
+
     @Override
     public String toString() {
         return ToString.builder("AwsCredentials")
                        .add("accessKeyId", accessKeyId)
                        .add("providerName", providerName)
+                       .add("accountId", accountId)
                        .build();
     }
 
@@ -137,21 +151,24 @@ public boolean equals(Object o) {
         }
         AwsBasicCredentials that = (AwsBasicCredentials) o;
         return Objects.equals(accessKeyId, that.accessKeyId) &&
-               Objects.equals(secretAccessKey, that.secretAccessKey);
+               Objects.equals(secretAccessKey, that.secretAccessKey) &&
+               Objects.equals(accountId, that.accountId().orElse(null));
     }
 
     @Override
     public int hashCode() {
         int hashCode = 1;
         hashCode = 31 * hashCode + Objects.hashCode(accessKeyId());
         hashCode = 31 * hashCode + Objects.hashCode(secretAccessKey());
+        hashCode = 31 * hashCode + Objects.hashCode(accountId);
         return hashCode;
     }
 
     @Override
     public Builder toBuilder() {
         return builder().accessKeyId(accessKeyId)
                         .secretAccessKey(secretAccessKey)
+                        .accountId(accountId)
                         .validateCredentials(validateCredentials)
                         .providerName(providerName);
     }
@@ -169,6 +186,7 @@ public static final class Builder implements CopyableBuilder<AwsBasicCredentials
         private String accessKeyId;
         private String secretAccessKey;
         private String providerName;
+        private String accountId;
         private boolean validateCredentials = true;
 
         private Builder() {
@@ -190,6 +208,14 @@ public Builder secretAccessKey(String secretAccessKey) {
             return this;
         }
 
+        /**
+         * The AWS account id associated with this credentials identity.
+         */
+        public Builder accountId(String accountId) {
+            this.accountId = accountId;
+            return this;
+        }
+
         /**
          * The name of the identity provider that created this credential identity.
          */

core/auth/src/main/java/software/amazon/awssdk/auth/credentials/AwsSessionCredentials.java

@@ -40,13 +40,15 @@ public final class AwsSessionCredentials implements AwsCredentials, AwsSessionCr
     private final String secretAccessKey;
     private final String sessionToken;
 
+    private final String accountId;
     private final Instant expirationTime;
     private final String providerName;
 
     private AwsSessionCredentials(Builder builder) {
         this.accessKeyId = Validate.paramNotNull(builder.accessKeyId, "accessKey");
         this.secretAccessKey = Validate.paramNotNull(builder.secretAccessKey, "secretKey");
         this.sessionToken = Validate.paramNotNull(builder.sessionToken, "sessionToken");
+        this.accountId = builder.accountId;
         this.expirationTime = builder.expirationTime;
         this.providerName = builder.providerName;
     }
@@ -111,11 +113,17 @@ public Optional<String> providerName() {
         return Optional.ofNullable(providerName);
     }
 
+    @Override
+    public Optional<String> accountId() {
+        return Optional.ofNullable(accountId);
+    }
+
     @Override
     public String toString() {
         return ToString.builder("AwsSessionCredentials")
                        .add("accessKeyId", accessKeyId())
                        .add("providerName", providerName)
+                       .add("accountId", accountId)
                        .build();
     }
 
@@ -132,6 +140,7 @@ public boolean equals(Object o) {
         return Objects.equals(accessKeyId, that.accessKeyId) &&
                Objects.equals(secretAccessKey, that.secretAccessKey) &&
                Objects.equals(sessionToken, that.sessionToken) &&
+               Objects.equals(accountId, that.accountId().orElse(null)) &&
                Objects.equals(expirationTime, that.expirationTime().orElse(null));
     }
 
@@ -141,6 +150,7 @@ public int hashCode() {
         hashCode = 31 * hashCode + Objects.hashCode(accessKeyId());
         hashCode = 31 * hashCode + Objects.hashCode(secretAccessKey());
         hashCode = 31 * hashCode + Objects.hashCode(sessionToken());
+        hashCode = 31 * hashCode + Objects.hashCode(accountId);
         hashCode = 31 * hashCode + Objects.hashCode(expirationTime);
         return hashCode;
     }
@@ -150,6 +160,7 @@ public Builder toBuilder() {
         return builder().accessKeyId(accessKeyId)
                         .secretAccessKey(secretAccessKey)
                         .sessionToken(sessionToken)
+                        .accountId(accountId)
                         .expirationTime(expirationTime)
                         .providerName(providerName);
     }
@@ -167,6 +178,7 @@ public static final class Builder implements CopyableBuilder<AwsSessionCredentia
         private String accessKeyId;
         private String secretAccessKey;
         private String sessionToken;
+        private String accountId;
         private Instant expirationTime;
         private String providerName;
 
@@ -195,6 +207,16 @@ public Builder sessionToken(String sessionToken) {
             return this;
         }
 
+        /**
+         * The AWS accountId
+         * @param accountId
+         * @return
+         */
+        public Builder accountId(String accountId) {
+            this.accountId = accountId;
+            return this;
+        }
+
         /**
          * The time after which this identity will no longer be valid. If this is empty,
          * an expiration time is not known (but the identity may still expire at some

core/auth/src/main/java/software/amazon/awssdk/auth/credentials/CredentialUtils.java

@@ -71,15 +71,21 @@ public static AwsCredentials toCredentials(AwsCredentialsIdentity awsCredentials
         // identity-spi defines 2 known types - AwsCredentialsIdentity and a sub-type AwsSessionCredentialsIdentity
         if (awsCredentialsIdentity instanceof AwsSessionCredentialsIdentity) {
             AwsSessionCredentialsIdentity awsSessionCredentialsIdentity = (AwsSessionCredentialsIdentity) awsCredentialsIdentity;
-            return AwsSessionCredentials.create(awsSessionCredentialsIdentity.accessKeyId(),
-                                                awsSessionCredentialsIdentity.secretAccessKey(),
-                                                awsSessionCredentialsIdentity.sessionToken());
+            return AwsSessionCredentials.builder()
+                                        .accessKeyId(awsSessionCredentialsIdentity.accessKeyId())
+                                        .secretAccessKey(awsSessionCredentialsIdentity.secretAccessKey())
+                                        .sessionToken(awsSessionCredentialsIdentity.sessionToken())
+                                        .accountId(awsSessionCredentialsIdentity.accountId().orElse(null))
+                                        .build();
         }
         if (isAnonymous(awsCredentialsIdentity)) {
             return AwsBasicCredentials.ANONYMOUS_CREDENTIALS;
         }
-        return AwsBasicCredentials.create(awsCredentialsIdentity.accessKeyId(),
-                                          awsCredentialsIdentity.secretAccessKey());
+        return AwsBasicCredentials.builder()
+                                  .accessKeyId(awsCredentialsIdentity.accessKeyId())
+                                  .secretAccessKey(awsCredentialsIdentity.secretAccessKey())
+                                  .accountId(awsCredentialsIdentity.accountId().orElse(null))
+                                  .build();
     }
 
     /**

core/auth/src/test/java/software/amazon/awssdk/auth/credentials/internal/AwsSessionCredentialsTest.java

@@ -16,7 +16,9 @@
 package software.amazon.awssdk.auth.credentials.internal;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import nl.jqno.equalsverifier.EqualsVerifier;
 import org.junit.jupiter.api.Test;
@@ -27,6 +29,7 @@ class AwsSessionCredentialsTest {
     private static final String ACCESS_KEY_ID = "accessKeyId";
     private static final String SECRET_ACCESS_KEY = "secretAccessKey";
     private static final String SESSION_TOKEN = "sessionToken";
+    private static final String ACCOUNT_ID = "accountId";
     private static final String PROVIDER_NAME = "StaticCredentialsProvider";
 
     @Test
@@ -65,6 +68,7 @@ void create_isSuccessful() {
         assertEquals(ACCESS_KEY_ID, identity.accessKeyId());
         assertEquals(SECRET_ACCESS_KEY, identity.secretAccessKey());
         assertEquals(SESSION_TOKEN, identity.sessionToken());
+        assertFalse(identity.accountId().isPresent());
     }
 
     @Test
@@ -73,10 +77,13 @@ void build_isSuccessful() {
                                                               .accessKeyId(ACCESS_KEY_ID)
                                                               .secretAccessKey(SECRET_ACCESS_KEY)
                                                               .sessionToken(SESSION_TOKEN)
+                                                              .accountId(ACCOUNT_ID)
                                                               .build();
         assertEquals(ACCESS_KEY_ID, identity.accessKeyId());
         assertEquals(SECRET_ACCESS_KEY, identity.secretAccessKey());
         assertEquals(SESSION_TOKEN, identity.sessionToken());
+        assertTrue(identity.accountId().isPresent());
+        assertEquals(ACCOUNT_ID, identity.accountId().get());
     }
 
     @Test
@@ -85,11 +92,14 @@ void copy_isSuccessful() {
                                                               .accessKeyId(ACCESS_KEY_ID)
                                                               .secretAccessKey(SECRET_ACCESS_KEY)
                                                               .sessionToken(SESSION_TOKEN)
+                                                              .accountId(ACCOUNT_ID)
                                                               .build();
         AwsSessionCredentials copy = identity.copy(c -> c.providerName(PROVIDER_NAME));
         assertEquals(ACCESS_KEY_ID, copy.accessKeyId());
         assertEquals(SECRET_ACCESS_KEY, copy.secretAccessKey());
         assertEquals(SESSION_TOKEN, copy.sessionToken());
+        assertTrue(identity.accountId().isPresent());
+        assertEquals(ACCOUNT_ID, identity.accountId().get());
         assertEquals(PROVIDER_NAME, copy.providerName().get());
     }
 }

core/identity-spi/src/main/java/software/amazon/awssdk/identity/spi/internal/DefaultAwsSessionCredentialsIdentity.java

@@ -94,7 +94,7 @@ public boolean equals(Object o) {
         return Objects.equals(accessKeyId, that.accessKeyId()) &&
                Objects.equals(secretAccessKey, that.secretAccessKey()) &&
                Objects.equals(sessionToken, that.sessionToken()) &&
-                Objects.equals(accountId, that.accountId().orElse(null));
+               Objects.equals(accountId, that.accountId().orElse(null));
     }
 
     @Override

services/sts/src/it/java/software/amazon/awssdk/services/sts/AssumeRoleIntegrationTest.java

@@ -178,6 +178,7 @@ public void profileCredentialsProviderCanAssumeRoles() throws InterruptedExcepti
 
         assertThat(awsCredentials.accessKeyId()).isNotBlank();
         assertThat(awsCredentials.secretAccessKey()).isNotBlank();
+        assertThat(awsCredentials.accountId()).isPresent();
         ((SdkAutoCloseable) awsCredentialsProvider).close();
     }
 
@@ -210,6 +211,7 @@ public void profileCredentialProviderCanAssumeRolesWithEnvironmentCredentialSour
 
             assertThat(awsCredentials.accessKeyId()).isNotBlank();
             assertThat(awsCredentials.secretAccessKey()).isNotBlank();
+            assertThat(awsCredentials.accountId()).isPresent();
             ((SdkAutoCloseable) awsCredentialsProvider).close();
         });
     }
@@ -247,6 +249,7 @@ public void profileCredentialProviderWithEnvironmentCredentialSourceAndSystemPro
 
                 assertThat(awsCredentials.accessKeyId()).isNotBlank();
                 assertThat(awsCredentials.secretAccessKey()).isNotBlank();
+                assertThat(awsCredentials.accountId()).isPresent();
                 ((SdkAutoCloseable) awsCredentialsProvider).close();
             });
         } finally {

services/sts/src/main/java/software/amazon/awssdk/services/sts/auth/SessionCredentialsHolder.java

@@ -0,0 +1,74 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.services.sts.auth;
+
+import java.time.Instant;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.annotations.ThreadSafe;
+import software.amazon.awssdk.auth.credentials.AwsSessionCredentials;
+import software.amazon.awssdk.services.sts.model.Credentials;
+
+/**
+ * Holder class used to atomically store a session with its expiration time.
+ */
+@SdkInternalApi
+@ThreadSafe
+final class SessionCredentialsHolder {
+    private final AwsSessionCredentials sessionCredentials;
+    private final Instant sessionCredentialsExpiration;
+
+    SessionCredentialsHolder(Builder b) {
+        Credentials credentials = b.credentials;
+        this.sessionCredentials = AwsSessionCredentials.builder()
+                                                       .accessKeyId(credentials.accessKeyId())
+                                                       .secretAccessKey(credentials.secretAccessKey())
+                                                       .sessionToken(credentials.sessionToken())
+                                                       .accountId(b.accountId)
+                                                       .build();
+        this.sessionCredentialsExpiration = credentials.expiration();
+    }
+
+    public static Builder builder() {
+        return new Builder();
+    }
+
+    public AwsSessionCredentials getSessionCredentials() {
+        return sessionCredentials;
+    }
+
+    public Instant getSessionCredentialsExpiration() {
+        return sessionCredentialsExpiration;
+    }
+
+    public static class Builder {
+        private Credentials credentials;
+        private String accountId;
+
+        public Builder credentials(Credentials credentials) {
+            this.credentials = credentials;
+            return this;
+        }
+
+        public Builder accountId(String accountId) {
+            this.accountId = accountId;
+            return this;
+        }
+
+        public SessionCredentialsHolder build() {
+            return new SessionCredentialsHolder(this);
+        }
+    }
+}