From 3d3b31b7ede90ed48230fa33146d44eba7959680 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Manuel=20Sugawara=20=28=E2=88=A9=EF=BD=80-=C2=B4=29?= =?UTF-8?q?=E2=8A=83=E2=94=81=E7=82=8E=E7=82=8E=E7=82=8E=E7=82=8E=E7=82=8E?= Date: Wed, 17 Jan 2024 14:14:28 -0800 Subject: [PATCH] Endpoint based auth scheme resolver should honor endpoint overrides --- .../tasks/AuthSchemeGeneratorTasks.java | 7 + .../scheme/AuthSchemeInterceptorSpec.java | 12 ++ .../poet/auth/scheme/AuthSchemeSpecUtils.java | 4 + .../scheme/DefaultAuthSchemeParamsSpec.java | 25 +++ .../EndpointAwareAuthSchemeParamsSpec.java | 75 +++++++ .../EndpointBasedAuthSchemeProviderSpec.java | 22 +- ...-scheme-default-params-with-allowlist.java | 22 +- ...heme-default-params-without-allowlist.java | 22 +- ...e-endpoint-provider-without-allowlist.java | 13 +- ...-params-auth-scheme-endpoint-provider.java | 13 +- ...ith-allowlist-auth-scheme-interceptor.java | 17 +- ...out-allowlist-auth-scheme-interceptor.java | 17 +- .../s3/EndpointOverrideConfigurationTest.java | 201 ++++++++++++++++++ .../S3ExpressAuthSchemeProviderTest.java | 3 + 14 files changed, 434 insertions(+), 19 deletions(-) create mode 100644 codegen/src/main/java/software/amazon/awssdk/codegen/poet/auth/scheme/EndpointAwareAuthSchemeParamsSpec.java create mode 100644 services/s3/src/test/java/software/amazon/awssdk/services/s3/EndpointOverrideConfigurationTest.java diff --git a/codegen/src/main/java/software/amazon/awssdk/codegen/emitters/tasks/AuthSchemeGeneratorTasks.java b/codegen/src/main/java/software/amazon/awssdk/codegen/emitters/tasks/AuthSchemeGeneratorTasks.java index 6f6ed28820a2..fbcec7931bd8 100644 --- a/codegen/src/main/java/software/amazon/awssdk/codegen/emitters/tasks/AuthSchemeGeneratorTasks.java +++ b/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 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)); } diff --git a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/auth/scheme/AuthSchemeInterceptorSpec.java b/codegen/src/main/java/software/amazon/awssdk/codegen/poet/auth/scheme/AuthSchemeInterceptorSpec.java index e464a05babcd..026e91a36742 100644 --- a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/auth/scheme/AuthSchemeInterceptorSpec.java +++ b/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(); } diff --git a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/auth/scheme/AuthSchemeSpecUtils.java b/codegen/src/main/java/software/amazon/awssdk/codegen/poet/auth/scheme/AuthSchemeSpecUtils.java index a9e26356b919..4c997473960f 100644 --- a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/auth/scheme/AuthSchemeSpecUtils.java +++ b/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"); } diff --git a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/auth/scheme/DefaultAuthSchemeParamsSpec.java b/codegen/src/main/java/software/amazon/awssdk/codegen/poet/auth/scheme/DefaultAuthSchemeParamsSpec.java index c2617f0b27d4..3592f0dbf5a9 100644 --- a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/auth/scheme/DefaultAuthSchemeParamsSpec.java +++ b/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())); } } diff --git a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/auth/scheme/EndpointAwareAuthSchemeParamsSpec.java b/codegen/src/main/java/software/amazon/awssdk/codegen/poet/auth/scheme/EndpointAwareAuthSchemeParamsSpec.java new file mode 100644 index 000000000000..4dee2a0c196f --- /dev/null +++ b/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(); + } +} diff --git a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/auth/scheme/EndpointBasedAuthSchemeProviderSpec.java b/codegen/src/main/java/software/amazon/awssdk/codegen/poet/auth/scheme/EndpointBasedAuthSchemeProviderSpec.java index a30dd68bbcce..f680c8307936 100644 --- a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/auth/scheme/EndpointBasedAuthSchemeProviderSpec.java +++ b/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); diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/auth/scheme/query-endpoint-auth-params-auth-scheme-default-params-with-allowlist.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/auth/scheme/query-endpoint-auth-params-auth-scheme-default-params-with-allowlist.java index e9fb9b7cfaeb..9bc1b2065fe1 100644 --- a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/auth/scheme/query-endpoint-auth-params-auth-scheme-default-params-with-allowlist.java +++ b/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); 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 index fc20bf433c72..078729faa1b7 100644 --- 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 @@ -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; @@ -46,6 +47,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; @@ -58,6 +61,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() { @@ -120,12 +124,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; @@ -148,6 +157,8 @@ private static final class Builder implements QueryAuthSchemeParams.Builder { private String operationContextParam; + private QueryEndpointProvider endpointProvider; + Builder() { } @@ -163,6 +174,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 @@ -238,6 +250,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); 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 index d22e9c1ca151..8aa44f160393 100644 --- 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 @@ -45,7 +45,7 @@ public List resolveAuthScheme(QueryAuthSchemeParams params) { .defaultStringParam(params.defaultStringParam()).deprecatedParam(params.deprecatedParam()) .booleanContextParam(params.booleanContextParam()).stringContextParam(params.stringContextParam()) .operationContextParam(params.operationContextParam()).build(); - Endpoint endpoint = CompletableFutureUtils.joinLikeSync(DELEGATE.resolveEndpoint(endpointParameters)); + Endpoint endpoint = CompletableFutureUtils.joinLikeSync(endpointProvider(params).resolveEndpoint(endpointParameters)); List authSchemes = endpoint.attribute(AwsEndpointAttribute.AUTH_SCHEMES); if (authSchemes == null) { return MODELED_RESOLVER.resolveAuthScheme(params); @@ -79,4 +79,15 @@ public List resolveAuthScheme(QueryAuthSchemeParams params) { } return Collections.unmodifiableList(options); } + + private QueryEndpointProvider endpointProvider(QueryAuthSchemeParams params) { + if (params instanceof QueryEndpointResolverAware) { + QueryEndpointResolverAware endpointAwareParams = (QueryEndpointResolverAware) params; + QueryEndpointProvider endpointProvider = endpointAwareParams.endpointProvider(); + if (endpointProvider != null) { + return endpointProvider; + } + } + return DELEGATE; + } } diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/auth/scheme/query-endpoint-auth-params-auth-scheme-endpoint-provider.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/auth/scheme/query-endpoint-auth-params-auth-scheme-endpoint-provider.java index 0fd273b82882..589880dbef0d 100644 --- a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/auth/scheme/query-endpoint-auth-params-auth-scheme-endpoint-provider.java +++ b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/auth/scheme/query-endpoint-auth-params-auth-scheme-endpoint-provider.java @@ -43,7 +43,7 @@ public List resolveAuthScheme(QueryAuthSchemeParams params) { .defaultTrueParam(params.defaultTrueParam()).defaultStringParam(params.defaultStringParam()) .deprecatedParam(params.deprecatedParam()).booleanContextParam(params.booleanContextParam()) .stringContextParam(params.stringContextParam()).operationContextParam(params.operationContextParam()).build(); - Endpoint endpoint = CompletableFutureUtils.joinLikeSync(DELEGATE.resolveEndpoint(endpointParameters)); + Endpoint endpoint = CompletableFutureUtils.joinLikeSync(endpointProvider(params).resolveEndpoint(endpointParameters)); List authSchemes = endpoint.attribute(AwsEndpointAttribute.AUTH_SCHEMES); if (authSchemes == null) { return MODELED_RESOLVER.resolveAuthScheme(params); @@ -77,4 +77,15 @@ public List resolveAuthScheme(QueryAuthSchemeParams params) { } return Collections.unmodifiableList(options); } + + private QueryEndpointProvider endpointProvider(QueryAuthSchemeParams params) { + if (params instanceof QueryEndpointResolverAware) { + QueryEndpointResolverAware endpointAwareParams = (QueryEndpointResolverAware) params; + QueryEndpointProvider endpointProvider = endpointAwareParams.endpointProvider(); + if (endpointProvider != null) { + return endpointProvider; + } + } + return DELEGATE; + } } diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/auth/scheme/query-endpoint-auth-params-with-allowlist-auth-scheme-interceptor.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/auth/scheme/query-endpoint-auth-params-with-allowlist-auth-scheme-interceptor.java index 5abf214378a0..aee41328cd1f 100644 --- a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/auth/scheme/query-endpoint-auth-params-with-allowlist-auth-scheme-interceptor.java +++ b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/auth/scheme/query-endpoint-auth-params-with-allowlist-auth-scheme-interceptor.java @@ -19,6 +19,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; @@ -32,6 +33,7 @@ import software.amazon.awssdk.services.query.auth.scheme.QueryAuthSchemeParams; import software.amazon.awssdk.services.query.auth.scheme.QueryAuthSchemeProvider; import software.amazon.awssdk.services.query.endpoints.QueryEndpointParams; +import software.amazon.awssdk.services.query.endpoints.QueryEndpointProvider; import software.amazon.awssdk.services.query.endpoints.internal.QueryResolveEndpointInterceptor; import software.amazon.awssdk.utils.Logger; import software.amazon.awssdk.utils.Validate; @@ -60,14 +62,12 @@ private SelectedAuthScheme selectAuthScheme(List> authSchemes = executionAttributes.getAttribute(SdkInternalExecutionAttribute.AUTH_SCHEMES); - IdentityProviders identityProviders = executionAttributes - .getAttribute(SdkInternalExecutionAttribute.IDENTITY_PROVIDERS); + IdentityProviders identityProviders = executionAttributes.getAttribute(SdkInternalExecutionAttribute.IDENTITY_PROVIDERS); List> discardedReasons = new ArrayList<>(); for (AuthSchemeOption authOption : authOptions) { AuthScheme authScheme = authSchemes.get(authOption.schemeId()); SelectedAuthScheme selectedAuthScheme = trySelectAuthScheme(authOption, authScheme, - identityProviders, discardedReasons, - metricCollector, executionAttributes); + identityProviders, discardedReasons, metricCollector, executionAttributes); if (selectedAuthScheme != null) { if (!discardedReasons.isEmpty()) { LOG.debug(() -> String.format("%s auth will be used, discarded: '%s'", authOption.schemeId(), @@ -95,12 +95,17 @@ private QueryAuthSchemeParams authSchemeParams(SdkRequest request, ExecutionAttr builder.operationContextParam(endpointParams.operationContextParam()); String operation = executionAttributes.getAttribute(SdkExecutionAttribute.OPERATION_NAME); builder.operation(operation); + if (builder instanceof QueryEndpointResolverAware.Builder) { + EndpointProvider endpointProvider = executionAttributes.getAttribute(SdkInternalExecutionAttribute.ENDPOINT_PROVIDER); + if (endpointProvider instanceof QueryEndpointProvider) { + ((QueryEndpointResolverAware.Builder) builder).endpointProvider((QueryEndpointProvider) endpointProvider); + } + } return builder.build(); } private SelectedAuthScheme trySelectAuthScheme(AuthSchemeOption authOption, AuthScheme authScheme, - IdentityProviders identityProviders, List> discardedReasons, - MetricCollector metricCollector, + IdentityProviders identityProviders, List> discardedReasons, MetricCollector metricCollector, ExecutionAttributes executionAttributes) { if (authScheme == null) { discardedReasons.add(() -> String.format("'%s' is not enabled for this request.", authOption.schemeId())); 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 index 378032221bc1..f95cc1366cd3 100644 --- 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 @@ -19,6 +19,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; @@ -32,6 +33,7 @@ import software.amazon.awssdk.services.query.auth.scheme.QueryAuthSchemeParams; import software.amazon.awssdk.services.query.auth.scheme.QueryAuthSchemeProvider; import software.amazon.awssdk.services.query.endpoints.QueryEndpointParams; +import software.amazon.awssdk.services.query.endpoints.QueryEndpointProvider; import software.amazon.awssdk.services.query.endpoints.internal.QueryResolveEndpointInterceptor; import software.amazon.awssdk.utils.Logger; import software.amazon.awssdk.utils.Validate; @@ -60,14 +62,12 @@ private SelectedAuthScheme selectAuthScheme(List> authSchemes = executionAttributes.getAttribute(SdkInternalExecutionAttribute.AUTH_SCHEMES); - IdentityProviders identityProviders = executionAttributes - .getAttribute(SdkInternalExecutionAttribute.IDENTITY_PROVIDERS); + IdentityProviders identityProviders = executionAttributes.getAttribute(SdkInternalExecutionAttribute.IDENTITY_PROVIDERS); List> discardedReasons = new ArrayList<>(); for (AuthSchemeOption authOption : authOptions) { AuthScheme authScheme = authSchemes.get(authOption.schemeId()); SelectedAuthScheme selectedAuthScheme = trySelectAuthScheme(authOption, authScheme, - identityProviders, discardedReasons, - metricCollector, executionAttributes); + identityProviders, discardedReasons, metricCollector, executionAttributes); if (selectedAuthScheme != null) { if (!discardedReasons.isEmpty()) { LOG.debug(() -> String.format("%s auth will be used, discarded: '%s'", authOption.schemeId(), @@ -98,12 +98,17 @@ private QueryAuthSchemeParams authSchemeParams(SdkRequest request, ExecutionAttr builder.operationContextParam(endpointParams.operationContextParam()); String operation = executionAttributes.getAttribute(SdkExecutionAttribute.OPERATION_NAME); builder.operation(operation); + if (builder instanceof QueryEndpointResolverAware.Builder) { + EndpointProvider endpointProvider = executionAttributes.getAttribute(SdkInternalExecutionAttribute.ENDPOINT_PROVIDER); + if (endpointProvider instanceof QueryEndpointProvider) { + ((QueryEndpointResolverAware.Builder) builder).endpointProvider((QueryEndpointProvider) endpointProvider); + } + } return builder.build(); } private SelectedAuthScheme trySelectAuthScheme(AuthSchemeOption authOption, AuthScheme authScheme, - IdentityProviders identityProviders, List> discardedReasons, - MetricCollector metricCollector, + IdentityProviders identityProviders, List> discardedReasons, MetricCollector metricCollector, ExecutionAttributes executionAttributes) { if (authScheme == null) { discardedReasons.add(() -> String.format("'%s' is not enabled for this request.", authOption.schemeId())); diff --git a/services/s3/src/test/java/software/amazon/awssdk/services/s3/EndpointOverrideConfigurationTest.java b/services/s3/src/test/java/software/amazon/awssdk/services/s3/EndpointOverrideConfigurationTest.java new file mode 100644 index 000000000000..d32a7653c2c1 --- /dev/null +++ b/services/s3/src/test/java/software/amazon/awssdk/services/s3/EndpointOverrideConfigurationTest.java @@ -0,0 +1,201 @@ +/* + * 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.s3; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +import java.util.concurrent.CompletableFuture; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import software.amazon.awssdk.auth.credentials.AwsBasicCredentials; +import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider; +import software.amazon.awssdk.core.SdkPlugin; +import software.amazon.awssdk.core.async.AsyncResponseTransformer; +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.endpoints.Endpoint; +import software.amazon.awssdk.endpoints.EndpointProvider; +import software.amazon.awssdk.regions.Region; +import software.amazon.awssdk.services.s3.endpoints.S3EndpointParams; +import software.amazon.awssdk.services.s3.endpoints.S3EndpointProvider; +import software.amazon.awssdk.utils.Validate; + +class EndpointOverrideConfigurationTest { + + private CapturingInterceptor interceptor; + private CallCountS3EndpointProvider callCountS3EndpointProvider; + + + @BeforeEach + void setup() { + this.interceptor = new CapturingInterceptor(); + this.callCountS3EndpointProvider = null; + } + + @Test + void testClientEndpointOverrideIsUsed() { + S3Client syncClient = syncClientBuilder() + .addPlugin(config -> config.overrideConfiguration(c -> c.addExecutionInterceptor(interceptor))) + .addPlugin(overrideEndpointPlugin()) + .build(); + assertThatThrownBy(() -> syncClient.getObject(r -> { + r.key("key").bucket("bucket"); + })) + .hasMessageContaining("boom!"); + + assertThat(callCountS3EndpointProvider).isNotNull(); + assertThat(callCountS3EndpointProvider.getCallCount()).isEqualTo(2); + } + + @Test + void testRequestEndpointOverrideIsUsed() { + S3Client syncClient = syncClientBuilder() + .addPlugin(config -> config.overrideConfiguration(c -> c.addExecutionInterceptor(interceptor))) + .build(); + assertThatThrownBy(() -> syncClient.getObject(r -> { + r.overrideConfiguration(c -> c.addPlugin(overrideEndpointPlugin())); + r.key("key").bucket("bucket"); + })) + .hasMessageContaining("boom!"); + + assertThat(callCountS3EndpointProvider).isNotNull(); + assertThat(callCountS3EndpointProvider.getCallCount()).isEqualTo(2); + } + + @Test + void testAsyncClientEndpointOverrideIsUsed() { + S3AsyncClient syncClient = asyncClientBuilder() + .addPlugin(config -> config.overrideConfiguration(c -> c.addExecutionInterceptor(interceptor))) + .addPlugin(overrideEndpointPlugin()) + .build(); + assertThatThrownBy(() -> syncClient.getObject(r -> { + r.key("key").bucket("bucket"); + }, AsyncResponseTransformer.toBytes()).join()) + .hasMessageContaining("boom!"); + + assertThat(callCountS3EndpointProvider).isNotNull(); + assertThat(callCountS3EndpointProvider.getCallCount()).isEqualTo(2); + } + + @Test + void testAsyncRequestEndpointOverrideIsUsed() { + S3AsyncClient syncClient = asyncClientBuilder() + .addPlugin(config -> config.overrideConfiguration(c -> c.addExecutionInterceptor(interceptor))) + .addPlugin(overrideEndpointPlugin()) + .build(); + assertThatThrownBy(() -> syncClient.getObject(r -> { + r.overrideConfiguration(c -> c.addPlugin(overrideEndpointPlugin())); + r.key("key").bucket("bucket"); + }, AsyncResponseTransformer.toBytes()).join()) + .hasMessageContaining("boom!"); + + assertThat(callCountS3EndpointProvider).isNotNull(); + assertThat(callCountS3EndpointProvider.getCallCount()).isEqualTo(2); + } + + private SdkPlugin overrideEndpointPlugin() { + return config -> { + EndpointProvider provider = config.endpointProvider(); + if (provider instanceof S3EndpointProvider) { + S3EndpointProvider s3EndpointProvider = (S3EndpointProvider) provider; + config.endpointProvider(endpointProvider(s3EndpointProvider)); + } + }; + } + + private S3ClientBuilder syncClientBuilder() { + return S3Client.builder() + .addPlugin(c -> { + S3ServiceClientConfiguration.Builder config = + Validate.isInstanceOf(S3ServiceClientConfiguration.Builder.class, c, + "\uD83E\uDD14"); + config.region(Region.US_WEST_2) + .credentialsProvider( + StaticCredentialsProvider.create( + AwsBasicCredentials.create("akid", "skid"))) + .overrideConfiguration(oc -> oc.addExecutionInterceptor(interceptor)); + }); + + } + + private S3AsyncClientBuilder asyncClientBuilder() { + return S3AsyncClient.builder() + .addPlugin(c -> { + S3ServiceClientConfiguration.Builder config = + Validate.isInstanceOf(S3ServiceClientConfiguration.Builder.class, c, + "\uD83E\uDD14"); + config.region(Region.US_WEST_2) + .credentialsProvider( + StaticCredentialsProvider.create( + AwsBasicCredentials.create("akid", "skid"))) + .overrideConfiguration(oc -> oc.addExecutionInterceptor(interceptor)); + }); + + } + + public S3EndpointProvider endpointProvider(S3EndpointProvider delegate) { + this.callCountS3EndpointProvider = new CallCountS3EndpointProvider(delegate); + return this.callCountS3EndpointProvider; + } + + static class CallCountS3EndpointProvider implements S3EndpointProvider { + private final S3EndpointProvider delegate; + private int callCount; + + CallCountS3EndpointProvider(S3EndpointProvider delegate) { + this.delegate = Validate.paramNotNull(delegate, "delegate"); + } + + public int getCallCount() { + return callCount; + } + + @Override + public CompletableFuture resolveEndpoint(S3EndpointParams endpointParams) { + callCount++; + return delegate.resolveEndpoint(endpointParams); + } + } + + public static class CapturingInterceptor implements ExecutionInterceptor { + private Context.BeforeTransmission context; + private ExecutionAttributes executionAttributes; + + @Override + public void beforeTransmission(Context.BeforeTransmission context, ExecutionAttributes executionAttributes) { + this.context = context; + this.executionAttributes = executionAttributes; + throw new RuntimeException("boom!"); + } + + public ExecutionAttributes executionAttributes() { + return executionAttributes; + } + + public Context.BeforeTransmission context() { + return context; + } + + public class CaptureCompletedException extends RuntimeException { + CaptureCompletedException(String message) { + super(message); + } + } + } + +} diff --git a/services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/s3express/S3ExpressAuthSchemeProviderTest.java b/services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/s3express/S3ExpressAuthSchemeProviderTest.java index d427e5b8b7de..3c0ea297b530 100644 --- a/services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/s3express/S3ExpressAuthSchemeProviderTest.java +++ b/services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/s3express/S3ExpressAuthSchemeProviderTest.java @@ -38,6 +38,7 @@ import software.amazon.awssdk.services.s3.auth.scheme.internal.DefaultS3AuthSchemeProvider; import software.amazon.awssdk.services.s3.auth.scheme.internal.S3AuthSchemeInterceptor; import software.amazon.awssdk.services.s3.endpoints.S3ClientContextParams; +import software.amazon.awssdk.services.s3.endpoints.S3EndpointProvider; import software.amazon.awssdk.services.s3.model.PutObjectRequest; import software.amazon.awssdk.services.s3.s3express.S3ExpressAuthScheme; import software.amazon.awssdk.services.s3.s3express.S3ExpressSessionCredentials; @@ -71,6 +72,8 @@ private ExecutionAttributes requiredExecutionAttributes(AttributeMap clientConte DefaultIdentityProviders.builder() .putIdentityProvider(DefaultCredentialsProvider.create()) .build()); + executionAttributes.putAttribute(SdkInternalExecutionAttribute.ENDPOINT_PROVIDER, + S3EndpointProvider.defaultProvider()); return executionAttributes; }