From ab5d0d7ebe1958936f65bffb5d8a05bacb20f0b8 Mon Sep 17 00:00:00 2001
From: Ian Botsford <83236726+ianbotsf@users.noreply.github.com>
Date: Fri, 12 Jul 2024 21:30:03 +0000
Subject: [PATCH] fix: correctly redact sensitive data in lists/maps

---
 .../30110eeb-0708-45b6-973f-960dc6062ce6.json |  5 +++
 .../codegen/rendering/StructureGenerator.kt   | 22 +++++++----
 .../rendering/StructureGeneratorTest.kt       | 38 +++++++++++++++++--
 3 files changed, 55 insertions(+), 10 deletions(-)
 create mode 100644 .changes/30110eeb-0708-45b6-973f-960dc6062ce6.json

diff --git a/.changes/30110eeb-0708-45b6-973f-960dc6062ce6.json b/.changes/30110eeb-0708-45b6-973f-960dc6062ce6.json
new file mode 100644
index 000000000..20f4712b5
--- /dev/null
+++ b/.changes/30110eeb-0708-45b6-973f-960dc6062ce6.json
@@ -0,0 +1,5 @@
+{
+    "id": "30110eeb-0708-45b6-973f-960dc6062ce6",
+    "type": "bugfix",
+    "description": "Correctly redact sensitive data in lists and maps"
+}
\ No newline at end of file
diff --git a/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/StructureGenerator.kt b/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/StructureGenerator.kt
index 7614072f1..c34eb5ab9 100644
--- a/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/StructureGenerator.kt
+++ b/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/StructureGenerator.kt
@@ -10,9 +10,12 @@ import software.amazon.smithy.kotlin.codegen.core.*
 import software.amazon.smithy.kotlin.codegen.lang.KotlinTypes
 import software.amazon.smithy.kotlin.codegen.model.*
 import software.amazon.smithy.kotlin.codegen.rendering.serde.ClientErrorCorrection
+import software.amazon.smithy.model.Model
 import software.amazon.smithy.model.shapes.*
 import software.amazon.smithy.model.traits.*
 
+private const val REDACTED_VALUE = "*** Sensitive Data Redacted ***"
+
 /**
  * Renders Smithy structure shapes
  */
@@ -105,18 +108,15 @@ class StructureGenerator(
             write("append(\"#T(\")", symbol)
 
             when {
-                shape.hasTrait<SensitiveTrait>() -> write("append(#S)", "*** Sensitive Data Redacted ***")
+                shape.isSensitive(model) -> write("append(#S)", REDACTED_VALUE)
                 else -> {
                     sortedMembers.forEachIndexed { index, memberShape ->
                         val (memberName, _) = memberNameSymbolIndex[memberShape]!!
+                        val isSensitive = memberShape.isSensitive(model)
+                        val value = if (isSensitive) REDACTED_VALUE else "\$$memberName"
                         val separator = if (index < sortedMembers.size - 1) "," else ""
 
-                        val targetShape = model.expectShape(memberShape.target)
-                        if (targetShape.hasTrait<SensitiveTrait>()) {
-                            write("append(\"#1L=*** Sensitive Data Redacted ***$separator\")", memberName)
-                        } else {
-                            write("append(\"#1L=\$#2L$separator\")", memberShape.defaultName(), memberName)
-                        }
+                        write("append(\"#L=#L#L\")", memberShape.defaultName(), value, separator)
                     }
                 }
             }
@@ -385,3 +385,11 @@ class StructureGenerator(
         if (hasConflictWithBaseClass) throw CodegenException("`sdkErrorMetadata` conflicts with property of same name inherited from SdkBaseException. Apply a rename customization/projection to fix.")
     }
 }
+
+private fun Shape.isSensitive(model: Model): Boolean = when {
+    this is MemberShape -> model.expectShape(target).isSensitive(model)
+    hasTrait<SensitiveTrait>() -> true
+    this is ListShape -> member.isSensitive(model)
+    this is MapShape -> key.isSensitive(model) || value.isSensitive(model)
+    else -> false
+}
diff --git a/codegen/smithy-kotlin-codegen/src/test/kotlin/software/amazon/smithy/kotlin/codegen/rendering/StructureGeneratorTest.kt b/codegen/smithy-kotlin-codegen/src/test/kotlin/software/amazon/smithy/kotlin/codegen/rendering/StructureGeneratorTest.kt
index 59a6a00cf..797e6f30e 100644
--- a/codegen/smithy-kotlin-codegen/src/test/kotlin/software/amazon/smithy/kotlin/codegen/rendering/StructureGeneratorTest.kt
+++ b/codegen/smithy-kotlin-codegen/src/test/kotlin/software/amazon/smithy/kotlin/codegen/rendering/StructureGeneratorTest.kt
@@ -285,11 +285,38 @@ class StructureGeneratorTest {
             @sensitive
             string Baz
             
+            list BazList {
+                member: Baz
+            }
+            
+            map BazToStringMap {
+                key: Baz
+                value: String
+            }
+            
+            map StringToBazMap {
+                key: String
+                value: Baz
+            }
+            
+            map StringToBazList {
+                key: String
+                value: BazList
+            }
+            
+            list StringToBazListList {
+                member: StringToBazList
+            }
+            
             structure Foo {
                 bar: Baz,
-                @documentation("Member documentation")
                 baz: Baz,
-                qux: String
+                qux: String,
+                quux: BazList,
+                corge: BazToStringMap,
+                grault: StringToBazMap,
+                garply: StringToBazList,
+                waldo: StringToBazListList,
             }
             
         """.prependNamespaceAndService().toSmithyModel()
@@ -306,7 +333,12 @@ class StructureGeneratorTest {
                 append("Foo(")
                 append("bar=*** Sensitive Data Redacted ***,")
                 append("baz=*** Sensitive Data Redacted ***,")
-                append("qux=${'$'}qux")
+                append("corge=*** Sensitive Data Redacted ***,")
+                append("garply=*** Sensitive Data Redacted ***,")
+                append("grault=*** Sensitive Data Redacted ***,")
+                append("quux=*** Sensitive Data Redacted ***,")
+                append("qux=${'$'}qux,")
+                append("waldo=*** Sensitive Data Redacted ***")
                 append(")")
             }
         """.formatForTest()