Support HostPrefix in E2.0, enable E2.0 for all services (#3564)

* Support HostPrefix in E2.0, enable E2.0 for all services

This changes fixes the previously incorrect handling of `hostPrefix` in
services, causing requests for some services to have the wrong URL. This change
also undoes the suppression of rules-based endpoint generation for all services
apart from EventBridge, S3, and S3 Control, which was implemented in #3520
because of the previously mentioned `hostPrefix` issue.

* Review comments

* Set tokenProvider() for Bearer auth services

.changes/next-release/bugfix-AWSSDKforJavav2-4ded69a.json

@@ -0,0 +1,6 @@
+{
+    "type": "bugfix",
+    "category": "AWS SDK for Java v2",
+    "contributor": "",
+    "description": "This changes fixes the previously incorrect handling of `hostPrefix` in services, causing requests for some services to have the wrong URL. This change also undoes the suppression of rules-based endpoint generation for all services apart from EventBridge, S3, and S3 Control, which was implemented in [#3520](https://github.com/aws/aws-sdk-java-v2/issues/3520) because of the previously mentioned `hostPrefix` issue."
+}

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

@@ -18,7 +18,6 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
-import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import software.amazon.awssdk.codegen.emitters.GeneratorTask;
@@ -46,10 +45,6 @@ public EndpointProviderTasks(GeneratorTaskParams dependencies) {
 
     @Override
     protected List<GeneratorTask> createTasks() throws Exception {
-        if (!generatorTaskParams.getModel().getCustomizationConfig().useRuleBasedEndpoints()) {
-            return Collections.emptyList();
-        }
-
         List<GeneratorTask> tasks = new ArrayList<>();
         tasks.add(generateInterface());
         tasks.add(generateParams());

codegen/src/main/java/software/amazon/awssdk/codegen/model/config/customization/CustomizationConfig.java

@@ -213,11 +213,6 @@ public class CustomizationConfig {
 
     private boolean useGlobalEndpoint;
 
-    /**
-     * Whether Endpoints 2.0/rule based endpoints should be used for endpoint resolution.
-     */
-    private boolean useRuleBasedEndpoints = false;
-
     private List<String> interceptors = new ArrayList<>();
 
     private CustomizationConfig() {
@@ -557,14 +552,6 @@ public void setSkipEndpointTests(Map<String, String> skipEndpointTests) {
         this.skipEndpointTests = skipEndpointTests;
     }
 
-    public boolean useRuleBasedEndpoints() {
-        return useRuleBasedEndpoints;
-    }
-
-    public void setUseRuleBasedEndpoints(boolean useRuleBasedEndpoints) {
-        this.useRuleBasedEndpoints = useRuleBasedEndpoints;
-    }
-
     public List<String> getInterceptors() {
         return interceptors;
     }

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

@@ -69,9 +69,7 @@ public TypeSpec poetSpec() {
             }
         }
 
-        if (endpointRulesSpecUtils.isEndpointRulesEnabled()) {
-            builder.addMethod(endpointProviderMethod());
-        }
+        builder.addMethod(endpointProviderMethod());
 
         if (AuthUtils.usesBearerAuth(model)) {
             builder.addMethod(bearerTokenProviderMethod());

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

@@ -111,11 +111,9 @@ public TypeSpec poetSpec() {
         builder.addMethod(finalizeServiceConfigurationMethod());
         defaultAwsAuthSignerMethod().ifPresent(builder::addMethod);
         builder.addMethod(signingNameMethod());
-        if (endpointRulesSpecUtils.isEndpointRulesEnabled()) {
-            builder.addMethod(defaultEndpointProviderMethod());
-        }
+        builder.addMethod(defaultEndpointProviderMethod());
 
-        if (hasClientContextParams() && endpointRulesSpecUtils.isEndpointRulesEnabled()) {
+        if (hasClientContextParams()) {
             model.getClientContextParams().forEach((n, m) -> {
                 builder.addMethod(clientContextParamSetter(n, m));
             });
@@ -191,9 +189,8 @@ private MethodSpec mergeServiceDefaultsMethod() {
                                                .addParameter(SdkClientConfiguration.class, "config")
                                                .addCode("return config.merge(c -> c");
 
-        if (endpointRulesSpecUtils.isEndpointRulesEnabled()) {
-            builder.addCode(".option($T.ENDPOINT_PROVIDER, defaultEndpointProvider())", SdkClientOption.class);
-        }
+        builder.addCode(".option($T.ENDPOINT_PROVIDER, defaultEndpointProvider())", SdkClientOption.class);
+
 
         if (defaultAwsAuthSignerMethod().isPresent()) {
             builder.addCode(".option($T.SIGNER, defaultSigner())\n", SdkAdvancedClientOption.class);
@@ -261,11 +258,9 @@ private MethodSpec finalizeServiceConfigurationMethod() {
 
         List<ClassName> builtInInterceptors = new ArrayList<>();
 
-        if (endpointRulesSpecUtils.isEndpointRulesEnabled()) {
-            builtInInterceptors.add(endpointRulesSpecUtils.resolverInterceptorName());
-            builtInInterceptors.add(endpointRulesSpecUtils.authSchemesInterceptorName());
-            builtInInterceptors.add(endpointRulesSpecUtils.requestModifierInterceptorName());
-        }
+        builtInInterceptors.add(endpointRulesSpecUtils.resolverInterceptorName());
+        builtInInterceptors.add(endpointRulesSpecUtils.authSchemesInterceptorName());
+        builtInInterceptors.add(endpointRulesSpecUtils.requestModifierInterceptorName());
 
         for (String interceptor : model.getCustomizationConfig().getInterceptors()) {
             builtInInterceptors.add(ClassName.bestGuess(interceptor));

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

@@ -38,7 +38,6 @@
 import software.amazon.awssdk.core.client.config.SdkAdvancedClientOption;
 import software.amazon.awssdk.utils.internal.CodegenNamingUtils;
 
-
 public class BaseClientBuilderInterface implements ClassSpec {
     private final IntermediateModel model;
     private final String basePackage;
@@ -73,14 +72,12 @@ public TypeSpec poetSpec() {
             builder.addMethod(serviceConfigurationConsumerBuilderMethod());
         }
 
-        if (endpointRulesSpecUtils.isEndpointRulesEnabled()) {
-            builder.addMethod(endpointProviderMethod());
+        builder.addMethod(endpointProviderMethod());
 
-            if (hasClientContextParams()) {
-                model.getClientContextParams().forEach((n, m) -> {
-                    builder.addMethod(clientContextParamSetter(n, m));
-                });
-            }
+        if (hasClientContextParams()) {
+            model.getClientContextParams().forEach((n, m) -> {
+                builder.addMethod(clientContextParamSetter(n, m));
+            });
         }
 
         if (generateTokenProviderMethod()) {

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

@@ -69,9 +69,7 @@ public TypeSpec poetSpec() {
             }
         }
 
-        if (endpointRulesSpecUtils.isEndpointRulesEnabled()) {
-            builder.addMethod(endpointProviderMethod());
-        }
+        builder.addMethod(endpointProviderMethod());
 
         if (AuthUtils.usesBearerAuth(model)) {
             builder.addMethod(tokenProviderMethodImpl());

codegen/src/main/java/software/amazon/awssdk/codegen/poet/rules/EndpointProviderTestSpec.java

@@ -88,7 +88,7 @@ private MethodSpec testsCasesMethod() {
             b.addStatement("testCases.add(new $T($L, $L))",
                            EndpointProviderTestCase.class,
                            createTestCase(test),
-                           TestGeneratorUtils.createExpect(test.getExpect()));
+                           TestGeneratorUtils.createExpect(test.getExpect(), null, null));
         });
 
         b.addStatement("return testCases");

codegen/src/main/java/software/amazon/awssdk/codegen/poet/rules/EndpointResolverInterceptorSpec.java

@@ -20,7 +20,9 @@
 import com.fasterxml.jackson.jr.stree.JrsString;
 import com.squareup.javapoet.ClassName;
 import com.squareup.javapoet.MethodSpec;
+import com.squareup.javapoet.ParameterizedTypeName;
 import com.squareup.javapoet.TypeSpec;
+import java.util.Iterator;
 import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.CompletionException;
@@ -32,6 +34,8 @@
 import software.amazon.awssdk.codegen.model.rules.endpoints.ParameterModel;
 import software.amazon.awssdk.codegen.model.service.ClientContextParam;
 import software.amazon.awssdk.codegen.model.service.ContextParam;
+import software.amazon.awssdk.codegen.model.service.EndpointTrait;
+import software.amazon.awssdk.codegen.model.service.HostPrefixProcessor;
 import software.amazon.awssdk.codegen.model.service.StaticContextParam;
 import software.amazon.awssdk.codegen.poet.ClassSpec;
 import software.amazon.awssdk.codegen.poet.PoetExtension;
@@ -41,9 +45,12 @@
 import software.amazon.awssdk.core.interceptor.Context;
 import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
 import software.amazon.awssdk.core.interceptor.ExecutionInterceptor;
+import software.amazon.awssdk.core.interceptor.SdkExecutionAttribute;
 import software.amazon.awssdk.core.interceptor.SdkInternalExecutionAttribute;
 import software.amazon.awssdk.endpoints.Endpoint;
 import software.amazon.awssdk.utils.AttributeMap;
+import software.amazon.awssdk.utils.HostnameValidator;
+import software.amazon.awssdk.utils.StringUtils;
 
 public class EndpointResolverInterceptorSpec implements ClassSpec {
     private final IntermediateModel model;
@@ -76,6 +83,8 @@ public TypeSpec poetSpec() {
             b.addMethod(setClientContextParamsMethod());
         }
 
+        b.addMethod(hostPrefixMethod());
+
         return b.build();
     }
 
@@ -106,6 +115,15 @@ private MethodSpec modifyRequestMethod() {
         b.beginControlFlow("try");
         b.addStatement("$T result = $N.resolveEndpoint(ruleParams(context, executionAttributes)).join()", Endpoint.class,
                        providerVar);
+        b.beginControlFlow("if (!$T.disableHostPrefixInjection(executionAttributes))",
+                           endpointRulesSpecUtils.rulesRuntimeClassName("AwsEndpointProviderUtils"));
+        b.addStatement("$T hostPrefix = hostPrefix(executionAttributes.getAttribute($T.OPERATION_NAME), context.request())",
+                       ParameterizedTypeName.get(Optional.class, String.class), SdkExecutionAttribute.class);
+        b.beginControlFlow("if (hostPrefix.isPresent())");
+        b.addStatement("result = $T.addHostPrefix(result, hostPrefix.get())",
+                       endpointRulesSpecUtils.rulesRuntimeClassName("AwsEndpointProviderUtils"));
+        b.endControlFlow();
+        b.endControlFlow();
         b.addStatement("executionAttributes.putAttribute(SdkInternalExecutionAttribute.RESOLVED_ENDPOINT, result)");
         b.addStatement("return context.request()");
         b.endControlFlow();
@@ -171,7 +189,7 @@ private MethodSpec ruleParams() {
                 case AWS_S3_FORCE_PATH_STYLE:
                 case AWS_S3_USE_ARN_REGION:
                 case AWS_S3_CONTROL_USE_ARN_REGION:
-                // end of S3 specific builtins
+                    // end of S3 specific builtins
                 case AWS_STS_USE_GLOBAL_ENDPOINT:
                     // V2 doesn't support this, only regional endpoints
                     return;
@@ -361,4 +379,66 @@ private MethodSpec setClientContextParamsMethod() {
 
         return b.build();
     }
+
+
+    private MethodSpec hostPrefixMethod() {
+        MethodSpec.Builder builder = MethodSpec.methodBuilder("hostPrefix")
+                                               .returns(ParameterizedTypeName.get(Optional.class, String.class))
+                                               .addParameter(String.class, "operationName")
+                                               .addParameter(SdkRequest.class, "request")
+                                               .addModifiers(Modifier.PRIVATE, Modifier.STATIC);
+
+        builder.beginControlFlow("switch (operationName)");
+
+        model.getOperations().forEach((name, opModel) -> {
+            String hostPrefix = getHostPrefix(opModel);
+            if (StringUtils.isBlank(hostPrefix)) {
+                return;
+            }
+
+            builder.beginControlFlow("case $S:", name);
+            HostPrefixProcessor processor = new HostPrefixProcessor(hostPrefix);
+
+            if (processor.c2jNames().isEmpty()) {
+                builder.addStatement("return $T.of($S)", Optional.class, processor.hostWithStringSpecifier());
+            } else {
+                String requestVar = opModel.getInput().getVariableName();
+                processor.c2jNames().forEach(c2jName -> {
+                    builder.addStatement("$1T.validateHostnameCompliant(request.getValueForField($2S, $3T.class).orElse(null), "
+                                         + "$2S, $4S)",
+                                         HostnameValidator.class,
+                                         c2jName,
+                                         String.class,
+                                         requestVar);
+                });
+
+                builder.addCode("return $T.of($T.format($S, ", Optional.class, String.class,
+                                processor.hostWithStringSpecifier());
+                Iterator<String> c2jNamesIter = processor.c2jNames().listIterator();
+                while (c2jNamesIter.hasNext()) {
+                    builder.addCode("request.getValueForField($S, $T.class).get()", c2jNamesIter.next(), String.class);
+                    if (c2jNamesIter.hasNext()) {
+                        builder.addCode(",");
+                    }
+                }
+                builder.addStatement("))");
+            }
+            builder.endControlFlow();
+        });
+
+        builder.addCode("default:");
+        builder.addStatement("return $T.empty()", Optional.class);
+        builder.endControlFlow();
+
+        return builder.build();
+    }
+
+    private String getHostPrefix(OperationModel opModel) {
+        EndpointTrait endpointTrait = opModel.getEndpointTrait();
+        if (endpointTrait == null) {
+            return null;
+        }
+
+        return endpointTrait.getHostPrefix();
+    }
 }

codegen/src/main/java/software/amazon/awssdk/codegen/poet/rules/EndpointRulesClientTestSpec.java

@@ -15,6 +15,8 @@
 
 package software.amazon.awssdk.codegen.poet.rules;
 
+import static software.amazon.awssdk.codegen.poet.rules.TestGeneratorUtils.getHostPrefixTemplate;
+
 import com.fasterxml.jackson.core.TreeNode;
 import com.fasterxml.jackson.jr.stree.JrsArray;
 import com.fasterxml.jackson.jr.stree.JrsObject;
@@ -56,6 +58,7 @@
 import software.amazon.awssdk.codegen.poet.ClassSpec;
 import software.amazon.awssdk.codegen.poet.PoetExtension;
 import software.amazon.awssdk.codegen.poet.PoetUtils;
+import software.amazon.awssdk.codegen.utils.AuthUtils;
 import software.amazon.awssdk.core.SdkSystemSetting;
 import software.amazon.awssdk.core.async.AsyncRequestBody;
 import software.amazon.awssdk.core.rules.testing.AsyncTestCase;
@@ -217,7 +220,8 @@ private MethodSpec syncTestsSourceMethod() {
                               SyncTestCase.class,
                               test.getDocumentation(),
                               syncOperationCallLambda(opModel, test.getParams(), opInput.getOperationParams()),
-                              TestGeneratorUtils.createExpect(test.getExpect()), getSkipReasonBlock(test.getDocumentation()));
+                              TestGeneratorUtils.createExpect(test.getExpect(), opModel, opInput.getOperationParams()),
+                              getSkipReasonBlock(test.getDocumentation()));
 
                     if (operationInputsIter.hasNext()) {
                         b.addCode(",");
@@ -228,7 +232,8 @@ private MethodSpec syncTestsSourceMethod() {
                           SyncTestCase.class,
                           test.getDocumentation(),
                           syncOperationCallLambda(defaultOpModel, test.getParams(), Collections.emptyMap()),
-                          TestGeneratorUtils.createExpect(test.getExpect()), getSkipReasonBlock(test.getDocumentation()));
+                          TestGeneratorUtils.createExpect(test.getExpect(), defaultOpModel, null),
+                          getSkipReasonBlock(test.getDocumentation()));
             }
 
             if (testIter.hasNext()) {
@@ -248,6 +253,9 @@ private CodeBlock syncOperationCallLambda(OperationModel opModel, Map<String, Tr
         b.beginControlFlow("() -> ");
         b.addStatement("$T builder = $T.builder()", syncClientBuilder(), syncClientClass());
         b.addStatement("builder.credentialsProvider($T.CREDENTIALS_PROVIDER)", BaseRuleSetClientTest.class);
+        if (AuthUtils.usesBearerAuth(model)) {
+            b.addStatement("builder.tokenProvider($T.TOKEN_PROVIDER)", BaseRuleSetClientTest.class);
+        }
         b.addStatement("builder.httpClient(getSyncHttpClient())");
 
         if (params != null) {
@@ -272,6 +280,9 @@ private CodeBlock asyncOperationCallLambda(OperationModel opModel, Map<String, T
         b.beginControlFlow("() -> ");
         b.addStatement("$T builder = $T.builder()", asyncClientBuilder(), asyncClientClass());
         b.addStatement("builder.credentialsProvider($T.CREDENTIALS_PROVIDER)", BaseRuleSetClientTest.class);
+        if (AuthUtils.usesBearerAuth(model)) {
+            b.addStatement("builder.tokenProvider($T.TOKEN_PROVIDER)", BaseRuleSetClientTest.class);
+        }
         b.addStatement("builder.httpClient(getAsyncHttpClient())");
 
         if (params != null) {
@@ -355,7 +366,8 @@ private MethodSpec asyncTestsSourceMethod() {
                               AsyncTestCase.class,
                               test.getDocumentation(),
                               asyncOperationCallLambda(opModel, test.getParams(), opInput.getOperationParams()),
-                              TestGeneratorUtils.createExpect(test.getExpect()), getSkipReasonBlock(test.getDocumentation()));
+                              TestGeneratorUtils.createExpect(test.getExpect(), opModel, opInput.getOperationParams()),
+                              getSkipReasonBlock(test.getDocumentation()));
 
                     if (operationInputsIter.hasNext()) {
                         b.addCode(",");
@@ -366,7 +378,8 @@ private MethodSpec asyncTestsSourceMethod() {
                           AsyncTestCase.class,
                           test.getDocumentation(),
                           asyncOperationCallLambda(defaultOpModel, test.getParams(), Collections.emptyMap()),
-                          TestGeneratorUtils.createExpect(test.getExpect()), getSkipReasonBlock(test.getDocumentation()));
+                          TestGeneratorUtils.createExpect(test.getExpect(), defaultOpModel, null),
+                          getSkipReasonBlock(test.getDocumentation()));
             }
 
             if (testIter.hasNext()) {
@@ -399,8 +412,10 @@ private CodeBlock requestCreation(OperationModel opModel, Map<String, TreeNode>
             return b.add(".build()").build();
         }
 
+        String hostPrefix = getHostPrefixTemplate(opModel).orElse("");
+
         inputShape.getMembers().forEach(m -> {
-            if (!boundToPath(m)) {
+            if (!boundToPath(m) && !hostPrefix.contains("{" + m.getC2jName() + "}")) {
                 return;
             }
 
@@ -445,6 +460,12 @@ private static boolean canBeEmpty(OperationModel opModel) {
             return false;
         }
 
+        String hostPrefix = getHostPrefixTemplate(opModel).orElse("");
+
+        if (hostPrefix.contains("{") && hostPrefix.contains("}")) {
+            return false;
+        }
+
         Optional<MemberModel> pathMemberOrStreaming = members.stream()
                                                              .filter(EndpointRulesClientTestSpec::boundToPath)
                                                              .findFirst();

codegen/src/main/java/software/amazon/awssdk/codegen/poet/rules/EndpointRulesSpecUtils.java

@@ -182,8 +182,4 @@ public boolean isS3Control() {
     public TypeName resolverReturnType() {
         return ParameterizedTypeName.get(CompletableFuture.class, Endpoint.class);
     }
-
-    public boolean isEndpointRulesEnabled() {
-        return intermediateModel.getCustomizationConfig().useRuleBasedEndpoints();
-    }
 }