Skip to content

Commit cf1e783

Browse files
authored
Add support for Unframing SigV4 signed messages for Servers (#4356)
# SigV4 Event Stream Support for Server SDK ## Problem Clients wrap event stream messages in SigV4 envelopes with signature headers (`:chunk-signature`, `:date`), but servers couldn't parse these signed messages because they expected the raw event shape, not the envelope. ## Solution Added server-side SigV4 event stream unsigning support that automatically extracts inner messages from signed envelopes while maintaining compatibility with unsigned messages. ## Implementation ### Type System Changes Event stream types are wrapped to handle both signed and unsigned messages: - `Receiver<Events, Error>` → `Receiver<SignedEvent<Events>, SignedEventError<Error>>` - `SignedEvent<T>` provides access to both the inner message and signature information - `SignedEventError<E>` wraps both extraction errors and underlying event errors ### Runtime Components **SigV4Unmarshaller**: Wraps the base event stream unmarshaller to handle SigV4 extraction: ```rust impl<T: UnmarshallMessage> UnmarshallMessage for SigV4Unmarshaller<T> { type Output = SignedEvent<T::Output>; type Error = SignedEventError<T::Error>; fn unmarshall(&self, message: &Message) -> Result<...> { match extract_signed_message(message) { Ok(Signed { message: inner, signature }) => { // Process inner message with base unmarshaller self.inner.unmarshall(&inner).map(|event| SignedEvent { message: event, signature: Some(signature), }) } Ok(Unsigned) => { // Process unsigned message directly self.inner.unmarshall(message).map(|event| SignedEvent { message: event, signature: None, }) } Err(err) => Ok(SignedEventError::InvalidSignedEvent(err)) } } } ``` ### Code Generation Integration **SigV4EventStreamDecorator**: - Detects services with `@sigv4` trait and event streams - Wraps event stream types using `SigV4EventStreamSymbolProvider` - Generates support structures (`SignedEvent`, `SigV4Unmarshaller`, etc.) - Injects unmarshaller wrapping via HTTP binding customizations **HTTP Binding Customization**: - Added `BeforeCreatingEventStreamReceiver` section to `HttpBindingGenerator` - Allows decorators to wrap the unmarshaller before `Receiver` creation - Generates: `let unmarshaller = SigV4Unmarshaller::new(unmarshaller);` ### Usage Server handlers receive `SignedEvent<T>` and extract the inner message: ```rust async fn streaming_operation_handler(input: StreamingOperationInput) -> Result<...> { let event = input.events.recv().await?; if let Some(signed_event) = event { let actual_event = &signed_event.message; // Extract inner message let signature_info = &signed_event.signature; // Access signature if present // Process actual_event... } } ``` ## Testing - Added `test_sigv4_signed_event_stream` that sends SigV4-wrapped events - Verifies both signed and unsigned messages work correctly - All existing event stream tests continue to pass ## Architecture Benefits - **Backward Compatible**: Unsigned messages work unchanged - **Type Safe**: Compile-time guarantees about message structure - **Extensible**: Pattern can be applied to other authentication schemes - **Minimal Impact**: Only affects services with `@sigv4` trait and event streams ## Checklist <!--- If a checkbox below is not applicable, then please DELETE it rather than leaving it unchecked --> - [ ] For changes to the smithy-rs codegen or runtime crates, I have created a changelog entry Markdown file in the `.changelog` directory, specifying "client," "server," or both in the `applies_to` key. - [ ] For changes to the AWS SDK, generated SDK code, or SDK runtime crates, I have created a changelog entry Markdown file in the `.changelog` directory, specifying "aws-sdk-rust" in the `applies_to` key. ---- _By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice._
1 parent 1de7020 commit cf1e783

File tree

21 files changed

+750
-76
lines changed

21 files changed

+750
-76
lines changed

.changelog/1761080224.md

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
---
2+
applies_to:
3+
- server
4+
authors:
5+
- rcoh
6+
references: ["smithy-rs#4356"]
7+
breaking: true
8+
new_feature: true
9+
bug_fix: false
10+
---
11+
Parse EventStream signed-frames for servers marked with `@sigv4`.
12+
13+
This is a breaking change, because events from SigV4 services are wrapped in a SignedEvent frame.

AGENTS.md

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,26 @@ operation MyOperation {
4848
- **`codegen-core/common-test-models/constraints.smithy`** - Constraint validation tests with restJson1
4949
- **`codegen-client-test/model/main.smithy`** - awsJson1_1 protocol tests
5050

51+
### httpQueryParams Bug Investigation
52+
53+
When investigating the `@httpQueryParams` bug (where query parameters weren't appearing in requests), the issue was in `RequestBindingGenerator.kt` line 173. The bug occurred when:
54+
55+
1. An operation had ONLY `@httpQueryParams` (no regular `@httpQuery` parameters)
56+
2. The condition `if (dynamicParams.isEmpty() && literalParams.isEmpty() && mapParams.isEmpty())` would skip generating the `uri_query` function
57+
58+
The fix was to ensure `mapParams.isEmpty()` was included in the condition check. The current implementation correctly generates query parameters for `@httpQueryParams` even when no other query parameters exist.
59+
60+
**Testing httpQueryParams**: Create operations with only `@httpQueryParams` to ensure they generate proper query strings in requests.
61+
62+
## rustTemplate Formatting
63+
64+
**CRITICAL**: Because `#` is the formatting character in `rustTemplate`, Rust attributes must be escaped:
65+
66+
❌ Wrong: `#[derive(Debug)]`
67+
✅ Correct: `##[derive(Debug)]`
68+
69+
This applies to ALL Rust attributes: `##[non_exhaustive]`, `##[derive(...)]`, `##[cfg(...)]`, etc.
70+
5171
## preludeScope: Rust Prelude Types
5272

5373
**Always use `preludeScope` for Rust prelude types:**

codegen-core/common-test-models/rpcv2Cbor-extras.smithy

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,10 @@ use smithy.framework#ValidationException
66
use smithy.protocols#rpcv2Cbor
77
use smithy.test#httpResponseTests
88
use smithy.test#httpMalformedRequestTests
9+
use aws.auth#sigv4
910

1011
@rpcv2Cbor
12+
@sigv4(name: "rpcv2-cbor")
1113
service RpcV2CborService {
1214
operations: [
1315
SimpleStructOperation

codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/SymbolExt.kt

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -83,9 +83,14 @@ fun Symbol.makeMaybeConstrained(): Symbol =
8383
* WARNING: This function does not update any symbol references (e.g., `symbol.addReference()`) on the
8484
* returned symbol. You will have to add those yourself if your logic relies on them.
8585
**/
86-
fun Symbol.mapRustType(f: (RustType) -> RustType): Symbol {
86+
fun Symbol.mapRustType(
87+
vararg dependencies: RuntimeType,
88+
f: (RustType) -> RustType,
89+
): Symbol {
8790
val newType = f(this.rustType())
88-
return Symbol.builder().rustType(newType)
91+
val builder = this.toBuilder()
92+
dependencies.forEach { builder.addReference(it.toSymbol()) }
93+
return builder.rustType(newType)
8994
.name(newType.name)
9095
.build()
9196
}

codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/generators/http/HttpBindingGenerator.kt

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,12 @@ sealed class HttpBindingSection(name: String) : Section(name) {
9696

9797
data class AfterDeserializingIntoADateTimeOfHttpHeaders(val memberShape: MemberShape) :
9898
HttpBindingSection("AfterDeserializingIntoADateTimeOfHttpHeaders")
99+
100+
data class BeforeCreatingEventStreamReceiver(
101+
val operationShape: OperationShape,
102+
val unionShape: UnionShape,
103+
val unmarshallerVariableName: String,
104+
) : HttpBindingSection("BeforeCreatingEventStreamReceiver")
99105
}
100106

101107
typealias HttpBindingCustomization = NamedCustomization<HttpBindingSection>
@@ -272,11 +278,27 @@ class HttpBindingGenerator(
272278
rustTemplate(
273279
"""
274280
let unmarshaller = #{unmarshallerConstructorFn}();
281+
""",
282+
"unmarshallerConstructorFn" to unmarshallerConstructorFn,
283+
)
284+
285+
// Allow customizations to wrap the unmarshaller
286+
for (customization in customizations) {
287+
customization.section(
288+
HttpBindingSection.BeforeCreatingEventStreamReceiver(
289+
operationShape,
290+
targetShape,
291+
"unmarshaller",
292+
),
293+
)(this)
294+
}
295+
296+
rustTemplate(
297+
"""
275298
let body = std::mem::replace(body, #{SdkBody}::taken());
276299
Ok(#{receiver:W})
277300
""",
278301
"SdkBody" to RuntimeType.sdkBody(runtimeConfig),
279-
"unmarshallerConstructorFn" to unmarshallerConstructorFn,
280302
"receiver" to
281303
writable {
282304
if (codegenTarget == CodegenTarget.SERVER) {

codegen-server-test/integration-tests/Cargo.lock

Lines changed: 9 additions & 9 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

codegen-server-test/integration-tests/eventstreams/tests/structured_eventstream_tests.rs

Lines changed: 60 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -143,13 +143,15 @@ async fn streaming_operation_handler(
143143
state.lock().unwrap().streaming_operation.num_calls += 1;
144144
let ev = input.events.recv().await;
145145

146-
if let Ok(Some(event)) = &ev {
146+
if let Ok(Some(signed_event)) = &ev {
147+
// Extract the actual event from the SignedEvent wrapper
148+
let actual_event = &signed_event.message;
147149
state
148150
.lock()
149151
.unwrap()
150152
.streaming_operation
151153
.events
152-
.push(event.clone());
154+
.push(actual_event.clone());
153155
}
154156

155157
Ok(output::StreamingOperationOutput::builder()
@@ -174,13 +176,15 @@ async fn streaming_operation_with_initial_data_handler(
174176

175177
let ev = input.events.recv().await;
176178

177-
if let Ok(Some(event)) = &ev {
179+
if let Ok(Some(signed_event)) = &ev {
180+
// Extract the actual event from the SignedEvent wrapper
181+
let actual_event = &signed_event.message;
178182
state
179183
.lock()
180184
.unwrap()
181185
.streaming_operation_with_initial_data
182186
.events
183-
.push(event.clone());
187+
.push(actual_event.clone());
184188
}
185189

186190
Ok(output::StreamingOperationWithInitialDataOutput::builder()
@@ -229,7 +233,7 @@ async fn streaming_operation_with_optional_data_handler(
229233
.unwrap()
230234
.streaming_operation_with_optional_data
231235
.events
232-
.push(event.clone());
236+
.push(event.message.clone());
233237
}
234238

235239
Ok(output::StreamingOperationWithOptionalDataOutput::builder()
@@ -348,6 +352,39 @@ fn build_event(event_type: &str) -> Message {
348352
Message::new_from_parts(headers, empty_cbor)
349353
}
350354

355+
fn build_sigv4_signed_event(event_type: &str) -> Message {
356+
use aws_smithy_eventstream::frame::write_message_to;
357+
use std::time::{SystemTime, UNIX_EPOCH};
358+
359+
// Build the inner event message
360+
let inner_event = build_event(event_type);
361+
362+
// Serialize the inner message to bytes
363+
let mut inner_bytes = Vec::new();
364+
write_message_to(&inner_event, &mut inner_bytes).unwrap();
365+
366+
// Create the SigV4 envelope with signature headers
367+
let timestamp = SystemTime::now()
368+
.duration_since(UNIX_EPOCH)
369+
.unwrap()
370+
.as_secs();
371+
372+
let headers = vec![
373+
Header::new(
374+
":chunk-signature",
375+
HeaderValue::ByteArray(Bytes::from(
376+
"example298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
377+
)),
378+
),
379+
Header::new(
380+
":date",
381+
HeaderValue::Timestamp(aws_smithy_types::DateTime::from_secs(timestamp as i64)),
382+
),
383+
];
384+
385+
Message::new_from_parts(headers, Bytes::from(inner_bytes))
386+
}
387+
351388
fn get_event_type(msg: &Message) -> &str {
352389
msg.headers()
353390
.iter()
@@ -439,6 +476,24 @@ async fn test_streaming_operation_with_initial_data_missing() {
439476
);
440477
}
441478

479+
/// Test that the server can handle SigV4 signed event stream messages.
480+
/// The client wraps the actual event in a SigV4 envelope with signature headers.
481+
#[tokio::test]
482+
async fn test_sigv4_signed_event_stream() {
483+
let mut harness = TestHarness::new("StreamingOperation").await;
484+
485+
// Send a SigV4 signed event - the inner message is wrapped in an envelope
486+
let signed_event = build_sigv4_signed_event("A");
487+
harness.client.send(signed_event).await.unwrap();
488+
489+
let resp = harness.expect_message().await;
490+
assert_eq!(get_event_type(&resp), "A");
491+
assert_eq!(
492+
harness.server.streaming_operation_events(),
493+
vec![Events::A(Event {})]
494+
);
495+
}
496+
442497
/// Test that when alwaysSendEventStreamInitialResponse is disabled, no initial-response is sent
443498
#[tokio::test]
444499
async fn test_server_no_initial_response_when_disabled() {

codegen-server/codegen-server-typescript/src/main/kotlin/software/amazon/smithy/rust/codegen/server/typescript/smithy/TsServerCodegenVisitor.kt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,7 @@ class TsServerCodegenVisitor(
6262
ServerProtocolLoader(
6363
codegenDecorator.protocols(
6464
service.id,
65-
ServerProtocolLoader.DefaultProtocols,
65+
ServerProtocolLoader.defaultProtocols(),
6666
),
6767
)
6868
.protocolFor(context.model, service)

codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/RustServerCodegenPlugin.kt

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import software.amazon.smithy.rust.codegen.core.smithy.StreamingShapeSymbolProvi
1919
import software.amazon.smithy.rust.codegen.core.smithy.SymbolVisitor
2020
import software.amazon.smithy.rust.codegen.server.smithy.customizations.CustomValidationExceptionWithReasonDecorator
2121
import software.amazon.smithy.rust.codegen.server.smithy.customizations.ServerRequiredCustomizations
22+
import software.amazon.smithy.rust.codegen.server.smithy.customizations.SigV4EventStreamDecorator
2223
import software.amazon.smithy.rust.codegen.server.smithy.customizations.SmithyValidationExceptionDecorator
2324
import software.amazon.smithy.rust.codegen.server.smithy.customizations.UserProvidedValidationExceptionDecorator
2425
import software.amazon.smithy.rust.codegen.server.smithy.customize.CombinedServerCodegenDecorator
@@ -54,6 +55,7 @@ class RustServerCodegenPlugin : ServerDecoratableBuildPlugin() {
5455
UserProvidedValidationExceptionDecorator(),
5556
SmithyValidationExceptionDecorator(),
5657
CustomValidationExceptionWithReasonDecorator(),
58+
SigV4EventStreamDecorator(),
5759
*decorator,
5860
)
5961
logger.info("Loaded plugin to generate pure Rust bindings for the server SDK")

codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/ServerCodegenVisitor.kt

Lines changed: 14 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -124,18 +124,7 @@ open class ServerCodegenVisitor(
124124

125125
val baseModel = baselineTransform(context.model)
126126
val service = settings.getService(baseModel)
127-
val (protocolShape, protocolGeneratorFactory) =
128-
ServerProtocolLoader(
129-
codegenDecorator.protocols(
130-
service.id,
131-
ServerProtocolLoader.DefaultProtocols,
132-
),
133-
)
134-
.protocolFor(context.model, service)
135-
this.protocolGeneratorFactory = protocolGeneratorFactory
136-
137127
model = codegenDecorator.transformModel(service, baseModel, settings)
138-
139128
val serverSymbolProviders =
140129
ServerSymbolProviders.from(
141130
settings,
@@ -146,7 +135,19 @@ open class ServerCodegenVisitor(
146135
codegenDecorator,
147136
RustServerCodegenPlugin::baseSymbolProvider,
148137
)
149-
138+
val (protocolShape, protocolGeneratorFactory) =
139+
ServerProtocolLoader(
140+
codegenDecorator.protocols(
141+
service.id,
142+
ServerProtocolLoader.defaultProtocols { it ->
143+
codegenDecorator.httpCustomizations(
144+
serverSymbolProviders.symbolProvider,
145+
it,
146+
)
147+
},
148+
),
149+
)
150+
.protocolFor(context.model, service)
150151
codegenContext =
151152
ServerCodegenContext(
152153
model,
@@ -160,6 +161,7 @@ open class ServerCodegenVisitor(
160161
serverSymbolProviders.constraintViolationSymbolProvider,
161162
serverSymbolProviders.pubCrateConstrainedShapeSymbolProvider,
162163
)
164+
this.protocolGeneratorFactory = protocolGeneratorFactory
163165

164166
// We can use a not-null assertion because [CombinedServerCodegenDecorator] returns a not null value.
165167
validationExceptionConversionGenerator = codegenDecorator.validationExceptionConversion(codegenContext)!!

0 commit comments

Comments
 (0)