Endpoint based auth scheme resolver should honor endpoint overrides

codegen/src/main/java/software/amazon/awssdk/codegen/emitters/tasks/AuthSchemeGeneratorTasks.java

@@ -25,6 +25,7 @@
 import software.amazon.awssdk.codegen.poet.auth.scheme.AuthSchemeProviderSpec;
 import software.amazon.awssdk.codegen.poet.auth.scheme.AuthSchemeSpecUtils;
 import software.amazon.awssdk.codegen.poet.auth.scheme.DefaultAuthSchemeParamsSpec;
+import software.amazon.awssdk.codegen.poet.auth.scheme.EndpointAwareAuthSchemeParamsSpec;
 import software.amazon.awssdk.codegen.poet.auth.scheme.EndpointBasedAuthSchemeProviderSpec;
 import software.amazon.awssdk.codegen.poet.auth.scheme.ModelBasedAuthSchemeProviderSpec;
 
@@ -47,6 +48,7 @@ protected List<GeneratorTask> createTasks() {
         tasks.add(generateAuthSchemeInterceptor());
         if (authSchemeSpecUtils.useEndpointBasedAuthProvider()) {
             tasks.add(generateEndpointBasedProvider());
+            tasks.add(generateEndpointAwareAuthSchemeParams());
         }
         return tasks;
     }
@@ -72,6 +74,11 @@ private GeneratorTask generateEndpointBasedProvider() {
                                      new EndpointBasedAuthSchemeProviderSpec(model));
     }
 
+    private GeneratorTask generateEndpointAwareAuthSchemeParams() {
+        return new PoetGeneratorTask(authSchemeDir(), model.getFileHeader(), new EndpointAwareAuthSchemeParamsSpec(model));
+
+    }
+
     private GeneratorTask generateAuthSchemeInterceptor() {
         return new PoetGeneratorTask(authSchemeInternalDir(), model.getFileHeader(), new AuthSchemeInterceptorSpec(model));
     }

codegen/src/main/java/software/amazon/awssdk/codegen/poet/auth/scheme/AuthSchemeInterceptorSpec.java

@@ -49,6 +49,7 @@
 import software.amazon.awssdk.core.interceptor.SdkInternalExecutionAttribute;
 import software.amazon.awssdk.core.internal.util.MetricUtils;
 import software.amazon.awssdk.core.metrics.CoreMetric;
+import software.amazon.awssdk.endpoints.EndpointProvider;
 import software.amazon.awssdk.http.auth.spi.scheme.AuthScheme;
 import software.amazon.awssdk.http.auth.spi.scheme.AuthSchemeOption;
 import software.amazon.awssdk.identity.spi.AwsCredentialsIdentity;
@@ -185,6 +186,17 @@ private MethodSpec generateAuthSchemeParams() {
                                  AwsExecutionAttribute.class);
             builder.addStatement("builder.region(region)");
         }
+        ClassName paramsBuilderClass = authSchemeSpecUtils.parametersEndpointAwareDefaultImplName().nestedClass("Builder");
+        builder.beginControlFlow("if (builder instanceof $T)",
+                                 paramsBuilderClass);
+        ClassName endpointProviderClass = endpointRulesSpecUtils.providerInterfaceName();
+        builder.addStatement("$T endpointProvider = executionAttributes.getAttribute($T.ENDPOINT_PROVIDER)",
+                             EndpointProvider.class,
+                             SdkInternalExecutionAttribute.class);
+        builder.beginControlFlow("if (endpointProvider instanceof $T)", endpointProviderClass);
+        builder.addStatement("(($T)builder).endpointProvider(($T)endpointProvider)", paramsBuilderClass, endpointProviderClass);
+        builder.endControlFlow();
+        builder.endControlFlow();
         builder.addStatement("return builder.build()");
         return builder.build();
     }

codegen/src/main/java/software/amazon/awssdk/codegen/poet/auth/scheme/AuthSchemeSpecUtils.java

