From fb721c152967e7e5c447cbd327cb157d3dbc9d6d Mon Sep 17 00:00:00 2001 From: Aaron Todd Date: Mon, 24 Mar 2025 13:56:15 -0400 Subject: [PATCH] fix event stream signing to include resource operations (#4055) ## Motivation and Context fixes https://github.com/smithy-lang/smithy-rs/issues/4054 ## Description `serviceShape.operations` only includes operations directly bound to the service and not operations bound via resources. This caused a bug with event stream signing because the operation came from a resource and didn't trigger the codegen path to enable the required feature flag in the runtime (see the ticket for linked code). This PR address the original issue and adds a test. I also searched for other places this may be happening as it's a common bug in Smithy codegen. ## Checklist - [x] 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. ---- _By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice._ --- .changelog/1742827344.md | 11 +++ .../smithy/rustsdk/SigV4AuthDecorator.kt | 2 +- .../codegen/client/smithy/ClientRustModule.kt | 6 +- .../generators/client/FluentClientDocs.kt | 13 ++- .../customizations/SmithyTypesPubUseExtra.kt | 12 ++- .../smithy/rust/codegen/core/util/Smithy.kt | 9 +- .../rust/codegen/core/util/ExtensionsTest.kt | 85 +++++++++++++++++++ 7 files changed, 122 insertions(+), 16 deletions(-) create mode 100644 .changelog/1742827344.md create mode 100644 codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/util/ExtensionsTest.kt diff --git a/.changelog/1742827344.md b/.changelog/1742827344.md new file mode 100644 index 000000000..cdb21004c --- /dev/null +++ b/.changelog/1742827344.md @@ -0,0 +1,11 @@ +--- +applies_to: +- client +authors: ["aajtodd"] +references: +- smithy-rs#4054 +breaking: false +new_feature: false +bug_fix: true +--- +Fix traversal of operations bound to resources in several places including logic to determine if an event stream exists diff --git a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/SigV4AuthDecorator.kt b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/SigV4AuthDecorator.kt index 6153123c9..e661377dc 100644 --- a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/SigV4AuthDecorator.kt +++ b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/SigV4AuthDecorator.kt @@ -205,7 +205,7 @@ private class AuthServiceRuntimePluginCustomization(private val codegenContext: val serviceHasEventStream = codegenContext.serviceShape.hasEventStreamOperations(codegenContext.model) if (serviceHasEventStream) { - // enable the aws-runtime `sign-eventstream` feature + // enable the aws-runtime `event-stream` feature addDependency( AwsCargoDependency.awsRuntime(runtimeConfig).withFeature("event-stream").toType() .toSymbol(), diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/ClientRustModule.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/ClientRustModule.kt index 64197ddc7..4dc936983 100644 --- a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/ClientRustModule.kt +++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/ClientRustModule.kt @@ -7,6 +7,7 @@ package software.amazon.smithy.rust.codegen.client.smithy import software.amazon.smithy.codegen.core.Symbol import software.amazon.smithy.model.Model +import software.amazon.smithy.model.knowledge.TopDownIndex import software.amazon.smithy.model.shapes.EnumShape import software.amazon.smithy.model.shapes.OperationShape import software.amazon.smithy.model.shapes.Shape @@ -156,11 +157,12 @@ class ClientModuleDocProvider( private fun customizeModuleDoc(): Writable = writable { val model = codegenContext.model + val operations = TopDownIndex.of(model).getContainedOperations(codegenContext.serviceShape) docs("Operation customization and supporting types.\n") - if (codegenContext.serviceShape.operations.isNotEmpty()) { + if (operations.isNotEmpty()) { val opFnName = FluentClientGenerator.clientOperationFnName( - codegenContext.serviceShape.operations.minOf { it } + operations.minOf { it.id } .let { model.expectShape(it, OperationShape::class.java) }, codegenContext.symbolProvider, ) diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/FluentClientDocs.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/FluentClientDocs.kt index 7570a7ce3..b1ac9feed 100644 --- a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/FluentClientDocs.kt +++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/FluentClientDocs.kt @@ -6,7 +6,6 @@ package software.amazon.smithy.rust.codegen.client.smithy.generators.client import software.amazon.smithy.model.knowledge.TopDownIndex -import software.amazon.smithy.model.shapes.OperationShape import software.amazon.smithy.model.shapes.StringShape import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext import software.amazon.smithy.rust.codegen.core.rustlang.docs @@ -59,17 +58,17 @@ object FluentClientDocs { writable { val model = codegenContext.model val symbolProvider = codegenContext.symbolProvider - if (model.operationShapes.isNotEmpty()) { + val operations = TopDownIndex.of(model).getContainedOperations(codegenContext.serviceShape) + if (operations.isNotEmpty()) { // Find an operation with a simple string member shape val (operation, member) = - codegenContext.serviceShape.operations - .map { id -> - val operationShape = model.expectShape(id, OperationShape::class.java) + operations + .map { op -> val member = - operationShape.inputShape(model) + op.inputShape(model) .members() .firstOrNull { model.expectShape(it.target) is StringShape } - operationShape to member + op to member } .sortedBy { it.first.id } .firstOrNull { (_, member) -> member != null } ?: (null to null) diff --git a/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/customizations/SmithyTypesPubUseExtra.kt b/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/customizations/SmithyTypesPubUseExtra.kt index aa1e1d40a..36e05d165 100644 --- a/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/customizations/SmithyTypesPubUseExtra.kt +++ b/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/customizations/SmithyTypesPubUseExtra.kt @@ -6,6 +6,8 @@ package software.amazon.smithy.rust.codegen.core.smithy.customizations import software.amazon.smithy.model.Model +import software.amazon.smithy.model.knowledge.TopDownIndex +import software.amazon.smithy.model.shapes.ServiceShape import software.amazon.smithy.model.shapes.Shape import software.amazon.smithy.model.shapes.StructureShape import software.amazon.smithy.rust.codegen.core.rustlang.Feature @@ -21,8 +23,12 @@ import software.amazon.smithy.rust.codegen.core.util.hasEventStreamOperations import software.amazon.smithy.rust.codegen.core.util.hasStreamingMember /** Returns true if the model has normal streaming operations (excluding event streams) */ -private fun hasStreamingOperations(model: Model): Boolean { - return model.operationShapes.any { operation -> +private fun hasStreamingOperations( + model: Model, + serviceShape: ServiceShape, +): Boolean { + val operations = TopDownIndex.of(model).getContainedOperations(serviceShape) + return operations.any { operation -> val input = model.expectShape(operation.inputShape, StructureShape::class.java) val output = model.expectShape(operation.outputShape, StructureShape::class.java) (input.hasStreamingMember(model) && !input.hasEventStreamMember(model)) || @@ -69,7 +75,7 @@ fun pubUseSmithyPrimitives( "Format" to RuntimeType.format(rc), ) } - if (hasStreamingOperations(model)) { + if (hasStreamingOperations(model, codegenContext.serviceShape)) { rustCrate.mergeFeature( Feature( "rt-tokio", diff --git a/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/util/Smithy.kt b/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/util/Smithy.kt index 32a6874fb..4616829de 100644 --- a/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/util/Smithy.kt +++ b/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/util/Smithy.kt @@ -8,6 +8,7 @@ package software.amazon.smithy.rust.codegen.core.util import software.amazon.smithy.aws.traits.ServiceTrait import software.amazon.smithy.codegen.core.CodegenException import software.amazon.smithy.model.Model +import software.amazon.smithy.model.knowledge.TopDownIndex import software.amazon.smithy.model.shapes.BooleanShape import software.amazon.smithy.model.shapes.ListShape import software.amazon.smithy.model.shapes.MapShape @@ -81,9 +82,11 @@ fun OperationShape.isOutputEventStream(model: Model): Boolean = fun OperationShape.isEventStream(model: Model): Boolean = isInputEventStream(model) || isOutputEventStream(model) fun ServiceShape.hasEventStreamOperations(model: Model): Boolean = - operations.any { id -> - model.expectShape(id, OperationShape::class.java).isEventStream(model) - } + TopDownIndex.of(model) + .getContainedOperations(this) + .any { op -> + op.isEventStream(model) + } fun Shape.shouldRedact(model: Model): Boolean = when { diff --git a/codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/util/ExtensionsTest.kt b/codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/util/ExtensionsTest.kt new file mode 100644 index 000000000..7c438c4ad --- /dev/null +++ b/codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/util/ExtensionsTest.kt @@ -0,0 +1,85 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package software.amazon.smithy.rust.codegen.core.util + +import io.kotest.matchers.shouldBe +import org.junit.jupiter.api.Test +import software.amazon.smithy.model.shapes.ServiceShape +import software.amazon.smithy.model.shapes.ShapeId +import software.amazon.smithy.rust.codegen.core.testutil.asSmithyModel + +class ExtensionsTest { + @Test + fun `it should find event streams on normal operations`() { + val model = + """ + namespace test + service TestService { + operations: [ EventStreamOp ] + } + + operation EventStreamOp { + input := { + events: Events + } + } + + @streaming + union Events { + foo: Foo + bar: Bar, + } + + structure Foo { + foo: String + } + + structure Bar { + bar: Long + } + """.asSmithyModel(smithyVersion = "2.0") + + val service = model.expectShape(ShapeId.from("test#TestService"), ServiceShape::class.java) + service.hasEventStreamOperations(model) shouldBe true + } + + @Test + fun `it should find event streams on resource operations`() { + val model = + """ + namespace test + service TestService { + resources: [ TestResource ] + } + + resource TestResource { + operations: [ EventStreamOp ] + } + + operation EventStreamOp { + input := { + events: Events + } + } + + @streaming + union Events { + foo: Foo + bar: Bar, + } + + structure Foo { + foo: String + } + + structure Bar { + bar: Long + } + """.asSmithyModel(smithyVersion = "2.0") + + val service = model.expectShape(ShapeId.from("test#TestService"), ServiceShape::class.java) + service.hasEventStreamOperations(model) shouldBe true + } +} -- GitLab