Fix GreedyLabel logic

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<Int>, 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<Int>, 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<Int> {
-        val uriLabels: Map<String, Int> = uriPattern
+    /** Constructs `LabelSensitivity` of a `Shape` */
+    internal fun findLabelSensitivity(uriPattern: UriPattern, rootShape: Shape): LabelSensitivity {
+        val sensitiveLabels = findSensitiveBound<HttpLabelTrait>(rootShape)
+
+        val labelMap: Map<String, Int> = uriPattern
             .segments
             .withIndex()
             .filter { (_, segment) -> segment.isLabel }.associate { (index, segment) -> Pair(segment.content, index) }
-        return findSensitiveBound<HttpLabelTrait>(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]

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..])?;