diff options
| author | Nikolay Govorov <me@govorov.online> | 2026-01-17 09:08:08 +0000 |
|---|---|---|
| committer | Nikolay Govorov <me@govorov.online> | 2026-01-17 09:08:08 +0000 |
| commit | 90dc56d85410d9ee6a51c10df5d496600d993351 (patch) | |
| tree | 2ddc93b00ec76228d40462184c7d4150c265efd1 | |
| parent | 6a9491a38e9895a0be5e136ca8d1f9481f0d95c3 (diff) | |
| download | tar tar.gz tar.bz2 tar.lz tar.xz tar.zst zip | |
Graceful shutdown and rate limits for server
Diffstat
| -rw-r--r-- | Cargo.lock | 253 | +245 −8 |
| -rw-r--r-- | Cargo.toml | 7 | +6 −1 |
| -rw-r--r-- | pkg/zorian.toml | 2 | +1 −1 |
| -rw-r--r-- | src/backends/zig.rs | 1 | +1 −0 |
| -rw-r--r-- | src/config.rs | 26 | +24 −2 |
| -rw-r--r-- | src/main.rs | 138 | +131 −7 |
6 files changed, 408 insertions, 19 deletions
diff --git a/Cargo.lock b/Cargo.lock index a0b66b6..f0a534c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -125,6 +125,17 @@ dependencies = [ ] [[package]] +name = "async-trait" +version = "0.1.89" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9035ad2d096bed7955a320ee7e2230574d28fd3c3a0f186cbea1ff3c7eed5dbb" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] name = "atoi" version = "2.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -759,6 +770,20 @@ dependencies = [ ] [[package]] +name = "dashmap" +version = "6.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5041cc499144891f3790297212f32a74fb938e5136a14943f338ef9e0ae276cf" +dependencies = [ + "cfg-if", + "crossbeam-utils", + "hashbrown 0.14.5", + "lock_api", + "once_cell", + "parking_lot_core", +] + +[[package]] name = "der" version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -994,6 +1019,16 @@ dependencies = [ ] [[package]] +name = "forwarded-header-value" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8835f84f38484cc86f110a805655697908257fb9a7af005234060891557198e9" +dependencies = [ + "nonempty", + "thiserror 1.0.69", +] + +[[package]] name = "fs-err" version = "3.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -1065,6 +1100,12 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f90f7dce0722e95104fcb095585910c0977252f286e354b5e3bd38902cd99988" [[package]] +name = "futures-timer" +version = "3.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f288b0a4f20f9a56b5d1da57e2227c661b7b16168e2f72365f57b63326e29b24" + +[[package]] name = "futures-util" version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -1108,9 +1149,11 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" dependencies = [ "cfg-if", + "js-sys", "libc", "r-efi", "wasip2", + "wasm-bindgen", ] [[package]] @@ -1904,6 +1947,29 @@ dependencies = [ ] [[package]] +name = "governor" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9efcab3c1958580ff1f25a2a41be1668f7603d849bb63af523b208a3cc1223b8" +dependencies = [ + "cfg-if", + "dashmap", + "futures-sink", + "futures-timer", + "futures-util", + "getrandom 0.3.4", + "hashbrown 0.16.1", + "nonzero_ext", + "parking_lot", + "portable-atomic", + "quanta", + "rand 0.9.2", + "smallvec", + "spinning_top", + "web-time", +] + +[[package]] name = "h2" version = "0.4.13" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -1933,6 +1999,12 @@ dependencies = [ [[package]] name = "hashbrown" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" + +[[package]] +name = "hashbrown" version = "0.15.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" @@ -2136,6 +2208,19 @@ dependencies = [ ] [[package]] +name = "hyper-timeout" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b90d566bffbce6a75bd8b09a05aa8c2cb1fabb6cb348f8840c9e4c90a0d83b0" +dependencies = [ + "hyper", + "hyper-util", + "pin-project-lite", + "tokio", + "tower-service", +] + +[[package]] name = "hyper-tls" version = "0.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -2687,6 +2772,18 @@ dependencies = [ ] [[package]] +name = "nonempty" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e9e591e719385e6ebaeb5ce5d3887f7d5676fceca6411d1925ccc95745f3d6f7" + +[[package]] +name = "nonzero_ext" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38bf9645c8b145698bb0b18a4637dcacbc421ea49bef2317e4fd8065a387cf21" + +[[package]] name = "normpath" version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -2715,7 +2812,7 @@ dependencies = [ "num-integer", "num-iter", "num-traits", - "rand", + "rand 0.8.5", "smallvec", "zeroize", ] @@ -2895,6 +2992,26 @@ dependencies = [ ] [[package]] +name = "pin-project" +version = "1.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "677f1add503faace112b9f1373e43e9e054bfdd22ff1a63c1bc485eaec6a6a8a" +dependencies = [ + "pin-project-internal", +] + +[[package]] +name = "pin-project-internal" +version = "1.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e918e4ff8c4549eb882f14b3a4bc8c8bc93de829416eacf579f1207a8fbf861" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] name = "pin-project-lite" version = "0.2.16" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -3029,6 +3146,21 @@ dependencies = [ ] [[package]] +name = "quanta" +version = "0.12.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3ab5a9d756f0d97bdc89019bd2e4ea098cf9cde50ee7564dde6b81ccc8f06c7" +dependencies = [ + "crossbeam-utils", + "libc", + "once_cell", + "raw-cpuid", + "wasi", + "web-sys", + "winapi", +] + +[[package]] name = "quick-xml" version = "0.37.5" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -3065,8 +3197,18 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" dependencies = [ "libc", - "rand_chacha", - "rand_core", + "rand_chacha 0.3.1", + "rand_core 0.6.4", +] + +[[package]] +name = "rand" +version = "0.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6db2770f06117d490610c7488547d543617b21bfa07796d7a12f6f1bd53850d1" +dependencies = [ + "rand_chacha 0.9.0", + "rand_core 0.9.5", ] [[package]] @@ -3076,7 +3218,17 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" dependencies = [ "ppv-lite86", - "rand_core", + "rand_core 0.6.4", +] + +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core 0.9.5", ] [[package]] @@ -3089,6 +3241,24 @@ dependencies = [ ] [[package]] +name = "rand_core" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +dependencies = [ + "getrandom 0.3.4", +] + +[[package]] +name = "raw-cpuid" +version = "11.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "498cd0dc59d73224351ee52a95fee0f1a617a2eae0e7d9d720cc622c73a54186" +dependencies = [ + "bitflags", +] + +[[package]] name = "rayon" version = "1.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -3221,7 +3391,7 @@ dependencies = [ "num-traits", "pkcs1", "pkcs8", - "rand_core", + "rand_core 0.6.4", "signature", "spki", "subtle", @@ -3657,7 +3827,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "77549399552de45a898a580c1b41d445bf730df867cc44e6c0233bbc4b8329de" dependencies = [ "digest", - "rand_core", + "rand_core 0.6.4", ] [[package]] @@ -3724,6 +3894,15 @@ dependencies = [ ] [[package]] +name = "spinning_top" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d96d2d1d716fb500937168cc09353ffdc7a012be8475ac7308e1bdf0e3923300" +dependencies = [ + "lock_api", +] + +[[package]] name = "spki" version = "0.7.3" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -3849,7 +4028,7 @@ dependencies = [ "memchr", "once_cell", "percent-encoding", - "rand", + "rand 0.8.5", "rsa", "serde", "sha1", @@ -3888,7 +4067,7 @@ dependencies = [ "md-5", "memchr", "once_cell", - "rand", + "rand 0.8.5", "serde", "serde_json", "sha2", @@ -4332,6 +4511,35 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ab16f14aed21ee8bfd8ec22513f7287cd4a91aa92e44edfe2c17ddd004e92607" [[package]] +name = "tonic" +version = "0.14.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eb7613188ce9f7df5bfe185db26c5814347d110db17920415cf2fbcad85e7203" +dependencies = [ + "async-trait", + "axum", + "base64", + "bytes", + "h2", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-timeout", + "hyper-util", + "percent-encoding", + "pin-project", + "socket2", + "sync_wrapper", + "tokio", + "tokio-stream", + "tower", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] name = "tower" version = "0.5.3" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -4395,6 +4603,23 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3" [[package]] +name = "tower_governor" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44de9b94d849d3c46e06a883d72d408c2de6403367b39df2b1c9d9e7b6736fe6" +dependencies = [ + "axum", + "forwarded-header-value", + "governor", + "http", + "pin-project", + "thiserror 2.0.17", + "tonic", + "tower", + "tracing", +] + +[[package]] name = "tracing" version = "0.1.44" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -4713,6 +4938,16 @@ dependencies = [ ] [[package]] +name = "web-time" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a6580f308b1fad9207618087a65c04e7a10bc77e02c8e84e9b00dd4b12fa0bb" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] name = "webpki-root-certs" version = "1.0.5" source = "registry+https://github.com/rust-lang/crates.io-index" @@ -5302,8 +5537,10 @@ dependencies = [ "toml", "tower", "tower-http", + "tower_governor", "tracing", "tracing-subscriber", + "url", "uuid", ] diff --git a/Cargo.toml b/Cargo.toml index b082481..6ce0a29 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -54,11 +54,16 @@ sqlx = { version = "0.8", features = [ thiserror = "2.0" tokio = { version = "1.49", features = ["full"] } toml = "0.9" -tokio-rustls = { version = "0.26", default-features = false, features = ["ring", "tls12"] } +tokio-rustls = { version = "0.26", default-features = false, features = [ + "ring", + "tls12", +] } tower = { version = "0.5", features = ["full"] } +tower_governor = { version = "0.8", features = ["axum", "tracing"] } tower-http = { version = "0.6", features = ["full"] } tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter", "json"] } +url = "2.5" uuid = { version = "1.19", features = ["v4", "v5"] } [dev-dependencies] diff --git a/pkg/zorian.toml b/pkg/zorian.toml index e290f13..3704f52 100644 --- a/pkg/zorian.toml +++ b/pkg/zorian.toml @@ -9,7 +9,7 @@ appname="zorian" # Empty hostnames means accept all requests (catch-all). [[listen]] addr = "0.0.0.0:3000" -hostnames = ["localhost", "127.0.0.1"] +hostnames = ["localhost", "127.0.0.1", "::1"] # For HTTPS, set tls_cert and tls_key to PEM file paths. # tls_crt = "/etc/zorian/cert.pem" diff --git a/src/backends/zig.rs b/src/backends/zig.rs index ed22326..5920645 100644 --- a/src/backends/zig.rs +++ b/src/backends/zig.rs @@ -212,6 +212,7 @@ impl ZigController { response::Response::builder() .status(status) .header(http::header::CONTENT_TYPE, "application/octet-stream") + .header(http::header::CONTENT_LENGTH, bytes.len()) .body(body::Body::from(bytes)) .unwrap() } diff --git a/src/config.rs b/src/config.rs index e8c6fcb..71bbe4f 100644 --- a/src/config.rs +++ b/src/config.rs @@ -7,6 +7,24 @@ use std::path::{Path, PathBuf}; use serde::Deserialize; use thiserror::Error; +/// When receiving a SIGINT/SIGTERM signal, we will wait for the proposed timeout before terminating workers +pub const SHUTDOWN_TIMEOUT: std::time::Duration = std::time::Duration::from_mins(1); + +/// Request timeout - maximum time to process a request (protects against Slowloris) +pub const REQUEST_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30); + +/// Maximum request body size in bytes (64 MB) +pub const MAX_BODY_SIZE: usize = 64 * 1024 * 1024; + +/// Maximum number of concurrent requests across all clients +pub const MAX_CONCURRENT_REQUESTS: usize = 512; + +/// Rate limit: requests per second per client IP +pub const RATE_LIMIT_PER_SECOND: u64 = 10; + +/// Rate limit: burst size (max requests allowed in a burst) per client IP +pub const RATE_LIMIT_BURST_SIZE: u32 = 50; + #[derive(Debug, Error)] pub enum ConfigError { #[error("failed to read config file: {0}")] @@ -58,8 +76,12 @@ pub struct ListenerConfig { impl Default for ListenerConfig { fn default() -> Self { Self { - addr: "0.0.0.0:3000".to_string(), - hostnames: Vec::new(), + addr: "127.0.0.1:3000".to_string(), + hostnames: vec![ + String::from("[::1]"), + String::from("127.0.0.1"), + String::from("localhost"), + ], tls_crt: None, tls_key: None, } diff --git a/src/main.rs b/src/main.rs index 34d89a9..a184a7f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -19,11 +19,16 @@ use axum::{ extract::{self, connect_info::Connected}, http::{self, Request, Response}, }; +use tokio::signal; use tokio_rustls::TlsAcceptor; use tracing::{error, info}; use tracing_subscriber::{layer::SubscriberExt, registry::LookupSpan, util::SubscriberInitExt}; use crate::backends::zig::ZigController; +use crate::config::{ + MAX_BODY_SIZE, MAX_CONCURRENT_REQUESTS, RATE_LIMIT_BURST_SIZE, RATE_LIMIT_PER_SECOND, + REQUEST_TIMEOUT, SHUTDOWN_TIMEOUT, +}; use crate::web::WebController; const VERSION: &str = env!("CARGO_PKG_VERSION"); @@ -67,6 +72,25 @@ struct RequestInfo { user_agent: Option<String>, } +/// Key extractor for tower-governor that uses ClientInfo to get the client IP. +#[derive(Clone)] +struct ClientIpKeyExtractor; + +impl tower_governor::key_extractor::KeyExtractor for ClientIpKeyExtractor { + type Key = std::net::IpAddr; + + fn name(&self) -> &'static str { + "ClientIpKeyExtractor" + } + + fn extract<T>(&self, req: &Request<T>) -> Result<Self::Key, tower_governor::GovernorError> { + req.extensions() + .get::<extract::ConnectInfo<ClientInfo>>() + .map(|ci| ci.0.0.ip()) + .ok_or(tower_governor::GovernorError::UnableToExtractKey) + } +} + #[tokio::main] async fn main() { let mut config_path = None; @@ -218,6 +242,13 @@ async fn main() { ) .on_failure(tower_http::trace::DefaultOnFailure::new().level(tracing::Level::ERROR)); + let governor_config = tower_governor::governor::GovernorConfigBuilder::default() + .per_second(RATE_LIMIT_PER_SECOND) + .burst_size(RATE_LIMIT_BURST_SIZE) + .key_extractor(ClientIpKeyExtractor) + .finish() + .unwrap(); + let app = axum::Router::new() .merge(web_controller.router()) .merge(zig_controller.router()) @@ -229,10 +260,24 @@ async fn main() { .layer(tower_http::request_id::SetRequestIdLayer::new( REQUEST_ID_HEADER.clone(), tower_http::request_id::MakeRequestUuid, + )) + .layer(tower_http::timeout::TimeoutLayer::with_status_code( + http::StatusCode::REQUEST_TIMEOUT, + REQUEST_TIMEOUT, + )) + .layer(tower_http::limit::RequestBodyLimitLayer::new(MAX_BODY_SIZE)) + .layer(tower_governor::GovernorLayer::new(Arc::new( + governor_config, + ))) + .layer(tower::limit::ConcurrencyLimitLayer::new( + MAX_CONCURRENT_REQUESTS, )); let mut tasks = tokio::task::JoinSet::new(); + // The channel is used to broadcast SIGINT/SIGTERM to all listeners. + let (shutdown_tx, _) = tokio::sync::broadcast::channel::<()>(1); + for listener_config in config.listeners() { let tcp_listener = tokio::net::TcpListener::bind(&listener_config.addr) .await @@ -259,22 +304,87 @@ async fn main() { })) .into_make_service_with_connect_info::<ClientInfo>(); + let mut shutdown_rx = shutdown_tx.subscribe(); + let shutdown_signal = async move { + shutdown_rx.recv().await.ok(); + }; + if let (Some(crt_path), Some(key_path)) = (&listener_config.tls_crt, &listener_config.tls_key) { let tls_listener = TlsListener::new(tcp_listener, crt_path, key_path); tasks.spawn(async move { - axum::serve(tls_listener, app).await.unwrap(); + axum::serve(tls_listener, app) + .with_graceful_shutdown(shutdown_signal) + .await + .unwrap(); }); } else { tasks.spawn(async move { - axum::serve(tcp_listener, app).await.unwrap(); + axum::serve(tcp_listener, app) + .with_graceful_shutdown(shutdown_signal) + .await + .unwrap(); }); } } - if let Some(result) = tasks.join_next().await { - result.unwrap(); + // Graceful shutdown + let sigint = async { + signal::ctrl_c() + .await + .expect("failed to install Ctrl+C handler"); + }; + let sigterm = { + #[cfg(unix)] + let terminate = async { + signal::unix::signal(signal::unix::SignalKind::terminate()) + .expect("failed to install signal handler") + .recv() + .await; + }; + + #[cfg(not(unix))] + let terminate = std::future::pending::<()>(); + + terminate + }; + tokio::select! { + _ = sigint => info!("received SIGINT, shutting down"), + _ = sigterm => info!("received SIGTERM, shutting down"), + result = tasks.join_next() => { + match result { + Some(Ok(())) => error!("listener exited unexpectedly, shutting down"), + Some(Err(e)) => error!("listener failed: {e}, shutting down"), + None => { + error!("no listeners running"); + return; + } + } + } + } + + drop(shutdown_tx); // broadcast + + // Wait for all listeners to finish with timeout + let shutdown_result = tokio::time::timeout(SHUTDOWN_TIMEOUT, async { + while let Some(result) = tasks.join_next().await { + if let Err(e) = result { + error!("listener task failed: {e}"); + } + } + }) + .await; + + if shutdown_result.is_err() { + error!( + "shutdown timeout after {:?}, aborting {} remaining tasks", + SHUTDOWN_TIMEOUT, + tasks.len() + ); + tasks.abort_all(); + } else { + info!("shutdown complete"); } } @@ -362,7 +472,6 @@ where fn call(&mut self, req: Request<Body>) -> Self::Future { let interface = req.extensions().get::<ListenerInfo>().cloned(); - // Check if hostname validation is needed if let Some(iface) = interface && !iface.hosts.is_empty() { @@ -370,10 +479,25 @@ where .headers() .get(http::header::HOST) .and_then(|v| v.to_str().ok()) - .map(|h| h.split(':').next().unwrap_or(h)); // Strip port if present + .and_then(|raw| { + let without_port = if let Some((host, port)) = raw.rsplit_once(':') + && port.parse::<u16>().is_ok() + && (host.ends_with(']') || !host.contains('[')) + { + host + } else { + raw + }; + url::Host::parse(without_port).ok().map(|h| h.to_string()) + }); let is_valid = host - .map(|h| iface.hosts.iter().any(|allowed| allowed == h)) + .map(|h| { + iface + .hosts + .iter() + .any(|allowed| allowed.eq_ignore_ascii_case(&h)) + }) .unwrap_or(false); if !is_valid { |