Something went wrong. Try again.
A downstream fork of Tranquil PDS for use in margin.cafe
Something went wrong. Try again.
25 kB · 669 lines
Rust
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670use std::net::SocketAddr;use std::sync::Arc;use std::time::Duration;
use axum::Router;use axum::body::Body;use axum::extract::ConnectInfo;use bytes::{Buf, Bytes};use futures_util::StreamExt;use http::header::ALT_SVC;use http::{HeaderValue, Request, Response, StatusCode};use quinn::crypto::rustls::QuicServerConfig;use quinn::{Endpoint, Incoming, ServerConfig, TransportConfig, VarInt};use tokio::sync::Semaphore;use tokio_util::sync::CancellationToken;use tokio_util::task::TaskTracker;use tower::ServiceExt;use tracing::debug;
use crate::tls::{ReloadableCertResolver, TlsError};
const MAX_CONCURRENT_BIDI_STREAMS: u32 = 256;const MAX_CONCURRENT_CONNECTIONS: usize = 512;const IDLE_TIMEOUT: Duration = Duration::from_secs(30);const SHUTDOWN_GRACE: Duration = Duration::from_secs(10);const ALT_SVC_MAX_AGE_SECS: u32 = 86_400;
pub fn build_quic_server_config( resolver: Arc<ReloadableCertResolver>,) -> Result<ServerConfig, TlsError> { let provider = Arc::new(rustls::crypto::ring::default_provider()); let mut crypto = rustls::ServerConfig::builder_with_provider(provider) .with_protocol_versions(&[&rustls::version::TLS13]) .map_err(|e| TlsError::Config(e.to_string()))? .with_no_client_auth() .with_cert_resolver(resolver); crypto.alpn_protocols = vec![b"h3".to_vec()];
let quic_crypto = QuicServerConfig::try_from(crypto).map_err(|e| TlsError::Config(e.to_string()))?; let mut config = ServerConfig::with_crypto(Arc::new(quic_crypto));
let mut transport = TransportConfig::default(); transport.max_concurrent_bidi_streams(VarInt::from_u32(MAX_CONCURRENT_BIDI_STREAMS)); transport.max_idle_timeout(Some( IDLE_TIMEOUT .try_into() .expect("idle timeout fits in varint"), )); config.transport_config(Arc::new(transport)); Ok(config)}
pub fn alt_svc_header(port: u16) -> HeaderValue { HeaderValue::from_str(&format!("h3=\":{port}\"; ma={ALT_SVC_MAX_AGE_SECS}")) .expect("alt-svc header value is valid ascii")}
pub fn with_alt_svc(app: Router, port: u16) -> Router { let value = alt_svc_header(port); app.layer(axum::middleware::map_response( move |mut response: Response<Body>| { let value = value.clone(); async move { if response.status() != StatusCode::SWITCHING_PROTOCOLS { response.headers_mut().insert(ALT_SVC, value); } response } }, ))}
pub fn with_host_from_authority(app: Router) -> Router { app.layer(axum::middleware::map_request( |mut request: Request<Body>| async move { let authority = request .uri() .authority() .map(|a| HeaderValue::from_str(a.as_str())); match ( request.headers().contains_key(http::header::HOST), authority, ) { (false, Some(Ok(value))) => { request.headers_mut().insert(http::header::HOST, value); request } _ => request, } }, ))}
pub async fn serve_http3(endpoint: Endpoint, app: Router, shutdown: CancellationToken) { let tracker = TaskTracker::new(); let conn_limiter = Arc::new(Semaphore::new(MAX_CONCURRENT_CONNECTIONS)); loop { tokio::select! { _ = shutdown.cancelled() => break, incoming = endpoint.accept() => { let Some(incoming) = incoming else { break }; let Ok(permit) = conn_limiter.clone().try_acquire_owned() else { debug!( peer = %incoming.remote_address(), max = MAX_CONCURRENT_CONNECTIONS, "refusing h3 connection: limit reached" ); incoming.refuse(); continue; }; let app = app.clone(); let conn_shutdown = shutdown.clone(); let conn_tracker = tracker.clone(); tracker.spawn(async move { let _permit = permit; if let Err(e) = serve_connection(incoming, app, conn_shutdown, conn_tracker).await { debug!("h3 connection ended: {e}"); } }); } } } tracker.close(); if tokio::time::timeout(SHUTDOWN_GRACE, tracker.wait()) .await .is_err() { debug!("h3 connections did not drain within grace, closing"); } endpoint.close(0u32.into(), b"shutdown"); endpoint.wait_idle().await;}
async fn serve_connection( incoming: Incoming, app: Router, shutdown: CancellationToken, tracker: TaskTracker,) -> Result<(), Box<dyn std::error::Error + Send + Sync>> { let conn = incoming.await?; let remote = conn.remote_address(); let mut h3_conn = h3::server::Connection::<_, Bytes>::new(h3_quinn::Connection::new(conn)).await?;
let mut draining = false; loop { tokio::select! { _ = shutdown.cancelled(), if !draining => { draining = true; let _ = h3_conn.shutdown(0).await; } resolved = h3_conn.accept() => match resolved { Ok(Some(resolver)) => { let app = app.clone(); tracker.spawn(async move { if let Err(e) = serve_request(resolver, app, remote).await { debug!(peer = %remote, "h3 request failed: {e}"); } }); } Ok(None) => break, Err(e) => { debug!(peer = %remote, "h3 accept error: {e}"); break; } } } } Ok(())}
async fn serve_request( resolver: h3::server::RequestResolver<h3_quinn::Connection, Bytes>, app: Router, remote: SocketAddr,) -> Result<(), Box<dyn std::error::Error + Send + Sync>> { let (req, stream) = resolver.resolve_request().await?; let (mut send, recv) = stream.split();
let (mut parts, ()) = req.into_parts(); parts.extensions.insert(ConnectInfo(remote)); let request = Request::from_parts(parts, request_body(recv));
let response = match app.oneshot(request).await { Ok(response) => response, Err(infallible) => match infallible {}, };
let (parts, body) = response.into_parts(); send.send_response(Response::from_parts(parts, ())).await?;
let mut data = body.into_data_stream(); while let Some(chunk) = data.next().await { match chunk { Ok(bytes) if bytes.has_remaining() => send.send_data(bytes).await?, Ok(_) => {} Err(e) => { debug!(peer = %remote, "h3 response body error: {e}"); send.stop_stream(h3::error::Code::H3_INTERNAL_ERROR); return Ok(()); } } } send.finish().await?; Ok(())}
struct RecvGuard { stream: h3::server::RequestStream<h3_quinn::RecvStream, Bytes>, ended: bool,}
impl Drop for RecvGuard { fn drop(&mut self) { if !self.ended { self.stream.stop_sending(h3::error::Code::H3_NO_ERROR); } }}
fn request_body(recv: h3::server::RequestStream<h3_quinn::RecvStream, Bytes>) -> Body { let guard = RecvGuard { stream: recv, ended: false, }; let stream = futures_util::stream::unfold(Some(guard), |state| async move { let mut guard = state?; match guard.stream.recv_data().await { Ok(Some(mut buf)) => { let bytes = buf.copy_to_bytes(buf.remaining()); Some((Ok::<Bytes, std::io::Error>(bytes), Some(guard))) } Ok(None) => { guard.ended = true; None } Err(e) => { guard.ended = true; Some((Err(std::io::Error::other(e.to_string())), None)) } } }); Body::from_stream(stream)}
#[cfg(test)]mod tests { use super::*; use axum::routing::get; use rustls::pki_types::{ CertificateDer, PrivateKeyDer, PrivatePkcs8KeyDer, ServerName, UnixTime, };
fn self_signed_resolver() -> Arc<ReloadableCertResolver> { let cert = rcgen::generate_simple_self_signed(vec!["localhost".to_string()]).unwrap(); let cert_der = cert.cert.der().clone(); let key_der = PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(cert.signing_key.serialize_der())); let signing_key = rustls::crypto::ring::sign::any_supported_type(&key_der).unwrap(); let certified = rustls::sign::CertifiedKey::new(vec![cert_der], signing_key); Arc::new(ReloadableCertResolver::new(certified)) }
#[derive(Debug)] struct SkipServerVerification(Arc<rustls::crypto::CryptoProvider>);
impl rustls::client::danger::ServerCertVerifier for SkipServerVerification { fn verify_server_cert( &self, _end_entity: &CertificateDer<'_>, _intermediates: &[CertificateDer<'_>], _server_name: &ServerName<'_>, _ocsp_response: &[u8], _now: UnixTime, ) -> Result<rustls::client::danger::ServerCertVerified, rustls::Error> { Ok(rustls::client::danger::ServerCertVerified::assertion()) }
fn verify_tls12_signature( &self, message: &[u8], cert: &CertificateDer<'_>, dss: &rustls::DigitallySignedStruct, ) -> Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> { rustls::crypto::verify_tls12_signature( message, cert, dss, &self.0.signature_verification_algorithms, ) }
fn verify_tls13_signature( &self, message: &[u8], cert: &CertificateDer<'_>, dss: &rustls::DigitallySignedStruct, ) -> Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> { rustls::crypto::verify_tls13_signature( message, cert, dss, &self.0.signature_verification_algorithms, ) }
fn supported_verify_schemes(&self) -> Vec<rustls::SignatureScheme> { self.0.signature_verification_algorithms.supported_schemes() } }
fn client_endpoint() -> Endpoint { let provider = Arc::new(rustls::crypto::ring::default_provider()); let mut crypto = rustls::ClientConfig::builder_with_provider(provider.clone()) .with_protocol_versions(&[&rustls::version::TLS13]) .unwrap() .dangerous() .with_custom_certificate_verifier(Arc::new(SkipServerVerification(provider))) .with_no_client_auth(); crypto.alpn_protocols = vec![b"h3".to_vec()]; let quic = quinn::crypto::rustls::QuicClientConfig::try_from(crypto).unwrap(); let mut endpoint = Endpoint::client("0.0.0.0:0".parse().unwrap()).unwrap(); endpoint.set_default_client_config(quinn::ClientConfig::new(Arc::new(quic))); endpoint }
#[tokio::test] async fn h3_get_roundtrips_through_router() { let app = Router::new().route("/", get(|| async { "ok" })); let server_config = build_quic_server_config(self_signed_resolver()).unwrap(); let server = Endpoint::server(server_config, "127.0.0.1:0".parse().unwrap()).unwrap(); let addr = server.local_addr().unwrap(); let shutdown = CancellationToken::new(); tokio::spawn(serve_http3(server, app, shutdown.clone()));
let client = client_endpoint(); let conn = client.connect(addr, "localhost").unwrap().await.unwrap(); let (mut driver, mut send_request) = h3::client::new(h3_quinn::Connection::new(conn)) .await .unwrap(); let drive = tokio::spawn(async move { std::future::poll_fn(|cx| driver.poll_close(cx)).await });
let req = Request::get("https://localhost/").body(()).unwrap(); let mut stream = send_request.send_request(req).await.unwrap(); stream.finish().await.unwrap();
let response = stream.recv_response().await.unwrap(); assert_eq!(response.status(), 200);
let mut body = Vec::new(); while let Some(mut chunk) = stream.recv_data().await.unwrap() { let bytes = chunk.copy_to_bytes(chunk.remaining()); body.extend_from_slice(&bytes); } assert_eq!(body.as_slice(), b"ok");
shutdown.cancel(); drive.abort(); }
#[test] fn alt_svc_header_advertises_h3() { assert_eq!( alt_svc_header(443).to_str().unwrap(), "h3=\":443\"; ma=86400" ); }
#[tokio::test] async fn alt_svc_added_to_responses_except_switching_protocols() { let app = with_alt_svc( Router::new().route("/ok", get(|| async { "ok" })).route( "/upgrade", get(|| async { Response::builder() .status(StatusCode::SWITCHING_PROTOCOLS) .body(Body::empty()) .unwrap() }), ), 443, );
let normal = app .clone() .oneshot(Request::get("/ok").body(Body::empty()).unwrap()) .await .unwrap(); assert_eq!( normal.headers().get(ALT_SVC).and_then(|v| v.to_str().ok()), Some("h3=\":443\"; ma=86400") );
let upgrade = app .oneshot(Request::get("/upgrade").body(Body::empty()).unwrap()) .await .unwrap(); assert!( upgrade.headers().get(ALT_SVC).is_none(), "101 responses must not carry Alt-Svc" ); }
fn make_cert(dns: &str) -> (rustls::sign::CertifiedKey, CertificateDer<'static>) { let cert = rcgen::generate_simple_self_signed(vec![dns.to_string()]).unwrap(); let cert_der = cert.cert.der().clone(); let key_der = PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(cert.signing_key.serialize_der())); let signing_key = rustls::crypto::ring::sign::any_supported_type(&key_der).unwrap(); let certified = rustls::sign::CertifiedKey::new(vec![cert_der.clone()], signing_key); (certified, cert_der) }
#[derive(Debug)] struct RecordingVerifier { provider: Arc<rustls::crypto::CryptoProvider>, seen: Arc<std::sync::Mutex<Vec<u8>>>, }
impl rustls::client::danger::ServerCertVerifier for RecordingVerifier { fn verify_server_cert( &self, end_entity: &CertificateDer<'_>, _intermediates: &[CertificateDer<'_>], _server_name: &ServerName<'_>, _ocsp_response: &[u8], _now: UnixTime, ) -> Result<rustls::client::danger::ServerCertVerified, rustls::Error> { *self.seen.lock().unwrap() = end_entity.as_ref().to_vec(); Ok(rustls::client::danger::ServerCertVerified::assertion()) }
fn verify_tls12_signature( &self, message: &[u8], cert: &CertificateDer<'_>, dss: &rustls::DigitallySignedStruct, ) -> Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> { rustls::crypto::verify_tls12_signature( message, cert, dss, &self.provider.signature_verification_algorithms, ) }
fn verify_tls13_signature( &self, message: &[u8], cert: &CertificateDer<'_>, dss: &rustls::DigitallySignedStruct, ) -> Result<rustls::client::danger::HandshakeSignatureValid, rustls::Error> { rustls::crypto::verify_tls13_signature( message, cert, dss, &self.provider.signature_verification_algorithms, ) }
fn supported_verify_schemes(&self) -> Vec<rustls::SignatureScheme> { self.provider .signature_verification_algorithms .supported_schemes() } }
async fn observe_server_cert(addr: SocketAddr) -> Vec<u8> { let provider = Arc::new(rustls::crypto::ring::default_provider()); let seen = Arc::new(std::sync::Mutex::new(Vec::new())); let verifier = Arc::new(RecordingVerifier { provider: provider.clone(), seen: seen.clone(), }); let mut crypto = rustls::ClientConfig::builder_with_provider(provider) .with_protocol_versions(&[&rustls::version::TLS13]) .unwrap() .dangerous() .with_custom_certificate_verifier(verifier) .with_no_client_auth(); crypto.alpn_protocols = vec![b"h3".to_vec()]; let quic = quinn::crypto::rustls::QuicClientConfig::try_from(crypto).unwrap(); let mut endpoint = Endpoint::client("127.0.0.1:0".parse().unwrap()).unwrap(); endpoint.set_default_client_config(quinn::ClientConfig::new(Arc::new(quic))); let conn = endpoint.connect(addr, "localhost").unwrap().await.unwrap(); conn.close(0u32.into(), b"done"); endpoint.wait_idle().await; seen.lock().unwrap().clone() }
#[tokio::test] async fn quic_handshake_observes_reloaded_certificate() { let (cert_a, der_a) = make_cert("localhost"); let (cert_b, der_b) = make_cert("localhost"); assert_ne!(der_a, der_b, "test must use two distinct certs");
let resolver = Arc::new(ReloadableCertResolver::new(cert_a)); let server_config = build_quic_server_config(resolver.clone()).unwrap(); let server = Endpoint::server(server_config, "127.0.0.1:0".parse().unwrap()).unwrap(); let addr = server.local_addr().unwrap(); let shutdown = CancellationToken::new(); let app = Router::new().route("/", get(|| async { "ok" })); tokio::spawn(serve_http3(server, app, shutdown.clone()));
let before = observe_server_cert(addr).await; assert_eq!( before.as_slice(), der_a.as_ref(), "first handshake must present the original cert" );
resolver.store(cert_b);
let after = observe_server_cert(addr).await; assert_eq!( after.as_slice(), der_b.as_ref(), "handshake after reload must present the new cert" ); assert_ne!(before, after, "reload must change the presented cert");
shutdown.cancel(); }
#[tokio::test] async fn h3_requests_carry_remote_addr_connect_info() { let app = Router::new().route( "/", get(|ConnectInfo(addr): ConnectInfo<SocketAddr>| async move { addr.to_string() }), ); let server_config = build_quic_server_config(self_signed_resolver()).unwrap(); let server = Endpoint::server(server_config, "127.0.0.1:0".parse().unwrap()).unwrap(); let addr = server.local_addr().unwrap(); let shutdown = CancellationToken::new(); tokio::spawn(serve_http3(server, app, shutdown.clone()));
let client = client_endpoint(); let client_port = client.local_addr().unwrap().port(); let conn = client.connect(addr, "localhost").unwrap().await.unwrap(); let (mut driver, mut send_request) = h3::client::new(h3_quinn::Connection::new(conn)) .await .unwrap(); let drive = tokio::spawn(async move { std::future::poll_fn(|cx| driver.poll_close(cx)).await });
let req = Request::get("https://localhost/").body(()).unwrap(); let mut stream = send_request.send_request(req).await.unwrap(); stream.finish().await.unwrap(); assert_eq!(stream.recv_response().await.unwrap().status(), 200);
let mut body = Vec::new(); while let Some(mut chunk) = stream.recv_data().await.unwrap() { let bytes = chunk.copy_to_bytes(chunk.remaining()); body.extend_from_slice(&bytes); } let reported: SocketAddr = String::from_utf8(body).unwrap().parse().unwrap(); assert!(reported.ip().is_loopback()); assert_eq!( reported.port(), client_port, "handlers must see the QUIC remote address via ConnectInfo" );
shutdown.cancel(); drive.abort(); }
#[tokio::test] async fn host_header_filled_from_authority() { let app = with_host_from_authority(Router::new().route( "/", get(|headers: http::HeaderMap| async move { headers .get(http::header::HOST) .and_then(|v| v.to_str().ok()) .map(str::to_owned) .unwrap_or_default() }), )); let server_config = build_quic_server_config(self_signed_resolver()).unwrap(); let server = Endpoint::server(server_config, "127.0.0.1:0".parse().unwrap()).unwrap(); let addr = server.local_addr().unwrap(); let shutdown = CancellationToken::new(); tokio::spawn(serve_http3(server, app, shutdown.clone()));
let client = client_endpoint(); let conn = client.connect(addr, "localhost").unwrap().await.unwrap(); let (mut driver, mut send_request) = h3::client::new(h3_quinn::Connection::new(conn)) .await .unwrap(); let drive = tokio::spawn(async move { std::future::poll_fn(|cx| driver.poll_close(cx)).await });
let req = Request::get("https://localhost/").body(()).unwrap(); let mut stream = send_request.send_request(req).await.unwrap(); stream.finish().await.unwrap(); assert_eq!(stream.recv_response().await.unwrap().status(), 200);
let mut body = Vec::new(); while let Some(mut chunk) = stream.recv_data().await.unwrap() { let bytes = chunk.copy_to_bytes(chunk.remaining()); body.extend_from_slice(&bytes); } assert_eq!( String::from_utf8(body).unwrap(), "localhost", "handlers must see the authority as the Host header" );
shutdown.cancel(); drive.abort(); }
#[tokio::test] async fn h3_body_error_resets_stream_instead_of_truncating() { let app = Router::new().route( "/", get(|| async { Body::from_stream(futures_util::stream::iter(vec![ Ok::<Bytes, std::io::Error>(Bytes::from_static(b"partial")), Err(std::io::Error::other("body source failed")), ])) }), ); let server_config = build_quic_server_config(self_signed_resolver()).unwrap(); let server = Endpoint::server(server_config, "127.0.0.1:0".parse().unwrap()).unwrap(); let addr = server.local_addr().unwrap(); let shutdown = CancellationToken::new(); tokio::spawn(serve_http3(server, app, shutdown.clone()));
let client = client_endpoint(); let conn = client.connect(addr, "localhost").unwrap().await.unwrap(); let (mut driver, mut send_request) = h3::client::new(h3_quinn::Connection::new(conn)) .await .unwrap(); let drive = tokio::spawn(async move { std::future::poll_fn(|cx| driver.poll_close(cx)).await });
let req = Request::get("https://localhost/").body(()).unwrap(); let mut stream = send_request.send_request(req).await.unwrap(); stream.finish().await.unwrap();
let outcome = async { stream.recv_response().await?; let mut body = Vec::new(); loop { match stream.recv_data().await { Ok(Some(mut chunk)) => { let bytes = chunk.copy_to_bytes(chunk.remaining()); body.extend_from_slice(&bytes); } Ok(None) => return Ok(body), Err(e) => return Err(e), } } } .await; assert!( outcome.is_err(), "a mid-body error must reset the stream, not end the body cleanly after {} bytes", outcome.map(|b| b.len()).unwrap_or(0) );
shutdown.cancel(); drive.abort(); }}