alice.pds.demo.boltless.dev
This repository has no description
did:plc:3fwecdnvtcscjnrx2p4n7a…
1use std::net::SocketAddr;
2use std::sync::Arc;
3use std::time::Duration;
4
5use axum::Router;
6use axum::body::Body;
7use hyper::body::Incoming;
8use hyper::service::service_fn;
9use hyper_util::rt::{TokioExecutor, TokioIo, TokioTimer};
10use hyper_util::server::conn::auto::Builder;
11use rustls::ServerConfig;
12use tokio::io::{AsyncRead, AsyncWrite};
13use tokio::net::{TcpListener, TcpStream};
14use tokio::sync::Semaphore;
15use tokio_rustls::TlsAcceptor;
16use tokio_util::sync::CancellationToken;
17use tokio_util::task::TaskTracker;
18use tower::{Service, ServiceExt};
19
20use crate::limits::ListenLimits;
21
22const HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(10);
23const CONNECTION_DRAIN_GRACE: Duration = Duration::from_secs(10);
24const LISTENER_DRAIN_GRACE: Duration = Duration::from_secs(30);
25const ACCEPT_BACKOFF: Duration = Duration::from_millis(250);
26
27pub async fn serve_plaintext(
28 listener: TcpListener,
29 router: Router,
30 limits: ListenLimits,
31 shutdown: CancellationToken,
32) -> std::io::Result<()> {
33 run_listener(listener, router, limits, shutdown, |stream| async move {
34 Some(stream)
35 })
36 .await
37}
38
39pub async fn serve_tls(
40 listener: TcpListener,
41 router: Router,
42 server_config: Arc<ServerConfig>,
43 limits: ListenLimits,
44 shutdown: CancellationToken,
45) -> std::io::Result<()> {
46 let acceptor = TlsAcceptor::from(server_config);
47 run_listener(listener, router, limits, shutdown, move |stream| {
48 let acceptor = acceptor.clone();
49 async move {
50 match tokio::time::timeout(HANDSHAKE_TIMEOUT, acceptor.accept(stream)).await {
51 Ok(Ok(tls_stream)) => Some(tls_stream),
52 Ok(Err(error)) => {
53 tracing::debug!("tls handshake failed: {error}");
54 None
55 }
56 Err(_) => {
57 tracing::debug!("tls handshake timed out after {HANDSHAKE_TIMEOUT:?}");
58 None
59 }
60 }
61 }
62 })
63 .await
64}
65
66async fn run_listener<IO, Upgrade, Fut>(
67 listener: TcpListener,
68 router: Router,
69 limits: ListenLimits,
70 shutdown: CancellationToken,
71 upgrade: Upgrade,
72) -> std::io::Result<()>
73where
74 IO: AsyncRead + AsyncWrite + Unpin + Send + 'static,
75 Upgrade: Fn(TcpStream) -> Fut + Send + Sync + 'static,
76 Fut: std::future::Future<Output = Option<IO>> + Send + 'static,
77{
78 let slots = Arc::new(Semaphore::new(limits.max_connections()));
79 let mut make_service = router.into_make_service_with_connect_info::<SocketAddr>();
80 let tracker = TaskTracker::new();
81 let upgrade = Arc::new(upgrade);
82
83 loop {
84 let accepted = tokio::select! {
85 () = shutdown.cancelled() => break,
86 accepted = listener.accept() => accepted,
87 };
88 let (stream, peer) = match accepted {
89 Ok(pair) => pair,
90 Err(error) if is_connection_error(&error) => continue,
91 Err(error) => {
92 tracing::warn!("tcp accept failed, backing off: {error}");
93 tokio::select! {
94 () = shutdown.cancelled() => break,
95 () = tokio::time::sleep(ACCEPT_BACKOFF) => {}
96 }
97 continue;
98 }
99 };
100 let Ok(slot) = Arc::clone(&slots).try_acquire_owned() else {
101 continue;
102 };
103 let service = match make_service.call(peer).await {
104 Ok(service) => service,
105 Err(never) => match never {},
106 };
107 let header_timeout = limits.header_timeout().get();
108 let conn_shutdown = shutdown.clone();
109 let upgrade = Arc::clone(&upgrade);
110 tracker.spawn(async move {
111 let _slot = slot;
112 let Some(io) = upgrade(stream).await else {
113 return;
114 };
115 let hyper_service = service_fn(move |request: hyper::Request<Incoming>| {
116 service.clone().oneshot(request.map(Body::new))
117 });
118 let budget = limits.connection_budget();
119 let mut builder = Builder::new(TokioExecutor::new());
120 builder
121 .http1()
122 .timer(TokioTimer::new())
123 .header_read_timeout(header_timeout);
124 builder
125 .http2()
126 .timer(TokioTimer::new())
127 .max_concurrent_streams(budget.max_concurrent_streams().get())
128 .initial_stream_window_size(budget.stream_receive_window())
129 .initial_connection_window_size(budget.connection_receive_window())
130 .keep_alive_interval(Some(header_timeout))
131 .keep_alive_timeout(header_timeout);
132 let connection =
133 builder.serve_connection_with_upgrades(TokioIo::new(io), hyper_service);
134 tokio::pin!(connection);
135
136 tokio::select! {
137 served = connection.as_mut() => drop(served),
138 () = conn_shutdown.cancelled() => {
139 connection.as_mut().graceful_shutdown();
140 let _ = tokio::time::timeout(CONNECTION_DRAIN_GRACE, connection.as_mut()).await;
141 }
142 }
143 });
144 }
145
146 tracker.close();
147 let _ = tokio::time::timeout(LISTENER_DRAIN_GRACE, tracker.wait()).await;
148 Ok(())
149}
150
151fn is_connection_error(error: &std::io::Error) -> bool {
152 matches!(
153 error.kind(),
154 std::io::ErrorKind::ConnectionRefused
155 | std::io::ErrorKind::ConnectionAborted
156 | std::io::ErrorKind::ConnectionReset
157 )
158}
159
160#[cfg(test)]
161mod tests {
162 use super::*;
163
164 use std::num::{NonZeroU32, NonZeroU64};
165
166 use axum::routing::get;
167 use tokio::io::{AsyncReadExt, AsyncWriteExt};
168 use tokio::net::TcpStream;
169
170 fn limits(header_timeout_ms: u64, max_connections: u32) -> ListenLimits {
171 ListenLimits::new(
172 crate::limits::HeaderTimeout::from_millis(NonZeroU64::new(header_timeout_ms).unwrap()),
173 crate::limits::IdleTimeout::from_millis(NonZeroU64::new(30_000).unwrap()),
174 NonZeroU32::new(max_connections).unwrap(),
175 )
176 }
177
178 async fn bind_and_serve(limits: ListenLimits) -> (SocketAddr, CancellationToken) {
179 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
180 let addr = listener.local_addr().unwrap();
181 let router = Router::new().route("/", get(|| async { "ok" }));
182 let shutdown = CancellationToken::new();
183 tokio::spawn(serve_plaintext(listener, router, limits, shutdown.clone()));
184 (addr, shutdown)
185 }
186
187 async fn read_to_close(stream: &mut TcpStream) -> Vec<u8> {
188 let mut collected = Vec::new();
189 stream.read_to_end(&mut collected).await.unwrap();
190 collected
191 }
192
193 #[tokio::test]
194 async fn a_full_request_is_answered() {
195 let (addr, _shutdown) = bind_and_serve(limits(5_000, 4)).await;
196 let mut stream = TcpStream::connect(addr).await.unwrap();
197 stream
198 .write_all(b"GET / HTTP/1.1\r\nhost: oyster.cafe\r\nconnection: close\r\n\r\n")
199 .await
200 .unwrap();
201 let answer = read_to_close(&mut stream).await;
202 let head = String::from_utf8_lossy(&answer);
203 assert!(head.starts_with("HTTP/1.1 200"), "got: {head}");
204 }
205
206 #[tokio::test]
207 async fn a_slowloris_connection_is_cut_at_the_header_timeout() {
208 let (addr, _shutdown) = bind_and_serve(limits(200, 4)).await;
209 let mut stream = TcpStream::connect(addr).await.unwrap();
210 stream.write_all(b"GET / HTT").await.unwrap();
211 let closed = tokio::time::timeout(Duration::from_secs(5), read_to_close(&mut stream))
212 .await
213 .expect("server cuts connection instead of waiting forever");
214 let head = String::from_utf8_lossy(&closed);
215 assert!(
216 !head.contains("200"),
217 "half-sent request must never be answered, got: {head}"
218 );
219 }
220
221 #[tokio::test]
222 async fn an_idle_keep_alive_connection_is_cut_at_the_header_timeout() {
223 let (addr, _shutdown) = bind_and_serve(limits(200, 4)).await;
224 let mut stream = TcpStream::connect(addr).await.unwrap();
225 stream
226 .write_all(b"GET / HTTP/1.1\r\nhost: oyster.cafe\r\n\r\n")
227 .await
228 .unwrap();
229 let answer = tokio::time::timeout(Duration::from_secs(5), read_to_close(&mut stream))
230 .await
231 .expect("idle keep-alive connection is cut after answered request");
232 let head = String::from_utf8_lossy(&answer);
233 assert!(head.starts_with("HTTP/1.1 200"), "got: {head}");
234 }
235
236 #[tokio::test]
237 async fn a_connection_beyond_the_limit_is_refused() {
238 let (addr, _shutdown) = bind_and_serve(limits(5_000, 1)).await;
239 let mut held = TcpStream::connect(addr).await.unwrap();
240 held.write_all(b"GET / HTT").await.unwrap();
241 tokio::time::sleep(Duration::from_millis(100)).await;
242
243 let mut refused = TcpStream::connect(addr).await.unwrap();
244 refused.write_all(b"GET / HTT").await.unwrap();
245 let mut answer = Vec::new();
246 let outcome =
247 tokio::time::timeout(Duration::from_secs(1), refused.read_to_end(&mut answer))
248 .await
249 .expect("over-limit connection is dropped at accept instead of held to timeout");
250 match outcome {
251 Ok(_) => assert!(
252 answer.is_empty(),
253 "over-limit connection gets no bytes, got: {}",
254 String::from_utf8_lossy(&answer)
255 ),
256 Err(reset) => assert_eq!(reset.kind(), std::io::ErrorKind::ConnectionReset),
257 }
258 }
259
260 #[tokio::test]
261 async fn an_in_flight_request_finishes_during_drain() {
262 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
263 let addr = listener.local_addr().unwrap();
264 let router = Router::new().route(
265 "/slow",
266 get(|| async {
267 tokio::time::sleep(Duration::from_millis(300)).await;
268 "drained-clean"
269 }),
270 );
271 let shutdown = CancellationToken::new();
272 tokio::spawn(serve_plaintext(
273 listener,
274 router,
275 limits(5_000, 4),
276 shutdown.clone(),
277 ));
278
279 let mut stream = TcpStream::connect(addr).await.unwrap();
280 stream
281 .write_all(b"GET /slow HTTP/1.1\r\nhost: oyster.cafe\r\nconnection: close\r\n\r\n")
282 .await
283 .unwrap();
284
285 tokio::time::sleep(Duration::from_millis(100)).await;
286 shutdown.cancel();
287
288 let answer = tokio::time::timeout(Duration::from_secs(5), read_to_close(&mut stream))
289 .await
290 .expect("an in-flight request is answered through the graceful drain");
291 let text = String::from_utf8_lossy(&answer);
292 assert!(text.starts_with("HTTP/1.1 200"), "got: {text}");
293 assert!(
294 text.trim_end().ends_with("drained-clean"),
295 "the in-flight response must complete during drain, got: {text}"
296 );
297 }
298
299 #[tokio::test]
300 async fn cancelling_the_token_stops_the_listener() {
301 let (addr, shutdown) = bind_and_serve(limits(5_000, 4)).await;
302 let mut stream = TcpStream::connect(addr).await.unwrap();
303 stream
304 .write_all(b"GET / HTTP/1.1\r\nhost: oyster.cafe\r\nconnection: close\r\n\r\n")
305 .await
306 .unwrap();
307 let _ = read_to_close(&mut stream).await;
308
309 shutdown.cancel();
310 tokio::time::sleep(Duration::from_millis(100)).await;
311 let refused = TcpStream::connect(addr).await;
312 if let Ok(mut late) = refused {
313 late.write_all(b"GET / HTTP/1.1\r\nhost: oyster.cafe\r\nconnection: close\r\n\r\n")
314 .await
315 .ok();
316 let mut answer = Vec::new();
317 let _ =
318 tokio::time::timeout(Duration::from_secs(1), late.read_to_end(&mut answer)).await;
319 assert!(
320 !String::from_utf8_lossy(&answer).contains("200"),
321 "a drained listener mustn't answer new requests"
322 );
323 }
324 }
325}
326
327#[cfg(test)]
328mod tls_tests {
329 use super::*;
330
331 use std::num::{NonZeroU32, NonZeroU64};
332
333 use axum::routing::get;
334 use futures::StreamExt;
335 use rustls::NamedGroup;
336 use rustls::crypto::aws_lc_rs;
337 use rustls::pki_types::ServerName;
338 use tokio::io::{AsyncReadExt, AsyncWriteExt};
339 use tokio::net::TcpStream;
340 use tokio_rustls::TlsConnector;
341
342 use crate::tls;
343
344 fn client(alpn: &[&[u8]]) -> TlsConnector {
345 client_with_provider(aws_lc_rs::default_provider(), alpn)
346 }
347
348 fn client_with_provider(
349 provider: rustls::crypto::CryptoProvider,
350 alpn: &[&[u8]],
351 ) -> TlsConnector {
352 let mut config = rustls::ClientConfig::builder_with_provider(Arc::new(provider))
353 .with_safe_default_protocol_versions()
354 .unwrap()
355 .dangerous()
356 .with_custom_certificate_verifier(Arc::new(tls::test_support::AcceptAnyServerCert))
357 .with_no_client_auth();
358 config.alpn_protocols = alpn.iter().map(|p| p.to_vec()).collect();
359 TlsConnector::from(Arc::new(config))
360 }
361
362 async fn spawn_tls_server(
363 resolver: Arc<tls::ReloadableCertResolver>,
364 limits: ListenLimits,
365 ) -> (SocketAddr, CancellationToken) {
366 let server_config = Arc::new(tls::build_tls_server_config(resolver, &[]).unwrap());
367 let listener = TcpListener::bind("[::1]:0").await.unwrap();
368 let addr = listener.local_addr().unwrap();
369 let router = Router::new().route("/", get(|| async { "ok" }));
370 let shutdown = CancellationToken::new();
371 tokio::spawn(serve_tls(
372 listener,
373 router,
374 server_config,
375 limits,
376 shutdown.clone(),
377 ));
378 (addr, shutdown)
379 }
380
381 async fn serve(
382 alpn_offer: &[&[u8]],
383 ) -> (
384 SocketAddr,
385 CancellationToken,
386 tokio_rustls::client::TlsStream<TcpStream>,
387 ) {
388 let (addr, shutdown) =
389 spawn_tls_server(tls::test_support::resolver(), tls::test_support::limits()).await;
390 let tcp = TcpStream::connect(addr).await.unwrap();
391 let name = ServerName::try_from("localhost").unwrap();
392 let tls = client(alpn_offer).connect(name, tcp).await.unwrap();
393 (addr, shutdown, tls)
394 }
395
396 #[tokio::test]
397 async fn it_terminates_tls_over_http1_and_negotiates_post_quantum() {
398 let (_addr, shutdown, mut tls) = serve(&[b"http/1.1"]).await;
399
400 let group = tls.get_ref().1.negotiated_key_exchange_group().unwrap();
401 assert_eq!(
402 group.name(),
403 NamedGroup::X25519MLKEM768,
404 "prefer-post-quantum must select X25519MLKEM768 with a capable client"
405 );
406 assert_eq!(
407 tls.get_ref().1.alpn_protocol(),
408 Some(b"http/1.1".as_slice())
409 );
410
411 tls.write_all(b"GET / HTTP/1.1\r\nhost: localhost\r\nconnection: close\r\n\r\n")
412 .await
413 .unwrap();
414 let mut response = Vec::new();
415 tls.read_to_end(&mut response).await.unwrap();
416 let text = String::from_utf8_lossy(&response);
417 assert!(text.starts_with("HTTP/1.1 200"), "got: {text}");
418 assert!(text.trim_end().ends_with("ok"), "got: {text}");
419
420 shutdown.cancel();
421 }
422
423 #[tokio::test]
424 async fn it_negotiates_h2_when_the_client_offers_only_h2() {
425 let (_addr, shutdown, tls) = serve(&[b"h2"]).await;
426 assert_eq!(tls.get_ref().1.alpn_protocol(), Some(b"h2".as_slice()));
427 shutdown.cancel();
428 }
429
430 #[tokio::test]
431 async fn a_classical_only_client_completes_over_x25519() {
432 let (addr, shutdown) =
433 spawn_tls_server(tls::test_support::resolver(), tls::test_support::limits()).await;
434
435 let tcp = TcpStream::connect(addr).await.unwrap();
436 let name = ServerName::try_from("localhost").unwrap();
437 let tls =
438 client_with_provider(tls::test_support::classical_only_provider(), &[b"http/1.1"])
439 .connect(name, tcp)
440 .await
441 .unwrap();
442
443 let group = tls.get_ref().1.negotiated_key_exchange_group().unwrap();
444 assert_eq!(
445 group.name(),
446 NamedGroup::X25519,
447 "a client without ML-KEM must still complete the handshake over classical X25519"
448 );
449
450 shutdown.cancel();
451 }
452
453 #[tokio::test(flavor = "multi_thread")]
454 async fn concurrent_handshakes_survive_a_cert_reload_under_load() {
455 let high_limit = ListenLimits::new(
456 crate::limits::HeaderTimeout::from_millis(NonZeroU64::new(5_000).unwrap()),
457 crate::limits::IdleTimeout::from_millis(NonZeroU64::new(30_000).unwrap()),
458 NonZeroU32::new(512).unwrap(),
459 );
460 let resolver = tls::test_support::resolver();
461 let (addr, shutdown) = spawn_tls_server(resolver.clone(), high_limit).await;
462
463 let churn = {
464 let resolver = resolver.clone();
465 let stop = shutdown.clone();
466 tokio::spawn(async move {
467 futures::stream::unfold(0u32, move |swaps| {
468 let resolver = resolver.clone();
469 let stop = stop.clone();
470 async move {
471 match stop.is_cancelled() {
472 true => None,
473 false => {
474 resolver.store(tls::test_support::self_signed());
475 tokio::time::sleep(Duration::from_millis(1)).await;
476 Some((swaps + 1, swaps + 1))
477 }
478 }
479 }
480 })
481 .fold(0u32, |_, swaps| async move { swaps })
482 .await
483 })
484 };
485
486 let clients: Vec<_> = (0..48)
487 .map(|_| {
488 tokio::spawn(async move {
489 let tcp = TcpStream::connect(addr).await.unwrap();
490 let name = ServerName::try_from("localhost").unwrap();
491 let mut tls = client(&[b"http/1.1"]).connect(name, tcp).await.unwrap();
492 tls.write_all(
493 b"GET / HTTP/1.1\r\nhost: localhost\r\nconnection: close\r\n\r\n",
494 )
495 .await
496 .unwrap();
497 let mut response = Vec::new();
498 tls.read_to_end(&mut response).await.unwrap();
499 let text = String::from_utf8_lossy(&response).into_owned();
500 text.starts_with("HTTP/1.1 200") && text.trim_end().ends_with("ok")
501 })
502 })
503 .collect();
504
505 let outcomes: Vec<bool> = futures::future::join_all(clients)
506 .await
507 .into_iter()
508 .map(|joined| joined.unwrap())
509 .collect();
510 assert!(
511 outcomes.iter().all(|served| *served),
512 "every handshake racing a cert reload must complete and serve the router uncorrupted"
513 );
514
515 shutdown.cancel();
516 assert!(
517 churn.await.unwrap() > 1,
518 "the resolver must have reloaded the certificate during the load"
519 );
520 }
521}