Vend tracing instrumentation as a Plugin (#1738)

* Add InstrumentLayer and InstrumentPlugin

* Add Sensitivity codegen

codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/ServerOperationGenerator.kt

@@ -45,6 +45,10 @@ class ServerOperationGenerator(
     fun render(writer: RustWriter) {
         writer.documentShape(operation, model)
 
+        val generator = ServerHttpSensitivityGenerator(model, operation, runtimeConfig)
+        val requestFmt = generator.requestFmt()
+        val responseFmt = generator.responseFmt()
+
         writer.rustTemplate(
             """
             pub struct $operationName;
@@ -56,8 +60,25 @@ class ServerOperationGenerator(
                 type Output = crate::output::${operationName}Output;
                 type Error = #{Error:W};
             }
+
+            impl #{SmithyHttpServer}::logging::sensitivity::Sensitivity for $operationName {
+                type RequestFmt = #{RequestType:W};
+                type ResponseFmt = #{ResponseType:W};
+
+                fn request_fmt() -> Self::RequestFmt {
+                    #{RequestValue:W}
+                }
+
+                fn response_fmt() -> Self::ResponseFmt {
+                    #{ResponseValue:W}
+                }
+            }
             """,
             "Error" to operationError(),
+            "RequestValue" to requestFmt.value,
+            "RequestType" to requestFmt.type,
+            "ResponseValue" to responseFmt.value,
+            "ResponseType" to responseFmt.type,
             *codegenScope,
         )
         // Adds newline to end of render

rust-runtime/aws-smithy-http-server/src/logging/layer.rs

+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+use tower::Layer;
+
+use super::{InstrumentOperation, MakeIdentity};
+
+/// A [`Layer`] used to apply [`InstrumentOperation`].
+#[derive(Debug)]
+pub struct InstrumentLayer<RequestMakeFmt = MakeIdentity, ResponseMakeFmt = MakeIdentity> {
+    operation_name: &'static str,
+    make_request: RequestMakeFmt,
+    make_response: ResponseMakeFmt,
+}
+
+impl InstrumentLayer {
+    /// Constructs a new [`InstrumentLayer`] with no data redacted.
+    pub fn new(operation_name: &'static str) -> Self {
+        Self {
+            operation_name,
+            make_request: MakeIdentity,
+            make_response: MakeIdentity,
+        }
+    }
+}
+
+impl<RequestMakeFmt, ResponseMakeFmt> InstrumentLayer<RequestMakeFmt, ResponseMakeFmt> {
+    /// Configures the request format.
+    ///
+    /// The argument is typically [`RequestFmt`](super::sensitivity::RequestFmt).
+    pub fn request_fmt<R>(self, make_request: R) -> InstrumentLayer<R, ResponseMakeFmt> {
+        InstrumentLayer {
+            operation_name: self.operation_name,
+            make_request,
+            make_response: self.make_response,
+        }
+    }
+
+    /// Configures the response format.
+    ///
+    /// The argument is typically [`ResponseFmt`](super::sensitivity::ResponseFmt).
+    pub fn response_fmt<R>(self, make_response: R) -> InstrumentLayer<RequestMakeFmt, R> {
+        InstrumentLayer {
+            operation_name: self.operation_name,
+            make_request: self.make_request,
+            make_response,
+        }
+    }
+}
+
+impl<S, RequestMakeFmt, ResponseMakeFmt> Layer<S> for InstrumentLayer<RequestMakeFmt, ResponseMakeFmt>
+where
+    RequestMakeFmt: Clone,
+    ResponseMakeFmt: Clone,
+{
+    type Service = InstrumentOperation<S, RequestMakeFmt, ResponseMakeFmt>;
+
+    fn layer(&self, service: S) -> Self::Service {
+        InstrumentOperation::new(service, self.operation_name)
+            .request_fmt(self.make_request.clone())
+            .response_fmt(self.make_response.clone())
+    }
+}

rust-runtime/aws-smithy-http-server/src/logging/mod.rs

@@ -57,11 +57,15 @@
 //!
 //! [sensitive trait]: https://awslabs.github.io/smithy/1.0/spec/core/documentation-traits.html?highlight=sensitive%20trait#sensitive-trait
 
+mod layer;
+mod plugin;
 pub mod sensitivity;
 mod service;
 
 use std::fmt::{Debug, Display};
 
+pub use layer::*;
+pub use plugin::*;
 pub use service::*;
 
 /// A standard interface for taking some component of the HTTP request/response and transforming it into new struct

rust-runtime/aws-smithy-http-server/src/logging/plugin.rs

+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+use tower::layer::util::Stack;
+
+use crate::{
+    operation::{Operation, OperationShape},
+    plugin::{Pluggable, Plugin},
+};
+
+use super::{layer::InstrumentLayer, sensitivity::Sensitivity};
+
+/// An [`Plugin`] which applies [`InstrumentLayer`] to all operations in the builder.
+#[derive(Debug)]
+pub struct InstrumentPlugin;
+
+impl<P, Op, S, L> Plugin<P, Op, S, L> for InstrumentPlugin
+where
+    Op: OperationShape,
+    Op: Sensitivity,
+{
+    type Service = S;
+    type Layer = Stack<L, InstrumentLayer<Op::RequestFmt, Op::ResponseFmt>>;
+
+    fn map(&self, operation: Operation<S, L>) -> Operation<Self::Service, Self::Layer> {
+        let layer = InstrumentLayer::new(Op::NAME)
+            .request_fmt(Op::request_fmt())
+            .response_fmt(Op::response_fmt());
+        operation.layer(layer)
+    }
+}
+
+/// An extension trait for applying [`InstrumentLayer`] to all operations.
+pub trait InstrumentExt: Pluggable<InstrumentPlugin> {
+    /// Applies an [`InstrumentLayer`] to all operations which respects the [@sensitive] trait given on the input and
+    /// output models. See [`InstrumentOperation`](super::InstrumentOperation) for more information.
+    ///
+    /// [@sensitive]: https://awslabs.github.io/smithy/2.0/spec/documentation-traits.html#sensitive-trait
+    fn instrument(self) -> Self::Output
+    where
+        Self: Sized,
+    {
+        self.apply(InstrumentPlugin)
+    }
+}
+
+impl<Builder> InstrumentExt for Builder where Builder: Pluggable<InstrumentPlugin> {}

rust-runtime/aws-smithy-http-server/src/logging/sensitivity/mod.rs

@@ -15,9 +15,27 @@ mod response;
 mod sensitive;
 pub mod uri;
 
+use http::{HeaderMap, StatusCode, Uri};
 pub use request::*;
 pub use response::*;
 pub use sensitive::*;
 
+use super::{MakeDebug, MakeDisplay};
+
 /// The string placeholder for redacted data.
 pub const REDACTED: &str = "{redacted}";
+
+/// An interface for providing [`MakeDebug`] and [`MakeDisplay`] for [`Request`](http::Request) and
+/// [`Response`](http::Response).
+pub trait Sensitivity {
+    /// The [`MakeDebug`] and [`MakeDisplay`] for the request [`HeaderMap`] and [`Uri`].
+    type RequestFmt: for<'a> MakeDebug<&'a HeaderMap> + for<'a> MakeDisplay<&'a Uri>;
+    /// The [`MakeDebug`] and [`MakeDisplay`] for the response [`HeaderMap`] and [`StatusCode`].
+    type ResponseFmt: for<'a> MakeDebug<&'a HeaderMap> + MakeDisplay<StatusCode>;
+
+    /// Returns the [`RequestFmt`](Sensitivity::RequestFmt).
+    fn request_fmt() -> Self::RequestFmt;
+
+    /// Returns the [`ResponseFmt`](Sensitivity::ResponseFmt).
+    fn response_fmt() -> Self::ResponseFmt;
+}