Configure the (in|ex)clusion of the Debug trait for containers per me…

…mbers' sensitive trait (#2029)

* Removes Debug for shape with sensitive trait

This commit centralizes a place where the codegen excludes the Debug
trait if a shape has the sensitive trait. Previously the exclusion was
handled locally in each shape, e.g. StructureGenerator and EnumGenerator.
However, that approach may overlook a certain shape we also need to treat
as such. We now handle the exclusion of the Debug trait in one place,
BaseSymbolMetadataProvider.

* Stop excluding the Debug trait locally

This commit updates EnumGenerator and StructureGenerator based on the
change made to BaseSymbolMetadataProvider in the previous commit.
Now that the exclusion of the Debug trait was centralized, those
classes in question no longer need to do so individually.

* Implement a custom Debug trait in BuilderGenerator

This commit implements a custom Debug trait in BuilderGenerator now that
the derived Debug trait is excluded from BaseSymbolMetadataProvider for
the structure container. The implementation of the custom Debug trait
pretty much follows that of StructureGenerator.

* Implement a custom Debug trait in UnionGenerator

This commit implements a custom Debug trait in BuilderGenerator now that
the derived Debug trait is excluded from BaseSymbolMetadataProvider for
the union container. The implementation of the custom Debug trait pretty
much follows that of EnumGenerator.

* Implement a custom Debug trait in ServerBuilderGenerator

This commit implements a custom Debug trait in ServerBuilderGenerator now
that the derived Debug trait is excluded from BaseSymbolMetadataProvider
for the structure container. The implementation of the custom Debug trait
pretty much follows that of StructureGenerator.

* Add the Copyright header

* Update CHANGELOG.next.toml

* Update Debug impl for UnionGenerator

This commit updates the implementation of a custom Debug trait impl for
UnionGenerator. Turns out that in a Union, a member target can be marked
as sensitive separately outside the Union. Therefore, the implementation
of a custom Debug trait has two cases depending on where the sensitive
trait appears, either it is applied to the whole Union or to a member
target.

* Peek at member sensitivity for Debug trait (in|ex)clusion

This commit addresses #2029 (comment).
With this change, structure shapes no longer need to exclude the Debug
trait unconditionally. The upshot is that we may be able to avoid a
custom Debug impl for a structure where the derived Debug will do, i.e.
when there is no sensitive trait either at a container level or at a
member level.

* Remove statement that does not seem to take effect

This commit addresses #2029 (comment)

* Rename renderDebugImplForUnion -> renderFullyRedactedDebugImpl

This commit addresses #2029 (comment)

* Rename renderDebugImplForUnionMemberWise -> renderDebugImpl

This commit addresses #2029 (comment)

Co-authored-by: Yuki Saito <[email protected]>

CHANGELOG.next.toml

@@ -487,13 +487,13 @@ meta = { "breaking" = false, "tada" = false, "bug" = false, "target" = "all" }
 author = "lsr0"
 
 [[aws-sdk-rust]]
-message = "Implementation of the Debug trait for enums can redact what is printed per the sensitive trait."
-references = ["smithy-rs#1983"]
+message = "Implementation of the Debug trait for container shapes can redact what is printed per the sensitive trait."
+references = ["smithy-rs#1983", "smithy-rs#2029"]
 meta = { "breaking" = true, "tada" = false, "bug" = false }
 author = "ysaito1001"
 
 [[smithy-rs]]
-message = "Implementation of the Debug trait for enums can redact what is printed per the sensitive trait."
-references = ["smithy-rs#1983"]
+message = "Implementation of the Debug trait for container shapes can redact what is printed per the sensitive trait."
+references = ["smithy-rs#1983", "smithy-rs#2029"]
 meta = { "breaking" = true, "tada" = false, "bug" = false, "target" = "all" }
 author = "ysaito1001"

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

@@ -15,6 +15,7 @@ import software.amazon.smithy.model.shapes.StructureShape
 import software.amazon.smithy.model.shapes.UnionShape
 import software.amazon.smithy.model.traits.EnumDefinition
 import software.amazon.smithy.model.traits.EnumTrait
