1use 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>;
37const ALPN_H2: &[u8] = b"h2";
39const 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 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 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 pub fn local_addr(&self) -> &A {
183 &self.0.local_addr
184 }
185
186 pub fn trigger_shutdown(&self) {
189 self.0.graceful_shutdown_token.cancel();
190 }
191
192 pub async fn wait_for_shutdown(&self) {
198 self.0.watch_sender.closed().await
199 }
200
201 pub async fn shutdown(&self) {
204 self.trigger_shutdown();
205 self.wait_for_shutdown().await;
206 }
207
208 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 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 _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 (io, remote_addr) = self.listener.accept(), if self.accepts_more_connections() => {
264 self.handle_incoming(io, remote_addr);
265 },
266 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 self.shutdown().await;
296
297 Ok(())
298 }
299
300 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 _ = 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 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 const CONNECTION_SHUTDOWN_GRACE_PERIOD: Duration = Duration::from_secs(1);
405
406 self.graceful_shutdown_token.cancel();
408
409 self.pending_connections.shutdown().await;
411
412 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
434async 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
448fn 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 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 assert_eq!(handle.connections().len(), 0);
514 }
515
516 const SERVER_NAME: &str = "iota-http-test";
517
518 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 #[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 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 #[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 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 #[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 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 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 #[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 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 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 #[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 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 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 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 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 #[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
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 #[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 .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 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 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 #[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 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 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 #[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 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 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 #[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}