Fix HttpPrefixHeader and QueryParams detection

smithy-lang · Sep 16, 2022 · 3f78f00 · 3f78f00
1 parent dd4b292
commit 3f78f00
Show file tree

Hide file tree

Showing 2 changed files with 78 additions and 23 deletions.
diff --git a/...are/amazon/smithy/rust/codegen/server/smithy/generators/ServerHttpSensitivityGenerator.kt b/...are/amazon/smithy/rust/codegen/server/smithy/generators/ServerHttpSensitivityGenerator.kt
@@ -153,15 +153,9 @@ sealed class HeaderSensitivity(
     ) : HeaderSensitivity(headerKeys, runtimeConfig)
 
     /** Is there anything to redact? */
-    private fun hasRedactions(): Boolean {
-        return when (this) {
-            is NotSensitiveMapValue -> {
-                prefixHeader != null || headerKeys.isNotEmpty()
-            }
-            is SensitiveMapValue -> {
-                true
-            }
-        }
+    internal fun hasRedactions(): Boolean = headerKeys.isNotEmpty() || when (this) {
+        is NotSensitiveMapValue -> prefixHeader != null
+        is SensitiveMapValue -> true
     }
 
     /** Returns the type of the `MakeDebug`. */
@@ -244,14 +238,12 @@ sealed class QuerySensitivity(
     class SensitiveMapValue(allKeysSensitive: Boolean, runtimeConfig: RuntimeConfig) : QuerySensitivity(allKeysSensitive, runtimeConfig)
 
     /** Is there anything to redact? */
-    private fun hasRedactions(): Boolean {
-        return when (this) {
-            is NotSensitiveMapValue -> {
-                allKeysSensitive || queryKeys.isNotEmpty()
-            }
-            is SensitiveMapValue -> {
-                true
-            }
+    internal fun hasRedactions(): Boolean = when (this) {
+        is NotSensitiveMapValue -> {
+            allKeysSensitive || queryKeys.isNotEmpty()
+        }
+        is SensitiveMapValue -> {
+            true
         }
     }
 
@@ -391,8 +383,10 @@ class ServerHttpSensitivityGenerator(
         return if (valuesSensitive) {
             // All values are sensitive
             HeaderSensitivity.SensitiveMapValue(headerKeys, keySensitive, httpPrefixName, runtimeConfig)
-        } else {
+        } else if (keySensitive) {
             HeaderSensitivity.NotSensitiveMapValue(headerKeys, httpPrefixName, runtimeConfig)
+        } else {
+            HeaderSensitivity.NotSensitiveMapValue(headerKeys, null, runtimeConfig)
         }
     }
 

diff --git a/...amazon/smithy/rust/codegen/server/smithy/generators/ServerHttpSensitivityGeneratorTest.kt b/...amazon/smithy/rust/codegen/server/smithy/generators/ServerHttpSensitivityGeneratorTest.kt
@@ -128,7 +128,7 @@ class ServerHttpSensitivityGeneratorTest {
 
         val input = generator.input()!!
         val querySensitivity = generator.findQuerySensitivity(input)
-        assertEquals(querySensitivity.allKeysSensitive, false)
+        assert(!querySensitivity.allKeysSensitive)
         assertEquals((querySensitivity as QuerySensitivity.NotSensitiveMapValue).queryKeys, listOf("query_c"))
 
         val testProject = TestWorkspace.testProject(serverTestSymbolProvider(model))
@@ -176,7 +176,7 @@ class ServerHttpSensitivityGeneratorTest {
         val input = generator.input()!!
         val querySensitivity = generator.findQuerySensitivity(input)
 
-        assertEquals(querySensitivity.allKeysSensitive, true)
+        assert(querySensitivity.allKeysSensitive)
         querySensitivity as QuerySensitivity.SensitiveMapValue
 
         val testProject = TestWorkspace.testProject(serverTestSymbolProvider(model))
@@ -225,7 +225,7 @@ class ServerHttpSensitivityGeneratorTest {
 
         val input = generator.input()!!
         val querySensitivity = generator.findQuerySensitivity(input)
-        assertEquals(querySensitivity.allKeysSensitive, true)
+        assert(querySensitivity.allKeysSensitive)
         assert((querySensitivity as QuerySensitivity.NotSensitiveMapValue).queryKeys.isEmpty())
 
         val testProject = TestWorkspace.testProject(serverTestSymbolProvider(model))
@@ -273,7 +273,7 @@ class ServerHttpSensitivityGeneratorTest {
 
         val input = generator.input()!!
         val querySensitivity = generator.findQuerySensitivity(input)
-        assertEquals(querySensitivity.allKeysSensitive, false)
+        assert(!querySensitivity.allKeysSensitive)
         querySensitivity as QuerySensitivity.SensitiveMapValue
 
         val testProject = TestWorkspace.testProject(serverTestSymbolProvider(model))
@@ -292,6 +292,37 @@ class ServerHttpSensitivityGeneratorTest {
         testProject.compileAndTest()
     }
 
+    @Test
+    fun `query params none`() {
+        val model = """
+            namespace test
+
+            operation Secret {
+                input: Input,
+            }
+
+            structure Input {
+                @required
+                @httpQueryParams()
+                queryMap: QueryMap,
+            }
+
+            map QueryMap {
+                key: String,
+                value: String
+            }
+
+        """.asSmithyModel()
+        val operation = model.operationShapes.toList()[0]
+        val generator = ServerHttpSensitivityGenerator(model, operation, TestRuntimeConfig)
+
+        val input = generator.input()!!
+        val querySensitivity = generator.findQuerySensitivity(input)
+        assert(!querySensitivity.allKeysSensitive)
+        querySensitivity as QuerySensitivity.NotSensitiveMapValue
+        assert(!querySensitivity.hasRedactions())
+    }
+
     @Test
     fun `header closure`() {
         val model = """
@@ -392,6 +423,36 @@ class ServerHttpSensitivityGeneratorTest {
         testProject.compileAndTest()
     }
 
+    @Test
+    fun `prefix header none`() {
+        val model = """
+            namespace test
+
+            operation Secret {
+                input: Input,
+            }
+
+            structure Input {
+                @required
+                @httpPrefixHeaders("prefix-")
+                prefixMap: PrefixMap,
+            }
+
+            map PrefixMap {
+                key: String,
+                value: String
+            }
+
+        """.asSmithyModel()
+        val operation = model.operationShapes.toList()[0]
+        val generator = ServerHttpSensitivityGenerator(model, operation, TestRuntimeConfig)
+
+        val inputShape = operation.inputShape(model)
+        val headerData = generator.findHeaderSensitivity(inputShape)
+        headerData as HeaderSensitivity.NotSensitiveMapValue
+        assert(!headerData.hasRedactions())
+    }
+
     @Test
     fun `prefix headers key closure`() {
         val model = """
@@ -477,7 +538,7 @@ class ServerHttpSensitivityGeneratorTest {
         assert(headerData.headerKeys.isEmpty())
         val asSensitiveMapValue = (headerData as HeaderSensitivity.SensitiveMapValue)
         assertEquals(asSensitiveMapValue.prefixHeader, "prefix-")
-        assertEquals(asSensitiveMapValue.keySensitive, false)
+        assert(!asSensitiveMapValue.keySensitive)
 
         val testProject = TestWorkspace.testProject(serverTestSymbolProvider(model))
         testProject.lib { writer ->