Unverified Commit 7cb9d6af authored by ysaito1001's avatar ysaito1001 Committed by GitHub
Browse files

Add DefaultEndpointResolver to the orchestrator (#2577) 

## Motivation and Context
This PR adds `DefaultEndpointResolver`, a default implementer of the
[EndpointResolver](https://github.com/awslabs/smithy-rs/blob/1e27efe05fe7b991c9f9bbf3d63a297b2dded334/rust-runtime/aws-smithy-runtime-api/src/client/orchestrator.rs#L182-L184

)
trait, to the orchestrator.

## Description
Roughly speaking, the endpoint-related work is currently done by
`make_operation` and `SmithyEndpointStage`, where `make_operation`
constructs endpoint params and applies an endpoint resolver and the
`SmithyEndpointStage` later updates the request header based on the
resolved endpoint. In the orchestrator world, that work is handled by
`DefaultEndpointResolver::resolve_and_apply_endpoint`.

The way endpoint parameters and an endpoint prefix are made available to
`DefaultEndpointResolver::resolve_and_apply_endpoint` is done by
interceptors because they _may_ require pieces of information only
available in an operation input struct. The place that has access to
both an operation input struct and a config bag happens to be an
interceptor.

There are two interceptors `<Operation>EndpointParamsInterceptor` and
`<Operation>EndpointParamsFinalizerInterceptor` per operation. We pass
around endpoint params _builder_ through interceptors to allow it to be
configured with more information at a later point; An end point
parameters builder is first initialized within the
`ServiceRuntimePlugin` with the field values obtained from the service
config, and is stored in a config bag. The finalizer interceptor "seals"
the builder and creates the actual endpoint params before it is passed
to `DefaultEndpointResolver::resolve_and_apply_endpoint`. These
interceptors implement `read_before_execution` because they need to
access an operation input before it gets serialized (otherwise,
`context.input()` would return a `None`).

## Testing
Replaced `StaticUriEndpointResolver` with `DefaultEndpointResolver` in
`sra_test` and verified the test continued passing.

UPDATE: The test currently fails in the `main` branch (it returns a
`None` when extracting `SigV4OperationSigningConfig` from the config bag
in `OverrideSigningTimeInterceptor`, hence `unwrap` fails), and by
merging the `main` branch this PR no longer passes the test, but it does
not add new failures either.

----

_By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice._

---------

Co-authored-by: default avatarYuki Saito <awsaito@amazon.com>
parent a2d37ad2

aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/SigV4AuthDecorator.kt

+2 −0
Original line number Diff line number Diff line
@@ -8,6 +8,7 @@ package software.amazon.smithy.rustsdk 
import software.amazon.smithy.aws.traits.auth.SigV4Trait

import software.amazon.smithy.aws.traits.auth.UnsignedPayloadTrait

import software.amazon.smithy.model.knowledge.ServiceIndex

import software.amazon.smithy.model.shapes.OperationShape

import software.amazon.smithy.model.traits.OptionalAuthTrait

import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext

import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegenDecorator
@@ -38,6 +39,7 @@ class SigV4AuthDecorator : ClientCodegenDecorator { 


    override fun operationRuntimePluginCustomizations(

        codegenContext: ClientCodegenContext,

        operation: OperationShape,

        baseCustomizations: List<OperationRuntimePluginCustomization>,

    ): List<OperationRuntimePluginCustomization> =

        baseCustomizations.letIf(codegenContext.settings.codegenConfig.enableNewSmithyRuntime) {

aws/sra-test/integration-tests/aws-sdk-s3/Cargo.toml

+12 −12
Original line number Diff line number Diff line
@@ -6,18 +6,18 @@ publish = false 
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html



[dependencies]

aws-credential-types = { path = "../../../rust-runtime/aws-credential-types", features = ["test-util"] }

aws-http = { path = "../../../rust-runtime/aws-http" }

aws-runtime = { path = "../../../rust-runtime/aws-runtime" }

aws-sdk-s3 = { path = "../../build/sdk/aws-sdk-s3", features = ["test-util"] }

aws-sigv4 = { path = "../../../rust-runtime/aws-sigv4" }

aws-types = { path = "../../../rust-runtime/aws-types" }

aws-smithy-async = { path = "../../../../rust-runtime/aws-smithy-async", features = ["rt-tokio"] }

aws-smithy-client = { path = "../../../../rust-runtime/aws-smithy-client", features = ["test-util"] }

aws-smithy-types = { path = "../../../../rust-runtime/aws-smithy-types" }

aws-smithy-http = { path = "../../../../rust-runtime/aws-smithy-http" }

aws-smithy-runtime = { path = "../../../../rust-runtime/aws-smithy-runtime", features = ["test-util"] }

aws-smithy-runtime-api = { path = "../../../../rust-runtime/aws-smithy-runtime-api" }

aws-credential-types = { path = "../../../sdk/build/aws-sdk/sdk/aws-credential-types", features = ["test-util"] }

aws-http = { path = "../../../sdk/build/aws-sdk/sdk/aws-http" }

aws-runtime = { path = "../../../sdk/build/aws-sdk/sdk/aws-runtime" }

aws-sdk-s3 = { path = "../../../sdk/build/aws-sdk/sdk/s3/", features = ["test-util"] }

aws-sigv4 = { path = "../../../sdk/build/aws-sdk/sdk/aws-sigv4" }

aws-types = { path = "../../../sdk/build/aws-sdk/sdk/aws-types" }

aws-smithy-async = { path = "../../../sdk/build/aws-sdk/sdk/aws-smithy-async", features = ["rt-tokio"] }

aws-smithy-client = { path = "../../../sdk/build/aws-sdk/sdk/aws-smithy-client", features = ["test-util"] }

aws-smithy-types = { path = "../../../sdk/build/aws-sdk/sdk/aws-smithy-types" }

aws-smithy-http = { path = "../../../sdk/build/aws-sdk/sdk/aws-smithy-http" }

aws-smithy-runtime = { path = "../../../sdk/build/aws-sdk/sdk/aws-smithy-runtime", features = ["test-util"] }

aws-smithy-runtime-api = { path = "../../../sdk/build/aws-sdk/sdk/aws-smithy-runtime-api" }

tokio = { version = "1.23.1", features = ["macros", "test-util", "rt-multi-thread"] }

tracing = "0.1.37"

tracing-subscriber = { version = "0.3.15", features = ["env-filter", "json"] }

aws/sra-test/integration-tests/aws-sdk-s3/tests/sra_test.rs

+87 −7
Original line number Diff line number Diff line
@@ -17,8 +17,10 @@ use aws_sdk_s3::primitives::SdkBody; 
use aws_smithy_client::erase::DynConnector;

use aws_smithy_client::test_connection::TestConnection;

use aws_smithy_runtime::client::connections::adapter::DynConnectorAdapter;

use aws_smithy_runtime_api::client::endpoints::StaticUriEndpointResolver;

use aws_smithy_runtime_api::client::interceptors::{Interceptor, InterceptorContext, Interceptors};

use aws_smithy_runtime_api::client::endpoints::DefaultEndpointResolver;

use aws_smithy_runtime_api::client::interceptors::{

    Interceptor, InterceptorContext, InterceptorError, Interceptors,

};

use aws_smithy_runtime_api::client::orchestrator::{

    BoxError, ConfigBagAccessors, Connection, HttpRequest, HttpResponse, TraceProbe,

};
@@ -27,7 +29,6 @@ use aws_smithy_runtime_api::config_bag::ConfigBag; 
use aws_smithy_runtime_api::type_erasure::TypedBox;

use aws_types::region::SigningRegion;

use aws_types::SigningService;

use http::Uri;

use std::sync::Arc;

use std::time::{Duration, UNIX_EPOCH};
@@ -108,9 +109,16 @@ async fn sra_manual_test() { 
                ),

            );



            cfg.set_endpoint_resolver(StaticUriEndpointResolver::uri(Uri::from_static(

                "https://test-bucket.s3.us-east-1.amazonaws.com/",

            )));

            cfg.set_endpoint_resolver(DefaultEndpointResolver::new(

                aws_smithy_http::endpoint::SharedEndpointResolver::new(

                    aws_sdk_s3::endpoint::DefaultResolver::new(),

                ),

            ));



            let params_builder = aws_sdk_s3::endpoint::Params::builder()

                .set_region(Some("us-east-1".to_owned()))

                .set_endpoint(Some("https://s3.us-east-1.amazonaws.com/".to_owned()));

            cfg.put(params_builder);



            cfg.set_retry_strategy(

                aws_smithy_runtime_api::client::retries::NeverRetryStrategy::new(),
@@ -162,6 +170,77 @@ async fn sra_manual_test() { 
        }

    }



    // This is a temporary operation runtime plugin until <Operation>EndpointParamsInterceptor and

    // <Operation>EndpointParamsFinalizerInterceptor have been fully implemented, in which case

    // `.with_operation_plugin(ManualOperationRuntimePlugin)` can be removed.

    struct ManualOperationRuntimePlugin;



    impl RuntimePlugin for ManualOperationRuntimePlugin {

        fn configure(&self, cfg: &mut ConfigBag) -> Result<(), BoxError> {

            #[derive(Debug)]

            struct ListObjectsV2EndpointParamsInterceptor;

            impl Interceptor<HttpRequest, HttpResponse> for ListObjectsV2EndpointParamsInterceptor {

                fn read_before_execution(

                    &self,

                    context: &InterceptorContext<HttpRequest, HttpResponse>,

                    cfg: &mut ConfigBag,

                ) -> Result<(), BoxError> {

                    let input = context.input()?;

                    let input = input

                        .downcast_ref::<ListObjectsV2Input>()

                        .ok_or_else(|| InterceptorError::invalid_input_access())?;

                    let mut params_builder = cfg

                        .get::<aws_sdk_s3::endpoint::ParamsBuilder>()

                        .ok_or(InterceptorError::read_before_execution(

                            "missing endpoint params builder",

                        ))?

                        .clone();

                    params_builder = params_builder.set_bucket(input.bucket.clone());

                    cfg.put(params_builder);



                    Ok(())

                }

            }



            #[derive(Debug)]

            struct ListObjectsV2EndpointParamsFinalizerInterceptor;

            impl Interceptor<HttpRequest, HttpResponse> for ListObjectsV2EndpointParamsFinalizerInterceptor {

                fn read_before_execution(

                    &self,

                    _context: &InterceptorContext<HttpRequest, HttpResponse>,

                    cfg: &mut ConfigBag,

                ) -> Result<(), BoxError> {

                    let params_builder = cfg

                        .get::<aws_sdk_s3::endpoint::ParamsBuilder>()

                        .ok_or(InterceptorError::read_before_execution(

                            "missing endpoint params builder",

                        ))?

                        .clone();

                    let params = params_builder

                        .build()

                        .map_err(InterceptorError::read_before_execution)?;

                    cfg.put(

                        aws_smithy_runtime_api::client::orchestrator::EndpointResolverParams::new(

                            params,

                        ),

                    );



                    Ok(())

                }

            }



            cfg.get::<Interceptors<HttpRequest, HttpResponse>>()

                .expect("interceptors set")

                .register_operation_interceptor(

                    Arc::new(ListObjectsV2EndpointParamsInterceptor) as _

                )

                .register_operation_interceptor(Arc::new(

                    ListObjectsV2EndpointParamsFinalizerInterceptor,

                ) as _);

            Ok(())

        }

    }



    let conn = TestConnection::new(vec![(

                http::Request::builder()

                    .header("authorization", "AWS4-HMAC-SHA256 Credential=ANOTREAL/20210618/us-east-1/s3/aws4_request, SignedHeaders=host;x-amz-content-sha256;x-amz-date;x-amz-security-token;x-amz-user-agent, Signature=ae78f74d26b6b0c3a403d9e8cc7ec3829d6264a2b33db672bf2b151bbb901786")
@@ -187,7 +266,8 @@ async fn sra_manual_test() { 


    let runtime_plugins = aws_smithy_runtime_api::client::runtime_plugin::RuntimePlugins::new()

        .with_client_plugin(ManualServiceRuntimePlugin(conn.clone()))

        .with_operation_plugin(aws_sdk_s3::operation::list_objects_v2::ListObjectsV2::new());

        .with_operation_plugin(aws_sdk_s3::operation::list_objects_v2::ListObjectsV2::new())

        .with_operation_plugin(ManualOperationRuntimePlugin);



    let input = ListObjectsV2Input::builder()

        .bucket("test-bucket")

codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/RustClientCodegenPlugin.kt

+2 −0
Original line number Diff line number Diff line
@@ -15,6 +15,7 @@ import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegen 
import software.amazon.smithy.rust.codegen.client.smithy.customize.CombinedClientCodegenDecorator

import software.amazon.smithy.rust.codegen.client.smithy.customize.NoOpEventStreamSigningDecorator

import software.amazon.smithy.rust.codegen.client.smithy.customize.RequiredCustomizations

import software.amazon.smithy.rust.codegen.client.smithy.endpoint.EndpointParamsDecorator

import software.amazon.smithy.rust.codegen.client.smithy.endpoint.EndpointsDecorator

import software.amazon.smithy.rust.codegen.client.smithy.generators.client.FluentClientDecorator

import software.amazon.smithy.rust.codegen.client.testutil.ClientDecoratableBuildPlugin
@@ -58,6 +59,7 @@ class RustClientCodegenPlugin : ClientDecoratableBuildPlugin() { 
                RequiredCustomizations(),

                FluentClientDecorator(),

                EndpointsDecorator(),

                EndpointParamsDecorator(),

                NoOpEventStreamSigningDecorator(),

                ApiKeyAuthDecorator(),

                *decorator,

codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/customizations/EndpointPrefixGenerator.kt

+8 −3
Original line number Diff line number Diff line
@@ -7,16 +7,16 @@ package software.amazon.smithy.rust.codegen.client.smithy.customizations 


import software.amazon.smithy.model.shapes.OperationShape

import software.amazon.smithy.model.traits.EndpointTrait

import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext

import software.amazon.smithy.rust.codegen.client.smithy.generators.EndpointTraitBindings

import software.amazon.smithy.rust.codegen.core.rustlang.Writable

import software.amazon.smithy.rust.codegen.core.rustlang.rust

import software.amazon.smithy.rust.codegen.core.rustlang.withBlock

import software.amazon.smithy.rust.codegen.core.rustlang.writable

import software.amazon.smithy.rust.codegen.core.smithy.CodegenContext

import software.amazon.smithy.rust.codegen.core.smithy.customize.OperationCustomization

import software.amazon.smithy.rust.codegen.core.smithy.customize.OperationSection



class EndpointPrefixGenerator(private val codegenContext: CodegenContext, private val shape: OperationShape) :

class EndpointPrefixGenerator(private val codegenContext: ClientCodegenContext, private val shape: OperationShape) :

    OperationCustomization() {

    override fun section(section: OperationSection): Writable = when (section) {

        is OperationSection.MutateRequest -> writable {
@@ -29,11 +29,16 @@ class EndpointPrefixGenerator(private val codegenContext: CodegenContext, privat 
                    epTrait,

                )

                withBlock("let endpoint_prefix = ", "?;") {

                    endpointTraitBindings.render(this, "self")

                    endpointTraitBindings.render(

                        this,

                        "self",

                        codegenContext.settings.codegenConfig.enableNewSmithyRuntime,

                    )

                }

                rust("request.properties_mut().insert(endpoint_prefix);")

            }

        }



        else -> emptySection

    }

}