Sensitive trait (#229)

* Add support for the Sensitive Trait

* Add kms integration test for sensitive trait

* Add additional test

* Always generate a custom debug impl

This actually causes a reduction in llvm-lines & apparently can improve compile performance. It also simplifies the code.

.gitignore

@@ -43,3 +43,6 @@ gradle-app.setting
 
 # MacOS
 .DS_Store
+
+# Rust build artifacts
+target/

aws/sdk/build.gradle.kts

@@ -105,6 +105,11 @@ task("relocateServices") {
                 from("$buildDir/smithyprojections/sdk/${it.module}/rust-codegen")
                 into(sdkOutputDir.resolve(it.module))
             }
+
+            copy {
+                from(projectDir.resolve("integration-tests/${it.module}/tests"))
+                into(sdkOutputDir.resolve(it.module).resolve("tests"))
+            }
         }
     }
     outputs.upToDateWhen { false }

aws/sdk/integration-tests/.gitignore

@@ -0,0 +1,2 @@
+Cargo.lock
+target/

aws/sdk/integration-tests/README.md

@@ -0,0 +1,8 @@
+# Handwritten Integration Test Root
+
+This folder contains hand-written integration tests that are specific to individual services. In order for your test to be merged into the final artifact:
+
+- The crate name must match the generated crate name, eg. `kms`, `dynamodb`
+- Your test must be placed into the `tests` folder. **Everything else in your test crate is ignored.**
+
+The contents of the `test` folder will be combined with codegenerated integration tests & inserted into the `tests` folder of the final generated service crate.

aws/sdk/integration-tests/kms/Cargo.toml

@@ -0,0 +1,11 @@
+# This Cargo.toml is unused in generated code. It exists solely to enable these tests to compile in-situ
+[package]
+name = "kms-tests"
+version = "0.1.0"
+authors = ["Russell Cohen <[email protected]>"]
+edition = "2018"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+kms = { path = "../../build/aws-sdk/kms" }

aws/sdk/integration-tests/kms/tests/sensitive-it.rs

@@ -0,0 +1,12 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0.
+ */
+
+use kms::output::GenerateRandomOutput;
+use kms::Blob;
+#[test]
+fn validate_sensitive_trait() {
+    let output = GenerateRandomOutput::builder().plaintext(Blob::new("some output")).build();
+    assert_eq!(format!("{:?}", output), "GenerateRandomOutput { plaintext: \"*** Sensitive Data Redacted ***\" }");
+}

codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/smithy/CodegenVisitor.kt

@@ -25,6 +25,7 @@ import software.amazon.smithy.rust.codegen.smithy.generators.ModelBuilderGenerat
 import software.amazon.smithy.rust.codegen.smithy.generators.ProtocolConfig
 import software.amazon.smithy.rust.codegen.smithy.generators.ProtocolGeneratorFactory
 import software.amazon.smithy.rust.codegen.smithy.generators.ServiceGenerator
+import software.amazon.smithy.rust.codegen.smithy.generators.SmithyTypesPubUseGenerator
 import software.amazon.smithy.rust.codegen.smithy.generators.StructureGenerator
 import software.amazon.smithy.rust.codegen.smithy.generators.UnionGenerator
 import software.amazon.smithy.rust.codegen.smithy.generators.implBlock
@@ -35,7 +36,8 @@ import software.amazon.smithy.rust.codegen.util.CommandFailed
 import software.amazon.smithy.rust.codegen.util.runCommand
 import java.util.logging.Logger
 
