Fix bug with @sensitive structs

smithy-lang · Jul 16, 2024 · 4b398c7 · 4b398c7
1 parent 5a65100
commit 4b398c7
Show file tree

Hide file tree

Showing 4 changed files with 43 additions and 16 deletions.
diff --git a/...ain/kotlin/software/amazon/smithy/rust/codegen/core/smithy/generators/BuilderGenerator.kt b/...ain/kotlin/software/amazon/smithy/rust/codegen/core/smithy/generators/BuilderGenerator.kt
@@ -45,10 +45,12 @@ import software.amazon.smithy.rust.codegen.core.smithy.isOptional
 import software.amazon.smithy.rust.codegen.core.smithy.makeOptional
 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.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.letIf
 import software.amazon.smithy.rust.codegen.core.util.redactIfNecessary
+import software.amazon.smithy.rust.codegen.core.util.shouldRedact
 import software.amazon.smithy.rust.codegen.core.util.toSnakeCase
 
 // TODO(https://github.com/smithy-lang/smithy-rs/issues/1401) This builder generator is only used by the client.
@@ -353,7 +355,16 @@ class BuilderGenerator(
                 rust("""let mut formatter = f.debug_struct(${builderName.dq()});""")
                 members.forEach { member ->
                     val memberName = symbolProvider.toMemberName(member)
-                    val fieldValue = member.redactIfNecessary(model, "self.$memberName")
+                    // If the struct is marked sensitive all fields get redacted, otherwise each field is determined on its own
+                    val fieldValue =
+                        if (shape.shouldRedact(model)) {
+                            REDACTION
+                        } else {
+                            member.redactIfNecessary(
+                                model,
+                                "self.$memberName",
+                            )
+                        }
 
                     rust(
                         "formatter.field(${memberName.dq()}, &$fieldValue);",

diff --git a/...n/kotlin/software/amazon/smithy/rust/codegen/core/smithy/generators/StructureGenerator.kt b/...n/kotlin/software/amazon/smithy/rust/codegen/core/smithy/generators/StructureGenerator.kt
@@ -33,9 +33,11 @@ import software.amazon.smithy.rust.codegen.core.smithy.customize.writeCustomizat
 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.smithy.rustType
+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.getTrait
 import software.amazon.smithy.rust.codegen.core.util.redactIfNecessary
+import software.amazon.smithy.rust.codegen.core.util.shouldRedact
 
 /** StructureGenerator customization sections */
 sealed class StructureSection(name: String) : Section(name) {
@@ -102,9 +104,19 @@ open class StructureGenerator(
         ) {
             writer.rustBlock("fn fmt(&self, f: &mut #1T::Formatter<'_>) -> #1T::Result", RuntimeType.stdFmt) {
                 rust("""let mut formatter = f.debug_struct(${name.dq()});""")
+
                 members.forEach { member ->
                     val memberName = symbolProvider.toMemberName(member)
-                    val fieldValue = member.redactIfNecessary(model, "self.$memberName")
+                    // If the struct is marked sensitive all fields get redacted, otherwise each field is determined on its own
+                    val fieldValue =
+                        if (shape.shouldRedact(model)) {
+                            REDACTION
+                        } else {
+                            member.redactIfNecessary(
+                                model,
+                                "self.$memberName",
+                            )
+                        }
 
                     rust(
                         "formatter.field(${memberName.dq()}, &$fieldValue);",

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
@@ -102,8 +102,8 @@ fun ServiceShape.hasEventStreamOperations(model: Model): Boolean =
 
 fun Shape.shouldRedact(model: Model): Boolean =
     when {
-        this is MemberShape -> model.expectShape(target).shouldRedact(model)
         hasTrait<SensitiveTrait>() -> true
+        this is MemberShape -> model.expectShape(target).shouldRedact(model)
         this is ListShape -> member.shouldRedact(model)
         this is MapShape -> key.shouldRedact(model) || value.shouldRedact(model)
         else -> false

diff --git a/...kotlin/software/amazon/smithy/rust/codegen/core/smithy/generators/BuilderGeneratorTest.kt b/...kotlin/software/amazon/smithy/rust/codegen/core/smithy/generators/BuilderGeneratorTest.kt
@@ -124,7 +124,8 @@ internal class BuilderGeneratorTest {
                         .username("admin")
                         .password("pswd")
                         .secret_key("12345");
-                         assert_eq!(format!("{:?}", builder), "Builder { username: Some(\"admin\"), password: \"*** Sensitive Data Redacted ***\", secret_key: \"*** Sensitive Data Redacted ***\" }");
+                         assert_eq!(format!("{:?}", builder),
+                         "Builder { username: Some(\"admin\"), password: \"*** Sensitive Data Redacted ***\", secret_key: \"*** Sensitive Data Redacted ***\", secret_value_map: \"*** Sensitive Data Redacted ***\", secret_key_map: \"*** Sensitive Data Redacted ***\", secret_list: \"*** Sensitive Data Redacted ***\" }");
                     """,
                 )
             }
@@ -278,18 +279,21 @@ internal class BuilderGeneratorTest {
             implBlock(provider.toSymbol(inner)) {
                 BuilderGenerator.renderConvenienceMethod(this, provider, inner)
             }
-            unitTest("test", additionalAttributes = listOf(Attribute.DenyDeprecated), block = {
-                rust(
-                    // Notice that the builder is instantiated directly, and not through the Inner::builder() method.
-                    // This is because Inner is marked with `deprecated`, so any usage of `Inner` inside the test will
-                    // fail the compilation.
-                    //
-                    // This piece of code would fail though if the Builder inherits the attributes from Inner.
-                    """
-                    let _ = test_inner::Builder::default();
-                    """,
-                )
-            })
+            unitTest(
+                "test", additionalAttributes = listOf(Attribute.DenyDeprecated),
+                block = {
+                    rust(
+                        // Notice that the builder is instantiated directly, and not through the Inner::builder() method.
+                        // This is because Inner is marked with `deprecated`, so any usage of `Inner` inside the test will
+                        // fail the compilation.
+                        //
+                        // This piece of code would fail though if the Builder inherits the attributes from Inner.
+                        """
+                        let _ = test_inner::Builder::default();
+                        """,
+                    )
+                },
+            )
         }
         project.withModule(provider.moduleForBuilder(inner)) {
             BuilderGenerator(model, provider, inner, emptyList()).render(this)