+import software.amazon.smithy.model.traits.SensitiveTrait
 import software.amazon.smithy.model.traits.StreamingTrait
 import software.amazon.smithy.rust.codegen.core.rustlang.Attribute
 import software.amazon.smithy.rust.codegen.core.rustlang.RustMetadata
@@ -77,13 +78,27 @@ abstract class SymbolMetadataProvider(private val base: RustSymbolProvider) : Wr
 class BaseSymbolMetadataProvider(
     base: RustSymbolProvider,
     private val model: Model,
-    additionalAttributes: List<Attribute>,
+    private val additionalAttributes: List<Attribute>,
 ) : SymbolMetadataProvider(base) {
-    private val containerDefault = RustMetadata(
-        Attribute.Derives(defaultDerives.toSet()),
-        additionalAttributes = additionalAttributes,
-        visibility = Visibility.PUBLIC,
-    )
+    private fun containerDefault(shape: Shape): RustMetadata {
+        val isSensitive = shape.hasTrait<SensitiveTrait>() ||
+            // Checking the shape's direct members for the sensitive trait should suffice.
+            // Whether their descendants, i.e. a member's member, is sensitive does not
+            // affect the inclusion/exclusion of the derived Debug trait of _this_ container
+            // shape; any sensitive descendant should still be printed as redacted.
+            shape.members().any { it.getMemberTrait(model, SensitiveTrait::class.java).isPresent }
+
+        val setOfDerives = if (isSensitive) {
+            defaultDerives.toSet() - RuntimeType.Debug
+        } else {
+            defaultDerives.toSet()
+        }
+        return RustMetadata(
+            Attribute.Derives(setOfDerives),
+            additionalAttributes = additionalAttributes,
+            visibility = Visibility.PUBLIC,
+        )
+    }
 
     override fun memberMeta(memberShape: MemberShape): RustMetadata {
         val container = model.expectShape(memberShape.container)
@@ -113,15 +128,15 @@ class BaseSymbolMetadataProvider(
     }
 
     override fun structureMeta(structureShape: StructureShape): RustMetadata {
-        return containerDefault
+        return containerDefault(structureShape)
     }
 
     override fun unionMeta(unionShape: UnionShape): RustMetadata {
-        return containerDefault
+        return containerDefault(unionShape)
     }
 
     override fun enumMeta(stringShape: StringShape): RustMetadata {
-        return containerDefault.withDerives(
+        return containerDefault(stringShape).withDerives(
             RuntimeType.std.member("hash::Hash"),
         ).withDerives(
             // enums can be eq because they can only contain ints and strings

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

@@ -44,6 +44,7 @@ import software.amazon.smithy.rust.codegen.core.smithy.rustType
 import software.amazon.smithy.rust.codegen.core.smithy.traits.SyntheticInputTrait
 import software.amazon.smithy.rust.codegen.core.util.dq
 import software.amazon.smithy.rust.codegen.core.util.hasTrait
+import software.amazon.smithy.rust.codegen.core.util.redactIfNecessary
 import software.amazon.smithy.rust.codegen.core.util.toSnakeCase
 
 // TODO(https://github.com/awslabs/smithy-rs/issues/1401) This builder generator is only used by the client.
@@ -110,13 +111,20 @@ class BuilderGenerator(
     private val runtimeConfig = symbolProvider.config().runtimeConfig
     private val members: List<MemberShape> = shape.allMembers.values.toList()
     private val structureSymbol = symbolProvider.toSymbol(shape)
+    private val builderSymbol = shape.builderSymbol(symbolProvider)
+    private val baseDerives = structureSymbol.expectRustMetadata().derives
+    private val builderDerives = baseDerives.derives.intersect(setOf(RuntimeType.Debug, RuntimeType.PartialEq, RuntimeType.Clone)) + RuntimeType.Default
+    private val builderName = "Builder"
 
     fun render(writer: RustWriter) {
         val symbol = symbolProvider.toSymbol(shape)
         writer.docs("See #D.", symbol)
-        val segments = shape.builderSymbol(symbolProvider).namespace.split("::")
         writer.withInlineModule(shape.builderSymbol(symbolProvider).module()) {
+            // Matching derives to the main structure + `Default` since we are a builder and everything is optional.
             renderBuilder(this)
+            if (!builderDerives.contains(RuntimeType.Debug)) {
+                renderDebugImpl(this)
+            }
         }
     }
 
@@ -142,7 +150,6 @@ class BuilderGenerator(
     }
 
     fun renderConvenienceMethod(implBlock: RustWriter) {
-        val builderSymbol = shape.builderSymbol(symbolProvider)
         implBlock.docs("Creates a new builder-style object to manufacture #D.", structureSymbol)
         implBlock.rustBlock("pub fn builder() -> #T", builderSymbol) {
             write("#T::default()", builderSymbol)
@@ -195,13 +202,8 @@ class BuilderGenerator(
     }
 
     private fun renderBuilder(writer: RustWriter) {
-        val builderName = "Builder"
-
         writer.docs("A builder for #D.", structureSymbol)
-        // Matching derives to the main structure + `Default` since we are a builder and everything is optional.
-        val baseDerives = structureSymbol.expectRustMetadata().derives
-        val derives = baseDerives.derives.intersect(setOf(RuntimeType.Debug, RuntimeType.PartialEq, RuntimeType.Clone)) + RuntimeType.Default
-        baseDerives.copy(derives = derives).render(writer)
+        baseDerives.copy(derives = builderDerives).render(writer)
         writer.rustBlock("pub struct $builderName") {
             for (member in members) {
                 val memberName = symbolProvider.toMemberName(member)
@@ -232,6 +234,23 @@ class BuilderGenerator(
         }
     }
 
+    private fun renderDebugImpl(writer: RustWriter) {
+        writer.rustBlock("impl #T for $builderName", RuntimeType.Debug) {
+            writer.rustBlock("fn fmt(&self, f: &mut #1T::Formatter<'_>) -> #1T::Result", RuntimeType.stdfmt) {
+                rust("""let mut formatter = f.debug_struct(${builderName.dq()});""")
+                members.forEach { member ->
+                    val memberName = symbolProvider.toMemberName(member)
+                    val fieldValue = member.redactIfNecessary(model, "self.$memberName")
+
+                    rust(
+                        "formatter.field(${memberName.dq()}, &$fieldValue);",
+                    )
+                }
+                rust("formatter.finish()")
+            }
+        }
+    }
+
     private fun RustWriter.renderVecHelper(member: MemberShape, memberName: String, coreType: RustType.Vec) {
         docs("Appends an item to `$memberName`.")
         rust("///")

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

@@ -94,12 +94,7 @@ open class EnumGenerator(
 ) {
     protected val symbol: Symbol = symbolProvider.toSymbol(shape)
     protected val enumName: String = symbol.name
-    private val isSensitive = shape.shouldRedact(model)
-    protected val meta = if (isSensitive) {
-        symbol.expectRustMetadata().withoutDerives(RuntimeType.Debug)
-    } else {
-        symbol.expectRustMetadata()
-    }
+    protected val meta = symbol.expectRustMetadata()
     protected val sortedMembers: List<EnumMemberModel> =
         enumTrait.values.sortedBy { it.value }.map { EnumMemberModel(it, symbolProvider) }
     protected open var target: CodegenTarget = CodegenTarget.CLIENT
@@ -136,7 +131,7 @@ open class EnumGenerator(
             renderUnnamedEnum()
         }
 
-        if (isSensitive) {
+        if (shape.shouldRedact(model)) {
             renderDebugImplForSensitiveEnum()
         }
     }

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

@@ -144,8 +144,7 @@ open class StructureGenerator(
         val containerMeta = symbol.expectRustMetadata()
         writer.documentShape(shape, model)
         writer.deprecatedShape(shape)
-        val withoutDebug = containerMeta.derives.copy(derives = containerMeta.derives.derives - RuntimeType.Debug)
-        containerMeta.copy(derives = withoutDebug).render(writer)
+        containerMeta.render(writer)
 
         writer.rustBlock("struct $name ${lifetimeDeclaration()}") {
             writer.forEachMember(members) { member, memberName, memberSymbol ->
@@ -154,7 +153,9 @@ open class StructureGenerator(
         }
 
         renderStructureImpl()
-        renderDebugImpl()
+        if (!containerMeta.derives.derives.contains(RuntimeType.Debug)) {
+            renderDebugImpl()
+        }
     }
 
     protected fun RustWriter.forEachMember(

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

@@ -9,16 +9,23 @@ import software.amazon.smithy.codegen.core.SymbolProvider
 import software.amazon.smithy.model.Model
 import software.amazon.smithy.model.shapes.MemberShape
 import software.amazon.smithy.model.shapes.UnionShape
+import software.amazon.smithy.model.traits.SensitiveTrait
 import software.amazon.smithy.rust.codegen.core.rustlang.Attribute
 import software.amazon.smithy.rust.codegen.core.rustlang.RustWriter
 import software.amazon.smithy.rust.codegen.core.rustlang.deprecatedShape
 import software.amazon.smithy.rust.codegen.core.rustlang.docs
 import software.amazon.smithy.rust.codegen.core.rustlang.documentShape
 import software.amazon.smithy.rust.codegen.core.rustlang.rust
 import software.amazon.smithy.rust.codegen.core.rustlang.rustBlock
+import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
 import software.amazon.smithy.rust.codegen.core.smithy.CodegenTarget
+import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
 import software.amazon.smithy.rust.codegen.core.smithy.expectRustMetadata
 import software.amazon.smithy.rust.codegen.core.smithy.renamedFrom
+import software.amazon.smithy.rust.codegen.core.util.REDACTION
+import software.amazon.smithy.rust.codegen.core.util.dq
+import software.amazon.smithy.rust.codegen.core.util.hasTrait
+import software.amazon.smithy.rust.codegen.core.util.shouldRedact
 import software.amazon.smithy.rust.codegen.core.util.toSnakeCase
 
 fun CodegenTarget.renderUnknownVariant() = when (this) {
@@ -47,16 +54,24 @@ class UnionGenerator(
     private val renderUnknownVariant: Boolean = true,
 ) {
     private val sortedMembers: List<MemberShape> = shape.allMembers.values.sortedBy { symbolProvider.toMemberName(it) }
+    private val unionSymbol = symbolProvider.toSymbol(shape)
 
     fun render() {
         renderUnion()
+        renderImplBlock()
+        if (!unionSymbol.expectRustMetadata().derives.derives.contains(RuntimeType.Debug)) {
+            if (shape.hasTrait<SensitiveTrait>()) {
+                renderFullyRedactedDebugImpl()
+            } else {
+                renderDebugImpl()
+            }
+        }
     }
 
     private fun renderUnion() {
         writer.documentShape(shape, model)
         writer.deprecatedShape(shape)
 
-        val unionSymbol = symbolProvider.toSymbol(shape)
         val containerMeta = unionSymbol.expectRustMetadata()
         containerMeta.render(writer)
         writer.rustBlock("enum ${unionSymbol.name}") {
@@ -82,6 +97,9 @@ class UnionGenerator(
                 rust("Unknown,")
             }
         }
+    }
+
+    private fun renderImplBlock() {
         writer.rustBlock("impl ${unionSymbol.name}") {
             sortedMembers.forEach { member ->
                 val memberSymbol = symbolProvider.toSymbol(member)
@@ -91,7 +109,11 @@ class UnionGenerator(
                 if (sortedMembers.size == 1) {
                     Attribute.Custom("allow(irrefutable_let_patterns)").render(this)
                 }
-                rust("/// Tries to convert the enum instance into [`$variantName`](#T::$variantName), extracting the inner #D.", unionSymbol, memberSymbol)
+                rust(
+                    "/// Tries to convert the enum instance into [`$variantName`](#T::$variantName), extracting the inner #D.",
+                    unionSymbol,
+                    memberSymbol,
+                )
                 rust("/// Returns `Err(&Self)` if it can't be converted.")
                 rustBlock("pub fn as_$funcNamePart(&self) -> std::result::Result<&#T, &Self>", memberSymbol) {
                     rust("if let ${unionSymbol.name}::$variantName(val) = &self { Ok(val) } else { Err(self) }")
@@ -110,10 +132,45 @@ class UnionGenerator(
         }
     }
 
+    private fun renderFullyRedactedDebugImpl() {
+        writer.rustTemplate(
+            """
+            impl #{Debug} for ${unionSymbol.name} {
+                fn fmt(&self, f: &mut #{StdFmt}::Formatter<'_>) -> #{StdFmt}::Result {
+                    write!(f, $REDACTION)
+                }
+            }
+            """,
+            "Debug" to RuntimeType.Debug,
+            "StdFmt" to RuntimeType.stdfmt,
+        )
+    }
+
+    private fun renderDebugImpl() {
+        writer.rustBlock("impl #T for ${unionSymbol.name}", RuntimeType.Debug) {
+            writer.rustBlock("fn fmt(&self, f: &mut #1T::Formatter<'_>) -> #1T::Result", RuntimeType.stdfmt) {
+                rustBlock("match self") {
+                    sortedMembers.forEach { member ->
+                        val memberName = symbolProvider.toMemberName(member)
+                        if (member.shouldRedact(model)) {
+                            rust("${unionSymbol.name}::$memberName(_) => f.debug_tuple($REDACTION).finish(),")
+                        } else {
+                            rust("${unionSymbol.name}::$memberName(val) => f.debug_tuple(${memberName.dq()}).field(&val).finish(),")
+                        }
+                    }
+                    if (renderUnknownVariant) {
+                        rust("${unionSymbol.name}::$UnknownVariantName => f.debug_tuple(${UnknownVariantName.dq()}).finish(),")
+                    }
+                }
+            }
+        }
+    }
+
     companion object {
         const val UnknownVariantName = "Unknown"
     }
 }
+
 fun unknownVariantError(union: String) =
     "Cannot serialize `$union::${UnionGenerator.UnknownVariantName}` for the request. " +
         "The `Unknown` variant is intended for responses only. " +

codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/smithy/generators/BuilderGeneratorTest.kt

@@ -24,6 +24,7 @@ internal class BuilderGeneratorTest {
     private val model = StructureGeneratorTest.model
     private val inner = StructureGeneratorTest.inner
     private val struct = StructureGeneratorTest.struct
+    private val credentials = StructureGeneratorTest.credentials
 
     @Test
     fun `generate builders`() {
@@ -94,4 +95,27 @@ internal class BuilderGeneratorTest {
             """,
         )
     }
+
+    @Test
+    fun `builder for a struct with sensitive fields should implement the debug trait as such`() {
+        val provider = testSymbolProvider(model)
+        val writer = RustWriter.forModule("model")
+        val credsGenerator = StructureGenerator(model, provider, writer, credentials)
+        val builderGenerator = BuilderGenerator(model, provider, credentials)
+        credsGenerator.render()
+        builderGenerator.render(writer)
+        writer.implBlock(credentials, provider) {
+            builderGenerator.renderConvenienceMethod(this)
+        }
+        writer.compileAndTest(
+            """
+            use super::*;
+            let builder = Credentials::builder()
+                .username("admin")
+                .password("pswd")
+                .secret_key("12345");
+                 assert_eq!(format!("{:?}", builder), "Builder { username: Some(\"admin\"), password: \"*** Sensitive Data Redacted ***\", secret_key: \"*** Sensitive Data Redacted ***\" }");
+            """,
+        )
+    }
 }