@@ -79,6 +79,10 @@ public ClassName parametersInterfaceName() {
         return ClassName.get(basePackage(), intermediateModel.getMetadata().getServiceName() + "AuthSchemeParams");
     }
 
+    public ClassName parametersEndpointAwareDefaultImplName() {
+        return ClassName.get(internalPackage(), intermediateModel.getMetadata().getServiceName() + "EndpointResolverAware");
+    }
+
     public ClassName parametersInterfaceBuilderInterfaceName() {
         return parametersInterfaceName().nestedClass("Builder");
     }

codegen/src/main/java/software/amazon/awssdk/codegen/poet/auth/scheme/DefaultAuthSchemeParamsSpec.java

@@ -59,6 +59,10 @@ public TypeSpec poetSpec() {
                                       .addMethod(builderMethod())
                                       .addType(builderImplSpec());
 
+        if (authSchemeSpecUtils.useEndpointBasedAuthProvider()) {
+            b.addSuperinterface(authSchemeSpecUtils.parametersEndpointAwareDefaultImplName());
+        }
+
         addFieldsAndAccessors(b);
         addToBuilder(b);
         return b.build();
@@ -89,6 +93,7 @@ private MethodSpec constructor() {
 
                 }
             });
+            b.addStatement("this.endpointProvider = builder.endpointProvider");
         }
 
         return b.build();
@@ -112,6 +117,10 @@ private TypeSpec builderImplSpec() {
                                      .addModifiers(Modifier.PRIVATE, Modifier.STATIC, Modifier.FINAL)
                                      .addSuperinterface(authSchemeSpecUtils.parametersInterfaceBuilderInterfaceName());
 
+        if (authSchemeSpecUtils.useEndpointBasedAuthProvider()) {
+            b.addSuperinterface(authSchemeSpecUtils.parametersEndpointAwareDefaultImplName().nestedClass("Builder"));
+        }
+
         addBuilderConstructors(b);
         addBuilderFieldsAndSetter(b);
 
@@ -142,6 +151,7 @@ private void addBuilderConstructors(TypeSpec.Builder b) {
                     builderFromInstance.addStatement("this.$1N = params.$1N", endpointRulesSpecUtils.variableName(name));
                 }
             });
+            builderFromInstance.addStatement("this.endpointProvider = params.endpointProvider");
         }
         b.addMethod(builderFromInstance.build());
     }
@@ -181,6 +191,17 @@ private void addFieldsAndAccessors(TypeSpec.Builder b) {
                                                       .build());
                 }
             });
+            ClassName endpointProvider = endpointRulesSpecUtils.providerInterfaceName();
+            b.addField(FieldSpec.builder(endpointProvider, "endpointProvider")
+                                .addModifiers(Modifier.PRIVATE, Modifier.FINAL)
+                                .build());
+            b.addMethod(MethodSpec.methodBuilder("endpointProvider")
+                                  .addModifiers(Modifier.PUBLIC)
+                                  .addAnnotation(Override.class)
+                                  .returns(endpointProvider)
+                                  .addStatement("return endpointProvider")
+                                  .build());
+
         }
     }
 
@@ -213,6 +234,10 @@ private void addBuilderFieldsAndSetter(TypeSpec.Builder b) {
                     b.addMethod(endpointRulesSpecUtils.parameterBuilderSetterMethod(className(), name, model));
                 }
             });
+            b.addField(FieldSpec.builder(endpointRulesSpecUtils.providerInterfaceName(), "endpointProvider")
+                                .addModifiers(Modifier.PRIVATE)
+                                .build());
+            b.addMethod(builderSetterMethod("endpointProvider", endpointRulesSpecUtils.providerInterfaceName()));
         }
     }
 

codegen/src/main/java/software/amazon/awssdk/codegen/poet/auth/scheme/EndpointAwareAuthSchemeParamsSpec.java

