Accept and use the new TokenIdentity interfaces

codegen/src/main/java/software/amazon/awssdk/codegen/poet/builder/AsyncClientBuilderClass.java

@@ -19,10 +19,10 @@
 import com.squareup.javapoet.MethodSpec;
 import com.squareup.javapoet.ParameterizedTypeName;
 import com.squareup.javapoet.TypeSpec;
+import com.squareup.javapoet.WildcardTypeName;
 import java.net.URI;
 import javax.lang.model.element.Modifier;
 import software.amazon.awssdk.annotations.SdkInternalApi;
-import software.amazon.awssdk.auth.token.credentials.SdkTokenProvider;
 import software.amazon.awssdk.awscore.client.config.AwsClientOption;
 import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
 import software.amazon.awssdk.codegen.poet.ClassSpec;
@@ -32,6 +32,8 @@
 import software.amazon.awssdk.codegen.utils.AuthUtils;
 import software.amazon.awssdk.core.client.config.SdkClientConfiguration;
 import software.amazon.awssdk.core.client.config.SdkClientOption;
+import software.amazon.awssdk.identity.spi.IdentityProvider;
+import software.amazon.awssdk.identity.spi.TokenIdentity;
 
 public class AsyncClientBuilderClass implements ClassSpec {
     private final IntermediateModel model;
@@ -76,7 +78,7 @@ public TypeSpec poetSpec() {
         builder.addMethod(endpointProviderMethod());
 
         if (AuthUtils.usesBearerAuth(model)) {
-            builder.addMethod(bearerTokenProviderMethod());
+            builder.addMethod(tokenProviderMethod());
         }
 
         return builder.addMethod(buildClientMethod()).build();
@@ -141,12 +143,14 @@ private MethodSpec buildClientMethod() {
                          .build();
     }
 
-    private MethodSpec bearerTokenProviderMethod() {
+    private MethodSpec tokenProviderMethod() {
+        ParameterizedTypeName tokenProviderTypeName = ParameterizedTypeName.get(ClassName.get(IdentityProvider.class),
+                                                                                WildcardTypeName.subtypeOf(TokenIdentity.class));
         return MethodSpec.methodBuilder("tokenProvider").addModifiers(Modifier.PUBLIC)
                          .addAnnotation(Override.class)
-                         .addParameter(SdkTokenProvider.class, "tokenProvider")
+                         .addParameter(tokenProviderTypeName, "tokenProvider")
                          .returns(builderClassName)
-                         .addStatement("clientConfiguration.option($T.TOKEN_PROVIDER, tokenProvider)",
+                         .addStatement("clientConfiguration.option($T.TOKEN_IDENTITY_PROVIDER, tokenProvider)",
                                        AwsClientOption.class)
                          .addStatement("return this")
                          .build();

codegen/src/main/java/software/amazon/awssdk/codegen/poet/builder/BaseClientBuilderClass.java

@@ -27,6 +27,7 @@
 import com.squareup.javapoet.TypeName;
 import com.squareup.javapoet.TypeSpec;
 import com.squareup.javapoet.TypeVariableName;
+import com.squareup.javapoet.WildcardTypeName;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
@@ -35,7 +36,6 @@
 import javax.lang.model.element.Modifier;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.auth.signer.Aws4Signer;
-import software.amazon.awssdk.auth.token.credentials.SdkTokenProvider;
 import software.amazon.awssdk.auth.token.credentials.aws.DefaultAwsTokenProvider;
 import software.amazon.awssdk.auth.token.signer.aws.BearerTokenSigner;
 import software.amazon.awssdk.awscore.client.builder.AwsDefaultClientBuilder;
@@ -59,6 +59,8 @@
 import software.amazon.awssdk.core.signer.Signer;
 import software.amazon.awssdk.http.Protocol;
 import software.amazon.awssdk.http.SdkHttpConfigurationOption;
+import software.amazon.awssdk.identity.spi.IdentityProvider;
+import software.amazon.awssdk.identity.spi.TokenIdentity;
 import software.amazon.awssdk.protocols.query.interceptor.QueryParametersToBodyInterceptor;
 import software.amazon.awssdk.utils.AttributeMap;
 import software.amazon.awssdk.utils.CollectionUtils;
@@ -205,7 +207,7 @@ private MethodSpec mergeServiceDefaultsMethod() {
         }
 
         if (AuthUtils.usesBearerAuth(model)) {
-            builder.addCode(".option($T.TOKEN_PROVIDER, defaultTokenProvider())\n", AwsClientOption.class);
+            builder.addCode(".option($T.TOKEN_IDENTITY_PROVIDER, defaultTokenProvider())\n", AwsClientOption.class);
             builder.addCode(".option($T.TOKEN_SIGNER, defaultTokenSigner())", SdkAdvancedClientOption.class);
         }
 
@@ -595,7 +597,8 @@ private MethodSpec clientContextParamSetter(String name, ClientContextParam para
 
     private MethodSpec defaultBearerTokenProviderMethod() {
         return MethodSpec.methodBuilder("defaultTokenProvider")
-                         .returns(SdkTokenProvider.class)
+                         .returns(ParameterizedTypeName.get(ClassName.get(IdentityProvider.class),
+                                                            WildcardTypeName.subtypeOf(TokenIdentity.class)))
                          .addModifiers(PRIVATE)
                          .addStatement("return $T.create()", DefaultAwsTokenProvider.class)
                          .build();
@@ -638,11 +641,10 @@ private MethodSpec validateClientOptionsMethod() {
                                  SdkAdvancedClientOption.class,
                                  "The 'overrideConfiguration.advancedOption[TOKEN_SIGNER]' "
                                  + "must be configured in the client builder.");
-            builder.addStatement("$T.notNull(c.option($T.TOKEN_PROVIDER), $S)",
+            builder.addStatement("$T.notNull(c.option($T.TOKEN_IDENTITY_PROVIDER), $S)",
                                  Validate.class,
                                  AwsClientOption.class,
-                                 "The 'overrideConfiguration.advancedOption[TOKEN_PROVIDER]' "
-                                 + "must be configured in the client builder.");
+                                 "The 'tokenProvider' must be configured in the client builder.");
         }
 
         return builder.build();

codegen/src/main/java/software/amazon/awssdk/codegen/poet/builder/BaseClientBuilderInterface.java

@@ -22,6 +22,7 @@
 import com.squareup.javapoet.TypeName;
 import com.squareup.javapoet.TypeSpec;
 import com.squareup.javapoet.TypeVariableName;
+import com.squareup.javapoet.WildcardTypeName;
 import java.util.function.Consumer;
 import javax.lang.model.element.Modifier;
 import software.amazon.awssdk.auth.token.credentials.SdkTokenProvider;
@@ -36,9 +37,14 @@
 import software.amazon.awssdk.codegen.poet.rules.EndpointRulesSpecUtils;
 import software.amazon.awssdk.codegen.utils.AuthUtils;
 import software.amazon.awssdk.core.client.config.SdkAdvancedClientOption;
+import software.amazon.awssdk.identity.spi.IdentityProvider;
+import software.amazon.awssdk.identity.spi.TokenIdentity;
 import software.amazon.awssdk.utils.internal.CodegenNamingUtils;
 
 public class BaseClientBuilderInterface implements ClassSpec {
+    private static final ParameterizedTypeName TOKEN_IDENTITY_PROVIDER_TYPE_NAME =
+        ParameterizedTypeName.get(ClassName.get(IdentityProvider.class), WildcardTypeName.subtypeOf(TokenIdentity.class));
+
     private final IntermediateModel model;
     private final String basePackage;
     private final ClassName builderInterfaceName;
@@ -82,6 +88,7 @@ public TypeSpec poetSpec() {
 
         if (generateTokenProviderMethod()) {
             builder.addMethod(tokenProviderMethod());
+            builder.addMethod(tokenIdentityProviderMethod());
         }
 
         return builder.build();
@@ -168,7 +175,7 @@ private boolean generateTokenProviderMethod() {
 
     private MethodSpec tokenProviderMethod() {
         return MethodSpec.methodBuilder("tokenProvider")
-                         .addModifiers(Modifier.PUBLIC, Modifier.ABSTRACT)
+                         .addModifiers(Modifier.PUBLIC, Modifier.DEFAULT)
                          .returns(TypeVariableName.get("B"))
                          .addParameter(SdkTokenProvider.class, "tokenProvider")
                          .addJavadoc("Set the token provider to use for bearer token authorization. This is optional, if none "
@@ -182,6 +189,27 @@ private MethodSpec tokenProviderMethod() {
                                      DefaultAwsTokenProvider.class,
                                      SdkAdvancedClientOption.class,
                                      BearerTokenSigner.class)
+                         .addStatement("return tokenProvider(($T) tokenProvider)", TOKEN_IDENTITY_PROVIDER_TYPE_NAME)
+                         .build();
+    }
+
+    private MethodSpec tokenIdentityProviderMethod() {
+        return MethodSpec.methodBuilder("tokenProvider")
+                         .addModifiers(Modifier.PUBLIC, Modifier.DEFAULT)
+                         .returns(TypeVariableName.get("B"))
+                         .addParameter(TOKEN_IDENTITY_PROVIDER_TYPE_NAME, "tokenProvider")
+                         .addJavadoc("Set the token provider to use for bearer token authorization. This is optional, if none "
+                                     + "is provided, the SDK will use {@link $T}.\n"
+                                     + "<p>\n"
+                                     + "If the service, or any of its operations require Bearer Token Authorization, then the "
+                                     + "SDK will default to this token provider to retrieve the token to use for authorization.\n"
+                                     + "<p>\n"
+                                     + "This provider works in conjunction with the {@code $T.TOKEN_SIGNER} set on the client. "
+                                     + "By default it is {@link $T}.",
+                                     DefaultAwsTokenProvider.class,
+                                     SdkAdvancedClientOption.class,
+                                     BearerTokenSigner.class)
+                         .addStatement("throw new $T()", UnsupportedOperationException.class)
                          .build();
     }
 

codegen/src/main/java/software/amazon/awssdk/codegen/poet/builder/SyncClientBuilderClass.java

@@ -19,10 +19,10 @@
 import com.squareup.javapoet.MethodSpec;
 import com.squareup.javapoet.ParameterizedTypeName;
 import com.squareup.javapoet.TypeSpec;
+import com.squareup.javapoet.WildcardTypeName;
 import java.net.URI;
 import javax.lang.model.element.Modifier;
 import software.amazon.awssdk.annotations.SdkInternalApi;
-import software.amazon.awssdk.auth.token.credentials.SdkTokenProvider;
 import software.amazon.awssdk.awscore.client.config.AwsClientOption;
 import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
 import software.amazon.awssdk.codegen.poet.ClassSpec;
@@ -32,6 +32,8 @@
 import software.amazon.awssdk.codegen.utils.AuthUtils;
 import software.amazon.awssdk.core.client.config.SdkClientConfiguration;
 import software.amazon.awssdk.core.client.config.SdkClientOption;
+import software.amazon.awssdk.identity.spi.IdentityProvider;
+import software.amazon.awssdk.identity.spi.TokenIdentity;
 
 public class SyncClientBuilderClass implements ClassSpec {
     private final IntermediateModel model;
@@ -142,11 +144,13 @@ private MethodSpec buildClientMethod() {
     }
 
     private MethodSpec tokenProviderMethodImpl() {
+        ParameterizedTypeName tokenProviderTypeName = ParameterizedTypeName.get(ClassName.get(IdentityProvider.class),
+                                                                                WildcardTypeName.subtypeOf(TokenIdentity.class));
         return MethodSpec.methodBuilder("tokenProvider").addModifiers(Modifier.PUBLIC)
                          .addAnnotation(Override.class)
-                         .addParameter(SdkTokenProvider.class, "tokenProvider")
+                         .addParameter(tokenProviderTypeName, "tokenProvider")
                          .returns(builderClassName)
-                         .addStatement("clientConfiguration.option($T.TOKEN_PROVIDER, tokenProvider)",
+                         .addStatement("clientConfiguration.option($T.TOKEN_IDENTITY_PROVIDER, tokenProvider)",
                                        AwsClientOption.class)
                          .addStatement("return this")
                          .build();

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/builder/test-async-client-builder-class.java

@@ -3,10 +3,11 @@
 import java.net.URI;
 import software.amazon.awssdk.annotations.Generated;
 import software.amazon.awssdk.annotations.SdkInternalApi;
-import software.amazon.awssdk.auth.token.credentials.SdkTokenProvider;
 import software.amazon.awssdk.awscore.client.config.AwsClientOption;
 import software.amazon.awssdk.core.client.config.SdkClientConfiguration;
 import software.amazon.awssdk.core.client.config.SdkClientOption;
+import software.amazon.awssdk.identity.spi.IdentityProvider;
+import software.amazon.awssdk.identity.spi.TokenIdentity;
 import software.amazon.awssdk.services.json.endpoints.JsonEndpointProvider;
 
 /**
@@ -23,8 +24,8 @@ public DefaultJsonAsyncClientBuilder endpointProvider(JsonEndpointProvider endpo
     }
 
     @Override
-    public DefaultJsonAsyncClientBuilder tokenProvider(SdkTokenProvider tokenProvider) {
-        clientConfiguration.option(AwsClientOption.TOKEN_PROVIDER, tokenProvider);
+    public DefaultJsonAsyncClientBuilder tokenProvider(IdentityProvider<? extends TokenIdentity> tokenProvider) {
+        clientConfiguration.option(AwsClientOption.TOKEN_IDENTITY_PROVIDER, tokenProvider);
         return this;
     }
 

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/builder/test-bearer-auth-client-builder-class.java

@@ -4,7 +4,6 @@
 import java.util.List;
 import software.amazon.awssdk.annotations.Generated;
 import software.amazon.awssdk.annotations.SdkInternalApi;
-import software.amazon.awssdk.auth.token.credentials.SdkTokenProvider;
 import software.amazon.awssdk.auth.token.credentials.aws.DefaultAwsTokenProvider;
 import software.amazon.awssdk.auth.token.signer.aws.BearerTokenSigner;
 import software.amazon.awssdk.awscore.client.builder.AwsDefaultClientBuilder;
@@ -15,6 +14,8 @@
 import software.amazon.awssdk.core.interceptor.ClasspathInterceptorChainFactory;
 import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
 import software.amazon.awssdk.core.signer.Signer;
+import software.amazon.awssdk.identity.spi.IdentityProvider;
+import software.amazon.awssdk.identity.spi.TokenIdentity;
 import software.amazon.awssdk.services.json.endpoints.JsonEndpointProvider;
 import software.amazon.awssdk.services.json.endpoints.internal.JsonEndpointAuthSchemeInterceptor;
 import software.amazon.awssdk.services.json.endpoints.internal.JsonRequestSetEndpointInterceptor;
@@ -42,7 +43,7 @@ protected final String serviceName() {
     protected final SdkClientConfiguration mergeServiceDefaults(SdkClientConfiguration config) {
         return config.merge(c -> c.option(SdkClientOption.ENDPOINT_PROVIDER, defaultEndpointProvider())
                                   .option(SdkClientOption.CRC32_FROM_COMPRESSED_DATA_ENABLED, false)
-                                  .option(AwsClientOption.TOKEN_PROVIDER, defaultTokenProvider())
+                                  .option(AwsClientOption.TOKEN_IDENTITY_PROVIDER, defaultTokenProvider())
                                   .option(SdkAdvancedClientOption.TOKEN_SIGNER, defaultTokenSigner()));
     }
 
@@ -71,7 +72,7 @@ private JsonEndpointProvider defaultEndpointProvider() {
         return JsonEndpointProvider.defaultProvider();
     }
 
-    private SdkTokenProvider defaultTokenProvider() {
+    private IdentityProvider<? extends TokenIdentity> defaultTokenProvider() {
         return DefaultAwsTokenProvider.create();
     }
 
@@ -82,7 +83,7 @@ private Signer defaultTokenSigner() {
     protected static void validateClientOptions(SdkClientConfiguration c) {
         Validate.notNull(c.option(SdkAdvancedClientOption.TOKEN_SIGNER),
                          "The 'overrideConfiguration.advancedOption[TOKEN_SIGNER]' must be configured in the client builder.");
-        Validate.notNull(c.option(AwsClientOption.TOKEN_PROVIDER),
-                         "The 'overrideConfiguration.advancedOption[TOKEN_PROVIDER]' must be configured in the client builder.");
+        Validate.notNull(c.option(AwsClientOption.TOKEN_IDENTITY_PROVIDER),
+                         "The 'tokenProvider' must be configured in the client builder.");
     }
 }

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/builder/test-bearer-auth-client-builder-interface.java

@@ -3,6 +3,8 @@
 import software.amazon.awssdk.annotations.Generated;
 import software.amazon.awssdk.auth.token.credentials.SdkTokenProvider;
 import software.amazon.awssdk.awscore.client.builder.AwsClientBuilder;
+import software.amazon.awssdk.identity.spi.IdentityProvider;
+import software.amazon.awssdk.identity.spi.TokenIdentity;
 import software.amazon.awssdk.services.json.endpoints.JsonEndpointProvider;
 
 /**
@@ -30,5 +32,22 @@ default B endpointProvider(JsonEndpointProvider endpointProvider) {
      * {@code software.amazon.awssdk.core.client.config.SdkAdvancedClientOption.TOKEN_SIGNER} set on the client. By
      * default it is {@link software.amazon.awssdk.auth.token.signer.aws.BearerTokenSigner}.
      */
-    B tokenProvider(SdkTokenProvider tokenProvider);
+    default B tokenProvider(SdkTokenProvider tokenProvider) {
+        return tokenProvider((IdentityProvider<? extends TokenIdentity>) tokenProvider);
+    }
+
+    /**
+     * Set the token provider to use for bearer token authorization. This is optional, if none is provided, the SDK will
+     * use {@link software.amazon.awssdk.auth.token.credentials.aws.DefaultAwsTokenProvider}.
+     * <p>
+     * If the service, or any of its operations require Bearer Token Authorization, then the SDK will default to this
+     * token provider to retrieve the token to use for authorization.
+     * <p>
+     * This provider works in conjunction with the
+     * {@code software.amazon.awssdk.core.client.config.SdkAdvancedClientOption.TOKEN_SIGNER} set on the client. By
+     * default it is {@link software.amazon.awssdk.auth.token.signer.aws.BearerTokenSigner}.
+     */
+    default B tokenProvider(IdentityProvider<? extends TokenIdentity> tokenProvider) {
+        throw new UnsupportedOperationException();
+    }
 }