diff --git a/.github/FUNDING.yml b/.github/FUNDING.yml new file mode 100644 index 0000000..3832f20 --- /dev/null +++ b/.github/FUNDING.yml @@ -0,0 +1,10 @@ +# Arvik — Funding options +# This file adds the "Sponsor" button to the GitHub repository. +# Docs: https://docs.github.com/en/repositories/managing-your-repositorys-settings-and-features/customizing-your-repository/displaying-a-sponsor-button-in-your-repository + +github: + - aarambh-darshan + +custom: + - "https://buymeacoffee.com/aarambhdevhub" + - "https://razorpay.me/@aarambhdevhub" diff --git a/.github/ISSUE_TEMPLATE/bug_report.md b/.github/ISSUE_TEMPLATE/bug_report.md new file mode 100644 index 0000000..09839e1 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/bug_report.md @@ -0,0 +1,49 @@ +--- +name: Bug Report +about: Something is broken or behaving unexpectedly +title: "fix: " +labels: bug +assignees: "" +--- + +## Description + +A clear, one-paragraph description of the bug. + +## Minimal Reproduction + +```rust +// The smallest possible code that demonstrates the bug. +// Remove everything unrelated. +// This must compile (or explain why it doesn't). +use arvik::{Router, get, serve_app}; + +#[tokio::main] +async fn main() { + let app = Router::new() + .route("/", get(|| async { "hello" })); + // reproduce the bug here + serve_app("127.0.0.1:8080", app).await.unwrap(); +} +``` + +## Expected Behavior + +What you expected to happen. + +## Actual Behavior + +What actually happened. Include the full error message, panic output, or unexpected value. + +## Environment + +| | | +|---|---| +| Arvik version | `0.9.x` | +| Rust version | `rustc --version` output | +| OS | e.g. Pop OS 22.04 / macOS 14 / Windows 11 | +| Active features | e.g. `default`, `tls`, `ws`, `config`, `macros` | + +## Additional Context + +Anything else that might be relevant — related issues, workarounds you've tried, etc. diff --git a/.github/ISSUE_TEMPLATE/feature_request.md b/.github/ISSUE_TEMPLATE/feature_request.md new file mode 100644 index 0000000..44b5981 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/feature_request.md @@ -0,0 +1,35 @@ +--- +name: Feature Request +about: Suggest a new feature or API improvement +title: "feat: " +labels: enhancement +assignees: "" +--- + +## Use Case + +What are you trying to accomplish? Describe the problem you are solving, not the solution you want. +A good use case sounds like: "I need to stream server-sent events with per-client backpressure but the current API requires Y workaround." + +## Proposed API + +Show what you would want to write. Code is clearer than prose. + +```rust +// Example of the API you'd like to exist: +let app = Router::new() + .route("/events", get(events_handler)) + .layer(YourNewLayer::new()); +``` + +## Alternatives Considered + +What can you do today instead? Why is it insufficient? + +## Affected Crate(s) + +Which crate(s) would this change? e.g. `arvik-core`, `arvik-router`, `arvik-middleware`, `arvik-hyper`, `arvik-extract`, `arvik-ws`, `arvik-sse`, `arvik-static`, `arvik-tls`, `arvik-config`, `arvik-observe` + +## Additional Context + +Links to prior art (axum, actix-web, tower-http…), related issues, or any other context that helps evaluate the request. diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md new file mode 100644 index 0000000..4f14d39 --- /dev/null +++ b/.github/PULL_REQUEST_TEMPLATE.md @@ -0,0 +1,37 @@ +## Summary + +Brief description of what this PR does and why. + +Closes # + +--- + +## Type of Change + +- [ ] Bug fix — non-breaking, fixes incorrect behavior +- [ ] New feature — non-breaking, adds functionality +- [ ] Breaking change — changes existing behavior or API +- [ ] Documentation — no behavior change +- [ ] Performance — no behavior change, improves speed or memory +- [ ] Refactor — no behavior change, improves code structure +- [ ] Test — adds or improves test coverage +- [ ] CI / build — changes to workflow or build configuration + +--- + +## Checklist + +- [ ] `cargo fmt` has been run — code is formatted +- [ ] `cargo clippy --workspace --all-features -- -D warnings` is clean +- [ ] `cargo test --workspace --all-features` passes +- [ ] `cargo check --workspace` passes without default features +- [ ] New or changed public APIs have `///` doc comments with examples +- [ ] `cargo doc --workspace --all-features --no-deps` builds without warnings +- [ ] `CHANGELOG.md` has been updated under `[Unreleased]` +- [ ] Tests are included for new behavior or bug fixes + +--- + +## Notes for Reviewers + +Anything the reviewer should pay particular attention to, tricky parts of the implementation, or decisions you were uncertain about. diff --git a/.github/release-notes/0.9.5.md b/.github/release-notes/0.9.5.md new file mode 100644 index 0000000..f0d24da --- /dev/null +++ b/.github/release-notes/0.9.5.md @@ -0,0 +1,83 @@ +## Arvik v0.9.5 — Performance Sprint + +Arvik v0.9.5 is the performance-focused release: low-level socket and listener +tuning, HTTP/2 presets, lower-copy `Body` constructors, and a manual benchmark +harness for measuring all of it. + +### Highlights + +- Low-level socket/listener tuning on `ServerConfig`: TCP_NODELAY, keepalive, + reuse port/address, backlog, socket buffers, and accept worker settings. +- HTTP/2 performance presets: `ServerConfig::http2_low_latency()` and + `ServerConfig::http2_high_throughput()`. +- `RuntimeConfig` for tuned Tokio runtime construction. +- Optional `runtime-metrics` facade feature forwarding to + `arvik-hyper/runtime-metrics`. +- Lower-copy `Body` constructors: `from_static`, `from_chunks`, and + `is_empty_hint`. +- Benchmark examples under `examples/benchmarks` and a tuning example under + `examples/performance_tuning`. +- Manual benchmark workflow that uploads benchmark artifacts. +- `scripts/bench/wrk.sh`, `scripts/bench/hey.sh`, and + `scripts/bench/run-all.sh`. + +### Changed + +- Router matched paths now use interned `Arc` metadata. +- Path parameter keys use stack-backed `SmallVec` storage with interned names. +- JSON responses serialize into `BytesMut` and freeze into `Bytes`. +- Arvik config maps socket/listener settings into `ServerConfig`. +- Workspace version bumped to `0.9.5`. + +### Notes + +- The planned full Linux `io_uring` backend is deferred as future experimental + work. The 0.9.x completion path is maintainable Tokio/socket tuning plus + optional runtime metrics. +- Benchmark CI is manual and artifact-based, not PR-commenting or release + gating. + +### Compatibility + +No public Rust API was removed or renamed. New `ServerConfig` methods are +additive; existing configurations continue to work unchanged. + +### Install + +```toml +[dependencies] +arvik = "0.9.5" +``` + +With TLS, WebSocket, and configuration support: + +```toml +[dependencies] +arvik = { version = "0.9.5", features = ["tls", "ws", "config"] } +``` + +### Crates published to crates.io + +| Crate | Version | Description | +|-------|---------|-------------| +| `arvik-core` | `0.9.5` | Core traits, request/response types, extractors | +| `arvik-router` | `0.9.5` | Radix-trie router and method dispatch | +| `arvik-hyper` | `0.9.5` | Hyper-based server with graceful shutdown | +| `arvik-extract` | `0.9.5` | Body/multipart/cookie extractors | +| `arvik-middleware` | `0.9.5` | CORS, auth, rate limiting, compression | +| `arvik-ws` | `0.9.5` | WebSocket upgrades and sockets | +| `arvik-sse` | `0.9.5` | Server-sent events | +| `arvik-static` | `0.9.5` | Static file and embedded asset serving | +| `arvik-tls` | `0.9.5` | rustls + native-tls backends | +| `arvik-macros` | `0.9.5` | Route and handler proc macros | +| `arvik-test` | `0.9.5` | Test client utilities | +| `arvik-config` | `0.9.5` | File/env runtime configuration | +| `arvik-observe` | `0.9.5` | Metrics, logging, tracing, health checks | +| `arvik` | `0.9.5` | Stable facade crate | + +### Documentation + +- [README](../../README.md) +- [Architecture](../../ARCHITECTURE.md) +- [Roadmap](../../ROADMAP.md) +- [Changelog](../../CHANGELOG.md) diff --git a/Cargo.lock b/Cargo.lock index 57c2922..25645c5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -204,6 +204,7 @@ dependencies = [ "async-compression", "base64", "cookie", + "futures-util", "http", "http-body", "mime", @@ -312,6 +313,7 @@ dependencies = [ "cookie", "futures-util", "http", + "httpdate", "serde", "serde_json", "serde_urlencoded", diff --git a/arvik-config/src/lib.rs b/arvik-config/src/lib.rs index ebdef92..d632a43 100644 --- a/arvik-config/src/lib.rs +++ b/arvik-config/src/lib.rs @@ -248,6 +248,28 @@ impl ArvikConfig { self.http2.initial_connection_window_size, "http2.initial_connection_window_size", )?; + // h2 asserts `size <= (1 << 31) - 1` when the connection is built; + // reject oversized values here with a clear message instead of + // panicking at runtime on the first HTTP/2 request. + const H2_MAX_WINDOW: u64 = (1 << 31) - 1; + for (field, value) in [ + ( + "http2.initial_stream_window_size", + self.http2.initial_stream_window_size, + ), + ( + "http2.initial_connection_window_size", + self.http2.initial_connection_window_size, + ), + ] { + if let Some(size) = value + && size as u64 > H2_MAX_WINDOW + { + return Err(ConfigError::Validation(format!( + "{field} must be at most {H2_MAX_WINDOW} (h2 protocol maximum), got {size}" + ))); + } + } validate_nonzero( self.http2.max_concurrent_streams, "http2.max_concurrent_streams", @@ -418,7 +440,7 @@ impl Default for Http2Section { } /// Optional TLS file settings. -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[derive(Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(default, deny_unknown_fields)] pub struct TlsSection { /// TLS backend. @@ -433,6 +455,23 @@ pub struct TlsSection { pub pkcs12_password: Option, } +impl std::fmt::Debug for TlsSection { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + // Never render the keystore password: any `{:?}` of the config would + // otherwise leak it to logs and crash dumps. + f.debug_struct("TlsSection") + .field("backend", &self.backend) + .field("cert_path", &self.cert_path) + .field("key_path", &self.key_path) + .field("pkcs12_path", &self.pkcs12_path) + .field( + "pkcs12_password", + &self.pkcs12_password.as_ref().map(|_| "***REDACTED***"), + ) + .finish() + } +} + impl Default for TlsSection { fn default() -> Self { Self { @@ -599,16 +638,39 @@ fn merge_values(base: &mut Value, overlay: Value) { fn apply_env(root: &mut Value, prefix: &str) -> Result<()> { let prefix = format!("{}_", prefix.trim_end_matches('_')); - for (var, value) in std::env::vars() { + let mut unknown: Vec = Vec::new(); + for (raw_var, raw_value) in std::env::vars_os() { + // Skip non-UTF-8 entries instead of panicking on unrelated variables + // somewhere in the process environment. + let (Ok(var), Ok(value)) = (raw_var.into_string(), raw_value.into_string()) else { + continue; + }; if !var.starts_with(&prefix) { continue; } let suffix = &var[prefix.len()..]; - let Some((path, value)) = parse_env_override(suffix, &var, &value)? else { - continue; - }; - set_path(root, path, value); + match parse_env_override(suffix, &var, &value)? { + Some((path, value)) => set_path(root, path, value), + None if suffix_starts_with_known_section(suffix) => { + // A typo inside a known section must not silently fall back to + // file/default values — the same mistake in a TOML file fails + // loudly. + unknown.push(var); + } + // Variables outside the config schema (other tooling sharing the + // prefix) are ignored, as they always were. + None => {} + } + } + + if !unknown.is_empty() { + return Err(ConfigError::Validation(format!( + "unknown configuration environment variable(s): {}. Expected names prefixed with `{}` (e.g. {}SERVER_PORT); check for typos.", + unknown.join(", "), + prefix, + prefix + ))); } Ok(()) @@ -709,6 +771,17 @@ fn parse_env_override( Ok(parsed) } +/// True when the variable's first path segment names a config section, so an +/// unmatched remainder is almost certainly a typo (e.g. `SERVER_MAXCONNECTIONS` +/// instead of `SERVER_MAX_CONNECTIONS`) rather than another tool's variable. +fn suffix_starts_with_known_section(suffix: &str) -> bool { + let section = suffix.split('_').next().unwrap_or_default(); + matches!( + section.to_ascii_lowercase().as_str(), + "server" | "http1" | "http2" | "shutdown" | "tls" + ) +} + fn parse_bool(var: &str, value: &str) -> Result { match value.to_ascii_lowercase().as_str() { "true" | "1" | "yes" | "on" => Ok(true), @@ -779,10 +852,21 @@ impl ConfigWatcher { }) .map_err(|err| ConfigError::Watch(err.to_string()))?; + // Watch the parent *directories*, not the files themselves: editors + // and deploy tooling replace configs via temp-file+rename, which + // silently orphans an inode watch. Reloads read by path, so they pick + // up the replacement on the next event regardless. + let mut watched_dirs = std::collections::HashSet::new(); for file in &builder.files { - watcher - .watch(file, notify::RecursiveMode::NonRecursive) - .map_err(|err| ConfigError::Watch(err.to_string()))?; + let dir = match file.parent() { + Some(dir) => dir.to_path_buf(), + None => continue, + }; + if watched_dirs.insert(dir.clone()) { + watcher + .watch(&dir, notify::RecursiveMode::NonRecursive) + .map_err(|err| ConfigError::Watch(err.to_string()))?; + } } let current_task = std::sync::Arc::clone(¤t); @@ -794,7 +878,13 @@ impl ConfigWatcher { } tokio::time::sleep(Duration::from_millis(200)).await; - while event_rx.try_recv().is_ok() {} + // Coalesce bursts of events, but forward any watcher errors + // that arrived during the debounce window instead of dropping. + while let Ok(pending) = event_rx.try_recv() { + if let Err(err) = pending { + let _ = update_tx.send(Err(ConfigError::Watch(err.to_string()))); + } + } match builder.build_inner() { Ok(config) => { @@ -937,6 +1027,60 @@ mod tests { assert_eq!(config.server.port, 9090); } + #[test] + fn typo_in_known_section_fails_loudly() { + let prefix = unique_prefix(); + let key = format!("{prefix}_SERVER_MAXCONNECTIONS"); + + unsafe { + std::env::set_var(&key, "50"); + } + let result = ArvikConfig::builder().env_prefix(&prefix).build(); + unsafe { + std::env::remove_var(&key); + } + + let err = result.expect_err("typo'd env var in a known section must fail"); + let text = err.to_string(); + assert!(text.contains("SERVER_MAXCONNECTIONS"), "got: {text}"); + } + + #[test] + fn foreign_vars_in_prefix_namespace_are_ignored() { + // Another tool sharing the prefix must not break config loading. + let prefix = unique_prefix(); + let key = format!("{prefix}_FOREIGN_TOOL_FLAG"); + + unsafe { + std::env::set_var(&key, "1"); + } + let result = ArvikConfig::builder().env_prefix(&prefix).build(); + unsafe { + std::env::remove_var(&key); + } + + result.unwrap(); + } + + #[test] + fn h2_window_above_protocol_maximum_is_rejected() { + let config = ArvikConfig { + http2: Http2Section { + initial_stream_window_size: Some((1 << 31) as u32), // h2 max is (1<<31)-1 + ..Http2Section::default() + }, + ..ArvikConfig::default() + }; + + let err = config + .validate() + .expect_err("oversized window must fail validation"); + assert!( + err.to_string().contains("initial_stream_window_size"), + "got: {err}" + ); + } + #[test] fn validation_reports_bad_limits() { let dir = tempfile::tempdir().unwrap(); diff --git a/arvik-core/src/accept.rs b/arvik-core/src/accept.rs new file mode 100644 index 0000000..d872239 --- /dev/null +++ b/arvik-core/src/accept.rs @@ -0,0 +1,104 @@ +//! Shared `Accept-Encoding` negotiation used by the compression middleware +//! and static precompressed asset selection. +//! +//! Parses full RFC 9110 §12.5.3 semantics the naive `contains("br")` check +//! missed: per-entry `q`-values, explicit `;q=0` refusals overriding the +//! `*` wildcard, and multi-line header values. + +/// Parse one header value into `(coding, quality)` pairs, lowercased. +/// +/// `q=0` entries are preserved so callers can distinguish an explicit +/// refusal from mere absence. +pub fn parse_accept_encodings(value: &str) -> Vec<(String, f32)> { + value + .split(',') + .filter_map(|entry| { + let mut parts = entry.trim().split(';'); + let coding = parts.next()?.trim().to_ascii_lowercase(); + if coding.is_empty() { + return None; + } + let mut quality = 1.0_f32; + for parameter in parts { + let parameter = parameter.trim(); + if let Some(q) = parameter + .strip_prefix("q=") + .or_else(|| parameter.strip_prefix("Q=")) + { + quality = q.trim().parse().unwrap_or(1.0); + } + } + Some((coding, quality.clamp(0.0, 1.0))) + }) + .collect() +} + +/// Pick the first entry of `supported` (in priority order) the client accepts. +/// +/// An exact `;q=0` entry refuses that coding even when `*` is present; codings +/// not mentioned at all fall back to the `*` wildcard's weight. +pub fn negotiate<'a>(supported: &[&'a str], accept_encoding: &str) -> Option<&'a str> { + let entries = parse_accept_encodings(accept_encoding); + let wildcard_q = entries + .iter() + .rev() + .find(|(coding, _)| coding == "*") + .map(|(_, q)| *q) + .unwrap_or(0.0); + + supported.iter().copied().find(|candidate| { + let candidate = candidate.to_ascii_lowercase(); + match entries + .iter() + .rev() + .find(|(coding, _)| *coding == candidate) + { + Some((_, q)) => *q > 0.0, + None => wildcard_q > 0.0, + } + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_quality_values() { + let parsed = parse_accept_encodings("gzip;q=0.5, BR, zstd ; Q=0 , *;q=0.1"); + assert_eq!( + parsed, + vec![ + ("gzip".to_string(), 0.5), + ("br".to_string(), 1.0), + ("zstd".to_string(), 0.0), + ("*".to_string(), 0.1), + ] + ); + } + + #[test] + fn explicit_refusal_beats_wildcard() { + // The client refused brotli even though * is allowed. + assert_eq!(negotiate(&["br", "gzip"], "gzip, br;q=0, *"), Some("gzip")); + assert_eq!(negotiate(&["br"], "*, br;q=0"), None); + } + + #[test] + fn unlisted_codings_follow_the_wildcard() { + assert_eq!(negotiate(&["zstd"], "*"), Some("zstd")); + assert_eq!(negotiate(&["zstd"], "gzip"), None); + } + + #[test] + fn priority_order_is_preserved() { + // Server preference wins among acceptable codings (tower-http-style); + // only explicit refusal or absence redirects the choice. + assert_eq!(negotiate(&["br", "gzip"], "gzip, br"), Some("br")); + assert_eq!( + negotiate(&["br", "gzip"], "br;q=0.2, gzip"), + Some("br"), + "any q > 0 keeps the server-preferred coding" + ); + } +} diff --git a/arvik-core/src/body.rs b/arvik-core/src/body.rs index 307f0e9..620048e 100644 --- a/arvik-core/src/body.rs +++ b/arvik-core/src/body.rs @@ -121,6 +121,42 @@ impl Body { Ok(collected.to_bytes()) } + /// Collect the body into [`Bytes`], enforcing a maximum size. + /// + /// Rejects early (before reading the stream) when a known + /// `Content-Length` exceeds `limit`; otherwise stops reading and returns + /// `BodyLimitError::TooLarge` as soon as the accumulated data would + /// grow past `limit`. + pub async fn to_bytes_limited( + self, + limit: usize, + ) -> Result { + use crate::body_limit::BodyLimitError; + + // Fast path: declared length already over the limit. + if let Some(upper) = http_body::Body::size_hint(&self).upper() { + if upper > limit as u64 { + return Err(BodyLimitError::TooLarge); + } + } + + let mut buf = bytes::BytesMut::new(); + let mut body = self; + while let Some(frame) = BodyExt::frame(&mut body).await { + let data = frame?.into_data().map_err(|_| { + Box::new(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "unexpected body trailer", + )) as BoxError + })?; + if buf.len() + data.len() > limit { + return Err(BodyLimitError::TooLarge); + } + buf.extend_from_slice(&data); + } + Ok(buf.freeze()) + } + /// Collect the entire body into a UTF-8 [`String`]. /// /// Returns an error if the body is not valid UTF-8 or if @@ -360,6 +396,37 @@ mod tests { assert_eq!(body.to_bytes().await.unwrap(), Bytes::new()); } + #[tokio::test] + async fn to_bytes_limited_rejects_oversized_full_body() { + let body = Body::from_bytes(vec![0u8; 128].into()); + match body.to_bytes_limited(64).await { + Err(crate::body_limit::BodyLimitError::TooLarge) => {} + other => panic!("expected TooLarge, got {:?}", other.map(|b| b.len())), + } + } + + #[tokio::test] + async fn to_bytes_limited_accepts_bodies_under_limit() { + let body = Body::from_bytes(vec![7u8; 10].into()); + let bytes = body.to_bytes_limited(64).await.unwrap(); + assert_eq!(bytes.len(), 10); + } + + #[tokio::test] + async fn to_bytes_limited_rejects_streaming_body_without_content_length() { + let chunks: Vec> = vec![ + Ok(Bytes::from_static(&[0u8; 40])), + Ok(Bytes::from_static(&[0u8; 40])), + ]; + let body = Body::from_stream(futures_util::stream::iter(chunks)); + // No size-hint upper bound — the cap must trigger mid-stream. + assert!(http_body::Body::size_hint(&body).upper().is_none()); + match body.to_bytes_limited(64).await { + Err(crate::body_limit::BodyLimitError::TooLarge) => {} + other => panic!("expected TooLarge, got {:?}", other.map(|b| b.len())), + } + } + #[tokio::test] async fn static_body_is_zero_copy_bytes() { let body = Body::from_static(b"hello"); diff --git a/arvik-core/src/body_limit.rs b/arvik-core/src/body_limit.rs new file mode 100644 index 0000000..c334850 --- /dev/null +++ b/arvik-core/src/body_limit.rs @@ -0,0 +1,136 @@ +//! Default request-body size limit for buffering extractors. +//! +//! Extractors that buffer the whole request body ([`Bytes`](bytes::Bytes), +//! [`String`], [`Json`](https://docs.rs/arvik/latest/arvik/struct.Json.html), +//! [`Form`](https://docs.rs/arvik/latest/arvik/struct.Form.html)) enforce a +//! default limit of [`DEFAULT_BODY_LIMIT`] (2 MiB). Requests exceeding the +//! limit are rejected with `413 Payload Too Large` instead of being buffered +//! into memory. +//! +//! # Raising or disabling the limit +//! +//! Insert a [`DefaultBodyLimit`] extension into the request to override the +//! default — e.g. with the convenience +//! [`DefaultBodyLimitLayer`](https://docs.rs/arvik-middleware) from +//! `arvik-middleware`, or any middleware/mapping that inserts the extension: +//! +//! ```rust,ignore +//! use arvik_core::body_limit::DefaultBodyLimit; +//! +//! // Allow 16 MiB bodies on this route. +//! Router::new() +//! .route("/upload", post(upload)) +//! .layer(DefaultBodyLimitLayer::max(16 * 1024 * 1024)); +//! ``` + +use std::fmt; + +/// Extension holding the maximum buffered-body size for this request. +/// +/// Extractors resolve the effective limit as: the `DefaultBodyLimit` +/// extension if present, otherwise [`DEFAULT_BODY_LIMIT`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct DefaultBodyLimit(pub usize); + +impl DefaultBodyLimit { + /// The limit applied when no extension is present (2 MiB). + pub const DEFAULT: Self = Self(DEFAULT_BODY_LIMIT); + + /// Create a limit override of `limit` bytes. + pub fn max(limit: usize) -> Self { + Self(limit) + } + + /// Disable the buffered-body limit entirely. + /// + /// Only sensible for routes behind other protections (e.g. an upstream + /// proxy enforcing a size cap). + pub fn disabled() -> Self { + Self(usize::MAX) + } +} + +/// The default maximum request-body size buffering extractors accept: 2 MiB, +/// matching common framework defaults. +pub const DEFAULT_BODY_LIMIT: usize = 2 * 1024 * 1024; + +/// Canonical error carried by streaming bodies that exceed their size limit +/// while being read (e.g. by [`crate::body::Body`] wrappers enforcing a cap). +/// +/// Readers detect it via [`BodyLimitError::is_payload_too_large`] and reject +/// with `413`. +#[derive(Debug)] +pub struct BodyTooLarge; + +impl fmt::Display for BodyTooLarge { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "request body exceeds the maximum allowed size") + } +} + +impl std::error::Error for BodyTooLarge {} + +/// Error produced when reading a body under a size limit. +#[derive(Debug)] +pub enum BodyLimitError { + /// The body exceeded the configured limit. + TooLarge, + /// The underlying body stream failed. + Read(crate::body::BoxError), +} + +impl BodyLimitError { + /// True when this failure represents an over-limit body — either caught + /// directly or surfaced by the underlying stream as [`BodyTooLarge`]. + pub fn is_payload_too_large(&self) -> bool { + match self { + Self::TooLarge => true, + Self::Read(e) => e.is::(), + } + } +} + +impl fmt::Display for BodyLimitError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::TooLarge => write!(f, "request body exceeds the maximum allowed size"), + Self::Read(e) => write!(f, "failed to read request body: {e}"), + } + } +} + +impl std::error::Error for BodyLimitError {} + +impl From for BodyLimitError { + fn from(e: crate::body::BoxError) -> Self { + Self::Read(e) + } +} + +/// Resolve the effective body limit from request extensions. +/// +/// Returns the [`DefaultBodyLimit`] extension value if present, otherwise +/// [`DEFAULT_BODY_LIMIT`]. +pub fn resolve_limit(extensions: &http::Extensions) -> usize { + extensions + .get::() + .map(|l| l.0) + .unwrap_or(DEFAULT_BODY_LIMIT) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn resolves_extension_over_default() { + let mut ext = http::Extensions::new(); + assert_eq!(resolve_limit(&ext), DEFAULT_BODY_LIMIT); + + ext.insert(DefaultBodyLimit::max(64)); + assert_eq!(resolve_limit(&ext), 64); + + ext.insert(DefaultBodyLimit::disabled()); + assert_eq!(resolve_limit(&ext), usize::MAX); + } +} diff --git a/arvik-core/src/error.rs b/arvik-core/src/error.rs index 6638741..50647f7 100644 --- a/arvik-core/src/error.rs +++ b/arvik-core/src/error.rs @@ -184,14 +184,17 @@ impl From for Error { } impl From for Error { - fn from(msg: String) -> Self { - Self::new(msg.clone()).with_message(msg) + fn from(err: String) -> Self { + // Internal by default: interpolated strings frequently embed driver + // messages, SQL fragments, or paths. Call `.public()` on the error to + // intentionally expose text to clients. + Self::new(err) } } impl From<&'static str> for Error { - fn from(msg: &'static str) -> Self { - Self::new(msg).with_message(msg) + fn from(err: &'static str) -> Self { + Self::new(err) } } diff --git a/arvik-core/src/extract.rs b/arvik-core/src/extract.rs index b1de0f0..72bee81 100644 --- a/arvik-core/src/extract.rs +++ b/arvik-core/src/extract.rs @@ -261,25 +261,107 @@ impl FromRequest for Request { } /// Extract the raw body as [`bytes::Bytes`]. +/// +/// The body is buffered under the request's [`crate::body_limit::DefaultBodyLimit`] +/// (2 MiB by default); larger bodies are rejected with `413`. impl FromRequest for bytes::Bytes { type Rejection = crate::Error; async fn from_request(req: Request, _state: &S) -> Result { - req.into_body().to_bytes().await.map_err(crate::Error::new) + let limit = crate::body_limit::resolve_limit(req.extensions()); + req.into_body() + .to_bytes_limited(limit) + .await + .map_err(body_limit_rejection) } } /// Extract the raw body as a UTF-8 [`String`]. +/// +/// The body is buffered under the request's [`crate::body_limit::DefaultBodyLimit`] +/// (2 MiB by default); larger bodies are rejected with `413`. impl FromRequest for String { type Rejection = crate::Error; async fn from_request(req: Request, _state: &S) -> Result { + let limit = crate::body_limit::resolve_limit(req.extensions()); let bytes = req .into_body() - .to_bytes() + .to_bytes_limited(limit) .await - .map_err(crate::Error::new)?; + .map_err(body_limit_rejection)?; String::from_utf8(bytes.to_vec()) .map_err(|e| crate::Error::new(e).with_status(http::StatusCode::BAD_REQUEST)) } } + +/// Map a limited-body read failure onto an [`crate::Error`] response: +/// `413 Payload Too Large` when the limit was hit, `500` on read errors. +fn body_limit_rejection(e: crate::body_limit::BodyLimitError) -> crate::Error { + if e.is_payload_too_large() { + return crate::Error::new(e).with_status(http::StatusCode::PAYLOAD_TOO_LARGE); + } + crate::Error::new(e) +} + +#[cfg(test)] +mod body_limit_tests { + use super::*; + use crate::body::Body; + use crate::body_limit::{DEFAULT_BODY_LIMIT, DefaultBodyLimit}; + use http::StatusCode; + + fn post_request(body: Body) -> Request { + Request::new( + http::Request::builder() + .method(http::Method::POST) + .uri("/") + .body(body) + .unwrap(), + ) + } + + #[tokio::test] + async fn bytes_extractor_enforces_default_body_limit() { + let req = post_request(Body::from_bytes(vec![0u8; DEFAULT_BODY_LIMIT + 1].into())); + let err = >::from_request(req, &()) + .await + .unwrap_err(); + assert_eq!(err.into_response().status(), StatusCode::PAYLOAD_TOO_LARGE); + } + + #[tokio::test] + async fn bytes_extractor_accepts_body_at_limit() { + let req = post_request(Body::from_bytes(vec![0u8; DEFAULT_BODY_LIMIT].into())); + let bytes = >::from_request(req, &()) + .await + .unwrap(); + assert_eq!(bytes.len(), DEFAULT_BODY_LIMIT); + } + + #[tokio::test] + async fn default_body_limit_extension_overrides_the_default() { + let mut req = post_request(Body::from_bytes(vec![0u8; 16].into())); + req.extensions_mut().insert(DefaultBodyLimit::max(8)); + let err = >::from_request(req, &()) + .await + .unwrap_err(); + assert_eq!(err.into_response().status(), StatusCode::PAYLOAD_TOO_LARGE); + + let mut req = post_request(Body::from_bytes(vec![0u8; 16].into())); + req.extensions_mut().insert(DefaultBodyLimit::disabled()); + let bytes = >::from_request(req, &()) + .await + .unwrap(); + assert_eq!(bytes.len(), 16); + } + + #[tokio::test] + async fn string_extractor_rejects_oversized_body() { + let req = post_request(Body::from_bytes(vec![b'a'; DEFAULT_BODY_LIMIT + 1].into())); + let err = >::from_request(req, &()) + .await + .unwrap_err(); + assert_eq!(err.into_response().status(), StatusCode::PAYLOAD_TOO_LARGE); + } +} diff --git a/arvik-core/src/into_response.rs b/arvik-core/src/into_response.rs index 755e76e..fec122a 100644 --- a/arvik-core/src/into_response.rs +++ b/arvik-core/src/into_response.rs @@ -255,8 +255,19 @@ where let (headers, body) = self; let mut r = body.into_response(); for (key, value) in headers { - if let (Ok(name), Ok(val)) = (key.try_into(), value.try_into()) { - r.headers_mut().insert(name, val); + match (key.try_into(), value.try_into()) { + (Ok(name), Ok(val)) => { + r.headers_mut().insert(name, val); + } + (Err(e), _) => tracing::warn!( + header = %std::any::type_name::(), + error = ?e, + "response tuple: invalid header name; header dropped" + ), + (_, Err(e)) => tracing::warn!( + error = ?e, + "response tuple: invalid header value; header dropped" + ), } } r @@ -280,8 +291,19 @@ where let mut r = body.into_response(); *r.status_mut() = status; for (key, value) in headers { - if let (Ok(name), Ok(val)) = (key.try_into(), value.try_into()) { - r.headers_mut().insert(name, val); + match (key.try_into(), value.try_into()) { + (Ok(name), Ok(val)) => { + r.headers_mut().insert(name, val); + } + (Err(e), _) => tracing::warn!( + header = %std::any::type_name::(), + error = ?e, + "response tuple: invalid header name; header dropped" + ), + (_, Err(e)) => tracing::warn!( + error = ?e, + "response tuple: invalid header value; header dropped" + ), } } r @@ -318,59 +340,6 @@ where } } -#[cfg(test)] -mod tests { - use super::*; - - #[tokio::test] - async fn trailers_are_appended_after_body() { - let mut trailers = HeaderMap::new(); - trailers.insert("x-checksum", "abc123".parse().unwrap()); - - let response = Trailers::new("body", trailers).into_response(); - let collected = response.into_body().collect().await.unwrap(); - let collected_trailers = collected.trailers().cloned().unwrap(); - - assert_eq!(collected_trailers.get("x-checksum").unwrap(), "abc123"); - assert_eq!(collected.to_bytes(), Bytes::from_static(b"body")); - } - - #[tokio::test] - async fn json_response_serializes_to_bytes() { - let response = Json(serde_json::json!({ "name": "Arvik" })).into_response(); - assert_eq!(response.status(), StatusCode::OK); - assert_eq!( - response.headers().get(http::header::CONTENT_TYPE).unwrap(), - "application/json" - ); - assert_eq!( - response.into_body().to_bytes().await.unwrap(), - Bytes::from_static(br#"{"name":"Arvik"}"#) - ); - } - - #[tokio::test] - async fn json_response_handles_serialization_failure() { - struct Fails; - - impl serde::Serialize for Fails { - fn serialize(&self, _serializer: S) -> Result - where - S: serde::Serializer, - { - Err(serde::ser::Error::custom("nope")) - } - } - - let response = Json(Fails).into_response(); - assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR); - assert_eq!( - response.into_body().to_bytes().await.unwrap(), - Bytes::from_static(br#"{"error":"Serialization failed","code":500}"#) - ); - } -} - // --------------------------------------------------------------------------- // (P1, P2, R) — two IntoResponseParts sets + body (0.3.2) // @@ -490,3 +459,56 @@ impl> IntoResponse for Html { .body(Body::from(self.0.into())) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn trailers_are_appended_after_body() { + let mut trailers = HeaderMap::new(); + trailers.insert("x-checksum", "abc123".parse().unwrap()); + + let response = Trailers::new("body", trailers).into_response(); + let collected = response.into_body().collect().await.unwrap(); + let collected_trailers = collected.trailers().cloned().unwrap(); + + assert_eq!(collected_trailers.get("x-checksum").unwrap(), "abc123"); + assert_eq!(collected.to_bytes(), Bytes::from_static(b"body")); + } + + #[tokio::test] + async fn json_response_serializes_to_bytes() { + let response = Json(serde_json::json!({ "name": "Arvik" })).into_response(); + assert_eq!(response.status(), StatusCode::OK); + assert_eq!( + response.headers().get(http::header::CONTENT_TYPE).unwrap(), + "application/json" + ); + assert_eq!( + response.into_body().to_bytes().await.unwrap(), + Bytes::from_static(br#"{"name":"Arvik"}"#) + ); + } + + #[tokio::test] + async fn json_response_handles_serialization_failure() { + struct Fails; + + impl serde::Serialize for Fails { + fn serialize(&self, _serializer: S) -> Result + where + S: serde::Serializer, + { + Err(serde::ser::Error::custom("nope")) + } + } + + let response = Json(Fails).into_response(); + assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR); + assert_eq!( + response.into_body().to_bytes().await.unwrap(), + Bytes::from_static(br#"{"error":"Serialization failed","code":500}"#) + ); + } +} diff --git a/arvik-core/src/into_response_parts.rs b/arvik-core/src/into_response_parts.rs index 0ae4cee..d53787a 100644 --- a/arvik-core/src/into_response_parts.rs +++ b/arvik-core/src/into_response_parts.rs @@ -156,10 +156,22 @@ where fn into_response_parts(self, mut parts: ResponseParts) -> Result { for (key, value) in self.0 { - // Skip silently on invalid header name/value — mirrors - // the existing array-tuple behaviour. - let (Ok(name), Ok(val)) = (key.try_into(), value.try_into()) else { - continue; + let (name, val) = match (key.try_into(), value.try_into()) { + (Ok(name), Ok(val)) => (name, val), + (Err(e), _) => { + tracing::warn!( + error = %e, + "AppendHeaders: invalid header name; header dropped" + ); + continue; + } + (_, Err(e)) => { + tracing::warn!( + error = %e, + "AppendHeaders: invalid header value; header dropped" + ); + continue; + } }; parts.headers_mut().append(name, val); } diff --git a/arvik-core/src/lib.rs b/arvik-core/src/lib.rs index 77fdaa3..819b442 100644 --- a/arvik-core/src/lib.rs +++ b/arvik-core/src/lib.rs @@ -17,7 +17,9 @@ //! - [`Html`] — HTML response type //! - [`Error`] — Framework error type +pub mod accept; pub mod body; +pub mod body_limit; pub mod error; pub mod extract; pub mod handler; @@ -31,6 +33,7 @@ pub mod stream_body; // Re-exports pub use body::Body; +pub use body_limit::{DEFAULT_BODY_LIMIT, DefaultBodyLimit}; pub use error::{Error, ErrorResponse}; pub use extract::{FromRequest, FromRequestParts}; pub use handler::Handler; diff --git a/arvik-core/src/method_filter.rs b/arvik-core/src/method_filter.rs index 4d49643..ae0bfe6 100644 --- a/arvik-core/src/method_filter.rs +++ b/arvik-core/src/method_filter.rs @@ -35,8 +35,13 @@ impl MethodFilter { pub const OPTIONS: Self = Self(1 << 6); /// Matches HTTP TRACE requests. pub const TRACE: Self = Self(1 << 7); - /// Matches any HTTP method. - pub const ANY: Self = Self(0xFF); + /// Matches non-standard extension methods (e.g. `PURGE`, `PROPFIND`, + /// `CONNECT`). Produced by [`MethodFilter::from_method`] for any method + /// outside the eight standard ones, so extension requests fall through + /// to `405 Method Not Allowed` instead of matching a standard filter. + pub const EXTENSION: Self = Self(1 << 8); + /// Matches any HTTP method, including non-standard extension methods. + pub const ANY: Self = Self(0xFF | (1 << 8)); /// Matches no HTTP method (empty filter). pub const NONE: Self = Self(0); @@ -66,7 +71,11 @@ impl MethodFilter { http::Method::HEAD => Self::HEAD, http::Method::OPTIONS => Self::OPTIONS, http::Method::TRACE => Self::TRACE, - _ => Self(0), // Unknown methods match nothing + // Extension methods get their own bit so they never match a + // standard-method filter (an empty mask would be a subset of + // every filter and would dispatch to the first registered + // handler). + _ => Self::EXTENSION, } } } @@ -121,4 +130,17 @@ mod tests { MethodFilter::POST ); } + + #[test] + fn extension_methods_never_match_standard_filters() { + let purge = http::Method::from_bytes(b"PURGE").unwrap(); + assert_eq!(MethodFilter::from_method(&purge), MethodFilter::EXTENSION); + // An unknown method must not subset-match a specific-method filter. + assert!(!MethodFilter::GET.contains(MethodFilter::EXTENSION)); + assert!(!MethodFilter::DELETE.contains(MethodFilter::EXTENSION)); + assert!(!MethodFilter::DELETE.matches(&purge)); + // ANY remains a true catch-all, per its contract. + assert!(MethodFilter::ANY.contains(MethodFilter::EXTENSION)); + assert!(MethodFilter::ANY.matches(&purge)); + } } diff --git a/arvik-core/src/request_parts.rs b/arvik-core/src/request_parts.rs index 4bf1ea0..e545918 100644 --- a/arvik-core/src/request_parts.rs +++ b/arvik-core/src/request_parts.rs @@ -2,7 +2,7 @@ //! //! [`RequestParts`] combines the standard HTTP request parts with //! framework-specific extensions. This is the input type for -//! [`FromRequestParts`] extractors that need access to headers, +//! `FromRequestParts` extractors that need access to headers, //! URI, method, and framework extensions without consuming the body. use http::Extensions; diff --git a/arvik-core/src/response.rs b/arvik-core/src/response.rs index 830584b..632b0f1 100644 --- a/arvik-core/src/response.rs +++ b/arvik-core/src/response.rs @@ -74,17 +74,38 @@ impl ResponseBuilder { } /// Build the response with the given body. + /// + /// Errors deferred by the builder (e.g. a header value that failed + /// validation — a user-controlled redirect target containing control + /// characters, for instance) are surfaced as a plain 500 instead of + /// panicking the connection task. pub fn body(self, body: impl Into) -> Response { - self.inner.body(body.into()).expect("valid response") + match self.inner.body(body.into()) { + Ok(response) => response, + Err(err) => { + tracing::error!("invalid response construction: {err}"); + internal_error_response() + } + } } /// Build a JSON response. /// /// Serializes `data` as JSON, sets `Content-Type: application/json`, - /// and returns the response. + /// and returns the response. Serialization failures (non-finite floats, + /// non-string map keys, …) return a 500 like [`crate::Json`] instead of + /// panicking. pub fn json(self, data: &T) -> Response { let mut writer = BytesMut::with_capacity(128).writer(); - serde_json::to_writer(&mut writer, data).expect("valid JSON serialization"); + if let Err(err) = serde_json::to_writer(&mut writer, data) { + tracing::error!("response JSON serialization failed: {err}"); + return ResponseBuilder::new() + .status(StatusCode::INTERNAL_SERVER_ERROR) + .header(http::header::CONTENT_TYPE, "application/json") + .body(Body::from( + r#"{"error":"Internal Server Error","code":500}"#, + )); + } let json_bytes = writer.into_inner().freeze(); self.header(http::header::CONTENT_TYPE, "application/json") .body(Body::from_bytes(json_bytes)) @@ -118,6 +139,15 @@ impl Default for ResponseBuilder { } } +/// Static 500 used when a response could not be constructed. +fn internal_error_response() -> Response { + http::Response::builder() + .status(StatusCode::INTERNAL_SERVER_ERROR) + .header(http::header::CONTENT_TYPE, "text/plain; charset=utf-8") + .body(Body::from("Internal Server Error")) + .expect("static fallback response is valid") +} + /// Create a redirect response. /// /// # Variants diff --git a/arvik-extract/src/form.rs b/arvik-extract/src/form.rs index 1dfb471..bd963ad 100644 --- a/arvik-extract/src/form.rs +++ b/arvik-extract/src/form.rs @@ -44,12 +44,13 @@ where return Err(FormRejection::InvalidContentType); } - // Read body - let body_bytes = req - .into_body() - .to_bytes() - .await - .map_err(|e| FormRejection::BodyReadFailed(e.to_string()))?; + // Read body under the configured size limit (2 MiB by default) + let limit = arvik_core::body_limit::resolve_limit(req.extensions()); + let body_bytes = match req.into_body().to_bytes_limited(limit).await { + Ok(bytes) => bytes, + Err(e) if e.is_payload_too_large() => return Err(FormRejection::PayloadTooLarge), + Err(e) => return Err(FormRejection::BodyReadFailed(e.to_string())), + }; // Deserialize let value = serde_urlencoded::from_bytes(&body_bytes) @@ -71,7 +72,60 @@ fn form_content_type(headers: &http::HeaderMap) -> bool { Err(_) => return false, }; - content_type - .to_lowercase() - .starts_with("application/x-www-form-urlencoded") + // Parse as a media type so parameter suffixes (`;charset=utf-8`) pass but + // bogus extensions (`…urlencodedfoo`) do not. + match content_type.parse::() { + Ok(mime) => mime.type_() == mime::APPLICATION && mime.subtype() == "x-www-form-urlencoded", + Err(_) => false, + } +} + +#[cfg(test)] +mod body_limit_tests { + use super::*; + use arvik_core::Body; + use arvik_core::body_limit::{DEFAULT_BODY_LIMIT, DefaultBodyLimit}; + + fn form_request(body: Body) -> Request { + Request::new( + http::Request::builder() + .method(http::Method::POST) + .uri("/") + .header( + http::header::CONTENT_TYPE, + "application/x-www-form-urlencoded", + ) + .body(body) + .unwrap(), + ) + } + + #[tokio::test] + async fn oversized_form_is_rejected_with_413() { + let big = "k=v&".repeat(DEFAULT_BODY_LIMIT / 2); + let req = form_request(Body::from_bytes(big.into_bytes().into())); + match Form::>::from_request(req, &()).await { + Err(FormRejection::PayloadTooLarge) => {} + other => panic!("expected PayloadTooLarge, got {:?}", other.err()), + } + } + + #[tokio::test] + async fn form_within_limit_extracts() { + let req = form_request(Body::from_bytes(b"user=darshan".to_vec().into())); + let Form(value) = Form::>::from_request(req, &()) + .await + .unwrap(); + assert_eq!(value["user"], "darshan"); + } + + #[tokio::test] + async fn default_body_limit_extension_tightens_form_limit() { + let mut req = form_request(Body::from_bytes(b"user=darshan".to_vec().into())); + req.extensions_mut().insert(DefaultBodyLimit::max(4)); + match Form::>::from_request(req, &()).await { + Err(FormRejection::PayloadTooLarge) => {} + other => panic!("expected PayloadTooLarge, got {:?}", other.err()), + } + } } diff --git a/arvik-extract/src/json.rs b/arvik-extract/src/json.rs index 62a7dca..7fa8f98 100644 --- a/arvik-extract/src/json.rs +++ b/arvik-extract/src/json.rs @@ -55,12 +55,13 @@ where return Err(JsonRejection::MissingJsonContentType); } - // Read body - let body_bytes = req - .into_body() - .to_bytes() - .await - .map_err(|e| JsonRejection::BodyReadFailed(e.to_string()))?; + // Read body under the configured size limit (2 MiB by default) + let limit = arvik_core::body_limit::resolve_limit(req.extensions()); + let body_bytes = match req.into_body().to_bytes_limited(limit).await { + Ok(bytes) => bytes, + Err(e) if e.is_payload_too_large() => return Err(JsonRejection::PayloadTooLarge), + Err(e) => return Err(JsonRejection::BodyReadFailed(e.to_string())), + }; // Deserialize let value = serde_json::from_slice(&body_bytes) @@ -110,3 +111,50 @@ fn json_content_type(headers: &http::HeaderMap) -> bool { && (mime.subtype() == mime::JSON || mime.suffix().is_some_and(|suffix| suffix == mime::JSON)) } + +#[cfg(test)] +mod body_limit_tests { + use super::*; + use arvik_core::Body; + use arvik_core::body_limit::{DEFAULT_BODY_LIMIT, DefaultBodyLimit}; + + fn json_request(body: Body) -> Request { + Request::new( + http::Request::builder() + .method(http::Method::POST) + .uri("/") + .header(http::header::CONTENT_TYPE, "application/json") + .body(body) + .unwrap(), + ) + } + + #[tokio::test] + async fn oversized_json_is_rejected_with_413() { + let big = format!("[{}]", "0,".repeat(DEFAULT_BODY_LIMIT / 2)); + let req = json_request(Body::from_bytes(big.into_bytes().into())); + match Json::::from_request(req, &()).await { + Err(JsonRejection::PayloadTooLarge) => {} + other => panic!("expected PayloadTooLarge, got {:?}", other.err()), + } + } + + #[tokio::test] + async fn json_within_limit_extracts() { + let req = json_request(Body::from_bytes(br#"{"ok":true}"#.to_vec().into())); + let Json(value) = Json::::from_request(req, &()) + .await + .unwrap(); + assert_eq!(value["ok"], serde_json::Value::Bool(true)); + } + + #[tokio::test] + async fn default_body_limit_extension_raises_json_limit() { + let mut req = json_request(Body::from_bytes(br#"{"ok":1}"#.to_vec().into())); + req.extensions_mut().insert(DefaultBodyLimit::max(4)); + match Json::::from_request(req, &()).await { + Err(JsonRejection::PayloadTooLarge) => {} + other => panic!("expected PayloadTooLarge, got {:?}", other.err()), + } + } +} diff --git a/arvik-extract/src/lib.rs b/arvik-extract/src/lib.rs index 58bf1f9..ceff7ac 100644 --- a/arvik-extract/src/lib.rs +++ b/arvik-extract/src/lib.rs @@ -3,7 +3,7 @@ //! Request extractors for the Arvik web framework. //! //! This crate provides all built-in extractors that implement -//! [`FromRequestParts`] and [`FromRequest`] from `arvik-core`: +//! `FromRequestParts` and `FromRequest` from `arvik-core`: //! //! ## Parts-Only Extractors (`FromRequestParts`) //! diff --git a/arvik-extract/src/path.rs b/arvik-extract/src/path.rs index ebd4e7f..3620152 100644 --- a/arvik-extract/src/path.rs +++ b/arvik-extract/src/path.rs @@ -62,10 +62,8 @@ where .get::() .ok_or(PathRejection::MissingPathParams)?; - let pairs: Vec<(String, String)> = path_params - .iter() - .map(|(k, v)| (k.to_string(), v.to_string())) - .collect(); + // Borrow the params: no per-key/per-value String clones here. + let pairs: Vec<(&str, &str)> = path_params.iter().collect(); let deserializer = PathDeserializer::new(&pairs); let value = T::deserialize(deserializer) diff --git a/arvik-extract/src/path_de.rs b/arvik-extract/src/path_de.rs index 9a99508..812fcb4 100644 --- a/arvik-extract/src/path_de.rs +++ b/arvik-extract/src/path_de.rs @@ -13,11 +13,11 @@ use std::fmt; /// A serde deserializer for path parameters. pub(crate) struct PathDeserializer<'de> { - params: &'de [(String, String)], + params: &'de [(&'de str, &'de str)], } impl<'de> PathDeserializer<'de> { - pub(crate) fn new(params: &'de [(String, String)]) -> Self { + pub(crate) fn new(params: &'de [(&'de str, &'de str)]) -> Self { Self { params } } } @@ -27,7 +27,7 @@ macro_rules! delegate_single { $( fn $method>(self, visitor: V) -> Result { if self.params.len() == 1 { - ValueDeserializer(self.params[0].1.clone()).$method(visitor) + ValueDeserializer(self.params[0].1).$method(visitor) } else { Err(PathDeserializeError::custom(concat!( "expected single path parameter for ", @@ -44,7 +44,7 @@ impl<'de> Deserializer<'de> for PathDeserializer<'de> { fn deserialize_any>(self, visitor: V) -> Result { if self.params.len() == 1 { - visitor.visit_string(self.params[0].1.clone()) + visitor.visit_str(self.params[0].1) } else { self.deserialize_map(visitor) } @@ -99,7 +99,7 @@ impl<'de> Deserializer<'de> for PathDeserializer<'de> { visitor: V, ) -> Result { if self.params.len() == 1 { - visitor.visit_enum(self.params[0].1.clone().into_enum_deserializer()) + visitor.visit_enum(StringEnumDeserializer(self.params[0].1)) } else { Err(PathDeserializeError::custom( "enums can only be extracted from a single path parameter", @@ -109,7 +109,7 @@ impl<'de> Deserializer<'de> for PathDeserializer<'de> { fn deserialize_option>(self, visitor: V) -> Result { if self.params.len() == 1 { - ValueDeserializer(self.params[0].1.clone()).deserialize_option(visitor) + ValueDeserializer(self.params[0].1).deserialize_option(visitor) } else { Err(PathDeserializeError::custom( "expected single path parameter for option", @@ -134,12 +134,12 @@ impl<'de> Deserializer<'de> for PathDeserializer<'de> { // --------------------------------------------------------------------------- struct PathMapAccess<'de> { - params: &'de [(String, String)], + params: &'de [(&'de str, &'de str)], index: usize, } impl<'de> PathMapAccess<'de> { - fn new(params: &'de [(String, String)]) -> Self { + fn new(params: &'de [(&'de str, &'de str)]) -> Self { Self { params, index: 0 } } } @@ -155,7 +155,7 @@ impl<'de> MapAccess<'de> for PathMapAccess<'de> { return Ok(None); } let key = &self.params[self.index].0; - seed.deserialize(key.as_str().into_deserializer()).map(Some) + seed.deserialize(key.into_deserializer()).map(Some) } fn next_value_seed>( @@ -164,7 +164,7 @@ impl<'de> MapAccess<'de> for PathMapAccess<'de> { ) -> Result { let value = &self.params[self.index].1; self.index += 1; - seed.deserialize(ValueDeserializer(value.clone())) + seed.deserialize(ValueDeserializer(value)) } } @@ -173,12 +173,12 @@ impl<'de> MapAccess<'de> for PathMapAccess<'de> { // --------------------------------------------------------------------------- struct PathSeqAccess<'de> { - params: &'de [(String, String)], + params: &'de [(&'de str, &'de str)], index: usize, } impl<'de> PathSeqAccess<'de> { - fn new(params: &'de [(String, String)]) -> Self { + fn new(params: &'de [(&'de str, &'de str)]) -> Self { Self { params, index: 0 } } } @@ -195,7 +195,7 @@ impl<'de> SeqAccess<'de> for PathSeqAccess<'de> { } let value = &self.params[self.index].1; self.index += 1; - seed.deserialize(ValueDeserializer(value.clone())).map(Some) + seed.deserialize(ValueDeserializer(value)).map(Some) } } @@ -203,17 +203,17 @@ impl<'de> SeqAccess<'de> for PathSeqAccess<'de> { // Value deserializer (parses a single string value into typed values) // --------------------------------------------------------------------------- -struct ValueDeserializer(String); +struct ValueDeserializer<'de>(&'de str); -impl<'de> Deserializer<'de> for ValueDeserializer { +impl<'de> Deserializer<'de> for ValueDeserializer<'de> { type Error = PathDeserializeError; fn deserialize_any>(self, visitor: V) -> Result { - visitor.visit_string(self.0) + visitor.visit_str(self.0) } fn deserialize_bool>(self, visitor: V) -> Result { - match self.0.as_str() { + match self.0 { "true" | "1" => visitor.visit_bool(true), "false" | "0" => visitor.visit_bool(false), _ => Err(PathDeserializeError::custom(format!( @@ -304,11 +304,11 @@ impl<'de> Deserializer<'de> for ValueDeserializer { } fn deserialize_string>(self, visitor: V) -> Result { - visitor.visit_string(self.0) + visitor.visit_str(self.0) } fn deserialize_str>(self, visitor: V) -> Result { - visitor.visit_string(self.0) + visitor.visit_str(self.0) } fn deserialize_char>(self, visitor: V) -> Result { @@ -340,7 +340,7 @@ impl<'de> Deserializer<'de> for ValueDeserializer { _variants: &'static [&'static str], visitor: V, ) -> Result { - visitor.visit_enum(self.0.into_enum_deserializer()) + visitor.visit_enum(StringEnumDeserializer(self.0)) } forward_to_deserialize_any! { @@ -353,9 +353,9 @@ impl<'de> Deserializer<'de> for ValueDeserializer { // String-as-enum deserializer // --------------------------------------------------------------------------- -struct StringEnumDeserializer(String); +struct StringEnumDeserializer<'de>(&'de str); -impl<'de> de::EnumAccess<'de> for StringEnumDeserializer { +impl<'de> de::EnumAccess<'de> for StringEnumDeserializer<'de> { type Error = PathDeserializeError; type Variant = UnitVariant; @@ -405,20 +405,6 @@ impl<'de> de::VariantAccess<'de> for UnitVariant { } } -// --------------------------------------------------------------------------- -// IntoDeserializer helpers -// --------------------------------------------------------------------------- - -trait IntoEnumDeserializer { - fn into_enum_deserializer(self) -> StringEnumDeserializer; -} - -impl IntoEnumDeserializer for String { - fn into_enum_deserializer(self) -> StringEnumDeserializer { - StringEnumDeserializer(self) - } -} - // Simple string deserializer for map keys struct StrDeserializer<'a>(&'a str); diff --git a/arvik-extract/src/rejection.rs b/arvik-extract/src/rejection.rs index c98cdd6..efd5a85 100644 --- a/arvik-extract/src/rejection.rs +++ b/arvik-extract/src/rejection.rs @@ -36,8 +36,15 @@ impl std::fmt::Display for PathRejection { impl IntoResponse for PathRejection { fn into_response(self) -> Response { + // MissingPathParams means routing did not attach params — a server-side + // wiring mistake (fallback handler / middleware outside routing), not a + // client error. + let status = match &self { + Self::MissingPathParams => StatusCode::INTERNAL_SERVER_ERROR, + Self::DeserializationFailed(_) => StatusCode::BAD_REQUEST, + }; ResponseBuilder::new() - .status(StatusCode::BAD_REQUEST) + .status(status) .header(http::header::CONTENT_TYPE, "text/plain; charset=utf-8") .body(arvik_core::Body::from(self.to_string())) } @@ -87,6 +94,9 @@ pub enum JsonRejection { MissingJsonContentType, /// JSON deserialization failed. DeserializationFailed(String), + /// Request body exceeded the configured size limit + /// ([`DefaultBodyLimit`](arvik_core::body_limit::DefaultBodyLimit), 2 MiB by default). + PayloadTooLarge, } impl std::fmt::Display for JsonRejection { @@ -97,6 +107,9 @@ impl std::fmt::Display for JsonRejection { write!(f, "Expected Content-Type: application/json") } Self::DeserializationFailed(msg) => write!(f, "Invalid JSON: {msg}"), + Self::PayloadTooLarge => { + write!(f, "Request payload exceeds the maximum allowed size") + } } } } @@ -105,6 +118,7 @@ impl IntoResponse for JsonRejection { fn into_response(self) -> Response { let status = match &self { JsonRejection::MissingJsonContentType => StatusCode::UNSUPPORTED_MEDIA_TYPE, + JsonRejection::PayloadTooLarge => StatusCode::PAYLOAD_TOO_LARGE, _ => StatusCode::BAD_REQUEST, }; ResponseBuilder::new() @@ -127,6 +141,9 @@ pub enum FormRejection { InvalidContentType, /// Form deserialization failed. DeserializationFailed(String), + /// Request body exceeded the configured size limit + /// ([`DefaultBodyLimit`](arvik_core::body_limit::DefaultBodyLimit), 2 MiB by default). + PayloadTooLarge, } impl std::fmt::Display for FormRejection { @@ -140,6 +157,9 @@ impl std::fmt::Display for FormRejection { ) } Self::DeserializationFailed(msg) => write!(f, "Invalid form data: {msg}"), + Self::PayloadTooLarge => { + write!(f, "Request payload exceeds the maximum allowed size") + } } } } @@ -148,6 +168,7 @@ impl IntoResponse for FormRejection { fn into_response(self) -> Response { let status = match &self { FormRejection::InvalidContentType => StatusCode::UNSUPPORTED_MEDIA_TYPE, + FormRejection::PayloadTooLarge => StatusCode::PAYLOAD_TOO_LARGE, _ => StatusCode::BAD_REQUEST, }; ResponseBuilder::new() diff --git a/arvik-hyper/src/config.rs b/arvik-hyper/src/config.rs index eabd39f..66afa68 100644 --- a/arvik-hyper/src/config.rs +++ b/arvik-hyper/src/config.rs @@ -27,6 +27,11 @@ pub struct ServerConfig { http1_title_case_headers: bool, http1_preserve_header_case: bool, http1_pipeline_flush: bool, + /// Upper bound on reading request headers (slowloris hardening). + http1_header_read_timeout: Option, + http1_max_buf_size: Option, + /// Upper bound on the TLS handshake phase (`None` disables the guard). + handshake_timeout: Option, http2_only: bool, http2_adaptive_window: bool, http2_initial_stream_window_size: Option, @@ -43,7 +48,10 @@ impl Default for ServerConfig { fn default() -> Self { Self { max_connections: None, - tcp_nodelay: None, + // Nagle × delayed-ACK adds tens-of-ms stalls per small response; + // virtually every HTTP server enables NODELAY unconditionally. + // Opt out with .tcp_nodelay(false). + tcp_nodelay: Some(true), tcp_keepalive: None, tcp_keepalive_interval: None, tcp_keepalive_retries: None, @@ -58,6 +66,9 @@ impl Default for ServerConfig { http1_title_case_headers: false, http1_preserve_header_case: false, http1_pipeline_flush: false, + http1_header_read_timeout: None, + http1_max_buf_size: None, + handshake_timeout: Some(Duration::from_secs(10)), http2_only: false, http2_adaptive_window: true, http2_initial_stream_window_size: None, @@ -311,6 +322,34 @@ impl ServerConfig { self } + /// Upper bound on reading request headers. + /// + /// Hardens against slow-loris clients that dribble headers forever + /// (hyper's own default is 30 s; `None` keeps that default). + #[must_use] + pub fn http1_header_read_timeout(mut self, timeout: impl Into>) -> Self { + self.http1_header_read_timeout = timeout.into(); + self + } + + /// Cap the per-connection HTTP/1 read buffer size in bytes. + #[must_use] + pub fn http1_max_buf_size(mut self, size: usize) -> Self { + self.http1_max_buf_size = Some(size); + self + } + + /// Upper bound on the TLS handshake phase (default 10 s). + /// + /// Clients that open a connection but never complete the handshake are + /// disconnected after this duration instead of holding a task and socket + /// indefinitely. Pass `None` to disable the guard entirely. + #[must_use] + pub fn handshake_timeout(mut self, timeout: impl Into>) -> Self { + self.handshake_timeout = timeout.into(); + self + } + /// Accept HTTP/2 prior-knowledge connections only. #[must_use] pub fn http2_only(mut self, enabled: bool) -> Self { @@ -399,6 +438,19 @@ impl ServerConfig { self.http2_only } + pub(crate) fn http1_header_read_timeout_value(&self) -> Option { + self.http1_header_read_timeout + } + + pub(crate) fn http1_max_buf_size_value(&self) -> Option { + self.http1_max_buf_size + } + + #[cfg_attr(not(any(feature = "tls", feature = "native-tls")), allow(dead_code))] + pub(crate) fn handshake_timeout_value(&self) -> Option { + self.handshake_timeout + } + pub(crate) fn needs_tuned_listener(&self) -> bool { self.reuse_port || self.reuse_address @@ -418,6 +470,12 @@ impl ServerConfig { http1.title_case_headers(self.http1_title_case_headers); http1.preserve_header_case(self.http1_preserve_header_case); http1.pipeline_flush(self.http1_pipeline_flush); + if let Some(timeout) = self.http1_header_read_timeout_value() { + http1.header_read_timeout(timeout); + } + if let Some(size) = self.http1_max_buf_size_value() { + http1.max_buf_size(size); + } http1.timer(TokioTimer::new()); } @@ -512,4 +570,33 @@ mod tests { assert_eq!(high.http2_max_concurrent_streams, Some(2_000)); assert_eq!(high.http2_initial_connection_window_size, Some(8_388_608)); } + + #[test] + fn handshake_and_http1_timeout_knobs() { + // Defaults: 10 s handshake guard, hyper defaults for HTTP/1 knobs. + let config = ServerConfig::new(); + assert_eq!( + config.handshake_timeout_value(), + Some(Duration::from_secs(10)) + ); + assert_eq!(config.http1_header_read_timeout_value(), None); + assert_eq!(config.http1_max_buf_size_value(), None); + + let tuned = ServerConfig::new() + .handshake_timeout(Duration::from_secs(3)) + .http1_header_read_timeout(Duration::from_secs(15)) + .http1_max_buf_size(64 * 1024); + assert_eq!( + tuned.handshake_timeout_value(), + Some(Duration::from_secs(3)) + ); + assert_eq!( + tuned.http1_header_read_timeout_value(), + Some(Duration::from_secs(15)) + ); + assert_eq!(tuned.http1_max_buf_size_value(), Some(64 * 1024)); + + let disabled = ServerConfig::new().handshake_timeout(None); + assert_eq!(disabled.handshake_timeout_value(), None); + } } diff --git a/arvik-hyper/src/graceful.rs b/arvik-hyper/src/graceful.rs index 2b06010..f960c1e 100644 --- a/arvik-hyper/src/graceful.rs +++ b/arvik-hyper/src/graceful.rs @@ -74,14 +74,17 @@ impl ShutdownConfig { } pub(crate) fn call_connected(&self, info: ConnectionInfo) { - if let Some(hook) = &self.on_connected { - hook(info); + if let Some(hook) = self.on_connected.clone() { + // Fire-and-forget: hooks are arbitrary user closures, and a + // blocking one must not stall the single accept loop. + tokio::spawn(async move { hook(info) }); } } pub(crate) fn call_disconnected(&self, info: ConnectionInfo) { - if let Some(hook) = &self.on_disconnected { - hook(info); + if let Some(hook) = self.on_disconnected.clone() { + // Same as above — invoked from ConnectionGuard::drop. + tokio::spawn(async move { hook(info) }); } } } diff --git a/arvik-hyper/src/lib.rs b/arvik-hyper/src/lib.rs index c00dace..03e9f48 100644 --- a/arvik-hyper/src/lib.rs +++ b/arvik-hyper/src/lib.rs @@ -5,7 +5,7 @@ //! This crate provides: //! - TCP listener and connection management //! - Hyper service integration -//! - Handler-based serving ([`serve`]) +//! - Handler-based serving ([`serve()`](crate::serve())) //! - Method router-based serving ([`serve_router`]) //! - Graceful shutdown support //! diff --git a/arvik-hyper/src/server.rs b/arvik-hyper/src/server.rs index 61f6c58..e3c4697 100644 --- a/arvik-hyper/src/server.rs +++ b/arvik-hyper/src/server.rs @@ -9,6 +9,7 @@ use std::pin::Pin; use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; use std::task::{Context, Poll}; +use std::time::Duration; use arvik_core::Body; use arvik_core::handler::Handler; @@ -20,7 +21,7 @@ use hyper::service::service_fn; use hyper_util::rt::TokioIo; use socket2::{Domain, Protocol, Socket, TcpKeepalive, Type}; use tokio::net::{TcpListener, TcpStream, lookup_host}; -use tokio::sync::mpsc; +use tokio::sync::{mpsc, watch}; use tokio::task::JoinSet; use tower_service::Service as _; @@ -310,6 +311,7 @@ impl Server { let mut accept_workers = spawn_accept_workers(self.listeners); let mut accepted = accept_workers.receiver; let active = Arc::new(AtomicUsize::new(0)); + let (conn_shutdown_tx, conn_shutdown_rx) = watch::channel(()); loop { tokio::select! { @@ -317,6 +319,9 @@ impl Server { _ = &mut signal => { tracing::info!("Shutdown signal received; stopping listener on {}", self.addr); + // Tell every established connection to finish its current + // request and close instead of serving new ones. + let _ = conn_shutdown_tx.send(()); break; } joined = connections.join_next(), if !connections.is_empty() => { @@ -326,11 +331,19 @@ impl Server { let Some(next) = next else { break; }; - let (stream, peer_addr) = next?; - apply_stream_options(&stream, &self.config)?; + let Some((stream, peer_addr)) = handle_accept_result(next, self.addr).await? else { + continue; + }; + if let Err(err) = apply_stream_options(&stream, &self.config) { + tracing::warn!(peer = %peer_addr, error = %err, + "dropping connection: failed to apply TCP options"); + continue; + } let info = ConnectionInfo { local_addr: self.addr, peer_addr }; if !try_admit_connection(&self.config, &active, info) { - drop(stream); + // Best-effort 503 so clients and load balancers see an + // HTTP signal instead of a bare connection reset. + reject_over_capacity(stream).await; continue; } @@ -343,9 +356,10 @@ impl Server { tracing::debug!("Accepted connection from {}", peer_addr); shutdown_config.call_connected(info); + let shutdown_rx = conn_shutdown_rx.clone(); connections.spawn(async move { let _guard = ConnectionGuard::new(active, shutdown_config, info); - run_connection(io, service, config, peer_addr).await; + run_connection(io, service, config, peer_addr, shutdown_rx).await; }); } } @@ -375,6 +389,8 @@ impl Server { let mut accept_workers = spawn_accept_workers(self.listeners); let mut accepted = accept_workers.receiver; let active = Arc::new(AtomicUsize::new(0)); + let handshake_timeout = self.config.handshake_timeout_value(); + let (conn_shutdown_tx, conn_shutdown_rx) = watch::channel(()); loop { tokio::select! { @@ -382,6 +398,9 @@ impl Server { _ = &mut signal => { tracing::info!("Shutdown signal received; stopping TLS listener on {}", self.addr); + // Tell every established connection to finish its current + // request and close instead of serving new ones. + let _ = conn_shutdown_tx.send(()); break; } joined = connections.join_next(), if !connections.is_empty() => { @@ -391,11 +410,19 @@ impl Server { let Some(next) = next else { break; }; - let (stream, peer_addr) = next?; - apply_stream_options(&stream, &self.config)?; + let Some((stream, peer_addr)) = handle_accept_result(next, self.addr).await? else { + continue; + }; + if let Err(err) = apply_stream_options(&stream, &self.config) { + tracing::warn!(peer = %peer_addr, error = %err, + "dropping connection: failed to apply TCP options"); + continue; + } let info = ConnectionInfo { local_addr: self.addr, peer_addr }; if !try_admit_connection(&self.config, &active, info) { - drop(stream); + // Best-effort 503 so clients and load balancers see an + // HTTP signal instead of a bare connection reset. + reject_over_capacity(stream).await; continue; } @@ -408,13 +435,28 @@ impl Server { tracing::debug!("Accepted TLS connection from {}", peer_addr); shutdown_config.call_connected(info); + let shutdown_rx = conn_shutdown_rx.clone(); connections.spawn(async move { let _guard = ConnectionGuard::new(active, shutdown_config, info); - match tls_config.accept(stream).await { - Ok(tls_stream) => { - run_connection(TokioIo::new(tls_stream), service, config, peer_addr).await; + let handshake = async { + match tls_config.accept(stream).await { + Ok(tls_stream) => { + run_connection(TokioIo::new(tls_stream), service, config, peer_addr, shutdown_rx).await; + } + Err(err) => tracing::warn!("TLS handshake failed from {}: {}", peer_addr, err), + } + }; + match handshake_timeout { + Some(limit) => { + if tokio::time::timeout(limit, handshake).await.is_err() { + tracing::warn!( + peer = %peer_addr, + timeout = ?limit, + "TLS handshake timed out; dropping connection" + ); + } } - Err(err) => tracing::warn!("TLS handshake failed from {}: {}", peer_addr, err), + None => handshake.await, } }); } @@ -445,6 +487,8 @@ impl Server { let mut accept_workers = spawn_accept_workers(self.listeners); let mut accepted = accept_workers.receiver; let active = Arc::new(AtomicUsize::new(0)); + let handshake_timeout = self.config.handshake_timeout_value(); + let (conn_shutdown_tx, conn_shutdown_rx) = watch::channel(()); loop { tokio::select! { @@ -452,6 +496,9 @@ impl Server { _ = &mut signal => { tracing::info!("Shutdown signal received; stopping native-tls listener on {}", self.addr); + // Tell every established connection to finish its current + // request and close instead of serving new ones. + let _ = conn_shutdown_tx.send(()); break; } joined = connections.join_next(), if !connections.is_empty() => { @@ -461,11 +508,19 @@ impl Server { let Some(next) = next else { break; }; - let (stream, peer_addr) = next?; - apply_stream_options(&stream, &self.config)?; + let Some((stream, peer_addr)) = handle_accept_result(next, self.addr).await? else { + continue; + }; + if let Err(err) = apply_stream_options(&stream, &self.config) { + tracing::warn!(peer = %peer_addr, error = %err, + "dropping connection: failed to apply TCP options"); + continue; + } let info = ConnectionInfo { local_addr: self.addr, peer_addr }; if !try_admit_connection(&self.config, &active, info) { - drop(stream); + // Best-effort 503 so clients and load balancers see an + // HTTP signal instead of a bare connection reset. + reject_over_capacity(stream).await; continue; } @@ -478,15 +533,30 @@ impl Server { tracing::debug!("Accepted native-tls connection from {}", peer_addr); shutdown_config.call_connected(info); + let shutdown_rx = conn_shutdown_rx.clone(); connections.spawn(async move { let _guard = ConnectionGuard::new(active, shutdown_config, info); - match tls_config.accept(stream).await { - Ok(tls_stream) => { - run_connection(TokioIo::new(tls_stream), service, config, peer_addr).await; + let handshake = async { + match tls_config.accept(stream).await { + Ok(tls_stream) => { + run_connection(TokioIo::new(tls_stream), service, config, peer_addr, shutdown_rx).await; + } + Err(err) => { + tracing::warn!("native-tls handshake failed from {}: {}", peer_addr, err) + } } - Err(err) => { - tracing::warn!("native-tls handshake failed from {}: {}", peer_addr, err) + }; + match handshake_timeout { + Some(limit) => { + if tokio::time::timeout(limit, handshake).await.is_err() { + tracing::warn!( + peer = %peer_addr, + timeout = ?limit, + "native-tls handshake timed out; dropping connection" + ); + } } + None => handshake.await, } }); } @@ -589,9 +659,20 @@ fn spawn_accept_workers(listeners: Vec) -> AcceptWorkers { let sender = sender.clone(); tasks.spawn(async move { loop { - let accepted = listener.accept().await; - if sender.send(accepted).await.is_err() { - break; + match listener.accept().await { + Ok(pair) => { + if sender.send(Ok(pair)).await.is_err() { + break; + } + } + Err(err) => { + // Pace persistent failures instead of spinning the + // channel full of errors; the main loop classifies them. + tokio::time::sleep(Duration::from_millis(10)).await; + if sender.send(Err(err)).await.is_err() { + break; + } + } } } }); @@ -601,6 +682,101 @@ fn spawn_accept_workers(listeners: Vec) -> AcceptWorkers { AcceptWorkers { receiver, tasks } } +/// Classify one accept result from the worker channel. +/// +/// Transient per-connection failures (`ECONNABORTED`, `EINTR`, `EWOULDBLOCK`) +/// and resource-exhaustion conditions (fd/memory pressure) are logged and +/// skipped — a single failed accept must never stop the server. Only genuinely +/// fatal socket errors (bad descriptor, invalid argument, permission denied) +/// are propagated to the caller. +async fn handle_accept_result( + result: io::Result<(TcpStream, SocketAddr)>, + listener_addr: SocketAddr, +) -> Result, BoxError> { + let (stream, peer_addr) = match result { + Ok(pair) => pair, + Err(err) => { + match err.kind() { + // Transient per-connection failures — safe to continue. + io::ErrorKind::ConnectionAborted + | io::ErrorKind::WouldBlock + | io::ErrorKind::Interrupted => { + tracing::debug!(listener = %listener_addr, error = %err, + "transient accept error; continuing"); + return Ok(None); + } + _ if is_fatal_accept_error(&err) => return Err(err.into()), + _ => { + // Anything else is treated as (possibly) transient. + // Resource exhaustion gets a short backoff so the OS can + // drain descriptors before the next attempt. + if is_resource_exhaustion(&err) { + tracing::warn!(listener = %listener_addr, error = %err, + "resource exhaustion on accept; backing off 100ms"); + tokio::time::sleep(Duration::from_millis(100)).await; + } else { + tracing::warn!(listener = %listener_addr, error = %err, + "accept error; continuing"); + } + return Ok(None); + } + } + } + }; + Ok(Some((stream, peer_addr))) +} + +/// Errors that indicate the listener itself is broken — serving cannot continue. +fn is_fatal_accept_error(err: &io::Error) -> bool { + const EBADF: i32 = 9; // listener closed / not a socket + const EINVAL: i32 = 22; // listener in an unusable state + + matches!(err.kind(), io::ErrorKind::PermissionDenied) + || matches!(err.raw_os_error(), Some(EBADF) | Some(EINVAL)) +} + +/// fd/memory exhaustion conditions that occur under load (EMFILE/ENFILE/ +/// ENOMEM/ENOBUFS). EMFILE=24 and ENFILE=23 share values on Linux and macOS; +/// ENOBUFS differs (Linux 105, macOS 55). +fn is_resource_exhaustion(err: &io::Error) -> bool { + const EMFILE: i32 = 24; + const ENFILE: i32 = 23; + const ENOMEM: i32 = 12; + + #[cfg(target_os = "linux")] + const ENOBUFS: i32 = 105; + #[cfg(not(target_os = "linux"))] + const ENOBUFS: i32 = 55; + + matches!(err.kind(), io::ErrorKind::OutOfMemory) + || matches!( + err.raw_os_error(), + Some(EMFILE) | Some(ENFILE) | Some(ENOMEM) | Some(ENOBUFS) + ) +} + +/// Write a minimal `503 Service Unavailable` and close. +/// +/// Used when the configured connection limit is exhausted. Never panics and +/// never blocks longer than two seconds — if the write fails, the socket is +/// dropped as before. +async fn reject_over_capacity(mut stream: TcpStream) { + use tokio::io::AsyncWriteExt as _; + + let body = "{\"error\":\"Service Unavailable\",\"code\":503}"; + let response = format!( + "HTTP/1.1 503 Service Unavailable\r\nContent-Type: application/json\r\n Content-Length: {}\r\nConnection: close\r\n\r\n{}", + body.len(), + body + ); + + let _ = tokio::time::timeout(Duration::from_secs(2), async move { + let _ = stream.write_all(response.as_bytes()).await; + let _ = stream.shutdown().await; + }) + .await; +} + fn apply_stream_options(stream: &TcpStream, config: &ServerConfig) -> io::Result<()> { if let Some(enabled) = config.tcp_nodelay_setting() { stream.set_nodelay(enabled)?; @@ -746,6 +922,7 @@ async fn run_connection( service: BoxCloneService, config: ServerConfig, peer_addr: SocketAddr, + mut shutdown_rx: watch::Receiver<()>, ) where I: hyper::rt::Read + hyper::rt::Write + Unpin + Send + 'static, { @@ -763,10 +940,36 @@ async fn run_connection( }); let builder = config.auto_builder(); + + // Serve requests until shutdown is signalled. On signal, tell hyper to + // stop taking NEW requests: idle keep-alive connections close immediately + // and in-flight responses are allowed to finish — instead of serving new + // requests for the whole drain window and then being hard-aborted. + // + // The plain and upgrade-capable connections are distinct types, so each + // branch runs its own shutdown-aware select. let result = if config.is_http2_only() { - builder.serve_connection(io, hyper_svc).await + let conn = builder.serve_connection(io, hyper_svc); + tokio::pin!(conn); + tokio::select! { + biased; + res = &mut conn => res, + _ = shutdown_rx.changed() => { + conn.as_mut().graceful_shutdown(); + (&mut conn).await + } + } } else { - builder.serve_connection_with_upgrades(io, hyper_svc).await + let conn = builder.serve_connection_with_upgrades(io, hyper_svc); + tokio::pin!(conn); + tokio::select! { + biased; + res = &mut conn => res, + _ = shutdown_rx.changed() => { + conn.as_mut().graceful_shutdown(); + (&mut conn).await + } + } }; if let Err(err) = result { @@ -838,6 +1041,48 @@ impl tower_service::Service for MethodRouterService { mod tests { use super::*; + #[tokio::test] + async fn connection_aborted_is_skipped_not_fatal() { + let err = io::Error::new( + io::ErrorKind::ConnectionAborted, + "client reset before accept", + ); + let out = handle_accept_result(Err(err), "127.0.0.1:8080".parse().unwrap()) + .await + .unwrap(); + assert!(out.is_none()); + } + + #[tokio::test] + async fn fd_exhaustion_backs_off_and_skips_instead_of_dying() { + // EMFILE — the classic under-load accept failure. + let err = io::Error::from_raw_os_error(24); + assert!(is_resource_exhaustion(&err)); + let out = handle_accept_result(Err(err), "127.0.0.1:8080".parse().unwrap()) + .await + .unwrap(); + assert!(out.is_none()); + } + + #[tokio::test] + async fn broken_listener_still_stops_the_server() { + // EBADF — the listener itself is gone; serving cannot continue. + let err = io::Error::from_raw_os_error(9); + assert!(is_fatal_accept_error(&err)); + let out = handle_accept_result(Err(err), "127.0.0.1:8080".parse().unwrap()).await; + assert!(out.is_err()); + } + + #[test] + fn fatal_classification_matches_only_broken_listeners() { + assert!(is_fatal_accept_error(&io::Error::from_raw_os_error(22))); // EINVAL + assert!(!is_fatal_accept_error(&io::Error::new( + io::ErrorKind::ConnectionAborted, + "x" + ))); + assert!(!is_fatal_accept_error(&io::Error::other("plain"))); + } + #[tokio::test] async fn default_bind_path_works() { let Some(server) = bind_or_skip(Server::bind("127.0.0.1:0").await) else { diff --git a/arvik-hyper/tests/graceful.rs b/arvik-hyper/tests/graceful.rs index 1ed6bdb..a1d03a3 100644 --- a/arvik-hyper/tests/graceful.rs +++ b/arvik-hyper/tests/graceful.rs @@ -168,3 +168,90 @@ fn is_permission_denied(err: &(dyn std::error::Error + Send + Sync + 'static)) - false } + +async fn read_one_response(stream: &mut tokio::net::TcpStream) -> String { + // Read until end of headers, then exactly Content-Length bytes of body. + let mut buf: Vec = Vec::new(); + let header_end; + loop { + let mut chunk = [0u8; 1024]; + let n = stream.read(&mut chunk).await.unwrap(); + assert!(n > 0, "connection closed before response completed"); + buf.extend_from_slice(&chunk[..n]); + if let Some(pos) = buf.windows(4).position(|w| w == b"\r\n\r\n") { + header_end = pos + 4; + break; + } + } + let headers = String::from_utf8_lossy(&buf[..header_end]).to_string(); + let content_length: usize = headers + .lines() + .find_map(|l| { + let (name, value) = l.split_once(':')?; + name.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse().ok())? + }) + .unwrap_or(0); + while buf.len() < header_end + content_length { + let mut chunk = [0u8; 4096]; + let n = stream.read(&mut chunk).await.unwrap(); + assert!(n > 0, "connection closed mid-body"); + buf.extend_from_slice(&chunk[..n]); + } + String::from_utf8_lossy(&buf).to_string() +} + +#[tokio::test] +async fn idle_keepalive_connection_closes_promptly_on_shutdown() { + let app = Router::new().route("/", get(|| async { "ok" })); + let Some(server) = bind_test_server().await else { + return; + }; + let addr = server.local_addr(); + let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel(); + + // Generous drain window: before the fix, an open idle keep-alive + // connection held `serve_app_with_graceful_shutdown` for this entire + // duration because nothing ever told connections to stop serving. + let task = tokio::spawn(async move { + server + .serve_app_with_graceful_shutdown( + app, + async { + let _ = shutdown_rx.await; + }, + ShutdownConfig::default().drain_timeout(Duration::from_secs(30)), + ) + .await + }); + + // Open an explicit keep-alive connection and complete one request over + // it, then leave it idle. + let mut stream = tokio::net::TcpStream::connect(addr).await.unwrap(); + stream + .write_all(b"GET / HTTP/1.1\r\nHost: localhost\r\nConnection: keep-alive\r\n\r\n") + .await + .unwrap(); + let response = read_one_response(&mut stream).await; + assert!(response.contains("200 OK")); + assert!(response.contains("ok")); + + shutdown_tx.send(()).unwrap(); + let started = std::time::Instant::now(); + + // The idle connection must see EOF promptly... + let mut eof_buf = Vec::new(); + let eof = tokio::time::timeout(Duration::from_secs(5), stream.read_to_end(&mut eof_buf)).await; + assert!( + eof.is_ok(), + "server did not close the idle connection within 5s" + ); + + // ...and the whole drain must finish long before the 30s window. + task.await.unwrap().unwrap(); + assert!( + started.elapsed() < Duration::from_secs(10), + "idle keep-alive connections held the drain for {:?}", + started.elapsed() + ); +} diff --git a/arvik-macros/src/lib.rs b/arvik-macros/src/lib.rs index b94c59a..aadd9bb 100644 --- a/arvik-macros/src/lib.rs +++ b/arvik-macros/src/lib.rs @@ -662,7 +662,9 @@ fn normalize_route_path(path: &str) -> std::result::Result { fn arvik_path() -> TokenStream2 { match proc_macro_crate::crate_name("arvik") { - Ok(proc_macro_crate::FoundCrate::Itself) => quote!(::arvik), + // proc-macro-crate docs: from the defining crate itself, refer via + // `crate` — `::arvik` does not resolve there. + Ok(proc_macro_crate::FoundCrate::Itself) => quote!(crate), Ok(proc_macro_crate::FoundCrate::Name(name)) => { let ident = Ident::new(&name, Span::call_site()); quote!(::#ident) diff --git a/arvik-middleware/Cargo.toml b/arvik-middleware/Cargo.toml index 8b6e22f..f1c4578 100644 --- a/arvik-middleware/Cargo.toml +++ b/arvik-middleware/Cargo.toml @@ -21,6 +21,7 @@ tokio-util = { workspace = true } tokio = { workspace = true } mime = { workspace = true } uuid = { workspace = true } +futures-util = { workspace = true } parking_lot = { workspace = true } base64 = { workspace = true } cookie = { workspace = true } diff --git a/arvik-middleware/src/auth.rs b/arvik-middleware/src/auth.rs index e8988dc..ecefeed 100644 --- a/arvik-middleware/src/auth.rs +++ b/arvik-middleware/src/auth.rs @@ -40,6 +40,25 @@ use http::{HeaderValue, StatusCode, header::AUTHORIZATION}; use tower_layer::Layer; use tower_service::Service; +/// Constant-time equality for secrets (tokens, passwords, CSRF tokens). +/// +/// Accumulates an XOR of every byte pair, so where a mismatch occurs does not +/// change timing. Only the *length* remains observable, which is treated as +/// public information (scheme names like `Bearer ` and standard token formats +/// are not secret). +pub(crate) fn secret_eq(a: &str, b: &str) -> bool { + let a = a.as_bytes(); + let b = b.as_bytes(); + if a.len() != b.len() { + return false; + } + let mut diff = 0u8; + for (x, y) in a.iter().zip(b.iter()) { + diff |= x ^ y; + } + diff == 0 +} + // ── AuthStrategy ────────────────────────────────────────────────────────────── enum AuthStrategy { @@ -56,7 +75,7 @@ impl AuthStrategy { req.headers() .get(AUTHORIZATION) .and_then(|v| v.to_str().ok()) - .map(|v| v == expected.as_str()) + .map(|v| crate::auth::secret_eq(v, &expected)) .unwrap_or(false) } AuthStrategy::Basic { username, password } => { @@ -66,7 +85,7 @@ impl AuthStrategy { req.headers() .get(AUTHORIZATION) .and_then(|v| v.to_str().ok()) - .map(|v| v == expected.as_str()) + .map(|v| crate::auth::secret_eq(v, &expected)) .unwrap_or(false) } AuthStrategy::Custom(f) => f(req), @@ -174,3 +193,45 @@ where }) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn secret_eq_compares_every_byte_without_short_circuit() { + assert!(secret_eq("", "")); + assert!(secret_eq("Bearer abc123", "Bearer abc123")); + // Same length, single byte differs. + assert!(!secret_eq("Bearer abc123", "Bearer abc124")); + assert!(!secret_eq("Bearer abc123", "Xearer abc123")); + } + + #[test] + fn secret_eq_is_length_aware() { + assert!(!secret_eq("abc", "abcd")); + assert!(!secret_eq("abcd", "abc")); + } + + #[test] + fn bearer_strategy_uses_constant_time_compare() { + let strategy = AuthStrategy::Bearer("s3cret-token".to_string()); + let mut req = Request::new( + http::Request::builder() + .uri("/") + .header(AUTHORIZATION, "Bearer s3cret-token") + .body(arvik_core::Body::empty()) + .unwrap(), + ); + assert!(strategy.is_authorized(&req)); + + req = Request::new( + http::Request::builder() + .uri("/") + .header(AUTHORIZATION, "Bearer wrong-token") + .body(arvik_core::Body::empty()) + .unwrap(), + ); + assert!(!strategy.is_authorized(&req)); + } +} diff --git a/arvik-middleware/src/body_limit.rs b/arvik-middleware/src/body_limit.rs index 1c6b36f..afb7bd2 100644 --- a/arvik-middleware/src/body_limit.rs +++ b/arvik-middleware/src/body_limit.rs @@ -1,10 +1,13 @@ //! Request body size limit middleware. //! -//! Enforces a maximum request body size. Checks the `Content-Length` -//! header immediately; for streaming bodies without `Content-Length`, -//! enforces the limit during body collection. +//! Enforces a maximum request body size **without buffering**: requests whose +//! `Content-Length` exceeds the limit are rejected immediately; streaming +//! bodies are passed through wrapped in a counting frame filter that errors +//! once more than `limit` bytes have been read (readers surface that as +//! `413 Payload Too Large`). //! -//! Returns `413 Payload Too Large` when the limit is exceeded. +//! To configure the *extractor-side* limit instead (`Bytes`, `String`, `Json`, +//! `Form`), see [`DefaultBodyLimitLayer`]. //! //! # Example //! @@ -21,9 +24,11 @@ use std::future::Future; use std::pin::Pin; use std::task::{Context, Poll}; +use arvik_core::body_limit::DefaultBodyLimit; use arvik_core::{Body, Request, Response}; use http::StatusCode; -use tokio_util::bytes::{Bytes, BytesMut}; +use http_body::Frame; +use tokio_util::bytes::Bytes; use tower_layer::Layer; use tower_service::Service; @@ -95,58 +100,62 @@ where } } - // For streaming bodies: collect up to limit bytes. + // For streaming bodies: pass the body through unbuffered with a + // counting cap — readers get an error once the limit is crossed. let (parts, body) = req.into_request_parts(); - - match collect_limited(body, limit).await { - Ok(bytes) => { - let req = Request::from_request_parts(parts, Body::from_bytes(bytes)); - inner.call(req).await - } - Err(()) => { - tracing::warn!( - limit = limit, - "Request body exceeds size limit (stream check)" - ); - Ok::(payload_too_large(limit)) - } - } + let limited = Body::new(LimitedBody { + inner: body, + limit: limit as u64, + read: 0, + }); + let req = Request::from_request_parts(parts, limited); + inner.call(req).await }) } } -/// Collect up to `limit` bytes from `body`. -/// -/// Returns `Err(())` if the body exceeds the limit. -/// -/// FIX: replaced the previous `futures_util::future::poll_fn` + manual -/// pin projection with a straightforward loop using `std::future::poll_fn` -/// and `Pin::new_unchecked` — body is `Unpin` so this is safe — or simply -/// box the body to get a stable address. -async fn collect_limited(body: Body, limit: usize) -> Result { - use http_body::Body as HttpBodyTrait; - - // Box the body so we can pin it without needing the caller to be Unpin. - let mut pinned = Box::pin(body); - let mut buf = BytesMut::new(); - - loop { - match std::future::poll_fn(|cx| pinned.as_mut().poll_frame(cx)).await { - Some(Ok(frame)) => { - if let Ok(data) = frame.into_data() { - buf.extend_from_slice(&data); - if buf.len() > limit { - return Err(()); +/// Streaming frame filter that errors once more than `limit` bytes have been +/// read. Frames themselves are passed through untouched — nothing is buffered. +struct LimitedBody { + inner: Body, + limit: u64, + read: u64, +} + +impl http_body::Body for LimitedBody { + type Data = Bytes; + type Error = arvik_core::body::BoxError; + + fn poll_frame( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + let this = &mut *self; + match Pin::new(&mut this.inner).poll_frame(cx) { + Poll::Ready(Some(Ok(frame))) => match frame.into_data() { + Ok(data) => { + this.read += data.len() as u64; + if this.read > this.limit { + return Poll::Ready(Some(Err(Box::new( + arvik_core::body_limit::BodyTooLarge, + )))); } + Poll::Ready(Some(Ok(Frame::data(data)))) } - // Trailers frames are silently skipped. - } - Some(Err(_)) => return Err(()), - None => break, + // Trailers frames pass through unchanged. + Err(frame) => Poll::Ready(Some(Ok(frame))), + }, + other => other, } } - Ok(buf.freeze()) + fn is_end_stream(&self) -> bool { + self.inner.is_end_stream() + } + + fn size_hint(&self) -> http_body::SizeHint { + self.inner.size_hint() + } } fn payload_too_large(limit: usize) -> Response { @@ -159,3 +168,170 @@ fn payload_too_large(limit: usize) -> Response { ))) .unwrap() } + +// ── DefaultBodyLimitLayer ──────────────────────────────────────────────────── +// +// Sets the size limit that *buffering extractors* (Bytes, String, Json, Form) +// enforce, by inserting a `DefaultBodyLimit` extension into each request. +// Unlike `RequestBodyLimitLayer`, this layer reads nothing itself — the +// extractors enforce the limit while collecting. + +/// Tower layer that configures the extractor-side body limit. +/// +/// Buffering extractors accept up to `limit` bytes instead of the 2 MiB +/// default ([`arvik_core::DEFAULT_BODY_LIMIT`]); larger requests are rejected +/// with `413 Payload Too Large` by the extractor itself. +/// +/// ```rust,ignore +/// use arvik_middleware::body_limit::DefaultBodyLimitLayer; +/// +/// // Allow 16 MiB JSON bodies on this route. +/// Router::new() +/// .route("/upload", post(upload)) +/// .layer(DefaultBodyLimitLayer::max(16 * 1024 * 1024)); +/// ``` +#[derive(Debug, Clone, Copy)] +pub struct DefaultBodyLimitLayer { + limit: usize, +} + +impl DefaultBodyLimitLayer { + /// Allow buffering extractors to accept up to `limit` bytes. + pub fn max(limit: usize) -> Self { + Self { limit } + } + + /// Disable the extractor-side limit for wrapped routes. + /// + /// Only sensible behind another protection (e.g. an upstream proxy + /// enforcing its own size cap). + pub fn disabled() -> Self { + Self { limit: usize::MAX } + } +} + +impl Layer for DefaultBodyLimitLayer { + type Service = DefaultBodyLimitService; + + fn layer(&self, inner: S) -> Self::Service { + DefaultBodyLimitService { + inner, + limit: self.limit, + } + } +} + +/// Tower service produced by [`DefaultBodyLimitLayer`]. +#[derive(Clone)] +pub struct DefaultBodyLimitService { + inner: S, + limit: usize, +} + +impl Service for DefaultBodyLimitService +where + S: Service + Clone + Send + 'static, + S::Future: Send + 'static, +{ + type Response = Response; + type Error = Infallible; + type Future = Pin> + Send + 'static>>; + + fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { + self.inner.poll_ready(cx) + } + + fn call(&mut self, mut req: Request) -> Self::Future { + let cloned = self.inner.clone(); + let mut inner = std::mem::replace(&mut self.inner, cloned); + + req.extensions_mut().insert(DefaultBodyLimit(self.limit)); + Box::pin(inner.call(req)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arvik_core::IntoResponse; + use arvik_core::body_limit::DefaultBodyLimit; + use arvik_core::extract::FromRequest; + + /// Inner service that reads the whole body via the `Bytes` extractor and + /// reports its length — with the extractor-side default disabled so this + /// suite exercises *this layer's* enforcement. + #[derive(Clone)] + struct ReadBodyService; + + impl Service for ReadBodyService { + type Response = Response; + type Error = Infallible; + type Future = Pin> + Send + 'static>>; + + fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn call(&mut self, mut req: Request) -> Self::Future { + Box::pin(async move { + req.extensions_mut().insert(DefaultBodyLimit::disabled()); + let resp = match >::from_request(req, &()).await { + Ok(b) => http::Response::builder() + .status(StatusCode::OK) + .body(Body::from(format!("len={}", b.len()))) + .unwrap(), + Err(err) => err.into_response(), + }; + Ok(resp) + }) + } + } + + fn chunked_request(chunks: &[&[u8]]) -> Request { + Request::new( + http::Request::builder() + .method(http::Method::POST) + .uri("/") + .header(http::header::TRANSFER_ENCODING, "chunked") + .body(Body::from_chunks(chunks.iter().map(|c| c.to_vec()))) + .unwrap(), + ) + } + + #[tokio::test] + async fn streams_body_through_intact_under_the_limit() { + let mut svc = RequestBodyLimitLayer::new(64).layer(ReadBodyService); + // Two chunks prove pass-through: nothing is reassembled by the layer. + let res = svc + .call(chunked_request(&[b"hello ", b"world"])) + .await + .unwrap(); + assert_eq!(res.status(), StatusCode::OK); + assert_eq!(Body::to_string(res.into_body()).await.unwrap(), "len=11"); + } + + #[tokio::test] + async fn streaming_body_over_limit_yields_413_from_reader() { + let mut svc = RequestBodyLimitLayer::new(8).layer(ReadBodyService); + let res = svc + .call(chunked_request(&[b"0123456789", b"abcdefghij"])) + .await + .unwrap(); + assert_eq!(res.status(), StatusCode::PAYLOAD_TOO_LARGE); + } + + #[tokio::test] + async fn oversized_content_length_is_rejected_before_reading() { + let mut svc = RequestBodyLimitLayer::new(8).layer(ReadBodyService); + let req = Request::new( + http::Request::builder() + .method(http::Method::POST) + .uri("/") + .header(http::header::CONTENT_LENGTH, "1000") + .body(Body::from(Bytes::from_static(b"tiny"))) + .unwrap(), + ); + let res = svc.call(req).await.unwrap(); + assert_eq!(res.status(), StatusCode::PAYLOAD_TOO_LARGE); + } +} diff --git a/arvik-middleware/src/catch_panic.rs b/arvik-middleware/src/catch_panic.rs index d469389..5291e34 100644 --- a/arvik-middleware/src/catch_panic.rs +++ b/arvik-middleware/src/catch_panic.rs @@ -1,10 +1,12 @@ //! Panic recovery middleware. //! //! Catches panics in handlers and returns a `500 Internal Server Error` -//! response instead of crashing the server task. +//! response instead of crashing the connection task. //! -//! Uses `tokio::task::spawn` under the hood — if the spawned task panics, -//! the `JoinError` is caught and converted to a 500 response. +//! Wraps the handler future in `catch_unwind` directly — zero per-request +//! overhead. Note: panics raised later, while *polling the response body +//! stream*, happen outside this wrapper (the body is driven by hyper after +//! the response is returned). //! //! # Example //! @@ -23,6 +25,7 @@ //! })); //! ``` +use futures_util::FutureExt as _; use std::convert::Infallible; use std::future::Future; use std::pin::Pin; @@ -111,15 +114,19 @@ where let panic_handler = self.panic_handler.clone(); Box::pin(async move { - // Spawn onto a new task so we can catch panics via JoinError - let handle = tokio::task::spawn(async move { inner.call(req).await }); - - match handle.await { + // catch_unwind on the future directly — no per-request task spawn, + // which cost an allocation + scheduling round-trip on 100% of + // requests to handle the rare panic case. + // + // AssertUnwindSafe: the tower service contract (clone-per-call, + // infallible) already treats inner state as recoverable. + match std::panic::AssertUnwindSafe(inner.call(req)) + .catch_unwind() + .await + { Ok(Ok(response)) => Ok(response), Ok(Err(infallible)) => match infallible {}, - Err(join_error) => { - let panic_payload = join_error.into_panic(); - + Err(panic_payload) => { let panic_msg = panic_payload .downcast_ref::<&str>() .copied() diff --git a/arvik-middleware/src/compression.rs b/arvik-middleware/src/compression.rs index 16d25fa..33833f8 100644 --- a/arvik-middleware/src/compression.rs +++ b/arvik-middleware/src/compression.rs @@ -143,21 +143,29 @@ impl CompressionLayer { } fn preferred_encoding(&self, accept_encoding: &str) -> Option { - // Simple preference: zstd > br > gzip > deflate - let lower = accept_encoding.to_lowercase(); - if self.zstd && lower.contains("zstd") { - return Some(Encoding::Zstd); + // Priority order zstd > br > gzip > deflate, negotiated with real + // q-value semantics (substring matching used to serve brotli to + // clients that explicitly refused it with `br;q=0`). + let mut candidates = Vec::with_capacity(4); + if self.zstd { + candidates.push(("zstd", Encoding::Zstd)); } - if self.br && (lower.contains("br") || lower.contains("brotli")) { - return Some(Encoding::Br); + if self.br { + candidates.push(("br", Encoding::Br)); } - if self.gzip && lower.contains("gzip") { - return Some(Encoding::Gzip); + if self.gzip { + candidates.push(("gzip", Encoding::Gzip)); } - if self.deflate && lower.contains("deflate") { - return Some(Encoding::Deflate); + if self.deflate { + candidates.push(("deflate", Encoding::Deflate)); } - None + let names: Vec<&str> = candidates.iter().map(|(name, _)| *name).collect(); + arvik_core::accept::negotiate(&names, accept_encoding).and_then(|chosen| { + candidates + .into_iter() + .find(|(name, _)| *name == chosen) + .map(|(_, encoding)| encoding) + }) } } @@ -204,6 +212,8 @@ where let cloned = self.inner.clone(); let mut inner = std::mem::replace(&mut self.inner, cloned); + let is_head = req.method() == http::Method::HEAD; + Box::pin(async move { let response = inner.call(req).await?; @@ -212,6 +222,11 @@ where return Ok(response); } + // HEAD has no body to compress; keep Content-Length semantics intact. + if is_head { + return Ok(response); + } + // Don't compress non-compressible content types. if !should_compress(response.headers()) { return Ok(response); @@ -234,31 +249,40 @@ async fn compress_response( ) -> Response { let (mut parts, body) = response.into_parts(); - let body_bytes: Bytes = match body.to_bytes().await { - Ok(b) => b, - Err(_) => { - return http::Response::builder() - .status(StatusCode::INTERNAL_SERVER_ERROR) - .body(Body::from("Compression error")) - .unwrap(); - } - }; - - // Skip compression for small bodies. - if body_bytes.len() < config.min_size { + // Skip compression when the size hint proves the body is smaller than + // the threshold — without buffering anything. + if let Some(upper) = http_body::Body::size_hint(&body).upper() + && upper < config.min_size as u64 + { parts .headers .insert(VARY, HeaderValue::from_static("Accept-Encoding")); - return http::Response::from_parts(parts, Body::from_bytes(body_bytes)); + return http::Response::from_parts(parts, body); } - // FIX: `compress_bytes` always returns `Vec` (either compressed or - // a copy of the original bytes as Vec). Previously the fallback arm - // returned `body_bytes` (type `Bytes`) while the Ok arm returned - // `Vec`, causing a type mismatch. Now both arms produce `Vec`. - let compressed: Vec = match compress_bytes(&body_bytes, encoding).await { - Ok(b) => b, - Err(_) => body_bytes.to_vec(), // FIX: .to_vec() converts Bytes → Vec + // Stream the body through the encoder instead of collecting it: time to + // first byte stays live for large responses and peak memory is one chunk + // rather than uncompressed+compressed copies of the whole payload. + let reader = tokio_util::io::StreamReader::new(body_byte_stream(body)); + let level = match config.level { + CompressionLevel::Fastest => async_compression::Level::Fastest, + CompressionLevel::Best => async_compression::Level::Best, + CompressionLevel::Default => async_compression::Level::Default, + }; + + let compressed: Body = match encoding { + Encoding::Gzip => Body::from_stream(tokio_util::io::ReaderStream::new( + GzipEncoder::with_quality(reader, level), + )), + Encoding::Br => Body::from_stream(tokio_util::io::ReaderStream::new( + BrotliEncoder::with_quality(reader, level), + )), + Encoding::Zstd => Body::from_stream(tokio_util::io::ReaderStream::new( + ZstdEncoder::with_quality(reader, level), + )), + Encoding::Deflate => Body::from_stream(tokio_util::io::ReaderStream::new( + DeflateEncoder::with_quality(reader, level), + )), }; parts.headers.insert( @@ -268,43 +292,38 @@ async fn compress_response( parts .headers .insert(VARY, HeaderValue::from_static("Accept-Encoding")); - // Remove Content-Length since compressed size differs. + // Compressed length is not known up front while streaming — hyper will + // use chunked transfer encoding. parts.headers.remove(CONTENT_LENGTH); - let final_body = Body::from_bytes(Bytes::from(compressed)); - http::Response::from_parts(parts, final_body) + http::Response::from_parts(parts, compressed) } -/// Compress `data` using the given encoding. Returns `Vec`. -async fn compress_bytes(data: &[u8], encoding: Encoding) -> std::io::Result> { - // `std::io::Cursor<&[u8]>` is `Unpin` and satisfies `AsyncRead`. - let cursor = std::io::Cursor::new(data); - match encoding { - Encoding::Gzip => { - let mut enc = GzipEncoder::new(cursor); - let mut out = Vec::new(); - enc.read_to_end(&mut out).await?; - Ok(out) - } - Encoding::Br => { - let mut enc = BrotliEncoder::new(cursor); - let mut out = Vec::new(); - enc.read_to_end(&mut out).await?; - Ok(out) - } - Encoding::Zstd => { - let mut enc = ZstdEncoder::new(cursor); - let mut out = Vec::new(); - enc.read_to_end(&mut out).await?; - Ok(out) - } - Encoding::Deflate => { - let mut enc = DeflateEncoder::new(cursor); - let mut out = Vec::new(); - enc.read_to_end(&mut out).await?; - Ok(out) - } - } +/// Adapt an HTTP body into a plain byte stream for `StreamReader`. +/// +/// Trailer frames terminate the data stream; stream errors surface as +/// `io::Error` mid-stream (the connection is aborted by hyper at that point). +type ByteStream = + futures_util::stream::BoxStream<'static, std::io::Result>; + +fn body_byte_stream(body: Body) -> ByteStream { + use futures_util::StreamExt as _; + + let mut body = Box::pin(body); + futures_util::stream::poll_fn( + move |cx| match http_body::Body::poll_frame(body.as_mut(), cx) { + std::task::Poll::Ready(Some(Ok(frame))) => match frame.into_data() { + Ok(data) => std::task::Poll::Ready(Some(Ok(data))), + Err(_trailer) => std::task::Poll::Ready(None), + }, + std::task::Poll::Ready(Some(Err(err))) => { + std::task::Poll::Ready(Some(Err(std::io::Error::other(err.to_string())))) + } + std::task::Poll::Ready(None) => std::task::Poll::Ready(None), + std::task::Poll::Pending => std::task::Poll::Pending, + }, + ) + .boxed() } fn should_compress(headers: &http::HeaderMap) -> bool { @@ -422,8 +441,14 @@ where let new_body = match decompressed { Some(data) => { - // Remove Content-Encoding since the body is now decoded. + // The body is now decoded: drop Content-Encoding and fix + // Content-Length, which still described the compressed + // wire size and would defeat downstream size checks. parts.headers_mut().remove(CONTENT_ENCODING); + parts.headers_mut().insert( + CONTENT_LENGTH, + HeaderValue::from_str(&data.len().to_string()).expect("valid CL"), + ); Body::from_bytes(Bytes::from(data)) } None => Body::from_bytes(body_bytes), // pass through as-is @@ -434,3 +459,167 @@ where }) } } + +#[cfg(test)] +mod tests { + use super::*; + + /// Inner service that reports the Content-Length header it received and + /// the actual body length, so the test can compare them. + #[derive(Clone)] + struct ReportService; + + impl Service for ReportService { + type Response = Response; + type Error = Infallible; + type Future = Pin> + Send + 'static>>; + + fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn call(&mut self, req: Request) -> Self::Future { + Box::pin(async move { + let cl = req + .headers() + .get(CONTENT_LENGTH) + .and_then(|v| v.to_str().ok()) + .unwrap_or("absent") + .to_string(); + let body = Body::to_string(req.into_body()).await.unwrap_or_default(); + let report = format!("cl={cl}|len={}", body.len()); + Ok(http::Response::builder() + .status(StatusCode::OK) + .body(Body::from(report)) + .unwrap()) + }) + } + } + + async fn gzip_compress(data: &[u8]) -> Bytes { + let mut encoder = GzipEncoder::new(std::io::Cursor::new(data)); + let mut out = Vec::new(); + tokio::io::copy(&mut encoder, &mut out).await.unwrap(); + Bytes::from(out) + } + + /// Inner service returning a fixed body larger than `min_size`. + #[derive(Clone)] + struct FixedBodyService { + body: String, + } + + impl Service for FixedBodyService { + type Response = Response; + type Error = Infallible; + type Future = Pin> + Send + 'static>>; + + fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn call(&mut self, _: Request) -> Self::Future { + let body = self.body.clone(); + Box::pin(async move { + Ok(http::Response::builder() + .status(StatusCode::OK) + .header(http::header::CONTENT_TYPE, "text/plain") + .body(Body::from(body)) + .unwrap()) + }) + } + } + + #[tokio::test] + async fn large_responses_are_streamed_compressed() { + // Larger than the default min_size so compression applies. + let payload = "arvik-streaming-compression-".repeat(200); + + let mut svc = CompressionLayer::new() + .gzip(true) + .br(false) + .zstd(false) + .deflate(false) + .layer(FixedBodyService { + body: payload.clone(), + }); + + let res = svc + .call(Request::new( + http::Request::builder() + .method(http::Method::GET) + .uri("/") + .header(ACCEPT_ENCODING, "gzip") + .body(Body::empty()) + .unwrap(), + )) + .await + .unwrap(); + + assert_eq!( + res.headers() + .get(CONTENT_ENCODING) + .and_then(|v| v.to_str().ok()), + Some("gzip") + ); + // Streaming: no Content-Length up front. + assert!(res.headers().get(CONTENT_LENGTH).is_none()); + + // The streamed body must decode back to the original payload. + let raw = res.into_body().to_bytes().await.unwrap(); + let mut decoder = GzipDecoder::new(std::io::Cursor::new(raw.as_ref())); + let mut decoded = Vec::new(); + decoder.read_to_end(&mut decoded).await.unwrap(); + assert_eq!(decoded, payload.as_bytes()); + } + + #[tokio::test] + async fn head_responses_skip_compression() { + let mut svc = CompressionLayer::new().layer(FixedBodyService { + body: "y".repeat(4096), + }); + + let res = svc + .call(Request::new( + http::Request::builder() + .method(http::Method::HEAD) + .uri("/") + .header(ACCEPT_ENCODING, "gzip") + .body(Body::empty()) + .unwrap(), + )) + .await + .unwrap(); + + assert!(res.headers().get(CONTENT_ENCODING).is_none()); + } + + #[tokio::test] + async fn decompression_fixes_content_length_to_decoded_size() { + let original = b"hello world, this payload is much larger once decompressed!"; + let compressed = gzip_compress(original).await; + assert!(compressed.len() != original.len()); + + let mut svc = DecompressionLayer::new().layer(ReportService); + let res = svc + .call(Request::new( + http::Request::builder() + .method(http::Method::POST) + .uri("/") + .header(CONTENT_ENCODING, "gzip") + .header(CONTENT_LENGTH, http::HeaderValue::from(compressed.len())) + .body(Body::from_bytes(compressed)) + .unwrap(), + )) + .await + .unwrap(); + + let text = Body::to_string(res.into_body()).await.unwrap(); + // The stale compressed length must not survive; both views of size + // must agree on the decoded payload. + assert_eq!( + text, + format!("cl={}|len={}", original.len(), original.len()) + ); + } +} diff --git a/arvik-middleware/src/cors.rs b/arvik-middleware/src/cors.rs index fee7cbe..c1b5d3d 100644 --- a/arvik-middleware/src/cors.rs +++ b/arvik-middleware/src/cors.rs @@ -42,9 +42,10 @@ //! //! # Credential support //! -//! When `allow_credentials(true)` is set, the `Access-Control-Allow-Origin` -//! header **cannot** be `*` per the CORS spec. Arvik automatically switches -//! from `*` to reflecting the request `Origin`. +//! `allow_credentials(true)` requires an explicit origin allowlist — the CORS +//! spec forbids `Access-Control-Allow-Origin: *` with credentials, and +//! mirroring arbitrary origins would make every site a trusted credentialed +//! origin. That combination panics at construction. use std::convert::Infallible; use std::future::Future; @@ -154,10 +155,14 @@ impl CorsLayer { } } - /// Allow all origins with credentials. + /// Allow all origins, methods, and headers — for **local development**. /// - /// Mirrors the request `Origin` (CORS spec forbids `*` with credentials). - /// Suitable for trusted cross-origin applications that need cookies or auth. + /// Mirrors the request `Origin`. Credentials are NOT enabled: mirroring + /// arbitrary origins together with `Access-Control-Allow-Credentials: + /// true` lets any website read authenticated responses from your users. + /// For credentialed cross-origin apps in production, list explicit + /// origins with [`CorsLayer::allow_origin`] and then enable + /// [`CorsLayer::allow_credentials`]. pub fn very_permissive() -> Self { Self { config: Arc::new(CorsConfig { @@ -165,7 +170,7 @@ impl CorsLayer { allow_methods: AllowMethods::Any, allow_headers: AllowHeaders::Mirror, expose_headers: Vec::new(), - allow_credentials: true, + allow_credentials: false, max_age_secs: Some(86400), }), } @@ -176,6 +181,7 @@ impl CorsLayer { fn mutate(self, f: F) -> Self { let mut cfg = Arc::try_unwrap(self.config).unwrap_or_else(|arc| (*arc).clone()); f(&mut cfg); + assert_credentials_safe(&cfg); Self { config: Arc::new(cfg), } @@ -232,19 +238,17 @@ impl CorsLayer { /// Allow or deny cookies, HTTP auth, and TLS client certificates. /// - /// When set to `true`: - /// - `Access-Control-Allow-Credentials: true` is added to responses. - /// - `Access-Control-Allow-Origin` cannot be `*`; the layer automatically - /// switches to mirroring the request origin. + /// When set to `true`, `Access-Control-Allow-Credentials: true` is added + /// to responses and the allowed origins MUST be an explicit list — + /// wildcard or mirrored origins would make every site a trusted + /// credentialed origin, so that combination is rejected at construction. + /// + /// # Panics + /// + /// Panics when credentials are enabled while origins are wildcarded or + /// mirrored. List explicit origins with [`CorsLayer::allow_origin`]. pub fn allow_credentials(self, allow: bool) -> Self { - self.mutate(|c| { - c.allow_credentials = allow; - if allow { - if let AllowOrigin::Any = c.allow_origins { - c.allow_origins = AllowOrigin::Mirror; - } - } - }) + self.mutate(|c| c.allow_credentials = allow) } /// Set the preflight response cache duration (`Access-Control-Max-Age`). @@ -259,9 +263,27 @@ impl Default for CorsLayer { } } +// ── Safety invariant ───────────────────────────────────────────────────────── + +/// Reject the one CORS combination that disables the browser's protection +/// entirely: reflecting arbitrary origins while allowing credentials. That +/// lets any website attach the victim's cookies/auth and read responses. +/// Fail loudly at construction instead of silently in production. +fn assert_credentials_safe(cfg: &CorsConfig) { + if cfg.allow_credentials && matches!(cfg.allow_origins, AllowOrigin::Mirror | AllowOrigin::Any) + { + panic!( + "CORS misconfiguration: allow_credentials(true) combined with wildcard or \ + mirrored origins trusts EVERY site with the user's cookies and auth headers. \ + List explicit origins via .allow_origin([\"https://app.example.com\", ...]) \ + before enabling credentials." + ); + } +} + // ── IntoAllowOrigin conversions ─────────────────────────────────────────────── -/// Types that can be converted into an [`AllowOrigin`] configuration. +/// Types that can be converted into an `AllowOrigin` configuration. pub trait IntoAllowOrigin: sealed::Sealed { #[doc(hidden)] #[allow(private_interfaces)] @@ -451,10 +473,13 @@ impl CorsConfig { } } - // Vary: Origin — required when response differs per origin - if !matches!(&self.allow_origins, AllowOrigin::Any) || self.allow_credentials { - builder = builder.header(VARY, "origin"); - } + // Preflight replies are cacheable and differ per Origin, per requested + // method, and per requested headers — declare all three so a shared + // cache/CDN cannot serve one page's preflight to another. + builder = builder.header( + VARY, + "Origin, Access-Control-Request-Method, Access-Control-Request-Headers", + ); builder.body(Body::empty()) } @@ -497,3 +522,102 @@ impl CorsConfig { } } } + +#[cfg(test)] +mod tests { + use super::*; + use arvik_core::{Body, Request}; + use std::convert::Infallible; + use std::pin::Pin; + + fn preflight(origin: &str, method: &str) -> Request { + Request::new( + http::Request::builder() + .method(Method::OPTIONS) + .uri("/") + .header(ORIGIN, origin) + .header(ACCESS_CONTROL_REQUEST_METHOD, method) + .body(Body::empty()) + .unwrap(), + ) + } + + #[derive(Clone)] + struct OkService; + + impl Service for OkService { + type Response = Response; + type Error = Infallible; + type Future = Pin> + Send + 'static>>; + + fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn call(&mut self, _: Request) -> Self::Future { + Box::pin(async { + Ok(http::Response::builder() + .status(StatusCode::OK) + .body(Body::empty()) + .unwrap()) + }) + } + } + + #[test] + #[should_panic(expected = "CORS misconfiguration")] + fn mirrored_origins_with_credentials_are_rejected_at_construction() { + let _ = CorsLayer::very_permissive().allow_credentials(true); + } + + #[test] + #[should_panic(expected = "CORS misconfiguration")] + fn permissive_plus_credentials_is_rejected_at_construction() { + let _ = CorsLayer::permissive().allow_credentials(true); + } + + #[tokio::test] + async fn explicit_origins_with_credentials_are_allowed() { + let mut svc = CorsLayer::new() + .allow_origin(["https://app.example.com".parse::().unwrap()]) + .allow_credentials(true) + .layer(OkService); + + let res = svc + .call(preflight("https://app.example.com", "GET")) + .await + .unwrap(); + assert_eq!(res.status(), StatusCode::NO_CONTENT); + assert_eq!( + res.headers() + .get(ACCESS_CONTROL_ALLOW_CREDENTIALS) + .and_then(|v| v.to_str().ok()), + Some("true") + ); + } + + #[tokio::test] + async fn preflight_vary_names_all_request_dependent_dimensions() { + let mut svc = CorsLayer::very_permissive().layer(OkService); + + let res = svc + .call(preflight("https://dev.local", "POST")) + .await + .unwrap(); + let vary = res + .headers() + .get(VARY) + .and_then(|v| v.to_str().ok()) + .unwrap_or_default(); + for expected in [ + "Origin", + "Access-Control-Request-Method", + "Access-Control-Request-Headers", + ] { + assert!( + vary.contains(expected), + "Vary `{vary}` must contain `{expected}`" + ); + } + } +} diff --git a/arvik-middleware/src/csrf.rs b/arvik-middleware/src/csrf.rs index a80dbf3..b71ecb6 100644 --- a/arvik-middleware/src/csrf.rs +++ b/arvik-middleware/src/csrf.rs @@ -198,7 +198,7 @@ where .map(|s| s.to_string()); match header_token { - Some(h) if h == csrf_token.as_str() => { + Some(h) if crate::auth::secret_eq(&h, csrf_token.as_str()) => { // Valid — proceed } Some(bad) => { diff --git a/arvik-middleware/src/from_fn.rs b/arvik-middleware/src/from_fn.rs index 87d01c8..9adf944 100644 --- a/arvik-middleware/src/from_fn.rs +++ b/arvik-middleware/src/from_fn.rs @@ -2,7 +2,7 @@ //! //! Tower [`Layer`] and [`Service`] produced by [`from_fn`] and [`from_fn_with_state`]. //! -//! Apply it to a [`Router`] via `.layer(from_fn(my_fn))`. +//! Apply it to a `Router` via `.layer(from_fn(my_fn))`. //! You typically do not interact with these types directly. use std::marker::PhantomData; @@ -19,7 +19,7 @@ use tower_service::Service; /// Create a Tower middleware layer from a plain async function. /// -/// The function may accept any number of [`FromRequestParts`] extractors +/// The function may accept any number of `FromRequestParts` extractors /// as leading parameters, followed by `Request` and `Next`: /// /// ```text @@ -27,12 +27,12 @@ use tower_service::Service; /// ``` /// /// All extractors that work with state `()` are available without any -/// additional setup: [`Method`], [`Uri`], [`Version`], [`HeaderMap`], -/// [`Path`], [`Query`], [`MatchedPath`], [`RawPathParams`], -/// [`ConnectInfo`], [`Extension`], [`TypedHeader`], [`CookieJar`]. +/// additional setup: `Method`, `Uri`, `Version`, `HeaderMap`, +/// `Path`, `Query`, `MatchedPath`, `RawPathParams`, +/// `ConnectInfo`, `Extension`, `TypedHeader`, `CookieJar`. /// -/// For extractors that need shared application state — [`State`], -/// [`SignedCookieJar`], [`PrivateCookieJar`] — use [`from_fn_with_state`]. +/// For extractors that need shared application state — `State`, +/// `SignedCookieJar`, `PrivateCookieJar` — use [`from_fn_with_state`]. /// /// ### Examples /// @@ -101,13 +101,13 @@ where /// Create a Tower middleware layer from an async function that has access to /// shared application state through extractors. /// -/// The `state` provided here is used to resolve [`FromRequestParts`] -/// extractors inside the middleware function — just as [`Router::with_state`] +/// The `state` provided here is used to resolve `FromRequestParts` +/// extractors inside the middleware function — just as `Router::with_state` /// provides state to handlers. /// /// **Pass the same state object here as you pass to `.with_state()`** so -/// extractors like [`State`], [`SignedCookieJar`], and -/// [`PrivateCookieJar`] work correctly. +/// extractors like `State`, `SignedCookieJar`, and +/// `PrivateCookieJar` work correctly. /// /// ### Signature /// @@ -216,7 +216,7 @@ where /// Tower [`Layer`] produced by [`from_fn`] or [`from_fn_with_state`]. /// -/// Apply it to a [`Router`] via `.layer(from_fn(my_fn))`. +/// Apply it to a `Router` via `.layer(from_fn(my_fn))`. /// You typically do not interact with this type directly. pub struct FromFnLayer { f: F, @@ -250,6 +250,7 @@ where Svc: Service + Clone + Send + + Sync + 'static, Svc::Future: Send + 'static, { @@ -306,6 +307,7 @@ where Svc: Service + Clone + Send + + Sync + 'static, Svc::Future: Send + 'static, { diff --git a/arvik-middleware/src/lib.rs b/arvik-middleware/src/lib.rs index 24a03b6..d5db936 100644 --- a/arvik-middleware/src/lib.rs +++ b/arvik-middleware/src/lib.rs @@ -22,10 +22,10 @@ //! //! | Layer | Description | Version | //! |---|---|---| -//! | [`from_fn`] | Middleware from async fn with extractor support | ✅ 0.4.1 | +//! | [from_fn()](crate::from_fn()) | Middleware from async fn with extractor support | ✅ 0.4.1 | //! | [`from_fn_with_state`] | Same, with router state access | ✅ 0.4.1 | -//! | [`map_request`] | Transform request only | ✅ 0.4.1 | -//! | [`map_response`] | Transform response only | ✅ 0.4.1 | +//! | [map_request()](crate::map_request()) | Transform request only | ✅ 0.4.1 | +//! | [map_response()](crate::map_response()) | Transform response only | ✅ 0.4.1 | //! | [`cors::CorsLayer`] | Full CORS spec | ✅ 0.4.1 | //! | [`compression::CompressionLayer`] | gzip / br / zstd / deflate | ✅ 0.4.2 | //! | [`compression::DecompressionLayer`] | Decompress request bodies | ✅ 0.4.2 | diff --git a/arvik-middleware/src/map_body.rs b/arvik-middleware/src/map_body.rs index 2371677..34f437f 100644 --- a/arvik-middleware/src/map_body.rs +++ b/arvik-middleware/src/map_body.rs @@ -3,7 +3,7 @@ //! [`MapRequestBodyLayer`] transforms the request body before it reaches the handler. //! [`MapResponseBodyLayer`] transforms the response body after the handler runs. //! -//! These are lower-level than [`map_request`]/[`map_response`] — they operate +//! These are lower-level than `map_request`/`map_response` — they operate //! on the raw body bytes rather than the full request/response. //! //! # Example @@ -95,7 +95,7 @@ where let mut inner = std::mem::replace(&mut self.inner, cloned); Box::pin(async move { - let (parts, body) = req.into_request_parts(); + let (mut parts, body) = req.into_request_parts(); let bytes = match body.to_bytes().await { Ok(b) => b, @@ -103,6 +103,18 @@ where }; let transformed = f(bytes).await; + // The transform may change the size — refresh Content-Length so + // downstream length checks describe the real body. + if parts + .headers_mut() + .contains_key(http::header::CONTENT_LENGTH) + { + let value = http::HeaderValue::from_str(&transformed.len().to_string()) + .expect("Content-Length from usize is a valid header value"); + parts + .headers_mut() + .insert(http::header::CONTENT_LENGTH, value); + } let new_req = Request::from_request_parts(parts, Body::from_bytes(transformed)); inner.call(new_req).await }) @@ -170,7 +182,7 @@ where Box::pin(async move { let response = inner.call(req).await?; - let (parts, body) = response.into_parts(); + let (mut parts, body) = response.into_parts(); let bytes = match body.to_bytes().await { Ok(b) => b, @@ -178,8 +190,132 @@ where }; let transformed = f(bytes).await; + // The transform may change the size — a stale Content-Length makes + // hyper frame the response by the old length (truncated bodies). + if parts.headers.contains_key(http::header::CONTENT_LENGTH) { + let value = http::HeaderValue::from_str(&transformed.len().to_string()) + .expect("Content-Length from usize is a valid header value"); + parts.headers.insert(http::header::CONTENT_LENGTH, value); + } let new_body = Body::from_bytes(transformed); Ok(http::Response::from_parts(parts, new_body)) }) } } + +#[cfg(test)] +mod tests { + use super::*; + use arvik_core::extract::FromRequest; + + #[derive(Clone)] + struct FixedResponseService; + + impl Service for FixedResponseService { + type Response = Response; + type Error = Infallible; + type Future = Pin> + Send + 'static>>; + + fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn call(&mut self, _: Request) -> Self::Future { + Box::pin(async move { + // "hello" with an accurate Content-Length. + let res = http::Response::builder() + .status(http::StatusCode::OK) + .header(http::header::CONTENT_LENGTH, "5") + .body(Body::from(Bytes::from_static(b"hello"))) + .unwrap(); + Ok(res) + }) + } + } + + #[derive(Clone)] + struct EchoBodyService; + + impl Service for EchoBodyService { + type Response = Response; + type Error = Infallible; + type Future = Pin> + Send + 'static>>; + + fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn call(&mut self, req: Request) -> Self::Future { + Box::pin(async move { + let cl = req + .headers() + .get(http::header::CONTENT_LENGTH) + .and_then(|v| v.to_str().ok()) + .unwrap_or("absent") + .to_string(); + let bytes = >::from_request(req, &()) + .await + .unwrap(); + let body = format!("len={cl}|body={}", String::from_utf8_lossy(&bytes)); + Ok(http::Response::builder() + .status(http::StatusCode::OK) + .body(Body::from(body)) + .unwrap()) + }) + } + } + + /// Size-changing transform from the module docs' own example. + fn append_footer(bytes: Bytes) -> Pin + Send>> { + Box::pin(async move { + let mut out = bytes.to_vec(); + out.extend_from_slice(b"!!"); + Bytes::from(out) + }) + } + + #[tokio::test] + async fn response_content_length_tracks_size_changing_transform() { + let mut svc = MapResponseBodyLayer::new(append_footer).layer(FixedResponseService); + let res = svc + .call(Request::new( + http::Request::builder() + .uri("/") + .body(Body::empty()) + .unwrap(), + )) + .await + .unwrap(); + + assert_eq!(res.status(), http::StatusCode::OK); + let cl: usize = res + .headers() + .get(http::header::CONTENT_LENGTH) + .unwrap() + .to_str() + .unwrap() + .parse() + .unwrap(); + assert_eq!(cl, 7, "Content-Length must describe the transformed body"); + assert_eq!(Body::to_string(res.into_body()).await.unwrap(), "hello!!"); + } + + #[tokio::test] + async fn request_content_length_tracks_size_changing_transform() { + let mut svc = MapRequestBodyLayer::new(append_footer).layer(EchoBodyService); + let res = svc + .call(Request::new( + http::Request::builder() + .method(http::Method::POST) + .uri("/") + .header(http::header::CONTENT_LENGTH, "3") + .body(Body::from(Bytes::from_static(b"abc"))) + .unwrap(), + )) + .await + .unwrap(); + + let text = Body::to_string(res.into_body()).await.unwrap(); + assert_eq!(text, "len=5|body=abc!!"); + } +} diff --git a/arvik-middleware/src/map_request.rs b/arvik-middleware/src/map_request.rs index 15f46ac..0e88072 100644 --- a/arvik-middleware/src/map_request.rs +++ b/arvik-middleware/src/map_request.rs @@ -3,7 +3,7 @@ //! Create middleware that **transforms only the request** before it reaches //! the handler. The response is returned unchanged. //! -//! Slightly more efficient than [`from_fn`] when you only need to mutate +//! Slightly more efficient than `from_fn` when you only need to mutate //! the request (no response inspection). use std::convert::Infallible; @@ -21,7 +21,7 @@ use tower_service::Service; /// Create middleware that **transforms only the request** before it reaches /// the handler. The response is returned unchanged. /// -/// Slightly more efficient than [`from_fn`] when you only need to mutate +/// Slightly more efficient than `from_fn` when you only need to mutate /// the request (no response inspection). /// /// ### Example diff --git a/arvik-middleware/src/map_response.rs b/arvik-middleware/src/map_response.rs index 629523f..40870bb 100644 --- a/arvik-middleware/src/map_response.rs +++ b/arvik-middleware/src/map_response.rs @@ -3,7 +3,7 @@ //! Create middleware that **transforms only the response** after the handler //! runs. The request is passed through unchanged. //! -//! Slightly more efficient than [`from_fn`] when you only need to mutate +//! Slightly more efficient than `from_fn` when you only need to mutate //! the response (no request inspection). use std::convert::Infallible; @@ -21,7 +21,7 @@ use tower_service::Service; /// Create middleware that **transforms only the response** after the handler /// runs. The request is passed through unchanged. /// -/// Slightly more efficient than [`from_fn`] when you only need to mutate +/// Slightly more efficient than `from_fn` when you only need to mutate /// the response (no request inspection). /// /// ### Example diff --git a/arvik-middleware/src/rate_limit.rs b/arvik-middleware/src/rate_limit.rs index 6927fb8..71954db 100644 --- a/arvik-middleware/src/rate_limit.rs +++ b/arvik-middleware/src/rate_limit.rs @@ -4,6 +4,15 @@ //! `429 Too Many Requests` with a `Retry-After` header when the //! limit is exceeded. //! +//! # Client identification +//! +//! By default clients are keyed by **socket address** — client-supplied +//! `X-Forwarded-For` / `X-Real-IP` headers are ignored, so a client cannot +//! mint fresh buckets by rotating header values. When deployed behind a +//! trusted reverse proxy, declare it with [`RateLimitLayer::trust_proxies`]; +//! the limiter then applies the standard *rightmost-untrusted-hop* rule to +//! `X-Forwarded-For`. +//! //! # Example //! //! ```rust,ignore @@ -14,13 +23,22 @@ //! Router::new() //! .route("/api", get(handler)) //! .layer(RateLimitLayer::new(100, Duration::from_secs(1))); +//! +//! // 600 requests per minute per IP, behind one known proxy +//! Router::new() +//! .route("/api", get(handler)) +//! .layer(RateLimitLayer::new(600, Duration::from_secs(60)) +//! .trust_proxies(["10.0.0.9".parse::().unwrap()])); //! ``` use std::collections::HashMap; +use std::collections::hash_map::RandomState; use std::convert::Infallible; use std::future::Future; +use std::hash::BuildHasher; use std::pin::Pin; use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; use std::task::{Context, Poll}; use std::time::{Duration, Instant}; @@ -30,22 +48,39 @@ use parking_lot::Mutex; use tower_layer::Layer; use tower_service::Service; +/// Number of independently locked bucket shards. +const SHARDS: usize = 32; + +/// Run an idle-bucket sweep roughly once per this many requests. +const SWEEP_EVERY: u64 = 4096; + +/// Buckets idle longer than this are evicted from the shared state. +const IDLE_TTL: Duration = Duration::from_secs(300); + +/// Upper bound on custom-header key length (bytes). +const MAX_HEADER_KEY_LEN: usize = 128; + // ── Token bucket ───────────────────────────────────────────────────────────── struct Bucket { tokens: f64, + #[allow(dead_code)] capacity: f64, - refill_rate: f64, // tokens per second + refill_rate: f64, // tokens per second = capacity / window last_refill: Instant, + last_seen: Instant, } impl Bucket { - fn new(capacity: u64) -> Self { + fn new(capacity: u64, window: Duration) -> Self { + let secs = window.as_secs_f64(); + debug_assert!(secs > 0.0, "rate limit window must be greater than zero"); Self { tokens: capacity as f64, capacity: capacity as f64, - refill_rate: capacity as f64, + refill_rate: capacity as f64 / secs.max(f64::MIN_POSITIVE), last_refill: Instant::now(), + last_seen: Instant::now(), } } @@ -58,6 +93,7 @@ impl Bucket { let elapsed = now.duration_since(self.last_refill).as_secs_f64(); self.tokens = (self.tokens + elapsed * self.refill_rate).min(self.capacity); self.last_refill = now; + self.last_seen = now; if self.tokens >= 1.0 { self.tokens -= 1.0; @@ -70,57 +106,157 @@ impl Bucket { } } +// ── Shared, sharded bucket state ───────────────────────────────────────────── + +/// Sharded bucket store: each shard has its own lock so concurrent requests +/// with different keys never serialize behind one mutex. +struct BucketStore { + shards: Vec>>, + hasher: RandomState, + calls: AtomicU64, +} + +impl BucketStore { + fn new(_window: Duration) -> Self { + Self { + shards: (0..SHARDS).map(|_| Mutex::new(HashMap::new())).collect(), + hasher: RandomState::new(), + calls: AtomicU64::new(0), + } + } + + fn shard_for(&self, key: &str) -> &Mutex> { + let idx = (self.hasher.hash_one(key) % self.shards.len() as u64) as usize; + &self.shards[idx] + } + + /// Consume one token for `key`, creating the bucket on first sight. + /// + /// Returns `None` when allowed, or `Some(retry_after_secs)` when limited. + fn try_consume(&self, key: &str, capacity: u64, window: Duration) -> Option { + // Amortized idle-bucket sweep keeps memory bounded under key churn. + let calls = self.calls.fetch_add(1, Ordering::Relaxed); + if calls % SWEEP_EVERY == 0 { + self.evict_idle(); + } + + let mut map = self.shard_for(key).lock(); + // Hot path: existing bucket, no allocation. + if let Some(bucket) = map.get_mut(key) { + return bucket.try_consume(); + } + let mut bucket = Bucket::new(capacity, window); + let outcome = bucket.try_consume(); + map.insert(key.to_string(), bucket); + outcome + } + + /// Drop buckets untouched for longer than [`IDLE_TTL`]. + fn evict_idle(&self) { + for shard in &self.shards { + shard + .lock() + .retain(|_, bucket| bucket.last_seen.elapsed() < IDLE_TTL); + } + } +} + // ── Key extraction ──────────────────────────────────────────────────────────── +/// How much the limiter trusts client-supplied IP headers. +/// +/// Defaults to [`ProxyTrust::Never`] — clients cannot mint new rate-limit +/// buckets by sending different `X-Forwarded-For` values. +#[derive(Debug, Clone, Default)] +pub enum ProxyTrust { + /// Ignore `X-Forwarded-For` / `X-Real-IP`; key by socket address. + #[default] + Never, + /// Honor proxy headers **only** when the immediate peer is one of these + /// addresses. The client is the rightmost `X-Forwarded-For` entry that is + /// not itself a trusted proxy (the standard rightmost-untrusted-hop rule). + Trusted(Vec), +} + /// Strategy for extracting a rate-limit key from a request. #[derive(Debug, Clone)] pub enum KeyExtractor { - /// Rate limit per client IP address (default). - /// Reads `X-Forwarded-For` first, falls back to socket addr extension. + /// Rate limit per client IP address (default). Identification honors the + /// layer's [`ProxyTrust`] setting. IpAddress, - /// Custom header value as the key. + /// Custom header value as the key. Values longer than 128 bytes are + /// truncated (long unique values must not grow the bucket map). Header(String), /// Rate limit globally (all requests share one bucket). Global, } impl KeyExtractor { - fn extract_key(&self, req: &Request) -> String { + fn extract_key(&self, req: &Request, trust: &ProxyTrust) -> String { match self { - KeyExtractor::IpAddress => extract_ip(req), - KeyExtractor::Header(name) => req - .headers() - .get(name.as_str()) - .and_then(|v| v.to_str().ok()) - .unwrap_or("unknown") - .to_string(), + KeyExtractor::IpAddress => extract_ip(req, trust), + KeyExtractor::Header(name) => { + let raw = req + .headers() + .get(name.as_str()) + .and_then(|v| v.to_str().ok()) + .unwrap_or("unknown"); + let mut key = raw.trim().to_string(); + key.truncate(MAX_HEADER_KEY_LEN); + key + } KeyExtractor::Global => "__global__".to_string(), } } } -fn extract_ip(req: &Request) -> String { - // Try X-Forwarded-For first (behind proxy). +/// Resolve the client IP for rate limiting. +/// +/// With [`ProxyTrust::Never`] this is always the socket address. With +/// [`ProxyTrust::Trusted`], header values are considered only when the +/// immediate peer is a declared proxy, and only well-formed IPs are accepted — +/// garbage entries fall back to the peer address rather than becoming keys. +fn extract_ip(req: &Request, trust: &ProxyTrust) -> String { + let peer = req + .extensions() + .get::() + .map(|addr| addr.ip()); + + let Some(peer_ip) = peer else { + // No connection info available (e.g. direct-call test paths). + return "unknown".to_string(); + }; + + let ProxyTrust::Trusted(proxies) = trust else { + return peer_ip.to_string(); + }; + if !proxies.contains(&peer_ip) { + return peer_ip.to_string(); + } + + // Peer is a trusted proxy: walk XFF right-to-left for the first entry + // that is not itself a trusted proxy address. if let Some(forwarded) = req .headers() .get("x-forwarded-for") .and_then(|v| v.to_str().ok()) { - if let Some(ip) = forwarded.split(',').next() { - return ip.trim().to_string(); + for candidate in forwarded.split(',').rev() { + let candidate = candidate.trim(); + if candidate.is_empty() { + continue; + } + match candidate.parse::() { + Ok(ip) if !proxies.contains(&ip) => return ip.to_string(), + Ok(_) => continue, // another trusted hop; keep walking left + Err(_) => break, // malformed entry — do not trust the rest + } } } - // Try X-Real-IP. - if let Some(real_ip) = req.headers().get("x-real-ip").and_then(|v| v.to_str().ok()) { - return real_ip.trim().to_string(); - } - - // Fall back to ConnectInfo extension. - req.extensions() - .get::() - .map(|addr| addr.ip().to_string()) - .unwrap_or_else(|| "unknown".to_string()) + // Every reported hop was trusted (or the header was absent/malformed): + // attribute to the proxy itself rather than to attacker-controlled text. + peer_ip.to_string() } // ── RateLimitLayer ──────────────────────────────────────────────────────────── @@ -129,26 +265,44 @@ fn extract_ip(req: &Request) -> String { #[derive(Clone)] pub struct RateLimitLayer { capacity: u64, + window: Duration, extractor: KeyExtractor, - /// Shared bucket state — all clones of this layer share the same map. - state: Arc>>, + proxy_trust: ProxyTrust, + store: Arc, } impl RateLimitLayer { - /// Create a new rate limiter allowing `capacity` requests per `window`. + /// Create a rate limiter allowing `capacity` requests per `window`, + /// refilling continuously (e.g. `new(600, 60s)` admits bursts up to 600 + /// and sustains 10 requests/second). + /// + /// # Panics /// - /// Keys by client IP address by default. The `window` parameter is - /// kept for API symmetry but the token bucket naturally distributes - /// `capacity` tokens over a one-second replenishment cycle; adjust - /// `capacity` to match your `window` (e.g. 10 req / 10s → capacity=1). - pub fn new(capacity: u64, _window: Duration) -> Self { + /// Panics if `window` is zero. + pub fn new(capacity: u64, window: Duration) -> Self { + assert!( + !window.is_zero(), + "RateLimitLayer window must be greater than zero" + ); Self { capacity, + window, extractor: KeyExtractor::IpAddress, - state: Arc::new(Mutex::new(HashMap::new())), + proxy_trust: ProxyTrust::default(), + store: Arc::new(BucketStore::new(window)), } } + /// Declare trusted reverse proxies for client-IP resolution. + /// + /// Only these peers may set `X-Forwarded-For`; the client becomes the + /// rightmost non-trusted entry. Without this, IP keys come from the + /// socket address and proxy headers are ignored. + pub fn trust_proxies(mut self, proxies: impl IntoIterator) -> Self { + self.proxy_trust = ProxyTrust::Trusted(proxies.into_iter().collect()); + self + } + /// Rate limit by a custom request header value. pub fn by_header(mut self, header_name: impl Into) -> Self { self.extractor = KeyExtractor::Header(header_name.into()); @@ -175,8 +329,10 @@ impl Layer for RateLimitLayer { RateLimitService { inner, capacity: self.capacity, + window: self.window, extractor: self.extractor.clone(), - state: Arc::clone(&self.state), + proxy_trust: self.proxy_trust.clone(), + store: Arc::clone(&self.store), } } } @@ -186,8 +342,10 @@ impl Layer for RateLimitLayer { pub struct RateLimitService { inner: S, capacity: u64, + window: Duration, extractor: KeyExtractor, - state: Arc>>, + proxy_trust: ProxyTrust, + store: Arc, } impl Service for RateLimitService @@ -204,21 +362,16 @@ where } fn call(&mut self, req: Request) -> Self::Future { - let key = self.extractor.extract_key(&req); + let key = self.extractor.extract_key(&req, &self.proxy_trust); let capacity = self.capacity; - let state = Arc::clone(&self.state); + let window = self.window; + let store = Arc::clone(&self.store); let cloned = self.inner.clone(); let mut inner = std::mem::replace(&mut self.inner, cloned); Box::pin(async move { - // Check the bucket inside a short synchronous critical section. - let retry_after: Option = { - let mut buckets = state.lock(); - let bucket = buckets - .entry(key.clone()) - .or_insert_with(|| Bucket::new(capacity)); - bucket.try_consume() - }; + // Short synchronous critical section on one shard only. + let retry_after: Option = store.try_consume(&key, capacity, window); if let Some(retry_secs) = retry_after { let retry_secs_ceil = retry_secs.ceil() as u64; @@ -228,8 +381,6 @@ where "Rate limit exceeded" ); - // FIX: explicit Ok::(...) so Rust can infer the - // error type of this early-return future independently of `inner.call`. return Ok::( http::Response::builder() .status(StatusCode::TOO_MANY_REQUESTS) @@ -248,3 +399,243 @@ where }) } } + +#[cfg(test)] +mod tests { + use super::*; + use std::net::SocketAddr; + + fn request_with_peer_and_xff(peer: SocketAddr, xff: Option<&str>) -> Request { + let mut builder = http::Request::builder().method(http::Method::GET).uri("/"); + if let Some(xff) = xff { + builder = builder.header("x-forwarded-for", xff); + } + let mut req = Request::new(builder.body(Body::empty()).unwrap()); + req.extensions_mut().insert(peer); + req + } + + fn peer_addr(ip: [u8; 4], port: u16) -> SocketAddr { + SocketAddr::from((ip, port)) + } + + // ── H6: client identification ──────────────────────────────────────────── + + #[test] + fn spoofed_forwarded_headers_are_ignored_by_default() { + let trust = ProxyTrust::default(); + let req = request_with_peer_and_xff( + peer_addr([203, 0, 113, 7], 50_000), + Some("1.2.3.4, 5.6.7.8"), + ); + assert_eq!(extract_ip(&req, &trust), "203.0.113.7"); + } + + #[test] + fn missing_connection_info_falls_back_to_unknown() { + let req = Request::new( + http::Request::builder() + .method(http::Method::GET) + .uri("/") + .body(Body::empty()) + .unwrap(), + ); + assert_eq!(extract_ip(&req, &ProxyTrust::default()), "unknown"); + } + + #[test] + fn untrusted_proxy_cannot_set_client_identity() { + let trust = ProxyTrust::Trusted(vec!["10.0.0.9".parse().unwrap()]); + let req = request_with_peer_and_xff( + peer_addr([198, 51, 100, 4], 40_000), // NOT a trusted proxy + Some("1.2.3.4"), + ); + assert_eq!(extract_ip(&req, &trust), "198.51.100.4"); + } + + #[test] + fn trusted_proxy_yields_rightmost_untrusted_hop() { + let trust = ProxyTrust::Trusted(vec!["10.0.0.9".parse().unwrap()]); + let req = request_with_peer_and_xff( + peer_addr([10, 0, 0, 9], 40_000), // trusted proxy + Some("203.0.113.5, 10.0.0.9"), // client, then the proxy chain + ); + assert_eq!(extract_ip(&req, &trust), "203.0.113.5"); + } + + #[test] + fn fully_trusted_chain_attributes_to_the_proxy() { + let trust = ProxyTrust::Trusted(vec!["10.0.0.9".parse().unwrap()]); + let req = + request_with_peer_and_xff(peer_addr([10, 0, 0, 9], 40_000), Some("10.0.0.9, 10.0.0.9")); + assert_eq!(extract_ip(&req, &trust), "10.0.0.9"); + } + + #[test] + fn malformed_forwarded_entries_are_not_keys() { + let trust = ProxyTrust::Trusted(vec!["10.0.0.9".parse().unwrap()]); + // Attacker-controlled garbage must not become a bucket key. + let req = request_with_peer_and_xff( + peer_addr([10, 0, 0, 9], 40_000), + Some(", 10.0.0.9"), + ); + assert_eq!(extract_ip(&req, &trust), "10.0.0.9"); + } + + #[test] + fn header_keys_are_trimmed_and_bounded() { + let extractor = KeyExtractor::Header("x-api-key".to_string()); + let long_value = "k".repeat(500); + let mut req = Request::new( + http::Request::builder() + .method(http::Method::GET) + .uri("/") + .header("x-api-key", long_value.as_str()) + .body(Body::empty()) + .unwrap(), + ); + req.extensions_mut().insert(peer_addr([1, 2, 3, 4], 80)); + assert_eq!( + extractor.extract_key(&req, &ProxyTrust::default()).len(), + 128 + ); + } + + // ── M8: window drives refill rate ──────────────────────────────────────── + + #[test] + fn refill_rate_matches_capacity_per_window() { + let minute = Bucket::new(600, Duration::from_secs(60)); + assert!((minute.refill_rate - 10.0).abs() < 1e-9); // 10/s sustained + + let second = Bucket::new(100, Duration::from_secs(1)); + assert!((second.refill_rate - 100.0).abs() < 1e-9); + } + + #[tokio::test] + async fn burst_cap_and_retry_after_reflect_the_window() { + let mut svc = RateLimitLayer::new(5, Duration::from_secs(50)) + .global() + .layer(PassthroughService); + + // Burn the burst allowance. + for _ in 0..5 { + let res = svc + .call(Request::new( + http::Request::builder() + .uri("/") + .body(Body::empty()) + .unwrap(), + )) + .await + .unwrap(); + assert_eq!(res.status(), StatusCode::OK); + } + + // Next request is limited; Retry-After must be on the order of the + // window (≈50s), not the old hardcoded 1-second refill. + let res = svc + .call(Request::new( + http::Request::builder() + .uri("/") + .body(Body::empty()) + .unwrap(), + )) + .await + .unwrap(); + assert_eq!(res.status(), StatusCode::TOO_MANY_REQUESTS); + let retry_after: u64 = res + .headers() + .get("retry-after") + .unwrap() + .to_str() + .unwrap() + .parse() + .unwrap(); + // Time to the next token is window/capacity = 50/5 = 10s — an order of + // magnitude beyond the old hardcoded 1-second refill, proving the + // window parameter now drives the refill rate. + assert!( + (9..=50).contains(&retry_after), + "retry-after={retry_after}s should be ~window/capacity = 10s" + ); + } + + // ── M9: sharding + eviction ────────────────────────────────────────────── + + #[test] + fn idle_buckets_are_evicted() { + let store = BucketStore::new(Duration::from_secs(1)); + assert_eq!( + store.try_consume("stale-key", 10, Duration::from_secs(1)), + None + ); + + // Age every bucket past the TTL. + for shard in &store.shards { + for bucket in shard.lock().values_mut() { + bucket.last_seen -= IDLE_TTL + Duration::from_secs(1); + } + } + + store.evict_idle(); + let live: usize = store.shards.iter().map(|s| s.lock().len()).sum(); + assert_eq!(live, 0, "idle buckets should be gone"); + + // A fresh bucket is created for the next request. + assert_eq!( + store.try_consume("stale-key", 10, Duration::from_secs(1)), + None + ); + let live: usize = store.shards.iter().map(|s| s.lock().len()).sum(); + assert_eq!(live, 1); + } + + #[test] + fn distinct_keys_land_on_independent_locks_without_error() { + let store = Arc::new(BucketStore::new(Duration::from_secs(1))); + let handles: Vec<_> = (0..8) + .map(|i| { + let store = Arc::clone(&store); + std::thread::spawn(move || { + for j in 0..250 { + store.try_consume( + format!("key-{i}-{j}").as_str(), + 10, + Duration::from_secs(1), + ); + } + }) + }) + .collect(); + for handle in handles { + handle.join().unwrap(); + } + let total: usize = store.shards.iter().map(|s| s.lock().len()).sum(); + assert_eq!(total, 8 * 250); + } + + // ── test double ────────────────────────────────────────────────────────── + + #[derive(Clone)] + struct PassthroughService; + + impl Service for PassthroughService { + type Response = Response; + type Error = Infallible; + type Future = Pin> + Send + 'static>>; + + fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn call(&mut self, _: Request) -> Self::Future { + Box::pin(async { + Ok(http::Response::builder() + .status(StatusCode::OK) + .body(Body::empty()) + .unwrap()) + }) + } + } +} diff --git a/arvik-middleware/src/security_headers.rs b/arvik-middleware/src/security_headers.rs index 11aa7e1..92109a1 100644 --- a/arvik-middleware/src/security_headers.rs +++ b/arvik-middleware/src/security_headers.rs @@ -26,6 +26,7 @@ use std::convert::Infallible; use std::future::Future; use std::pin::Pin; +use std::sync::Arc; use std::task::{Context, Poll}; use arvik_core::{Request, Response}; @@ -416,9 +417,34 @@ impl Layer for SecurityHeadersLayer { type Service = SecurityHeadersService; fn layer(&self, inner: S) -> Self::Service { + // Every value is constant for the layer's lifetime — build the + // HeaderValues once here instead of cloning the whole config and + // re-formatting the HSTS string on every request. + let mut headers: Vec<(HeaderName, HeaderValue)> = Vec::with_capacity(7); + let mut add = |name: &'static str, value: &str| { + if let Ok(value) = HeaderValue::from_str(value) { + headers.push((HeaderName::from_static(name), value)); + } + }; + add("x-frame-options", self.frame_options); + add("x-content-type-options", "nosniff"); + add("x-xss-protection", "1; mode=block"); + add( + "strict-transport-security", + &format!("max-age={}; includeSubDomains", self.hsts_max_age), + ); + add("referrer-policy", "strict-origin-when-cross-origin"); + add( + "permissions-policy", + "geolocation=(), microphone=(), camera=(), payment=()", + ); + if let Some(csp) = &self.csp { + add("content-security-policy", csp); + } + SecurityHeadersService { inner, - config: self.clone(), + headers: headers.into(), } } } @@ -427,7 +453,7 @@ impl Layer for SecurityHeadersLayer { #[derive(Clone)] pub struct SecurityHeadersService { inner: S, - config: SecurityHeadersLayer, + headers: Arc<[(HeaderName, HeaderValue)]>, } impl Service for SecurityHeadersService @@ -446,41 +472,16 @@ where fn call(&mut self, req: Request) -> Self::Future { let cloned = self.inner.clone(); let mut inner = std::mem::replace(&mut self.inner, cloned); - let config = self.config.clone(); + let headers = Arc::clone(&self.headers); Box::pin(async move { let mut response = inner.call(req).await?; let h = response.headers_mut(); - - set_header(h, "x-frame-options", config.frame_options); - set_header(h, "x-content-type-options", "nosniff"); - set_header(h, "x-xss-protection", "1; mode=block"); - set_header( - h, - "strict-transport-security", - &format!("max-age={}; includeSubDomains", config.hsts_max_age), - ); - set_header(h, "referrer-policy", "strict-origin-when-cross-origin"); - set_header( - h, - "permissions-policy", - "geolocation=(), microphone=(), camera=(), payment=()", - ); - - if let Some(csp) = &config.csp { - set_header(h, "content-security-policy", csp); + // or_insert preserves handler-set headers, as before. + for (name, value) in headers.iter() { + h.entry(name.clone()).or_insert_with(|| value.clone()); } - Ok(response) }) } } - -fn set_header(headers: &mut http::HeaderMap, name: &str, value: &str) { - if let (Ok(n), Ok(v)) = ( - HeaderName::from_bytes(name.as_bytes()), - HeaderValue::from_str(value), - ) { - headers.entry(n).or_insert(v); - } -} diff --git a/arvik-middleware/src/trace.rs b/arvik-middleware/src/trace.rs index 7d2ab6c..9e8a651 100644 --- a/arvik-middleware/src/trace.rs +++ b/arvik-middleware/src/trace.rs @@ -32,7 +32,7 @@ use std::time::Instant; use arvik_core::{Request, Response}; use tower_layer::Layer; use tower_service::Service; -use tracing::{Level, Span}; +use tracing::{Instrument, Level, Span}; // ── LatencyUnit ─────────────────────────────────────────────────────────────── @@ -99,13 +99,14 @@ impl DefaultMakeSpan { let path = req.uri().path().to_owned(); let version = format!("{:?}", req.version()); - match self.level { + let span = match self.level { Level::ERROR => tracing::error_span!( "request", http.method = %method, http.path = %path, http.version = %version, http.status_code = tracing::field::Empty, + http.request_headers = tracing::field::Empty, latency = tracing::field::Empty, ), Level::WARN => tracing::warn_span!( @@ -114,6 +115,7 @@ impl DefaultMakeSpan { http.path = %path, http.version = %version, http.status_code = tracing::field::Empty, + http.request_headers = tracing::field::Empty, latency = tracing::field::Empty, ), Level::DEBUG => tracing::debug_span!( @@ -122,6 +124,7 @@ impl DefaultMakeSpan { http.path = %path, http.version = %version, http.status_code = tracing::field::Empty, + http.request_headers = tracing::field::Empty, latency = tracing::field::Empty, ), Level::TRACE => tracing::trace_span!( @@ -130,6 +133,7 @@ impl DefaultMakeSpan { http.path = %path, http.version = %version, http.status_code = tracing::field::Empty, + http.request_headers = tracing::field::Empty, latency = tracing::field::Empty, ), _ => tracing::info_span!( @@ -138,9 +142,23 @@ impl DefaultMakeSpan { http.path = %path, http.version = %version, http.status_code = tracing::field::Empty, + http.request_headers = tracing::field::Empty, latency = tracing::field::Empty, ), + }; + + if self.include_headers { + // Recorded as one compact field; values are echoed as received. + let headers = req + .headers() + .iter() + .map(|(name, value)| format!("{}={}", name, value.to_str().unwrap_or("?"))) + .collect::>() + .join("; "); + span.record("http.request_headers", headers.as_str()); } + + span } } @@ -233,10 +251,10 @@ where Box::pin(async move { let start = Instant::now(); - let response = { - let _guard = span.enter(); - inner.call(req).await? - }; + // .instrument ties the span to this future only — unlike + // span.enter(), it cannot corrupt the thread-local span stack + // while the task awaits across worker threads. + let response = inner.call(req).instrument(span.clone()).await?; let elapsed = start.elapsed(); let status = response.status(); diff --git a/arvik-observe/Cargo.toml b/arvik-observe/Cargo.toml index df5b50e..2ace570 100644 --- a/arvik-observe/Cargo.toml +++ b/arvik-observe/Cargo.toml @@ -37,7 +37,7 @@ uuid = { workspace = true, optional = true } [features] default = [] -metrics = ["dep:once_cell", "dep:prometheus"] +metrics = ["dep:once_cell", "dep:prometheus", "dep:tokio"] opentelemetry = [ "dep:opentelemetry", "dep:opentelemetry-otlp", diff --git a/arvik-observe/src/health.rs b/arvik-observe/src/health.rs index 36c7f19..71ffd98 100644 --- a/arvik-observe/src/health.rs +++ b/arvik-observe/src/health.rs @@ -10,7 +10,7 @@ use arvik_core::{Response, ResponseBuilder}; use futures_util::future::join_all; use http::StatusCode; use once_cell::sync::Lazy; -use parking_lot::RwLock; +use parking_lot::{Mutex, RwLock}; use serde::Serialize; const DEFAULT_CHECK_TIMEOUT: Duration = Duration::from_secs(2); @@ -29,6 +29,8 @@ struct HealthInner { checks: RwLock>, check_timeout: RwLock, startup_complete: AtomicBool, + readiness_cache_ttl: RwLock, + readiness_cache: parking_lot::Mutex)>>, } #[derive(Clone)] @@ -94,7 +96,12 @@ where fn into_health_check_result(self) -> HealthCheckResult { match self { Ok(_) => HealthCheckResult::ok(), - Err(err) => HealthCheckResult::unhealthy(err.to_string()), + Err(err) => { + // Driver/DSN text must not reach the public probe body during + // an outage; keep the detail in logs for operators. + tracing::error!(error = %err, "health check failed"); + HealthCheckResult::unhealthy("check failed") + } } } } @@ -107,7 +114,7 @@ struct ProbeResponse { checks: Vec, } -#[derive(Debug, Serialize)] +#[derive(Debug, Clone, Serialize)] struct CheckReport { name: String, status: &'static str, @@ -116,6 +123,10 @@ struct CheckReport { message: Option, } +/// How long a readiness probe result is reused before checks run again. +/// Bursts of kubelet/LB probes then cost one dependency round-trip instead of N. +const DEFAULT_READINESS_CACHE_TTL: Duration = Duration::from_secs(1); + impl HealthRegistry { /// Create a health registry. Startup is considered complete by default. pub fn new() -> Self { @@ -125,6 +136,8 @@ impl HealthRegistry { checks: RwLock::new(Vec::new()), check_timeout: RwLock::new(DEFAULT_CHECK_TIMEOUT), startup_complete: AtomicBool::new(true), + readiness_cache_ttl: RwLock::new(DEFAULT_READINESS_CACHE_TTL), + readiness_cache: Mutex::new(None), }), } } @@ -205,9 +218,34 @@ impl HealthRegistry { ) } + /// Set how long a readiness result is reused before checks run again + /// (default 1 s). A zero TTL disables caching. + pub fn set_readiness_cache_ttl(&self, ttl: Duration) { + *self.inner.readiness_cache_ttl.write() = ttl; + } + /// Build the `/health/ready` response. pub async fn readiness_response(&self) -> Response { - let checks = self.run_checks().await; + let ttl = *self.inner.readiness_cache_ttl.read(); + + let checks = if ttl.is_zero() { + self.run_checks().await + } else { + let cached = self + .inner + .readiness_cache + .lock() + .clone() + .filter(|(at, _)| at.elapsed() < ttl); + match cached { + Some((_, checks)) => checks, + None => { + let checks = self.run_checks().await; + *self.inner.readiness_cache.lock() = Some((Instant::now(), checks.clone())); + checks + } + } + }; let healthy = checks.iter().all(|check| check.status == "ok"); let status = if healthy { StatusCode::OK diff --git a/arvik-observe/src/logging.rs b/arvik-observe/src/logging.rs index 90574c4..4d6e198 100644 --- a/arvik-observe/src/logging.rs +++ b/arvik-observe/src/logging.rs @@ -107,6 +107,77 @@ impl ArvikLoggerBuilder { .try_init(), LogFormat::Auto => unreachable!("LogFormat::resolve never returns Auto"), } + // try_init already returns our exact LoggerError box type. + .map_err(|err| err as LoggerError) + } + + /// Install logging and OpenTelemetry as **one** subscriber. + /// + /// [`ArvikLogger::init`] and [`OtelConfig::install`](crate::trace::OtelConfig::install) + /// each try to claim the global subscriber, so calling both always fails + /// for whichever runs second. This method composes the log layer and the + /// OTel span layer into a single subscriber instead. The returned guard + /// shuts the tracer provider down on drop; on failure nothing is left + /// installed. + #[cfg(all(feature = "logging", feature = "opentelemetry"))] + pub fn init_with_otel( + self, + otel: crate::trace::OtelConfig, + ) -> Result { + use opentelemetry::global; + use opentelemetry::trace::TracerProvider as _; + use tracing_subscriber::layer::SubscriberExt; + use tracing_subscriber::util::SubscriberInitExt; + + let provider = otel + .build_provider() + .map_err(|err| -> LoggerError { Box::new(err) })?; + + let format = self.format.resolve(); + let env_filter = EnvFilter::try_new(self.env_filter_directive())?; + + let fmt_layer = tracing_subscriber::fmt::layer() + .with_target(self.with_target) + .with_thread_ids(self.with_thread_ids); + + // The Json/Pretty layers have distinct types, so each arm composes and + // installs its own subscriber. + let install_result = match format { + LogFormat::Json => { + let otel_layer = tracing_opentelemetry::layer() + .with_tracer(provider.tracer(otel.service_name().to_string())); + tracing_subscriber::registry() + .with(env_filter) + .with(fmt_layer.json()) + .with(otel_layer) + .try_init() + } + LogFormat::Pretty => { + let otel_layer = tracing_opentelemetry::layer() + .with_tracer(provider.tracer(otel.service_name().to_string())); + tracing_subscriber::registry() + .with(env_filter) + .with(fmt_layer.pretty()) + .with(otel_layer) + .try_init() + } + LogFormat::Auto => unreachable!("LogFormat::resolve never returns Auto"), + }; + + if let Err(err) = install_result { + // Mirror OtelConfig::install's failure hygiene: release exporter + // threads and drop the global reference before surfacing the error. + let _ = provider.shutdown(); + global::set_tracer_provider( + opentelemetry_sdk::trace::SdkTracerProvider::builder().build(), + ); + return Err(Box::new(crate::trace::OtelError::Subscriber( + err.to_string(), + ))); + } + + global::set_tracer_provider(provider.clone()); + Ok(crate::trace::OtelGuard::from_provider(provider)) } fn env_filter_directive(&self) -> String { @@ -239,8 +310,17 @@ where } fn call(&mut self, mut req: Request) -> Self::Future { - let method = req.method().as_str().to_string(); - let uri = req.uri().to_string(); + // Request-ID plumbing runs unconditionally (it is functional), but the + // span field strings are only built when INFO spans are actually + // enabled — otherwise every filtered-out request still paid several + // heap allocations. + let span_enabled = tracing::enabled!(tracing::Level::INFO); + + let (method, uri) = if span_enabled { + (req.method().as_str().to_string(), req.uri().to_string()) + } else { + (String::new(), String::new()) + }; let request_id = request_id(&req, &self.request_id_header); req.extensions_mut() @@ -250,13 +330,17 @@ where .insert(self.request_id_header.clone(), value); } - let span = request_span( - &method, - &uri, - &request_id, - self.include_headers - .then(|| masked_headers(req.headers(), &self.sensitive_headers)), - ); + let span = if span_enabled { + request_span( + &method, + &uri, + &request_id, + self.include_headers + .then(|| masked_headers(req.headers(), &self.sensitive_headers)), + ) + } else { + tracing::Span::none() + }; let request_id_header = self.request_id_header.clone(); let request_id_for_response = request_id.clone(); @@ -313,15 +397,30 @@ where } } +/// Longest accepted client-supplied request id. +const MAX_REQUEST_ID_LEN: usize = 128; + +/// True when the id is safe to echo into headers and every log line: +/// printable ASCII, no whitespace/control characters, bounded length. +/// Anything else gets replaced with a fresh UUID so attackers can neither +/// forge correlation keys nor inflate log volume. +fn is_valid_request_id(id: &str) -> bool { + !id.is_empty() + && id.len() <= MAX_REQUEST_ID_LEN + && id.bytes().all(|b| (0x21..=0x7e).contains(&b)) +} + fn request_id(req: &Request, header: &HeaderName) -> String { req.headers() .get(header) .and_then(|value| value.to_str().ok()) + .filter(|id| is_valid_request_id(id)) .map(str::to_owned) .or_else(|| { req.extensions() .get::() .map(|id| id.as_str().to_string()) + .filter(|id| is_valid_request_id(id)) }) .unwrap_or_else(|| Uuid::new_v4().to_string()) } diff --git a/arvik-observe/src/metrics.rs b/arvik-observe/src/metrics.rs index 4d62ed5..828cdaf 100644 --- a/arvik-observe/src/metrics.rs +++ b/arvik-observe/src/metrics.rs @@ -377,9 +377,19 @@ pub async fn metrics_handler() -> Response { /// Prometheus scrape handler for a specific registry. pub async fn metrics_handler_with_registry(registry: MetricsRegistry) -> Response { + // gather() locks every collector and deep-clones metric families, then + // the encoder walks them — CPU-bound work that must not run inline on + // the tokio worker handling the scrape. + let encoded = tokio::task::spawn_blocking(move || registry.encode_text()) + .await + .unwrap_or_else(|err| { + tracing::error!("metrics encode task failed: {}", err); + String::new() + }); + ResponseBuilder::new() .header(http::header::CONTENT_TYPE, MetricsRegistry::content_type()) - .body(registry.encode_text()) + .body(encoded) } type BodyCallback = dyn Fn(u64) + Send + Sync + 'static; diff --git a/arvik-observe/src/trace.rs b/arvik-observe/src/trace.rs index 8685847..3592c5f 100644 --- a/arvik-observe/src/trace.rs +++ b/arvik-observe/src/trace.rs @@ -97,20 +97,28 @@ where } fn call(&mut self, req: Request) -> Self::Future { - let method = req.method().as_str().to_owned(); - let url = req.uri().to_string(); - let user_agent = req - .headers() - .get(http::header::USER_AGENT) - .and_then(|value| value.to_str().ok()) - .unwrap_or("") - .to_owned(); - let parent = extract_parent(req.headers(), &self.propagators); - let span = make_http_span(&self.service_name, &method, &url, &user_agent); - - if let Some(parent) = parent { - let _ = span.set_parent(parent); - } + // Skip all span construction (including the field strings and the + // propagator extraction) when INFO spans are filtered out anyway. + let span = if !tracing::enabled!(tracing::Level::INFO) { + tracing::Span::none() + } else { + let method = req.method().as_str().to_owned(); + // Path only — query strings routinely carry credentials and must + // not land in traces by default. + let url = req.uri().path().to_string(); + let user_agent = req + .headers() + .get(http::header::USER_AGENT) + .and_then(|value| value.to_str().ok()) + .unwrap_or("") + .to_owned(); + let span = make_http_span(&self.service_name, &method, &url, &user_agent); + + if let Some(parent) = extract_parent(req.headers(), &self.propagators) { + let _ = span.set_parent(parent); + } + span + }; let cloned = self.inner.clone(); let mut inner = std::mem::replace(&mut self.inner, cloned); @@ -234,9 +242,14 @@ impl OtelConfig { let builder = SdkTracerProvider::builder().with_resource(resource); match &self.exporter { - Exporter::Stdout => Ok(builder - .with_simple_exporter(opentelemetry_stdout::SpanExporter::default()) - .build()), + Exporter::Stdout => { + // Batch, not simple: the simple processor serializes and + // writes each span inline on whichever task ends it — a + // synchronous stdout write on the request hot path. + Ok(builder + .with_batch_exporter(opentelemetry_stdout::SpanExporter::default()) + .build()) + } Exporter::OtlpGrpc(endpoint) => { let mut exporter = opentelemetry_otlp::SpanExporter::builder().with_tonic(); if let Some(endpoint) = endpoint { @@ -262,16 +275,38 @@ impl OtelConfig { } } + /// Return the configured service name. + pub fn service_name(&self) -> &str { + &self.service_name + } + /// Install the tracer provider and a `tracing-opentelemetry` subscriber layer. + /// + /// The global subscriber is installed first; if that fails (typically + /// because another subscriber — e.g. [`crate::logging::ArvikLogger`] — is + /// already installed), the freshly built provider is shut down and the + /// global reference released instead of leaking its exporter task. + /// + /// To combine structured logging with OpenTelemetry in one subscriber, + /// use `crate::logging::ArvikLogger::init_with_otel`. pub fn install(self) -> Result { let provider = self.build_provider()?; let tracer = provider.tracer(self.service_name.clone()); - global::set_tracer_provider(provider.clone()); - tracing_subscriber::registry() + let install_result = tracing_subscriber::registry() .with(tracing_opentelemetry::layer().with_tracer(tracer)) - .try_init() - .map_err(|err| OtelError::Subscriber(err.to_string()))?; + .try_init(); + + if let Err(err) = install_result { + // Stop exporter threads, then overwrite the global reference so + // nothing keeps the provider alive for the rest of the process. + let _ = provider.shutdown(); + global::set_tracer_provider(SdkTracerProvider::builder().build()); + return Err(OtelError::Subscriber(err.to_string())); + } + + // Only publish globally once installation actually succeeded. + global::set_tracer_provider(provider.clone()); Ok(OtelGuard { provider: Some(provider), @@ -284,6 +319,15 @@ pub struct OtelGuard { provider: Option, } +impl OtelGuard { + /// Wrap an externally built provider so its lifetime is managed here. + pub(crate) fn from_provider(provider: SdkTracerProvider) -> Self { + Self { + provider: Some(provider), + } + } +} + impl OtelGuard { /// Flush and shut down the tracer provider now. pub fn shutdown(mut self) -> Result<(), OtelError> { diff --git a/arvik-router/src/layer.rs b/arvik-router/src/layer.rs index e3074bf..1e6d520 100644 --- a/arvik-router/src/layer.rs +++ b/arvik-router/src/layer.rs @@ -31,7 +31,7 @@ pub struct BoxCloneService(Box); /// A closure that applies a Tower layer to a [`BoxCloneService`] and returns a new one. /// -/// Stored inside [`Router`] and [`MethodRouter`] to defer layer application until +/// Stored inside `Router` and `MethodRouter` to defer layer application until /// the service is actually built at serve time. /// /// ```rust,ignore @@ -43,7 +43,7 @@ pub type LayerFn = Arc BoxCloneService + Send + Sync // ── BoxCloneService ───────────────────────────────────────────────────────── /// Object-safe inner trait. -trait ErasedSvc: Send { +trait ErasedSvc: Send + Sync { fn call_erased( &mut self, req: Request, @@ -56,7 +56,7 @@ trait ErasedSvc: Send { impl ErasedSvc for S where - S: Service + Clone + Send + 'static, + S: Service + Clone + Send + Sync + 'static, S::Future: Send + 'static, { fn call_erased( @@ -79,7 +79,11 @@ impl BoxCloneService { /// Wrap a concrete service in a `BoxCloneService`. pub fn new(svc: S) -> Self where - S: Service + Clone + Send + 'static, + S: Service + + Clone + + Send + + Sync + + 'static, S::Future: Send + 'static, { Self(Box::new(svc)) @@ -116,18 +120,19 @@ impl std::fmt::Debug for BoxCloneService { /// Convert any Tower [`Layer`] into a [`LayerFn`]. /// -/// Used by [`Router::layer`], [`Router::route_layer`], and [`MethodRouter::layer`]. +/// Used by `Router::layer`, `Router::route_layer`, and `MethodRouter::layer`. /// /// # Bounds /// /// The resulting service `L::Service` must be: /// - `Service` -/// - `Clone + Send + 'static` +/// - `Clone + Send + Sync + 'static` /// - Its future must be `Send + 'static` pub fn into_layer_fn(layer: L) -> LayerFn where L: Layer + Clone + Send + Sync + 'static, - L::Service: Service + Clone + Send + 'static, + L::Service: + Service + Clone + Send + Sync + 'static, >::Future: Send + 'static, { Arc::new(move |svc: BoxCloneService| -> BoxCloneService { diff --git a/arvik-router/src/method_router.rs b/arvik-router/src/method_router.rs index 52028ab..cf0c18e 100644 --- a/arvik-router/src/method_router.rs +++ b/arvik-router/src/method_router.rs @@ -30,6 +30,9 @@ use crate::layer::{BoxCloneService, LayerFn, apply_layers, into_layer_fn, onesho // ── MethodRouter ──────────────────────────────────────────────────────────── +/// Shared, type-erased handler list. +type HandlerList = std::sync::Arc>)>>; + /// Stores one handler per HTTP method for a single route. /// /// Created via the top-level constructor functions [`get`], [`post`], etc. @@ -42,19 +45,42 @@ use crate::layer::{BoxCloneService, LayerFn, apply_layers, into_layer_fn, onesho /// .layer(RequireAuthLayer::new()); /// ``` pub struct MethodRouter { - /// (method_filter, type-erased handler) pairs. - handlers: Vec<(MethodFilter, Box>)>, + /// (method_filter, type-erased handler) pairs, shared behind an Arc so + /// cloning a router for per-request dispatch is an O(1) refcount bump + /// instead of deep-cloning every boxed handler. + handlers: HandlerList, /// Bitmask of all registered methods — used to build the `Allow` header. allow_methods: MethodFilter, /// Tower layers applied to each matched handler (innermost = first in vec). layers: Vec, } +/// Mutable access to a handler list that may be shared behind an Arc. +/// +/// During router construction the Arc is almost always unique, so this is a +/// no-op there; sharing only appears once built routers get cloned. +fn handlers_mut( + handlers: &mut HandlerList, +) -> &mut Vec<(MethodFilter, Box>)> +where + S: Clone + Send + Sync + 'static, +{ + if std::sync::Arc::get_mut(handlers).is_none() { + *handlers = std::sync::Arc::new( + handlers + .iter() + .map(|(filter, handler)| (*filter, handler.clone_box())) + .collect(), + ); + } + std::sync::Arc::get_mut(handlers).expect("just made unique") +} + impl MethodRouter { /// Create an empty `MethodRouter` with no handlers. pub fn new() -> Self { Self { - handlers: Vec::new(), + handlers: std::sync::Arc::new(Vec::new()), allow_methods: MethodFilter::NONE, layers: Vec::new(), } @@ -67,7 +93,7 @@ impl MethodRouter { T: 'static, { self.allow_methods |= filter; - self.handlers.push((filter, into_erased(handler))); + handlers_mut(&mut self.handlers).push((filter, into_erased(handler))); self } @@ -81,7 +107,8 @@ impl MethodRouter { } self.allow_methods |= other.allow_methods; - self.handlers.extend(other.handlers); + handlers_mut(&mut self.handlers) + .extend(other.handlers.iter().map(|(f, h)| (*f, h.clone_box()))); self.layers.extend(other.layers); self } @@ -173,8 +200,11 @@ impl MethodRouter { pub fn layer(mut self, layer: L) -> Self where L: Layer + Clone + Send + Sync + 'static, - L::Service: - Service + Clone + Send + 'static, + L::Service: Service + + Clone + + Send + + Sync + + 'static, >::Future: Send + 'static, { self.layers.push(into_layer_fn(layer)); @@ -184,14 +214,27 @@ impl MethodRouter { /// Dispatch the request to the matching method handler, applying any /// configured layers. /// + /// HEAD requests are served by the GET handler when no dedicated HEAD + /// handler is registered (RFC 9110 §9.3.2 — hyper strips the response + /// body for HEAD). + /// /// Returns `405 Method Not Allowed` (with an `Allow` header) if no handler - /// matches the request method. + /// matches the request method — including non-standard extension methods + /// (e.g. `PURGE`) unless the route registered [`MethodFilter::EXTENSION`] + /// or [`MethodFilter::ANY`]. pub async fn call(&self, req: Request, state: S) -> Response { let method = req.method().clone(); let method_filter = MethodFilter::from_method(&method); - for (filter, handler) in &self.handlers { - if filter.contains(method_filter) { + // RFC 9110 §9.3.2: a HEAD response is a GET response without a body, + // so fall back to the GET handler unless HEAD was explicitly bound. + let head_falls_back_to_get = + method == http::Method::HEAD && !self.allow_methods.contains(MethodFilter::HEAD); + + for (filter, handler) in self.handlers.iter() { + let matched = filter.contains(method_filter) + || (head_falls_back_to_get && filter.contains(MethodFilter::GET)); + if matched { if self.layers.is_empty() { // ── Fast path: no per-route layers ────────────────────── let h = handler.clone_box(); @@ -219,10 +262,20 @@ impl MethodRouter { .text("Method Not Allowed") } + /// Append already-erased layers to this router's per-route stack. + /// + /// Used by [`crate::Router::nest`] / [`crate::Router::merge`] to preserve + /// middleware attached inside a sub-router when its routes are flattened + /// into the parent: appended layers wrap the ones registered earlier, + /// matching the router-level (`layer` outside `route_layer`) ordering. + pub(crate) fn extend_layers(&mut self, extra: impl IntoIterator) { + self.layers.extend(extra); + } + /// Bind application state, converting `MethodRouter` → `MethodRouter<()>`. /// /// After calling this, the method router is ready to be served directly - /// or attached to a [`Router`] that is itself bound with [`Router::with_state`]. + /// or attached to a `Router` that is itself bound with `Router::with_state`. /// /// # Example /// @@ -235,17 +288,18 @@ impl MethodRouter { /// ``` pub fn with_state(self, state: S) -> MethodRouter<()> { let state = Arc::new(state); - let handlers = self - .handlers - .into_iter() - .map(|(filter, handler)| { - let bound: Box> = Box::new(StateBound { - inner: handler, - state: Arc::clone(&state), - }); - (filter, bound) - }) - .collect(); + let handlers = std::sync::Arc::new( + self.handlers + .iter() + .map(|(filter, handler)| { + let bound: Box> = Box::new(StateBound { + inner: handler.clone_box(), + state: Arc::clone(&state), + }); + (*filter, bound) + }) + .collect(), + ); MethodRouter { handlers, @@ -257,14 +311,11 @@ impl MethodRouter { impl Clone for MethodRouter { fn clone(&self) -> Self { + // handlers live behind an Arc: no per-handler clone_box here. Self { - handlers: self - .handlers - .iter() - .map(|(f, h)| (*f, h.clone_box())) - .collect(), + handlers: std::sync::Arc::clone(&self.handlers), allow_methods: self.allow_methods, - layers: self.layers.clone(), // Arc — O(n) cheap clone + layers: self.layers.clone(), // LayerFn is Arc-backed — O(n) cheap } } } @@ -436,3 +487,79 @@ fn build_allow_header(filter: MethodFilter) -> String { .collect::>() .join(", ") } + +#[cfg(test)] +mod tests { + use super::*; + use arvik_core::Body; + + async fn get_handler() -> &'static str { + "get-body" + } + + async fn delete_handler() -> &'static str { + "deleted" + } + + async fn head_handler() -> &'static str { + "head-body" + } + + async fn any_handler() -> &'static str { + "any-body" + } + + fn request(method: &str) -> Request { + Request::new( + http::Request::builder() + .method(method) + .uri("/") + .body(Body::empty()) + .unwrap(), + ) + } + + async fn body_string(res: Response) -> String { + res.into_body().to_string().await.unwrap() + } + + #[tokio::test] + async fn unknown_method_returns_405_not_the_registered_handler() { + let router = MethodRouter::new().delete(delete_handler); + let res = router.call(request("PURGE"), ()).await; + assert_eq!(res.status(), StatusCode::METHOD_NOT_ALLOWED); + assert_eq!(res.headers().get(http::header::ALLOW).unwrap(), "DELETE"); + assert_eq!(body_string(res).await, "Method Not Allowed"); + } + + #[tokio::test] + async fn known_method_mismatch_still_returns_405() { + let router = MethodRouter::new().delete(delete_handler); + let res = router.call(request("POST"), ()).await; + assert_eq!(res.status(), StatusCode::METHOD_NOT_ALLOWED); + } + + #[tokio::test] + async fn any_route_still_matches_extension_methods() { + let router = MethodRouter::new().on(MethodFilter::ANY, any_handler); + let res = router.call(request("PURGE"), ()).await; + assert_eq!(res.status(), StatusCode::OK); + assert_eq!(body_string(res).await, "any-body"); + } + + #[tokio::test] + async fn head_falls_back_to_get_handler() { + let router = MethodRouter::new().get(get_handler); + let res = router.call(request("HEAD"), ()).await; + assert_eq!(res.status(), StatusCode::OK); + assert_eq!(body_string(res).await, "get-body"); + } + + #[tokio::test] + async fn dedicated_head_handler_takes_precedence_over_get() { + let router = MethodRouter::new().get(get_handler).head(head_handler); + let res = router.call(request("HEAD"), ()).await; + assert_eq!(res.status(), StatusCode::OK); + assert_eq!(body_string(res).await, "head-body"); + } +} diff --git a/arvik-router/src/router.rs b/arvik-router/src/router.rs index 6a4df8c..f2e65b1 100644 --- a/arvik-router/src/router.rs +++ b/arvik-router/src/router.rs @@ -13,9 +13,9 @@ //! //! After calling [`Router::with_state`] you receive a `Router<()>`. Call //! [`Router::into_service`] to obtain a [`BoxCloneService`] with all layers -//! baked in, ready to be passed to [`arvik_hyper::Server::serve_service`]. +//! baked in, ready to be passed to `arvik_hyper::Server::serve_service`. //! -//! [`arvik_hyper::serve_app`] calls `into_service()` internally. +//! `arvik_hyper::serve_app` calls `into_service()` internally. use std::convert::Infallible; use std::future::Future; @@ -160,7 +160,6 @@ pub struct Router { struct RouteEntry { pattern: Arc, method_router: MethodRouter, - param_names: Arc<[Arc]>, } impl RouteEntry { @@ -168,7 +167,6 @@ impl RouteEntry { Self { pattern: Arc::from(path), method_router, - param_names: extract_param_names(path).into(), } } } @@ -224,8 +222,32 @@ impl Router { } /// Mount a sub-router under a path prefix (flatten strategy). + /// + /// Middleware attached to the sub-router (`.layer()` / `.route_layer()`) + /// is preserved by folding it onto each flattened route. Note the + /// flattening caveat: those layers run only on **matched** sub-router + /// routes — unmatched paths in the prefix fall through to this router's + /// fallback, like any other unmatched path. + /// + /// # Panics + /// + /// Panics if the sub-router has its own `.fallback()`: flattened routing + /// cannot scope a fallback to the prefix. Set `.fallback()` on the outer + /// router instead. pub fn nest(mut self, prefix: &str, other: Router) -> Self { let prefix = prefix.trim_end_matches('/'); + if other.fallback.is_some() { + panic!( + "nest(\"{prefix}\"): the sub-router has a `.fallback()`, which cannot be \ + scoped under a prefix when routes are flattened. Set `.fallback()` on the \ + outer router instead." + ); + } + // Preserve middleware attached inside the sub-router by folding it + // onto every flattened route's method-router layer stack: + // handler ← method layers ← sub route_layers ← sub layers. + let mut sub_layers = other.route_layers; + sub_layers.extend(other.layers); for entry in other.routes { let path = entry.pattern.as_ref(); let full = if path == "/" { @@ -237,20 +259,31 @@ impl Router { if let Err(e) = self.trie.insert(&full, idx) { panic!("Nested route conflict for `{full}`: {e}"); } - self.routes - .push(RouteEntry::new(&full, entry.method_router)); + let mut method_router = entry.method_router; + method_router.extend_layers(sub_layers.iter().cloned()); + self.routes.push(RouteEntry::new(&full, method_router)); } self } /// Merge all routes from another router into this one. + /// + /// Middleware attached to the merged router is preserved by folding it + /// onto each merged route; a merged `.fallback()` is adopted when this + /// router has none. pub fn merge(mut self, other: Router) -> Self { + let mut sub_layers = other.route_layers; + sub_layers.extend(other.layers); for entry in other.routes { let path = entry.pattern.as_ref(); let idx = self.routes.len(); if let Err(e) = self.trie.insert(path, idx) { panic!("Merge conflict for `{path}`: {e}"); } + let mut entry = entry; + entry + .method_router + .extend_layers(sub_layers.iter().cloned()); self.routes.push(entry); } if self.fallback.is_none() { @@ -322,8 +355,11 @@ impl Router { pub fn layer(mut self, layer: L) -> Self where L: Layer + Clone + Send + Sync + 'static, - L::Service: - Service + Clone + Send + 'static, + L::Service: Service + + Clone + + Send + + Sync + + 'static, >::Future: Send + 'static, { self.layers.push(into_layer_fn(layer)); @@ -332,7 +368,7 @@ impl Router { /// Apply a Tower middleware layer to **matched routes only**. /// - /// Unlike [`layer`], this does not run on 404/405 responses. + /// Unlike `layer`, this does not run on 404/405 responses. /// Processes the request **after** outer layers but **before** per-method layers. /// /// ```rust,ignore @@ -344,8 +380,11 @@ impl Router { pub fn route_layer(mut self, layer: L) -> Self where L: Layer + Clone + Send + Sync + 'static, - L::Service: - Service + Clone + Send + 'static, + L::Service: Service + + Clone + + Send + + Sync + + 'static, >::Future: Send + 'static, { self.route_layers.push(into_layer_fn(layer)); @@ -356,8 +395,8 @@ impl Router { /// Bind application state, converting `Router` → `Router<()>`. /// - /// Call this once before passing your router to [`serve_app`] or - /// [`into_service`]. + /// Call this once before passing your router to `serve_app` or + /// `into_service`. pub fn with_state(self, state: S) -> Router<()> { let state = Arc::new(state); @@ -367,7 +406,6 @@ impl Router { .map(|entry| RouteEntry { pattern: entry.pattern, method_router: entry.method_router.with_state((*state).clone()), - param_names: entry.param_names, }) .collect(); @@ -392,7 +430,7 @@ impl Router { /// Dispatch a request directly, **without** applying outer layers. /// /// This method exists for backward compatibility and internal testing. - /// For production serving use [`into_service`] (via [`serve_app`]). + /// For production serving use `into_service` (via `serve_app`). pub async fn call(&self, mut req: Request, state: S) -> Response { let path = req.uri().path(); @@ -404,13 +442,11 @@ impl Router { if !matched.params.is_empty() { let mut pp = PathParams::new(); - for (idx, (k, v)) in matched.params.iter().enumerate() { - let key = entry - .param_names - .get(idx) - .cloned() - .unwrap_or_else(|| Arc::from(k)); - pp.push(key, percent_decode(v)); + // matchit's keys are the declared names — including + // mid-segment params like `{id}.png` that positional + // pattern parsing cannot align with. + for (k, v) in matched.params.iter() { + pp.push(Arc::from(k), percent_decode(v).into_owned()); } req.extensions_mut().insert(pp); } @@ -459,12 +495,29 @@ impl Router<()> { let route_layers = self.route_layers; let outer_layers = self.layers; - // Inner core: route dispatch + route_layers + // Pre-fold the router-level route_layers into each entry ONCE, so the + // per-request path clones a ready-made service instead of deep-cloning + // the MethodRouter and re-wrapping every layer on every request. + let layered_routes: Vec> = self + .routes + .iter() + .map(|entry| { + if route_layers.is_empty() { + None + } else { + let base = + BoxCloneService::new(MethodRouterService(entry.method_router.clone())); + Some(apply_layers(base, &route_layers)) + } + }) + .collect(); + + // Inner core: route dispatch + prebuilt layered services let inner = Arc::new(RouterInner { trie: self.trie, routes: self.routes, fallback: self.fallback, - route_layers, + layered_routes, }); // Wrap in a clone-friendly Tower service @@ -489,7 +542,9 @@ struct RouterInner { trie: matchit::Router, routes: Vec>, fallback: Option>>, - route_layers: Vec, + /// Pre-folded `route_layer(method_router)` services, index-aligned with + /// `routes`; `None` when the route has no router-level layers. + layered_routes: Vec>, } impl RouterInner { @@ -504,25 +559,21 @@ impl RouterInner { if !matched.params.is_empty() { let mut pp = PathParams::new(); - for (idx, (k, v)) in matched.params.iter().enumerate() { - let key = entry - .param_names - .get(idx) - .cloned() - .unwrap_or_else(|| Arc::from(k)); - pp.push(key, percent_decode(v)); + // matchit's keys are the declared names — including + // mid-segment params like `{id}.png` that positional + // pattern parsing cannot align with. + for (k, v) in matched.params.iter() { + pp.push(Arc::from(k), percent_decode(v).into_owned()); } req.extensions_mut().insert(pp); } req.extensions_mut().insert(MatchedPathExt(pattern.clone())); - let mut response = if self.route_layers.is_empty() { - entry.method_router.call(req, ()).await - } else { - let base = - BoxCloneService::new(MethodRouterService(entry.method_router.clone())); - let svc = apply_layers(base, &self.route_layers); - oneshot(svc, req).await + // Cloning the prebuilt stack is an O(1)-per-wrapper bump now + // that MethodRouter shares its handlers behind an Arc. + let mut response = match self.layered_routes[idx].as_ref() { + Some(svc) => oneshot(svc.clone(), req).await, + None => entry.method_router.call(req, ()).await, }; response.extensions_mut().insert(MatchedPathExt(pattern)); response @@ -614,28 +665,32 @@ fn normalize_service_prefix(prefix: &str) -> String { } } -fn extract_param_names(path: &str) -> Vec> { - path.split('/') - .filter_map(|segment| { - let inner = segment.strip_prefix('{')?.strip_suffix('}')?; - let name = inner.strip_prefix('*').unwrap_or(inner); - (!name.is_empty()).then(|| Arc::from(name)) - }) - .collect() -} - /// Percent-decode a URL path segment in-place (ASCII-only fast path). -fn percent_decode(input: &str) -> String { +/// Percent-decode a URL path segment (ASCII-only fast path). +/// +/// Decoding happens *after* route matching, so decoded values are handed to +/// handlers verbatim. Sequences that would change segment structure or smuggle +/// control characters are therefore kept encoded: `%2F` (`/`) and `%5C` (`\`) +/// cannot introduce separators into captured params, and `%00`, `%0A`, `%0D` +/// stay escaped so NUL/CRLF cannot reach handlers, filesystems, or logs +/// through decoded values. Literal separators in the URL are unaffected. +fn percent_decode(input: &str) -> std::borrow::Cow<'_, str> { let bytes = input.as_bytes(); if !bytes.contains(&b'%') { - return input.to_string(); // fast path: nothing to decode + return std::borrow::Cow::Borrowed(input); // fast path: nothing to decode } let mut out = Vec::with_capacity(bytes.len()); let mut i = 0; while i < bytes.len() { if bytes[i] == b'%' && i + 2 < bytes.len() { if let (Some(hi), Some(lo)) = (hex_nibble(bytes[i + 1]), hex_nibble(bytes[i + 2])) { - out.push(hi << 4 | lo); + let decoded = hi << 4 | lo; + if matches!(decoded, b'/' | b'\\' | 0x00 | b'\n' | b'\r') { + // Preserve the escape — see doc comment. + out.extend_from_slice(&bytes[i..i + 3]); + } else { + out.push(decoded); + } i += 3; continue; } @@ -643,7 +698,7 @@ fn percent_decode(input: &str) -> String { out.push(bytes[i]); i += 1; } - String::from_utf8(out).unwrap_or_else(|_| input.to_string()) + std::borrow::Cow::Owned(String::from_utf8(out).unwrap_or_else(|_| input.to_string())) } #[inline] @@ -665,16 +720,23 @@ mod tests { #[test] fn test_percent_decode() { assert_eq!(percent_decode("hello%20world"), "hello world"); - assert_eq!(percent_decode("foo%2Fbar"), "foo/bar"); assert_eq!(percent_decode("normal"), "normal"); assert_eq!(percent_decode("100%25"), "100%"); } #[test] - fn extracts_param_names_at_registration_time() { - let names = extract_param_names("/users/{id}/files/{*path}"); - let names: Vec<_> = names.iter().map(|name| name.as_ref()).collect(); - assert_eq!(names, vec!["id", "path"]); + fn percent_decode_keeps_structure_changing_escapes_encoded() { + // Encoded separators must not become real separators post-match. + assert_eq!(percent_decode("foo%2Fbar"), "foo%2Fbar"); + assert_eq!(percent_decode("a%5Cb"), "a%5Cb"); + // Control characters stay escaped. + assert_eq!(percent_decode("a%00b"), "a%00b"); + assert_eq!(percent_decode("a%0d%0Ab"), "a%0d%0Ab"); + // Ordinary decodes still work. + assert_eq!(percent_decode("a%2Eb"), "a.b"); + assert_eq!(percent_decode("%e2%82%ac"), "\u{20ac}"); + // Invalid UTF-8 falls back to the input unchanged. + assert_eq!(percent_decode("%ff"), "%ff"); } #[test] @@ -857,7 +919,35 @@ mod tests { } #[tokio::test] - async fn param_route_uses_interned_names_and_decodes_values() { + async fn mid_segment_params_expose_their_declared_names() { + async fn h(req: arvik_core::Request) -> String { + let params = req.extensions().get::().unwrap(); + format!( + "{}:{}", + params.get("id").unwrap_or("MISSING"), + params.get("tag").unwrap_or("MISSING") + ) + } + + // matchit supports mid-segment params; positional pattern parsing used + // to misalign these (audit L2). + let mut app = Router::new() + .route("/images/img{id}.png/{tag}", crate::get(h)) + .into_service(); + let req = Request::new( + http::Request::builder() + .uri("/images/img7.png/nature") + .body(arvik_core::Body::empty()) + .unwrap(), + ); + + let res = app.call(req).await.unwrap(); + let text = String::from_utf8(res.into_body().to_bytes().await.unwrap().to_vec()).unwrap(); + assert_eq!(text, "7:nature"); + } + + #[tokio::test] + async fn param_route_decodes_values() { async fn h(req: arvik_core::Request) -> String { let params = req.extensions().get::().unwrap(); format!( @@ -901,4 +991,132 @@ mod tests { let b = Router::new().route("/users", crate::get(h)); let _: Router<()> = a.merge(b); } + + // ── nest/merge middleware preservation (audit H3) ──────────────────────── + + /// Layer that appends its tag to the `x-mark` response header. + #[derive(Clone)] + struct MarkLayer(&'static str); + + impl Layer for MarkLayer { + type Service = MarkService; + + fn layer(&self, inner: BoxCloneService) -> Self::Service { + MarkService { inner, tag: self.0 } + } + } + + #[derive(Clone)] + struct MarkService { + inner: BoxCloneService, + tag: &'static str, + } + + impl Service for MarkService { + type Response = Response; + type Error = Infallible; + type Future = Pin> + Send + 'static>>; + + fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn call(&mut self, req: Request) -> Self::Future { + let cloned = self.inner.clone(); + let mut inner = std::mem::replace(&mut self.inner, cloned); + let value = self.tag.parse().unwrap(); + Box::pin(async move { + let mut res = inner + .call(req) + .await + .unwrap_or_else(|infallible| match infallible {}); + res.headers_mut().append("x-mark", value); + Ok(res) + }) + } + } + + async fn marked_handler() -> &'static str { + "inner" + } + + fn get_req(path: &str) -> Request { + Request::new( + http::Request::builder() + .method(http::Method::GET) + .uri(path) + .body(arvik_core::Body::empty()) + .unwrap(), + ) + } + + #[tokio::test] + async fn nest_preserves_subrouter_layers_and_route_layers() { + let sub = Router::new() + .route("/inner", crate::get(marked_handler)) + .route_layer(MarkLayer("route")) + .layer(MarkLayer("wide")); + let app: Router<()> = Router::new().nest("/sub", sub); + + // Matched nested route carries both sub-router layers (outer first). + let res = app.call(get_req("/sub/inner"), ()).await; + assert_eq!(res.status(), StatusCode::OK); + let marks: Vec<_> = res + .headers() + .get_all("x-mark") + .iter() + .filter_map(|v| v.to_str().ok()) + .collect(); + assert_eq!(marks, vec!["route", "wide"]); + } + + #[tokio::test] + async fn nested_layers_do_not_leak_outside_the_prefix_or_to_unmatched_paths() { + async fn outer_handler() -> &'static str { + "outer" + } + let sub = Router::new() + .route("/inner", crate::get(marked_handler)) + .layer(MarkLayer("sub")); + let app: Router<()> = Router::new() + .route("/own", crate::get(outer_handler)) + .nest("/sub", sub); + + // Parent's own route is untouched by sub-router layers. + let res = app.call(get_req("/own"), ()).await; + assert!(res.headers().get("x-mark").is_none()); + + // Unmatched path under the prefix: no sub-router layer runs (404). + let res = app.call(get_req("/sub/missing"), ()).await; + assert_eq!(res.status(), StatusCode::NOT_FOUND); + assert!(res.headers().get("x-mark").is_none()); + } + + #[test] + #[should_panic(expected = "cannot be scoped")] + fn nest_rejects_subrouter_fallback_instead_of_dropping_it() { + async fn h() -> &'static str { + "ok" + } + let sub: Router<()> = Router::new().route("/inner", crate::get(h)).fallback(h); + let _: Router<()> = Router::new().nest("/sub", sub); + } + + #[tokio::test] + async fn merge_preserves_merged_router_layers() { + let sub = Router::new() + .route("/a", crate::get(marked_handler)) + .route_layer(MarkLayer("merged")); + let app: Router<()> = Router::new().merge(sub); + + let res = app.call(get_req("/a"), ()).await; + assert_eq!(res.status(), StatusCode::OK); + let marks: Vec<&str> = res + .headers() + .get_all("x-mark") + .iter() + .filter_map(|v| v.to_str().ok()) + .collect(); + assert_eq!(marks, vec!["merged"]); + } } diff --git a/arvik-router/src/service.rs b/arvik-router/src/service.rs index eff68a0..1ff0355 100644 --- a/arvik-router/src/service.rs +++ b/arvik-router/src/service.rs @@ -1,9 +1,9 @@ //! Tower service adapter for the router. //! //! Provides [`ServiceHandler`] which wraps any Tower `Service` -//! into an Arvik [`Handler`], enabling services to be mounted -//! inside the router via [`Router::route_service`] and -//! [`Router::nest_service`]. +//! into an Arvik `Handler`, enabling services to be mounted +//! inside the router via `Router::route_service` and +//! `Router::nest_service`. use std::convert::Infallible; use std::future::Future; @@ -15,7 +15,7 @@ use arvik_core::response::Response; use http::Uri; use tower_service::Service; -/// Wraps a Tower [`Service`] to implement Arvik's [`Handler`] trait. +/// Wraps a Tower [`Service`] to implement Arvik's `Handler` trait. /// /// This adapter allows any compatible Tower service to be used /// as a route handler within the router. @@ -121,7 +121,9 @@ fn strip_uri_prefix(uri: &Uri, prefix: &str) -> Option { let path = uri.path(); let stripped_path = if prefix == "/" { path - } else if path == prefix || path == format!("{prefix}/") { + } else if path == prefix + || (path.len() == prefix.len() + 1 && path.starts_with(prefix) && path.ends_with('/')) + { "/" } else { path.strip_prefix(prefix) diff --git a/arvik-sse/src/event.rs b/arvik-sse/src/event.rs index 5829379..d99f569 100644 --- a/arvik-sse/src/event.rs +++ b/arvik-sse/src/event.rs @@ -83,10 +83,18 @@ impl Event { /// The browser uses this to send `Last-Event-ID` on reconnect so you /// can resume from where the stream left off. /// - /// `U+0000 NULL` bytes are stripped — browsers silently ignore ids - /// that contain null (SSE spec §9.2). + /// `U+0000 NULL` bytes are stripped, and the value is truncated at the + /// first CR/LF: an id containing `\n\n` would terminate the event framing + /// and silently drop every field serialized after it — an injection vector + /// whenever ids derive from untrusted input (e.g. an echoed + /// `Last-Event-ID` header). pub fn id(mut self, id: impl Into) -> Self { - let s: String = id.into().chars().filter(|&c| c != '\0').collect(); + let s: String = id + .into() + .chars() + .take_while(|&c| c != '\n' && c != '\r') + .filter(|&c| c != '\0') + .collect(); self.id = Some(s); self } diff --git a/arvik-static/src/lib.rs b/arvik-static/src/lib.rs index 4292e22..291473e 100644 --- a/arvik-static/src/lib.rs +++ b/arvik-static/src/lib.rs @@ -42,6 +42,9 @@ pub struct ServeDir { directory_listing: bool, chunk_size: usize, cache_control: Option, + /// When false (default), requests resolving through a symlink to outside + /// the (canonicalized) root are rejected. + follow_symlinks: bool, } #[cfg(feature = "fs")] @@ -58,9 +61,20 @@ impl ServeDir { directory_listing: false, chunk_size: DEFAULT_CHUNK_SIZE, cache_control: None, + follow_symlinks: false, } } + /// Allow serving paths that resolve through symlinks. + /// + /// By default, a request whose resolved location escapes the serve root + /// via a symlink is treated as not found. Enable this only when the tree + /// is trusted (e.g. intentional mounts). + pub fn follow_symlinks(mut self, enabled: bool) -> Self { + self.follow_symlinks = enabled; + self + } + /// Set a service called when a file is not found. pub fn not_found_service(mut self, service: S) -> Self where @@ -139,6 +153,14 @@ impl ServeDir { }; let candidate = self.root.join(&relative.fs_path); + + // Symlink containment: a planted link must not serve bytes from + // outside the (canonicalized) root. Checked on the canonical path so + // escapes through intermediate directories are caught too. + if !self.follow_symlinks && !path_contained(&self.root, &candidate).await { + return self.call_fallback(req).await; + } + let metadata = match tokio::fs::metadata(&candidate).await { Ok(metadata) => metadata, Err(_) => return self.call_fallback(req).await, @@ -217,50 +239,48 @@ impl ServeDir { asset_path: &str, accepts_br: bool, accepts_gzip: bool, - ) -> Result { + ) -> Result { let content_type = content_type(asset_path); if accepts_br { let br_path = append_suffix(&path, ".br"); - if let Ok(metadata) = tokio::fs::metadata(&br_path).await - && metadata.is_file() + if tokio::fs::metadata(&br_path) + .await + .is_ok_and(|m| m.is_file()) { - return Ok(FsAsset::new( - br_path, - metadata, + return Ok(FsCandidate { + path: br_path, content_type, - Some(HeaderValue::from_static("br")), - self.precompressed_gzip || self.precompressed_br, - )); + content_encoding: Some(HeaderValue::from_static("br")), + vary_accept_encoding: self.precompressed_gzip || self.precompressed_br, + }); } } if accepts_gzip { let gz_path = append_suffix(&path, ".gz"); - if let Ok(metadata) = tokio::fs::metadata(&gz_path).await - && metadata.is_file() + if tokio::fs::metadata(&gz_path) + .await + .is_ok_and(|m| m.is_file()) { - return Ok(FsAsset::new( - gz_path, - metadata, + return Ok(FsCandidate { + path: gz_path, content_type, - Some(HeaderValue::from_static("gzip")), - self.precompressed_gzip || self.precompressed_br, - )); + content_encoding: Some(HeaderValue::from_static("gzip")), + vary_accept_encoding: self.precompressed_gzip || self.precompressed_br, + }); } } - let metadata = tokio::fs::metadata(&path).await.map_err(|_| ())?; - if !metadata.is_file() { + if !tokio::fs::metadata(&path).await.is_ok_and(|m| m.is_file()) { return Err(()); } - Ok(FsAsset::new( + Ok(FsCandidate { path, - metadata, content_type, - None, - self.precompressed_gzip || self.precompressed_br, - )) + content_encoding: None, + vary_accept_encoding: self.precompressed_gzip || self.precompressed_br, + }) } async fn call_fallback(&self, req: Request) -> Response { @@ -374,48 +394,49 @@ impl ServeFile { } } - async fn select_file(&self, accepts_br: bool, accepts_gzip: bool) -> Result { + async fn select_file(&self, accepts_br: bool, accepts_gzip: bool) -> Result { if accepts_br { let br_path = append_suffix(&self.path, ".br"); - if let Ok(metadata) = tokio::fs::metadata(&br_path).await - && metadata.is_file() + if tokio::fs::metadata(&br_path) + .await + .is_ok_and(|m| m.is_file()) { - return Ok(FsAsset::new( - br_path, - metadata, - self.mime.clone(), - Some(HeaderValue::from_static("br")), - self.precompressed_gzip || self.precompressed_br, - )); + return Ok(FsCandidate { + path: br_path, + content_type: self.mime.clone(), + content_encoding: Some(HeaderValue::from_static("br")), + vary_accept_encoding: self.precompressed_gzip || self.precompressed_br, + }); } } if accepts_gzip { let gz_path = append_suffix(&self.path, ".gz"); - if let Ok(metadata) = tokio::fs::metadata(&gz_path).await - && metadata.is_file() + if tokio::fs::metadata(&gz_path) + .await + .is_ok_and(|m| m.is_file()) { - return Ok(FsAsset::new( - gz_path, - metadata, - self.mime.clone(), - Some(HeaderValue::from_static("gzip")), - self.precompressed_gzip || self.precompressed_br, - )); + return Ok(FsCandidate { + path: gz_path, + content_type: self.mime.clone(), + content_encoding: Some(HeaderValue::from_static("gzip")), + vary_accept_encoding: self.precompressed_gzip || self.precompressed_br, + }); } } - let metadata = tokio::fs::metadata(&*self.path).await.map_err(|_| ())?; - if !metadata.is_file() { + if !tokio::fs::metadata(&*self.path) + .await + .is_ok_and(|m| m.is_file()) + { return Err(()); } - Ok(FsAsset::new( - (*self.path).clone(), - metadata, - self.mime.clone(), - None, - self.precompressed_gzip || self.precompressed_br, - )) + Ok(FsCandidate { + path: (*self.path).clone(), + content_type: self.mime.clone(), + content_encoding: None, + vary_accept_encoding: self.precompressed_gzip || self.precompressed_br, + }) } } @@ -719,67 +740,68 @@ async fn call_boxed(mut service: BoxCloneService, req: Request) -> Response { .unwrap_or_else(|infallible| match infallible {}) } +/// A filesystem candidate selected for serving — everything except metadata. +/// +/// Metadata (length, mtime, ETag) is read from the *open* file at serving +/// time, closing the stat/open window where a deploy could swap or truncate +/// the file between selecting it and streaming it. #[cfg(feature = "fs")] -struct FsAsset { +struct FsCandidate { path: PathBuf, - len: u64, - modified: Option, - etag: HeaderValue, content_type: HeaderValue, content_encoding: Option, vary_accept_encoding: bool, } -#[cfg(feature = "fs")] -impl FsAsset { - fn new( - path: PathBuf, - metadata: std::fs::Metadata, - content_type: HeaderValue, - content_encoding: Option, - vary_accept_encoding: bool, - ) -> Self { - let modified = metadata.modified().ok(); - let etag = fs_etag(metadata.len(), modified); - Self { - path, - len: metadata.len(), - modified, - etag, - content_type, - content_encoding, - vary_accept_encoding, - } - } -} - #[cfg(feature = "fs")] async fn serve_fs_asset( - asset: FsAsset, + candidate: FsCandidate, req: Request, head: bool, chunk_size: usize, cache_control: Option, ) -> Response { - if is_not_modified(&req, &asset.etag, asset.modified) { + // Open first, then stat the open descriptor: the validators below are + // guaranteed to describe exactly the bytes this call will stream. + let file = match tokio::fs::File::open(&candidate.path).await { + Ok(file) => file, + Err(_) => return not_found(), + }; + let metadata = match file.metadata().await { + Ok(metadata) => metadata, + Err(_) => return not_found(), + }; + + let len = metadata.len(); + let modified = metadata.modified().ok(); + let etag = fs_etag(len, modified); + + if is_not_modified(&req, &etag, modified) { return not_modified( - &asset.etag, - asset.modified, + &etag, + modified, cache_control, - asset.vary_accept_encoding, + candidate.vary_accept_encoding, ); } - let range = match parse_range(req.headers().get(header::RANGE), asset.len) { - RangeDecision::Full => None, - RangeDecision::Partial(range) => Some(range), - RangeDecision::Unsatisfiable => return range_not_satisfiable(asset.len), + // RFC 9110 §13.1.5: If-Range must gate the range — a stale validator + // means the client's cached copy predates this file, so only a full 200 + // response can be safely stitched onto it. + let range = if if_range_allows_partial(&req, &etag, modified) { + match parse_range(req.headers().get(header::RANGE), len) { + RangeDecision::Full => None, + RangeDecision::Partial(range) => Some(range), + RangeDecision::Unsatisfiable => return range_not_satisfiable(len), + } + } else { + None }; let body_len = range .as_ref() .map(|range| range.end - range.start + 1) - .unwrap_or(asset.len); + .unwrap_or(len); let mut builder = ResponseBuilder::new() .status(if range.is_some() { @@ -787,27 +809,27 @@ async fn serve_fs_asset( } else { StatusCode::OK }) - .header(header::CONTENT_TYPE, asset.content_type) + .header(header::CONTENT_TYPE, candidate.content_type) .header(header::CONTENT_LENGTH, body_len.to_string()) .header(header::ACCEPT_RANGES, "bytes") - .header(header::ETAG, asset.etag.clone()); + .header(header::ETAG, etag.clone()); - if let Some(modified) = asset.modified { + if let Some(modified) = modified { builder = builder.header(header::LAST_MODIFIED, httpdate::fmt_http_date(modified)); } if let Some(cache_control) = cache_control { builder = builder.header(header::CACHE_CONTROL, cache_control); } - if let Some(content_encoding) = asset.content_encoding { + if let Some(content_encoding) = candidate.content_encoding { builder = builder.header(header::CONTENT_ENCODING, content_encoding); } - if asset.vary_accept_encoding { + if candidate.vary_accept_encoding { builder = builder.header(header::VARY, "Accept-Encoding"); } if let Some(range) = &range { builder = builder.header( header::CONTENT_RANGE, - format!("bytes {}-{}/{}", range.start, range.end, asset.len), + format!("bytes {}-{}/{}", range.start, range.end, len), ); } @@ -815,11 +837,6 @@ async fn serve_fs_asset( return builder.empty(); } - let file = match tokio::fs::File::open(&asset.path).await { - Ok(file) => file, - Err(_) => return not_found(), - }; - let body: Body = match range { Some(range) => ranged_file_body(file, range, chunk_size).await, None => file_body(file, chunk_size), @@ -828,6 +845,33 @@ async fn serve_fs_asset( builder.body(body) } +/// Evaluate `If-Range`: `true` means a partial response may be served. +/// +/// ETag comparators use strong comparison; HTTP-date comparators must match +/// `Last-Modified` exactly. A missing header always allows ranges. +fn if_range_allows_partial( + req: &Request, + etag: &HeaderValue, + modified: Option, +) -> bool { + let Some(value) = req.headers().get(header::IF_RANGE) else { + return true; + }; + let Ok(value) = value.to_str() else { + return false; + }; + + if value.starts_with('"') || value.starts_with("W/") { + // Weak validators never satisfy If-Range's strong comparison. + return !value.starts_with("W/") && value.as_bytes() == etag.as_bytes(); + } + + match modified { + Some(m) => httpdate::parse_http_date(value).is_ok_and(|date| date == m), + None => false, + } +} + #[cfg(feature = "fs")] fn file_body(file: tokio::fs::File, chunk_size: usize) -> Body { let stream = tokio_util::io::ReaderStream::with_capacity(file, chunk_size); @@ -910,10 +954,14 @@ fn serve_embedded_asset( ); } - let range = match parse_range(req.headers().get(header::RANGE), asset.len) { - RangeDecision::Full => None, - RangeDecision::Partial(range) => Some(range), - RangeDecision::Unsatisfiable => return range_not_satisfiable(asset.len), + let range = if if_range_allows_partial(&req, &asset.etag, asset.modified) { + match parse_range(req.headers().get(header::RANGE), asset.len) { + RangeDecision::Full => None, + RangeDecision::Partial(range) => Some(range), + RangeDecision::Unsatisfiable => return range_not_satisfiable(asset.len), + } + } else { + None }; let body_len = range @@ -965,6 +1013,22 @@ fn serve_embedded_asset( builder.body(body) } +/// Verify that `path`, after resolving symlinks, still lives inside `root`. +/// +/// Both sides are canonicalized so escapes through intermediate directory +/// links are caught as well as direct file links. Unresolvable paths are +/// treated as not contained. +#[cfg(feature = "fs")] +async fn path_contained(root: &Path, path: &Path) -> bool { + let Ok(canonical_root) = tokio::fs::canonicalize(root).await else { + return false; + }; + match tokio::fs::canonicalize(path).await { + Ok(canonical_path) => canonical_path.starts_with(&canonical_root), + Err(_) => false, + } +} + #[derive(Debug)] struct RelativePath { fs_path: PathBuf, @@ -1057,24 +1121,18 @@ fn join_asset_path(base: &str, file: &str) -> String { } fn accepts_encoding(req: &Request, encoding: &str) -> bool { - req.headers() - .get(header::ACCEPT_ENCODING) - .and_then(|value| value.to_str().ok()) - .map(|header| { - header.split(',').any(|part| { - let mut pieces = part.trim().split(';'); - let token = pieces.next().unwrap_or("").trim(); - let mut q = 1.0_f32; - for param in pieces { - let param = param.trim(); - if let Some(value) = param.strip_prefix("q=") { - q = value.parse::().unwrap_or(0.0); - } - } - q > 0.0 && (token.eq_ignore_ascii_case(encoding) || token == "*") - }) - }) - .unwrap_or(false) + // Combine every header line (RFC 9110 §5.2 allows multi-line lists) and + // negotiate with q-value semantics — the old single-line `any()` check let + // `*;q=1` override an explicit `br;q=0` refusal. + let combined = req + .headers() + .get_all(header::ACCEPT_ENCODING) + .iter() + .filter_map(|value| value.to_str().ok()) + .collect::>() + .join(", "); + + arvik_core::accept::negotiate(&[encoding], &combined).is_some() } #[derive(Clone, Copy)] @@ -1094,11 +1152,16 @@ fn parse_range(value: Option<&HeaderValue>, len: u64) -> RangeDecision { return RangeDecision::Full; }; + // RFC 9110 §14.2: an unrecognized range unit MUST be ignored (serve 200 + // full), not answered with 416. let Some(spec) = value.strip_prefix("bytes=") else { - return RangeDecision::Unsatisfiable; + return RangeDecision::Full; }; + // Multi-range requests would need multipart/byteranges bodies; serving a + // single span would corrupt clients expecting all ranges. Full 200 is the + // safe, spec-permitted answer. if spec.contains(',') || spec.is_empty() { - return RangeDecision::Unsatisfiable; + return RangeDecision::Full; } let Some((start, end)) = spec.split_once('-') else { @@ -1181,8 +1244,11 @@ fn fs_etag(len: u64, modified: Option) -> HeaderValue { let modified = modified .and_then(|time| time.duration_since(UNIX_EPOCH).ok()) .unwrap_or_default(); + // Strong validator (nginx-style mtime+len). Strong comparison is required + // for `If-Range` to allow partial responses at all, and lets range-aware + // clients resume downloads against this file. HeaderValue::from_str(&format!( - "W/\"{:x}-{:x}-{:x}\"", + "\"{:x}-{:x}-{:x}\"", len, modified.as_secs(), modified.subsec_nanos() @@ -1578,4 +1644,91 @@ mod tests { let res = service.call(request("/dir/")).await.unwrap(); assert!(embedded_text(res).await.contains("file.txt")); } + // ── audit group 7: If-Range, symlink containment ───────────────────────── + + #[test] + fn if_range_strong_etag_comparison() { + let etag = HeaderValue::from_static("\"abc123\""); + let modified = None; + + // No header → ranges allowed. + let req = request("/f"); + assert!(if_range_allows_partial(&req, &etag, modified)); + + // Exact strong match → allowed. + let mut req = request("/f"); + req.headers_mut() + .insert(header::IF_RANGE, HeaderValue::from_static("\"abc123\"")); + assert!(if_range_allows_partial(&req, &etag, modified)); + + // Different validator → not allowed. + let mut req = request("/f"); + req.headers_mut() + .insert(header::IF_RANGE, HeaderValue::from_static("\"other\"")); + assert!(!if_range_allows_partial(&req, &etag, modified)); + + // Weak validators never satisfy strong comparison. + let mut req = request("/f"); + req.headers_mut() + .insert(header::IF_RANGE, HeaderValue::from_static("W/\"abc123\"")); + assert!(!if_range_allows_partial(&req, &etag, modified)); + } + + #[tokio::test] + async fn stale_if_range_downgrades_range_to_full_response() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("data.bin"); + write(&path, b"0123456789"); + + let mut service = ServeFile::new(&path); + + // Discover the current ETag. + let res = service.call(request("/ignored")).await.unwrap(); + let etag = res.headers()[header::ETAG].clone(); + + // Matching If-Range → partial content. + let mut req = request("/ignored"); + req.headers_mut() + .insert(header::RANGE, HeaderValue::from_static("bytes=0-3")); + req.headers_mut().insert(header::IF_RANGE, etag.clone()); + let res = service.call(req).await.unwrap(); + assert_eq!(res.status(), StatusCode::PARTIAL_CONTENT); + assert_eq!(body_text(res).await, "0123"); + + // Stale If-Range → full 200 (client must discard its cached version). + let mut req = request("/ignored"); + req.headers_mut() + .insert(header::RANGE, HeaderValue::from_static("bytes=0-3")); + req.headers_mut() + .insert(header::IF_RANGE, HeaderValue::from_static("\"stale\"")); + let res = service.call(req).await.unwrap(); + assert_eq!(res.status(), StatusCode::OK); + assert_eq!(body_text(res).await, "0123456789"); + } + + #[cfg(unix)] + #[tokio::test] + async fn symlink_escape_is_blocked_by_default() { + let outside = tempfile::tempdir().unwrap(); + let secret = outside.path().join("secret.txt"); + write(&secret, b"top secret"); + + let root = tempfile::tempdir().unwrap(); + #[cfg(unix)] + std::os::unix::fs::symlink(&secret, root.path().join("leak.txt")).unwrap(); + + let mut blocked = ServeDir::new(root.path()); + let res = blocked.call(request("/leak.txt")).await.unwrap(); + assert_ne!( + body_text(res).await, + "top secret", + "symlink escape must not serve outside the root by default" + ); + + // Opt-in restores link traversal. + let mut allowed = ServeDir::new(root.path()).follow_symlinks(true); + let res = allowed.call(request("/leak.txt")).await.unwrap(); + assert_eq!(res.status(), StatusCode::OK); + assert_eq!(body_text(res).await, "top secret"); + } } diff --git a/arvik-test/Cargo.toml b/arvik-test/Cargo.toml index 67f01ac..c6454c9 100644 --- a/arvik-test/Cargo.toml +++ b/arvik-test/Cargo.toml @@ -10,6 +10,7 @@ readme.workspace = true description = "Testing utilities for the Arvik web framework" [dependencies] +httpdate = { workspace = true } arvik-core = { workspace = true } arvik-router = { workspace = true } arvik-hyper = { workspace = true, optional = true } diff --git a/arvik-test/src/lib.rs b/arvik-test/src/lib.rs index 8e1687c..3bd7ae8 100644 --- a/arvik-test/src/lib.rs +++ b/arvik-test/src/lib.rs @@ -97,15 +97,56 @@ impl TestClient { fn store_response_cookies(&self, headers: &HeaderMap) { let mut jar = self.cookies.lock().expect("test cookie jar poisoned"); for value in headers.get_all(SET_COOKIE) { - if let Ok(value) = value.to_str() { - if let Ok(cookie) = Cookie::parse(value.to_owned()) { - jar.add(cookie.into_owned()); - } + let Some(value) = value.to_str().ok() else { + continue; + }; + let Ok(cookie) = Cookie::parse(value.to_owned()) else { + continue; + }; + + // Honor deletion so logout/expiry assertions behave like real + // browsers: `Max-Age<=0` or a past `Expires` removes the cookie. + if set_cookie_is_expired(value) { + jar.remove(cookie.into_owned()); + } else { + jar.add(cookie.into_owned()); } } } } +/// RFC 6265 §5.3: `Max-Age` takes precedence over `Expires`; both zero/past +/// values delete the cookie. +fn set_cookie_is_expired(raw: &str) -> bool { + let attributes = raw.split(';').skip(1).map(str::trim); + + for attribute in attributes.clone() { + if let Some(v) = attribute + .strip_prefix("Max-Age=") + .or_else(|| attribute.strip_prefix("max-age=")) + { + return v + .trim() + .parse::() + .map(|secs| secs <= 0) + .unwrap_or(false); + } + } + + for attribute in attributes { + if let Some(v) = attribute + .strip_prefix("Expires=") + .or_else(|| attribute.strip_prefix("expires=")) + { + return httpdate::parse_http_date(v.trim()) + .map(|time| time <= std::time::SystemTime::now()) + .unwrap_or(false); + } + } + + false +} + /// Builder for test requests. #[derive(Debug)] pub struct TestRequestBuilder { @@ -186,8 +227,20 @@ impl TestRequestBuilder { /// Dispatch the request and return the response. pub async fn send(mut self) -> TestResponse { - if let Some(cookie_header) = self.client.cookie_header() { - self.headers.insert(COOKIE, cookie_header); + let jar_header = self + .client + .cookie_header() + .and_then(|value| value.to_str().ok().map(str::to_owned)); + if let Some(jar_header) = jar_header { + // Merge with any explicitly set Cookie header instead of + // clobbering it. + let merged = match self.headers.get(COOKIE).and_then(|v| v.to_str().ok()) { + Some(existing) => format!("{existing}; {jar_header}"), + None => jar_header, + }; + if let Ok(value) = HeaderValue::from_str(&merged) { + self.headers.insert(COOKIE, value); + } } if !self.extra_cookies.is_empty() { diff --git a/arvik-tls/src/lib.rs b/arvik-tls/src/lib.rs index 544659e..d0fb583 100644 --- a/arvik-tls/src/lib.rs +++ b/arvik-tls/src/lib.rs @@ -156,6 +156,19 @@ pub mod rustls { Arc::clone(&self.inner.read()) } + /// Restrict ALPN advertisement to HTTP/2 only. + /// + /// Use when the server is configured with `http2_only(true)`: otherwise + /// the default offer invites clients to negotiate `http/1.1` while the + /// server answers with HTTP/2 frames, which compliant clients reject. + pub fn with_alpn_http2_only(self) -> Self { + let current = self.inner.read().clone(); + let mut updated = (*current).clone(); + updated.alpn_protocols = vec![ALPN_HTTP2.to_vec()]; + *self.inner.write() = Arc::new(updated); + self + } + /// Return a TLS acceptor using the currently active config. pub fn acceptor(&self) -> TlsAcceptor { TlsAcceptor::from(self.current_config()) @@ -257,26 +270,69 @@ pub mod rustls { key_path: std::path::PathBuf, debounce: Duration, ) -> Result { - use notify::{RecursiveMode, Watcher as _}; + use notify::{EventKind, RecursiveMode, Watcher as _}; let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<()>(); + + // Watch the *parent directories*, not the file inodes: deploys + // replace certificates with temp-file+rename and certbot re-points + // `live/` symlinks — an inode watch goes silent after both. + let cert_dir = cert_path.parent().map(Path::to_path_buf); + let key_dir = key_path.parent().map(Path::to_path_buf); + let cert_name = cert_path.file_name().map(std::ffi::OsStr::to_owned); + let key_name = key_path.file_name().map(std::ffi::OsStr::to_owned); + + let notify_cert_path = cert_path.clone(); + let notify_key_path = key_path.clone(); + let notify_cert_dir = cert_dir.clone(); + let notify_key_dir = key_dir.clone(); let mut watcher = notify::recommended_watcher( move |event: notify::Result| match event { - Ok(_) => { - let _ = tx.send(()); + Ok(event) => { + // Only content-relevant changes should trigger a + // reload; access/metadata churn would otherwise cause + // constant re-parsing. + if !matches!( + event.kind, + EventKind::Create(_) | EventKind::Modify(_) | EventKind::Remove(_) + ) { + return; + } + // The watched directories see every sibling file too — + // react only to our own file names, or to a + // bare-directory event (how some backends report renames). + let relevant = event.paths.iter().any(|p| { + p == ¬ify_cert_path + || p == ¬ify_key_path + || p.file_name().is_some_and(|name| { + Some(name) == cert_name.as_deref() + || Some(name) == key_name.as_deref() + }) + || Some(p.as_path()) == notify_cert_dir.as_deref() + || Some(p.as_path()) == notify_key_dir.as_deref() + }); + if relevant { + let _ = tx.send(()); + } } Err(err) => tracing::warn!("TLS certificate watch error: {}", err), }, )?; - watcher.watch(&cert_path, RecursiveMode::NonRecursive)?; - watcher.watch(&key_path, RecursiveMode::NonRecursive)?; + let mut seen_dirs = std::collections::HashSet::new(); + for dir in [&cert_dir, &key_dir].into_iter().flatten() { + if seen_dirs.insert(dir.clone()) { + watcher.watch(dir, RecursiveMode::NonRecursive)?; + } + } let task = tokio::spawn(async move { while rx.recv().await.is_some() { tokio::time::sleep(debounce).await; while rx.try_recv().is_ok() {} + // Path-based reload resolves symlinks and replacements + // freshly on every attempt, so post-rename events land. if let Err(err) = config.reload_from_pem_file(&cert_path, &key_path).await { tracing::warn!( "TLS certificate reload failed; keeping previous config active: {}", @@ -355,6 +411,8 @@ pub mod native { pub fn from_pkcs12(data: impl AsRef<[u8]>, password: &str) -> Result { let identity = native_tls::Identity::from_pkcs12(data.as_ref(), password)?; let mut builder = native_tls::TlsAcceptor::builder(identity); + // native-tls 0.2.18: accept_alpn is infallible (returns the + // builder), so ALPN setup cannot silently fail on this version. builder.accept_alpn(&["h2", "http/1.1"]); let acceptor = builder.build()?; Ok(Self { diff --git a/arvik-ws/src/lib.rs b/arvik-ws/src/lib.rs index 44287af..f798ade 100644 --- a/arvik-ws/src/lib.rs +++ b/arvik-ws/src/lib.rs @@ -52,6 +52,12 @@ //! .on_upgrade(handler) //! ``` +// `tungstenite::Error` is this crate's public error contract, returned by +// value from every socket operation. It exceeds clippy's result_large_err +// threshold on newer toolchains; boxing it would allocate on each call and +// break the API. The large-variant cases are rare runtime failures. +#![allow(clippy::result_large_err)] + pub mod message; pub mod rejection; pub mod socket; diff --git a/arvik-ws/src/message.rs b/arvik-ws/src/message.rs index 463b99c..dc38c90 100644 --- a/arvik-ws/src/message.rs +++ b/arvik-ws/src/message.rs @@ -181,7 +181,7 @@ pub enum Message { Text(String), /// A binary message. Binary(Vec), - /// A ping frame (server should auto-pong — see [`WebSocket::recv`]). + /// A ping frame (server should auto-pong — see `WebSocket::recv`). Ping(Vec), /// A pong frame (response to ping). Pong(Vec), diff --git a/arvik-ws/src/rejection.rs b/arvik-ws/src/rejection.rs index 28f4c22..832c8d2 100644 --- a/arvik-ws/src/rejection.rs +++ b/arvik-ws/src/rejection.rs @@ -54,7 +54,10 @@ impl IntoResponse for WebSocketUpgradeRejection { fn into_response(self) -> Response { let status = match &self { Self::MethodNotGet => StatusCode::METHOD_NOT_ALLOWED, - Self::InvalidWebSocketVersionHeader => StatusCode::BAD_REQUEST, + // RFC 6455 §4.4/§4.2.2: an unsupported version gets 426 Upgrade + // Required plus the version(s) this server speaks, so capable + // clients can retry instead of failing unrecoverably. + Self::InvalidWebSocketVersionHeader => StatusCode::UPGRADE_REQUIRED, Self::ConnectionNotUpgradable => StatusCode::INTERNAL_SERVER_ERROR, _ => StatusCode::BAD_REQUEST, }; @@ -65,9 +68,14 @@ impl IntoResponse for WebSocketUpgradeRejection { status.as_u16() ); - ResponseBuilder::new() + let mut builder = ResponseBuilder::new() .status(status) - .header(http::header::CONTENT_TYPE, "application/json") - .body(arvik_core::Body::from(body)) + .header(http::header::CONTENT_TYPE, "application/json"); + + if matches!(self, Self::InvalidWebSocketVersionHeader) { + builder = builder.header("sec-websocket-version", "13"); + } + + builder.body(arvik_core::Body::from(body)) } } diff --git a/arvik-ws/src/split.rs b/arvik-ws/src/split.rs index 5375fa8..ebd176b 100644 --- a/arvik-ws/src/split.rs +++ b/arvik-ws/src/split.rs @@ -102,7 +102,7 @@ impl std::fmt::Debug for Sender { /// /// `Ping` frames received here are surfaced as [`Message::Ping`] so that /// your receive task can forward them to the `Sender` for an explicit pong. -/// Unlike [`WebSocket::recv`], the split `Receiver` does **not** auto-pong +/// Unlike `WebSocket::recv`, the split `Receiver` does **not** auto-pong /// (there is no access to the sink half to do so). Handle them manually: /// /// ```rust,ignore diff --git a/arvik/src/lib.rs b/arvik/src/lib.rs index 186c156..6ee498a 100644 --- a/arvik/src/lib.rs +++ b/arvik/src/lib.rs @@ -87,6 +87,10 @@ pub use arvik_core::ResponseParts; // 0.3.2 pub use arvik_core::FromRequest; pub use arvik_core::FromRequestParts; +// Body size limiting (0.9.6) +pub use arvik_core::body_limit::DEFAULT_BODY_LIMIT; +pub use arvik_core::body_limit::DefaultBodyLimit; + // --------------------------------------------------------------------------- // Router types // --------------------------------------------------------------------------- @@ -303,14 +307,14 @@ pub use arvik_static::{ServeDir, ServeFile}; /// /// | Function | Use case | /// |---|---| -/// | [`from_fn`] | Middleware with extractors, no shared state | -/// | [`from_fn_with_state`] | Middleware with extractors + shared app state | -/// | [`map_request`] | Transform the incoming request only | -/// | [`map_request_with_state`] | Transform request with state access | -/// | [`map_response`] | Transform the outgoing response only | -/// | [`map_response_with_state`] | Transform response with state access | +/// | `from_fn` | Middleware with extractors, no shared state | +/// | `from_fn_with_state` | Middleware with extractors + shared app state | +/// | `map_request` | Transform the incoming request only | +/// | `map_request_with_state` | Transform request with state access | +/// | `map_response` | Transform the outgoing response only | +/// | `map_response_with_state` | Transform response with state access | /// -/// See [`arvik_middleware::from_fn`] for detailed documentation and examples. +/// See [`arvik_middleware::from_fn()`] for detailed documentation and examples. pub mod middleware { pub use arvik_middleware::{ auth::RequireAuthorizationLayer, @@ -345,7 +349,7 @@ pub mod middleware { // Add at the top-level of lib.rs (outside the middleware module): pub use arvik_middleware::CorsLayer; pub use arvik_middleware::auth::RequireAuthorizationLayer; -pub use arvik_middleware::body_limit::RequestBodyLimitLayer; +pub use arvik_middleware::body_limit::{DefaultBodyLimitLayer, RequestBodyLimitLayer}; pub use arvik_middleware::catch_panic::CatchPanicLayer; pub use arvik_middleware::compression::{CompressionLayer, CompressionLevel, DecompressionLayer}; pub use arvik_middleware::csrf::{CsrfLayer, CsrfToken};