diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/StructureGenerator.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/StructureGenerator.kt index c1b93c7243..91dd67d871 100644 --- a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/StructureGenerator.kt +++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/StructureGenerator.kt @@ -36,6 +36,7 @@ import software.amazon.smithy.rust.codegen.core.smithy.traits.SyntheticInputTrai 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.hasTrait +import software.amazon.smithy.rust.codegen.core.util.redactIfNecessary fun RustWriter.implBlock(structureShape: Shape, symbolProvider: SymbolProvider, block: RustWriter.() -> Unit) { rustBlock("impl ${symbolProvider.toSymbol(structureShape).name}") { @@ -43,14 +44,6 @@ fun RustWriter.implBlock(structureShape: Shape, symbolProvider: SymbolProvider, } } -fun redactIfNecessary(member: MemberShape, model: Model, safeToPrint: String): String { - return if (member.getMemberTrait(model, SensitiveTrait::class.java).isPresent) { - "*** Sensitive Data Redacted ***".dq() - } else { - safeToPrint - } -} - open class StructureGenerator( val model: Model, private val symbolProvider: RustSymbolProvider, @@ -114,9 +107,7 @@ open class StructureGenerator( rust("""let mut formatter = f.debug_struct(${name.dq()});""") members.forEach { member -> val memberName = symbolProvider.toMemberName(member) - val fieldValue = redactIfNecessary( - member, model, "self.$memberName", - ) + val fieldValue = member.redactIfNecessary(model, "self.$memberName") rust( "formatter.field(${memberName.dq()}, &$fieldValue);", ) diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/http/HttpBindingGenerator.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/http/HttpBindingGenerator.kt index 2de37bbf81..c2590675e7 100644 --- a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/http/HttpBindingGenerator.kt +++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/http/HttpBindingGenerator.kt @@ -39,7 +39,6 @@ import software.amazon.smithy.rust.codegen.client.smithy.CoreCodegenContext import software.amazon.smithy.rust.codegen.client.smithy.RuntimeType import software.amazon.smithy.rust.codegen.client.smithy.generators.CodegenTarget import software.amazon.smithy.rust.codegen.client.smithy.generators.operationBuildError -import software.amazon.smithy.rust.codegen.client.smithy.generators.redactIfNecessary import software.amazon.smithy.rust.codegen.client.smithy.makeOptional import software.amazon.smithy.rust.codegen.client.smithy.mapRustType import software.amazon.smithy.rust.codegen.client.smithy.protocols.HttpBindingDescriptor @@ -54,6 +53,7 @@ import software.amazon.smithy.rust.codegen.core.util.inputShape import software.amazon.smithy.rust.codegen.core.util.isPrimitive import software.amazon.smithy.rust.codegen.core.util.isStreaming import software.amazon.smithy.rust.codegen.core.util.outputShape +import software.amazon.smithy.rust.codegen.core.util.redactIfNecessary import software.amazon.smithy.rust.codegen.core.util.toSnakeCase /** @@ -492,11 +492,7 @@ class HttpBindingGenerator( let header_value = $safeName; let header_value = http::header::HeaderValue::try_from(&*header_value).map_err(|err| { #{build_error}::InvalidField { field: "$memberName", details: format!("`{}` cannot be used as a header value: {}", &${ - redactIfNecessary( - memberShape, - model, - "header_value", - ) + memberShape.redactIfNecessary(model, "header_value") }, err)} })?; builder = builder.header("${httpBinding.locationName}", header_value); @@ -532,11 +528,7 @@ class HttpBindingGenerator( #{build_error}::InvalidField { field: "$memberName", details: format!("`{}` cannot be used as a header value: {}", ${ - redactIfNecessary( - memberShape, - model, - "v", - ) + memberShape.redactIfNecessary(model, "v") }, err)} })?; builder = builder.header(header_name, header_value); 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 index e7b7894478..c4b547658d 100644 --- 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 @@ -16,6 +16,7 @@ import software.amazon.smithy.model.shapes.Shape import software.amazon.smithy.model.shapes.ShapeId import software.amazon.smithy.model.shapes.StructureShape import software.amazon.smithy.model.shapes.UnionShape +import software.amazon.smithy.model.traits.SensitiveTrait import software.amazon.smithy.model.traits.StreamingTrait import software.amazon.smithy.model.traits.Trait import software.amazon.smithy.rust.codegen.core.smithy.traits.SyntheticInputTrait @@ -81,6 +82,16 @@ fun ServiceShape.hasEventStreamOperations(model: Model): Boolean = operations.an model.expectShape(id, OperationShape::class.java).isEventStream(model) } +fun Shape.redactIfNecessary(model: Model, safeToPrint: String): String = + when (this) { + is MemberShape -> model.expectShape(this.target).redactIfNecessary(model, safeToPrint) + else -> if (this.hasTrait()) { + "*** Sensitive Data Redacted ***".dq() + } else { + safeToPrint + } + } + /* * Returns the member of this structure targeted with streaming trait (if it exists). *