diff --git a/rust-runtime/aws-smithy-http-server/examples/pokemon-service/src/bin/pokemon-service.rs b/rust-runtime/aws-smithy-http-server/examples/pokemon-service/src/bin/pokemon-service.rs index 8c68643e3bfabcdae96b0b76cfea82621b587726..402e57555abc4497bc97c52f0c622ce765155064 100644 --- a/rust-runtime/aws-smithy-http-server/examples/pokemon-service/src/bin/pokemon-service.rs +++ b/rust-runtime/aws-smithy-http-server/examples/pokemon-service/src/bin/pokemon-service.rs @@ -6,7 +6,7 @@ // This program is exported as a binary named `pokemon-service`. use std::{net::SocketAddr, sync::Arc}; -use aws_smithy_http_server::{plugin::PluginPipeline, AddExtensionLayer}; +use aws_smithy_http_server::{extension::OperationExtensionExt, plugin::PluginPipeline, AddExtensionLayer}; use clap::Parser; use pokemon_service::{ capture_pokemon, check_health, do_nothing, get_pokemon_species, get_server_statistics, get_storage, @@ -29,8 +29,13 @@ struct Args { pub async fn main() { let args = Args::parse(); setup_tracing(); - // Apply the `PrintPlugin` defined in `plugin.rs` - let plugins = PluginPipeline::new().print(); + let plugins = PluginPipeline::new() + // Apply the `PrintPlugin` defined in `plugin.rs` + .print() + // Apply the `OperationExtensionPlugin` defined in `aws_smithy_http_server::extension`. This allows other + // plugins or tests to access a `aws_smithy_http_server::extension::OperationExtension` from + // `Response::extensions`, or infer routing failure when it's missing. + .insert_operation_extension(); let app = PokemonService::builder_with_plugins(plugins) // Build a registry containing implementations to all the operations in the service. These // are async functions or async closures that take as input the operation's input and diff --git a/rust-runtime/aws-smithy-http-server/src/extension.rs b/rust-runtime/aws-smithy-http-server/src/extension.rs index b2f1ba3191bca831c1ac8b6d31c48c189acc288d..5c302a757afd5776ecbf8a75e9a5812714acbdb7 100644 --- a/rust-runtime/aws-smithy-http-server/src/extension.rs +++ b/rust-runtime/aws-smithy-http-server/src/extension.rs @@ -19,27 +19,27 @@ //! //! [extensions]: https://docs.rs/http/latest/http/struct.Extensions.html -use std::ops::Deref; +use std::{fmt, future::Future, ops::Deref, pin::Pin, task::Context, task::Poll}; +use futures_util::ready; +use futures_util::TryFuture; use thiserror::Error; +use tower::{layer::util::Stack, Layer, Service}; +use crate::operation::{Operation, OperationShape}; +use crate::plugin::{plugin_from_operation_name_fn, OperationNameFn, Plugin, PluginPipeline, PluginStack}; #[allow(deprecated)] use crate::request::RequestParts; -pub use crate::request::extension::Extension; -pub use crate::request::extension::MissingExtension; +pub use crate::request::extension::{Extension, MissingExtension}; /// Extension type used to store information about Smithy operations in HTTP responses. -/// This extension type is set when it has been correctly determined that the request should be -/// routed to a particular operation. The operation handler might not even get invoked because the -/// request fails to deserialize into the modeled operation input. +/// This extension type is inserted, via the [`OperationExtensionPlugin`], whenever it has been correctly determined +/// that the request should be routed to a particular operation. The operation handler might not even get invoked +/// because the request fails to deserialize into the modeled operation input. /// /// The format given must be the absolute shape ID with `#` replaced with a `.`. #[derive(Debug, Clone)] -#[deprecated( - since = "0.52.0", - note = "This is no longer inserted by the new service builder. Layers should be constructed per operation using the plugin system." -)] pub struct OperationExtension { absolute: &'static str, @@ -85,6 +85,117 @@ impl OperationExtension { } } +pin_project_lite::pin_project! { + /// The [`Service::Future`] of [`OperationExtensionService`] - inserts an [`OperationExtension`] into the + /// [`http::Response]`. + pub struct OperationExtensionFuture { + #[pin] + inner: Fut, + operation_extension: Option + } +} + +impl Future for OperationExtensionFuture +where + Fut: TryFuture>, +{ + type Output = Result, Fut::Error>; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let this = self.project(); + let ext = this + .operation_extension + .take() + .expect("futures cannot be polled after completion"); + let resp = ready!(this.inner.try_poll(cx)); + Poll::Ready(resp.map(|mut resp| { + resp.extensions_mut().insert(ext); + resp + })) + } +} + +/// Inserts a [`OperationExtension`] into the extensions of the [`http::Response`]. +#[derive(Debug, Clone)] +pub struct OperationExtensionService { + inner: S, + operation_extension: OperationExtension, +} + +impl Service> for OperationExtensionService +where + S: Service, Response = http::Response>, +{ + type Response = http::Response; + type Error = S::Error; + type Future = OperationExtensionFuture; + + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { + self.inner.poll_ready(cx) + } + + fn call(&mut self, req: http::Request) -> Self::Future { + OperationExtensionFuture { + inner: self.inner.call(req), + operation_extension: Some(self.operation_extension.clone()), + } + } +} + +/// A [`Layer`] applying the [`OperationExtensionService`] to an inner [`Service`]. +#[derive(Debug, Clone)] +pub struct OperationExtensionLayer(OperationExtension); + +impl Layer for OperationExtensionLayer { + type Service = OperationExtensionService; + + fn layer(&self, inner: S) -> Self::Service { + OperationExtensionService { + inner, + operation_extension: self.0.clone(), + } + } +} + +/// A [`Plugin`] which applies [`OperationExtensionLayer`] to every operation. +pub struct OperationExtensionPlugin(OperationNameFn OperationExtensionLayer>); + +impl fmt::Debug for OperationExtensionPlugin { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_tuple("OperationExtensionPlugin").field(&"...").finish() + } +} + +impl Plugin for OperationExtensionPlugin +where + Op: OperationShape, +{ + type Service = S; + type Layer = Stack; + + fn map(&self, input: Operation) -> Operation { + OperationExtensionLayer> as Plugin>::map(&self.0, input) + } +} + +/// An extension trait on [`PluginPipeline`] allowing the application of [`OperationExtensionPlugin`]. +/// +/// See [`module`](crate::extension) documentation for more info. +pub trait OperationExtensionExt

{ + /// Apply the [`OperationExtensionPlugin`], which inserts the [`OperationExtension`] into every [`http::Response`]. + fn insert_operation_extension(self) -> PluginPipeline>; +} + +impl

OperationExtensionExt

for PluginPipeline

{ + fn insert_operation_extension(self) -> PluginPipeline> { + let plugin = OperationExtensionPlugin(plugin_from_operation_name_fn(|name| { + let operation_extension = OperationExtension::new(name).expect("Operation name is malformed, this should never happen. Please file an issue against https://github.com/awslabs/smithy-rs"); + OperationExtensionLayer(operation_extension) + })); + self.push(plugin) + } +} + /// Extension type used to store the type of user-modeled error returned by an operation handler. /// These are modeled errors, defined in the Smithy model. #[derive(Debug, Clone)] @@ -157,7 +268,6 @@ where } #[cfg(test)] -#[allow(deprecated)] mod tests { use super::*;