Skip to content

Commit

Permalink
Add protocol specific routers (#1666)
Browse files Browse the repository at this point in the history
* Add protocol specific routers

* Replace internals of `Router` with protocol specific routers
  • Loading branch information
hlbarber authored Aug 26, 2022
1 parent 30d985a commit 778539f
Show file tree
Hide file tree
Showing 8 changed files with 544 additions and 179 deletions.
12 changes: 12 additions & 0 deletions rust-runtime/aws-smithy-http-server/src/protocols.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,18 @@
use crate::rejection::MissingContentTypeReason;
use crate::request::RequestParts;

/// [AWS REST JSON 1.0 Protocol](https://awslabs.github.io/smithy/2.0/aws/protocols/aws-restjson1-protocol.html).
pub struct AwsRestJson1;

/// [AWS REST XML Protocol](https://awslabs.github.io/smithy/2.0/aws/protocols/aws-restxml-protocol.html).
pub struct AwsRestXml;

/// [AWS JSON 1.0 Protocol](https://awslabs.github.io/smithy/2.0/aws/protocols/aws-json-1_0-protocol.html).
pub struct AwsJson10;

/// [AWS JSON 1.1 Protocol](https://awslabs.github.io/smithy/2.0/aws/protocols/aws-json-1_1-protocol.html).
pub struct AwsJson11;

/// Supported protocols.
#[derive(Debug, Clone, Copy)]
pub enum Protocol {
Expand Down
6 changes: 6 additions & 0 deletions rust-runtime/aws-smithy-http-server/src/response.rs
Original file line number Diff line number Diff line change
Expand Up @@ -36,3 +36,9 @@ use crate::body::BoxBody;

#[doc(hidden)]
pub type Response<T = BoxBody> = http::Response<T>;

/// A protocol aware function taking `self` to [`http::Response`].
pub trait IntoResponse<Protocol> {
/// Performs a conversion into a [`http::Response`].
fn into_response(self) -> http::Response<BoxBody>;
}
22 changes: 2 additions & 20 deletions rust-runtime/aws-smithy-http-server/src/routing/future.rs
Original file line number Diff line number Diff line change
Expand Up @@ -33,29 +33,11 @@
*/

//! Future types.
use crate::body::BoxBody;
use futures_util::future::Either;
use http::{Request, Response};
use std::{convert::Infallible, future::ready};
use tower::util::Oneshot;
use super::Route;
pub use super::{into_make_service::IntoMakeService, route::RouteFuture};

type OneshotRoute<B> = Oneshot<super::Route<B>, Request<B>>;
type ReadyResponse = std::future::Ready<Result<Response<BoxBody>, Infallible>>;

opaque_future! {
/// Response future for [`Router`](super::Router).
pub type RouterFuture<B> =
futures_util::future::Either<OneshotRoute<B>, ReadyResponse>;
}

impl<B> RouterFuture<B> {
pub(super) fn from_oneshot(future: Oneshot<super::Route<B>, Request<B>>) -> Self {
Self::new(Either::Left(future))
}

pub(super) fn from_response(response: Response<BoxBody>) -> Self {
Self::new(Either::Right(ready(Ok(response))))
}
pub type RouterFuture<B> = super::routers::RoutingFuture<Route<B>, B>;
}
213 changes: 58 additions & 155 deletions rust-runtime/aws-smithy-http-server/src/routing/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,18 +8,17 @@
//! [Smithy specification]: https://awslabs.github.io/smithy/1.0/spec/core/http-traits.html
use self::request_spec::RequestSpec;
use self::tiny_map::TinyMap;
use self::routers::{aws_json::AwsJsonRouter, rest::RestRouter, RoutingService};
use crate::body::{boxed, Body, BoxBody, HttpBody};
use crate::error::BoxError;
use crate::protocols::Protocol;
use crate::runtime_error::{RuntimeError, RuntimeErrorKind};
use http::{Request, Response, StatusCode};
use crate::protocols::{AwsJson10, AwsJson11, AwsRestJson1, AwsRestXml};

use http::{Request, Response};
use std::{
convert::Infallible,
task::{Context, Poll},
};
use tower::layer::Layer;
use tower::util::ServiceExt;
use tower::{Service, ServiceBuilder};
use tower_http::map_response_body::MapResponseBodyLayer;

Expand All @@ -31,6 +30,7 @@ mod lambda_handler;
pub mod request_spec;

mod route;
mod routers;
mod tiny_map;

pub use self::lambda_handler::LambdaHandler;
Expand Down Expand Up @@ -61,11 +61,6 @@ pub struct Router<B = Body> {
routes: Routes<B>,
}

// This constant determines when the `TinyMap` implementation switches from being a `Vec` to a
// `HashMap`. This is chosen to be 15 as a result of the discussion around
// https://github.com/awslabs/smithy-rs/pull/1429#issuecomment-1147516546
const ROUTE_CUTOFF: usize = 15;

/// Protocol-aware routes types.
///
/// RestJson1 and RestXml routes are stored in a `Vec` because there can be multiple matches on the
Expand All @@ -75,10 +70,10 @@ const ROUTE_CUTOFF: usize = 15;
/// directly found in the `X-Amz-Target` HTTP header.
#[derive(Debug)]
enum Routes<B = Body> {
RestXml(Vec<(Route<B>, RequestSpec)>),
RestJson1(Vec<(Route<B>, RequestSpec)>),
AwsJson10(TinyMap<String, Route<B>, ROUTE_CUTOFF>),
AwsJson11(TinyMap<String, Route<B>, ROUTE_CUTOFF>),
RestXml(RoutingService<RestRouter<Route<B>>, AwsRestXml>),
RestJson1(RoutingService<RestRouter<Route<B>>, AwsRestJson1>),
AwsJson10(RoutingService<AwsJsonRouter<Route<B>>, AwsJson10>),
AwsJson11(RoutingService<AwsJsonRouter<Route<B>>, AwsJson11>),
}

impl<B> Clone for Router<B> {
Expand All @@ -104,29 +99,6 @@ impl<B> Router<B>
where
B: Send + 'static,
{
/// Return the correct, protocol-specific "Not Found" response for an unknown operation.
fn unknown_operation(&self) -> RouterFuture<B> {
let protocol = match &self.routes {
Routes::RestJson1(_) => Protocol::RestJson1,
Routes::RestXml(_) => Protocol::RestXml,
Routes::AwsJson10(_) => Protocol::AwsJson10,
Routes::AwsJson11(_) => Protocol::AwsJson11,
};
let error = RuntimeError {
protocol,
kind: RuntimeErrorKind::UnknownOperation,
};
RouterFuture::from_response(error.into_response())
}

/// Return the HTTP error response for non allowed method.
fn method_not_allowed(&self) -> RouterFuture<B> {
RouterFuture::from_response({
let mut res = Response::new(crate::body::empty());
*res.status_mut() = StatusCode::METHOD_NOT_ALLOWED;
res
})
}
/// Convert this router into a [`MakeService`], that is a [`Service`] whose
/// response is another service.
///
Expand Down Expand Up @@ -155,46 +127,21 @@ where
NewResBody::Error: Into<BoxError>,
{
let layer = ServiceBuilder::new()
.layer_fn(Route::new)
.layer(MapResponseBodyLayer::new(boxed))
.layer(layer);
match self.routes {
Routes::RestJson1(routes) => {
let routes = routes
.into_iter()
.map(|(route, request_spec)| (Layer::layer(&layer, route), request_spec))
.collect();
Router {
routes: Routes::RestJson1(routes),
}
}
Routes::RestXml(routes) => {
let routes = routes
.into_iter()
.map(|(route, request_spec)| (Layer::layer(&layer, route), request_spec))
.collect();
Router {
routes: Routes::RestXml(routes),
}
}
Routes::AwsJson10(routes) => {
let routes = routes
.into_iter()
.map(|(operation, route)| (operation, Layer::layer(&layer, route)))
.collect();
Router {
routes: Routes::AwsJson10(routes),
}
}
Routes::AwsJson11(routes) => {
let routes = routes
.into_iter()
.map(|(operation, route)| (operation, Layer::layer(&layer, route)))
.collect();
Router {
routes: Routes::AwsJson11(routes),
}
}
Routes::RestJson1(routes) => Router {
routes: Routes::RestJson1(routes.map(|router| router.layer(layer).boxed())),
},
Routes::RestXml(routes) => Router {
routes: Routes::RestXml(routes.map(|router| router.layer(layer).boxed())),
},
Routes::AwsJson10(routes) => Router {
routes: Routes::AwsJson10(routes.map(|router| router.layer(layer).boxed())),
},
Routes::AwsJson11(routes) => Router {
routes: Routes::AwsJson11(routes.map(|router| router.layer(layer).boxed())),
},
}
}

Expand All @@ -211,18 +158,14 @@ where
),
>,
{
let mut routes: Vec<(Route<B>, RequestSpec)> = routes
.into_iter()
.map(|(svc, request_spec)| (Route::from_box_clone_service(svc), request_spec))
.collect();

// Sort them once by specifity, with the more specific routes sorted before the less
// specific ones, so that when routing a request we can simply iterate through the routes
// and pick the first one that matches.
routes.sort_by_key(|(_route, request_spec)| std::cmp::Reverse(request_spec.rank()));

let svc = RoutingService::new(
routes
.into_iter()
.map(|(svc, request_spec)| (request_spec, Route::from_box_clone_service(svc)))
.collect(),
);
Self {
routes: Routes::RestJson1(routes),
routes: Routes::RestJson1(svc),
}
}

Expand All @@ -239,18 +182,14 @@ where
),
>,
{
let mut routes: Vec<(Route<B>, RequestSpec)> = routes
.into_iter()
.map(|(svc, request_spec)| (Route::from_box_clone_service(svc), request_spec))
.collect();

// Sort them once by specifity, with the more specific routes sorted before the less
// specific ones, so that when routing a request we can simply iterate through the routes
// and pick the first one that matches.
routes.sort_by_key(|(_route, request_spec)| std::cmp::Reverse(request_spec.rank()));

let svc = RoutingService::new(
routes
.into_iter()
.map(|(svc, request_spec)| (request_spec, Route::from_box_clone_service(svc)))
.collect(),
);
Self {
routes: Routes::RestXml(routes),
routes: Routes::RestXml(svc),
}
}

Expand All @@ -267,13 +206,15 @@ where
),
>,
{
let routes = routes
.into_iter()
.map(|(svc, operation)| (operation, Route::from_box_clone_service(svc)))
.collect();
let svc = RoutingService::new(
routes
.into_iter()
.map(|(svc, operation)| (operation, Route::from_box_clone_service(svc)))
.collect(),
);

Self {
routes: Routes::AwsJson10(routes),
routes: Routes::AwsJson10(svc),
}
}

Expand All @@ -290,13 +231,15 @@ where
),
>,
{
let routes = routes
.into_iter()
.map(|(svc, operation)| (operation, Route::from_box_clone_service(svc)))
.collect();
let svc = RoutingService::new(
routes
.into_iter()
.map(|(svc, operation)| (operation, Route::from_box_clone_service(svc)))
.collect(),
);

Self {
routes: Routes::AwsJson11(routes),
routes: Routes::AwsJson11(svc),
}
}
}
Expand All @@ -316,55 +259,15 @@ where

