Skip to content

Commit 696c69d

Browse files
aajtoddianbotsf
aajtodd
and
ianbotsf
authored
fix(codegen): respect sensitive trait when applied to container shape (#764)
Co-authored-by: Ian Smith Botsford <[email protected]>
1 parent cee4330 commit 696c69d

File tree

5 files changed

+181
-21
lines changed

5 files changed

+181
-21
lines changed

‎.changes/51ec552f-1a2f-49d2-8e87-a421ad1dec27.json‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
{
2+
"id": "51ec552f-1a2f-49d2-8e87-a421ad1dec27",
3+
"type": "bugfix",
4+
"description": "Respect @sensitive trait when applied to container shape",
5+
"issues": [
6+
"awslabs/smithy-kotlin#763"
7+
]
8+
}

‎smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/StructureGenerator.kt‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -128,11 +128,11 @@ class StructureGenerator(
128128
write("append(\"#T(\")", symbol)
129129

130130
when {
131-
sortedMembers.isEmpty() -> write("append(\")\")")
131+
shape.hasTrait<SensitiveTrait>() -> write("append(#S)", "*** Sensitive Data Redacted ***")
132132
else -> {
133133
sortedMembers.forEachIndexed { index, memberShape ->
134134
val (memberName, _) = memberNameSymbolIndex[memberShape]!!
135-
val separator = if (index < sortedMembers.size - 1) "," else ")"
135+
val separator = if (index < sortedMembers.size - 1) "," else ""
136136

137137
val targetShape = model.expectShape(memberShape.target)
138138
if (targetShape.hasTrait<SensitiveTrait>()) {
@@ -143,6 +143,8 @@ class StructureGenerator(
143143
}
144144
}
145145
}
146+
147+
write("append(\")\")")
146148
}
147149
}
148150

‎smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/UnionGenerator.kt‎

Lines changed: 23 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
*/
55
package software.amazon.smithy.kotlin.codegen.rendering
66

7+
import software.amazon.smithy.codegen.core.Symbol
78
import software.amazon.smithy.codegen.core.SymbolProvider
89
import software.amazon.smithy.kotlin.codegen.core.*
910
import software.amazon.smithy.kotlin.codegen.lang.KotlinTypes
@@ -12,6 +13,7 @@ import software.amazon.smithy.kotlin.codegen.model.hasTrait
1213
import software.amazon.smithy.kotlin.codegen.model.isBoxed
1314
import software.amazon.smithy.model.Model
1415
import software.amazon.smithy.model.shapes.*
16+
import software.amazon.smithy.model.traits.SensitiveTrait
1517
import software.amazon.smithy.model.traits.StreamingTrait
1618

1719
/**
@@ -23,13 +25,13 @@ class UnionGenerator(
2325
private val writer: KotlinWriter,
2426
private val shape: UnionShape,
2527
) {
28+
val symbol: Symbol = symbolProvider.toSymbol(shape)
2629

2730
/**
2831
* Renders a Smithy union to a Kotlin sealed class
2932
*/
3033
fun render() {
3134
check(!shape.allMembers.values.any { memberShape -> memberShape.memberName.equals("SdkUnknown", true) }) { "generating SdkUnknown would cause duplicate variant for union shape: $shape" }
32-
val symbol = symbolProvider.toSymbol(shape)
3335
writer.renderDocumentation(shape)
3436
writer.renderAnnotations(shape)
3537
writer.openBlock("public sealed class #T {", symbol)
@@ -44,20 +46,21 @@ class UnionGenerator(
4446
writer.renderAnnotations(it)
4547
val variantName = it.unionVariantName()
4648
val variantSymbol = symbolProvider.toSymbol(it)
47-
writer.writeInline("public data class #L(val value: #Q) : #Q()", variantName, variantSymbol, symbol)
48-
when (model.expectShape(it.target).type) {
49-
ShapeType.BLOB -> {
50-
writer.withBlock(" {", "}") {
51-
renderHashCode(model, listOf(it), symbolProvider, this)
52-
renderEquals(model, listOf(it), variantName, this)
53-
}
49+
writer.withBlock("public data class #L(val value: #Q) : #Q() {", "}", variantName, variantSymbol, symbol) {
50+
if (model.expectShape(it.target).type == ShapeType.BLOB) {
51+
renderHashCode(model, listOf(it), symbolProvider, this)
52+
renderEquals(model, listOf(it), variantName, this)
5453
}
55-
else -> writer.write("")
54+
55+
renderToString()
5656
}
57+
writer.write("")
5758
}
5859

5960
// generate the unknown which will always be last
60-
writer.write("public object SdkUnknown : #Q()", symbol)
61+
writer.withBlock("public object SdkUnknown : #Q() {", "}", symbol) {
62+
renderToString()
63+
}
6164

6265
members.sortedBy { it.memberName }.forEach {
6366
val variantName = it.unionVariantName()
@@ -172,4 +175,14 @@ class UnionGenerator(
172175
write("return true")
173176
}
174177
}
178+
179+
private fun renderToString() {
180+
if (shape.hasTrait<SensitiveTrait>()) {
181+
writer.write(
182+
"""override fun toString(): #Q = "#T(*** Sensitive Data Redacted ***)"""",
183+
KotlinTypes.String,
184+
symbol,
185+
)
186+
} // else just use regular data class toString implementation
187+
}
175188
}

‎smithy-kotlin-codegen/src/test/kotlin/software/amazon/smithy/kotlin/codegen/rendering/StructureGeneratorTest.kt‎

Lines changed: 53 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -107,7 +107,8 @@ class StructureGeneratorTest {
107107
append("byteValue=${'$'}byteValue,")
108108
append("foo=${'$'}foo,")
109109
append("object=${'$'}`object`,")
110-
append("quux=${'$'}quux)")
110+
append("quux=${'$'}quux")
111+
append(")")
111112
}
112113
""".formatForTest()
113114

@@ -273,12 +274,62 @@ class StructureGeneratorTest {
273274
append("Foo(")
274275
append("bar=*** Sensitive Data Redacted ***,")
275276
append("baz=*** Sensitive Data Redacted ***,")
276-
append("qux=${'$'}qux)")
277+
append("qux=${'$'}qux")
278+
append(")")
277279
}
278280
""".formatForTest()
279281
generated.shouldContainOnlyOnceWithDiff(expected)
280282
}
281283

284+
@Test
285+
fun testSensitiveTraitOnParent() {
286+
val model = """
287+
string Baz
288+
289+
@sensitive
290+
structure Foo {
291+
bar: Baz,
292+
@documentation("Member documentation")
293+
baz: Baz,
294+
qux: String
295+
}
296+
297+
structure Parent {
298+
foo: Foo
299+
}
300+
""".prependNamespaceAndService().toSmithyModel()
301+
302+
val provider: SymbolProvider = KotlinCodegenPlugin.createSymbolProvider(model)
303+
val writer = KotlinWriter(TestModelDefault.NAMESPACE)
304+
val struct = model.expectShape<StructureShape>("com.test#Foo")
305+
val renderingCtx = RenderingContext(writer, struct, model, provider, model.defaultSettings())
306+
StructureGenerator(renderingCtx).render()
307+
308+
val generated = writer.toString()
309+
val expected = """
310+
override fun toString(): kotlin.String = buildString {
311+
append("Foo(")
312+
append("*** Sensitive Data Redacted ***")
313+
append(")")
314+
}
315+
""".formatForTest()
316+
generated.shouldContainOnlyOnceWithDiff(expected)
317+
318+
val writer2 = KotlinWriter(TestModelDefault.NAMESPACE)
319+
val struct2 = model.expectShape<StructureShape>("com.test#Parent")
320+
val renderingCtx2 = RenderingContext(writer2, struct2, model, provider, model.defaultSettings())
321+
StructureGenerator(renderingCtx2).render()
322+
val generated2 = writer2.toString()
323+
val expected2 = """
324+
override fun toString(): kotlin.String = buildString {
325+
append("Parent(")
326+
append("foo=*** Sensitive Data Redacted ***")
327+
append(")")
328+
}
329+
""".formatForTest()
330+
generated2.shouldContainOnlyOnceWithDiff(expected2)
331+
}
332+
282333
@Test
283334
fun `it handles required HTTP fields in initializers`() {
284335
val model = """

‎smithy-kotlin-codegen/src/test/kotlin/software/amazon/smithy/kotlin/codegen/rendering/UnionGeneratorTest.kt‎

Lines changed: 93 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -47,8 +47,12 @@ class UnionGeneratorTest {
4747
* Documentation for MyUnion
4848
*/
4949
public sealed class MyUnion {
50-
public data class Bar(val value: kotlin.Int) : test.model.MyUnion()
51-
public data class Baz(val value: kotlin.Int) : test.model.MyUnion()
50+
public data class Bar(val value: kotlin.Int) : test.model.MyUnion() {
51+
}
52+
53+
public data class Baz(val value: kotlin.Int) : test.model.MyUnion() {
54+
}
55+
5256
public data class Blz(val value: kotlin.ByteArray) : test.model.MyUnion() {
5357
5458
override fun hashCode(): kotlin.Int {
@@ -66,12 +70,18 @@ class UnionGeneratorTest {
6670
return true
6771
}
6872
}
73+
6974
/**
7075
* Documentation for foo
7176
*/
72-
public data class Foo(val value: kotlin.String) : test.model.MyUnion()
73-
public data class MyStruct(val value: test.model.MyStruct) : test.model.MyUnion()
74-
public object SdkUnknown : test.model.MyUnion()
77+
public data class Foo(val value: kotlin.String) : test.model.MyUnion() {
78+
}
79+
80+
public data class MyStruct(val value: test.model.MyStruct) : test.model.MyUnion() {
81+
}
82+
83+
public object SdkUnknown : test.model.MyUnion() {
84+
}
7585
7686
/**
7787
* Casts this [MyUnion] as a [Bar] and retrieves its [kotlin.Int] value. Throws an exception if the [MyUnion] is not a
@@ -219,8 +229,11 @@ class UnionGeneratorTest {
219229

220230
val expectedClassDecl = """
221231
public sealed class MyUnion {
222-
public data class Foo(val value: test.model.MyStruct) : test.model.MyUnion()
223-
public object SdkUnknown : test.model.MyUnion()
232+
public data class Foo(val value: test.model.MyStruct) : test.model.MyUnion() {
233+
}
234+
235+
public object SdkUnknown : test.model.MyUnion() {
236+
}
224237
225238
/**
226239
* Casts this [MyUnion] as a [Foo] and retrieves its [test.model.MyStruct] value. Throws an exception if the [MyUnion] is not a
@@ -238,6 +251,79 @@ class UnionGeneratorTest {
238251
contents.shouldContainWithDiff(expectedClassDecl)
239252
}
240253

254+
@Test
255+
fun `it generates sensitive unions`() {
256+
val contents = generateUnion(
257+
"""
258+
@sensitive
259+
union MyUnion {
260+
foo: String,
261+
bar: Integer,
262+
baz: MyStruct
263+
}
264+
265+
structure MyStruct {
266+
qux: String
267+
}
268+
""",
269+
)
270+
271+
val expectedClassDecl = """
272+
public sealed class MyUnion {
273+
public data class Bar(val value: kotlin.Int) : test.model.MyUnion() {
274+
override fun toString(): kotlin.String = "MyUnion(*** Sensitive Data Redacted ***)"
275+
}
276+
277+
public data class Baz(val value: test.model.MyStruct) : test.model.MyUnion() {
278+
override fun toString(): kotlin.String = "MyUnion(*** Sensitive Data Redacted ***)"
279+
}
280+
281+
public data class Foo(val value: kotlin.String) : test.model.MyUnion() {
282+
override fun toString(): kotlin.String = "MyUnion(*** Sensitive Data Redacted ***)"
283+
}
284+
285+
public object SdkUnknown : test.model.MyUnion() {
286+
override fun toString(): kotlin.String = "MyUnion(*** Sensitive Data Redacted ***)"
287+
}
288+
289+
/**
290+
* Casts this [MyUnion] as a [Bar] and retrieves its [kotlin.Int] value. Throws an exception if the [MyUnion] is not a
291+
* [Bar].
292+
*/
293+
public fun asBar(): kotlin.Int = (this as MyUnion.Bar).value
294+
295+
/**
296+
* Casts this [MyUnion] as a [Bar] and retrieves its [kotlin.Int] value. Returns null if the [MyUnion] is not a [Bar].
297+
*/
298+
public fun asBarOrNull(): kotlin.Int? = (this as? MyUnion.Bar)?.value
299+
300+
/**
301+
* Casts this [MyUnion] as a [Baz] and retrieves its [test.model.MyStruct] value. Throws an exception if the [MyUnion] is not a
302+
* [Baz].
303+
*/
304+
public fun asBaz(): test.model.MyStruct = (this as MyUnion.Baz).value
305+
306+
/**
307+
* Casts this [MyUnion] as a [Baz] and retrieves its [test.model.MyStruct] value. Returns null if the [MyUnion] is not a [Baz].
308+
*/
309+
public fun asBazOrNull(): test.model.MyStruct? = (this as? MyUnion.Baz)?.value
310+
311+
/**
312+
* Casts this [MyUnion] as a [Foo] and retrieves its [kotlin.String] value. Throws an exception if the [MyUnion] is not a
313+
* [Foo].
314+
*/
315+
public fun asFoo(): kotlin.String = (this as MyUnion.Foo).value
316+
317+
/**
318+
* Casts this [MyUnion] as a [Foo] and retrieves its [kotlin.String] value. Returns null if the [MyUnion] is not a [Foo].
319+
*/
320+
public fun asFooOrNull(): kotlin.String? = (this as? MyUnion.Foo)?.value
321+
}
322+
""".trimIndent()
323+
324+
contents.shouldContainWithDiff(expectedClassDecl)
325+
}
326+
241327
private fun generateUnion(model: String): String {
242328
val fullModel = model.prependNamespaceAndService(namespace = "test").toSmithyModel()
243329

0 commit comments

Comments
 (0)