diff --git a/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/ServerHttpSensitivityGenerator.kt b/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/ServerHttpSensitivityGenerator.kt index a6d0d8b3e2..41365b7c15 100644 --- a/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/ServerHttpSensitivityGenerator.kt +++ b/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/ServerHttpSensitivityGenerator.kt @@ -61,42 +61,32 @@ class StatusCodeSensitivity(private val sensitive: Boolean, runtimeConfig: Runti } } +/** Represents the information needed to specify the position of a greedy label. */ data class GreedyLabel( - val index: Int, - val suffix: String, + // The segment index the greedy label. + val segmentIndex: Int, + // The number of characters from the end of the URI the greedy label terminates. + val endOffset: Int, ) -internal fun findGreedyLabel(uriPattern: UriPattern): GreedyLabel? = uriPattern - .segments - .asIterable() - .withIndex() - .find { (_, segment) -> - segment.isGreedyLabel - } - ?.let { (index, segment) -> - val remainingSegments = uriPattern.segments.asIterable().drop(index + 1) - val suffix = if (remainingSegments.isNotEmpty()) { - remainingSegments.joinToString(prefix = "/", separator = "/") - } else { - "" - } - GreedyLabel(index, suffix) - } - /** Models the ways labels can be bound and sensitive. */ class LabelSensitivity(private val labelIndexes: List, private val greedyLabel: GreedyLabel?, runtimeConfig: RuntimeConfig) { private val codegenScope = arrayOf("SmithyHttpServer" to ServerCargoDependency.SmithyHttpServer(runtimeConfig).asType()) /** Returns the closure used during construction. */ fun closure(): Writable = writable { - rustTemplate( - """ - { - |index: usize| matches!(index, ${labelIndexes.joinToString("|")}) - } as fn(_) -> _ - """, - *codegenScope, - ) + if (labelIndexes.isNotEmpty()) { + rustTemplate( + """ + { + |index: usize| matches!(index, ${labelIndexes.joinToString("|")}) + } as fn(_) -> _ + """, + *codegenScope, + ) + } else { + rust("{ |_index: usize| false } as fn(_) -> _") + } } private fun hasRedactions(): Boolean = labelIndexes.isNotEmpty() || greedyLabel != null @@ -112,7 +102,7 @@ class LabelSensitivity(private val labelIndexes: List, private val greedyLa if (greedyLabel != null) { rustTemplate( """ - Some(#{SmithyHttpServer}::logging::sensitivity::uri::GreedyLabel::new(${greedyLabel.index}, "${greedyLabel.suffix}"))""", + Some(#{SmithyHttpServer}::logging::sensitivity::uri::GreedyLabel::new(${greedyLabel.segmentIndex}, ${greedyLabel.endOffset}))""", *codegenScope, ) } else { @@ -428,20 +418,41 @@ class ServerHttpSensitivityGenerator( return QuerySensitivity.NotSensitiveMapValue(queries, keysSensitive, runtimeConfig) } - internal fun findUriLabelIndexes(uriPattern: UriPattern, rootShape: Shape): List { - val uriLabels: Map = uriPattern + /** Constructs `LabelSensitivity` of a `Shape` */ + internal fun findLabelSensitivity(uriPattern: UriPattern, rootShape: Shape): LabelSensitivity { + val sensitiveLabels = findSensitiveBound(rootShape) + + val labelMap: Map = uriPattern .segments .withIndex() .filter { (_, segment) -> segment.isLabel }.associate { (index, segment) -> Pair(segment.content, index) } - return findSensitiveBound(rootShape).mapNotNull { uriLabels[it.memberName] } - } + val labelsIndex = sensitiveLabels.mapNotNull { labelMap[it.memberName] } - /** Constructs `LabelSensitivity` of a `Shape` */ - internal fun findLabelSensitivity(uriPattern: UriPattern, rootShape: Shape): LabelSensitivity = LabelSensitivity( - findUriLabelIndexes(uriPattern, rootShape), - findGreedyLabel(uriPattern), - runtimeConfig, - ) + val greedyLabel = uriPattern + .segments + .asIterable() + .withIndex() + .find { (_, segment) -> + segment.isGreedyLabel + }?.let { (index, segment) -> + // Check if sensitive + if (sensitiveLabels.find { it.memberName == segment.content } != null) { + index + } else { + null + } + }?.let { index -> + val remainingSegments = uriPattern.segments.asIterable().drop(index + 1) + val suffix = if (remainingSegments.isNotEmpty()) { + remainingSegments.joinToString(prefix = "/", separator = "/") + } else { + "" + } + GreedyLabel(index, suffix.length) + } + + return LabelSensitivity(labelsIndex, greedyLabel, runtimeConfig) + } // Find member shapes with trait `B` contained in a shape enjoying `A`. // [trait|A] ~> [trait|B] diff --git a/rust-runtime/aws-smithy-http-server/src/logging/sensitivity/uri/label.rs b/rust-runtime/aws-smithy-http-server/src/logging/sensitivity/uri/label.rs index cd48908ae1..7845f8cc84 100644 --- a/rust-runtime/aws-smithy-http-server/src/logging/sensitivity/uri/label.rs +++ b/rust-runtime/aws-smithy-http-server/src/logging/sensitivity/uri/label.rs @@ -102,10 +102,9 @@ where if hit_greedy { if let Some(end_index) = self.path.len().checked_sub(greedy_label.end_offset) { - if greedy_start + 1 <= end_index { + if greedy_start < end_index { let greedy_redaction = Sensitive(&self.path[greedy_start + 1..end_index]); let remainder = &self.path[end_index..]; - println!("GREEDY: {greedy_redaction} vs {}", &self.path[greedy_start..end_index]); write!(f, "/{greedy_redaction}{remainder}")?; } else { write!(f, "{}", &self.path[greedy_start..])?;