use futures_core::ready;
use http::{Request, Response};
use pin_project_lite::pin_project;
use std::{
future::Future,
pin::Pin,
task::{Context, Poll},
};
use tower_layer::Layer;
use tower_service::Service;
#[derive(Debug, Clone)]
pub struct AsyncRequireAuthorizationLayer<T> {
auth: T,
}
impl<T> AsyncRequireAuthorizationLayer<T> {
pub fn new(auth: T) -> AsyncRequireAuthorizationLayer<T> {
Self { auth }
}
}
impl<S, T> Layer<S> for AsyncRequireAuthorizationLayer<T>
where
T: Clone,
{
type Service = AsyncRequireAuthorization<S, T>;
fn layer(&self, inner: S) -> Self::Service {
AsyncRequireAuthorization::new(inner, self.auth.clone())
}
}
#[derive(Clone, Debug)]
pub struct AsyncRequireAuthorization<S, T> {
inner: S,
auth: T,
}
impl<S, T> AsyncRequireAuthorization<S, T> {
define_inner_service_accessors!();
}
impl<S, T> AsyncRequireAuthorization<S, T> {
pub fn new(inner: S, auth: T) -> AsyncRequireAuthorization<S, T> {
Self { inner, auth }
}
pub fn layer(auth: T) -> AsyncRequireAuthorizationLayer<T> {
AsyncRequireAuthorizationLayer::new(auth)
}
}
impl<ReqBody, ResBody, S, Auth> Service<Request<ReqBody>> for AsyncRequireAuthorization<S, Auth>
where
Auth: AsyncAuthorizeRequest<ReqBody, ResponseBody = ResBody>,
S: Service<Request<Auth::RequestBody>, Response = Response<ResBody>> + Clone,
{
type Response = Response<ResBody>;
type Error = S::Error;
type Future = ResponseFuture<Auth, S, ReqBody>;
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.inner.poll_ready(cx)
}
fn call(&mut self, req: Request<ReqBody>) -> Self::Future {
let inner = self.inner.clone();
let authorize = self.auth.authorize(req);
ResponseFuture {
state: State::Authorize { authorize },
service: inner,
}
}
}
pin_project! {
pub struct ResponseFuture<Auth, S, ReqBody>
where
Auth: AsyncAuthorizeRequest<ReqBody>,
S: Service<Request<Auth::RequestBody>>,
{
#[pin]
state: State<Auth::Future, S::Future>,
service: S,
}
}
pin_project! {
#[project = StateProj]
enum State<A, SFut> {
Authorize {
#[pin]
authorize: A,
},
Authorized {
#[pin]
fut: SFut,
},
}
}
impl<Auth, S, ReqBody, B> Future for ResponseFuture<Auth, S, ReqBody>
where
Auth: AsyncAuthorizeRequest<ReqBody, ResponseBody = B>,
S: Service<Request<Auth::RequestBody>, Response = Response<B>>,
{
type Output = Result<Response<B>, S::Error>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let mut this = self.project();
loop {
match this.state.as_mut().project() {
StateProj::Authorize { authorize } => {
let auth = ready!(authorize.poll(cx));
match auth {
Ok(req) => {
let fut = this.service.call(req);
this.state.set(State::Authorized { fut })
}
Err(res) => {
return Poll::Ready(Ok(res));
}
};
}
StateProj::Authorized { fut } => {
return fut.poll(cx);
}
}
}
}
}
pub trait AsyncAuthorizeRequest<B> {
type RequestBody;
type ResponseBody;
type Future: Future<Output = Result<Request<Self::RequestBody>, Response<Self::ResponseBody>>>;
fn authorize(&mut self, request: Request<B>) -> Self::Future;
}
impl<B, F, Fut, ReqBody, ResBody> AsyncAuthorizeRequest<B> for F
where
F: FnMut(Request<B>) -> Fut,
Fut: Future<Output = Result<Request<ReqBody>, Response<ResBody>>>,
{
type RequestBody = ReqBody;
type ResponseBody = ResBody;
type Future = Fut;
fn authorize(&mut self, request: Request<B>) -> Self::Future {
self(request)
}
}
#[cfg(test)]
mod tests {
#[allow(unused_imports)]
use super::*;
use futures_util::future::BoxFuture;
use http::{header, StatusCode};
use hyper::Body;
use tower::{BoxError, ServiceBuilder, ServiceExt};
#[derive(Clone, Copy)]
struct MyAuth;
impl<B> AsyncAuthorizeRequest<B> for MyAuth
where
B: Send + 'static,
{
type RequestBody = B;
type ResponseBody = Body;
type Future = BoxFuture<'static, Result<Request<B>, Response<Self::ResponseBody>>>;
fn authorize(&mut self, mut request: Request<B>) -> Self::Future {
Box::pin(async move {
let authorized = request
.headers()
.get(header::AUTHORIZATION)
.and_then(|it| it.to_str().ok())
.and_then(|it| it.strip_prefix("Bearer "))
.map(|it| it == "69420")
.unwrap_or(false);
if authorized {
let user_id = UserId("6969".to_owned());
request.extensions_mut().insert(user_id);
Ok(request)
} else {
Err(Response::builder()
.status(StatusCode::UNAUTHORIZED)
.body(Body::empty())
.unwrap())
}
})
}
}
#[derive(Debug)]
struct UserId(String);
#[tokio::test]
async fn require_async_auth_works() {
let mut service = ServiceBuilder::new()
.layer(AsyncRequireAuthorizationLayer::new(MyAuth))
.service_fn(echo);
let request = Request::get("/")
.header(header::AUTHORIZATION, "Bearer 69420")
.body(Body::empty())
.unwrap();
let res = service.ready().await.unwrap().call(request).await.unwrap();
assert_eq!(res.status(), StatusCode::OK);
}
#[tokio::test]
async fn require_async_auth_401() {
let mut service = ServiceBuilder::new()
.layer(AsyncRequireAuthorizationLayer::new(MyAuth))
.service_fn(echo);
let request = Request::get("/")
.header(header::AUTHORIZATION, "Bearer deez")
.body(Body::empty())
.unwrap();
let res = service.ready().await.unwrap().call(request).await.unwrap();
assert_eq!(res.status(), StatusCode::UNAUTHORIZED);
}
async fn echo(req: Request<Body>) -> Result<Response<Body>, BoxError> {
Ok(Response::new(req.into_body()))
}
}