structured_proxy/
serve.rs

1//! Running a [`ProxyService`] on a TCP listener: HTTP/1.1 and HTTP/2 on one
2//! port, optionally behind TLS, with an optional cap on open connections.
3
4use std::pin::Pin;
5use std::sync::Arc;
6use std::task::{Context, Poll};
7use std::time::Duration;
8
9use hyper_util::rt::{TokioExecutor, TokioIo, TokioTimer};
10use hyper_util::server::conn::auto::Builder;
11use hyper_util::service::TowerToHyperService;
12use pin_project_lite::pin_project;
13use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
14use tokio::net::TcpListener;
15use tokio::sync::{OwnedSemaphorePermit, Semaphore};
16use tonic::transport::server::Connected;
17
18use crate::service::{ConnectionInfo, ProxyService};
19use crate::upstream::Upstream;
20
21mod idle;
22
23/// How [`serve_with`] runs its listener.
24///
25/// # Examples
26///
27/// ```
28/// use std::time::Duration;
29/// use structured_proxy::ServeOptions;
30///
31/// let options = ServeOptions::new()
32///     .max_connections(10_000)
33///     .idle_timeout(Some(Duration::from_secs(120)));
34/// # let _ = options;
35/// ```
36#[derive(Clone, Debug)]
37pub struct ServeOptions {
38    max_connections: Option<usize>,
39    tls: Option<Arc<rustls::ServerConfig>>,
40    idle_timeout: Option<Duration>,
41    header_read_timeout: Duration,
42    tls_handshake_timeout: Duration,
43}
44
45impl Default for ServeOptions {
46    fn default() -> Self {
47        Self {
48            max_connections: None,
49            tls: None,
50            idle_timeout: Some(Duration::from_secs(60)),
51            header_read_timeout: Duration::from_secs(30),
52            tls_handshake_timeout: Duration::from_secs(10),
53        }
54    }
55}
56
57impl ServeOptions {
58    /// Cleartext, with no limit on connections; a connection idle for 60 s is
59    /// closed, a client gets 30 s to send the headers of an HTTP/1.1 request
60    /// and 10 s to finish a TLS handshake.
61    pub fn new() -> Self {
62        Self::default()
63    }
64
65    /// Close a connection that has had no request in flight for `timeout`
66    /// (gracefully: HTTP/2 gets a GOAWAY); `None` keeps idle connections
67    /// open. A request counts until its response body ends, so a stream keeps
68    /// its connection. Without it an idle client holds a
69    /// [`max_connections`](Self::max_connections) slot for as long as it likes.
70    /// A connection upgraded to another protocol (a WebSocket in a fallback)
71    /// leaves HTTP and this timeout with it; it keeps its slot until it
72    /// closes.
73    #[must_use]
74    pub fn idle_timeout(mut self, timeout: Option<Duration>) -> Self {
75        self.idle_timeout = timeout;
76        self
77    }
78
79    /// How long an HTTP/1.1 client has to send a request's headers, so a
80    /// client that trickles them in holds no connection for long.
81    #[must_use]
82    pub fn header_read_timeout(mut self, timeout: Duration) -> Self {
83        self.header_read_timeout = timeout;
84        self
85    }
86
87    /// How long a client has to finish its TLS handshake before the
88    /// connection is dropped, so a stalled client holds neither a task nor a
89    /// connection slot.
90    #[must_use]
91    pub fn tls_handshake_timeout(mut self, timeout: Duration) -> Self {
92        self.tls_handshake_timeout = timeout;
93        self
94    }
95
96    /// Serve at most `max` connections at once: past it, the next connection
97    /// is accepted when one closes, and waits in the listen backlog until then.
98    ///
99    /// # Panics
100    ///
101    /// `max` is zero, which would never accept a connection.
102    #[must_use]
103    pub fn max_connections(mut self, max: usize) -> Self {
104        assert!(max > 0, "max_connections must be at least 1");
105        self.max_connections = Some(max.min(Semaphore::MAX_PERMITS));
106        self
107    }
108
109    /// Terminate TLS with `config`. Set its `alpn_protocols` to `h2` and
110    /// `http/1.1` for gRPC clients to get HTTP/2; an empty list is filled in
111    /// with those two. A client certificate the config verifies reaches a
112    /// tonic upstream as `Request::peer_certs`.
113    #[must_use]
114    pub fn tls(mut self, mut config: rustls::ServerConfig) -> Self {
115        if config.alpn_protocols.is_empty() {
116            config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()];
117        }
118        self.tls = Some(Arc::new(config));
119        self
120    }
121}
122
123/// Serve `service` on `listener` until the listener fails: cleartext HTTP/1.1
124/// and HTTP/2 on the same port, so REST clients and native gRPC clients share
125/// it. [`serve_with`] adds TLS and a connection limit.
126///
127/// # Errors
128///
129/// None so far: a failed accept is logged and retried, as a full file
130/// descriptor table recovers when connections close.
131///
132/// # Examples
133///
134/// ```no_run
135/// use structured_proxy::ProxyServer;
136///
137/// # async fn run() -> anyhow::Result<()> {
138/// let grpc = tonic::service::Routes::default();
139/// let service = ProxyServer::from_yaml_str("service:\n  name: demo\n")?.service(grpc)?;
140/// let listener = tokio::net::TcpListener::bind("0.0.0.0:8080").await?;
141/// structured_proxy::serve(listener, service).await?;
142/// # Ok(())
143/// # }
144/// ```
145pub async fn serve<U: Upstream>(
146    listener: TcpListener,
147    service: ProxyService<U>,
148) -> std::io::Result<()> {
149    serve_with(listener, service, ServeOptions::new()).await
150}
151
152/// [`serve`] with `options`: TLS termination (mTLS when the config verifies
153/// client certificates) and a cap on open connections. Each connection's
154/// service gets its peer, and behind TLS its client certificates, through
155/// [`ProxyService::for_connection`]. The TLS handshake runs in the
156/// connection's own task, so a slow client does not hold up the others.
157///
158/// # Errors
159///
160/// See [`serve`].
161///
162/// # Examples
163///
164/// ```no_run
165/// use structured_proxy::{ProxyServer, ServeOptions};
166///
167/// # async fn run(tls: rustls::ServerConfig) -> anyhow::Result<()> {
168/// let service = ProxyServer::from_yaml_str("service:\n  name: demo\n")?
169///     .service(tonic::service::Routes::default())?;
170/// let listener = tokio::net::TcpListener::bind("0.0.0.0:8443").await?;
171/// let options = ServeOptions::new().tls(tls).max_connections(10_000);
172/// structured_proxy::serve_with(listener, service, options).await?;
173/// # Ok(())
174/// # }
175/// ```
176pub async fn serve_with<U: Upstream>(
177    listener: TcpListener,
178    service: ProxyService<U>,
179    options: ServeOptions,
180) -> std::io::Result<()> {
181    let slots = options
182        .max_connections
183        .map(|max| Arc::new(Semaphore::new(max)));
184    let acceptor = options.tls.map(tokio_rustls::TlsAcceptor::from);
185    let handshake_timeout = options.tls_handshake_timeout;
186    let limits = ConnectionLimits {
187        idle_timeout: options.idle_timeout,
188        header_read_timeout: options.header_read_timeout,
189    };
190    loop {
191        // The slot is taken before the accept, so a full server leaves new
192        // connections in the kernel's backlog instead of accepting and
193        // dropping them.
194        let slot = match &slots {
195            Some(slots) => Some(
196                slots
197                    .clone()
198                    .acquire_owned()
199                    .await
200                    .expect("the connection semaphore is never closed"),
201            ),
202            None => None,
203        };
204        let tcp = match listener.accept().await {
205            Ok((tcp, _)) => tcp,
206            Err(error) => {
207                accept_failed(error).await;
208                continue;
209            }
210        };
211        let service = service.clone();
212        let acceptor = acceptor.clone();
213        tokio::spawn(async move {
214            // Small gRPC frames and REST answers are latency-bound.
215            if let Err(error) = tcp.set_nodelay(true) {
216                tracing::debug!(%error, "cannot set TCP_NODELAY");
217            }
218            match acceptor {
219                None => {
220                    let service = service.for_connection(tcp.connect_info());
221                    serve_connection(SlotIo { io: tcp, slot }, service, limits).await;
222                }
223                Some(acceptor) => {
224                    // The slot is held through the handshake, then by the
225                    // stream.
226                    let stream =
227                        match tokio::time::timeout(handshake_timeout, acceptor.accept(tcp)).await {
228                            Ok(Ok(stream)) => stream,
229                            Ok(Err(error)) => {
230                                tracing::debug!(%error, "TLS handshake failed");
231                                return;
232                            }
233                            Err(_) => {
234                                tracing::debug!("TLS handshake timed out");
235                                return;
236                            }
237                        };
238                    let service =
239                        service.for_connection(ConnectionInfo::tls(stream.connect_info()));
240                    serve_connection(SlotIo { io: stream, slot }, service, limits).await;
241                }
242            }
243        });
244    }
245}
246
247pin_project! {
248    /// A connection's IO holding its `max_connections` slot for as long as
249    /// the socket is open. hyper hands the IO of an upgraded connection (a
250    /// WebSocket in a fallback) to the upgrading service and finishes the
251    /// connection future, so the slot has to live with the IO, not the task.
252    struct SlotIo<I> {
253        #[pin]
254        io: I,
255        slot: Option<OwnedSemaphorePermit>,
256    }
257}
258
259impl<I: AsyncRead> AsyncRead for SlotIo<I> {
260    #[inline]
261    fn poll_read(
262        self: Pin<&mut Self>,
263        cx: &mut Context<'_>,
264        buf: &mut ReadBuf<'_>,
265    ) -> Poll<std::io::Result<()>> {
266        self.project().io.poll_read(cx, buf)
267    }
268}
269
270impl<I: AsyncWrite> AsyncWrite for SlotIo<I> {
271    #[inline]
272    fn poll_write(
273        self: Pin<&mut Self>,
274        cx: &mut Context<'_>,
275        buf: &[u8],
276    ) -> Poll<std::io::Result<usize>> {
277        self.project().io.poll_write(cx, buf)
278    }
279
280    #[inline]
281    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
282        self.project().io.poll_flush(cx)
283    }
284
285    #[inline]
286    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
287        self.project().io.poll_shutdown(cx)
288    }
289
290    #[inline]
291    fn poll_write_vectored(
292        self: Pin<&mut Self>,
293        cx: &mut Context<'_>,
294        bufs: &[std::io::IoSlice<'_>],
295    ) -> Poll<std::io::Result<usize>> {
296        self.project().io.poll_write_vectored(cx, bufs)
297    }
298
299    #[inline]
300    fn is_write_vectored(&self) -> bool {
301        self.io.is_write_vectored()
302    }
303}
304
305/// The timeouts every connection is served under.
306#[derive(Clone, Copy, Debug)]
307struct ConnectionLimits {
308    idle_timeout: Option<Duration>,
309    header_read_timeout: Duration,
310}
311
312/// HTTP/1.1 or HTTP/2, whichever the client speaks, on one connection.
313async fn serve_connection<U, I>(io: I, service: ProxyService<U>, limits: ConnectionLimits)
314where
315    U: Upstream,
316    I: AsyncRead + AsyncWrite + Unpin + Send + 'static,
317{
318    let mut builder = Builder::new(TokioExecutor::new());
319    // hyper times nothing without a timer: its own default header read
320    // timeout is dropped with a warning.
321    builder
322        .http1()
323        .timer(TokioTimer::new())
324        .header_read_timeout(limits.header_read_timeout);
325    builder.http2().timer(TokioTimer::new());
326    let served = match limits.idle_timeout {
327        None => {
328            builder
329                .serve_connection_with_upgrades(TokioIo::new(io), TowerToHyperService::new(service))
330                .await
331        }
332        Some(timeout) => {
333            let activity = Arc::new(idle::Activity::default());
334            let service = idle::Tracked {
335                inner: service,
336                activity: Arc::clone(&activity),
337            };
338            let connection = builder.serve_connection_with_upgrades(
339                TokioIo::new(io),
340                TowerToHyperService::new(service),
341            );
342            tokio::pin!(connection);
343            tokio::select! {
344                served = connection.as_mut() => served,
345                () = activity.idle_for(timeout) => {
346                    // HTTP/2 gets a GOAWAY, HTTP/1.1 closes after the request
347                    // it is reading, if any.
348                    connection.as_mut().graceful_shutdown();
349                    connection.await
350                }
351            }
352        }
353    };
354    if let Err(error) = served {
355        tracing::debug!(%error, "connection ended");
356    }
357}
358
359/// A failed accept: a connection the client gave up on is skipped; anything
360/// else (a full file descriptor table) is logged and retried a second later,
361/// the pause that lets open connections close.
362async fn accept_failed(error: std::io::Error) {
363    use std::io::ErrorKind;
364    if matches!(
365        error.kind(),
366        ErrorKind::ConnectionRefused | ErrorKind::ConnectionAborted | ErrorKind::ConnectionReset
367    ) {
368        return;
369    }
370    tracing::error!(%error, "accepting a connection failed; retrying in 1s");
371    tokio::time::sleep(Duration::from_secs(1)).await;
372}
373
374#[cfg(test)]
375mod tests;
structured_proxy/serve.rs

structured_proxy/
serve.rs
