Skip to main content

iota_http/
lib.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2025 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::{collections::HashMap, sync::Arc, time::Duration};
6
7use connection_handler::{OnConnectionClose, sleep_or_pending};
8pub use http;
9use http::{Request, Response};
10use hyper_util::service::TowerToHyperService;
11use io::ServerIo;
12use metrics::{HandshakeResult, TrackRequests, TrackedIo};
13use tokio::task::JoinSet;
14use tokio_rustls::{TlsAcceptor, rustls};
15use tower::{Service, ServiceBuilder, ServiceExt};
16use tracing::trace;
17
18use self::{
19    body::BoxBody,
20    connection_info::{ActiveConnections, PeerConnectionCounts},
21};
22
23pub mod body;
24mod config;
25mod connection_handler;
26mod connection_info;
27mod fuse;
28mod io;
29mod listener;
30pub mod metrics;
31
32pub use config::{Config, PeerConnectionEvent};
33pub use connection_info::{ConnectInfo, ConnectionId, ConnectionInfo, PeerCertificates};
34pub use listener::{Listener, ListenerExt};
35
36pub(crate) type BoxError = Box<dyn std::error::Error + Send + Sync>;
37/// h2 alpn in plain format for rustls.
38const ALPN_H2: &[u8] = b"h2";
39/// h1 alpn in plain format for rustls.
40const ALPN_H1: &[u8] = b"http/1.1";
41
42#[derive(Default)]
43pub struct Builder {
44    config: Config,
45    tls_config: Option<rustls::ServerConfig>,
46}
47
48impl Builder {
49    pub fn new() -> Self {
50        Self::default()
51    }
52
53    pub fn config(mut self, config: Config) -> Self {
54        self.config = config;
55        self
56    }
57
58    // Convenience method for configuring TLS with a single server cert
59    //
60    // Attempts to load PEM formatted files for the certificate chain and private
61    // key material from the provided file system paths.
62    pub fn tls_single_cert(
63        self,
64        cert_file: impl AsRef<std::path::Path>,
65        private_key_file: impl AsRef<std::path::Path>,
66    ) -> Result<Self, BoxError> {
67        let tls_config =
68            iota_tls::create_rustls_server_config_from_pem(cert_file, private_key_file)?;
69        Ok(self.tls_config(tls_config))
70    }
71
72    pub fn tls_config(mut self, tls_config: rustls::ServerConfig) -> Self {
73        self.tls_config = Some(tls_config);
74        self
75    }
76
77    pub fn serve<A, S, ResponseBody>(
78        self,
79        addr: A,
80        service: S,
81    ) -> Result<ServerHandle<std::net::SocketAddr>, BoxError>
82    where
83        A: std::net::ToSocketAddrs,
84        S: Service<
85                Request<BoxBody>,
86                Response = Response<ResponseBody>,
87                Error: Into<BoxError>,
88                Future: Send,
89            > + Clone
90            + Send
91            + 'static,
92        ResponseBody: http_body::Body<Data = bytes::Bytes, Error: Into<BoxError>> + Send + 'static,
93    {
94        let listener = listener::TcpListenerWithOptions::new(addr, self.config.tcp_keepalive)?;
95
96        Self::serve_with_listener(self, listener, service)
97    }
98
99    fn serve_with_listener<L, S, ResponseBody>(
100        self,
101        listener: L,
102        service: S,
103    ) -> Result<ServerHandle<L::Addr>, BoxError>
104    where
105        L: Listener<Addr = std::net::SocketAddr>,
106        S: Service<
107                Request<BoxBody>,
108                Response = Response<ResponseBody>,
109                Error: Into<BoxError>,
110                Future: Send,
111            > + Clone
112            + Send
113            + 'static,
114        ResponseBody: http_body::Body<Data = bytes::Bytes, Error: Into<BoxError>> + Send + 'static,
115    {
116        self.config.validate()?;
117
118        let local_addr = listener.local_addr()?;
119        let graceful_shutdown_token = tokio_util::sync::CancellationToken::new();
120        let connections = ActiveConnections::default();
121
122        let tls_config = self.tls_config.map(|mut tls| {
123            tls.alpn_protocols.push(ALPN_H2.into());
124            if self.config.accept_http1 {
125                tls.alpn_protocols.push(ALPN_H1.into());
126            }
127            Arc::new(tls)
128        });
129
130        if let (Some(metrics), Some(_)) = (&self.config.metrics, &tls_config) {
131            metrics.enable_tls();
132        }
133
134        let (watch_sender, watch_receiver) = tokio::sync::watch::channel(());
135        let peer_connection_counts =
136            PeerConnectionCounts::new(self.config.on_peer_connection_event.clone());
137        let server = Server {
138            config: self.config,
139            tls_config,
140            listener,
141            local_addr,
142            service: ServiceBuilder::new()
143                .layer(tower::util::BoxCloneService::layer())
144                .map_response(|response: Response<ResponseBody>| response.map(body::boxed))
145                .map_err(Into::into)
146                .service(service),
147            pending_connections: JoinSet::new(),
148            connection_handlers: JoinSet::new(),
149            connections: connections.clone(),
150            graceful_shutdown_token: graceful_shutdown_token.clone(),
151            _watch_receiver: watch_receiver,
152            peer_connection_counts,
153        };
154
155        let handle = ServerHandle(Arc::new(HandleInner {
156            local_addr,
157            connections,
158            graceful_shutdown_token,
159            watch_sender,
160        }));
161
162        tokio::spawn(server.serve());
163
164        Ok(handle)
165    }
166}
167
168#[derive(Debug)]
169pub struct ServerHandle<A = std::net::SocketAddr>(Arc<HandleInner<A>>);
170
171#[derive(Debug)]
172struct HandleInner<A = std::net::SocketAddr> {
173    /// The local address of the server.
174    local_addr: A,
175    connections: ActiveConnections<A>,
176    graceful_shutdown_token: tokio_util::sync::CancellationToken,
177    watch_sender: tokio::sync::watch::Sender<()>,
178}
179
180impl<A> ServerHandle<A> {
181    /// Returns the local address of the server
182    pub fn local_addr(&self) -> &A {
183        &self.0.local_addr
184    }
185
186    /// Trigger a graceful shutdown of the server, but don't wait till the
187    /// server has completed shutting down
188    pub fn trigger_shutdown(&self) {
189        self.0.graceful_shutdown_token.cancel();
190    }
191
192    /// Completes once the network has been shutdown.
193    ///
194    /// This explicitly *does not* trigger the network to shutdown, see
195    /// `trigger_shutdown` or `shutdown` if you want to trigger shutting
196    /// down the server.
197    pub async fn wait_for_shutdown(&self) {
198        self.0.watch_sender.closed().await
199    }
200
201    /// Triggers a shutdown of the server and waits for it to complete shutting
202    /// down.
203    pub async fn shutdown(&self) {
204        self.trigger_shutdown();
205        self.wait_for_shutdown().await;
206    }
207
208    /// Checks if the Server has been shutdown.
209    pub fn is_shutdown(&self) -> bool {
210        self.0.watch_sender.is_closed()
211    }
212
213    pub fn connections(
214        &self,
215    ) -> std::sync::RwLockReadGuard<'_, HashMap<ConnectionId, ConnectionInfo<A>>> {
216        self.0.connections.read().unwrap()
217    }
218
219    /// Returns the number of active connections the server is handling
220    pub fn number_of_connections(&self) -> usize {
221        self.connections().len()
222    }
223}
224
225impl<A> Clone for ServerHandle<A> {
226    fn clone(&self) -> Self {
227        Self(self.0.clone())
228    }
229}
230
231type ConnectingOutput<Io, Addr> = Result<(ServerIo<TrackedIo<Io>>, Addr), crate::BoxError>;
232
233struct Server<L: Listener> {
234    config: Config,
235    tls_config: Option<Arc<rustls::ServerConfig>>,
236
237    listener: L,
238    local_addr: L::Addr,
239    service: tower::util::BoxCloneService<Request<BoxBody>, Response<BoxBody>, crate::BoxError>,
240
241    pending_connections: JoinSet<ConnectingOutput<L::Io, L::Addr>>,
242    connection_handlers: JoinSet<()>,
243    connections: ActiveConnections<L::Addr>,
244    graceful_shutdown_token: tokio_util::sync::CancellationToken,
245    // Used to signal to a ServerHandle when the server has completed shutting down
246    _watch_receiver: tokio::sync::watch::Receiver<()>,
247    peer_connection_counts: PeerConnectionCounts,
248}
249
250impl<L> Server<L>
251where
252    L: Listener<Addr = std::net::SocketAddr>,
253{
254    async fn serve(mut self) -> Result<(), BoxError> {
255        loop {
256            tokio::select! {
257                _ = self.graceful_shutdown_token.cancelled() => {
258                    trace!("signal received, shutting down");
259                    break;
260                },
261                // While the handshake limit is reached, leave new connections in the kernel
262                // backlog rather than accepting them into an unbounded set of pending tasks.
263                (io, remote_addr) = self.listener.accept(), if self.accepts_more_connections() => {
264                    self.handle_incoming(io, remote_addr);
265                },
266                // A failed task affects only its own connection, so the loop keeps serving
267                // the others.
268                Some(maybe_connection) = self.pending_connections.join_next() => {
269                    let (io, remote_addr) = match maybe_connection {
270                        Ok(Ok((io, remote_addr))) => {
271                            (io, remote_addr)
272                        }
273                        Ok(Err(e)) => {
274                            tracing::debug!(error = %e, "error accepting connection");
275                            continue;
276                        }
277                        Err(e) => {
278                            tracing::error!(error = %e, "connection handshake task failed");
279                            continue;
280                        }
281                    };
282
283                    trace!("connection accepted");
284                    self.handle_connection(io, remote_addr);
285                },
286                Some(connection_handler_output) = self.connection_handlers.join_next() => {
287                    if let Err(e) = connection_handler_output {
288                        tracing::error!(error = %e, "connection task failed");
289                    }
290                },
291            }
292        }
293
294        // Shutting down, wait for all connection handlers to finish
295        self.shutdown().await;
296
297        Ok(())
298    }
299
300    /// Whether another connection may be accepted, or the limit on concurrent
301    /// TLS handshakes is currently reached.
302    fn accepts_more_connections(&self) -> bool {
303        self.config
304            .max_pending_connections
305            .is_none_or(|max| self.pending_connections.len() < max)
306    }
307
308    fn handle_incoming(&mut self, io: L::Io, remote_addr: L::Addr) {
309        let io = TrackedIo::new(io, self.config.metrics.as_ref(), remote_addr);
310        if let Some(tls) = self.tls_config.clone() {
311            let tls_acceptor = TlsAcceptor::from(tls);
312            let handshake_timeout = self.config.handshake_timeout;
313            let handshake_guard = io.connection().and_then(|c| c.track_handshake());
314            self.pending_connections.spawn(async move {
315                let result = tokio::select! {
316                    result = handshake(io, remote_addr, tls_acceptor) => Some(result),
317                    // Dropping the handshake closes the connection, releasing its file descriptor.
318                    _ = sleep_or_pending(handshake_timeout) => None,
319                };
320                if let Some(guard) = handshake_guard {
321                    guard.record_result(match &result {
322                        Some(Ok(_)) => HandshakeResult::Completed,
323                        Some(Err(_)) => HandshakeResult::Failed,
324                        None => HandshakeResult::TimedOut,
325                    });
326                }
327                result.unwrap_or_else(|| {
328                    Err(std::io::Error::new(
329                        std::io::ErrorKind::TimedOut,
330                        "TLS handshake timed out",
331                    )
332                    .into())
333                })
334            });
335        } else {
336            self.handle_connection(ServerIo::new_io(io), remote_addr);
337        }
338    }
339
340    fn handle_connection(&mut self, io: ServerIo<TrackedIo<L::Io>>, remote_addr: L::Addr) {
341        let mut peer_connection_guard = None;
342        if let (Some(max), Some(peer)) =
343            (self.config.max_connections_per_peer, peer_public_key(&io))
344        {
345            let Some(guard) = self.peer_connection_counts.register(&peer, max) else {
346                // Dropping the connection closes it, releasing its file descriptor.
347                trace!("peer already holds {max} connections, closing the new one");
348                return;
349            };
350            peer_connection_guard = Some(guard);
351        }
352
353        let connection_shutdown_token = self.graceful_shutdown_token.child_token();
354        let connection_info = ConnectionInfo::new(
355            remote_addr,
356            io.peer_certs(),
357            connection_shutdown_token.clone(),
358        );
359        let connection_id = connection_info.id();
360        let connect_info = connection_info::ConnectInfo {
361            local_addr: self.local_addr,
362            remote_addr: *connection_info.remote_address(),
363        };
364        let peer_certificates = connection_info.peer_certificates().cloned();
365        let connection_metrics = io.connection().map(|connection| connection.metrics());
366        let hyper_io = hyper_util::rt::TokioIo::new(io);
367
368        let service = TrackRequests::new(self.service.clone(), connection_metrics);
369        let hyper_svc = TowerToHyperService::new(service.map_request(
370            move |mut request: Request<hyper::body::Incoming>| {
371                request.extensions_mut().insert(connect_info.clone());
372                if let Some(peer_certificates) = peer_certificates.clone() {
373                    request.extensions_mut().insert(peer_certificates);
374                }
375
376                request.map(body::boxed)
377            },
378        ));
379
380        self.connections
381            .write()
382            .unwrap()
383            .insert(connection_id, connection_info);
384        let on_connection_close = OnConnectionClose::new(
385            connection_id,
386            self.connections.clone(),
387            peer_connection_guard,
388        );
389
390        self.connection_handlers
391            .spawn(connection_handler::serve_connection(
392                hyper_io,
393                hyper_svc,
394                self.config.connection_builder(),
395                connection_shutdown_token,
396                self.config.max_connection_age,
397                on_connection_close,
398            ));
399    }
400
401    async fn shutdown(mut self) {
402        // The time we are willing to wait for a connection to get gracefully shutdown
403        // before we attempt to forcefully shutdown all active connections
404        const CONNECTION_SHUTDOWN_GRACE_PERIOD: Duration = Duration::from_secs(1);
405
406        // Just to be careful make sure the token is canceled
407        self.graceful_shutdown_token.cancel();
408
409        // Terminate any in-progress pending connections
410        self.pending_connections.shutdown().await;
411
412        // Wait for all connection handlers to terminate
413        trace!(
414            "waiting for {} connections to close",
415            self.connection_handlers.len()
416        );
417
418        let graceful_shutdown =
419            async { while self.connection_handlers.join_next().await.is_some() {} };
420
421        if tokio::time::timeout(CONNECTION_SHUTDOWN_GRACE_PERIOD, graceful_shutdown)
422            .await
423            .is_err()
424        {
425            tracing::warn!(
426                "Failed to stop all connection handlers in {:?}. Forcing shutdown.",
427                CONNECTION_SHUTDOWN_GRACE_PERIOD
428            );
429            self.connection_handlers.shutdown().await;
430        }
431    }
432}
433
434/// Runs the TLS handshake to completion.
435async fn handshake<Io, Addr>(
436    io: TrackedIo<Io>,
437    remote_addr: Addr,
438    tls_acceptor: TlsAcceptor,
439) -> ConnectingOutput<Io, Addr>
440where
441    Io: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static,
442{
443    tracing::trace!("accepting TLS connection");
444    let io = tls_acceptor.accept(io).await?;
445    Ok((ServerIo::new_tls_io(io), remote_addr))
446}
447
448/// Identifies the peer by the public key of the single certificate it
449/// authenticated with, or `None` if it presented no certificate.
450fn peer_public_key<Io>(io: &ServerIo<Io>) -> Option<Vec<u8>> {
451    let certs = io.peer_certs()?;
452    let [certificate] = certs.as_slice() else {
453        trace!("unexpected number of peer certificates: {}", certs.len());
454        return None;
455    };
456
457    match iota_tls::public_key_from_certificate(certificate) {
458        Ok(public_key) => Some(AsRef::<[u8]>::as_ref(&public_key).to_vec()),
459        Err(e) => {
460            trace!("failed to extract the public key from the peer certificate: {e:?}");
461            None
462        }
463    }
464}
465
466#[cfg(test)]
467mod tests {
468    use axum::Router;
469
470    use super::*;
471
472    #[tokio::test]
473    async fn simple() {
474        const MESSAGE: &str = "Hello, World!";
475
476        let app = Router::new().route("/", axum::routing::get(|| async { MESSAGE }));
477
478        let handle = Builder::new().serve(("localhost", 0), app).unwrap();
479
480        let url = format!("http://{}", handle.local_addr());
481
482        let _ = rustls::crypto::ring::default_provider().install_default();
483        let response = reqwest::get(url).await.unwrap().bytes().await.unwrap();
484
485        assert_eq!(response, MESSAGE.as_bytes());
486    }
487
488    #[tokio::test]
489    async fn shutdown() {
490        const MESSAGE: &str = "Hello, World!";
491
492        let app = Router::new().route("/", axum::routing::get(|| async { MESSAGE }));
493
494        let handle = Builder::new().serve(("localhost", 0), app).unwrap();
495
496        let url = format!("http://{}", handle.local_addr());
497
498        let _ = rustls::crypto::ring::default_provider().install_default();
499        let response = reqwest::get(url).await.unwrap().bytes().await.unwrap();
500
501        // a request was just made so we should have 1 active connection
502        assert_eq!(handle.connections().len(), 1);
503
504        assert_eq!(response, MESSAGE.as_bytes());
505
506        assert!(!handle.is_shutdown());
507
508        handle.shutdown().await;
509
510        assert!(handle.is_shutdown());
511
512        // Now that the network has been shutdown there should be zero connections
513        assert_eq!(handle.connections().len(), 0);
514    }
515
516    const SERVER_NAME: &str = "iota-http-test";
517
518    /// A server config and a client config that trusts it, without client
519    /// authentication.
520    fn test_tls_configs() -> (rustls::ServerConfig, rustls::ClientConfig) {
521        use fastcrypto::{
522            ed25519::{Ed25519KeyPair, Ed25519PrivateKey},
523            traits::{KeyPair, ToFromBytes},
524        };
525
526        let keypair = Ed25519KeyPair::from(Ed25519PrivateKey::from_bytes(&[42; 32]).unwrap());
527        let public_key = keypair.public().to_owned();
528        (
529            iota_tls::create_rustls_server_config(keypair.private(), SERVER_NAME.to_string()),
530            iota_tls::create_rustls_client_config(public_key, SERVER_NAME.to_string(), None),
531        )
532    }
533
534    /// The peer is unauthenticated until its handshake completes, so a peer
535    /// that never starts one must not hold the connection open indefinitely.
536    #[tokio::test]
537    async fn silent_peer_is_dropped_after_the_handshake_timeout() {
538        use tokio::io::AsyncReadExt as _;
539
540        const HANDSHAKE_TIMEOUT: Duration = Duration::from_millis(200);
541
542        let (server_tls_config, _) = test_tls_configs();
543        let handle = Builder::new()
544            .config(Config::default().handshake_timeout(Some(HANDSHAKE_TIMEOUT)))
545            .tls_config(server_tls_config)
546            .serve(("localhost", 0), Router::new())
547            .unwrap();
548
549        // Connect, then never send a ClientHello.
550        let mut connection = tokio::net::TcpStream::connect(handle.local_addr())
551            .await
552            .unwrap();
553
554        let mut buf = [0u8; 1];
555        let read = tokio::time::timeout(HANDSHAKE_TIMEOUT * 25, connection.read(&mut buf))
556            .await
557            .expect("the server must not wait for the handshake past the deadline");
558
559        assert!(
560            matches!(read, Ok(0) | Err(_)),
561            "the server must close the connection, got {read:?}"
562        );
563    }
564
565    /// A TLS-configured listener only speaks TLS: a peer that starts with a
566    /// plaintext HTTP/2 preface is closed before any request is dispatched.
567    #[tokio::test]
568    async fn plaintext_peer_is_refused_by_a_tls_listener() {
569        use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
570
571        let (server_tls_config, _) = test_tls_configs();
572        let served = Arc::new(std::sync::atomic::AtomicBool::new(false));
573        let app = Router::new().route(
574            "/",
575            axum::routing::get({
576                let served = served.clone();
577                move || async move {
578                    served.store(true, std::sync::atomic::Ordering::SeqCst);
579                    "served"
580                }
581            }),
582        );
583        let handle = Builder::new()
584            .tls_config(server_tls_config)
585            .serve(("localhost", 0), app)
586            .unwrap();
587
588        let mut connection = tokio::net::TcpStream::connect(handle.local_addr())
589            .await
590            .unwrap();
591        connection
592            .write_all(b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n")
593            .await
594            .unwrap();
595        connection
596            .write_all(&[0, 0, 0, 0x4, 0, 0, 0, 0, 0])
597            .await
598            .unwrap();
599
600        // The TLS acceptor answers garbage with an alert record and closes;
601        // anything else would mean the plaintext bytes were served.
602        let received = tokio::time::timeout(Duration::from_secs(5), async {
603            let mut received = Vec::new();
604            let mut buf = [0u8; 64];
605            loop {
606                match connection.read(&mut buf).await {
607                    Ok(0) | Err(_) => break received,
608                    Ok(n) => received.extend_from_slice(&buf[..n]),
609                }
610            }
611        })
612        .await
613        .expect("the server must close a plaintext connection promptly");
614        const TLS_ALERT_RECORD: u8 = 0x15;
615        assert!(
616            received
617                .first()
618                .is_none_or(|first| *first == TLS_ALERT_RECORD),
619            "the server must answer plaintext with a TLS alert at most, got {received:?}"
620        );
621        assert!(
622            !served.load(std::sync::atomic::Ordering::SeqCst),
623            "no request may reach the service over plaintext"
624        );
625    }
626
627    /// An HTTP/1 peer that never finishes sending its request headers is
628    /// closed once the header deadline passes.
629    #[tokio::test]
630    async fn http1_peer_that_never_finishes_its_headers_is_closed() {
631        use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
632
633        const HEADER_DEADLINE: Duration = Duration::from_millis(200);
634
635        let handle = Builder::new()
636            .config(Config::default().http1_header_read_timeout(Some(HEADER_DEADLINE)))
637            .serve(("localhost", 0), Router::new())
638            .unwrap();
639
640        let mut connection = tokio::net::TcpStream::connect(handle.local_addr())
641            .await
642            .unwrap();
643        // A request line and one header, but never the blank line that ends
644        // the header block.
645        connection
646            .write_all(b"GET / HTTP/1.1\r\nHost: localhost\r\n")
647            .await
648            .unwrap();
649        let sent_at = tokio::time::Instant::now();
650
651        // Without the deadline the peer is held forever, so a wide window here
652        // only guards against a slow machine, it does not weaken the assertion.
653        let closed = tokio::time::timeout(HEADER_DEADLINE * 25, async {
654            let mut buf = [0u8; 1024];
655            loop {
656                match connection.read(&mut buf).await {
657                    Ok(0) | Err(_) => break,
658                    Ok(_) => continue,
659                }
660            }
661        })
662        .await;
663        assert!(
664            closed.is_ok(),
665            "the server must close a peer that stalls its request headers"
666        );
667        assert!(
668            sent_at.elapsed() >= HEADER_DEADLINE,
669            "the peer must be given the whole deadline before it is closed"
670        );
671    }
672
673    /// A server that does not accept HTTP/1 refuses it on plaintext, and on
674    /// TLS whether or not the client chose HTTP/2 through ALPN. A plaintext
675    /// HTTP/2 request is still served.
676    #[tokio::test]
677    async fn a_server_without_http1_does_not_serve_an_http1_request() {
678        use std::sync::atomic::{AtomicUsize, Ordering};
679
680        let (server_tls_config, client_tls_config) = test_tls_configs();
681        let served = Arc::new(AtomicUsize::new(0));
682        let app = Router::new().route(
683            "/",
684            axum::routing::get({
685                let served = served.clone();
686                move || async move {
687                    served.fetch_add(1, Ordering::SeqCst);
688                    "served"
689                }
690            }),
691        );
692        let config = Config::default().accept_http1(false);
693        let plaintext = Builder::new()
694            .config(config.clone())
695            .serve(("localhost", 0), app.clone())
696            .unwrap();
697        let tls = Builder::new()
698            .config(config)
699            .tls_config(server_tls_config)
700            .serve(("localhost", 0), app)
701            .unwrap();
702
703        let connection = tokio::net::TcpStream::connect(plaintext.local_addr())
704            .await
705            .unwrap();
706        assert_http1_request_is_refused(connection).await;
707
708        for alpn_protocols in [vec![], vec![ALPN_H2.to_vec()]] {
709            let mut client_tls_config = client_tls_config.clone();
710            client_tls_config.alpn_protocols = alpn_protocols;
711            let connection = tokio::net::TcpStream::connect(tls.local_addr())
712                .await
713                .unwrap();
714            let connection = tokio_rustls::TlsConnector::from(Arc::new(client_tls_config))
715                .connect(
716                    rustls::pki_types::ServerName::try_from(SERVER_NAME).unwrap(),
717                    connection,
718                )
719                .await
720                .unwrap();
721            assert_http1_request_is_refused(connection).await;
722        }
723
724        assert_eq!(
725            served.load(Ordering::SeqCst),
726            0,
727            "no HTTP/1 request may reach the service"
728        );
729
730        let _ = rustls::crypto::ring::default_provider().install_default();
731        let response = reqwest::Client::builder()
732            .http2_prior_knowledge()
733            .build()
734            .unwrap()
735            .get(format!("http://{}", plaintext.local_addr()))
736            .send()
737            .await
738            .unwrap();
739        assert_eq!(response.text().await.unwrap(), "served");
740    }
741
742    /// Sends an HTTP/1 request and asserts that the server answers as HTTP/2
743    /// and closes the connection.
744    async fn assert_http1_request_is_refused(
745        mut connection: impl tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
746    ) {
747        use tokio::io::AsyncWriteExt as _;
748
749        // The type of an HTTP/2 SETTINGS frame (RFC 9113, section 6.5).
750        const HTTP2_SETTINGS_FRAME: u8 = 0x4;
751
752        connection
753            .write_all(b"GET / HTTP/1.1\r\nHost: localhost\r\n\r\n")
754            .await
755            .unwrap();
756
757        let mut received = Vec::new();
758        let closed = read_until_done_or_closed(&mut connection, &mut received, |_| false).await;
759        assert_eq!(
760            frame_types(&received).first(),
761            Some(&HTTP2_SETTINGS_FRAME),
762            "the server must answer as HTTP/2, got {:?}",
763            String::from_utf8_lossy(&received)
764        );
765        assert!(closed, "the server must close the connection");
766    }
767
768    /// A served HTTP/2 connection gets a GOAWAY frame when the server shuts
769    /// down and when the connection reaches its maximum age, whether or not
770    /// the server accepts HTTP/1.
771    #[tokio::test]
772    async fn an_http2_connection_gets_goaway_on_shutdown_and_at_its_maximum_age() {
773        use tokio::io::AsyncWriteExt as _;
774
775        // The client connection preface: the fixed string and an empty
776        // SETTINGS frame (RFC 9113, section 3.4).
777        const HTTP2_CLIENT_PREFACE: &[u8] = b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n\0\0\0\x04\0\0\0\0\0";
778        // The types of HTTP/2 SETTINGS and GOAWAY frames (RFC 9113, sections
779        // 6.5 and 6.8).
780        const HTTP2_SETTINGS_FRAME: u8 = 0x4;
781        const HTTP2_GOAWAY_FRAME: u8 = 0x7;
782
783        for accept_http1 in [true, false] {
784            for max_connection_age in [None, Some(Duration::from_millis(500))] {
785                let mut config = Config::default().accept_http1(accept_http1);
786                if let Some(max_connection_age) = max_connection_age {
787                    config = config.max_connection_age(max_connection_age);
788                }
789                let handle = Builder::new()
790                    .config(config)
791                    .serve(("localhost", 0), Router::new())
792                    .unwrap();
793                let mut connection = tokio::net::TcpStream::connect(handle.local_addr())
794                    .await
795                    .unwrap();
796                connection.write_all(HTTP2_CLIENT_PREFACE).await.unwrap();
797
798                let mut received = Vec::new();
799                read_until_done_or_closed(&mut connection, &mut received, |received| {
800                    !frame_types(received).is_empty()
801                })
802                .await;
803                assert_eq!(
804                    frame_types(&received).first(),
805                    Some(&HTTP2_SETTINGS_FRAME),
806                    "the server must serve the connection as HTTP/2"
807                );
808
809                if max_connection_age.is_none() {
810                    handle.trigger_shutdown();
811                }
812                read_until_done_or_closed(&mut connection, &mut received, |received| {
813                    frame_types(received).contains(&HTTP2_GOAWAY_FRAME)
814                })
815                .await;
816                assert!(
817                    frame_types(&received).contains(&HTTP2_GOAWAY_FRAME),
818                    "accept_http1: {accept_http1}, max_connection_age: {max_connection_age:?}, \
819                     frames: {:?}",
820                    frame_types(&received)
821                );
822            }
823        }
824    }
825
826    /// Reads from `connection` into `received` until `done(received)` holds or
827    /// the peer closes the connection. Returns `false` if neither happens
828    /// within 5 s.
829    async fn read_until_done_or_closed(
830        connection: &mut (impl tokio::io::AsyncRead + Unpin),
831        received: &mut Vec<u8>,
832        done: impl Fn(&[u8]) -> bool,
833    ) -> bool {
834        use tokio::io::AsyncReadExt as _;
835
836        tokio::time::timeout(Duration::from_secs(5), async {
837            let mut buf = [0u8; 1024];
838            while !done(received) {
839                match connection.read(&mut buf).await {
840                    Ok(0) | Err(_) => break,
841                    Ok(n) => received.extend_from_slice(&buf[..n]),
842                }
843            }
844        })
845        .await
846        .is_ok()
847    }
848
849    /// The types of the HTTP/2 frames whose header is in `bytes`.
850    fn frame_types(bytes: &[u8]) -> Vec<u8> {
851        let mut types = Vec::new();
852        let mut rest = bytes;
853        while let Some((header, payload)) = rest.split_first_chunk::<9>() {
854            types.push(header[3]);
855            let length = u32::from_be_bytes([0, header[0], header[1], header[2]]) as usize;
856            rest = payload.get(length..).unwrap_or_default();
857        }
858        types
859    }
860
861    /// A peer that completes the HTTP/2 preface and then falls silent is
862    /// closed once it fails to answer the keepalive ping.
863    #[tokio::test]
864    async fn silent_http2_peer_is_closed_by_keepalive() {
865        use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
866
867        const KEEPALIVE: Duration = Duration::from_millis(100);
868
869        let handle = Builder::new()
870            .config(
871                Config::default()
872                    .http2_keepalive_interval(Some(KEEPALIVE))
873                    .http2_keepalive_timeout(Some(KEEPALIVE)),
874            )
875            .serve(("localhost", 0), Router::new())
876            .unwrap();
877
878        let mut connection = tokio::net::TcpStream::connect(handle.local_addr())
879            .await
880            .unwrap();
881        // Connection preface and an empty SETTINGS frame, then nothing more.
882        connection
883            .write_all(b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n")
884            .await
885            .unwrap();
886        connection
887            .write_all(&[0, 0, 0, 0x4, 0, 0, 0, 0, 0])
888            .await
889            .unwrap();
890
891        let closed = tokio::time::timeout(KEEPALIVE * 50, async {
892            let mut buf = [0u8; 1024];
893            loop {
894                match connection.read(&mut buf).await {
895                    Ok(0) | Err(_) => break,
896                    Ok(_) => continue,
897                }
898            }
899        })
900        .await;
901        assert!(
902            closed.is_ok(),
903            "the server must close a peer that ignores keepalive pings"
904        );
905    }
906
907    /// While the pending limit is reached the server stops accepting, and
908    /// resumes once a slot frees.
909    #[tokio::test]
910    async fn pending_connections_are_capped() {
911        let (server_tls_config, client_tls_config) = test_tls_configs();
912        let handle = Builder::new()
913            .config(
914                Config::default()
915                    // Long enough that the stalled peer below is released by the test
916                    // rather than by the deadline.
917                    .handshake_timeout(Some(Duration::from_secs(60)))
918                    .max_pending_connections(Some(1)),
919            )
920            .tls_config(server_tls_config)
921            .serve(("localhost", 0), Router::new())
922            .unwrap();
923        let addr = *handle.local_addr();
924
925        // Occupy the only handshake slot with a peer that never speaks.
926        let stalled = tokio::net::TcpStream::connect(addr).await.unwrap();
927        tokio::time::sleep(Duration::from_millis(200)).await;
928
929        let connector = tokio_rustls::TlsConnector::from(Arc::new(client_tls_config));
930        let server_name = rustls::pki_types::ServerName::try_from(SERVER_NAME).unwrap();
931        let handshake = async {
932            let io = tokio::net::TcpStream::connect(addr).await.unwrap();
933            connector.connect(server_name, io).await
934        };
935        tokio::pin!(handshake);
936
937        // The connection reaches the kernel backlog but is not accepted. Without the
938        // limit the handshake completes in milliseconds, so a wide window here only
939        // guards against a slow machine, it does not weaken the assertion.
940        assert!(
941            tokio::time::timeout(Duration::from_secs(2), &mut handshake)
942                .await
943                .is_err(),
944            "the server must not handshake while the pending limit is reached"
945        );
946
947        drop(stalled);
948
949        tokio::time::timeout(Duration::from_secs(10), &mut handshake)
950            .await
951            .expect("the server must resume accepting once a slot frees")
952            .expect("the handshake must succeed");
953    }
954
955    /// A peer at its connection limit gets no further connections, while every
956    /// other peer is served as usual and the peer itself recovers a slot as
957    /// soon as one of its connections closes.
958    #[tokio::test]
959    async fn connections_per_peer_are_capped() {
960        use fastcrypto::{
961            ed25519::{Ed25519KeyPair, Ed25519PrivateKey},
962            traits::{KeyPair, ToFromBytes},
963        };
964        use tokio::io::AsyncReadExt as _;
965
966        const MAX_PER_PEER: usize = 2;
967
968        let client_key =
969            |seed: u8| Ed25519KeyPair::from(Ed25519PrivateKey::from_bytes(&[seed; 32]).unwrap());
970        let server_keypair = client_key(1);
971        let server_public_key = server_keypair.public().to_owned();
972        let server_config = iota_tls::create_rustls_server_config_with_client_verifier(
973            server_keypair.private(),
974            SERVER_NAME.to_string(),
975            iota_tls::AllowPublicKeys::new(
976                [
977                    client_key(2).public().to_owned(),
978                    client_key(3).public().to_owned(),
979                ]
980                .into(),
981            ),
982        );
983
984        let events = Arc::new(std::sync::Mutex::new(Vec::new()));
985        let handle = Builder::new()
986            .config(
987                Config::default()
988                    .max_connections_per_peer(Some(MAX_PER_PEER))
989                    .on_peer_connection_event({
990                        let events = events.clone();
991                        move |peer, event| events.lock().unwrap().push((peer.to_vec(), event))
992                    }),
993            )
994            .tls_config(server_config)
995            .serve(("localhost", 0), Router::new())
996            .unwrap();
997        let addr = *handle.local_addr();
998
999        let server_name = rustls::pki_types::ServerName::try_from(SERVER_NAME).unwrap();
1000        let connect = |seed: u8| {
1001            let connector =
1002                tokio_rustls::TlsConnector::from(Arc::new(iota_tls::create_rustls_client_config(
1003                    server_public_key.clone(),
1004                    SERVER_NAME.to_string(),
1005                    Some(client_key(seed).private()),
1006                )));
1007            let server_name = server_name.clone();
1008            async move {
1009                let io = tokio::net::TcpStream::connect(addr).await.unwrap();
1010                connector.connect(server_name, io).await.unwrap()
1011            }
1012        };
1013
1014        let events_of = |seed: u8| -> Vec<PeerConnectionEvent> {
1015            let peer = client_key(seed).public().as_ref().to_vec();
1016            events
1017                .lock()
1018                .unwrap()
1019                .iter()
1020                .filter(|(reported, _)| *reported == peer)
1021                .map(|(_, event)| *event)
1022                .collect()
1023        };
1024
1025        // The connection is established before the server has necessarily
1026        // registered or dropped it, so settle on the count rather than assert
1027        // on it immediately.
1028        let established = async |expected: usize| {
1029            for _ in 0..100 {
1030                if handle.number_of_connections() == expected {
1031                    return;
1032                }
1033                tokio::time::sleep(Duration::from_millis(20)).await;
1034            }
1035            panic!(
1036                "expected {expected} connections, the server holds {}",
1037                handle.number_of_connections()
1038            );
1039        };
1040
1041        let mut peer = Vec::new();
1042        for _ in 0..MAX_PER_PEER {
1043            peer.push(connect(2).await);
1044        }
1045        established(MAX_PER_PEER).await;
1046
1047        let mut refused = connect(2).await;
1048        let mut buf = [0u8; 1];
1049        let read = tokio::time::timeout(Duration::from_secs(10), refused.read(&mut buf))
1050            .await
1051            .expect("the server must close a connection past the limit");
1052        assert!(
1053            matches!(read, Ok(0) | Err(_)),
1054            "the server must close the connection, got {read:?}"
1055        );
1056        established(MAX_PER_PEER).await;
1057        assert_eq!(
1058            events_of(2),
1059            [
1060                PeerConnectionEvent::Established { held: 1 },
1061                PeerConnectionEvent::Established { held: 2 },
1062                PeerConnectionEvent::RefusedAtLimit { held: 2 },
1063            ],
1064            "the refusal must be reported against the peer that caused it"
1065        );
1066
1067        // A peer sitting at its limit must not consume anyone else's budget.
1068        let other_peer = connect(3).await;
1069        established(MAX_PER_PEER + 1).await;
1070
1071        peer.pop();
1072        established(MAX_PER_PEER).await;
1073        let reconnected = connect(2).await;
1074        established(MAX_PER_PEER + 1).await;
1075        assert_eq!(
1076            events_of(2),
1077            [
1078                PeerConnectionEvent::Established { held: 1 },
1079                PeerConnectionEvent::Established { held: 2 },
1080                PeerConnectionEvent::RefusedAtLimit { held: 2 },
1081                PeerConnectionEvent::Closed { held: 1 },
1082                PeerConnectionEvent::Established { held: 2 },
1083            ],
1084            "closing and reconnecting must be reported with the count held afterwards"
1085        );
1086        assert_eq!(
1087            events_of(3),
1088            [PeerConnectionEvent::Established { held: 1 }],
1089            "another peer's events must not leak into this one's"
1090        );
1091
1092        drop((peer, other_peer, reconnected));
1093    }
1094
1095    /// The cap identifies a peer by its single certificate, so a peer sending
1096    /// a longer chain is refused at the handshake and never holds a
1097    /// connection.
1098    #[tokio::test]
1099    async fn peer_with_extra_certificates_is_refused_at_the_handshake() {
1100        use fastcrypto::{
1101            ed25519::{Ed25519KeyPair, Ed25519PrivateKey},
1102            traits::{KeyPair, ToFromBytes},
1103        };
1104        use tokio::io::AsyncReadExt as _;
1105
1106        let key =
1107            |seed: u8| Ed25519KeyPair::from(Ed25519PrivateKey::from_bytes(&[seed; 32]).unwrap());
1108        let server_public_key = key(1).public().to_owned();
1109        let client_public_key = key(2).public().to_owned();
1110        let server_config = iota_tls::create_rustls_server_config_with_client_verifier(
1111            key(1).private(),
1112            SERVER_NAME.to_string(),
1113            iota_tls::AllowPublicKeys::new([client_public_key].into()),
1114        );
1115        let handle = Builder::new()
1116            .config(Config::default().max_connections_per_peer(Some(1)))
1117            .tls_config(server_config)
1118            .serve(("localhost", 0), Router::new())
1119            .unwrap();
1120
1121        // The allowed certificate followed by an unrelated one.
1122        let client_certificate =
1123            iota_tls::SelfSignedCertificate::new(key(2).private(), SERVER_NAME);
1124        let extra_certificate = iota_tls::SelfSignedCertificate::new(key(3).private(), SERVER_NAME);
1125        let client_config =
1126            iota_tls::ServerCertVerifier::new(server_public_key, SERVER_NAME.to_string())
1127                .rustls_client_config_with_client_auth(
1128                    vec![
1129                        client_certificate.rustls_certificate(),
1130                        extra_certificate.rustls_certificate(),
1131                    ],
1132                    client_certificate.rustls_private_key(),
1133                )
1134                .unwrap();
1135
1136        // In TLS 1.3 the client may consider the handshake done before the
1137        // server's rejection arrives, so the refusal can surface on the first
1138        // read instead.
1139        let connector = tokio_rustls::TlsConnector::from(Arc::new(client_config));
1140        let io = tokio::net::TcpStream::connect(handle.local_addr())
1141            .await
1142            .unwrap();
1143        let server_name = rustls::pki_types::ServerName::try_from(SERVER_NAME).unwrap();
1144        if let Ok(mut connection) = connector.connect(server_name, io).await {
1145            let mut buf = [0u8; 1];
1146            let read = tokio::time::timeout(Duration::from_secs(10), connection.read(&mut buf))
1147                .await
1148                .expect("the server must close a connection with more than one certificate");
1149            assert!(
1150                matches!(read, Ok(0) | Err(_)),
1151                "the server must close the connection, got {read:?}"
1152            );
1153        }
1154        assert_eq!(handle.number_of_connections(), 0);
1155    }
1156
1157    /// A limit the accept loop can never fall below, and a limit whose slots
1158    /// nothing releases, both leave the server unable to accept.
1159    #[tokio::test]
1160    async fn unrecoverable_limits_are_rejected() {
1161        let served = |config| {
1162            Builder::new()
1163                .config(config)
1164                .tls_config(test_tls_configs().0)
1165                .serve(("localhost", 0), Router::new())
1166        };
1167
1168        assert!(
1169            served(Config::default().max_pending_connections(Some(0))).is_err(),
1170            "a zero limit must be rejected"
1171        );
1172        assert!(
1173            served(
1174                Config::default()
1175                    .handshake_timeout(None)
1176                    .max_pending_connections(Some(8))
1177            )
1178            .is_err(),
1179            "a limit without a handshake deadline must be rejected"
1180        );
1181        assert!(
1182            served(
1183                Config::default()
1184                    .handshake_timeout(None)
1185                    .max_pending_connections(None)
1186            )
1187            .is_ok(),
1188            "removing both bounds stays allowed"
1189        );
1190    }
1191}