Skip to content

Commit

Permalink
Fix label sensitivity
Browse files Browse the repository at this point in the history
  • Loading branch information
Harry Barber committed Sep 16, 2022
1 parent 3f78f00 commit 9883173
Show file tree
Hide file tree
Showing 8 changed files with 272 additions and 253 deletions.
109 changes: 58 additions & 51 deletions ...are/amazon/smithy/rust/codegen/server/smithy/generators/ServerHttpSensitivityGenerator.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,6 @@ import software.amazon.smithy.rust.codegen.client.rustlang.Writable
import software.amazon.smithy.rust.codegen.client.rustlang.asType
import software.amazon.smithy.rust.codegen.client.rustlang.plus
import software.amazon.smithy.rust.codegen.client.rustlang.rust
import software.amazon.smithy.rust.codegen.client.rustlang.rustBlock
import software.amazon.smithy.rust.codegen.client.rustlang.rustBlockTemplate
import software.amazon.smithy.rust.codegen.client.rustlang.rustTemplate
import software.amazon.smithy.rust.codegen.client.rustlang.withBlock
Expand All @@ -39,13 +38,6 @@ import software.amazon.smithy.rust.codegen.core.util.hasTrait
import software.amazon.smithy.rust.codegen.server.smithy.ServerCargoDependency
import java.util.*

internal fun findUriGreedyLabelPosition(uriPattern: UriPattern): Int? {
return uriPattern
.greedyLabel
.orElse(null)
?.let { uriPattern.toString().indexOf("$it") }
}

