use crate::app::{AppBuilder, Plugin};
use crate::body::Body;
use crate::call::Call;
use crate::pipeline::{Middleware, Next, Phase};
use crate::response::{IntoResponse, Response};
use async_trait::async_trait;
use std::sync::Arc;
use tower_service::Service;
tokio::task_local! {
static IN_FLIGHT: std::cell::RefCell<Option<(Call, Next)>>;
}
#[derive(Clone, Copy, Default)]
pub struct NextService;
impl Service<http::Request<Body>> for NextService {
type Response = http::Response<Body>;
type Error = crate::Error;
type Future = std::pin::Pin<
Box<dyn std::future::Future<Output = Result<Self::Response, Self::Error>> + Send>,
>;
fn poll_ready(
&mut self,
_cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Result<(), Self::Error>> {
std::task::Poll::Ready(Ok(()))
}
fn call(&mut self, req: http::Request<Body>) -> Self::Future {
let taken = match IN_FLIGHT.try_with(|slot| slot.borrow_mut().take()) {
Ok(taken) => taken,
Err(_) => {
return Box::pin(async {
Err(crate::Error::internal(
"this tower layer drove the inner service from another task; \
the Churust adapter cannot cross task boundaries",
))
})
}
};
Box::pin(async move {
let Some((mut call, next)) = taken else {
return Err(crate::Error::internal(
"tower layer called the inner service more than once per request",
));
};
*call.headers_mut() = req.headers().clone();
call.set_uri(req.uri().clone());
let res = next.run(call).await;
let mut out = http::Response::builder().status(res.status);
if let Some(h) = out.headers_mut() {
*h = res.headers;
}
Ok(out.body(res.body).expect("response build is infallible"))
})
}
}
pub struct TowerMiddleware<S> {
service: std::sync::Mutex<S>,
phase: Phase,
}
impl<S> TowerMiddleware<S> {
pub fn new<L>(layer: L) -> Self
where
L: tower_layer::Layer<NextService, Service = S>,
{
TowerMiddleware {
service: std::sync::Mutex::new(layer.layer(NextService)),
phase: Phase::Plugins,
}
}
pub fn in_phase(mut self, phase: Phase) -> Self {
self.phase = phase;
self
}
}
#[async_trait]
impl<S, ResBody> Middleware for TowerMiddleware<S>
where
S: Service<http::Request<Body>, Response = http::Response<ResBody>> + Clone + Send + 'static,
S::Error: std::fmt::Display + Send,
S::Future: Send,
ResBody: http_body::Body<Data = bytes::Bytes> + Send + 'static,
ResBody::Error: std::fmt::Display,
{
async fn handle(&self, call: Call, next: Next) -> Response {
let mut service = match self.service.lock() {
Ok(s) => s.clone(),
Err(poisoned) => poisoned.into_inner().clone(),
};
let req = {
let mut b = http::Request::builder()
.method(call.method().clone())
.uri(call.uri().clone());
if let Some(h) = b.headers_mut() {
*h = call.headers().clone();
}
b.body(Body::empty()).expect("request build is infallible")
};
let slot = std::cell::RefCell::new(Some((call, next)));
let outcome = IN_FLIGHT
.scope(slot, async move {
futures_util::future::poll_fn(|cx| service.poll_ready(cx))
.await
.map_err(|e| {
tracing::error!(error = %e, "tower service not ready");
})?;
service.call(req).await.map_err(|e| {
tracing::error!(error = %e, "tower service failed");
})
})
.await;
match outcome {
Ok(res) => {
let (parts, body) = res.into_parts();
let mut out = Response::new(parts.status);
out.headers = parts.headers;
let exact = http_body::Body::size_hint(&body).exact();
out.body = Body::from_stream(http_body_util::BodyDataStream::new(body));
if let Some(len) = exact {
if !out.headers.contains_key(http::header::CONTENT_LENGTH) {
if let Ok(value) = http::HeaderValue::from_str(&len.to_string()) {
out.headers.insert(http::header::CONTENT_LENGTH, value);
}
}
}
out
}
Err(()) => crate::Error::internal("Internal Server Error").into_response(),
}
}
}
impl<S, ResBody> Plugin for TowerMiddleware<S>
where
S: Service<http::Request<Body>, Response = http::Response<ResBody>> + Clone + Send + 'static,
S::Error: std::fmt::Display + Send,
S::Future: Send,
ResBody: http_body::Body<Data = bytes::Bytes> + Send + 'static,
ResBody::Error: std::fmt::Display,
{
fn install(self: Box<Self>, app: &mut AppBuilder) {
let phase = self.phase;
app.add_middleware_in(phase, Arc::new(*self));
}
}