#[inline]
fn call(&mut self, req: Request<B>) -> Self::Future {
match &self.routes {
let fut = match &mut self.routes {
// REST routes.
Routes::RestJson1(routes) | Routes::RestXml(routes) => {
let mut method_not_allowed = false;

// Loop through all the routes and validate if any of them matches. Routes are already ranked.
for (route, request_spec) in routes {
match request_spec.matches(&req) {
request_spec::Match::Yes => {
return RouterFuture::from_oneshot(route.clone().oneshot(req));
}
request_spec::Match::MethodNotAllowed => method_not_allowed = true,
// Continue looping to see if another route matches.
request_spec::Match::No => continue,
}
}

if method_not_allowed {
// The HTTP method is not correct.
self.method_not_allowed()
} else {
// In any other case return the `RuntimeError::UnknownOperation`.
self.unknown_operation()
}
}
Routes::RestJson1(routes) => routes.call(req),
Routes::RestXml(routes) => routes.call(req),
// AwsJson routes.
Routes::AwsJson10(routes) | Routes::AwsJson11(routes) => {
if req.uri() == "/" {
// Check the request method for POST.
if req.method() == http::Method::POST {
// Find the `x-amz-target` header.
if let Some(target) = req.headers().get("x-amz-target") {
if let Ok(target) = target.to_str() {
// Lookup in the `TinyMap` for a route for the target.
let route = routes.get(target);
if let Some(route) = route {
return RouterFuture::from_oneshot(route.clone().oneshot(req));
}
}
}
} else {
// The HTTP method is not POST.
return self.method_not_allowed();
}
}
// In any other case return the `RuntimeError::UnknownOperation`.
self.unknown_operation()
}
}
Routes::AwsJson10(routes) => routes.call(req),
Routes::AwsJson11(routes) => routes.call(req),
};
RouterFuture::new(fut)
}
}

Expand All @@ -376,7 +279,7 @@ mod rest_tests {
routing::request_spec::*,
};
use futures_util::Future;
use http::{HeaderMap, Method};
use http::{HeaderMap, Method, StatusCode};
use std::pin::Pin;

/// Helper function to build a `Request`. Used in other test modules.
Expand Down Expand Up @@ -601,7 +504,7 @@ mod awsjson_tests {
use super::*;
use crate::body::boxed;
use futures_util::Future;
use http::{HeaderMap, HeaderValue, Method};
use http::{HeaderMap, HeaderValue, Method, StatusCode};
use pretty_assertions::assert_eq;
use std::pin::Pin;

Expand Down
Loading

0 comments on commit 778539f

Please sign in to comment.