use std::sync::Arc;
use crate::axum::body::Body;
use crate::axum::extract::Request;
use crate::axum::response::Response;
#[derive(Clone)]
pub struct ErrorMapFn(Arc<dyn Fn(Response) -> Response + Send + Sync + 'static>);
impl ErrorMapFn {
#[must_use]
pub fn new<F>(f: F) -> Self
where
F: Fn(Response) -> Response + Send + Sync + 'static,
{
Self(Arc::new(f))
}
#[must_use]
pub fn map(&self, response: Response) -> Response {
(self.0)(response)
}
}
impl std::fmt::Debug for ErrorMapFn {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_tuple("ErrorMapFn").finish_non_exhaustive()
}
}
#[derive(Clone)]
pub struct ErrorMappingLayer {
map: Option<ErrorMapFn>,
}
impl ErrorMappingLayer {
#[must_use]
pub fn new(map: Option<ErrorMapFn>) -> Self {
Self { map }
}
}
impl<S> tower::Layer<S> for ErrorMappingLayer {
type Service = ErrorMappingService<S>;
fn layer(&self, inner: S) -> Self::Service {
ErrorMappingService {
inner,
map: self.map.clone(),
}
}
}
#[derive(Clone)]
pub struct ErrorMappingService<S> {
inner: S,
map: Option<ErrorMapFn>,
}
impl<S> tower::Service<Request<Body>> for ErrorMappingService<S>
where
S: tower::Service<Request<Body>, Response = Response, Error = std::convert::Infallible>
+ Clone
+ Send
+ 'static,
S::Future: Send + 'static,
{
type Response = Response;
type Error = std::convert::Infallible;
type Future = std::pin::Pin<
Box<dyn std::future::Future<Output = Result<Response, std::convert::Infallible>> + Send>,
>;
fn poll_ready(
&mut self,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), Self::Error>> {
self.inner.poll_ready(cx)
}
fn call(&mut self, req: Request<Body>) -> Self::Future {
let clone = self.inner.clone();
let mut inner = std::mem::replace(&mut self.inner, clone);
let map = self.map.clone();
Box::pin(async move {
let response = inner.call(req).await.unwrap_or_else(|inf| match inf {});
match map {
Some(f) => Ok(f.map(response)),
None => Ok(response),
}
})
}
}