/** Models the ways status codes can be bound and sensitive. */
class StatusCodeSensitivity(private val sensitive: Boolean, runtimeConfig: RuntimeConfig) {
private val codegenScope = arrayOf(
Expand All @@ -69,58 +61,69 @@ class StatusCodeSensitivity(private val sensitive: Boolean, runtimeConfig: Runti
}
}

/** Models the ways labels can be bound and sensitive. */
sealed class LabelSensitivity(runtimeConfig: RuntimeConfig) {
private val codegenScope = arrayOf("SmithyHttpServer" to ServerCargoDependency.SmithyHttpServer(runtimeConfig).asType())
data class GreedyLabel(
val index: Int,
val suffix: String,
)

class Normal(val indexes: List<Int>, runtimeConfig: RuntimeConfig) : LabelSensitivity(runtimeConfig) {
/** Returns the closure used during construction. */
fun closure(): Writable = writable {
withBlock("{", "} as fn(_) -> _") {
rustBlock("|index: usize|") {
if (indexes.isNotEmpty()) {
withBlock("matches!(index,", ")") {
val matches = indexes.joinToString("|") { "$it" }
rust(matches)
}
} else {
rust("{_ = index; false}")
}
}
}
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)
}
class Greedy(val suffixPosition: Int, runtimeConfig: RuntimeConfig) : LabelSensitivity(runtimeConfig)

fun hasRedactions(): Boolean = when (this) {
is Normal -> indexes.isNotEmpty()
is Greedy -> true
/** 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,
)
}
private fun hasRedactions(): Boolean = labelIndexes.isNotEmpty() || greedyLabel != null

/** Returns the type of the `MakeFmt`. */
fun type(): Writable = if (hasRedactions()) when (this) {
is Normal -> writable {
rustTemplate("#{SmithyHttpServer}::logging::sensitivity::uri::MakeLabel<fn(usize) -> bool>", *codegenScope)
}
is Greedy -> writable {
rustTemplate("#{SmithyHttpServer}::logging::sensitivity::uri::MakeGreedyLabel", *codegenScope)
}
fun type(): Writable = if (hasRedactions()) writable {
rustTemplate("#{SmithyHttpServer}::logging::sensitivity::uri::MakeLabel<fn(usize) -> bool>", *codegenScope)
} else writable {
rustTemplate("#{SmithyHttpServer}::logging::MakeIdentity", *codegenScope)
}

/** Returns the setter enclosing the closure or suffix position. */
fun setter(): Writable = if (hasRedactions()) when (this) {
is Normal -> writable {
rustTemplate(".label(#{Closure:W})", "Closure" to closure())
}
is Greedy -> {
val suffixPosition = suffixPosition
writable {
rust(".greedy_label($suffixPosition)")
}
/** Returns the value of the `GreedyLabel`. */
private fun greedyLabelStruct(): Writable = writable {
if (greedyLabel != null) {
rustTemplate(
"""
Some(#{SmithyHttpServer}::logging::sensitivity::uri::GreedyLabel::new(${greedyLabel.index}, "${greedyLabel.suffix}"))""",
*codegenScope,
)
} else {
rust("None")
}
} else writable {}
}

/** Returns the setter enclosing the closure or suffix position. */
fun setter(): Writable = if (hasRedactions()) writable {
rustTemplate(".label(#{Closure:W}, #{GreedyLabel:W})", "Closure" to closure(), "GreedyLabel" to greedyLabelStruct())
} else writable { }
}

/** Models the ways headers can be bound and sensitive */
Expand Down Expand Up @@ -384,8 +387,10 @@ class ServerHttpSensitivityGenerator(
// All values are sensitive
HeaderSensitivity.SensitiveMapValue(headerKeys, keySensitive, httpPrefixName, runtimeConfig)
} else if (keySensitive) {
// Only keys are sensitive
HeaderSensitivity.NotSensitiveMapValue(headerKeys, httpPrefixName, runtimeConfig)
} else {
// No values are sensitive
HeaderSensitivity.NotSensitiveMapValue(headerKeys, null, runtimeConfig)
}
}
Expand Down Expand Up @@ -432,9 +437,11 @@ class ServerHttpSensitivityGenerator(
}

/** Constructs `LabelSensitivity` of a `Shape` */
internal fun findLabelSensitivity(uriPattern: UriPattern, rootShape: Shape): LabelSensitivity {
return findUriGreedyLabelPosition(uriPattern)?.let { LabelSensitivity.Greedy(it, runtimeConfig) } ?: findUriLabelIndexes(uriPattern, rootShape).let { LabelSensitivity.Normal(it, runtimeConfig) }
}
internal fun findLabelSensitivity(uriPattern: UriPattern, rootShape: Shape): LabelSensitivity = LabelSensitivity(
findUriLabelIndexes(uriPattern, rootShape),
findGreedyLabel(uriPattern),
runtimeConfig,
)

// Find member shapes with trait `B` contained in a shape enjoying `A`.
// [trait|A] ~> [trait|B]
Expand Down
49 changes: 26 additions & 23 deletions ...amazon/smithy/rust/codegen/server/smithy/generators/ServerHttpSensitivityGeneratorTest.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@

package software.amazon.smithy.rust.codegen.server.smithy.generators

import io.kotest.assertions.fail
import org.junit.jupiter.api.Assertions.assertEquals
import org.junit.jupiter.api.Test
import software.amazon.smithy.model.pattern.UriPattern
Expand Down Expand Up @@ -34,11 +33,19 @@ class ServerHttpSensitivityGeneratorTest {
)

@Test
fun `find greedy label`() {
fun `find greedy label end`() {
val uri = "/pokemon-species/{name+}"
val pattern = UriPattern.parse(uri)
val position = findUriGreedyLabelPosition(pattern)!!
assertEquals(position, 17)
val value = findGreedyLabel(pattern)!!
assertEquals(value, GreedyLabel(1, ""))
}

@Test
fun `find greedy label`() {
val uri = "/pokemon-species/{name+}/ash/ketchum"
val pattern = UriPattern.parse(uri)
val value = findGreedyLabel(pattern)!!
assertEquals(value, GreedyLabel(1, "/ash/ketchum"))
}

@Test
Expand Down Expand Up @@ -591,26 +598,22 @@ class ServerHttpSensitivityGeneratorTest {
val labeledUriIndexes = generator.findUriLabelIndexes(uri, input)
assertEquals(labeledUriIndexes, listOf(2, 1))

when (val labelData = generator.findLabelSensitivity(uri, input)) {
is LabelSensitivity.Greedy -> fail("expected normal http label")
is LabelSensitivity.Normal -> {
val testProject = TestWorkspace.testProject(serverTestSymbolProvider(model))
testProject.lib { writer ->
writer.unitTest("uri_closure") {
rustTemplate(
"""
let closure = #{Closure:W};
assert_eq!(closure(0), false);
assert_eq!(closure(1), true);
assert_eq!(closure(2), true);
""",
"Closure" to labelData.closure(),
*codegenScope,
)
}
}
testProject.compileAndTest()
val labelData = generator.findLabelSensitivity(uri, input)
val testProject = TestWorkspace.testProject(serverTestSymbolProvider(model))
testProject.lib { writer ->
writer.unitTest("uri_closure") {
rustTemplate(
"""
let closure = #{Closure:W};
assert_eq!(closure(0), false);
assert_eq!(closure(1), true);
assert_eq!(closure(2), true);
""",
"Closure" to labelData.closure(),
*codegenScope,
)
}
}
testProject.compileAndTest()
}
}
2 changes: 1 addition & 1 deletion rust-runtime/aws-smithy-http-server/src/logging/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@
//! key_suffix: None,
//! })
//! .query(|name| QueryMarker { key: false, value: name == "bar" })
//! .label(|index| index % 2 == 0);
//! .label(|index| index % 2 == 0, None);
//! let response_fmt = ResponseFmt::new()
//! .header(|name| {
//! if name.as_str().starts_with("prefix-") {
Expand Down
26 changes: 10 additions & 16 deletions rust-runtime/aws-smithy-http-server/src/logging/sensitivity/request.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ use crate::logging::{MakeFmt, MakeIdentity};

use super::{
headers::{HeaderMarker, MakeHeaders},
uri::{MakeGreedyLabel, MakeLabel, MakeQuery, MakeUri, QueryMarker},
uri::{GreedyLabel, MakeLabel, MakeQuery, MakeUri, QueryMarker},
};

/// Allows the modification the requests URIs [`Display`](std::fmt::Display) and headers
Expand Down Expand Up @@ -70,27 +70,21 @@ impl<Header, P, Q> RequestFmt<Header, MakeUri<P, Q>> {
/// Marks parts of the URI as sensitive.
///
/// See [`Label`](super::uri::Label) for more info.
pub fn label<F>(self, label: F) -> RequestFmt<Header, MakeUri<MakeLabel<F>, Q>>
pub fn label<F>(
self,
label_marker: F,
greedy_label: Option<GreedyLabel>,
) -> RequestFmt<Header, MakeUri<MakeLabel<F>, Q>>
where
F: Fn(usize) -> bool,
{
RequestFmt {
headers: self.headers,
uri: MakeUri {
make_path: MakeLabel(label),
make_query: self.uri.make_query,
},
}
}

/// Marks parts of the URI as sensitive.
///
/// See [`GreedyLabel`](super::uri::GreedyLabel) for more info.
pub fn greedy_label(self, position: usize) -> RequestFmt<Header, MakeUri<MakeGreedyLabel, Q>> {
RequestFmt {
headers: self.headers,
uri: MakeUri {
make_path: MakeGreedyLabel(position),
make_path: MakeLabel {
label_marker,
greedy_label,
},
make_query: self.uri.make_query,
},
}
Expand Down
113 changes: 0 additions & 113 deletions rust-runtime/aws-smithy-http-server/src/logging/sensitivity/uri/greedy_label.rs
View file

This file was deleted.

Loading

0 comments on commit 9883173

Please sign in to comment.