Simplify event stream message signer configuration

This PR creates a `DeferredSigner` implementation that allows for the event stream message signer to be wired up by the signing implementation later in the request lifecycle rather than by adding an event stream signer method to the config. Refactoring this brings the middleware client implementation closer to how the orchestrator implementation will work, which unblocks the work required to make event streams work in the orchestrator.
smithy-lang · May 4, 2023 · 6c7eb17 · 6c7eb17
1 parent beedd2c
commit 6c7eb17
Show file tree

Hide file tree

Showing 17 changed files with 330 additions and 380 deletions.
diff --git a/aws/rust-runtime/aws-sig-auth/src/event_stream.rs b/aws/rust-runtime/aws-sig-auth/src/event_stream.rs
@@ -3,114 +3,98 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-use crate::middleware::Signature;
 use aws_credential_types::Credentials;
 use aws_sigv4::event_stream::{sign_empty_message, sign_message};
 use aws_sigv4::SigningParams;
 use aws_smithy_eventstream::frame::{Message, SignMessage, SignMessageError};
-use aws_smithy_http::property_bag::{PropertyBag, SharedPropertyBag};
 use aws_types::region::SigningRegion;
 use aws_types::SigningService;
 use std::time::SystemTime;
 
 /// Event Stream SigV4 signing implementation.
 #[derive(Debug)]
 pub struct SigV4Signer {
-    properties: SharedPropertyBag,
-    last_signature: Option<String>,
+    last_signature: String,
+    credentials: Credentials,
+    signing_region: SigningRegion,
+    signing_service: SigningService,
+    time: Option<SystemTime>,
 }
 
 impl SigV4Signer {
-    pub fn new(properties: SharedPropertyBag) -> Self {
+    pub fn new(
+        last_signature: String,
+        credentials: Credentials,
+        signing_region: SigningRegion,
+        signing_service: SigningService,
+        time: Option<SystemTime>,
+    ) -> Self {
         Self {
-            properties,
-            last_signature: None,
+            last_signature,
+            credentials,
+            signing_region,
+            signing_service,
+            time,
         }
     }
 
-    fn signing_params(properties: &PropertyBag) -> SigningParams<()> {
-        // Every single one of these values would have been retrieved during the initial request,
-        // so we can safely assume they all exist in the property bag at this point.
-        let credentials = properties.get::<Credentials>().unwrap();
-        let region = properties.get::<SigningRegion>().unwrap();
-        let signing_service = properties.get::<SigningService>().unwrap();
-        let time = properties
-            .get::<SystemTime>()
-            .copied()
-            .unwrap_or_else(SystemTime::now);
+    fn signing_params(&self) -> SigningParams<()> {
         let mut builder = SigningParams::builder()
-            .access_key(credentials.access_key_id())
-            .secret_key(credentials.secret_access_key())
-            .region(region.as_ref())
-            .service_name(signing_service.as_ref())
-            .time(time)
+            .access_key(self.credentials.access_key_id())
+            .secret_key(self.credentials.secret_access_key())
+            .region(self.signing_region.as_ref())
+            .service_name(self.signing_service.as_ref())
+            .time(self.time.unwrap_or_else(SystemTime::now))
             .settings(());
-        builder.set_security_token(credentials.session_token());
+        builder.set_security_token(self.credentials.session_token());
         builder.build().unwrap()
     }
 }
 
 impl SignMessage for SigV4Signer {
     fn sign(&mut self, message: Message) -> Result<Message, SignMessageError> {
-        let properties = self.properties.acquire();
-        if self.last_signature.is_none() {
-            // The Signature property should exist in the property bag for all Event Stream requests.
-            self.last_signature = Some(
-                properties
-                    .get::<Signature>()
-                    .expect("property bag contains initial Signature")
-                    .as_ref()
-                    .into(),
-            )
-        }
-
         let (signed_message, signature) = {
-            let params = Self::signing_params(&properties);
-            sign_message(&message, self.last_signature.as_ref().unwrap(), &params).into_parts()
+            let params = self.signing_params();
+            sign_message(&message, &self.last_signature, &params).into_parts()
         };
-        self.last_signature = Some(signature);
+        self.last_signature = signature;
         Ok(signed_message)
     }
 
     fn sign_empty(&mut self) -> Option<Result<Message, SignMessageError>> {
-        let properties = self.properties.acquire();
-        if self.last_signature.is_none() {
-            // The Signature property should exist in the property bag for all Event Stream requests.
-            self.last_signature = Some(properties.get::<Signature>().unwrap().as_ref().into())
-        }
         let (signed_message, signature) = {
-            let params = Self::signing_params(&properties);
-            sign_empty_message(self.last_signature.as_ref().unwrap(), &params).into_parts()
+            let params = self.signing_params();
+            sign_empty_message(&self.last_signature, &params).into_parts()
         };
-        self.last_signature = Some(signature);
+        self.last_signature = signature;
         Some(Ok(signed_message))
     }
 }
 
 #[cfg(test)]
 mod tests {
     use crate::event_stream::SigV4Signer;
-    use crate::middleware::Signature;
     use aws_credential_types::Credentials;
     use aws_smithy_eventstream::frame::{HeaderValue, Message, SignMessage};
-    use aws_smithy_http::property_bag::PropertyBag;
     use aws_types::region::Region;
     use aws_types::region::SigningRegion;
     use aws_types::SigningService;
     use std::time::{Duration, UNIX_EPOCH};
 
+    fn check_send_sync<T: Send + Sync>(value: T) -> T {
+        value
+    }
+
     #[test]
     fn sign_message() {
         let region = Region::new("us-east-1");
-        let mut properties = PropertyBag::new();
-        properties.insert(region.clone());
-        properties.insert(UNIX_EPOCH + Duration::new(1611160427, 0));
-        properties.insert(SigningService::from_static("transcribe"));
-        properties.insert(Credentials::for_tests());
-        properties.insert(SigningRegion::from(region));
-        properties.insert(Signature::new("initial-signature".into()));
-
-        let mut signer = SigV4Signer::new(properties.into());
+        let mut signer = check_send_sync(SigV4Signer::new(
+            "initial-signature".into(),
+            Credentials::for_tests(),
+            SigningRegion::from(region),
+            SigningService::from_static("transcribe"),
+            Some(UNIX_EPOCH + Duration::new(1611160427, 0)),
+        ));
         let mut signatures = Vec::new();
         for _ in 0..5 {
             let signed = signer

diff --git a/aws/rust-runtime/aws-sig-auth/src/middleware.rs b/aws/rust-runtime/aws-sig-auth/src/middleware.rs
@@ -20,21 +20,10 @@ use crate::signer::{
     OperationSigningConfig, RequestConfig, SigV4Signer, SigningError, SigningRequirements,
 };
 
-/// Container for the request signature for use in the property bag.
-#[non_exhaustive]
-pub struct Signature(String);
-
-impl Signature {
-    pub fn new(signature: String) -> Self {
-        Self(signature)
-    }
-}
-
-impl AsRef<str> for Signature {
-    fn as_ref(&self) -> &str {
-        &self.0
-    }
-}
+#[cfg(feature = "sign-eventstream")]
+use crate::event_stream::SigV4Signer as EventStreamSigV4Signer;
+#[cfg(feature = "sign-eventstream")]
+use aws_smithy_eventstream::frame::DeferredSignerSender;
 
 /// Middleware stage to sign requests with SigV4
 ///
@@ -177,11 +166,26 @@ impl MapRequest for SigV4SigningStage {
                     SigningRequirements::Required => signing_config(config)?,
                 };
 
-            let signature = self
+            let _signature = self
                 .signer
                 .sign(operation_config, &request_config, &creds, &mut req)
                 .map_err(SigningStageErrorKind::SigningFailure)?;
-            config.insert(signature);
+
+            // If this is an event stream operation, set up the event stream signer
+            #[cfg(feature = "sign-eventstream")]
+            if let Some(signer_sender) = config.get::<DeferredSignerSender>() {
+                let time_override = config.get::<SystemTime>().copied();
+                signer_sender
+                    .send(Box::new(EventStreamSigV4Signer::new(
+                        _signature,
+                        creds,
+                        request_config.region.clone(),
+                        request_config.service.clone(),
+                        time_override,
+                    )) as _)
+                    .expect("failed to send deferred signer");
+            }
+
             Ok(req)
         })
     }
@@ -202,13 +206,17 @@ mod test {
     use aws_types::region::{Region, SigningRegion};
     use aws_types::SigningService;
 
-    use crate::middleware::{
-        SigV4SigningStage, Signature, SigningStageError, SigningStageErrorKind,
-    };
+    use crate::middleware::{SigV4SigningStage, SigningStageError, SigningStageErrorKind};
     use crate::signer::{OperationSigningConfig, SigV4Signer};
 
+    #[cfg(feature = "sign-eventstream")]
     #[test]
-    fn places_signature_in_property_bag() {
+    fn sends_event_stream_signer_for_event_stream_operations() {
+        use crate::event_stream::SigV4Signer as EventStreamSigV4Signer;
+        use aws_smithy_eventstream::frame::{DeferredSigner, SignMessage};
+        use std::time::SystemTime;
+
+        let (mut deferred_signer, deferred_signer_sender) = DeferredSigner::new();
         let req = http::Request::builder()
             .uri("https://test-service.test-region.amazonaws.com/")
             .body(SdkBody::from(""))
@@ -217,21 +225,31 @@ mod test {
         let req = operation::Request::new(req)
             .augment(|req, properties| {
                 properties.insert(region.clone());
-                properties.insert(UNIX_EPOCH + Duration::new(1611160427, 0));
+                properties.insert::<SystemTime>(UNIX_EPOCH + Duration::new(1611160427, 0));
                 properties.insert(SigningService::from_static("kinesis"));
                 properties.insert(OperationSigningConfig::default_config());
                 properties.insert(Credentials::for_tests());
-                properties.insert(SigningRegion::from(region));
+                properties.insert(SigningRegion::from(region.clone()));
+                properties.insert(deferred_signer_sender);
                 Result::<_, Infallible>::Ok(req)
             })
             .expect("succeeds");
 
         let signer = SigV4SigningStage::new(SigV4Signer::new());
-        let req = signer.apply(req).unwrap();
+        let _ = signer.apply(req).unwrap();
+
+        let mut signer_for_comparison = EventStreamSigV4Signer::new(
+            // This is the expected SigV4 signature for the HTTP request above
+            "abac477b4afabf5651079e7b9a0aa6a1a3e356a7418a81d974cdae9d4c8e5441".into(),
+            Credentials::for_tests(),
+            SigningRegion::from(region),
+            SigningService::from_static("kinesis"),
+            Some(UNIX_EPOCH + Duration::new(1611160427, 0)),
+        );
 
-        let property_bag = req.properties();
-        let signature = property_bag.get::<Signature>();
-        assert!(signature.is_some());
+        let expected_signed_empty = signer_for_comparison.sign_empty().unwrap().unwrap();
+        let actual_signed_empty = deferred_signer.sign_empty().unwrap().unwrap();
+        assert_eq!(expected_signed_empty, actual_signed_empty);
     }
 
     // check that the endpoint middleware followed by signing middleware produce the expected result

diff --git a/aws/rust-runtime/aws-sig-auth/src/signer.rs b/aws/rust-runtime/aws-sig-auth/src/signer.rs
@@ -3,7 +3,6 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-use crate::middleware::Signature;
 use aws_credential_types::Credentials;
 use aws_sigv4::http_request::{
     sign, PayloadChecksumKind, PercentEncodingMode, SessionTokenMode, SignableRequest,
@@ -191,7 +190,7 @@ impl SigV4Signer {
         request_config: &RequestConfig<'_>,
         credentials: &Credentials,
         request: &mut http::Request<SdkBody>,
-    ) -> Result<Signature, SigningError> {
+    ) -> Result<String, SigningError> {
         let settings = Self::settings(operation_config);
         let signing_params = Self::signing_params(settings, credentials, request_config);
 
@@ -223,7 +222,7 @@ impl SigV4Signer {
 
         signing_instructions.apply_to_request(request);
 
-        Ok(Signature::new(signature))
+        Ok(signature)
     }
 }
 

diff --git a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsPresigningDecorator.kt b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsPresigningDecorator.kt
@@ -21,6 +21,7 @@ import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegen
 import software.amazon.smithy.rust.codegen.client.smithy.generators.client.FluentClientCustomization
 import software.amazon.smithy.rust.codegen.client.smithy.generators.client.FluentClientSection
 import software.amazon.smithy.rust.codegen.client.smithy.generators.protocol.MakeOperationGenerator
+import software.amazon.smithy.rust.codegen.client.smithy.protocols.ClientHttpBoundProtocolPayloadGenerator
 import software.amazon.smithy.rust.codegen.core.rustlang.RustWriter
 import software.amazon.smithy.rust.codegen.core.rustlang.Writable
 import software.amazon.smithy.rust.codegen.core.rustlang.docs
@@ -34,7 +35,6 @@ import software.amazon.smithy.rust.codegen.core.smithy.RuntimeConfig
 import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
 import software.amazon.smithy.rust.codegen.core.smithy.customize.OperationCustomization
 import software.amazon.smithy.rust.codegen.core.smithy.customize.OperationSection
-import software.amazon.smithy.rust.codegen.core.smithy.protocols.HttpBoundProtocolPayloadGenerator
 import software.amazon.smithy.rust.codegen.core.util.cloneOperation
 import software.amazon.smithy.rust.codegen.core.util.expectTrait
 import software.amazon.smithy.rust.codegen.core.util.hasTrait
@@ -173,7 +173,7 @@ class AwsInputPresignedMethod(
         MakeOperationGenerator(
             codegenContext,
             protocol,
-            HttpBoundProtocolPayloadGenerator(codegenContext, protocol),
+            ClientHttpBoundProtocolPayloadGenerator(codegenContext, protocol),
             // Prefixed with underscore to avoid colliding with modeled functions
             functionName = makeOperationFn,
             public = false,

diff --git a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/SigV4SigningDecorator.kt b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/SigV4SigningDecorator.kt
@@ -16,7 +16,7 @@ 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
 import software.amazon.smithy.rust.codegen.client.smithy.generators.config.ConfigCustomization
-import software.amazon.smithy.rust.codegen.client.smithy.generators.config.EventStreamSigningConfig
+import software.amazon.smithy.rust.codegen.client.smithy.generators.config.ServiceConfig
 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.rustTemplate
@@ -77,17 +77,17 @@ class SigV4SigningDecorator : ClientCodegenDecorator {
 }
 
 class SigV4SigningConfig(
-    runtimeConfig: RuntimeConfig,
+    private val runtimeConfig: RuntimeConfig,
     private val serviceHasEventStream: Boolean,
     private val sigV4Trait: SigV4Trait,
-) : EventStreamSigningConfig(runtimeConfig) {
-    private val codegenScope = arrayOf(
-        "SigV4Signer" to AwsRuntimeType.awsSigAuthEventStream(runtimeConfig).resolve("event_stream::SigV4Signer"),
-    )
-
-    override fun configImplSection(): Writable {
-        return writable {
-            rustTemplate(
+) : ConfigCustomization() {
+    override fun section(section: ServiceConfig): Writable = writable {
+        if (section is ServiceConfig.ConfigImpl) {
+            if (serviceHasEventStream) {
+                // enable the aws-sig-auth `sign-eventstream` feature
+                addDependency(AwsRuntimeType.awsSigAuthEventStream(runtimeConfig).toSymbol())
+            }
+            rust(
                 """
                 /// The signature version 4 service signing name to use in the credential scope when signing requests.
                 ///
@@ -97,24 +97,7 @@ class SigV4SigningConfig(
                     ${sigV4Trait.name.dq()}
                 }
                 """,
-                *codegenScope,
             )
-            if (serviceHasEventStream) {
-                rustTemplate(
-                    "#{signerFn:W}",
-                    "signerFn" to
-                        renderEventStreamSignerFn { propertiesName ->
-                            writable {
-                                rustTemplate(
-                                    """
-                                    #{SigV4Signer}::new($propertiesName)
-                                    """,
-                                    *codegenScope,
-                                )
-                            }
-                        },
-                )
-            }
         }
     }
 }
@@ -209,5 +192,3 @@ class SigV4SigningFeature(
         }
     }
 }
-
-fun RuntimeConfig.sigAuth() = awsRuntimeCrate("aws-sig-auth")
diff --git a/.../rustsdk/SigV4SigningCustomizationTest.kt → ...ithy/rustsdk/SigV4SigningDecoratorTest.kt b/.../rustsdk/SigV4SigningCustomizationTest.kt → ...ithy/rustsdk/SigV4SigningDecoratorTest.kt
@@ -12,7 +12,7 @@ import software.amazon.smithy.rust.codegen.core.testutil.TestWorkspace
 import software.amazon.smithy.rust.codegen.core.testutil.compileAndTest
 import software.amazon.smithy.rust.codegen.core.testutil.unitTest
 
-internal class SigV4SigningCustomizationTest {
+internal class SigV4SigningDecoratorTest {
     @Test
     fun `generates a valid config`() {
         val project = stubConfigProject(