use std::future::Future;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};
use http::{Request, Response};
use tower::Service;
use crate::http::{Body, BoxError, HttpService};
pub struct Next {
inner: Arc<Mutex<HttpService>>,
}
impl Next {
pub async fn run(self, req: Request<Body>) -> Result<Response<Body>, BoxError> {
std::future::poll_fn(|cx| self.inner.lock().unwrap().poll_ready(cx)).await?;
let fut = self.inner.lock().unwrap().call(req);
fut.await
}
}
pub fn from_fn<F, Fut>(f: F) -> MiddlewareFnLayer<F>
where
F: Fn(Request<Body>, Next) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<Response<Body>, BoxError>> + Send + 'static,
{
MiddlewareFnLayer { f: Arc::new(f) }
}
pub struct MiddlewareFnLayer<F> {
f: Arc<F>,
}
impl<F, Fut> tower::Layer<HttpService> for MiddlewareFnLayer<F>
where
F: Fn(Request<Body>, Next) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<Response<Body>, BoxError>> + Send + 'static,
{
type Service = MiddlewareFnService<F>;
fn layer(&self, inner: HttpService) -> Self::Service {
MiddlewareFnService {
f: self.f.clone(),
inner: Arc::new(Mutex::new(inner)),
}
}
}
pub struct MiddlewareFnService<F> {
f: Arc<F>,
inner: Arc<Mutex<HttpService>>,
}
impl<F, Fut> Service<Request<Body>> for MiddlewareFnService<F>
where
F: Fn(Request<Body>, Next) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Result<Response<Body>, BoxError>> + Send + 'static,
{
type Response = Response<Body>;
type Error = BoxError;
type Future = Pin<Box<dyn Future<Output = Result<Response<Body>, BoxError>> + Send>>;
fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn call(&mut self, req: Request<Body>) -> Self::Future {
let next = Next {
inner: self.inner.clone(),
};
let f = self.f.clone();
Box::pin(async move { f(req, next).await })
}
}