diff --git a/Cargo.lock b/Cargo.lock index f782807..8be9ebf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1769,12 +1769,13 @@ version = "2.0.0" dependencies = [ "anyhow", "axum", - "axum-server", "bytes", "clap", "governor", + "h2", "http", "httpdate", + "hyper-util", "mainline", "pkarr", "reqwest", @@ -2868,6 +2869,7 @@ dependencies = [ "http-body", "percent-encoding", "pin-project-lite", + "tokio", "tower-layer", "tower-service", "tracing", diff --git a/relay/Cargo.toml b/relay/Cargo.toml index 1932bcb..b927b9c 100644 --- a/relay/Cargo.toml +++ b/relay/Cargo.toml @@ -25,14 +25,16 @@ mainline = { workspace = true, optional = true } tokio = { version = "1", features = [ "fs", "macros", + "net", "rt-multi-thread", "signal", + "sync", "time", ] } -tower-http = { version = "0.7", features = ["cors", "trace"] } +hyper-util = { version = "0.1", features = ["server-auto", "service", "tokio"] } +tower-http = { version = "0.7", features = ["cors", "timeout", "trace"] } tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter"] } -axum-server = "0.8" http = "1" bytes = "1" tower_governor = { version = "0.8", default-features = false, features = ["axum"] } @@ -48,6 +50,7 @@ pkarr = { path = "../pkarr", version = "8", default-features = false, features = ] } [dev-dependencies] +h2 = "0.4" mainline = { workspace = true } reqwest = { workspace = true } diff --git a/relay/src/config.rs b/relay/src/config.rs index 184bdef..1a95dbf 100644 --- a/relay/src/config.rs +++ b/relay/src/config.rs @@ -7,11 +7,14 @@ use std::{ net::Ipv4Addr, path::{Path, PathBuf}, }; +use tokio::sync::Semaphore; use crate::rate_limiting::RateLimiterConfig; pub const DEFAULT_CACHE_SIZE: usize = 1_000_000; pub const DEFAULT_HTTP_PORT: u16 = 6881; +pub const DEFAULT_HTTP_MAX_CONNECTIONS: usize = 1_024; +pub const DEFAULT_HTTP_MAX_CONNECTION_AGE_SECONDS: u64 = 5 * 60; pub const CACHE_DIR: &str = "pkarr-cache"; #[derive(Serialize, Deserialize, Debug)] @@ -40,12 +43,18 @@ impl Default for RelayConfig { pub struct HttpConfig { #[serde(default = "default_http_port")] pub port: u16, + #[serde(default = "default_http_max_connections")] + pub max_connections: usize, + #[serde(default = "default_http_max_connection_age_seconds")] + pub max_connection_age_seconds: u64, } impl Default for HttpConfig { fn default() -> Self { Self { port: default_http_port(), + max_connections: default_http_max_connections(), + max_connection_age_seconds: default_http_max_connection_age_seconds(), } } } @@ -91,6 +100,7 @@ impl RelayConfig { .path .map(|cache_path| resolve_cache_path(cache_path, path.as_ref())) .transpose()?; + validate_http(&config.http)?; validate_cache_ttl(&config.cache)?; Ok(config) @@ -101,6 +111,14 @@ fn default_http_port() -> u16 { DEFAULT_HTTP_PORT } +fn default_http_max_connections() -> usize { + DEFAULT_HTTP_MAX_CONNECTIONS +} + +fn default_http_max_connection_age_seconds() -> u64 { + DEFAULT_HTTP_MAX_CONNECTION_AGE_SECONDS +} + fn default_cache_size() -> usize { DEFAULT_CACHE_SIZE } @@ -113,6 +131,24 @@ fn default_maximum_ttl() -> u32 { DEFAULT_MAXIMUM_TTL } +fn validate_http(http: &HttpConfig) -> Result<()> { + anyhow::ensure!( + http.max_connections > 0, + "http.max_connections must be greater than zero" + ); + anyhow::ensure!( + http.max_connections <= Semaphore::MAX_PERMITS, + "http.max_connections must be less than or equal to {}", + Semaphore::MAX_PERMITS + ); + anyhow::ensure!( + http.max_connection_age_seconds > 0, + "http.max_connection_age_seconds must be greater than zero" + ); + + Ok(()) +} + fn validate_cache_ttl(cache: &CacheConfig) -> Result<()> { anyhow::ensure!( cache.minimum_ttl <= cache.maximum_ttl, @@ -142,13 +178,64 @@ fn resolve_cache_path(cache_path: PathBuf, config_file_path: &Path) -> Result Self { + Self::from(&HttpConfig::default()) + } +} + +impl From<&HttpConfig> for Limits { + fn from(config: &HttpConfig) -> Self { + Self { + max_connections: config.max_connections, + initial_request_header_timeout: INITIAL_REQUEST_HEADER_TIMEOUT, + max_connection_age: Duration::from_secs(config.max_connection_age_seconds), + drain_timeout: CONNECTION_DRAIN_TIMEOUT, + http1_header_read_timeout: HTTP1_HEADER_READ_TIMEOUT, + } + } +} + +pub(crate) struct HttpServer { + shutdown: watch::Sender<()>, +} + +impl HttpServer { + pub(crate) fn spawn( + listener: TcpListener, + app: Router, + config: &HttpConfig, + ) -> io::Result { + Self::spawn_with_limits(listener, app, Limits::from(config)) + } + + fn spawn_with_limits(listener: TcpListener, app: Router, limits: Limits) -> io::Result { + let listener = tokio::net::TcpListener::from_std(listener)?; + let (shutdown, shutdown_receiver) = watch::channel(()); + + tokio::spawn(serve(listener, app, limits, shutdown_receiver)); + + Ok(Self { shutdown }) + } + + pub(crate) fn shutdown(&self) { + let _ = self.shutdown.send(()); + } +} + +async fn serve( + listener: tokio::net::TcpListener, + app: Router, + limits: Limits, + mut shutdown: watch::Receiver<()>, +) { + let connection_slots = Arc::new(Semaphore::new(limits.max_connections)); + let mut builder = auto::Builder::new(TokioExecutor::new()); + builder + .http1() + .timer(TokioTimer::new()) + .header_read_timeout(limits.http1_header_read_timeout); + builder + .http2() + .max_concurrent_streams(HTTP2_MAX_CONCURRENT_STREAMS) + .max_header_list_size(HTTP2_MAX_HEADER_LIST_SIZE); + let builder = Arc::new(builder); + + loop { + let accepted = tokio::select! { + biased; + _ = shutdown.changed() => return, + accepted = listener.accept() => accepted, + }; + + let (stream, peer_address) = match accepted { + Ok(connection) => connection, + Err(error) => { + tracing::warn!(%error, "failed to accept HTTP connection; retrying"); + tokio::select! { + biased; + _ = shutdown.changed() => return, + _ = sleep(ACCEPT_ERROR_BACKOFF) => {} + } + continue; + } + }; + + let Ok(connection_slot) = Arc::clone(&connection_slots).try_acquire_owned() else { + // Refuse excess connections immediately instead of allocating another task. + tracing::debug!( + %peer_address, + max_connections = limits.max_connections, + "HTTP connection refused because the connection limit was reached" + ); + drop(stream); + continue; + }; + + tokio::spawn(serve_connection( + Arc::clone(&builder), + app.clone(), + stream, + peer_address, + limits, + shutdown.clone(), + connection_slot, + )); + } +} + +async fn serve_connection( + builder: Arc>, + app: Router, + stream: TcpStream, + peer_address: SocketAddr, + limits: Limits, + mut shutdown: watch::Receiver<()>, + _connection_slot: OwnedSemaphorePermit, +) { + let first_request_headers_received = Arc::new(Notify::new()); + let request_headers_notification = Arc::clone(&first_request_headers_received); + // This middleware runs only after Hyper has parsed a complete request + // header block, keeping the deadline armed through negotiation and parsing. + let app = app + .layer(Extension(ConnectInfo(peer_address))) + .layer(map_request(move |request: Request| { + request_headers_notification.notify_one(); + std::future::ready(request) + })); + let service = TowerToHyperService::new(app); + let connection = builder.serve_connection(TokioIo::new(stream), service); + let maximum_age = sleep(limits.max_connection_age); + let first_request_headers = first_request_headers_received.notified(); + let initial_request_header_deadline = sleep(limits.initial_request_header_timeout); + tokio::pin!( + connection, + maximum_age, + first_request_headers, + initial_request_header_deadline + ); + let mut awaiting_first_request_headers = true; + + loop { + tokio::select! { + biased; + _ = shutdown.changed() => return, + result = connection.as_mut() => { + if let Err(error) = result { + tracing::debug!(%error, %peer_address, "HTTP connection closed with an error"); + } + return; + } + _ = maximum_age.as_mut() => break, + _ = first_request_headers.as_mut(), if awaiting_first_request_headers => { + awaiting_first_request_headers = false; + } + _ = initial_request_header_deadline.as_mut(), if awaiting_first_request_headers => { + tracing::debug!(%peer_address, "HTTP connection did not produce initial request headers in time"); + return; + } + } + } + + // For HTTP/2 this sends GOAWAY and refuses new streams. HTTP/1 stops accepting + // new requests on this connection. Already accepted requests may finish. + connection.as_mut().graceful_shutdown(); + + tokio::select! { + biased; + _ = shutdown.changed() => { + tracing::debug!(%peer_address, "HTTP connection drain interrupted by server shutdown"); + } + result = connection.as_mut() => { + if let Err(error) = result { + tracing::debug!(%error, %peer_address, "HTTP connection closed with an error while draining"); + } + } + _ = sleep(limits.drain_timeout) => { + tracing::debug!(%peer_address, "HTTP connection exceeded its drain timeout"); + } + } +} + +#[cfg(test)] +mod tests; diff --git a/relay/src/http_server/tests.rs b/relay/src/http_server/tests.rs new file mode 100644 index 0000000..c27abab --- /dev/null +++ b/relay/src/http_server/tests.rs @@ -0,0 +1,354 @@ +use std::{ + io, + net::{Ipv4Addr, SocketAddr, TcpListener}, + time::Duration, +}; + +use axum::Router; +use h2::Ping; +use http::StatusCode; +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::TcpStream, + time::{sleep, timeout}, +}; + +use super::{HttpServer, Limits}; + +#[tokio::test] +async fn connection_without_request_hits_initial_request_header_timeout() { + let server = TestServer::spawn(Limits { + initial_request_header_timeout: Duration::from_millis(200), + max_connection_age: Duration::from_secs(1), + ..test_limits() + }); + let mut stream = TcpStream::connect(server.address).await.unwrap(); + + timeout(Duration::from_secs(1), wait_for_disconnect(&mut stream)) + .await + .expect("connection should close at its initial request header deadline") + .unwrap(); +} + +#[tokio::test] +async fn partial_protocol_preface_hits_initial_request_header_timeout() { + let server = TestServer::spawn(Limits { + initial_request_header_timeout: Duration::from_millis(200), + max_connection_age: Duration::from_secs(1), + ..test_limits() + }); + let mut stream = TcpStream::connect(server.address).await.unwrap(); + stream.write_all(b"P").await.unwrap(); + + timeout(Duration::from_secs(1), wait_for_disconnect(&mut stream)) + .await + .expect("partial protocol detection should reach the initial request header deadline") + .unwrap(); +} + +#[tokio::test] +async fn incomplete_http2_handshake_hits_initial_request_header_timeout() { + let server = TestServer::spawn(Limits { + initial_request_header_timeout: Duration::from_millis(200), + max_connection_age: Duration::from_secs(1), + ..test_limits() + }); + let mut stream = TcpStream::connect(server.address).await.unwrap(); + stream + .write_all(b"PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n") + .await + .unwrap(); + + timeout(Duration::from_secs(1), wait_for_disconnect(&mut stream)) + .await + .expect("an incomplete HTTP/2 handshake should reach the initial request header deadline") + .unwrap(); +} + +#[tokio::test] +async fn complete_request_headers_disarm_initial_request_header_timeout() { + let server = TestServer::spawn(Limits { + initial_request_header_timeout: Duration::from_millis(200), + max_connection_age: Duration::from_secs(1), + ..test_limits() + }); + let stream = TcpStream::connect(server.address).await.unwrap(); + let (mut client, connection) = h2::client::handshake(stream).await.unwrap(); + let connection_task = tokio::spawn(connection); + + let request = http::Request::get(format!("http://{}/", server.address)) + .body(()) + .unwrap(); + let (response, _) = client.send_request(request, true).unwrap(); + assert_eq!(response.await.unwrap().status(), StatusCode::NOT_FOUND); + + sleep(Duration::from_millis(300)).await; + let request = http::Request::get(format!("http://{}/", server.address)) + .body(()) + .unwrap(); + let (response, _) = client.send_request(request, true).unwrap(); + assert_eq!(response.await.unwrap().status(), StatusCode::NOT_FOUND); + + server.server.shutdown(); + timeout(Duration::from_secs(1), connection_task) + .await + .expect("connection should close on shutdown") + .unwrap() + .unwrap(); +} + +#[tokio::test] +async fn idle_connection_is_closed_at_maximum_age() { + let server = TestServer::spawn(Limits { + max_connection_age: Duration::from_millis(200), + ..test_limits() + }); + let mut stream = TcpStream::connect(server.address).await.unwrap(); + + timeout(Duration::from_secs(1), wait_for_disconnect(&mut stream)) + .await + .expect("idle connection should reach its maximum age") + .unwrap(); +} + +#[tokio::test] +async fn incomplete_http1_headers_hit_header_timeout() { + let server = TestServer::spawn(Limits { + max_connection_age: Duration::from_secs(1), + http1_header_read_timeout: Duration::from_millis(200), + ..test_limits() + }); + let mut stream = TcpStream::connect(server.address).await.unwrap(); + stream.write_all(b"GET / HTTP/1.1\r\nHost:").await.unwrap(); + + timeout(Duration::from_secs(1), wait_for_disconnect(&mut stream)) + .await + .expect("incomplete headers should time out") + .unwrap(); +} + +#[tokio::test] +async fn active_http2_stream_finishes_during_grace_period() { + let app = Router::new().route( + "/", + axum::routing::get(|| async { + sleep(Duration::from_millis(300)).await; + StatusCode::NO_CONTENT + }), + ); + let server = TestServer::spawn_with_app( + Limits { + max_connection_age: Duration::from_millis(200), + drain_timeout: Duration::from_millis(500), + ..test_limits() + }, + app, + ); + let stream = TcpStream::connect(server.address).await.unwrap(); + let (mut client, connection) = h2::client::handshake(stream).await.unwrap(); + let connection_task = tokio::spawn(connection); + let request = http::Request::get(format!("http://{}/", server.address)) + .body(()) + .unwrap(); + let (response, _) = client.send_request(request, true).unwrap(); + + let response = timeout(Duration::from_secs(1), response) + .await + .expect("request should finish during the grace period") + .unwrap(); + assert_eq!(response.status(), StatusCode::NO_CONTENT); + + timeout(Duration::from_secs(1), connection_task) + .await + .expect("connection should close after draining") + .unwrap() + .unwrap(); +} + +#[tokio::test] +async fn http2_stream_exceeding_grace_period_is_interrupted() { + let app = Router::new().route( + "/", + axum::routing::get(|| async { + sleep(Duration::from_secs(1)).await; + StatusCode::NO_CONTENT + }), + ); + let server = TestServer::spawn_with_app( + Limits { + max_connection_age: Duration::from_millis(200), + drain_timeout: Duration::from_millis(200), + ..test_limits() + }, + app, + ); + let stream = TcpStream::connect(server.address).await.unwrap(); + let (mut client, connection) = h2::client::handshake(stream).await.unwrap(); + let connection_task = tokio::spawn(connection); + let request = http::Request::get(format!("http://{}/", server.address)) + .body(()) + .unwrap(); + let (response, _) = client.send_request(request, true).unwrap(); + + assert!(timeout(Duration::from_secs(1), response) + .await + .expect("connection should close after the grace period") + .is_err()); + + timeout(Duration::from_secs(1), connection_task) + .await + .expect("client connection task should finish") + .unwrap() + .unwrap(); +} + +#[tokio::test] +async fn http2_ping_traffic_does_not_extend_maximum_age() { + let server = TestServer::spawn(Limits { + max_connection_age: Duration::from_millis(200), + ..test_limits() + }); + let stream = TcpStream::connect(server.address).await.unwrap(); + let (_client, mut connection) = h2::client::handshake(stream).await.unwrap(); + let mut ping_pong = connection.ping_pong().unwrap(); + let connection_task = tokio::spawn(connection); + let ping_task = tokio::spawn(async move { + loop { + sleep(Duration::from_millis(50)).await; + if ping_pong.ping(Ping::opaque()).await.is_err() { + break; + } + } + }); + + timeout(Duration::from_secs(1), connection_task) + .await + .expect("PING traffic should not prevent graceful shutdown") + .unwrap() + .unwrap(); + + ping_task.abort(); +} + +#[tokio::test] +async fn excess_connection_is_refused() { + let server = TestServer::spawn(Limits { + max_connections: 1, + max_connection_age: Duration::from_secs(1), + ..test_limits() + }); + let _first = TcpStream::connect(server.address).await.unwrap(); + sleep(Duration::from_millis(100)).await; + let mut excess = TcpStream::connect(server.address).await.unwrap(); + + timeout(Duration::from_secs(1), wait_for_disconnect(&mut excess)) + .await + .expect("connection above the limit should be refused") + .unwrap(); +} + +#[tokio::test] +async fn shutdown_closes_existing_connections() { + let server = TestServer::spawn(test_limits()); + let mut stream = TcpStream::connect(server.address).await.unwrap(); + + server.server.shutdown(); + + timeout(Duration::from_secs(1), wait_for_disconnect(&mut stream)) + .await + .expect("server shutdown should close existing connections") + .unwrap(); +} + +#[tokio::test] +async fn shutdown_interrupts_connection_drain() { + let app = Router::new().route( + "/", + axum::routing::get(|| async { + sleep(Duration::from_secs(5)).await; + StatusCode::NO_CONTENT + }), + ); + let server = TestServer::spawn_with_app( + Limits { + max_connection_age: Duration::from_millis(200), + drain_timeout: Duration::from_secs(2), + ..test_limits() + }, + app, + ); + let stream = TcpStream::connect(server.address).await.unwrap(); + let (mut client, connection) = h2::client::handshake(stream).await.unwrap(); + let connection_task = tokio::spawn(connection); + let request = http::Request::get(format!("http://{}/", server.address)) + .body(()) + .unwrap(); + let (_response, _) = client.send_request(request, true).unwrap(); + + sleep(Duration::from_millis(300)).await; + server.server.shutdown(); + + timeout(Duration::from_millis(500), connection_task) + .await + .expect("shutdown should interrupt the connection drain") + .unwrap() + .unwrap(); +} + +fn test_limits() -> Limits { + Limits { + max_connections: 16, + initial_request_header_timeout: Duration::from_secs(1), + max_connection_age: Duration::from_secs(1), + drain_timeout: Duration::from_millis(100), + http1_header_read_timeout: Duration::from_secs(1), + } +} + +struct TestServer { + address: SocketAddr, + server: HttpServer, +} + +impl TestServer { + fn spawn(limits: Limits) -> Self { + Self::spawn_with_app(limits, Router::new()) + } + + fn spawn_with_app(limits: Limits, app: Router) -> Self { + let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).unwrap(); + listener.set_nonblocking(true).unwrap(); + let address = listener.local_addr().unwrap(); + let server = HttpServer::spawn_with_limits(listener, app, limits).unwrap(); + + Self { address, server } + } +} + +impl Drop for TestServer { + fn drop(&mut self) { + self.server.shutdown(); + } +} + +async fn wait_for_disconnect(stream: &mut TcpStream) -> io::Result<()> { + let mut buffer = [0; 1024]; + + loop { + match stream.read(&mut buffer).await { + Ok(0) => return Ok(()), + Ok(_) => {} + Err(error) + if matches!( + error.kind(), + io::ErrorKind::ConnectionAborted + | io::ErrorKind::ConnectionReset + | io::ErrorKind::NotConnected + ) => + { + return Ok(()); + } + Err(error) => return Err(error), + } + } +} diff --git a/relay/src/lib.rs b/relay/src/lib.rs index be52b08..46008b7 100644 --- a/relay/src/lib.rs +++ b/relay/src/lib.rs @@ -12,6 +12,7 @@ mod dht_service; mod error; mod extractors; mod handlers; +mod http_server; mod index; mod quota_config; mod rate_limiting; @@ -27,9 +28,8 @@ use std::{ use anyhow::anyhow; use axum::{extract::DefaultBodyLimit, http::HeaderName, Router}; -use axum_server::Handle; -use tower_http::{cors::CorsLayer, trace::TraceLayer}; +use tower_http::{cors::CorsLayer, timeout::RequestBodyDeadlineLayer, trace::TraceLayer}; use tracing::info; use pkarr::{ @@ -41,10 +41,13 @@ use url::Url; use config::{RelayConfig, CACHE_DIR}; use dht_service::DhtService; +use http_server::HttpServer; pub use quota_config::{RequestCountQuota, TimeUnit}; pub use rate_limiting::RateLimiterConfig; +const HTTP_REQUEST_BODY_TIMEOUT: Duration = Duration::from_secs(30); + /// A builder for Pkarr [Relay] pub struct RelayBuilder { config: RelayConfig, @@ -140,7 +143,7 @@ impl RelayBuilder { /// This struct represents a running relay server and provides methods to interact with it, /// such as retrieving the server's address or shutting it down. pub struct Relay { - handle: Handle, + http_server: HttpServer, relay_address: SocketAddr, } @@ -186,8 +189,7 @@ impl Relay { let dht_client = DhtClient::build(dht_config)?; let listener = TcpListener::bind(SocketAddr::from(([0, 0, 0, 0], config.http.port)))?; - // On axum-server 0.8.0 the `.set_nonblocking(true)` call does not take place internally anymore - // See open issue https://github.com/programatik29/axum-server/issues/181 + // Tokio requires a non-blocking listener when converting from std. listener.set_nonblocking(true)?; let node_address = dht_client.info().await.local_addr(); @@ -207,17 +209,10 @@ impl Relay { ); let state = AppState { dht }; let app = create_app(state, rate_limiters.http, rate_limiters.behind_proxy); - - let handle = Handle::new(); - - let task = axum_server::from_tcp(listener)? - .handle(handle.clone()) - .serve(app.into_make_service_with_connect_info::()); - - tokio::spawn(task); + let http_server = HttpServer::spawn(listener, app, &config.http)?; Ok(Relay { - handle, + http_server, relay_address, }) } @@ -252,7 +247,10 @@ impl Relay { /// because the possible Undefined Behavior (UB) if the lock file is broken. pub async fn run_test(bootstrap: &[T]) -> anyhow::Result { let config = RelayConfig { - http: config::HttpConfig { port: 0 }, + http: config::HttpConfig { + port: 0, + ..Default::default() + }, rate_limiter: None, ..Default::default() }; @@ -281,7 +279,10 @@ impl Relay { } let config = RelayConfig { - http: config::HttpConfig { port: 15411 }, + http: config::HttpConfig { + port: 15411, + ..Default::default() + }, rate_limiter: None, ..Default::default() }; @@ -308,7 +309,7 @@ impl Relay { /// Shutdown the relay server. pub fn shutdown(&self) { - self.handle.shutdown(); + self.http_server.shutdown(); } } @@ -369,6 +370,7 @@ fn create_app( .route("/", axum::routing::get(crate::handlers::index)) .with_state(state) .layer(DefaultBodyLimit::max(1104)) + .layer(RequestBodyDeadlineLayer::new(HTTP_REQUEST_BODY_TIMEOUT)) .layer(cors_layer()) .layer(TraceLayer::new_for_http());