From 46ffa8fe5c58d4f29422598a60fc28d4fc0e388c Mon Sep 17 00:00:00 2001 From: Harry Barber Date: Tue, 20 Sep 2022 17:22:37 +0000 Subject: [PATCH] Add InstrumentLayer and InstrumentPlugin --- .../src/logging/layer.rs | 65 +++++++++++++++++++ .../aws-smithy-http-server/src/logging/mod.rs | 4 ++ .../src/logging/plugin.rs | 47 ++++++++++++++ .../src/logging/sensitivity/mod.rs | 18 +++++ 4 files changed, 134 insertions(+) create mode 100644 rust-runtime/aws-smithy-http-server/src/logging/layer.rs create mode 100644 rust-runtime/aws-smithy-http-server/src/logging/plugin.rs diff --git a/rust-runtime/aws-smithy-http-server/src/logging/layer.rs b/rust-runtime/aws-smithy-http-server/src/logging/layer.rs new file mode 100644 index 00000000000..956fd6b24dc --- /dev/null +++ b/rust-runtime/aws-smithy-http-server/src/logging/layer.rs @@ -0,0 +1,65 @@ +/* + * 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 { + 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 InstrumentLayer { + /// Configures the request format. + /// + /// The argument is typically [`RequestFmt`](super::sensitivity::RequestFmt). + pub fn request_fmt(self, make_request: R) -> InstrumentLayer { + 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(self, make_response: R) -> InstrumentLayer { + InstrumentLayer { + operation_name: self.operation_name, + make_request: self.make_request, + make_response, + } + } +} + +impl Layer for InstrumentLayer +where + RequestMakeFmt: Clone, + ResponseMakeFmt: Clone, +{ + type Service = InstrumentOperation; + + 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()) + } +} diff --git a/rust-runtime/aws-smithy-http-server/src/logging/mod.rs b/rust-runtime/aws-smithy-http-server/src/logging/mod.rs index 646179730dd..99f6f94869e 100644 --- a/rust-runtime/aws-smithy-http-server/src/logging/mod.rs +++ b/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 diff --git a/rust-runtime/aws-smithy-http-server/src/logging/plugin.rs b/rust-runtime/aws-smithy-http-server/src/logging/plugin.rs new file mode 100644 index 00000000000..494964ab3d5 --- /dev/null +++ b/rust-runtime/aws-smithy-http-server/src/logging/plugin.rs @@ -0,0 +1,47 @@ +/* + * 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 Plugin for InstrumentPlugin +where + Op: OperationShape, + Op: Sensitivity, +{ + type Service = S; + type Layer = Stack>; + + fn map(&self, operation: Operation) -> Operation { + 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 { + /// Applies [`InstrumentLayer`] to all operations. See [`InstrumentOperation`](super::InstrumentOperation) for more + /// information. + fn instrument(self) -> Self::Output + where + Self: Sized, + { + self.apply(InstrumentPlugin) + } +} + +impl InstrumentExt for Builder where Builder: Pluggable {} diff --git a/rust-runtime/aws-smithy-http-server/src/logging/sensitivity/mod.rs b/rust-runtime/aws-smithy-http-server/src/logging/sensitivity/mod.rs index ca1485fbe43..fff18b3d5c8 100644 --- a/rust-runtime/aws-smithy-http-server/src/logging/sensitivity/mod.rs +++ b/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 [`Uri`]. + type ResponseFmt: for<'a> MakeDebug<&'a HeaderMap> + MakeDisplay; + + /// Returns the [`RequestFmt`](Sensitivity::RequestFmt). + fn request_fmt() -> Self::RequestFmt; + + /// Returns the [`ResponseFmt`](Sensitivity::ResponseFmt). + fn response_fmt() -> Self::ResponseFmt; +}