structured_proxy/
service.rs

1//! The proxy as one tower service: native gRPC to the upstream, everything
2//! else to the proxy's routes, and what no route answers to a fallback.
3
4use std::convert::Infallible;
5use std::future::Future;
6use std::net::SocketAddr;
7use std::pin::Pin;
8use std::sync::Arc;
9use std::task::{Context, Poll};
10
11use axum::extract::connect_info::ConnectInfo;
12use axum::response::IntoResponse;
13use axum::routing::future::RouteFuture;
14use bytes::Bytes;
15use pin_project_lite::pin_project;
16use rustls::pki_types::CertificateDer;
17use tonic::transport::server::{Connected, TcpConnectInfo};
18use tower::{Layer, Service, ServiceExt};
19use tower_http::cors::{Cors, CorsLayer};
20
21use crate::guard::{BoxedService, Class, GrpcRejections, Guards};
22
23/// tonic's `TlsConnectInfo<TcpConnectInfo>`, the record its TLS server puts on
24/// a request and `Request::peer_certs` reads. tonic exports the name only with
25/// a TLS backend feature, so it is reached through the stream it describes.
26pub type TlsConnectInfo =
27    <tokio_rustls::server::TlsStream<tokio::net::TcpStream> as Connected>::ConnectInfo;
28
29use crate::upstream::{
30    grpc_protocol, is_grpc_web_preflight, BoxError, GrpcProtocol, PassThrough, Upstream,
31};
32
33/// The proxy as a tower service, built by
34/// [`ProxyServer::service`](crate::ProxyServer::service).
35///
36/// A request with a gRPC or gRPC-Web content type goes to the upstream as it
37/// arrived, so one listener carries REST and native gRPC. gRPC-Web is the
38/// upstream's to translate (tonic-web's `GrpcWebLayer` around its services),
39/// or the proxy's with `grpc_web.translate` for an upstream that speaks only
40/// gRPC. A gRPC-Web
41/// answer gets the proxy's CORS policy (`cors.grpc_web`), since the proxy
42/// answers the browser's preflight for it. Every other request goes to the
43/// proxy's routes (transcoded RPCs, health, metrics, OpenAPI, OIDC,
44/// forward-auth, extra routes) behind the proxy's middleware. A request no
45/// route matches is answered `404`, or handed to the service set with
46/// [`with_fallback`](Self::with_fallback). Native gRPC and the fallback pass
47/// only the guards whose scope names them (`grpc`, `fallback`); a guard's
48/// rejection of a gRPC call is a gRPC status.
49///
50/// Serve it with [`serve`](crate::serve) or
51/// [`serve_with`](crate::serve_with) (TLS, a connection limit), or hand it to
52/// any server that takes a tower service of `http` types: a Unix socket, an
53/// existing hyper or axum server. Native gRPC needs HTTP/2 on that server
54/// (ALPN `h2` next to `http/1.1` behind TLS). Such a server tells the proxy
55/// which connection a request came on with
56/// [`for_connection`](Self::for_connection).
57///
58/// # Examples
59///
60/// ```
61/// use structured_proxy::ProxyServer;
62///
63/// # fn build() -> anyhow::Result<()> {
64/// // The embedder's own tonic services, called in process.
65/// let grpc = tonic::service::Routes::default();
66/// let service = ProxyServer::from_yaml_str("service:\n  name: demo\n")?.service(grpc)?;
67/// # let _ = service;
68/// # Ok(())
69/// # }
70/// # build().unwrap();
71/// ```
72#[derive(Clone)]
73pub struct ProxyService<U> {
74    upstream: U,
75    routes: axum::Router,
76    /// Binary and text gRPC-Web to the upstream under the CORS policy, each
77    /// built once; `None` when the upstream owns CORS for gRPC-Web.
78    grpc_web: Option<GrpcWebCors<U>>,
79    /// The gRPC paths behind guards or through gRPC-Web translation; a
80    /// protocol that needs neither passes through with nothing in between.
81    boxed: BoxedGrpc,
82    /// The guards, for the fallback an embedder sets later.
83    guards: Arc<Guards>,
84    /// The connection the requests arrive on, set per connection by the
85    /// server.
86    connection: Option<ConnectionInfo>,
87}
88
89impl<U> std::fmt::Debug for ProxyService<U> {
90    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
91        f.debug_struct("ProxyService")
92            .field("boxed_grpc", &self.boxed.grpc.is_some())
93            .field("boxed_grpc_web", &self.boxed.web.is_some())
94            .field("connection", &self.connection)
95            .finish_non_exhaustive()
96    }
97}
98
99/// The gRPC paths that need more than a pass-through (guards, gRPC-Web
100/// translation), each protocol's stack built once; `None` for a protocol that
101/// passes through as it is.
102#[derive(Clone, Default)]
103struct BoxedGrpc {
104    grpc: Option<BoxedService>,
105    web: Option<BoxedService>,
106    web_text: Option<BoxedService>,
107}
108
109/// Hands a gRPC-Web call, translated to gRPC by tonic-web, to an upstream
110/// that speaks only gRPC.
111#[derive(Clone, Debug)]
112struct Translated<U> {
113    upstream: U,
114}
115
116impl<U: Upstream> Service<http::Request<tonic::body::Body>> for Translated<U> {
117    type Response = http::Response<axum::body::Body>;
118    type Error = Infallible;
119    type Future = PassThrough<U>;
120
121    #[inline]
122    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
123        Poll::Ready(Ok(()))
124    }
125
126    fn call(&mut self, request: http::Request<tonic::body::Body>) -> Self::Future {
127        let (mut parts, body) = request.into_parts();
128        // A browser's call arrives over HTTP/1.1; as gRPC it is an HTTP/2
129        // request (gRPC PROTOCOL-HTTP2).
130        parts.version = http::Version::HTTP_2;
131        // tonic-web reports the size of the base64 text body for the decoded
132        // one, and an HTTP/2 client would announce that as its length; the
133        // gRPC body goes without one.
134        let body = tonic::body::Body::new(Unsized { body });
135        PassThrough::new(
136            self.upstream.clone(),
137            http::Request::from_parts(parts, body),
138            GrpcProtocol::Grpc,
139        )
140    }
141}
142
143pin_project! {
144    /// `body` without its size hint.
145    struct Unsized {
146        #[pin]
147        body: tonic::body::Body,
148    }
149}
150
151impl http_body::Body for Unsized {
152    type Data = Bytes;
153    type Error = tonic::Status;
154
155    #[inline]
156    fn poll_frame(
157        self: Pin<&mut Self>,
158        cx: &mut Context<'_>,
159    ) -> Poll<Option<Result<http_body::Frame<Bytes>, tonic::Status>>> {
160        self.project().body.poll_frame(cx)
161    }
162
163    #[inline]
164    fn is_end_stream(&self) -> bool {
165        self.body.is_end_stream()
166    }
167}
168
169/// gRPC-Web pass-through, one CORS-wrapped path per encoding.
170#[derive(Clone, Debug)]
171struct GrpcWebCors<U> {
172    web: Cors<Forward<U>>,
173    web_text: Cors<Forward<U>>,
174}
175
176/// Hands a request to the upstream in a known protocol, as a tower service so
177/// a layer can wrap it.
178#[derive(Clone, Debug)]
179struct Forward<U> {
180    upstream: U,
181    protocol: GrpcProtocol,
182}
183
184impl<U: Upstream> Service<http::Request<tonic::body::Body>> for Forward<U> {
185    type Response = http::Response<axum::body::Body>;
186    type Error = Infallible;
187    type Future = PassThrough<U>;
188
189    #[inline]
190    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
191        // The upstream's readiness is waited for per request, on its clone.
192        Poll::Ready(Ok(()))
193    }
194
195    fn call(&mut self, request: http::Request<tonic::body::Body>) -> Self::Future {
196        PassThrough::new(self.upstream.clone(), request, self.protocol)
197    }
198}
199
200/// The guarded path's end: the guards run on axum's request type.
201impl<U: Upstream> Service<axum::extract::Request> for Forward<U> {
202    type Response = http::Response<axum::body::Body>;
203    type Error = Infallible;
204    type Future = PassThrough<U>;
205
206    #[inline]
207    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
208        Poll::Ready(Ok(()))
209    }
210
211    fn call(&mut self, request: axum::extract::Request) -> Self::Future {
212        let request = request.map(tonic::body::Body::new);
213        PassThrough::new(self.upstream.clone(), request, self.protocol)
214    }
215}
216
217/// The connection requests arrive on, in the form a tonic server records it:
218/// the TCP ends, and behind TLS the client's certificate chain.
219///
220/// Built from what [`Connected::connect_info`] returns for a
221/// `tokio::net::TcpStream` (`ConnectionInfo::from`) or a
222/// `tokio_rustls::server::TlsStream<TcpStream>` ([`ConnectionInfo::tls`]), so a
223/// tonic handler behind the proxy reads it as it would behind tonic's own
224/// server: `Request::remote_addr`, and `Request::peer_certs` for mTLS.
225///
226/// [`Connected::connect_info`]: tonic::transport::server::Connected::connect_info
227#[derive(Clone, Debug)]
228pub struct ConnectionInfo {
229    tcp: TcpConnectInfo,
230    tls: Option<TlsConnectInfo>,
231}
232
233impl ConnectionInfo {
234    /// The client's address.
235    pub fn remote_addr(&self) -> Option<SocketAddr> {
236        self.tcp.remote_addr
237    }
238
239    /// The certificate chain the client presented over TLS, if it did.
240    pub fn peer_certs(&self) -> Option<Arc<Vec<CertificateDer<'static>>>> {
241        self.tls.as_ref().and_then(|tls| tls.peer_certs())
242    }
243
244    /// Put the connection where tonic reads it: [`TcpConnectInfo`] for
245    /// `Request::remote_addr` and, behind TLS, [`TlsConnectInfo`] for
246    /// `Request::peer_certs`.
247    pub(crate) fn into_tonic_extensions(self, extensions: &mut http::Extensions) {
248        extensions.insert(self.tcp);
249        if let Some(tls) = self.tls {
250            extensions.insert(tls);
251        }
252    }
253}
254
255impl From<TcpConnectInfo> for ConnectionInfo {
256    fn from(tcp: TcpConnectInfo) -> Self {
257        Self { tcp, tls: None }
258    }
259}
260
261impl ConnectionInfo {
262    /// A TLS connection, from what `connect_info` reports for a
263    /// `tokio_rustls::server::TlsStream<TcpStream>`: the TCP ends and the
264    /// client's certificate chain.
265    pub fn tls(tls: TlsConnectInfo) -> Self {
266        Self {
267            tcp: tls.get_ref().clone(),
268            tls: Some(tls),
269        }
270    }
271}
272
273impl<U: Upstream> ProxyService<U> {
274    /// The service over `upstream` and `routes`; gRPC-Web answers carry
275    /// `grpc_web_cors` when set, and gRPC-Web calls are translated to gRPC
276    /// when `translate_grpc_web`.
277    pub(crate) fn new(
278        upstream: U,
279        routes: axum::Router,
280        grpc_web_cors: Option<CorsLayer>,
281        guards: Arc<Guards>,
282        translate_grpc_web: bool,
283    ) -> Self {
284        let forward = |protocol| Forward {
285            upstream: upstream.clone(),
286            protocol,
287        };
288        let guard_grpc = guards.cover(Class::Grpc);
289        let stack = |protocol: GrpcProtocol| {
290            let web = protocol != GrpcProtocol::Grpc;
291            let mut service = if web && translate_grpc_web {
292                let translator = tonic_web::GrpcWebLayer::new().layer(Translated {
293                    upstream: upstream.clone(),
294                });
295                BoxedService::new(ServiceExt::<axum::extract::Request>::map_response(
296                    translator,
297                    |response| response.map(axum::body::Body::new),
298                ))
299            } else {
300                BoxedService::new(forward(protocol))
301            };
302            if guard_grpc {
303                service = BoxedService::new(GrpcRejections {
304                    inner: guards.service(service, Class::Grpc),
305                    protocol,
306                });
307            }
308            // CORS outermost, so a rejected gRPC-Web call still carries it.
309            match &grpc_web_cors {
310                Some(cors) if web => BoxedService::new(cors.layer(service)),
311                _ => service,
312            }
313        };
314        let boxed_web = guard_grpc || translate_grpc_web;
315        let boxed = BoxedGrpc {
316            grpc: guard_grpc.then(|| stack(GrpcProtocol::Grpc)),
317            web: boxed_web.then(|| stack(GrpcProtocol::Web)),
318            web_text: boxed_web.then(|| stack(GrpcProtocol::WebText)),
319        };
320        // The pass-through of gRPC-Web, and the answer to its preflights.
321        let grpc_web = grpc_web_cors.map(|cors| GrpcWebCors {
322            web: cors.layer(forward(GrpcProtocol::Web)),
323            web_text: cors.layer(forward(GrpcProtocol::WebText)),
324        });
325        Self {
326            upstream,
327            routes,
328            grpc_web,
329            boxed,
330            guards,
331            connection: None,
332        }
333    }
334
335    /// Hand the requests no route matches to `fallback` instead of answering
336    /// `404`: an embedder's own REST routes, a static site, anything that is a
337    /// tower service. Only the guards whose scope names `fallback` traffic see
338    /// them; CORS and tracing are the fallback's own.
339    ///
340    /// A request whose path a route answers but not with its method stays with
341    /// the proxy (`405`), as does every gRPC request and every browser
342    /// preflight for a gRPC-Web call, which follows the call it announces.
343    #[must_use]
344    pub fn with_fallback<F>(mut self, fallback: F) -> Self
345    where
346        F: Service<axum::extract::Request, Error = Infallible> + Clone + Send + Sync + 'static,
347        F::Response: IntoResponse,
348        F::Future: Send + 'static,
349    {
350        self.routes = if self.guards.cover(Class::Fallback) {
351            let fallback = BoxedService::new(fallback.map_response(IntoResponse::into_response));
352            self.routes
353                .fallback_service(self.guards.service(fallback, Class::Fallback))
354        } else {
355            self.routes.fallback_service(fallback)
356        };
357        self
358    }
359
360    /// This service for the requests of one connection: the proxy's
361    /// middleware sees its client's address (rate limits by IP, the auth
362    /// decider), and a tonic upstream in process reads it with
363    /// `Request::remote_addr`, and its TLS client certificates with
364    /// `Request::peer_certs`, for native and transcoded calls alike.
365    ///
366    /// A server of your own calls it once per accepted connection, with what
367    /// tonic's [`Connected`] trait
368    /// reports for the stream. [`serve_with`](crate::serve_with) does this
369    /// itself.
370    ///
371    /// # Examples
372    ///
373    /// ```no_run
374    /// use structured_proxy::ProxyServer;
375    /// use tonic::transport::server::Connected;
376    ///
377    /// # async fn run() -> anyhow::Result<()> {
378    /// let proxy = ProxyServer::from_yaml_str("service:\n  name: demo\n")?
379    ///     .service(tonic::service::Routes::default())?;
380    /// let listener = tokio::net::TcpListener::bind("0.0.0.0:8443").await?;
381    /// let (tcp, _) = listener.accept().await?;
382    /// // After a TLS handshake, the `TlsStream` reports the client's certificates too.
383    /// let service = proxy.for_connection(tcp.connect_info());
384    /// # let _ = service;
385    /// # Ok(())
386    /// # }
387    /// ```
388    ///
389    /// Without it, a request's peer comes from the `ConnectInfo` an outer
390    /// axum server recorded, when there is one.
391    #[must_use]
392    pub fn for_connection(&self, connection: impl Into<ConnectionInfo>) -> Self {
393        Self {
394            upstream: self.upstream.clone(),
395            routes: self.routes.clone(),
396            grpc_web: self.grpc_web.clone(),
397            boxed: self.boxed.clone(),
398            guards: self.guards.clone(),
399            connection: Some(connection.into()),
400        }
401    }
402}
403
404/// The connection `request` came on: `connection`, the one given to
405/// [`ProxyService::for_connection`], else the peer an outer axum server
406/// recorded as `ConnectInfo`, so the upstream sees the same client the proxy's
407/// middleware does. A free function, so the caller keeps its other fields.
408fn connection_of<B>(
409    connection: Option<&ConnectionInfo>,
410    request: &http::Request<B>,
411) -> Option<ConnectionInfo> {
412    if let Some(connection) = connection {
413        return Some(connection.clone());
414    }
415    let ConnectInfo(remote) = request.extensions().get::<ConnectInfo<SocketAddr>>()?;
416    Some(ConnectionInfo::from(TcpConnectInfo {
417        local_addr: None,
418        remote_addr: Some(*remote),
419    }))
420}
421
422impl<U, B> Service<http::Request<B>> for ProxyService<U>
423where
424    U: Upstream,
425    B: http_body::Body<Data = Bytes> + Send + 'static,
426    B::Error: Into<BoxError>,
427{
428    type Response = http::Response<axum::body::Body>;
429    type Error = Infallible;
430    type Future = ResponseFuture<U>;
431
432    #[inline]
433    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
434        // Readiness of the upstream is waited for per request, on its clone.
435        Poll::Ready(Ok(()))
436    }
437
438    fn call(&mut self, mut request: http::Request<B>) -> Self::Future {
439        // A gRPC-Web preflight goes where the call it announces goes, so both
440        // get one CORS policy: the proxy's, or the upstream's when it owns
441        // CORS. A fallback in between would answer it with neither.
442        let call = grpc_protocol(request.headers());
443        let protocol = call.or_else(|| {
444            is_grpc_web_preflight(request.method(), request.headers()).then_some(GrpcProtocol::Web)
445        });
446        // A preflight is not a call: it skips the guards and the translation.
447        let boxed = match call {
448            Some(GrpcProtocol::Grpc) => self.boxed.grpc.as_mut(),
449            Some(GrpcProtocol::Web) => self.boxed.web.as_mut(),
450            Some(GrpcProtocol::WebText) => self.boxed.web_text.as_mut(),
451            None => None,
452        };
453        let inner = if let Some(service) = boxed {
454            // The guards read the peer as axum's `ConnectInfo`, the upstream
455            // as tonic records it.
456            if let Some(connection) = connection_of(self.connection.as_ref(), &request) {
457                let extensions = request.extensions_mut();
458                if let Some(remote) = connection.remote_addr() {
459                    extensions.insert(ConnectInfo(remote));
460                }
461                connection.into_tonic_extensions(extensions);
462            }
463            // Every service on these paths is ready at once: guards and the
464            // translator are middleware, and `Forward` waits for the upstream
465            // per request.
466            Inner::Boxed {
467                future: service.call(request.map(axum::body::Body::new)),
468            }
469        } else if let Some(protocol) = protocol {
470            // A native call carries its connection the way tonic's server
471            // hands it to a handler.
472            if let Some(connection) = connection_of(self.connection.as_ref(), &request) {
473                connection.into_tonic_extensions(request.extensions_mut());
474            }
475            let request = request.map(tonic::body::Body::new);
476            // A browser only speaks gRPC-Web, and the routes answered its
477            // preflight: its call gets the same CORS policy.
478            let cors = match (&mut self.grpc_web, protocol) {
479                (Some(cors), GrpcProtocol::Web) => Some(&mut cors.web),
480                (Some(cors), GrpcProtocol::WebText) => Some(&mut cors.web_text),
481                _ => None,
482            };
483            match cors {
484                // `Forward` is always ready, and so is CORS around it.
485                Some(cors) => Inner::GrpcWeb {
486                    future: cors.call(request),
487                },
488                None => Inner::Grpc {
489                    call: PassThrough::new(self.upstream.clone(), request, protocol),
490                },
491            }
492        } else {
493            // The middleware reads the peer as axum's `ConnectInfo`; the
494            // transcoder passes the whole connection on to the upstream.
495            if let Some(connection) = &self.connection {
496                let extensions = request.extensions_mut();
497                if let Some(remote) = connection.remote_addr() {
498                    extensions.insert(ConnectInfo(remote));
499                }
500                extensions.insert(connection.clone());
501            } else if let Some(connection) = connection_of(None, &request) {
502                request.extensions_mut().insert(connection);
503            }
504            Inner::Routes {
505                future: self.routes.call(request),
506            }
507        };
508        ResponseFuture { inner }
509    }
510}
511
512pin_project! {
513    /// The response future of [`ProxyService`].
514    pub struct ResponseFuture<U: Upstream> {
515        #[pin]
516        inner: Inner<U>,
517    }
518}
519
520pin_project! {
521    #[project = InnerProj]
522    enum Inner<U: Upstream> {
523        Routes {
524            #[pin]
525            future: RouteFuture<Infallible>,
526        },
527        Grpc {
528            #[pin]
529            call: PassThrough<U>,
530        },
531        GrpcWeb {
532            #[pin]
533            future: tower_http::cors::ResponseFuture<PassThrough<U>>,
534        },
535        Boxed {
536            #[pin]
537            future: <BoxedService as Service<axum::extract::Request>>::Future,
538        },
539    }
540}
541
542impl<U: Upstream> Future for ResponseFuture<U> {
543    type Output = Result<http::Response<axum::body::Body>, Infallible>;
544
545    #[inline]
546    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
547        match self.project().inner.project() {
548            InnerProj::Routes { future } => future.poll(cx),
549            InnerProj::Grpc { call } => call.poll(cx),
550            InnerProj::GrpcWeb { future } => future.poll(cx),
551            InnerProj::Boxed { future } => future.poll(cx),
552        }
553    }
554}
555
556#[cfg(test)]
557mod tests;
structured_proxy/service.rs

structured_proxy/
service.rs