-class CodegenVisitor(context: PluginContext, private val codegenDecorator: RustCodegenDecorator) : ShapeVisitor.Default<Unit>() {
+class CodegenVisitor(context: PluginContext, private val codegenDecorator: RustCodegenDecorator) :
+    ShapeVisitor.Default<Unit>() {
 
     private val logger = Logger.getLogger(javaClass.name)
     private val settings = RustSettings.from(context.model, context.settings)
@@ -53,13 +55,19 @@ class CodegenVisitor(context: PluginContext, private val codegenDecorator: RustC
             SymbolVisitorConfig(runtimeConfig = settings.runtimeConfig, codegenConfig = settings.codegenConfig)
         val baseModel = baselineTransform(context.model)
         val service = settings.getService(baseModel)
-        val (protocol, generator) = ProtocolLoader(codegenDecorator.protocols(service.id, ProtocolLoader.DefaultProtocols)).protocolFor(context.model, service)
+        val (protocol, generator) = ProtocolLoader(
+            codegenDecorator.protocols(
+                service.id,
+                ProtocolLoader.DefaultProtocols
+            )
+        ).protocolFor(context.model, service)
         protocolGenerator = generator
         model = generator.transformModel(baseModel)
         val baseProvider = RustCodegenPlugin.BaseSymbolProvider(model, symbolVisitorConfig)
         symbolProvider = codegenDecorator.symbolProvider(generator.symbolProvider(model, baseProvider))
 
-        protocolConfig = ProtocolConfig(model, symbolProvider, settings.runtimeConfig, service, protocol, settings.moduleName)
+        protocolConfig =
+            ProtocolConfig(model, symbolProvider, settings.runtimeConfig, service, protocol, settings.moduleName)
         writers = CodegenWriterDelegator(
             context.fileManifest,
             symbolProvider,
@@ -76,7 +84,13 @@ class CodegenVisitor(context: PluginContext, private val codegenDecorator: RustC
         val serviceShapes = Walker(model).walkShapes(service)
         serviceShapes.forEach { it.accept(this) }
         // TODO: if we end up with a lot of these on-by-default customizations, we may want to refactor them somewhere
-        writers.finalize(settings, codegenDecorator.libRsCustomizations(protocolConfig, listOf(CrateVersionGenerator())))
+        writers.finalize(
+            settings,
+            codegenDecorator.libRsCustomizations(
+                protocolConfig,
+                listOf(CrateVersionGenerator(), SmithyTypesPubUseGenerator(protocolConfig.runtimeConfig))
+            )
+        )
         try {
             "cargo fmt".runCommand(fileManifest.baseDir)
         } catch (_: CommandFailed) {

codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/smithy/RuntimeTypes.kt

@@ -59,12 +59,15 @@ data class RuntimeType(val name: String?, val dependency: RustDependency?, val n
         // val Blob = RuntimeType("Blob", RustDependency.IO_CORE, "blob")
         val From = RuntimeType("From", dependency = null, namespace = "std::convert")
         val AsRef = RuntimeType("AsRef", dependency = null, namespace = "std::convert")
-        fun StdFmt(member: String) = RuntimeType("fmt::$member", dependency = null, namespace = "std")
+        fun StdFmt(member: String?) = RuntimeType(member, dependency = null, namespace = "std::fmt")
         fun Std(member: String) = RuntimeType(member, dependency = null, namespace = "std")
         val StdError = RuntimeType("Error", dependency = null, namespace = "std::error")
         val HashSet = RuntimeType(RustType.SetType, dependency = null, namespace = "std::collections")
         val HashMap = RuntimeType("HashMap", dependency = null, namespace = "std::collections")
         val ByteSlab = RuntimeType("Vec<u8>", dependency = null, namespace = "std::vec")
+        val Debug = StdFmt("Debug")
+        val PartialEq = Std("cmp::PartialEq")
+        val Clone = Std("clone::Clone")
 
         fun Instant(runtimeConfig: RuntimeConfig) =
             RuntimeType("Instant", CargoDependency.SmithyTypes(runtimeConfig), "${runtimeConfig.cratePrefix}_types")

codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/smithy/SymbolMetadataProvider.kt

@@ -91,8 +91,9 @@ class BaseSymbolMetadataProvider(base: RustSymbolProvider) : SymbolMetadataProvi
     }
 
     companion object {
-        private val defaultDerives =
-            listOf(RuntimeType.StdFmt("Debug"), RuntimeType.Std("cmp::PartialEq"), RuntimeType.Std("clone::Clone"))
+        private val defaultDerives = with(RuntimeType) {
+            listOf(Debug, PartialEq, Clone)
+        }
     }
 }
 

codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/smithy/generators/SmithyTypesPubUseGenerator.kt

@@ -0,0 +1,19 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0.
+ */
+
+package software.amazon.smithy.rust.codegen.smithy.generators
+
+import software.amazon.smithy.rust.codegen.rustlang.rust
+import software.amazon.smithy.rust.codegen.rustlang.writable
+import software.amazon.smithy.rust.codegen.smithy.RuntimeConfig
+import software.amazon.smithy.rust.codegen.smithy.RuntimeType
+
+class SmithyTypesPubUseGenerator(private val runtimeConfig: RuntimeConfig) : LibRsCustomization() {
+    override fun section(section: LibRsSection) = writable {
+        when (section) {
+            LibRsSection.Body -> rust("pub use #T;", RuntimeType.Blob(runtimeConfig))
+        }
+    }
+}

codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/smithy/generators/StructureGenerator.kt

@@ -11,29 +11,37 @@ import software.amazon.smithy.model.shapes.MemberShape
 import software.amazon.smithy.model.shapes.Shape
 import software.amazon.smithy.model.shapes.StructureShape
 import software.amazon.smithy.model.traits.ErrorTrait
+import software.amazon.smithy.model.traits.SensitiveTrait
 import software.amazon.smithy.rust.codegen.rustlang.RustType
 import software.amazon.smithy.rust.codegen.rustlang.RustWriter
 import software.amazon.smithy.rust.codegen.rustlang.documentShape
+import software.amazon.smithy.rust.codegen.rustlang.rust
 import software.amazon.smithy.rust.codegen.rustlang.rustBlock
+import software.amazon.smithy.rust.codegen.smithy.RuntimeType
 import software.amazon.smithy.rust.codegen.smithy.RustSymbolProvider
 import software.amazon.smithy.rust.codegen.smithy.canUseDefault
 import software.amazon.smithy.rust.codegen.smithy.expectRustMetadata
 import software.amazon.smithy.rust.codegen.smithy.isOptional
 import software.amazon.smithy.rust.codegen.smithy.rustType
+import software.amazon.smithy.rust.codegen.util.dq
 
 fun RustWriter.implBlock(structureShape: Shape, symbolProvider: SymbolProvider, block: RustWriter.() -> Unit) {
     rustBlock("impl ${symbolProvider.toSymbol(structureShape).name}") {
         block(this)
     }
 }
 
+fun StructureShape.hasSensitiveMember(model: Model) =
+    this.members().any { it.getMemberTrait(model, SensitiveTrait::class.java).isPresent }
+
 class StructureGenerator(
     val model: Model,
     private val symbolProvider: RustSymbolProvider,
     private val writer: RustWriter,
     private val shape: StructureShape
 ) {
     private val members: List<MemberShape> = shape.allMembers.values.toList()
+    private val name = symbolProvider.toSymbol(shape).name
 
     fun render() {
         renderStructure()
@@ -72,19 +80,42 @@ class StructureGenerator(
         } else ""
     }
 
+    /** Render a custom debug implementation
+     * When [SensitiveTrait] support is required, render a custom debug implementation to redact sensitive data
+     */
+    private fun renderDebugImpl() {
+        writer.rustBlock("impl ${lifetimeDeclaration()} #T for $name ${lifetimeDeclaration()}", RuntimeType.Debug) {
+            writer.rustBlock("fn fmt(&self, f: &mut #1T::Formatter<'_>) -> #1T::Result", RuntimeType.StdFmt(null)) {
+                rust("""let mut formatter = f.debug_struct(${name.dq()});""")
+                members.forEach { member ->
+                    val memberName = symbolProvider.toMemberName(member)
+                    if (member.getMemberTrait(model, SensitiveTrait::class.java).isPresent) {
+                        rust("""formatter.field(${memberName.dq()}, &"*** Sensitive Data Redacted ***");""")
+                    } else {
+                        rust("formatter.field(${memberName.dq()}, &self.$memberName);")
+                    }
+                }
+                rust("formatter.finish()")
+            }
+        }
+    }
+
     private fun renderStructure() {
         val symbol = symbolProvider.toSymbol(shape)
         val containerMeta = symbol.expectRustMetadata()
         writer.documentShape(shape, model)
-        containerMeta.render(writer)
+        val withoutDebug = containerMeta.derives.copy(derives = containerMeta.derives.derives - RuntimeType.Debug)
+        containerMeta.copy(derives = withoutDebug).render(writer)
 
-        writer.rustBlock("struct ${symbol.name} ${lifetimeDeclaration()}") {
+        writer.rustBlock("struct $name ${lifetimeDeclaration()}") {
             members.forEach { member ->
                 val memberName = symbolProvider.toMemberName(member)
                 writer.documentShape(member, model)
                 symbolProvider.toSymbol(member).expectRustMetadata().render(this)
                 write("$memberName: #T,", symbolProvider.toSymbol(member))
             }
         }
+
+        renderDebugImpl()
     }
 }

codegen/src/test/kotlin/software/amazon/smithy/rust/codegen/generators/StructureGeneratorTest.kt

@@ -7,7 +7,6 @@ package software.amazon.smithy.rust.codegen.generators
 
 import io.kotest.matchers.string.shouldContainInOrder
 import org.junit.jupiter.api.Test
-import software.amazon.smithy.model.shapes.ShapeId
 import software.amazon.smithy.model.shapes.StructureShape
 import software.amazon.smithy.rust.codegen.rustlang.Custom
 import software.amazon.smithy.rust.codegen.rustlang.RustMetadata
@@ -47,10 +46,23 @@ class StructureGeneratorTest {
         structure MyError {
             message: String
         }
+
+        @sensitive
+        string SecretKey
+
+        structure Credentials {
+            username: String,
+            @sensitive
+            password: String,
+
+            // test that sensitive can be applied directly to a member or to the shape
+            secretKey: SecretKey
+        }
         """.asSmithyModel()
-        val struct = model.expectShape(ShapeId.from("com.test#MyStruct"), StructureShape::class.java)
-        val inner = model.expectShape(ShapeId.from("com.test#Inner"), StructureShape::class.java)
-        val error = model.expectShape(ShapeId.from("com.test#MyError"), StructureShape::class.java)
+        val struct = model.lookup<StructureShape>("com.test#MyStruct")
+        val inner = model.lookup<StructureShape>("com.test#Inner")
+        val credentials = model.lookup<StructureShape>("com.test#Credentials")
+        val error = model.lookup<StructureShape>("com.test#MyError")
     }
 
     @Test
@@ -115,6 +127,25 @@ class StructureGeneratorTest {
         )
     }
 
+    @Test
+    fun `generate a custom debug implementation when the sensitive trait is present`() {
+        val provider = testSymbolProvider(model)
+        val writer = RustWriter.forModule("lib")
+        val generator = StructureGenerator(model, provider, writer, credentials)
+        generator.render()
+        writer.unitTest(
+            """
+            let creds = Credentials {
+                username: Some("not_redacted".to_owned()),
+                password: Some("don't leak me".to_owned()),
+                secret_key: Some("don't leak me".to_owned())
+            };
+            assert_eq!(format!("{:?}", creds), "Credentials { username: Some(\"not_redacted\"), password: \"*** Sensitive Data Redacted ***\", secret_key: \"*** Sensitive Data Redacted ***\" }");
+        """
+        )
+        writer.compileAndTest()
+    }
+
     @Test
     fun `attach docs to everything`() {
         val model = """