Skip to main content

tachyon_web/routing/
tower_compat.rs

1//! Optional interop with the `tower` ecosystem (`tower::Service` / `tower::Layer`).
2//!
3//! This module exists so real Axum apps that mount a pre-built Tower service
4//! (e.g. `tower_http::services::ServeDir`, a `tonic` gRPC service) or apply a
5//! `tower::Layer` (tracing, compression, timeouts, concurrency limits) can still
6//! be ported without a rewrite.
7//!
8//! **This is deliberately not the recommended path.** Every call into a Tower
9//! `Service` goes through an extra `Box<dyn Future>` indirection and (for
10//! layers) a fresh `Service` value per request — overhead the native
11//! `.hoop()`/`.hoop_at()` middleware system is built to avoid. Prefer native
12//! middleware; reach for this only to bridge in existing Tower/tower-http code.
13
14use crate::http::error::Error;
15use crate::http::response::{Body, IntoResponse};
16use crate::routing::handler::{BoxedFuture, Handler, ResponseFuture};
17use crate::routing::middleware::Next;
18use bytes::Bytes;
19use hyper::{Request, Response};
20use std::future::Future;
21use std::pin::Pin;
22use std::sync::Arc;
23use std::task::{Context, Poll};
24use tower::{Layer, Service, ServiceExt};
25
26/// Marker type parameter for [`Handler`] impls backed by a raw `tower::Service`.
27#[derive(Debug)]
28pub struct TowerServiceMarker;
29
30/// Wraps a `tower::Service` so it can be registered as a route handler.
31///
32/// Used by [`Router::route_service`](crate::routing::Router::route_service),
33/// [`Router::nest_service`](crate::routing::Router::nest_service), and
34/// [`Router::fallback_service`](crate::routing::Router::fallback_service).
35#[derive(Debug, Clone)]
36pub struct ServiceHandler<Svc> {
37    pub(crate) service: Svc,
38    /// When `Some(prefix)`, the request URI's path is rewritten via
39    /// [`crate::routing::strip_uri_prefix`] to strip this exact leading byte
40    /// sequence (used by `nest_service`, matching Axum's nested-service path
41    /// rewriting — see that function's docs for the percent-encoding caveat).
42    pub(crate) strip_prefix: Option<Arc<str>>,
43}
44
45impl<Svc, RespBody, S> Handler<TowerServiceMarker, S> for ServiceHandler<Svc>
46where
47    S: Send + Sync + 'static,
48    Svc: Service<Request<Bytes>, Response = Response<RespBody>> + Clone + Send + Sync + 'static,
49    Svc::Future: Send + 'static,
50    Svc::Error: Into<Error> + Send,
51    RespBody: hyper::body::Body<Data = Bytes> + Send + 'static,
52    RespBody::Error: Into<Error>,
53{
54    fn call(self, mut req: Request<Body>, _state: Arc<S>) -> BoxedFuture {
55        if let Some(prefix) = &self.strip_prefix {
56            crate::routing::strip_uri_prefix(&mut req, prefix);
57        }
58
59        let mut service = self.service;
60        ResponseFuture::Boxed(Box::pin(async move {
61            // Tower services conventionally expect a fully-buffered body; only
62            // native tachyon handlers (via `BodyStream`/`Request<Body>`) get the
63            // option of true streaming.
64            let limit = crate::routing::extract::max_body_size(req.extensions());
65            let (parts, body) = req.into_parts();
66            let bytes = match body.collect_bytes(limit).await {
67                Ok(b) => b,
68                Err(e) => return e.into_response(),
69            };
70            let req = Request::from_parts(parts, bytes);
71            match service.ready().await {
72                Ok(ready) => match ready.call(req).await {
73                    Ok(resp) => {
74                        let (parts, body) = resp.into_parts();
75                        Response::from_parts(parts, Body::stream(body))
76                    }
77                    Err(e) => Into::<Error>::into(e).into_response(),
78                },
79                Err(e) => Into::<Error>::into(e).into_response(),
80            }
81        }))
82    }
83}
84
85/// A one-shot `tower::Service` adapter over a [`Next`] continuation — lets a
86/// `tower::Layer` wrap "the rest of the tachyon pipeline" for this request.
87pub struct NextService<S> {
88    next: Option<Next<S>>,
89}
90
91impl<S> std::fmt::Debug for NextService<S> {
92    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
93        f.debug_struct("NextService").finish_non_exhaustive()
94    }
95}
96
97impl<S> NextService<S> {
98    pub(crate) const fn new(next: Next<S>) -> Self {
99        Self { next: Some(next) }
100    }
101}
102
103impl<S: Send + Sync + 'static> Service<Request<Bytes>> for NextService<S> {
104    type Response = Response<Body>;
105    type Error = std::convert::Infallible;
106    type Future = Pin<Box<dyn Future<Output = Result<Response<Body>, Self::Error>> + Send>>;
107
108    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
109        Poll::Ready(Ok(()))
110    }
111
112    fn call(&mut self, req: Request<Bytes>) -> Self::Future {
113        let next = self.next.take();
114        Box::pin(async move {
115            let Some(next) = next else {
116                unreachable!("NextService called more than once for the same request")
117            };
118            let (parts, bytes) = req.into_parts();
119            Ok(next
120                .run(Request::from_parts(parts, Body::full(bytes)))
121                .await)
122        })
123    }
124}
125
126/// Adapts a `tower::Layer` into a native tachyon middleware closure, so it can
127/// be installed via the same `.hoop_at()` machinery as any other middleware.
128pub(crate) fn from_tower_layer<L, S, RespBody>(
129    layer: L,
130) -> impl Fn(Request<Body>, Next<S>) -> ResponseFuture + Clone + Send + Sync + 'static
131where
132    S: Send + Sync + 'static,
133    L: Layer<NextService<S>> + Clone + Send + Sync + 'static,
134    L::Service: Service<Request<Bytes>, Response = Response<RespBody>> + Send + 'static,
135    <L::Service as Service<Request<Bytes>>>::Future: Send + 'static,
136    <L::Service as Service<Request<Bytes>>>::Error: Into<Error> + Send,
137    RespBody: hyper::body::Body<Data = Bytes> + Send + 'static,
138    RespBody::Error: Into<Error>,
139{
140    move |req, next| {
141        let layer = layer.clone();
142        ResponseFuture::Boxed(Box::pin(async move {
143            // As with `ServiceHandler`, the Tower side of the bridge always sees a
144            // fully-buffered body.
145            let limit = crate::routing::extract::max_body_size(req.extensions());
146            let (parts, body) = req.into_parts();
147            let bytes = match body.collect_bytes(limit).await {
148                Ok(b) => b,
149                Err(e) => return e.into_response(),
150            };
151            let req = Request::from_parts(parts, bytes);
152            let mut layered = layer.layer(NextService::new(next));
153            match layered.ready().await {
154                Ok(ready) => match ready.call(req).await {
155                    Ok(resp) => {
156                        let (parts, body) = resp.into_parts();
157                        Response::from_parts(parts, Body::stream(body))
158                    }
159                    Err(e) => Into::<Error>::into(e).into_response(),
160                },
161                Err(e) => Into::<Error>::into(e).into_response(),
162            }
163        }))
164    }
165}
166
167/// Lets a compiled router be driven directly as a `tower::Service` — the
168/// idiomatic Axum testing pattern (`app.oneshot(request).await`, via
169/// `tower::ServiceExt`) works unchanged against a `CompiledRouter`, and a
170/// `CompiledRouter` can be handed to any other Tower/Hyper API expecting a
171/// `Service<Request<B>>`.
172///
173/// Unlike Axum (which only implements this for `Router<()>`, since a
174/// `Router<S>` for `S != ()` hasn't been given its state yet), this is
175/// implemented for `CompiledRouter<S>` for **any** state type: a compiled
176/// router is always fully self-contained (its state was bound at `compile()`
177/// time), so there's no equivalent "not runnable yet" state to restrict this to.
178impl<S, B> Service<Request<B>> for crate::routing::CompiledRouter<S>
179where
180    S: Clone + Send + Sync + 'static,
181    B: hyper::body::Body<Data = Bytes> + Send + 'static,
182    B::Error: Into<Error>,
183{
184    type Response = Response<Body>;
185    type Error = std::convert::Infallible;
186    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
187
188    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
189        Poll::Ready(Ok(()))
190    }
191
192    fn call(&mut self, req: Request<B>) -> Self::Future {
193        let (parts, body) = req.into_parts();
194        let req = Request::from_parts(parts, Body::stream(body));
195        let this = self.clone();
196        Box::pin(async move { Ok(this.handle_request(req).await) })
197    }
198}
199
200/// Lets an uncompiled, stateless [`crate::routing::Router`] be driven
201/// directly as a `tower::Service` — e.g. `Router::new().route(...).oneshot(req)`
202/// — with no separate `.compile()` call, matching `axum::Router`'s drop-in
203/// ergonomics exactly: build with `.route()`/`.nest()`/`.merge()`/`.hoop()`/
204/// etc., then hand the same value straight to `.oneshot()`, a `tower`
205/// server, or anything else expecting a `Service`.
206///
207/// Internally this compiles the `matchit` route tree **once**, the first
208/// time `call()` runs, and caches the result in `Router`'s private
209/// `compiled` field — every builder method that mutates the route table
210/// resets that cache, so it's impossible to silently dispatch against a
211/// stale tree. Every call after the first is exactly as cheap as calling
212/// the already-`CompiledRouter` directly; the split from `axum::Router`
213/// (which has no separate compiled form at all) is now purely internal.
214///
215/// Deliberately restricted to `Router<()>`, matching Axum exactly (Axum only
216/// implements `Service` for `Router<()>` too — a `Router<S>` for `S != ()`
217/// hasn't been given its state yet, so there's nothing meaningful to serve).
218/// This is what makes plain `Router::new().route(...).oneshot(req)` type-check
219/// with no turbofish: `Service` has exactly one impl to unify against, same
220/// as in Axum. A router built with real shared state (`State<T>` extractors)
221/// still needs `.with_state(actual_state)` first, same as Axum — at that
222/// point it's already a `Router<()>` too. For testing a still-generic
223/// `Router<S>`/`CompiledRouter<S>` for `S != ()` directly, use
224/// [`CompiledRouter`](crate::routing::CompiledRouter)'s broader impl above
225/// via an explicit `.compile()`.
226///
227/// ```rust,no_run
228/// use tachyon_web::{Router, get};
229/// use tower::ServiceExt;
230///
231/// async fn handler() -> &'static str { "hi" }
232///
233/// # async fn build() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
234/// let app = Router::new().route("/", get(handler));
235/// let req = hyper::Request::builder().uri("/").body(http_body_util::Full::new(bytes::Bytes::new()))?;
236/// let resp = app.oneshot(req).await?;
237/// # let _ = resp;
238/// # Ok(())
239/// # }
240/// ```
241impl<B> Service<Request<B>> for crate::routing::Router<()>
242where
243    B: hyper::body::Body<Data = Bytes> + Send + 'static,
244    B::Error: Into<Error>,
245{
246    type Response = Response<Body>;
247    type Error = std::convert::Infallible;
248    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
249
250    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
251        Poll::Ready(Ok(()))
252    }
253
254    #[allow(clippy::expect_used)]
255    fn call(&mut self, req: Request<B>) -> Self::Future {
256        if self.compiled.is_none() {
257            let built = std::mem::take(self);
258            self.compiled = Some(
259                built
260                    .compile()
261                    .expect("Router compilation failed (e.g. an overlapping/duplicate route)"),
262            );
263        }
264        let compiled = self.compiled.as_mut().expect("just populated above");
265        Service::call(compiled, req)
266    }
267}