Skip to content

Commit efbad71

Browse files
sichanyooSichan Yoo
sichanyoo
and
Sichan Yoo
authored
fix: Widen sensitive trait check to nested members of list and map members. (#783)
* Widen sensitive trait check to nested members of list and map members. * Simplified code using Kotlin's partition method. * Narrow redaction to only redact keys or values in map shapes, if only values / keys are sensitive. * ktlint * Clean up code & add codegen tests. * Add missing optional chaining operator for accessing keys / values of map. --------- Co-authored-by: Sichan Yoo <[email protected]>
1 parent 054e265 commit efbad71

File tree

4 files changed

+164
-11
lines changed

4 files changed

+164
-11
lines changed

‎smithy-swift-codegen/src/main/kotlin/software/amazon/smithy/swift/codegen/DirectedSwiftCodegen.kt‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ import software.amazon.smithy.model.shapes.ServiceShape
2222
import software.amazon.smithy.model.traits.SensitiveTrait
2323
import software.amazon.smithy.swift.codegen.core.GenerationContext
2424
import software.amazon.smithy.swift.codegen.integration.CustomDebugStringConvertibleGenerator
25+
import software.amazon.smithy.swift.codegen.integration.CustomDebugStringConvertibleGenerator.Companion.isSensitive
2526
import software.amazon.smithy.swift.codegen.integration.ProtocolGenerator
2627
import software.amazon.smithy.swift.codegen.integration.SwiftIntegration
2728
import software.amazon.smithy.swift.codegen.model.hasTrait
@@ -114,7 +115,7 @@ class DirectedSwiftCodegen(val context: PluginContext) :
114115
StructureGenerator(model, symbolProvider, writer, shape, settings, protocolGenerator?.serviceErrorProtocolSymbol).render()
115116
}
116117

