diff --git a/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/customizations/SigV4EventStreamDecorator.kt b/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/customizations/SigV4EventStreamDecorator.kt index 0cf966749e..f8be36aea9 100644 --- a/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/customizations/SigV4EventStreamDecorator.kt +++ b/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/customizations/SigV4EventStreamDecorator.kt @@ -5,6 +5,7 @@ package software.amazon.smithy.rust.codegen.server.smithy.customizations +import software.amazon.smithy.aws.traits.auth.SigV4ATrait import software.amazon.smithy.aws.traits.auth.SigV4Trait import software.amazon.smithy.codegen.core.Symbol import software.amazon.smithy.model.knowledge.ServiceIndex @@ -50,8 +51,10 @@ class SigV4EventStreamDecorator : ServerCodegenDecorator { } } -internal fun RustSymbolProvider.usesSigAuth(): Boolean = - ServiceIndex.of(model).getAuthSchemes(moduleProviderContext.serviceShape!!).containsKey(SigV4Trait.ID) +internal fun RustSymbolProvider.usesSigAuth(): Boolean { + val authSchemes = ServiceIndex.of(model).getAuthSchemes(moduleProviderContext.serviceShape!!) + return authSchemes.containsKey(SigV4Trait.ID) || authSchemes.containsKey(SigV4ATrait.ID) +} // Goes from `T` to `SignedEvent` fun wrapInSignedEvent( diff --git a/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/customizations/SigV4EventStreamSupportStructures.kt b/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/customizations/SigV4EventStreamSupportStructures.kt index 81f67f1dd4..781f16752f 100644 --- a/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/customizations/SigV4EventStreamSupportStructures.kt +++ b/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/customizations/SigV4EventStreamSupportStructures.kt @@ -12,6 +12,7 @@ import software.amazon.smithy.rust.codegen.core.rustlang.RustType import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate 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.RuntimeType.Companion.preludeScope import software.amazon.smithy.rust.codegen.core.smithy.mapRustType import software.amazon.smithy.rust.codegen.core.smithy.rustType import software.amazon.smithy.rust.codegen.core.util.PANIC @@ -77,11 +78,12 @@ object SigV4EventStreamSupportStructures { ##[derive(Debug, Clone)] pub struct SignatureInfo { /// The chunk signature bytes from the `:chunk-signature` header - pub chunk_signature: Vec, + pub chunk_signature: #{Vec}, /// The timestamp from the `:date` header pub timestamp: #{SystemTime}, } """, + *preludeScope, "SystemTime" to RuntimeType.std.resolve("time::SystemTime"), ) } @@ -116,18 +118,25 @@ object SigV4EventStreamSupportStructures { ##[derive(Debug)] pub enum SignedEventError { /// Error from the underlying event stream - Event(E), + Event { + /// The error from the underlying event stream + error: E, + /// Signature information if the message was signed + signature: #{Option}<#{SignatureInfo}>, + }, /// Error extracting signed message InvalidSignedEvent(#{ExtractionError}), } impl From for SignedEventError { fn from(err: E) -> Self { - SignedEventError::Event(err) + SignedEventError::Event { error: err, signature: None } } } """, "ExtractionError" to extractionError(runtimeConfig), + "SignatureInfo" to signatureInfo(), + *preludeScope, ) } @@ -182,11 +191,14 @@ object SigV4EventStreamSupportStructures { #{UnmarshalledMessage}::Event(event) => { Ok(#{UnmarshalledMessage}::Event(#{SignedEvent} { message: event, - signature: Some(signature), + signature: Some(signature.clone()), })) } #{UnmarshalledMessage}::Error(err) => { - Ok(#{UnmarshalledMessage}::Error(#{SignedEventError}::Event(err))) + Ok(#{UnmarshalledMessage}::Error(#{SignedEventError}::Event { + error: err, + signature: Some(signature), + })) } }, Err(err) => Err(err), @@ -203,7 +215,10 @@ object SigV4EventStreamSupportStructures { })) } #{UnmarshalledMessage}::Error(err) => { - Ok(#{UnmarshalledMessage}::Error(#{SignedEventError}::Event(err))) + Ok(#{UnmarshalledMessage}::Error(#{SignedEventError}::Event { + error: err, + signature: None, + })) } }, Err(err) => Err(err),