Adds accountId support to process credentials provider (#4332)

cenedhryn · web-flow · commit ec2bfca64e1c · 2023-08-23T14:39:21.000-07:00
* Adds accountId support to process credentials provider

* Updates codegen tests from upstream branch to take account ID into account
diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/auth/scheme/query-endpoint-auth-params-auth-scheme-default-params-without-allowlist.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/auth/scheme/query-endpoint-auth-params-auth-scheme-default-params-without-allowlist.java
@@ -32,6 +32,9 @@ public final class DefaultQueryAuthSchemeParams implements QueryAuthSchemeParams
 
     private final Boolean useFIPSEndpoint;
 
+
+    private final String awsAccountId;
+
     private final String endpointId;
 
     private final Boolean defaultTrueParam;
@@ -51,6 +54,7 @@ private DefaultQueryAuthSchemeParams(Builder builder) {
         this.region = builder.region;
         this.useDualStackEndpoint = builder.useDualStackEndpoint;
         this.useFIPSEndpoint = builder.useFIPSEndpoint;
+        this.awsAccountId = builder.awsAccountId;
         this.endpointId = builder.endpointId;
         this.defaultTrueParam = Validate.paramNotNull(builder.defaultTrueParam, "defaultTrueParam");
         this.defaultStringParam = Validate.paramNotNull(builder.defaultStringParam, "defaultStringParam");
@@ -84,6 +88,11 @@ public Boolean useFipsEndpoint() {
         return useFIPSEndpoint;
     }
 
+    @Override
+    public String awsAccountId() {
+        return awsAccountId;
+    }
+
     @Override
     public String endpointId() {
         return endpointId;
@@ -134,6 +143,8 @@ private static final class Builder implements QueryAuthSchemeParams.Builder {
 
         private Boolean useFIPSEndpoint;
 
+        private String awsAccountId;
+
         private String endpointId;
 
         private Boolean defaultTrueParam = true;
@@ -156,6 +167,7 @@ private static final class Builder implements QueryAuthSchemeParams.Builder {
             this.region = params.region;
             this.useDualStackEndpoint = params.useDualStackEndpoint;
             this.useFIPSEndpoint = params.useFIPSEndpoint;
+            this.awsAccountId = params.awsAccountId;
             this.endpointId = params.endpointId;
             this.defaultTrueParam = params.defaultTrueParam;
             this.defaultStringParam = params.defaultStringParam;
@@ -189,6 +201,12 @@ public Builder useFipsEndpoint(Boolean useFIPSEndpoint) {
             return this;
         }
 
+        @Override
+        public Builder awsAccountId(String awsAccountId) {
+            this.awsAccountId = awsAccountId;
+            return this;
+        }
+
         @Override
         public Builder endpointId(String endpointId) {
             this.endpointId = endpointId;
diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/auth/scheme/query-endpoint-auth-params-auth-scheme-endpoint-provider-without-allowlist.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/auth/scheme/query-endpoint-auth-params-auth-scheme-endpoint-provider-without-allowlist.java
@@ -55,7 +55,7 @@ public static QueryAuthSchemeProvider create() {
     public List<AuthSchemeOption> resolveAuthScheme(QueryAuthSchemeParams params) {
         QueryEndpointParams endpointParameters = QueryEndpointParams.builder().region(params.region())
                 .useDualStackEndpoint(params.useDualStackEndpoint()).useFipsEndpoint(params.useFipsEndpoint())
-                .endpointId(params.endpointId()).defaultTrueParam(params.defaultTrueParam())
+                 .awsAccountId(params.awsAccountId()).endpointId(params.endpointId()).defaultTrueParam(params.defaultTrueParam())
                 .defaultStringParam(params.defaultStringParam()).deprecatedParam(params.deprecatedParam())
                 .booleanContextParam(params.booleanContextParam()).stringContextParam(params.stringContextParam())
                 .operationContextParam(params.operationContextParam()).build();
diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/auth/scheme/query-endpoint-auth-params-auth-scheme-params-without-allowlist.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/auth/scheme/query-endpoint-auth-params-auth-scheme-params-without-allowlist.java
@@ -49,6 +49,8 @@ static Builder builder() {
 
     Boolean useFipsEndpoint();
 
+    String awsAccountId();
+
     String endpointId();
 
     /**
@@ -90,6 +92,8 @@ interface Builder extends CopyableBuilder<Builder, QueryAuthSchemeParams> {
 
         Builder useFipsEndpoint(Boolean useFIPSEndpoint);
 
+        Builder awsAccountId(String awsAccountId);
+
         Builder endpointId(String endpointId);
 
         /**
diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/auth/scheme/query-endpoint-auth-params-without-allowlist-auth-scheme-interceptor.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/auth/scheme/query-endpoint-auth-params-without-allowlist-auth-scheme-interceptor.java
@@ -102,6 +102,7 @@ private QueryAuthSchemeParams authSchemeParams(SdkRequest request, ExecutionAttr
         builder.region(endpointParams.region());
         builder.useDualStackEndpoint(endpointParams.useDualStackEndpoint());
         builder.useFipsEndpoint(endpointParams.useFipsEndpoint());
+        builder.awsAccountId(endpointParams.awsAccountId());
         builder.endpointId(endpointParams.endpointId());
         builder.defaultTrueParam(endpointParams.defaultTrueParam());
         builder.defaultStringParam(endpointParams.defaultStringParam());
diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/rules/endpoint-parameters.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/rules/endpoint-parameters.java
@@ -158,6 +158,7 @@ private BuilderImpl(QueryEndpointParams builder) {
             this.region = builder.region;
             this.useDualStackEndpoint = builder.useDualStackEndpoint;
             this.useFIPSEndpoint = builder.useFIPSEndpoint;
+            this.awsAccountId = builder.awsAccountId;
             this.endpointId = builder.endpointId;
             this.defaultTrueParam = builder.defaultTrueParam;
             this.defaultStringParam = builder.defaultStringParam;
diff --git a/core/auth/src/main/java/software/amazon/awssdk/auth/credentials/ProcessCredentialsProvider.java b/core/auth/src/main/java/software/amazon/awssdk/auth/credentials/ProcessCredentialsProvider.java
@@ -129,11 +129,10 @@ private RefreshResult<AwsCredentials> refreshCredentials() {
             JsonNode credentialsJson = parseProcessOutput(processOutput);
 
             AwsCredentials credentials = credentials(credentialsJson);
-            Instant credentialExpirationTime = credentialExpirationTime(credentialsJson);
-
+            Instant expirationTime = credentials.expirationTime().orElse(Instant.MAX);
             return RefreshResult.builder(credentials)
-                                .staleTime(credentialExpirationTime)
-                                .prefetchTime(credentialExpirationTime.minusMillis(credentialRefreshThreshold.toMillis()))
+                                .staleTime(expirationTime)
+                                .prefetchTime(expirationTime.minusMillis(credentialRefreshThreshold.toMillis()))
                                 .build();
         } catch (InterruptedException e) {
             throw new IllegalStateException("Process-based credential refreshing has been interrupted.", e);
@@ -166,28 +165,33 @@ private AwsCredentials credentials(JsonNode credentialsJson) {
         String accessKeyId = getText(credentialsJson, "AccessKeyId");
         String secretAccessKey = getText(credentialsJson, "SecretAccessKey");
         String sessionToken = getText(credentialsJson, "SessionToken");
+        String accountId = getText(credentialsJson, "AccountId");
 
         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);
+            return AwsSessionCredentials.builder()
+                                        .accessKeyId(accessKeyId)
+                                        .secretAccessKey(secretAccessKey)
+                                        .sessionToken(sessionToken)
+                                        .expirationTime(credentialExpirationTime(credentialsJson))
+                                        .accountId(accountId)
+                                        .build();
         }
+        return AwsBasicCredentials.builder()
+                                  .accessKeyId(accessKeyId)
+                                  .secretAccessKey(secretAccessKey)
+                                  .accountId(accountId)
+                                  .build();
     }
 
     /**
      * Parse the process output to retrieve the expiration date and time.
      */
     private Instant credentialExpirationTime(JsonNode credentialsJson) {
         String expiration = getText(credentialsJson, "Expiration");
-
-        if (expiration != null) {
-            return DateUtils.parseIso8601Date(expiration);
-        } else {
-            return Instant.MAX;
-        }
+        return expiration != null ? DateUtils.parseIso8601Date(expiration) : null;
     }
 
     /**
diff --git a/core/auth/src/test/java/software/amazon/awssdk/auth/credentials/ProcessCredentialsProviderTest.java b/core/auth/src/test/java/software/amazon/awssdk/auth/credentials/ProcessCredentialsProviderTest.java
diff --git a/core/auth/src/test/resources/resources/process/linux-credentials-script.sh b/core/auth/src/test/resources/resources/process/linux-credentials-script.sh
diff --git a/core/auth/src/test/resources/resources/process/windows-credentials-script.bat b/core/auth/src/test/resources/resources/process/windows-credentials-script.bat