@@ -0,0 +1,75 @@
+/*
+ * 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.codegen.poet.auth.scheme;
+
+import com.squareup.javapoet.ClassName;
+import com.squareup.javapoet.MethodSpec;
+import com.squareup.javapoet.TypeSpec;
+import javax.lang.model.element.Modifier;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
+import software.amazon.awssdk.codegen.poet.ClassSpec;
+import software.amazon.awssdk.codegen.poet.PoetUtils;
+import software.amazon.awssdk.codegen.poet.rules.EndpointRulesSpecUtils;
+
+public class EndpointAwareAuthSchemeParamsSpec implements ClassSpec {
+
+    private final AuthSchemeSpecUtils authSchemeSpecUtils;
+    private final EndpointRulesSpecUtils endpointRulesSpecUtils;
+
+    public EndpointAwareAuthSchemeParamsSpec(IntermediateModel intermediateModel) {
+        this.authSchemeSpecUtils = new AuthSchemeSpecUtils(intermediateModel);
+        this.endpointRulesSpecUtils = new EndpointRulesSpecUtils(intermediateModel);
+    }
+
+    @Override
+    public ClassName className() {
+        return authSchemeSpecUtils.parametersEndpointAwareDefaultImplName();
+    }
+
+    @Override
+    public TypeSpec poetSpec() {
+        TypeSpec.Builder b = PoetUtils.createInterfaceBuilder(className())
+                                      .addAnnotation(SdkInternalApi.class)
+                                      .addMethod(endpointProviderMethod())
+                                      .addType(builderSpec());
+        return b.build();
+    }
+
+    private ClassName builderClassName() {
+        return className().nestedClass("Builder");
+    }
+
+    private MethodSpec endpointProviderMethod() {
+        return MethodSpec.methodBuilder("endpointProvider")
+                         .addModifiers(Modifier.ABSTRACT, Modifier.PUBLIC)
+                         .returns(endpointRulesSpecUtils.providerInterfaceName())
+                         .build();
+    }
+
+    private TypeSpec builderSpec() {
+        ClassName builderClassName = builderClassName();
+        TypeSpec.Builder b = TypeSpec.interfaceBuilder(builderClassName)
+                                     .addModifiers(Modifier.PUBLIC, Modifier.STATIC);
+        b.addMethod(MethodSpec.methodBuilder("endpointProvider")
+                              .addModifiers(Modifier.ABSTRACT, Modifier.PUBLIC)
+                              .addParameter(endpointRulesSpecUtils.providerInterfaceName(), "endpointProvider")
+                              .returns(builderClassName)
+                              .build());
+
+        return b.build();
+    }
+}

codegen/src/main/java/software/amazon/awssdk/codegen/poet/auth/scheme/EndpointBasedAuthSchemeProviderSpec.java

@@ -72,6 +72,7 @@ public TypeSpec poetSpec() {
                         .addField(endpointDelegateInstance())
                         .addMethod(createMethod())
                         .addMethod(resolveAuthSchemeMethod())
+                        .addMethod(endpointProvider())
                         .build();
     }
 
@@ -93,6 +94,25 @@ private FieldSpec endpointDelegateInstance() {
                         .build();
     }
 
+    private MethodSpec endpointProvider() {
+        ClassName endpointProviderClass = endpointRulesSpecUtils.providerInterfaceName();
+        MethodSpec.Builder builder = MethodSpec.methodBuilder("endpointProvider")
+                                               .addModifiers(Modifier.PRIVATE)
+                                               .returns(endpointProviderClass)
+                                               .addParameter(authSchemeSpecUtils.parametersInterfaceName(), "params");
+
+        ClassName endpointAwareParams = authSchemeSpecUtils.parametersEndpointAwareDefaultImplName();
+        builder.beginControlFlow("if (params instanceof $T)", endpointAwareParams);
+        builder.addStatement("$1T endpointAwareParams = ($1T) params", endpointAwareParams);
+        builder.addStatement("$T endpointProvider = endpointAwareParams.endpointProvider()", endpointProviderClass);
+        builder.beginControlFlow("if (endpointProvider != null)");
+        builder.addStatement("return endpointProvider");
+        builder.endControlFlow();
+        builder.endControlFlow();
+        builder.addStatement("return DELEGATE");
+        return builder.build();
+    }
+
     private FieldSpec modeledResolverInstance() {
         return FieldSpec.builder(authSchemeSpecUtils.providerInterfaceName(), "MODELED_RESOLVER")
                         .addModifiers(Modifier.PRIVATE, Modifier.STATIC, Modifier.FINAL)
@@ -124,7 +144,7 @@ private MethodSpec resolveAuthSchemeMethod() {
             }
         });
         spec.addStatement(".build()");
-        spec.addStatement("$T endpoint = $T.joinLikeSync(DELEGATE.resolveEndpoint(endpointParameters))",
+        spec.addStatement("$T endpoint = $T.joinLikeSync(endpointProvider(params).resolveEndpoint(endpointParameters))",
                           Endpoint.class, CompletableFutureUtils.class);
         spec.addStatement("$T authSchemes = endpoint.attribute($T.AUTH_SCHEMES)",
                           ParameterizedTypeName.get(List.class, EndpointAuthScheme.class), AwsEndpointAttribute.class);

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/auth/scheme/query-endpoint-auth-params-auth-scheme-default-params-with-allowlist.java

@@ -19,11 +19,12 @@
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.regions.Region;
 import software.amazon.awssdk.services.query.auth.scheme.QueryAuthSchemeParams;
+import software.amazon.awssdk.services.query.endpoints.QueryEndpointProvider;
 import software.amazon.awssdk.utils.Validate;
 
 @Generated("software.amazon.awssdk:codegen")
 @SdkInternalApi
-public final class DefaultQueryAuthSchemeParams implements QueryAuthSchemeParams {
+public final class DefaultQueryAuthSchemeParams implements QueryAuthSchemeParams, QueryEndpointResolverAware {
     private final String operation;
 
     private final Region region;
@@ -40,6 +41,8 @@ public final class DefaultQueryAuthSchemeParams implements QueryAuthSchemeParams
 
     private final String operationContextParam;
 
+    private final QueryEndpointProvider endpointProvider;
+
     private DefaultQueryAuthSchemeParams(Builder builder) {
         this.operation = Validate.paramNotNull(builder.operation, "operation");
         this.region = builder.region;
@@ -49,6 +52,7 @@ private DefaultQueryAuthSchemeParams(Builder builder) {
         this.booleanContextParam = builder.booleanContextParam;
         this.stringContextParam = builder.stringContextParam;
         this.operationContextParam = builder.operationContextParam;
+        this.endpointProvider = builder.endpointProvider;
     }
 
     public static QueryAuthSchemeParams.Builder builder() {
@@ -96,12 +100,17 @@ public String operationContextParam() {
         return operationContextParam;
     }
 
+    @Override
+    public QueryEndpointProvider endpointProvider() {
+        return endpointProvider;
+    }
+
     @Override
     public QueryAuthSchemeParams.Builder toBuilder() {
         return new Builder(this);
     }
 
-    private static final class Builder implements QueryAuthSchemeParams.Builder {
+    private static final class Builder implements QueryAuthSchemeParams.Builder, QueryEndpointResolverAware.Builder {
         private String operation;
 
         private Region region;
@@ -118,6 +127,8 @@ private static final class Builder implements QueryAuthSchemeParams.Builder {
 
         private String operationContextParam;
 
+        private QueryEndpointProvider endpointProvider;
+
         Builder() {
         }
 
@@ -130,6 +141,7 @@ private static final class Builder implements QueryAuthSchemeParams.Builder {
             this.booleanContextParam = params.booleanContextParam;
             this.stringContextParam = params.stringContextParam;
             this.operationContextParam = params.operationContextParam;
+            this.endpointProvider = params.endpointProvider;
         }
 
         @Override
@@ -187,6 +199,12 @@ public Builder operationContextParam(String operationContextParam) {
             return this;
         }
 
+        @Override
+        public Builder endpointProvider(QueryEndpointProvider endpointProvider) {
+            this.endpointProvider = endpointProvider;
+            return this;
+        }
+
         @Override
         public QueryAuthSchemeParams build() {
             return new DefaultQueryAuthSchemeParams(this);