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

structured_proxy/
service.rs