117-
if (shape.hasTrait<SensitiveTrait>() || shape.members().any { it.hasTrait<SensitiveTrait>() || model.expectShape(it.target).hasTrait<SensitiveTrait>() }) {
118+
if (shape.hasTrait<SensitiveTrait>() || shape.members().any { it.isSensitive(model) }) {
118119
writers.useShapeExtensionWriter(shape, "CustomDebugStringConvertible") { writer: SwiftWriter ->
119120
CustomDebugStringConvertibleGenerator(symbolProvider, writer, shape, model).render()
120121
}
@@ -134,7 +135,7 @@ class DirectedSwiftCodegen(val context: PluginContext) :
134135
StructureGenerator(model, symbolProvider, writer, shape, settings, protocolGenerator?.serviceErrorProtocolSymbol).renderErrors()
135136
}
136137

137-
if (shape.hasTrait<SensitiveTrait>() || shape.members().any { it.hasTrait<SensitiveTrait>() || model.expectShape(it.target).hasTrait<SensitiveTrait>() }) {
138+
if (shape.hasTrait<SensitiveTrait>() || shape.members().any { it.isSensitive(model) }) {
138139
writers.useShapeExtensionWriter(shape, "CustomDebugStringConvertible") { writer: SwiftWriter ->
139140
CustomDebugStringConvertibleGenerator(symbolProvider, writer, shape, model).render()
140141
}

‎smithy-swift-codegen/src/main/kotlin/software/amazon/smithy/swift/codegen/integration/CustomDebugStringConvertibleGenerator.kt‎

Lines changed: 26 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,10 @@ package software.amazon.smithy.swift.codegen.integration
88
import software.amazon.smithy.codegen.core.Symbol
99
import software.amazon.smithy.codegen.core.SymbolProvider
1010
import software.amazon.smithy.model.Model
11+
import software.amazon.smithy.model.shapes.ListShape
12+
import software.amazon.smithy.model.shapes.MapShape
1113
import software.amazon.smithy.model.shapes.MemberShape
14+
import software.amazon.smithy.model.shapes.Shape
1215
import software.amazon.smithy.model.shapes.StructureShape
1316
import software.amazon.smithy.model.traits.ErrorTrait
1417
import software.amazon.smithy.model.traits.SensitiveTrait
@@ -25,6 +28,14 @@ class CustomDebugStringConvertibleGenerator(
2528
) {
2629
companion object {
2730
const val REDACT_STRING = "CONTENT_REDACTED"
31+
32+
fun Shape.isSensitive(model: Model): Boolean = when {
33+
this is MemberShape -> model.expectShape(target).isSensitive(model)
34+
hasTrait<SensitiveTrait>() -> true
35+
this is ListShape -> member.isSensitive(model)
36+
this is MapShape -> key.isSensitive(model) || value.isSensitive(model)
37+
else -> false
38+
}
2839
}
2940

3041
private val structSymbol: Symbol by lazy {
@@ -48,14 +59,7 @@ class CustomDebugStringConvertibleGenerator(
4859
private fun renderDescription() {
4960
val symbolName = structSymbol.name
5061
writer.writeInline("\"$symbolName(")
51-
val membersWithoutSensitiveTrait = membersSortedByName
52-
.filterNot { it.hasTrait<SensitiveTrait>() || model.expectShape(it.target).hasTrait<SensitiveTrait>() }
53-
.sortedBy { it.memberName }
54-
.toList()
55-
val membersWithSensitiveTrait = membersSortedByName
56-
.filter { it.hasTrait<SensitiveTrait>() || model.expectShape(it.target).hasTrait<SensitiveTrait>() }
57-
.sortedBy { it.memberName }
58-
.toList()
62+
val (membersWithSensitiveTrait, membersWithoutSensitiveTrait) = membersSortedByName.partition { it.isSensitive(model) }
5963
for (member in membersWithoutSensitiveTrait) {
6064
renderMemberDescription(writer, member, false)
6165
renderComma(writer, member != membersWithoutSensitiveTrait.last())
@@ -77,7 +81,12 @@ class CustomDebugStringConvertibleGenerator(
7781
) {
7882
val memberNames = symbolProvider.toMemberNames(member)
7983
val path = "properties.".takeIf { shape.hasTrait<ErrorTrait>() } ?: ""
80-
val description = if (isRedacted) "\\\"$REDACT_STRING\\\"" else "\\(${SwiftTypes.String}(describing: $path${memberNames.first}))"
84+
var description = ""
85+
if (model.expectShape(member.target).isMapShape) {
86+
description = getStringForLoggingMapShape(model.expectShape(member.target).asMapShape().get(), path, memberNames)
87+
} else {
88+
description = if (isRedacted) "\\\"$REDACT_STRING\\\"" else "\\(${SwiftTypes.String}(describing: $path${memberNames.first}))"
89+
}
8190
writer.writeInline("${memberNames.second}: $description")
8291
}
8392

@@ -86,4 +95,12 @@ class CustomDebugStringConvertibleGenerator(
8695
writer.writeInline(", ")
8796
}
8897
}
98+
99+
private fun getStringForLoggingMapShape(member: MapShape, path: String, memberNames: Pair<String, String>): String {
100+
if (member.hasTrait<SensitiveTrait>()) return "\\\"$REDACT_STRING\\\""
101+
if (member.key.isSensitive(model) && member.value.isSensitive(model)) return "\\\"$REDACT_STRING\\\""
102+
if (member.key.isSensitive(model)) return "[keys: \\\"$REDACT_STRING\\\", values: \\(${SwiftTypes.String}(describing: $path${memberNames.first}?.values))]"
103+
if (member.value.isSensitive(model)) return "[keys: \\(${SwiftTypes.String}(describing: $path${memberNames.first}?.keys)), values: \\\"$REDACT_STRING\\\"]"
104+
return "\\(${SwiftTypes.String}(describing: $path${memberNames.first}))"
105+
}
89106
}

‎smithy-swift-codegen/src/test/kotlin/CustomDebugStringConvertibleGeneratorTests.kt‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
import io.kotest.matchers.string.shouldContainOnlyOnce
2+
import org.junit.jupiter.api.Test
3+
import software.amazon.smithy.swift.codegen.DefaultClientConfigurationIntegration
4+
5+
class CustomDebugStringConvertibleGeneratorTests {
6+
@Test
7+
fun `list with sensitive trait gets redacted correctly`() {
8+
val context = setupTests("custom-debug-string-convertible-generator-test.smithy", "com.test#Example")
9+
val contents = getFileContents(context.manifest, "Sources/RestJson/models/GetFooOutput+CustomDebugStringConvertible.swift")
10+
contents.shouldSyntacticSanityCheck()
11+
contents.shouldContainOnlyOnce("listWithSensitiveTrait: \\\"CONTENT_REDACTED\\\"")
12+
}
13+
14+
@Test
15+
fun `list with member shape that targets a shape with sensitive trait gets redacted correctly`() {
16+
val context = setupTests("custom-debug-string-convertible-generator-test.smithy", "com.test#Example")
17+
val contents = getFileContents(context.manifest, "Sources/RestJson/models/GetFooOutput+CustomDebugStringConvertible.swift")
18+
contents.shouldSyntacticSanityCheck()
19+
contents.shouldContainOnlyOnce("listWithSensitiveTargetedByMember: \\\"CONTENT_REDACTED\\\"")
20+
}
21+
22+
@Test
23+
fun `map with sensitive trait gets redacted correctly`() {
24+
val context = setupTests("custom-debug-string-convertible-generator-test.smithy", "com.test#Example")
25+
val contents = getFileContents(context.manifest, "Sources/RestJson/models/GetFooOutput+CustomDebugStringConvertible.swift")
26+
contents.shouldSyntacticSanityCheck()
27+
contents.shouldContainOnlyOnce("mapWithSensitiveTrait: \\\"CONTENT_REDACTED\\\"")
28+
}
29+
30+
@Test
31+
fun `map with key that targets a shape with sensitive trait gets redacted correctly`() {
32+
val context = setupTests("custom-debug-string-convertible-generator-test.smithy", "com.test#Example")
33+
val contents = getFileContents(context.manifest, "Sources/RestJson/models/GetFooOutput+CustomDebugStringConvertible.swift")
34+
contents.shouldSyntacticSanityCheck()
35+
contents.shouldContainOnlyOnce("mapWithSensitiveTargetedByKey: [keys: \\\"CONTENT_REDACTED\\\", values: \\(Swift.String(describing: mapWithSensitiveTargetedByKey?.values))]")
36+
}
37+
38+
@Test
39+
fun `map with value that targets a shape with sensitive trait gets redacted correctly`() {
40+
val context = setupTests("custom-debug-string-convertible-generator-test.smithy", "com.test#Example")
41+
val contents = getFileContents(context.manifest, "Sources/RestJson/models/GetFooOutput+CustomDebugStringConvertible.swift")
42+
contents.shouldSyntacticSanityCheck()
43+
contents.shouldContainOnlyOnce("mapWithSensitiveTargetedByValue: [keys: \\(Swift.String(describing: mapWithSensitiveTargetedByValue?.keys)), values: \\\"CONTENT_REDACTED\\\"]")
44+
}
45+
46+
@Test
47+
fun `map with key and value that target shapes with sensitive trait gets redacted correctly`() {
48+
val context = setupTests("custom-debug-string-convertible-generator-test.smithy", "com.test#Example")
49+
val contents = getFileContents(context.manifest, "Sources/RestJson/models/GetFooOutput+CustomDebugStringConvertible.swift")
50+
contents.shouldSyntacticSanityCheck()
51+
contents.shouldContainOnlyOnce("mapWithSensitiveTargetedByBoth: \\\"CONTENT_REDACTED\\\"")
52+
}
53+
54+
private fun setupTests(smithyFile: String, serviceShapeId: String): TestContext {
55+
val context = TestContext.initContextFrom(
56+
listOf(smithyFile),
57+
serviceShapeId,
58+
MockHTTPRestJsonProtocolGenerator(),
59+
{ model -> model.defaultSettings(serviceShapeId, "RestJson", "2019-12-16", "Rest Json Protocol") },
60+
listOf(DefaultClientConfigurationIntegration())
61+
)
62+
63+
context.generator.initializeMiddleware(context.generationCtx)
64+
context.generator.generateProtocolClient(context.generationCtx)
65+
context.generator.generateSerializers(context.generationCtx)
66+
context.generationCtx.delegator.flushWriters()
67+
return context
68+
}
69+
}

‎smithy-swift-codegen/src/test/resources/custom-debug-string-convertible-generator-test.smithy‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
$version: "1.0"
2+
namespace com.test
3+
4+
use aws.protocols#awsJson1_1
5+
6+
@awsJson1_1
7+
service Example {
8+
version: "1.0.0",
9+
operations: [
10+
GetFoo
11+
]
12+
}
13+
14+
@http(method: "GET", uri: "/foo")
15+
operation GetFoo {
16+
input: GetFooRequest,
17+
output: GetFooResponse
18+
}
19+
20+
structure GetFooRequest {}
21+
22+
structure GetFooResponse {
23+
listWithSensitiveTrait: SensitiveList
24+
listWithSensitiveTargetedByMember: SensitiveMemberTargetList
25+
26+
mapWithSensitiveTrait: SensitiveMap
27+
mapWithSensitiveTargetedByKey: SensitiveKeyTargetMap
28+
mapWithSensitiveTargetedByValue: SensitiveValueTargetMap
29+
mapWithSensitiveTargetedByBoth: SensitiveTargetMap
30+
}
31+
32+
@sensitive
33+
list SensitiveList {
34+
member: String
35+
}
36+
37+
list SensitiveMemberTargetList {
38+
member: SensitiveTargetStruct
39+
}
40+
41+
@sensitive
42+
map SensitiveMap {
43+
key: String
44+
value: String
45+
}
46+
47+
map SensitiveKeyTargetMap {
48+
key: SensitiveTargetString
49+
value: String
50+
}
51+
52+
map SensitiveValueTargetMap {
53+
key: String
54+
value: SensitiveTargetStruct
55+
}
56+
57+
map SensitiveTargetMap {
58+
key: SensitiveTargetString
59+
value: SensitiveTargetStruct
60+
}
61+
62+
@sensitive
63+
structure SensitiveTargetStruct {}
64+
65+
@sensitive
66+
string SensitiveTargetString

0 commit comments

Comments
 (0)