diff --git a/Cargo.lock b/Cargo.lock index 8d0b790..847d676 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2780,7 +2780,7 @@ dependencies = [ [[package]] name = "telemt" -version = "3.5.7" +version = "3.5.8" dependencies = [ "aes", "anyhow", diff --git a/Cargo.toml b/Cargo.toml index 6c17a13..d30b391 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "telemt" -version = "3.5.7" +version = "3.5.8" edition = "2024" rust-version = "1.88" @@ -36,6 +36,7 @@ nix = { version = "0.31.3", default-features = false, features = [ "net", "user", "process", + "dir", "fs", "signal", ] } diff --git a/src/api/config_edit.rs b/src/api/config_edit.rs index 7b087c6..b1e8719 100644 --- a/src/api/config_edit.rs +++ b/src/api/config_edit.rs @@ -6,10 +6,13 @@ use serde_json::Value as Json; use toml::Value as Toml; use super::ApiShared; +#[cfg(test)] +use super::config_store::write_atomic; use super::config_store::{ EDITABLE_SECTIONS, EDITABLE_SERVER_FIELDS, compute_snapshot_revision, is_editable_section, load_candidate_snapshot, load_config_snapshot, render_server_listeners, - render_top_level_section, resolve_single_source_owner, upsert_toml_table, write_atomic, + render_top_level_section, resolve_single_source_owner, upsert_toml_table, + write_atomic_if_unchanged, }; use super::model::ApiFailure; use crate::config::ProxyConfig; @@ -43,7 +46,10 @@ pub(super) struct PatchConfigResponse { } struct PreparedConfigPatch { + config_path: PathBuf, + expected_revision: String, owner_path: PathBuf, + expected_owner_contents: String, owner_contents: String, desired_config: Arc, response: PatchConfigResponse, @@ -56,6 +62,21 @@ pub(super) async fn patch_config( expected_revision: Option, reload_request: Option, shared: &ApiShared, +) -> Result { + let shared = shared.clone(); + shared + .clone() + .run_mutation_completion(async move { + patch_config_to_completion(patch_json, expected_revision, reload_request, &shared).await + }) + .await +} + +async fn patch_config_to_completion( + patch_json: Json, + expected_revision: Option, + reload_request: Option, + shared: &ApiShared, ) -> Result { let _guard = shared.mutation_lock.lock().await; let active_config = shared.active_runtime.load_full().config(); @@ -77,7 +98,14 @@ pub(super) async fn patch_config( } else { None }; - write_atomic(prepared.owner_path, prepared.owner_contents).await?; + prepared.response.revision = write_atomic_if_unchanged( + prepared.config_path, + prepared.expected_revision, + prepared.owner_path, + prepared.expected_owner_contents, + prepared.owner_contents, + ) + .await?; if let Some(reservation) = reservation { prepared.response.reload = Some(reservation.enqueue(prepared.desired_config)); } @@ -110,8 +138,16 @@ pub(super) async fn apply_patch_to_path( patch_json: &Json, expected_revision: Option, ) -> Result { - let prepared = prepare_patch_to_path(config_path, patch_json, expected_revision).await?; - write_atomic(prepared.owner_path, prepared.owner_contents).await?; + let mut prepared = prepare_patch_to_path(config_path, patch_json, expected_revision).await?; + let revision = write_atomic_if_unchanged( + prepared.config_path, + prepared.expected_revision, + prepared.owner_path, + prepared.expected_owner_contents, + prepared.owner_contents, + ) + .await?; + prepared.response.revision = revision; Ok(prepared.response) } @@ -197,6 +233,7 @@ async fn prepare_patch_to_path( .get(&owner_path) .cloned() .ok_or_else(|| ApiFailure::internal("config source owner is missing from snapshot"))?; + let expected_owner_contents = owner_contents.clone(); for section in &touched { if *section == "server" { let rendered = render_server_listeners(&requested_cfg)?; @@ -233,7 +270,10 @@ async fn prepare_patch_to_path( deferred_process_fields(&old_cfg, &new_cfg).map_err(ApiFailure::bad_request)?; Ok(PreparedConfigPatch { + config_path: config_path.to_path_buf(), + expected_revision: current, owner_path, + expected_owner_contents, owner_contents, desired_config: Arc::new(new_cfg), response: PatchConfigResponse { @@ -407,6 +447,9 @@ fn deep_merge(base: &mut Toml, patch: &Toml) { } } +#[cfg(test)] +#[path = "config_edit/base_path_tests.rs"] +mod base_path_tests; #[cfg(test)] #[path = "config_edit/tests.rs"] mod tests; diff --git a/src/api/config_edit/base_path_tests.rs b/src/api/config_edit/base_path_tests.rs new file mode 100644 index 0000000..5a2bbfb --- /dev/null +++ b/src/api/config_edit/base_path_tests.rs @@ -0,0 +1,88 @@ +use super::*; + +fn web_config() -> &'static str { + r#" +[access.users] +alice = "000102030405060708090a0b0c0d0e0f" + +[[server.listeners]] +ip = "127.0.0.1" +port = 18080 +transport = "web" +proxy_protocol = false +web_client_ip_source = "x_forwarded_for" +web_trusted_proxy_cidrs = ["127.0.0.1/32"] + +[web] +enabled = true + +[[web.vhosts]] +host = "proxy.example.com" +public_addr = "203.0.113.10:443" + +[web.vhosts.decoy] +mode = "http_upstream" +upstream = "http://127.0.0.1:18081" + +[[web.vhosts.profiles]] +user = "alice" +secret_mode = "plain" +"# +} + +fn vhosts_patch(base_path: &str) -> Json { + serde_json::json!({ + "web": { + "vhosts": [{ + "host": "proxy.example.com", + "base_path": base_path, + "public_addr": "203.0.113.10:443", + "decoy": { + "mode": "http_upstream", + "upstream": "http://127.0.0.1:18081" + }, + "profiles": [{ + "user": "alice", + "secret_mode": "plain" + }] + }] + } + }) +} + +#[tokio::test] +async fn config_api_applies_valid_base_path_and_preserves_source_on_invalid_patch() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("config.toml"); + std::fs::write(&path, web_config()).unwrap(); + let active = ProxyConfig::load(&path).unwrap(); + + let mut response = apply_patch_to_path(&path, &vhosts_patch("MixedCase/path"), None) + .await + .unwrap(); + let desired = ProxyConfig::load(&path).unwrap(); + let resolved = reconcile_runtime_effect(&mut response, &active, &desired).unwrap(); + assert!(!response.restart_required); + assert!(response.runtime_reload_required); + assert!(!response.process_restart_required); + assert!(response.deferred_process_fields.is_empty()); + assert!(resolved.runtime_changed); + assert_eq!(desired.web.vhosts[0].base_path, "MixedCase/path"); + assert_eq!( + resolved.effective.web.runtime.as_ref().unwrap().vhosts["proxy.example.com"].base, + "/MixedCase/path/" + ); + + let (managed, _revision) = read_managed_config(&path).await.unwrap(); + let vhosts = managed["web"]["vhosts"].as_array().unwrap(); + assert_eq!(vhosts[0]["base_path"].as_str(), Some("MixedCase/path")); + assert!(managed["web"].get("runtime").is_none()); + assert!(!managed.as_table().unwrap().contains_key("access")); + + let before_invalid = std::fs::read(&path).unwrap(); + let error = apply_patch_to_path(&path, &vhosts_patch("/invalid"), None) + .await + .unwrap_err(); + assert_eq!(error.status, hyper::StatusCode::BAD_REQUEST); + assert_eq!(std::fs::read(&path).unwrap(), before_invalid); +} diff --git a/src/api/config_edit/tests.rs b/src/api/config_edit/tests.rs index 1143a31..fb2f9a4 100644 --- a/src/api/config_edit/tests.rs +++ b/src/api/config_edit/tests.rs @@ -107,14 +107,24 @@ async fn read_managed_config_exposes_web_without_runtime_or_access_secrets() { #[tokio::test] async fn patch_web_debug_is_hot_and_limits_are_process_deferred() { let (path, _directory) = temp_config("[web]\nenabled = false\n"); + let active = ProxyConfig::load(&path).unwrap(); let debug_patch: Json = serde_json::json!({ - "web": {"debug": {"enabled": true, "capture_headers": false}} + "web": {"debug": { + "enabled": true, + "sideband": true, + "capture_headers": false + }} }); - let debug = apply_patch_to_path(&path, &debug_patch, None) + let mut debug = apply_patch_to_path(&path, &debug_patch, None) .await .unwrap(); + let desired = ProxyConfig::load(&path).unwrap(); + reconcile_runtime_effect(&mut debug, &active, &desired).unwrap(); + assert!(!debug.restart_required); + assert!(debug.runtime_reload_required); assert!(!debug.process_restart_required); assert!(debug.changed.iter().any(|section| section == "web")); + assert!(desired.web.debug.sideband); let limits_patch: Json = serde_json::json!({ "web": {"limits": {"max_http_connections": 2049}} @@ -313,6 +323,30 @@ async fn patch_writes_the_included_section_owner_only() { ); } +#[tokio::test] +async fn prepared_patch_rejects_external_edit_before_commit() { + let (path, _directory) = temp_config("[censorship]\ntls_domain = \"old.example\"\n"); + let patch: Json = serde_json::json!({ + "censorship": {"tls_domain": "api.example"} + }); + let prepared = prepare_patch_to_path(&path, &patch, None).await.unwrap(); + let external = "[censorship]\ntls_domain = \"external.example\"\n"; + tokio::fs::write(&path, external).await.unwrap(); + + let error = write_atomic_if_unchanged( + prepared.config_path, + prepared.expected_revision, + prepared.owner_path, + prepared.expected_owner_contents, + prepared.owner_contents, + ) + .await + .unwrap_err(); + + assert_eq!(error.code, "revision_conflict"); + assert_eq!(tokio::fs::read_to_string(&path).await.unwrap(), external); +} + #[tokio::test] async fn patch_rejects_multiple_source_owners_without_writing() { let dir = tempfile::tempdir().unwrap(); diff --git a/src/api/config_store.rs b/src/api/config_store.rs index 7e97ec7..5bb7869 100644 --- a/src/api/config_store.rs +++ b/src/api/config_store.rs @@ -10,12 +10,15 @@ use super::model::ApiFailure; // Source-preserving TOML rendering and atomic persistence helpers. mod persistence; +// Compare-and-replace file persistence and metadata preservation. +mod atomic; +pub(in crate::api) use atomic::{write_atomic, write_atomic_if_unchanged}; #[cfg(test)] use persistence::{find_toml_table_bounds, render_access_section, save_sections_to_disk}; pub(in crate::api) use persistence::{ render_server_listeners, render_top_level_section, save_access_sections_to_disk, - upsert_toml_table, write_atomic, + save_access_sections_to_disk_if_revision, upsert_toml_table, }; #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -54,22 +57,21 @@ pub(super) fn parse_if_match(headers: &hyper::HeaderMap) -> Option { .map(|value| value.trim_matches('"').to_string()) } -pub(super) async fn ensure_expected_revision( +/// Loads one mutation base and validates its revision from the same source snapshot. +pub(super) async fn load_config_for_mutation( config_path: &Path, expected_revision: Option<&str>, -) -> Result<(), ApiFailure> { - let Some(expected) = expected_revision else { - return Ok(()); - }; - let current = current_revision(config_path).await?; - if current != expected { +) -> Result<(ProxyConfig, String), ApiFailure> { + let loaded = load_config_snapshot(config_path, false).await?; + let revision = compute_snapshot_revision(&loaded); + if expected_revision.is_some_and(|expected| expected != revision) { return Err(ApiFailure::new( hyper::StatusCode::CONFLICT, "revision_conflict", "Config revision mismatch", )); } - Ok(()) + Ok((loaded.config, revision)) } pub(super) async fn current_revision(config_path: &Path) -> Result { @@ -243,15 +245,24 @@ pub(super) async fn load_candidate_snapshot( } fn normalize_source_path(path: &Path) -> PathBuf { - path.canonicalize().unwrap_or_else(|_| { - if path.is_absolute() { - path.to_path_buf() - } else { - std::env::current_dir() - .map(|cwd| cwd.join(path)) - .unwrap_or_else(|_| path.to_path_buf()) + let absolute = if path.is_absolute() { + path.to_path_buf() + } else { + std::env::current_dir() + .map(|cwd| cwd.join(path)) + .unwrap_or_else(|_| path.to_path_buf()) + }; + let mut normalized = PathBuf::new(); + for component in absolute.components() { + match component { + std::path::Component::CurDir => {} + std::path::Component::ParentDir => { + normalized.pop(); + } + component => normalized.push(component.as_os_str()), } - }) + } + normalized } pub(super) async fn load_config_from_disk(config_path: &Path) -> Result { diff --git a/src/api/config_store/atomic.rs b/src/api/config_store/atomic.rs new file mode 100644 index 0000000..221abcc --- /dev/null +++ b/src/api/config_store/atomic.rs @@ -0,0 +1,432 @@ +use std::fs::File; +use std::io::{Read, Write}; +use std::path::{Path, PathBuf}; + +#[cfg(unix)] +use std::os::unix::fs::{MetadataExt, PermissionsExt}; + +#[cfg(unix)] +use nix::fcntl::{Flock, FlockArg, OFlag, openat, renameat}; +#[cfg(unix)] +use nix::sys::stat::Mode; +#[cfg(unix)] +use nix::unistd::{UnlinkatFlags, fsync, unlinkat}; +#[cfg(unix)] +use tracing::warn; + +use super::compute_source_revision; +use crate::api::model::ApiFailure; +use crate::config::ProxyConfig; +#[cfg(unix)] +use crate::util::secure_fs::AnchoredPath; + +const MAX_CONFIG_SOURCE_BYTES: u64 = 8 * 1024 * 1024; + +enum AtomicWriteError { + Conflict, + ReadGraph(String), + Io(std::io::Error), +} + +struct ExistingTarget { + contents: String, + metadata: std::fs::Metadata, +} + +struct GraphFence<'a> { + config_path: &'a Path, + expected_revision: &'a str, +} + +struct ConfigWriteLock { + #[cfg(unix)] + _file: Flock, +} + +impl ConfigWriteLock { + fn acquire(path: &Path) -> std::io::Result { + let path = normalize_path(path); + #[cfg(unix)] + { + let lock_path = sibling_lock_path(&path); + let anchored = AnchoredPath::open_creating_parents(&lock_path, 0o750)?; + let descriptor = openat( + anchored.parent(), + anchored.name(), + OFlag::O_RDWR | OFlag::O_CREAT | OFlag::O_NOFOLLOW | OFlag::O_CLOEXEC, + Mode::from_bits_truncate(0o600), + ) + .map_err(errno_to_io)?; + let file = File::from(descriptor); + let metadata = file.metadata()?; + if !metadata.is_file() || metadata.nlink() != 1 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "config lock must be a regular file with one directory entry", + )); + } + let file = Flock::lock(file, FlockArg::LockExclusive) + .map_err(|(_, error)| errno_to_io(error))?; + Ok(Self { _file: file }) + } + #[cfg(not(unix))] + { + let _ = path; + Ok(Self {}) + } + } +} + +/// Replaces one config source through a same-directory rename after syncing file data. +pub(in crate::api) async fn write_atomic( + path: PathBuf, + contents: String, +) -> Result<(), ApiFailure> { + tokio::task::spawn_blocking(move || { + let _lock = ConfigWriteLock::acquire(&path)?; + write_atomic_sync(&path, None, &contents, None).map(|_| ()) + }) + .await + .map_err(|error| ApiFailure::internal(format!("failed to join writer: {error}")))? + .map_err(|error| ApiFailure::internal(format!("failed to write config: {error}"))) +} + +/// Replaces one source only if both its graph revision and owner contents are unchanged. +pub(in crate::api) async fn write_atomic_if_unchanged( + config_path: PathBuf, + expected_revision: String, + path: PathBuf, + expected_contents: String, + contents: String, +) -> Result { + tokio::task::spawn_blocking(move || { + let config_path = normalize_path(&config_path); + let path = normalize_path(&path); + // Every API mutation locks the root source so writes to different includes serialize. + let _lock = ConfigWriteLock::acquire(&config_path).map_err(AtomicWriteError::Io)?; + let graph = ProxyConfig::read_source_graph(&config_path) + .map_err(|error| AtomicWriteError::ReadGraph(error.to_string()))?; + if compute_source_revision(&graph) != expected_revision { + return Err(AtomicWriteError::Conflict); + } + write_atomic_sync( + &path, + Some(&expected_contents), + &contents, + Some(GraphFence { + config_path: &config_path, + expected_revision: &expected_revision, + }), + ) + .map_err(|error| { + if error.kind() == std::io::ErrorKind::AlreadyExists { + AtomicWriteError::Conflict + } else { + AtomicWriteError::Io(error) + } + })? + .ok_or_else(|| { + AtomicWriteError::Io(std::io::Error::other( + "config graph fence did not produce a committed revision", + )) + }) + }) + .await + .map_err(|error| ApiFailure::internal(format!("failed to join writer: {error}")))? + .map_err(|error| match error { + AtomicWriteError::Conflict => revision_conflict(), + AtomicWriteError::ReadGraph(error) => { + ApiFailure::internal(format!("failed to verify config graph: {error}")) + } + AtomicWriteError::Io(error) => { + ApiFailure::internal(format!("failed to write config: {error}")) + } + }) +} + +fn revision_conflict() -> ApiFailure { + ApiFailure::new( + hyper::StatusCode::CONFLICT, + "revision_conflict", + "Config revision changed before persistence", + ) +} + +fn sibling_lock_path(path: &Path) -> PathBuf { + let mut name = path + .file_name() + .unwrap_or_else(|| std::ffi::OsStr::new("config.toml")) + .to_os_string(); + name.push(".lock"); + path.parent().unwrap_or_else(|| Path::new(".")).join(name) +} + +fn normalize_path(path: &Path) -> PathBuf { + let absolute = if path.is_absolute() { + path.to_path_buf() + } else { + std::env::current_dir() + .map(|current| current.join(path)) + .unwrap_or_else(|_| path.to_path_buf()) + }; + let mut normalized = PathBuf::new(); + for component in absolute.components() { + match component { + std::path::Component::CurDir => {} + std::path::Component::ParentDir => { + normalized.pop(); + } + component => normalized.push(component.as_os_str()), + } + } + normalized +} + +fn fenced_post_commit_revision( + fence: GraphFence<'_>, + path: &Path, + contents: &str, +) -> std::io::Result { + let mut graph = ProxyConfig::read_source_graph(fence.config_path) + .map_err(|error| std::io::Error::other(error.to_string()))?; + if compute_source_revision(&graph) != fence.expected_revision { + return Err(std::io::Error::new( + std::io::ErrorKind::AlreadyExists, + "config graph changed during persistence", + )); + } + let path = normalize_path(path); + let Some(owner) = graph.source_contents.get_mut(&path) else { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "config source owner left the source graph during persistence", + )); + }; + *owner = contents.to_string(); + Ok(compute_source_revision(&graph)) +} + +#[cfg(unix)] +fn open_existing_target(anchored: &AnchoredPath) -> std::io::Result> { + let descriptor = match openat( + anchored.parent(), + anchored.name(), + OFlag::O_RDONLY | OFlag::O_NONBLOCK | OFlag::O_NOFOLLOW | OFlag::O_CLOEXEC, + Mode::empty(), + ) { + Ok(descriptor) => descriptor, + Err(nix::errno::Errno::ENOENT) => return Ok(None), + Err(error) => return Err(errno_to_io(error)), + }; + let mut file = File::from(descriptor); + let metadata = file.metadata()?; + if !metadata.is_file() || metadata.nlink() != 1 || metadata.len() > MAX_CONFIG_SOURCE_BYTES { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "config target must be a bounded regular file with one directory entry", + )); + } + let mut contents = String::with_capacity(metadata.len() as usize); + Read::take(&mut file, MAX_CONFIG_SOURCE_BYTES + 1).read_to_string(&mut contents)?; + if contents.len() as u64 > MAX_CONFIG_SOURCE_BYTES { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "config target exceeds the source size limit", + )); + } + let completed = file.metadata()?; + if !same_target(&metadata, &completed) || metadata.len() != completed.len() { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "config target changed while it was read", + )); + } + Ok(Some(ExistingTarget { contents, metadata })) +} + +#[cfg(not(unix))] +fn open_existing_target(path: &Path) -> std::io::Result> { + let mut file = match File::open(path) { + Ok(file) => file, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None), + Err(error) => return Err(error), + }; + let metadata = file.metadata()?; + if !metadata.is_file() || metadata.len() > MAX_CONFIG_SOURCE_BYTES { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "config target must be a bounded regular file", + )); + } + let mut contents = String::new(); + file.read_to_string(&mut contents)?; + Ok(Some(ExistingTarget { contents, metadata })) +} + +fn same_target(left: &std::fs::Metadata, right: &std::fs::Metadata) -> bool { + #[cfg(unix)] + { + left.dev() == right.dev() && left.ino() == right.ino() + } + #[cfg(not(unix))] + { + left.len() == right.len() && left.modified().ok() == right.modified().ok() + } +} + +#[cfg(unix)] +fn write_atomic_sync( + path: &Path, + expected_contents: Option<&str>, + contents: &str, + graph_fence: Option>, +) -> std::io::Result> { + let anchored = AnchoredPath::open_creating_parents(path, 0o750)?; + let existing = open_existing_target(&anchored)?; + validate_expected_contents(existing.as_ref(), expected_contents)?; + let temp_name = format!( + ".{}.tmp-{}", + path.file_name() + .and_then(|name| name.to_str()) + .unwrap_or("config.toml"), + rand::random::() + ); + let descriptor = openat( + anchored.parent(), + temp_name.as_str(), + OFlag::O_WRONLY | OFlag::O_CREAT | OFlag::O_EXCL | OFlag::O_NOFOLLOW | OFlag::O_CLOEXEC, + Mode::from_bits_truncate(0o600), + ) + .map_err(errno_to_io)?; + let write_result = write_and_publish( + descriptor, + path, + &anchored, + &temp_name, + existing.as_ref(), + contents, + graph_fence, + ); + if write_result.is_err() { + let _ = unlinkat( + anchored.parent(), + temp_name.as_str(), + UnlinkatFlags::NoRemoveDir, + ); + } + write_result +} + +#[cfg(unix)] +fn write_and_publish( + descriptor: std::os::fd::OwnedFd, + path: &Path, + anchored: &AnchoredPath, + temp_name: &str, + existing: Option<&ExistingTarget>, + contents: &str, + graph_fence: Option>, +) -> std::io::Result> { + let mut file = File::from(descriptor); + if let Some(existing) = existing { + use nix::unistd::{Gid, Uid, fchown}; + + fchown( + &file, + Some(Uid::from_raw(existing.metadata.uid())), + Some(Gid::from_raw(existing.metadata.gid())), + ) + .map_err(errno_to_io)?; + file.set_permissions(std::fs::Permissions::from_mode( + existing.metadata.mode() & 0o7777, + ))?; + } + file.write_all(contents.as_bytes())?; + file.sync_all()?; + let current = open_existing_target(anchored)?; + if !target_unchanged(existing, current.as_ref()) { + return Err(std::io::Error::new( + std::io::ErrorKind::AlreadyExists, + "config target changed during persistence", + )); + } + let committed_revision = graph_fence + .map(|fence| fenced_post_commit_revision(fence, path, contents)) + .transpose()?; + renameat( + anchored.parent(), + temp_name, + anchored.parent(), + anchored.name(), + ) + .map_err(errno_to_io)?; + // Rename is the commit boundary. A later directory-sync error cannot be reported as an + // uncommitted mutation because mandatory in-process publication must still run. + if let Err(error) = fsync(anchored.parent()).map_err(errno_to_io) { + warn!( + path = %path.display(), + %error, + "Config rename committed but directory durability sync failed" + ); + } + Ok(committed_revision) +} + +#[cfg(not(unix))] +fn write_atomic_sync( + path: &Path, + expected_contents: Option<&str>, + contents: &str, + graph_fence: Option>, +) -> std::io::Result> { + let parent = path.parent().unwrap_or_else(|| Path::new(".")); + std::fs::create_dir_all(parent)?; + let existing = open_existing_target(path)?; + validate_expected_contents(existing.as_ref(), expected_contents)?; + let temp = parent.join(format!(".telemt.tmp-{}", rand::random::())); + std::fs::write(&temp, contents)?; + let current = open_existing_target(path)?; + if !target_unchanged(existing.as_ref(), current.as_ref()) { + let _ = std::fs::remove_file(&temp); + return Err(std::io::Error::new( + std::io::ErrorKind::AlreadyExists, + "config target changed during persistence", + )); + } + let committed_revision = graph_fence + .map(|fence| fenced_post_commit_revision(fence, path, contents)) + .transpose()?; + std::fs::rename(temp, path)?; + Ok(committed_revision) +} + +fn validate_expected_contents( + existing: Option<&ExistingTarget>, + expected_contents: Option<&str>, +) -> std::io::Result<()> { + if expected_contents + .is_some_and(|expected| existing.is_none_or(|target| target.contents != expected)) + { + return Err(std::io::Error::new( + std::io::ErrorKind::AlreadyExists, + "config source changed before persistence", + )); + } + Ok(()) +} + +fn target_unchanged(existing: Option<&ExistingTarget>, current: Option<&ExistingTarget>) -> bool { + match (existing, current) { + (Some(expected), Some(current)) => { + same_target(&expected.metadata, ¤t.metadata) + && expected.contents == current.contents + } + (None, None) => true, + _ => false, + } +} + +#[cfg(unix)] +fn errno_to_io(error: nix::errno::Errno) -> std::io::Error { + std::io::Error::from_raw_os_error(error as i32) +} diff --git a/src/api/config_store/persistence.rs b/src/api/config_store/persistence.rs index 8893e9e..83c392a 100644 --- a/src/api/config_store/persistence.rs +++ b/src/api/config_store/persistence.rs @@ -1,12 +1,14 @@ use std::collections::BTreeMap; -use std::io::Write; -use std::path::{Path, PathBuf}; +use std::path::Path; use chrono::{DateTime, Utc}; use serde::Serialize; use crate::config::{ProxyConfig, RateLimitBps}; +#[cfg(test)] +use super::atomic::write_atomic; +use super::atomic::write_atomic_if_unchanged; #[cfg(test)] use super::compute_revision; use super::{ @@ -99,8 +101,22 @@ pub(in crate::api) async fn save_access_sections_to_disk( config_path: &Path, cfg: &ProxyConfig, sections: &[AccessSection], +) -> Result { + save_access_sections_to_disk_if_revision(config_path, cfg, sections, None).await +} + +/// Persists access tables only while the complete source graph remains unchanged. +pub(in crate::api) async fn save_access_sections_to_disk_if_revision( + config_path: &Path, + cfg: &ProxyConfig, + sections: &[AccessSection], + expected_revision: Option<&str>, ) -> Result { let loaded = load_config_snapshot(config_path, false).await?; + let loaded_revision = compute_snapshot_revision(&loaded); + if expected_revision.is_some_and(|expected| expected != loaded_revision) { + return Err(revision_conflict()); + } let mut applied = Vec::new(); for section in sections { if applied.contains(section) { @@ -117,7 +133,7 @@ pub(in crate::api) async fn save_access_sections_to_disk( }) }); if applied.is_empty() { - return Ok(compute_snapshot_revision(&loaded)); + return Ok(loaded_revision); } let targets = applied @@ -130,6 +146,7 @@ pub(in crate::api) async fn save_access_sections_to_disk( .get(&owner_path) .cloned() .ok_or_else(|| ApiFailure::internal("config source owner is missing from snapshot"))?; + let expected_owner_contents = owner_contents.clone(); for section in applied { let rendered = render_access_section(cfg, section)?; owner_contents = upsert_toml_table(&owner_contents, section.table_name(), &rendered); @@ -142,8 +159,15 @@ pub(in crate::api) async fn save_access_sections_to_disk( owner_contents.clone(), ) .await?; - let revision = compute_snapshot_revision(&candidate); - write_atomic(owner_path, owner_contents).await?; + let _candidate_revision = compute_snapshot_revision(&candidate); + let revision = write_atomic_if_unchanged( + config_path.to_path_buf(), + loaded_revision, + owner_path, + expected_owner_contents, + owner_contents, + ) + .await?; Ok(revision) } @@ -373,46 +397,10 @@ fn find_all_table_blocks(source: &str, table_name: &str) -> Vec<(usize, usize)> blocks } -/// Replaces one config source through a durable same-directory rename. -pub(in crate::api) async fn write_atomic( - path: PathBuf, - contents: String, -) -> Result<(), ApiFailure> { - tokio::task::spawn_blocking(move || write_atomic_sync(&path, &contents)) - .await - .map_err(|e| ApiFailure::internal(format!("failed to join writer: {}", e)))? - .map_err(|e| ApiFailure::internal(format!("failed to write config: {}", e))) -} - -fn write_atomic_sync(path: &Path, contents: &str) -> std::io::Result<()> { - let parent = path.parent().unwrap_or_else(|| Path::new(".")); - std::fs::create_dir_all(parent)?; - - let tmp_name = format!( - ".{}.tmp-{}", - path.file_name() - .and_then(|s| s.to_str()) - .unwrap_or("config.toml"), - rand::random::() - ); - let tmp_path = parent.join(tmp_name); - - let write_result = (|| { - let mut file = std::fs::OpenOptions::new() - .create_new(true) - .write(true) - .open(&tmp_path)?; - file.write_all(contents.as_bytes())?; - file.sync_all()?; - std::fs::rename(&tmp_path, path)?; - if let Ok(dir) = std::fs::File::open(parent) { - let _ = dir.sync_all(); - } - Ok(()) - })(); - - if write_result.is_err() { - let _ = std::fs::remove_file(&tmp_path); - } - write_result +fn revision_conflict() -> ApiFailure { + ApiFailure::new( + hyper::StatusCode::CONFLICT, + "revision_conflict", + "Config revision changed before persistence", + ) } diff --git a/src/api/config_store/tests.rs b/src/api/config_store/tests.rs index ef35909..6635201 100644 --- a/src/api/config_store/tests.rs +++ b/src/api/config_store/tests.rs @@ -260,6 +260,104 @@ async fn access_mutation_writes_only_the_single_included_owner() { assert_eq!(revision, current_revision(&root).await.unwrap()); } +#[tokio::test] +async fn access_mutation_rejects_source_graph_change_after_snapshot() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path().join("config.toml"); + let included = dir.path().join("users.toml"); + let root_body = "include = \"users.toml\"\n[censorship]\ntls_domain = \"one.example\"\n"; + let external_root = "include = \"users.toml\"\n[censorship]\ntls_domain = \"two.example\"\n"; + let included_body = "[access.users]\nalice = \"00000000000000000000000000000000\"\n"; + tokio::fs::write(&root, root_body).await.unwrap(); + tokio::fs::write(&included, included_body).await.unwrap(); + let (mut cfg, revision) = load_config_for_mutation(&root, None).await.unwrap(); + cfg.access.users.insert( + "bob".to_string(), + "11111111111111111111111111111111".to_string(), + ); + tokio::fs::write(&root, external_root).await.unwrap(); + + let error = save_access_sections_to_disk_if_revision( + &root, + &cfg, + &[AccessSection::Users], + Some(&revision), + ) + .await + .unwrap_err(); + + assert_eq!(error.code, "revision_conflict"); + assert_eq!( + tokio::fs::read_to_string(&root).await.unwrap(), + external_root + ); + assert_eq!( + tokio::fs::read_to_string(&included).await.unwrap(), + included_body + ); +} + +#[cfg(unix)] +#[tokio::test] +async fn atomic_write_preserves_existing_file_mode() { + use std::os::unix::fs::{MetadataExt, PermissionsExt}; + + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("config.toml"); + tokio::fs::write(&path, "old").await.unwrap(); + std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o640)).unwrap(); + let before = std::fs::metadata(&path).unwrap(); + + write_atomic(path.clone(), "new".to_string()).await.unwrap(); + + let after = std::fs::metadata(&path).unwrap(); + assert_eq!(after.mode() & 0o7777, 0o640); + assert_eq!(after.uid(), before.uid()); + assert_eq!(after.gid(), before.gid()); +} + +#[tokio::test] +async fn config_sidecar_lock_serializes_competing_revision_writers() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("config.toml"); + let original = concat!( + "[censorship]\n", + "tls_domain = \"original.example\"\n", + "[access.users]\n", + "alice = \"00000000000000000000000000000000\"\n" + ); + tokio::fs::write(&path, original).await.unwrap(); + let graph = ProxyConfig::read_source_graph(&path).unwrap(); + let revision = compute_source_revision(&graph); + + let first = tokio::spawn(write_atomic_if_unchanged( + path.clone(), + revision.clone(), + path.clone(), + original.to_string(), + original.replace("original.example", "first.example"), + )); + let second = tokio::spawn(write_atomic_if_unchanged( + path.clone(), + revision, + path.clone(), + original.to_string(), + original.replace("original.example", "second.example"), + )); + let first = first.await.unwrap(); + let second = second.await.unwrap(); + + assert_ne!(first.is_ok(), second.is_ok()); + let (winner_revision, conflict) = match (first, second) { + (Ok(revision), Err(error)) | (Err(error), Ok(revision)) => (revision, error), + _ => unreachable!("exactly one cooperative writer must commit"), + }; + assert_eq!(conflict.code, "revision_conflict"); + assert_eq!(winner_revision, current_revision(&path).await.unwrap()); + let persisted = tokio::fs::read_to_string(&path).await.unwrap(); + assert!(persisted.contains("first.example") || persisted.contains("second.example")); +} + #[tokio::test] async fn access_mutation_rejects_sections_with_different_source_owners() { let dir = tempfile::tempdir().unwrap(); @@ -307,3 +405,17 @@ fn render_user_rate_limits_section() { assert!(rendered.starts_with("[access.user_rate_limits]\n")); assert!(rendered.contains("alice = { up_bps = 1024, down_bps = 2048 }")); } + +#[cfg(unix)] +#[test] +fn source_owner_normalization_preserves_symlinks() { + use std::os::unix::fs::symlink; + + let dir = tempfile::tempdir().unwrap(); + let real = dir.path().join("real.toml"); + let linked = dir.path().join("linked.toml"); + std::fs::write(&real, "").unwrap(); + symlink(&real, &linked).unwrap(); + + assert_eq!(normalize_source_path(&linked), linked); +} diff --git a/src/api/handler/fixed_routes.rs b/src/api/handler/fixed_routes.rs index 57cdf99..e6e55bb 100644 --- a/src/api/handler/fixed_routes.rs +++ b/src/api/handler/fixed_routes.rs @@ -21,7 +21,6 @@ pub(super) async fn create_user_route( } let expected_revision = parse_if_match(req.headers()); let body = read_json::(req.into_body(), body_limit).await?; - let requested_enabled = body.enabled; let result = create_user(body, expected_revision, shared).await; let (mut data, revision) = match result { Ok(ok) => ok, @@ -34,25 +33,6 @@ pub(super) async fn create_user_route( }; let runtime_cfg = config_rx.borrow().clone(); data.user.in_runtime = runtime_cfg.access.users.contains_key(&data.user.username); - if let Some(enabled) = requested_enabled { - shared - .proxy_shared - .set_user_enabled(&data.user.username, enabled); - if !enabled { - let cancelled = shared - .proxy_shared - .cancel_user_sessions(&data.user.username); - if cancelled > 0 { - shared.runtime_events.record( - "api.user.disable.runtime", - format!( - "username={} cancelled_sessions={}", - data.user.username, cancelled - ), - ); - } - } - } shared.runtime_events.record( "api.user.create.ok", format!("username={}", data.user.username), diff --git a/src/api/handler/user_routes.rs b/src/api/handler/user_routes.rs index 004f01d..a5dfe3b 100644 --- a/src/api/handler/user_routes.rs +++ b/src/api/handler/user_routes.rs @@ -61,7 +61,6 @@ pub(super) async fn handle( }; let runtime_cfg = config_rx.borrow().clone(); data.in_runtime = runtime_cfg.access.users.contains_key(&data.username); - shared.proxy_shared.set_user_enabled(base_user, true); shared .runtime_events .record("api.user.enable.ok", format!("username={}", base_user)); @@ -104,15 +103,9 @@ pub(super) async fn handle( }; let runtime_cfg = config_rx.borrow().clone(); data.in_runtime = runtime_cfg.access.users.contains_key(&data.username); - let newly_disabled = shared.proxy_shared.set_user_enabled(base_user, false); - let cancelled = shared.proxy_shared.cancel_user_sessions(base_user); - shared.runtime_events.record( - "api.user.disable.ok", - format!( - "username={} newly_disabled={} cancelled_sessions={}", - base_user, newly_disabled, cancelled - ), - ); + shared + .runtime_events + .record("api.user.disable.ok", format!("username={}", base_user)); let status = if data.in_runtime { StatusCode::OK } else { @@ -139,38 +132,55 @@ pub(super) async fn handle( )); } let expected_revision = parse_if_match(req.headers()); - let _mutation_guard = shared.mutation_lock.lock().await; - let disk_cfg = load_config_from_disk(&shared.config_path).await?; - ensure_expected_revision(&shared.config_path, expected_revision.as_deref()).await?; - if !disk_cfg.access.users.contains_key(user) { - return Ok(error_response( - request_id, - ApiFailure::new(StatusCode::NOT_FOUND, "not_found", "User not found"), - )); - } - let configured_users = disk_cfg - .access - .users - .keys() - .cloned() - .collect::>(); - let snapshot = match shared.quota_state.reset_user(&configured_users, user).await { - Ok(snapshot) => snapshot, - Err(error) => { - shared.runtime_events.record( - "api.user.reset_quota.failed", - format!("username={} error={}", user, error), + let completion_shared = shared.as_ref().clone(); + let user_owned = user.to_string(); + let completion = shared + .run_mutation_completion(async move { + let _mutation_guard = completion_shared.mutation_lock.lock().await; + let (disk_cfg, _) = load_config_for_mutation( + &completion_shared.config_path, + expected_revision.as_deref(), + ) + .await?; + if !disk_cfg.access.users.contains_key(&user_owned) { + return Err(ApiFailure::new( + StatusCode::NOT_FOUND, + "not_found", + "User not found", + )); + } + let configured_users = disk_cfg + .access + .users + .keys() + .cloned() + .collect::>(); + let snapshot = completion_shared + .quota_state + .reset_user(&configured_users, &user_owned) + .await + .map_err(|error| { + completion_shared.runtime_events.record( + "api.user.reset_quota.failed", + format!("username={} error={}", user_owned, error), + ); + ApiFailure::internal(format!("Failed to reset user quota: {}", error)) + })?; + completion_shared.runtime_events.record( + "api.user.reset_quota.ok", + format!("username={}", user_owned), ); - return Err(ApiFailure::internal(format!( - "Failed to reset user quota: {}", - error - ))); + let revision = current_revision(&completion_shared.config_path).await?; + Ok((snapshot, revision)) + }) + .await; + let (snapshot, revision) = match completion { + Ok(result) => result, + Err(error) if error.code == "not_found" => { + return Ok(error_response(request_id, error)); } + Err(error) => return Err(error), }; - shared - .runtime_events - .record("api.user.reset_quota.ok", format!("username={}", user)); - let revision = current_revision(&shared.config_path).await?; return Ok(success_response( StatusCode::OK, ResetUserQuotaResponse { @@ -271,11 +281,6 @@ pub(super) async fn handle( } let expected_revision = parse_if_match(req.headers()); let body = read_json::(req.into_body(), body_limit).await?; - let enabled_update = match &body.enabled { - Patch::Unchanged => None, - Patch::Remove => Some(true), - Patch::Set(enabled) => Some(*enabled), - }; let result = patch_user(user, body, expected_revision, shared).await; let (mut data, revision) = match result { Ok(ok) => ok, @@ -289,21 +294,6 @@ pub(super) async fn handle( }; let runtime_cfg = config_rx.borrow().clone(); data.in_runtime = runtime_cfg.access.users.contains_key(&data.username); - if let Some(enabled) = enabled_update { - shared - .proxy_shared - .set_user_enabled(&data.username, enabled); - if !enabled { - let cancelled = shared.proxy_shared.cancel_user_sessions(&data.username); - shared.runtime_events.record( - "api.user.disable.runtime", - format!( - "username={} cancelled_sessions={}", - data.username, cancelled - ), - ); - } - } shared .runtime_events .record("api.user.patch.ok", format!("username={}", data.username)); @@ -337,12 +327,9 @@ pub(super) async fn handle( return Err(error); } }; - shared.proxy_shared.set_user_enabled(&deleted_user, true); - let cancelled = shared.proxy_shared.cancel_user_sessions(&deleted_user); - shared.runtime_events.record( - "api.user.delete.ok", - format!("username={} cancelled_sessions={}", deleted_user, cancelled), - ); + shared + .runtime_events + .record("api.user.delete.ok", format!("username={}", deleted_user)); let runtime_cfg = config_rx.borrow().clone(); let in_runtime = runtime_cfg.access.users.contains_key(&deleted_user); let response = DeleteUserResponse { diff --git a/src/api/mod.rs b/src/api/mod.rs index e4aad93..a17b6cf 100644 --- a/src/api/mod.rs +++ b/src/api/mod.rs @@ -17,7 +17,7 @@ use hyper::service::service_fn; use hyper::{Method, Request, Response, StatusCode}; use subtle::ConstantTimeEq; use tokio::net::TcpListener; -use tokio::sync::{Mutex, RwLock, Semaphore, watch}; +use tokio::sync::{Mutex, RwLock, Semaphore, oneshot, watch}; use tokio::time::timeout; use tracing::{debug, info, warn}; @@ -59,7 +59,7 @@ mod web_runtime; mod web_status; use config_store::{ - current_revision, ensure_expected_revision, load_config_for_reload, load_config_from_disk, + current_revision, load_config_for_mutation, load_config_for_reload, load_config_from_disk, parse_if_match, }; use events::ApiEventStore; @@ -70,7 +70,6 @@ use model::{ PatchUserRequest, ResetUserQuotaResponse, RotateSecretRequest, SummaryData, UserActiveIps, is_valid_username, }; -use patch::Patch; use runtime_edge::{ EdgeConnectionsCacheEntry, build_runtime_connections_summary_data, build_runtime_events_recent_data, build_runtime_tls_fingerprints_data, @@ -135,6 +134,7 @@ pub(super) struct ApiShared { pub(super) active_runtime: Arc>, pub(super) web_trace: Arc, pub(super) web_runtime_rx: watch::Receiver, + pub(super) control_plane: ProcessControlPlane, } impl ApiShared { @@ -170,8 +170,32 @@ impl ApiShared { active_runtime: self.active_runtime.clone(), web_trace: self.web_trace.clone(), web_runtime_rx: self.web_runtime_rx.clone(), + control_plane: self.control_plane.clone(), } } + + /// Keeps an accepted mutation alive until persistence and mandatory publication finish. + async fn run_mutation_completion(&self, future: F) -> Result + where + T: Send + 'static, + F: std::future::Future> + Send + 'static, + { + let (result_tx, result_rx) = oneshot::channel(); + self.control_plane + .spawn_completion(async move { + let _ = result_tx.send(future.await); + }) + .map_err(|_| { + ApiFailure::new( + StatusCode::SERVICE_UNAVAILABLE, + "control_plane_shutting_down", + "Control plane is shutting down", + ) + })?; + result_rx.await.map_err(|_| { + ApiFailure::internal("accepted config mutation did not report completion") + })? + } } fn auth_header_matches(actual: &str, expected: &str) -> bool { @@ -365,6 +389,7 @@ pub(crate) async fn serve( active_runtime, web_trace, web_runtime_rx, + control_plane: control_plane.clone(), }); spawn_runtime_watchers( diff --git a/src/api/model.rs b/src/api/model.rs index 3c7497d..127c348 100644 --- a/src/api/model.rs +++ b/src/api/model.rs @@ -124,6 +124,10 @@ pub(super) struct ZeroCoreData { pub(super) conntrack_pressure_active: bool, pub(super) conntrack_event_queue_depth: u64, pub(super) conntrack_rule_apply_ok: bool, + pub(super) conntrack_rule_reconcile_success_total: u64, + pub(super) conntrack_rule_reconcile_error_total: u64, + pub(super) conntrack_rule_rollback_success_total: u64, + pub(super) conntrack_rule_rollback_error_total: u64, pub(super) conntrack_delete_attempt_total: u64, pub(super) conntrack_delete_success_total: u64, pub(super) conntrack_delete_not_found_total: u64, diff --git a/src/api/runtime_edge.rs b/src/api/runtime_edge.rs index 639dbe0..31235bc 100644 --- a/src/api/runtime_edge.rs +++ b/src/api/runtime_edge.rs @@ -314,9 +314,7 @@ async fn recompute_connections_payload( let mut active_users = 0usize; for entry in shared.stats.iter_user_stats() { let user_stats = entry.value(); - let current_connections = user_stats - .curr_connects - .load(std::sync::atomic::Ordering::Relaxed); + let current_connections = shared.stats.get_process_user_curr_connects(entry.key()); let total_octets = user_stats .octets_from_client .load(std::sync::atomic::Ordering::Relaxed) diff --git a/src/api/runtime_min.rs b/src/api/runtime_min.rs index 5338baf..a89f5bc 100644 --- a/src/api/runtime_min.rs +++ b/src/api/runtime_min.rs @@ -1,4 +1,3 @@ -use std::collections::BTreeSet; use std::time::{SystemTime, UNIX_EPOCH}; use serde::Serialize; @@ -17,82 +16,6 @@ pub(super) struct SecurityWhitelistData { pub(super) entries: Vec, } -#[derive(Serialize)] -pub(super) struct RuntimeMePoolStateGenerationData { - pub(super) active_generation: u64, - pub(super) warm_generation: u64, - pub(super) warm_generations: Vec, - pub(super) pending_hardswap_generation: u64, - pub(super) pending_hardswap_age_secs: Option, - pub(super) reinit_inflight: usize, - pub(super) reinit_max_concurrency_effective: usize, - pub(super) draining_generations: Vec, -} - -#[derive(Serialize)] -pub(super) struct RuntimeMePoolStateHardswapData { - pub(super) enabled: bool, - pub(super) pending: bool, -} - -#[derive(Serialize)] -pub(super) struct RuntimeMePoolStateWriterContourData { - pub(super) warm: usize, - pub(super) active: usize, - pub(super) draining: usize, -} - -#[derive(Serialize)] -pub(super) struct RuntimeMePoolStateWriterHealthData { - pub(super) healthy: usize, - pub(super) degraded: usize, - pub(super) draining: usize, -} - -#[derive(Serialize)] -pub(super) struct RuntimeMePoolStateWriterData { - pub(super) total: usize, - pub(super) alive_non_draining: usize, - pub(super) draining: usize, - pub(super) degraded: usize, - pub(super) contour: RuntimeMePoolStateWriterContourData, - pub(super) health: RuntimeMePoolStateWriterHealthData, -} - -#[derive(Serialize)] -pub(super) struct RuntimeMePoolStateRefillDcData { - pub(super) dc: i16, - pub(super) family: &'static str, - pub(super) inflight: usize, -} - -#[derive(Serialize)] -pub(super) struct RuntimeMePoolStateRefillData { - pub(super) inflight_endpoints_total: usize, - pub(super) inflight_dc_total: usize, - pub(super) running_dc_total: usize, - pub(super) pending_dc_total: usize, - pub(super) by_dc: Vec, -} - -#[derive(Serialize)] -pub(super) struct RuntimeMePoolStatePayload { - pub(super) generations: RuntimeMePoolStateGenerationData, - pub(super) hardswap: RuntimeMePoolStateHardswapData, - pub(super) writers: RuntimeMePoolStateWriterData, - pub(super) refill: RuntimeMePoolStateRefillData, -} - -#[derive(Serialize)] -pub(super) struct RuntimeMePoolStateData { - pub(super) enabled: bool, - #[serde(skip_serializing_if = "Option::is_none")] - pub(super) reason: Option<&'static str>, - pub(super) generated_at_epoch_secs: u64, - #[serde(skip_serializing_if = "Option::is_none")] - pub(super) data: Option, -} - #[derive(Serialize)] pub(super) struct RuntimeMeQualityCountersData { pub(super) idle_close_by_peer_total: u64, @@ -285,100 +208,6 @@ pub(super) fn build_security_whitelist_data(cfg: &ProxyConfig) -> SecurityWhitel } } -pub(super) async fn build_runtime_me_pool_state_data(shared: &ApiShared) -> RuntimeMePoolStateData { - let now_epoch_secs = now_epoch_secs(); - let Some(pool) = shared.me_pool.read().await.clone() else { - return RuntimeMePoolStateData { - enabled: false, - reason: Some(SOURCE_UNAVAILABLE_REASON), - generated_at_epoch_secs: now_epoch_secs, - data: None, - }; - }; - - let (status, runtime) = pool.api_coherent_snapshots().await; - let refill = pool.api_refill_snapshot().await; - - let mut draining_generations = BTreeSet::::new(); - let mut contour_warm = 0usize; - let mut contour_active = 0usize; - let mut contour_draining = 0usize; - let mut draining = 0usize; - let mut degraded = 0usize; - let mut healthy = 0usize; - - for writer in &status.writers { - if writer.draining { - draining_generations.insert(writer.generation); - draining += 1; - } - if writer.degraded && !writer.draining { - degraded += 1; - } - if !writer.degraded && !writer.draining { - healthy += 1; - } - match writer.state { - "warm" => contour_warm += 1, - "active" => contour_active += 1, - _ => contour_draining += 1, - } - } - - RuntimeMePoolStateData { - enabled: true, - reason: None, - generated_at_epoch_secs: status.generated_at_epoch_secs, - data: Some(RuntimeMePoolStatePayload { - generations: RuntimeMePoolStateGenerationData { - active_generation: runtime.active_generation, - warm_generation: runtime.warm_generation, - warm_generations: runtime.warm_generations, - pending_hardswap_generation: runtime.pending_hardswap_generation, - pending_hardswap_age_secs: runtime.pending_hardswap_age_secs, - reinit_inflight: runtime.reinit_inflight, - reinit_max_concurrency_effective: runtime.reinit_max_concurrency_effective, - draining_generations: draining_generations.into_iter().collect(), - }, - hardswap: RuntimeMePoolStateHardswapData { - enabled: runtime.hardswap_enabled, - pending: runtime.pending_hardswap_generation != 0, - }, - writers: RuntimeMePoolStateWriterData { - total: status.writers.len(), - alive_non_draining: status.writers.len().saturating_sub(draining), - draining, - degraded, - contour: RuntimeMePoolStateWriterContourData { - warm: contour_warm, - active: contour_active, - draining: contour_draining, - }, - health: RuntimeMePoolStateWriterHealthData { - healthy, - degraded, - draining, - }, - }, - refill: RuntimeMePoolStateRefillData { - inflight_endpoints_total: refill.inflight_endpoints_total, - inflight_dc_total: refill.inflight_dc_total, - running_dc_total: refill.running_dc_total, - pending_dc_total: refill.pending_dc_total, - by_dc: refill - .by_dc - .into_iter() - .map(|entry| RuntimeMePoolStateRefillDcData { - dc: entry.dc, - family: entry.family, - inflight: entry.inflight, - }) - .collect(), - }, - }), - } -} - pub(super) async fn build_runtime_me_quality_data(shared: &ApiShared) -> RuntimeMeQualityData { let now_epoch_secs = now_epoch_secs(); let Some(pool) = shared.me_pool.read().await.clone() else { @@ -541,7 +370,10 @@ pub(super) async fn build_runtime_upstream_quality_data( } } +// ME pool runtime-state projection. +mod me_pool; // NAT/STUN runtime projection and timestamping. mod nat; +pub(super) use me_pool::build_runtime_me_pool_state_data; pub(super) use nat::build_runtime_nat_stun_data; use nat::now_epoch_secs; diff --git a/src/api/runtime_min/me_pool.rs b/src/api/runtime_min/me_pool.rs new file mode 100644 index 0000000..f85e77e --- /dev/null +++ b/src/api/runtime_min/me_pool.rs @@ -0,0 +1,192 @@ +//! ME pool runtime-state projection. + +use std::collections::BTreeSet; + +use serde::Serialize; + +use super::{ApiShared, SOURCE_UNAVAILABLE_REASON, now_epoch_secs}; + +#[derive(Serialize)] +struct RuntimeMePoolStateGenerationData { + active_generation: u64, + warm_generation: u64, + warm_generations: Vec, + pending_hardswap_generation: u64, + pending_hardswap_age_secs: Option, + reinit_inflight: usize, + reinit_max_concurrency_effective: usize, + draining_generations: Vec, +} + +#[derive(Serialize)] +struct RuntimeMePoolStateHardswapData { + enabled: bool, + pending: bool, + pending_writers_current: usize, + pending_writer_deficit: usize, + pending_missing_dc_groups: usize, + pending_map_current: Option, + orphan_warm_writers_current: usize, + replacement_preparing_current: usize, + replacement_retiring_current: usize, +} + +#[derive(Serialize)] +struct RuntimeMePoolStateWriterContourData { + warm: usize, + active: usize, + draining: usize, +} + +#[derive(Serialize)] +struct RuntimeMePoolStateWriterHealthData { + healthy: usize, + degraded: usize, + draining: usize, +} + +#[derive(Serialize)] +struct RuntimeMePoolStateWriterData { + total: usize, + alive_non_draining: usize, + draining: usize, + degraded: usize, + contour: RuntimeMePoolStateWriterContourData, + health: RuntimeMePoolStateWriterHealthData, +} + +#[derive(Serialize)] +struct RuntimeMePoolStateRefillDcData { + dc: i16, + family: &'static str, + inflight: usize, +} + +#[derive(Serialize)] +struct RuntimeMePoolStateRefillData { + inflight_endpoints_total: usize, + inflight_dc_total: usize, + running_dc_total: usize, + pending_dc_total: usize, + by_dc: Vec, +} + +#[derive(Serialize)] +struct RuntimeMePoolStatePayload { + generations: RuntimeMePoolStateGenerationData, + hardswap: RuntimeMePoolStateHardswapData, + writers: RuntimeMePoolStateWriterData, + refill: RuntimeMePoolStateRefillData, +} + +#[derive(Serialize)] +struct RuntimeMePoolStateData { + enabled: bool, + #[serde(skip_serializing_if = "Option::is_none")] + reason: Option<&'static str>, + generated_at_epoch_secs: u64, + #[serde(skip_serializing_if = "Option::is_none")] + data: Option, +} + +/// Builds the bounded runtime ME pool response projection. +pub(in crate::api) async fn build_runtime_me_pool_state_data(shared: &ApiShared) -> impl Serialize { + let now_epoch_secs = now_epoch_secs(); + let Some(pool) = shared.me_pool.read().await.clone() else { + return RuntimeMePoolStateData { + enabled: false, + reason: Some(SOURCE_UNAVAILABLE_REASON), + generated_at_epoch_secs: now_epoch_secs, + data: None, + }; + }; + + let (status, runtime) = pool.api_coherent_snapshots().await; + let refill = pool.api_refill_snapshot().await; + + let mut draining_generations = BTreeSet::::new(); + let mut contour_warm = 0usize; + let mut contour_active = 0usize; + let mut contour_draining = 0usize; + let mut draining = 0usize; + let mut degraded = 0usize; + let mut healthy = 0usize; + + for writer in &status.writers { + if writer.draining { + draining_generations.insert(writer.generation); + draining += 1; + } + if writer.degraded && !writer.draining { + degraded += 1; + } + if !writer.degraded && !writer.draining { + healthy += 1; + } + match writer.state { + "warm" => contour_warm += 1, + "active" => contour_active += 1, + _ => contour_draining += 1, + } + } + + RuntimeMePoolStateData { + enabled: true, + reason: None, + generated_at_epoch_secs: status.generated_at_epoch_secs, + data: Some(RuntimeMePoolStatePayload { + generations: RuntimeMePoolStateGenerationData { + active_generation: runtime.active_generation, + warm_generation: runtime.warm_generation, + warm_generations: runtime.warm_generations, + pending_hardswap_generation: runtime.pending_hardswap_generation, + pending_hardswap_age_secs: runtime.pending_hardswap_age_secs, + reinit_inflight: runtime.reinit_inflight, + reinit_max_concurrency_effective: runtime.reinit_max_concurrency_effective, + draining_generations: draining_generations.into_iter().collect(), + }, + hardswap: RuntimeMePoolStateHardswapData { + enabled: runtime.hardswap_enabled, + pending: runtime.pending_hardswap_generation != 0, + pending_writers_current: runtime.pending_writers_current, + pending_writer_deficit: runtime.pending_writer_deficit, + pending_missing_dc_groups: runtime.pending_missing_dc_groups, + pending_map_current: runtime.pending_map_current, + orphan_warm_writers_current: runtime.orphan_warm_writers_current, + replacement_preparing_current: runtime.replacement_preparing_current, + replacement_retiring_current: runtime.replacement_retiring_current, + }, + writers: RuntimeMePoolStateWriterData { + total: status.writers.len(), + alive_non_draining: status.writers.len().saturating_sub(draining), + draining, + degraded, + contour: RuntimeMePoolStateWriterContourData { + warm: contour_warm, + active: contour_active, + draining: contour_draining, + }, + health: RuntimeMePoolStateWriterHealthData { + healthy, + degraded, + draining, + }, + }, + refill: RuntimeMePoolStateRefillData { + inflight_endpoints_total: refill.inflight_endpoints_total, + inflight_dc_total: refill.inflight_dc_total, + running_dc_total: refill.running_dc_total, + pending_dc_total: refill.pending_dc_total, + by_dc: refill + .by_dc + .into_iter() + .map(|entry| RuntimeMePoolStateRefillDcData { + dc: entry.dc, + family: entry.family, + inflight: entry.inflight, + }) + .collect(), + }, + }), + } +} diff --git a/src/api/runtime_stats.rs b/src/api/runtime_stats.rs index d9536f1..1563df2 100644 --- a/src/api/runtime_stats.rs +++ b/src/api/runtime_stats.rs @@ -61,6 +61,12 @@ pub(super) fn build_zero_all_data(stats: &Stats, configured_users: usize) -> Zer conntrack_pressure_active: stats.get_conntrack_pressure_active(), conntrack_event_queue_depth: stats.get_conntrack_event_queue_depth(), conntrack_rule_apply_ok: stats.get_conntrack_rule_apply_ok(), + conntrack_rule_reconcile_success_total: stats + .get_conntrack_rule_reconcile_success_total(), + conntrack_rule_reconcile_error_total: stats.get_conntrack_rule_reconcile_error_total(), + conntrack_rule_rollback_success_total: stats + .get_conntrack_rule_rollback_success_total(), + conntrack_rule_rollback_error_total: stats.get_conntrack_rule_rollback_error_total(), conntrack_delete_attempt_total: stats.get_conntrack_delete_attempt_total(), conntrack_delete_success_total: stats.get_conntrack_delete_success_total(), conntrack_delete_not_found_total: stats.get_conntrack_delete_not_found_total(), diff --git a/src/api/runtime_zero.rs b/src/api/runtime_zero.rs index e9dce7f..98491af 100644 --- a/src/api/runtime_zero.rs +++ b/src/api/runtime_zero.rs @@ -200,7 +200,7 @@ pub(super) async fn build_runtime_gates_data( && cfg.general.me2dc_fallback && matches!(route_state.mode, RelayRouteMode::Direct); let reroute_to_direct_at_epoch_secs = if reroute_active { - shared.route_runtime.direct_since_epoch_secs() + route_state.direct_since_epoch_secs } else { None }; diff --git a/src/api/users.rs b/src/api/users.rs index 5228277..b16f4bb 100644 --- a/src/api/users.rs +++ b/src/api/users.rs @@ -5,12 +5,13 @@ use hyper::StatusCode; use crate::config::ProxyConfig; use crate::config::RateLimitBps; use crate::ip_tracker::UserIpTracker; +use crate::proxy::user_admission::credential_id_from_hex; use crate::stats::Stats; use super::ApiShared; use super::config_store::{ - AccessSection, current_revision, ensure_expected_revision, load_config_from_disk, - save_access_sections_to_disk, + AccessSection, current_revision, load_config_for_mutation, + save_access_sections_to_disk_if_revision, }; use super::model::{ ApiFailure, CreateUserRequest, CreateUserResponse, PatchUserRequest, RotateSecretRequest, diff --git a/src/api/users/create.rs b/src/api/users/create.rs index a46ed59..509b1ca 100644 --- a/src/api/users/create.rs +++ b/src/api/users/create.rs @@ -4,6 +4,20 @@ pub(in crate::api) async fn create_user( body: CreateUserRequest, expected_revision: Option, shared: &ApiShared, +) -> Result<(CreateUserResponse, String), ApiFailure> { + let shared = shared.clone(); + shared + .clone() + .run_mutation_completion(async move { + create_user_to_completion(body, expected_revision, &shared).await + }) + .await +} + +async fn create_user_to_completion( + body: CreateUserRequest, + expected_revision: Option, + shared: &ApiShared, ) -> Result<(CreateUserResponse, String), ApiFailure> { let touches_user_ad_tags = body.user_ad_tag.is_some(); let touches_user_max_tcp_conns = body.max_tcp_conns.is_some(); @@ -41,9 +55,11 @@ pub(in crate::api) async fn create_user( } let expiration = parse_optional_expiration(body.expiration_rfc3339.as_deref())?; + let credential_id = credential_id_from_hex(&secret) + .ok_or_else(|| ApiFailure::internal("validated user secret could not be decoded"))?; let _guard = shared.mutation_lock.lock().await; - let mut cfg = load_config_from_disk(&shared.config_path).await?; - ensure_expected_revision(&shared.config_path, expected_revision.as_deref()).await?; + let (mut cfg, base_revision) = + load_config_for_mutation(&shared.config_path, expected_revision.as_deref()).await?; if cfg.access.users.contains_key(&body.username) { return Err(ApiFailure::new( @@ -122,9 +138,18 @@ pub(in crate::api) async fn create_user( touched_sections.push(AccessSection::UserEnabled); } - let revision = - save_access_sections_to_disk(&shared.config_path, &cfg, &touched_sections).await?; - drop(_guard); + let revision = save_access_sections_to_disk_if_revision( + &shared.config_path, + &cfg, + &touched_sections, + Some(&base_revision), + ) + .await?; + shared.proxy_shared.stage_user_credential( + &body.username, + credential_id, + cfg.access.is_user_enabled(&body.username), + ); if let Some(limit) = updated_limit { shared @@ -132,6 +157,7 @@ pub(in crate::api) async fn create_user( .set_user_limit(&body.username, limit) .await; } + drop(_guard); let (detected_ip_v4, detected_ip_v6) = shared.detected_link_ips(); let users = users_from_config( diff --git a/src/api/users/lifecycle.rs b/src/api/users/lifecycle.rs index a68850e..b5a595a 100644 --- a/src/api/users/lifecycle.rs +++ b/src/api/users/lifecycle.rs @@ -6,6 +6,22 @@ pub(in crate::api) async fn rotate_secret( body: RotateSecretRequest, expected_revision: Option, shared: &ApiShared, +) -> Result<(CreateUserResponse, String), ApiFailure> { + let shared = shared.clone(); + let user = user.to_string(); + shared + .clone() + .run_mutation_completion(async move { + rotate_secret_to_completion(&user, body, expected_revision, &shared).await + }) + .await +} + +async fn rotate_secret_to_completion( + user: &str, + body: RotateSecretRequest, + expected_revision: Option, + shared: &ApiShared, ) -> Result<(CreateUserResponse, String), ApiFailure> { let secret = body.secret.unwrap_or_else(random_user_secret); if !is_valid_user_secret(&secret) { @@ -13,10 +29,12 @@ pub(in crate::api) async fn rotate_secret( "secret must be exactly 32 hex characters", )); } + let credential_id = credential_id_from_hex(&secret) + .ok_or_else(|| ApiFailure::internal("validated user secret could not be decoded"))?; let _guard = shared.mutation_lock.lock().await; - let mut cfg = load_config_from_disk(&shared.config_path).await?; - ensure_expected_revision(&shared.config_path, expected_revision.as_deref()).await?; + let (mut cfg, base_revision) = + load_config_for_mutation(&shared.config_path, expected_revision.as_deref()).await?; if !cfg.access.users.contains_key(user) { return Err(ApiFailure::new( @@ -29,8 +47,18 @@ pub(in crate::api) async fn rotate_secret( cfg.access.users.insert(user.to_string(), secret.clone()); cfg.validate() .map_err(|e| ApiFailure::bad_request(format!("config validation failed: {}", e)))?; - let revision = - save_access_sections_to_disk(&shared.config_path, &cfg, &[AccessSection::Users]).await?; + let revision = save_access_sections_to_disk_if_revision( + &shared.config_path, + &cfg, + &[AccessSection::Users], + Some(&base_revision), + ) + .await?; + shared.proxy_shared.stage_user_credential( + user, + credential_id, + cfg.access.is_user_enabled(user), + ); drop(_guard); let (detected_ip_v4, detected_ip_v6) = shared.detected_link_ips(); @@ -61,10 +89,25 @@ pub(in crate::api) async fn delete_user( user: &str, expected_revision: Option, shared: &ApiShared, +) -> Result<(String, String), ApiFailure> { + let shared = shared.clone(); + let user = user.to_string(); + shared + .clone() + .run_mutation_completion(async move { + delete_user_to_completion(&user, expected_revision, &shared).await + }) + .await +} + +async fn delete_user_to_completion( + user: &str, + expected_revision: Option, + shared: &ApiShared, ) -> Result<(String, String), ApiFailure> { let _guard = shared.mutation_lock.lock().await; - let mut cfg = load_config_from_disk(&shared.config_path).await?; - ensure_expected_revision(&shared.config_path, expected_revision.as_deref()).await?; + let (mut cfg, base_revision) = + load_config_for_mutation(&shared.config_path, expected_revision.as_deref()).await?; if !cfg.access.users.contains_key(user) { return Err(ApiFailure::new( @@ -107,8 +150,14 @@ pub(in crate::api) async fn delete_user( cfg.validate() .map_err(|e| ApiFailure::bad_request(format!("config validation failed: {}", e)))?; - let revision = - save_access_sections_to_disk(&shared.config_path, &cfg, &touched_sections).await?; + let revision = save_access_sections_to_disk_if_revision( + &shared.config_path, + &cfg, + &touched_sections, + Some(&base_revision), + ) + .await?; + let deleted_incarnation = shared.proxy_shared.delete_user(user).incarnation; let configured_users = cfg.access.users.keys().cloned().collect(); if let Err(error) = shared .quota_state @@ -121,9 +170,12 @@ pub(in crate::api) async fn delete_user( "Deleted user quota checkpoint cleanup will be reconciled on restart" ); } - drop(_guard); shared.ip_tracker.remove_user_limit(user).await; - shared.ip_tracker.clear_user_ips(user).await; + shared + .ip_tracker + .clear_user_ips_if_not_newer(user, deleted_incarnation) + .await; + drop(_guard); Ok((user.to_string(), revision)) } diff --git a/src/api/users/update.rs b/src/api/users/update.rs index bfd1af2..713dd53 100644 --- a/src/api/users/update.rs +++ b/src/api/users/update.rs @@ -5,6 +5,22 @@ pub(in crate::api) async fn patch_user( body: PatchUserRequest, expected_revision: Option, shared: &ApiShared, +) -> Result<(UserInfo, String), ApiFailure> { + let shared = shared.clone(); + let user = user.to_string(); + shared + .clone() + .run_mutation_completion(async move { + patch_user_to_completion(&user, body, expected_revision, &shared).await + }) + .await +} + +async fn patch_user_to_completion( + user: &str, + body: PatchUserRequest, + expected_revision: Option, + shared: &ApiShared, ) -> Result<(UserInfo, String), ApiFailure> { let touches_users = body.secret.is_some(); let touches_user_ad_tags = !matches!(&body.user_ad_tag, Patch::Unchanged); @@ -32,8 +48,8 @@ pub(in crate::api) async fn patch_user( } let expiration = parse_patch_expiration(&body.expiration_rfc3339)?; let _guard = shared.mutation_lock.lock().await; - let mut cfg = load_config_from_disk(&shared.config_path).await?; - ensure_expected_revision(&shared.config_path, expected_revision.as_deref()).await?; + let (mut cfg, base_revision) = + load_config_for_mutation(&shared.config_path, expected_revision.as_deref()).await?; if !cfg.access.users.contains_key(user) { return Err(ApiFailure::new( @@ -138,6 +154,19 @@ pub(in crate::api) async fn patch_user( cfg.validate() .map_err(|e| ApiFailure::bad_request(format!("config validation failed: {}", e)))?; + let staged_credential = + if touches_users || touches_user_enabled { + let secret = cfg + .access + .users + .get(user) + .ok_or_else(|| ApiFailure::internal("updated user secret is missing"))?; + Some(credential_id_from_hex(secret).ok_or_else(|| { + ApiFailure::internal("validated user secret could not be decoded") + })?) + } else { + None + }; let mut touched_sections = Vec::new(); if touches_users { @@ -168,14 +197,27 @@ pub(in crate::api) async fn patch_user( let revision = if touched_sections.is_empty() { current_revision(&shared.config_path).await? } else { - save_access_sections_to_disk(&shared.config_path, &cfg, &touched_sections).await? + save_access_sections_to_disk_if_revision( + &shared.config_path, + &cfg, + &touched_sections, + Some(&base_revision), + ) + .await? }; - drop(_guard); + if let Some(credential_id) = staged_credential { + shared.proxy_shared.stage_user_credential( + user, + credential_id, + cfg.access.is_user_enabled(user), + ); + } match max_unique_ips_change { Some(Some(limit)) => shared.ip_tracker.set_user_limit(user, limit).await, Some(None) => shared.ip_tracker.remove_user_limit(user).await, None => {} } + drop(_guard); let (detected_ip_v4, detected_ip_v6) = shared.detected_link_ips(); let users = users_from_config( &cfg, @@ -199,10 +241,26 @@ pub(in crate::api) async fn set_user_enabled( enabled: bool, expected_revision: Option, shared: &ApiShared, +) -> Result<(UserInfo, String), ApiFailure> { + let shared = shared.clone(); + let user = user.to_string(); + shared + .clone() + .run_mutation_completion(async move { + set_user_enabled_to_completion(&user, enabled, expected_revision, &shared).await + }) + .await +} + +async fn set_user_enabled_to_completion( + user: &str, + enabled: bool, + expected_revision: Option, + shared: &ApiShared, ) -> Result<(UserInfo, String), ApiFailure> { let _guard = shared.mutation_lock.lock().await; - let mut cfg = load_config_from_disk(&shared.config_path).await?; - ensure_expected_revision(&shared.config_path, expected_revision.as_deref()).await?; + let (mut cfg, base_revision) = + load_config_for_mutation(&shared.config_path, expected_revision.as_deref()).await?; if !cfg.access.users.contains_key(user) { return Err(ApiFailure::new( @@ -220,9 +278,22 @@ pub(in crate::api) async fn set_user_enabled( cfg.validate() .map_err(|e| ApiFailure::bad_request(format!("config validation failed: {}", e)))?; - let revision = - save_access_sections_to_disk(&shared.config_path, &cfg, &[AccessSection::UserEnabled]) - .await?; + let credential_id = cfg + .access + .users + .get(user) + .and_then(|secret| credential_id_from_hex(secret)) + .ok_or_else(|| ApiFailure::internal("validated user secret could not be decoded"))?; + let revision = save_access_sections_to_disk_if_revision( + &shared.config_path, + &cfg, + &[AccessSection::UserEnabled], + Some(&base_revision), + ) + .await?; + shared + .proxy_shared + .stage_user_credential(user, credential_id, enabled); drop(_guard); let (detected_ip_v4, detected_ip_v6) = shared.detected_link_ips(); diff --git a/src/api/users/view.rs b/src/api/users/view.rs index f38ce75..4c89d4d 100644 --- a/src/api/users/view.rs +++ b/src/api/users/view.rs @@ -71,7 +71,7 @@ pub(in crate::api) async fn users_from_config( .filter(|limit| *limit > 0) .or((cfg.access.user_max_unique_ips_global_each > 0) .then_some(cfg.access.user_max_unique_ips_global_each)), - current_connections: stats.get_user_curr_connects(&username), + current_connections: stats.get_process_user_curr_connects(&username), active_unique_ips: active_ip_list.len(), active_unique_ips_list: active_ip_list, recent_unique_ips: recent_ip_list.len(), diff --git a/src/api/web_status.rs b/src/api/web_status.rs index d162e0d..e11d1a4 100644 --- a/src/api/web_status.rs +++ b/src/api/web_status.rs @@ -67,6 +67,11 @@ pub(super) async fn render( push_filter_form(&mut html, &query); html.push_str("

Store

"); summary_row(&mut html, "debug enabled", yes_no(status.policy.enabled)); + summary_row( + &mut html, + "sideband", + yes_no(status.policy.bridge_diagnostics_enabled()), + ); summary_row(&mut html, "body capture", body_mode(&status.policy)); summary_row(&mut html, "window seconds", &query.window_secs.to_string()); summary_row( diff --git a/src/api/web_status/tests.rs b/src/api/web_status/tests.rs index 52ef582..1314104 100644 --- a/src/api/web_status/tests.rs +++ b/src/api/web_status/tests.rs @@ -22,6 +22,7 @@ fn html_escaping_covers_active_markup_characters() { async fn renderer_filters_groups_and_sets_control_plane_security_headers() { let policy = WebDebugConfig { enabled: true, + sideband: true, ..Default::default() }; let limits = crate::config::WebLimitsConfig { @@ -58,6 +59,7 @@ async fn renderer_filters_groups_and_sets_control_plane_security_headers() { let body = response.into_body().collect().await.unwrap().to_bytes(); let body = std::str::from_utf8(&body).unwrap(); assert!(body.contains("session_created")); + assert!(body.contains("")); assert!(body.contains("0123456789abcdef")); assert!(body.contains("192.0.2.40")); } diff --git a/src/cli.rs b/src/cli.rs index e740b2b..d9a4e15 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -8,15 +8,22 @@ //! - `run [OPTIONS] [config.toml]` - Run in foreground (default behavior) //! - `healthcheck [OPTIONS] [config.toml]` - Run control-plane health probe -use rand::RngExt; -use std::fs; -use std::path::{Path, PathBuf}; -use std::process::Command; +use std::path::PathBuf; use crate::healthcheck::{self, HealthcheckMode}; #[cfg(unix)] -use crate::daemon::{self, DEFAULT_PID_FILE, DaemonOptions}; +use crate::daemon::{DEFAULT_PID_FILE, DaemonOptions}; + +// Unix daemon control and argument parsing. +#[cfg(unix)] +mod daemon_commands; +// Fire-and-forget installation workflow. +mod init; + +#[cfg(unix)] +pub use daemon_commands::parse_daemon_args; +pub use init::{InitOptions, parse_init_args, run_init}; /// CLI subcommand to execute. #[derive(Debug, Clone, PartialEq, Eq)] @@ -40,13 +47,20 @@ pub enum Subcommand { /// Parsed subcommand with its options. #[derive(Debug)] pub struct ParsedCommand { + /// Selected command mode. pub subcommand: Subcommand, + /// PID file used by daemon-control commands. pub pid_file: PathBuf, + /// Configuration file passed to runtime or healthcheck. pub config_path: String, + /// Requested healthcheck mode. pub healthcheck_mode: HealthcheckMode, + /// Invalid healthcheck mode retained for command diagnostics. pub healthcheck_mode_invalid: Option, #[cfg(unix)] + /// Unix daemon lifecycle options. pub daemon_opts: DaemonOptions, + /// Fire-and-forget initialization options. pub init_opts: Option, } @@ -79,7 +93,6 @@ pub fn parse_command(args: &[String]) -> ParsedCommand { return cmd; } - // Check for subcommand as first argument if let Some(first) = args.first() { match first.as_str() { "start" => { @@ -120,11 +133,9 @@ pub fn parse_command(args: &[String]) -> ParsedCommand { } } - // Parse remaining options let mut i = 0; while i < args.len() { match args[i].as_str() { - // Skip subcommand names "start" | "stop" | "reload" | "status" | "run" | "healthcheck" => {} "--mode" => { i += 1; @@ -154,7 +165,6 @@ pub fn parse_command(args: &[String]) -> ParsedCommand { } } } - // PID file option (for stop/reload/status) "--pid-file" => { i += 1; if i < args.len() { @@ -189,9 +199,9 @@ pub fn parse_command(args: &[String]) -> ParsedCommand { #[cfg(unix)] pub fn execute_subcommand(cmd: &ParsedCommand) -> Option { match cmd.subcommand { - Subcommand::Stop => Some(cmd_stop(&cmd.pid_file)), - Subcommand::Reload => Some(cmd_reload(&cmd.pid_file)), - Subcommand::Status => Some(cmd_status(&cmd.pid_file)), + Subcommand::Stop => Some(daemon_commands::stop(&cmd.pid_file)), + Subcommand::Reload => Some(daemon_commands::reload(&cmd.pid_file)), + Subcommand::Status => Some(daemon_commands::status(&cmd.pid_file)), Subcommand::Healthcheck => { if let Some(invalid_mode) = cmd.healthcheck_mode_invalid.as_ref() { if invalid_mode.is_empty() { @@ -224,6 +234,7 @@ pub fn execute_subcommand(cmd: &ParsedCommand) -> Option { } } +/// Executes a non-server subcommand on platforms without daemon support. #[cfg(not(unix))] pub fn execute_subcommand(cmd: &ParsedCommand) -> Option { match cmd.subcommand { @@ -261,487 +272,3 @@ pub fn execute_subcommand(cmd: &ParsedCommand) -> Option { Subcommand::Run | Subcommand::Start => None, } } - -/// Stop command: send SIGTERM to the running daemon. -#[cfg(unix)] -fn cmd_stop(pid_file: &Path) -> i32 { - use nix::sys::signal::Signal; - - println!("Stopping telemt daemon..."); - - match daemon::signal_pid_file(pid_file, Signal::SIGTERM) { - Ok(()) => { - println!("Stop signal sent successfully"); - - // Wait for process to exit (up to 10 seconds) - for _ in 0..20 { - std::thread::sleep(std::time::Duration::from_millis(500)); - if let daemon::DaemonStatus::NotRunning = daemon::check_status(pid_file) { - println!("Daemon stopped"); - return 0; - } - } - println!("Daemon may still be shutting down"); - 0 - } - Err(e) => { - eprintln!("Failed to stop daemon: {}", e); - 1 - } - } -} - -/// Reload command: send SIGHUP to trigger config reload. -#[cfg(unix)] -fn cmd_reload(pid_file: &Path) -> i32 { - use nix::sys::signal::Signal; - - println!("Reloading telemt configuration..."); - - match daemon::signal_pid_file(pid_file, Signal::SIGHUP) { - Ok(()) => { - println!("Reload signal sent successfully"); - 0 - } - Err(e) => { - eprintln!("Failed to reload daemon: {}", e); - 1 - } - } -} - -/// Status command: check if daemon is running. -#[cfg(unix)] -fn cmd_status(pid_file: &Path) -> i32 { - match daemon::check_status(pid_file) { - daemon::DaemonStatus::Running(pid) => { - println!("telemt is running (pid {})", pid); - 0 - } - daemon::DaemonStatus::Stale(pid) => { - println!("telemt is not running (stale pid file, was pid {})", pid); - // Clean up stale PID file - let _ = std::fs::remove_file(pid_file); - 1 - } - daemon::DaemonStatus::NotRunning => { - println!("telemt is not running"); - 1 - } - } -} - -/// Options for the init command -#[derive(Debug, Clone)] -pub struct InitOptions { - pub port: u16, - pub domain: String, - pub secret: Option, - pub username: String, - pub config_dir: PathBuf, - pub no_start: bool, -} - -/// Parse daemon-related options from CLI args. -#[cfg(unix)] -pub fn parse_daemon_args(args: &[String]) -> DaemonOptions { - let mut opts = DaemonOptions::default(); - let mut i = 0; - - while i < args.len() { - match args[i].as_str() { - "--daemon" | "-d" => { - opts.daemonize = true; - } - "--foreground" | "-f" => { - opts.foreground = true; - } - "--pid-file" => { - i += 1; - if i < args.len() { - opts.pid_file = Some(PathBuf::from(&args[i])); - } - } - s if s.starts_with("--pid-file=") => { - opts.pid_file = Some(PathBuf::from(s.trim_start_matches("--pid-file="))); - } - "--run-as-user" => { - i += 1; - if i < args.len() { - opts.user = Some(args[i].clone()); - } - } - s if s.starts_with("--run-as-user=") => { - opts.user = Some(s.trim_start_matches("--run-as-user=").to_string()); - } - "--run-as-group" => { - i += 1; - if i < args.len() { - opts.group = Some(args[i].clone()); - } - } - s if s.starts_with("--run-as-group=") => { - opts.group = Some(s.trim_start_matches("--run-as-group=").to_string()); - } - "--working-dir" => { - i += 1; - if i < args.len() { - opts.working_dir = Some(PathBuf::from(&args[i])); - } - } - s if s.starts_with("--working-dir=") => { - opts.working_dir = Some(PathBuf::from(s.trim_start_matches("--working-dir="))); - } - _ => {} - } - i += 1; - } - - opts -} - -impl Default for InitOptions { - fn default() -> Self { - Self { - port: 443, - domain: "www.google.com".to_string(), - secret: None, - username: "user".to_string(), - config_dir: PathBuf::from("/etc/telemt"), - no_start: false, - } - } -} - -/// Parse --init subcommand options from CLI args. -/// -/// Returns `Some(InitOptions)` if `--init` was found, `None` otherwise. -pub fn parse_init_args(args: &[String]) -> Option { - if !args.iter().any(|a| a == "--init") { - return None; - } - - let mut opts = InitOptions::default(); - let mut i = 0; - - while i < args.len() { - match args[i].as_str() { - "--port" => { - i += 1; - if i < args.len() { - opts.port = args[i].parse().unwrap_or(443); - } - } - "--domain" => { - i += 1; - if i < args.len() { - opts.domain = args[i].clone(); - } - } - "--secret" => { - i += 1; - if i < args.len() { - opts.secret = Some(args[i].clone()); - } - } - "--user" => { - i += 1; - if i < args.len() { - opts.username = args[i].clone(); - } - } - "--config-dir" => { - i += 1; - if i < args.len() { - opts.config_dir = PathBuf::from(&args[i]); - } - } - "--no-start" => { - opts.no_start = true; - } - _ => {} - } - i += 1; - } - - Some(opts) -} - -/// Run the fire-and-forget setup. -pub fn run_init(opts: InitOptions) -> Result<(), Box> { - use crate::service::{self, InitSystem, ServiceOptions}; - - eprintln!("[telemt] Fire-and-forget setup"); - eprintln!(); - - // 1. Detect init system - let init_system = service::detect_init_system(); - eprintln!("[+] Detected init system: {}", init_system); - - // 2. Generate or validate secret - let secret = match opts.secret { - Some(s) => { - if s.len() != 32 || !s.chars().all(|c| c.is_ascii_hexdigit()) { - eprintln!("[error] Secret must be exactly 32 hex characters"); - std::process::exit(1); - } - s - } - None => generate_secret(), - }; - - eprintln!("[+] Secret: {}", secret); - eprintln!("[+] User: {}", opts.username); - eprintln!("[+] Port: {}", opts.port); - eprintln!("[+] Domain: {}", opts.domain); - - // 3. Create config directory - fs::create_dir_all(&opts.config_dir)?; - let config_path = opts.config_dir.join("config.toml"); - - // 4. Write config - let config_content = generate_config(&opts.username, &secret, opts.port, &opts.domain); - fs::write(&config_path, &config_content)?; - eprintln!("[+] Config written to {}", config_path.display()); - - // 5. Generate and write service file - let exe_path = - std::env::current_exe().unwrap_or_else(|_| PathBuf::from("/usr/local/bin/telemt")); - - let service_opts = ServiceOptions { - exe_path: &exe_path, - config_path: &config_path, - user: None, // Let systemd/init handle user - group: None, - pid_file: "/var/run/telemt.pid", - working_dir: Some("/var/lib/telemt"), - description: "Telemt MTProxy - Telegram MTProto Proxy", - }; - - let service_path = service::service_file_path(init_system); - let service_content = service::generate_service_file(init_system, &service_opts); - - // Ensure parent directory exists - if let Some(parent) = Path::new(service_path).parent() { - let _ = fs::create_dir_all(parent); - } - - match fs::write(service_path, &service_content) { - Ok(()) => { - eprintln!("[+] Service file written to {}", service_path); - - // Make script executable for OpenRC/FreeBSD - #[cfg(unix)] - if init_system == InitSystem::OpenRC || init_system == InitSystem::FreeBSDRc { - use std::os::unix::fs::PermissionsExt; - let mut perms = fs::metadata(service_path)?.permissions(); - perms.set_mode(0o755); - fs::set_permissions(service_path, perms)?; - } - } - Err(e) => { - eprintln!("[!] Cannot write service file (run as root?): {}", e); - eprintln!("[!] Manual service file content:"); - eprintln!("{}", service_content); - - // Still print links and installation instructions - eprintln!(); - eprintln!("{}", service::installation_instructions(init_system)); - print_links(&opts.username, &secret, opts.port, &opts.domain); - return Ok(()); - } - } - - // 6. Install and enable service based on init system - match init_system { - InitSystem::Systemd => { - run_cmd("systemctl", &["daemon-reload"]); - run_cmd("systemctl", &["enable", "telemt.service"]); - eprintln!("[+] Service enabled"); - - if !opts.no_start { - run_cmd("systemctl", &["start", "telemt.service"]); - eprintln!("[+] Service started"); - - std::thread::sleep(std::time::Duration::from_secs(1)); - let status = Command::new("systemctl") - .args(["is-active", "telemt.service"]) - .output(); - - match status { - Ok(out) if out.status.success() => { - eprintln!("[+] Service is running"); - } - _ => { - eprintln!("[!] Service may not have started correctly"); - eprintln!("[!] Check: journalctl -u telemt.service -n 20"); - } - } - } else { - eprintln!("[+] Service not started (--no-start)"); - eprintln!("[+] Start manually: systemctl start telemt.service"); - } - } - InitSystem::OpenRC => { - run_cmd("rc-update", &["add", "telemt", "default"]); - eprintln!("[+] Service enabled"); - - if !opts.no_start { - run_cmd("rc-service", &["telemt", "start"]); - eprintln!("[+] Service started"); - } else { - eprintln!("[+] Service not started (--no-start)"); - eprintln!("[+] Start manually: rc-service telemt start"); - } - } - InitSystem::FreeBSDRc => { - run_cmd("sysrc", &["telemt_enable=YES"]); - eprintln!("[+] Service enabled"); - - if !opts.no_start { - run_cmd("service", &["telemt", "start"]); - eprintln!("[+] Service started"); - } else { - eprintln!("[+] Service not started (--no-start)"); - eprintln!("[+] Start manually: service telemt start"); - } - } - InitSystem::Unknown => { - eprintln!("[!] Unknown init system - service file written but not installed"); - eprintln!("[!] You may need to install it manually"); - } - } - - eprintln!(); - - // 7. Print links - print_links(&opts.username, &secret, opts.port, &opts.domain); - - Ok(()) -} - -fn generate_secret() -> String { - let mut rng = rand::rng(); - let bytes: Vec = (0..16).map(|_| rng.random::()).collect(); - hex::encode(bytes) -} - -fn generate_config(username: &str, secret: &str, port: u16, domain: &str) -> String { - format!( - r#"# Telemt MTProxy — auto-generated config -# Re-run `telemt --init` to regenerate - -show_link = ["{username}"] - -[general] -# prefer_ipv6 is deprecated; use [network].prefer -prefer_ipv6 = false -fast_mode = true -use_middle_proxy = false -log_level = "normal" -desync_all_full = false -update_every = 43200 -hardswap = false -me_pool_drain_ttl_secs = 90 -me_instadrain = false -me_pool_drain_threshold = 32 -me_pool_drain_soft_evict_grace_secs = 10 -me_pool_drain_soft_evict_per_writer = 2 -me_pool_drain_soft_evict_budget_per_core = 16 -me_pool_drain_soft_evict_cooldown_ms = 1000 -me_bind_stale_mode = "never" -me_pool_min_fresh_ratio = 0.8 -me_reinit_drain_timeout_secs = 90 -tg_connect = 10 - -[network] -ipv4 = true -ipv6 = true -prefer = 4 -multipath = false - -[general.modes] -classic = false -secure = false -tls = true - -[server] -listen_addr_ipv4 = "0.0.0.0" -listen_addr_ipv6 = "::" - -[[server.listeners]] -ip = "0.0.0.0" -port = {port} -# reuse_allow = false # Set true only when intentionally running multiple telemt instances on same port - -[[server.listeners]] -ip = "::" -port = {port} - -[timeouts] -client_first_byte_idle_secs = 300 -client_handshake = 60 -client_keepalive = 60 -client_ack = 300 - -[censorship] -tls_domain = "{domain}" -mask = true -mask_port = 443 -fake_cert_len = 2048 -serverhello_compact = false -tls_full_cert_ttl_secs = 90 - -[access] -user_max_tcp_conns_global_each = 0 -replay_check_len = 65536 -replay_window_secs = 120 -ignore_time_skew = false - -[access.users] -{username} = "{secret}" - -[[upstreams]] -type = "direct" -enabled = true -weight = 10 -# Optional per-upstream DC family policy: -# ipv6 = true -# prefer = 6 -"#, - username = username, - secret = secret, - port = port, - domain = domain, - ) -} - -fn run_cmd(cmd: &str, args: &[&str]) { - match Command::new(cmd).args(args).output() { - Ok(output) => { - if !output.status.success() { - let stderr = String::from_utf8_lossy(&output.stderr); - eprintln!("[!] {} {} failed: {}", cmd, args.join(" "), stderr.trim()); - } - } - Err(e) => { - eprintln!("[!] Failed to run {} {}: {}", cmd, args.join(" "), e); - } - } -} - -fn print_links(username: &str, secret: &str, port: u16, domain: &str) { - let domain_hex = hex::encode(domain); - - println!("=== Proxy Links ==="); - println!("[{}]", username); - println!( - " EE-TLS: tg://proxy?server=YOUR_SERVER_IP&port={}&secret=ee{}{}", - port, secret, domain_hex - ); - println!(); - println!("Replace YOUR_SERVER_IP with your server's public IP."); - println!("The proxy will auto-detect and display the correct link on startup."); - println!("Check: journalctl -u telemt.service | head -30"); - println!("==================="); -} diff --git a/src/cli/daemon_commands.rs b/src/cli/daemon_commands.rs new file mode 100644 index 0000000..ee6fcca --- /dev/null +++ b/src/cli/daemon_commands.rs @@ -0,0 +1,141 @@ +use std::path::{Path, PathBuf}; + +use crate::daemon::{self, DaemonOptions}; + +/// Parses daemon-related options from CLI arguments. +pub fn parse_daemon_args(args: &[String]) -> DaemonOptions { + let mut opts = DaemonOptions::default(); + let mut i = 0; + + while i < args.len() { + match args[i].as_str() { + "--daemon" | "-d" => { + opts.daemonize = true; + } + "--foreground" | "-f" => { + opts.foreground = true; + } + "--pid-file" => { + i += 1; + if i < args.len() { + opts.pid_file = Some(PathBuf::from(&args[i])); + } + } + s if s.starts_with("--pid-file=") => { + opts.pid_file = Some(PathBuf::from(s.trim_start_matches("--pid-file="))); + } + "--run-as-user" => { + i += 1; + if i < args.len() { + opts.user = Some(args[i].clone()); + } + } + s if s.starts_with("--run-as-user=") => { + opts.user = Some(s.trim_start_matches("--run-as-user=").to_string()); + } + "--run-as-group" => { + i += 1; + if i < args.len() { + opts.group = Some(args[i].clone()); + } + } + s if s.starts_with("--run-as-group=") => { + opts.group = Some(s.trim_start_matches("--run-as-group=").to_string()); + } + "--working-dir" => { + i += 1; + if i < args.len() { + opts.working_dir = Some(PathBuf::from(&args[i])); + } + } + s if s.starts_with("--working-dir=") => { + opts.working_dir = Some(PathBuf::from(s.trim_start_matches("--working-dir="))); + } + _ => {} + } + i += 1; + } + + opts +} + +/// Sends SIGTERM and waits briefly for graceful PID-file cleanup. +pub(super) fn stop(pid_file: &Path) -> i32 { + use nix::sys::signal::Signal; + + println!("Stopping telemt daemon..."); + + match daemon::signal_pid_file(pid_file, Signal::SIGTERM) { + Ok(()) => { + println!("Stop signal sent successfully"); + + // Wait for process to exit for up to ten seconds. + for _ in 0..20 { + std::thread::sleep(std::time::Duration::from_millis(500)); + if let daemon::DaemonStatus::NotRunning = daemon::check_status(pid_file) { + println!("Daemon stopped"); + return 0; + } + } + println!("Daemon may still be shutting down"); + 0 + } + Err(e) => { + eprintln!("Failed to stop daemon: {}", e); + 1 + } + } +} + +/// Sends SIGHUP to trigger configuration reload. +pub(super) fn reload(pid_file: &Path) -> i32 { + use nix::sys::signal::Signal; + + println!("Reloading telemt configuration..."); + + match daemon::signal_pid_file(pid_file, Signal::SIGHUP) { + Ok(()) => { + println!("Reload signal sent successfully"); + 0 + } + Err(e) => { + eprintln!("Failed to reload daemon: {}", e); + 1 + } + } +} + +/// Reports daemon status without mutating PID lifecycle state. +pub(super) fn status(pid_file: &Path) -> i32 { + match daemon::check_status(pid_file) { + daemon::DaemonStatus::Running(pid) => { + println!("telemt is running (pid {})", pid); + 0 + } + daemon::DaemonStatus::Stale(pid) => { + println!("telemt is not running (stale pid file, was pid {})", pid); + 1 + } + daemon::DaemonStatus::NotRunning => { + println!("telemt is not running"); + 1 + } + } +} + +#[cfg(test)] +mod tests { + use std::fs; + + use super::*; + + #[test] + fn status_does_not_remove_stale_pid_file() { + let directory = tempfile::tempdir().unwrap(); + let pid_file = directory.path().join("telemt.pid"); + fs::write(&pid_file, b"2000000000\n").unwrap(); + + assert_eq!(status(&pid_file), 1); + assert!(pid_file.exists()); + } +} diff --git a/src/cli/init.rs b/src/cli/init.rs new file mode 100644 index 0000000..43282c3 --- /dev/null +++ b/src/cli/init.rs @@ -0,0 +1,368 @@ +use std::path::{Path, PathBuf}; +use std::process::Command; + +use rand::RngExt; + +use crate::util::trusted_command::resolve_trusted_helper; + +/// Options for the fire-and-forget init command. +#[derive(Debug, Clone)] +pub struct InitOptions { + /// Public listener port. + pub port: u16, + /// TLS camouflage domain. + pub domain: String, + /// Optional pre-generated proxy secret. + pub secret: Option, + /// Initial access username. + pub username: String, + /// Destination directory for generated configuration. + pub config_dir: PathBuf, + /// Generate service files without starting the service. + pub no_start: bool, +} + +impl Default for InitOptions { + fn default() -> Self { + Self { + port: 443, + domain: "www.google.com".to_string(), + secret: None, + username: "user".to_string(), + config_dir: PathBuf::from("/etc/telemt"), + no_start: false, + } + } +} + +/// Parse --init subcommand options from CLI args. +/// +/// Returns `Some(InitOptions)` if `--init` was found, `None` otherwise. +pub fn parse_init_args(args: &[String]) -> Option { + if !args.iter().any(|a| a == "--init") { + return None; + } + + let mut opts = InitOptions::default(); + let mut i = 0; + + while i < args.len() { + match args[i].as_str() { + "--port" => { + i += 1; + if i < args.len() { + opts.port = args[i].parse().unwrap_or(443); + } + } + "--domain" => { + i += 1; + if i < args.len() { + opts.domain = args[i].clone(); + } + } + "--secret" => { + i += 1; + if i < args.len() { + opts.secret = Some(args[i].clone()); + } + } + "--user" => { + i += 1; + if i < args.len() { + opts.username = args[i].clone(); + } + } + "--config-dir" => { + i += 1; + if i < args.len() { + opts.config_dir = PathBuf::from(&args[i]); + } + } + "--no-start" => { + opts.no_start = true; + } + _ => {} + } + i += 1; + } + + Some(opts) +} + +/// Run the fire-and-forget setup. +pub fn run_init(opts: InitOptions) -> Result<(), Box> { + use crate::service::{self, InitSystem, ServiceOptions}; + + eprintln!("[telemt] Fire-and-forget setup"); + eprintln!(); + + let init_system = service::detect_init_system(); + eprintln!("[+] Detected init system: {}", init_system); + + let secret = match opts.secret { + Some(s) => { + if s.len() != 32 || !s.chars().all(|c| c.is_ascii_hexdigit()) { + eprintln!("[error] Secret must be exactly 32 hex characters"); + std::process::exit(1); + } + s + } + None => generate_secret(), + }; + + eprintln!("[+] Secret: {}", secret); + eprintln!("[+] User: {}", opts.username); + eprintln!("[+] Port: {}", opts.port); + eprintln!("[+] Domain: {}", opts.domain); + + let config_path = opts.config_dir.join("config.toml"); + let config_content = generate_config(&opts.username, &secret, opts.port, &opts.domain); + write_init_file(&config_path, &config_content, 0o600)?; + eprintln!("[+] Config written to {}", config_path.display()); + + let exe_path = + std::env::current_exe().unwrap_or_else(|_| PathBuf::from("/usr/local/bin/telemt")); + let service_opts = ServiceOptions { + exe_path: &exe_path, + config_path: &config_path, + // Let the selected init system manage process identity. + user: None, + group: None, + pid_file: "/var/run/telemt.pid", + working_dir: Some("/var/lib/telemt"), + description: "Telemt MTProxy - Telegram MTProto Proxy", + }; + + let service_path = service::service_file_path(init_system); + let service_content = service::generate_service_file(init_system, &service_opts); + let service_mode = if init_system == InitSystem::OpenRC || init_system == InitSystem::FreeBSDRc + { + 0o755 + } else { + 0o644 + }; + match write_init_file(Path::new(service_path), &service_content, service_mode) { + Ok(()) => { + eprintln!("[+] Service file written to {}", service_path); + } + Err(e) => { + eprintln!("[!] Cannot write service file (run as root?): {}", e); + eprintln!("[!] Manual service file content:"); + eprintln!("{}", service_content); + eprintln!(); + eprintln!("{}", service::installation_instructions(init_system)); + print_links(&opts.username, &secret, opts.port, &opts.domain); + return Ok(()); + } + } + + match init_system { + InitSystem::Systemd => { + run_cmd("systemctl", &["daemon-reload"]); + run_cmd("systemctl", &["enable", "telemt.service"]); + eprintln!("[+] Service enabled"); + + if !opts.no_start { + run_cmd("systemctl", &["start", "telemt.service"]); + eprintln!("[+] Service started"); + + std::thread::sleep(std::time::Duration::from_secs(1)); + let status = resolve_trusted_helper("systemctl").and_then(|command_path| { + Command::new(command_path) + .args(["is-active", "telemt.service"]) + .output() + .ok() + }); + match status { + Some(out) if out.status.success() => { + eprintln!("[+] Service is running"); + } + _ => { + eprintln!("[!] Service may not have started correctly"); + eprintln!("[!] Check: journalctl -u telemt.service -n 20"); + } + } + } else { + eprintln!("[+] Service not started (--no-start)"); + eprintln!("[+] Start manually: systemctl start telemt.service"); + } + } + InitSystem::OpenRC => { + run_cmd("rc-update", &["add", "telemt", "default"]); + eprintln!("[+] Service enabled"); + + if !opts.no_start { + run_cmd("rc-service", &["telemt", "start"]); + eprintln!("[+] Service started"); + } else { + eprintln!("[+] Service not started (--no-start)"); + eprintln!("[+] Start manually: rc-service telemt start"); + } + } + InitSystem::FreeBSDRc => { + run_cmd("sysrc", &["telemt_enable=YES"]); + eprintln!("[+] Service enabled"); + + if !opts.no_start { + run_cmd("service", &["telemt", "start"]); + eprintln!("[+] Service started"); + } else { + eprintln!("[+] Service not started (--no-start)"); + eprintln!("[+] Start manually: service telemt start"); + } + } + InitSystem::Unknown => { + eprintln!("[!] Unknown init system - service file written but not installed"); + eprintln!("[!] You may need to install it manually"); + } + } + + eprintln!(); + print_links(&opts.username, &secret, opts.port, &opts.domain); + Ok(()) +} + +fn write_init_file(path: &Path, contents: &str, mode: u32) -> std::io::Result<()> { + #[cfg(unix)] + { + crate::util::secure_fs::atomic_replace(path, contents.as_bytes(), mode) + } + #[cfg(not(unix))] + { + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent)?; + } + let _ = mode; + std::fs::write(path, contents) + } +} + +fn generate_secret() -> String { + let mut rng = rand::rng(); + let bytes: Vec = (0..16).map(|_| rng.random::()).collect(); + hex::encode(bytes) +} + +fn generate_config(username: &str, secret: &str, port: u16, domain: &str) -> String { + format!( + r#"# Telemt MTProxy — auto-generated config +# Re-run `telemt --init` to regenerate + +show_link = ["{username}"] + +[general] +# prefer_ipv6 is deprecated; use [network].prefer +prefer_ipv6 = false +fast_mode = true +use_middle_proxy = false +log_level = "normal" +desync_all_full = false +update_every = 43200 +hardswap = false +me_pool_drain_ttl_secs = 90 +me_instadrain = false +me_pool_drain_threshold = 32 +me_pool_drain_soft_evict_grace_secs = 10 +me_pool_drain_soft_evict_per_writer = 2 +me_pool_drain_soft_evict_budget_per_core = 16 +me_pool_drain_soft_evict_cooldown_ms = 1000 +me_bind_stale_mode = "never" +me_pool_min_fresh_ratio = 0.8 +me_reinit_drain_timeout_secs = 90 +tg_connect = 10 + +[network] +ipv4 = true +ipv6 = true +prefer = 4 +multipath = false + +[general.modes] +classic = false +secure = false +tls = true + +[server] +listen_addr_ipv4 = "0.0.0.0" +listen_addr_ipv6 = "::" + +[[server.listeners]] +ip = "0.0.0.0" +port = {port} +# reuse_allow = false # Set true only when intentionally running multiple telemt instances on same port + +[[server.listeners]] +ip = "::" +port = {port} + +[timeouts] +client_first_byte_idle_secs = 300 +client_handshake = 60 +client_keepalive = 60 +client_ack = 300 + +[censorship] +tls_domain = "{domain}" +mask = true +mask_port = 443 +fake_cert_len = 2048 +serverhello_compact = false +tls_full_cert_ttl_secs = 90 + +[access] +user_max_tcp_conns_global_each = 0 +replay_check_len = 65536 +replay_window_secs = 120 +ignore_time_skew = false + +[access.users] +{username} = "{secret}" + +[[upstreams]] +type = "direct" +enabled = true +weight = 10 +# Optional per-upstream DC family policy: +# ipv6 = true +# prefer = 6 +"#, + username = username, + secret = secret, + port = port, + domain = domain, + ) +} + +fn run_cmd(cmd: &str, args: &[&str]) { + let Some(command_path) = resolve_trusted_helper(cmd) else { + eprintln!("[!] Refusing unavailable or untrusted command: {}", cmd); + return; + }; + match Command::new(command_path).args(args).output() { + Ok(output) => { + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr); + eprintln!("[!] {} {} failed: {}", cmd, args.join(" "), stderr.trim()); + } + } + Err(e) => { + eprintln!("[!] Failed to run {} {}: {}", cmd, args.join(" "), e); + } + } +} + +fn print_links(username: &str, secret: &str, port: u16, domain: &str) { + let domain_hex = hex::encode(domain); + + println!("=== Proxy Links ==="); + println!("[{}]", username); + println!( + " EE-TLS: tg://proxy?server=YOUR_SERVER_IP&port={}&secret=ee{}{}", + port, secret, domain_hex + ); + println!(); + println!("Replace YOUR_SERVER_IP with your server's public IP."); + println!("The proxy will auto-detect and display the correct link on startup."); + println!("Check: journalctl -u telemt.service | head -30"); + println!("==================="); +} diff --git a/src/config/hot_reload.rs b/src/config/hot_reload.rs index 1a5f511..b01ac54 100644 --- a/src/config/hot_reload.rs +++ b/src/config/hot_reload.rs @@ -62,5 +62,8 @@ use reporting::log_changes; #[cfg(test)] use watcher::{ReloadState, reload_config}; +#[cfg(test)] +#[path = "hot_reload/base_path_tests.rs"] +mod base_path_tests; #[cfg(test)] mod tests; diff --git a/src/config/hot_reload/base_path_tests.rs b/src/config/hot_reload/base_path_tests.rs new file mode 100644 index 0000000..8d06772 --- /dev/null +++ b/src/config/hot_reload/base_path_tests.rs @@ -0,0 +1,81 @@ +use base64::Engine as _; + +use super::*; + +fn write_base_path_config(path: &Path, base_path: &str) { + let base_path = if base_path.is_empty() { + String::new() + } else { + format!("base_path = \"{base_path}\"\n") + }; + let config = format!( + r#" +[access.users] +alice = "000102030405060708090a0b0c0d0e0f" + +[[server.listeners]] +ip = "127.0.0.1" +port = 18080 +transport = "web" +proxy_protocol = false +web_client_ip_source = "x_forwarded_for" +web_trusted_proxy_cidrs = ["127.0.0.1/32"] + +[web] +enabled = true + +[[web.vhosts]] +host = "proxy.example.com" +{base_path}public_addr = "203.0.113.10:443" + +[web.vhosts.decoy] +mode = "http_upstream" +upstream = "http://127.0.0.1:18081" + +[[web.vhosts.profiles]] +user = "alice" +secret_mode = "plain" +"#, + ); + std::fs::write(path, config).unwrap(); +} + +#[test] +fn reload_rejects_invalid_base_then_publishes_route_identity_together() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("config.toml"); + write_base_path_config(&path, ""); + let initial = Arc::new(ProxyConfig::load(&path).unwrap()); + let initial_hash = ProxyConfig::load_with_metadata(&path) + .unwrap() + .rendered_hash; + let initial_capability = initial.web.runtime.as_ref().unwrap().capabilities[0]; + let (config_tx, _config_rx) = watch::channel(Arc::clone(&initial)); + let (log_tx, _log_rx) = watch::channel(initial.general.log_level.clone()); + let mut reload_state = ReloadState::new(Some(initial_hash)); + + write_base_path_config(&path, "/invalid"); + reload_config(&path, &config_tx, &log_tx, None, None, &mut reload_state); + let unchanged = config_tx.borrow().clone(); + assert!(Arc::ptr_eq(&unchanged, &initial)); + assert_eq!(unchanged.web.vhosts[0].base_path, ""); + assert_eq!( + unchanged.web.runtime.as_ref().unwrap().capabilities[0], + initial_capability + ); + + write_base_path_config(&path, "dobry-cola-super-app"); + reload_config(&path, &config_tx, &log_tx, None, None, &mut reload_state); + let applied = config_tx.borrow().clone(); + let runtime = applied.web.runtime.as_ref().unwrap(); + let vhost = &runtime.vhosts["proxy.example.com"]; + assert_eq!(applied.web.vhosts[0].base_path, "dobry-cola-super-app"); + assert_eq!(vhost.base, "/dobry-cola-super-app/"); + assert_eq!(vhost.capabilities[0], vhost.profiles[0].capability); + assert_eq!(runtime.capabilities.as_ref(), vhost.capabilities.as_ref()); + assert!(!runtime.capabilities.contains(&initial_capability)); + assert_eq!( + base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(vhost.capabilities[0]), + "hHz99Xs93EN1j91G9gpNepXwGNNt5YdAFkEVk_LlqdQ" + ); +} diff --git a/src/config/hot_reload/tests.rs b/src/config/hot_reload/tests.rs index 4741f10..f81bbd7 100644 --- a/src/config/hot_reload/tests.rs +++ b/src/config/hot_reload/tests.rs @@ -144,11 +144,13 @@ fn web_debug_policy_is_hot_while_debug_capacity_is_process_owned() { let old = sample_config(); let mut new = old.clone(); new.web.debug.enabled = true; + new.web.debug.sideband = true; new.web.debug.default_window_secs = 60; new.web.limits.debug_records_capacity += 1; let applied = overlay_hot_fields(&old, &new); assert!(applied.web.debug.enabled); + assert!(applied.web.debug.sideband); assert_eq!(applied.web.debug.default_window_secs, 60); assert_eq!( applied.web.limits.debug_records_capacity, diff --git a/src/config/hot_reload/watcher.rs b/src/config/hot_reload/watcher.rs index 9197d2b..178691a 100644 --- a/src/config/hot_reload/watcher.rs +++ b/src/config/hot_reload/watcher.rs @@ -52,15 +52,24 @@ impl ReloadState { } fn normalize_watch_path(path: &Path) -> PathBuf { - path.canonicalize().unwrap_or_else(|_| { - if path.is_absolute() { - path.to_path_buf() - } else { - std::env::current_dir() - .map(|cwd| cwd.join(path)) - .unwrap_or_else(|_| path.to_path_buf()) + let absolute = if path.is_absolute() { + path.to_path_buf() + } else { + std::env::current_dir() + .map(|cwd| cwd.join(path)) + .unwrap_or_else(|_| path.to_path_buf()) + }; + let mut normalized = PathBuf::new(); + for component in absolute.components() { + match component { + std::path::Component::CurDir => {} + std::path::Component::ParentDir => { + normalized.pop(); + } + component => normalized.push(component.as_os_str()), } - }) + } + normalized } fn sync_watch_paths( @@ -433,3 +442,21 @@ pub fn spawn_config_watcher( (config_rx, log_rx, task) } + +#[cfg(all(test, unix))] +mod path_tests { + use std::os::unix::fs::symlink; + + use super::normalize_watch_path; + + #[test] + fn watch_path_normalization_preserves_symlinks() { + let dir = tempfile::tempdir().unwrap(); + let real = dir.path().join("real.toml"); + let linked = dir.path().join("linked.toml"); + std::fs::write(&real, "").unwrap(); + symlink(&real, &linked).unwrap(); + + assert_eq!(normalize_watch_path(&linked), linked); + } +} diff --git a/src/config/load.rs b/src/config/load.rs index e6560a2..2ec5bd5 100644 --- a/src/config/load.rs +++ b/src/config/load.rs @@ -9,6 +9,7 @@ use rand::RngExt; use serde::{Deserialize, Serialize}; use tracing::warn; +use crate::crypto::sha256; use crate::error::{ProxyError, Result}; use super::defaults::*; @@ -35,7 +36,9 @@ mod validate_server; mod validate_web; mod validation; -use self::includes::{hash_rendered_snapshot, normalize_config_path, preprocess_includes}; +use self::includes::{ + hash_rendered_snapshot, normalize_config_path, preprocess_includes, read_config_source, +}; use self::normalize::{ is_valid_ad_tag, is_valid_tls_domain_name, normalize_domain_to_ascii, normalize_exclusive_mask_target, normalize_mask_host_to_ascii, parse_exclusive_mask_target, @@ -63,9 +66,9 @@ const MAX_API_REQUEST_BODY_LIMIT_BYTES: usize = 1024 * 1024; pub(crate) struct LoadedConfig { /// Validated and normalized effective configuration. pub(crate) config: ProxyConfig, - /// Canonical paths participating in the recursive include graph. + /// Normalized absolute paths participating in the recursive include graph. pub(crate) source_files: Vec, - /// Raw source bytes keyed by canonical source path. + /// Raw source bytes keyed by normalized absolute source path. pub(crate) source_contents: BTreeMap, /// Legacy hash of the include-expanded rendered snapshot. pub(crate) rendered_hash: u64, @@ -74,7 +77,7 @@ pub(crate) struct LoadedConfig { /// Raw recursive source graph captured before typed deserialization. #[derive(Debug, Clone)] pub(crate) struct ConfigSourceGraph { - /// Raw source bytes keyed by canonical source path. + /// Raw source bytes keyed by normalized absolute source path. pub(crate) source_contents: BTreeMap, /// Include-expanded TOML used for typed deserialization. pub(crate) rendered: String, @@ -174,21 +177,42 @@ impl ProxyConfig { source_overrides: &BTreeMap, ) -> Result { let path = path.as_ref(); - let normalized_path = normalize_config_path(path); + let mut previous = Self::capture_source_graph(path, source_overrides)?; + for _ in 0..2 { + let current = Self::capture_source_graph(path, source_overrides)?; + if current.source_contents == previous.source_contents + && current.rendered == previous.rendered + { + return Ok(current); + } + previous = current; + } + Err(ProxyError::Config( + "config source graph changed repeatedly while it was read".to_string(), + )) + } + + fn capture_source_graph( + path: &Path, + source_overrides: &BTreeMap, + ) -> Result { + let path = path.as_ref(); + let (normalized_path, disk_content) = read_config_source(path)?; let content = source_overrides .get(&normalized_path) .cloned() - .map(Ok) - .unwrap_or_else(|| std::fs::read_to_string(path)) - .map_err(|e| ProxyError::Config(e.to_string()))?; - let base_dir = path.parent().unwrap_or(Path::new(".")); + .unwrap_or(disk_content); + let base_dir = normalized_path + .parent() + .unwrap_or(Path::new(".")) + .to_path_buf(); let mut source_files = BTreeSet::new(); source_files.insert(normalized_path.clone()); let mut source_contents = BTreeMap::new(); source_contents.insert(normalized_path, content.clone()); let processed = preprocess_includes( &content, - base_dir, + &base_dir, 0, &mut source_files, &mut source_contents, @@ -230,6 +254,25 @@ impl ProxyConfig { self.runtime_user_auth.as_deref() } + /// Returns the credential identity frozen into this runtime snapshot. + pub(crate) fn runtime_user_credential_id(&self, user: &str) -> Option<[u8; 16]> { + self.runtime_user_auth() + .and_then(|snapshot| snapshot.credential_id_by_name(user)) + .or_else(|| { + self.access + .users + .get(user) + .and_then(|secret| hex::decode(secret).ok()) + .and_then(|secret| <[u8; 16]>::try_from(secret).ok()) + .map(|secret| { + let digest = sha256(&secret); + let mut credential_id = [0; 16]; + credential_id.copy_from_slice(&digest[..16]); + credential_id + }) + }) + } + /// Validates cross-field configuration invariants after deserialization. pub fn validate(&self) -> Result<()> { if self.access.users.is_empty() { diff --git a/src/config/load/includes.rs b/src/config/load/includes.rs index 89780b2..b964157 100644 --- a/src/config/load/includes.rs +++ b/src/config/load/includes.rs @@ -4,16 +4,27 @@ use std::path::{Path, PathBuf}; use crate::error::{ProxyError, Result}; +const MAX_CONFIG_SOURCE_BYTES: usize = 8 * 1024 * 1024; + pub(super) fn normalize_config_path(path: &Path) -> PathBuf { - path.canonicalize().unwrap_or_else(|_| { - if path.is_absolute() { - path.to_path_buf() - } else { - std::env::current_dir() - .map(|cwd| cwd.join(path)) - .unwrap_or_else(|_| path.to_path_buf()) + let absolute = if path.is_absolute() { + path.to_path_buf() + } else { + std::env::current_dir() + .map(|cwd| cwd.join(path)) + .unwrap_or_else(|_| path.to_path_buf()) + }; + let mut normalized = PathBuf::new(); + for component in absolute.components() { + match component { + std::path::Component::CurDir => {} + std::path::Component::ParentDir => { + normalized.pop(); + } + component => normalized.push(component.as_os_str()), } - }) + } + normalized } pub(super) fn hash_rendered_snapshot(rendered: &str) -> u64 { @@ -22,6 +33,29 @@ pub(super) fn hash_rendered_snapshot(rendered: &str) -> u64 { hasher.finish() } +pub(super) fn read_config_source(path: &Path) -> Result<(PathBuf, String)> { + #[cfg(unix)] + let bytes = crate::util::secure_fs::read_regular_limited(path, MAX_CONFIG_SOURCE_BYTES) + .map_err(|error| ProxyError::Config(error.to_string()))?; + #[cfg(not(unix))] + let bytes = std::fs::read(path).map_err(|error| ProxyError::Config(error.to_string()))?; + if bytes.len() > MAX_CONFIG_SOURCE_BYTES { + return Err(ProxyError::Config(format!( + "config source `{}` exceeds {} bytes", + path.display(), + MAX_CONFIG_SOURCE_BYTES + ))); + } + let contents = String::from_utf8(bytes).map_err(|error| { + ProxyError::Config(format!( + "config source `{}` is not valid UTF-8: {error}", + path.display() + )) + })?; + let normalized = normalize_config_path(path); + Ok((normalized, contents)) +} + pub(super) fn preprocess_includes( content: &str, base_dir: &Path, @@ -41,16 +75,17 @@ pub(super) fn preprocess_includes( if let Some(rest) = rest.strip_prefix('=') { let path_str = rest.trim().trim_matches('"'); let resolved = base_dir.join(path_str); - let normalized = normalize_config_path(&resolved); - source_files.insert(normalized.clone()); + let (normalized, disk_contents) = read_config_source(&resolved)?; let included = source_overrides .get(&normalized) .cloned() - .map(Ok) - .unwrap_or_else(|| std::fs::read_to_string(&resolved)) - .map_err(|e| ProxyError::Config(e.to_string()))?; - source_contents.insert(normalized, included.clone()); - let included_dir = resolved.parent().unwrap_or(base_dir); + .or_else(|| source_contents.get(&normalized).cloned()) + .unwrap_or(disk_contents); + source_files.insert(normalized.clone()); + source_contents + .entry(normalized.clone()) + .or_insert_with(|| included.clone()); + let included_dir = normalized.parent().unwrap_or(base_dir); output.push_str(&preprocess_includes( &included, included_dir, diff --git a/src/config/load/runtime_auth.rs b/src/config/load/runtime_auth.rs index 01d1c4b..c7e5ec1 100644 --- a/src/config/load/runtime_auth.rs +++ b/src/config/load/runtime_auth.rs @@ -2,6 +2,7 @@ use std::collections::HashMap; use std::collections::hash_map::DefaultHasher; use std::hash::Hasher; +use crate::crypto::sha256; use crate::error::{ProxyError, Result}; const ACCESS_SECRET_BYTES: usize = 16; @@ -11,6 +12,7 @@ const ACCESS_SECRET_BYTES: usize = 16; pub(crate) struct UserAuthSnapshot { entries: Vec, by_name: HashMap, + by_hint_key: HashMap>, sni_index: HashMap>, sni_initial_index: HashMap>, } @@ -19,16 +21,23 @@ pub(crate) struct UserAuthSnapshot { pub(crate) struct UserAuthEntry { pub(crate) user: String, pub(crate) secret: [u8; ACCESS_SECRET_BYTES], + /// Stable secret identity used by process-wide admission fencing. + pub(crate) credential_id: [u8; 16], + /// Stable compact key used only to resolve bounded authentication hints. + pub(crate) hint_key: u64, } impl UserAuthSnapshot { pub(super) fn from_users(users: &HashMap) -> Result { let mut entries = Vec::with_capacity(users.len()); let mut by_name = HashMap::with_capacity(users.len()); + let mut by_hint_key = HashMap::with_capacity(users.len()); let mut sni_index = HashMap::with_capacity(users.len()); let mut sni_initial_index = HashMap::with_capacity(users.len()); - for (user, secret_hex) in users { + let mut ordered_users = users.iter().collect::>(); + ordered_users.sort_unstable_by(|(left, _), (right, _)| left.cmp(right)); + for (user, secret_hex) in ordered_users { let decoded = hex::decode(secret_hex).map_err(|_| ProxyError::InvalidSecret { user: user.clone(), reason: "Must be 32 hex characters".to_string(), @@ -46,11 +55,30 @@ impl UserAuthSnapshot { let mut secret = [0u8; ACCESS_SECRET_BYTES]; secret.copy_from_slice(&decoded); + let digest = sha256(&secret); + let mut credential_id = [0; 16]; + credential_id.copy_from_slice(&digest[..16]); + let hint_key = u64::from_le_bytes([ + credential_id[0], + credential_id[1], + credential_id[2], + credential_id[3], + credential_id[4], + credential_id[5], + credential_id[6], + credential_id[7], + ]) | 1; entries.push(UserAuthEntry { user: user.clone(), secret, + credential_id, + hint_key, }); by_name.insert(user.clone(), user_id); + by_hint_key + .entry(hint_key) + .or_insert_with(Vec::new) + .push(user_id); sni_index .entry(Self::sni_lookup_hash(user)) .or_insert_with(Vec::new) @@ -70,6 +98,7 @@ impl UserAuthSnapshot { Ok(Self { entries, by_name, + by_hint_key, sni_index, sni_initial_index, }) @@ -88,6 +117,18 @@ impl UserAuthSnapshot { self.entries.get(idx) } + /// Returns the stable credential identity for an exact configured username. + pub(crate) fn credential_id_by_name(&self, user: &str) -> Option<[u8; 16]> { + self.user_id_by_name(user) + .and_then(|user_id| self.entry_by_id(user_id)) + .map(|entry| entry.credential_id) + } + + /// Returns every bounded authentication candidate sharing a stable hint key. + pub(crate) fn candidate_ids_by_hint_key(&self, hint_key: u64) -> Option<&[u32]> { + self.by_hint_key.get(&hint_key).map(Vec::as_slice) + } + pub(crate) fn sni_candidates(&self, sni: &str) -> Option<&[u32]> { self.sni_index .get(&Self::sni_lookup_hash(sni)) @@ -110,3 +151,48 @@ impl UserAuthSnapshot { hasher.finish() } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn credential_hint_survives_positional_id_shift() { + let mut initial = HashMap::new(); + initial.insert( + "alice".to_string(), + "11111111111111111111111111111111".to_string(), + ); + initial.insert( + "bob".to_string(), + "22222222222222222222222222222222".to_string(), + ); + let initial = UserAuthSnapshot::from_users(&initial).unwrap(); + let initial_id = initial.user_id_by_name("alice").unwrap(); + let hint_key = initial.entry_by_id(initial_id).unwrap().hint_key; + + let mut reloaded = HashMap::new(); + reloaded.insert( + "aaron".to_string(), + "33333333333333333333333333333333".to_string(), + ); + reloaded.insert( + "alice".to_string(), + "11111111111111111111111111111111".to_string(), + ); + reloaded.insert( + "bob".to_string(), + "22222222222222222222222222222222".to_string(), + ); + let reloaded = UserAuthSnapshot::from_users(&reloaded).unwrap(); + let reloaded_id = reloaded.user_id_by_name("alice").unwrap(); + + assert_ne!(initial_id, reloaded_id); + assert!( + reloaded + .candidate_ids_by_hint_key(hint_key) + .unwrap() + .contains(&reloaded_id) + ); + } +} diff --git a/src/config/load/runtime_web.rs b/src/config/load/runtime_web.rs index 159a5fa..90698cb 100644 --- a/src/config/load/runtime_web.rs +++ b/src/config/load/runtime_web.rs @@ -5,15 +5,33 @@ use std::path::Path; use std::sync::Arc; #[cfg(unix)] -use std::os::unix::fs::OpenOptionsExt; +use std::ffi::OsString; +#[cfg(unix)] +use std::os::unix::ffi::OsStringExt; +#[cfg(unix)] +use std::os::unix::fs::MetadataExt; + +#[cfg(unix)] +use nix::dir::Dir; +#[cfg(unix)] +use nix::fcntl::{OFlag, openat}; +#[cfg(unix)] +use nix::sys::stat::Mode; use bytes::Bytes; use hmac::{Hmac, Mac}; use sha2::{Digest, Sha256}; use super::*; +#[cfg(unix)] +use crate::util::secure_fs::open_dir_nofollow; -const WEB_CAPABILITY_CONTEXT: &[u8] = b"tdesktop-web-proxy-bridge-v1\n"; +// Path-based static snapshot fallback for platforms without directory descriptors. +#[cfg(not(unix))] +mod static_site_fallback; + +const WEB_CAPABILITY_CONTEXT_V1: &[u8] = b"tdesktop-web-proxy-bridge-v1\n"; +const WEB_CAPABILITY_CONTEXT_V2: &[u8] = b"tdesktop-web-proxy-bridge-v2\n"; const WEB_DEBUG_FINGERPRINT_CONTEXT: &[u8] = b"telemt-web-debug-key-fingerprint-v1\0"; const MAX_WEB_STATIC_DEPTH: usize = 64; @@ -24,6 +42,7 @@ pub(super) fn rebuild(config: &mut ProxyConfig) -> Result<()> { })?; let mut runtime_vhosts = BTreeMap::new(); let mut runtime_profiles = Vec::new(); + let mut runtime_capabilities = Vec::new(); let mut static_files = 0usize; let mut static_bytes = 0usize; @@ -50,8 +69,11 @@ pub(super) fn rebuild(config: &mut ProxyConfig) -> Result<()> { })?; let (client_secret, client_secret_len) = client_secret(auth_entry.secret, profile.secret_mode); - let capability = - derive_web_capability(&client_secret[..client_secret_len], vhost.host.as_bytes())?; + let capability = derive_web_capability( + &client_secret[..client_secret_len], + vhost.host.as_bytes(), + vhost.base_path.as_bytes(), + )?; let key_fingerprint = debug_key_fingerprint(&client_secret[..client_secret_len]); if !capabilities.insert(capability) { return Err(ProxyError::Config(format!( @@ -63,6 +85,7 @@ pub(super) fn rebuild(config: &mut ProxyConfig) -> Result<()> { host: vhost.host.clone(), public_addr: vhost.public_addr, user: profile.user.clone(), + credential_id: auth_entry.credential_id, secret_mode: profile.secret_mode, carrier: config.web.carrier, carrier_negotiation_enabled: config.web.carrier_negotiation_enabled(), @@ -86,6 +109,7 @@ pub(super) fn rebuild(config: &mut ProxyConfig) -> Result<()> { .unwrap_or(config.web.limits.max_streams_per_session), }); capability_table.push(capability); + runtime_capabilities.push(capability); profiles.push(Arc::clone(&runtime_profile)); runtime_profiles.push(runtime_profile); } @@ -93,6 +117,11 @@ pub(super) fn rebuild(config: &mut ProxyConfig) -> Result<()> { vhost.host.clone(), Arc::new(WebRuntimeVhost { host: vhost.host.clone(), + base: if vhost.base_path.is_empty() { + "/".to_string() + } else { + format!("/{}/", vhost.base_path) + }, decoy_fasttrack_mode: config.web.decoy_fasttrack_mode, decoy, decoy_header_secs: config.web.timeouts.decoy_header_secs, @@ -105,6 +134,7 @@ pub(super) fn rebuild(config: &mut ProxyConfig) -> Result<()> { config.web.runtime = Some(Arc::new(WebRuntimeConfig { vhosts: runtime_vhosts, profiles: runtime_profiles, + capabilities: runtime_capabilities.into_boxed_slice(), })); Ok(()) } @@ -116,12 +146,23 @@ fn debug_key_fingerprint(secret: &[u8]) -> String { hex::encode(&digest.finalize()[..8]) } -/// Derives the Telegram Desktop WEB capability for one exact secret and host. -pub(crate) fn derive_web_capability(secret: &[u8], host: &[u8]) -> Result<[u8; 32]> { +/// Derives the Telegram Desktop WEB capability for one exact secret, host, and base path. +pub(crate) fn derive_web_capability( + secret: &[u8], + host: &[u8], + base_path: &[u8], +) -> Result<[u8; 32]> { let mut mac = Hmac::::new_from_slice(secret) .map_err(|_| ProxyError::Config("WEB capability secret must not be empty".to_string()))?; - mac.update(WEB_CAPABILITY_CONTEXT); - mac.update(host); + if base_path.is_empty() { + mac.update(WEB_CAPABILITY_CONTEXT_V1); + mac.update(host); + } else { + mac.update(WEB_CAPABILITY_CONTEXT_V2); + mac.update(host); + mac.update(b"\n"); + mac.update(base_path); + } Ok(mac.finalize().into_bytes().into()) } @@ -192,34 +233,31 @@ fn load_static_site( total_files: &mut usize, total_bytes: &mut usize, ) -> Result { - let root_metadata = fs::symlink_metadata(root).map_err(|error| { - ProxyError::Config(format!( - "failed to inspect WEB static directory `{}`: {error}", - root.display() - )) - })?; - if root_metadata.file_type().is_symlink() || !root_metadata.is_dir() { - return Err(ProxyError::Config(format!( - "WEB static directory `{}` must be a real directory, not a symlink", - root.display() - ))); - } - let canonical_root = fs::canonicalize(root).map_err(|error| { - ProxyError::Config(format!( - "failed to canonicalize WEB static directory `{}`: {error}", - root.display() - )) - })?; let mut assets = BTreeMap::new(); - load_static_directory( - &canonical_root, - &canonical_root, - &mut assets, - total_files, - total_bytes, - limits, - 0, - )?; + #[cfg(unix)] + { + let directory = open_static_root(root)?; + load_static_directory( + directory, + Path::new(""), + root, + &mut assets, + total_files, + total_bytes, + limits, + 0, + )?; + } + #[cfg(not(unix))] + { + static_site_fallback::load_static_site_by_path( + root, + limits, + &mut assets, + total_files, + total_bytes, + )?; + } if !assets.contains_key(&format!("/{index}")) { return Err(ProxyError::Config(format!( "WEB static directory `{}` does not contain index `{index}`", @@ -232,54 +270,95 @@ fn load_static_site( }) } +#[cfg(unix)] +fn open_static_root(root: &Path) -> Result { + let descriptor = open_dir_nofollow(root).map_err(|error| { + ProxyError::Config(format!( + "WEB static directory `{}` must be a real directory, not a symlink: {error}", + root.display() + )) + })?; + Dir::from_fd(descriptor).map_err(|error| { + ProxyError::Config(format!( + "failed to read WEB static directory `{}`: {error}", + root.display() + )) + }) +} + +#[cfg(unix)] fn load_static_directory( + mut directory: Dir, + relative: &Path, root: &Path, - directory: &Path, assets: &mut BTreeMap, total_files: &mut usize, total_bytes: &mut usize, limits: &WebLimitsConfig, depth: usize, ) -> Result<()> { - let entries = fs::read_dir(directory).map_err(|error| { - ProxyError::Config(format!( - "failed to read WEB static directory `{}`: {error}", - directory.display() - )) - })?; - for entry in entries { + let mut entries = Vec::new(); + for entry in directory.iter() { let entry = entry.map_err(|error| { - ProxyError::Config(format!("failed to read WEB static entry: {error}")) + ProxyError::Config(format!( + "failed to read WEB static directory `{}`: {error}", + root.join(relative).display() + )) })?; + let name = entry.file_name().to_bytes(); + if name == b"." || name == b".." { + continue; + } if *total_files >= limits.max_static_files { return Err(ProxyError::Config( "WEB static entries exceed process-wide web.limits.max_static_files".to_string(), )); } *total_files += 1; - let path = entry.path(); - let file_type = entry.file_type().map_err(|error| { + entries.push(OsString::from_vec(name.to_vec())); + } + entries.sort_unstable(); + + for name in entries { + let relative_path = relative.join(&name); + let display_path = root.join(&relative_path); + let descriptor = openat( + &directory, + name.as_os_str(), + OFlag::O_RDONLY | OFlag::O_NOFOLLOW | OFlag::O_CLOEXEC, + Mode::empty(), + ) + .map_err(|error| { ProxyError::Config(format!( - "failed to inspect WEB static entry `{}`: {error}", - path.display() + "failed to open WEB static entry `{}` without following symlinks: {error}", + display_path.display() )) })?; - if file_type.is_symlink() { - return Err(ProxyError::Config(format!( - "WEB static entry `{}` must not be a symlink", - path.display() - ))); - } - if file_type.is_dir() { + let file = fs::File::from(descriptor); + let metadata = file.metadata().map_err(|error| { + ProxyError::Config(format!( + "failed to inspect WEB static entry `{}`: {error}", + display_path.display() + )) + })?; + if metadata.is_dir() { if depth >= MAX_WEB_STATIC_DEPTH { return Err(ProxyError::Config(format!( "WEB static directory `{}` exceeds the maximum nesting depth", - path.display() + display_path.display() ))); } + let descriptor = file.into(); + let child = Dir::from_fd(descriptor).map_err(|error| { + ProxyError::Config(format!( + "failed to open WEB static directory `{}`: {error}", + display_path.display() + )) + })?; load_static_directory( + child, + &relative_path, root, - &path, assets, total_files, total_bytes, @@ -288,83 +367,107 @@ fn load_static_directory( )?; continue; } - if !file_type.is_file() { - return Err(ProxyError::Config(format!( - "WEB static entry `{}` must be a regular file", - path.display() - ))); - } - let mut options = fs::OpenOptions::new(); - options.read(true); - #[cfg(unix)] - options.custom_flags(libc::O_CLOEXEC | libc::O_NOFOLLOW); - let file = options.open(&path).map_err(|error| { - ProxyError::Config(format!( - "failed to open WEB static file `{}`: {error}", - path.display() - )) - })?; - let metadata = file.metadata().map_err(|error| { - ProxyError::Config(format!( - "failed to inspect WEB static file `{}`: {error}", - path.display() - )) - })?; if !metadata.is_file() { return Err(ProxyError::Config(format!( - "WEB static entry `{}` changed before it was opened", - path.display() + "WEB static entry `{}` must be a regular file", + display_path.display() ))); } - let file_len = usize::try_from(metadata.len()).map_err(|_| { - ProxyError::Config(format!("WEB static file `{}` is too large", path.display())) - })?; - if file_len > limits.max_static_file_bytes { - return Err(ProxyError::Config(format!( - "WEB static file `{}` exceeds web.limits.max_static_file_bytes", - path.display() - ))); - } - *total_bytes = total_bytes.checked_add(file_len).ok_or_else(|| { - ProxyError::Config("WEB static snapshot byte count overflowed usize".to_string()) - })?; - if *total_bytes > limits.max_static_bytes { - return Err(ProxyError::Config( - "WEB static snapshots exceed process-wide web.limits.max_static_bytes".to_string(), - )); - } - let relative = path.strip_prefix(root).map_err(|_| { - ProxyError::Config("WEB static path escaped its configured root".to_string()) - })?; - let route = static_route(relative)?; - let mut body = Vec::with_capacity(file_len); - file.take(limits.max_static_file_bytes as u64 + 1) - .read_to_end(&mut body) - .map_err(|error| { - ProxyError::Config(format!( - "failed to read WEB static file `{}`: {error}", - path.display() - )) - })?; - if body.len() != file_len { - return Err(ProxyError::Config(format!( - "WEB static file `{}` changed while its snapshot was built", - path.display() - ))); - } - let etag = format!("\"{}\"", hex::encode(Sha256::digest(&body))); - assets.insert( - route, - WebStaticAsset { - body: Bytes::from(body), - content_type: static_content_type(&path), - etag, - }, - ); + load_static_file( + file, + &metadata, + &relative_path, + &display_path, + assets, + total_bytes, + limits, + )?; } Ok(()) } +fn load_static_file( + mut file: fs::File, + metadata: &fs::Metadata, + relative: &Path, + display_path: &Path, + assets: &mut BTreeMap, + total_bytes: &mut usize, + limits: &WebLimitsConfig, +) -> Result<()> { + let file_len = usize::try_from(metadata.len()).map_err(|_| { + ProxyError::Config(format!( + "WEB static file `{}` is too large", + display_path.display() + )) + })?; + if file_len > limits.max_static_file_bytes { + return Err(ProxyError::Config(format!( + "WEB static file `{}` exceeds web.limits.max_static_file_bytes", + display_path.display() + ))); + } + *total_bytes = total_bytes.checked_add(file_len).ok_or_else(|| { + ProxyError::Config("WEB static snapshot byte count overflowed usize".to_string()) + })?; + if *total_bytes > limits.max_static_bytes { + return Err(ProxyError::Config( + "WEB static snapshots exceed process-wide web.limits.max_static_bytes".to_string(), + )); + } + let route = static_route(relative)?; + let mut body = Vec::with_capacity(file_len); + file.by_ref() + .take(limits.max_static_file_bytes as u64 + 1) + .read_to_end(&mut body) + .map_err(|error| { + ProxyError::Config(format!( + "failed to read WEB static file `{}`: {error}", + display_path.display() + )) + })?; + let final_metadata = file.metadata().map_err(|error| { + ProxyError::Config(format!( + "failed to recheck WEB static file `{}`: {error}", + display_path.display() + )) + })?; + if body.len() != file_len || !static_file_version_matches(metadata, &final_metadata) { + return Err(ProxyError::Config(format!( + "WEB static file `{}` changed while its snapshot was built", + display_path.display() + ))); + } + let etag = format!("\"{}\"", hex::encode(Sha256::digest(&body))); + assets.insert( + route, + WebStaticAsset { + body: Bytes::from(body), + content_type: static_content_type(relative), + etag, + }, + ); + Ok(()) +} + +#[cfg(unix)] +fn static_file_version_matches(before: &fs::Metadata, after: &fs::Metadata) -> bool { + before.dev() == after.dev() + && before.ino() == after.ino() + && before.len() == after.len() + && before.mtime() == after.mtime() + && before.mtime_nsec() == after.mtime_nsec() + && before.ctime() == after.ctime() + && before.ctime_nsec() == after.ctime_nsec() +} + +#[cfg(not(unix))] +fn static_file_version_matches(before: &fs::Metadata, after: &fs::Metadata) -> bool { + before.len() == after.len() + && before.modified().ok() == after.modified().ok() + && before.created().ok() == after.created().ok() +} + fn static_route(relative: &Path) -> Result { let mut route = String::new(); for component in relative.components() { @@ -402,26 +505,7 @@ fn static_content_type(path: &Path) -> &'static str { } } +// Runtime WEB construction tests remain separate from the production loader. #[cfg(test)] -mod tests { - use base64::Engine as _; - - use super::*; - - #[test] - fn capability_matches_reference_vectors() { - let secret = hex::decode("000102030405060708090a0b0c0d0e0f").unwrap(); - let plain = derive_web_capability(&secret, b"proxy.example.com").unwrap(); - assert_eq!( - base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(plain), - "MHLEY5PmW1GWqJkSrlmJpvJUiLhBH_QKy6yKg8a0JPk" - ); - let mut dd_secret = vec![0xdd]; - dd_secret.extend_from_slice(&secret); - let dd = derive_web_capability(&dd_secret, b"proxy.example.com").unwrap(); - assert_eq!( - base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(dd), - "IpJrt3e7sKtzPyoXy6w-Zj6GGEvsvclN66JzQEfPYLA" - ); - } -} +#[path = "runtime_web/tests.rs"] +mod tests; diff --git a/src/config/load/runtime_web/static_site_fallback.rs b/src/config/load/runtime_web/static_site_fallback.rs new file mode 100644 index 0000000..ba33da2 --- /dev/null +++ b/src/config/load/runtime_web/static_site_fallback.rs @@ -0,0 +1,138 @@ +use std::collections::BTreeMap; +use std::fs; +use std::path::Path; + +use super::*; + +/// Builds a bounded static-site snapshot on platforms without directory descriptors. +pub(super) fn load_static_site_by_path( + root: &Path, + limits: &WebLimitsConfig, + assets: &mut BTreeMap, + total_files: &mut usize, + total_bytes: &mut usize, +) -> Result<()> { + let root_metadata = fs::symlink_metadata(root).map_err(|error| { + ProxyError::Config(format!( + "failed to inspect WEB static directory `{}`: {error}", + root.display() + )) + })?; + if root_metadata.file_type().is_symlink() || !root_metadata.is_dir() { + return Err(ProxyError::Config(format!( + "WEB static directory `{}` must be a real directory, not a symlink", + root.display() + ))); + } + let canonical_root = fs::canonicalize(root).map_err(|error| { + ProxyError::Config(format!( + "failed to canonicalize WEB static directory `{}`: {error}", + root.display() + )) + })?; + load_static_directory( + &canonical_root, + &canonical_root, + assets, + total_files, + total_bytes, + limits, + 0, + ) +} + +fn load_static_directory( + root: &Path, + directory: &Path, + assets: &mut BTreeMap, + total_files: &mut usize, + total_bytes: &mut usize, + limits: &WebLimitsConfig, + depth: usize, +) -> Result<()> { + let entries = fs::read_dir(directory).map_err(|error| { + ProxyError::Config(format!( + "failed to read WEB static directory `{}`: {error}", + directory.display() + )) + })?; + for entry in entries { + let entry = entry.map_err(|error| { + ProxyError::Config(format!("failed to read WEB static entry: {error}")) + })?; + if *total_files >= limits.max_static_files { + return Err(ProxyError::Config( + "WEB static entries exceed process-wide web.limits.max_static_files".to_string(), + )); + } + *total_files += 1; + let path = entry.path(); + let file_type = entry.file_type().map_err(|error| { + ProxyError::Config(format!( + "failed to inspect WEB static entry `{}`: {error}", + path.display() + )) + })?; + if file_type.is_symlink() { + return Err(ProxyError::Config(format!( + "WEB static entry `{}` must not be a symlink", + path.display() + ))); + } + if file_type.is_dir() { + if depth >= MAX_WEB_STATIC_DEPTH { + return Err(ProxyError::Config(format!( + "WEB static directory `{}` exceeds the maximum nesting depth", + path.display() + ))); + } + load_static_directory( + root, + &path, + assets, + total_files, + total_bytes, + limits, + depth + 1, + )?; + continue; + } + if !file_type.is_file() { + return Err(ProxyError::Config(format!( + "WEB static entry `{}` must be a regular file", + path.display() + ))); + } + let file = fs::File::open(&path).map_err(|error| { + ProxyError::Config(format!( + "failed to open WEB static file `{}`: {error}", + path.display() + )) + })?; + let metadata = file.metadata().map_err(|error| { + ProxyError::Config(format!( + "failed to inspect WEB static file `{}`: {error}", + path.display() + )) + })?; + if !metadata.is_file() { + return Err(ProxyError::Config(format!( + "WEB static entry `{}` changed before it was opened", + path.display() + ))); + } + let relative = path.strip_prefix(root).map_err(|_| { + ProxyError::Config("WEB static path escaped its configured root".to_string()) + })?; + load_static_file( + file, + &metadata, + relative, + &path, + assets, + total_bytes, + limits, + )?; + } + Ok(()) +} diff --git a/src/config/load/runtime_web/tests.rs b/src/config/load/runtime_web/tests.rs new file mode 100644 index 0000000..5d3f5b2 --- /dev/null +++ b/src/config/load/runtime_web/tests.rs @@ -0,0 +1,92 @@ +use base64::Engine as _; + +use super::*; + +#[test] +fn capability_matches_reference_vectors() { + let secret = hex::decode("000102030405060708090a0b0c0d0e0f").unwrap(); + let mut dd_secret = vec![0xdd]; + dd_secret.extend_from_slice(&secret); + for (client_secret, base_path, expected) in [ + ( + secret.as_slice(), + b"".as_slice(), + "MHLEY5PmW1GWqJkSrlmJpvJUiLhBH_QKy6yKg8a0JPk", + ), + ( + dd_secret.as_slice(), + b"".as_slice(), + "IpJrt3e7sKtzPyoXy6w-Zj6GGEvsvclN66JzQEfPYLA", + ), + ( + secret.as_slice(), + b"dobry-cola-super-app".as_slice(), + "hHz99Xs93EN1j91G9gpNepXwGNNt5YdAFkEVk_LlqdQ", + ), + ( + dd_secret.as_slice(), + b"dobry-cola-super-app".as_slice(), + "TGUkZaevsavLbHvlNWipnRoYxgzZ51ioWvbxgGT3wHo", + ), + ] { + let capability = + derive_web_capability(client_secret, b"proxy.example.com", base_path).unwrap(); + assert_eq!( + base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(capability), + expected + ); + } +} + +#[test] +fn capability_binds_the_exact_host_and_base_path_identity() { + let secret = hex::decode("000102030405060708090a0b0c0d0e0f").unwrap(); + let root = derive_web_capability(&secret, b"proxy.example.com", b"").unwrap(); + let mixed = derive_web_capability(&secret, b"proxy.example.com", b"MixedCase/path").unwrap(); + let lower = derive_web_capability(&secret, b"proxy.example.com", b"mixedcase/path").unwrap(); + let other_path = + derive_web_capability(&secret, b"proxy.example.com", b"MixedCase/other").unwrap(); + let other_host = + derive_web_capability(&secret, b"other.example.com", b"MixedCase/path").unwrap(); + + let identities = [root, mixed, lower, other_path, other_host] + .into_iter() + .collect::>(); + assert_eq!(identities.len(), 5); +} + +#[cfg(unix)] +#[test] +fn static_snapshot_remains_anchored_after_root_path_replacement() { + use std::os::unix::fs::symlink; + + let temp = tempfile::tempdir().unwrap(); + let root = temp.path().join("site"); + let detached = temp.path().join("detached"); + let replacement = temp.path().join("replacement"); + fs::create_dir(&root).unwrap(); + fs::write(root.join("index.html"), b"original").unwrap(); + fs::create_dir(&replacement).unwrap(); + fs::write(replacement.join("index.html"), b"replacement").unwrap(); + + let directory = open_static_root(&root).unwrap(); + fs::rename(&root, &detached).unwrap(); + symlink(&replacement, &root).unwrap(); + + let mut assets = BTreeMap::new(); + let mut total_files = 0; + let mut total_bytes = 0; + load_static_directory( + directory, + Path::new(""), + &root, + &mut assets, + &mut total_files, + &mut total_bytes, + &WebLimitsConfig::default(), + 0, + ) + .unwrap(); + + assert_eq!(assets["/index.html"].body.as_ref(), b"original"); +} diff --git a/src/config/load/strict_keys.rs b/src/config/load/strict_keys.rs index 69dc383..3040518 100644 --- a/src/config/load/strict_keys.rs +++ b/src/config/load/strict_keys.rs @@ -326,6 +326,7 @@ const WEB_LIMITS_CONFIG_KEYS: &[&str] = &[ const WEB_DEBUG_CONFIG_KEYS: &[&str] = &[ "enabled", + "sideband", "capture_lifecycle", "capture_headers", "capture_timings", @@ -364,7 +365,7 @@ const WEB_TIMEOUTS_CONFIG_KEYS: &[&str] = &[ "decoy_header_secs", ]; -const WEB_VHOST_CONFIG_KEYS: &[&str] = &["host", "public_addr", "decoy", "profiles"]; +const WEB_VHOST_CONFIG_KEYS: &[&str] = &["host", "base_path", "public_addr", "decoy", "profiles"]; const WEB_DECOY_CONFIG_KEYS: &[&str] = &["mode", "upstream", "directory", "index"]; const WEB_PROFILE_CONFIG_KEYS: &[&str] = &[ "user", diff --git a/src/config/load/validate_runtime.rs b/src/config/load/validate_runtime.rs index ae05757..6cae0ed 100644 --- a/src/config/load/validate_runtime.rs +++ b/src/config/load/validate_runtime.rs @@ -196,6 +196,13 @@ pub(super) fn validate(config: &mut ProxyConfig) -> Result<()> { "access.user_rate_limits.{user} must set at least one non-zero direction" ))); } + for (direction, value) in [("up_bps", limit.up_bps), ("down_bps", limit.down_bps)] { + if value > MAX_RATE_LIMIT_BPS { + return Err(ProxyError::Config(format!( + "access.user_rate_limits.{user}.{direction} must be within [0, {MAX_RATE_LIMIT_BPS}]" + ))); + } + } } for (cidr, limit) in &config.access.cidr_rate_limits { @@ -204,6 +211,13 @@ pub(super) fn validate(config: &mut ProxyConfig) -> Result<()> { "access.cidr_rate_limits.{cidr} must set at least one non-zero direction" ))); } + for (direction, value) in [("up_bps", limit.up_bps), ("down_bps", limit.down_bps)] { + if value > MAX_RATE_LIMIT_BPS { + return Err(ProxyError::Config(format!( + "access.cidr_rate_limits.{cidr}.{direction} must be within [0, {MAX_RATE_LIMIT_BPS}]" + ))); + } + } } let mut cidr_auto_templates = HashSet::new(); for cidr in config.access.cidr_rate_limits.keys() { diff --git a/src/config/load/validate_web/memory.rs b/src/config/load/validate_web/memory.rs index 7dca7a6..3a8d161 100644 --- a/src/config/load/validate_web/memory.rs +++ b/src/config/load/validate_web/memory.rs @@ -6,7 +6,8 @@ const WEB_DEBUG_GROUP_SCRATCH_BYTES: usize = 4 * 1024 * 1024; const WEB_CARRIER_LEARNING_ENTRY_BYTES: usize = 512; const WEB_LANE_STATE_BYTES: usize = 512; const WEB_OVERLOAD_CONNECTION_BYTES: usize = 4 * 1024; -const WEB_CAPABILITY_INDEX_ENTRY_BYTES: usize = 32; +// Each profile capability is stored in its vhost and in the global containment table. +const WEB_CAPABILITY_INDEX_ENTRY_BYTES: usize = 64; /// Validates process-wide body, header, queue, static, and debug reservations. pub(super) fn validate(limits: &WebLimitsConfig) -> Result<()> { diff --git a/src/config/load/validate_web/vhosts.rs b/src/config/load/validate_web/vhosts.rs index ec82852..571d4fa 100644 --- a/src/config/load/validate_web/vhosts.rs +++ b/src/config/load/validate_web/vhosts.rs @@ -9,6 +9,10 @@ pub(super) fn validate_vhosts(config: &mut ProxyConfig) -> Result<()> { let mut profile_count = 0usize; for (vhost_idx, vhost) in config.web.vhosts.iter_mut().enumerate() { vhost.host = normalize_web_host(&vhost.host, &format!("web.vhosts[{vhost_idx}].host"))?; + validate_web_base_path( + &vhost.base_path, + &format!("web.vhosts[{vhost_idx}].base_path"), + )?; if !hosts.insert(vhost.host.clone()) { return config_error(&format!("duplicate WEB vhost host `{}`", vhost.host)); } @@ -75,6 +79,25 @@ pub(super) fn validate_vhosts(config: &mut ProxyConfig) -> Result<()> { Ok(()) } +fn validate_web_base_path(value: &str, field: &str) -> Result<()> { + let valid = value.len() <= 128 + && !value.starts_with('/') + && !value.ends_with('/') + && value.split('/').all(|segment| { + let mut bytes = segment.bytes(); + bytes + .next() + .is_some_and(|byte| byte.is_ascii_alphanumeric()) + && bytes.all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_')) + }); + if value.is_empty() || valid { + return Ok(()); + } + config_error(&format!( + "{field} must be empty or contain at most 128 ASCII bytes in slash-separated [A-Za-z0-9][A-Za-z0-9_-]* segments" + )) +} + pub(super) fn normalize_web_host(value: &str, field: &str) -> Result { let input = value.trim(); if input.is_empty() diff --git a/src/config/tests/load_basic_tests.rs b/src/config/tests/load_basic_tests.rs index e24285f..82c7da4 100644 --- a/src/config/tests/load_basic_tests.rs +++ b/src/config/tests/load_basic_tests.rs @@ -46,6 +46,8 @@ mod legacy_policy_tests; mod me_route_tests; #[path = "load_basic_tests/me_startup_tests.rs"] mod me_startup_tests; +#[path = "load_basic_tests/source_security_tests.rs"] +mod source_security_tests; #[path = "load_basic_tests/synlimit_mss_tests.rs"] mod synlimit_mss_tests; #[path = "load_basic_tests/tls_fetch_tests.rs"] diff --git a/src/config/tests/load_basic_tests/defaults_access_tests.rs b/src/config/tests/load_basic_tests/defaults_access_tests.rs index 7d3d8de..4f6908e 100644 --- a/src/config/tests/load_basic_tests/defaults_access_tests.rs +++ b/src/config/tests/load_basic_tests/defaults_access_tests.rs @@ -288,6 +288,68 @@ fn cidr_rate_limits_reject_duplicate_normalized_auto_templates() { assert!(error.contains("duplicates normalized auto-template *6/128")); } +#[test] +fn rate_limits_accept_the_packed_counter_maximum() { + let cfg = load_config_from_temp_toml( + r#" + [censorship] + tls_domain = "example.com" + + [access.users] + user = "00000000000000000000000000000000" + + [access.user_rate_limits] + user = { up_bps = 100000000000, down_bps = 0 } + + [access.cidr_rate_limits] + "203.0.113.0/24" = { up_bps = 0, down_bps = 100000000000 } + "#, + ); + + assert_eq!(cfg.access.user_rate_limits["user"].up_bps, 100_000_000_000); + assert_eq!( + cfg.access.cidr_rate_limits[&CidrRateLimitKey::Network("203.0.113.0/24".parse().unwrap())] + .down_bps, + 100_000_000_000 + ); +} + +#[test] +fn user_rate_limits_reject_values_above_the_packed_counter_maximum() { + let error = load_config_error_from_temp_toml( + r#" + [censorship] + tls_domain = "example.com" + + [access.users] + user = "00000000000000000000000000000000" + + [access.user_rate_limits] + user = { up_bps = 100000000001, down_bps = 0 } + "#, + ); + + assert!(error.contains("access.user_rate_limits.user.up_bps must be within")); +} + +#[test] +fn cidr_rate_limits_reject_values_above_the_packed_counter_maximum() { + let error = load_config_error_from_temp_toml( + r#" + [censorship] + tls_domain = "example.com" + + [access.users] + user = "00000000000000000000000000000000" + + [access.cidr_rate_limits] + "203.0.113.0/24" = { up_bps = 0, down_bps = 100000000001 } + "#, + ); + + assert!(error.contains("access.cidr_rate_limits.203.0.113.0/24.down_bps must be within")); +} + #[test] fn file_logging_requires_path() { let error = load_config_error_from_temp_toml( diff --git a/src/config/tests/load_basic_tests/source_security_tests.rs b/src/config/tests/load_basic_tests/source_security_tests.rs new file mode 100644 index 0000000..86706eb --- /dev/null +++ b/src/config/tests/load_basic_tests/source_security_tests.rs @@ -0,0 +1,42 @@ +#[cfg(unix)] +use std::os::unix::fs::symlink; + +use super::*; + +#[cfg(unix)] +#[test] +fn config_loader_rejects_final_and_intermediate_symlinks() { + let directory = tempfile::tempdir().unwrap(); + let real_directory = directory.path().join("real"); + let linked_directory = directory.path().join("linked"); + std::fs::create_dir(&real_directory).unwrap(); + let real_config = real_directory.join("config.toml"); + let final_link = directory.path().join("config.toml"); + std::fs::write(&real_config, "[general]\n").unwrap(); + symlink(&real_config, &final_link).unwrap(); + symlink(&real_directory, &linked_directory).unwrap(); + + assert!(ProxyConfig::load(&final_link).is_err()); + assert!(ProxyConfig::load(linked_directory.join("config.toml")).is_err()); +} + +#[test] +fn config_loader_rejects_oversized_source() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("config.toml"); + std::fs::write(&path, vec![b' '; 8 * 1024 * 1024 + 1]).unwrap(); + + let error = ProxyConfig::load(&path).unwrap_err().to_string(); + + assert!(error.contains("size limit") || error.contains("exceeds")); +} + +#[cfg(unix)] +#[test] +fn config_loader_rejects_fifo_without_blocking() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("config.toml"); + nix::unistd::mkfifo(&path, nix::sys::stat::Mode::S_IRUSR).unwrap(); + + assert!(ProxyConfig::load(&path).is_err()); +} diff --git a/src/config/tests/load_basic_tests/web_tests.rs b/src/config/tests/load_basic_tests/web_tests.rs index b18d004..b0a071d 100644 --- a/src/config/tests/load_basic_tests/web_tests.rs +++ b/src/config/tests/load_basic_tests/web_tests.rs @@ -1,5 +1,8 @@ use super::*; +#[path = "web_tests/base_path_tests.rs"] +mod base_path_tests; + const WEB_CONFIG: &str = r#" [access.users] alice = "000102030405060708090a0b0c0d0e0f" @@ -40,9 +43,11 @@ fn web_config_builds_canonical_runtime_snapshot() { .vhosts .get("proxy.example.com") .expect("canonical WEB vhost"); + assert_eq!(vhost.base, "/"); assert_eq!(vhost.profiles.len(), 1); assert_eq!(vhost.capabilities.len(), vhost.profiles.len()); assert_eq!(vhost.capabilities[0], vhost.profiles[0].capability); + assert_eq!(runtime.capabilities.as_ref(), vhost.capabilities.as_ref()); assert_eq!(vhost.decoy_fasttrack_mode, WebDecoyFastTrackMode::Off); assert_eq!(vhost.profiles[0].user, "alice"); assert_eq!(vhost.profiles[0].secret_mode, WebSecretMode::Dd); @@ -59,6 +64,66 @@ fn web_config_builds_canonical_runtime_snapshot() { ); } +#[test] +fn web_base_path_is_canonical_and_precomputed() { + let maximum = "a".repeat(128); + for base_path in ["a", "a/b/c9_x-y", "Dobry-Cola/super_app", maximum.as_str()] { + let configured = WEB_CONFIG.replace( + "host = \"Proxy.Example.COM\"", + &format!("host = \"Proxy.Example.COM\"\nbase_path = \"{base_path}\""), + ); + let config = load_config_from_temp_toml(&configured); + assert_eq!(config.web.vhosts[0].base_path, base_path); + assert_eq!( + config.web.runtime.as_ref().unwrap().vhosts["proxy.example.com"].base, + format!("/{base_path}/") + ); + } + + let strict = format!( + "[general]\nconfig_strict = true\n{}", + WEB_CONFIG.replace( + "host = \"Proxy.Example.COM\"", + "host = \"Proxy.Example.COM\"\nbase_path = \"relay\"", + ) + ); + assert_eq!( + load_config_from_temp_toml(&strict).web.vhosts[0].base_path, + "relay" + ); +} + +#[test] +fn web_base_path_rejects_noncanonical_forms() { + let oversized = "a".repeat(129); + for base_path in [ + "/relay", + "relay/", + "relay//nested", + "-relay", + "_relay", + "a/-lead", + "a/_lead", + "dot.ted", + "..", + "/", + "relay/.hidden", + "relay/%2fhidden", + "relay path", + "relay/тест", + oversized.as_str(), + ] { + let invalid = WEB_CONFIG.replace( + "host = \"Proxy.Example.COM\"", + &format!("host = \"Proxy.Example.COM\"\nbase_path = \"{base_path}\""), + ); + assert!( + load_config_error_from_temp_toml(&invalid).contains("web.vhosts[0].base_path"), + "base path {base_path:?} was accepted" + ); + } +} + #[test] fn web_decoy_fasttrack_mode_is_typed_and_defaults_off() { let defaults = ProxyConfig::default(); @@ -301,17 +366,31 @@ fn web_carrier_learning_capacity_must_remain_nonzero() { #[test] fn web_debug_table_uses_debug_name_and_bounded_defaults() { + assert!(!crate::config::WebDebugConfig::default().sideband); + let mut ineffective = crate::config::WebDebugConfig { + enabled: true, + sideband: true, + ..Default::default() + }; + ineffective.capture_lifecycle = false; + assert!(!ineffective.bridge_diagnostics_enabled()); + let configured = WEB_CONFIG.replace( "[[web.vhosts]]", - "[web.debug]\nenabled = true\nbody_capture = \"prefix\"\nbody_prefix_bytes = 2048\ndefault_window_secs = 180\nmax_window_secs = 900\n\n[[web.vhosts]]", + "[web.debug]\nenabled = true\nsideband = true\nbody_capture = \"prefix\"\nbody_prefix_bytes = 2048\ndefault_window_secs = 180\nmax_window_secs = 900\n\n[[web.vhosts]]", ); let config = load_config_from_temp_toml(&configured); assert!(config.web.debug.enabled); + assert!(config.web.debug.sideband); + assert!(config.web.debug.bridge_diagnostics_enabled()); assert_eq!(config.web.debug.body_capture, WebDebugBodyCapture::Prefix); assert_eq!(config.web.debug.body_prefix_bytes, 2048); assert_eq!(config.web.debug.default_window_secs, 180); assert_eq!(config.web.debug.max_window_secs, 900); + let strict = format!("[general]\nconfig_strict = true\n{configured}"); + assert!(load_config_from_temp_toml(&strict).web.debug.sideband); + let old_name = format!( "[general]\nconfig_strict = true\n{}", WEB_CONFIG.replace( @@ -321,6 +400,16 @@ fn web_debug_table_uses_debug_name_and_bounded_defaults() { ); let error = load_config_error_from_temp_toml(&old_name); assert!(error.contains("web.trace")); + + let old_parameter = format!( + "[general]\nconfig_strict = true\n{}", + WEB_CONFIG.replace( + "[[web.vhosts]]", + "[web.debug]\nbridge_diagnostics = true\n\n[[web.vhosts]]", + ) + ); + let error = load_config_error_from_temp_toml(&old_parameter); + assert!(error.contains("web.debug.bridge_diagnostics")); } #[test] diff --git a/src/config/tests/load_basic_tests/web_tests/base_path_tests.rs b/src/config/tests/load_basic_tests/web_tests/base_path_tests.rs new file mode 100644 index 0000000..37417cc --- /dev/null +++ b/src/config/tests/load_basic_tests/web_tests/base_path_tests.rs @@ -0,0 +1,33 @@ +use super::*; + +#[test] +fn web_runtime_collects_every_vhost_capability() { + let configured = format!( + "{WEB_CONFIG}\n{}", + r#" +[[web.vhosts]] +host = "Other.Example.COM" +base_path = "other/path" +public_addr = "203.0.113.11:443" + +[web.vhosts.decoy] +mode = "http_upstream" +upstream = "http://127.0.0.1:18082" + +[[web.vhosts.profiles]] +user = "alice" +secret_mode = "dd" +"# + ); + let config = load_config_from_temp_toml(&configured); + let runtime = config.web.runtime.as_ref().unwrap(); + let first = &runtime.vhosts["proxy.example.com"]; + let second = &runtime.vhosts["other.example.com"]; + + assert_eq!(runtime.capabilities.len(), 2); + assert_eq!(first.capabilities[0], first.profiles[0].capability); + assert_eq!(second.capabilities[0], second.profiles[0].capability); + assert_ne!(first.capabilities[0], second.capabilities[0]); + assert!(runtime.capabilities.contains(&first.capabilities[0])); + assert!(runtime.capabilities.contains(&second.capabilities[0])); +} diff --git a/src/config/types.rs b/src/config/types.rs index dfc0026..e769ef4 100644 --- a/src/config/types.rs +++ b/src/config/types.rs @@ -31,7 +31,7 @@ mod web_debug; pub use access::{AccessConfig, CidrRateLimitKey, RateLimitBps}; #[allow(unused_imports)] -pub(crate) use access::{CidrAutoTemplate, CidrAutoTemplateFamily}; +pub(crate) use access::{CidrAutoTemplate, CidrAutoTemplateFamily, MAX_RATE_LIMIT_BPS}; pub use api::{ApiConfig, ApiGrayAction}; pub use censorship::{ AntiCensorshipConfig, ExclusiveMaskTarget, TlsFetchConfig, TlsFetchProfile, UnknownSniAction, diff --git a/src/config/types/access.rs b/src/config/types/access.rs index 3f68d60..ac35dff 100644 --- a/src/config/types/access.rs +++ b/src/config/types/access.rs @@ -1,5 +1,8 @@ use super::*; +/// Highest rate that fits one packed 20 ms shaping epoch. +pub(crate) const MAX_RATE_LIMIT_BPS: u64 = 100_000_000_000; + #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct AccessConfig { #[serde(default = "default_access_users")] @@ -260,10 +263,10 @@ fn parse_cidr_auto_prefix( /// Transport rate limit in bits-per-second. #[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] pub struct RateLimitBps { - /// Upload direction limit in bits-per-second; `0` means unlimited. + /// Upload limit in bits-per-second within `0..=100_000_000_000`; `0` means unlimited. #[serde(default)] pub up_bps: u64, - /// Download direction limit in bits-per-second; `0` means unlimited. + /// Download limit in bits-per-second within `0..=100_000_000_000`; `0` means unlimited. #[serde(default)] pub down_bps: u64, } diff --git a/src/config/types/web.rs b/src/config/types/web.rs index 70fb062..3f88ce5 100644 --- a/src/config/types/web.rs +++ b/src/config/types/web.rs @@ -71,6 +71,9 @@ pub enum WebDecoyConfig { pub struct WebVhostConfig { /// Canonical lowercase ACE hostname used by Telegram Desktop. pub host: String, + /// Optional canonical WEB endpoint prefix without surrounding slashes. + #[serde(default)] + pub base_path: String, /// Stable public destination tuple used by inner relay routing and KDF metadata. pub public_addr: SocketAddr, /// Ordinary-site fallback for this hostname. diff --git a/src/config/types/web/runtime.rs b/src/config/types/web/runtime.rs index 0485ee2..58b79ff 100644 --- a/src/config/types/web/runtime.rs +++ b/src/config/types/web/runtime.rs @@ -7,6 +7,8 @@ pub(crate) struct WebRuntimeConfig { pub(crate) vhosts: BTreeMap>, /// Flat profile inventory used by startup link emission. pub(crate) profiles: Vec>, + /// Complete active capability table used to contain misplaced credentials. + pub(crate) capabilities: Box<[[u8; 32]]>, } /// Precomputed immutable virtual-host data. @@ -14,6 +16,8 @@ pub(crate) struct WebRuntimeConfig { pub(crate) struct WebRuntimeVhost { /// Canonical lowercase ACE hostname. pub(crate) host: String, + /// Exact slash-delimited endpoint base, including the trailing slash. + pub(crate) base: String, /// Restart-frozen decoy capability-scan policy. pub(crate) decoy_fasttrack_mode: WebDecoyFastTrackMode, /// Immutable ordinary-site fallback snapshot. @@ -35,6 +39,8 @@ pub(crate) struct WebRuntimeProfile { pub(crate) public_addr: SocketAddr, /// Exact access user authenticated by logical streams. pub(crate) user: String, + /// Stable credential identity used by process-wide admission fencing. + pub(crate) credential_id: [u8; 16], /// Client secret representation and inner protocol policy. pub(crate) secret_mode: WebSecretMode, /// Sole carrier or final fallback frozen into the issued bridge policy. diff --git a/src/config/types/web_debug.rs b/src/config/types/web_debug.rs index 4bc0df9..63cf7b2 100644 --- a/src/config/types/web_debug.rs +++ b/src/config/types/web_debug.rs @@ -23,6 +23,9 @@ pub struct WebDebugConfig { /// Enables process-owned WEB debug collection. #[serde(default)] pub enabled: bool, + /// Enables generated-bridge diagnostic reports over the HTTPS sideband. + #[serde(default)] + pub sideband: bool, /// Records typed bridge, session, stream, handshake, and relay events. #[serde(default = "default_true")] pub capture_lifecycle: bool, @@ -56,6 +59,7 @@ impl Default for WebDebugConfig { fn default() -> Self { Self { enabled: false, + sideband: false, capture_lifecycle: true, capture_headers: true, capture_timings: true, @@ -69,6 +73,13 @@ impl Default for WebDebugConfig { } } +impl WebDebugConfig { + /// Returns whether newly issued bridges may report diagnostic lifecycle events. + pub(crate) const fn bridge_diagnostics_enabled(&self) -> bool { + self.enabled && self.sideband && self.capture_lifecycle + } +} + fn default_true() -> bool { true } diff --git a/src/conntrack_control.rs b/src/conntrack_control.rs index 1cb5b7d..2f650d9 100644 --- a/src/conntrack_control.rs +++ b/src/conntrack_control.rs @@ -1,20 +1,23 @@ -use std::collections::BTreeSet; -use std::net::IpAddr; -use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; -use tokio::io::AsyncWriteExt; -use tokio::process::Command; use tokio::sync::{mpsc, watch}; use tokio_util::sync::CancellationToken; -use tracing::{debug, info, warn}; +use tracing::{info, warn}; -use crate::config::{ConntrackBackend, ConntrackMode, ProxyConfig}; +use crate::config::ProxyConfig; use crate::proxy::middle_relay::note_global_relay_pressure; use crate::proxy::shared_state::{ConntrackCloseEvent, ConntrackCloseReason, ProxySharedState}; use crate::stats::Stats; +// Privileged netfilter rule and conntrack helper execution. +mod firewall; + +pub(crate) use firewall::FirewallAuthority; +use firewall::{ + DeleteOutcome, delete_conntrack_entry, effective_conntrack_enabled, probe_runtime_support, +}; + const CONNTRACK_EVENT_QUEUE_CAPACITY: usize = 32_768; const PRESSURE_RELEASE_TICKS: u8 = 3; const PRESSURE_SAMPLE_INTERVAL: Duration = Duration::from_secs(1); @@ -112,7 +115,6 @@ async fn run_conntrack_controller_worker( runtime_support, false, ); - reconcile_rules(&cfg, runtime_support, stats.as_ref()).await; loop { tokio::select! { @@ -126,7 +128,6 @@ async fn run_conntrack_controller_worker( effective_enabled = effective_conntrack_enabled(&cfg, runtime_support); delete_budget_tokens = cfg.server.conntrack_control.delete_budget_per_sec; apply_runtime_state(stats.as_ref(), shared.as_ref(), &cfg, runtime_support, pressure_state.active); - reconcile_rules(&cfg, runtime_support, stats.as_ref()).await; } event = close_rx.recv() => { let Some(event) = event else { @@ -310,383 +311,6 @@ fn update_pressure_state( state.low_streak = 0; } -async fn reconcile_rules( - cfg: &ProxyConfig, - runtime_support: ConntrackRuntimeSupport, - stats: &Stats, -) { - if !cfg.server.conntrack_control.inline_conntrack_control { - clear_notrack_rules_all_backends().await; - stats.set_conntrack_rule_apply_ok(true); - return; - } - - if !effective_conntrack_enabled(cfg, runtime_support) { - clear_notrack_rules_all_backends().await; - stats.set_conntrack_rule_apply_ok(false); - return; - } - - let backend = runtime_support - .netfilter_backend - .expect("netfilter backend must be available for effective conntrack control"); - - let apply_result = match backend { - NetfilterBackend::Nftables => apply_nft_rules(cfg).await, - NetfilterBackend::Iptables => apply_iptables_rules(cfg).await, - }; - - if let Err(error) = apply_result { - warn!(error = %error, "Failed to reconcile conntrack/notrack rules"); - stats.set_conntrack_rule_apply_ok(false); - } else { - stats.set_conntrack_rule_apply_ok(true); - } -} - -fn probe_runtime_support(configured_backend: ConntrackBackend) -> ConntrackRuntimeSupport { - ConntrackRuntimeSupport { - netfilter_backend: pick_backend(configured_backend), - has_cap_net_admin: has_cap_net_admin(), - has_conntrack_binary: command_exists("conntrack"), - } -} - -fn effective_conntrack_enabled( - cfg: &ProxyConfig, - runtime_support: ConntrackRuntimeSupport, -) -> bool { - cfg.server.conntrack_control.inline_conntrack_control - && runtime_support.has_cap_net_admin - && runtime_support.netfilter_backend.is_some() - && runtime_support.has_conntrack_binary -} - -fn pick_backend(configured: ConntrackBackend) -> Option { - match configured { - ConntrackBackend::Auto => { - if command_exists("nft") { - Some(NetfilterBackend::Nftables) - } else if command_exists("iptables") { - Some(NetfilterBackend::Iptables) - } else { - None - } - } - ConntrackBackend::Nftables => command_exists("nft").then_some(NetfilterBackend::Nftables), - ConntrackBackend::Iptables => { - command_exists("iptables").then_some(NetfilterBackend::Iptables) - } - } -} - -fn command_exists(binary: &str) -> bool { - let Some(path_var) = std::env::var_os("PATH") else { - return false; - }; - std::env::split_paths(&path_var).any(|dir| { - let candidate: PathBuf = dir.join(binary); - candidate.exists() && candidate.is_file() - }) -} - -fn listener_port_set(cfg: &ProxyConfig) -> Vec { - let mut ports: BTreeSet = BTreeSet::new(); - if cfg.server.listeners.is_empty() { - ports.insert(cfg.server.port); - } else { - for listener in &cfg.server.listeners { - ports.insert(listener.port.unwrap_or(cfg.server.port)); - } - } - ports.into_iter().collect() -} - -fn notrack_targets(cfg: &ProxyConfig) -> (Vec<(Option, u16)>, Vec<(Option, u16)>) { - let mode = cfg.server.conntrack_control.mode; - let mut v4_targets: BTreeSet<(Option, u16)> = BTreeSet::new(); - let mut v6_targets: BTreeSet<(Option, u16)> = BTreeSet::new(); - - match mode { - ConntrackMode::Tracked => {} - ConntrackMode::Notrack => { - if cfg.server.listeners.is_empty() { - let port = cfg.server.port; - if let Some(ipv4) = cfg - .server - .listen_addr_ipv4 - .as_ref() - .and_then(|s| s.parse::().ok()) - { - if ipv4.is_unspecified() { - v4_targets.insert((None, port)); - } else { - v4_targets.insert((Some(ipv4), port)); - } - } - if let Some(ipv6) = cfg - .server - .listen_addr_ipv6 - .as_ref() - .and_then(|s| s.parse::().ok()) - { - if ipv6.is_unspecified() { - v6_targets.insert((None, port)); - } else { - v6_targets.insert((Some(ipv6), port)); - } - } - } else { - for listener in &cfg.server.listeners { - let port = listener.port.unwrap_or(cfg.server.port); - if listener.ip.is_ipv4() { - if listener.ip.is_unspecified() { - v4_targets.insert((None, port)); - } else { - v4_targets.insert((Some(listener.ip), port)); - } - } else if listener.ip.is_unspecified() { - v6_targets.insert((None, port)); - } else { - v6_targets.insert((Some(listener.ip), port)); - } - } - } - } - ConntrackMode::Hybrid => { - let ports = listener_port_set(cfg); - for ip in &cfg.server.conntrack_control.hybrid_listener_ips { - if ip.is_ipv4() { - for port in &ports { - v4_targets.insert((Some(*ip), *port)); - } - } else { - for port in &ports { - v6_targets.insert((Some(*ip), *port)); - } - } - } - } - } - - ( - v4_targets.into_iter().collect(), - v6_targets.into_iter().collect(), - ) -} - -async fn apply_nft_rules(cfg: &ProxyConfig) -> Result<(), String> { - let _ = run_command( - "nft", - &["delete", "table", "inet", "telemt_conntrack"], - None, - ) - .await; - if matches!(cfg.server.conntrack_control.mode, ConntrackMode::Tracked) { - return Ok(()); - } - - let (v4_targets, v6_targets) = notrack_targets(cfg); - let mut rules = Vec::new(); - for (ip, port) in v4_targets { - let rule = if let Some(ip) = ip { - format!("tcp dport {} ip daddr {} notrack", port, ip) - } else { - format!("tcp dport {} notrack", port) - }; - rules.push(rule); - } - for (ip, port) in v6_targets { - let rule = if let Some(ip) = ip { - format!("tcp dport {} ip6 daddr {} notrack", port, ip) - } else { - format!("tcp dport {} notrack", port) - }; - rules.push(rule); - } - - let rule_blob = if rules.is_empty() { - String::new() - } else { - format!(" {}\n", rules.join("\n ")) - }; - let script = format!( - "table inet telemt_conntrack {{\n chain preraw {{\n type filter hook prerouting priority raw; policy accept;\n{rule_blob} }}\n}}\n" - ); - run_command("nft", &["-f", "-"], Some(script)).await -} - -async fn apply_iptables_rules(cfg: &ProxyConfig) -> Result<(), String> { - apply_iptables_rules_for_binary("iptables", cfg, true).await?; - apply_iptables_rules_for_binary("ip6tables", cfg, false).await?; - Ok(()) -} - -async fn apply_iptables_rules_for_binary( - binary: &str, - cfg: &ProxyConfig, - ipv4: bool, -) -> Result<(), String> { - if !command_exists(binary) { - return Ok(()); - } - let chain = "TELEMT_NOTRACK"; - let _ = run_command( - binary, - &["-t", "raw", "-D", "PREROUTING", "-j", chain], - None, - ) - .await; - let _ = run_command(binary, &["-t", "raw", "-F", chain], None).await; - let _ = run_command(binary, &["-t", "raw", "-X", chain], None).await; - - if matches!(cfg.server.conntrack_control.mode, ConntrackMode::Tracked) { - return Ok(()); - } - - run_command(binary, &["-t", "raw", "-N", chain], None).await?; - run_command(binary, &["-t", "raw", "-F", chain], None).await?; - if run_command( - binary, - &["-t", "raw", "-C", "PREROUTING", "-j", chain], - None, - ) - .await - .is_err() - { - run_command( - binary, - &["-t", "raw", "-I", "PREROUTING", "1", "-j", chain], - None, - ) - .await?; - } - - let (v4_targets, v6_targets) = notrack_targets(cfg); - let selected = if ipv4 { v4_targets } else { v6_targets }; - for (ip, port) in selected { - let mut args = vec![ - "-t".to_string(), - "raw".to_string(), - "-A".to_string(), - chain.to_string(), - "-p".to_string(), - "tcp".to_string(), - "--dport".to_string(), - port.to_string(), - ]; - if let Some(ip) = ip { - args.push("-d".to_string()); - args.push(ip.to_string()); - } - args.push("-j".to_string()); - args.push("CT".to_string()); - args.push("--notrack".to_string()); - let arg_refs: Vec<&str> = args.iter().map(String::as_str).collect(); - run_command(binary, &arg_refs, None).await?; - } - Ok(()) -} - -async fn clear_notrack_rules_all_backends() { - let _ = run_command( - "nft", - &["delete", "table", "inet", "telemt_conntrack"], - None, - ) - .await; - let _ = run_command( - "iptables", - &["-t", "raw", "-D", "PREROUTING", "-j", "TELEMT_NOTRACK"], - None, - ) - .await; - let _ = run_command("iptables", &["-t", "raw", "-F", "TELEMT_NOTRACK"], None).await; - let _ = run_command("iptables", &["-t", "raw", "-X", "TELEMT_NOTRACK"], None).await; - let _ = run_command( - "ip6tables", - &["-t", "raw", "-D", "PREROUTING", "-j", "TELEMT_NOTRACK"], - None, - ) - .await; - let _ = run_command("ip6tables", &["-t", "raw", "-F", "TELEMT_NOTRACK"], None).await; - let _ = run_command("ip6tables", &["-t", "raw", "-X", "TELEMT_NOTRACK"], None).await; -} - -enum DeleteOutcome { - Deleted, - NotFound, - Error, -} - -async fn delete_conntrack_entry(event: ConntrackCloseEvent) -> DeleteOutcome { - if !command_exists("conntrack") { - return DeleteOutcome::Error; - } - let args = vec![ - "-D".to_string(), - "-p".to_string(), - "tcp".to_string(), - "-s".to_string(), - event.src.ip().to_string(), - "--sport".to_string(), - event.src.port().to_string(), - "-d".to_string(), - event.dst.ip().to_string(), - "--dport".to_string(), - event.dst.port().to_string(), - ]; - let arg_refs: Vec<&str> = args.iter().map(String::as_str).collect(); - match run_command("conntrack", &arg_refs, None).await { - Ok(()) => DeleteOutcome::Deleted, - Err(error) => { - if error.contains("0 flow entries have been deleted") { - DeleteOutcome::NotFound - } else { - debug!(error = %error, "conntrack delete failed"); - DeleteOutcome::Error - } - } - } -} - -async fn run_command(binary: &str, args: &[&str], stdin: Option) -> Result<(), String> { - if !command_exists(binary) { - return Err(format!("{binary} is not available")); - } - let mut command = Command::new(binary); - command.args(args); - if stdin.is_some() { - command.stdin(std::process::Stdio::piped()); - } - command.stdout(std::process::Stdio::null()); - command.stderr(std::process::Stdio::piped()); - let mut child = command - .spawn() - .map_err(|e| format!("spawn {binary} failed: {e}"))?; - if let Some(blob) = stdin - && let Some(mut writer) = child.stdin.take() - { - writer - .write_all(blob.as_bytes()) - .await - .map_err(|e| format!("stdin write {binary} failed: {e}"))?; - } - let output = child - .wait_with_output() - .await - .map_err(|e| format!("wait {binary} failed: {e}"))?; - if output.status.success() { - return Ok(()); - } - let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string(); - Err(if stderr.is_empty() { - format!("{binary} exited with status {}", output.status) - } else { - stderr - }) -} - fn fd_usage_pct() -> Option { let soft_limit = nofile_soft_limit()?; if soft_limit == 0 { @@ -715,29 +339,6 @@ fn nofile_soft_limit() -> Option { } } -fn has_cap_net_admin() -> bool { - #[cfg(target_os = "linux")] - { - let Ok(status) = std::fs::read_to_string("/proc/self/status") else { - return false; - }; - for line in status.lines() { - if let Some(raw) = line.strip_prefix("CapEff:") { - let caps = raw.trim(); - if let Ok(bits) = u64::from_str_radix(caps, 16) { - const CAP_NET_ADMIN_BIT: u64 = 12; - return (bits & (1u64 << CAP_NET_ADMIN_BIT)) != 0; - } - } - } - false - } - #[cfg(not(target_os = "linux"))] - { - false - } -} - #[cfg(test)] mod tests { use super::*; diff --git a/src/conntrack_control/firewall.rs b/src/conntrack_control/firewall.rs new file mode 100644 index 0000000..f8a6a3a --- /dev/null +++ b/src/conntrack_control/firewall.rs @@ -0,0 +1,18 @@ +// Process-owned firewall reconciliation and privileged conntrack helpers. + +mod actor; +mod command; +mod iptables; +mod model; +mod nftables; +mod runtime; +mod transaction; + +pub(crate) use actor::FirewallAuthority; +pub(super) use runtime::{ + DeleteOutcome, delete_conntrack_entry, effective_conntrack_enabled, probe_runtime_support, +}; + +#[cfg(test)] +#[path = "firewall/tests.rs"] +mod tests; diff --git a/src/conntrack_control/firewall/actor.rs b/src/conntrack_control/firewall/actor.rs new file mode 100644 index 0000000..13d6d3f --- /dev/null +++ b/src/conntrack_control/firewall/actor.rs @@ -0,0 +1,371 @@ +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::Duration; + +use tokio::sync::{Notify, watch}; +use tokio_util::sync::CancellationToken; +use tracing::warn; + +use crate::config::ProxyConfig; +use crate::maestro::control_plane::ProcessControlPlane; +use crate::stats::Stats; + +use super::command::{CommandError, FirewallCommandRunner, SystemCommandRunner}; +use super::model::{AppliedPlan, AppliedState, DesiredPolicy, DesiredState}; +use super::transaction::{InterruptibleRunner, reconcile_once, recover_to_empty}; + +const SHUTDOWN_CLEANUP_TIMEOUT: Duration = Duration::from_secs(30); +const SHUTDOWN_WAIT_TIMEOUT: Duration = Duration::from_secs(35); +const INITIAL_RECONCILE_TIMEOUT: Duration = Duration::from_secs(65); +const RETRY_DELAYS: [Duration; 6] = [ + Duration::from_secs(1), + Duration::from_secs(2), + Duration::from_secs(4), + Duration::from_secs(8), + Duration::from_secs(16), + Duration::from_secs(30), +]; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(super) enum ReconcileOutcome { + Applied, + Failed, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(super) struct ReconcileStatus { + pub(super) generation: u64, + pub(super) outcome: ReconcileOutcome, +} + +/// Process-owned publisher and shutdown owner for conntrack firewall policy. +#[derive(Clone)] +pub(crate) struct FirewallAuthority { + desired_tx: watch::Sender>, + status_rx: watch::Receiver>, + terminal: CancellationToken, + closed: Arc, + completed_flag: Arc, + cleanup_succeeded: Arc, + completed: Arc, +} + +impl FirewallAuthority { + /// Starts the single process-owned firewall reconciler. + pub(crate) fn spawn(control_plane: &ProcessControlPlane) -> Result { + let (desired_tx, desired_rx) = watch::channel(None); + let (status_tx, status_rx) = watch::channel(None); + let terminal = CancellationToken::new(); + let closed = Arc::new(AtomicBool::new(false)); + let completed_flag = Arc::new(AtomicBool::new(false)); + let cleanup_succeeded = Arc::new(AtomicBool::new(false)); + let completed = Arc::new(Notify::new()); + let actor = FirewallReconciler::new( + SystemCommandRunner, + desired_rx, + status_tx, + terminal.clone(), + Arc::clone(&closed), + Arc::clone(&completed_flag), + Arc::clone(&cleanup_succeeded), + Arc::clone(&completed), + ); + control_plane + .spawn_cooperative(move |process_cancellation| async move { + actor.run(process_cancellation).await; + }) + .map_err(|_| { + "process control-plane admission closed before conntrack firewall startup" + .to_string() + })?; + Ok(Self { + desired_tx, + status_rx, + terminal, + closed, + completed_flag, + cleanup_succeeded, + completed, + }) + } + + /// Publishes policy only after its runtime generation becomes active. + pub(crate) fn publish( + &self, + generation: u64, + config: Arc, + stats: Arc, + ) -> bool { + if self.closed.load(Ordering::Acquire) { + stats.set_conntrack_rule_apply_ok(false); + return false; + } + stats.set_conntrack_rule_apply_ok(false); + self.desired_tx.send_replace(Some(DesiredState { + generation, + policy: DesiredPolicy::from_config(config.as_ref()), + stats, + })); + true + } + + /// Publishes startup policy and waits for its first bounded attempt. + pub(crate) async fn publish_initial( + &self, + generation: u64, + config: Arc, + stats: Arc, + ) -> bool { + let mut status_rx = self.status_rx.clone(); + if !self.publish(generation, config, stats) { + return false; + } + tokio::time::timeout(INITIAL_RECONCILE_TIMEOUT, async move { + loop { + if let Some(status) = *status_rx.borrow_and_update() + && status.generation == generation + { + return status.outcome == ReconcileOutcome::Applied; + } + if status_rx.changed().await.is_err() { + return false; + } + } + }) + .await + .unwrap_or(false) + } + + /// Stops policy admission and waits for bounded terminal cleanup. + pub(crate) async fn shutdown_and_clear(&self) -> bool { + let completed = self.completed.notified(); + tokio::pin!(completed); + completed.as_mut().enable(); + if !self.closed.swap(true, Ordering::AcqRel) { + if let Some(desired) = self.desired_tx.borrow().as_ref() { + desired.stats.set_conntrack_rule_apply_ok(false); + } + self.terminal.cancel(); + } + let finished = self.completed_flag.load(Ordering::Acquire) + || tokio::time::timeout(SHUTDOWN_WAIT_TIMEOUT, completed) + .await + .is_ok(); + finished && self.cleanup_succeeded.load(Ordering::Acquire) + } +} + +struct CompletionGuard { + closed: Arc, + completed_flag: Arc, + completed: Arc, +} + +impl Drop for CompletionGuard { + fn drop(&mut self) { + self.closed.store(true, Ordering::Release); + self.completed_flag.store(true, Ordering::Release); + self.completed.notify_waiters(); + } +} + +pub(super) struct FirewallReconciler { + runner: R, + desired_rx: watch::Receiver>, + status_tx: watch::Sender>, + terminal: CancellationToken, + completion: CompletionGuard, + cleanup_succeeded: Arc, + applied: AppliedState, + last_generation: u64, + last_policy: Option, + last_stats: Option>, +} + +impl FirewallReconciler +where + R: FirewallCommandRunner + 'static, +{ + pub(super) fn new( + runner: R, + desired_rx: watch::Receiver>, + status_tx: watch::Sender>, + terminal: CancellationToken, + closed: Arc, + completed_flag: Arc, + cleanup_succeeded: Arc, + completed: Arc, + ) -> Self { + Self { + runner, + desired_rx, + status_tx, + terminal, + completion: CompletionGuard { + closed, + completed_flag, + completed, + }, + cleanup_succeeded, + applied: AppliedState::Unknown, + last_generation: 0, + last_policy: None, + last_stats: None, + } + } + + pub(super) async fn run(mut self, process_cancellation: CancellationToken) { + let mut current = None; + let mut retry_index = 0usize; + 'run: loop { + if current.is_none() { + let changed = tokio::select! { + biased; + _ = self.terminal.cancelled() => false, + _ = process_cancellation.cancelled() => false, + changed = self.desired_rx.changed() => changed.is_ok(), + }; + if !changed { + break; + } + current = self.take_latest_desired(); + retry_index = 0; + if current.is_none() { + continue; + } + } + + let desired = current.as_ref().expect("desired state is present").clone(); + let interruptible = + InterruptibleRunner::new(&self.runner, &self.terminal, &process_cancellation); + let result = + reconcile_once(&interruptible, &interruptible, &mut self.applied, &desired).await; + match result { + Ok(()) => { + desired + .stats + .increment_conntrack_rule_reconcile_success_total(); + desired.stats.set_conntrack_rule_apply_ok(true); + self.status_tx.send_replace(Some(ReconcileStatus { + generation: desired.generation, + outcome: ReconcileOutcome::Applied, + })); + if self.terminal.is_cancelled() || process_cancellation.is_cancelled() { + desired.stats.set_conntrack_rule_apply_ok(false); + break; + } + current = None; + retry_index = 0; + } + Err(failure) if failure.cancelled => break, + Err(failure) => { + desired + .stats + .increment_conntrack_rule_reconcile_error_total(); + desired.stats.set_conntrack_rule_apply_ok(false); + if let Some(rollback_succeeded) = failure.rollback_succeeded { + if rollback_succeeded { + desired + .stats + .increment_conntrack_rule_rollback_success_total(); + } else { + desired + .stats + .increment_conntrack_rule_rollback_error_total(); + } + } + self.status_tx.send_replace(Some(ReconcileStatus { + generation: desired.generation, + outcome: ReconcileOutcome::Failed, + })); + warn!( + generation = desired.generation, + error = %failure.message, + "Failed to reconcile conntrack firewall policy" + ); + + let delay = RETRY_DELAYS[retry_index.min(RETRY_DELAYS.len() - 1)]; + retry_index = retry_index.saturating_add(1); + let retry_deadline = tokio::time::Instant::now() + delay; + loop { + tokio::select! { + biased; + _ = self.terminal.cancelled() => break 'run, + _ = process_cancellation.cancelled() => break 'run, + changed = self.desired_rx.changed() => { + if changed.is_err() { + break 'run; + } + if let Some(next) = self.take_latest_desired() { + current = Some(next); + retry_index = 0; + break; + } + } + _ = tokio::time::sleep_until(retry_deadline) => break, + } + } + } + } + + if self.terminal.is_cancelled() || process_cancellation.is_cancelled() { + break; + } + if self.desired_rx.has_changed().unwrap_or(false) { + if let Some(next) = self.take_latest_desired() { + current = Some(next); + retry_index = 0; + } + } + } + + if let Some(stats) = &self.last_stats { + stats.set_conntrack_rule_apply_ok(false); + } + if let Err(error) = + tokio::time::timeout(SHUTDOWN_CLEANUP_TIMEOUT, recover_to_empty(&self.runner)) + .await + .unwrap_or_else(|_| { + Err(CommandError::failed("firewall shutdown cleanup timed out")) + }) + { + warn!(error = %error, "Failed to clear conntrack firewall policy during shutdown"); + } else { + self.applied = AppliedState::Known(AppliedPlan::Empty); + self.cleanup_succeeded.store(true, Ordering::Release); + } + let _completion = &self.completion; + } + + pub(super) fn take_latest_desired(&mut self) -> Option { + let next = self.desired_rx.borrow_and_update().clone()?; + if next.generation < self.last_generation { + warn!( + generation = next.generation, + active_generation = self.last_generation, + "Ignored stale conntrack firewall policy publication" + ); + return None; + } + if next.generation == self.last_generation { + if self.last_policy.as_ref() != Some(&next.policy) { + warn!( + generation = next.generation, + "Ignored conflicting conntrack firewall policy for active generation" + ); + return None; + } + self.last_stats = Some(next.stats.clone()); + if let AppliedState::Known(applied) = &self.applied + && applied.matches_policy(&next.policy) + { + next.stats.set_conntrack_rule_apply_ok(true); + } + return Some(next); + } + self.last_generation = next.generation; + self.last_policy = Some(next.policy.clone()); + self.last_stats = Some(next.stats.clone()); + Some(next) + } +} diff --git a/src/conntrack_control/firewall/command.rs b/src/conntrack_control/firewall/command.rs new file mode 100644 index 0000000..c9fb7d5 --- /dev/null +++ b/src/conntrack_control/firewall/command.rs @@ -0,0 +1,204 @@ +use std::time::Duration; + +#[cfg(unix)] +use tokio::io::AsyncWriteExt; +#[cfg(unix)] +use tokio::process::Command; + +#[cfg(unix)] +use crate::util::trusted_command::resolve_trusted_helper; + +const COMMAND_TIMEOUT: Duration = Duration::from_secs(30); + +#[derive(Clone, Debug, Eq, PartialEq)] +pub(super) struct CommandSpec { + pub(super) binary: &'static str, + pub(super) args: Vec, + pub(super) stdin: Option, +} + +impl CommandSpec { + pub(super) fn new(binary: &'static str, args: impl IntoIterator) -> Self { + Self { + binary, + args: args.into_iter().map(str::to_string).collect(), + stdin: None, + } + } + + pub(super) fn with_stdin( + binary: &'static str, + args: impl IntoIterator, + stdin: String, + ) -> Self { + Self { + binary, + args: args.into_iter().map(str::to_string).collect(), + stdin: Some(stdin), + } + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(super) enum CommandErrorKind { + Missing, + NotFound, + Cancelled, + Timeout, + Failed, +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub(super) struct CommandError { + pub(super) kind: CommandErrorKind, + pub(super) message: String, +} + +impl CommandError { + pub(super) fn cancelled() -> Self { + Self { + kind: CommandErrorKind::Cancelled, + message: "firewall transaction cancelled".to_string(), + } + } + + pub(super) fn failed(message: impl Into) -> Self { + Self { + kind: CommandErrorKind::Failed, + message: message.into(), + } + } +} + +impl std::fmt::Display for CommandError { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str(&self.message) + } +} + +pub(super) trait FirewallCommandRunner: Send + Sync { + fn available(&self, binary: &str) -> bool; + + fn has_cap_net_admin(&self) -> bool; + + async fn run(&self, spec: CommandSpec) -> Result<(), CommandError>; +} + +#[derive(Clone, Copy, Default)] +pub(super) struct SystemCommandRunner; + +impl FirewallCommandRunner for SystemCommandRunner { + fn available(&self, binary: &str) -> bool { + #[cfg(unix)] + { + resolve_trusted_helper(binary).is_some() + } + #[cfg(not(unix))] + { + let _ = binary; + false + } + } + + fn has_cap_net_admin(&self) -> bool { + #[cfg(target_os = "linux")] + { + let Ok(status) = std::fs::read_to_string("/proc/self/status") else { + return false; + }; + for line in status.lines() { + if let Some(raw) = line.strip_prefix("CapEff:") { + let caps = raw.trim(); + if let Ok(bits) = u64::from_str_radix(caps, 16) { + const CAP_NET_ADMIN_BIT: u64 = 12; + return (bits & (1u64 << CAP_NET_ADMIN_BIT)) != 0; + } + } + } + false + } + #[cfg(not(target_os = "linux"))] + { + false + } + } + + async fn run(&self, spec: CommandSpec) -> Result<(), CommandError> { + #[cfg(not(unix))] + { + Err(CommandError { + kind: CommandErrorKind::Missing, + message: format!("{} is not available", spec.binary), + }) + } + #[cfg(unix)] + { + let Some(command_path) = resolve_trusted_helper(spec.binary) else { + return Err(CommandError { + kind: CommandErrorKind::Missing, + message: format!("{} is not available", spec.binary), + }); + }; + let mut command = Command::new(command_path); + command.args(&spec.args); + command.env("LC_ALL", "C"); + if spec.stdin.is_some() { + command.stdin(std::process::Stdio::piped()); + } + command.stdout(std::process::Stdio::null()); + command.stderr(std::process::Stdio::piped()); + command.kill_on_drop(true); + let mut child = command.spawn().map_err(|error| CommandError { + kind: CommandErrorKind::Failed, + message: format!("spawn {} failed: {error}", spec.binary), + })?; + let binary = spec.binary; + let output = tokio::time::timeout(COMMAND_TIMEOUT, async move { + if let Some(blob) = spec.stdin + && let Some(mut writer) = child.stdin.take() + { + writer + .write_all(blob.as_bytes()) + .await + .map_err(|error| CommandError { + kind: CommandErrorKind::Failed, + message: format!("stdin write {binary} failed: {error}"), + })?; + } + child + .wait_with_output() + .await + .map_err(|error| CommandError { + kind: CommandErrorKind::Failed, + message: format!("wait {binary} failed: {error}"), + }) + }) + .await + .map_err(|_| CommandError { + kind: CommandErrorKind::Timeout, + message: format!("{binary} timed out after {}s", COMMAND_TIMEOUT.as_secs()), + })??; + if output.status.success() { + return Ok(()); + } + let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string(); + let message = if stderr.is_empty() { + format!("{binary} exited with status {}", output.status) + } else { + stderr + }; + let kind = if is_not_found_error(&message) { + CommandErrorKind::NotFound + } else { + CommandErrorKind::Failed + }; + Err(CommandError { kind, message }) + } + } +} + +pub(super) fn is_not_found_error(message: &str) -> bool { + message.contains("No chain/target/match by that name") + || message.contains("Bad rule (does a matching rule exist in that chain?)") + || message.contains("Could not process rule: No such file or directory") +} diff --git a/src/conntrack_control/firewall/iptables.rs b/src/conntrack_control/firewall/iptables.rs new file mode 100644 index 0000000..6e06c1b --- /dev/null +++ b/src/conntrack_control/firewall/iptables.rs @@ -0,0 +1,240 @@ +use super::command::{CommandError, CommandErrorKind, CommandSpec, FirewallCommandRunner}; +use super::model::{NotrackTarget, ShadowSlot}; + +const DISPATCH_CHAIN: &str = "TELEMT_NOTRACK"; +const SHADOW_CHAIN_A: &str = "TELEMT_NT_A"; +const SHADOW_CHAIN_B: &str = "TELEMT_NT_B"; +const MAX_OWNED_JUMPS: usize = 8; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(super) enum IpFamily { + V4, + V6, +} + +impl IpFamily { + fn command_binary(self) -> &'static str { + match self { + Self::V4 => "iptables", + Self::V6 => "ip6tables", + } + } + + fn restore_binary(self) -> &'static str { + match self { + Self::V4 => "iptables-restore", + Self::V6 => "ip6tables-restore", + } + } +} + +pub(super) fn family_available(runner: &R, family: IpFamily) -> bool { + runner.available(family.command_binary()) && runner.available(family.restore_binary()) +} + +fn shadow_chain(slot: ShadowSlot) -> &'static str { + match slot { + ShadowSlot::A => SHADOW_CHAIN_A, + ShadowSlot::B => SHADOW_CHAIN_B, + } +} + +pub(super) async fn stage_family( + runner: &R, + family: IpFamily, + slot: ShadowSlot, + targets: &[NotrackTarget], +) -> Result<(), CommandError> { + if targets.is_empty() { + return Ok(()); + } + require_family(runner, family)?; + ensure_owned_chains(runner, family).await?; + let script = render_stage_script(slot, targets); + runner + .run(CommandSpec::with_stdin( + family.restore_binary(), + ["--noflush"], + script, + )) + .await +} + +pub(super) async fn activate_family( + runner: &R, + family: IpFamily, + slot: Option, +) -> Result<(), CommandError> { + require_family(runner, family)?; + if slot.is_some() { + ensure_prerouting_jump(runner, family).await?; + } + runner + .run(CommandSpec::with_stdin( + family.restore_binary(), + ["--noflush"], + render_dispatch_script(slot), + )) + .await +} + +pub(super) async fn cleanup_all(runner: &R) -> Result<(), CommandError> { + let mut errors = Vec::new(); + for family in [IpFamily::V4, IpFamily::V6] { + if !runner.available(family.command_binary()) { + continue; + } + if let Err(error) = cleanup_family(runner, family).await { + errors.push(error.message); + } + } + if errors.is_empty() { + Ok(()) + } else { + Err(CommandError::failed(errors.join("; "))) + } +} + +async fn cleanup_family( + runner: &R, + family: IpFamily, +) -> Result<(), CommandError> { + let binary = family.command_binary(); + let mut errors = Vec::new(); + for _ in 0..MAX_OWNED_JUMPS { + let result = runner + .run(CommandSpec::new( + binary, + ["-t", "raw", "-D", "PREROUTING", "-j", DISPATCH_CHAIN], + )) + .await; + match result { + Ok(()) => {} + Err(error) + if matches!( + error.kind, + CommandErrorKind::NotFound | CommandErrorKind::Missing + ) => + { + break; + } + Err(error) => { + errors.push(error.message); + break; + } + } + } + for chain in [DISPATCH_CHAIN, SHADOW_CHAIN_A, SHADOW_CHAIN_B] { + for operation in ["-F", "-X"] { + let result = runner + .run(CommandSpec::new(binary, ["-t", "raw", operation, chain])) + .await; + if let Err(error) = result + && !matches!( + error.kind, + CommandErrorKind::NotFound | CommandErrorKind::Missing + ) + { + errors.push(error.message); + } + } + } + if errors.is_empty() { + Ok(()) + } else { + Err(CommandError::failed(errors.join("; "))) + } +} + +async fn ensure_prerouting_jump( + runner: &R, + family: IpFamily, +) -> Result<(), CommandError> { + let binary = family.command_binary(); + match runner + .run(CommandSpec::new( + binary, + ["-t", "raw", "-C", "PREROUTING", "-j", DISPATCH_CHAIN], + )) + .await + { + Ok(()) => Ok(()), + Err(error) if error.kind == CommandErrorKind::NotFound => { + runner + .run(CommandSpec::new( + binary, + ["-t", "raw", "-I", "PREROUTING", "1", "-j", DISPATCH_CHAIN], + )) + .await + } + Err(error) => Err(error), + } +} + +async fn ensure_owned_chains( + runner: &R, + family: IpFamily, +) -> Result<(), CommandError> { + let binary = family.command_binary(); + for chain in [DISPATCH_CHAIN, SHADOW_CHAIN_A, SHADOW_CHAIN_B] { + match runner + .run(CommandSpec::new(binary, ["-t", "raw", "-N", chain])) + .await + { + Ok(()) => {} + Err(error) if is_chain_exists_error(&error.message) => {} + Err(error) => return Err(error), + } + } + Ok(()) +} + +pub(super) fn is_chain_exists_error(message: &str) -> bool { + message.contains("Chain already exists") +} + +fn require_family( + runner: &R, + family: IpFamily, +) -> Result<(), CommandError> { + for binary in [family.command_binary(), family.restore_binary()] { + if !runner.available(binary) { + return Err(CommandError { + kind: CommandErrorKind::Missing, + message: format!("{binary} is required for conntrack firewall reconciliation"), + }); + } + } + Ok(()) +} + +pub(super) fn render_stage_script(slot: ShadowSlot, targets: &[NotrackTarget]) -> String { + let chain = shadow_chain(slot); + let mut script = format!("*raw\n-F {chain}\n"); + for target in targets { + script.push_str("-A "); + script.push_str(chain); + script.push_str(" -p tcp --dport "); + script.push_str(&target.port.to_string()); + if let Some(ip) = target.ip { + script.push_str(" -d "); + script.push_str(&ip.to_string()); + } + script.push_str(" -j CT --notrack\n"); + } + script.push_str("COMMIT\n"); + script +} + +pub(super) fn render_dispatch_script(slot: Option) -> String { + let mut script = format!("*raw\n-F {DISPATCH_CHAIN}\n"); + if let Some(slot) = slot { + script.push_str("-A "); + script.push_str(DISPATCH_CHAIN); + script.push_str(" -j "); + script.push_str(shadow_chain(slot)); + script.push('\n'); + } + script.push_str("COMMIT\n"); + script +} diff --git a/src/conntrack_control/firewall/model.rs b/src/conntrack_control/firewall/model.rs new file mode 100644 index 0000000..8aea825 --- /dev/null +++ b/src/conntrack_control/firewall/model.rs @@ -0,0 +1,204 @@ +use std::collections::BTreeSet; +use std::net::IpAddr; +use std::sync::Arc; + +use crate::config::{ConntrackBackend, ConntrackMode, ProxyConfig}; +use crate::stats::Stats; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(super) enum ShadowSlot { + A, + B, +} + +impl ShadowSlot { + pub(super) fn other(self) -> Self { + match self { + Self::A => Self::B, + Self::B => Self::A, + } + } +} + +#[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd)] +pub(super) struct NotrackTarget { + pub(super) ip: Option, + pub(super) port: u16, +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub(super) enum DesiredPolicy { + Empty, + Rules { + configured_backend: ConntrackBackend, + v4: Vec, + v6: Vec, + }, +} + +impl DesiredPolicy { + pub(super) fn from_config(cfg: &ProxyConfig) -> Self { + if !cfg.server.conntrack_control.inline_conntrack_control + || matches!(cfg.server.conntrack_control.mode, ConntrackMode::Tracked) + { + return Self::Empty; + } + let (v4, v6) = notrack_targets(cfg); + if v4.is_empty() && v6.is_empty() { + Self::Empty + } else { + Self::Rules { + configured_backend: cfg.server.conntrack_control.backend, + v4, + v6, + } + } + } +} + +#[derive(Clone)] +pub(super) struct DesiredState { + pub(super) generation: u64, + pub(super) policy: DesiredPolicy, + pub(super) stats: Arc, +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub(super) enum AppliedPlan { + Empty, + Iptables { + slot: ShadowSlot, + v4: Vec, + v6: Vec, + }, + Nftables { + slot: ShadowSlot, + v4: Vec, + v6: Vec, + }, +} + +impl AppliedPlan { + pub(super) fn slot(&self) -> Option { + match self { + Self::Empty => None, + Self::Iptables { slot, .. } | Self::Nftables { slot, .. } => Some(*slot), + } + } + + pub(super) fn matches_policy(&self, policy: &DesiredPolicy) -> bool { + match (self, policy) { + (Self::Empty, DesiredPolicy::Empty) => true, + ( + Self::Iptables { v4, v6, .. }, + DesiredPolicy::Rules { + configured_backend, + v4: desired_v4, + v6: desired_v6, + }, + ) => { + matches!( + configured_backend, + ConntrackBackend::Auto | ConntrackBackend::Iptables + ) && v4 == desired_v4 + && v6 == desired_v6 + } + ( + Self::Nftables { v4, v6, .. }, + DesiredPolicy::Rules { + configured_backend, + v4: desired_v4, + v6: desired_v6, + }, + ) => { + matches!( + configured_backend, + ConntrackBackend::Auto | ConntrackBackend::Nftables + ) && v4 == desired_v4 + && v6 == desired_v6 + } + _ => false, + } + } +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub(super) enum AppliedState { + Known(AppliedPlan), + Unknown, +} + +fn listener_port_set(cfg: &ProxyConfig) -> Vec { + let mut ports = BTreeSet::new(); + if cfg.server.listeners.is_empty() { + ports.insert(cfg.server.port); + } else { + for listener in &cfg.server.listeners { + ports.insert(listener.port.unwrap_or(cfg.server.port)); + } + } + ports.into_iter().collect() +} + +fn notrack_targets(cfg: &ProxyConfig) -> (Vec, Vec) { + let mut v4_targets = BTreeSet::new(); + let mut v6_targets = BTreeSet::new(); + match cfg.server.conntrack_control.mode { + ConntrackMode::Tracked => {} + ConntrackMode::Notrack => { + if cfg.server.listeners.is_empty() { + let port = cfg.server.port; + for raw in [ + cfg.server.listen_addr_ipv4.as_deref(), + cfg.server.listen_addr_ipv6.as_deref(), + ] + .into_iter() + .flatten() + { + if let Ok(ip) = raw.parse::() { + let target = NotrackTarget { + ip: (!ip.is_unspecified()).then_some(ip), + port, + }; + if ip.is_ipv4() { + v4_targets.insert(target); + } else { + v6_targets.insert(target); + } + } + } + } else { + for listener in &cfg.server.listeners { + let target = NotrackTarget { + ip: (!listener.ip.is_unspecified()).then_some(listener.ip), + port: listener.port.unwrap_or(cfg.server.port), + }; + if listener.ip.is_ipv4() { + v4_targets.insert(target); + } else { + v6_targets.insert(target); + } + } + } + } + ConntrackMode::Hybrid => { + for ip in &cfg.server.conntrack_control.hybrid_listener_ips { + for port in listener_port_set(cfg) { + let target = NotrackTarget { + ip: Some(*ip), + port, + }; + if ip.is_ipv4() { + v4_targets.insert(target); + } else { + v6_targets.insert(target); + } + } + } + } + } + ( + v4_targets.into_iter().collect(), + v6_targets.into_iter().collect(), + ) +} diff --git a/src/conntrack_control/firewall/nftables.rs b/src/conntrack_control/firewall/nftables.rs new file mode 100644 index 0000000..370d8f9 --- /dev/null +++ b/src/conntrack_control/firewall/nftables.rs @@ -0,0 +1,146 @@ +use super::command::{CommandError, CommandErrorKind, CommandSpec, FirewallCommandRunner}; +use super::model::{NotrackTarget, ShadowSlot}; + +const LEGACY_TABLE: &str = "telemt_conntrack"; +const TABLE_A: &str = "telemt_conntrack_a"; +const TABLE_B: &str = "telemt_conntrack_b"; + +pub(super) fn available(runner: &R) -> bool { + runner.available("nft") +} + +fn table(slot: ShadowSlot) -> &'static str { + match slot { + ShadowSlot::A => TABLE_A, + ShadowSlot::B => TABLE_B, + } +} + +pub(super) async fn stage( + runner: &R, + slot: ShadowSlot, + v4: &[NotrackTarget], + v6: &[NotrackTarget], +) -> Result<(), CommandError> { + require_nft(runner)?; + delete_table_if_present(runner, table(slot)).await?; + runner + .run(CommandSpec::with_stdin( + "nft", + ["-f", "-"], + render_stage_script(slot, v4, v6), + )) + .await +} + +pub(super) async fn activate( + runner: &R, + slot: ShadowSlot, +) -> Result<(), CommandError> { + require_nft(runner)?; + runner + .run(CommandSpec::with_stdin( + "nft", + ["-f", "-"], + render_activate_script(slot), + )) + .await +} + +pub(super) async fn deactivate( + runner: &R, + slot: ShadowSlot, +) -> Result<(), CommandError> { + delete_table_if_present(runner, table(slot)).await +} + +pub(super) async fn cleanup_all(runner: &R) -> Result<(), CommandError> { + if !runner.available("nft") { + return Ok(()); + } + let mut errors = Vec::new(); + for table_name in [LEGACY_TABLE, TABLE_A, TABLE_B] { + if let Err(error) = delete_table_if_present(runner, table_name).await { + errors.push(error.message); + } + } + if errors.is_empty() { + Ok(()) + } else { + Err(CommandError::failed(errors.join("; "))) + } +} + +async fn delete_table_if_present( + runner: &R, + table_name: &'static str, +) -> Result<(), CommandError> { + match runner + .run(CommandSpec::new( + "nft", + ["delete", "table", "inet", table_name], + )) + .await + { + Ok(()) => Ok(()), + Err(error) + if matches!( + error.kind, + CommandErrorKind::NotFound | CommandErrorKind::Missing + ) => + { + Ok(()) + } + Err(error) => Err(error), + } +} + +fn require_nft(runner: &R) -> Result<(), CommandError> { + if runner.available("nft") { + Ok(()) + } else { + Err(CommandError { + kind: CommandErrorKind::Missing, + message: "nft is required for conntrack firewall reconciliation".to_string(), + }) + } +} + +pub(super) fn render_stage_script( + slot: ShadowSlot, + v4: &[NotrackTarget], + v6: &[NotrackTarget], +) -> String { + let table = table(slot); + let mut script = format!("add table inet {table}\nadd chain inet {table} rules\n"); + for target in v4 { + script.push_str("add rule inet "); + script.push_str(table); + script.push_str(" rules tcp dport "); + script.push_str(&target.port.to_string()); + if let Some(ip) = target.ip { + script.push_str(" ip daddr "); + script.push_str(&ip.to_string()); + } + script.push_str(" notrack\n"); + } + for target in v6 { + script.push_str("add rule inet "); + script.push_str(table); + script.push_str(" rules tcp dport "); + script.push_str(&target.port.to_string()); + if let Some(ip) = target.ip { + script.push_str(" ip6 daddr "); + script.push_str(&ip.to_string()); + } + script.push_str(" notrack\n"); + } + script +} + +pub(super) fn render_activate_script(slot: ShadowSlot) -> String { + let table = table(slot); + format!( + "add chain inet {table} preraw {{ type filter hook prerouting priority raw; policy accept; }}\nadd rule inet {table} preraw jump rules\n" + ) +} diff --git a/src/conntrack_control/firewall/runtime.rs b/src/conntrack_control/firewall/runtime.rs new file mode 100644 index 0000000..706e007 --- /dev/null +++ b/src/conntrack_control/firewall/runtime.rs @@ -0,0 +1,91 @@ +use tracing::debug; + +use crate::config::{ConntrackBackend, ProxyConfig}; +use crate::conntrack_control::{ConntrackRuntimeSupport, NetfilterBackend}; +use crate::proxy::shared_state::ConntrackCloseEvent; + +use super::command::{CommandSpec, FirewallCommandRunner, SystemCommandRunner}; +use super::nftables; + +/// Probes the effective firewall backend and conntrack deletion capability. +pub(in crate::conntrack_control) fn probe_runtime_support( + configured_backend: ConntrackBackend, +) -> ConntrackRuntimeSupport { + let runner = SystemCommandRunner; + let iptables_available = runner.available("iptables"); + let netfilter_backend = match configured_backend { + ConntrackBackend::Auto if nftables::available(&runner) => Some(NetfilterBackend::Nftables), + ConntrackBackend::Auto if iptables_available => Some(NetfilterBackend::Iptables), + ConntrackBackend::Nftables if nftables::available(&runner) => { + Some(NetfilterBackend::Nftables) + } + ConntrackBackend::Iptables if iptables_available => Some(NetfilterBackend::Iptables), + _ => None, + }; + ConntrackRuntimeSupport { + netfilter_backend, + has_cap_net_admin: runner.has_cap_net_admin(), + has_conntrack_binary: runner.available("conntrack"), + } +} + +/// Resolves whether conntrack close publication is usable for this runtime. +pub(in crate::conntrack_control) fn effective_conntrack_enabled( + cfg: &ProxyConfig, + runtime_support: ConntrackRuntimeSupport, +) -> bool { + cfg.server.conntrack_control.inline_conntrack_control + && runtime_support.has_cap_net_admin + && runtime_support.netfilter_backend.is_some() + && runtime_support.has_conntrack_binary +} + +/// Result of one best-effort kernel conntrack deletion. +pub(in crate::conntrack_control) enum DeleteOutcome { + /// The kernel reported successful deletion. + Deleted, + /// No matching conntrack entry existed. + NotFound, + /// The helper was unavailable or returned an unexpected failure. + Error, +} + +/// Deletes the exact TCP tuple represented by one close event. +pub(in crate::conntrack_control) async fn delete_conntrack_entry( + event: ConntrackCloseEvent, +) -> DeleteOutcome { + let runner = SystemCommandRunner; + if !runner.available("conntrack") { + return DeleteOutcome::Error; + } + let args = vec![ + "-D".to_string(), + "-p".to_string(), + "tcp".to_string(), + "-s".to_string(), + event.src.ip().to_string(), + "--sport".to_string(), + event.src.port().to_string(), + "-d".to_string(), + event.dst.ip().to_string(), + "--dport".to_string(), + event.dst.port().to_string(), + ]; + match runner + .run(CommandSpec { + binary: "conntrack", + args, + stdin: None, + }) + .await + { + Ok(()) => DeleteOutcome::Deleted, + Err(error) if error.message.contains("0 flow entries have been deleted") => { + DeleteOutcome::NotFound + } + Err(error) => { + debug!(error = %error, "conntrack delete failed"); + DeleteOutcome::Error + } + } +} diff --git a/src/conntrack_control/firewall/tests.rs b/src/conntrack_control/firewall/tests.rs new file mode 100644 index 0000000..ddd1239 --- /dev/null +++ b/src/conntrack_control/firewall/tests.rs @@ -0,0 +1,488 @@ +use std::collections::{BTreeMap, BTreeSet}; +use std::future::pending; +use std::sync::atomic::Ordering; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use tokio::sync::{Notify, watch}; +use tokio_util::sync::CancellationToken; + +use crate::config::ConntrackBackend; +use crate::stats::Stats; + +use super::actor::{FirewallReconciler, ReconcileOutcome}; +use super::command::{CommandError, CommandErrorKind, CommandSpec, FirewallCommandRunner}; +use super::model::{ + AppliedPlan, AppliedState, DesiredPolicy, DesiredState, NotrackTarget, ShadowSlot, +}; +use super::transaction::{InterruptibleRunner, reconcile_once}; + +#[path = "tests/model_tests.rs"] +mod model_tests; + +#[derive(Clone)] +struct FailureRule { + binary: &'static str, + occurrence: usize, +} + +#[derive(Default)] +struct FakeState { + calls: Vec, + binary_calls: BTreeMap<&'static str, usize>, + failures: Vec, +} + +#[derive(Clone)] +struct FakeRunner { + available: Arc>, + has_cap_net_admin: bool, + state: Arc>, +} + +impl FakeRunner { + fn all_available() -> Self { + Self { + available: Arc::new(BTreeSet::from([ + "conntrack", + "ip6tables", + "ip6tables-restore", + "iptables", + "iptables-restore", + "nft", + ])), + has_cap_net_admin: true, + state: Arc::new(Mutex::new(FakeState::default())), + } + } + + fn with_failure(self, binary: &'static str, occurrence: usize) -> Self { + self.state + .lock() + .unwrap() + .failures + .push(FailureRule { binary, occurrence }); + self + } + + fn calls(&self) -> Vec { + self.state.lock().unwrap().calls.clone() + } +} + +impl FirewallCommandRunner for FakeRunner { + fn available(&self, binary: &str) -> bool { + self.available.contains(binary) + } + + fn has_cap_net_admin(&self) -> bool { + self.has_cap_net_admin + } + + async fn run(&self, spec: CommandSpec) -> Result<(), CommandError> { + let mut state = self.state.lock().unwrap(); + let occurrence = { + let count = state.binary_calls.entry(spec.binary).or_default(); + *count += 1; + *count + }; + state.calls.push(spec.clone()); + if state + .failures + .iter() + .any(|failure| failure.binary == spec.binary && failure.occurrence == occurrence) + { + return Err(CommandError::failed(format!( + "injected {} failure at occurrence {}", + spec.binary, occurrence + ))); + } + drop(state); + + let operation = spec.args.get(2).map(String::as_str); + if (matches!(spec.binary, "iptables" | "ip6tables") + && matches!(operation, Some("-C" | "-D" | "-F" | "-X"))) + || (spec.binary == "nft" && spec.args.first().map(String::as_str) == Some("delete")) + { + return Err(CommandError { + kind: CommandErrorKind::NotFound, + message: "injected object not found".to_string(), + }); + } + Ok(()) + } +} + +struct BlockingRunner { + entered: Arc, +} + +impl FirewallCommandRunner for BlockingRunner { + fn available(&self, _binary: &str) -> bool { + true + } + + fn has_cap_net_admin(&self) -> bool { + true + } + + async fn run(&self, _spec: CommandSpec) -> Result<(), CommandError> { + self.entered.notify_one(); + pending().await + } +} + +fn target(ip: Option<&str>, port: u16) -> NotrackTarget { + NotrackTarget { + ip: ip.map(|value| value.parse().unwrap()), + port, + } +} + +fn desired(generation: u64, policy: DesiredPolicy) -> DesiredState { + DesiredState { + generation, + policy, + stats: Arc::new(Stats::new()), + } +} + +fn dual_stack_policy(port: u16) -> DesiredPolicy { + DesiredPolicy::Rules { + configured_backend: ConntrackBackend::Iptables, + v4: vec![target(Some("192.0.2.10"), port)], + v6: vec![target(Some("2001:db8::10"), port)], + } +} + +fn nft_dual_stack_policy(port: u16) -> DesiredPolicy { + DesiredPolicy::Rules { + configured_backend: ConntrackBackend::Nftables, + v4: vec![target(Some("192.0.2.10"), port)], + v6: vec![target(Some("2001:db8::10"), port)], + } +} + +#[tokio::test] +async fn successful_reconcile_flips_between_shadow_slots() { + let runner = FakeRunner::all_available(); + let mut applied = AppliedState::Known(AppliedPlan::Empty); + let first = desired(1, dual_stack_policy(443)); + + reconcile_once(&runner, &runner, &mut applied, &first) + .await + .unwrap(); + assert!(matches!( + applied, + AppliedState::Known(AppliedPlan::Iptables { + slot: ShadowSlot::A, + .. + }) + )); + + let second = desired(2, dual_stack_policy(8443)); + reconcile_once(&runner, &runner, &mut applied, &second) + .await + .unwrap(); + assert!(matches!( + applied, + AppliedState::Known(AppliedPlan::Iptables { + slot: ShadowSlot::B, + .. + }) + )); + assert!(runner.calls().iter().any(|call| { + call.stdin + .as_deref() + .is_some_and(|script| script.contains("-A TELEMT_NOTRACK -j TELEMT_NT_B")) + })); +} + +#[tokio::test] +async fn identical_policy_is_command_free_after_convergence() { + let runner = FakeRunner::all_available(); + let mut applied = AppliedState::Known(AppliedPlan::Empty); + reconcile_once( + &runner, + &runner, + &mut applied, + &desired(1, dual_stack_policy(443)), + ) + .await + .unwrap(); + let calls_after_convergence = runner.calls().len(); + + reconcile_once( + &runner, + &runner, + &mut applied, + &desired(2, dual_stack_policy(443)), + ) + .await + .unwrap(); + + assert_eq!(runner.calls().len(), calls_after_convergence); +} + +#[tokio::test] +async fn backend_migration_failure_restores_previous_backend() { + let runner = FakeRunner::all_available().with_failure("ip6tables-restore", 3); + let mut applied = AppliedState::Known(AppliedPlan::Empty); + reconcile_once( + &runner, + &runner, + &mut applied, + &desired(1, dual_stack_policy(443)), + ) + .await + .unwrap(); + let previous = applied.clone(); + + let failure = reconcile_once( + &runner, + &runner, + &mut applied, + &desired(2, nft_dual_stack_policy(8443)), + ) + .await + .unwrap_err(); + + assert_eq!(failure.rollback_succeeded, Some(true)); + assert_eq!(applied, previous); +} + +#[tokio::test] +async fn unavailable_target_does_not_clear_confirmed_applied_policy() { + let runner = FakeRunner::all_available(); + let mut applied = AppliedState::Known(AppliedPlan::Empty); + reconcile_once( + &runner, + &runner, + &mut applied, + &desired(1, dual_stack_policy(443)), + ) + .await + .unwrap(); + let previous = applied.clone(); + let unavailable = FakeRunner { + available: Arc::new(BTreeSet::new()), + has_cap_net_admin: false, + state: Arc::new(Mutex::new(FakeState::default())), + }; + + let failure = reconcile_once( + &unavailable, + &unavailable, + &mut applied, + &desired(2, nft_dual_stack_policy(8443)), + ) + .await + .unwrap_err(); + + assert_eq!(failure.rollback_succeeded, None); + assert_eq!(applied, previous); + assert!(unavailable.calls().is_empty()); +} + +#[tokio::test] +async fn partial_dual_stack_failure_restores_previous_applied_plan() { + let runner = FakeRunner::all_available().with_failure("ip6tables-restore", 4); + let mut applied = AppliedState::Known(AppliedPlan::Empty); + let first = desired(1, dual_stack_policy(443)); + reconcile_once(&runner, &runner, &mut applied, &first) + .await + .unwrap(); + let previous = applied.clone(); + + let failure = reconcile_once( + &runner, + &runner, + &mut applied, + &desired(2, dual_stack_policy(8443)), + ) + .await + .unwrap_err(); + + assert_eq!(failure.rollback_succeeded, Some(true)); + assert_eq!(applied, previous); +} + +#[tokio::test] +async fn rollback_failure_marks_applied_state_unknown() { + let runner = FakeRunner::all_available() + .with_failure("ip6tables-restore", 4) + .with_failure("nft", 2); + let mut applied = AppliedState::Known(AppliedPlan::Empty); + reconcile_once( + &runner, + &runner, + &mut applied, + &desired(1, dual_stack_policy(443)), + ) + .await + .unwrap(); + + let failure = reconcile_once( + &runner, + &runner, + &mut applied, + &desired(2, dual_stack_policy(8443)), + ) + .await + .unwrap_err(); + + assert_eq!(failure.rollback_succeeded, Some(false)); + assert_eq!(applied, AppliedState::Unknown); +} + +#[tokio::test] +async fn transaction_cancellation_does_not_claim_a_new_applied_plan() { + let entered = Arc::new(Notify::new()); + let runner = BlockingRunner { + entered: entered.clone(), + }; + let terminal = CancellationToken::new(); + let process_cancellation = CancellationToken::new(); + let interruptible = InterruptibleRunner::new(&runner, &terminal, &process_cancellation); + let mut applied = AppliedState::Known(AppliedPlan::Empty); + let desired = desired(1, dual_stack_policy(443)); + let failure = { + let transaction = reconcile_once(&interruptible, &runner, &mut applied, &desired); + tokio::pin!(transaction); + + tokio::select! { + _ = entered.notified() => terminal.cancel(), + _ = &mut transaction => panic!("transaction completed before injected cancellation"), + } + transaction.await.unwrap_err() + }; + + assert!(failure.cancelled); + assert_eq!(applied, AppliedState::Known(AppliedPlan::Empty)); +} + +#[test] +fn desired_watch_coalesces_and_rejects_stale_or_conflicting_generations() { + let runner = FakeRunner::all_available(); + let (desired_tx, desired_rx) = watch::channel(None); + let (status_tx, _status_rx) = watch::channel(None); + let terminal = CancellationToken::new(); + let closed = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let completed = Arc::new(Notify::new()); + let mut reconciler = FirewallReconciler::new( + runner, + desired_rx, + status_tx, + terminal, + closed, + Arc::new(std::sync::atomic::AtomicBool::new(false)), + Arc::new(std::sync::atomic::AtomicBool::new(false)), + completed, + ); + + desired_tx.send_replace(Some(desired(1, dual_stack_policy(443)))); + desired_tx.send_replace(Some(desired(2, dual_stack_policy(8443)))); + assert_eq!(reconciler.take_latest_desired().unwrap().generation, 2); + + desired_tx.send_replace(Some(desired(1, dual_stack_policy(443)))); + assert!(reconciler.take_latest_desired().is_none()); + + desired_tx.send_replace(Some(desired(2, dual_stack_policy(9443)))); + assert!(reconciler.take_latest_desired().is_none()); +} + +#[tokio::test(start_paused = true)] +async fn actor_retries_with_backoff_and_cleans_owned_rules_on_shutdown() { + let runner = FakeRunner::all_available().with_failure("iptables-restore", 1); + let observed_runner = runner.clone(); + let (desired_tx, desired_rx) = watch::channel(None); + let (status_tx, mut status_rx) = watch::channel(None); + let terminal = CancellationToken::new(); + let closed = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let completed_flag = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let cleanup_succeeded = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let completed = Arc::new(Notify::new()); + let reconciler = FirewallReconciler::new( + runner, + desired_rx, + status_tx, + terminal.clone(), + closed.clone(), + completed_flag.clone(), + cleanup_succeeded.clone(), + completed, + ); + let process_cancellation = CancellationToken::new(); + let task = tokio::spawn(reconciler.run(process_cancellation)); + + let initial_desired = desired(2, dual_stack_policy(443)); + let desired_stats = initial_desired.stats.clone(); + desired_tx.send_replace(Some(initial_desired)); + status_rx.changed().await.unwrap(); + assert_eq!( + status_rx.borrow().as_ref().unwrap().outcome, + ReconcileOutcome::Failed + ); + let calls_after_failure = observed_runner.calls().len(); + + desired_tx.send_replace(Some(desired(1, dual_stack_policy(7443)))); + tokio::task::yield_now().await; + desired_tx.send_replace(Some(desired(2, dual_stack_policy(9443)))); + tokio::task::yield_now().await; + tokio::time::advance(Duration::from_millis(999)).await; + tokio::task::yield_now().await; + assert_eq!(observed_runner.calls().len(), calls_after_failure); + + tokio::time::advance(Duration::from_millis(1)).await; + status_rx.changed().await.unwrap(); + assert_eq!( + status_rx.borrow().as_ref().unwrap().outcome, + ReconcileOutcome::Applied + ); + assert_eq!(status_rx.borrow().as_ref().unwrap().generation, 2); + assert!(observed_runner.calls().iter().all(|call| { + call.stdin + .as_ref() + .is_none_or(|script| !script.contains("7443") && !script.contains("9443")) + })); + let calls_before_shutdown = observed_runner.calls().len(); + + terminal.cancel(); + tokio::time::timeout(Duration::from_secs(1), task) + .await + .unwrap() + .unwrap(); + assert!(closed.load(Ordering::Acquire)); + assert!(completed_flag.load(Ordering::Acquire)); + assert!(cleanup_succeeded.load(Ordering::Acquire)); + assert!(!desired_stats.get_conntrack_rule_apply_ok()); + assert!(observed_runner.calls().len() > calls_before_shutdown); +} + +#[tokio::test] +async fn actor_reports_terminal_cleanup_failure_separately_from_completion() { + let runner = FakeRunner::all_available().with_failure("nft", 1); + let (_desired_tx, desired_rx) = watch::channel(None); + let (status_tx, _status_rx) = watch::channel(None); + let terminal = CancellationToken::new(); + terminal.cancel(); + let closed = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let completed_flag = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let cleanup_succeeded = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let reconciler = FirewallReconciler::new( + runner, + desired_rx, + status_tx, + terminal, + closed.clone(), + completed_flag.clone(), + cleanup_succeeded.clone(), + Arc::new(Notify::new()), + ); + + reconciler.run(CancellationToken::new()).await; + + assert!(closed.load(Ordering::Acquire)); + assert!(completed_flag.load(Ordering::Acquire)); + assert!(!cleanup_succeeded.load(Ordering::Acquire)); +} diff --git a/src/conntrack_control/firewall/tests/model_tests.rs b/src/conntrack_control/firewall/tests/model_tests.rs new file mode 100644 index 0000000..d1cc49a --- /dev/null +++ b/src/conntrack_control/firewall/tests/model_tests.rs @@ -0,0 +1,118 @@ +use crate::config::{ConntrackBackend, ConntrackMode, ProxyConfig}; + +use super::super::command::is_not_found_error; +use super::super::iptables::{self, is_chain_exists_error}; +use super::super::model::{DesiredPolicy, ShadowSlot}; +use super::super::nftables; +use super::target; + +#[test] +fn desired_policy_derives_exact_listener_targets() { + let mut config = ProxyConfig::default(); + config.server.port = 8443; + config.server.listen_addr_ipv4 = Some("0.0.0.0".to_string()); + config.server.listen_addr_ipv6 = Some("2001:db8::10".to_string()); + config.server.conntrack_control.inline_conntrack_control = true; + config.server.conntrack_control.mode = ConntrackMode::Notrack; + config.server.conntrack_control.backend = ConntrackBackend::Iptables; + + assert_eq!( + DesiredPolicy::from_config(&config), + DesiredPolicy::Rules { + configured_backend: ConntrackBackend::Iptables, + v4: vec![target(None, 8443)], + v6: vec![target(Some("2001:db8::10"), 8443)], + } + ); + + config.server.conntrack_control.mode = ConntrackMode::Tracked; + assert_eq!(DesiredPolicy::from_config(&config), DesiredPolicy::Empty); + config.server.conntrack_control.mode = ConntrackMode::Notrack; + config.server.conntrack_control.inline_conntrack_control = false; + assert_eq!(DesiredPolicy::from_config(&config), DesiredPolicy::Empty); +} + +#[test] +fn hybrid_policy_is_a_sorted_deduplicated_address_port_product() { + let mut config = ProxyConfig::default(); + config.server.conntrack_control.inline_conntrack_control = true; + config.server.conntrack_control.mode = ConntrackMode::Hybrid; + config.server.conntrack_control.hybrid_listener_ips = vec![ + "2001:db8::10".parse().unwrap(), + "192.0.2.10".parse().unwrap(), + "192.0.2.10".parse().unwrap(), + ]; + config.server.listeners = vec![ + serde_json::from_value(serde_json::json!({ + "ip": "0.0.0.0", + "port": 8443 + })) + .unwrap(), + serde_json::from_value(serde_json::json!({ + "ip": "::", + "port": 443 + })) + .unwrap(), + ]; + + assert_eq!( + DesiredPolicy::from_config(&config), + DesiredPolicy::Rules { + configured_backend: ConntrackBackend::Auto, + v4: vec![ + target(Some("192.0.2.10"), 443), + target(Some("192.0.2.10"), 8443), + ], + v6: vec![ + target(Some("2001:db8::10"), 443), + target(Some("2001:db8::10"), 8443), + ], + } + ); +} + +#[test] +fn restore_renderers_keep_staging_detached_from_activation() { + let stage = iptables::render_stage_script(ShadowSlot::B, &[target(Some("192.0.2.20"), 443)]); + assert!(stage.contains("-F TELEMT_NT_B\n")); + assert!(stage.contains("-A TELEMT_NT_B -p tcp --dport 443 -d 192.0.2.20 -j CT --notrack\n")); + assert!(!stage.contains("-A TELEMT_NOTRACK -j TELEMT_NT_B")); + assert!(!stage.contains(":TELEMT_")); + + let activation = iptables::render_dispatch_script(Some(ShadowSlot::B)); + assert!(activation.contains("-F TELEMT_NOTRACK\n")); + assert!(activation.contains("-A TELEMT_NOTRACK -j TELEMT_NT_B\n")); + assert!(!activation.contains(":TELEMT_")); + + let nft_stage = nftables::render_stage_script( + ShadowSlot::A, + &[target(None, 443)], + &[target(Some("2001:db8::20"), 8443)], + ); + assert!(!nft_stage.contains("hook prerouting")); + assert!(nft_stage.contains("tcp dport 443 notrack\n")); + assert!(nft_stage.contains("tcp dport 8443 ip6 daddr 2001:db8::20 notrack\n")); + assert!(nftables::render_activate_script(ShadowSlot::A).contains("hook prerouting")); +} + +#[test] +fn command_error_classification_only_accepts_absent_owned_objects() { + assert!(is_not_found_error( + "iptables: No chain/target/match by that name." + )); + assert!(is_not_found_error( + "Bad rule (does a matching rule exist in that chain?)." + )); + assert!(is_not_found_error( + "Error: Could not process rule: No such file or directory" + )); + assert!(!is_not_found_error("Permission denied")); + assert!(!is_not_found_error( + "can't initialize iptables table `raw': Table does not exist" + )); + assert!(!is_not_found_error( + "Another app is currently holding the xtables lock" + )); + assert!(is_chain_exists_error("iptables: Chain already exists.")); + assert!(!is_chain_exists_error("Permission denied")); +} diff --git a/src/conntrack_control/firewall/transaction.rs b/src/conntrack_control/firewall/transaction.rs new file mode 100644 index 0000000..0c7ff92 --- /dev/null +++ b/src/conntrack_control/firewall/transaction.rs @@ -0,0 +1,335 @@ +use std::time::Duration; + +use tokio_util::sync::CancellationToken; + +use crate::config::ConntrackBackend; + +use super::command::{CommandError, CommandErrorKind, CommandSpec, FirewallCommandRunner}; +use super::iptables::{self, IpFamily}; +use super::model::{AppliedPlan, AppliedState, DesiredPolicy, DesiredState, ShadowSlot}; +use super::nftables; + +const TRANSACTION_TIMEOUT: Duration = Duration::from_secs(30); + +pub(super) struct InterruptibleRunner<'a, R> { + inner: &'a R, + terminal: &'a CancellationToken, + process_cancellation: &'a CancellationToken, +} + +impl<'a, R> InterruptibleRunner<'a, R> { + pub(super) fn new( + inner: &'a R, + terminal: &'a CancellationToken, + process_cancellation: &'a CancellationToken, + ) -> Self { + Self { + inner, + terminal, + process_cancellation, + } + } +} + +impl FirewallCommandRunner for InterruptibleRunner<'_, R> { + fn available(&self, binary: &str) -> bool { + self.inner.available(binary) + } + + fn has_cap_net_admin(&self) -> bool { + self.inner.has_cap_net_admin() + } + + async fn run(&self, spec: CommandSpec) -> Result<(), CommandError> { + tokio::select! { + biased; + _ = self.terminal.cancelled() => Err(CommandError::cancelled()), + _ = self.process_cancellation.cancelled() => Err(CommandError::cancelled()), + result = self.inner.run(spec) => result, + } + } +} + +#[derive(Debug)] +pub(super) struct ReconcileFailure { + pub(super) message: String, + pub(super) rollback_succeeded: Option, + pub(super) cancelled: bool, +} + +pub(super) async fn reconcile_once( + interruptible: &I, + recovery_runner: &R, + applied: &mut AppliedState, + desired: &DesiredState, +) -> Result<(), ReconcileFailure> +where + I: FirewallCommandRunner, + R: FirewallCommandRunner, +{ + if let AppliedState::Known(current) = applied + && current.matches_policy(&desired.policy) + { + return Ok(()); + } + + if matches!(applied, AppliedState::Unknown) { + match tokio::time::timeout(TRANSACTION_TIMEOUT, recover_to_empty(interruptible)).await { + Ok(Ok(())) => *applied = AppliedState::Known(AppliedPlan::Empty), + Ok(Err(error)) if error.kind == CommandErrorKind::Cancelled => { + return Err(cancelled_failure()); + } + Ok(Err(error)) => { + return Err(ReconcileFailure { + message: format!("startup recovery failed: {error}"), + rollback_succeeded: None, + cancelled: false, + }); + } + Err(_) => { + return Err(ReconcileFailure { + message: "startup recovery timed out".to_string(), + rollback_succeeded: None, + cancelled: false, + }); + } + } + } + + let previous = match applied { + AppliedState::Known(plan) => plan.clone(), + AppliedState::Unknown => unreachable!("unknown state was recovered above"), + }; + let target = match resolve_target(interruptible, &desired.policy, &previous) { + Ok(target) => target, + Err(error) => { + return Err(ReconcileFailure { + message: error.message, + rollback_succeeded: None, + cancelled: false, + }); + } + }; + + let transition = tokio::time::timeout( + TRANSACTION_TIMEOUT, + transition_plan(interruptible, &previous, &target), + ) + .await; + match transition { + Ok(Ok(())) => { + *applied = AppliedState::Known(target); + Ok(()) + } + Ok(Err(error)) if error.kind == CommandErrorKind::Cancelled => Err(cancelled_failure()), + Ok(Err(error)) => { + rollback_after_failure(recovery_runner, applied, previous, error.message).await + } + Err(_) => { + rollback_after_failure( + recovery_runner, + applied, + previous, + "firewall apply transaction timed out".to_string(), + ) + .await + } + } +} + +async fn rollback_after_failure( + runner: &R, + applied: &mut AppliedState, + previous: AppliedPlan, + apply_error: String, +) -> Result<(), ReconcileFailure> { + let rollback = tokio::time::timeout(TRANSACTION_TIMEOUT, restore_plan(runner, &previous)).await; + let rollback_result = match rollback { + Ok(result) => result, + Err(_) => Err(CommandError::failed("firewall rollback timed out")), + }; + match rollback_result { + Ok(()) => { + *applied = AppliedState::Known(previous); + Err(ReconcileFailure { + message: apply_error, + rollback_succeeded: Some(true), + cancelled: false, + }) + } + Err(rollback_error) => { + *applied = AppliedState::Unknown; + if rollback_error.kind == CommandErrorKind::Cancelled { + return Err(cancelled_failure()); + } + Err(ReconcileFailure { + message: format!("{apply_error}; rollback failed: {rollback_error}"), + rollback_succeeded: Some(false), + cancelled: false, + }) + } + } +} + +fn cancelled_failure() -> ReconcileFailure { + ReconcileFailure { + message: "firewall transaction cancelled for process shutdown".to_string(), + rollback_succeeded: None, + cancelled: true, + } +} + +pub(super) fn resolve_target( + runner: &R, + desired: &DesiredPolicy, + previous: &AppliedPlan, +) -> Result { + let DesiredPolicy::Rules { + configured_backend, + v4, + v6, + } = desired + else { + return Ok(AppliedPlan::Empty); + }; + if !runner.has_cap_net_admin() { + return Err(CommandError::failed( + "CAP_NET_ADMIN is required for conntrack firewall reconciliation", + )); + } + let next_slot = previous.slot().map_or(ShadowSlot::A, ShadowSlot::other); + let iptables_available = (v4.is_empty() || iptables::family_available(runner, IpFamily::V4)) + && (v6.is_empty() || iptables::family_available(runner, IpFamily::V6)); + match configured_backend { + ConntrackBackend::Nftables if nftables::available(runner) => Ok(AppliedPlan::Nftables { + slot: next_slot, + v4: v4.clone(), + v6: v6.clone(), + }), + ConntrackBackend::Iptables if iptables_available => Ok(AppliedPlan::Iptables { + slot: next_slot, + v4: v4.clone(), + v6: v6.clone(), + }), + ConntrackBackend::Auto if nftables::available(runner) => Ok(AppliedPlan::Nftables { + slot: next_slot, + v4: v4.clone(), + v6: v6.clone(), + }), + ConntrackBackend::Auto if iptables_available => Ok(AppliedPlan::Iptables { + slot: next_slot, + v4: v4.clone(), + v6: v6.clone(), + }), + backend => Err(CommandError::failed(format!( + "configured conntrack firewall backend {backend:?} is unavailable" + ))), + } +} + +pub(super) async fn transition_plan( + runner: &R, + previous: &AppliedPlan, + target: &AppliedPlan, +) -> Result<(), CommandError> { + install_plan(runner, target).await?; + match (previous, target) { + ( + AppliedPlan::Iptables { + v4: previous_v4, + v6: previous_v6, + .. + }, + AppliedPlan::Iptables { v4, v6, .. }, + ) => { + if !previous_v4.is_empty() && v4.is_empty() { + iptables::activate_family(runner, IpFamily::V4, None).await?; + } + if !previous_v6.is_empty() && v6.is_empty() { + iptables::activate_family(runner, IpFamily::V6, None).await?; + } + } + ( + AppliedPlan::Nftables { + slot: previous_slot, + .. + }, + AppliedPlan::Nftables { slot, .. }, + ) if previous_slot != slot => { + nftables::deactivate(runner, *previous_slot).await?; + } + (_, _) if previous != target => clear_plan(runner, previous).await?, + _ => {} + } + Ok(()) +} + +async fn install_plan( + runner: &R, + plan: &AppliedPlan, +) -> Result<(), CommandError> { + match plan { + AppliedPlan::Empty => Ok(()), + AppliedPlan::Iptables { slot, v4, v6 } => { + if !v4.is_empty() { + iptables::stage_family(runner, IpFamily::V4, *slot, v4).await?; + } + if !v6.is_empty() { + iptables::stage_family(runner, IpFamily::V6, *slot, v6).await?; + } + if !v4.is_empty() { + iptables::activate_family(runner, IpFamily::V4, Some(*slot)).await?; + } + if !v6.is_empty() { + iptables::activate_family(runner, IpFamily::V6, Some(*slot)).await?; + } + Ok(()) + } + AppliedPlan::Nftables { slot, v4, v6 } => { + nftables::stage(runner, *slot, v4, v6).await?; + nftables::activate(runner, *slot).await + } + } +} + +async fn clear_plan( + runner: &R, + plan: &AppliedPlan, +) -> Result<(), CommandError> { + match plan { + AppliedPlan::Empty => Ok(()), + AppliedPlan::Iptables { v4, v6, .. } => { + if !v4.is_empty() { + iptables::activate_family(runner, IpFamily::V4, None).await?; + } + if !v6.is_empty() { + iptables::activate_family(runner, IpFamily::V6, None).await?; + } + Ok(()) + } + AppliedPlan::Nftables { slot, .. } => nftables::deactivate(runner, *slot).await, + } +} + +async fn restore_plan( + runner: &R, + plan: &AppliedPlan, +) -> Result<(), CommandError> { + recover_to_empty(runner).await?; + install_plan(runner, plan).await +} + +pub(super) async fn recover_to_empty( + runner: &R, +) -> Result<(), CommandError> { + let nft_result = nftables::cleanup_all(runner).await; + let iptables_result = iptables::cleanup_all(runner).await; + match (nft_result, iptables_result) { + (Ok(()), Ok(())) => Ok(()), + (Err(first), Ok(())) | (Ok(()), Err(first)) => Err(first), + (Err(first), Err(second)) => Err(CommandError::failed(format!( + "{}; {}", + first.message, second.message + ))), + } +} diff --git a/src/daemon/mod.rs b/src/daemon/mod.rs index 66fca0f..15d274c 100644 --- a/src/daemon/mod.rs +++ b/src/daemon/mod.rs @@ -4,14 +4,20 @@ //! and privilege dropping for running telemt as a background service. use std::fs::{self, File, OpenOptions}; -use std::io::{self, Read, Write}; +use std::io; use std::os::unix::fs::OpenOptionsExt; use std::path::{Path, PathBuf}; use nix::errno::Errno; -use nix::fcntl::{Flock, FlockArg}; -use nix::unistd::{self, ForkResult, Gid, Pid, Uid, chdir, close, fork, getpid, setsid}; -use tracing::{debug, info, warn}; +use nix::unistd::{self, ForkResult, Gid, Uid, chdir, close, fork, getpid, setsid}; +use tracing::info; + +// PID file ownership and process-control helpers. +mod pid_file; + +pub use pid_file::DaemonStatus; +#[allow(unused_imports)] +pub use pid_file::{PidFile, check_status, read_pid_file, signal_pid_file}; /// Default PID file location. pub const DEFAULT_PID_FILE: &str = "/var/run/telemt.pid"; @@ -51,36 +57,47 @@ impl DaemonOptions { /// Error types for daemon operations. #[derive(Debug, thiserror::Error)] pub enum DaemonError { + /// A daemonization fork failed. #[error("fork failed: {0}")] ForkFailed(#[source] nix::Error), + /// Creation of the detached session failed. #[error("setsid failed: {0}")] SetsidFailed(#[source] nix::Error), + /// Switching to the configured working directory failed. #[error("chdir failed: {0}")] ChdirFailed(#[source] nix::Error), + /// Opening `/dev/null` for standard-stream redirection failed. #[error("failed to open /dev/null: {0}")] DevNullFailed(#[source] io::Error), + /// Redirecting a standard file descriptor failed. #[error("failed to redirect stdio: {0}")] RedirectFailed(#[source] nix::Error), + /// A PID lifecycle operation failed. #[error("PID file error: {0}")] PidFile(String), + /// Another process owns the daemon PID lifecycle. #[error("another instance is already running (pid {0})")] AlreadyRunning(i32), + /// The configured runtime user does not exist. #[error("user '{0}' not found")] UserNotFound(String), + /// The configured runtime group does not exist. #[error("group '{0}' not found")] GroupNotFound(String), + /// Applying the configured runtime identity failed. #[error("failed to set uid/gid: {0}")] PrivilegeDrop(#[source] nix::Error), + /// An underlying filesystem operation failed. #[error("io error: {0}")] Io(#[from] io::Error), } @@ -106,38 +123,28 @@ pub enum DaemonizeResult { /// Returns `DaemonizeResult::Parent` in the original parent (which should exit), /// or `DaemonizeResult::Child` in the final daemon child. pub fn daemonize(working_dir: Option<&Path>) -> Result { - // First fork match unsafe { fork() } { Ok(ForkResult::Parent { .. }) => { - // Parent exits return Ok(DaemonizeResult::Parent); } - Ok(ForkResult::Child) => { - // Child continues - } + Ok(ForkResult::Child) => {} Err(e) => return Err(DaemonError::ForkFailed(e)), } - // Create new session, become session leader setsid().map_err(DaemonError::SetsidFailed)?; // Second fork to ensure we can never acquire a controlling terminal match unsafe { fork() } { Ok(ForkResult::Parent { .. }) => { - // Intermediate parent exits std::process::exit(0); } - Ok(ForkResult::Child) => { - // Final daemon child continues - } + Ok(ForkResult::Child) => {} Err(e) => return Err(DaemonError::ForkFailed(e)), } - // Change working directory let target_dir = working_dir.unwrap_or(Path::new("/")); chdir(target_dir).map_err(DaemonError::ChdirFailed)?; - // Redirect stdin, stdout, stderr to /dev/null redirect_stdio_to_devnull()?; Ok(DaemonizeResult::Child) @@ -156,21 +163,17 @@ fn redirect_stdio_to_devnull() -> Result<(), DaemonError> { // Use libc::dup2 directly for redirecting standard file descriptors // nix 0.31's dup2 requires OwnedFd which doesn't work well with stdio fds unsafe { - // Redirect stdin (fd 0) if libc::dup2(devnull_fd, 0) < 0 { return Err(DaemonError::RedirectFailed(Errno::last())); } - // Redirect stdout (fd 1) if libc::dup2(devnull_fd, 1) < 0 { return Err(DaemonError::RedirectFailed(Errno::last())); } - // Redirect stderr (fd 2) if libc::dup2(devnull_fd, 2) < 0 { return Err(DaemonError::RedirectFailed(Errno::last())); } } - // Close original devnull fd if it's not one of the standard fds if devnull_fd > 2 { let _ = close(devnull_fd); } @@ -178,166 +181,6 @@ fn redirect_stdio_to_devnull() -> Result<(), DaemonError> { Ok(()) } -/// PID file manager with flock-based locking. -pub struct PidFile { - path: PathBuf, - file: Option, - locked: bool, -} - -impl PidFile { - /// Creates a new PID file manager for the given path. - pub fn new>(path: P) -> Self { - Self { - path: path.as_ref().to_path_buf(), - file: None, - locked: false, - } - } - - /// Checks if another instance is already running. - /// - /// Returns the PID of the running instance if one exists. - pub fn check_running(&self) -> Result, DaemonError> { - if !self.path.exists() { - return Ok(None); - } - - // Try to read existing PID - let mut contents = String::new(); - File::open(&self.path) - .and_then(|mut f| f.read_to_string(&mut contents)) - .map_err(|e| { - DaemonError::PidFile(format!("cannot read {}: {}", self.path.display(), e)) - })?; - - let pid: i32 = contents - .trim() - .parse() - .map_err(|_| DaemonError::PidFile(format!("invalid PID in {}", self.path.display())))?; - - // Check if process is still running - if is_process_running(pid) { - Ok(Some(pid)) - } else { - // Stale PID file - debug!(pid, path = %self.path.display(), "Removing stale PID file"); - let _ = fs::remove_file(&self.path); - Ok(None) - } - } - - /// Acquires the PID file lock and writes the current PID. - /// - /// Fails if another instance is already running. - pub fn acquire(&mut self) -> Result<(), DaemonError> { - // Check for running instance first - if let Some(pid) = self.check_running()? { - return Err(DaemonError::AlreadyRunning(pid)); - } - - // Ensure parent directory exists - if let Some(parent) = self.path.parent() { - if !parent.exists() { - fs::create_dir_all(parent).map_err(|e| { - DaemonError::PidFile(format!( - "cannot create directory {}: {}", - parent.display(), - e - )) - })?; - } - } - - // Open/create PID file with exclusive lock - let file = OpenOptions::new() - .write(true) - .create(true) - .truncate(true) - .mode(0o644) - .open(&self.path) - .map_err(|e| { - DaemonError::PidFile(format!("cannot open {}: {}", self.path.display(), e)) - })?; - - // Try to acquire exclusive lock (non-blocking) - let flock = Flock::lock(file, FlockArg::LockExclusiveNonblock).map_err(|(_, errno)| { - // Check if another instance grabbed the lock - if let Some(pid) = self.check_running().ok().flatten() { - DaemonError::AlreadyRunning(pid) - } else { - DaemonError::PidFile(format!("cannot lock {}: {}", self.path.display(), errno)) - } - })?; - - // Write our PID - let pid = getpid(); - let mut file = flock - .unlock() - .map_err(|(_, errno)| DaemonError::PidFile(format!("unlock failed: {}", errno)))?; - - writeln!(file, "{}", pid).map_err(|e| { - DaemonError::PidFile(format!( - "cannot write PID to {}: {}", - self.path.display(), - e - )) - })?; - - // Re-acquire lock and keep it - let flock = Flock::lock(file, FlockArg::LockExclusiveNonblock).map_err(|(_, errno)| { - DaemonError::PidFile(format!("cannot re-lock {}: {}", self.path.display(), errno)) - })?; - - self.file = Some(flock.unlock().map_err(|(_, errno)| { - DaemonError::PidFile(format!("unlock for storage failed: {}", errno)) - })?); - self.locked = true; - - info!(pid = pid.as_raw(), path = %self.path.display(), "PID file created"); - Ok(()) - } - - /// Releases the PID file lock and removes the file. - pub fn release(&mut self) -> Result<(), DaemonError> { - if let Some(file) = self.file.take() { - drop(file); - } - self.locked = false; - - if self.path.exists() { - fs::remove_file(&self.path).map_err(|e| { - DaemonError::PidFile(format!("cannot remove {}: {}", self.path.display(), e)) - })?; - debug!(path = %self.path.display(), "PID file removed"); - } - - Ok(()) - } - - /// Returns the path to this PID file. - #[allow(dead_code)] - pub fn path(&self) -> &Path { - &self.path - } -} - -impl Drop for PidFile { - fn drop(&mut self) { - if self.locked { - if let Err(e) = self.release() { - warn!(error = %e, "Failed to clean up PID file on drop"); - } - } - } -} - -/// Checks if a process with the given PID is running. -fn is_process_running(pid: i32) -> bool { - // kill(pid, 0) checks if process exists without sending a signal - nix::sys::signal::kill(Pid::from_raw(pid), None).is_ok() -} - // macOS gates nix::unistd::setgroups differently in the current dependency set, // so call libc directly there while preserving the original nix path elsewhere. fn set_supplementary_groups(gid: Gid) -> Result<(), nix::Error> { @@ -383,9 +226,11 @@ pub fn drop_privileges( }; if (target_uid.is_some() || target_gid.is_some()) - && let Some(file) = pid_file.and_then(|pid| pid.file.as_ref()) + && let Some(pid_file) = pid_file { - unistd::fchown(file, target_uid, target_gid).map_err(DaemonError::PrivilegeDrop)?; + for file in pid_file.ownership_file_handles().into_iter().flatten() { + unistd::fchown(file, target_uid, target_gid).map_err(DaemonError::PrivilegeDrop)?; + } } if let Some(gid) = target_gid { @@ -401,7 +246,7 @@ pub fn drop_privileges( if uid.as_raw() != 0 && let Some(pid) = pid_file { - let parent = pid.path.parent().unwrap_or(Path::new(".")); + let parent = pid.path().parent().unwrap_or(Path::new(".")); let probe_path = parent.join(format!( ".telemt_pid_probe_{}_{}", std::process::id(), @@ -436,7 +281,6 @@ pub fn drop_privileges( /// Looks up a user by name and returns their UID. fn lookup_user(name: &str) -> Result { - // Use libc getpwnam let c_name = std::ffi::CString::new(name).map_err(|_| DaemonError::UserNotFound(name.to_string()))?; @@ -480,76 +324,6 @@ fn lookup_group(name: &str) -> Result { } } -/// Reads PID from a PID file. -#[allow(dead_code)] -pub fn read_pid_file>(path: P) -> Result { - let path = path.as_ref(); - let mut contents = String::new(); - File::open(path) - .and_then(|mut f| f.read_to_string(&mut contents)) - .map_err(|e| DaemonError::PidFile(format!("cannot read {}: {}", path.display(), e)))?; - - contents - .trim() - .parse() - .map_err(|_| DaemonError::PidFile(format!("invalid PID in {}", path.display()))) -} - -/// Sends a signal to the process specified in a PID file. -#[allow(dead_code)] -pub fn signal_pid_file>( - path: P, - signal: nix::sys::signal::Signal, -) -> Result<(), DaemonError> { - let pid = read_pid_file(&path)?; - - if !is_process_running(pid) { - return Err(DaemonError::PidFile(format!( - "process {} from {} is not running", - pid, - path.as_ref().display() - ))); - } - - nix::sys::signal::kill(Pid::from_raw(pid), signal) - .map_err(|e| DaemonError::PidFile(format!("cannot signal process {}: {}", pid, e)))?; - - Ok(()) -} - -/// Returns the status of the daemon based on PID file. -#[allow(dead_code)] -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum DaemonStatus { - /// Daemon is running with the given PID. - Running(i32), - /// PID file exists but process is not running. - Stale(i32), - /// No PID file exists. - NotRunning, -} - -/// Checks the daemon status from a PID file. -#[allow(dead_code)] -pub fn check_status>(path: P) -> DaemonStatus { - let path = path.as_ref(); - - if !path.exists() { - return DaemonStatus::NotRunning; - } - - match read_pid_file(path) { - Ok(pid) => { - if is_process_running(pid) { - DaemonStatus::Running(pid) - } else { - DaemonStatus::Stale(pid) - } - } - Err(_) => DaemonStatus::NotRunning, - } -} - #[cfg(test)] mod tests { use super::*; @@ -571,29 +345,4 @@ mod tests { }; assert!(!opts.should_daemonize()); } - - #[test] - fn test_check_status_not_running() { - let path = "/tmp/telemt_test_nonexistent.pid"; - assert_eq!(check_status(path), DaemonStatus::NotRunning); - } - - #[test] - fn test_pid_file_basic() { - let path = "/tmp/telemt_test_pidfile.pid"; - let _ = fs::remove_file(path); - - let mut pf = PidFile::new(path); - assert!(pf.check_running().unwrap().is_none()); - - pf.acquire().unwrap(); - assert!(Path::new(path).exists()); - - // Read it back - let pid = read_pid_file(path).unwrap(); - assert_eq!(pid, std::process::id() as i32); - - pf.release().unwrap(); - assert!(!Path::new(path).exists()); - } } diff --git a/src/daemon/pid_file.rs b/src/daemon/pid_file.rs new file mode 100644 index 0000000..0eb70f0 --- /dev/null +++ b/src/daemon/pid_file.rs @@ -0,0 +1,511 @@ +use std::ffi::OsStr; +use std::fs::{self, File}; +use std::io::{self, ErrorKind, Read, Write}; +#[cfg(target_os = "linux")] +use std::os::fd::{FromRawFd, OwnedFd}; +use std::os::unix::fs::{MetadataExt, PermissionsExt}; +use std::path::{Path, PathBuf}; + +use nix::fcntl::{Flock, FlockArg, OFlag, openat}; +use nix::sys::stat::Mode; +use nix::unistd::{Pid, UnlinkatFlags, getpid, unlinkat}; +use tracing::{debug, info, warn}; + +use super::DaemonError; +use crate::util::secure_fs::AnchoredPath; + +/// PID file manager backed by a persistent sibling lock file. +pub struct PidFile { + path: PathBuf, + lock_path: PathBuf, + pid_file: Option, + pid_identity: Option, + lock_file: Option>, + anchor: Option, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +struct FileIdentity { + device: u64, + inode: u64, +} + +impl FileIdentity { + fn from_metadata(metadata: &fs::Metadata) -> Self { + Self { + device: metadata.dev(), + inode: metadata.ino(), + } + } +} + +impl PidFile { + /// Creates a new PID file manager for the given path. + pub fn new>(path: P) -> Self { + let path = normalize_pid_path(path.as_ref()); + let lock_path = sibling_lock_path(&path); + Self { + path, + lock_path, + pid_file: None, + pid_identity: None, + lock_file: None, + anchor: None, + } + } + + /// Checks whether the PID file names a running process without modifying either file. + pub fn check_running(&self) -> Result, DaemonError> { + let Some(pid) = read_pid_file_if_exists(&self.path)? else { + return Ok(None); + }; + Ok(is_process_running(pid).then_some(pid)) + } + + /// Acquires the persistent sibling lock and writes the current PID. + /// + /// Fails if another owner holds the lock or the existing PID names a running process. + pub fn acquire(&mut self) -> Result<(), DaemonError> { + let anchor = + AnchoredPath::open_trusted_parent_or_create(&self.path, 0o755).map_err(|error| { + DaemonError::PidFile(format!( + "cannot open trusted parent for {}: {}", + self.path.display(), + error + )) + })?; + let lock_name = self.lock_path.file_name().ok_or_else(|| { + DaemonError::PidFile(format!( + "lock path {} has no file name", + self.lock_path.display() + )) + })?; + let lock_file = open_file_at( + &anchor, + lock_name, + OFlag::O_RDWR | OFlag::O_CREAT | OFlag::O_CLOEXEC | OFlag::O_NOFOLLOW, + 0o644, + ) + .map_err(|error| { + DaemonError::PidFile(format!( + "cannot open lock file {}: {}", + self.lock_path.display(), + error + )) + })?; + validate_regular_single_link(&lock_file, &self.lock_path)?; + let lock_file = + Flock::lock(lock_file, FlockArg::LockExclusiveNonblock).map_err(|(_, errno)| { + if let Some(pid) = read_pid_file_at(&anchor, &self.path) + .ok() + .flatten() + .filter(|pid| is_process_running(*pid)) + { + DaemonError::AlreadyRunning(pid) + } else { + DaemonError::PidFile(format!( + "cannot lock {}: {}", + self.lock_path.display(), + errno + )) + } + })?; + + if let Some(pid) = read_pid_file_at(&anchor, &self.path)? + && is_process_running(pid) + { + return Err(DaemonError::AlreadyRunning(pid)); + } + + let mut pid_file = open_file_at( + &anchor, + anchor.name(), + OFlag::O_RDWR | OFlag::O_CREAT | OFlag::O_CLOEXEC | OFlag::O_NOFOLLOW, + 0o644, + ) + .map_err(|error| { + DaemonError::PidFile(format!("cannot open {}: {}", self.path.display(), error)) + })?; + let pid_metadata = validate_regular_single_link(&pid_file, &self.path)?; + let pid_identity = FileIdentity::from_metadata(&pid_metadata); + // Validate the opened inode before modifying it so a hard-link substitution + // cannot turn PID publication into truncation of an unrelated file. + pid_file.set_len(0).map_err(|error| { + DaemonError::PidFile(format!( + "cannot truncate {}: {}", + self.path.display(), + error + )) + })?; + let pid = getpid(); + writeln!(pid_file, "{}", pid).map_err(|error| { + DaemonError::PidFile(format!( + "cannot write PID to {}: {}", + self.path.display(), + error + )) + })?; + pid_file.sync_data().map_err(|error| { + DaemonError::PidFile(format!( + "cannot sync PID file {}: {}", + self.path.display(), + error + )) + })?; + + self.pid_file = Some(pid_file); + self.pid_identity = Some(pid_identity); + self.lock_file = Some(lock_file); + self.anchor = Some(anchor); + info!(pid = pid.as_raw(), path = %self.path.display(), "PID file created"); + Ok(()) + } + + /// Removes the PID file while retaining exclusive lock ownership until cleanup completes. + pub fn release(&mut self) -> Result<(), DaemonError> { + if self.lock_file.is_none() { + self.pid_file = None; + self.pid_identity = None; + self.anchor = None; + return Ok(()); + } + + let removal = match self.anchor.as_ref() { + Some(anchor) => remove_owned_pid_file(anchor, &self.path, self.pid_identity), + None => Err(DaemonError::PidFile( + "PID file lock is held without a directory anchor".to_string(), + )), + }; + self.pid_file = None; + self.pid_identity = None; + self.lock_file = None; + self.anchor = None; + removal?; + debug!(path = %self.path.display(), "PID file removed"); + Ok(()) + } + + /// Returns the path to this PID file. + pub fn path(&self) -> &Path { + &self.path + } + + /// Returns open files whose ownership must follow the target runtime identity. + pub(super) fn ownership_file_handles(&self) -> [Option<&File>; 2] { + [self.pid_file.as_ref(), self.lock_file.as_deref()] + } +} + +impl Drop for PidFile { + fn drop(&mut self) { + if self.lock_file.is_some() + && let Err(error) = self.release() + { + warn!(error = %error, "Failed to clean up PID file on drop"); + } + } +} + +fn sibling_lock_path(path: &Path) -> PathBuf { + let mut lock_path = path.as_os_str().to_os_string(); + lock_path.push(".lock"); + lock_path.into() +} + +fn normalize_pid_path(path: &Path) -> PathBuf { + let legacy_run = Path::new("/var/run"); + let Ok(remainder) = path.strip_prefix(legacy_run) else { + return path.to_path_buf(); + }; + let Ok(var_metadata) = fs::metadata("/var") else { + return path.to_path_buf(); + }; + let Ok(link_metadata) = fs::symlink_metadata(legacy_run) else { + return path.to_path_buf(); + }; + let Ok(target) = fs::read_link(legacy_run) else { + return path.to_path_buf(); + }; + let trusted_var = var_metadata.is_dir() + && var_metadata.uid() == 0 + && var_metadata.permissions().mode() & 0o022 == 0; + let trusted_alias = link_metadata.file_type().is_symlink() + && link_metadata.uid() == 0 + && (target == Path::new("/run") || target == Path::new("../run")); + if trusted_var && trusted_alias { + Path::new("/run").join(remainder) + } else { + path.to_path_buf() + } +} + +fn open_file_at(anchor: &AnchoredPath, name: &OsStr, flags: OFlag, mode: u32) -> io::Result { + let descriptor = openat(anchor.parent(), name, flags, Mode::from_bits_truncate(mode)) + .map_err(|error| io::Error::from_raw_os_error(error as i32))?; + Ok(File::from(descriptor)) +} + +fn read_pid_file_if_exists(path: &Path) -> Result, DaemonError> { + let anchor = match AnchoredPath::open_trusted_parent(path) { + Ok(anchor) => anchor, + Err(error) if error.kind() == ErrorKind::NotFound => return Ok(None), + Err(error) => { + return Err(DaemonError::PidFile(format!( + "cannot open trusted parent for {}: {}", + path.display(), + error + ))); + } + }; + read_pid_file_at(&anchor, path) +} + +fn read_pid_file_at(anchor: &AnchoredPath, path: &Path) -> Result, DaemonError> { + let mut file = match open_file_at( + anchor, + anchor.name(), + OFlag::O_RDONLY | OFlag::O_CLOEXEC | OFlag::O_NOFOLLOW, + 0, + ) { + Ok(file) => file, + Err(error) if error.kind() == ErrorKind::NotFound => return Ok(None), + Err(error) => { + return Err(DaemonError::PidFile(format!( + "cannot read {}: {}", + path.display(), + error + ))); + } + }; + let metadata = validate_regular_single_link(&file, path)?; + if metadata.len() > 64 { + return Err(DaemonError::PidFile(format!( + "invalid PID in {}", + path.display() + ))); + } + let mut contents = String::new(); + file.read_to_string(&mut contents).map_err(|error| { + DaemonError::PidFile(format!("cannot read {}: {}", path.display(), error)) + })?; + let pid: i32 = contents + .trim() + .parse() + .map_err(|_| DaemonError::PidFile(format!("invalid PID in {}", path.display())))?; + if pid <= 1 { + return Err(DaemonError::PidFile(format!( + "invalid PID in {}", + path.display() + ))); + } + Ok(Some(pid)) +} + +fn remove_owned_pid_file( + anchor: &AnchoredPath, + path: &Path, + expected: Option, +) -> Result<(), DaemonError> { + let file = match open_file_at( + anchor, + anchor.name(), + OFlag::O_RDONLY | OFlag::O_CLOEXEC | OFlag::O_NOFOLLOW, + 0, + ) { + Ok(file) => file, + Err(error) if error.kind() == ErrorKind::NotFound => return Ok(()), + Err(error) => { + return Err(DaemonError::PidFile(format!( + "cannot inspect {} before removal: {}", + path.display(), + error + ))); + } + }; + let metadata = validate_regular_single_link(&file, path)?; + if expected != Some(FileIdentity::from_metadata(&metadata)) { + return Err(DaemonError::PidFile(format!( + "refusing to remove replaced PID file {}", + path.display() + ))); + } + drop(file); + unlinkat(anchor.parent(), anchor.name(), UnlinkatFlags::NoRemoveDir).map_err(|error| { + DaemonError::PidFile(format!( + "cannot remove {}: {}", + path.display(), + io::Error::from_raw_os_error(error as i32) + )) + }) +} + +fn validate_regular_single_link(file: &File, path: &Path) -> Result { + let metadata = file.metadata().map_err(|error| { + DaemonError::PidFile(format!("cannot inspect {}: {}", path.display(), error)) + })?; + if !metadata.is_file() || metadata.nlink() != 1 { + return Err(DaemonError::PidFile(format!( + "{} must be a regular file with one directory entry", + path.display() + ))); + } + Ok(metadata) +} + +/// Reads a PID from a PID file. +#[allow(dead_code)] +pub fn read_pid_file>(path: P) -> Result { + let path = normalize_pid_path(path.as_ref()); + read_pid_file_if_exists(&path)?.ok_or_else(|| { + DaemonError::PidFile(format!( + "cannot read {}: file does not exist", + path.display() + )) + }) +} + +/// Sends a signal to the process specified in a PID file. +#[allow(dead_code)] +pub fn signal_pid_file>( + path: P, + signal: nix::sys::signal::Signal, +) -> Result<(), DaemonError> { + let path = normalize_pid_path(path.as_ref()); + let pid = read_pid_file(&path)?; + #[cfg(target_os = "linux")] + let pidfd = open_pidfd(pid)?; + if !daemon_lock_is_held(&path)? { + return Err(DaemonError::PidFile(format!( + "refusing to signal unlocked or stale PID file {}", + path.display() + ))); + } + #[cfg(target_os = "linux")] + return signal_pidfd(&pidfd, pid, signal); + #[cfg(not(target_os = "linux"))] + nix::sys::signal::kill(Pid::from_raw(pid), signal) + .map_err(|error| DaemonError::PidFile(format!("cannot signal process {}: {}", pid, error))) +} + +/// Daemon state derived from the PID file. +#[allow(dead_code)] +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum DaemonStatus { + /// Daemon is running with the given PID. + Running(i32), + /// PID file exists but the named process is not running. + Stale(i32), + /// No readable PID file exists. + NotRunning, +} + +/// Checks daemon status without modifying the PID or lock file. +#[allow(dead_code)] +pub fn check_status>(path: P) -> DaemonStatus { + let path = normalize_pid_path(path.as_ref()); + match read_pid_file_if_exists(&path) { + Ok(Some(pid)) if daemon_lock_is_held(&path).unwrap_or(false) && is_process_running(pid) => { + DaemonStatus::Running(pid) + } + Ok(Some(pid)) => DaemonStatus::Stale(pid), + Ok(None) | Err(_) => DaemonStatus::NotRunning, + } +} + +fn daemon_lock_is_held(path: &Path) -> Result { + let lock_path = sibling_lock_path(path); + let anchor = match AnchoredPath::open_trusted_parent(path) { + Ok(anchor) => anchor, + Err(error) if error.kind() == ErrorKind::NotFound => return Ok(false), + Err(error) => { + return Err(DaemonError::PidFile(format!( + "cannot open trusted parent for {}: {}", + path.display(), + error + ))); + } + }; + let lock_name = lock_path.file_name().ok_or_else(|| { + DaemonError::PidFile(format!( + "lock path {} has no file name", + lock_path.display() + )) + })?; + let file = match open_file_at( + &anchor, + lock_name, + OFlag::O_RDWR | OFlag::O_CLOEXEC | OFlag::O_NOFOLLOW, + 0, + ) { + Ok(file) => file, + Err(error) if error.kind() == ErrorKind::NotFound => return Ok(false), + Err(error) => { + return Err(DaemonError::PidFile(format!( + "cannot inspect lock {}: {}", + lock_path.display(), + error + ))); + } + }; + validate_regular_single_link(&file, &lock_path)?; + match Flock::lock(file, FlockArg::LockExclusiveNonblock) { + Ok(_available) => Ok(false), + Err((_file, nix::errno::Errno::EWOULDBLOCK)) => Ok(true), + Err((_file, error)) => Err(DaemonError::PidFile(format!( + "cannot inspect lock ownership for {}: {}", + lock_path.display(), + error + ))), + } +} + +#[cfg(target_os = "linux")] +fn open_pidfd(pid: i32) -> Result { + // SAFETY: `pidfd_open` receives a validated positive PID and no pointer arguments. + let descriptor = unsafe { libc::syscall(libc::SYS_pidfd_open, pid, 0) }; + if descriptor < 0 { + return Err(DaemonError::PidFile(format!( + "cannot open stable process handle for {}: {}", + pid, + std::io::Error::last_os_error() + ))); + } + // SAFETY: a successful `pidfd_open` returns one newly owned descriptor. + Ok(unsafe { OwnedFd::from_raw_fd(descriptor as i32) }) +} + +#[cfg(target_os = "linux")] +fn signal_pidfd( + pidfd: &OwnedFd, + pid: i32, + signal: nix::sys::signal::Signal, +) -> Result<(), DaemonError> { + use std::os::fd::AsRawFd; + + // SAFETY: the pidfd is owned and valid, and both optional pointer arguments are null. + let result = unsafe { + libc::syscall( + libc::SYS_pidfd_send_signal, + pidfd.as_raw_fd(), + signal as libc::c_int, + std::ptr::null::(), + 0, + ) + }; + if result == 0 { + Ok(()) + } else { + Err(DaemonError::PidFile(format!( + "cannot signal process {} through stable handle: {}", + pid, + std::io::Error::last_os_error() + ))) + } +} + +fn is_process_running(pid: i32) -> bool { + nix::sys::signal::kill(Pid::from_raw(pid), None).is_ok() +} + +#[cfg(test)] +mod tests; diff --git a/src/daemon/pid_file/tests.rs b/src/daemon/pid_file/tests.rs new file mode 100644 index 0000000..d4cbed1 --- /dev/null +++ b/src/daemon/pid_file/tests.rs @@ -0,0 +1,283 @@ +use std::os::unix::fs::{MetadataExt, symlink}; +use std::process::{Child, Command, Stdio}; +use std::thread; +use std::time::{Duration, Instant}; + +use super::*; + +const HELPER_PID_PATH: &str = "TELEMT_PID_LOCK_HELPER_PATH"; +const HELPER_READY_PATH: &str = "TELEMT_PID_LOCK_HELPER_READY"; +const HELPER_STOP_PATH: &str = "TELEMT_PID_LOCK_HELPER_STOP"; + +fn wait_for_path(path: &Path, timeout: Duration) -> bool { + let deadline = Instant::now() + timeout; + while Instant::now() < deadline { + if path.exists() { + return true; + } + thread::sleep(Duration::from_millis(10)); + } + false +} + +fn wait_for_child(child: &mut Child, timeout: Duration) -> Option { + let deadline = Instant::now() + timeout; + while Instant::now() < deadline { + if let Some(status) = child.try_wait().unwrap() { + return Some(status); + } + thread::sleep(Duration::from_millis(10)); + } + None +} + +#[test] +fn pid_file_remains_send_and_sync() { + fn assert_send_sync() {} + + assert_send_sync::(); +} + +#[test] +fn system_var_run_alias_keeps_the_default_pid_path_usable() { + let Ok(metadata) = fs::symlink_metadata("/var/run") else { + return; + }; + let Ok(target) = fs::read_link("/var/run") else { + return; + }; + if !metadata.file_type().is_symlink() + || (target != Path::new("/run") && target != Path::new("../run")) + { + return; + } + + let pid_file = PidFile::new("/var/run/telemt.pid"); + + assert_eq!(pid_file.path(), Path::new("/run/telemt.pid")); +} + +#[test] +fn lock_holder_subprocess() { + let Some(pid_path) = std::env::var_os(HELPER_PID_PATH) else { + return; + }; + let ready_path = PathBuf::from(std::env::var_os(HELPER_READY_PATH).unwrap()); + let stop_path = PathBuf::from(std::env::var_os(HELPER_STOP_PATH).unwrap()); + let mut pid_file = PidFile::new(PathBuf::from(pid_path)); + pid_file.acquire().unwrap(); + fs::write(&ready_path, b"ready").unwrap(); + assert!(wait_for_path(&stop_path, Duration::from_secs(10))); + pid_file.release().unwrap(); +} + +#[test] +fn persistent_sibling_lock_serializes_processes_after_pid_unlink() { + let directory = tempfile::tempdir().unwrap(); + let pid_path = directory.path().join("telemt.pid"); + let lock_path = sibling_lock_path(&pid_path); + let ready_path = directory.path().join("ready"); + let stop_path = directory.path().join("stop"); + let mut child = Command::new(std::env::current_exe().unwrap()) + .args([ + "--exact", + "daemon::pid_file::tests::lock_holder_subprocess", + "--nocapture", + ]) + .env(HELPER_PID_PATH, &pid_path) + .env(HELPER_READY_PATH, &ready_path) + .env(HELPER_STOP_PATH, &stop_path) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn() + .unwrap(); + + if !wait_for_path(&ready_path, Duration::from_secs(5)) { + let _ = child.kill(); + let _ = child.wait(); + panic!("PID lock holder did not become ready"); + } + let lock_inode = fs::metadata(&lock_path).unwrap().ino(); + fs::remove_file(&pid_path).unwrap(); + + let mut contender = PidFile::new(&pid_path); + assert!(contender.acquire().is_err()); + + fs::write(&stop_path, b"stop").unwrap(); + let status = wait_for_child(&mut child, Duration::from_secs(5)).unwrap_or_else(|| { + let _ = child.kill(); + child.wait().unwrap() + }); + assert!(status.success()); + assert_eq!(fs::metadata(&lock_path).unwrap().ino(), lock_inode); + + contender.acquire().unwrap(); + assert_eq!(fs::metadata(&lock_path).unwrap().ino(), lock_inode); + contender.release().unwrap(); + assert!(!pid_path.exists()); + assert!(lock_path.exists()); +} + +#[test] +fn stale_pid_checks_are_read_only() { + let directory = tempfile::tempdir().unwrap(); + let pid_path = directory.path().join("telemt.pid"); + fs::write(&pid_path, b"2000000000\n").unwrap(); + let pid_file = PidFile::new(&pid_path); + + assert_eq!(pid_file.check_running().unwrap(), None); + assert_eq!(check_status(&pid_path), DaemonStatus::Stale(2_000_000_000)); + assert!(pid_path.exists()); +} + +#[test] +fn status_requires_live_lock_ownership() { + let directory = tempfile::tempdir().unwrap(); + let pid_path = directory.path().join("telemt.pid"); + fs::write(&pid_path, format!("{}\n", std::process::id())).unwrap(); + assert_eq!( + check_status(&pid_path), + DaemonStatus::Stale(std::process::id() as i32) + ); + + fs::remove_file(&pid_path).unwrap(); + let mut owner = PidFile::new(&pid_path); + owner.acquire().unwrap(); + assert_eq!( + check_status(&pid_path), + DaemonStatus::Running(std::process::id() as i32) + ); + owner.release().unwrap(); +} + +#[test] +fn unowned_release_does_not_remove_pid_file() { + let directory = tempfile::tempdir().unwrap(); + let pid_path = directory.path().join("telemt.pid"); + fs::write(&pid_path, b"2000000000\n").unwrap(); + let mut pid_file = PidFile::new(&pid_path); + + pid_file.release().unwrap(); + + assert!(pid_path.exists()); +} + +#[test] +fn acquire_rejects_pid_symlink_without_truncating_target() { + let directory = tempfile::tempdir().unwrap(); + let pid_path = directory.path().join("telemt.pid"); + let target_path = directory.path().join("target"); + fs::write(&target_path, b"preserve\n").unwrap(); + symlink(&target_path, &pid_path).unwrap(); + let mut pid_file = PidFile::new(&pid_path); + + assert!(pid_file.acquire().is_err()); + assert_eq!(fs::read(&target_path).unwrap(), b"preserve\n"); +} + +#[test] +fn acquire_rejects_pid_hard_link_without_truncating_target() { + let directory = tempfile::tempdir().unwrap(); + let pid_path = directory.path().join("telemt.pid"); + let target_path = directory.path().join("target"); + fs::write(&target_path, b"preserve\n").unwrap(); + fs::hard_link(&target_path, &pid_path).unwrap(); + let mut pid_file = PidFile::new(&pid_path); + + assert!(pid_file.acquire().is_err()); + assert_eq!(fs::read(&target_path).unwrap(), b"preserve\n"); +} + +#[test] +fn acquire_rejects_symlinked_parent_without_publishing_outside() { + let directory = tempfile::tempdir().unwrap(); + let real_parent = directory.path().join("real"); + let linked_parent = directory.path().join("linked"); + fs::create_dir(&real_parent).unwrap(); + symlink(&real_parent, &linked_parent).unwrap(); + let pid_path = linked_parent.join("telemt.pid"); + let mut pid_file = PidFile::new(&pid_path); + + assert!(pid_file.acquire().is_err()); + assert!(!real_parent.join("telemt.pid").exists()); + assert!(!real_parent.join("telemt.pid.lock").exists()); +} + +#[test] +fn release_remains_anchored_after_parent_path_replacement() { + let directory = tempfile::tempdir().unwrap(); + let active_parent = directory.path().join("active"); + let moved_parent = directory.path().join("moved"); + fs::create_dir(&active_parent).unwrap(); + let pid_path = active_parent.join("telemt.pid"); + let mut pid_file = PidFile::new(&pid_path); + pid_file.acquire().unwrap(); + + fs::rename(&active_parent, &moved_parent).unwrap(); + fs::create_dir(&active_parent).unwrap(); + fs::write(active_parent.join("telemt.pid"), b"replacement\n").unwrap(); + + pid_file.release().unwrap(); + + assert!(!moved_parent.join("telemt.pid").exists()); + assert_eq!( + fs::read(active_parent.join("telemt.pid")).unwrap(), + b"replacement\n" + ); +} + +#[test] +fn release_does_not_remove_replacement_path() { + let directory = tempfile::tempdir().unwrap(); + let pid_path = directory.path().join("telemt.pid"); + let owned_path = directory.path().join("owned.pid"); + let mut pid_file = PidFile::new(&pid_path); + pid_file.acquire().unwrap(); + fs::rename(&pid_path, &owned_path).unwrap(); + fs::write(&pid_path, b"replacement\n").unwrap(); + + let error = pid_file.release().unwrap_err(); + + assert!( + error + .to_string() + .contains("refusing to remove replaced PID file") + ); + assert_eq!(fs::read(&pid_path).unwrap(), b"replacement\n"); +} + +#[test] +fn pid_parser_rejects_process_group_values() { + let directory = tempfile::tempdir().unwrap(); + let pid_path = directory.path().join("telemt.pid"); + + for value in ["-1\n", "0\n", "1\n"] { + fs::write(&pid_path, value).unwrap(); + assert!(read_pid_file(&pid_path).is_err()); + } +} + +#[test] +fn pid_file_release_keeps_lock_inode() { + let directory = tempfile::tempdir().unwrap(); + let pid_path = directory.path().join("telemt.pid"); + let lock_path = sibling_lock_path(&pid_path); + let mut pid_file = PidFile::new(&pid_path); + + pid_file.acquire().unwrap(); + assert!( + pid_file + .ownership_file_handles() + .into_iter() + .all(|file| file.is_some()) + ); + assert_eq!(read_pid_file(&pid_path).unwrap(), std::process::id() as i32); + let lock_inode = fs::metadata(&lock_path).unwrap().ino(); + pid_file.release().unwrap(); + + assert!(!pid_path.exists()); + assert!(lock_path.exists()); + pid_file.acquire().unwrap(); + assert_eq!(fs::metadata(&lock_path).unwrap().ino(), lock_inode); + pid_file.release().unwrap(); +} diff --git a/src/ip_tracker.rs b/src/ip_tracker.rs index a54410f..f2a71f6 100644 --- a/src/ip_tracker.rs +++ b/src/ip_tracker.rs @@ -15,6 +15,7 @@ use arc_swap::ArcSwap; use tokio::sync::{Mutex as AsyncMutex, RwLock}; use crate::config::UserMaxUniqueIpsMode; +use crate::proxy::user_admission::UserIncarnation; const CLEANUP_DRAIN_BATCH_LIMIT: usize = 1024; const MAX_ACTIVE_IP_ENTRIES: u64 = 131_072; @@ -32,15 +33,20 @@ mod tests; struct UserIpShard { active_ips: HashMap>, recent_ips: HashMap>, + incarnations: HashMap, } #[derive(Debug, Default)] struct CleanupShard { - queue: Mutex>>, + queue: Mutex, } +type CleanupQueue = HashMap>>; +type CleanupBatch = HashMap<(String, UserIncarnation, IpAddr), usize>; + #[derive(Debug, Clone)] struct UserIpLimitPolicy { + source_generation: u64, max_ips: Arc>, default_max_ips: usize, mode: UserMaxUniqueIpsMode, @@ -50,6 +56,7 @@ struct UserIpLimitPolicy { impl Default for UserIpLimitPolicy { fn default() -> Self { Self { + source_generation: 0, max_ips: Arc::new(HashMap::new()), default_max_ips: 0, mode: UserMaxUniqueIpsMode::ActiveWindow, @@ -68,6 +75,7 @@ pub struct UserIpTracker { recent_cap_rejects: Arc, cleanup_deferred_releases: Arc, limit_policy: Arc>, + policy_update: Arc>, last_compact_epoch_secs: Arc, cleanup_queue_len: Arc, cleanup_shards: Arc>, @@ -119,6 +127,7 @@ impl UserIpTracker { recent_cap_rejects: Arc::new(AtomicU64::new(0)), cleanup_deferred_releases: Arc::new(AtomicU64::new(0)), limit_policy: Arc::new(ArcSwap::from_pointee(UserIpLimitPolicy::default())), + policy_update: Arc::new(Mutex::new(())), last_compact_epoch_secs: Arc::new(AtomicU64::new(0)), cleanup_queue_len: Arc::new(AtomicU64::new(0)), cleanup_shards: Arc::new(cleanup_shards), @@ -194,19 +203,26 @@ impl UserIpTracker { } pub(super) fn pop_one_cleanup( - queue: &mut HashMap>, - ) -> Option<(String, IpAddr, usize)> { + queue: &mut CleanupQueue, + ) -> Option<(String, UserIncarnation, IpAddr, usize)> { let user = queue.keys().next().cloned()?; - let ip = queue.get(&user)?.keys().next().copied()?; - let count = queue.get_mut(&user)?.remove(&ip)?; - let remove_user = queue - .get(&user) - .map(|user_queue| user_queue.is_empty()) - .unwrap_or(false); - if remove_user { + let incarnation = queue.get(&user)?.keys().next().copied()?; + let ip = queue + .get(&user)? + .get(&incarnation)? + .keys() + .next() + .copied()?; + let incarnations = queue.get_mut(&user)?; + let ips = incarnations.get_mut(&incarnation)?; + let count = ips.remove(&ip)?; + if ips.is_empty() { + incarnations.remove(&incarnation); + } + if incarnations.is_empty() { queue.remove(&user); } - Some((user, ip, count)) + Some((user, incarnation, ip, count)) } #[cfg(test)] @@ -222,6 +238,19 @@ impl UserIpTracker { #[cfg(not(test))] pub(super) fn observe_cleanup_poison_for_tests(&self) {} + #[cfg(test)] + pub(crate) async fn hold_user_shard_for_tests( + &self, + user: &str, + entered: tokio::sync::oneshot::Sender<()>, + release: tokio::sync::oneshot::Receiver<()>, + ) { + let shard_idx = Self::shard_idx(user); + let _guard = self.shards[shard_idx].write().await; + let _ = entered.send(()); + let _ = release.await; + } + pub(super) fn now_epoch_secs() -> u64 { std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) diff --git a/src/ip_tracker/admission.rs b/src/ip_tracker/admission.rs index df3bd98..8fb39a1 100644 --- a/src/ip_tracker/admission.rs +++ b/src/ip_tracker/admission.rs @@ -2,47 +2,84 @@ use super::*; impl UserIpTracker { pub async fn set_limit_policy(&self, mode: UserMaxUniqueIpsMode, window_secs: u64) { - self.limit_policy.rcu(|current| { - Arc::new(UserIpLimitPolicy { - mode, - window_secs: window_secs.max(1), - ..(**current).clone() - }) + let _policy_update = self.policy_update.lock().unwrap_or_else(|poisoned| { + self.policy_update.clear_poison(); + poisoned.into_inner() }); + let current = self.limit_policy.load_full(); + self.limit_policy.store(Arc::new(UserIpLimitPolicy { + mode, + window_secs: window_secs.max(1), + ..(*current).clone() + })); } pub async fn set_user_limit(&self, username: &str, max_ips: usize) { - let username = username.to_string(); - self.limit_policy.rcu(|current| { - let mut limits = current.max_ips.as_ref().clone(); - limits.insert(username.clone(), max_ips); - Arc::new(UserIpLimitPolicy { - max_ips: Arc::new(limits), - ..(**current).clone() - }) + let _policy_update = self.policy_update.lock().unwrap_or_else(|poisoned| { + self.policy_update.clear_poison(); + poisoned.into_inner() }); + let current = self.limit_policy.load_full(); + let mut limits = current.max_ips.as_ref().clone(); + limits.insert(username.to_string(), max_ips); + self.limit_policy.store(Arc::new(UserIpLimitPolicy { + max_ips: Arc::new(limits), + ..(*current).clone() + })); } pub async fn remove_user_limit(&self, username: &str) { - self.limit_policy.rcu(|current| { - let mut limits = current.max_ips.as_ref().clone(); - limits.remove(username); - Arc::new(UserIpLimitPolicy { - max_ips: Arc::new(limits), - ..(**current).clone() - }) + let _policy_update = self.policy_update.lock().unwrap_or_else(|poisoned| { + self.policy_update.clear_poison(); + poisoned.into_inner() }); + let current = self.limit_policy.load_full(); + let mut limits = current.max_ips.as_ref().clone(); + limits.remove(username); + self.limit_policy.store(Arc::new(UserIpLimitPolicy { + max_ips: Arc::new(limits), + ..(*current).clone() + })); } pub async fn load_limits(&self, default_limit: usize, limits: &HashMap) { - let limits = Arc::new(limits.clone()); - self.limit_policy.rcu(|current| { - Arc::new(UserIpLimitPolicy { - max_ips: Arc::clone(&limits), - default_max_ips: default_limit, - ..(**current).clone() - }) + let _policy_update = self.policy_update.lock().unwrap_or_else(|poisoned| { + self.policy_update.clear_poison(); + poisoned.into_inner() }); + let current = self.limit_policy.load_full(); + self.limit_policy.store(Arc::new(UserIpLimitPolicy { + max_ips: Arc::new(limits.clone()), + default_max_ips: default_limit, + ..(*current).clone() + })); + } + + /// Atomically publishes one coherent policy from the active runtime generation. + pub(crate) async fn apply_policy_from_source( + &self, + source_generation: u64, + default_limit: usize, + limits: &HashMap, + mode: UserMaxUniqueIpsMode, + window_secs: u64, + ) -> bool { + let _policy_update = self.policy_update.lock().unwrap_or_else(|poisoned| { + self.policy_update.clear_poison(); + poisoned.into_inner() + }); + let current = self.limit_policy.load_full(); + if source_generation < current.source_generation { + return false; + } + self.limit_policy.store(Arc::new(UserIpLimitPolicy { + source_generation, + max_ips: Arc::new(limits.clone()), + default_max_ips: default_limit, + mode, + window_secs: window_secs.max(1), + })); + true } pub(super) fn prune_recent( @@ -59,8 +96,17 @@ impl UserIpTracker { } pub async fn check_and_add(&self, username: &str, ip: IpAddr) -> Result<(), String> { + self.check_and_add_for_incarnation(username, 0, ip).await + } + + /// Reserves an IP slot for one exact user incarnation. + pub(crate) async fn check_and_add_for_incarnation( + &self, + username: &str, + incarnation: UserIncarnation, + ip: IpAddr, + ) -> Result<(), String> { self.drain_cleanup_for_user(username).await; - self.maybe_compact_empty_users().await; let policy = self.limit_policy.load(); let limit = Self::user_limit(&policy, username); let mode = policy.mode; @@ -69,6 +115,30 @@ impl UserIpTracker { let shard_idx = Self::shard_idx(username); let mut shard = self.shards[shard_idx].write().await; + if let Some(current) = shard.incarnations.get(username).copied() { + if current > incarnation { + return Err(format!( + "IP tracker rejected stale user incarnation for '{username}'" + )); + } + if current < incarnation { + let removed_active = shard + .active_ips + .remove(username) + .map(|ips| ips.len()) + .unwrap_or(0); + let removed_recent = shard + .recent_ips + .remove(username) + .map(|ips| ips.len()) + .unwrap_or(0); + Self::decrement_counter(&self.active_entry_count, removed_active); + Self::decrement_counter(&self.recent_entry_count, removed_recent); + shard.incarnations.insert(username.to_string(), incarnation); + } + } else { + shard.incarnations.insert(username.to_string(), incarnation); + } let user_active = shard.active_ips.entry(username.to_string()).or_default(); let active_contains_ip = user_active.contains_key(&ip); let active_len = user_active.len(); @@ -174,9 +244,21 @@ impl UserIpTracker { } pub async fn remove_ip(&self, username: &str, ip: IpAddr) { - self.maybe_compact_empty_users().await; + self.remove_ip_for_incarnation(username, 0, ip).await; + } + + /// Releases an IP slot only from the incarnation that acquired it. + pub(crate) async fn remove_ip_for_incarnation( + &self, + username: &str, + incarnation: UserIncarnation, + ip: IpAddr, + ) { let shard_idx = Self::shard_idx(username); let mut shard = self.shards[shard_idx].write().await; + if shard.incarnations.get(username).copied() != Some(incarnation) { + return; + } let mut removed_active_entries = 0usize; if let Some(user_ips) = shard.active_ips.get_mut(username) { if let Some(count) = user_ips.get_mut(&ip) { diff --git a/src/ip_tracker/cleanup.rs b/src/ip_tracker/cleanup.rs index 86b4f8b..fd5ab0c 100644 --- a/src/ip_tracker/cleanup.rs +++ b/src/ip_tracker/cleanup.rs @@ -1,15 +1,90 @@ use super::*; +struct DetachedCleanupBatch<'a> { + tracker: &'a UserIpTracker, + shard_idx: usize, + entries: CleanupBatch, +} + +impl DetachedCleanupBatch<'_> { + fn is_empty(&self) -> bool { + self.entries.is_empty() + } + + fn entries(&self) -> impl Iterator { + self.entries.iter() + } + + fn commit(mut self) { + let committed = self.entries.len(); + self.entries.clear(); + UserIpTracker::decrement_counter(&self.tracker.cleanup_queue_len, committed); + } +} + +impl Drop for DetachedCleanupBatch<'_> { + fn drop(&mut self) { + if self.entries.is_empty() { + return; + } + + let cleanup_shard = &self.tracker.cleanup_shards[self.shard_idx]; + let mut duplicate_entries = 0usize; + let mut restore = |queue: &mut CleanupQueue| { + for ((user, incarnation, ip), count) in self.entries.drain() { + let queued = queue + .entry(user) + .or_default() + .entry(incarnation) + .or_default() + .entry(ip) + .or_insert(0); + if *queued != 0 { + duplicate_entries = duplicate_entries.saturating_add(1); + } + *queued = queued.saturating_add(count); + } + }; + match cleanup_shard.queue.lock() { + Ok(mut queue) => restore(&mut queue), + Err(poisoned) => { + let mut queue = poisoned.into_inner(); + restore(&mut queue); + cleanup_shard.queue.clear_poison(); + tracing::warn!( + "UserIpTracker cleanup_queue lock poisoned while restoring a cancelled cleanup batch" + ); + } + } + UserIpTracker::decrement_counter(&self.tracker.cleanup_queue_len, duplicate_entries); + } +} + impl UserIpTracker { /// Queues a deferred active IP cleanup for a later async drain. pub fn enqueue_cleanup(&self, user: String, ip: IpAddr) { + self.enqueue_cleanup_for_incarnation(user, 0, ip); + } + + /// Queues cleanup for the exact user incarnation that owns the reservation. + pub(crate) fn enqueue_cleanup_for_incarnation( + &self, + user: String, + incarnation: UserIncarnation, + ip: IpAddr, + ) { self.observe_cleanup_poison_for_tests(); let shard_idx = Self::shard_idx(&user); let cleanup_shard = &self.cleanup_shards[shard_idx]; match cleanup_shard.queue.lock() { Ok(mut queue) => { - let user_queue = queue.entry(user).or_default(); - let count = user_queue.entry(ip).or_insert(0); + let count = queue + .entry(user) + .or_default() + .entry(incarnation) + .or_default() + .entry(ip) + .or_insert(0); if *count == 0 { self.cleanup_queue_len.fetch_add(1, Ordering::Relaxed); } @@ -19,8 +94,13 @@ impl UserIpTracker { } Err(poisoned) => { let mut queue = poisoned.into_inner(); - let user_queue = queue.entry(user.clone()).or_default(); - let count = user_queue.entry(ip).or_insert(0); + let count = queue + .entry(user.clone()) + .or_default() + .entry(incarnation) + .or_default() + .entry(ip) + .or_insert(0); if *count == 0 { self.cleanup_queue_len.fetch_add(1, Ordering::Relaxed); } @@ -42,6 +122,26 @@ impl UserIpTracker { self.cleanup_queue_len.load(Ordering::Relaxed) as usize } + #[cfg(test)] + pub(crate) fn cleanup_queue_physical_entries_for_tests(&self) -> usize { + self.cleanup_shards + .iter() + .map(|cleanup_shard| { + let count = |queue: &CleanupQueue| { + queue + .values() + .flat_map(HashMap::values) + .map(HashMap::len) + .sum::() + }; + match cleanup_shard.queue.lock() { + Ok(queue) => count(&queue), + Err(poisoned) => count(&poisoned.into_inner()), + } + }) + .sum() + } + #[cfg(test)] pub(crate) fn cleanup_queue_mutex_for_tests( &self, @@ -63,12 +163,13 @@ impl UserIpTracker { return; } let shard_idx = Self::shard_idx(user); + let _drain_guard = self.cleanup_drain_locks[shard_idx].lock().await; let cleanup_shard = &self.cleanup_shards[shard_idx]; let to_remove = match cleanup_shard.queue.lock() { - Ok(mut queue) => queue.remove(user).unwrap_or_default(), + Ok(mut queue) => detach_user_cleanup(self, shard_idx, &mut queue, user), Err(poisoned) => { let mut queue = poisoned.into_inner(); - let drained = queue.remove(user).unwrap_or_default(); + let drained = detach_user_cleanup(self, shard_idx, &mut queue, user); cleanup_shard.queue.clear_poison(); drained } @@ -76,16 +177,19 @@ impl UserIpTracker { if to_remove.is_empty() { return; } - self.cleanup_queue_len - .fetch_sub(to_remove.len() as u64, Ordering::Relaxed); let mut shard = self.shards[shard_idx].write().await; let mut removed_active_entries = 0usize; - for (ip, pending_count) in to_remove { + for ((queued_user, incarnation, ip), pending_count) in to_remove.entries() { + if shard.incarnations.get(queued_user).copied() != Some(*incarnation) { + continue; + } removed_active_entries = removed_active_entries.saturating_add( - Self::apply_active_cleanup(&mut shard.active_ips, user, ip, pending_count), + Self::apply_active_cleanup(&mut shard.active_ips, queued_user, *ip, *pending_count), ); } Self::decrement_counter(&self.active_entry_count, removed_active_entries); + drop(shard); + to_remove.commit(); } pub(super) async fn drain_cleanup_shard(&self, shard_idx: usize) { @@ -100,16 +204,20 @@ impl UserIpTracker { if queue.is_empty() { return; } - let mut drained = - HashMap::with_capacity(queue.len().min(CLEANUP_DRAIN_BATCH_LIMIT)); + let mut drained = HashMap::with_capacity(CLEANUP_DRAIN_BATCH_LIMIT); for _ in 0..CLEANUP_DRAIN_BATCH_LIMIT { - let Some((user, ip, count)) = Self::pop_one_cleanup(&mut queue) else { + let Some((user, incarnation, ip, count)) = + Self::pop_one_cleanup(&mut queue) + else { break; }; - self.cleanup_queue_len.fetch_sub(1, Ordering::Relaxed); - drained.insert((user, ip), count); + drained.insert((user, incarnation, ip), count); + } + DetachedCleanupBatch { + tracker: self, + shard_idx, + entries: drained, } - drained } Err(poisoned) => { let mut queue = poisoned.into_inner(); @@ -117,32 +225,61 @@ impl UserIpTracker { cleanup_shard.queue.clear_poison(); return; } - let mut drained = - HashMap::with_capacity(queue.len().min(CLEANUP_DRAIN_BATCH_LIMIT)); + let mut drained = HashMap::with_capacity(CLEANUP_DRAIN_BATCH_LIMIT); for _ in 0..CLEANUP_DRAIN_BATCH_LIMIT { - let Some((user, ip, count)) = Self::pop_one_cleanup(&mut queue) else { + let Some((user, incarnation, ip, count)) = + Self::pop_one_cleanup(&mut queue) + else { break; }; - self.cleanup_queue_len.fetch_sub(1, Ordering::Relaxed); - drained.insert((user, ip), count); + drained.insert((user, incarnation, ip), count); } cleanup_shard.queue.clear_poison(); - drained + DetachedCleanupBatch { + tracker: self, + shard_idx, + entries: drained, + } } } }; - drop(_drain_guard); if to_remove.is_empty() { return; } let mut shard = self.shards[shard_idx].write().await; let mut removed_active_entries = 0usize; - for ((user, ip), pending_count) in to_remove { + for ((user, incarnation, ip), pending_count) in to_remove.entries() { + if shard.incarnations.get(user).copied() != Some(*incarnation) { + continue; + } removed_active_entries = removed_active_entries.saturating_add( - Self::apply_active_cleanup(&mut shard.active_ips, &user, ip, pending_count), + Self::apply_active_cleanup(&mut shard.active_ips, user, *ip, *pending_count), ); } Self::decrement_counter(&self.active_entry_count, removed_active_entries); + drop(shard); + to_remove.commit(); + } +} + +fn detach_user_cleanup<'a>( + tracker: &'a UserIpTracker, + shard_idx: usize, + queue: &mut CleanupQueue, + user: &str, +) -> DetachedCleanupBatch<'a> { + let mut entries = CleanupBatch::new(); + if let Some(incarnations) = queue.remove(user) { + for (incarnation, ips) in incarnations { + for (ip, count) in ips { + entries.insert((user.to_string(), incarnation, ip), count); + } + } + } + DetachedCleanupBatch { + tracker, + shard_idx, + entries, } } diff --git a/src/ip_tracker/snapshot.rs b/src/ip_tracker/snapshot.rs index f696c4e..7e89ee1 100644 --- a/src/ip_tracker/snapshot.rs +++ b/src/ip_tracker/snapshot.rs @@ -68,10 +68,13 @@ impl UserIpTracker { } } - pub async fn run_periodic_maintenance(self: Arc) { + pub async fn run_periodic_maintenance(self: Arc, source_generation: u64) { let mut interval = tokio::time::interval(Duration::from_secs(1)); loop { interval.tick().await; + if self.limit_policy.load().source_generation != source_generation { + continue; + } self.drain_cleanup_queue().await; self.maybe_compact_empty_users().await; } @@ -228,8 +231,25 @@ impl UserIpTracker { } pub async fn clear_user_ips(&self, username: &str) { + self.clear_user_ips_if_not_newer(username, 0).await; + } + + /// Clears state while advancing the username fence to a newer incarnation. + pub(crate) async fn clear_user_ips_if_not_newer( + &self, + username: &str, + incarnation: UserIncarnation, + ) { let shard_idx = Self::shard_idx(username); let mut shard = self.shards[shard_idx].write().await; + if shard + .incarnations + .get(username) + .is_some_and(|current| *current > incarnation) + { + return; + } + shard.incarnations.insert(username.to_string(), incarnation); let removed_active_entries = shard .active_ips .remove(username) @@ -246,23 +266,36 @@ impl UserIpTracker { } pub async fn clear_all(&self) { + let mut cleanup_drain_guards = Vec::with_capacity(USER_IP_TRACKER_SHARDS); + for drain_lock in self.cleanup_drain_locks.iter() { + cleanup_drain_guards.push(drain_lock.lock().await); + } for shard_lock in self.shards.iter() { let mut shard = shard_lock.write().await; shard.active_ips.clear(); shard.recent_ips.clear(); + shard.incarnations.clear(); } self.active_entry_count.store(0, Ordering::Relaxed); self.recent_entry_count.store(0, Ordering::Relaxed); + let mut cleanup_queue_guards = Vec::with_capacity(USER_IP_TRACKER_SHARDS); for cleanup_shard in self.cleanup_shards.iter() { - match cleanup_shard.queue.lock() { - Ok(mut queue) => queue.clear(), + let queue = match cleanup_shard.queue.lock() { + Ok(queue) => queue, Err(poisoned) => { - poisoned.into_inner().clear(); + let queue = poisoned.into_inner(); cleanup_shard.queue.clear_poison(); + queue } - } + }; + cleanup_queue_guards.push(queue); + } + for queue in cleanup_queue_guards.iter_mut() { + queue.clear(); } self.cleanup_queue_len.store(0, Ordering::Relaxed); + drop(cleanup_queue_guards); + drop(cleanup_drain_guards); } pub async fn is_ip_active(&self, username: &str, ip: IpAddr) -> bool { diff --git a/src/ip_tracker/tests.rs b/src/ip_tracker/tests.rs index bf2925b..913aeaf 100644 --- a/src/ip_tracker/tests.rs +++ b/src/ip_tracker/tests.rs @@ -3,6 +3,8 @@ use std::net::{IpAddr, Ipv4Addr, Ipv6Addr}; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; +mod cleanup_invariants; + fn test_ipv4(oct1: u8, oct2: u8, oct3: u8, oct4: u8) -> IpAddr { IpAddr::V4(Ipv4Addr::new(oct1, oct2, oct3, oct4)) } @@ -189,6 +191,37 @@ async fn test_clear_user_ips() { assert_eq!(tracker.get_active_ip_count("test_user").await, 0); } +#[tokio::test] +async fn stale_incarnation_cleanup_cannot_release_recreated_user_ip() { + let tracker = UserIpTracker::new(); + tracker.set_user_limit("test_user", 1).await; + let old_ip = test_ipv4(192, 168, 2, 1); + let current_ip = test_ipv4(192, 168, 2, 2); + let rejected_ip = test_ipv4(192, 168, 2, 3); + + tracker + .check_and_add_for_incarnation("test_user", 1, old_ip) + .await + .unwrap(); + tracker.clear_user_ips_if_not_newer("test_user", 2).await; + tracker + .check_and_add_for_incarnation("test_user", 3, current_ip) + .await + .unwrap(); + + tracker + .remove_ip_for_incarnation("test_user", 1, old_ip) + .await; + + assert!(tracker.is_ip_active("test_user", current_ip).await); + assert!( + tracker + .check_and_add_for_incarnation("test_user", 3, rejected_ip) + .await + .is_err() + ); +} + #[tokio::test] async fn test_is_ip_active() { let tracker = UserIpTracker::new(); @@ -233,6 +266,64 @@ async fn test_load_limits_replaces_previous_map() { assert_eq!(tracker.get_user_limit("user2").await, Some(5)); } +#[tokio::test] +async fn stale_runtime_cannot_overwrite_newer_ip_policy() { + let tracker = UserIpTracker::new(); + let mut newer = HashMap::new(); + newer.insert("alice".to_string(), 5); + assert!( + tracker + .apply_policy_from_source(2, 7, &newer, UserMaxUniqueIpsMode::Combined, 90,) + .await + ); + + let mut stale = HashMap::new(); + stale.insert("alice".to_string(), 1); + assert!( + !tracker + .apply_policy_from_source(1, 1, &stale, UserMaxUniqueIpsMode::ActiveWindow, 1,) + .await + ); + + let policy = tracker.limit_policy.load_full(); + assert_eq!(policy.source_generation, 2); + assert_eq!(policy.default_max_ips, 7); + assert_eq!(policy.max_ips["alice"], 5); + assert_eq!(policy.mode, UserMaxUniqueIpsMode::Combined); + assert_eq!(policy.window_secs, 90); +} + +#[tokio::test] +async fn active_runtime_can_publish_coherent_same_generation_ip_policy() { + let tracker = UserIpTracker::new(); + assert!( + tracker + .apply_policy_from_source( + 3, + 1, + &HashMap::new(), + UserMaxUniqueIpsMode::ActiveWindow, + 10, + ) + .await + ); + let mut limits = HashMap::new(); + limits.insert("alice".to_string(), 4); + + assert!( + tracker + .apply_policy_from_source(3, 6, &limits, UserMaxUniqueIpsMode::TimeWindow, 30,) + .await + ); + + let policy = tracker.limit_policy.load_full(); + assert_eq!(policy.source_generation, 3); + assert_eq!(policy.default_max_ips, 6); + assert_eq!(policy.max_ips["alice"], 4); + assert_eq!(policy.mode, UserMaxUniqueIpsMode::TimeWindow); + assert_eq!(policy.window_secs, 30); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn concurrent_policy_replacement_never_exposes_partial_limit_map() { const USER_COUNT: usize = 4_096; @@ -404,10 +495,7 @@ async fn test_compact_prunes_stale_recent_entries() { } tracker.last_compact_epoch_secs.store(0, Ordering::Relaxed); - tracker - .check_and_add("trigger-user", test_ipv4(10, 3, 0, 2)) - .await - .unwrap(); + tracker.maybe_compact_empty_users().await; let shard_idx = UserIpTracker::shard_idx(&stale_user); let shard = tracker.shards[shard_idx].read().await; diff --git a/src/ip_tracker/tests/cleanup_invariants.rs b/src/ip_tracker/tests/cleanup_invariants.rs new file mode 100644 index 0000000..830fa60 --- /dev/null +++ b/src/ip_tracker/tests/cleanup_invariants.rs @@ -0,0 +1,109 @@ +use super::super::*; +use std::net::{IpAddr, Ipv4Addr}; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +fn test_ipv4(oct1: u8, oct2: u8, oct3: u8, oct4: u8) -> IpAddr { + IpAddr::V4(Ipv4Addr::new(oct1, oct2, oct3, oct4)) +} + +#[tokio::test] +async fn cancelled_cleanup_drain_restores_detached_batch() { + let tracker = Arc::new(UserIpTracker::new()); + let user = "cancelled-cleanup-user"; + let ip = test_ipv4(10, 2, 1, 1); + tracker.check_and_add(user, ip).await.unwrap(); + tracker.enqueue_cleanup(user.to_string(), ip); + + let (entered_tx, entered_rx) = tokio::sync::oneshot::channel(); + let (release_tx, release_rx) = tokio::sync::oneshot::channel(); + let held_tracker = Arc::clone(&tracker); + let held_user = user.to_string(); + let holder = tokio::spawn(async move { + held_tracker + .hold_user_shard_for_tests(&held_user, entered_tx, release_rx) + .await; + }); + entered_rx.await.unwrap(); + + let shard_idx = UserIpTracker::shard_idx(user); + let drain_tracker = Arc::clone(&tracker); + let drain = tokio::spawn(async move { + drain_tracker.drain_cleanup_shard(shard_idx).await; + }); + tokio::time::timeout(Duration::from_secs(1), async { + while tracker.cleanup_queue_physical_entries_for_tests() != 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("cleanup batch must detach before waiting for the IP shard"); + + drain.abort(); + assert!(drain.await.unwrap_err().is_cancelled()); + assert_eq!(tracker.cleanup_queue_len_for_tests(), 1); + assert_eq!(tracker.cleanup_queue_physical_entries_for_tests(), 1); + + let _ = release_tx.send(()); + holder.await.unwrap(); + tracker.drain_cleanup_queue().await; + assert_eq!(tracker.cleanup_queue_len_for_tests(), 0); + assert_eq!(tracker.get_active_ip_count(user).await, 0); +} + +#[test] +fn clear_all_serializes_queue_reset_with_concurrent_enqueue() { + let tracker = Arc::new(UserIpTracker::new()); + let first_shard_user = (0u64..) + .map(|index| format!("clear-race-{index}")) + .find(|user| UserIpTracker::shard_idx(user) == 0) + .unwrap(); + let last_shard = USER_IP_TRACKER_SHARDS - 1; + let last_queue_guard = tracker.cleanup_shards[last_shard].queue.lock().unwrap(); + + let clear_tracker = Arc::clone(&tracker); + let clear = std::thread::spawn(move || { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + runtime.block_on(clear_tracker.clear_all()); + }); + + let wait_deadline = Instant::now() + Duration::from_secs(1); + loop { + match tracker.cleanup_shards[0].queue.try_lock() { + Ok(queue) => drop(queue), + Err(std::sync::TryLockError::WouldBlock) => break, + Err(std::sync::TryLockError::Poisoned(_)) => panic!("cleanup queue lock poisoned"), + } + assert!( + Instant::now() < wait_deadline, + "clear_all did not reach queue reset" + ); + std::thread::yield_now(); + } + + let (started_tx, started_rx) = std::sync::mpsc::channel(); + let (completed_tx, completed_rx) = std::sync::mpsc::channel(); + let enqueue_tracker = Arc::clone(&tracker); + let enqueue = std::thread::spawn(move || { + started_tx.send(()).unwrap(); + enqueue_tracker.enqueue_cleanup(first_shard_user, test_ipv4(10, 2, 2, 1)); + completed_tx.send(()).unwrap(); + }); + started_rx.recv().unwrap(); + assert!( + completed_rx + .recv_timeout(Duration::from_millis(50)) + .is_err() + ); + + drop(last_queue_guard); + clear.join().unwrap(); + completed_rx.recv_timeout(Duration::from_secs(1)).unwrap(); + enqueue.join().unwrap(); + + assert_eq!(tracker.cleanup_queue_len_for_tests(), 1); + assert_eq!(tracker.cleanup_queue_physical_entries_for_tests(), 1); +} diff --git a/src/logging.rs b/src/logging.rs index 08f1131..9567fa1 100644 --- a/src/logging.rs +++ b/src/logging.rs @@ -8,8 +8,6 @@ // Infrastructure module used via CLI flags. #![allow(dead_code)] -use std::path::Path; - use crate::config::{LogRotation, LoggingConfig, LoggingDestination}; use tracing_subscriber::layer::SubscriberExt; @@ -144,31 +142,9 @@ pub fn init_logging( } LogDestination::File { options } => { - let (non_blocking, guard) = if options.max_size_bytes > 0 - || options.max_files > 0 - || options.max_age_secs > 0 - { - let file_appender = file::BoundedFileAppender::new(options.clone()) - .expect("Failed to open log file"); - tracing_appender::non_blocking(file_appender) - } else if !matches!(options.rotation, LogRotation::Never) { - let path = Path::new(&options.path); - let dir = log_file_dir(path); - let prefix = log_file_name(path); - let file_appender = tracing_appender::rolling::RollingFileAppender::builder() - .rotation(to_tracing_rotation(options.rotation)) - .filename_prefix(prefix) - .build(dir) - .expect("Failed to open log file"); - tracing_appender::non_blocking(file_appender) - } else { - let file = std::fs::OpenOptions::new() - .create(true) - .append(true) - .open(&options.path) - .expect("Failed to open log file"); - tracing_appender::non_blocking(file) - }; + let file_appender = + file::BoundedFileAppender::new(options.clone()).expect("Failed to open log file"); + let (non_blocking, guard) = tracing_appender::non_blocking(file_appender); let fmt_layer = fmt::Layer::default() .with_ansi(false) @@ -185,28 +161,6 @@ pub fn init_logging( } } -fn log_file_dir(path: &Path) -> &Path { - path.parent() - .filter(|parent| !parent.as_os_str().is_empty()) - .unwrap_or_else(|| Path::new(".")) -} - -fn log_file_name(path: &Path) -> &str { - path.file_name() - .and_then(|s| s.to_str()) - .unwrap_or("telemt") -} - -fn to_tracing_rotation(rotation: LogRotation) -> tracing_appender::rolling::Rotation { - match rotation { - LogRotation::Never => tracing_appender::rolling::Rotation::NEVER, - LogRotation::Minutely => tracing_appender::rolling::Rotation::MINUTELY, - LogRotation::Hourly => tracing_appender::rolling::Rotation::HOURLY, - LogRotation::Daily => tracing_appender::rolling::Rotation::DAILY, - LogRotation::Weekly => tracing_appender::rolling::Rotation::WEEKLY, - } -} - /// Syslog writer for tracing. #[cfg(unix)] #[derive(Clone, Copy)] diff --git a/src/logging/file.rs b/src/logging/file.rs index 3b96903..f14c24f 100644 --- a/src/logging/file.rs +++ b/src/logging/file.rs @@ -1,8 +1,26 @@ -use std::fs::{self, File, OpenOptions}; +#[cfg(not(unix))] +use std::fs::OpenOptions; +use std::fs::{self, File}; use std::io::{self, Write}; use std::path::{Path, PathBuf}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; +#[cfg(unix)] +use std::ffi::OsString; +#[cfg(unix)] +use std::os::fd::OwnedFd; +#[cfg(unix)] +use std::os::unix::ffi::OsStringExt; + +#[cfg(unix)] +use nix::dir::Dir; +#[cfg(unix)] +use nix::fcntl::{OFlag, openat, renameat}; +#[cfg(unix)] +use nix::sys::stat::Mode; +#[cfg(unix)] +use nix::unistd::{UnlinkatFlags, dup, unlinkat}; + use chrono::{DateTime, Datelike, Duration as ChronoDuration, Utc}; use crate::config::LogRotation; @@ -20,6 +38,8 @@ pub(crate) struct BoundedFileAppender { current_size: u64, last_cleanup: DateTime, file: Option, + #[cfg(unix)] + dir_fd: OwnedFd, now: Box DateTime + Send + Sync>, } @@ -46,6 +66,11 @@ impl BoundedFileAppender { let start = now(); let current_path = active_path_for(&dir, &base_name, options.rotation, &start); + #[cfg(unix)] + let dir_fd = crate::util::secure_fs::open_trusted_dir_nofollow_or_create(&dir, 0o750)?; + #[cfg(unix)] + let (file, current_size) = open_append_file(&dir_fd, ¤t_path)?; + #[cfg(not(unix))] let (file, current_size) = open_append_file(¤t_path)?; let mut appender = Self { options, @@ -55,6 +80,8 @@ impl BoundedFileAppender { current_size, last_cleanup: start, file: Some(file), + #[cfg(unix)] + dir_fd, now, }; appender.cleanup(&start); @@ -79,10 +106,8 @@ impl BoundedFileAppender { fn rotate_for_size(&mut self, now: &DateTime) -> io::Result<()> { self.close_current()?; - if self.current_path.exists() { - let archive_path = self.archive_path(now); - fs::rename(&self.current_path, archive_path)?; - } + let archive_path = self.archive_path(now); + self.rename_current_if_present(&archive_path)?; self.open_current() } @@ -95,7 +120,7 @@ impl BoundedFileAppender { let stamp = now.format("%Y%m%d%H%M%S"); for seq in 0..1000 { let candidate = self.dir.join(format!("{file_name}.{stamp}.{seq}")); - if !candidate.exists() { + if !self.path_exists(&candidate) { return candidate; } } @@ -103,6 +128,9 @@ impl BoundedFileAppender { } fn open_current(&mut self) -> io::Result<()> { + #[cfg(unix)] + let (file, current_size) = open_append_file(&self.dir_fd, &self.current_path)?; + #[cfg(not(unix))] let (file, current_size) = open_append_file(&self.current_path)?; self.file = Some(file); self.current_size = current_size; @@ -130,40 +158,10 @@ impl BoundedFileAppender { fn cleanup(&mut self, now: &DateTime) { self.last_cleanup = now.clone(); - let Ok(entries) = fs::read_dir(&self.dir) else { + let Ok(mut candidates) = self.collect_candidates() else { return; }; - let mut candidates = Vec::new(); - let prefix = format!("{}.", self.base_name); - for entry in entries.flatten() { - let path = entry.path(); - let Ok(file_type) = entry.file_type() else { - continue; - }; - if !file_type.is_file() { - continue; - } - - let is_current = path == self.current_path; - let Some(name) = entry.file_name().to_str().map(|name| name.to_string()) else { - continue; - }; - if !is_current && !name.starts_with(&prefix) { - continue; - } - - let Ok(metadata) = entry.metadata() else { - continue; - }; - let modified = metadata.modified().unwrap_or(UNIX_EPOCH); - candidates.push(LogFileCandidate { - path, - modified, - is_current, - }); - } - if self.options.max_age_secs > 0 { let cutoff = system_time_from_utc(now) .checked_sub(Duration::from_secs(self.options.max_age_secs)) @@ -172,7 +170,7 @@ impl BoundedFileAppender { if candidate.is_current || candidate.modified >= cutoff { true } else { - let _ = fs::remove_file(&candidate.path); + self.remove_candidate(candidate); false } }); @@ -189,11 +187,147 @@ impl BoundedFileAppender { if total <= self.options.max_files { break; } - let _ = fs::remove_file(candidate.path); + self.remove_candidate(&candidate); total -= 1; } } } + + #[cfg(unix)] + fn path_exists(&self, path: &Path) -> bool { + let Some(name) = path.file_name() else { + return true; + }; + match openat( + &self.dir_fd, + name, + OFlag::O_RDONLY | OFlag::O_NONBLOCK | OFlag::O_NOFOLLOW | OFlag::O_CLOEXEC, + Mode::empty(), + ) { + Ok(_) => true, + Err(nix::errno::Errno::ENOENT) => false, + Err(_) => true, + } + } + + #[cfg(not(unix))] + fn path_exists(&self, path: &Path) -> bool { + path.exists() + } + + #[cfg(unix)] + fn rename_current_if_present(&self, archive_path: &Path) -> io::Result<()> { + let current_name = self.current_path.file_name().ok_or_else(|| { + io::Error::new(io::ErrorKind::InvalidInput, "log path has no file name") + })?; + let archive_name = archive_path.file_name().ok_or_else(|| { + io::Error::new(io::ErrorKind::InvalidInput, "archive path has no file name") + })?; + match renameat(&self.dir_fd, current_name, &self.dir_fd, archive_name) { + Ok(()) => Ok(()), + Err(nix::errno::Errno::ENOENT) => Ok(()), + Err(error) => Err(io::Error::from_raw_os_error(error as i32)), + } + } + + #[cfg(not(unix))] + fn rename_current_if_present(&self, archive_path: &Path) -> io::Result<()> { + if self.current_path.exists() { + fs::rename(&self.current_path, archive_path)?; + } + Ok(()) + } + + #[cfg(unix)] + fn collect_candidates(&self) -> io::Result> { + use std::os::unix::fs::MetadataExt; + + let descriptor = + dup(&self.dir_fd).map_err(|error| io::Error::from_raw_os_error(error as i32))?; + let mut directory = + Dir::from_fd(descriptor).map_err(|error| io::Error::from_raw_os_error(error as i32))?; + let mut candidates = Vec::new(); + let prefix = format!("{}.", self.base_name); + for entry in directory.iter().flatten() { + let bytes = entry.file_name().to_bytes(); + if bytes == b"." || bytes == b".." { + continue; + } + let name = OsString::from_vec(bytes.to_vec()); + let path = self.dir.join(&name); + let is_current = path == self.current_path; + let Some(name_text) = name.to_str() else { + continue; + }; + if !is_current && !name_text.starts_with(&prefix) { + continue; + } + let Ok(descriptor) = openat( + &self.dir_fd, + name.as_os_str(), + OFlag::O_RDONLY | OFlag::O_NONBLOCK | OFlag::O_NOFOLLOW | OFlag::O_CLOEXEC, + Mode::empty(), + ) else { + continue; + }; + let file = File::from(descriptor); + let Ok(metadata) = file.metadata() else { + continue; + }; + if !metadata.is_file() || metadata.nlink() != 1 { + continue; + } + candidates.push(LogFileCandidate { + path, + modified: metadata.modified().unwrap_or(UNIX_EPOCH), + is_current, + }); + } + Ok(candidates) + } + + #[cfg(not(unix))] + fn collect_candidates(&self) -> io::Result> { + let mut candidates = Vec::new(); + let prefix = format!("{}.", self.base_name); + for entry in fs::read_dir(&self.dir)?.flatten() { + let path = entry.path(); + let Ok(file_type) = entry.file_type() else { + continue; + }; + if !file_type.is_file() { + continue; + } + let is_current = path == self.current_path; + let Some(name) = entry.file_name().to_str().map(str::to_string) else { + continue; + }; + if !is_current && !name.starts_with(&prefix) { + continue; + } + let Ok(metadata) = entry.metadata() else { + continue; + }; + candidates.push(LogFileCandidate { + path, + modified: metadata.modified().unwrap_or(UNIX_EPOCH), + is_current, + }); + } + Ok(candidates) + } + + #[cfg(unix)] + fn remove_candidate(&self, candidate: &LogFileCandidate) { + if let Some(name) = candidate.path.file_name() { + let _ = unlinkat(&self.dir_fd, name, UnlinkatFlags::NoRemoveDir); + } + } + + #[cfg(not(unix))] + fn remove_candidate(&self, candidate: &LogFileCandidate) { + let _ = fs::remove_file(&candidate.path); + } } impl Write for BoundedFileAppender { @@ -233,6 +367,17 @@ struct LogFileCandidate { is_current: bool, } +#[cfg(unix)] +fn open_append_file(dir_fd: &OwnedFd, path: &Path) -> io::Result<(File, u64)> { + let name = path + .file_name() + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "log path has no file name"))?; + let file = crate::util::secure_fs::open_append_regular_at(dir_fd, name, 0o640)?; + let current_size = file.metadata()?.len(); + Ok((file, current_size)) +} + +#[cfg(not(unix))] fn open_append_file(path: &Path) -> io::Result<(File, u64)> { let mut options = OpenOptions::new(); options.create(true).append(true); @@ -291,105 +436,4 @@ fn system_time_from_utc(now: &DateTime) -> SystemTime { } #[cfg(test)] -mod tests { - use std::io::Write; - - use tempfile::tempdir; - - use super::*; - - fn fixed_now() -> DateTime { - DateTime::::from(UNIX_EPOCH + Duration::from_secs(10)) - } - - fn options(path: PathBuf) -> FileLogOptions { - FileLogOptions { - path: path.to_string_lossy().to_string(), - rotation: LogRotation::Never, - max_size_bytes: 0, - max_files: 0, - max_age_secs: 0, - } - } - - fn matching_logs(dir: &Path) -> Vec { - let mut files: Vec<_> = fs::read_dir(dir) - .unwrap() - .flatten() - .map(|entry| entry.path()) - .filter(|path| { - path.file_name() - .and_then(|name| name.to_str()) - .map(|name| name.starts_with("telemt.log")) - .unwrap_or(false) - }) - .collect(); - files.sort(); - files - } - - #[test] - fn size_rotation_keeps_latest_write_in_active_file() { - let dir = tempdir().unwrap(); - let path = dir.path().join("telemt.log"); - let mut options = options(path.clone()); - options.max_size_bytes = 6; - - let mut appender = BoundedFileAppender::with_now(options, Box::new(fixed_now)).unwrap(); - appender.write_all(b"abc\n").unwrap(); - appender.write_all(b"def\n").unwrap(); - appender.flush().unwrap(); - - assert_eq!(fs::read_to_string(path).unwrap(), "def\n"); - assert_eq!(matching_logs(dir.path()).len(), 2); - } - - #[test] - fn max_files_retention_removes_oldest_archives() { - let dir = tempdir().unwrap(); - let path = dir.path().join("telemt.log"); - let mut options = options(path); - options.max_size_bytes = 4; - options.max_files = 2; - - let mut appender = BoundedFileAppender::with_now(options, Box::new(fixed_now)).unwrap(); - for line in [b"aa\n", b"bb\n", b"cc\n", b"dd\n"] { - appender.write_all(line).unwrap(); - } - appender.flush().unwrap(); - - assert!(matching_logs(dir.path()).len() <= 2); - } - - #[cfg(unix)] - #[test] - fn max_age_retention_removes_old_archives() { - use std::ffi::CString; - use std::os::unix::ffi::OsStrExt; - - let dir = tempdir().unwrap(); - let path = dir.path().join("telemt.log"); - let old_archive = dir.path().join("telemt.log.20000101000000.0"); - fs::write(&old_archive, "old").unwrap(); - - let c_path = CString::new(old_archive.as_os_str().as_bytes()).unwrap(); - let times = [ - libc::timespec { - tv_sec: 0, - tv_nsec: 0, - }, - libc::timespec { - tv_sec: 0, - tv_nsec: 0, - }, - ]; - let rc = unsafe { libc::utimensat(libc::AT_FDCWD, c_path.as_ptr(), times.as_ptr(), 0) }; - assert_eq!(rc, 0); - - let mut options = options(path); - options.max_age_secs = 1; - let _appender = BoundedFileAppender::with_now(options, Box::new(fixed_now)).unwrap(); - - assert!(!old_archive.exists()); - } -} +mod tests; diff --git a/src/logging/file/tests.rs b/src/logging/file/tests.rs new file mode 100644 index 0000000..bd2656c --- /dev/null +++ b/src/logging/file/tests.rs @@ -0,0 +1,143 @@ +use std::io::Write; + +use tempfile::tempdir; + +use super::*; + +fn fixed_now() -> DateTime { + DateTime::::from(UNIX_EPOCH + Duration::from_secs(10)) +} + +fn options(path: PathBuf) -> FileLogOptions { + FileLogOptions { + path: path.to_string_lossy().to_string(), + rotation: LogRotation::Never, + max_size_bytes: 0, + max_files: 0, + max_age_secs: 0, + } +} + +fn matching_logs(dir: &Path) -> Vec { + let mut files: Vec<_> = fs::read_dir(dir) + .unwrap() + .flatten() + .map(|entry| entry.path()) + .filter(|path| { + path.file_name() + .and_then(|name| name.to_str()) + .map(|name| name.starts_with("telemt.log")) + .unwrap_or(false) + }) + .collect(); + files.sort(); + files +} + +#[test] +fn size_rotation_keeps_latest_write_in_active_file() { + let dir = tempdir().unwrap(); + let path = dir.path().join("telemt.log"); + let mut options = options(path.clone()); + options.max_size_bytes = 6; + + let mut appender = BoundedFileAppender::with_now(options, Box::new(fixed_now)).unwrap(); + appender.write_all(b"abc\n").unwrap(); + appender.write_all(b"def\n").unwrap(); + appender.flush().unwrap(); + + assert_eq!(fs::read_to_string(path).unwrap(), "def\n"); + assert_eq!(matching_logs(dir.path()).len(), 2); +} + +#[test] +fn max_files_retention_removes_oldest_archives() { + let dir = tempdir().unwrap(); + let path = dir.path().join("telemt.log"); + let mut options = options(path); + options.max_size_bytes = 4; + options.max_files = 2; + + let mut appender = BoundedFileAppender::with_now(options, Box::new(fixed_now)).unwrap(); + for line in [b"aa\n", b"bb\n", b"cc\n", b"dd\n"] { + appender.write_all(line).unwrap(); + } + appender.flush().unwrap(); + + assert!(matching_logs(dir.path()).len() <= 2); +} + +#[cfg(unix)] +#[test] +fn max_age_retention_removes_old_archives() { + use std::ffi::CString; + use std::os::unix::ffi::OsStrExt; + + let dir = tempdir().unwrap(); + let path = dir.path().join("telemt.log"); + let old_archive = dir.path().join("telemt.log.20000101000000.0"); + fs::write(&old_archive, "old").unwrap(); + + let c_path = CString::new(old_archive.as_os_str().as_bytes()).unwrap(); + let times = [ + libc::timespec { + tv_sec: 0, + tv_nsec: 0, + }, + libc::timespec { + tv_sec: 0, + tv_nsec: 0, + }, + ]; + let rc = unsafe { libc::utimensat(libc::AT_FDCWD, c_path.as_ptr(), times.as_ptr(), 0) }; + assert_eq!(rc, 0); + + let mut options = options(path); + options.max_age_secs = 1; + let _appender = BoundedFileAppender::with_now(options, Box::new(fixed_now)).unwrap(); + + assert!(!old_archive.exists()); +} + +#[cfg(unix)] +#[test] +fn rotation_stays_bound_to_opened_directory_after_path_replacement() { + use std::os::unix::fs::symlink; + + let root = tempdir().unwrap(); + let original = root.path().join("logs"); + let moved = root.path().join("logs-moved"); + let redirect = root.path().join("redirect"); + fs::create_dir(&original).unwrap(); + fs::create_dir(&redirect).unwrap(); + let mut options = options(original.join("telemt.log")); + options.max_size_bytes = 4; + let mut appender = BoundedFileAppender::with_now(options, Box::new(fixed_now)).unwrap(); + appender.write_all(b"aa\n").unwrap(); + fs::rename(&original, &moved).unwrap(); + symlink(&redirect, &original).unwrap(); + + appender.write_all(b"bb\n").unwrap(); + appender.flush().unwrap(); + + assert!(!matching_logs(&moved).is_empty()); + assert!(matching_logs(&redirect).is_empty()); +} + +#[cfg(unix)] +#[test] +fn appender_rejects_group_writable_log_directory() { + use std::os::unix::fs::PermissionsExt; + + let current = std::env::current_dir().unwrap(); + let dir = tempfile::Builder::new() + .prefix("telemt-untrusted-log-") + .tempdir_in(current) + .unwrap(); + fs::set_permissions(dir.path(), fs::Permissions::from_mode(0o770)).unwrap(); + + assert!( + BoundedFileAppender::with_now(options(dir.path().join("telemt.log")), Box::new(fixed_now),) + .is_err() + ); +} diff --git a/src/maestro/bootstrap.rs b/src/maestro/bootstrap.rs index a89ee39..fd9abb0 100644 --- a/src/maestro/bootstrap.rs +++ b/src/maestro/bootstrap.rs @@ -74,26 +74,7 @@ pub(super) async fn bootstrap( data_path.as_deref(), ); - if !runtime_base_dir.exists() - && let Err(e) = std::fs::create_dir_all(&runtime_base_dir) - { - eprintln!( - "[telemt] Can't create runtime directory {}: {}", - runtime_base_dir.display(), - e - ); - std::process::exit(1); - } - - if !runtime_base_dir.is_dir() { - eprintln!( - "[telemt] Runtime path exists but is not a directory: {}", - runtime_base_dir.display() - ); - std::process::exit(1); - } - - if let Err(e) = std::env::set_current_dir(&runtime_base_dir) { + if let Err(e) = enter_runtime_directory(&runtime_base_dir) { eprintln!( "[telemt] Can't use runtime directory {}: {}", runtime_base_dir.display(), @@ -125,7 +106,7 @@ pub(super) async fn bootstrap( if config_path_explicit { if let Some(serialized) = serialized.as_ref() { - if let Err(write_error) = std::fs::write(&config_path, serialized) { + if let Err(write_error) = write_private_file(&config_path, serialized) { eprintln!( "[telemt] Error: failed to create explicit config at {}: {}", config_path.display(), @@ -149,7 +130,7 @@ pub(super) async fn bootstrap( if let Some(serialized) = serialized.as_ref() { match std::fs::create_dir_all(&runtime_base_dir) { - Ok(()) => match std::fs::write(&runtime_config_path, serialized) { + Ok(()) => match write_private_file(&runtime_config_path, serialized) { Ok(()) => { config_path = runtime_config_path; eprintln!( @@ -176,7 +157,7 @@ pub(super) async fn bootstrap( } if !persisted { - match std::fs::write(&fallback_config_path, serialized) { + match write_private_file(&fallback_config_path, serialized) { Ok(()) => { config_path = fallback_config_path; eprintln!( @@ -226,24 +207,7 @@ pub(super) async fn bootstrap( std::process::exit(1); } - if data_path.exists() { - if !data_path.is_dir() { - eprintln!( - "[telemt] data_path exists but is not a directory: {}", - data_path.display() - ); - std::process::exit(1); - } - } else if let Err(e) = std::fs::create_dir_all(data_path) { - eprintln!( - "[telemt] Can't create data_path {}: {}", - data_path.display(), - e - ); - std::process::exit(1); - } - - if let Err(e) = std::env::set_current_dir(data_path) { + if let Err(e) = enter_runtime_directory(data_path) { eprintln!( "[telemt] Can't use data_path {}: {}", data_path.display(), @@ -376,3 +340,26 @@ pub(super) async fn bootstrap( logging_guard, }) } + +fn enter_runtime_directory(path: &std::path::Path) -> std::io::Result<()> { + #[cfg(unix)] + { + crate::util::secure_fs::chdir_nofollow_or_create(path, 0o750) + } + #[cfg(not(unix))] + { + std::fs::create_dir_all(path)?; + std::env::set_current_dir(path) + } +} + +fn write_private_file(path: &std::path::Path, contents: &str) -> std::io::Result<()> { + #[cfg(unix)] + { + crate::util::secure_fs::atomic_replace(path, contents.as_bytes(), 0o600) + } + #[cfg(not(unix))] + { + std::fs::write(path, contents) + } +} diff --git a/src/maestro/control_plane.rs b/src/maestro/control_plane.rs index 89cda3a..5d4c60d 100644 --- a/src/maestro/control_plane.rs +++ b/src/maestro/control_plane.rs @@ -120,6 +120,34 @@ impl ProcessControlPlane { Ok(()) } + /// Registers work that must finish once accepted, even after shutdown cancellation starts. + pub(crate) fn spawn_completion(&self, future: F) -> Result<(), F> + where + F: Future + Send + 'static, + { + let Some(registration) = self.inner.admission.try_register() else { + return Err(future); + }; + self.inner.tasks.spawn(future); + drop(registration); + Ok(()) + } + + /// Registers a cooperatively cancelled task whose cleanup future must finish. + pub(crate) fn spawn_cooperative(&self, spawn: S) -> Result<(), S> + where + S: FnOnce(CancellationToken) -> F, + F: Future + Send + 'static, + { + let Some(registration) = self.inner.admission.try_register() else { + return Err(spawn); + }; + let cancellation = self.inner.cancellation.clone(); + self.inner.tasks.spawn(spawn(cancellation)); + drop(registration); + Ok(()) + } + /// Closes task admission, cancels all owned work, and joins it within the deadline. pub(crate) async fn shutdown(&self, timeout: Duration) -> bool { let deadline = tokio::time::Instant::now() + timeout; @@ -214,4 +242,49 @@ mod tests { assert!(scope.shutdown(Duration::from_secs(1)).await); } + + #[tokio::test] + async fn shutdown_waits_for_accepted_completion_without_cancelling_it() { + let scope = ProcessControlPlane::new(); + let (release_tx, release_rx) = tokio::sync::oneshot::channel(); + let completed = Arc::new(AtomicBool::new(false)); + let completed_task = completed.clone(); + assert!( + scope + .spawn_completion(async move { + let _ = release_rx.await; + completed_task.store(true, Ordering::Release); + }) + .is_ok() + ); + + let shutdown_scope = scope.clone(); + let shutdown = + tokio::spawn(async move { shutdown_scope.shutdown(Duration::from_secs(1)).await }); + tokio::task::yield_now().await; + assert!(!shutdown.is_finished()); + assert!(!completed.load(Ordering::Acquire)); + + release_tx.send(()).unwrap(); + assert!(shutdown.await.unwrap()); + assert!(completed.load(Ordering::Acquire)); + } + + #[tokio::test] + async fn cooperative_task_observes_cancellation_and_finishes_cleanup() { + let scope = ProcessControlPlane::new(); + let completed = Arc::new(AtomicBool::new(false)); + let completed_task = Arc::clone(&completed); + assert!( + scope + .spawn_cooperative(move |cancellation| async move { + cancellation.cancelled().await; + completed_task.store(true, Ordering::Release); + }) + .is_ok() + ); + + assert!(scope.shutdown(Duration::from_secs(1)).await); + assert!(completed.load(Ordering::Acquire)); + } } diff --git a/src/maestro/generation.rs b/src/maestro/generation.rs index 870b683..10da12f 100644 --- a/src/maestro/generation.rs +++ b/src/maestro/generation.rs @@ -22,6 +22,11 @@ use crate::tls_front::TlsFrontCache; use crate::transport::UpstreamManager; use crate::transport::middle_proxy::MePool; +// Cancellation guards preserve runtime ownership across preparation and drain futures. +mod lifecycle; +pub(crate) use lifecycle::RuntimeTaskScopePreparationGuard; +use lifecycle::SessionDrainCancellationGuard; + const SESSION_STOP_TIMEOUT: Duration = Duration::from_secs(5); const BACKGROUND_STOP_TIMEOUT: Duration = Duration::from_secs(5); const SESSION_ADMISSION_CLOSED: usize = 1 << (usize::BITS - 1); @@ -145,12 +150,17 @@ impl RuntimeTaskScope { self.cancel.clone() } - /// Cancels the scope and waits within the bounded background-task budget. - pub(crate) async fn stop(&self) { + /// Synchronously closes task admission and signals every tracked task. + pub(crate) fn begin_stop(&self) { self.admission.close(); - self.admission.wait_for_registrations().await; self.cancel.cancel(); self.tracker.close(); + } + + /// Cancels the scope and waits within the bounded background-task budget. + pub(crate) async fn stop(&self) { + self.begin_stop(); + self.admission.wait_for_registrations().await; let _ = tokio::time::timeout(BACKGROUND_STOP_TIMEOUT, self.tracker.wait()).await; } } @@ -298,24 +308,31 @@ impl RuntimeGeneration { /// Waits for registered sessions and cancels them when the deadline expires. pub(crate) async fn drain_sessions(&self, timeout: Duration) -> bool { self.stop_accepting_sessions(); + let mut cancellation_guard = SessionDrainCancellationGuard::new(self); self.session_admission.wait_for_registrations().await; self.sessions.close(); if tokio::time::timeout(timeout, self.sessions.wait()) .await .is_ok() { + cancellation_guard.disarm(); return true; } + cancellation_guard.disarm(); self.stop_sessions().await; false } - /// Cancels all sessions and waits within the bounded session-stop budget. - pub(crate) async fn stop_sessions(&self) { + fn begin_stop_sessions(&self) { self.stop_accepting_sessions(); - self.session_admission.wait_for_registrations().await; self.session_cancel.cancel(); self.sessions.close(); + } + + /// Cancels all sessions and waits within the bounded session-stop budget. + pub(crate) async fn stop_sessions(&self) { + self.begin_stop_sessions(); + self.session_admission.wait_for_registrations().await; let _ = tokio::time::timeout(SESSION_STOP_TIMEOUT, self.sessions.wait()).await; } @@ -335,6 +352,8 @@ impl RuntimeGeneration { impl Drop for RuntimeGeneration { fn drop(&mut self) { + self.background_tasks.begin_stop(); + self.begin_stop_sessions(); if let Some(pool) = self.me_pool.as_ref() { pool.begin_shutdown(); } diff --git a/src/maestro/generation/lifecycle.rs b/src/maestro/generation/lifecycle.rs new file mode 100644 index 0000000..4cbf0f3 --- /dev/null +++ b/src/maestro/generation/lifecycle.rs @@ -0,0 +1,158 @@ +use super::*; + +/// Cancels tasks if runtime preparation exits before ownership reaches a generation. +#[must_use = "runtime preparation guards must be disarmed after ownership transfer"] +pub(crate) struct RuntimeTaskScopePreparationGuard { + scope: RuntimeTaskScope, + armed: bool, +} + +impl RuntimeTaskScopePreparationGuard { + /// Arms cancellation for a newly created, not-yet-published task scope. + pub(crate) fn new(scope: RuntimeTaskScope) -> Self { + Self { scope, armed: true } + } + + /// Confirms that a runtime generation now owns the task scope. + pub(crate) fn disarm(mut self) { + self.armed = false; + } +} + +impl Drop for RuntimeTaskScopePreparationGuard { + fn drop(&mut self) { + if self.armed { + self.scope.begin_stop(); + } + } +} + +pub(super) struct SessionDrainCancellationGuard<'a> { + generation: &'a RuntimeGeneration, + armed: bool, +} + +impl<'a> SessionDrainCancellationGuard<'a> { + pub(super) fn new(generation: &'a RuntimeGeneration) -> Self { + Self { + generation, + armed: true, + } + } + + pub(super) fn disarm(&mut self) { + self.armed = false; + } +} + +impl Drop for SessionDrainCancellationGuard<'_> { + fn drop(&mut self) { + if self.armed { + self.generation.begin_stop_sessions(); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + struct NotifyOnDrop(Arc); + + impl Drop for NotifyOnDrop { + fn drop(&mut self) { + self.0.notify_one(); + } + } + + #[tokio::test] + async fn preparation_guard_drop_cancels_scope_and_rejects_late_spawn() { + let scope = RuntimeTaskScope::new(); + let guard = RuntimeTaskScopePreparationGuard::new(scope.clone()); + drop(guard); + + assert!(scope.cancellation_token().is_cancelled()); + let ran = Arc::new(AtomicUsize::new(0)); + let ran_task = Arc::clone(&ran); + scope.spawn(async move { + ran_task.fetch_add(1, Ordering::AcqRel); + }); + tokio::task::yield_now().await; + assert_eq!(ran.load(Ordering::Acquire), 0); + } + + #[tokio::test] + async fn preparation_guard_disarm_transfers_ownership() { + let scope = RuntimeTaskScope::new(); + RuntimeTaskScopePreparationGuard::new(scope.clone()).disarm(); + + assert!(!scope.cancellation_token().is_cancelled()); + scope.stop().await; + } + + #[tokio::test] + async fn aborted_scope_stop_still_cancels_children() { + let scope = RuntimeTaskScope::new(); + let registration = scope.admission.try_register().unwrap(); + let started = Arc::new(Notify::new()); + let dropped = Arc::new(Notify::new()); + let started_task = Arc::clone(&started); + let dropped_task = Arc::clone(&dropped); + scope.spawn(async move { + let _drop_signal = NotifyOnDrop(dropped_task); + started_task.notify_one(); + std::future::pending::<()>().await; + }); + started.notified().await; + + let stop_scope = scope.clone(); + let stop = tokio::spawn(async move { + stop_scope.stop().await; + }); + scope.cancellation_token().cancelled().await; + stop.abort(); + assert!(stop.await.unwrap_err().is_cancelled()); + drop(registration); + + tokio::time::timeout(Duration::from_secs(1), dropped.notified()) + .await + .expect("scope cancellation must drop the tracked child"); + assert!(scope.admission.try_register().is_none()); + } + + #[tokio::test] + async fn aborted_graceful_drain_forces_session_cancellation() { + let generation = test_runtime_generation(1, ProxyConfig::default()); + let registration = generation.session_admission.try_register().unwrap(); + let started = Arc::new(Notify::new()); + let dropped = Arc::new(Notify::new()); + let started_task = Arc::clone(&started); + let dropped_task = Arc::clone(&dropped); + assert!(generation.spawn_session(async move { + let _drop_signal = NotifyOnDrop(dropped_task); + started_task.notify_one(); + std::future::pending::<()>().await; + })); + started.notified().await; + + let drain_generation = Arc::clone(&generation); + let drain = tokio::spawn(async move { + drain_generation + .drain_sessions(Duration::from_secs(60)) + .await + }); + while generation.session_admission.state.load(Ordering::Acquire) & SESSION_ADMISSION_CLOSED + == 0 + { + tokio::task::yield_now().await; + } + drain.abort(); + assert!(drain.await.unwrap_err().is_cancelled()); + assert!(generation.session_cancel.is_cancelled()); + drop(registration); + + tokio::time::timeout(Duration::from_secs(1), dropped.notified()) + .await + .expect("aborted graceful drain must cancel existing sessions"); + } +} diff --git a/src/maestro/helpers.rs b/src/maestro/helpers.rs index 9a50f16..57a38e3 100644 --- a/src/maestro/helpers.rs +++ b/src/maestro/helpers.rs @@ -1,20 +1,8 @@ -#![allow(clippy::items_after_test_module)] - use std::path::{Path, PathBuf}; use std::sync::atomic::{AtomicBool, Ordering}; -use std::time::Duration; - -use tokio::sync::watch; -use tracing::{debug, error, info, warn}; use crate::cli; -use crate::config::ProxyConfig; use crate::logging::LogCliOptions; -use crate::transport::UpstreamManager; -use crate::transport::middle_proxy::{ - ProxyConfigData, fetch_proxy_config_with_raw_via_upstream, load_proxy_config_cache, - save_proxy_config_cache, -}; const MAESTRO_COLOR: &str = "\x1b[92m"; const COLOR_RESET: &str = "\x1b[0m"; @@ -47,6 +35,20 @@ pub(crate) fn resolve_runtime_config_path( startup_cwd: &Path, config_path_explicit: bool, ) -> PathBuf { + let normalize = |path: PathBuf| { + let mut normalized = PathBuf::new(); + for component in path.components() { + match component { + std::path::Component::CurDir => {} + std::path::Component::ParentDir => { + normalized.pop(); + } + component => normalized.push(component.as_os_str()), + } + } + normalized + }; + if config_path_explicit { let raw = PathBuf::from(config_path_cli); let absolute = if raw.is_absolute() { @@ -54,7 +56,7 @@ pub(crate) fn resolve_runtime_config_path( } else { startup_cwd.join(raw) }; - return absolute.canonicalize().unwrap_or(absolute); + return normalize(absolute); } let etc_telemt = std::path::Path::new("/etc/telemt"); @@ -66,7 +68,7 @@ pub(crate) fn resolve_runtime_config_path( ]; for candidate in candidates { if candidate.is_file() { - return candidate.canonicalize().unwrap_or(candidate); + return normalize(candidate); } } @@ -103,7 +105,17 @@ fn normalize_runtime_dir(path: &Path, startup_cwd: &Path) -> PathBuf { } else { startup_cwd.join(path) }; - absolute.canonicalize().unwrap_or(absolute) + let mut normalized = PathBuf::new(); + for component in absolute.components() { + match component { + std::path::Component::CurDir => {} + std::path::Component::ParentDir => { + normalized.pop(); + } + component => normalized.push(component.as_os_str()), + } + } + normalized } /// Parsed CLI arguments. @@ -314,588 +326,10 @@ fn print_help() { } } +// Runtime reporting and startup snapshot helpers. +mod runtime; + +pub(crate) use runtime::*; + #[cfg(test)] -mod tests { - use std::path::{Path, PathBuf}; - - use super::{ - expected_handshake_close_description, format_maestro_line, is_expected_handshake_eof, - peer_close_description, resolve_runtime_base_dir, resolve_runtime_config_path, - }; - use crate::error::{ProxyError, StreamError}; - - #[test] - fn maestro_line_formatter_respects_disabled_colors() { - let plain = format_maestro_line("boot", false); - assert_eq!(plain, "MAESTRO: boot"); - assert!(!plain.contains('\x1b')); - } - - #[test] - fn maestro_line_formatter_keeps_color_when_enabled() { - let colored = format_maestro_line("boot", true); - assert!(colored.contains("\x1b[92mMAESTRO\x1b[0m")); - } - - #[test] - fn resolve_runtime_config_path_anchors_relative_to_startup_cwd() { - let nonce = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_nanos(); - let startup_cwd = std::env::temp_dir().join(format!("telemt_cfg_path_{nonce}")); - std::fs::create_dir_all(&startup_cwd).unwrap(); - let target = startup_cwd.join("config.toml"); - std::fs::write(&target, " ").unwrap(); - - let resolved = resolve_runtime_config_path("config.toml", &startup_cwd, true); - assert_eq!(resolved, target.canonicalize().unwrap()); - - let _ = std::fs::remove_file(&target); - let _ = std::fs::remove_dir(&startup_cwd); - } - - #[test] - fn resolve_runtime_config_path_keeps_absolute_for_missing_file() { - let nonce = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_nanos(); - let startup_cwd = std::env::temp_dir().join(format!("telemt_cfg_path_missing_{nonce}")); - std::fs::create_dir_all(&startup_cwd).unwrap(); - - let resolved = resolve_runtime_config_path("missing.toml", &startup_cwd, true); - assert_eq!(resolved, startup_cwd.join("missing.toml")); - - let _ = std::fs::remove_dir(&startup_cwd); - } - - #[test] - fn resolve_runtime_config_path_uses_startup_candidates_when_not_explicit() { - let nonce = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_nanos(); - let startup_cwd = - std::env::temp_dir().join(format!("telemt_cfg_startup_candidates_{nonce}")); - std::fs::create_dir_all(&startup_cwd).unwrap(); - let telemt = startup_cwd.join("telemt.toml"); - std::fs::write(&telemt, " ").unwrap(); - - let resolved = resolve_runtime_config_path("config.toml", &startup_cwd, false); - assert_eq!(resolved, telemt.canonicalize().unwrap()); - - let _ = std::fs::remove_file(&telemt); - let _ = std::fs::remove_dir(&startup_cwd); - } - - #[test] - fn resolve_runtime_config_path_defaults_to_startup_config_when_none_found() { - let nonce = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_nanos(); - let startup_cwd = std::env::temp_dir().join(format!("telemt_cfg_startup_default_{nonce}")); - std::fs::create_dir_all(&startup_cwd).unwrap(); - - let resolved = resolve_runtime_config_path("config.toml", &startup_cwd, false); - assert_eq!(resolved, startup_cwd.join("config.toml")); - - let _ = std::fs::remove_dir(&startup_cwd); - } - - #[test] - fn resolve_runtime_base_dir_prefers_cli_data_path() { - let nonce = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_nanos(); - let startup_cwd = std::env::temp_dir().join(format!("telemt_runtime_base_cwd_{nonce}")); - let data_path = std::env::temp_dir().join(format!("telemt_runtime_base_data_{nonce}")); - std::fs::create_dir_all(&startup_cwd).unwrap(); - std::fs::create_dir_all(&data_path).unwrap(); - - let resolved = resolve_runtime_base_dir( - &startup_cwd.join("config.toml"), - &startup_cwd, - true, - Some(&data_path), - ); - assert_eq!(resolved, data_path.canonicalize().unwrap()); - - let _ = std::fs::remove_dir(&data_path); - let _ = std::fs::remove_dir(&startup_cwd); - } - - #[test] - fn resolve_runtime_base_dir_uses_working_directory_before_explicit_config_parent() { - let nonce = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_nanos(); - let startup_cwd = std::env::temp_dir().join(format!("telemt_runtime_base_start_{nonce}")); - let config_dir = std::env::temp_dir().join(format!("telemt_runtime_base_cfg_{nonce}")); - std::fs::create_dir_all(&startup_cwd).unwrap(); - std::fs::create_dir_all(&config_dir).unwrap(); - - let resolved = - resolve_runtime_base_dir(&config_dir.join("telemt.toml"), &startup_cwd, true, None); - assert_eq!(resolved, startup_cwd.canonicalize().unwrap()); - - let _ = std::fs::remove_dir(&config_dir); - let _ = std::fs::remove_dir(&startup_cwd); - } - - #[test] - fn resolve_runtime_base_dir_uses_explicit_config_parent_from_root() { - let nonce = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_nanos(); - let config_dir = std::env::temp_dir().join(format!("telemt_runtime_base_root_cfg_{nonce}")); - std::fs::create_dir_all(&config_dir).unwrap(); - - let resolved = - resolve_runtime_base_dir(&config_dir.join("telemt.toml"), Path::new("/"), true, None); - assert_eq!(resolved, config_dir.canonicalize().unwrap()); - - let _ = std::fs::remove_dir(&config_dir); - } - - #[test] - fn resolve_runtime_base_dir_uses_systemd_working_directory_before_etc() { - let nonce = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() - .as_nanos(); - let startup_cwd = std::env::temp_dir().join(format!("telemt_runtime_base_systemd_{nonce}")); - std::fs::create_dir_all(&startup_cwd).unwrap(); - - let resolved = - resolve_runtime_base_dir(&startup_cwd.join("config.toml"), &startup_cwd, false, None); - assert_eq!(resolved, startup_cwd.canonicalize().unwrap()); - - let _ = std::fs::remove_dir(&startup_cwd); - } - - #[test] - fn resolve_runtime_base_dir_falls_back_to_etc_from_root() { - let resolved = resolve_runtime_base_dir( - Path::new("/etc/telemt/config.toml"), - Path::new("/"), - false, - None, - ); - assert_eq!(resolved, PathBuf::from("/etc/telemt")); - } - - #[test] - fn expected_handshake_eof_matches_connection_reset() { - let err = ProxyError::Io(std::io::Error::from(std::io::ErrorKind::ConnectionReset)); - assert!(is_expected_handshake_eof(&err)); - } - - #[test] - fn expected_handshake_eof_matches_stream_io_unexpected_eof() { - let err = ProxyError::Stream(StreamError::Io(std::io::Error::from( - std::io::ErrorKind::UnexpectedEof, - ))); - assert!(is_expected_handshake_eof(&err)); - } - - #[test] - fn peer_close_description_is_human_readable_for_all_peer_close_kinds() { - let cases = [ - ( - std::io::ErrorKind::ConnectionReset, - "Peer reset TCP connection (RST)", - ), - ( - std::io::ErrorKind::ConnectionAborted, - "Peer aborted TCP connection during transport", - ), - ( - std::io::ErrorKind::BrokenPipe, - "Peer closed write side (broken pipe)", - ), - ( - std::io::ErrorKind::NotConnected, - "Socket was already closed by peer", - ), - ]; - - for (kind, expected) in cases { - let err = ProxyError::Io(std::io::Error::from(kind)); - assert_eq!(peer_close_description(&err), Some(expected)); - } - } - - #[test] - fn handshake_close_description_is_human_readable_for_all_expected_kinds() { - let cases = [ - ( - ProxyError::Io(std::io::Error::from(std::io::ErrorKind::UnexpectedEof)), - "Peer closed before sending full 64-byte MTProto handshake", - ), - ( - ProxyError::Io(std::io::Error::from(std::io::ErrorKind::ConnectionReset)), - "Peer reset TCP connection during initial MTProto handshake", - ), - ( - ProxyError::Io(std::io::Error::from(std::io::ErrorKind::ConnectionAborted)), - "Peer aborted TCP connection during initial MTProto handshake", - ), - ( - ProxyError::Io(std::io::Error::from(std::io::ErrorKind::BrokenPipe)), - "Peer closed write side before MTProto handshake completed", - ), - ( - ProxyError::Io(std::io::Error::from(std::io::ErrorKind::NotConnected)), - "Handshake socket was already closed by peer", - ), - ( - ProxyError::Stream(StreamError::UnexpectedEof), - "Peer closed before sending full 64-byte MTProto handshake", - ), - ]; - - for (err, expected) in cases { - assert_eq!(expected_handshake_close_description(&err), Some(expected)); - } - } -} - -pub(crate) fn print_proxy_links(host: &str, port: u16, config: &ProxyConfig) { - print_maestro_line(format!("Proxy links ({host})")); - for user_name in config - .general - .links - .show - .resolve_users(&config.access.users) - { - if let Some(secret) = config.access.users.get(user_name) { - print_maestro_line(format!("User: {user_name}")); - if config.general.modes.classic { - print_maestro_line(format!( - "Classic: tg://proxy?server={host}&port={port}&secret={secret}" - )); - } - if config.general.modes.secure { - print_maestro_line(format!( - "DD: tg://proxy?server={host}&port={port}&secret=dd{secret}" - )); - } - if config.general.modes.tls { - let mut domains = Vec::with_capacity(1 + config.censorship.tls_domains.len()); - domains.push(config.censorship.tls_domain.clone()); - for d in &config.censorship.tls_domains { - if !domains.contains(d) { - domains.push(d.clone()); - } - } - - for domain in domains { - let domain_hex = hex::encode(&domain); - print_maestro_line(format!( - "EE-TLS: tg://proxy?server={host}&port={port}&secret=ee{secret}{domain_hex}" - )); - } - } - } else { - warn!(target: "telemt::links", "User '{}' in show_link not found", user_name); - } - } -} - -/// Prints WEB links only for profiles selected by the existing link policy. -pub(crate) fn print_web_proxy_links(config: &ProxyConfig) { - if !config.web.enabled || config.general.links.show.is_empty() { - return; - } - let Some(runtime) = config.web.runtime.as_ref() else { - return; - }; - let shown = config - .general - .links - .show - .resolve_users(&config.access.users); - let mut heading_printed = false; - for profile in &runtime.profiles { - if !shown.iter().any(|user| user.as_str() == profile.user) { - continue; - } - if !heading_printed { - print_maestro_line("WEB proxy links"); - heading_printed = true; - } - let Some(secret) = config.access.users.get(&profile.user) else { - continue; - }; - let prefix = match profile.secret_mode { - crate::config::WebSecretMode::Plain => "", - crate::config::WebSecretMode::Dd => "dd", - }; - print_maestro_line(format!( - "User: {} ({:?})", - profile.user, profile.secret_mode - )); - print_maestro_line(format!( - "WEB: tg://webproxy?server={}&secret={prefix}{secret}", - profile.host, - )); - } -} - -pub(crate) async fn write_beobachten_snapshot(path: &str, payload: &str) -> std::io::Result<()> { - if let Some(parent) = std::path::Path::new(path).parent() - && !parent.as_os_str().is_empty() - { - tokio::fs::create_dir_all(parent).await?; - } - tokio::fs::write(path, payload).await -} - -pub(crate) fn unit_label(value: u64, singular: &'static str, plural: &'static str) -> &'static str { - if value == 1 { singular } else { plural } -} - -pub(crate) fn format_uptime(total_secs: u64) -> String { - const SECS_PER_MINUTE: u64 = 60; - const SECS_PER_HOUR: u64 = 60 * SECS_PER_MINUTE; - const SECS_PER_DAY: u64 = 24 * SECS_PER_HOUR; - const SECS_PER_MONTH: u64 = 30 * SECS_PER_DAY; - const SECS_PER_YEAR: u64 = 12 * SECS_PER_MONTH; - - let mut remaining = total_secs; - let years = remaining / SECS_PER_YEAR; - remaining %= SECS_PER_YEAR; - let months = remaining / SECS_PER_MONTH; - remaining %= SECS_PER_MONTH; - let days = remaining / SECS_PER_DAY; - remaining %= SECS_PER_DAY; - let hours = remaining / SECS_PER_HOUR; - remaining %= SECS_PER_HOUR; - let minutes = remaining / SECS_PER_MINUTE; - let seconds = remaining % SECS_PER_MINUTE; - - let mut parts = Vec::new(); - if total_secs > SECS_PER_YEAR { - parts.push(format!("{} {}", years, unit_label(years, "year", "years"))); - } - if total_secs > SECS_PER_MONTH { - parts.push(format!( - "{} {}", - months, - unit_label(months, "month", "months") - )); - } - if total_secs > SECS_PER_DAY { - parts.push(format!("{} {}", days, unit_label(days, "day", "days"))); - } - if total_secs > SECS_PER_HOUR { - parts.push(format!("{} {}", hours, unit_label(hours, "hour", "hours"))); - } - if total_secs > SECS_PER_MINUTE { - parts.push(format!( - "{} {}", - minutes, - unit_label(minutes, "minute", "minutes") - )); - } - parts.push(format!( - "{} {}", - seconds, - unit_label(seconds, "second", "seconds") - )); - - format!("{} / {} seconds", parts.join(", "), total_secs) -} - -#[allow(dead_code)] -pub(crate) async fn wait_until_admission_open(admission_rx: &mut watch::Receiver) -> bool { - loop { - if *admission_rx.borrow() { - return true; - } - if admission_rx.changed().await.is_err() { - return *admission_rx.borrow(); - } - } -} - -pub(crate) fn is_expected_handshake_eof(err: &crate::error::ProxyError) -> bool { - expected_handshake_close_description(err).is_some() -} - -pub(crate) fn peer_close_description(err: &crate::error::ProxyError) -> Option<&'static str> { - fn from_kind(kind: std::io::ErrorKind) -> Option<&'static str> { - match kind { - std::io::ErrorKind::ConnectionReset => Some("Peer reset TCP connection (RST)"), - std::io::ErrorKind::ConnectionAborted => { - Some("Peer aborted TCP connection during transport") - } - std::io::ErrorKind::BrokenPipe => Some("Peer closed write side (broken pipe)"), - std::io::ErrorKind::NotConnected => Some("Socket was already closed by peer"), - _ => None, - } - } - - match err { - crate::error::ProxyError::Io(ioe) => from_kind(ioe.kind()), - crate::error::ProxyError::Stream(crate::error::StreamError::Io(ioe)) => { - from_kind(ioe.kind()) - } - _ => None, - } -} - -pub(crate) fn expected_handshake_close_description( - err: &crate::error::ProxyError, -) -> Option<&'static str> { - fn from_kind(kind: std::io::ErrorKind) -> Option<&'static str> { - match kind { - std::io::ErrorKind::UnexpectedEof => { - Some("Peer closed before sending full 64-byte MTProto handshake") - } - std::io::ErrorKind::ConnectionReset => { - Some("Peer reset TCP connection during initial MTProto handshake") - } - std::io::ErrorKind::ConnectionAborted => { - Some("Peer aborted TCP connection during initial MTProto handshake") - } - std::io::ErrorKind::BrokenPipe => { - Some("Peer closed write side before MTProto handshake completed") - } - std::io::ErrorKind::NotConnected => Some("Handshake socket was already closed by peer"), - _ => None, - } - } - - match err { - crate::error::ProxyError::Io(ioe) => from_kind(ioe.kind()), - crate::error::ProxyError::Stream(crate::error::StreamError::UnexpectedEof) => { - Some("Peer closed before sending full 64-byte MTProto handshake") - } - crate::error::ProxyError::Stream(crate::error::StreamError::Io(ioe)) => { - from_kind(ioe.kind()) - } - _ => None, - } -} - -pub(crate) async fn load_startup_proxy_config_snapshot( - url: &str, - cache_path: Option<&str>, - me2dc_fallback: bool, - label: &'static str, - upstream: Option>, -) -> Option { - loop { - match fetch_proxy_config_with_raw_via_upstream(url, upstream.clone()).await { - Ok((cfg, raw)) => { - if !cfg.map.is_empty() { - if let Some(path) = cache_path - && let Err(e) = save_proxy_config_cache(path, &raw).await - { - warn!(error = %e, path, snapshot = label, "Failed to store startup proxy-config cache"); - } - return Some(cfg); - } - - warn!( - snapshot = label, - url, "Startup proxy-config is empty; trying disk cache" - ); - if let Some(path) = cache_path { - match load_proxy_config_cache(path).await { - Ok(cached) if !cached.map.is_empty() => { - info!( - snapshot = label, - path, - proxy_for_lines = cached.proxy_for_lines, - "Loaded startup proxy-config from disk cache" - ); - return Some(cached); - } - Ok(_) => { - warn!( - snapshot = label, - path, "Startup proxy-config cache is empty; ignoring cache file" - ); - } - Err(cache_err) => { - debug!( - snapshot = label, - path, - error = %cache_err, - "Startup proxy-config cache unavailable" - ); - } - } - } - - if me2dc_fallback { - error!( - snapshot = label, - "Startup proxy-config unavailable and no saved config found; falling back to direct mode" - ); - return None; - } - - warn!( - snapshot = label, - retry_in_secs = 2, - "Startup proxy-config unavailable and no saved config found; retrying because me2dc_fallback=false" - ); - tokio::time::sleep(Duration::from_secs(2)).await; - } - Err(fetch_err) => { - if let Some(path) = cache_path { - match load_proxy_config_cache(path).await { - Ok(cached) if !cached.map.is_empty() => { - info!( - snapshot = label, - path, - proxy_for_lines = cached.proxy_for_lines, - "Loaded startup proxy-config from disk cache" - ); - return Some(cached); - } - Ok(_) => { - warn!( - snapshot = label, - path, "Startup proxy-config cache is empty; ignoring cache file" - ); - } - Err(cache_err) => { - debug!( - snapshot = label, - path, - error = %cache_err, - "Startup proxy-config cache unavailable" - ); - } - } - } - - if me2dc_fallback { - error!( - snapshot = label, - error = %fetch_err, - "Startup proxy-config unavailable and no cached data; falling back to direct mode" - ); - return None; - } - - warn!( - snapshot = label, - error = %fetch_err, - retry_in_secs = 2, - "Startup proxy-config unavailable; retrying because me2dc_fallback=false" - ); - tokio::time::sleep(Duration::from_secs(2)).await; - } - } - } -} +mod tests; diff --git a/src/maestro/helpers/runtime.rs b/src/maestro/helpers/runtime.rs new file mode 100644 index 0000000..fcd19d1 --- /dev/null +++ b/src/maestro/helpers/runtime.rs @@ -0,0 +1,476 @@ +use std::time::Duration; + +use base64::Engine as _; +use tokio::sync::watch; +use tracing::{debug, error, info, warn}; + +use crate::config::ProxyConfig; +use crate::transport::UpstreamManager; +use crate::transport::middle_proxy::{ + ProxyConfigData, fetch_proxy_config_with_raw_via_upstream, load_proxy_config_cache, + save_proxy_config_cache, +}; + +use super::print_maestro_line; + +/// Prints configured MTProxy links through the direct MAESTRO output channel. +pub(crate) fn print_proxy_links(host: &str, port: u16, config: &ProxyConfig) { + print_maestro_line(format!("Proxy links ({host})")); + for user_name in config + .general + .links + .show + .resolve_users(&config.access.users) + { + if let Some(secret) = config.access.users.get(user_name) { + print_maestro_line(format!("User: {user_name}")); + if config.general.modes.classic { + print_maestro_line(format!( + "Classic: tg://proxy?server={host}&port={port}&secret={secret}" + )); + } + if config.general.modes.secure { + print_maestro_line(format!( + "DD: tg://proxy?server={host}&port={port}&secret=dd{secret}" + )); + } + if config.general.modes.tls { + let mut domains = Vec::with_capacity(1 + config.censorship.tls_domains.len()); + domains.push(config.censorship.tls_domain.clone()); + for d in &config.censorship.tls_domains { + if !domains.contains(d) { + domains.push(d.clone()); + } + } + + for domain in domains { + let domain_hex = hex::encode(&domain); + print_maestro_line(format!( + "EE-TLS: tg://proxy?server={host}&port={port}&secret=ee{secret}{domain_hex}" + )); + } + } + } else { + warn!(target: "telemt::links", "User '{}' in show_link not found", user_name); + } + } +} + +/// Prints WEB links only for profiles selected by the existing link policy. +pub(crate) fn print_web_proxy_links(config: &ProxyConfig) { + if !config.web.enabled || config.general.links.show.is_empty() { + return; + } + let Some(runtime) = config.web.runtime.as_ref() else { + return; + }; + let shown = config + .general + .links + .show + .resolve_users(&config.access.users); + let mut heading_printed = false; + for profile in &runtime.profiles { + if !shown.iter().any(|user| user.as_str() == profile.user) { + continue; + } + if !heading_printed { + print_maestro_line("WEB proxy links"); + heading_printed = true; + } + let Some(secret) = config.access.users.get(&profile.user) else { + continue; + }; + let Some(vhost) = config + .web + .vhosts + .iter() + .find(|vhost| vhost.host == profile.host) + else { + continue; + }; + print_maestro_line(format!( + "User: {} ({:?})", + profile.user, profile.secret_mode + )); + if let Some(link) = + format_web_proxy_link(&profile.host, &vhost.base_path, secret, profile.secret_mode) + { + print_maestro_line(format!("WEB: {link}")); + } + } +} + +fn format_web_proxy_link( + host: &str, + base_path: &str, + secret: &str, + mode: crate::config::WebSecretMode, +) -> Option { + if base_path.is_empty() { + let prefix = match mode { + crate::config::WebSecretMode::Plain => "", + crate::config::WebSecretMode::Dd => "dd", + }; + return Some(format!( + "tg://webproxy?server={host}&secret={prefix}{secret}" + )); + } + let decoded = hex::decode(secret).ok()?; + let mut marked = Vec::with_capacity(decoded.len() + 2); + marked.push(0x70); + if mode == crate::config::WebSecretMode::Dd { + marked.push(0xdd); + } + marked.extend_from_slice(&decoded); + let server = url::form_urlencoded::byte_serialize(format!("{host}/{base_path}").as_bytes()) + .collect::(); + let marked = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(marked); + Some(format!("tg://webproxy?server={server}&secret={marked}")) +} + +/// Durably replaces one Beobachten snapshot without following Unix symlinks. +pub(crate) async fn write_beobachten_snapshot(path: &str, payload: &str) -> std::io::Result<()> { + #[cfg(unix)] + { + crate::util::secure_fs::atomic_replace_async( + std::path::PathBuf::from(path), + payload.as_bytes().to_vec(), + 0o600, + ) + .await + } + #[cfg(not(unix))] + { + if let Some(parent) = std::path::Path::new(path).parent() + && !parent.as_os_str().is_empty() + { + tokio::fs::create_dir_all(parent).await?; + } + tokio::fs::write(path, payload).await + } +} + +/// Selects a singular or plural display label for one integer value. +pub(crate) fn unit_label(value: u64, singular: &'static str, plural: &'static str) -> &'static str { + if value == 1 { singular } else { plural } +} + +/// Formats process uptime into bounded human-readable units and exact seconds. +pub(crate) fn format_uptime(total_secs: u64) -> String { + const SECS_PER_MINUTE: u64 = 60; + const SECS_PER_HOUR: u64 = 60 * SECS_PER_MINUTE; + const SECS_PER_DAY: u64 = 24 * SECS_PER_HOUR; + const SECS_PER_MONTH: u64 = 30 * SECS_PER_DAY; + const SECS_PER_YEAR: u64 = 12 * SECS_PER_MONTH; + + let mut remaining = total_secs; + let years = remaining / SECS_PER_YEAR; + remaining %= SECS_PER_YEAR; + let months = remaining / SECS_PER_MONTH; + remaining %= SECS_PER_MONTH; + let days = remaining / SECS_PER_DAY; + remaining %= SECS_PER_DAY; + let hours = remaining / SECS_PER_HOUR; + remaining %= SECS_PER_HOUR; + let minutes = remaining / SECS_PER_MINUTE; + let seconds = remaining % SECS_PER_MINUTE; + + let mut parts = Vec::new(); + if total_secs > SECS_PER_YEAR { + parts.push(format!("{} {}", years, unit_label(years, "year", "years"))); + } + if total_secs > SECS_PER_MONTH { + parts.push(format!( + "{} {}", + months, + unit_label(months, "month", "months") + )); + } + if total_secs > SECS_PER_DAY { + parts.push(format!("{} {}", days, unit_label(days, "day", "days"))); + } + if total_secs > SECS_PER_HOUR { + parts.push(format!("{} {}", hours, unit_label(hours, "hour", "hours"))); + } + if total_secs > SECS_PER_MINUTE { + parts.push(format!( + "{} {}", + minutes, + unit_label(minutes, "minute", "minutes") + )); + } + parts.push(format!( + "{} {}", + seconds, + unit_label(seconds, "second", "seconds") + )); + + format!("{} / {} seconds", parts.join(", "), total_secs) +} + +#[allow(dead_code)] +/// Waits until admission opens or its watch channel closes. +pub(crate) async fn wait_until_admission_open(admission_rx: &mut watch::Receiver) -> bool { + loop { + if *admission_rx.borrow() { + return true; + } + if admission_rx.changed().await.is_err() { + return *admission_rx.borrow(); + } + } +} + +/// Classifies peer closure that is expected during an incomplete handshake. +pub(crate) fn is_expected_handshake_eof(err: &crate::error::ProxyError) -> bool { + expected_handshake_close_description(err).is_some() +} + +/// Returns a stable diagnostic description for transport-level peer closure. +pub(crate) fn peer_close_description(err: &crate::error::ProxyError) -> Option<&'static str> { + fn from_kind(kind: std::io::ErrorKind) -> Option<&'static str> { + match kind { + std::io::ErrorKind::ConnectionReset => Some("Peer reset TCP connection (RST)"), + std::io::ErrorKind::ConnectionAborted => { + Some("Peer aborted TCP connection during transport") + } + std::io::ErrorKind::BrokenPipe => Some("Peer closed write side (broken pipe)"), + std::io::ErrorKind::NotConnected => Some("Socket was already closed by peer"), + _ => None, + } + } + + match err { + crate::error::ProxyError::Io(ioe) => from_kind(ioe.kind()), + crate::error::ProxyError::Stream(crate::error::StreamError::Io(ioe)) => { + from_kind(ioe.kind()) + } + _ => None, + } +} + +/// Returns a stable diagnostic description for expected handshake closure. +pub(crate) fn expected_handshake_close_description( + err: &crate::error::ProxyError, +) -> Option<&'static str> { + fn from_kind(kind: std::io::ErrorKind) -> Option<&'static str> { + match kind { + std::io::ErrorKind::UnexpectedEof => { + Some("Peer closed before sending full 64-byte MTProto handshake") + } + std::io::ErrorKind::ConnectionReset => { + Some("Peer reset TCP connection during initial MTProto handshake") + } + std::io::ErrorKind::ConnectionAborted => { + Some("Peer aborted TCP connection during initial MTProto handshake") + } + std::io::ErrorKind::BrokenPipe => { + Some("Peer closed write side before MTProto handshake completed") + } + std::io::ErrorKind::NotConnected => Some("Handshake socket was already closed by peer"), + _ => None, + } + } + + match err { + crate::error::ProxyError::Io(ioe) => from_kind(ioe.kind()), + crate::error::ProxyError::Stream(crate::error::StreamError::UnexpectedEof) => { + Some("Peer closed before sending full 64-byte MTProto handshake") + } + crate::error::ProxyError::Stream(crate::error::StreamError::Io(ioe)) => { + from_kind(ioe.kind()) + } + _ => None, + } +} + +/// Loads a non-empty startup endpoint snapshot with bounded cache fallback. +pub(crate) async fn load_startup_proxy_config_snapshot( + url: &str, + cache_path: Option<&str>, + me2dc_fallback: bool, + label: &'static str, + upstream: Option>, +) -> Option { + loop { + match fetch_proxy_config_with_raw_via_upstream(url, upstream.clone()).await { + Ok((cfg, raw)) => { + if !cfg.map.is_empty() { + if let Some(path) = cache_path + && let Err(e) = save_proxy_config_cache(path, &raw).await + { + warn!(error = %e, path, snapshot = label, "Failed to store startup proxy-config cache"); + } + return Some(cfg); + } + + warn!( + snapshot = label, + url, "Startup proxy-config is empty; trying disk cache" + ); + if let Some(path) = cache_path { + match load_proxy_config_cache(path).await { + Ok(cached) if !cached.map.is_empty() => { + info!( + snapshot = label, + path, + proxy_for_lines = cached.proxy_for_lines, + "Loaded startup proxy-config from disk cache" + ); + return Some(cached); + } + Ok(_) => { + warn!( + snapshot = label, + path, "Startup proxy-config cache is empty; ignoring cache file" + ); + } + Err(cache_err) => { + debug!( + snapshot = label, + path, + error = %cache_err, + "Startup proxy-config cache unavailable" + ); + } + } + } + + if me2dc_fallback { + error!( + snapshot = label, + "Startup proxy-config unavailable and no saved config found; falling back to direct mode" + ); + return None; + } + + warn!( + snapshot = label, + retry_in_secs = 2, + "Startup proxy-config unavailable and no saved config found; retrying because me2dc_fallback=false" + ); + tokio::time::sleep(Duration::from_secs(2)).await; + } + Err(fetch_err) => { + if let Some(path) = cache_path { + match load_proxy_config_cache(path).await { + Ok(cached) if !cached.map.is_empty() => { + info!( + snapshot = label, + path, + proxy_for_lines = cached.proxy_for_lines, + "Loaded startup proxy-config from disk cache" + ); + return Some(cached); + } + Ok(_) => { + warn!( + snapshot = label, + path, "Startup proxy-config cache is empty; ignoring cache file" + ); + } + Err(cache_err) => { + debug!( + snapshot = label, + path, + error = %cache_err, + "Startup proxy-config cache unavailable" + ); + } + } + } + + if me2dc_fallback { + error!( + snapshot = label, + error = %fetch_err, + "Startup proxy-config unavailable and no cached data; falling back to direct mode" + ); + return None; + } + + warn!( + snapshot = label, + error = %fetch_err, + retry_in_secs = 2, + "Startup proxy-config unavailable; retrying because me2dc_fallback=false" + ); + tokio::time::sleep(Duration::from_secs(2)).await; + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::WebSecretMode; + + const SECRET: &str = "000102030405060708090a0b0c0d0e0f"; + + #[test] + fn root_web_proxy_links_keep_the_legacy_secret_form() { + assert_eq!( + format_web_proxy_link("proxy.example.com", "", SECRET, WebSecretMode::Plain), + Some(format!( + "tg://webproxy?server=proxy.example.com&secret={SECRET}" + )) + ); + assert_eq!( + format_web_proxy_link("proxy.example.com", "", SECRET, WebSecretMode::Dd), + Some(format!( + "tg://webproxy?server=proxy.example.com&secret=dd{SECRET}" + )) + ); + } + + #[test] + fn path_web_proxy_links_use_the_tdesktop_marker() { + assert_eq!( + format_web_proxy_link( + "proxy.example.com", + "dobry-cola/super_app", + SECRET, + WebSecretMode::Plain, + ), + Some("tg://webproxy?server=proxy.example.com%2Fdobry-cola%2Fsuper_app&secret=cAABAgMEBQYHCAkKCwwNDg8".to_string()) + ); + assert_eq!( + format_web_proxy_link( + "proxy.example.com", + "dobry-cola/super_app", + SECRET, + WebSecretMode::Dd, + ), + Some("tg://webproxy?server=proxy.example.com%2Fdobry-cola%2Fsuper_app&secret=cN0AAQIDBAUGBwgJCgsMDQ4P".to_string()) + ); + } + + #[test] + fn path_web_proxy_link_round_trips_through_the_tdesktop_grammar() { + for (mode, expected_secret) in [ + (WebSecretMode::Plain, hex::decode(SECRET).unwrap()), + ( + WebSecretMode::Dd, + [vec![0xdd], hex::decode(SECRET).unwrap()].concat(), + ), + ] { + let link = format_web_proxy_link("proxy.example.com", "MixedCase/a_b-9", SECRET, mode) + .unwrap(); + let parsed = url::Url::parse(&link).unwrap(); + let query = parsed + .query_pairs() + .collect::>(); + assert_eq!( + query["server"].as_ref(), + "proxy.example.com/MixedCase/a_b-9" + ); + let marked = base64::engine::general_purpose::URL_SAFE_NO_PAD + .decode(query["secret"].as_bytes()) + .unwrap(); + assert_eq!(marked.first(), Some(&0x70)); + assert_eq!(&marked[1..], expected_secret.as_slice()); + } + } +} diff --git a/src/maestro/helpers/tests.rs b/src/maestro/helpers/tests.rs new file mode 100644 index 0000000..2dc5524 --- /dev/null +++ b/src/maestro/helpers/tests.rs @@ -0,0 +1,266 @@ +use std::path::{Path, PathBuf}; + +use super::{ + expected_handshake_close_description, format_maestro_line, is_expected_handshake_eof, + peer_close_description, resolve_runtime_base_dir, resolve_runtime_config_path, +}; +use crate::error::{ProxyError, StreamError}; + +#[test] +fn maestro_line_formatter_respects_disabled_colors() { + let plain = format_maestro_line("boot", false); + assert_eq!(plain, "MAESTRO: boot"); + assert!(!plain.contains('\x1b')); +} + +#[test] +fn maestro_line_formatter_keeps_color_when_enabled() { + let colored = format_maestro_line("boot", true); + assert!(colored.contains("\x1b[92mMAESTRO\x1b[0m")); +} + +#[test] +fn resolve_runtime_config_path_anchors_relative_to_startup_cwd() { + let nonce = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos(); + let startup_cwd = std::env::temp_dir().join(format!("telemt_cfg_path_{nonce}")); + std::fs::create_dir_all(&startup_cwd).unwrap(); + let target = startup_cwd.join("config.toml"); + std::fs::write(&target, " ").unwrap(); + + let resolved = resolve_runtime_config_path("config.toml", &startup_cwd, true); + assert_eq!(resolved, target.canonicalize().unwrap()); + + let _ = std::fs::remove_file(&target); + let _ = std::fs::remove_dir(&startup_cwd); +} + +#[test] +fn resolve_runtime_config_path_keeps_absolute_for_missing_file() { + let nonce = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos(); + let startup_cwd = std::env::temp_dir().join(format!("telemt_cfg_path_missing_{nonce}")); + std::fs::create_dir_all(&startup_cwd).unwrap(); + + let resolved = resolve_runtime_config_path("missing.toml", &startup_cwd, true); + assert_eq!(resolved, startup_cwd.join("missing.toml")); + + let _ = std::fs::remove_dir(&startup_cwd); +} + +#[cfg(unix)] +#[test] +fn runtime_paths_preserve_symlinks_for_descriptor_validation() { + use std::os::unix::fs::symlink; + + let dir = tempfile::tempdir().unwrap(); + let real_dir = dir.path().join("real"); + let linked_dir = dir.path().join("linked"); + std::fs::create_dir(&real_dir).unwrap(); + std::fs::write(real_dir.join("config.toml"), " ").unwrap(); + symlink(&real_dir, &linked_dir).unwrap(); + let linked_config = linked_dir.join("config.toml"); + + let config = resolve_runtime_config_path(linked_config.to_str().unwrap(), dir.path(), true); + let runtime = resolve_runtime_base_dir(&linked_config, dir.path(), true, Some(&linked_dir)); + + assert_eq!(config, linked_config); + assert_eq!(runtime, linked_dir); +} + +#[test] +fn resolve_runtime_config_path_uses_startup_candidates_when_not_explicit() { + let nonce = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos(); + let startup_cwd = std::env::temp_dir().join(format!("telemt_cfg_startup_candidates_{nonce}")); + std::fs::create_dir_all(&startup_cwd).unwrap(); + let telemt = startup_cwd.join("telemt.toml"); + std::fs::write(&telemt, " ").unwrap(); + + let resolved = resolve_runtime_config_path("config.toml", &startup_cwd, false); + assert_eq!(resolved, telemt.canonicalize().unwrap()); + + let _ = std::fs::remove_file(&telemt); + let _ = std::fs::remove_dir(&startup_cwd); +} + +#[test] +fn resolve_runtime_config_path_defaults_to_startup_config_when_none_found() { + let nonce = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos(); + let startup_cwd = std::env::temp_dir().join(format!("telemt_cfg_startup_default_{nonce}")); + std::fs::create_dir_all(&startup_cwd).unwrap(); + + let resolved = resolve_runtime_config_path("config.toml", &startup_cwd, false); + assert_eq!(resolved, startup_cwd.join("config.toml")); + + let _ = std::fs::remove_dir(&startup_cwd); +} + +#[test] +fn resolve_runtime_base_dir_prefers_cli_data_path() { + let nonce = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos(); + let startup_cwd = std::env::temp_dir().join(format!("telemt_runtime_base_cwd_{nonce}")); + let data_path = std::env::temp_dir().join(format!("telemt_runtime_base_data_{nonce}")); + std::fs::create_dir_all(&startup_cwd).unwrap(); + std::fs::create_dir_all(&data_path).unwrap(); + + let resolved = resolve_runtime_base_dir( + &startup_cwd.join("config.toml"), + &startup_cwd, + true, + Some(&data_path), + ); + assert_eq!(resolved, data_path.canonicalize().unwrap()); + + let _ = std::fs::remove_dir(&data_path); + let _ = std::fs::remove_dir(&startup_cwd); +} + +#[test] +fn resolve_runtime_base_dir_uses_working_directory_before_explicit_config_parent() { + let nonce = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos(); + let startup_cwd = std::env::temp_dir().join(format!("telemt_runtime_base_start_{nonce}")); + let config_dir = std::env::temp_dir().join(format!("telemt_runtime_base_cfg_{nonce}")); + std::fs::create_dir_all(&startup_cwd).unwrap(); + std::fs::create_dir_all(&config_dir).unwrap(); + + let resolved = + resolve_runtime_base_dir(&config_dir.join("telemt.toml"), &startup_cwd, true, None); + assert_eq!(resolved, startup_cwd.canonicalize().unwrap()); + + let _ = std::fs::remove_dir(&config_dir); + let _ = std::fs::remove_dir(&startup_cwd); +} + +#[test] +fn resolve_runtime_base_dir_uses_explicit_config_parent_from_root() { + let nonce = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos(); + let config_dir = std::env::temp_dir().join(format!("telemt_runtime_base_root_cfg_{nonce}")); + std::fs::create_dir_all(&config_dir).unwrap(); + + let resolved = + resolve_runtime_base_dir(&config_dir.join("telemt.toml"), Path::new("/"), true, None); + assert_eq!(resolved, config_dir.canonicalize().unwrap()); + + let _ = std::fs::remove_dir(&config_dir); +} + +#[test] +fn resolve_runtime_base_dir_uses_systemd_working_directory_before_etc() { + let nonce = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos(); + let startup_cwd = std::env::temp_dir().join(format!("telemt_runtime_base_systemd_{nonce}")); + std::fs::create_dir_all(&startup_cwd).unwrap(); + + let resolved = + resolve_runtime_base_dir(&startup_cwd.join("config.toml"), &startup_cwd, false, None); + assert_eq!(resolved, startup_cwd.canonicalize().unwrap()); + + let _ = std::fs::remove_dir(&startup_cwd); +} + +#[test] +fn resolve_runtime_base_dir_falls_back_to_etc_from_root() { + let resolved = resolve_runtime_base_dir( + Path::new("/etc/telemt/config.toml"), + Path::new("/"), + false, + None, + ); + assert_eq!(resolved, PathBuf::from("/etc/telemt")); +} + +#[test] +fn expected_handshake_eof_matches_connection_reset() { + let err = ProxyError::Io(std::io::Error::from(std::io::ErrorKind::ConnectionReset)); + assert!(is_expected_handshake_eof(&err)); +} + +#[test] +fn expected_handshake_eof_matches_stream_io_unexpected_eof() { + let err = ProxyError::Stream(StreamError::Io(std::io::Error::from( + std::io::ErrorKind::UnexpectedEof, + ))); + assert!(is_expected_handshake_eof(&err)); +} + +#[test] +fn peer_close_description_is_human_readable_for_all_peer_close_kinds() { + let cases = [ + ( + std::io::ErrorKind::ConnectionReset, + "Peer reset TCP connection (RST)", + ), + ( + std::io::ErrorKind::ConnectionAborted, + "Peer aborted TCP connection during transport", + ), + ( + std::io::ErrorKind::BrokenPipe, + "Peer closed write side (broken pipe)", + ), + ( + std::io::ErrorKind::NotConnected, + "Socket was already closed by peer", + ), + ]; + + for (kind, expected) in cases { + let err = ProxyError::Io(std::io::Error::from(kind)); + assert_eq!(peer_close_description(&err), Some(expected)); + } +} + +#[test] +fn handshake_close_description_is_human_readable_for_all_expected_kinds() { + let cases = [ + ( + ProxyError::Io(std::io::Error::from(std::io::ErrorKind::UnexpectedEof)), + "Peer closed before sending full 64-byte MTProto handshake", + ), + ( + ProxyError::Io(std::io::Error::from(std::io::ErrorKind::ConnectionReset)), + "Peer reset TCP connection during initial MTProto handshake", + ), + ( + ProxyError::Io(std::io::Error::from(std::io::ErrorKind::ConnectionAborted)), + "Peer aborted TCP connection during initial MTProto handshake", + ), + ( + ProxyError::Io(std::io::Error::from(std::io::ErrorKind::BrokenPipe)), + "Peer closed write side before MTProto handshake completed", + ), + ( + ProxyError::Io(std::io::Error::from(std::io::ErrorKind::NotConnected)), + "Handshake socket was already closed by peer", + ), + ( + ProxyError::Stream(StreamError::UnexpectedEof), + "Peer closed before sending full 64-byte MTProto handshake", + ), + ]; + + for (err, expected) in cases { + assert_eq!(expected_handshake_close_description(&err), Some(expected)); + } +} diff --git a/src/maestro/listeners/bind.rs b/src/maestro/listeners/bind.rs index 4efb527..524498b 100644 --- a/src/maestro/listeners/bind.rs +++ b/src/maestro/listeners/bind.rs @@ -1,5 +1,13 @@ use std::error::Error; +#[cfg(unix)] +use std::io::{Error as IoError, ErrorKind}; use std::net::{IpAddr, SocketAddr}; +#[cfg(unix)] +use std::os::unix::fs::{FileTypeExt, MetadataExt}; +#[cfg(unix)] +use std::os::unix::net::UnixStream as StdUnixStream; +#[cfg(unix)] +use std::path::Path; use std::sync::Arc; use socket2::Socket; @@ -12,6 +20,8 @@ use crate::config::{ListenerTransport, ProxyConfig}; use crate::startup::{COMPONENT_LISTENERS_BIND, StartupTracker}; use crate::transport::find_listener_processes; use crate::transport::socket::{activate_listener_socket, bind_listener_socket}; +#[cfg(unix)] +use crate::util::secure_fs::AnchoredPath; use super::plan::{ListenerBindSpec, listener_bind_plan}; use crate::maestro::helpers::{print_proxy_links, print_web_proxy_links}; @@ -187,6 +197,64 @@ fn print_configured_links( print_proxy_links(host, port, config); } +#[cfg(unix)] +fn unix_path_identity(metadata: &std::fs::Metadata) -> (u64, u64) { + (metadata.dev(), metadata.ino()) +} + +#[cfg(unix)] +fn remove_stale_unix_socket(path: &Path) -> std::io::Result<()> { + let metadata = match std::fs::symlink_metadata(path) { + Ok(metadata) => metadata, + Err(error) if error.kind() == ErrorKind::NotFound => return Ok(()), + Err(error) => return Err(error), + }; + if !metadata.file_type().is_socket() { + return Err(IoError::new( + ErrorKind::AlreadyExists, + format!( + "refusing to remove non-socket Unix listener path {}", + path.display() + ), + )); + } + match StdUnixStream::connect(path) { + Ok(_) => { + return Err(IoError::new( + ErrorKind::AddrInUse, + format!("Unix listener {} is already active", path.display()), + )); + } + Err(error) if error.kind() == ErrorKind::ConnectionRefused => {} + Err(error) => return Err(error), + } + let current = std::fs::symlink_metadata(path)?; + if !current.file_type().is_socket() + || unix_path_identity(¤t) != unix_path_identity(&metadata) + { + return Err(IoError::new( + ErrorKind::AlreadyExists, + format!( + "Unix listener path {} changed during cleanup", + path.display() + ), + )); + } + std::fs::remove_file(path) +} + +#[cfg(unix)] +fn verify_bound_unix_socket(path: &Path, expected: (u64, u64)) -> std::io::Result<()> { + let metadata = std::fs::symlink_metadata(path)?; + if metadata.file_type().is_socket() && unix_path_identity(&metadata) == expected { + return Ok(()); + } + Err(IoError::new( + ErrorKind::AlreadyExists, + format!("Unix listener path {} was replaced", path.display()), + )) +} + /// Binds every eligible configured listener or fails without a partial inventory. pub(crate) async fn bind_listeners( config: &Arc, @@ -217,27 +285,45 @@ pub(crate) async fn bind_listeners( let mut unix_listener_out = None; #[cfg(unix)] if let Some(unix_path) = &config.server.listen_unix_sock { - let _ = tokio::fs::remove_file(unix_path).await; + let unix_path = Path::new(unix_path); + let anchored_path = AnchoredPath::open_trusted_parent(unix_path)?; + remove_stale_unix_socket(unix_path)?; let unix_listener = UnixListener::bind(unix_path)?; + let socket_metadata = std::fs::symlink_metadata(unix_path)?; + if !socket_metadata.file_type().is_socket() { + return Err(IoError::new( + ErrorKind::AlreadyExists, + format!("Unix listener path {} was replaced", unix_path.display()), + ) + .into()); + } + let socket_identity = unix_path_identity(&socket_metadata); if let Some(perm_str) = &config.server.listen_unix_sock_perm { match u32::from_str_radix(perm_str.trim_start_matches('0'), 8) { Ok(mode) => { - use std::os::unix::fs::PermissionsExt; - let permissions = std::fs::Permissions::from_mode(mode); - if let Err(error_value) = std::fs::set_permissions(unix_path, permissions) { + use nix::sys::stat::{FchmodatFlags, Mode, fchmodat}; + + verify_bound_unix_socket(unix_path, socket_identity)?; + if let Err(error_value) = fchmodat( + anchored_path.parent(), + anchored_path.name(), + Mode::from_bits_truncate(mode), + FchmodatFlags::NoFollowSymlink, + ) { error!( - path = %unix_path, + path = %unix_path.display(), permissions = %perm_str, error = %error_value, "Failed to set Unix socket permissions" ); } else { - info!(path = %unix_path, permissions = %perm_str, "Listening on Unix socket"); + verify_bound_unix_socket(unix_path, socket_identity)?; + info!(path = %unix_path.display(), permissions = %perm_str, "Listening on Unix socket"); } } Err(error_value) => { warn!( - path = %unix_path, + path = %unix_path.display(), permissions = %perm_str, error = %error_value, "Invalid Unix socket permissions; keeping umask-derived mode" @@ -245,7 +331,7 @@ pub(crate) async fn bind_listeners( } } } else { - info!(path = %unix_path, "Listening on Unix socket"); + info!(path = %unix_path.display(), "Listening on Unix socket"); } unix_listener_out = Some(unix_listener); } @@ -271,3 +357,54 @@ pub(crate) async fn bind_listeners( unix_listener: unix_listener_out, }) } + +#[cfg(all(test, unix))] +mod tests { + use std::os::unix::fs::symlink; + use std::os::unix::net::UnixListener as StdUnixListener; + + use super::*; + + #[test] + fn unix_socket_cleanup_refuses_regular_file_and_symlink() { + let directory = tempfile::tempdir().unwrap(); + let regular = directory.path().join("regular"); + let link = directory.path().join("listener.sock"); + std::fs::write(®ular, b"preserve").unwrap(); + symlink(®ular, &link).unwrap(); + + assert!(remove_stale_unix_socket(®ular).is_err()); + assert!(remove_stale_unix_socket(&link).is_err()); + assert_eq!(std::fs::read(®ular).unwrap(), b"preserve"); + assert!( + std::fs::symlink_metadata(&link) + .unwrap() + .file_type() + .is_symlink() + ); + } + + #[test] + fn unix_socket_cleanup_removes_only_stale_socket() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("listener.sock"); + let listener = StdUnixListener::bind(&path).unwrap(); + drop(listener); + + remove_stale_unix_socket(&path).unwrap(); + + assert!(!path.exists()); + } + + #[test] + fn unix_socket_cleanup_preserves_live_listener() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("listener.sock"); + let _listener = StdUnixListener::bind(&path).unwrap(); + + let error = remove_stale_unix_socket(&path).unwrap_err(); + + assert_eq!(error.kind(), ErrorKind::AddrInUse); + assert!(path.exists()); + } +} diff --git a/src/maestro/orchestrator.rs b/src/maestro/orchestrator.rs index cf0fc61..32c7869 100644 --- a/src/maestro/orchestrator.rs +++ b/src/maestro/orchestrator.rs @@ -3,8 +3,8 @@ use std::net::{IpAddr, SocketAddr}; use std::sync::Arc; use arc_swap::ArcSwap; -use tokio::sync::{RwLock, watch}; -use tracing::{error, info}; +use tokio::sync::{RwLock, Semaphore, watch}; +use tracing::{error, info, warn}; use crate::api; use crate::ip_tracker::UserIpTracker; @@ -12,6 +12,9 @@ use crate::network::probe::{decide_network_capabilities, log_probe_result, run_p use crate::proxy::direct_buffer_budget::{DirectBufferBudget, resolve_direct_buffer_hard_limit}; use crate::proxy::route_mode::{RelayRouteMode, RouteRuntimeController}; use crate::proxy::shared_state::ProxySharedState; +use crate::proxy::traffic_limiter::TrafficLimiter; +use crate::proxy::user_admission::UserAdmissionAuthority; +use crate::proxy::user_connection_authority::UserConnectionAuthority; use crate::startup::{COMPONENT_API_BOOTSTRAP, COMPONENT_NETWORK_PROBE}; use crate::stats::telemetry::TelemetryPolicy; use crate::stats::{QuotaStore, Stats}; @@ -37,7 +40,7 @@ pub(super) async fn run_telemt_core( process_started_at, process_started_at_epoch_secs, startup_tracker, - config, + mut config, config_path, has_rust_log, effective_log_level, @@ -45,11 +48,22 @@ pub(super) async fn run_telemt_core( logging_guard: _logging_guard, } = bootstrap::bootstrap(privilege_drop_requested).await?; + if privilege_drop_requested && config.server.conntrack_control.inline_conntrack_control { + warn!("Inline conntrack control is disabled when process privileges are dropped"); + config.server.conntrack_control.inline_conntrack_control = false; + } + let quota_store = Arc::new(QuotaStore::default()); - let stats = Arc::new(Stats::with_quota_store(quota_store.clone())); + let connection_authority = Arc::new(UserConnectionAuthority::default()); + let stats = Arc::new(Stats::with_process_authorities( + quota_store.clone(), + connection_authority, + )); let tls_full_cert_budget = Arc::new(TlsFullCertBudget::new()); let process_control_plane = control_plane::ProcessControlPlane::new(); let runtime_task_scope = generation::RuntimeTaskScope::new(); + let runtime_task_scope_guard = + generation::RuntimeTaskScopePreparationGuard::new(runtime_task_scope.clone()); stats.apply_telemetry_policy(TelemetryPolicy::from_config(&config.general.telemetry)); let quota_state_path = config.general.quota_state_path.clone(); let quota_state = @@ -71,14 +85,11 @@ pub(super) async fn run_telemt_core( .with_dns_overrides(&config.network.dns_overrides)?, ); let ip_tracker = Arc::new(UserIpTracker::new()); - ip_tracker - .load_limits( + let _ = ip_tracker + .apply_policy_from_source( + 1, config.access.user_max_unique_ips_global_each, &config.access.user_max_unique_ips, - ) - .await; - ip_tracker - .set_limit_policy( config.access.user_max_unique_ips_mode, config.access.user_max_unique_ips_window_secs, ) @@ -101,18 +112,36 @@ pub(super) async fn run_telemt_core( let direct_buffer_hard_limit = resolve_direct_buffer_hard_limit(config.general.direct_relay_buffer_budget_max_bytes).await; let direct_buffer_budget = DirectBufferBudget::new(direct_buffer_hard_limit); + direct_buffer_budget.activate_controller(1); info!( hard_limit_bytes = direct_buffer_hard_limit, configured_override_bytes = config.general.direct_relay_buffer_budget_max_bytes, "Direct relay buffer budget initialized" ); - let shared_state = - ProxySharedState::new_with_direct_buffer_budget(direct_buffer_budget.clone()); - shared_state.apply_user_enabled_config(&config.access.user_enabled); - shared_state.traffic_limiter.apply_policy( + let user_admission = UserAdmissionAuthority::new_with_quota_store(quota_store.clone()); + let traffic_limiter = TrafficLimiter::new(); + let _ = traffic_limiter.apply_policy_from_source( + 1, config.access.user_rate_limits.clone(), config.access.cidr_rate_limits.clone(), ); + let shared_state = ProxySharedState::new_with_process_authorities( + direct_buffer_budget.clone(), + traffic_limiter, + user_admission, + ); + let _ = shared_state.activate_user_config_source( + 1, + None, + &config.access.users, + &config.access.user_enabled, + ); + let max_connections_limit = if config.server.max_connections == 0 { + Semaphore::MAX_PERMITS + } else { + config.server.max_connections as usize + }; + let max_connections = Arc::new(Semaphore::new(max_connections_limit)); let web_trace = WebTraceStore::new(config.web.debug.clone(), &config.web.limits); let web_runtime_control = WebRuntimeControl::new(); @@ -294,6 +323,7 @@ pub(super) async fn run_telemt_core( ip_tracker.clone(), shared_state.clone(), direct_buffer_budget, + max_connections, route_runtime.clone(), api_me_pool.clone(), runtime_task_scope.clone(), @@ -324,6 +354,7 @@ pub(super) async fn run_telemt_core( runtime.max_connections, runtime_task_scope, ); + runtime_task_scope_guard.disarm(); let active_runtime = Arc::new(ArcSwap::from(runtime_generation)); let bound = listeners::bind_listeners( &runtime.config, @@ -341,9 +372,26 @@ pub(super) async fn run_telemt_core( .await .map_err(std::io::Error::other)?; + #[cfg(target_os = "linux")] + let conntrack_firewall = { + let authority = crate::conntrack_control::FirewallAuthority::spawn(&process_control_plane) + .map_err(std::io::Error::other)?; + if !authority + .publish_initial(1, runtime.config.clone(), stats.clone()) + .await + { + warn!( + "Initial conntrack firewall reconciliation failed; background retries remain active" + ); + } + Some(authority) + }; + #[cfg(not(target_os = "linux"))] + let conntrack_firewall = None::; + drop_after_bind(); - runtime_tasks::spawn_metrics_if_configured( + if let Err(error) = runtime_tasks::spawn_metrics_if_configured( &runtime.config, &startup_tracker, active_runtime.clone(), @@ -351,7 +399,15 @@ pub(super) async fn run_telemt_core( tls_full_cert_budget.clone(), process_control_plane.clone(), ) - .await?; + .await + { + if let Some(conntrack_firewall) = &conntrack_firewall + && !conntrack_firewall.shutdown_and_clear().await + { + warn!("Conntrack firewall cleanup failed after metrics startup error"); + } + return Err(error.into()); + } runtime_watch_tx.send_replace(Some(active_runtime.load_full().watch_state())); active_runtime_tx.send_replace(Some(active_runtime.clone())); @@ -375,6 +431,7 @@ pub(super) async fn run_telemt_core( runtime_watch_tx, listener_manager, web_trace, + conntrack_firewall.clone(), ); shutdown::spawn_signal_handlers( @@ -387,6 +444,7 @@ pub(super) async fn run_telemt_core( active_runtime, quota_state, reload_supervisor, + conntrack_firewall, process_control_plane, ) .await; diff --git a/src/maestro/reload_supervisor.rs b/src/maestro/reload_supervisor.rs index d186745..d50a5cf 100644 --- a/src/maestro/reload_supervisor.rs +++ b/src/maestro/reload_supervisor.rs @@ -8,6 +8,7 @@ use tokio::sync::watch; use tokio_util::sync::CancellationToken; use tracing::{info, warn}; +use crate::conntrack_control::FirewallAuthority; use crate::stats::QuotaStore; use crate::tls_front::cache::TlsFullCertBudget; use crate::web::trace::WebTraceStore; @@ -33,6 +34,7 @@ pub(crate) struct ReloadSupervisor { runtime_watch_tx: watch::Sender>, listener_manager: Arc>, web_trace: Arc, + conntrack_firewall: Option, } /// Process-owned handle that quiesces reloads before shutdown snapshots the runtime. @@ -106,6 +108,7 @@ impl ReloadSupervisor { runtime_watch_tx: watch::Sender>, listener_manager: ListenerManager, web_trace: Arc, + conntrack_firewall: Option, ) -> ReloadSupervisorHandle { let listener_manager = Arc::new(Mutex::new(listener_manager)); let supervisor = Self { @@ -120,6 +123,7 @@ impl ReloadSupervisor { runtime_watch_tx, listener_manager: listener_manager.clone(), web_trace, + conntrack_firewall, }; let control = supervisor.control.clone(); let shutdown = CancellationToken::new(); @@ -175,8 +179,14 @@ impl ReloadSupervisor { resolved.effective, &self.config_path, self.quota_store.clone(), + old_runtime.stats.connection_authority(), self.runtime_log_filter.clone(), self.tls_full_cert_budget.clone(), + old_runtime.proxy_shared.user_admission(), + old_runtime.ip_tracker.clone(), + old_runtime.proxy_shared.traffic_limiter.clone(), + old_runtime.proxy_shared.direct_buffer_budget.clone(), + old_runtime.max_connections.clone(), ) .await { @@ -277,6 +287,7 @@ impl ReloadSupervisor { generation: new_runtime, detected_ips, config_watcher_activation, + user_admission_epoch, } = prepared; let pending_listener_transition = if let Some(listener_transition) = listener_transition { match self @@ -298,11 +309,48 @@ impl ReloadSupervisor { } else { None }; + let config = new_runtime.config(); + let _ = new_runtime.proxy_shared.activate_user_config_source( + new_runtime.id, + Some(user_admission_epoch), + &config.access.users, + &config.access.user_enabled, + ); + let _ = new_runtime + .ip_tracker + .apply_policy_from_source( + new_runtime.id, + config.access.user_max_unique_ips_global_each, + &config.access.user_max_unique_ips, + config.access.user_max_unique_ips_mode, + config.access.user_max_unique_ips_window_secs, + ) + .await; + let _ = new_runtime + .proxy_shared + .traffic_limiter + .apply_policy_from_source( + new_runtime.id, + config.access.user_rate_limits.clone(), + config.access.cidr_rate_limits.clone(), + ); + new_runtime + .proxy_shared + .direct_buffer_budget + .activate_controller(new_runtime.id); let replaced = { let listener_manager = self.listener_manager.lock().await; old_runtime.stop_accepting_sessions(); listener_manager.activate_runtime_generation(new_runtime.clone()) }; + let conntrack_firewall_published = match &self.conntrack_firewall { + Some(conntrack_firewall) => conntrack_firewall.publish( + new_runtime.id, + new_runtime.config(), + new_runtime.stats.clone(), + ), + None => true, + }; self.web_trace .apply_policy(new_runtime.id, &new_runtime.config().web.debug); config_watcher_activation.send_replace(true); @@ -317,6 +365,12 @@ impl ReloadSupervisor { .apply_reload(&new_runtime.config().general.log_level); self.runtime_watch_tx .send_replace(Some(new_runtime.watch_state())); + if !conntrack_firewall_published { + let warning = + "conntrack firewall reconciler is unavailable after runtime activation".to_string(); + warn!(reload_id = command.reload_id, warning = %warning); + self.control.add_warning(command.reload_id, warning).await; + } info!( reload_id = command.reload_id, diff --git a/src/maestro/reload_supervisor_tests.rs b/src/maestro/reload_supervisor_tests.rs index 27b6e4b..2842409 100644 --- a/src/maestro/reload_supervisor_tests.rs +++ b/src/maestro/reload_supervisor_tests.rs @@ -23,10 +23,12 @@ fn runtime_log_filter() -> RuntimeLogFilter { fn prepared_runtime(generation: Arc) -> PreparedRuntime { let (config_watcher_activation, _activation_rx) = watch::channel(false); + let user_admission_epoch = generation.proxy_shared.user_admission().epoch(); PreparedRuntime { generation, detected_ips: (None, None), config_watcher_activation, + user_admission_epoch, } } @@ -59,6 +61,7 @@ async fn fixture(request: ReloadRequest) -> ReloadFixture { runtime_watch_tx, listener_manager, web_trace, + conntrack_firewall: None, }); let command = ReloadCommand { reload_id: accepted.reload_id, @@ -273,6 +276,7 @@ async fn quiesce_joins_idle_supervisor_and_rejects_later_submissions() { runtime.config().web.debug.clone(), &runtime.config().web.limits, ), + None, ); tokio::time::timeout(Duration::from_secs(1), handle.quiesce()) diff --git a/src/maestro/runtime_build.rs b/src/maestro/runtime_build.rs index 638478f..dd6edc5 100644 --- a/src/maestro/runtime_build.rs +++ b/src/maestro/runtime_build.rs @@ -11,11 +11,12 @@ use crate::config::{ use crate::crypto::SecureRandom; use crate::ip_tracker::UserIpTracker; use crate::network::probe::{decide_network_capabilities, run_probe}; -use crate::proxy::direct_buffer_budget::{ - DirectBufferBudget, resolve_direct_buffer_hard_limit, run_direct_buffer_budget_controller, -}; +use crate::proxy::direct_buffer_budget::{DirectBufferBudget, run_direct_buffer_budget_controller}; use crate::proxy::route_mode::{RelayRouteMode, RouteRuntimeController}; use crate::proxy::shared_state::ProxySharedState; +use crate::proxy::traffic_limiter::TrafficLimiter; +use crate::proxy::user_admission::UserAdmissionAuthority; +use crate::proxy::user_connection_authority::UserConnectionAuthority; use crate::startup::StartupTracker; use crate::stats::beobachten::BeobachtenStore; use crate::stats::telemetry::TelemetryPolicy; @@ -26,7 +27,7 @@ use crate::transport::UpstreamManager; use crate::transport::middle_proxy::MePool; use super::admission; -use super::generation::{RuntimeGeneration, RuntimeTaskScope}; +use super::generation::{RuntimeGeneration, RuntimeTaskScope, RuntimeTaskScopePreparationGuard}; use super::listeners::listener_rebind_supported; use super::runtime_tasks::RuntimeLogFilter; use super::{me_startup, runtime_tasks, tls_bootstrap}; @@ -39,6 +40,8 @@ pub(crate) struct PreparedRuntime { pub(crate) detected_ips: (Option, Option), /// Gate opened only after the candidate becomes the active generation. pub(crate) config_watcher_activation: watch::Sender, + /// User-authority epoch captured before candidate construction. + pub(crate) user_admission_epoch: u64, } pub(crate) async fn prepare_runtime( @@ -46,9 +49,16 @@ pub(crate) async fn prepare_runtime( config: ProxyConfig, config_path: &Path, quota_store: Arc, + connection_authority: Arc, runtime_log_filter: RuntimeLogFilter, tls_full_cert_budget: Arc, + user_admission: Arc, + ip_tracker: Arc, + traffic_limiter: Arc, + direct_buffer_budget: Arc, + max_connections: Arc, ) -> Result { + let user_admission_epoch = user_admission.epoch(); config .validate_web_decoy_listener_separation() .map_err(|error| error.to_string())?; @@ -58,7 +68,11 @@ pub(crate) async fn prepare_runtime( .as_secs(); let startup_tracker = Arc::new(StartupTracker::new(started_at_epoch_secs)); let task_scope = RuntimeTaskScope::new(); - let stats = Arc::new(Stats::with_quota_store(quota_store)); + let task_scope_guard = RuntimeTaskScopePreparationGuard::new(task_scope.clone()); + let stats = Arc::new(Stats::with_process_authorities( + quota_store, + connection_authority, + )); stats.apply_telemetry_policy(TelemetryPolicy::from_config(&config.general.telemetry)); let upstream_manager = Arc::new( @@ -75,29 +89,10 @@ pub(crate) async fn prepare_runtime( .with_dns_overrides(&config.network.dns_overrides) .map_err(|error| format!("DNS override preparation failed: {}", error))?, ); - let ip_tracker = Arc::new(UserIpTracker::new()); - ip_tracker - .load_limits( - config.access.user_max_unique_ips_global_each, - &config.access.user_max_unique_ips, - ) - .await; - ip_tracker - .set_limit_policy( - config.access.user_max_unique_ips_mode, - config.access.user_max_unique_ips_window_secs, - ) - .await; - - let hard_limit = - resolve_direct_buffer_hard_limit(config.general.direct_relay_buffer_budget_max_bytes).await; - let direct_buffer_budget = DirectBufferBudget::new(hard_limit); - let proxy_shared = - ProxySharedState::new_with_direct_buffer_budget(direct_buffer_budget.clone()); - proxy_shared.apply_user_enabled_config(&config.access.user_enabled); - proxy_shared.traffic_limiter.apply_policy( - config.access.user_rate_limits.clone(), - config.access.cidr_rate_limits.clone(), + let proxy_shared = ProxySharedState::new_with_process_authorities( + direct_buffer_budget.clone(), + traffic_limiter, + user_admission, ); let probe = run_probe( @@ -178,14 +173,9 @@ pub(crate) async fn prepare_runtime( Duration::from_secs(config.access.replay_window_secs), )); let buffer_pool = Arc::new(BufferPool::with_config(64 * 1024, 4096)); - let max_connections_limit = if config.server.max_connections == 0 { - Semaphore::MAX_PERMITS - } else { - config.server.max_connections as usize - }; - let max_connections = Arc::new(Semaphore::new(max_connections_limit)); let (config_watcher_activation, config_watcher_activation_rx) = watch::channel(false); let watches = runtime_tasks::spawn_runtime_tasks( + generation_id, &config, config_path, &probe, @@ -281,10 +271,12 @@ pub(crate) async fn prepare_runtime( conntrack_scope.cancellation_token(), )); task_scope.spawn(run_direct_buffer_budget_controller( + generation_id, direct_buffer_budget, buffer_pool.clone(), stats.clone(), proxy_shared.clone(), + max_connections.clone(), config.server.max_connections, )); let generation = RuntimeGeneration::new( @@ -306,11 +298,13 @@ pub(crate) async fn prepare_runtime( max_connections, task_scope, ); + task_scope_guard.disarm(); drop(admission_tx); Ok(PreparedRuntime { generation, config_watcher_activation, + user_admission_epoch, detected_ips: ( probe.detected_ipv4.map(IpAddr::V4), probe.detected_ipv6.map(IpAddr::V6), @@ -406,6 +400,23 @@ pub(crate) fn resolve_reload_config( effective.server.metrics_listen = old.server.metrics_listen.clone(); effective.server.metrics_port = old.server.metrics_port; } + if old.server.max_connections != desired.server.max_connections { + fields.push("server.max_connections".to_string()); + effective.server.max_connections = old.server.max_connections; + } + if serde_json::to_value(&old.server.conntrack_control).ok() + != serde_json::to_value(&desired.server.conntrack_control).ok() + { + fields.push("server.conntrack_control".to_string()); + effective.server.conntrack_control = old.server.conntrack_control.clone(); + } + if old.general.direct_relay_buffer_budget_max_bytes + != desired.general.direct_relay_buffer_budget_max_bytes + { + fields.push("general.direct_relay_buffer_budget_max_bytes".to_string()); + effective.general.direct_relay_buffer_budget_max_bytes = + old.general.direct_relay_buffer_budget_max_bytes; + } if old.general.quota_state_path != desired.general.quota_state_path { fields.push("general.quota_state_path".to_string()); effective.general.quota_state_path = old.general.quota_state_path.clone(); diff --git a/src/maestro/runtime_build_tests.rs b/src/maestro/runtime_build_tests.rs index babaa54..1554ddc 100644 --- a/src/maestro/runtime_build_tests.rs +++ b/src/maestro/runtime_build_tests.rs @@ -94,6 +94,70 @@ fn global_mss_profiles_are_deferred_with_the_listener_socket_group() { assert!(!resolved.runtime_changed); } +#[test] +fn process_wide_connection_and_direct_buffer_envelopes_are_restart_only() { + let old = ProxyConfig::default(); + let mut desired = old.clone(); + desired.server.max_connections = old.server.max_connections.saturating_add(1); + desired.general.direct_relay_buffer_budget_max_bytes = old + .general + .direct_relay_buffer_budget_max_bytes + .saturating_add(4 * 1024); + + let resolved = resolve_reload_config(&old, &desired).unwrap(); + + assert_eq!( + resolved.deferred_process_fields, + vec![ + "server.max_connections".to_string(), + "general.direct_relay_buffer_budget_max_bytes".to_string(), + ] + ); + assert_eq!( + resolved.effective.server.max_connections, + old.server.max_connections + ); + assert_eq!( + resolved + .effective + .general + .direct_relay_buffer_budget_max_bytes, + old.general.direct_relay_buffer_budget_max_bytes + ); + assert!(!resolved.runtime_changed); +} + +#[test] +fn conntrack_control_policy_is_restart_only_as_one_process_owned_unit() { + let old = ProxyConfig::default(); + let mut desired = old.clone(); + desired.server.conntrack_control.inline_conntrack_control = + !old.server.conntrack_control.inline_conntrack_control; + desired.server.conntrack_control.mode = crate::config::ConntrackMode::Notrack; + desired.server.conntrack_control.backend = crate::config::ConntrackBackend::Iptables; + desired.server.conntrack_control.profile = crate::config::ConntrackPressureProfile::Aggressive; + desired.server.conntrack_control.hybrid_listener_ips = vec!["192.0.2.10".parse().unwrap()]; + desired.server.conntrack_control.pressure_high_watermark_pct = 90; + desired.server.conntrack_control.pressure_low_watermark_pct = 40; + desired.server.conntrack_control.delete_budget_per_sec = old + .server + .conntrack_control + .delete_budget_per_sec + .saturating_add(1); + + let resolved = resolve_reload_config(&old, &desired).unwrap(); + + assert_eq!( + resolved.deferred_process_fields, + vec!["server.conntrack_control".to_string()] + ); + assert_eq!( + serde_json::to_value(&resolved.effective.server.conntrack_control).unwrap(), + serde_json::to_value(&old.server.conntrack_control).unwrap() + ); + assert!(!resolved.runtime_changed); +} + #[test] fn mixed_reload_retains_process_state_and_applies_runtime_state() { let old = ProxyConfig::default(); @@ -231,6 +295,28 @@ fn web_decoy_fasttrack_mode_is_deferred_without_runtime_publication() { assert!(!resolved.runtime_changed); } +#[test] +fn base_path_change_is_runtime_owned_and_rebuilds_route_identity() { + let old = web_config_with_fasttrack("off"); + let old_runtime = old.web.runtime.as_ref().unwrap(); + let old_capability = old_runtime.vhosts["proxy.example.com"].capabilities[0]; + let mut desired = old.clone(); + desired.web.vhosts[0].base_path = "MixedCase/path".to_string(); + + let resolved = resolve_reload_config(&old, &desired).unwrap(); + + assert!(resolved.deferred_process_fields.is_empty()); + assert!(resolved.runtime_changed); + assert_eq!(resolved.effective.web.vhosts[0].base_path, "MixedCase/path"); + let runtime = resolved.effective.web.runtime.as_ref().unwrap(); + let vhost = &runtime.vhosts["proxy.example.com"]; + assert_eq!(vhost.base, "/MixedCase/path/"); + assert_ne!(vhost.capabilities[0], old_capability); + assert_eq!(vhost.capabilities[0], vhost.profiles[0].capability); + assert_eq!(runtime.capabilities.as_ref(), vhost.capabilities.as_ref()); + assert!(!runtime.capabilities.contains(&old_capability)); +} + #[test] fn enabling_learning_is_deferred_when_retained_capacity_is_too_small() { let mut old = ProxyConfig::default(); diff --git a/src/maestro/runtime_startup.rs b/src/maestro/runtime_startup.rs index bfa4485..d580fb0 100644 --- a/src/maestro/runtime_startup.rs +++ b/src/maestro/runtime_startup.rs @@ -56,6 +56,7 @@ pub(super) async fn prepare_runtime( ip_tracker: Arc, shared_state: Arc, direct_buffer_budget: Arc, + max_connections: Arc, route_runtime: Arc, api_me_pool: Arc>>>, runtime_task_scope: RuntimeTaskScope, @@ -69,13 +70,6 @@ pub(super) async fn prepare_runtime( let beobachten = Arc::new(BeobachtenStore::new()); let rng = Arc::new(SecureRandom::new()); - let max_connections_limit = if config.server.max_connections == 0 { - Semaphore::MAX_PERMITS - } else { - config.server.max_connections as usize - }; - let max_connections = Arc::new(Semaphore::new(max_connections_limit)); - let me2dc_fallback = config.general.me2dc_fallback; let me_init_retry_attempts = config.general.me_init_retry_attempts; if use_middle_proxy && !decision.ipv4_me && !decision.ipv6_me { @@ -229,6 +223,7 @@ pub(super) async fn prepare_runtime( } let runtime_watches = runtime_tasks::spawn_runtime_tasks( + 1, &config, config_path, probe, @@ -345,10 +340,12 @@ pub(super) async fn prepare_runtime( conntrack_scope.cancellation_token(), )); runtime_task_scope.spawn(run_direct_buffer_budget_controller( + 1, direct_buffer_budget, buffer_pool.clone(), stats, shared_state, + max_connections.clone(), config.server.max_connections, )); diff --git a/src/maestro/runtime_tasks.rs b/src/maestro/runtime_tasks.rs index cc2a124..891e6d6 100644 --- a/src/maestro/runtime_tasks.rs +++ b/src/maestro/runtime_tasks.rs @@ -90,6 +90,7 @@ impl RuntimeLogFilter { #[allow(clippy::too_many_arguments)] pub(crate) async fn spawn_runtime_tasks( + generation_id: u64, config: &Arc, config_path: &Path, probe: &NetworkProbe, @@ -137,7 +138,9 @@ pub(crate) async fn spawn_runtime_tasks( let ip_tracker_maintenance = ip_tracker.clone(); task_scope.spawn(async move { - ip_tracker_maintenance.run_periodic_maintenance().await; + ip_tracker_maintenance + .run_periodic_maintenance(generation_id) + .await; }); let detected_ip_v4: Option = probe.detected_ipv4.map(IpAddr::V4); @@ -196,20 +199,7 @@ pub(crate) async fn spawn_runtime_tasks( let ip_tracker_policy = ip_tracker.clone(); let mut config_rx_ip_limits = config_rx.clone(); task_scope.spawn(async move { - let mut prev_limits = config_rx_ip_limits - .borrow() - .access - .user_max_unique_ips - .clone(); - let mut prev_global_each = config_rx_ip_limits - .borrow() - .access - .user_max_unique_ips_global_each; - let mut prev_mode = config_rx_ip_limits.borrow().access.user_max_unique_ips_mode; - let mut prev_window = config_rx_ip_limits - .borrow() - .access - .user_max_unique_ips_window_secs; + let mut previous = config_rx_ip_limits.borrow().access.clone(); loop { if config_rx_ip_limits.changed().await.is_err() { @@ -217,39 +207,28 @@ pub(crate) async fn spawn_runtime_tasks( } let cfg = config_rx_ip_limits.borrow_and_update().clone(); - if prev_limits != cfg.access.user_max_unique_ips - || prev_global_each != cfg.access.user_max_unique_ips_global_each + if previous.user_max_unique_ips != cfg.access.user_max_unique_ips + || previous.user_max_unique_ips_global_each + != cfg.access.user_max_unique_ips_global_each + || previous.user_max_unique_ips_mode != cfg.access.user_max_unique_ips_mode + || previous.user_max_unique_ips_window_secs + != cfg.access.user_max_unique_ips_window_secs { - ip_tracker_policy - .load_limits( + let _ = ip_tracker_policy + .apply_policy_from_source( + generation_id, cfg.access.user_max_unique_ips_global_each, &cfg.access.user_max_unique_ips, - ) - .await; - prev_limits = cfg.access.user_max_unique_ips.clone(); - prev_global_each = cfg.access.user_max_unique_ips_global_each; - } - - if prev_mode != cfg.access.user_max_unique_ips_mode - || prev_window != cfg.access.user_max_unique_ips_window_secs - { - ip_tracker_policy - .set_limit_policy( cfg.access.user_max_unique_ips_mode, cfg.access.user_max_unique_ips_window_secs, ) .await; - prev_mode = cfg.access.user_max_unique_ips_mode; - prev_window = cfg.access.user_max_unique_ips_window_secs; + previous = cfg.access.clone(); } } }); let limiter = shared_state.traffic_limiter.clone(); - limiter.apply_policy( - config.access.user_rate_limits.clone(), - config.access.cidr_rate_limits.clone(), - ); let mut config_rx_rate_limits = config_rx.clone(); task_scope.spawn(async move { let mut prev_user_limits = config_rx_rate_limits @@ -270,7 +249,8 @@ pub(crate) async fn spawn_runtime_tasks( if prev_user_limits != cfg.access.user_rate_limits || prev_cidr_limits != cfg.access.cidr_rate_limits { - limiter.apply_policy( + let _ = limiter.apply_policy_from_source( + generation_id, cfg.access.user_rate_limits.clone(), cfg.access.cidr_rate_limits.clone(), ); @@ -288,8 +268,14 @@ pub(crate) async fn spawn_runtime_tasks( break; } let cfg = config_rx_user_enabled.borrow_and_update().clone(); - for user in shared_user_enabled.apply_user_enabled_config(&cfg.access.user_enabled) { - let cancelled = shared_user_enabled.cancel_user_sessions(&user); + let Some(cancelled_users) = shared_user_enabled.apply_user_config_from_source( + generation_id, + &cfg.access.users, + &cfg.access.user_enabled, + ) else { + continue; + }; + for (user, cancelled) in cancelled_users { if cancelled > 0 { info!( user = %user, diff --git a/src/maestro/shutdown.rs b/src/maestro/shutdown.rs index 226060e..6264d7d 100644 --- a/src/maestro/shutdown.rs +++ b/src/maestro/shutdown.rs @@ -23,6 +23,7 @@ use super::control_plane::ProcessControlPlane; use super::generation::RuntimeGeneration; use super::helpers::{format_uptime, unit_label}; use super::reload_supervisor::ReloadSupervisorHandle; +use crate::conntrack_control::FirewallAuthority; use crate::quota_state::QuotaStateOwner; use crate::stats::Stats; use crate::synlimit_control; @@ -54,6 +55,7 @@ pub(crate) async fn wait_for_shutdown( active_runtime: Arc>, quota_state: Arc, reload_supervisor: ReloadSupervisorHandle, + conntrack_firewall: Option, process_control_plane: ProcessControlPlane, ) { let signal = wait_for_shutdown_signal().await; @@ -63,6 +65,7 @@ pub(crate) async fn wait_for_shutdown( active_runtime, quota_state, reload_supervisor, + conntrack_firewall, process_control_plane, ) .await; @@ -95,6 +98,7 @@ async fn perform_shutdown( active_runtime: Arc>, quota_state: Arc, reload_supervisor: ReloadSupervisorHandle, + conntrack_firewall: Option, process_control_plane: ProcessControlPlane, ) { let shutdown_started_at = Instant::now(); @@ -126,6 +130,12 @@ async fn perform_shutdown( warn!("ME shutdown: pool lifecycle deadline expired"); } + if let Some(conntrack_firewall) = conntrack_firewall + && !conntrack_firewall.shutdown_and_clear().await + { + warn!("Conntrack firewall cleanup did not complete successfully"); + } + if let Err(error) = synlimit_control::clear_synlimit_rules_all_backends().await { warn!(error = %error, "Failed to clear SYN limiter rules during shutdown"); } diff --git a/src/main.rs b/src/main.rs index c0453d9..8e2f307 100644 --- a/src/main.rs +++ b/src/main.rs @@ -27,6 +27,7 @@ mod protocol; mod proxy; mod quota_state; mod service; +mod slot_budget; mod startup; mod stats; mod stream; diff --git a/src/metrics.rs b/src/metrics.rs index dfa69ba..74a7796 100644 --- a/src/metrics.rs +++ b/src/metrics.rs @@ -236,6 +236,10 @@ async fn handle( let config = runtime.config(); if req.uri().path() == "/metrics" { + let me_hardswap = match runtime.current_me_pool().await { + Some(pool) => Some(pool.api_hardswap_snapshot().await), + None => None, + }; let body = render_metrics( stats, shared_state, @@ -244,6 +248,7 @@ async fn handle( tls_cache, tls_full_cert_budget, web_publication, + me_hardswap.as_ref(), ) .await; let resp = Response::builder() diff --git a/src/metrics/render.rs b/src/metrics/render.rs index e7af27d..b900143 100644 --- a/src/metrics/render.rs +++ b/src/metrics/render.rs @@ -12,6 +12,8 @@ mod me_lifecycle; mod me_buffers; // ME writer selection, KDF, and hardswap metrics. mod me_policy; +// Live hardswap ownership and replacement progress metrics. +mod me_hardswap; // Adaptive-floor and writer-cap metrics. mod me_floor; // Desync, pool recovery, and refill metrics. @@ -27,6 +29,7 @@ pub(super) async fn render_metrics( tls_cache: Option<&TlsFrontCache>, tls_full_cert_budget: &TlsFullCertBudget, web_publication: &crate::web::control::WebRuntimePublication, + me_hardswap: Option<&crate::transport::middle_proxy::MeApiHardswapSnapshot>, ) -> String { let mut out = String::with_capacity(4096); let telemetry = stats.telemetry_policy(); @@ -62,6 +65,7 @@ pub(super) async fn render_metrics( me_allows_debug, ); me_policy::render(&mut out, stats, me_allows_normal, me_allows_debug); + me_hardswap::render(&mut out, me_hardswap, me_allows_normal); me_floor::render(&mut out, stats, config, me_allows_normal); me_recovery::render(&mut out, stats, me_allows_normal, me_allows_debug); users::render( diff --git a/src/metrics/render/connections.rs b/src/metrics/render/connections.rs index f7c2ae0..2285291 100644 --- a/src/metrics/render/connections.rs +++ b/src/metrics/render/connections.rs @@ -266,6 +266,53 @@ pub(super) fn render( } ); + let _ = writeln!( + out, + "# HELP telemt_conntrack_rule_reconcile_total Conntrack firewall reconciliations by result" + ); + let _ = writeln!(out, "# TYPE telemt_conntrack_rule_reconcile_total counter"); + let _ = writeln!( + out, + "telemt_conntrack_rule_reconcile_total{{result=\"success\"}} {}", + if core_enabled { + stats.get_conntrack_rule_reconcile_success_total() + } else { + 0 + } + ); + let _ = writeln!( + out, + "telemt_conntrack_rule_reconcile_total{{result=\"error\"}} {}", + if core_enabled { + stats.get_conntrack_rule_reconcile_error_total() + } else { + 0 + } + ); + let _ = writeln!( + out, + "# HELP telemt_conntrack_rule_rollback_total Conntrack firewall rollbacks by result" + ); + let _ = writeln!(out, "# TYPE telemt_conntrack_rule_rollback_total counter"); + let _ = writeln!( + out, + "telemt_conntrack_rule_rollback_total{{result=\"success\"}} {}", + if core_enabled { + stats.get_conntrack_rule_rollback_success_total() + } else { + 0 + } + ); + let _ = writeln!( + out, + "telemt_conntrack_rule_rollback_total{{result=\"error\"}} {}", + if core_enabled { + stats.get_conntrack_rule_rollback_error_total() + } else { + 0 + } + ); + let _ = writeln!( out, "# HELP telemt_conntrack_event_queue_depth Pending close events in conntrack control queue" diff --git a/src/metrics/render/me_hardswap.rs b/src/metrics/render/me_hardswap.rs new file mode 100644 index 0000000..75c66af --- /dev/null +++ b/src/metrics/render/me_hardswap.rs @@ -0,0 +1,149 @@ +use std::fmt::Write; + +use crate::transport::middle_proxy::MeApiHardswapSnapshot; + +/// Renders fixed-cardinality hardswap and writer-replacement gauges. +pub(super) fn render(out: &mut String, snapshot: Option<&MeApiHardswapSnapshot>, enabled: bool) { + let snapshot = enabled.then_some(snapshot).flatten(); + let pending = snapshot.is_some_and(|value| value.pending); + let pending_age_secs = snapshot + .and_then(|value| value.pending_age_secs) + .unwrap_or(0); + let pending_writers_current = snapshot + .map(|value| value.pending_writers_current) + .unwrap_or(0); + let pending_writer_deficit = snapshot + .map(|value| value.pending_writer_deficit) + .unwrap_or(0); + let pending_missing_dc_groups = snapshot + .map(|value| value.pending_missing_dc_groups) + .unwrap_or(0); + let pending_map_current = snapshot + .and_then(|value| value.pending_map_current) + .is_some_and(|value| value); + let orphan_warm_writers_current = snapshot + .map(|value| value.orphan_warm_writers_current) + .unwrap_or(0); + let replacement_preparing_current = snapshot + .map(|value| value.replacement_preparing_current) + .unwrap_or(0); + let replacement_retiring_current = snapshot + .map(|value| value.replacement_retiring_current) + .unwrap_or(0); + + let _ = writeln!( + out, + "# HELP telemt_me_hardswap_pending Whether an ME hardswap generation is pending" + ); + let _ = writeln!(out, "# TYPE telemt_me_hardswap_pending gauge"); + let _ = writeln!(out, "telemt_me_hardswap_pending {}", usize::from(pending)); + let _ = writeln!( + out, + "# HELP telemt_me_hardswap_pending_age_seconds Age of the pending ME hardswap generation" + ); + let _ = writeln!(out, "# TYPE telemt_me_hardswap_pending_age_seconds gauge"); + let _ = writeln!( + out, + "telemt_me_hardswap_pending_age_seconds {pending_age_secs}" + ); + let _ = writeln!( + out, + "# HELP telemt_me_hardswap_pending_writers_current Authoritative warm writers in the pending generation" + ); + let _ = writeln!( + out, + "# TYPE telemt_me_hardswap_pending_writers_current gauge" + ); + let _ = writeln!( + out, + "telemt_me_hardswap_pending_writers_current {pending_writers_current}" + ); + let _ = writeln!( + out, + "# HELP telemt_me_hardswap_pending_writer_deficit Writers missing from the pending generation floor" + ); + let _ = writeln!( + out, + "# TYPE telemt_me_hardswap_pending_writer_deficit gauge" + ); + let _ = writeln!( + out, + "telemt_me_hardswap_pending_writer_deficit {pending_writer_deficit}" + ); + let _ = writeln!( + out, + "# HELP telemt_me_hardswap_pending_missing_dc_groups Desired DC-family groups below the pending-generation floor" + ); + let _ = writeln!( + out, + "# TYPE telemt_me_hardswap_pending_missing_dc_groups gauge" + ); + let _ = writeln!( + out, + "telemt_me_hardswap_pending_missing_dc_groups {pending_missing_dc_groups}" + ); + let _ = writeln!( + out, + "# HELP telemt_me_hardswap_pending_map_current Whether the pending generation targets the current endpoint map" + ); + let _ = writeln!(out, "# TYPE telemt_me_hardswap_pending_map_current gauge"); + let _ = writeln!( + out, + "telemt_me_hardswap_pending_map_current {}", + usize::from(pending_map_current) + ); + let _ = writeln!( + out, + "# HELP telemt_me_hardswap_orphan_warm_writers_current Warm writers not owned by the pending hardswap generation" + ); + let _ = writeln!( + out, + "# TYPE telemt_me_hardswap_orphan_warm_writers_current gauge" + ); + let _ = writeln!( + out, + "telemt_me_hardswap_orphan_warm_writers_current {orphan_warm_writers_current}" + ); + let _ = writeln!( + out, + "# HELP telemt_me_writer_replacement_current ME writer replacements by transaction phase" + ); + let _ = writeln!(out, "# TYPE telemt_me_writer_replacement_current gauge"); + let _ = writeln!( + out, + "telemt_me_writer_replacement_current{{state=\"preparing\"}} {replacement_preparing_current}" + ); + let _ = writeln!( + out, + "telemt_me_writer_replacement_current{{state=\"retiring\"}} {replacement_retiring_current}" + ); +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn renders_bounded_hardswap_and_replacement_state() { + let snapshot = MeApiHardswapSnapshot { + pending: true, + pending_age_secs: Some(42), + pending_writers_current: 3, + pending_writer_deficit: 4, + pending_missing_dc_groups: 2, + pending_map_current: Some(true), + orphan_warm_writers_current: 1, + replacement_preparing_current: 5, + replacement_retiring_current: 6, + }; + let mut out = String::new(); + + render(&mut out, Some(&snapshot), true); + + assert!(out.contains("telemt_me_hardswap_pending 1")); + assert!(out.contains("telemt_me_hardswap_pending_age_seconds 42")); + assert!(out.contains("telemt_me_hardswap_pending_writer_deficit 4")); + assert!(out.contains("telemt_me_writer_replacement_current{state=\"preparing\"} 5")); + assert!(out.contains("telemt_me_writer_replacement_current{state=\"retiring\"} 6")); + } +} diff --git a/src/metrics/render/traffic.rs b/src/metrics/render/traffic.rs index 29636f6..136fd08 100644 --- a/src/metrics/render/traffic.rs +++ b/src/metrics/render/traffic.rs @@ -118,6 +118,49 @@ pub(super) fn render( } ); + let _ = writeln!( + out, + "# HELP telemt_rate_limiter_cas_retry_exhausted_total Traffic limiter operations that exhausted their bounded CAS attempt budget" + ); + let _ = writeln!( + out, + "# TYPE telemt_rate_limiter_cas_retry_exhausted_total counter" + ); + for (scope, direction, reserve, refund) in [ + ( + "user", + "up", + limiter_metrics.user_reserve_cas_retry_exhausted_up_total, + limiter_metrics.user_refund_cas_retry_exhausted_up_total, + ), + ( + "user", + "down", + limiter_metrics.user_reserve_cas_retry_exhausted_down_total, + limiter_metrics.user_refund_cas_retry_exhausted_down_total, + ), + ( + "cidr", + "up", + limiter_metrics.cidr_reserve_cas_retry_exhausted_up_total, + limiter_metrics.cidr_refund_cas_retry_exhausted_up_total, + ), + ( + "cidr", + "down", + limiter_metrics.cidr_reserve_cas_retry_exhausted_down_total, + limiter_metrics.cidr_refund_cas_retry_exhausted_down_total, + ), + ] { + for (operation, value) in [("reserve", reserve), ("refund", refund)] { + let _ = writeln!( + out, + "telemt_rate_limiter_cas_retry_exhausted_total{{scope=\"{scope}\",direction=\"{direction}\",operation=\"{operation}\"}} {}", + if core_enabled { value } else { 0 } + ); + } + } + let _ = writeln!( out, "# HELP telemt_rate_limiter_active_leases Active relay leases under rate limiting by scope" diff --git a/src/metrics/render/users.rs b/src/metrics/render/users.rs index bdfd5e8..87897e8 100644 --- a/src/metrics/render/users.rs +++ b/src/metrics/render/users.rs @@ -153,7 +153,7 @@ pub(super) async fn render( out, "telemt_user_connections_current{{user=\"{}\"}} {}", user, - s.curr_connects.load(std::sync::atomic::Ordering::Relaxed) + stats.get_process_user_curr_connects(user) ); let _ = writeln!( out, diff --git a/src/metrics/tests.rs b/src/metrics/tests.rs index 04a0293..5d974e9 100644 --- a/src/metrics/tests.rs +++ b/src/metrics/tests.rs @@ -3,10 +3,22 @@ use http_body_util::BodyExt; use std::net::IpAddr; use std::time::SystemTime; +use crate::stats::telemetry::TelemetryPolicy; use crate::tls_front::types::{ CachedTlsData, ParsedServerHello, TlsBehaviorProfile, TlsCertPayload, TlsProfileSource, }; +const CAS_CONTENTION_SERIES: [(&str, &str, &str, u64); 8] = [ + ("user", "up", "reserve", 1), + ("user", "down", "reserve", 2), + ("user", "up", "refund", 3), + ("user", "down", "refund", 4), + ("cidr", "up", "reserve", 5), + ("cidr", "down", "reserve", 6), + ("cidr", "up", "refund", 7), + ("cidr", "down", "refund", 8), +]; + fn test_web_publication() -> crate::web::control::WebRuntimePublication { let control = crate::web::control::WebRuntimeControl::new(); control.subscribe().borrow().clone() @@ -18,6 +30,9 @@ async fn test_render_metrics_format() { let shared_state = ProxySharedState::new(); let tracker = UserIpTracker::new(); let mut config = ProxyConfig::default(); + shared_state + .traffic_limiter + .set_cas_contention_metrics_for_test([1, 2, 3, 4, 5, 6, 7, 8]); config .access .user_max_unique_ips @@ -28,6 +43,10 @@ async fn test_render_metrics_format() { stats.increment_connects_bad_with_class("tls_handshake_bad_client"); stats.increment_handshake_timeouts(); stats.increment_handshake_failure_class("timeout"); + stats.increment_conntrack_rule_reconcile_success_total(); + stats.increment_conntrack_rule_reconcile_error_total(); + stats.increment_conntrack_rule_rollback_success_total(); + stats.increment_conntrack_rule_rollback_error_total(); shared_state .handshake .auth_expensive_checks_total @@ -69,6 +88,10 @@ async fn test_render_metrics_format() { stats.increment_me_endpoint_quarantine_draining_suppressed_total(); stats.increment_user_connects("alice"); stats.increment_user_curr_connects("alice"); + let _connection_permit = stats + .connection_authority() + .try_acquire("alice", None) + .unwrap(); stats.add_user_octets_from("alice", 1024); stats.add_user_octets_to("alice", 2048); stats.increment_user_msgs_from("alice"); @@ -87,6 +110,7 @@ async fn test_render_metrics_format() { None, &TlsFullCertBudget::new(), &test_web_publication(), + None, ) .await; @@ -103,6 +127,10 @@ async fn test_render_metrics_format() { ); assert!(output.contains("telemt_handshake_timeouts_total 1")); assert!(output.contains("telemt_handshake_failures_by_class_total{class=\"timeout\"} 1")); + assert!(output.contains("telemt_conntrack_rule_reconcile_total{result=\"success\"} 1")); + assert!(output.contains("telemt_conntrack_rule_reconcile_total{result=\"error\"} 1")); + assert!(output.contains("telemt_conntrack_rule_rollback_total{result=\"success\"} 1")); + assert!(output.contains("telemt_conntrack_rule_rollback_total{result=\"error\"} 1")); assert!(output.contains("telemt_auth_expensive_checks_total 9")); assert!(output.contains("telemt_auth_budget_exhausted_total 2")); assert!(output.contains("telemt_upstream_connect_attempt_total 2")); @@ -151,6 +179,17 @@ async fn test_render_metrics_format() { assert!(output.contains("telemt_ip_tracker_users{scope=\"active\"} 1")); assert!(output.contains("telemt_ip_tracker_entries{scope=\"active\"} 1")); assert!(output.contains("telemt_ip_tracker_cleanup_queue_len 0")); + for (scope, direction, operation, value) in CAS_CONTENTION_SERIES { + assert!(output.contains(&format!( + "telemt_rate_limiter_cas_retry_exhausted_total{{scope=\"{scope}\",direction=\"{direction}\",operation=\"{operation}\"}} {value}" + ))); + } + assert_eq!( + output + .matches("telemt_rate_limiter_cas_retry_exhausted_total{") + .count(), + 8 + ); } #[tokio::test] @@ -223,6 +262,7 @@ async fn test_render_tls_front_profile_health() { Some(&cache), &TlsFullCertBudget::new(), &test_web_publication(), + None, ) .await; @@ -293,6 +333,7 @@ async fn process_tls_budget_metrics_survive_a_generation_without_tls_cache() { None, budget.as_ref(), &test_web_publication(), + None, ) .await; @@ -303,6 +344,13 @@ async fn process_tls_budget_metrics_survive_a_generation_without_tls_cache() { async fn test_render_empty_stats() { let stats = Stats::new(); let shared_state = ProxySharedState::new(); + stats.apply_telemetry_policy(TelemetryPolicy { + core_enabled: false, + ..TelemetryPolicy::default() + }); + shared_state + .traffic_limiter + .set_cas_contention_metrics_for_test([1, 2, 3, 4, 5, 6, 7, 8]); let tracker = UserIpTracker::new(); let config = ProxyConfig::default(); let output = render_metrics( @@ -313,6 +361,7 @@ async fn test_render_empty_stats() { None, &TlsFullCertBudget::new(), &test_web_publication(), + None, ) .await; assert!(output.contains("telemt_connections_total 0")); @@ -322,6 +371,11 @@ async fn test_render_empty_stats() { assert!(output.contains("telemt_auth_budget_exhausted_total 0")); assert!(output.contains("telemt_user_unique_ips_current{user=")); assert!(output.contains("telemt_user_unique_ips_recent_window{user=")); + for (scope, direction, operation, _) in CAS_CONTENTION_SERIES { + assert!(output.contains(&format!( + "telemt_rate_limiter_cas_retry_exhausted_total{{scope=\"{scope}\",direction=\"{direction}\",operation=\"{operation}\"}} 0" + ))); + } } #[tokio::test] @@ -346,6 +400,7 @@ async fn test_render_uses_global_each_unique_ip_limit() { None, &TlsFullCertBudget::new(), &test_web_publication(), + None, ) .await; @@ -367,6 +422,7 @@ async fn test_render_has_type_annotations() { None, &TlsFullCertBudget::new(), &test_web_publication(), + None, ) .await; assert!(output.contains("# TYPE telemt_uptime_seconds gauge")); diff --git a/src/proxy/authenticated.rs b/src/proxy/authenticated.rs index 60315fb..a62ce25 100644 --- a/src/proxy/authenticated.rs +++ b/src/proxy/authenticated.rs @@ -14,7 +14,9 @@ use crate::proxy::handshake::HandshakeSuccess; use crate::proxy::middle_relay::{handle_via_middle_proxy, handle_via_middle_proxy_with_conntrack}; use crate::proxy::route_mode::{RelayRouteMode, RouteRuntimeController}; use crate::proxy::shared_state::{ConntrackClosePolicy, ProxySharedState}; -use crate::stats::Stats; +use crate::proxy::user_admission::UserIncarnation; +use crate::proxy::user_connection_authority::UserConnectionPermit; +use crate::stats::{Stats, UserConnectionObservation, UserQuotaHandle}; use crate::stream::{BufferPool, CryptoReader, CryptoWriter}; use crate::transport::UpstreamManager; use crate::transport::middle_proxy::MePool; @@ -59,13 +61,21 @@ where W: AsyncWrite + Unpin + Send + 'static, { let user = success.user.clone(); - if !deps.shared.is_user_enabled(&user) { + let Some(credential_id) = deps.config.runtime_user_credential_id(&user) else { + warn!(user = %user, "Authenticated user is absent from the runtime credential snapshot"); + return Err(ProxyError::UserDisabled { user }); + }; + let Some(user_incarnation) = deps + .shared + .authenticated_user_incarnation(&user, credential_id) + else { warn!(user = %user, "Disabled user rejected"); return Err(ProxyError::UserDisabled { user }); - } + }; - let user_reservation = acquire_user_connection_reservation( + let user_reservation = acquire_user_connection_reservation_for_incarnation( &user, + user_incarnation, &deps.config, Arc::clone(&deps.stats), peer_addr, @@ -76,10 +86,24 @@ where warn!(user = %user, error = %error, "User admission check failed"); error })?; + let quota_handle = user_reservation.quota_handle(); let route_snapshot = deps.route_runtime.snapshot(); let session_id = deps.rng.u64(); - let user_session = deps.shared.register_user_session(&user, session_id); + let Some(user_session) = deps + .shared + .register_authenticated_user_session(&user, credential_id) + else { + user_reservation.release_deferred(); + warn!(user = %user, "Disabled user rejected during final admission"); + return Err(ProxyError::UserDisabled { user }); + }; + if user_session.incarnation() != user_incarnation { + drop(user_session); + user_reservation.release_deferred(); + warn!(user = %user, "User incarnation changed during admission"); + return Err(ProxyError::UserDisabled { user }); + } let session_cancel = user_session.token(); let selected_me_pool = if deps.config.general.use_middle_proxy && matches!(route_snapshot.mode, RelayRouteMode::Middle) @@ -115,6 +139,7 @@ where session_id, session_cancel.clone(), Arc::clone(&deps.shared), + quota_handle.clone(), ) .await } else { @@ -134,6 +159,7 @@ where session_cancel.clone(), Arc::clone(&deps.shared), ConntrackClosePolicy::Suppress, + quota_handle.clone(), ) .await } @@ -149,6 +175,7 @@ where local_addr, session_cancel.clone(), conntrack_close_policy, + quota_handle.clone(), ) .await } @@ -163,6 +190,7 @@ where local_addr, session_cancel, conntrack_close_policy, + quota_handle, ) .await }; @@ -180,6 +208,7 @@ async fn run_direct( local_addr: SocketAddr, session_cancel: tokio_util::sync::CancellationToken, conntrack_close_policy: ConntrackClosePolicy, + quota_handle: UserQuotaHandle, ) -> Result<()> where R: AsyncRead + Unpin + Send + 'static, @@ -201,6 +230,7 @@ where session_cancel, Arc::clone(&deps.shared), conntrack_close_policy, + quota_handle, ) .await } @@ -209,11 +239,60 @@ where /// Owns one authenticated user's connection and source-IP admission slots. pub(crate) struct UserConnectionReservation { stats: Arc, - ip_tracker: Arc, + quota_handle: UserQuotaHandle, + _connection_permit: UserConnectionPermit, + _stats_observation: Option, + ip_permit: Option, + released: bool, +} + +struct UserIpPermit { + tracker: Arc, + owner: Option, +} + +struct UserIpOwner { user: String, + incarnation: UserIncarnation, ip: IpAddr, - tracks_ip: bool, - active: bool, +} + +impl UserIpPermit { + fn new( + tracker: Arc, + user: String, + incarnation: UserIncarnation, + ip: IpAddr, + ) -> Self { + Self { + tracker, + owner: Some(UserIpOwner { + user, + incarnation, + ip, + }), + } + } + + async fn release(mut self) { + let Some(owner) = self.owner.as_ref() else { + return; + }; + self.tracker + .remove_ip_for_incarnation(&owner.user, owner.incarnation, owner.ip) + .await; + self.owner = None; + } +} + +impl Drop for UserIpPermit { + fn drop(&mut self) { + let Some(owner) = self.owner.take() else { + return; + }; + self.tracker + .enqueue_cleanup_for_incarnation(owner.user, owner.incarnation, owner.ip); + } } impl UserConnectionReservation { @@ -225,40 +304,73 @@ impl UserConnectionReservation { ip: IpAddr, tracks_ip: bool, ) -> Self { - Self { + let quota_handle = stats.current_user_quota_handle(&user); + let connection_permit = stats + .connection_authority() + .try_acquire(&user, None) + .expect("unlimited test connection permit must be available"); + let stats_observation = stats.observe_user_current_connection(&user); + Self::new_for_incarnation( stats, ip_tracker, user, ip, + 0, + quota_handle, + connection_permit, + stats_observation, tracks_ip, - active: true, + ) + } + + /// Creates a reservation fenced to one authenticated user incarnation. + pub(crate) fn new_for_incarnation( + stats: Arc, + ip_tracker: Arc, + user: String, + ip: IpAddr, + incarnation: UserIncarnation, + quota_handle: UserQuotaHandle, + connection_permit: UserConnectionPermit, + stats_observation: Option, + tracks_ip: bool, + ) -> Self { + let ip_permit = tracks_ip.then(|| UserIpPermit::new(ip_tracker, user, incarnation, ip)); + Self { + stats, + quota_handle, + _connection_permit: connection_permit, + _stats_observation: stats_observation, + ip_permit, + released: false, } } + /// Returns quota ownership pinned to the authenticated user incarnation. + pub(crate) fn quota_handle(&self) -> UserQuotaHandle { + self.quota_handle.clone() + } + /// Releases both admission counters through the asynchronous cleanup path. pub(crate) async fn release(mut self) { - if !self.active { - return; + if let Some(ip_permit) = self.ip_permit.take() { + ip_permit.release().await; } - self.active = false; - if self.tracks_ip { - self.ip_tracker.remove_ip(&self.user, self.ip).await; - } - self.stats.decrement_user_curr_connects(&self.user); + self.released = true; + } + + /// Defers IP cleanup when admission fails after the asynchronous reservation step. + pub(crate) fn release_deferred(mut self) { + self.released = true; } } impl Drop for UserConnectionReservation { fn drop(&mut self) { - if !self.active { + if self.released { return; } - self.active = false; self.stats.increment_session_drop_fallback_total(); - self.stats.decrement_user_curr_connects(&self.user); - if self.tracks_ip { - self.ip_tracker.enqueue_cleanup(self.user.clone(), self.ip); - } } } @@ -269,6 +381,20 @@ pub(crate) async fn acquire_user_connection_reservation( stats: Arc, peer_addr: SocketAddr, ip_tracker: Arc, +) -> Result { + acquire_user_connection_reservation_for_incarnation( + user, 0, config, stats, peer_addr, ip_tracker, + ) + .await +} + +async fn acquire_user_connection_reservation_for_incarnation( + user: &str, + incarnation: UserIncarnation, + config: &ProxyConfig, + stats: Arc, + peer_addr: SocketAddr, + ip_tracker: Arc, ) -> Result { if let Some(expiration) = config.access.user_expirations.get(user) && chrono::Utc::now() > *expiration @@ -277,8 +403,13 @@ pub(crate) async fn acquire_user_connection_reservation( user: user.to_string(), }); } + let Some(quota_handle) = stats.quota_handle_for_incarnation(user, incarnation) else { + return Err(ProxyError::UserDisabled { + user: user.to_string(), + }); + }; if let Some(quota) = config.access.user_data_quota.get(user) - && stats.get_user_quota_used(user) >= *quota + && quota_handle.used() >= *quota { return Err(ProxyError::DataQuotaExceeded { user: user.to_string(), @@ -294,14 +425,17 @@ pub(crate) async fn acquire_user_connection_reservation( .or((config.access.user_max_tcp_conns_global_each > 0) .then_some(config.access.user_max_tcp_conns_global_each)) .map(|value| value as u64); - if !stats.try_acquire_user_curr_connects(user, limit) { + let Some(connection_permit) = stats.connection_authority().try_acquire(user, limit) else { return Err(ProxyError::ConnectionLimitExceeded { user: user.to_string(), }); - } + }; + let stats_observation = stats.observe_user_current_connection(user); - if let Err(reason) = ip_tracker.check_and_add(user, peer_addr.ip()).await { - stats.decrement_user_curr_connects(user); + if let Err(reason) = ip_tracker + .check_and_add_for_incarnation(user, incarnation, peer_addr.ip()) + .await + { warn!( user = %user, ip = %peer_addr.ip(), @@ -313,11 +447,15 @@ pub(crate) async fn acquire_user_connection_reservation( }); } - Ok(UserConnectionReservation::new( + Ok(UserConnectionReservation::new_for_incarnation( stats, ip_tracker, user.to_string(), peer_addr.ip(), + incarnation, + quota_handle, + connection_permit, + stats_observation, true, )) } diff --git a/src/proxy/client/authenticated.rs b/src/proxy/client/authenticated.rs index d131863..5374392 100644 --- a/src/proxy/client/authenticated.rs +++ b/src/proxy/client/authenticated.rs @@ -25,6 +25,28 @@ impl RunningClientHandler { R: AsyncRead + Unpin + Send + 'static, W: AsyncWrite + Unpin + Send + 'static, { + // Manually constructed handshake fixtures bypass credential validation, so + // materialize the process-authority state that a real handshake generation owns. + let config = if config.runtime_user_credential_id(&success.user).is_none() { + let mut test_config = (*config).clone(); + test_config.access.users.insert( + success.user.clone(), + "00000000000000000000000000000000".to_string(), + ); + test_config.rebuild_runtime_user_auth()?; + Arc::new(test_config) + } else { + config + }; + let shared = ProxySharedState::new_with_direct_buffer_budget_and_user_admission( + crate::proxy::direct_buffer_budget::DirectBufferBudget::new( + crate::proxy::direct_buffer_budget::fallback_direct_buffer_hard_limit(), + ), + crate::proxy::user_admission::UserAdmissionAuthority::new_with_quota_store( + stats.quota_store(), + ), + ); + shared.apply_user_config(&config.access.users, &config.access.user_enabled); Self::handle_authenticated_static_with_shared( client_reader, client_writer, @@ -40,7 +62,7 @@ impl RunningClientHandler { local_addr, peer_addr, ip_tracker, - ProxySharedState::new(), + shared, ) .await } @@ -133,18 +155,17 @@ impl RunningClientHandler { .or((config.access.user_max_tcp_conns_global_each > 0) .then_some(config.access.user_max_tcp_conns_global_each)) .map(|v| v as u64); - if !stats.try_acquire_user_curr_connects(user, limit) { + let Some(_connection_permit) = stats.connection_authority().try_acquire(user, limit) else { return Err(ProxyError::ConnectionLimitExceeded { user: user.to_string(), }); - } + }; match ip_tracker.check_and_add(user, peer_addr.ip()).await { Ok(()) => { ip_tracker.remove_ip(user, peer_addr.ip()).await; } Err(reason) => { - stats.decrement_user_curr_connects(user); warn!( user = %user, ip = %peer_addr.ip(), @@ -156,8 +177,6 @@ impl RunningClientHandler { }); } } - - stats.decrement_user_curr_connects(user); Ok(()) } } diff --git a/src/proxy/direct_buffer_budget.rs b/src/proxy/direct_buffer_budget.rs index 0242294..8204351 100644 --- a/src/proxy/direct_buffer_budget.rs +++ b/src/proxy/direct_buffer_budget.rs @@ -2,12 +2,16 @@ use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; use std::time::Duration; +use parking_lot::{Mutex as ParkingMutex, MutexGuard as ParkingMutexGuard}; use tokio::sync::watch; -use crate::stats::Stats; -use crate::stream::BufferPool; - -use super::shared_state::ProxySharedState; +// Process controller and system-memory sampling remain outside data-plane accounting. +mod controller; +#[cfg(test)] +use controller::connection_fill_pct; +pub(crate) use controller::{ + resolve_direct_buffer_hard_limit, run_direct_buffer_budget_controller, +}; /// Accounting granularity for process-wide Direct copy-buffer reservations. pub(crate) const DIRECT_BUFFER_UNIT_BYTES: usize = 4 * 1024; @@ -70,6 +74,8 @@ pub(crate) struct DirectBufferBudget { hard_limit_bytes: u64, target_bytes: AtomicU64, reserved_bytes: AtomicU64, + active_controller_generation: AtomicU64, + controller_update: ParkingMutex<()>, pressure_generation: AtomicU64, pressure_tx: watch::Sender, memory_total_bytes: AtomicU64, @@ -94,6 +100,8 @@ impl DirectBufferBudget { hard_limit_bytes, target_bytes: AtomicU64::new(hard_limit_bytes), reserved_bytes: AtomicU64::new(0), + active_controller_generation: AtomicU64::new(0), + controller_update: ParkingMutex::new(()), pressure_generation: AtomicU64::new(0), pressure_tx, memory_total_bytes: AtomicU64::new(0), @@ -120,6 +128,19 @@ impl DirectBufferBudget { self.pressure_tx.subscribe() } + /// Transfers adaptive-target writes to the active runtime generation. + pub(crate) fn activate_controller(&self, generation: u64) { + let _controller_update = self.controller_update.lock(); + self.active_controller_generation + .fetch_max(generation, Ordering::AcqRel); + } + + fn begin_controller_update(&self, generation: u64) -> Option> { + let controller_update = self.controller_update.lock(); + (self.active_controller_generation.load(Ordering::Acquire) == generation) + .then_some(controller_update) + } + /// Reserves bytes against either the adaptive target or the absolute ceiling. pub(crate) fn try_reserve( self: &Arc, @@ -322,222 +343,6 @@ impl Drop for DirectBufferLease { } } -/// Resolves the startup hard ceiling from config, cgroup, and host memory. -pub(crate) async fn resolve_direct_buffer_hard_limit(configured: usize) -> usize { - if configured != 0 { - return align_down(configured); - } - let sample = read_system_memory_sample().await; - if sample.total_bytes == 0 { - return AUTO_HARD_FALLBACK_BYTES; - } - let derived = (sample.total_bytes / 4) - .clamp(AUTO_HARD_MIN_BYTES as u64, AUTO_HARD_MAX_BYTES as u64) - .min(sample.total_bytes); - align_down(derived as usize).max(DIRECT_BUFFER_UNIT_BYTES) -} - -/// Runs the control-plane loop for Direct budget and shared pool pressure. -pub(crate) async fn run_direct_buffer_budget_controller( - budget: Arc, - buffer_pool: Arc, - stats: Arc, - shared: Arc, - max_connections: u32, -) { - let mut interval = tokio::time::interval(CONTROL_INTERVAL); - interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); - let mut healthy_streak = 0u8; - let mut previous_denied = 0u64; - let mut previous_fallback = 0u64; - let mut previous_rejected = 0u64; - let pool_trim_low = buffer_pool - .max_buffers() - .min(BUFFER_POOL_TRIM_LOW_WATERMARK); - let pool_trim_high = buffer_pool - .max_buffers() - .min(BUFFER_POOL_TRIM_HIGH_WATERMARK); - let mut pool_trim_armed = true; - - loop { - interval.tick().await; - let sample = read_system_memory_sample().await; - budget.update_system_sample(sample); - - let snapshot = budget.snapshot(); - let denied_delta = snapshot - .promotion_denied_total - .saturating_sub(previous_denied); - previous_denied = snapshot.promotion_denied_total; - let fallback_delta = snapshot - .minimum_fallback_total - .saturating_sub(previous_fallback); - previous_fallback = snapshot.minimum_fallback_total; - let rejected_delta = snapshot - .admission_rejected_total - .saturating_sub(previous_rejected); - previous_rejected = snapshot.admission_rejected_total; - - let connection_pct = connection_fill_pct(stats.as_ref(), max_connections); - let memory_available_pct = percentage(sample.available_bytes, sample.total_bytes); - let target_utilization_pct = percentage(snapshot.reserved_bytes, snapshot.target_bytes); - let pressure = shared.conntrack_pressure_active() - || connection_pct.is_some_and(|value| value >= 85) - || memory_available_pct.is_some_and(|value| value <= 15) - || target_utilization_pct.is_some_and(|value| value >= 90) - || denied_delta > 0 - || fallback_delta > 0 - || rejected_delta > 0; - - if !pressure { - pool_trim_armed = true; - } else if pool_trim_armed && buffer_pool.pooled() > pool_trim_high { - buffer_pool.trim_to(pool_trim_low); - pool_trim_armed = false; - } - - let pool_snapshot = buffer_pool.stats(); - stats.set_buffer_pool_gauges( - pool_snapshot.pooled, - pool_snapshot.allocated, - pool_snapshot.allocated.saturating_sub(pool_snapshot.pooled), - ); - stats.set_buffer_pool_replaced_nonstandard_total(pool_snapshot.replaced_nonstandard); - - let headroom_target = if sample.total_bytes == 0 { - snapshot.hard_limit_bytes - } else { - snapshot - .reserved_bytes - .saturating_add(sample.available_bytes / 4) - .min(snapshot.hard_limit_bytes) - }; - - if pressure { - healthy_streak = 0; - let reduced = snapshot.target_bytes.saturating_mul(3) / 4; - budget.set_target_bytes(reduced.min(headroom_target)); - continue; - } - - let healthy = memory_available_pct.is_none_or(|value| value >= 30) - && connection_pct.is_none_or(|value| value <= 70); - if !healthy { - healthy_streak = 0; - if headroom_target < snapshot.target_bytes { - budget.set_target_bytes(headroom_target); - } - continue; - } - - healthy_streak = healthy_streak.saturating_add(1); - if healthy_streak >= HEALTHY_RECOVERY_SAMPLES { - healthy_streak = 0; - let increment = (snapshot.target_bytes / 16).max(4 * 1024 * 1024); - budget.set_target_bytes( - snapshot - .target_bytes - .saturating_add(increment) - .min(headroom_target), - ); - } - } -} - -fn connection_fill_pct(stats: &Stats, max_connections: u32) -> Option { - if max_connections == 0 { - return None; - } - Some( - ((stats.get_current_connections_total().saturating_mul(100)) / u64::from(max_connections)) - .min(100) as u8, - ) -} - -fn percentage(value: u64, total: u64) -> Option { - if total == 0 { - return None; - } - Some(((value.saturating_mul(100)) / total).min(100) as u8) -} - -async fn read_system_memory_sample() -> SystemMemorySample { - #[cfg(target_os = "linux")] - { - let meminfo = tokio::fs::read_to_string("/proc/meminfo") - .await - .unwrap_or_default(); - let status = tokio::fs::read_to_string("/proc/self/status") - .await - .unwrap_or_default(); - let host_total = parse_kib_field(&meminfo, "MemTotal:"); - let host_available = parse_kib_field(&meminfo, "MemAvailable:"); - let process_rss = parse_kib_field(&status, "VmRSS:"); - - let cgroup_v2_max = read_cgroup_limit("/sys/fs/cgroup/memory.max").await; - let cgroup_v2_current = read_u64_file("/sys/fs/cgroup/memory.current").await; - let cgroup_v1_max = read_cgroup_limit("/sys/fs/cgroup/memory/memory.limit_in_bytes").await; - let cgroup_v1_current = read_u64_file("/sys/fs/cgroup/memory/memory.usage_in_bytes").await; - let cgroup_max = cgroup_v2_max.or(cgroup_v1_max); - let cgroup_current = cgroup_v2_current.or(cgroup_v1_current); - - let total = match (host_total, cgroup_max) { - (0, Some(limit)) => limit, - (host, Some(limit)) => host.min(limit), - (host, None) => host, - }; - let cgroup_available = cgroup_max - .zip(cgroup_current) - .map(|(limit, current)| limit.saturating_sub(current)); - let available = match (host_available, cgroup_available) { - (0, Some(value)) => value, - (host, Some(value)) => host.min(value), - (host, None) => host, - }; - return SystemMemorySample { - total_bytes: total, - available_bytes: available, - process_rss_bytes: process_rss, - }; - } - #[cfg(not(target_os = "linux"))] - { - SystemMemorySample::default() - } -} - -#[cfg(target_os = "linux")] -async fn read_cgroup_limit(path: &str) -> Option { - let raw = tokio::fs::read_to_string(path).await.ok()?; - let raw = raw.trim(); - if raw == "max" { - return None; - } - let value = raw.parse::().ok()?; - (value < (1u64 << 60)).then_some(value) -} - -#[cfg(target_os = "linux")] -async fn read_u64_file(path: &str) -> Option { - tokio::fs::read_to_string(path) - .await - .ok()? - .trim() - .parse() - .ok() -} - -#[cfg(target_os = "linux")] -fn parse_kib_field(raw: &str, key: &str) -> u64 { - raw.lines() - .find_map(|line| { - let value = line.strip_prefix(key)?.split_whitespace().next()?; - value.parse::().ok() - }) - .unwrap_or(0) - .saturating_mul(1024) -} - fn align_up(bytes: usize) -> usize { bytes .div_ceil(DIRECT_BUFFER_UNIT_BYTES) diff --git a/src/proxy/direct_buffer_budget/controller.rs b/src/proxy/direct_buffer_budget/controller.rs new file mode 100644 index 0000000..589f00e --- /dev/null +++ b/src/proxy/direct_buffer_budget/controller.rs @@ -0,0 +1,235 @@ +use std::sync::Arc; + +use tokio::sync::Semaphore; + +use super::*; +use crate::proxy::shared_state::ProxySharedState; +use crate::stats::Stats; +use crate::stream::BufferPool; + +/// Resolves the startup hard ceiling from config, cgroup, and host memory. +pub(crate) async fn resolve_direct_buffer_hard_limit(configured: usize) -> usize { + if configured != 0 { + return align_down(configured); + } + let sample = read_system_memory_sample().await; + if sample.total_bytes == 0 { + return AUTO_HARD_FALLBACK_BYTES; + } + let derived = (sample.total_bytes / 4) + .clamp(AUTO_HARD_MIN_BYTES as u64, AUTO_HARD_MAX_BYTES as u64) + .min(sample.total_bytes); + align_down(derived as usize).max(DIRECT_BUFFER_UNIT_BYTES) +} + +/// Runs the control-plane loop for Direct budget and shared pool pressure. +pub(crate) async fn run_direct_buffer_budget_controller( + source_generation: u64, + budget: Arc, + buffer_pool: Arc, + stats: Arc, + shared: Arc, + connection_slots: Arc, + max_connections: u32, +) { + let mut interval = tokio::time::interval(CONTROL_INTERVAL); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + let mut healthy_streak = 0u8; + let mut previous_denied = 0u64; + let mut previous_fallback = 0u64; + let mut previous_rejected = 0u64; + let pool_trim_low = buffer_pool + .max_buffers() + .min(BUFFER_POOL_TRIM_LOW_WATERMARK); + let pool_trim_high = buffer_pool + .max_buffers() + .min(BUFFER_POOL_TRIM_HIGH_WATERMARK); + let mut pool_trim_armed = true; + + loop { + interval.tick().await; + if budget.active_controller_generation.load(Ordering::Acquire) != source_generation { + continue; + } + let sample = read_system_memory_sample().await; + let Some(_controller_update) = budget.begin_controller_update(source_generation) else { + continue; + }; + budget.update_system_sample(sample); + + let snapshot = budget.snapshot(); + let denied_delta = snapshot + .promotion_denied_total + .saturating_sub(previous_denied); + previous_denied = snapshot.promotion_denied_total; + let fallback_delta = snapshot + .minimum_fallback_total + .saturating_sub(previous_fallback); + previous_fallback = snapshot.minimum_fallback_total; + let rejected_delta = snapshot + .admission_rejected_total + .saturating_sub(previous_rejected); + previous_rejected = snapshot.admission_rejected_total; + + let connection_pct = connection_fill_pct(connection_slots.as_ref(), max_connections); + let memory_available_pct = percentage(sample.available_bytes, sample.total_bytes); + let target_utilization_pct = percentage(snapshot.reserved_bytes, snapshot.target_bytes); + let pressure = shared.conntrack_pressure_active() + || connection_pct.is_some_and(|value| value >= 85) + || memory_available_pct.is_some_and(|value| value <= 15) + || target_utilization_pct.is_some_and(|value| value >= 90) + || denied_delta > 0 + || fallback_delta > 0 + || rejected_delta > 0; + + if !pressure { + pool_trim_armed = true; + } else if pool_trim_armed && buffer_pool.pooled() > pool_trim_high { + buffer_pool.trim_to(pool_trim_low); + pool_trim_armed = false; + } + + let pool_snapshot = buffer_pool.stats(); + stats.set_buffer_pool_gauges( + pool_snapshot.pooled, + pool_snapshot.allocated, + pool_snapshot.allocated.saturating_sub(pool_snapshot.pooled), + ); + stats.set_buffer_pool_replaced_nonstandard_total(pool_snapshot.replaced_nonstandard); + + let headroom_target = if sample.total_bytes == 0 { + snapshot.hard_limit_bytes + } else { + snapshot + .reserved_bytes + .saturating_add(sample.available_bytes / 4) + .min(snapshot.hard_limit_bytes) + }; + + if pressure { + healthy_streak = 0; + let reduced = snapshot.target_bytes.saturating_mul(3) / 4; + budget.set_target_bytes(reduced.min(headroom_target)); + continue; + } + + let healthy = memory_available_pct.is_none_or(|value| value >= 30) + && connection_pct.is_none_or(|value| value <= 70); + if !healthy { + healthy_streak = 0; + if headroom_target < snapshot.target_bytes { + budget.set_target_bytes(headroom_target); + } + continue; + } + + healthy_streak = healthy_streak.saturating_add(1); + if healthy_streak >= HEALTHY_RECOVERY_SAMPLES { + healthy_streak = 0; + let increment = (snapshot.target_bytes / 16).max(4 * 1024 * 1024); + budget.set_target_bytes( + snapshot + .target_bytes + .saturating_add(increment) + .min(headroom_target), + ); + } + } +} + +pub(super) fn connection_fill_pct( + connection_slots: &Semaphore, + max_connections: u32, +) -> Option { + if max_connections == 0 { + return None; + } + let max_connections = max_connections as usize; + let active = + max_connections.saturating_sub(connection_slots.available_permits().min(max_connections)); + Some((active.saturating_mul(100) / max_connections).min(100) as u8) +} + +fn percentage(value: u64, total: u64) -> Option { + if total == 0 { + return None; + } + Some(((value.saturating_mul(100)) / total).min(100) as u8) +} + +async fn read_system_memory_sample() -> SystemMemorySample { + #[cfg(target_os = "linux")] + { + let meminfo = tokio::fs::read_to_string("/proc/meminfo") + .await + .unwrap_or_default(); + let status = tokio::fs::read_to_string("/proc/self/status") + .await + .unwrap_or_default(); + let host_total = parse_kib_field(&meminfo, "MemTotal:"); + let host_available = parse_kib_field(&meminfo, "MemAvailable:"); + let process_rss = parse_kib_field(&status, "VmRSS:"); + + let cgroup_v2_max = read_cgroup_limit("/sys/fs/cgroup/memory.max").await; + let cgroup_v2_current = read_u64_file("/sys/fs/cgroup/memory.current").await; + let cgroup_v1_max = read_cgroup_limit("/sys/fs/cgroup/memory/memory.limit_in_bytes").await; + let cgroup_v1_current = read_u64_file("/sys/fs/cgroup/memory/memory.usage_in_bytes").await; + let cgroup_max = cgroup_v2_max.or(cgroup_v1_max); + let cgroup_current = cgroup_v2_current.or(cgroup_v1_current); + + let total = match (host_total, cgroup_max) { + (0, Some(limit)) => limit, + (host, Some(limit)) => host.min(limit), + (host, None) => host, + }; + let cgroup_available = cgroup_max + .zip(cgroup_current) + .map(|(limit, current)| limit.saturating_sub(current)); + let available = match (host_available, cgroup_available) { + (0, Some(value)) => value, + (host, Some(value)) => host.min(value), + (host, None) => host, + }; + return SystemMemorySample { + total_bytes: total, + available_bytes: available, + process_rss_bytes: process_rss, + }; + } + #[cfg(not(target_os = "linux"))] + { + SystemMemorySample::default() + } +} + +#[cfg(target_os = "linux")] +async fn read_cgroup_limit(path: &str) -> Option { + let raw = tokio::fs::read_to_string(path).await.ok()?; + let raw = raw.trim(); + if raw == "max" { + return None; + } + let value = raw.parse::().ok()?; + (value < (1u64 << 60)).then_some(value) +} + +#[cfg(target_os = "linux")] +async fn read_u64_file(path: &str) -> Option { + tokio::fs::read_to_string(path) + .await + .ok()? + .trim() + .parse() + .ok() +} + +#[cfg(target_os = "linux")] +fn parse_kib_field(raw: &str, key: &str) -> u64 { + raw.lines() + .find_map(|line| { + let value = line.strip_prefix(key)?.split_whitespace().next()?; + value.parse::().ok() + }) + .unwrap_or(0) + .saturating_mul(1024) +} diff --git a/src/proxy/direct_relay.rs b/src/proxy/direct_relay.rs index ef65298..45a1b8b 100644 --- a/src/proxy/direct_relay.rs +++ b/src/proxy/direct_relay.rs @@ -1,5 +1,6 @@ use std::collections::HashSet; use std::ffi::OsString; +#[cfg(all(test, unix))] use std::fs::OpenOptions; use std::io::Write; use std::net::SocketAddr; @@ -26,6 +27,7 @@ use crate::proxy::shared_state::{ ProxySharedState, }; use crate::stats::Stats; +use crate::stats::UserQuotaHandle; use crate::stream::{BufferPool, CryptoReader, CryptoWriter}; use crate::transport::UpstreamManager; #[cfg(unix)] @@ -33,7 +35,7 @@ use nix::fcntl::{Flock, FlockArg, OFlag, openat}; #[cfg(unix)] use nix::sys::stat::Mode; -#[cfg(unix)] +#[cfg(all(test, unix))] use std::os::unix::fs::OpenOptionsExt; // Direct relay lifecycle and conntrack publication. @@ -178,10 +180,7 @@ fn open_unknown_dc_log_append_anchored( ) -> std::io::Result { #[cfg(unix)] { - let parent = OpenOptions::new() - .read(true) - .custom_flags(libc::O_DIRECTORY | libc::O_NOFOLLOW | libc::O_CLOEXEC) - .open(&path.allowed_parent)?; + let parent = crate::util::secure_fs::open_dir_nofollow(&path.allowed_parent)?; let oflags = OFlag::O_CREAT | OFlag::O_APPEND diff --git a/src/proxy/direct_relay/relay.rs b/src/proxy/direct_relay/relay.rs index b70420a..298ff5d 100644 --- a/src/proxy/direct_relay/relay.rs +++ b/src/proxy/direct_relay/relay.rs @@ -59,6 +59,7 @@ where R: AsyncRead + Unpin + Send + 'static, W: AsyncWrite + Unpin + Send + 'static, { + let quota_handle = stats.current_user_quota_handle(&success.user); handle_via_direct_with_shared_and_conntrack( client_reader, client_writer, @@ -75,6 +76,7 @@ where session_cancel, shared, ConntrackClosePolicy::Publish, + quota_handle, ) .await } @@ -96,6 +98,7 @@ pub(crate) async fn handle_via_direct_with_shared_and_conntrack( session_cancel: CancellationToken, shared: Arc, conntrack_close_policy: ConntrackClosePolicy, + quota_handle: UserQuotaHandle, ) -> Result<()> where R: AsyncRead + Unpin + Send + 'static, @@ -171,6 +174,7 @@ where config.server.max_connections, user, Arc::clone(&stats), + quota_handle, config.access.user_data_quota.get(user).copied(), traffic_lease, relay_activity_timeout, diff --git a/src/proxy/handshake.rs b/src/proxy/handshake.rs index 96cefd3..4431f55 100644 --- a/src/proxy/handshake.rs +++ b/src/proxy/handshake.rs @@ -74,9 +74,10 @@ pub(crate) use self::auth_probe::{ auth_probe_saturation_is_throttled_at_for_testing_in_shared, auth_probe_saturation_is_throttled_for_testing_in_shared, auth_probe_saturation_state_for_testing_in_shared, - auth_probe_saturation_state_lock_for_testing_in_shared, auth_probe_state_for_testing_in_shared, - clear_auth_probe_state_for_testing_in_shared, + auth_probe_saturation_state_lock_for_testing_in_shared, auth_probe_slots_for_testing_in_shared, + auth_probe_state_for_testing_in_shared, clear_auth_probe_state_for_testing_in_shared, clear_unknown_sni_warn_state_for_testing_in_shared, clear_warned_secrets_for_testing_in_shared, + insert_auth_probe_state_for_testing_in_shared, should_emit_unknown_sni_warn_for_testing_in_shared, warned_secrets_for_testing_in_shared, }; @@ -89,13 +90,16 @@ const WARNED_SECRET_MAX_ENTRIES: usize = 1_024; const AUTH_PROBE_TRACK_RETENTION_SECS: u64 = 10 * 60; #[cfg(test)] -const AUTH_PROBE_TRACK_MAX_ENTRIES: usize = 256; +pub(super) const AUTH_PROBE_TRACK_MAX_ENTRIES: usize = 256; #[cfg(not(test))] -const AUTH_PROBE_TRACK_MAX_ENTRIES: usize = 65_536; +pub(super) const AUTH_PROBE_TRACK_MAX_ENTRIES: usize = 65_536; const AUTH_PROBE_PRUNE_SCAN_LIMIT: usize = 1_024; const AUTH_PROBE_BACKOFF_START_FAILS: u32 = 4; const AUTH_PROBE_SATURATION_GRACE_FAILS: u32 = 2; -const STICKY_HINT_MAX_ENTRIES: usize = 65_536; +#[cfg(test)] +pub(super) const STICKY_HINT_MAX_ENTRIES: usize = 256; +#[cfg(not(test))] +pub(super) const STICKY_HINT_MAX_ENTRIES: usize = 65_536; const CANDIDATE_HINT_TRACK_CAP: usize = 64; const OVERLOAD_CANDIDATE_BUDGET_HINTED: usize = 16; const OVERLOAD_CANDIDATE_BUDGET_UNHINTED: usize = 8; diff --git a/src/proxy/handshake/auth_candidates.rs b/src/proxy/handshake/auth_candidates.rs index 48954ee..79db15c 100644 --- a/src/proxy/handshake/auth_candidates.rs +++ b/src/proxy/handshake/auth_candidates.rs @@ -42,7 +42,7 @@ pub(super) fn ip_prefix_hint_key(peer_ip: IpAddr) -> u64 { } } -pub(super) fn sticky_hint_get_by_ip(shared: &ProxySharedState, peer_ip: IpAddr) -> Option { +pub(super) fn sticky_hint_get_by_ip(shared: &ProxySharedState, peer_ip: IpAddr) -> Option { shared .handshake .sticky_user_by_ip @@ -53,7 +53,7 @@ pub(super) fn sticky_hint_get_by_ip(shared: &ProxySharedState, peer_ip: IpAddr) pub(super) fn sticky_hint_get_by_ip_prefix( shared: &ProxySharedState, peer_ip: IpAddr, -) -> Option { +) -> Option { shared .handshake .sticky_user_by_ip_prefix @@ -61,7 +61,7 @@ pub(super) fn sticky_hint_get_by_ip_prefix( .map(|entry| *entry) } -pub(super) fn sticky_hint_get_by_sni(shared: &ProxySharedState, sni: &str) -> Option { +pub(super) fn sticky_hint_get_by_sni(shared: &ProxySharedState, sni: &str) -> Option { let key = sni_hint_hash(sni); shared .handshake @@ -73,34 +73,76 @@ pub(super) fn sticky_hint_get_by_sni(shared: &ProxySharedState, sni: &str) -> Op pub(super) fn sticky_hint_record_success_in( shared: &ProxySharedState, peer_ip: IpAddr, - user_id: u32, + hint_key: u64, sni: Option<&str>, ) { - if shared.handshake.sticky_user_by_ip.len() > STICKY_HINT_MAX_ENTRIES { - shared.handshake.sticky_user_by_ip.clear(); - } - shared.handshake.sticky_user_by_ip.insert(peer_ip, user_id); - - if shared.handshake.sticky_user_by_ip_prefix.len() > STICKY_HINT_MAX_ENTRIES { - shared.handshake.sticky_user_by_ip_prefix.clear(); - } - shared - .handshake - .sticky_user_by_ip_prefix - .insert(ip_prefix_hint_key(peer_ip), user_id); + bounded_sticky_hint_upsert( + &shared.handshake.sticky_user_by_ip, + &shared.handshake.sticky_user_by_ip_slots, + peer_ip, + hint_key, + ); + bounded_sticky_hint_upsert( + &shared.handshake.sticky_user_by_ip_prefix, + &shared.handshake.sticky_user_by_ip_prefix_slots, + ip_prefix_hint_key(peer_ip), + hint_key, + ); if let Some(sni) = sni { - if shared.handshake.sticky_user_by_sni_hash.len() > STICKY_HINT_MAX_ENTRIES { - shared.handshake.sticky_user_by_sni_hash.clear(); - } - shared - .handshake - .sticky_user_by_sni_hash - .insert(sni_hint_hash(sni), user_id); + bounded_sticky_hint_upsert( + &shared.handshake.sticky_user_by_sni_hash, + &shared.handshake.sticky_user_by_sni_hash_slots, + sni_hint_hash(sni), + hint_key, + ); } } -pub(super) fn record_recent_user_success_in(shared: &ProxySharedState, user_id: u32) { +fn bounded_sticky_hint_upsert( + entries: &DashMap, + slots: &crate::slot_budget::SlotBudget, + key: K, + hint_key: u64, +) where + K: Clone + Eq + Hash, +{ + if let Some(mut existing) = entries.get_mut(&key) { + *existing = hint_key; + return; + } + + for _ in 0..2 { + if let Some(slot) = slots.try_acquire() { + match entries.entry(key.clone()) { + Entry::Occupied(mut entry) => { + entry.insert(hint_key); + } + Entry::Vacant(entry) => { + entry.insert(hint_key); + slot.commit(); + } + } + return; + } + + let Some((victim_key, victim_hint_key)) = entries + .iter() + .next() + .map(|entry| (entry.key().clone(), *entry.value())) + else { + return; + }; + if entries + .remove_if(&victim_key, |_, current| *current == victim_hint_key) + .is_some() + { + slots.release(); + } + } +} + +pub(super) fn record_recent_user_success_in(shared: &ProxySharedState, hint_key: u64) { let ring = &shared.handshake.recent_user_ring; if ring.is_empty() { return; @@ -110,7 +152,7 @@ pub(super) fn record_recent_user_success_in(shared: &ProxySharedState, user_id: .recent_user_ring_seq .fetch_add(1, Ordering::Relaxed); let idx = (seq as usize) % ring.len(); - ring[idx].store(user_id.saturating_add(1), Ordering::Relaxed); + ring[idx].store(hint_key, Ordering::Relaxed); } pub(super) fn mark_candidate_if_new( @@ -343,6 +385,67 @@ mod web_mode_tests { } } +#[cfg(test)] +mod bounded_registry_tests { + use std::sync::Arc; + + use super::*; + + #[test] + fn parallel_sticky_hints_never_exceed_their_hard_caps() { + const ATTEMPTS: usize = 10_000; + + let shared = ProxySharedState::new(); + std::thread::scope(|scope| { + for worker in 0..16 { + let shared = Arc::clone(&shared); + scope.spawn(move || { + for index in (worker..ATTEMPTS).step_by(16) { + let octets = (index as u32).to_be_bytes(); + let peer_ip = IpAddr::V4(std::net::Ipv4Addr::new( + octets[1], + octets[2], + octets[3], + worker as u8, + )); + sticky_hint_record_success_in( + shared.as_ref(), + peer_ip, + index as u64 | 1, + Some(&format!("host-{index}.example")), + ); + } + }); + } + }); + + assert_eq!( + shared.handshake.sticky_user_by_ip.len(), + STICKY_HINT_MAX_ENTRIES + ); + assert_eq!( + shared.handshake.sticky_user_by_ip_prefix.len(), + STICKY_HINT_MAX_ENTRIES + ); + assert_eq!( + shared.handshake.sticky_user_by_sni_hash.len(), + STICKY_HINT_MAX_ENTRIES + ); + assert_eq!( + shared.handshake.sticky_user_by_ip_slots.used(), + shared.handshake.sticky_user_by_ip.len() + ); + assert_eq!( + shared.handshake.sticky_user_by_ip_prefix_slots.used(), + shared.handshake.sticky_user_by_ip_prefix.len() + ); + assert_eq!( + shared.handshake.sticky_user_by_sni_hash_slots.used(), + shared.handshake.sticky_user_by_sni_hash.len() + ); + } +} + pub(super) fn decode_user_secrets_in( shared: &ProxySharedState, config: &ProxyConfig, diff --git a/src/proxy/handshake/auth_probe.rs b/src/proxy/handshake/auth_probe.rs index 13e0428..4cd2534 100644 --- a/src/proxy/handshake/auth_probe.rs +++ b/src/proxy/handshake/auth_probe.rs @@ -98,9 +98,14 @@ pub(super) fn auth_probe_is_throttled_in( }; if auth_probe_state_expired(&entry, now) { drop(entry); - state.remove_if(&peer_ip, |_, current| { - auth_probe_state_expired(current, now) - }); + if state + .remove_if(&peer_ip, |_, current| { + auth_probe_state_expired(current, now) + }) + .is_some() + { + shared.handshake.auth_probe_slots.release(); + } return false; } now < entry.blocked_until @@ -118,9 +123,14 @@ pub(super) fn auth_probe_saturation_grace_exhausted_in( }; if auth_probe_state_expired(&entry, now) { drop(entry); - state.remove_if(&peer_ip, |_, current| { - auth_probe_state_expired(current, now) - }); + if state + .remove_if(&peer_ip, |_, current| { + auth_probe_state_expired(current, now) + }) + .is_some() + { + shared.handshake.auth_probe_slots.release(); + } return false; } @@ -216,7 +226,13 @@ pub(super) fn auth_probe_record_failure_in( ) { let peer_ip = normalize_auth_probe_ip(peer_ip); let state = &shared.handshake.auth_probe; - auth_probe_record_failure_with_state_in(shared, state, peer_ip, now); + auth_probe_record_failure_with_state_and_budget_in( + shared, + state, + Some(&shared.handshake.auth_probe_slots), + peer_ip, + now, + ); } pub(super) fn auth_probe_record_failure_with_state_in( @@ -224,6 +240,16 @@ pub(super) fn auth_probe_record_failure_with_state_in( state: &DashMap, peer_ip: IpAddr, now: Instant, +) { + auth_probe_record_failure_with_state_and_budget_in(shared, state, None, peer_ip, now); +} + +fn auth_probe_record_failure_with_state_and_budget_in( + shared: &ProxySharedState, + state: &DashMap, + slots: Option<&crate::slot_budget::SlotBudget>, + peer_ip: IpAddr, + now: Instant, ) { let make_new_state = || AuthProbeState { fail_streak: 1, @@ -279,6 +305,9 @@ pub(super) fn auth_probe_record_failure_with_state_in( }) .is_some() { + if let Some(slots) = slots { + slots.release(); + } break; } continue; @@ -347,9 +376,15 @@ pub(super) fn auth_probe_record_failure_with_state_in( } for stale_key in stale_keys { - state.remove_if(&stale_key, |_, current| { - auth_probe_state_expired(current, now) - }); + if state + .remove_if(&stale_key, |_, current| { + auth_probe_state_expired(current, now) + }) + .is_some() + && let Some(slots) = slots + { + slots.release(); + } } if state.len() < AUTH_PROBE_TRACK_MAX_ENTRIES { @@ -360,19 +395,37 @@ pub(super) fn auth_probe_record_failure_with_state_in( auth_probe_note_saturation_in(shared, now); return; }; - state.remove_if(&evict_key, |_, current| { - current.fail_streak == evict_fail_streak && current.last_seen == evict_last_seen - }); + if state + .remove_if(&evict_key, |_, current| { + current.fail_streak == evict_fail_streak && current.last_seen == evict_last_seen + }) + .is_some() + && let Some(slots) = slots + { + slots.release(); + } auth_probe_note_saturation_in(shared, now); } } + let slot = if let Some(slots) = slots { + let Some(slot) = slots.try_acquire() else { + auth_probe_note_saturation_in(shared, now); + return; + }; + Some(slot) + } else { + None + }; match state.entry(peer_ip) { Entry::Occupied(mut entry) => { update_existing(entry.get_mut()); } Entry::Vacant(entry) => { entry.insert(make_new_state()); + if let Some(slot) = slot { + slot.commit(); + } } } } @@ -380,125 +433,15 @@ pub(super) fn auth_probe_record_failure_with_state_in( pub(super) fn auth_probe_record_success_in(shared: &ProxySharedState, peer_ip: IpAddr) { let peer_ip = normalize_auth_probe_ip(peer_ip); let state = &shared.handshake.auth_probe; - state.remove(&peer_ip); -} - -#[cfg(test)] -pub(crate) fn auth_probe_record_failure_for_testing( - shared: &ProxySharedState, - peer_ip: IpAddr, - now: Instant, -) { - auth_probe_record_failure_in(shared, peer_ip, now); -} - -#[cfg(test)] -pub(crate) fn auth_probe_fail_streak_for_testing_in_shared( - shared: &ProxySharedState, - peer_ip: IpAddr, -) -> Option { - let peer_ip = normalize_auth_probe_ip(peer_ip); - shared - .handshake - .auth_probe - .get(&peer_ip) - .map(|entry| entry.fail_streak) -} - -#[cfg(test)] -pub(crate) fn clear_auth_probe_state_for_testing_in_shared(shared: &ProxySharedState) { - shared.handshake.auth_probe.clear(); - match shared.handshake.auth_probe_saturation.lock() { - Ok(mut saturation) => { - *saturation = None; - } - Err(poisoned) => { - let mut saturation = poisoned.into_inner(); - *saturation = None; - shared.handshake.auth_probe_saturation.clear_poison(); - } + if state.remove(&peer_ip).is_some() { + shared.handshake.auth_probe_slots.release(); } } #[cfg(test)] -pub(crate) fn auth_probe_state_for_testing_in_shared( - shared: &ProxySharedState, -) -> &DashMap { - &shared.handshake.auth_probe -} - +mod testing; #[cfg(test)] -pub(crate) fn auth_probe_saturation_state_for_testing_in_shared( - shared: &ProxySharedState, -) -> &Mutex> { - &shared.handshake.auth_probe_saturation -} - -#[cfg(test)] -pub(crate) fn auth_probe_saturation_state_lock_for_testing_in_shared( - shared: &ProxySharedState, -) -> std::sync::MutexGuard<'_, Option> { - shared - .handshake - .auth_probe_saturation - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()) -} - -#[cfg(test)] -pub(crate) fn clear_unknown_sni_warn_state_for_testing_in_shared(shared: &ProxySharedState) { - let mut guard = shared - .handshake - .unknown_sni_warn_next_allowed - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()); - *guard = None; -} - -#[cfg(test)] -pub(crate) fn should_emit_unknown_sni_warn_for_testing_in_shared( - shared: &ProxySharedState, - now: Instant, -) -> bool { - should_emit_unknown_sni_warn_in(shared, now) -} - -#[cfg(test)] -pub(crate) fn clear_warned_secrets_for_testing_in_shared(shared: &ProxySharedState) { - if let Ok(mut guard) = shared.handshake.invalid_secret_warned.lock() { - guard.clear(); - } -} - -#[cfg(test)] -pub(crate) fn warned_secrets_for_testing_in_shared( - shared: &ProxySharedState, -) -> &Mutex> { - &shared.handshake.invalid_secret_warned -} - -#[cfg(test)] -pub(crate) fn auth_probe_is_throttled_for_testing_in_shared( - shared: &ProxySharedState, - peer_ip: IpAddr, -) -> bool { - auth_probe_is_throttled_in(shared, peer_ip, Instant::now()) -} - -#[cfg(test)] -pub(crate) fn auth_probe_saturation_is_throttled_for_testing_in_shared( - shared: &ProxySharedState, -) -> bool { - auth_probe_saturation_is_throttled_in(shared, Instant::now()) -} - -#[cfg(test)] -pub(crate) fn auth_probe_saturation_is_throttled_at_for_testing_in_shared( - shared: &ProxySharedState, - now: Instant, -) -> bool { - auth_probe_saturation_is_throttled_in(shared, now) -} +pub(crate) use testing::*; #[inline] pub(super) fn find_matching_tls_domain<'a>(config: &'a ProxyConfig, sni: &str) -> Option<&'a str> { diff --git a/src/proxy/handshake/auth_probe/testing.rs b/src/proxy/handshake/auth_probe/testing.rs new file mode 100644 index 0000000..24eb86d --- /dev/null +++ b/src/proxy/handshake/auth_probe/testing.rs @@ -0,0 +1,182 @@ +use super::*; + +/// Records one deterministic authentication failure against an isolated shared state. +pub(crate) fn auth_probe_record_failure_for_testing( + shared: &ProxySharedState, + peer_ip: IpAddr, + now: Instant, +) { + auth_probe_record_failure_in(shared, peer_ip, now); +} + +/// Returns the normalized peer failure streak from an isolated shared state. +pub(crate) fn auth_probe_fail_streak_for_testing_in_shared( + shared: &ProxySharedState, + peer_ip: IpAddr, +) -> Option { + let peer_ip = normalize_auth_probe_ip(peer_ip); + shared + .handshake + .auth_probe + .get(&peer_ip) + .map(|entry| entry.fail_streak) +} + +/// Clears probe entries, exact capacity accounting, and saturation state together. +pub(crate) fn clear_auth_probe_state_for_testing_in_shared(shared: &ProxySharedState) { + let removed = shared.handshake.auth_probe.len(); + assert_eq!(shared.handshake.auth_probe_slots.used(), removed); + shared.handshake.auth_probe.clear(); + shared.handshake.auth_probe_slots.release_many(removed); + match shared.handshake.auth_probe_saturation.lock() { + Ok(mut saturation) => { + *saturation = None; + } + Err(poisoned) => { + let mut saturation = poisoned.into_inner(); + *saturation = None; + shared.handshake.auth_probe_saturation.clear_poison(); + } + } +} + +/// Inserts one fixture entry while preserving exact registry capacity accounting. +pub(crate) fn insert_auth_probe_state_for_testing_in_shared( + shared: &ProxySharedState, + peer_ip: IpAddr, + state: AuthProbeState, +) { + let peer_ip = normalize_auth_probe_ip(peer_ip); + let slot = shared + .handshake + .auth_probe_slots + .try_acquire() + .expect("test auth-probe registry capacity must be available"); + match shared.handshake.auth_probe.entry(peer_ip) { + Entry::Occupied(mut entry) => { + entry.insert(state); + } + Entry::Vacant(entry) => { + entry.insert(state); + slot.commit(); + } + } +} + +/// Exposes the isolated probe registry to adversarial tests. +pub(crate) fn auth_probe_state_for_testing_in_shared( + shared: &ProxySharedState, +) -> &DashMap { + &shared.handshake.auth_probe +} + +/// Returns exact committed probe slots for capacity assertions. +pub(crate) fn auth_probe_slots_for_testing_in_shared(shared: &ProxySharedState) -> usize { + shared.handshake.auth_probe_slots.used() +} + +/// Exposes the isolated saturation state mutex to tests. +pub(crate) fn auth_probe_saturation_state_for_testing_in_shared( + shared: &ProxySharedState, +) -> &Mutex> { + &shared.handshake.auth_probe_saturation +} + +/// Locks isolated saturation state while recovering poisoned test fixtures. +pub(crate) fn auth_probe_saturation_state_lock_for_testing_in_shared( + shared: &ProxySharedState, +) -> std::sync::MutexGuard<'_, Option> { + shared + .handshake + .auth_probe_saturation + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) +} + +/// Resets the isolated unknown-SNI warning rate limiter. +pub(crate) fn clear_unknown_sni_warn_state_for_testing_in_shared(shared: &ProxySharedState) { + let mut guard = shared + .handshake + .unknown_sni_warn_next_allowed + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + *guard = None; +} + +/// Evaluates unknown-SNI warning admission at a deterministic instant. +pub(crate) fn should_emit_unknown_sni_warn_for_testing_in_shared( + shared: &ProxySharedState, + now: Instant, +) -> bool { + should_emit_unknown_sni_warn_in(shared, now) +} + +/// Clears the isolated invalid-secret warning deduplication set. +pub(crate) fn clear_warned_secrets_for_testing_in_shared(shared: &ProxySharedState) { + if let Ok(mut guard) = shared.handshake.invalid_secret_warned.lock() { + guard.clear(); + } +} + +/// Exposes the isolated invalid-secret warning set to tests. +pub(crate) fn warned_secrets_for_testing_in_shared( + shared: &ProxySharedState, +) -> &Mutex> { + &shared.handshake.invalid_secret_warned +} + +/// Evaluates peer throttling against the current test clock. +pub(crate) fn auth_probe_is_throttled_for_testing_in_shared( + shared: &ProxySharedState, + peer_ip: IpAddr, +) -> bool { + auth_probe_is_throttled_in(shared, peer_ip, Instant::now()) +} + +/// Evaluates global saturation throttling against the current test clock. +pub(crate) fn auth_probe_saturation_is_throttled_for_testing_in_shared( + shared: &ProxySharedState, +) -> bool { + auth_probe_saturation_is_throttled_in(shared, Instant::now()) +} + +/// Evaluates global saturation throttling at a deterministic instant. +pub(crate) fn auth_probe_saturation_is_throttled_at_for_testing_in_shared( + shared: &ProxySharedState, + now: Instant, +) -> bool { + auth_probe_saturation_is_throttled_in(shared, now) +} + +#[test] +fn parallel_distinct_failures_respect_exact_auth_probe_capacity() { + const ATTEMPTS: usize = 10_000; + + let shared = ProxySharedState::new(); + std::thread::scope(|scope| { + for worker in 0..16 { + let shared = Arc::clone(&shared); + scope.spawn(move || { + for index in (worker..ATTEMPTS).step_by(16) { + let octets = (index as u32).to_be_bytes(); + let peer_ip = IpAddr::V4(std::net::Ipv4Addr::new( + octets[1], + octets[2], + octets[3], + worker as u8, + )); + auth_probe_record_failure_in(shared.as_ref(), peer_ip, Instant::now()); + } + }); + } + }); + + assert_eq!( + shared.handshake.auth_probe.len(), + AUTH_PROBE_TRACK_MAX_ENTRIES + ); + assert_eq!( + auth_probe_slots_for_testing_in_shared(shared.as_ref()), + AUTH_PROBE_TRACK_MAX_ENTRIES + ); +} diff --git a/src/proxy/handshake/mtproto.rs b/src/proxy/handshake/mtproto.rs index da9b351..b5a4b18 100644 --- a/src/proxy/handshake/mtproto.rs +++ b/src/proxy/handshake/mtproto.rs @@ -151,10 +151,14 @@ where if let Some(snapshot) = config.runtime_user_auth() { let sticky_ip_hint = sticky_hint_get_by_ip(shared, peer.ip()); let sticky_prefix_hint = sticky_hint_get_by_ip_prefix(shared, peer.ip()); + let sticky_ip_candidates = + sticky_ip_hint.and_then(|hint_key| snapshot.candidate_ids_by_hint_key(hint_key)); + let sticky_prefix_candidates = + sticky_prefix_hint.and_then(|hint_key| snapshot.candidate_ids_by_hint_key(hint_key)); let preferred_user_id = preferred_user.and_then(|user| snapshot.user_id_by_name(user)); let exact_user_id = exact_user.and_then(|user| snapshot.user_id_by_name(user)); - let has_hint = sticky_ip_hint.is_some() - || sticky_prefix_hint.is_some() + let has_hint = sticky_ip_candidates.is_some_and(|ids| !ids.is_empty()) + || sticky_prefix_candidates.is_some_and(|ids| !ids.is_empty()) || preferred_user_id.is_some() || exact_user_id.is_some(); let overload = auth_probe_saturation_is_throttled_in(shared, Instant::now()); @@ -204,9 +208,17 @@ where let mut matched = exact_user_id.is_some_and(|user_id| try_user_id!(user_id)); if exact_user.is_none() - && let Some(user_id) = sticky_ip_hint + && let Some(candidate_ids) = sticky_ip_candidates { - matched = try_user_id!(user_id); + for &user_id in candidate_ids { + if try_user_id!(user_id) { + matched = true; + break; + } + if budget_exhausted { + break; + } + } } if exact_user.is_none() @@ -218,9 +230,17 @@ where if exact_user.is_none() && !matched - && let Some(user_id) = sticky_prefix_hint + && let Some(candidate_ids) = sticky_prefix_candidates { - matched = try_user_id!(user_id); + for &user_id in candidate_ids { + if try_user_id!(user_id) { + matched = true; + break; + } + if budget_exhausted { + break; + } + } } if exact_user.is_none() && !matched && !budget_exhausted { @@ -231,18 +251,22 @@ where .recent_user_ring_seq .load(Ordering::Relaxed); let scan_limit = ring.len().min(RECENT_USER_RING_SCAN_LIMIT); - for offset in 0..scan_limit { + 'recent_hints: for offset in 0..scan_limit { let idx = (next_seq as usize + ring.len() - 1 - offset) % ring.len(); - let encoded_user_id = ring[idx].load(Ordering::Relaxed); - if encoded_user_id == 0 { + let hint_key = ring[idx].load(Ordering::Relaxed); + if hint_key == 0 { continue; } - if try_user_id!(encoded_user_id - 1) { - matched = true; - break; - } - if budget_exhausted { - break; + if let Some(candidate_ids) = snapshot.candidate_ids_by_hint_key(hint_key) { + for &user_id in candidate_ids { + if try_user_id!(user_id) { + matched = true; + break 'recent_hints; + } + if budget_exhausted { + break 'recent_hints; + } + } } } } @@ -357,8 +381,10 @@ where auth_probe_record_success_in(shared, peer.ip()); if let Some(user_id) = matched_user_id { - sticky_hint_record_success_in(shared, peer.ip(), user_id, None); - record_recent_user_success_in(shared, user_id); + if let Some(entry) = snapshot.entry_by_id(user_id) { + sticky_hint_record_success_in(shared, peer.ip(), entry.hint_key, None); + record_recent_user_success_in(shared, entry.hint_key); + } } let max_pending = config.general.crypto_pending_buffer; diff --git a/src/proxy/handshake/tls_handshake.rs b/src/proxy/handshake/tls_handshake.rs index c98dc17..cde4325 100644 --- a/src/proxy/handshake/tls_handshake.rs +++ b/src/proxy/handshake/tls_handshake.rs @@ -257,14 +257,15 @@ where secret: validated_secret, user_id: validated_user_id, } = validation; - // Reject known replay digests before expensive cache/domain/ALPN policy work. + // Reserve the replay digest before any asynchronous policy work so a concurrent + // duplicate cannot pass the check-to-commit window. let digest_half = &validation_digest[..tls::TLS_DIGEST_HALF_LEN]; - if replay_checker.check_tls_digest(digest_half) { + let Some(replay_claim) = replay_checker.claim_tls_digest(digest_half) else { auth_probe_record_failure_in(shared, peer.ip(), Instant::now()); maybe_apply_server_hello_delay(config).await; warn!(peer = %peer, "TLS replay attack detected (duplicate digest)"); return HandshakeResult::BadClient { reader, writer }; - } + }; let selected_tls_domain = matched_tls_domain.unwrap_or(config.censorship.tls_domain.as_str()); let cached_entry = if config.censorship.tls_emulation { @@ -337,8 +338,12 @@ where None }; - // Add replay digest only for policy-valid handshakes. - replay_checker.add_tls_digest(digest_half); + // Commit only policy-valid handshakes; early returns release the pending claim. + if !replay_claim.commit() { + auth_probe_record_failure_in(shared, peer.ip(), Instant::now()); + warn!(peer = %peer, "TLS replay claim lost before commit"); + return HandshakeResult::BadClient { reader, writer }; + } let validation_session_id_slice = &validation_session_id[..validation_session_id_len]; @@ -391,8 +396,13 @@ where auth_probe_record_success_in(shared, peer.ip()); if let Some(user_id) = validated_user_id { - sticky_hint_record_success_in(shared, peer.ip(), user_id, client_sni.as_deref()); - record_recent_user_success_in(shared, user_id); + if let Some(entry) = config + .runtime_user_auth() + .and_then(|snapshot| snapshot.entry_by_id(user_id)) + { + sticky_hint_record_success_in(shared, peer.ip(), entry.hint_key, client_sni.as_deref()); + record_recent_user_success_in(shared, entry.hint_key); + } } HandshakeResult::Success(( diff --git a/src/proxy/handshake/tls_validation.rs b/src/proxy/handshake/tls_validation.rs index 4b83c08..f7c949c 100644 --- a/src/proxy/handshake/tls_validation.rs +++ b/src/proxy/handshake/tls_validation.rs @@ -40,11 +40,17 @@ pub(super) async fn validate_tls_client( }; let sticky_ip_hint = sticky_hint_get_by_ip(shared, peer.ip()); + let sticky_ip_candidates = + sticky_ip_hint.and_then(|hint_key| snapshot.candidate_ids_by_hint_key(hint_key)); let preferred_user_id = preferred_user_hint.and_then(|user| snapshot.user_id_by_name(user)); let sticky_sni_hint = client_sni .as_deref() .and_then(|sni| sticky_hint_get_by_sni(shared, sni)); + let sticky_sni_candidates = + sticky_sni_hint.and_then(|hint_key| snapshot.candidate_ids_by_hint_key(hint_key)); let sticky_prefix_hint = sticky_hint_get_by_ip_prefix(shared, peer.ip()); + let sticky_prefix_candidates = + sticky_prefix_hint.and_then(|hint_key| snapshot.candidate_ids_by_hint_key(hint_key)); let sni_candidates = client_sni .as_deref() .and_then(|sni| snapshot.sni_candidates(sni)); @@ -52,10 +58,10 @@ pub(super) async fn validate_tls_client( .as_deref() .and_then(|sni| snapshot.sni_initial_candidates(sni)); - let has_hint = sticky_ip_hint.is_some() + let has_hint = sticky_ip_candidates.is_some_and(|ids| !ids.is_empty()) || preferred_user_id.is_some() - || sticky_sni_hint.is_some() - || sticky_prefix_hint.is_some() + || sticky_sni_candidates.is_some_and(|ids| !ids.is_empty()) + || sticky_prefix_candidates.is_some_and(|ids| !ids.is_empty()) || sni_candidates.is_some_and(|ids| !ids.is_empty()) || sni_initial_candidates.is_some_and(|ids| !ids.is_empty()); let overload = auth_probe_saturation_is_throttled_in(shared, Instant::now()); @@ -95,20 +101,44 @@ pub(super) async fn validate_tls_client( } let mut matched = false; - if let Some(user_id) = sticky_ip_hint { - matched = try_user_id!(user_id); + if let Some(candidate_ids) = sticky_ip_candidates { + for &user_id in candidate_ids { + if try_user_id!(user_id) { + matched = true; + break; + } + if budget_exhausted { + break; + } + } } if !matched && let Some(user_id) = preferred_user_id { matched = try_user_id!(user_id); } - if !matched && let Some(user_id) = sticky_sni_hint { - matched = try_user_id!(user_id); + if !matched && let Some(candidate_ids) = sticky_sni_candidates { + for &user_id in candidate_ids { + if try_user_id!(user_id) { + matched = true; + break; + } + if budget_exhausted { + break; + } + } } - if !matched && let Some(user_id) = sticky_prefix_hint { - matched = try_user_id!(user_id); + if !matched && let Some(candidate_ids) = sticky_prefix_candidates { + for &user_id in candidate_ids { + if try_user_id!(user_id) { + matched = true; + break; + } + if budget_exhausted { + break; + } + } } if !matched @@ -149,18 +179,22 @@ pub(super) async fn validate_tls_client( .recent_user_ring_seq .load(Ordering::Relaxed); let scan_limit = ring.len().min(RECENT_USER_RING_SCAN_LIMIT); - for offset in 0..scan_limit { + 'recent_hints: for offset in 0..scan_limit { let idx = (next_seq as usize + ring.len() - 1 - offset) % ring.len(); - let encoded_user_id = ring[idx].load(Ordering::Relaxed); - if encoded_user_id == 0 { + let hint_key = ring[idx].load(Ordering::Relaxed); + if hint_key == 0 { continue; } - if try_user_id!(encoded_user_id - 1) { - matched = true; - break; - } - if budget_exhausted { - break; + if let Some(candidate_ids) = snapshot.candidate_ids_by_hint_key(hint_key) { + for &user_id in candidate_ids { + if try_user_id!(user_id) { + matched = true; + break 'recent_hints; + } + if budget_exhausted { + break 'recent_hints; + } + } } } } diff --git a/src/proxy/middle_relay.rs b/src/proxy/middle_relay.rs index 8182c03..56a4a5a 100644 --- a/src/proxy/middle_relay.rs +++ b/src/proxy/middle_relay.rs @@ -15,6 +15,7 @@ use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc, oneshot, watch}; use tokio::time::timeout; use tokio_util::sync::CancellationToken; +use tokio_util::task::AbortOnDropHandle; use tracing::{debug, info, trace, warn}; use crate::config::{ConntrackPressureProfile, ProxyConfig}; @@ -31,7 +32,8 @@ use crate::proxy::shared_state::{ }; use crate::proxy::traffic_limiter::{RateDirection, TrafficLease, next_refill_delay}; use crate::stats::{ - MeD2cFlushReason, MeD2cQuotaRejectStage, MeD2cWriteMode, QuotaReserveError, Stats, UserStats, + MeD2cFlushReason, MeD2cQuotaRejectStage, MeD2cWriteMode, QuotaReserveError, Stats, + UserQuotaHandle, UserStats, }; use crate::stream::{BufferPool, CryptoReader, CryptoWriter, PooledBuffer}; use crate::transport::middle_proxy::{ConnLease, MePool, MeResponse, proto_flags_for_tag}; @@ -108,6 +110,7 @@ pub(crate) async fn handle_via_middle_proxy( session_id: u64, session_cancel: CancellationToken, shared: Arc, + quota_handle: UserQuotaHandle, ) -> Result<()> where R: AsyncRead + Unpin + Send + 'static, @@ -129,6 +132,7 @@ where session_cancel, shared, ConntrackClosePolicy::Publish, + quota_handle, ) .await } @@ -152,7 +156,7 @@ const ME_D2C_FLUSH_BATCH_MAX_FRAMES_MIN: usize = 1; const ME_D2C_FLUSH_BATCH_MAX_BYTES_MIN: usize = 4096; const ME_D2C_FRAME_BUF_SHRINK_HYSTERESIS_FACTOR: usize = 2; const ME_D2C_SINGLE_WRITE_COALESCE_MAX_BYTES: usize = 128 * 1024; -const QUOTA_RESERVE_SPIN_RETRIES: usize = 32; +const QUOTA_RESERVE_ATTEMPTS_PER_ROUND: usize = 4; const QUOTA_RESERVE_BACKOFF_MIN_MS: u64 = 1; const QUOTA_RESERVE_BACKOFF_MAX_MS: u64 = 16; const QUOTA_RESERVE_MAX_BACKOFF_ROUNDS: usize = 16; diff --git a/src/proxy/middle_relay/d2c.rs b/src/proxy/middle_relay/d2c.rs index a1906fc..03f9185 100644 --- a/src/proxy/middle_relay/d2c.rs +++ b/src/proxy/middle_relay/d2c.rs @@ -135,6 +135,7 @@ pub(crate) async fn process_me_writer_response( where W: AsyncWrite + Unpin + Send + 'static, { + let quota_handle = quota_limit.map(|_| stats.current_user_quota_handle(user)); process_me_writer_response_with_traffic_lease( response, client_writer, @@ -144,6 +145,7 @@ where stats, user, quota_user_stats, + quota_handle.as_ref(), quota_limit, quota_soft_overshoot_bytes, None, @@ -165,6 +167,7 @@ pub(crate) async fn process_me_writer_response_with_traffic_lease( stats: &Stats, user: &str, quota_user_stats: Option<&UserStats>, + quota_handle: Option<&UserQuotaHandle>, quota_limit: Option, quota_soft_overshoot_bytes: u64, traffic_lease: Option<&Arc>, @@ -185,10 +188,15 @@ where trace!(conn_id, bytes = data.len(), flags, "ME->C data"); } let data_len = data.len() as u64; - if let (Some(limit), Some(user_stats)) = (quota_limit, quota_user_stats) { + if let (Some(limit), Some(quota_handle)) = (quota_limit, quota_handle) { let soft_limit = quota_soft_cap(limit, quota_soft_overshoot_bytes); match reserve_user_quota_with_yield( - user_stats, data_len, soft_limit, stats, cancel, None, + quota_handle, + data_len, + soft_limit, + stats, + cancel, + None, ) .await { diff --git a/src/proxy/middle_relay/quota.rs b/src/proxy/middle_relay/quota.rs index 3a04c00..e7a6eb4 100644 --- a/src/proxy/middle_relay/quota.rs +++ b/src/proxy/middle_relay/quota.rs @@ -12,7 +12,7 @@ pub(super) fn quota_soft_cap(limit: u64, overshoot: u64) -> u64 { } pub(super) async fn reserve_user_quota_with_yield( - user_stats: &UserStats, + quota_handle: &UserQuotaHandle, bytes: u64, limit: u64, stats: &Stats, @@ -22,15 +22,14 @@ pub(super) async fn reserve_user_quota_with_yield( let mut backoff_ms = QUOTA_RESERVE_BACKOFF_MIN_MS; let mut backoff_rounds = 0usize; loop { - for _ in 0..QUOTA_RESERVE_SPIN_RETRIES { - match user_stats.quota_try_reserve(bytes, limit) { - Ok(total) => return Ok(total), + for _ in 0..QUOTA_RESERVE_ATTEMPTS_PER_ROUND { + match quota_handle.try_reserve(bytes, limit) { + Ok(reservation) => return Ok(reservation.commit()), Err(QuotaReserveError::LimitExceeded) => { return Err(MiddleQuotaReserveError::LimitExceeded); } Err(QuotaReserveError::Contended) => { stats.increment_quota_contention_total(); - std::hint::spin_loop(); } } } diff --git a/src/proxy/middle_relay/session.rs b/src/proxy/middle_relay/session.rs index bee6675..3531cb7 100644 --- a/src/proxy/middle_relay/session.rs +++ b/src/proxy/middle_relay/session.rs @@ -2,11 +2,15 @@ use super::*; // Bounded C2ME sender and downstream writer tasks. mod tasks; +// Child-task ownership aborts relay tasks when the parent future is cancelled. +mod children; // Conntrack close classification. mod close_reason; +use children::RelayChildTasks; use close_reason::classify_conntrack_close_reason; use tasks::{run_c2me_sender, run_me_writer}; + struct RelayConnLease { connection: Option, conn_id: u64, @@ -61,6 +65,7 @@ pub(crate) async fn handle_via_middle_proxy_with_conntrack( session_cancel: CancellationToken, shared: Arc, conntrack_close_policy: ConntrackClosePolicy, + quota_handle: UserQuotaHandle, ) -> Result<()> where R: AsyncRead + Unpin + Send + 'static, @@ -73,6 +78,7 @@ where let quota_limit = config.access.user_data_quota.get(&user).copied(); let quota_user_stats = quota_limit.map(|_| stats.get_or_create_user_stats_handle(&user)); + let quota_handle = quota_limit.map(|_| quota_handle); let peer = success.peer; let traffic_lease = shared.traffic_limiter.acquire_lease(&user, peer.ip()); let proto_tag = success.proto_tag; @@ -183,7 +189,7 @@ where let c2me_byte_semaphore = Arc::new(Semaphore::new(c2me_byte_budget)); let (c2me_tx, c2me_rx) = mpsc::channel::(c2me_channel_capacity); let me_pool_c2me = me_pool.clone(); - let mut c2me_sender = tokio::spawn(run_c2me_sender( + let c2me_sender = AbortOnDropHandle::new(tokio::spawn(run_c2me_sender( c2me_rx, me_pool_c2me, conn_id, @@ -191,7 +197,7 @@ where peer, translated_local_addr, effective_tag_array, - )); + ))); let (stop_tx, stop_rx) = oneshot::channel::<()>(); let flow_cancel = CancellationToken::new(); @@ -200,18 +206,20 @@ where let rng_clone = rng.clone(); let user_clone = user.clone(); let quota_user_stats_me_writer = quota_user_stats.clone(); + let quota_handle_me_writer = quota_handle.clone(); let traffic_lease_me_writer = traffic_lease.clone(); let flow_cancel_me_writer = flow_cancel.clone(); let last_downstream_activity_ms_clone = last_downstream_activity_ms.clone(); let bytes_me2c_clone = bytes_me2c.clone(); let d2c_flush_policy = MeD2cFlushPolicy::from_config(&config); - let mut me_writer = tokio::spawn(run_me_writer( + let me_writer = AbortOnDropHandle::new(tokio::spawn(run_me_writer( crypto_writer, me_rx_task, stats_clone, rng_clone, user_clone, quota_user_stats_me_writer, + quota_handle_me_writer, quota_limit, traffic_lease_me_writer, flow_cancel_me_writer, @@ -222,7 +230,13 @@ where session_started_at, conn_id, stop_rx, - )); + ))); + let mut child_tasks = RelayChildTasks { + c2me_sender, + me_writer, + flow_cancel: flow_cancel.clone(), + stop_tx: Some(stop_tx), + }; let mut main_result: Result<()> = Ok(()); let mut client_closed = false; @@ -340,11 +354,11 @@ where forensics.bytes_c2me = forensics .bytes_c2me .saturating_add(payload.len() as u64); - if let (Some(limit), Some(user_stats)) = - (quota_limit, quota_user_stats.as_deref()) + if let (Some(limit), Some(quota_handle)) = + (quota_limit, quota_handle.as_ref()) { match reserve_user_quota_with_yield( - user_stats, + quota_handle, payload.len() as u64, limit, stats.as_ref(), @@ -379,7 +393,11 @@ where break; } } - stats.add_user_octets_from_handle(user_stats, payload.len() as u64); + if let Some(user_stats) = quota_user_stats.as_deref() { + stats.add_user_octets_from_handle(user_stats, payload.len() as u64); + } else { + stats.add_user_octets_from(&user, payload.len() as u64); + } } else { stats.add_user_octets_from(&user, payload.len() as u64); } @@ -446,28 +464,30 @@ where } drop(c2me_tx); - let c2me_result = match timeout(ME_CHILD_JOIN_TIMEOUT, &mut c2me_sender).await { + let c2me_result = match timeout(ME_CHILD_JOIN_TIMEOUT, &mut child_tasks.c2me_sender).await { Ok(joined) => { joined.unwrap_or_else(|e| Err(ProxyError::Proxy(format!("ME sender join error: {e}")))) } Err(_) => { stats.increment_me_child_join_timeout_total(); stats.increment_me_child_abort_total(); - c2me_sender.abort(); + child_tasks.c2me_sender.abort(); Err(ProxyError::Proxy("ME sender join timeout".into())) } }; flow_cancel.cancel(); - let _ = stop_tx.send(()); - let mut writer_result = match timeout(ME_CHILD_JOIN_TIMEOUT, &mut me_writer).await { + if let Some(stop_tx) = child_tasks.stop_tx.take() { + let _ = stop_tx.send(()); + } + let mut writer_result = match timeout(ME_CHILD_JOIN_TIMEOUT, &mut child_tasks.me_writer).await { Ok(joined) => { joined.unwrap_or_else(|e| Err(ProxyError::Proxy(format!("ME writer join error: {e}")))) } Err(_) => { stats.increment_me_child_join_timeout_total(); stats.increment_me_child_abort_total(); - me_writer.abort(); + child_tasks.me_writer.abort(); Err(ProxyError::Proxy("ME writer join timeout".into())) } }; diff --git a/src/proxy/middle_relay/session/children.rs b/src/proxy/middle_relay/session/children.rs new file mode 100644 index 0000000..1b1cd8c --- /dev/null +++ b/src/proxy/middle_relay/session/children.rs @@ -0,0 +1,65 @@ +use super::*; + +pub(super) struct RelayChildTasks { + pub(super) c2me_sender: AbortOnDropHandle>, + pub(super) me_writer: AbortOnDropHandle>, + pub(super) flow_cancel: CancellationToken, + pub(super) stop_tx: Option>, +} + +impl Drop for RelayChildTasks { + fn drop(&mut self) { + self.flow_cancel.cancel(); + if let Some(stop_tx) = self.stop_tx.take() { + let _ = stop_tx.send(()); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::{AtomicUsize, Ordering}; + + struct DropSignal(Arc); + + impl Drop for DropSignal { + fn drop(&mut self) { + self.0.fetch_add(1, Ordering::AcqRel); + } + } + + async fn pending_child(signal: DropSignal) -> Result<()> { + let _signal = signal; + std::future::pending::<()>().await; + Ok(()) + } + + #[tokio::test] + async fn relay_child_scope_drop_aborts_both_children() { + let dropped = Arc::new(AtomicUsize::new(0)); + let flow_cancel = CancellationToken::new(); + let (stop_tx, stop_rx) = oneshot::channel(); + let child_tasks = RelayChildTasks { + c2me_sender: AbortOnDropHandle::new(tokio::spawn(pending_child(DropSignal( + Arc::clone(&dropped), + )))), + me_writer: AbortOnDropHandle::new(tokio::spawn(pending_child(DropSignal(Arc::clone( + &dropped, + ))))), + flow_cancel: flow_cancel.clone(), + stop_tx: Some(stop_tx), + }; + + drop(child_tasks); + assert!(flow_cancel.is_cancelled()); + assert!(stop_rx.await.is_ok()); + tokio::time::timeout(Duration::from_secs(1), async { + while dropped.load(Ordering::Acquire) != 2 { + tokio::task::yield_now().await; + } + }) + .await + .expect("both relay child futures must be dropped after scope cancellation"); + } +} diff --git a/src/proxy/middle_relay/session/tasks.rs b/src/proxy/middle_relay/session/tasks.rs index 56f52dd..b2e747e 100644 --- a/src/proxy/middle_relay/session/tasks.rs +++ b/src/proxy/middle_relay/session/tasks.rs @@ -53,6 +53,7 @@ pub(super) async fn run_me_writer( rng_clone: Arc, user_clone: String, quota_user_stats_me_writer: Option>, + quota_handle_me_writer: Option, quota_limit: Option, traffic_lease_me_writer: Option>, flow_cancel_me_writer: CancellationToken, @@ -105,6 +106,7 @@ where stats_clone.as_ref(), &user_clone, quota_user_stats_me_writer.as_deref(), + quota_handle_me_writer.as_ref(), quota_limit, d2c_flush_policy.quota_soft_overshoot_bytes, traffic_lease_me_writer.as_ref(), @@ -167,6 +169,7 @@ where stats_clone.as_ref(), &user_clone, quota_user_stats_me_writer.as_deref(), + quota_handle_me_writer.as_ref(), quota_limit, d2c_flush_policy.quota_soft_overshoot_bytes, traffic_lease_me_writer.as_ref(), @@ -233,6 +236,7 @@ where stats_clone.as_ref(), &user_clone, quota_user_stats_me_writer.as_deref(), + quota_handle_me_writer.as_ref(), quota_limit, d2c_flush_policy.quota_soft_overshoot_bytes, traffic_lease_me_writer.as_ref(), @@ -304,6 +308,7 @@ where stats_clone.as_ref(), &user_clone, quota_user_stats_me_writer.as_deref(), + quota_handle_me_writer.as_ref(), quota_limit, d2c_flush_policy.quota_soft_overshoot_bytes, traffic_lease_me_writer.as_ref(), diff --git a/src/proxy/mod.rs b/src/proxy/mod.rs index b1b71bf..01b9f08 100644 --- a/src/proxy/mod.rs +++ b/src/proxy/mod.rs @@ -73,6 +73,9 @@ pub mod route_mode; pub mod session_eviction; pub mod shared_state; pub mod traffic_limiter; +pub(crate) mod user_admission; +// Process-wide per-user connection admission remains independent from telemetry. +pub(crate) mod user_connection_authority; pub use client::ClientHandler; #[allow(unused_imports)] diff --git a/src/proxy/relay.rs b/src/proxy/relay.rs index 6695026..0cc8cc5 100644 --- a/src/proxy/relay.rs +++ b/src/proxy/relay.rs @@ -290,6 +290,7 @@ where // ── Combine split halves into bidirectional streams ────────────── let client_combined = CombinedStream::new(client_reader, client_writer); let mut server = CombinedStream::new(server_reader, server_writer); + let quota_handle = stats.current_user_quota_handle(&user_owned); // Wrap client with stats/activity tracking let mut client = StatsIo::new_with_traffic_lease( @@ -297,6 +298,7 @@ where Arc::clone(&counters), Arc::clone(&stats), user_owned.clone(), + quota_handle, traffic_lease, quota_limit, Arc::clone("a_exceeded), diff --git a/src/proxy/relay/adaptive_copy.rs b/src/proxy/relay/adaptive_copy.rs index cae723c..ae6ca2a 100644 --- a/src/proxy/relay/adaptive_copy.rs +++ b/src/proxy/relay/adaptive_copy.rs @@ -19,7 +19,7 @@ use crate::proxy::direct_buffer_budget::{ DIRECT_BASE_C2S_BYTES, DIRECT_BASE_S2C_BYTES, DirectBufferBudget, DirectBufferLease, }; use crate::proxy::traffic_limiter::TrafficLease; -use crate::stats::Stats; +use crate::stats::{Stats, UserQuotaHandle}; use super::WATCHDOG_INTERVAL; use super::io::{SharedCounters, StatsIo, is_quota_io_error}; @@ -141,6 +141,7 @@ pub(crate) async fn relay_direct_adaptive( max_connections: u32, user: &str, stats: Arc, + quota_handle: UserQuotaHandle, quota_limit: Option, traffic_lease: Option>, activity_timeout: Duration, @@ -200,6 +201,7 @@ where Arc::clone(&counters), Arc::clone(&stats), user_owned.clone(), + quota_handle.clone(), traffic_lease.clone(), quota_limit, Arc::clone("a_exceeded), @@ -210,6 +212,7 @@ where Arc::clone(&counters), Arc::clone(&stats), user_owned.clone(), + quota_handle, traffic_lease, quota_limit, Arc::clone("a_exceeded), diff --git a/src/proxy/relay/io.rs b/src/proxy/relay/io.rs index fb30f3f..29332d2 100644 --- a/src/proxy/relay/io.rs +++ b/src/proxy/relay/io.rs @@ -1,5 +1,5 @@ use crate::proxy::traffic_limiter::{RateDirection, TrafficLease, next_refill_delay}; -use crate::stats::{Stats, UserStats}; +use crate::stats::{Stats, UserQuotaHandle, UserStats}; use std::io; use std::pin::Pin; use std::sync::Arc; @@ -16,10 +16,7 @@ mod quota; pub(super) use self::combined::CombinedStream; pub(super) use self::counters::SharedCounters; pub(super) use self::quota::is_quota_io_error; -use self::quota::{ - QUOTA_RESERVE_MAX_ROUNDS, QUOTA_RESERVE_SPIN_RETRIES, quota_io_error, - refund_reserved_quota_bytes, -}; +use self::quota::{QUOTA_RESERVE_MAX_ATTEMPTS_PER_POLL, quota_io_error}; pub(super) use self::quota::{quota_adaptive_interval_bytes, should_immediate_quota_check}; /// Transparent I/O wrapper that tracks per-user statistics and activity. @@ -43,6 +40,7 @@ pub(super) struct StatsIo { stats: Arc, user: String, user_stats: Arc, + quota_handle: UserQuotaHandle, traffic_lease: Option>, c2s_rate_debt_bytes: u64, c2s_wait: RateWaitState, @@ -74,11 +72,13 @@ impl StatsIo { quota_exceeded: Arc, epoch: Instant, ) -> Self { + let quota_handle = stats.current_user_quota_handle(&user); Self::new_with_traffic_lease( inner, counters, stats, user, + quota_handle, None, quota_limit, quota_exceeded, @@ -91,6 +91,7 @@ impl StatsIo { counters: Arc, stats: Arc, user: String, + quota_handle: UserQuotaHandle, traffic_lease: Option>, quota_limit: Option, quota_exceeded: Arc, @@ -105,6 +106,7 @@ impl StatsIo { stats, user, user_stats, + quota_handle, traffic_lease, c2s_rate_debt_bytes: 0, c2s_wait: RateWaitState::default(), @@ -213,51 +215,36 @@ impl AsyncRead for StatsIo { } let mut remaining_before = None; - let mut reserved_read_bytes = 0u64; + let mut quota_reservation = None; let mut read_limit = buf.remaining(); if let Some(limit) = this.quota_limit { - let used_before = this.user_stats.quota_used(); - let remaining = limit.saturating_sub(used_before); - if remaining == 0 { - this.quota_exceeded.store(true, Ordering::Release); - return Poll::Ready(Err(quota_io_error())); - } - remaining_before = Some(remaining); - read_limit = read_limit.min(remaining as usize); - if read_limit == 0 { - this.quota_exceeded.store(true, Ordering::Release); - return Poll::Ready(Err(quota_io_error())); - } - - let desired = read_limit as u64; - let mut reserve_rounds = 0usize; - while reserved_read_bytes == 0 { - for _ in 0..QUOTA_RESERVE_SPIN_RETRIES { - match this.user_stats.quota_try_reserve(desired, limit) { - Ok(_) => { - reserved_read_bytes = desired; - break; - } - Err(crate::stats::QuotaReserveError::LimitExceeded) => { - this.quota_exceeded.store(true, Ordering::Release); - return Poll::Ready(Err(quota_io_error())); - } - Err(crate::stats::QuotaReserveError::Contended) => { - this.stats.increment_quota_contention_total(); - } + for _ in 0..QUOTA_RESERVE_MAX_ATTEMPTS_PER_POLL { + let used_before = this.quota_handle.used(); + let remaining = limit.saturating_sub(used_before); + if remaining == 0 { + this.quota_exceeded.store(true, Ordering::Release); + return Poll::Ready(Err(quota_io_error())); + } + let desired = remaining.min(read_limit as u64); + match this.quota_handle.try_reserve(desired, limit) { + Ok(reservation) => { + remaining_before = Some(remaining); + read_limit = desired as usize; + quota_reservation = Some(reservation); + break; + } + Err(crate::stats::QuotaReserveError::LimitExceeded) + | Err(crate::stats::QuotaReserveError::Contended) => { + this.stats.increment_quota_contention_total(); } } - - if reserved_read_bytes == 0 { - reserve_rounds = reserve_rounds.saturating_add(1); - if reserve_rounds >= QUOTA_RESERVE_MAX_ROUNDS { - this.stats.increment_quota_contention_timeout_total(); - if this.arm_quota_wait(cx).is_pending() { - return Poll::Pending; - } - reserve_rounds = 0; - } + } + if quota_reservation.is_none() { + this.stats.increment_quota_contention_timeout_total(); + if this.arm_quota_wait(cx).is_ready() { + cx.waker().wake_by_ref(); } + return Poll::Pending; } } @@ -287,9 +274,9 @@ impl AsyncRead for StatsIo { match read_result { Poll::Ready(Ok(n)) => { - if reserved_read_bytes > n as u64 { - let refund_bytes = reserved_read_bytes - n as u64; - refund_reserved_quota_bytes(this.user_stats.as_ref(), refund_bytes); + if let Some(reservation) = quota_reservation.take() { + let refund_bytes = reservation.reserved_bytes().saturating_sub(n as u64); + reservation.settle(n as u64); this.stats.add_quota_refund_bytes_total(refund_bytes); } if n > 0 { @@ -308,7 +295,7 @@ impl AsyncRead for StatsIo { } } if let Some(limit) = this.quota_limit - && this.user_stats.quota_used() >= limit + && this.quota_handle.used() >= limit { this.quota_exceeded.store(true, Ordering::Release); } @@ -333,16 +320,16 @@ impl AsyncRead for StatsIo { Poll::Ready(Ok(())) } Poll::Pending => { - if reserved_read_bytes > 0 { - refund_reserved_quota_bytes(this.user_stats.as_ref(), reserved_read_bytes); - this.stats.add_quota_refund_bytes_total(reserved_read_bytes); + if let Some(reservation) = quota_reservation.take() { + this.stats + .add_quota_refund_bytes_total(reservation.reserved_bytes()); } Poll::Pending } Poll::Ready(Err(err)) => { - if reserved_read_bytes > 0 { - refund_reserved_quota_bytes(this.user_stats.as_ref(), reserved_read_bytes); - this.stats.add_quota_refund_bytes_total(reserved_read_bytes); + if let Some(reservation) = quota_reservation.take() { + this.stats + .add_quota_refund_bytes_total(reservation.reserved_bytes()); } Poll::Ready(Err(err)) } @@ -361,14 +348,15 @@ impl AsyncWrite for StatsIo { return Poll::Ready(Err(quota_io_error())); } - let mut shaper_reserved_bytes = 0u64; + let mut shaper_reservation = None; let mut write_buf = buf; if let Some(lease) = this.traffic_lease.as_ref() { if !buf.is_empty() { loop { - let consume = lease.try_consume(RateDirection::Down, buf.len() as u64); + let reservation = lease.try_reserve(RateDirection::Down, buf.len() as u64); + let consume = reservation.result(); if consume.granted > 0 { - shaper_reserved_bytes = consume.granted; + shaper_reservation = Some(reservation); if consume.granted < buf.len() as u64 { write_buf = &buf[..consume.granted as usize]; } @@ -398,62 +386,44 @@ impl AsyncWrite for StatsIo { } let mut remaining_before = None; - let mut reserved_bytes = 0u64; + let mut quota_reservation = None; if let Some(limit) = this.quota_limit { if !write_buf.is_empty() { - let mut reserve_rounds = 0usize; - while reserved_bytes == 0 { - let used_before = this.user_stats.quota_used(); + for _ in 0..QUOTA_RESERVE_MAX_ATTEMPTS_PER_POLL { + let used_before = this.quota_handle.used(); let remaining = limit.saturating_sub(used_before); if remaining == 0 { - if let Some(lease) = this.traffic_lease.as_ref() { - lease.refund(RateDirection::Down, shaper_reserved_bytes); - } this.quota_exceeded.store(true, Ordering::Release); return Poll::Ready(Err(quota_io_error())); } remaining_before = Some(remaining); let desired = remaining.min(write_buf.len() as u64); - let mut saw_contention = false; - for _ in 0..QUOTA_RESERVE_SPIN_RETRIES { - match this.user_stats.quota_try_reserve(desired, limit) { - Ok(_) => { - reserved_bytes = desired; - write_buf = &write_buf[..desired as usize]; - break; - } - Err(crate::stats::QuotaReserveError::LimitExceeded) => { - break; - } - Err(crate::stats::QuotaReserveError::Contended) => { - this.stats.increment_quota_contention_total(); - saw_contention = true; - } + match this.quota_handle.try_reserve(desired, limit) { + Ok(reservation) => { + quota_reservation = Some(reservation); + write_buf = &write_buf[..desired as usize]; + break; } - } - - if reserved_bytes == 0 { - reserve_rounds = reserve_rounds.saturating_add(1); - if reserve_rounds >= QUOTA_RESERVE_MAX_ROUNDS { - this.stats.increment_quota_contention_timeout_total(); - if let Some(lease) = this.traffic_lease.as_ref() { - lease.refund(RateDirection::Down, shaper_reserved_bytes); - } - let _ = this.arm_quota_wait(cx); - return Poll::Pending; - } else if saw_contention { - std::hint::spin_loop(); + Err(crate::stats::QuotaReserveError::LimitExceeded) + | Err(crate::stats::QuotaReserveError::Contended) => { + this.stats.increment_quota_contention_total(); } } } + if quota_reservation.is_none() { + this.stats.increment_quota_contention_timeout_total(); + Self::arm_wait(&mut this.quota_wait, false, false); + if Self::poll_wait(&mut this.quota_wait, cx, None, RateDirection::Up).is_ready() + { + cx.waker().wake_by_ref(); + } + return Poll::Pending; + } } else { - let used_before = this.user_stats.quota_used(); + let used_before = this.quota_handle.used(); let remaining = limit.saturating_sub(used_before); if remaining == 0 { - if let Some(lease) = this.traffic_lease.as_ref() { - lease.refund(RateDirection::Down, shaper_reserved_bytes); - } this.quota_exceeded.store(true, Ordering::Release); return Poll::Ready(Err(quota_io_error())); } @@ -463,15 +433,13 @@ impl AsyncWrite for StatsIo { match Pin::new(&mut this.inner).poll_write(cx, write_buf) { Poll::Ready(Ok(n)) => { - if reserved_bytes > n as u64 { - let refund_bytes = reserved_bytes - n as u64; - refund_reserved_quota_bytes(this.user_stats.as_ref(), refund_bytes); + if let Some(reservation) = quota_reservation.take() { + let refund_bytes = reservation.reserved_bytes().saturating_sub(n as u64); + reservation.settle(n as u64); this.stats.add_quota_refund_bytes_total(refund_bytes); } - if shaper_reserved_bytes > n as u64 - && let Some(lease) = this.traffic_lease.as_ref() - { - lease.refund(RateDirection::Down, shaper_reserved_bytes - n as u64); + if let Some(reservation) = shaper_reservation.take() { + reservation.settle_written(n as u64); } if n > 0 { if let Some(lease) = this.traffic_lease.as_ref() { @@ -492,7 +460,7 @@ impl AsyncWrite for StatsIo { if let (Some(limit), Some(remaining)) = (this.quota_limit, remaining_before) { if should_immediate_quota_check(remaining, n_to_charge) { this.quota_bytes_since_check = 0; - if this.user_stats.quota_used() >= limit { + if this.quota_handle.used() >= limit { this.quota_exceeded.store(true, Ordering::Release); } } else { @@ -501,7 +469,7 @@ impl AsyncWrite for StatsIo { let interval = quota_adaptive_interval_bytes(remaining); if this.quota_bytes_since_check >= interval { this.quota_bytes_since_check = 0; - if this.user_stats.quota_used() >= limit { + if this.quota_handle.used() >= limit { this.quota_exceeded.store(true, Ordering::Release); } } @@ -513,26 +481,16 @@ impl AsyncWrite for StatsIo { Poll::Ready(Ok(n)) } Poll::Ready(Err(err)) => { - if reserved_bytes > 0 { - refund_reserved_quota_bytes(this.user_stats.as_ref(), reserved_bytes); - this.stats.add_quota_refund_bytes_total(reserved_bytes); - } - if shaper_reserved_bytes > 0 - && let Some(lease) = this.traffic_lease.as_ref() - { - lease.refund(RateDirection::Down, shaper_reserved_bytes); + if let Some(reservation) = quota_reservation.take() { + this.stats + .add_quota_refund_bytes_total(reservation.reserved_bytes()); } Poll::Ready(Err(err)) } Poll::Pending => { - if reserved_bytes > 0 { - refund_reserved_quota_bytes(this.user_stats.as_ref(), reserved_bytes); - this.stats.add_quota_refund_bytes_total(reserved_bytes); - } - if shaper_reserved_bytes > 0 - && let Some(lease) = this.traffic_lease.as_ref() - { - lease.refund(RateDirection::Down, shaper_reserved_bytes); + if let Some(reservation) = quota_reservation.take() { + this.stats + .add_quota_refund_bytes_total(reservation.reserved_bytes()); } Poll::Pending } diff --git a/src/proxy/relay/io/quota.rs b/src/proxy/relay/io/quota.rs index 82e83f5..2b9c8c1 100644 --- a/src/proxy/relay/io/quota.rs +++ b/src/proxy/relay/io/quota.rs @@ -1,4 +1,3 @@ -use crate::stats::UserStats; use std::io; #[derive(Debug)] @@ -28,8 +27,7 @@ const QUOTA_NEAR_LIMIT_BYTES: u64 = 64 * 1024; const QUOTA_LARGE_CHARGE_BYTES: u64 = 16 * 1024; const QUOTA_ADAPTIVE_INTERVAL_MIN_BYTES: u64 = 4 * 1024; const QUOTA_ADAPTIVE_INTERVAL_MAX_BYTES: u64 = 64 * 1024; -pub(super) const QUOTA_RESERVE_SPIN_RETRIES: usize = 64; -pub(super) const QUOTA_RESERVE_MAX_ROUNDS: usize = 8; +pub(super) const QUOTA_RESERVE_MAX_ATTEMPTS_PER_POLL: usize = 4; #[inline] pub(in crate::proxy::relay) fn quota_adaptive_interval_bytes(remaining_before: u64) -> u64 { @@ -46,10 +44,3 @@ pub(in crate::proxy::relay) fn should_immediate_quota_check( ) -> bool { remaining_before <= QUOTA_NEAR_LIMIT_BYTES || charge_bytes >= QUOTA_LARGE_CHARGE_BYTES } - -pub(super) fn refund_reserved_quota_bytes(user_stats: &UserStats, reserved_bytes: u64) { - if reserved_bytes == 0 { - return; - } - user_stats.refund_quota(reserved_bytes); -} diff --git a/src/proxy/route_mode.rs b/src/proxy/route_mode.rs index a3a5d6c..d5324fd 100644 --- a/src/proxy/route_mode.rs +++ b/src/proxy/route_mode.rs @@ -1,5 +1,4 @@ use std::sync::Arc; -use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use tokio::sync::watch; @@ -24,11 +23,11 @@ impl RelayRouteMode { pub(crate) struct RouteCutoverState { pub mode: RelayRouteMode, pub generation: u64, + pub direct_since_epoch_secs: Option, } #[derive(Clone)] pub(crate) struct RouteRuntimeController { - direct_since_epoch_secs: Arc, tx: watch::Sender, } @@ -37,17 +36,11 @@ impl RouteRuntimeController { let initial = RouteCutoverState { mode: initial_mode, generation: 0, + direct_since_epoch_secs: matches!(initial_mode, RelayRouteMode::Direct) + .then(now_epoch_secs), }; let (tx, _rx) = watch::channel(initial); - let direct_since_epoch_secs = if matches!(initial_mode, RelayRouteMode::Direct) { - now_epoch_secs() - } else { - 0 - }; - Self { - direct_since_epoch_secs: Arc::new(AtomicU64::new(direct_since_epoch_secs)), - tx, - } + Self { tx } } pub(crate) fn snapshot(&self) -> RouteCutoverState { @@ -58,25 +51,16 @@ impl RouteRuntimeController { self.tx.subscribe() } - pub(crate) fn direct_since_epoch_secs(&self) -> Option { - let value = self.direct_since_epoch_secs.load(Ordering::Relaxed); - (value > 0).then_some(value) - } - pub(crate) fn set_mode(&self, mode: RelayRouteMode) -> Option { let mut next = None; let changed = self.tx.send_if_modified(|state| { if state.mode == mode { return false; } - if matches!(mode, RelayRouteMode::Direct) { - self.direct_since_epoch_secs - .store(now_epoch_secs(), Ordering::Relaxed); - } else { - self.direct_since_epoch_secs.store(0, Ordering::Relaxed); - } state.mode = mode; state.generation = state.generation.saturating_add(1); + state.direct_since_epoch_secs = + matches!(mode, RelayRouteMode::Direct).then(now_epoch_secs); next = Some(*state); true }); diff --git a/src/proxy/shared_state.rs b/src/proxy/shared_state.rs index 6a8774c..0bc500d 100644 --- a/src/proxy/shared_state.rs +++ b/src/proxy/shared_state.rs @@ -1,18 +1,22 @@ use std::collections::hash_map::RandomState; use std::collections::{HashMap, HashSet}; use std::net::{IpAddr, SocketAddr}; -use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::{Arc, Mutex}; use std::time::Instant; use dashmap::DashMap; use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc}; -use tokio_util::sync::CancellationToken; use crate::proxy::direct_buffer_budget::{DirectBufferBudget, fallback_direct_buffer_hard_limit}; use crate::proxy::handshake::{AuthProbeSaturationState, AuthProbeState}; use crate::proxy::middle_relay::{DesyncDedupRotationState, RelayIdleCandidateRegistry}; use crate::proxy::traffic_limiter::TrafficLimiter; +use crate::proxy::user_admission::{ + UserAdmissionAuthority, UserAdmissionPublication, UserCredentialId, UserIncarnation, + UserMutationResult, UserSessionRegistration, +}; +use crate::slot_budget::SlotBudget; const HANDSHAKE_RECENT_USER_RING_LEN: usize = 64; const MASKING_FALLBACK_MAX_CONCURRENT: usize = 512; @@ -52,14 +56,26 @@ pub(crate) enum ConntrackClosePolicy { pub(crate) struct HandshakeSharedState { pub(crate) auth_probe: DashMap, + /// Exact capacity authority for the authentication probe registry. + pub(crate) auth_probe_slots: SlotBudget, pub(crate) auth_probe_saturation: Mutex>, pub(crate) auth_probe_eviction_hasher: RandomState, pub(crate) invalid_secret_warned: Mutex>, pub(crate) unknown_sni_warn_next_allowed: Mutex>, - pub(crate) sticky_user_by_ip: DashMap, - pub(crate) sticky_user_by_ip_prefix: DashMap, - pub(crate) sticky_user_by_sni_hash: DashMap, - pub(crate) recent_user_ring: Box<[AtomicU32]>, + /// Stable credential hints keyed by exact peer IP. + pub(crate) sticky_user_by_ip: DashMap, + /// Exact capacity authority for peer-IP credential hints. + pub(crate) sticky_user_by_ip_slots: SlotBudget, + /// Stable credential hints keyed by bounded peer network prefix. + pub(crate) sticky_user_by_ip_prefix: DashMap, + /// Exact capacity authority for peer-prefix credential hints. + pub(crate) sticky_user_by_ip_prefix_slots: SlotBudget, + /// Stable credential hints keyed by normalized SNI hash. + pub(crate) sticky_user_by_sni_hash: DashMap, + /// Exact capacity authority for SNI credential hints. + pub(crate) sticky_user_by_sni_hash_slots: SlotBudget, + /// Bounded recent credential-hint ring used as an authentication fallback. + pub(crate) recent_user_ring: Box<[AtomicU64]>, pub(crate) recent_user_ring_seq: AtomicU64, pub(crate) auth_expensive_checks_total: AtomicU64, pub(crate) auth_budget_exhausted_total: AtomicU64, @@ -80,36 +96,12 @@ pub(crate) struct ProxySharedState { pub(crate) middle_relay: MiddleRelaySharedState, pub(crate) traffic_limiter: Arc, pub(crate) direct_buffer_budget: Arc, - disabled_users: DashMap, - active_user_sessions: DashMap<(String, u64), CancellationToken>, + user_admission: Arc, pub(crate) conntrack_pressure_active: AtomicBool, pub(crate) conntrack_close_tx: Mutex>>, masking_fallback_permits: Arc, } -#[must_use = "registered user sessions must be kept alive until relay completion"] -pub(crate) struct UserSessionRegistration { - token: CancellationToken, - _guard: UserSessionGuard, -} - -impl UserSessionRegistration { - pub(crate) fn token(&self) -> CancellationToken { - self.token.clone() - } -} - -struct UserSessionGuard { - shared: Arc, - key: (String, u64), -} - -impl Drop for UserSessionGuard { - fn drop(&mut self) { - self.shared.active_user_sessions.remove(&self.key); - } -} - impl ProxySharedState { pub(crate) fn new() -> Arc { Self::new_with_direct_buffer_budget(DirectBufferBudget::new( @@ -120,18 +112,54 @@ impl ProxySharedState { /// Creates process state with the startup-resolved Direct buffer envelope. pub(crate) fn new_with_direct_buffer_budget( direct_buffer_budget: Arc, + ) -> Arc { + Self::new_with_direct_buffer_budget_and_user_admission( + direct_buffer_budget, + UserAdmissionAuthority::new(), + ) + } + + /// Creates generation state around one process-owned user authority. + pub(crate) fn new_with_direct_buffer_budget_and_user_admission( + direct_buffer_budget: Arc, + user_admission: Arc, + ) -> Arc { + Self::new_with_process_authorities( + direct_buffer_budget, + TrafficLimiter::new(), + user_admission, + ) + } + + /// Creates generation-local caches around process-owned data-plane authorities. + pub(crate) fn new_with_process_authorities( + direct_buffer_budget: Arc, + traffic_limiter: Arc, + user_admission: Arc, ) -> Arc { Arc::new(Self { handshake: HandshakeSharedState { auth_probe: DashMap::new(), + auth_probe_slots: SlotBudget::new( + crate::proxy::handshake::AUTH_PROBE_TRACK_MAX_ENTRIES, + ), auth_probe_saturation: Mutex::new(None), auth_probe_eviction_hasher: RandomState::new(), invalid_secret_warned: Mutex::new(HashSet::new()), unknown_sni_warn_next_allowed: Mutex::new(None), sticky_user_by_ip: DashMap::new(), + sticky_user_by_ip_slots: SlotBudget::new( + crate::proxy::handshake::STICKY_HINT_MAX_ENTRIES, + ), sticky_user_by_ip_prefix: DashMap::new(), + sticky_user_by_ip_prefix_slots: SlotBudget::new( + crate::proxy::handshake::STICKY_HINT_MAX_ENTRIES, + ), sticky_user_by_sni_hash: DashMap::new(), - recent_user_ring: std::iter::repeat_with(|| AtomicU32::new(0)) + sticky_user_by_sni_hash_slots: SlotBudget::new( + crate::proxy::handshake::STICKY_HINT_MAX_ENTRIES, + ), + recent_user_ring: std::iter::repeat_with(|| AtomicU64::new(0)) .take(HANDSHAKE_RECENT_USER_RING_LEN) .collect::>() .into_boxed_slice(), @@ -148,10 +176,9 @@ impl ProxySharedState { relay_idle_registry: RelayIdleCandidateRegistry::default(), relay_idle_mark_seq: AtomicU64::new(0), }, - traffic_limiter: TrafficLimiter::new(), + traffic_limiter, direct_buffer_budget, - disabled_users: DashMap::new(), - active_user_sessions: DashMap::new(), + user_admission, conntrack_pressure_active: AtomicBool::new(false), conntrack_close_tx: Mutex::new(None), masking_fallback_permits: Arc::new(Semaphore::new(MASKING_FALLBACK_MAX_CONCURRENT)), @@ -167,72 +194,117 @@ impl ProxySharedState { } pub(crate) fn is_user_enabled(&self, user: &str) -> bool { - !self.disabled_users.contains_key(user) + self.user_admission.is_user_enabled(user) } - pub(crate) fn set_user_enabled(&self, user: &str, enabled: bool) -> bool { - if enabled { - self.disabled_users.remove(user); - false - } else { - self.disabled_users.insert(user.to_string(), ()).is_none() - } + /// Returns the process authority shared by every runtime generation. + pub(crate) fn user_admission(&self) -> Arc { + Arc::clone(&self.user_admission) } - pub(crate) fn apply_user_enabled_config( + /// Reconciles the complete user authentication policy from configuration. + pub(crate) fn apply_user_config( &self, + users: &HashMap, user_enabled: &HashMap, - ) -> Vec { - let desired_disabled = user_enabled - .iter() - .filter_map(|(user, enabled)| (!*enabled).then_some(user.clone())) - .collect::>(); - let current_disabled = self - .disabled_users - .iter() - .map(|entry| entry.key().clone()) - .collect::>(); + ) -> Vec<(String, usize)> { + self.user_admission.apply_config(users, user_enabled) + } - for user in current_disabled.difference(&desired_disabled) { - self.disabled_users.remove(user); - } - let newly_disabled = desired_disabled - .difference(¤t_disabled) - .cloned() - .collect::>(); - for user in desired_disabled { - self.disabled_users.insert(user, ()); - } - newly_disabled + /// Transfers user-policy ownership to one runtime generation. + pub(crate) fn activate_user_config_source( + &self, + source_generation: u64, + expected_epoch: Option, + users: &HashMap, + user_enabled: &HashMap, + ) -> Option> { + self.user_admission.activate_config_source( + source_generation, + expected_epoch, + users, + user_enabled, + ) + } + + /// Applies an update only from the active runtime generation. + pub(crate) fn apply_user_config_from_source( + &self, + source_generation: u64, + users: &HashMap, + user_enabled: &HashMap, + ) -> Option> { + self.user_admission + .apply_config_from_source(source_generation, users, user_enabled) + } + + /// Applies one persisted user mutation before asynchronous config reload. + pub(crate) fn stage_user( + &self, + user: &str, + secret: &str, + enabled: bool, + ) -> Option { + self.user_admission.stage_user(user, secret, enabled) + } + + /// Applies one prevalidated persisted credential before asynchronous reload. + pub(crate) fn stage_user_credential( + &self, + user: &str, + credential_id: UserCredentialId, + enabled: bool, + ) -> UserMutationResult { + self.user_admission + .stage_user_credential(user, credential_id, enabled) + } + + /// Installs a deletion tombstone and cancels every current owner. + pub(crate) fn delete_user(&self, user: &str) -> UserMutationResult { + self.user_admission.delete_user(user) + } + + /// Returns the current incarnation for an exact authenticated credential. + pub(crate) fn authenticated_user_incarnation( + &self, + user: &str, + credential_id: UserCredentialId, + ) -> Option { + self.user_admission + .authenticated_incarnation(user, credential_id) + } + + /// Starts an atomic publication boundary for an authenticated owner. + pub(crate) fn claim_authenticated_user( + self: &Arc, + user: &str, + credential_id: UserCredentialId, + ) -> Option> { + self.user_admission.claim_authenticated(user, credential_id) } pub(crate) fn register_user_session( self: &Arc, user: &str, - session_id: u64, - ) -> UserSessionRegistration { - let token = CancellationToken::new(); - let key = (user.to_string(), session_id); - self.active_user_sessions.insert(key.clone(), token.clone()); - UserSessionRegistration { - token, - _guard: UserSessionGuard { - shared: Arc::clone(self), - key, - }, - } + _session_id: u64, + ) -> Option { + self.user_admission.register_legacy(user) + } + + /// Registers a relay session against the exact credential that authenticated it. + pub(crate) fn register_authenticated_user_session( + self: &Arc, + user: &str, + credential_id: UserCredentialId, + ) -> Option { + let mut publication = self.claim_authenticated_user(user, credential_id)?; + let registration = publication.take_registration()?; + publication.commit(); + Some(registration) } pub(crate) fn cancel_user_sessions(&self, user: &str) -> usize { - let tokens = self - .active_user_sessions - .iter() - .filter_map(|entry| (entry.key().0 == user).then(|| entry.value().clone())) - .collect::>(); - for token in &tokens { - token.cancel(); - } - tokens.len() + self.user_admission.cancel_user_owners(user) } pub(crate) fn set_conntrack_close_sender(&self, tx: mpsc::Sender) { @@ -300,34 +372,56 @@ impl ProxySharedState { mod tests { use super::*; + const ALICE_SECRET: &str = "00112233445566778899aabbccddeeff"; + + fn configured_shared() -> Arc { + let shared = ProxySharedState::new(); + let users = HashMap::from([ + ("alice".to_string(), ALICE_SECRET.to_string()), + ( + "bob".to_string(), + "ffeeddccbbaa99887766554433221100".to_string(), + ), + ]); + shared.apply_user_config(&users, &HashMap::new()); + shared + } + #[test] fn user_enabled_config_sync_tracks_disabled_overrides() { - let shared = ProxySharedState::new(); + let shared = configured_shared(); assert!(shared.is_user_enabled("alice")); + let users = HashMap::from([ + ("alice".to_string(), ALICE_SECRET.to_string()), + ( + "bob".to_string(), + "ffeeddccbbaa99887766554433221100".to_string(), + ), + ]); let mut user_enabled = HashMap::new(); user_enabled.insert("alice".to_string(), false); user_enabled.insert("bob".to_string(), true); - let mut newly_disabled = shared.apply_user_enabled_config(&user_enabled); + let mut newly_disabled = shared.apply_user_config(&users, &user_enabled); newly_disabled.sort(); - assert_eq!(newly_disabled, vec!["alice".to_string()]); + assert_eq!(newly_disabled, vec![("alice".to_string(), 0)]); assert!(!shared.is_user_enabled("alice")); assert!(shared.is_user_enabled("bob")); - assert!(shared.apply_user_enabled_config(&user_enabled).is_empty()); + assert!(shared.apply_user_config(&users, &user_enabled).is_empty()); user_enabled.clear(); - assert!(shared.apply_user_enabled_config(&user_enabled).is_empty()); + assert!(shared.apply_user_config(&users, &user_enabled).is_empty()); assert!(shared.is_user_enabled("alice")); } #[test] fn cancel_user_sessions_cancels_only_registered_matching_user() { - let shared = ProxySharedState::new(); - let alice_1 = shared.register_user_session("alice", 1); - let alice_2 = shared.register_user_session("alice", 2); - let bob = shared.register_user_session("bob", 1); + let shared = configured_shared(); + let alice_1 = shared.register_user_session("alice", 1).unwrap(); + let alice_2 = shared.register_user_session("alice", 2).unwrap(); + let bob = shared.register_user_session("bob", 1).unwrap(); let alice_1_token = alice_1.token(); let alice_2_token = alice_2.token(); let bob_token = bob.token(); @@ -339,4 +433,68 @@ mod tests { assert!(alice_2_token.is_cancelled()); assert!(!bob_token.is_cancelled()); } + + #[test] + fn disabled_user_cannot_register_after_the_cancellation_snapshot() { + let shared = configured_shared(); + + let result = shared.stage_user("alice", ALICE_SECRET, false).unwrap(); + assert!(result.newly_disabled); + assert_eq!(result.cancelled, 0); + assert_eq!(shared.cancel_user_sessions("alice"), 0); + + let late = shared.register_user_session("alice", 1); + assert!( + late.is_none(), + "a session registered after disable returned must be rejected" + ); + } + + #[test] + fn disabling_user_cancels_existing_sessions_before_return() { + let shared = configured_shared(); + let registration = shared.register_user_session("alice", 1).unwrap(); + let token = registration.token(); + + let result = shared.stage_user("alice", ALICE_SECRET, false).unwrap(); + assert!(result.newly_disabled); + assert_eq!(result.cancelled, 1); + assert!(token.is_cancelled()); + assert!(shared.register_user_session("alice", 2).is_none()); + + let result = shared.stage_user("alice", ALICE_SECRET, true).unwrap(); + assert!(!result.newly_disabled); + assert_eq!(result.cancelled, 0); + assert!(shared.register_user_session("alice", 3).is_some()); + } + + #[test] + fn concurrent_disable_and_registration_never_leave_a_live_session() { + const ITERATIONS: usize = 10_000; + + let shared = ProxySharedState::new(); + let barrier = Arc::new(std::sync::Barrier::new(2)); + let register_shared = Arc::clone(&shared); + let register_barrier = Arc::clone(&barrier); + let register = std::thread::spawn(move || { + let mut registrations = Vec::with_capacity(ITERATIONS); + for session_id in 0..ITERATIONS as u64 { + let user = format!("user-{session_id}"); + register_barrier.wait(); + registrations.push(register_shared.register_user_session(&user, session_id)); + } + registrations + }); + + for session_id in 0..ITERATIONS as u64 { + let user = format!("user-{session_id}"); + barrier.wait(); + let secret = "00112233445566778899aabbccddeeff"; + shared.stage_user(&user, secret, false).unwrap(); + } + + for registration in register.join().unwrap().into_iter().flatten() { + assert!(registration.token().is_cancelled()); + } + } } diff --git a/src/proxy/tests/client_security_tests.rs b/src/proxy/tests/client_security_tests.rs index 338129d..5b132e4 100644 --- a/src/proxy/tests/client_security_tests.rs +++ b/src/proxy/tests/client_security_tests.rs @@ -276,14 +276,13 @@ async fn user_connection_reservation_drop_enqueues_cleanup_synchronously() { ip_tracker.set_user_limit(&user, 1).await; ip_tracker.check_and_add(&user, ip).await.unwrap(); - stats.increment_user_curr_connects(&user); - assert_eq!(ip_tracker.get_active_ip_count(&user).await, 1); - assert_eq!(stats.get_user_curr_connects(&user), 1); let reservation = UserConnectionReservation::new(stats.clone(), ip_tracker.clone(), user.clone(), ip, true); + assert_eq!(stats.get_user_curr_connects(&user), 1); + // Drop the reservation synchronously without any tokio::spawn/await yielding! drop(reservation); @@ -304,6 +303,111 @@ async fn user_connection_reservation_drop_enqueues_cleanup_synchronously() { assert_eq!(ip_tracker.get_active_ip_count(&user).await, 0); } +#[tokio::test] +async fn cancelled_ip_admission_releases_process_connection_permit() { + let ip_tracker = Arc::new(UserIpTracker::new()); + let stats = Arc::new(Stats::new()); + let user = "cancelled-admission-user"; + let peer_addr: SocketAddr = "198.51.100.210:50000".parse().unwrap(); + let mut config = ProxyConfig::default(); + config.access.user_max_tcp_conns.insert(user.to_string(), 1); + + let (entered_tx, entered_rx) = tokio::sync::oneshot::channel(); + let (release_tx, release_rx) = tokio::sync::oneshot::channel(); + let held_tracker = Arc::clone(&ip_tracker); + let held_user = user.to_string(); + let holder = tokio::spawn(async move { + held_tracker + .hold_user_shard_for_tests(&held_user, entered_tx, release_rx) + .await; + }); + entered_rx.await.unwrap(); + + let acquire_stats = Arc::clone(&stats); + let acquire_tracker = Arc::clone(&ip_tracker); + let acquire_config = config.clone(); + let acquire = tokio::spawn(async move { + acquire_user_connection_reservation( + user, + &acquire_config, + acquire_stats, + peer_addr, + acquire_tracker, + ) + .await + }); + + tokio::time::timeout(Duration::from_secs(1), async { + while stats.get_process_user_curr_connects(user) != 1 { + tokio::task::yield_now().await; + } + }) + .await + .expect("connection permit must be acquired before IP admission completes"); + + acquire.abort(); + let acquire_result = acquire.await; + assert!( + acquire_result + .as_ref() + .is_err_and(tokio::task::JoinError::is_cancelled) + ); + assert_eq!(stats.get_process_user_curr_connects(user), 0); + assert_eq!(stats.get_user_curr_connects(user), 0); + assert_eq!(ip_tracker.cleanup_queue_len_for_tests(), 0); + + let _ = release_tx.send(()); + holder.await.unwrap(); +} + +#[tokio::test] +async fn cancelled_async_release_preserves_ip_cleanup_ownership() { + let ip_tracker = Arc::new(UserIpTracker::new()); + let stats = Arc::new(Stats::new()); + let user = "cancelled-release-user"; + let peer_addr: SocketAddr = "198.51.100.211:50001".parse().unwrap(); + let mut config = ProxyConfig::default(); + config.access.user_max_tcp_conns.insert(user.to_string(), 1); + + let reservation = acquire_user_connection_reservation( + user, + &config, + Arc::clone(&stats), + peer_addr, + Arc::clone(&ip_tracker), + ) + .await + .unwrap(); + assert_eq!(stats.get_process_user_curr_connects(user), 1); + assert_eq!(ip_tracker.get_active_ip_count(user).await, 1); + + let (entered_tx, entered_rx) = tokio::sync::oneshot::channel(); + let (release_tx, release_rx) = tokio::sync::oneshot::channel(); + let held_tracker = Arc::clone(&ip_tracker); + let held_user = user.to_string(); + let holder = tokio::spawn(async move { + held_tracker + .hold_user_shard_for_tests(&held_user, entered_tx, release_rx) + .await; + }); + entered_rx.await.unwrap(); + + let release = tokio::spawn(reservation.release()); + tokio::task::yield_now().await; + release.abort(); + assert!(release.await.unwrap_err().is_cancelled()); + + assert_eq!(stats.get_process_user_curr_connects(user), 0); + assert_eq!(stats.get_user_curr_connects(user), 0); + assert_eq!(ip_tracker.cleanup_queue_len_for_tests(), 1); + + let _ = release_tx.send(()); + holder.await.unwrap(); + ip_tracker.drain_cleanup_queue().await; + assert_eq!(ip_tracker.get_active_ip_count(user).await, 0); + assert_eq!(ip_tracker.cleanup_queue_len_for_tests(), 0); +} + #[tokio::test] async fn relay_task_abort_releases_user_gate_and_ip_reservation() { let tg_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); @@ -2813,7 +2917,11 @@ async fn tcp_limit_rejection_does_not_reserve_ip_or_trigger_rollback() { .insert("user".to_string(), 1); let stats = Stats::new(); - stats.increment_user_curr_connects("user"); + let _existing_connection = stats + .connection_authority() + .try_acquire("user", Some(1)) + .expect("existing connection must occupy the process admission slot"); + let _existing_observation = stats.observe_user_current_connection("user"); let ip_tracker = UserIpTracker::new(); let peer_addr: SocketAddr = "198.51.100.210:50000".parse().unwrap(); @@ -2853,7 +2961,11 @@ async fn zero_tcp_limit_uses_global_fallback_and_rejects_without_side_effects() config.access.user_max_tcp_conns_global_each = 1; let stats = Stats::new(); - stats.increment_user_curr_connects("user"); + let _existing_connection = stats + .connection_authority() + .try_acquire("user", Some(1)) + .expect("existing connection must occupy the process admission slot"); + let _existing_observation = stats.observe_user_current_connection("user"); let ip_tracker = UserIpTracker::new(); let peer_addr: SocketAddr = "198.51.100.211:50001".parse().unwrap(); @@ -2914,7 +3026,11 @@ async fn global_tcp_fallback_applies_when_per_user_limit_is_missing() { config.access.user_max_tcp_conns_global_each = 1; let stats = Stats::new(); - stats.increment_user_curr_connects("user"); + let _existing_connection = stats + .connection_authority() + .try_acquire("user", Some(1)) + .expect("existing connection must occupy the process admission slot"); + let _existing_observation = stats.observe_user_current_connection("user"); let ip_tracker = UserIpTracker::new(); let peer_addr: SocketAddr = "198.51.100.213:50003".parse().unwrap(); @@ -4024,7 +4140,11 @@ async fn concurrent_limit_rejections_from_mixed_ips_leave_no_ip_footprint() { let config = Arc::new(config); let stats = Arc::new(Stats::new()); - stats.increment_user_curr_connects("user"); + let _existing_connection = stats + .connection_authority() + .try_acquire("user", Some(1)) + .expect("existing connection must occupy the process admission slot"); + let _existing_observation = stats.observe_user_current_connection("user"); let ip_tracker = Arc::new(UserIpTracker::new()); let mut tasks = tokio::task::JoinSet::new(); diff --git a/src/proxy/tests/direct_buffer_budget_tests.rs b/src/proxy/tests/direct_buffer_budget_tests.rs index 796e7fc..d75f905 100644 --- a/src/proxy/tests/direct_buffer_budget_tests.rs +++ b/src/proxy/tests/direct_buffer_budget_tests.rs @@ -1,4 +1,5 @@ use super::*; +use tokio::sync::Semaphore; #[test] fn lease_drop_releases_the_complete_reservation() { @@ -36,3 +37,64 @@ fn growth_and_shrink_keep_accounting_balanced() { drop(lease); assert_eq!(budget.snapshot().reserved_bytes, 0); } + +#[test] +fn runtime_generations_share_one_absolute_reservation_envelope() { + let first_generation = DirectBufferBudget::new(16 * 1024); + let second_generation = Arc::clone(&first_generation); + let first = first_generation + .try_reserve(12 * 1024, true) + .expect("first generation reservation must fit"); + + assert!(second_generation.try_reserve(8 * 1024, true).is_none()); + assert_eq!(second_generation.snapshot().reserved_bytes, 12 * 1024); + + drop(first); + assert!(second_generation.try_reserve(8 * 1024, true).is_some()); +} + +#[test] +fn stale_runtime_cannot_reclaim_direct_controller_ownership() { + let budget = DirectBufferBudget::new(16 * 1024); + budget.activate_controller(2); + budget.activate_controller(1); + + assert_eq!( + budget.active_controller_generation.load(Ordering::Acquire), + 2 + ); +} + +#[test] +fn controller_handoff_waits_for_inflight_update_and_fences_old_generation() { + let budget = DirectBufferBudget::new(16 * 1024); + budget.activate_controller(1); + let update = budget.begin_controller_update(1).unwrap(); + let (activated_tx, activated_rx) = std::sync::mpsc::channel(); + let next_budget = Arc::clone(&budget); + let activation = std::thread::spawn(move || { + next_budget.activate_controller(2); + activated_tx.send(()).unwrap(); + }); + + assert!( + activated_rx + .recv_timeout(Duration::from_millis(50)) + .is_err() + ); + drop(update); + activated_rx.recv_timeout(Duration::from_secs(1)).unwrap(); + activation.join().unwrap(); + + assert!(budget.begin_controller_update(1).is_none()); + assert!(budget.begin_controller_update(2).is_some()); +} + +#[test] +fn connection_pressure_uses_process_wide_slot_ownership() { + let slots = Arc::new(Semaphore::new(10)); + let _old_generation = Arc::clone(&slots).try_acquire_many_owned(3).unwrap(); + let _new_generation = Arc::clone(&slots).try_acquire_many_owned(2).unwrap(); + + assert_eq!(connection_fill_pct(slots.as_ref(), 10), Some(50)); +} diff --git a/src/proxy/tests/direct_relay_security_tests/anchored.rs b/src/proxy/tests/direct_relay_security_tests/anchored.rs index 6f0b242..2a8926f 100644 --- a/src/proxy/tests/direct_relay_security_tests/anchored.rs +++ b/src/proxy/tests/direct_relay_security_tests/anchored.rs @@ -72,6 +72,49 @@ fn adversarial_parent_swap_after_check_is_blocked_by_anchored_open() { ); } +#[cfg(unix)] +#[test] +fn adversarial_intermediate_parent_swap_is_blocked_by_component_walk() { + use std::os::unix::fs::symlink; + + let directory = tempfile::tempdir().expect("temporary directory must be creatable"); + let parent = directory.path().join("parent"); + let moved = directory.path().join("moved"); + let outside = directory.path().join("outside"); + fs::create_dir_all(parent.join("nested")).expect("original nested directory must be creatable"); + fs::create_dir_all(outside.join("nested")).expect("outside nested directory must be creatable"); + + let candidate = parent.join("nested/unknown-dc.log"); + let sanitized = sanitize_unknown_dc_log_path( + candidate + .to_str() + .expect("temporary path must be valid UTF-8"), + ) + .expect("candidate must sanitize before intermediate parent swap"); + assert!( + unknown_dc_log_path_is_still_safe(&sanitized), + "precondition: target should initially pass revalidation" + ); + + fs::rename(&parent, &moved).expect("intermediate parent must be movable"); + symlink(&outside, &parent).expect("intermediate parent symlink must be creatable"); + + let err = open_unknown_dc_log_append_anchored(&sanitized) + .expect_err("anchored open must reject a swapped intermediate component"); + let raw = err.raw_os_error(); + assert!( + matches!( + raw, + Some(libc::ELOOP) | Some(libc::ENOTDIR) | Some(libc::ENOENT) + ), + "component walk must fail closed on intermediate swap, got raw_os_error={raw:?}" + ); + assert!( + !outside.join("nested/unknown-dc.log").exists(), + "component walk must not create a log through a swapped intermediate directory" + ); +} + #[cfg(unix)] #[test] fn anchored_open_nix_path_writes_expected_lines() { diff --git a/src/proxy/tests/handshake_security_tests.rs b/src/proxy/tests/handshake_security_tests.rs index 8a21e58..48a3f3d 100644 --- a/src/proxy/tests/handshake_security_tests.rs +++ b/src/proxy/tests/handshake_security_tests.rs @@ -1115,7 +1115,8 @@ async fn tls_unknown_sni_reject_handshake_policy_emits_unrecognized_name_alert() // Drain what the server wrote. We expect exactly one TLS alert record: // 0x15 0x03 0x03 0x00 0x02 0x02 0x70 // (ContentType.alert, TLS 1.2, length=2, fatal, unrecognized_name) - drop(result); // drops the server-side writer so peer_side sees EOF + // Drop the server-side writer so `peer_side` observes EOF. + drop(result); let mut buf = Vec::new(); peer_side.read_to_end(&mut buf).await.unwrap(); assert_eq!( @@ -1241,7 +1242,10 @@ async fn tls_runtime_snapshot_updates_sticky_and_recent_hints() { .sticky_user_by_ip .get(&peer.ip()) .map(|entry| *entry), - Some(0), + config + .runtime_user_auth() + .and_then(|snapshot| snapshot.entry_by_id(0)) + .map(|entry| entry.hint_key), "successful runtime-snapshot auth must seed sticky ip cache" ); assert_eq!( @@ -3047,7 +3051,8 @@ async fn valid_tls_is_blocked_by_per_ip_preauth_throttle_without_saturation() { let rng = SecureRandom::new(); let peer: SocketAddr = "198.51.100.103:45103".parse().unwrap(); - auth_probe_state_for_testing_in_shared(shared.as_ref()).insert( + insert_auth_probe_state_for_testing_in_shared( + shared.as_ref(), normalize_auth_probe_ip(peer.ip()), AuthProbeState { fail_streak: AUTH_PROBE_BACKOFF_START_FAILS, @@ -3085,7 +3090,8 @@ async fn saturation_allows_valid_tls_even_when_peer_ip_is_currently_throttled() let peer: SocketAddr = "198.51.100.104:45104".parse().unwrap(); let now = Instant::now(); - auth_probe_state_for_testing_in_shared(shared.as_ref()).insert( + insert_auth_probe_state_for_testing_in_shared( + shared.as_ref(), normalize_auth_probe_ip(peer.ip()), AuthProbeState { fail_streak: AUTH_PROBE_BACKOFF_START_FAILS, @@ -3181,7 +3187,8 @@ async fn saturation_grace_exhaustion_preauth_throttles_repeated_invalid_tls_prob let rng = SecureRandom::new(); let peer: SocketAddr = "198.51.100.205:45205".parse().unwrap(); let now = Instant::now(); - auth_probe_state_for_testing_in_shared(shared.as_ref()).insert( + insert_auth_probe_state_for_testing_in_shared( + shared.as_ref(), normalize_auth_probe_ip(peer.ip()), AuthProbeState { fail_streak: AUTH_PROBE_BACKOFF_START_FAILS + AUTH_PROBE_SATURATION_GRACE_FAILS, @@ -3235,7 +3242,8 @@ async fn saturation_allows_valid_mtproto_even_when_peer_ip_is_currently_throttle let peer: SocketAddr = "198.51.100.106:45106".parse().unwrap(); let now = Instant::now(); - auth_probe_state_for_testing_in_shared(shared.as_ref()).insert( + insert_auth_probe_state_for_testing_in_shared( + shared.as_ref(), normalize_auth_probe_ip(peer.ip()), AuthProbeState { fail_streak: AUTH_PROBE_BACKOFF_START_FAILS, @@ -3328,7 +3336,8 @@ async fn saturation_grace_exhaustion_preauth_throttles_repeated_invalid_mtproto_ let replay_checker = ReplayChecker::new(128, Duration::from_secs(60)); let peer: SocketAddr = "198.51.100.206:45206".parse().unwrap(); let now = Instant::now(); - auth_probe_state_for_testing_in_shared(shared.as_ref()).insert( + insert_auth_probe_state_for_testing_in_shared( + shared.as_ref(), normalize_auth_probe_ip(peer.ip()), AuthProbeState { fail_streak: AUTH_PROBE_BACKOFF_START_FAILS + AUTH_PROBE_SATURATION_GRACE_FAILS, @@ -3378,7 +3387,8 @@ async fn saturation_grace_progression_tls_reaches_cap_then_stops_incrementing() let rng = SecureRandom::new(); let peer: SocketAddr = "198.51.100.207:45207".parse().unwrap(); let now = Instant::now(); - auth_probe_state_for_testing_in_shared(shared.as_ref()).insert( + insert_auth_probe_state_for_testing_in_shared( + shared.as_ref(), normalize_auth_probe_ip(peer.ip()), AuthProbeState { fail_streak: AUTH_PROBE_BACKOFF_START_FAILS, @@ -3461,7 +3471,8 @@ async fn saturation_grace_progression_mtproto_reaches_cap_then_stops_incrementin let replay_checker = ReplayChecker::new(128, Duration::from_secs(60)); let peer: SocketAddr = "198.51.100.208:45208".parse().unwrap(); let now = Instant::now(); - auth_probe_state_for_testing_in_shared(shared.as_ref()).insert( + insert_auth_probe_state_for_testing_in_shared( + shared.as_ref(), normalize_auth_probe_ip(peer.ip()), AuthProbeState { fail_streak: AUTH_PROBE_BACKOFF_START_FAILS, @@ -3545,7 +3556,8 @@ async fn saturation_grace_boundary_still_admits_valid_tls_before_exhaustion() { let rng = SecureRandom::new(); let peer: SocketAddr = "198.51.100.209:45209".parse().unwrap(); let now = Instant::now(); - auth_probe_state_for_testing_in_shared(shared.as_ref()).insert( + insert_auth_probe_state_for_testing_in_shared( + shared.as_ref(), normalize_auth_probe_ip(peer.ip()), AuthProbeState { fail_streak: AUTH_PROBE_BACKOFF_START_FAILS + AUTH_PROBE_SATURATION_GRACE_FAILS - 1, @@ -3599,7 +3611,8 @@ async fn saturation_grace_exhaustion_blocks_valid_tls_until_backoff_expires() { let rng = SecureRandom::new(); let peer: SocketAddr = "198.51.100.210:45210".parse().unwrap(); let now = Instant::now(); - auth_probe_state_for_testing_in_shared(shared.as_ref()).insert( + insert_auth_probe_state_for_testing_in_shared( + shared.as_ref(), normalize_auth_probe_ip(peer.ip()), AuthProbeState { fail_streak: AUTH_PROBE_BACKOFF_START_FAILS + AUTH_PROBE_SATURATION_GRACE_FAILS, @@ -3667,7 +3680,8 @@ async fn saturation_grace_exhaustion_is_shared_across_tls_and_mtproto_for_same_p let rng = SecureRandom::new(); let peer: SocketAddr = "198.51.100.211:45211".parse().unwrap(); let now = Instant::now(); - auth_probe_state_for_testing_in_shared(shared.as_ref()).insert( + insert_auth_probe_state_for_testing_in_shared( + shared.as_ref(), normalize_auth_probe_ip(peer.ip()), AuthProbeState { fail_streak: AUTH_PROBE_BACKOFF_START_FAILS + AUTH_PROBE_SATURATION_GRACE_FAILS, @@ -3735,7 +3749,8 @@ async fn adversarial_same_peer_invalid_tls_storm_does_not_bypass_saturation_grac let rng = Arc::new(SecureRandom::new()); let peer: SocketAddr = "198.51.100.212:45212".parse().unwrap(); let now = Instant::now(); - auth_probe_state_for_testing_in_shared(shared.as_ref()).insert( + insert_auth_probe_state_for_testing_in_shared( + shared.as_ref(), normalize_auth_probe_ip(peer.ip()), AuthProbeState { fail_streak: AUTH_PROBE_BACKOFF_START_FAILS + AUTH_PROBE_SATURATION_GRACE_FAILS, @@ -3801,7 +3816,8 @@ async fn light_fuzz_saturation_grace_tls_invalid_inputs_never_authenticate_or_pa let rng = SecureRandom::new(); let peer: SocketAddr = "198.51.100.213:45213".parse().unwrap(); let now = Instant::now(); - auth_probe_state_for_testing_in_shared(shared.as_ref()).insert( + insert_auth_probe_state_for_testing_in_shared( + shared.as_ref(), normalize_auth_probe_ip(peer.ip()), AuthProbeState { fail_streak: AUTH_PROBE_BACKOFF_START_FAILS + AUTH_PROBE_SATURATION_GRACE_FAILS, @@ -3986,7 +4002,8 @@ async fn expired_saturation_keeps_per_ip_throttle_enforced_for_valid_tls() { let peer: SocketAddr = "198.51.100.110:45110".parse().unwrap(); let now = Instant::now(); - auth_probe_state_for_testing_in_shared(shared.as_ref()).insert( + insert_auth_probe_state_for_testing_in_shared( + shared.as_ref(), normalize_auth_probe_ip(peer.ip()), AuthProbeState { fail_streak: AUTH_PROBE_BACKOFF_START_FAILS, diff --git a/src/proxy/tests/middle_relay_atomic_quota_invariant_tests.rs b/src/proxy/tests/middle_relay_atomic_quota_invariant_tests.rs index 5cf6ac5..a05c13e 100644 --- a/src/proxy/tests/middle_relay_atomic_quota_invariant_tests.rs +++ b/src/proxy/tests/middle_relay_atomic_quota_invariant_tests.rs @@ -240,6 +240,7 @@ async fn me_writer_data_write_obeys_flow_cancellation() { user, None, None, + None, 0, None, &cancel, diff --git a/src/proxy/tests/route_mode_coherence_adversarial_tests.rs b/src/proxy/tests/route_mode_coherence_adversarial_tests.rs index b7f816e..b96d44b 100644 --- a/src/proxy/tests/route_mode_coherence_adversarial_tests.rs +++ b/src/proxy/tests/route_mode_coherence_adversarial_tests.rs @@ -9,7 +9,7 @@ fn positive_direct_cutover_sets_timestamp_and_snapshot_coherently() { let rx = runtime.subscribe(); assert!( - runtime.direct_since_epoch_secs().is_none(), + runtime.snapshot().direct_since_epoch_secs.is_none(), "middle startup must not expose direct-since timestamp" ); @@ -24,7 +24,7 @@ fn positive_direct_cutover_sets_timestamp_and_snapshot_coherently() { ); assert_eq!(observed.mode, RelayRouteMode::Direct); assert!( - runtime.direct_since_epoch_secs().is_some(), + observed.direct_since_epoch_secs.is_some(), "direct cutover must publish a non-empty direct-since timestamp" ); } @@ -34,12 +34,12 @@ fn negative_idempotent_set_mode_does_not_mutate_timestamp_or_generation() { let runtime = RouteRuntimeController::new(RelayRouteMode::Direct); let before_state = runtime.snapshot(); - let before_ts = runtime.direct_since_epoch_secs(); + let before_ts = before_state.direct_since_epoch_secs; let changed = runtime.set_mode(RelayRouteMode::Direct); let after_state = runtime.snapshot(); - let after_ts = runtime.direct_since_epoch_secs(); + let after_ts = after_state.direct_since_epoch_secs; assert!(changed.is_none(), "idempotent set_mode must return None"); assert_eq!( @@ -58,7 +58,7 @@ fn edge_middle_cutover_clears_timestamp() { let rx = runtime.subscribe(); assert!( - runtime.direct_since_epoch_secs().is_some(), + runtime.snapshot().direct_since_epoch_secs.is_some(), "direct startup must expose direct-since timestamp" ); @@ -73,7 +73,7 @@ fn edge_middle_cutover_clears_timestamp() { ); assert_eq!(observed.mode, RelayRouteMode::Middle); assert!( - runtime.direct_since_epoch_secs().is_none(), + observed.direct_since_epoch_secs.is_none(), "middle cutover must clear direct-since timestamp" ); } @@ -90,7 +90,7 @@ fn adversarial_blackhat_probe_sequence_observes_consistent_mode_timestamp_pairs( let observed_direct = *rx.borrow(); assert_eq!(observed_direct, emitted_direct); assert!( - runtime.direct_since_epoch_secs().is_some(), + observed_direct.direct_since_epoch_secs.is_some(), "direct observation must never expose empty timestamp" ); @@ -100,7 +100,7 @@ fn adversarial_blackhat_probe_sequence_observes_consistent_mode_timestamp_pairs( let observed_middle = *rx.borrow(); assert_eq!(observed_middle, emitted_middle); assert!( - runtime.direct_since_epoch_secs().is_none(), + observed_middle.direct_since_epoch_secs.is_none(), "middle observation must never expose direct timestamp" ); } @@ -136,9 +136,9 @@ fn integration_subscriber_and_runtime_gates_stay_coherent_across_cutovers() { assert_eq!(snapshot, emitted); if matches!(mode, RelayRouteMode::Direct) { - assert!(runtime.direct_since_epoch_secs().is_some()); + assert!(snapshot.direct_since_epoch_secs.is_some()); } else { - assert!(runtime.direct_since_epoch_secs().is_none()); + assert!(snapshot.direct_since_epoch_secs.is_none()); } } } @@ -176,12 +176,12 @@ fn light_fuzz_random_mode_plan_preserves_timestamp_and_generation_invariants() { if matches!(snapshot.mode, RelayRouteMode::Direct) { assert!( - runtime.direct_since_epoch_secs().is_some(), + snapshot.direct_since_epoch_secs.is_some(), "direct fuzz state must expose timestamp" ); } else { assert!( - runtime.direct_since_epoch_secs().is_none(), + snapshot.direct_since_epoch_secs.is_none(), "middle fuzz state must clear timestamp" ); } @@ -199,12 +199,17 @@ fn stress_parallel_subscribers_never_observe_generation_regression() { let rx = runtime.subscribe(); let mut last = rx.borrow().generation; for _ in 0..10_000usize { - let current = rx.borrow().generation; + let current = *rx.borrow(); assert!( - current >= last, + current.generation >= last, "watch generation must be monotonic for every subscriber" ); - last = current; + assert_eq!( + matches!(current.mode, RelayRouteMode::Direct), + current.direct_since_epoch_secs.is_some(), + "one route snapshot must contain a coherent mode/timestamp pair" + ); + last = current.generation; std::thread::yield_now(); } })); @@ -227,8 +232,8 @@ fn stress_parallel_subscribers_never_observe_generation_regression() { let final_state = runtime.snapshot(); if matches!(final_state.mode, RelayRouteMode::Direct) { - assert!(runtime.direct_since_epoch_secs().is_some()); + assert!(final_state.direct_since_epoch_secs.is_some()); } else { - assert!(runtime.direct_since_epoch_secs().is_none()); + assert!(final_state.direct_since_epoch_secs.is_none()); } } diff --git a/src/proxy/tests/route_mode_security_tests.rs b/src/proxy/tests/route_mode_security_tests.rs index e5925fc..cc169ca 100644 --- a/src/proxy/tests/route_mode_security_tests.rs +++ b/src/proxy/tests/route_mode_security_tests.rs @@ -120,6 +120,7 @@ fn session_is_not_affected_when_mode_matches_even_if_generation_advanced() { let current = RouteCutoverState { mode: RelayRouteMode::Direct, generation: 2, + direct_since_epoch_secs: Some(1), }; let session_generation = 0; @@ -134,6 +135,7 @@ fn cutover_predicate_rejects_equal_generation_even_if_mode_differs() { let current = RouteCutoverState { mode: RelayRouteMode::Middle, generation: 77, + direct_since_epoch_secs: None, }; assert!( !is_session_affected_by_cutover(current, RelayRouteMode::Direct, 77), @@ -168,13 +170,15 @@ fn adversarial_route_oscillation_only_cuts_over_sessions_with_different_final_mo fn light_fuzz_cutover_predicate_matches_reference_oracle() { let mut rng = StdRng::seed_from_u64(0xC0DEC0DE5EED); for _ in 0..20_000 { + let mode = if rng.random::() { + RelayRouteMode::Direct + } else { + RelayRouteMode::Middle + }; let current = RouteCutoverState { - mode: if rng.random::() { - RelayRouteMode::Direct - } else { - RelayRouteMode::Middle - }, + mode, generation: rng.random_range(0u64..1_000_000), + direct_since_epoch_secs: matches!(mode, RelayRouteMode::Direct).then_some(1), }; let session_mode = if rng.random::() { RelayRouteMode::Direct diff --git a/src/proxy/traffic_limiter.rs b/src/proxy/traffic_limiter.rs index a7a6068..1b16d0e 100644 --- a/src/proxy/traffic_limiter.rs +++ b/src/proxy/traffic_limiter.rs @@ -6,6 +6,7 @@ use std::sync::atomic::{AtomicU64, Ordering}; use arc_swap::ArcSwap; use dashmap::DashMap; use ipnetwork::IpNetwork; +use parking_lot::Mutex as ParkingMutex; use crate::config::RateLimitBps; @@ -32,42 +33,91 @@ const REGISTRY_SHARDS: usize = 64; const FAIR_EPOCH_MS: u64 = 20; const MAX_BORROW_CHUNK_BYTES: u64 = 32 * 1024; const CLEANUP_INTERVAL_SECS: u64 = 60; +const RESERVE_CAS_ATTEMPT_LIMIT: usize = 16; +const REFUND_CAS_ATTEMPT_LIMIT: usize = 8; +const PACKED_USAGE_BITS: u32 = 28; +const PACKED_USAGE_MASK: u64 = (1u64 << PACKED_USAGE_BITS) - 1; +const PACKED_EPOCH_MAX: u64 = (1u64 << (u64::BITS - PACKED_USAGE_BITS)) - 1; +/// Traffic direction used by rate-limit accounting. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum RateDirection { + /// Traffic received from a client. Up, + /// Traffic sent to a client. Down, } +/// Result of an immediate traffic-budget request. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct TrafficConsumeResult { + /// Number of bytes granted by all applicable limits. pub granted: u64, + /// Whether the per-user limit prevented a grant. pub blocked_user: bool, + /// Whether the per-CIDR limit prevented a grant. pub blocked_cidr: bool, } +/// Process-wide traffic-limiter counters and gauges. #[derive(Debug, Clone, Copy)] pub struct TrafficLimiterMetricsSnapshot { + /// Per-user upload throttle events. pub user_throttle_up_total: u64, + /// Per-user download throttle events. pub user_throttle_down_total: u64, + /// Per-CIDR upload throttle events. pub cidr_throttle_up_total: u64, + /// Per-CIDR download throttle events. pub cidr_throttle_down_total: u64, + /// Per-user accumulated upload wait time in milliseconds. pub user_wait_up_ms_total: u64, + /// Per-user accumulated download wait time in milliseconds. pub user_wait_down_ms_total: u64, + /// Per-CIDR accumulated upload wait time in milliseconds. pub cidr_wait_up_ms_total: u64, + /// Per-CIDR accumulated download wait time in milliseconds. pub cidr_wait_down_ms_total: u64, + /// Per-user upload reservations that exhausted their CAS budget. + pub user_reserve_cas_retry_exhausted_up_total: u64, + /// Per-user download reservations that exhausted their CAS budget. + pub user_reserve_cas_retry_exhausted_down_total: u64, + /// Per-user upload refunds that exhausted their CAS budget. + pub user_refund_cas_retry_exhausted_up_total: u64, + /// Per-user download refunds that exhausted their CAS budget. + pub user_refund_cas_retry_exhausted_down_total: u64, + /// Per-CIDR upload reservations that exhausted their CAS budget. + pub cidr_reserve_cas_retry_exhausted_up_total: u64, + /// Per-CIDR download reservations that exhausted their CAS budget. + pub cidr_reserve_cas_retry_exhausted_down_total: u64, + /// Per-CIDR upload refunds that exhausted their CAS budget. + pub cidr_refund_cas_retry_exhausted_up_total: u64, + /// Per-CIDR download refunds that exhausted their CAS budget. + pub cidr_refund_cas_retry_exhausted_down_total: u64, + /// Active leases with a per-user rate limit. pub user_active_leases: u64, + /// Active leases with a per-CIDR rate limit. pub cidr_active_leases: u64, + /// Configured per-user rate-limit entries. pub user_policy_entries: u64, + /// Configured per-CIDR rate-limit entries. pub cidr_policy_entries: u64, } +#[derive(Default)] +struct CasContentionMetrics { + reserve_exhausted_total: AtomicU64, + refund_exhausted_total: AtomicU64, +} + #[derive(Default)] struct ScopeMetrics { throttle_up_total: AtomicU64, throttle_down_total: AtomicU64, wait_up_ms_total: AtomicU64, wait_down_ms_total: AtomicU64, + contention_up: Arc, + contention_down: Arc, active_leases: AtomicU64, policy_entries: AtomicU64, } @@ -76,32 +126,73 @@ struct ScopeMetrics { struct AtomicRatePair { up_bps: AtomicU64, down_bps: AtomicU64, + revision: ParkingMutex, } #[derive(Default)] struct DirectionBucket { - epoch: AtomicU64, - used: AtomicU64, + state: AtomicU64, + contention: Arc, + #[cfg(test)] + forced_reserve_failures: AtomicU64, + #[cfg(test)] + forced_refund_failures: AtomicU64, + #[cfg(test)] + reserve_cas_attempts: AtomicU64, + #[cfg(test)] + refund_cas_attempts: AtomicU64, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum BucketReserveError { + StaleEpoch, + Contended, + RefundContended, + ReserveAndRefundContended, + FairShareContended, +} + +impl BucketReserveError { + fn exhausted_reserve_budget(self) -> bool { + matches!(self, Self::Contended | Self::ReserveAndRefundContended) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum BucketRefundOutcome { + Complete, + Contended, +} + +struct ReserveCasBudget { + remaining: usize, +} + +struct BucketReservation { + granted: u64, + debit: Option, +} + +struct CidrReservation { + granted: u64, + aggregate_debit: Option, + user_debit: Option, } struct UserBucket { rates: AtomicRatePair, - up: DirectionBucket, - down: DirectionBucket, + up: Arc, + down: Arc, active_leases: AtomicU64, } -#[derive(Default)] struct CidrDirectionBucket { - epoch: AtomicU64, - used: AtomicU64, - active_users: AtomicU64, + used: Arc, + active_users: Arc, } -#[derive(Default)] struct CidrUserDirectionState { - epoch: AtomicU64, - used: AtomicU64, + used: Arc, } struct CidrUserShare { @@ -139,6 +230,8 @@ enum CidrPolicyMatch<'a> { #[derive(Default)] struct PolicySnapshot { + revision: u64, + source_generation: u64, user_limits: HashMap, cidr_rules_v4: Vec, cidr_rules_v6: Vec, @@ -152,19 +245,47 @@ struct ShardedRegistry { mask: usize, } -pub struct TrafficLease { +struct TrafficLeaseBinding { limiter: Arc, + revision: u64, user_bucket: Option>, cidr_bucket: Option>, cidr_user_key: Option, cidr_user_share: Option>, } +/// A live traffic-limiter binding for one authenticated client session. +pub struct TrafficLease { + limiter: Arc, + user: String, + client_ip: IpAddr, + binding: ArcSwap, + refresh: ParkingMutex<()>, +} + +/// Owns rate-limit policy, shared buckets, and limiter telemetry. pub struct TrafficLimiter { policy: ArcSwap, + policy_update: ParkingMutex<()>, user_buckets: ShardedRegistry, cidr_buckets: ShardedRegistry, user_scope: ScopeMetrics, cidr_scope: ScopeMetrics, last_cleanup_epoch_secs: AtomicU64, } + +struct DirectionDebit { + bucket: Arc, + epoch: u64, + refundable: u64, +} + +/// Refunds uncommitted shaping budget when an I/O attempt is cancelled. +#[must_use = "traffic reservations must be settled after the I/O attempt"] +pub(crate) struct TrafficReservation { + result: TrafficConsumeResult, + _binding: Arc, + user: Option, + cidr: Option, + cidr_user: Option, +} diff --git a/src/proxy/traffic_limiter/buckets.rs b/src/proxy/traffic_limiter/buckets.rs index f6fe592..3685ba8 100644 --- a/src/proxy/traffic_limiter/buckets.rs +++ b/src/proxy/traffic_limiter/buckets.rs @@ -1,262 +1,496 @@ use super::*; -impl ScopeMetrics { - pub(super) fn throttle(&self, direction: RateDirection) { - match direction { - RateDirection::Up => { - self.throttle_up_total.fetch_add(1, Ordering::Relaxed); - } - RateDirection::Down => { - self.throttle_down_total.fetch_add(1, Ordering::Relaxed); - } +impl ReserveCasBudget { + pub(super) fn new() -> Self { + Self { + remaining: RESERVE_CAS_ATTEMPT_LIMIT, } } - pub(super) fn wait_ms(&self, direction: RateDirection, wait_ms: u64) { - match direction { - RateDirection::Up => { - self.wait_up_ms_total.fetch_add(wait_ms, Ordering::Relaxed); - } - RateDirection::Down => { - self.wait_down_ms_total - .fetch_add(wait_ms, Ordering::Relaxed); - } + fn take(&mut self) -> bool { + if self.remaining == 0 { + return false; } + self.remaining -= 1; + true + } + + pub(super) fn is_exhausted(&self) -> bool { + self.remaining == 0 } } impl AtomicRatePair { - pub(super) fn set(&self, limits: RateLimitBps) { - self.up_bps.store(limits.up_bps, Ordering::Relaxed); - self.down_bps.store(limits.down_bps, Ordering::Relaxed); + pub(super) fn new(revision: u64, limits: RateLimitBps) -> Self { + let rates = Self::default(); + rates.set(revision, limits); + rates + } + + pub(super) fn set(&self, revision: u64, limits: RateLimitBps) { + let mut current_revision = self.revision.lock(); + if revision < *current_revision { + return; + } + self.up_bps.store(limits.up_bps, Ordering::Release); + self.down_bps.store(limits.down_bps, Ordering::Release); + *current_revision = revision; } pub(super) fn get(&self, direction: RateDirection) -> u64 { match direction { - RateDirection::Up => self.up_bps.load(Ordering::Relaxed), - RateDirection::Down => self.down_bps.load(Ordering::Relaxed), + RateDirection::Up => self.up_bps.load(Ordering::Acquire), + RateDirection::Down => self.down_bps.load(Ordering::Acquire), } } } impl DirectionBucket { - pub(super) fn sync_epoch(&self, epoch: u64) { - let current = self.epoch.load(Ordering::Relaxed); - if current == epoch { - return; - } - if current < epoch - && self - .epoch - .compare_exchange(current, epoch, Ordering::Relaxed, Ordering::Relaxed) - .is_ok() - { - self.used.store(0, Ordering::Relaxed); + fn new(contention: Arc) -> Self { + Self { + state: AtomicU64::new(0), + contention, + #[cfg(test)] + forced_reserve_failures: AtomicU64::new(0), + #[cfg(test)] + forced_refund_failures: AtomicU64::new(0), + #[cfg(test)] + reserve_cas_attempts: AtomicU64::new(0), + #[cfg(test)] + refund_cas_attempts: AtomicU64::new(0), } } - pub(super) fn try_consume(&self, cap_bps: u64, requested: u64) -> u64 { - if requested == 0 { - return 0; + #[inline(always)] + fn compare_exchange_reserve(&self, current: u64, next: u64) -> Result { + #[cfg(test)] + if self.should_force_reserve_failure() { + return Err(current); } - if cap_bps == 0 { - return requested; + self.state + .compare_exchange(current, next, Ordering::Relaxed, Ordering::Relaxed) + } + + #[inline(always)] + fn compare_exchange_refund(&self, current: u64, next: u64) -> Result { + #[cfg(test)] + if self.should_force_refund_failure() { + return Err(current); } + self.state + .compare_exchange(current, next, Ordering::Relaxed, Ordering::Relaxed) + } - let epoch = current_epoch(); - self.sync_epoch(epoch); - let cap_epoch = bytes_per_epoch(cap_bps); + fn unpack(state: u64) -> (u64, u64) { + (state >> PACKED_USAGE_BITS, state & PACKED_USAGE_MASK) + } + fn pack(epoch: u64, used: u64) -> Option { + if epoch > PACKED_EPOCH_MAX || used > PACKED_USAGE_MASK { + return None; + } + Some((epoch << PACKED_USAGE_BITS) | used) + } + + #[cfg(test)] + pub(super) fn used_at(&self, epoch: u64) -> Option { + if epoch > PACKED_EPOCH_MAX { + return None; + } + let (current_epoch, used) = Self::unpack(self.state.load(Ordering::Relaxed)); + (current_epoch == epoch).then_some(used) + } + + fn used_in_epoch(&self, epoch: u64) -> Result { + let (observed_epoch, used) = Self::unpack(self.state.load(Ordering::Relaxed)); + if observed_epoch != epoch { + return Err(BucketReserveError::StaleEpoch); + } + Ok(used) + } + + pub(super) fn try_reserve_at( + self: &Arc, + epoch: u64, + cap: u64, + requested: u64, + budget: &mut ReserveCasBudget, + ) -> Result, BucketReserveError> { + if requested == 0 || cap == 0 { + return Ok(None); + } + if epoch > PACKED_EPOCH_MAX { + return Err(BucketReserveError::StaleEpoch); + } + let cap = cap.min(PACKED_USAGE_MASK); + + let mut observed = self.state.load(Ordering::Relaxed); + // The transaction-owned budget bounds retries across every participating bucket. loop { - let used = self.used.load(Ordering::Relaxed); - if used >= cap_epoch { - return 0; + let (observed_epoch, observed_used) = Self::unpack(observed); + if observed_epoch > epoch { + return Err(BucketReserveError::StaleEpoch); } - let remaining = cap_epoch.saturating_sub(used); + let used = if observed_epoch == epoch { + observed_used + } else { + 0 + }; + if used >= cap { + return Ok(None); + } + let remaining = cap - used; let grant = requested.min(remaining); if grant == 0 { - return 0; + return Ok(None); } - let next = used.saturating_add(grant); - if self - .used - .compare_exchange_weak(used, next, Ordering::Relaxed, Ordering::Relaxed) - .is_ok() - { - return grant; + let Some(next) = Self::pack(epoch, used + grant) else { + return Err(BucketReserveError::StaleEpoch); + }; + if !budget.take() { + return Err(BucketReserveError::Contended); + } + match self.compare_exchange_reserve(observed, next) { + Ok(_) => { + return Ok(Some(DirectionDebit { + bucket: Arc::clone(self), + epoch, + refundable: grant, + })); + } + Err(actual) => observed = actual, } } } - pub(super) fn refund(&self, bytes: u64) { - if bytes == 0 { - return; + fn refund_at(&self, epoch: u64, bytes: u64) -> BucketRefundOutcome { + if bytes == 0 || epoch > PACKED_EPOCH_MAX { + return BucketRefundOutcome::Complete; } - decrement_atomic_saturating(&self.used, bytes); + let mut observed = self.state.load(Ordering::Relaxed); + for _ in 0..REFUND_CAS_ATTEMPT_LIMIT { + let (observed_epoch, used) = Self::unpack(observed); + if observed_epoch != epoch || used == 0 { + return BucketRefundOutcome::Complete; + } + let next = Self::pack(epoch, used.saturating_sub(bytes)).unwrap_or(observed); + match self.compare_exchange_refund(observed, next) { + Ok(_) => return BucketRefundOutcome::Complete, + Err(actual) => observed = actual, + } + } + let (observed_epoch, used) = Self::unpack(observed); + if observed_epoch != epoch || used == 0 { + return BucketRefundOutcome::Complete; + } + // Retain the charge after bounded contention so accounting cannot under-enforce. + self.contention + .refund_exhausted_total + .fetch_add(1, Ordering::Relaxed); + BucketRefundOutcome::Contended + } +} +impl DirectionDebit { + fn granted(&self) -> u64 { + self.refundable + } + + pub(super) fn shrink_to(&mut self, retained: u64) -> BucketRefundOutcome { + let retained = retained.min(self.refundable); + let outcome = self + .bucket + .refund_at(self.epoch, self.refundable - retained); + // Never retry the attempted delta in Drop; failed refunds remain charged. + self.refundable = retained; + outcome + } + + fn refund_all(&mut self) -> BucketRefundOutcome { + self.shrink_to(0) + } + + pub(super) fn settle(&mut self, committed: u64) { + let _ = self.shrink_to(committed); + self.refundable = 0; + } + + pub(super) fn commit_all(&mut self) -> u64 { + let committed = self.refundable; + self.refundable = 0; + committed + } +} + +impl Drop for DirectionDebit { + fn drop(&mut self) { + let _ = self.bucket.refund_at(self.epoch, self.refundable); } } impl UserBucket { - pub(super) fn new(limits: RateLimitBps) -> Self { - let rates = AtomicRatePair::default(); - rates.set(limits); + pub(super) fn new( + revision: u64, + limits: RateLimitBps, + up_contention: Arc, + down_contention: Arc, + ) -> Self { Self { - rates, - up: DirectionBucket::default(), - down: DirectionBucket::default(), + rates: AtomicRatePair::new(revision, limits), + up: Arc::new(DirectionBucket::new(up_contention)), + down: Arc::new(DirectionBucket::new(down_contention)), active_leases: AtomicU64::new(0), } } - pub(super) fn set_rates(&self, limits: RateLimitBps) { - self.rates.set(limits); + pub(super) fn set_rates(&self, revision: u64, limits: RateLimitBps) { + self.rates.set(revision, limits); } - pub(super) fn try_consume(&self, direction: RateDirection, requested: u64) -> u64 { + pub(super) fn try_reserve( + &self, + direction: RateDirection, + epoch: u64, + requested: u64, + budget: &mut ReserveCasBudget, + ) -> Result { let cap_bps = self.rates.get(direction); - match direction { - RateDirection::Up => self.up.try_consume(cap_bps, requested), - RateDirection::Down => self.down.try_consume(cap_bps, requested), - } - } - - pub(super) fn refund(&self, direction: RateDirection, bytes: u64) { - match direction { - RateDirection::Up => self.up.refund(bytes), - RateDirection::Down => self.down.refund(bytes), + if cap_bps == 0 { + return Ok(BucketReservation { + granted: requested, + debit: None, + }); } + let cap = bytes_per_epoch(cap_bps); + let debit = match direction { + RateDirection::Up => self.up.try_reserve_at(epoch, cap, requested, budget)?, + RateDirection::Down => self.down.try_reserve_at(epoch, cap, requested, budget)?, + }; + let granted = debit.as_ref().map(DirectionDebit::granted).unwrap_or(0); + Ok(BucketReservation { granted, debit }) } } impl CidrDirectionBucket { - pub(super) fn sync_epoch(&self, epoch: u64) { - let current = self.epoch.load(Ordering::Relaxed); - if current == epoch { - return; - } - if current < epoch - && self - .epoch - .compare_exchange(current, epoch, Ordering::Relaxed, Ordering::Relaxed) - .is_ok() - { - self.used.store(0, Ordering::Relaxed); - self.active_users.store(0, Ordering::Relaxed); + fn new(contention: Arc) -> Self { + Self { + used: Arc::new(DirectionBucket::new(Arc::clone(&contention))), + active_users: Arc::new(DirectionBucket::new(contention)), } } - pub(super) fn try_consume( + pub(super) fn try_reserve( &self, user_state: &CidrUserDirectionState, + epoch: u64, cap_epoch: u64, requested: u64, - ) -> u64 { + budget: &mut ReserveCasBudget, + ) -> Result { if requested == 0 || cap_epoch == 0 { - return 0; + return Ok(CidrReservation { + granted: 0, + aggregate_debit: None, + user_debit: None, + }); } - let epoch = current_epoch(); - self.sync_epoch(epoch); - user_state.sync_epoch_and_mark_active(epoch, &self.active_users); - let active_users = self.active_users.load(Ordering::Relaxed).max(1); + if !user_state.ensure_active(epoch, &self.active_users, budget)? { + return Ok(CidrReservation { + granted: 0, + aggregate_debit: None, + user_debit: None, + }); + } + let active_users = self.active_users.used_in_epoch(epoch)?; + let active_users = active_users.max(1); let fair_share = cap_epoch.saturating_div(active_users).max(1); - loop { - let total_used = self.used.load(Ordering::Relaxed); - if total_used >= cap_epoch { - return 0; - } - let total_remaining = cap_epoch.saturating_sub(total_used); - let user_used = user_state.used.load(Ordering::Relaxed); - let guaranteed_remaining = fair_share.saturating_sub(user_used); + let user_used = user_state.used.used_in_epoch(epoch)?; + let guaranteed_remaining = fair_share.saturating_sub(user_used); + let (user_cap, desired) = if guaranteed_remaining > 0 { + (fair_share, requested.min(guaranteed_remaining)) + } else { + (PACKED_USAGE_MASK, requested.min(MAX_BORROW_CHUNK_BYTES)) + }; + let mut user_debit = user_state + .used + .try_reserve_at(epoch, user_cap, desired, budget)?; - let grant = if guaranteed_remaining > 0 { - requested.min(guaranteed_remaining).min(total_remaining) + // A competing reservation can consume the guaranteed share between the snapshot and CAS. + if user_debit.is_none() && guaranteed_remaining > 0 { + let refreshed_used = user_state.used.used_in_epoch(epoch)?; + if refreshed_used >= fair_share { + user_debit = user_state.used.try_reserve_at( + epoch, + PACKED_USAGE_MASK, + requested.min(MAX_BORROW_CHUNK_BYTES), + budget, + )?; } else { - requested.min(total_remaining).min(MAX_BORROW_CHUNK_BYTES) - }; - - if grant == 0 { - return 0; + user_debit = user_state.used.try_reserve_at( + epoch, + fair_share, + requested.min(fair_share - refreshed_used), + budget, + )?; } - - let next_total = total_used.saturating_add(grant); - if self - .used - .compare_exchange_weak(total_used, next_total, Ordering::Relaxed, Ordering::Relaxed) - .is_ok() - { - user_state.used.fetch_add(grant, Ordering::Relaxed); - return grant; + if user_debit.is_none() { + return Err(BucketReserveError::FairShareContended); } } + + let Some(mut user_debit) = user_debit else { + return Ok(CidrReservation { + granted: 0, + aggregate_debit: None, + user_debit: None, + }); + }; + let user_granted = user_debit.granted(); + let Some(aggregate_debit) = + self.used + .try_reserve_at(epoch, cap_epoch, user_granted, budget)? + else { + return Ok(CidrReservation { + granted: 0, + aggregate_debit: None, + user_debit: None, + }); + }; + let granted = aggregate_debit.granted(); + if granted < user_granted { + let _ = user_debit.shrink_to(granted); + } + Ok(CidrReservation { + granted, + aggregate_debit: Some(aggregate_debit), + user_debit: Some(user_debit), + }) } +} - pub(super) fn refund(&self, bytes: u64) { - if bytes == 0 { - return; - } - decrement_atomic_saturating(&self.used, bytes); +#[cfg(test)] +impl Default for CidrDirectionBucket { + fn default() -> Self { + Self::new(Arc::new(CasContentionMetrics::default())) } } impl CidrUserDirectionState { - pub(super) fn sync_epoch_and_mark_active(&self, epoch: u64, active_users: &AtomicU64) { - let current = self.epoch.load(Ordering::Relaxed); - if current == epoch { - return; - } - if current < epoch - && self - .epoch - .compare_exchange(current, epoch, Ordering::Relaxed, Ordering::Relaxed) - .is_ok() - { - self.used.store(0, Ordering::Relaxed); - active_users.fetch_add(1, Ordering::Relaxed); + fn new(contention: Arc) -> Self { + Self { + used: Arc::new(DirectionBucket::new(contention)), } } - pub(super) fn refund(&self, bytes: u64) { - if bytes == 0 { - return; + pub(super) fn ensure_active( + &self, + epoch: u64, + active_users: &Arc, + budget: &mut ReserveCasBudget, + ) -> Result { + if epoch > PACKED_EPOCH_MAX { + return Err(BucketReserveError::StaleEpoch); } - decrement_atomic_saturating(&self.used, bytes); + let mut observed = self.used.state.load(Ordering::Relaxed); + // Every retry spends shared reserve budget before publishing the user epoch. + loop { + let (observed_epoch, _) = DirectionBucket::unpack(observed); + if observed_epoch == epoch { + return Ok(true); + } + if observed_epoch > epoch { + return Err(BucketReserveError::StaleEpoch); + } + let Some(mut active_debit) = + active_users.try_reserve_at(epoch, PACKED_USAGE_MASK, 1, budget)? + else { + return Ok(false); + }; + let Some(next) = DirectionBucket::pack(epoch, 0) else { + return Err(BucketReserveError::StaleEpoch); + }; + if !budget.take() { + if active_debit.refund_all() == BucketRefundOutcome::Contended { + return Err(BucketReserveError::ReserveAndRefundContended); + } + return Err(BucketReserveError::Contended); + } + match self.used.compare_exchange_reserve(observed, next) { + Ok(_) => { + active_debit.commit_all(); + return Ok(true); + } + Err(actual) => { + let refund = active_debit.refund_all(); + let (actual_epoch, _) = DirectionBucket::unpack(actual); + if refund == BucketRefundOutcome::Contended { + return Err(if budget.is_exhausted() { + BucketReserveError::ReserveAndRefundContended + } else { + BucketReserveError::RefundContended + }); + } + if actual_epoch == epoch { + return Ok(true); + } + observed = actual; + } + } + } + } +} + +#[cfg(test)] +impl Default for CidrUserDirectionState { + fn default() -> Self { + Self::new(Arc::new(CasContentionMetrics::default())) } } impl CidrUserShare { - pub(super) fn new() -> Self { + pub(super) fn new( + up_contention: Arc, + down_contention: Arc, + ) -> Self { Self { active_conns: AtomicU64::new(0), - up: CidrUserDirectionState::default(), - down: CidrUserDirectionState::default(), + up: CidrUserDirectionState::new(up_contention), + down: CidrUserDirectionState::new(down_contention), } } } impl CidrBucket { - pub(super) fn new(limits: RateLimitBps) -> Self { - let rates = AtomicRatePair::default(); - rates.set(limits); + pub(super) fn new( + revision: u64, + limits: RateLimitBps, + up_contention: Arc, + down_contention: Arc, + ) -> Self { Self { - rates, - up: CidrDirectionBucket::default(), - down: CidrDirectionBucket::default(), + rates: AtomicRatePair::new(revision, limits), + up: CidrDirectionBucket::new(up_contention), + down: CidrDirectionBucket::new(down_contention), users: ShardedRegistry::new(REGISTRY_SHARDS), active_leases: AtomicU64::new(0), } } - pub(super) fn set_rates(&self, limits: RateLimitBps) { - self.rates.set(limits); + pub(super) fn set_rates(&self, revision: u64, limits: RateLimitBps) { + self.rates.set(revision, limits); } pub(super) fn acquire_user_share(&self, user: &str) -> Arc { - self.users - .get_or_insert_with(user, CidrUserShare::new, |share| { + let up_contention = Arc::clone(&self.up.used.contention); + let down_contention = Arc::clone(&self.down.used.contention); + self.users.get_or_insert_with( + user, + || CidrUserShare::new(up_contention, down_contention), + |share| { share.active_conns.fetch_add(1, Ordering::Relaxed); - }) + }, + ) } pub(super) fn release_user_share(&self, user: &str, share: &Arc) { @@ -268,37 +502,30 @@ impl CidrBucket { }); } - pub(super) fn try_consume_for_user( + pub(super) fn try_reserve_for_user( &self, direction: RateDirection, share: &CidrUserShare, + epoch: u64, requested: u64, - ) -> u64 { + budget: &mut ReserveCasBudget, + ) -> Result { let cap_bps = self.rates.get(direction); if cap_bps == 0 { - return requested; + return Ok(CidrReservation { + granted: requested, + aggregate_debit: None, + user_debit: None, + }); } let cap_epoch = bytes_per_epoch(cap_bps); match direction { - RateDirection::Up => self.up.try_consume(&share.up, cap_epoch, requested), - RateDirection::Down => self.down.try_consume(&share.down, cap_epoch, requested), - } - } - - pub(super) fn refund_for_user( - &self, - direction: RateDirection, - share: &CidrUserShare, - bytes: u64, - ) { - match direction { - RateDirection::Up => { - self.up.refund(bytes); - share.up.refund(bytes); - } + RateDirection::Up => self + .up + .try_reserve(&share.up, epoch, cap_epoch, requested, budget), RateDirection::Down => { - self.down.refund(bytes); - share.down.refund(bytes); + self.down + .try_reserve(&share.down, epoch, cap_epoch, requested, budget) } } } diff --git a/src/proxy/traffic_limiter/helpers.rs b/src/proxy/traffic_limiter/helpers.rs index 272fe03..c869d85 100644 --- a/src/proxy/traffic_limiter/helpers.rs +++ b/src/proxy/traffic_limiter/helpers.rs @@ -2,6 +2,7 @@ use std::sync::OnceLock; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; use super::*; +/// Returns the delay until the next traffic-limiter refill boundary. pub fn next_refill_delay() -> Duration { let start = limiter_epoch_start(); let elapsed_ms = start.elapsed().as_millis() as u64; @@ -56,7 +57,7 @@ pub(super) fn auto_cidr_bucket_key(ip: IpAddr, prefix_len: u8) -> Option pub(super) fn current_epoch() -> u64 { let start = limiter_epoch_start(); let elapsed_ms = start.elapsed().as_millis() as u64; - elapsed_ms / FAIR_EPOCH_MS + elapsed_ms / FAIR_EPOCH_MS + 1 } pub(super) fn limiter_epoch_start() -> &'static Instant { diff --git a/src/proxy/traffic_limiter/lease.rs b/src/proxy/traffic_limiter/lease.rs index ad136af..98fc15f 100644 --- a/src/proxy/traffic_limiter/lease.rs +++ b/src/proxy/traffic_limiter/lease.rs @@ -1,72 +1,185 @@ use super::*; impl TrafficLease { - pub fn try_consume(&self, direction: RateDirection, requested: u64) -> TrafficConsumeResult { + fn current_binding(&self) -> Arc { + let policy = self.limiter.policy.load(); + let current = self.binding.load_full(); + if current.revision == policy.revision { + return current; + } + drop(policy); + + let _refresh = self.refresh.lock(); + let policy = self.limiter.policy.load_full(); + let current = self.binding.load_full(); + if current.revision == policy.revision { + return current; + } + let next = self + .limiter + .build_binding(&self.user, self.client_ip, &policy); + self.binding.store(Arc::clone(&next)); + next + } + + /// Reserves shaping budget until the associated I/O result is settled. + pub(crate) fn try_reserve( + &self, + direction: RateDirection, + requested: u64, + ) -> TrafficReservation { + let binding = self.current_binding(); if requested == 0 { - return TrafficConsumeResult { - granted: 0, - blocked_user: false, - blocked_cidr: false, + return TrafficReservation { + result: TrafficConsumeResult { + granted: 0, + blocked_user: false, + blocked_cidr: false, + }, + _binding: binding, + user: None, + cidr: None, + cidr_user: None, + }; + } + if binding.user_bucket.is_none() && binding.cidr_bucket.is_none() { + return TrafficReservation { + result: TrafficConsumeResult { + granted: requested, + blocked_user: false, + blocked_cidr: false, + }, + _binding: binding, + user: None, + cidr: None, + cidr_user: None, }; } + let epoch = current_epoch(); + let mut budget = ReserveCasBudget::new(); let mut granted = requested; - if let Some(user_bucket) = self.user_bucket.as_ref() { - let user_granted = user_bucket.try_consume(direction, granted); - if user_granted == 0 { + let mut user_debit = None; + if let Some(user_bucket) = binding.user_bucket.as_ref() { + let user_reservation = + match user_bucket.try_reserve(direction, epoch, granted, &mut budget) { + Ok(reservation) => reservation, + Err(error) => { + if error.exhausted_reserve_budget() { + self.limiter + .user_scope + .reserve_cas_retry_exhausted(direction); + } + return TrafficReservation { + result: TrafficConsumeResult { + granted: 0, + blocked_user: false, + blocked_cidr: false, + }, + _binding: binding, + user: None, + cidr: None, + cidr_user: None, + }; + } + }; + user_debit = user_reservation.debit; + if user_reservation.granted == 0 { self.limiter.observe_throttle(direction, true, false); - return TrafficConsumeResult { - granted: 0, - blocked_user: true, - blocked_cidr: false, + return TrafficReservation { + result: TrafficConsumeResult { + granted: 0, + blocked_user: true, + blocked_cidr: false, + }, + _binding: binding, + user: user_debit, + cidr: None, + cidr_user: None, }; } - granted = user_granted; + granted = user_reservation.granted; } - if let (Some(cidr_bucket), Some(cidr_user_share)) = - (self.cidr_bucket.as_ref(), self.cidr_user_share.as_ref()) - { - let cidr_granted = - cidr_bucket.try_consume_for_user(direction, cidr_user_share, granted); - if cidr_granted < granted - && let Some(user_bucket) = self.user_bucket.as_ref() + let mut cidr_debit = None; + let mut cidr_user_debit = None; + if let (Some(cidr_bucket), Some(cidr_user_share)) = ( + binding.cidr_bucket.as_ref(), + binding.cidr_user_share.as_ref(), + ) { + let cidr_reservation = match cidr_bucket.try_reserve_for_user( + direction, + cidr_user_share, + epoch, + granted, + &mut budget, + ) { + Ok(reservation) => reservation, + Err(error) => { + if error.exhausted_reserve_budget() { + self.limiter + .cidr_scope + .reserve_cas_retry_exhausted(direction); + } + return TrafficReservation { + result: TrafficConsumeResult { + granted: 0, + blocked_user: false, + blocked_cidr: false, + }, + _binding: binding, + user: user_debit, + cidr: None, + cidr_user: None, + }; + } + }; + cidr_debit = cidr_reservation.aggregate_debit; + cidr_user_debit = cidr_reservation.user_debit; + if cidr_reservation.granted < granted + && let Some(debit) = user_debit.as_mut() { - user_bucket.refund(direction, granted.saturating_sub(cidr_granted)); + let _ = debit.shrink_to(cidr_reservation.granted); } - if cidr_granted == 0 { + if cidr_reservation.granted == 0 { self.limiter.observe_throttle(direction, false, true); - return TrafficConsumeResult { - granted: 0, - blocked_user: false, - blocked_cidr: true, + return TrafficReservation { + result: TrafficConsumeResult { + granted: 0, + blocked_user: false, + blocked_cidr: true, + }, + _binding: binding, + user: user_debit, + cidr: cidr_debit, + cidr_user: cidr_user_debit, }; } - granted = cidr_granted; + granted = cidr_reservation.granted; } - TrafficConsumeResult { - granted, - blocked_user: false, - blocked_cidr: false, + TrafficReservation { + result: TrafficConsumeResult { + granted, + blocked_user: false, + blocked_cidr: false, + }, + _binding: binding, + user: user_debit, + cidr: cidr_debit, + cidr_user: cidr_user_debit, } } - pub fn refund(&self, direction: RateDirection, bytes: u64) { - if bytes == 0 { - return; - } - - if let Some(user_bucket) = self.user_bucket.as_ref() { - user_bucket.refund(direction, bytes); - } - if let (Some(cidr_bucket), Some(cidr_user_share)) = - (self.cidr_bucket.as_ref(), self.cidr_user_share.as_ref()) - { - cidr_bucket.refund_for_user(direction, cidr_user_share, bytes); - } + /// Immediately consumes the budget granted by all applicable limits. + pub fn try_consume(&self, direction: RateDirection, requested: u64) -> TrafficConsumeResult { + let reservation = self.try_reserve(direction, requested); + let result = reservation.result(); + reservation.settle_written(result.granted); + result } + /// Records caller-observed limiter wait time for the blocking scopes. pub fn observe_wait_ms( &self, direction: RateDirection, @@ -82,7 +195,28 @@ impl TrafficLease { } } -impl Drop for TrafficLease { +impl TrafficReservation { + /// Returns the shaping decision associated with this reservation. + pub(crate) fn result(&self) -> TrafficConsumeResult { + self.result + } + + /// Commits written bytes and refunds the uncommitted remainder. + pub(crate) fn settle_written(mut self, committed: u64) { + let committed = committed.min(self.result.granted); + if let Some(debit) = self.user.as_mut() { + debit.settle(committed); + } + if let Some(debit) = self.cidr.as_mut() { + debit.settle(committed); + } + if let Some(debit) = self.cidr_user.as_mut() { + debit.settle(committed); + } + } +} + +impl Drop for TrafficLeaseBinding { fn drop(&mut self) { if let Some(bucket) = self.user_bucket.as_ref() { decrement_atomic_saturating(&bucket.active_leases, 1); diff --git a/src/proxy/traffic_limiter/limiter.rs b/src/proxy/traffic_limiter/limiter.rs index 4aaf411..b75fdb5 100644 --- a/src/proxy/traffic_limiter/limiter.rs +++ b/src/proxy/traffic_limiter/limiter.rs @@ -1,10 +1,55 @@ use crate::config::CidrRateLimitKey; use super::*; + +impl ScopeMetrics { + pub(super) fn throttle(&self, direction: RateDirection) { + match direction { + RateDirection::Up => { + self.throttle_up_total.fetch_add(1, Ordering::Relaxed); + } + RateDirection::Down => { + self.throttle_down_total.fetch_add(1, Ordering::Relaxed); + } + } + } + + pub(super) fn wait_ms(&self, direction: RateDirection, wait_ms: u64) { + match direction { + RateDirection::Up => { + self.wait_up_ms_total.fetch_add(wait_ms, Ordering::Relaxed); + } + RateDirection::Down => { + self.wait_down_ms_total + .fetch_add(wait_ms, Ordering::Relaxed); + } + } + } + + pub(super) fn contention(&self, direction: RateDirection) -> Arc { + match direction { + RateDirection::Up => Arc::clone(&self.contention_up), + RateDirection::Down => Arc::clone(&self.contention_down), + } + } + + pub(super) fn reserve_cas_retry_exhausted(&self, direction: RateDirection) { + let metrics = match direction { + RateDirection::Up => self.contention_up.as_ref(), + RateDirection::Down => self.contention_down.as_ref(), + }; + metrics + .reserve_exhausted_total + .fetch_add(1, Ordering::Relaxed); + } +} + impl TrafficLimiter { + /// Creates an empty limiter with no active rate-limit policy. pub fn new() -> Arc { Arc::new(Self { policy: ArcSwap::from_pointee(PolicySnapshot::default()), + policy_update: ParkingMutex::new(()), user_buckets: ShardedRegistry::new(REGISTRY_SHARDS), cidr_buckets: ShardedRegistry::new(REGISTRY_SHARDS), user_scope: ScopeMetrics::default(), @@ -13,11 +58,32 @@ impl TrafficLimiter { }) } + #[cfg(test)] + /// Replaces the in-memory policy for isolated limiter tests. pub fn apply_policy( &self, user_limits: HashMap, cidr_limits: HashMap, ) { + let _ = self.apply_policy_inner(None, user_limits, cidr_limits); + } + + /// Publishes policy only when the source runtime is not older than the active source. + pub(crate) fn apply_policy_from_source( + &self, + source_generation: u64, + user_limits: HashMap, + cidr_limits: HashMap, + ) -> bool { + self.apply_policy_inner(Some(source_generation), user_limits, cidr_limits) + } + + fn apply_policy_inner( + &self, + source_generation: Option, + user_limits: HashMap, + cidr_limits: HashMap, + ) -> bool { let filtered_users = user_limits .into_iter() .filter(|(_, limit)| limit.up_bps > 0 || limit.down_bps > 0) @@ -70,6 +136,17 @@ impl TrafficLimiter { let cidr_policy_entries = cidr_rule_keys.len() + cidr_auto_rules_v4.len() + cidr_auto_rules_v6.len(); + let policy_update = self.policy_update.lock(); + let current = self.policy.load_full(); + if source_generation.is_some_and(|source| source < current.source_generation) { + return false; + } + // Revision wrap could otherwise let an old lease restore stale rates. + let Some(revision) = current.revision.checked_add(1) else { + return false; + }; + let source_generation = source_generation.unwrap_or(current.source_generation); + self.user_scope .policy_entries .store(filtered_users.len() as u64, Ordering::Relaxed); @@ -78,6 +155,8 @@ impl TrafficLimiter { .store(cidr_policy_entries as u64, Ordering::Relaxed); self.policy.store(Arc::new(PolicySnapshot { + revision, + source_generation, user_limits: filtered_users, cidr_rules_v4, cidr_rules_v6, @@ -86,25 +165,51 @@ impl TrafficLimiter { cidr_rule_keys, })); + drop(policy_update); self.maybe_cleanup(); + true } + /// Creates a lease that follows policy revisions for one client identity. pub fn acquire_lease( self: &Arc, user: &str, client_ip: IpAddr, ) -> Option> { let policy = self.policy.load_full(); + let binding = self.build_binding(user, client_ip, &policy); + Some(Arc::new(TrafficLease { + limiter: Arc::clone(self), + user: user.to_string(), + client_ip, + binding: ArcSwap::from(binding), + refresh: ParkingMutex::new(()), + })) + } + + pub(super) fn build_binding( + self: &Arc, + user: &str, + client_ip: IpAddr, + policy: &PolicySnapshot, + ) -> Arc { let mut user_bucket = None; if let Some(limit) = policy.user_limits.get(user).copied() { let bucket = self.user_buckets.get_or_insert_with( user, - || UserBucket::new(limit), + || { + UserBucket::new( + policy.revision, + limit, + self.user_scope.contention(RateDirection::Up), + self.user_scope.contention(RateDirection::Down), + ) + }, |bucket| { bucket.active_leases.fetch_add(1, Ordering::Relaxed); }, ); - bucket.set_rates(limit); + bucket.set_rates(policy.revision, limit); self.user_scope .active_leases .fetch_add(1, Ordering::Relaxed); @@ -121,12 +226,19 @@ impl TrafficLimiter { }; let bucket = self.cidr_buckets.get_or_insert_with( key, - || CidrBucket::new(limits), + || { + CidrBucket::new( + policy.revision, + limits, + self.cidr_scope.contention(RateDirection::Up), + self.cidr_scope.contention(RateDirection::Down), + ) + }, |bucket| { bucket.active_leases.fetch_add(1, Ordering::Relaxed); }, ); - bucket.set_rates(limits); + bucket.set_rates(policy.revision, limits); self.cidr_scope .active_leases .fetch_add(1, Ordering::Relaxed); @@ -136,20 +248,17 @@ impl TrafficLimiter { cidr_bucket = Some(bucket); } - if user_bucket.is_none() && cidr_bucket.is_none() { - return None; - } - - self.maybe_cleanup(); - Some(Arc::new(TrafficLease { + Arc::new(TrafficLeaseBinding { limiter: Arc::clone(self), + revision: policy.revision, user_bucket, cidr_bucket, cidr_user_key, cidr_user_share, - })) + }) } + /// Captures limiter telemetry without locking bucket registries. pub fn metrics_snapshot(&self) -> TrafficLimiterMetricsSnapshot { TrafficLimiterMetricsSnapshot { user_throttle_up_total: self.user_scope.throttle_up_total.load(Ordering::Relaxed), @@ -160,6 +269,46 @@ impl TrafficLimiter { user_wait_down_ms_total: self.user_scope.wait_down_ms_total.load(Ordering::Relaxed), cidr_wait_up_ms_total: self.cidr_scope.wait_up_ms_total.load(Ordering::Relaxed), cidr_wait_down_ms_total: self.cidr_scope.wait_down_ms_total.load(Ordering::Relaxed), + user_reserve_cas_retry_exhausted_up_total: self + .user_scope + .contention_up + .reserve_exhausted_total + .load(Ordering::Relaxed), + user_reserve_cas_retry_exhausted_down_total: self + .user_scope + .contention_down + .reserve_exhausted_total + .load(Ordering::Relaxed), + user_refund_cas_retry_exhausted_up_total: self + .user_scope + .contention_up + .refund_exhausted_total + .load(Ordering::Relaxed), + user_refund_cas_retry_exhausted_down_total: self + .user_scope + .contention_down + .refund_exhausted_total + .load(Ordering::Relaxed), + cidr_reserve_cas_retry_exhausted_up_total: self + .cidr_scope + .contention_up + .reserve_exhausted_total + .load(Ordering::Relaxed), + cidr_reserve_cas_retry_exhausted_down_total: self + .cidr_scope + .contention_down + .reserve_exhausted_total + .load(Ordering::Relaxed), + cidr_refund_cas_retry_exhausted_up_total: self + .cidr_scope + .contention_up + .refund_exhausted_total + .load(Ordering::Relaxed), + cidr_refund_cas_retry_exhausted_down_total: self + .cidr_scope + .contention_down + .refund_exhausted_total + .load(Ordering::Relaxed), user_active_leases: self.user_scope.active_leases.load(Ordering::Relaxed), cidr_active_leases: self.cidr_scope.active_leases.load(Ordering::Relaxed), user_policy_entries: self.user_scope.policy_entries.load(Ordering::Relaxed), @@ -167,6 +316,57 @@ impl TrafficLimiter { } } + /// Sets fixed contention counters for renderer mapping tests. + #[cfg(test)] + pub(crate) fn set_cas_contention_metrics_for_test(&self, values: [u64; 8]) { + let [ + user_reserve_up, + user_reserve_down, + user_refund_up, + user_refund_down, + cidr_reserve_up, + cidr_reserve_down, + cidr_refund_up, + cidr_refund_down, + ] = values; + for (counter, value) in [ + ( + &self.user_scope.contention_up.reserve_exhausted_total, + user_reserve_up, + ), + ( + &self.user_scope.contention_down.reserve_exhausted_total, + user_reserve_down, + ), + ( + &self.user_scope.contention_up.refund_exhausted_total, + user_refund_up, + ), + ( + &self.user_scope.contention_down.refund_exhausted_total, + user_refund_down, + ), + ( + &self.cidr_scope.contention_up.reserve_exhausted_total, + cidr_reserve_up, + ), + ( + &self.cidr_scope.contention_down.reserve_exhausted_total, + cidr_reserve_down, + ), + ( + &self.cidr_scope.contention_up.refund_exhausted_total, + cidr_refund_up, + ), + ( + &self.cidr_scope.contention_down.refund_exhausted_total, + cidr_refund_down, + ), + ] { + counter.store(value, Ordering::Relaxed); + } + } + pub(super) fn observe_throttle( &self, direction: RateDirection, diff --git a/src/proxy/traffic_limiter/tests.rs b/src/proxy/traffic_limiter/tests.rs index b9147da..dc62b41 100644 --- a/src/proxy/traffic_limiter/tests.rs +++ b/src/proxy/traffic_limiter/tests.rs @@ -1,10 +1,146 @@ use super::*; use crate::config::CidrRateLimitKey; +mod bucket_contention; + +impl DirectionBucket { + pub(crate) fn should_force_reserve_failure(&self) -> bool { + self.reserve_cas_attempts.fetch_add(1, Ordering::Relaxed); + let remaining = self.forced_reserve_failures.load(Ordering::Relaxed); + remaining > 0 + && self + .forced_reserve_failures + .compare_exchange( + remaining, + remaining - 1, + Ordering::Relaxed, + Ordering::Relaxed, + ) + .is_ok() + } + + pub(crate) fn should_force_refund_failure(&self) -> bool { + self.refund_cas_attempts.fetch_add(1, Ordering::Relaxed); + let remaining = self.forced_refund_failures.load(Ordering::Relaxed); + remaining > 0 + && self + .forced_refund_failures + .compare_exchange( + remaining, + remaining - 1, + Ordering::Relaxed, + Ordering::Relaxed, + ) + .is_ok() + } + + pub(crate) fn force_reserve_failures(&self, failures: usize) { + self.reserve_cas_attempts.store(0, Ordering::Relaxed); + self.forced_reserve_failures + .store(failures as u64, Ordering::Relaxed); + } + + pub(crate) fn force_refund_failures(&self, failures: usize) { + self.refund_cas_attempts.store(0, Ordering::Relaxed); + self.forced_refund_failures + .store(failures as u64, Ordering::Relaxed); + } + + pub(crate) fn reserve_cas_attempts(&self) -> u64 { + self.reserve_cas_attempts.load(Ordering::Relaxed) + } + + pub(crate) fn refund_cas_attempts(&self) -> u64 { + self.refund_cas_attempts.load(Ordering::Relaxed) + } +} + fn rate(up_bps: u64, down_bps: u64) -> RateLimitBps { RateLimitBps { up_bps, down_bps } } +fn reserve_at( + bucket: &Arc, + epoch: u64, + cap: u64, + requested: u64, +) -> Result, BucketReserveError> { + bucket.try_reserve_at(epoch, cap, requested, &mut ReserveCasBudget::new()) +} + +#[test] +fn stale_runtime_cannot_overwrite_newer_rate_policy() { + let limiter = TrafficLimiter::new(); + let mut newer = HashMap::new(); + newer.insert("alice".to_string(), rate(2_000, 3_000)); + assert!(limiter.apply_policy_from_source(2, newer, HashMap::new())); + + let mut stale = HashMap::new(); + stale.insert("alice".to_string(), rate(1_000, 1_000)); + assert!(!limiter.apply_policy_from_source(1, stale, HashMap::new())); + + let policy = limiter.policy.load_full(); + assert_eq!(policy.source_generation, 2); + assert_eq!(policy.user_limits["alice"].up_bps, 2_000); + assert_eq!(policy.user_limits["alice"].down_bps, 3_000); +} + +#[test] +fn active_runtime_can_publish_same_generation_rate_update() { + let limiter = TrafficLimiter::new(); + assert!(limiter.apply_policy_from_source(4, HashMap::new(), HashMap::new())); + let mut updated = HashMap::new(); + updated.insert("alice".to_string(), rate(4_000, 5_000)); + + assert!(limiter.apply_policy_from_source(4, updated, HashMap::new())); + + let policy = limiter.policy.load_full(); + assert_eq!(policy.source_generation, 4); + assert_eq!(policy.user_limits["alice"].up_bps, 4_000); +} + +#[test] +fn lease_acquisition_and_refresh_do_not_wait_for_policy_publication_lock() { + let limiter = TrafficLimiter::new(); + let mut initial = HashMap::new(); + initial.insert("alice".to_string(), rate(1_000, 1_000)); + limiter.apply_policy(initial, HashMap::new()); + let lease = limiter + .acquire_lease("alice", "203.0.113.7".parse().unwrap()) + .unwrap(); + + let mut updated = HashMap::new(); + updated.insert("alice".to_string(), rate(2_000, 2_000)); + limiter.apply_policy(updated, HashMap::new()); + + let publication = limiter.policy_update.lock(); + let (completed_tx, completed_rx) = std::sync::mpsc::channel(); + let acquire_limiter = Arc::clone(&limiter); + let acquire_tx = completed_tx.clone(); + let acquire = std::thread::spawn(move || { + let _lease = acquire_limiter + .acquire_lease("bob", "203.0.113.8".parse().unwrap()) + .unwrap(); + acquire_tx.send(()).unwrap(); + }); + let refresh = std::thread::spawn(move || { + let _ = lease.try_consume(RateDirection::Up, 1); + completed_tx.send(()).unwrap(); + }); + + let first_completed = completed_rx + .recv_timeout(std::time::Duration::from_secs(1)) + .is_ok(); + let second_completed = completed_rx + .recv_timeout(std::time::Duration::from_secs(1)) + .is_ok(); + drop(publication); + acquire.join().unwrap(); + refresh.join().unwrap(); + + assert!(first_completed && second_completed); +} + #[test] fn explicit_cidr_rule_wins_over_auto_template() { let limiter = TrafficLimiter::new(); @@ -74,3 +210,284 @@ fn auto_cidr_bucket_key_canonicalizes_network_address() { "auto:6:2001:db8::/64" ); } + +#[test] +fn refund_from_an_old_epoch_does_not_reduce_the_current_epoch() { + let bucket = Arc::new(DirectionBucket::default()); + let old_debit = reserve_at(&bucket, 7, 100, 80).unwrap().unwrap(); + let current_debit = reserve_at(&bucket, 8, 100, 60).unwrap().unwrap(); + + drop(old_debit); + + assert_eq!(bucket.used_at(8), Some(60)); + assert_eq!(bucket.refund_cas_attempts(), 0); + assert_eq!( + bucket + .contention + .refund_exhausted_total + .load(Ordering::Relaxed), + 0 + ); + drop(current_debit); +} + +#[test] +fn concurrent_rollover_cannot_publish_multiple_epoch_budgets() { + const CONTENDERS: usize = 32; + + let bucket = Arc::new(DirectionBucket::default()); + let barrier = Arc::new(std::sync::Barrier::new(CONTENDERS)); + let mut threads = Vec::with_capacity(CONTENDERS); + for _ in 0..CONTENDERS { + let bucket = Arc::clone(&bucket); + let barrier = Arc::clone(&barrier); + threads.push(std::thread::spawn(move || { + barrier.wait(); + reserve_at(&bucket, 9, 100, 100) + .unwrap() + .map(|mut debit| debit.commit_all()) + .unwrap_or(0) + })); + } + + let granted = threads + .into_iter() + .map(|thread| thread.join().unwrap()) + .sum::(); + assert_eq!(granted, 100); + assert_eq!(bucket.used_at(9), Some(100)); +} + +#[test] +fn scheduler_pressure_never_exceeds_a_packed_epoch_budget() { + const CONTENDERS: usize = 4; + const EPOCHS: usize = 10_000; + + let bucket = Arc::new(DirectionBucket::default()); + let barrier = Arc::new(std::sync::Barrier::new(CONTENDERS)); + let mut threads = Vec::with_capacity(CONTENDERS); + for _ in 0..CONTENDERS { + let bucket = Arc::clone(&bucket); + let barrier = Arc::clone(&barrier); + threads.push(std::thread::spawn(move || { + let mut grants = Vec::with_capacity(EPOCHS); + for epoch in 1..=EPOCHS as u64 { + barrier.wait(); + let granted = reserve_at(&bucket, epoch, 100, 100) + .unwrap() + .map(|mut debit| debit.commit_all()) + .unwrap_or(0); + grants.push(granted); + barrier.wait(); + } + grants + })); + } + + let grants = threads + .into_iter() + .map(|thread| thread.join().unwrap()) + .collect::>(); + for epoch_index in 0..EPOCHS { + let granted = grants + .iter() + .map(|thread_grants| thread_grants[epoch_index]) + .sum::(); + assert_eq!(granted, 100); + } +} + +#[test] +fn stale_policy_revision_cannot_restore_an_old_rate() { + let bucket = UserBucket::new( + 2, + rate(2_000, 3_000), + Arc::new(CasContentionMetrics::default()), + Arc::new(CasContentionMetrics::default()), + ); + + bucket.set_rates(3, rate(4_000, 5_000)); + bucket.set_rates(2, rate(6_000, 7_000)); + + assert_eq!(bucket.rates.get(RateDirection::Up), 4_000); + assert_eq!(bucket.rates.get(RateDirection::Down), 5_000); +} + +#[test] +fn dropped_debit_refunds_only_its_packed_epoch() { + let bucket = Arc::new(DirectionBucket::default()); + let debit = reserve_at(&bucket, 11, 100, 80).unwrap().unwrap(); + + drop(debit); + + assert_eq!(bucket.used_at(11), Some(0)); + assert!(matches!( + reserve_at(&bucket, PACKED_EPOCH_MAX + 1, 100, 1), + Err(BucketReserveError::StaleEpoch) + )); +} + +#[test] +fn concurrent_first_use_counts_one_active_cidr_user() { + const CONTENDERS: usize = 32; + + let bucket = Arc::new(CidrDirectionBucket::default()); + let user = Arc::new(CidrUserDirectionState::default()); + let barrier = Arc::new(std::sync::Barrier::new(CONTENDERS)); + let mut threads = Vec::with_capacity(CONTENDERS); + for _ in 0..CONTENDERS { + let bucket = Arc::clone(&bucket); + let user = Arc::clone(&user); + let barrier = Arc::clone(&barrier); + threads.push(std::thread::spawn(move || { + barrier.wait(); + user.ensure_active(13, &bucket.active_users, &mut ReserveCasBudget::new()) + })); + } + let results: Vec<_> = threads + .into_iter() + .map(|thread| thread.join().unwrap()) + .collect(); + + assert!(results.iter().any(|result| *result == Ok(true))); + assert!(results.iter().all(|result| matches!( + result, + Ok(true) + | Err(BucketReserveError::Contended) + | Err(BucketReserveError::RefundContended) + | Err(BucketReserveError::ReserveAndRefundContended) + ))); + let leaked = bucket + .active_users + .contention + .refund_exhausted_total + .load(Ordering::Relaxed); + assert_eq!(bucket.active_users.used_at(13), Some(1 + leaked)); + assert!(1 + leaked <= CONTENDERS as u64); + assert_eq!(user.used.used_at(13), Some(0)); +} + +#[test] +fn configured_rate_maximum_fits_the_packed_epoch_budget() { + assert_eq!(bytes_per_epoch(100_000_000_000), 250_000_000); + assert!(bytes_per_epoch(100_000_000_000) <= PACKED_USAGE_MASK); +} + +#[test] +fn dropped_traffic_reservation_refunds_user_and_cidr_debits() { + let limiter = TrafficLimiter::new(); + let mut user_limits = HashMap::new(); + user_limits.insert("alice".to_string(), rate(400_000, 400_000)); + let mut cidr_limits = HashMap::new(); + cidr_limits.insert( + CidrRateLimitKey::Network("203.0.113.0/24".parse().unwrap()), + rate(400_000, 400_000), + ); + limiter.apply_policy(user_limits, cidr_limits); + let lease = limiter + .acquire_lease("alice", "203.0.113.7".parse().unwrap()) + .unwrap(); + + let reservation = lease.try_reserve(RateDirection::Down, 800); + assert_eq!(reservation.result().granted, 800); + let epoch = reservation.user.as_ref().unwrap().epoch; + assert_eq!(reservation.cidr.as_ref().unwrap().epoch, epoch); + assert_eq!(reservation.cidr_user.as_ref().unwrap().epoch, epoch); + drop(reservation); + + let binding = lease.binding.load_full(); + let user_bucket = binding.user_bucket.as_ref().unwrap(); + let cidr_bucket = binding.cidr_bucket.as_ref().unwrap(); + let cidr_user = binding.cidr_user_share.as_ref().unwrap(); + assert_eq!(user_bucket.down.used_at(epoch), Some(0)); + assert_eq!(cidr_bucket.down.used.used_at(epoch), Some(0)); + assert_eq!(cidr_user.down.used.used_at(epoch), Some(0)); +} + +#[test] +fn partial_traffic_settlement_charges_only_committed_bytes() { + let limiter = TrafficLimiter::new(); + let mut user_limits = HashMap::new(); + user_limits.insert("alice".to_string(), rate(400_000, 400_000)); + limiter.apply_policy(user_limits, HashMap::new()); + let lease = limiter + .acquire_lease("alice", "203.0.113.7".parse().unwrap()) + .unwrap(); + + let reservation = lease.try_reserve(RateDirection::Down, 800); + let epoch = reservation.user.as_ref().unwrap().epoch; + reservation.settle_written(300); + + assert_eq!( + lease + .binding + .load_full() + .user_bucket + .as_ref() + .unwrap() + .down + .used_at(epoch), + Some(300) + ); +} + +#[test] +fn active_lease_observes_policy_removal() { + let limiter = TrafficLimiter::new(); + let mut user_limits = HashMap::new(); + user_limits.insert("alice".to_string(), rate(1, 0)); + limiter.apply_policy(user_limits, HashMap::new()); + let lease = limiter + .acquire_lease("alice", "203.0.113.7".parse().unwrap()) + .unwrap(); + assert!(lease.binding.load().user_bucket.is_some()); + + limiter.apply_policy(HashMap::new(), HashMap::new()); + + assert_eq!(lease.try_consume(RateDirection::Up, 2).granted, 2); + assert!(lease.binding.load().user_bucket.is_none()); +} + +#[test] +fn active_lease_observes_policy_addition() { + let limiter = TrafficLimiter::new(); + let lease = limiter + .acquire_lease("alice", "203.0.113.7".parse().unwrap()) + .unwrap(); + assert!(lease.binding.load().user_bucket.is_none()); + + let mut user_limits = HashMap::new(); + user_limits.insert("alice".to_string(), rate(1, 0)); + limiter.apply_policy(user_limits, HashMap::new()); + + assert_eq!(lease.try_consume(RateDirection::Up, 2).granted, 1); + assert!(lease.binding.load().user_bucket.is_some()); +} + +#[test] +fn reservation_refund_stays_with_retired_binding() { + let limiter = TrafficLimiter::new(); + let mut user_limits = HashMap::new(); + user_limits.insert("alice".to_string(), rate(400_000, 400_000)); + limiter.apply_policy(user_limits, HashMap::new()); + let lease = limiter + .acquire_lease("alice", "203.0.113.7".parse().unwrap()) + .unwrap(); + + let reservation = lease.try_reserve(RateDirection::Down, 800); + let old_bucket = Arc::clone( + reservation + ._binding + .user_bucket + .as_ref() + .expect("the original policy must bind a user bucket"), + ); + let epoch = reservation.user.as_ref().unwrap().epoch; + limiter.apply_policy(HashMap::new(), HashMap::new()); + assert_eq!(lease.try_consume(RateDirection::Down, 1).granted, 1); + assert!(lease.binding.load().user_bucket.is_none()); + + drop(reservation); + + assert_eq!(old_bucket.down.used_at(epoch), Some(0)); +} diff --git a/src/proxy/traffic_limiter/tests/bucket_contention.rs b/src/proxy/traffic_limiter/tests/bucket_contention.rs new file mode 100644 index 0000000..731a710 --- /dev/null +++ b/src/proxy/traffic_limiter/tests/bucket_contention.rs @@ -0,0 +1,471 @@ +use super::*; +use proptest::prelude::*; + +#[test] +fn reserve_stops_after_the_attempt_limit() { + let bucket = Arc::new(DirectionBucket::default()); + bucket.force_reserve_failures(RESERVE_CAS_ATTEMPT_LIMIT); + let mut budget = ReserveCasBudget::new(); + + let reservation = bucket.try_reserve_at(1, 100, 1, &mut budget); + + assert!(matches!(reservation, Err(BucketReserveError::Contended))); + assert_eq!( + bucket.reserve_cas_attempts(), + RESERVE_CAS_ATTEMPT_LIMIT as u64 + ); + assert_eq!(bucket.used_at(1), None); +} + +#[test] +fn reserve_succeeds_on_the_last_allowed_attempt() { + let bucket = Arc::new(DirectionBucket::default()); + bucket.force_reserve_failures(RESERVE_CAS_ATTEMPT_LIMIT - 1); + let mut budget = ReserveCasBudget::new(); + + let mut debit = bucket + .try_reserve_at(1, 100, 80, &mut budget) + .unwrap() + .unwrap(); + + assert_eq!(debit.commit_all(), 80); + assert_eq!( + bucket.reserve_cas_attempts(), + RESERVE_CAS_ATTEMPT_LIMIT as u64 + ); + assert!(budget.is_exhausted()); + assert_eq!(bucket.used_at(1), Some(80)); +} + +#[test] +fn refund_stops_after_the_attempt_limit() { + let bucket = Arc::new(DirectionBucket::default()); + let debit = reserve_at(&bucket, 1, 100, 80).unwrap().unwrap(); + bucket.force_refund_failures(REFUND_CAS_ATTEMPT_LIMIT); + + drop(debit); + + assert_eq!( + bucket.refund_cas_attempts(), + REFUND_CAS_ATTEMPT_LIMIT as u64 + ); + assert_eq!(bucket.used_at(1), Some(80)); + assert_eq!( + bucket + .contention + .refund_exhausted_total + .load(Ordering::Relaxed), + 1 + ); + + let mut tail = reserve_at(&bucket, 1, 100, 100).unwrap().unwrap(); + assert_eq!(tail.commit_all(), 20); + assert!(reserve_at(&bucket, 1, 100, 1).unwrap().is_none()); + assert_eq!(bucket.used_at(1), Some(100)); + + let mut next = reserve_at(&bucket, 2, 100, 100).unwrap().unwrap(); + assert_eq!(next.commit_all(), 100); +} + +#[test] +fn refund_succeeds_on_the_last_allowed_attempt() { + let bucket = Arc::new(DirectionBucket::default()); + let debit = reserve_at(&bucket, 1, 100, 80).unwrap().unwrap(); + bucket.force_refund_failures(REFUND_CAS_ATTEMPT_LIMIT - 1); + + drop(debit); + + assert_eq!( + bucket.refund_cas_attempts(), + REFUND_CAS_ATTEMPT_LIMIT as u64 + ); + assert_eq!(bucket.used_at(1), Some(0)); + assert_eq!( + bucket + .contention + .refund_exhausted_total + .load(Ordering::Relaxed), + 0 + ); +} + +#[test] +fn partial_refund_exhaustion_retains_the_full_charge() { + let bucket = Arc::new(DirectionBucket::default()); + let mut debit = reserve_at(&bucket, 1, 100, 80).unwrap().unwrap(); + bucket.force_refund_failures(REFUND_CAS_ATTEMPT_LIMIT); + + debit.settle(30); + drop(debit); + + assert_eq!(bucket.used_at(1), Some(80)); + assert_eq!( + bucket.refund_cas_attempts(), + REFUND_CAS_ATTEMPT_LIMIT as u64 + ); + assert_eq!( + bucket + .contention + .refund_exhausted_total + .load(Ordering::Relaxed), + 1 + ); +} + +#[test] +fn lease_contention_is_not_reported_as_throttling() { + let limiter = TrafficLimiter::new(); + let mut user_limits = HashMap::new(); + user_limits.insert("alice".to_string(), rate(400_000, 400_000)); + limiter.apply_policy(user_limits, HashMap::new()); + let lease = limiter + .acquire_lease("alice", "203.0.113.7".parse().unwrap()) + .unwrap(); + let bucket = Arc::clone(&lease.binding.load_full().user_bucket.as_ref().unwrap().down); + bucket.force_reserve_failures(RESERVE_CAS_ATTEMPT_LIMIT); + + let result = lease.try_consume(RateDirection::Down, 1); + let metrics = limiter.metrics_snapshot(); + + assert_eq!(result.granted, 0); + assert!(!result.blocked_user); + assert!(!result.blocked_cidr); + assert_eq!(metrics.user_reserve_cas_retry_exhausted_down_total, 1); + assert_eq!(metrics.user_throttle_down_total, 0); +} + +#[test] +fn cidr_contention_rolls_back_provisional_user_debits() { + let limiter = TrafficLimiter::new(); + let mut user_limits = HashMap::new(); + user_limits.insert("alice".to_string(), rate(400_000, 400_000)); + let mut cidr_limits = HashMap::new(); + cidr_limits.insert( + CidrRateLimitKey::Network("203.0.113.0/24".parse().unwrap()), + rate(400_000, 400_000), + ); + limiter.apply_policy(user_limits, cidr_limits); + let lease = limiter + .acquire_lease("alice", "203.0.113.7".parse().unwrap()) + .unwrap(); + let binding = lease.binding.load_full(); + let cidr_bucket = binding.cidr_bucket.as_ref().unwrap(); + cidr_bucket + .down + .used + .force_reserve_failures(RESERVE_CAS_ATTEMPT_LIMIT); + + let reservation = lease.try_reserve(RateDirection::Down, 800); + let result = reservation.result(); + let epoch = reservation.user.as_ref().unwrap().epoch; + let user_attempts = binding + .user_bucket + .as_ref() + .unwrap() + .down + .reserve_cas_attempts(); + let active_user_attempts = cidr_bucket.down.active_users.reserve_cas_attempts(); + let cidr_user_attempts = binding + .cidr_user_share + .as_ref() + .unwrap() + .down + .used + .reserve_cas_attempts(); + let aggregate_attempts = cidr_bucket.down.used.reserve_cas_attempts(); + drop(reservation); + + assert_eq!(result.granted, 0); + assert!(!result.blocked_user); + assert!(!result.blocked_cidr); + assert_eq!( + binding.user_bucket.as_ref().unwrap().down.used_at(epoch), + Some(0) + ); + assert_eq!(cidr_bucket.down.used.used_at(epoch), None); + assert_eq!( + user_attempts + active_user_attempts + cidr_user_attempts + aggregate_attempts, + RESERVE_CAS_ATTEMPT_LIMIT as u64 + ); + assert_eq!( + binding + .cidr_user_share + .as_ref() + .unwrap() + .down + .used + .used_at(epoch), + Some(0) + ); + let metrics = limiter.metrics_snapshot(); + assert_eq!(metrics.cidr_reserve_cas_retry_exhausted_down_total, 1); + assert_eq!(metrics.cidr_throttle_down_total, 0); +} + +#[test] +fn contention_snapshot_preserves_scope_direction_and_operation() { + let limiter = TrafficLimiter::new(); + limiter + .user_scope + .contention_up + .reserve_exhausted_total + .store(1, Ordering::Relaxed); + limiter + .user_scope + .contention_down + .reserve_exhausted_total + .store(2, Ordering::Relaxed); + limiter + .user_scope + .contention_up + .refund_exhausted_total + .store(3, Ordering::Relaxed); + limiter + .user_scope + .contention_down + .refund_exhausted_total + .store(4, Ordering::Relaxed); + limiter + .cidr_scope + .contention_up + .reserve_exhausted_total + .store(5, Ordering::Relaxed); + limiter + .cidr_scope + .contention_down + .reserve_exhausted_total + .store(6, Ordering::Relaxed); + limiter + .cidr_scope + .contention_up + .refund_exhausted_total + .store(7, Ordering::Relaxed); + limiter + .cidr_scope + .contention_down + .refund_exhausted_total + .store(8, Ordering::Relaxed); + + let snapshot = limiter.metrics_snapshot(); + + assert_eq!(snapshot.user_reserve_cas_retry_exhausted_up_total, 1); + assert_eq!(snapshot.user_reserve_cas_retry_exhausted_down_total, 2); + assert_eq!(snapshot.user_refund_cas_retry_exhausted_up_total, 3); + assert_eq!(snapshot.user_refund_cas_retry_exhausted_down_total, 4); + assert_eq!(snapshot.cidr_reserve_cas_retry_exhausted_up_total, 5); + assert_eq!(snapshot.cidr_reserve_cas_retry_exhausted_down_total, 6); + assert_eq!(snapshot.cidr_refund_cas_retry_exhausted_up_total, 7); + assert_eq!(snapshot.cidr_refund_cas_retry_exhausted_down_total, 8); +} + +#[test] +fn cidr_activation_consumes_one_shared_attempt_budget() { + let bucket = CidrDirectionBucket::default(); + let user = CidrUserDirectionState::default(); + user.used.force_reserve_failures(RESERVE_CAS_ATTEMPT_LIMIT); + let mut budget = ReserveCasBudget::new(); + + let activation = user.ensure_active(13, &bucket.active_users, &mut budget); + + assert_eq!(activation, Err(BucketReserveError::Contended)); + assert_eq!( + user.used.reserve_cas_attempts() + bucket.active_users.reserve_cas_attempts(), + RESERVE_CAS_ATTEMPT_LIMIT as u64 + ); + assert_eq!(bucket.active_users.used_at(13), Some(0)); + assert_eq!(user.used.used_at(13), None); + assert_eq!( + bucket + .active_users + .contention + .refund_exhausted_total + .load(Ordering::Relaxed), + 0 + ); +} + +#[test] +fn cidr_activation_reports_refund_contention_without_reserve_exhaustion() { + let bucket = CidrDirectionBucket::default(); + let user = CidrUserDirectionState::default(); + user.used.force_reserve_failures(1); + bucket + .active_users + .force_refund_failures(REFUND_CAS_ATTEMPT_LIMIT); + let mut budget = ReserveCasBudget::new(); + + let activation = user.ensure_active(13, &bucket.active_users, &mut budget); + + assert_eq!(activation, Err(BucketReserveError::RefundContended)); + assert!(!budget.is_exhausted()); + assert!(!activation.unwrap_err().exhausted_reserve_budget()); + assert_eq!(bucket.active_users.used_at(13), Some(1)); + assert_eq!(user.used.used_at(13), None); + assert_eq!( + bucket + .active_users + .contention + .refund_exhausted_total + .load(Ordering::Relaxed), + 1 + ); +} + +#[test] +fn cidr_activation_preserves_combined_contention_cause() { + let bucket = CidrDirectionBucket::default(); + let user = CidrUserDirectionState::default(); + bucket + .active_users + .force_refund_failures(REFUND_CAS_ATTEMPT_LIMIT); + let mut budget = ReserveCasBudget { remaining: 1 }; + + let activation = user.ensure_active(13, &bucket.active_users, &mut budget); + + assert_eq!( + activation, + Err(BucketReserveError::ReserveAndRefundContended) + ); + assert!(budget.is_exhausted()); + assert!(activation.unwrap_err().exhausted_reserve_budget()); + assert_eq!(bucket.active_users.used_at(13), Some(1)); + assert_eq!(user.used.used_at(13), None); + assert_eq!( + bucket + .active_users + .contention + .refund_exhausted_total + .load(Ordering::Relaxed), + 1 + ); +} + +#[test] +fn cidr_first_grants_preserve_the_current_soft_fair_share() { + let bucket = CidrDirectionBucket::default(); + let first = CidrUserDirectionState::default(); + let second = CidrUserDirectionState::default(); + assert_eq!( + first.ensure_active(17, &bucket.active_users, &mut ReserveCasBudget::new(),), + Ok(true) + ); + assert_eq!( + second.ensure_active(17, &bucket.active_users, &mut ReserveCasBudget::new(),), + Ok(true) + ); + + let mut first_reservation = bucket + .try_reserve(&first, 17, 100, 100, &mut ReserveCasBudget::new()) + .unwrap(); + let mut second_reservation = bucket + .try_reserve(&second, 17, 100, 100, &mut ReserveCasBudget::new()) + .unwrap(); + + assert_eq!(first_reservation.granted, 50); + assert_eq!(second_reservation.granted, 50); + assert_eq!(bucket.active_users.used_at(17), Some(2)); + first_reservation + .aggregate_debit + .as_mut() + .unwrap() + .commit_all(); + first_reservation.user_debit.as_mut().unwrap().commit_all(); + second_reservation + .aggregate_debit + .as_mut() + .unwrap() + .commit_all(); + second_reservation.user_debit.as_mut().unwrap().commit_all(); + assert_eq!(bucket.used.used_at(17), Some(100)); + assert_eq!(first.used.used_at(17), Some(50)); + assert_eq!(second.used.used_at(17), Some(50)); +} + +#[test] +fn concurrent_commit_and_refund_preserve_fail_closed_accounting() { + const WORKERS: usize = 32; + const OPERATIONS_PER_WORKER: usize = 1_024; + const EPOCH: u64 = 19; + const CAP: u64 = (WORKERS * OPERATIONS_PER_WORKER) as u64; + + let bucket = Arc::new(DirectionBucket::default()); + let committed = Arc::new(AtomicU64::new(0)); + let barrier = Arc::new(std::sync::Barrier::new(WORKERS)); + let mut threads = Vec::with_capacity(WORKERS); + for worker in 0..WORKERS { + let bucket = Arc::clone(&bucket); + let committed = Arc::clone(&committed); + let barrier = Arc::clone(&barrier); + threads.push(std::thread::spawn(move || { + barrier.wait(); + for operation in 0..OPERATIONS_PER_WORKER { + match reserve_at(&bucket, EPOCH, CAP, 1) { + Ok(Some(mut debit)) if (worker + operation) % 2 == 0 => { + committed.fetch_add(debit.commit_all(), Ordering::Relaxed); + } + Ok(Some(debit)) => drop(debit), + Err(BucketReserveError::Contended) => {} + Ok(None) => panic!("capacity exhausted before every operation ran"), + Err(error) => panic!("unexpected reservation error: {error:?}"), + } + } + })); + } + for thread in threads { + thread.join().unwrap(); + } + + let committed = committed.load(Ordering::Relaxed); + let leaked = bucket + .contention + .refund_exhausted_total + .load(Ordering::Relaxed); + let used = bucket.used_at(EPOCH).unwrap(); + assert_eq!(used, committed + leaked); + assert!(used <= CAP); +} + +proptest! { + #[test] + fn sequential_debit_lifecycle_matches_the_epoch_model( + operations in prop::collection::vec((0u8..4, 1u64..128), 1..128), + ) { + const CAP: u64 = 512; + + let bucket = Arc::new(DirectionBucket::default()); + let mut epoch = 1u64; + let mut model_used = 0u64; + for (operation, requested) in operations { + if operation == 3 { + epoch += 1; + model_used = 0; + let debit = reserve_at(&bucket, epoch, CAP, 1).unwrap().unwrap(); + drop(debit); + prop_assert_eq!(bucket.used_at(epoch), Some(0)); + continue; + } + + let reservation = reserve_at(&bucket, epoch, CAP, requested).unwrap(); + let expected_grant = requested.min(CAP.saturating_sub(model_used)); + if expected_grant == 0 { + prop_assert!(reservation.is_none()); + continue; + } + let mut debit = reservation.unwrap(); + prop_assert_eq!(debit.refundable, expected_grant); + match operation { + 0 => { + model_used += debit.commit_all(); + } + 1 => { + let committed = expected_grant / 2; + debit.settle(committed); + model_used += committed; + } + 2 => drop(debit), + _ => unreachable!(), + } + prop_assert_eq!(bucket.used_at(epoch), Some(model_used)); + } + } +} diff --git a/src/proxy/user_admission.rs b/src/proxy/user_admission.rs new file mode 100644 index 0000000..8e8724c --- /dev/null +++ b/src/proxy/user_admission.rs @@ -0,0 +1,535 @@ +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::atomic::AtomicU8; + +use parking_lot::Mutex; +use tokio_util::sync::CancellationToken; + +use crate::crypto::sha256; +use crate::stats::QuotaStore; + +const REGISTRATION_PENDING: u8 = 0; +const REGISTRATION_ACTIVE: u8 = 1; +const REGISTRATION_DROPPED: u8 = 2; + +// Authenticated-owner publication and RAII deregistration. +mod registration; +pub(crate) use registration::{UserAdmissionPublication, UserSessionRegistration}; + +/// Stable secret identity used to fence authentication across runtime generations. +pub(crate) type UserCredentialId = [u8; 16]; + +/// Monotonic identity of one configured username lifetime. +pub(crate) type UserIncarnation = u64; + +#[derive(Clone, Copy, PartialEq, Eq)] +struct EffectiveUser { + credential_id: UserCredentialId, + enabled: bool, +} + +#[derive(Clone, Copy, PartialEq, Eq)] +enum UserOverride { + Present(EffectiveUser), + Deleted, +} + +struct UserRecord { + configured: Option, + mutation_override: Option, + incarnation: UserIncarnation, +} + +impl UserRecord { + fn effective(&self) -> Option { + match self.mutation_override { + Some(UserOverride::Present(user)) => Some(user), + Some(UserOverride::Deleted) => None, + None => self.configured, + } + } +} + +struct RegisteredOwner { + token: CancellationToken, + incarnation: UserIncarnation, +} + +#[derive(Default)] +struct UserAdmissionState { + initialized: bool, + epoch: u64, + active_config_source: Option, + stale_config_source_rejections: u64, + next_incarnation: UserIncarnation, + next_registration_id: u64, + users: HashMap, + owners_by_user: HashMap>, +} + +impl UserAdmissionState { + fn allocate_incarnation(&mut self) -> UserIncarnation { + self.next_incarnation = self.next_incarnation.checked_add(1).unwrap_or(u64::MAX); + self.next_incarnation + } + + fn allocate_registration_id(&mut self) -> Option { + let next = self.next_registration_id.checked_add(1)?; + self.next_registration_id = next; + Some(next) + } + + fn bump_epoch(&mut self) { + self.epoch = self.epoch.checked_add(1).unwrap_or(u64::MAX); + } + + fn owner_tokens(&self, user: &str) -> Vec { + self.owners_by_user + .get(user) + .map(|owners| owners.values().map(|owner| owner.token.clone()).collect()) + .unwrap_or_default() + } +} + +/// Result of one durable user mutation applied to the process admission authority. +pub(crate) struct UserMutationResult { + /// Incarnation invalidated or created by the mutation. + pub(crate) incarnation: UserIncarnation, + /// Number of live owners cancelled by the mutation. + pub(crate) cancelled: usize, + /// Whether the effective enabled state changed from enabled to disabled. + pub(crate) newly_disabled: bool, +} + +/// Process-owned user authentication and live-owner authority. +pub(crate) struct UserAdmissionAuthority { + state: Mutex, + quota_store: Arc, +} + +impl UserAdmissionAuthority { + /// Creates an uninitialized authority for isolated tests and startup wiring. + pub(crate) fn new() -> Arc { + Self::new_with_quota_store(Arc::new(QuotaStore::default())) + } + + /// Creates an authority coupled to the process-scoped quota identity store. + pub(crate) fn new_with_quota_store(quota_store: Arc) -> Arc { + Arc::new(Self { + state: Mutex::new(UserAdmissionState::default()), + quota_store, + }) + } + + /// Returns the mutation epoch used to reject stale candidate configuration. + pub(crate) fn epoch(&self) -> u64 { + self.state.lock().epoch + } + + /// Reconciles the complete configured user set into the process authority. + pub(crate) fn apply_config( + &self, + users: &HashMap, + user_enabled: &HashMap, + ) -> Vec<(String, usize)> { + self.activate_config_source(0, None, users, user_enabled) + .unwrap_or_default() + } + + /// Transfers configuration ownership to one runtime generation. + pub(crate) fn activate_config_source( + &self, + source_generation: u64, + expected_epoch: Option, + users: &HashMap, + user_enabled: &HashMap, + ) -> Option> { + self.apply_config_locked(source_generation, expected_epoch, true, users, user_enabled) + } + + /// Reconciles an update only while its runtime generation owns configuration. + pub(crate) fn apply_config_from_source( + &self, + source_generation: u64, + users: &HashMap, + user_enabled: &HashMap, + ) -> Option> { + self.apply_config_locked(source_generation, None, false, users, user_enabled) + } + + /// Returns the number of rejected updates from non-owning generations. + pub(crate) fn stale_config_source_rejections(&self) -> u64 { + self.state.lock().stale_config_source_rejections + } + + fn apply_config_locked( + &self, + source_generation: u64, + expected_epoch: Option, + activate_source: bool, + users: &HashMap, + user_enabled: &HashMap, + ) -> Option> { + let configured = users + .iter() + .filter_map(|(user, secret)| { + credential_id_from_hex(secret).map(|credential_id| { + ( + user.clone(), + EffectiveUser { + credential_id, + enabled: user_enabled.get(user).copied().unwrap_or(true), + }, + ) + }) + }) + .collect::>(); + let cancellations = { + let mut state = self.state.lock(); + if activate_source { + if state + .active_config_source + .is_some_and(|active| source_generation < active) + { + state.stale_config_source_rejections = + state.stale_config_source_rejections.saturating_add(1); + return None; + } + state.active_config_source = Some(source_generation); + } else if state.active_config_source != Some(source_generation) { + state.stale_config_source_rejections = + state.stale_config_source_rejections.saturating_add(1); + return None; + } + if expected_epoch.is_some_and(|epoch| state.epoch != epoch) { + return None; + } + + let mut changed = !state.initialized; + state.initialized = true; + let existing_users = state.users.keys().cloned().collect::>(); + let mut cancellations = Vec::new(); + + for user in existing_users { + let desired = configured.get(&user).copied(); + let old_effective = state.users.get(&user).and_then(UserRecord::effective); + let override_matches = state.users.get(&user).is_some_and(|record| { + matches!( + (record.mutation_override, desired), + (Some(UserOverride::Present(current)), Some(next)) if current == next + ) || matches!( + (record.mutation_override, desired), + (Some(UserOverride::Deleted), None) + ) + }); + + if let Some(record) = state.users.get_mut(&user) { + if record.configured != desired || override_matches { + changed = true; + } + record.configured = desired; + if override_matches { + record.mutation_override = None; + } + } + + let new_effective = state.users.get(&user).and_then(UserRecord::effective); + if old_effective != new_effective { + let identity_changed = old_effective.map(|entry| entry.credential_id) + != new_effective.map(|entry| entry.credential_id); + if identity_changed { + let incarnation = state.allocate_incarnation(); + if let Some(record) = state.users.get_mut(&user) { + record.incarnation = incarnation; + } + match (old_effective, new_effective) { + (None, Some(_)) => { + self.quota_store.activate_fresh(&user, incarnation); + } + (Some(_), Some(_)) => { + self.quota_store + .advance_preserving_usage(&user, incarnation); + } + (Some(_), None) => { + self.quota_store.retire_through(&user, incarnation); + } + (None, None) => {} + } + } + if identity_changed + || old_effective.is_some_and(|entry| entry.enabled) + && new_effective.is_none_or(|entry| !entry.enabled) + { + let tokens = state.owner_tokens(&user); + cancellations.push((user, tokens)); + } + } + } + + for (user, desired) in configured { + if state.users.contains_key(&user) { + continue; + } + changed = true; + let incarnation = state.allocate_incarnation(); + state.users.insert( + user.clone(), + UserRecord { + configured: Some(desired), + mutation_override: None, + incarnation, + }, + ); + self.quota_store.activate_fresh(&user, incarnation); + } + + if changed { + state.bump_epoch(); + } + cancellations + }; + Some(cancel_owners(cancellations)) + } + + /// Applies one persisted user value ahead of asynchronous runtime reload. + pub(crate) fn stage_user( + &self, + user: &str, + secret: &str, + enabled: bool, + ) -> Option { + let credential_id = credential_id_from_hex(secret)?; + Some(self.stage_user_credential(user, credential_id, enabled)) + } + + /// Applies one already validated credential mutation ahead of runtime reload. + pub(crate) fn stage_user_credential( + &self, + user: &str, + credential_id: UserCredentialId, + enabled: bool, + ) -> UserMutationResult { + let desired = EffectiveUser { + credential_id, + enabled, + }; + let (incarnation, newly_disabled, tokens) = { + let mut state = self.state.lock(); + let previous = state.users.get(user).and_then(UserRecord::effective); + let identity_changed = previous.map(|entry| entry.credential_id) != Some(credential_id); + let incarnation = if identity_changed { + state.allocate_incarnation() + } else { + state + .users + .get(user) + .map(|record| record.incarnation) + .unwrap_or_else(|| state.allocate_incarnation()) + }; + let record = state.users.entry(user.to_string()).or_insert(UserRecord { + configured: None, + mutation_override: None, + incarnation, + }); + record.mutation_override = Some(UserOverride::Present(desired)); + record.incarnation = incarnation; + if identity_changed { + if previous.is_some() { + self.quota_store.advance_preserving_usage(user, incarnation); + } else { + self.quota_store.activate_fresh(user, incarnation); + } + } + state.initialized = true; + state.bump_epoch(); + let newly_disabled = previous.is_some_and(|entry| entry.enabled) && !enabled; + let tokens = if identity_changed || !enabled { + state.owner_tokens(user) + } else { + Vec::new() + }; + (incarnation, newly_disabled, tokens) + }; + let cancelled = tokens.len(); + for token in tokens { + token.cancel(); + } + UserMutationResult { + incarnation, + cancelled, + newly_disabled, + } + } + + /// Installs a deletion tombstone and cancels every owner of the old incarnation. + pub(crate) fn delete_user(&self, user: &str) -> UserMutationResult { + let (incarnation, newly_disabled, tokens) = { + let mut state = self.state.lock(); + let previous = state.users.get(user).and_then(UserRecord::effective); + let incarnation = state.allocate_incarnation(); + let record = state.users.entry(user.to_string()).or_insert(UserRecord { + configured: None, + mutation_override: None, + incarnation, + }); + record.mutation_override = Some(UserOverride::Deleted); + record.incarnation = incarnation; + self.quota_store.retire_through(user, incarnation); + state.initialized = true; + state.bump_epoch(); + ( + incarnation, + previous.is_some_and(|entry| entry.enabled), + state.owner_tokens(user), + ) + }; + let cancelled = tokens.len(); + for token in tokens { + token.cancel(); + } + UserMutationResult { + incarnation, + cancelled, + newly_disabled, + } + } + + /// Returns whether the effective process policy currently enables a user. + pub(crate) fn is_user_enabled(&self, user: &str) -> bool { + let state = self.state.lock(); + if !state.initialized { + return true; + } + state + .users + .get(user) + .and_then(UserRecord::effective) + .is_some_and(|entry| entry.enabled) + } + + /// Returns the authenticated incarnation for an exact current credential. + pub(crate) fn authenticated_incarnation( + &self, + user: &str, + credential_id: UserCredentialId, + ) -> Option { + let state = self.state.lock(); + if !state.initialized { + return Some(0); + } + let record = state.users.get(user)?; + let effective = record.effective()?; + (effective.enabled && effective.credential_id == credential_id) + .then_some(record.incarnation) + } + + /// Starts a short publication critical section for one authenticated owner. + pub(crate) fn claim_authenticated( + self: &Arc, + user: &str, + credential_id: UserCredentialId, + ) -> Option> { + let mut state = self.state.lock(); + let incarnation = if state.initialized { + let record = state.users.get(user)?; + let effective = record.effective()?; + if !effective.enabled || effective.credential_id != credential_id { + return None; + } + record.incarnation + } else { + 0 + }; + let registration_id = state.allocate_registration_id()?; + let token = CancellationToken::new(); + let active = Arc::new(AtomicU8::new(REGISTRATION_PENDING)); + Some(UserAdmissionPublication { + state, + authority: Arc::clone(self), + user: user.to_string(), + registration_id, + incarnation, + token, + active, + registration_taken: false, + }) + } + + /// Registers a legacy owner when no credential snapshot is available. + pub(crate) fn register_legacy(self: &Arc, user: &str) -> Option { + let credential_id = { + let state = self.state.lock(); + if !state.initialized { + [0; 16] + } else { + state.users.get(user)?.effective()?.credential_id + } + }; + let mut publication = self.claim_authenticated(user, credential_id)?; + let registration = publication.take_registration()?; + publication.commit(); + Some(registration) + } + + /// Cancels all current owners without changing admission policy. + pub(crate) fn cancel_user_owners(&self, user: &str) -> usize { + let tokens = self.state.lock().owner_tokens(user); + let count = tokens.len(); + for token in tokens { + token.cancel(); + } + count + } + + fn unregister(&self, user: &str, registration_id: u64, incarnation: UserIncarnation) { + let mut state = self.state.lock(); + let remove_user = state + .owners_by_user + .get_mut(user) + .map(|owners| { + if owners + .get(®istration_id) + .is_some_and(|owner| owner.incarnation == incarnation) + { + owners.remove(®istration_id); + } + owners.is_empty() + }) + .unwrap_or(false); + if remove_user { + state.owners_by_user.remove(user); + } + } +} + +/// Derives the stable credential identity from one decoded MTProxy secret. +pub(crate) fn credential_id(secret: &[u8; 16]) -> UserCredentialId { + let digest = sha256(secret); + let mut id = [0; 16]; + id.copy_from_slice(&digest[..16]); + id +} + +/// Decodes one configured secret and derives its credential identity. +pub(crate) fn credential_id_from_hex(secret: &str) -> Option { + let decoded = hex::decode(secret).ok()?; + let secret: [u8; 16] = decoded.try_into().ok()?; + Some(credential_id(&secret)) +} + +fn cancel_owners(cancellations: Vec<(String, Vec)>) -> Vec<(String, usize)> { + cancellations + .into_iter() + .map(|(user, tokens)| { + let count = tokens.len(); + for token in tokens { + token.cancel(); + } + (user, count) + }) + .collect() +} + +#[cfg(test)] +#[path = "user_admission/tests.rs"] +mod tests; diff --git a/src/proxy/user_admission/registration.rs b/src/proxy/user_admission/registration.rs new file mode 100644 index 0000000..b69a61d --- /dev/null +++ b/src/proxy/user_admission/registration.rs @@ -0,0 +1,112 @@ +use std::sync::Arc; +use std::sync::atomic::{AtomicU8, Ordering}; + +use parking_lot::MutexGuard; +use tokio_util::sync::CancellationToken; + +use super::*; + +/// Authority lock retained until the caller publishes its owned object. +pub(crate) struct UserAdmissionPublication<'a> { + /// Locked authority state that linearizes publication with policy mutation. + pub(super) state: MutexGuard<'a, UserAdmissionState>, + /// Process authority used by the eventual registration guard. + pub(super) authority: Arc, + /// Username owning the published lifecycle object. + pub(super) user: String, + /// Unique registration identity within the process authority. + pub(super) registration_id: u64, + /// User incarnation authenticated by this publication. + pub(super) incarnation: UserIncarnation, + /// Revocation signal shared with the lifecycle owner. + pub(super) token: CancellationToken, + /// Atomic publication state shared with the registration guard. + pub(super) active: Arc, + /// Whether ownership has already moved into a registration guard. + pub(super) registration_taken: bool, +} + +impl UserAdmissionPublication<'_> { + /// Moves the registered owner out while retaining the authority lock. + pub(crate) fn take_registration(&mut self) -> Option { + if self.registration_taken { + return None; + } + self.registration_taken = true; + Some(UserSessionRegistration { + authority: Arc::clone(&self.authority), + user: self.user.clone(), + registration_id: self.registration_id, + incarnation: self.incarnation, + token: self.token.clone(), + active: Arc::clone(&self.active), + }) + } + + /// Commits the owner record after the caller publishes its lifecycle object. + pub(crate) fn commit(mut self) { + if !self.registration_taken { + return; + } + if self + .active + .compare_exchange( + REGISTRATION_PENDING, + REGISTRATION_ACTIVE, + Ordering::AcqRel, + Ordering::Acquire, + ) + .is_err() + { + return; + } + self.state + .owners_by_user + .entry(self.user.clone()) + .or_default() + .insert( + self.registration_id, + RegisteredOwner { + token: self.token.clone(), + incarnation: self.incarnation, + }, + ); + } +} + +/// RAII ownership registered against one user incarnation. +#[must_use = "registered user ownership must be retained until lifecycle completion"] +pub(crate) struct UserSessionRegistration { + authority: Arc, + user: String, + registration_id: u64, + incarnation: UserIncarnation, + token: CancellationToken, + active: Arc, +} + +impl UserSessionRegistration { + /// Returns the cancellation signal for revocation or credential replacement. + pub(crate) fn token(&self) -> CancellationToken { + self.token.clone() + } + + /// Returns the immutable user incarnation owned by this registration. + pub(crate) fn incarnation(&self) -> UserIncarnation { + self.incarnation + } + + /// Returns whether revocation has cancelled this ownership. + pub(crate) fn is_cancelled(&self) -> bool { + self.token.is_cancelled() + } +} + +impl Drop for UserSessionRegistration { + fn drop(&mut self) { + if self.active.swap(REGISTRATION_DROPPED, Ordering::AcqRel) == REGISTRATION_ACTIVE { + self.authority + .unregister(&self.user, self.registration_id, self.incarnation); + } + } +} diff --git a/src/proxy/user_admission/tests.rs b/src/proxy/user_admission/tests.rs new file mode 100644 index 0000000..5144a33 --- /dev/null +++ b/src/proxy/user_admission/tests.rs @@ -0,0 +1,141 @@ +use super::*; + +fn users(secret: &str) -> HashMap { + HashMap::from([("alice".to_string(), secret.to_string())]) +} + +#[test] +fn shared_authority_rejects_registration_through_an_old_generation() { + let authority = UserAdmissionAuthority::new(); + let secret = "00112233445566778899aabbccddeeff"; + authority.apply_config(&users(secret), &HashMap::new()); + let credential = credential_id_from_hex(secret).unwrap(); + + assert!(authority.claim_authenticated("alice", credential).is_some()); + authority.stage_user("alice", secret, false).unwrap(); + + assert!(authority.claim_authenticated("alice", credential).is_none()); +} + +#[test] +fn stale_credential_cannot_cross_delete_and_recreate() { + let authority = UserAdmissionAuthority::new(); + let old_secret = "00112233445566778899aabbccddeeff"; + let new_secret = "ffeeddccbbaa99887766554433221100"; + authority.apply_config(&users(old_secret), &HashMap::new()); + let old_credential = credential_id_from_hex(old_secret).unwrap(); + let old_incarnation = authority + .authenticated_incarnation("alice", old_credential) + .unwrap(); + + authority.delete_user("alice"); + let recreated = authority.stage_user("alice", new_secret, true).unwrap(); + + assert!(recreated.incarnation > old_incarnation); + assert!( + authority + .authenticated_incarnation("alice", old_credential) + .is_none() + ); +} + +#[test] +fn stale_candidate_cannot_overwrite_newer_mutation() { + let authority = UserAdmissionAuthority::new(); + let secret = "00112233445566778899aabbccddeeff"; + authority + .activate_config_source(1, None, &users(secret), &HashMap::new()) + .unwrap(); + let candidate_epoch = authority.epoch(); + authority.stage_user("alice", secret, false).unwrap(); + + assert!( + authority + .activate_config_source(2, Some(candidate_epoch), &users(secret), &HashMap::new(),) + .is_none() + ); + assert!(!authority.is_user_enabled("alice")); + + let disabled = HashMap::from([("alice".to_string(), false)]); + assert!( + authority + .apply_config_from_source(1, &users(secret), &disabled) + .is_none() + ); + assert!( + authority + .apply_config_from_source(2, &users(secret), &disabled) + .is_some() + ); + assert!(!authority.is_user_enabled("alice")); +} + +#[test] +fn stale_generation_snapshot_cannot_reopen_reconciled_user() { + let authority = UserAdmissionAuthority::new(); + let secret = "00112233445566778899aabbccddeeff"; + authority + .activate_config_source(1, None, &users(secret), &HashMap::new()) + .unwrap(); + authority.stage_user("alice", secret, false).unwrap(); + + let disabled = HashMap::from([("alice".to_string(), false)]); + let epoch = authority.epoch(); + authority + .activate_config_source(2, Some(epoch), &users(secret), &disabled) + .unwrap(); + assert!( + authority + .apply_config_from_source(1, &users(secret), &HashMap::new()) + .is_none() + ); + + assert!(!authority.is_user_enabled("alice")); + assert_eq!(authority.stale_config_source_rejections(), 1); +} + +#[test] +fn registration_dropped_before_publication_cannot_leave_an_owner() { + let authority = UserAdmissionAuthority::new(); + let secret = "00112233445566778899aabbccddeeff"; + authority.apply_config(&users(secret), &HashMap::new()); + let credential = credential_id_from_hex(secret).unwrap(); + let mut publication = authority.claim_authenticated("alice", credential).unwrap(); + let registration = publication.take_registration().unwrap(); + + drop(registration); + publication.commit(); + + assert_eq!(authority.cancel_user_owners("alice"), 0); +} + +#[test] +fn quota_identity_follows_credential_rotation_and_recreation() { + let quota_store = Arc::new(QuotaStore::default()); + let authority = UserAdmissionAuthority::new_with_quota_store(quota_store.clone()); + let old_secret = "00112233445566778899aabbccddeeff"; + let new_secret = "ffeeddccbbaa99887766554433221100"; + authority.apply_config(&users(old_secret), &HashMap::new()); + let old_incarnation = authority + .authenticated_incarnation("alice", credential_id_from_hex(old_secret).unwrap()) + .unwrap(); + let old_quota = quota_store.handle_exact("alice", old_incarnation).unwrap(); + old_quota.charge(40); + + let rotated = authority.stage_user("alice", new_secret, true).unwrap(); + let rotated_quota = quota_store + .handle_exact("alice", rotated.incarnation) + .unwrap(); + old_quota.charge(20); + + assert_eq!(rotated_quota.used(), 40); + authority.delete_user("alice"); + let recreated = authority.stage_user("alice", old_secret, true).unwrap(); + assert_eq!( + quota_store + .handle_exact("alice", recreated.incarnation) + .unwrap() + .used(), + 0 + ); +} diff --git a/src/proxy/user_connection_authority.rs b/src/proxy/user_connection_authority.rs new file mode 100644 index 0000000..373f99d --- /dev/null +++ b/src/proxy/user_connection_authority.rs @@ -0,0 +1,123 @@ +use std::sync::Arc; + +use dashmap::DashMap; +use dashmap::mapref::entry::Entry; + +/// Process-wide per-user connection admission shared by runtime generations. +#[derive(Default)] +pub(crate) struct UserConnectionAuthority { + active: DashMap, +} + +/// Owns one exact connection slot until the authenticated connection exits. +#[must_use = "connection permits must be retained for the connection lifetime"] +pub(crate) struct UserConnectionPermit { + authority: Arc, + user: String, +} + +impl UserConnectionAuthority { + /// Acquires a slot without consulting optional telemetry state. + pub(crate) fn try_acquire( + self: &Arc, + user: &str, + limit: Option, + ) -> Option { + match self.active.entry(user.to_string()) { + Entry::Occupied(mut entry) => { + if limit.is_some_and(|max| *entry.get() >= max) { + return None; + } + let next = entry.get().checked_add(1)?; + *entry.get_mut() = next; + } + Entry::Vacant(entry) => { + if limit == Some(0) { + return None; + } + entry.insert(1); + } + } + Some(UserConnectionPermit { + authority: Arc::clone(self), + user: user.to_string(), + }) + } + + /// Returns the authoritative active connection count for one username. + pub(crate) fn active(&self, user: &str) -> u64 { + self.active.get(user).map(|entry| *entry).unwrap_or(0) + } + + #[cfg(test)] + fn tracked_users(&self) -> usize { + self.active.len() + } +} + +impl Drop for UserConnectionPermit { + fn drop(&mut self) { + let Entry::Occupied(mut entry) = self.authority.active.entry(self.user.clone()) else { + debug_assert!(false, "connection permit owner entry disappeared"); + return; + }; + debug_assert!(*entry.get() > 0, "connection permit counter underflow"); + if *entry.get() <= 1 { + entry.remove(); + } else { + *entry.get_mut() -= 1; + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Barrier; + + #[test] + fn permit_drop_releases_and_removes_zero_entry() { + let authority = Arc::new(UserConnectionAuthority::default()); + let permit = authority.try_acquire("alice", Some(1)).unwrap(); + assert_eq!(authority.active("alice"), 1); + assert!(authority.try_acquire("alice", Some(1)).is_none()); + + drop(permit); + + assert_eq!(authority.active("alice"), 0); + assert_eq!(authority.tracked_users(), 0); + } + + #[test] + fn concurrent_acquire_never_exceeds_limit() { + const CONTENDERS: usize = 64; + const LIMIT: u64 = 7; + + let authority = Arc::new(UserConnectionAuthority::default()); + let barrier = Arc::new(Barrier::new(CONTENDERS + 1)); + let (permit_tx, permit_rx) = std::sync::mpsc::channel(); + let mut threads = Vec::with_capacity(CONTENDERS); + for _ in 0..CONTENDERS { + let authority = Arc::clone(&authority); + let barrier = Arc::clone(&barrier); + let permit_tx = permit_tx.clone(); + threads.push(std::thread::spawn(move || { + barrier.wait(); + permit_tx + .send(authority.try_acquire("alice", Some(LIMIT))) + .unwrap(); + })); + } + drop(permit_tx); + barrier.wait(); + for thread in threads { + thread.join().unwrap(); + } + let permits = permit_rx.into_iter().flatten().collect::>(); + + assert_eq!(permits.len() as u64, LIMIT); + assert_eq!(authority.active("alice"), LIMIT); + drop(permits); + assert_eq!(authority.active("alice"), 0); + } +} diff --git a/src/quota_state.rs b/src/quota_state.rs index bf3149b..2db71f0 100644 --- a/src/quota_state.rs +++ b/src/quota_state.rs @@ -1,11 +1,9 @@ use std::collections::{BTreeMap, BTreeSet}; -use std::io::Write; use std::path::{Path, PathBuf}; use std::sync::Arc; use std::time::{SystemTime, UNIX_EPOCH}; use serde::{Deserialize, Serialize}; -use tokio::io::AsyncReadExt; use tokio::sync::Mutex; use tracing::{info, warn}; @@ -106,34 +104,30 @@ impl QuotaStateOwner { used_bytes: 0, last_reset_epoch_secs, }; + let reset_target = self.store.current_or_legacy_handle(user); let state = self.state_for_users(configured_users, Some((user, prospective.clone()))); let path = self.path.clone(); - let store = Arc::clone(&self.store); - let user = user.to_string(); let task = tokio::task::spawn_blocking(move || { let _guard = guard; write_state_file_blocking(&path, &state)?; - Ok(store.reset(&user, last_reset_epoch_secs)) + Ok(reset_target.reset(last_reset_epoch_secs)) }); wait_for_blocking_io(task).await } - /// Removes a deleted user's persisted and in-memory quota ownership. + /// Removes a deleted user's persisted quota checkpoint. pub(crate) async fn remove_user( &self, configured_users: &BTreeSet, user: &str, ) -> std::io::Result<()> { let guard = Arc::clone(&self.mutation).lock_owned().await; + debug_assert!(!configured_users.contains(user)); let state = self.state_for_users(configured_users, None); let path = self.path.clone(); - let store = Arc::clone(&self.store); - let user = user.to_string(); let task = tokio::task::spawn_blocking(move || { let _guard = guard; - let persisted = write_state_file_blocking(&path, &state); - store.remove(&user); - persisted + write_state_file_blocking(&path, &state) }); wait_for_blocking_io(task).await } @@ -173,27 +167,42 @@ fn now_epoch_secs() -> u64 { } async fn read_state_file(path: &Path) -> std::io::Result> { - let file = match tokio::fs::File::open(path).await { - Ok(file) => file, + #[cfg(unix)] + let payload = match crate::util::secure_fs::read_regular_limited_async( + path.to_path_buf(), + QUOTA_STATE_MAX_BYTES as usize, + ) + .await + { + Ok(payload) => payload, Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None), Err(error) => return Err(error), }; - if file.metadata().await?.len() > QUOTA_STATE_MAX_BYTES { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidData, - "quota state file exceeds the 16 MiB limit", - )); - } - let mut payload = Vec::new(); - file.take(QUOTA_STATE_MAX_BYTES.saturating_add(1)) - .read_to_end(&mut payload) - .await?; - if payload.len() as u64 > QUOTA_STATE_MAX_BYTES { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidData, - "quota state file grew beyond the 16 MiB limit while reading", - )); - } + #[cfg(not(unix))] + let payload = { + let file = match tokio::fs::File::open(path).await { + Ok(file) => file, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None), + Err(error) => return Err(error), + }; + if file.metadata().await?.len() > QUOTA_STATE_MAX_BYTES { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "quota state file exceeds the 16 MiB limit", + )); + } + let mut payload = Vec::new(); + file.take(QUOTA_STATE_MAX_BYTES.saturating_add(1)) + .read_to_end(&mut payload) + .await?; + if payload.len() as u64 > QUOTA_STATE_MAX_BYTES { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "quota state file grew beyond the 16 MiB limit while reading", + )); + } + payload + }; let state = serde_json::from_slice(&payload).map_err(|error| { std::io::Error::new( std::io::ErrorKind::InvalidData, @@ -217,11 +226,6 @@ async fn wait_for_blocking_io( } fn write_state_file_blocking(path: &Path, state: &QuotaStateFile) -> std::io::Result<()> { - let parent = path - .parent() - .filter(|parent| !parent.as_os_str().is_empty()) - .unwrap_or_else(|| Path::new(".")); - std::fs::create_dir_all(parent)?; let mut payload = serde_json::to_vec_pretty(state)?; payload.push(b'\n'); if payload.len() as u64 > QUOTA_STATE_MAX_BYTES { @@ -231,45 +235,60 @@ fn write_state_file_blocking(path: &Path, state: &QuotaStateFile) -> std::io::Re )); } - let mut last_collision = None; - for _ in 0..8 { - let tmp_path = path.with_extension(format!( - "tmp.{}.{}", - std::process::id(), - rand::random::() - )); - let mut file = match std::fs::OpenOptions::new() - .write(true) - .create_new(true) - .open(&tmp_path) - { - Ok(file) => file, - Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => { - last_collision = Some(error); - continue; - } - Err(error) => return Err(error), - }; - let result = (|| { - file.write_all(&payload)?; - file.sync_all()?; - drop(file); - std::fs::rename(&tmp_path, path)?; - #[cfg(unix)] - std::fs::File::open(parent)?.sync_all()?; - Ok(()) - })(); - if result.is_err() { - let _ = std::fs::remove_file(&tmp_path); - } - return result; + #[cfg(unix)] + { + return crate::util::secure_fs::atomic_replace(path, &payload, 0o600); + } + #[cfg(not(unix))] + { + use std::io::Write; + + let parent = path + .parent() + .filter(|parent| !parent.as_os_str().is_empty()) + .unwrap_or_else(|| Path::new(".")); + std::fs::create_dir_all(parent)?; + + let mut last_collision = None; + for _ in 0..8 { + let tmp_path = path.with_extension(format!( + "tmp.{}.{}", + std::process::id(), + rand::random::() + )); + let mut file = match std::fs::OpenOptions::new() + .write(true) + .create_new(true) + .open(&tmp_path) + { + Ok(file) => file, + Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => { + last_collision = Some(error); + continue; + } + Err(error) => return Err(error), + }; + let result = (|| { + file.write_all(&payload)?; + file.sync_all()?; + drop(file); + std::fs::rename(&tmp_path, path)?; + #[cfg(unix)] + std::fs::File::open(parent)?.sync_all()?; + Ok(()) + })(); + if result.is_err() { + let _ = std::fs::remove_file(&tmp_path); + } + return result; + } + Err(last_collision.unwrap_or_else(|| { + std::io::Error::new( + std::io::ErrorKind::AlreadyExists, + "failed to allocate a unique quota checkpoint temporary file", + ) + })) } - Err(last_collision.unwrap_or_else(|| { - std::io::Error::new( - std::io::ErrorKind::AlreadyExists, - "failed to allocate a unique quota checkpoint temporary file", - ) - })) } fn quota_user_state(quota: UserQuotaSnapshot) -> QuotaUserState { diff --git a/src/slot_budget.rs b/src/slot_budget.rs new file mode 100644 index 0000000..737b316 --- /dev/null +++ b/src/slot_budget.rs @@ -0,0 +1,131 @@ +use std::sync::atomic::{AtomicUsize, Ordering}; + +/// Exact lock-free admission budget for bounded concurrent registries. +pub(crate) struct SlotBudget { + capacity: usize, + used: AtomicUsize, +} + +impl SlotBudget { + /// Creates an empty budget with the given hard capacity. + pub(crate) const fn new(capacity: usize) -> Self { + Self { + capacity, + used: AtomicUsize::new(0), + } + } + + /// Reserves one slot or returns `None` when the hard capacity is exhausted. + pub(crate) fn try_acquire(&self) -> Option> { + let mut current = self.used.load(Ordering::Acquire); + loop { + if current >= self.capacity { + return None; + } + match self.used.compare_exchange_weak( + current, + current + 1, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => { + return Some(SlotLease { + budget: self, + armed: true, + }); + } + Err(actual) => current = actual, + } + } + } + + /// Releases one slot after its registry entry has been removed. + pub(crate) fn release(&self) { + self.release_many(1); + } + + /// Releases a known number of slots after a batch registry removal completes. + pub(crate) fn release_many(&self, amount: usize) { + if amount == 0 { + return; + } + let released = self + .used + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| { + current.checked_sub(amount) + }); + debug_assert!( + released.is_ok(), + "slot budget release must match acquisitions" + ); + } + + /// Returns the exact number of currently committed or reserved slots. + pub(crate) fn used(&self) -> usize { + self.used.load(Ordering::Acquire) + } +} + +/// Provisional slot ownership that rolls back unless committed to a registry entry. +pub(crate) struct SlotLease<'a> { + budget: &'a SlotBudget, + armed: bool, +} + +impl SlotLease<'_> { + /// Transfers the reserved slot to the registry entry being published. + pub(crate) fn commit(mut self) { + self.armed = false; + } +} + +impl Drop for SlotLease<'_> { + fn drop(&mut self) { + if self.armed { + self.budget.release(); + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + + use super::*; + + #[test] + fn parallel_reservations_never_exceed_the_hard_capacity() { + const CAPACITY: usize = 127; + const ATTEMPTS: usize = 10_000; + + let budget = Arc::new(SlotBudget::new(CAPACITY)); + let admitted = Arc::new(AtomicUsize::new(0)); + std::thread::scope(|scope| { + for worker in 0..16 { + let budget = Arc::clone(&budget); + let admitted = Arc::clone(&admitted); + scope.spawn(move || { + for attempt in (worker..ATTEMPTS).step_by(16) { + if let Some(lease) = budget.try_acquire() { + admitted.fetch_add(1, Ordering::Relaxed); + lease.commit(); + } + std::hint::black_box(attempt); + } + }); + } + }); + + assert_eq!(admitted.load(Ordering::Relaxed), CAPACITY); + assert_eq!(budget.used(), CAPACITY); + } + + #[test] + fn dropped_provisional_lease_returns_its_slot() { + let budget = SlotBudget::new(1); + drop(budget.try_acquire().unwrap()); + + assert!(budget.try_acquire().is_some()); + } +} diff --git a/src/stats/core_counters.rs b/src/stats/core_counters.rs index 9017270..819324a 100644 --- a/src/stats/core_counters.rs +++ b/src/stats/core_counters.rs @@ -135,6 +135,38 @@ impl Stats { .store(ok, Ordering::Relaxed); } + /// Records one successful process-owned firewall reconciliation. + pub fn increment_conntrack_rule_reconcile_success_total(&self) { + if self.telemetry_core_enabled() { + self.conntrack_rule_reconcile_success_total + .fetch_add(1, Ordering::Relaxed); + } + } + + /// Records one failed process-owned firewall reconciliation. + pub fn increment_conntrack_rule_reconcile_error_total(&self) { + if self.telemetry_core_enabled() { + self.conntrack_rule_reconcile_error_total + .fetch_add(1, Ordering::Relaxed); + } + } + + /// Records one successful firewall transaction rollback. + pub fn increment_conntrack_rule_rollback_success_total(&self) { + if self.telemetry_core_enabled() { + self.conntrack_rule_rollback_success_total + .fetch_add(1, Ordering::Relaxed); + } + } + + /// Records one rollback failure that leaves applied firewall state unknown. + pub fn increment_conntrack_rule_rollback_error_total(&self) { + if self.telemetry_core_enabled() { + self.conntrack_rule_rollback_error_total + .fetch_add(1, Ordering::Relaxed); + } + } + pub fn increment_conntrack_delete_attempt_total(&self) { if self.telemetry_core_enabled() { self.conntrack_delete_attempt_total diff --git a/src/stats/core_getters.rs b/src/stats/core_getters.rs index c1730c1..01d4437 100644 --- a/src/stats/core_getters.rs +++ b/src/stats/core_getters.rs @@ -84,6 +84,26 @@ impl Stats { pub fn get_conntrack_rule_apply_ok(&self) -> bool { self.conntrack_rule_apply_ok_gauge.load(Ordering::Relaxed) } + /// Returns successful process-owned firewall reconciliations. + pub fn get_conntrack_rule_reconcile_success_total(&self) -> u64 { + self.conntrack_rule_reconcile_success_total + .load(Ordering::Relaxed) + } + /// Returns failed process-owned firewall reconciliations. + pub fn get_conntrack_rule_reconcile_error_total(&self) -> u64 { + self.conntrack_rule_reconcile_error_total + .load(Ordering::Relaxed) + } + /// Returns successful firewall transaction rollbacks. + pub fn get_conntrack_rule_rollback_success_total(&self) -> u64 { + self.conntrack_rule_rollback_success_total + .load(Ordering::Relaxed) + } + /// Returns rollback failures that left applied firewall state unknown. + pub fn get_conntrack_rule_rollback_error_total(&self) -> u64 { + self.conntrack_rule_rollback_error_total + .load(Ordering::Relaxed) + } pub fn get_conntrack_delete_attempt_total(&self) -> u64 { self.conntrack_delete_attempt_total.load(Ordering::Relaxed) } diff --git a/src/stats/mod.rs b/src/stats/mod.rs index 170ad08..036bec9 100644 --- a/src/stats/mod.rs +++ b/src/stats/mod.rs @@ -21,12 +21,14 @@ use std::sync::Arc; use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU64, AtomicUsize, Ordering}; use std::time::Instant; -pub(crate) use self::quota_store::QuotaStore; +pub(crate) use self::quota_store::{QuotaReservation, QuotaStore, UserQuotaHandle}; #[allow(unused_imports)] pub use self::replay::{ReplayChecker, ReplayStats}; use self::telemetry::TelemetryPolicy; pub use self::tls_fingerprints::TlsFingerprintSnapshotRow; +pub(crate) use self::users::UserConnectionObservation; use crate::config::MeWriterPickMode; +use crate::proxy::user_connection_authority::UserConnectionAuthority; const ME_HANDSHAKE_ERROR_CODE_MAX: usize = 64; @@ -148,6 +150,10 @@ pub struct Stats { conntrack_pressure_active_gauge: AtomicBool, conntrack_event_queue_depth_gauge: AtomicU64, conntrack_rule_apply_ok_gauge: AtomicBool, + conntrack_rule_reconcile_success_total: AtomicU64, + conntrack_rule_reconcile_error_total: AtomicU64, + conntrack_rule_rollback_success_total: AtomicU64, + conntrack_rule_rollback_error_total: AtomicU64, conntrack_delete_attempt_total: AtomicU64, conntrack_delete_success_total: AtomicU64, conntrack_delete_not_found_total: AtomicU64, @@ -351,6 +357,7 @@ pub struct Stats { tls_fingerprints: tls_fingerprints::TlsFingerprintCollector, user_stats: DashMap>, quota_store: Arc, + connection_authority: Arc, user_stats_last_cleanup_epoch_secs: AtomicU64, start_time: parking_lot::RwLock>, } @@ -392,11 +399,6 @@ impl UserStats { self.quota.used() } - #[inline] - pub(crate) fn refund_quota(&self, bytes: u64) { - self.quota.refund(bytes); - } - /// Attempts one CAS reservation step against the quota counter. /// /// Callers control retry/yield policy. This primitive intentionally does @@ -404,18 +406,43 @@ impl UserStats { /// with their own contention strategy. #[inline] pub fn quota_try_reserve(&self, bytes: u64, limit: u64) -> Result { + self.quota + .try_reserve(bytes, limit) + .map(QuotaReservation::commit) + } + + /// Reserves quota until a direct I/O attempt is settled. + #[inline] + pub(crate) fn quota_reserve( + &self, + bytes: u64, + limit: u64, + ) -> Result { self.quota.try_reserve(bytes, limit) } } impl Stats { pub fn new() -> Self { - Self::with_quota_store(Arc::new(QuotaStore::default())) + Self::with_process_authorities( + Arc::new(QuotaStore::default()), + Arc::new(UserConnectionAuthority::default()), + ) } + #[cfg(test)] pub(crate) fn with_quota_store(quota_store: Arc) -> Self { + Self::with_process_authorities(quota_store, Arc::new(UserConnectionAuthority::default())) + } + + /// Creates generation telemetry around process-owned enforcement authorities. + pub(crate) fn with_process_authorities( + quota_store: Arc, + connection_authority: Arc, + ) -> Self { let stats = Self { quota_store, + connection_authority, ..Self::default() }; stats.apply_telemetry_policy(TelemetryPolicy::default()); @@ -423,6 +450,17 @@ impl Stats { *stats.start_time.write() = Some(Instant::now()); stats } + + /// Returns the process-scoped quota authority for test runtime construction. + #[cfg(test)] + pub(crate) fn quota_store(&self) -> Arc { + Arc::clone(&self.quota_store) + } + + /// Returns process-owned per-user connection admission. + pub(crate) fn connection_authority(&self) -> Arc { + Arc::clone(&self.connection_authority) + } } #[cfg(test)] diff --git a/src/stats/quota_store.rs b/src/stats/quota_store.rs index e338b76..48e95b0 100644 --- a/src/stats/quota_store.rs +++ b/src/stats/quota_store.rs @@ -2,139 +2,418 @@ use std::collections::HashMap; use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; +use arc_swap::ArcSwap; use dashmap::DashMap; +use parking_lot::Mutex; use super::{QuotaReserveError, UserQuotaSnapshot}; +use crate::proxy::user_admission::UserIncarnation; /// Process-scoped per-user quota accounting shared by runtime generations. #[derive(Default)] pub struct QuotaStore { - users: DashMap>, + users: DashMap>, +} + +struct QuotaUserSlot { + state: Mutex, } -/// Atomic quota state for one configured user. #[derive(Default)] +struct QuotaSlotState { + high_water: UserIncarnation, + current: Option, + startup_seed: Option, +} + +struct QuotaAccount { + incarnation: UserIncarnation, + counters: Arc, +} + +/// Exact quota ownership pinned to one authenticated user incarnation. +#[derive(Clone)] +pub(crate) struct UserQuotaHandle { + incarnation: UserIncarnation, + counters: Arc, +} + +/// Atomically replaceable quota state for one configured user. pub(crate) struct UserQuotaCounters { + generation: ArcSwap, +} + +struct QuotaGeneration { used_bytes: AtomicU64, - last_reset_epoch_secs: AtomicU64, + last_reset_epoch_secs: u64, +} + +/// Owns a quota debit until the corresponding I/O outcome is known. +#[must_use = "quota reservations must be committed or settled"] +pub(crate) struct QuotaReservation { + generation: Arc, + reserved_bytes: u64, + total_after_reserve: u64, } impl QuotaStore { - pub(crate) fn user(&self, user: &str) -> Arc { + fn slot(&self, user: &str) -> Arc { if let Some(existing) = self.users.get(user) { return Arc::clone(existing.value()); } Arc::clone( self.users .entry(user.to_string()) - .or_insert_with(|| Arc::new(UserQuotaCounters::default())) + .or_insert_with(|| { + Arc::new(QuotaUserSlot { + state: Mutex::new(QuotaSlotState::default()), + }) + }) .value(), ) } - pub(crate) fn used(&self, user: &str) -> u64 { - self.users.get(user).map(|state| state.used()).unwrap_or(0) - } - - pub(crate) fn load(&self, user: &str, used_bytes: u64, last_reset_epoch_secs: u64) { - let state = self.user(user); - state.used_bytes.store(used_bytes, Ordering::Relaxed); - state - .last_reset_epoch_secs - .store(last_reset_epoch_secs, Ordering::Relaxed); - } - - pub(crate) fn reset(&self, user: &str, now_epoch_secs: u64) -> UserQuotaSnapshot { - let state = self.user(user); - state.used_bytes.store(0, Ordering::Relaxed); - state - .last_reset_epoch_secs - .store(now_epoch_secs, Ordering::Relaxed); - UserQuotaSnapshot { + pub(crate) fn current_or_legacy_handle(&self, user: &str) -> UserQuotaHandle { + let slot = self.slot(user); + let mut state = slot.state.lock(); + if let Some(account) = &state.current { + return UserQuotaHandle { + incarnation: account.incarnation, + counters: Arc::clone(&account.counters), + }; + } + let seed = state.startup_seed.take().unwrap_or(UserQuotaSnapshot { used_bytes: 0, - last_reset_epoch_secs: now_epoch_secs, + last_reset_epoch_secs: 0, + }); + let counters = Arc::new(UserQuotaCounters::from_snapshot(&seed)); + let incarnation = state.high_water; + state.current = Some(QuotaAccount { + incarnation, + counters: Arc::clone(&counters), + }); + UserQuotaHandle { + incarnation, + counters, } } + pub(crate) fn user(&self, user: &str) -> Arc { + self.current_or_legacy_handle(user).counters + } + + /// Returns the quota account owned by the exact current incarnation. + pub(crate) fn handle_exact( + &self, + user: &str, + incarnation: UserIncarnation, + ) -> Option { + let slot = self.users.get(user)?; + let state = slot.state.lock(); + let account = state.current.as_ref()?; + (account.incarnation == incarnation).then(|| UserQuotaHandle { + incarnation, + counters: Arc::clone(&account.counters), + }) + } + + /// Creates a quota account for a new username lifetime without inheriting a retired account. + pub(crate) fn activate_fresh(&self, user: &str, incarnation: UserIncarnation) { + let slot = self.slot(user); + let mut state = slot.state.lock(); + if incarnation <= state.high_water { + return; + } + let snapshot = if state.high_water == 0 { + state + .current + .as_ref() + .map(|account| account.counters.snapshot()) + .or_else(|| state.startup_seed.take()) + } else { + state.startup_seed = None; + None + } + .unwrap_or(UserQuotaSnapshot { + used_bytes: 0, + last_reset_epoch_secs: 0, + }); + state.high_water = state.high_water.max(incarnation); + state.current = Some(QuotaAccount { + incarnation, + counters: Arc::new(UserQuotaCounters::from_snapshot(&snapshot)), + }); + } + + /// Advances a credential incarnation while preserving usage captured at the transition. + pub(crate) fn advance_preserving_usage(&self, user: &str, incarnation: UserIncarnation) { + let slot = self.slot(user); + let mut state = slot.state.lock(); + if incarnation <= state.high_water { + return; + } + let snapshot = state + .current + .as_ref() + .map(|account| account.counters.snapshot()) + .or_else(|| state.startup_seed.take()) + .unwrap_or(UserQuotaSnapshot { + used_bytes: 0, + last_reset_epoch_secs: 0, + }); + state.high_water = incarnation; + state.current = Some(QuotaAccount { + incarnation, + counters: Arc::new(UserQuotaCounters::from_snapshot(&snapshot)), + }); + } + + /// Retires quota ownership without affecting a newer incarnation. + pub(crate) fn retire_through(&self, user: &str, incarnation: UserIncarnation) { + let slot = self.slot(user); + let mut state = slot.state.lock(); + if incarnation < state.high_water { + return; + } + state.high_water = incarnation; + state.current = None; + state.startup_seed = None; + } + + pub(crate) fn used(&self, user: &str) -> u64 { + self.users + .get(user) + .and_then(|slot| { + let state = slot.state.lock(); + state + .current + .as_ref() + .map(|account| account.counters.used()) + .or_else(|| state.startup_seed.as_ref().map(|seed| seed.used_bytes)) + }) + .unwrap_or(0) + } + + pub(crate) fn load(&self, user: &str, used_bytes: u64, last_reset_epoch_secs: u64) { + let slot = self.slot(user); + let mut state = slot.state.lock(); + let snapshot = UserQuotaSnapshot { + used_bytes, + last_reset_epoch_secs, + }; + if let Some(account) = &state.current { + account.replace_from_snapshot(&snapshot); + } else { + state.startup_seed = Some(snapshot); + } + } + + pub(crate) fn reset(&self, user: &str, now_epoch_secs: u64) -> UserQuotaSnapshot { + self.current_or_legacy_handle(user).reset(now_epoch_secs) + } + pub(crate) fn remove(&self, user: &str) { - self.users.remove(user); + let slot = self.slot(user); + let mut state = slot.state.lock(); + state.high_water = state.high_water.saturating_add(1); + state.current = None; + state.startup_seed = None; } pub(crate) fn snapshot(&self) -> HashMap { let mut out = HashMap::new(); for entry in self.users.iter() { - let state = entry.value(); - let used_bytes = state.used(); - let last_reset_epoch_secs = state.last_reset_epoch_secs.load(Ordering::Relaxed); - if used_bytes == 0 && last_reset_epoch_secs == 0 { + let state = entry.value().state.lock(); + let snapshot = if let Some(account) = state.current.as_ref() { + account.counters.snapshot() + } else if let Some(seed) = state.startup_seed.as_ref() { + seed.clone() + } else { + continue; + }; + if snapshot.used_bytes == 0 && snapshot.last_reset_epoch_secs == 0 { continue; } - out.insert( - entry.key().clone(), - UserQuotaSnapshot { - used_bytes, - last_reset_epoch_secs, - }, - ); + out.insert(entry.key().clone(), snapshot); } out } } +impl Default for UserQuotaCounters { + fn default() -> Self { + Self { + generation: ArcSwap::from_pointee(QuotaGeneration { + used_bytes: AtomicU64::new(0), + last_reset_epoch_secs: 0, + }), + } + } +} + impl UserQuotaCounters { + fn from_snapshot(snapshot: &UserQuotaSnapshot) -> Self { + Self { + generation: ArcSwap::from_pointee(QuotaGeneration { + used_bytes: AtomicU64::new(snapshot.used_bytes), + last_reset_epoch_secs: snapshot.last_reset_epoch_secs, + }), + } + } + + fn snapshot(&self) -> UserQuotaSnapshot { + let generation = self.generation.load_full(); + UserQuotaSnapshot { + used_bytes: generation.used_bytes.load(Ordering::Relaxed), + last_reset_epoch_secs: generation.last_reset_epoch_secs, + } + } + + fn replace(&self, used_bytes: u64, last_reset_epoch_secs: u64) { + self.generation.store(Arc::new(QuotaGeneration { + used_bytes: AtomicU64::new(used_bytes), + last_reset_epoch_secs, + })); + } + #[inline] pub(crate) fn used(&self) -> u64 { - self.used_bytes.load(Ordering::Relaxed) + self.generation.load().used_bytes.load(Ordering::Relaxed) } #[inline] pub(crate) fn charge(&self, bytes: u64) -> u64 { - self.used_bytes + self.generation + .load_full() + .used_bytes .fetch_add(bytes, Ordering::Relaxed) .saturating_add(bytes) } #[inline] - pub(crate) fn refund(&self, bytes: u64) { - let mut current = self.used_bytes.load(Ordering::Relaxed); - loop { - let next = current.saturating_sub(bytes); - match self.used_bytes.compare_exchange_weak( - current, - next, - Ordering::Relaxed, - Ordering::Relaxed, - ) { - Ok(_) => return, - Err(observed) => current = observed, - } - } - } - - #[inline] - pub(crate) fn try_reserve(&self, bytes: u64, limit: u64) -> Result { - let current = self.used_bytes.load(Ordering::Relaxed); + pub(crate) fn try_reserve( + &self, + bytes: u64, + limit: u64, + ) -> Result { + let generation = self.generation.load_full(); + let current = generation.used_bytes.load(Ordering::Relaxed); if bytes > limit.saturating_sub(current) { return Err(QuotaReserveError::LimitExceeded); } let next = current.saturating_add(bytes); - match self.used_bytes.compare_exchange_weak( + match generation.used_bytes.compare_exchange_weak( current, next, Ordering::Relaxed, Ordering::Relaxed, ) { - Ok(_) => Ok(next), + Ok(_) => Ok(QuotaReservation { + generation, + reserved_bytes: bytes, + total_after_reserve: next, + }), Err(_) => Err(QuotaReserveError::Contended), } } } +impl QuotaAccount { + fn replace_from_snapshot(&self, snapshot: &UserQuotaSnapshot) { + self.counters + .replace(snapshot.used_bytes, snapshot.last_reset_epoch_secs); + } +} + +impl UserQuotaHandle { + /// Returns the immutable incarnation owned by this handle. + pub(crate) fn incarnation(&self) -> UserIncarnation { + self.incarnation + } + + #[inline] + pub(crate) fn used(&self) -> u64 { + self.counters.used() + } + + #[inline] + pub(crate) fn charge(&self, bytes: u64) -> u64 { + self.counters.charge(bytes) + } + + #[inline] + pub(crate) fn try_reserve( + &self, + bytes: u64, + limit: u64, + ) -> Result { + self.counters.try_reserve(bytes, limit) + } + + /// Resets only the quota incarnation captured by this handle. + pub(crate) fn reset(&self, now_epoch_secs: u64) -> UserQuotaSnapshot { + self.counters.replace(0, now_epoch_secs); + UserQuotaSnapshot { + used_bytes: 0, + last_reset_epoch_secs: now_epoch_secs, + } + } +} + +impl QuotaReservation { + /// Returns the number of bytes held by this reservation. + pub(crate) fn reserved_bytes(&self) -> u64 { + self.reserved_bytes + } + + /// Commits the complete reservation and returns the generation-local total. + pub(crate) fn commit(mut self) -> u64 { + self.reserved_bytes = 0; + self.total_after_reserve + } + + /// Commits part of the reservation and refunds the remainder. + pub(crate) fn settle(mut self, committed_bytes: u64) { + let committed_bytes = committed_bytes.min(self.reserved_bytes); + refund_generation( + self.generation.as_ref(), + self.reserved_bytes - committed_bytes, + ); + self.reserved_bytes = 0; + } +} + +impl Drop for QuotaReservation { + fn drop(&mut self) { + refund_generation(self.generation.as_ref(), self.reserved_bytes); + } +} + +fn refund_generation(generation: &QuotaGeneration, bytes: u64) { + if bytes == 0 { + return; + } + let mut current = generation.used_bytes.load(Ordering::Relaxed); + loop { + let next = current.saturating_sub(bytes); + match generation.used_bytes.compare_exchange_weak( + current, + next, + Ordering::Relaxed, + Ordering::Relaxed, + ) { + Ok(_) => return, + Err(observed) => current = observed, + } + } +} + #[cfg(test)] mod tests { use super::*; use crate::stats::Stats; + use std::sync::Barrier; #[test] fn quota_counters_are_shared_across_stats_generations() { @@ -148,4 +427,118 @@ mod tests { second.reset_user_quota("alice"); assert_eq!(first.get_user_quota_used("alice"), 0); } + + #[test] + fn quota_snapshot_never_combines_different_generations() { + const ITERATIONS: u64 = 10_000; + + let store = Arc::new(QuotaStore::default()); + store.load("alice", 1, 1); + let barrier = Arc::new(Barrier::new(2)); + let writer_store = Arc::clone(&store); + let writer_barrier = Arc::clone(&barrier); + let writer = std::thread::spawn(move || { + writer_barrier.wait(); + for generation in 2..=ITERATIONS { + writer_store.load("alice", generation, generation); + } + }); + + barrier.wait(); + for _ in 0..ITERATIONS { + let snapshot = store.snapshot().remove("alice").unwrap(); + assert_eq!(snapshot.used_bytes, snapshot.last_reset_epoch_secs); + } + writer.join().unwrap(); + } + + #[test] + fn repeated_old_generation_refunds_leave_new_usage_intact() { + const ITERATIONS: u64 = 10_000; + + let store = QuotaStore::default(); + let state = store.user("alice"); + for generation in 1..=ITERATIONS { + let reservation = state.try_reserve(80, 100).unwrap(); + store.load("alice", 40, generation); + drop(reservation); + assert_eq!(store.used("alice"), 40); + store.reset("alice", generation); + } + } + + #[test] + fn retired_incarnation_cannot_charge_recreated_username() { + let store = QuotaStore::default(); + store.activate_fresh("alice", 1); + let retired = store.handle_exact("alice", 1).unwrap(); + retired.charge(40); + store.retire_through("alice", 2); + store.activate_fresh("alice", 3); + + retired.charge(20); + + assert_eq!(retired.used(), 60); + assert_eq!(store.handle_exact("alice", 3).unwrap().used(), 0); + } + + #[test] + fn credential_rotation_preserves_usage_without_sharing_future_charges() { + let store = QuotaStore::default(); + store.activate_fresh("alice", 1); + let old = store.handle_exact("alice", 1).unwrap(); + old.charge(40); + store.advance_preserving_usage("alice", 2); + let current = store.handle_exact("alice", 2).unwrap(); + + old.charge(20); + + assert_eq!(old.used(), 60); + assert_eq!(current.used(), 40); + } + + #[test] + fn captured_reset_handle_cannot_reset_a_new_incarnation() { + let store = QuotaStore::default(); + store.activate_fresh("alice", 1); + let reset_target = store.handle_exact("alice", 1).unwrap(); + reset_target.charge(40); + store.advance_preserving_usage("alice", 2); + let current = store.handle_exact("alice", 2).unwrap(); + current.charge(20); + + reset_target.reset(7); + + assert_eq!(reset_target.used(), 0); + assert_eq!(current.used(), 60); + } + + #[test] + fn stale_retirement_cannot_remove_newer_quota_owner() { + let store = QuotaStore::default(); + store.activate_fresh("alice", 1); + store.retire_through("alice", 2); + store.activate_fresh("alice", 3); + + store.retire_through("alice", 2); + + assert!(store.handle_exact("alice", 3).is_some()); + } + + #[test] + fn old_reservation_refund_does_not_debit_recreated_username() { + let store = QuotaStore::default(); + store.activate_fresh("alice", 1); + let old = store.handle_exact("alice", 1).unwrap(); + let reservation = old.try_reserve(80, 100).unwrap(); + store.retire_through("alice", 2); + store.activate_fresh("alice", 3); + let current = store.handle_exact("alice", 3).unwrap(); + current.charge(50); + + drop(reservation); + + assert_eq!(old.used(), 0); + assert_eq!(current.used(), 50); + } } diff --git a/src/stats/replay.rs b/src/stats/replay.rs index c0e2257..e6fe20f 100644 --- a/src/stats/replay.rs +++ b/src/stats/replay.rs @@ -1,6 +1,6 @@ use std::borrow::Borrow; -use std::collections::VecDeque; use std::collections::hash_map::DefaultHasher; +use std::collections::{HashMap, VecDeque}; use std::hash::{Hash, Hasher}; use std::num::NonZeroUsize; use std::sync::Arc; @@ -73,7 +73,9 @@ pub struct ReplayChecker { checks: AtomicU64, hits: AtomicU64, additions: AtomicU64, + capacity_rejections: AtomicU64, cleanups: AtomicU64, + next_claim_token: AtomicU64, } struct ReplayEntry { @@ -82,15 +84,24 @@ struct ReplayEntry { struct ReplayShard { cache: LruCache, + pending: HashMap, queue: VecDeque<(Instant, ReplayKey, u64)>, seq_counter: u64, capacity: usize, } +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum ReplayClaimResult { + Claimed, + Duplicate, + Capacity, +} + impl ReplayShard { fn new(cap: NonZeroUsize) -> Self { Self { cache: LruCache::new(cap), + pending: HashMap::new(), queue: VecDeque::with_capacity(cap.get()), seq_counter: 0, capacity: cap.get(), @@ -135,28 +146,112 @@ impl ReplayShard { return false; } self.cleanup(now, window); - self.cache.get(key).is_some() + self.cache.get(key).is_some() || self.pending.contains_key(key) } - fn add_owned(&mut self, key: ReplayKey, now: Instant, window: Duration) { + fn add_owned(&mut self, key: ReplayKey, now: Instant, window: Duration) -> bool { if window.is_zero() { - return; + return true; } self.cleanup(now, window); - if self.cache.peek(key.as_slice()).is_some() { - return; + if self.cache.peek(key.as_slice()).is_some() || self.pending.contains_key(key.as_slice()) { + return true; } - while self.queue.len() >= self.capacity { + while self.cache.len().saturating_add(self.pending.len()) >= self.capacity { + if self.queue.is_empty() { + return false; + } self.evict_queue_front(); } let seq = self.next_seq(); self.cache.put(key.clone(), ReplayEntry { seq }); self.queue.push_back((now, key, seq)); + true + } + + fn claim_owned( + &mut self, + key: ReplayKey, + now: Instant, + window: Duration, + token: u64, + ) -> ReplayClaimResult { + if window.is_zero() { + return ReplayClaimResult::Claimed; + } + self.cleanup(now, window); + if self.cache.peek(key.as_slice()).is_some() || self.pending.contains_key(key.as_slice()) { + return ReplayClaimResult::Duplicate; + } + while self.cache.len().saturating_add(self.pending.len()) >= self.capacity { + if self.queue.is_empty() { + return ReplayClaimResult::Capacity; + } + self.evict_queue_front(); + } + self.pending.insert(key, token); + ReplayClaimResult::Claimed + } + + fn remove_pending(&mut self, key: &[u8], token: u64) -> bool { + if self.pending.get(key).copied() != Some(token) { + return false; + } + self.pending.remove(key); + true } fn len(&self) -> usize { - self.cache.len() + self.cache.len().saturating_add(self.pending.len()) + } +} + +/// Exclusive in-flight ownership of one TLS replay digest. +#[must_use = "the claim must be committed after all TLS policy checks"] +pub(crate) struct TlsReplayClaim<'a> { + checker: &'a ReplayChecker, + shard_idx: usize, + key: Option, + token: u64, + reserved: bool, +} + +impl TlsReplayClaim<'_> { + /// Commits the claimed digest into the replay window. + pub(crate) fn commit(mut self) -> bool { + if !self.reserved { + return true; + } + let Some(key) = self.key.take() else { + return false; + }; + let mut shard = self.checker.tls_shards[self.shard_idx].lock(); + if !shard.remove_pending(key.as_slice(), self.token) { + return false; + } + if !shard.add_owned(key, Instant::now(), self.checker.tls_window) { + self.checker + .capacity_rejections + .fetch_add(1, Ordering::Relaxed); + return false; + } + self.checker.additions.fetch_add(1, Ordering::Relaxed); + self.reserved = false; + true + } +} + +impl Drop for TlsReplayClaim<'_> { + fn drop(&mut self) { + if !self.reserved { + return; + } + if let Some(key) = self.key.as_ref() { + self.checker.tls_shards[self.shard_idx] + .lock() + .remove_pending(key.as_slice(), self.token); + } } } @@ -183,7 +278,27 @@ impl ReplayChecker { checks: AtomicU64::new(0), hits: AtomicU64::new(0), additions: AtomicU64::new(0), + capacity_rejections: AtomicU64::new(0), cleanups: AtomicU64::new(0), + next_claim_token: AtomicU64::new(1), + } + } + + fn reserve_claim_token(&self) -> Option { + let mut current = self.next_claim_token.load(Ordering::Relaxed); + loop { + if current == u64::MAX { + return None; + } + match self.next_claim_token.compare_exchange_weak( + current, + current + 1, + Ordering::Relaxed, + Ordering::Relaxed, + ) { + Ok(_) => return Some(current), + Err(actual) => current = actual, + } } } @@ -207,11 +322,15 @@ impl ReplayChecker { let found = shard.check(data, now, window); if found { self.hits.fetch_add(1, Ordering::Relaxed); - } else { - shard.add_owned(owned_key, now, window); - self.additions.fetch_add(1, Ordering::Relaxed); + return true; + } + if shard.add_owned(owned_key, now, window) { + self.additions.fetch_add(1, Ordering::Relaxed); + false + } else { + self.capacity_rejections.fetch_add(1, Ordering::Relaxed); + true } - found } fn check_only_internal( @@ -231,11 +350,14 @@ impl ReplayChecker { } fn add_only(&self, data: &[u8], shards: &[Mutex], window: Duration) { - self.additions.fetch_add(1, Ordering::Relaxed); let idx = self.get_shard_idx(data); let owned_key = ReplayKey::from_slice(data); let mut shard = shards[idx].lock(); - shard.add_owned(owned_key, Instant::now(), window); + if shard.add_owned(owned_key, Instant::now(), window) { + self.additions.fetch_add(1, Ordering::Relaxed); + } else { + self.capacity_rejections.fetch_add(1, Ordering::Relaxed); + } } pub fn check_and_add_handshake(&self, data: &[u8]) -> bool { @@ -246,6 +368,43 @@ impl ReplayChecker { self.check_and_add_internal(data, &self.tls_shards, self.tls_window) } + /// Claims one TLS digest until the policy-valid handshake commits or aborts. + pub(crate) fn claim_tls_digest(&self, data: &[u8]) -> Option> { + self.checks.fetch_add(1, Ordering::Relaxed); + let shard_idx = self.get_shard_idx(data); + let key = ReplayKey::from_slice(data); + if self.tls_window.is_zero() { + return Some(TlsReplayClaim { + checker: self, + shard_idx, + key: None, + token: 0, + reserved: false, + }); + } + let token = self.reserve_claim_token()?; + let mut shard = self.tls_shards[shard_idx].lock(); + match shard.claim_owned(key.clone(), Instant::now(), self.tls_window, token) { + ReplayClaimResult::Claimed => {} + ReplayClaimResult::Duplicate => { + self.hits.fetch_add(1, Ordering::Relaxed); + return None; + } + ReplayClaimResult::Capacity => { + self.capacity_rejections.fetch_add(1, Ordering::Relaxed); + return None; + } + } + drop(shard); + Some(TlsReplayClaim { + checker: self, + shard_idx, + key: Some(key), + token, + reserved: true, + }) + } + pub fn check_handshake(&self, data: &[u8]) -> bool { self.check_and_add_handshake(data) } @@ -267,12 +426,12 @@ impl ReplayChecker { let mut total_queue_len = 0; for shard in &self.handshake_shards { let s = shard.lock(); - total_entries += s.cache.len(); + total_entries += s.len(); total_queue_len += s.queue.len(); } for shard in &self.tls_shards { let s = shard.lock(); - total_entries += s.cache.len(); + total_entries += s.len(); total_queue_len += s.queue.len(); } @@ -282,6 +441,7 @@ impl ReplayChecker { total_checks: self.checks.load(Ordering::Relaxed), total_hits: self.hits.load(Ordering::Relaxed), total_additions: self.additions.load(Ordering::Relaxed), + total_capacity_rejections: self.capacity_rejections.load(Ordering::Relaxed), total_cleanups: self.cleanups.load(Ordering::Relaxed), num_shards: self.handshake_shards.len() + self.tls_shards.len(), window_secs: self.window.as_secs(), @@ -332,6 +492,8 @@ pub struct ReplayStats { pub total_checks: u64, pub total_hits: u64, pub total_additions: u64, + /// Claims rejected because committed and pending entries exhausted a shard. + pub total_capacity_rejections: u64, pub total_cleanups: u64, pub num_shards: usize, pub window_secs: u64, @@ -354,3 +516,7 @@ impl ReplayStats { } } } + +#[cfg(test)] +#[path = "replay/tests.rs"] +mod capacity_tests; diff --git a/src/stats/replay/tests.rs b/src/stats/replay/tests.rs new file mode 100644 index 0000000..c1eb199 --- /dev/null +++ b/src/stats/replay/tests.rs @@ -0,0 +1,47 @@ +use super::*; + +#[test] +fn committed_and_pending_entries_share_one_shard_capacity() { + let capacity = NonZeroUsize::new(2).unwrap(); + let mut shard = ReplayShard::new(capacity); + let now = Instant::now(); + let window = Duration::from_secs(60); + assert_eq!( + shard.claim_owned(ReplayKey::from_slice(b"pending-a"), now, window, 1), + ReplayClaimResult::Claimed + ); + assert_eq!( + shard.claim_owned(ReplayKey::from_slice(b"pending-b"), now, window, 2), + ReplayClaimResult::Claimed + ); + + assert!(!shard.add_owned(ReplayKey::from_slice(b"committed"), now, window)); + + assert_eq!(shard.len(), capacity.get()); + assert!(shard.pending.contains_key(b"pending-a".as_slice())); + assert!(shard.pending.contains_key(b"pending-b".as_slice())); +} + +#[test] +fn pending_capacity_rejection_is_not_reported_as_a_replay_hit() { + let checker = ReplayChecker::new(64, Duration::from_secs(60)); + let first = checker + .claim_tls_digest(b"same-shard-key") + .expect("first claim must reserve the shard"); + let shard_idx = checker.get_shard_idx(b"same-shard-key"); + let rejected = (0..10_000u64) + .map(u64::to_le_bytes) + .find(|candidate| checker.get_shard_idx(candidate) == shard_idx) + .expect("bounded search must find a second key in the selected shard"); + let before = checker.stats(); + + assert!(checker.claim_tls_digest(&rejected).is_none()); + + let after = checker.stats(); + assert_eq!(after.total_hits, before.total_hits); + assert_eq!( + after.total_capacity_rejections, + before.total_capacity_rejections + 1 + ); + drop(first); +} diff --git a/src/stats/tests.rs b/src/stats/tests.rs index 650b123..f2165db 100644 --- a/src/stats/tests.rs +++ b/src/stats/tests.rs @@ -13,6 +13,33 @@ fn test_stats_shared_counters() { assert_eq!(stats.get_connects_all(), 3); } +#[test] +fn runtime_stats_share_process_connection_admission_authority() { + let quota_store = Arc::new(QuotaStore::default()); + let authority = Arc::new(UserConnectionAuthority::default()); + let first = Stats::with_process_authorities(Arc::clone("a_store), Arc::clone(&authority)); + let second = Stats::with_process_authorities(quota_store, authority); + + let permit = first + .connection_authority() + .try_acquire("alice", Some(1)) + .unwrap(); + assert!( + second + .connection_authority() + .try_acquire("alice", Some(1)) + .is_none() + ); + + drop(permit); + assert!( + second + .connection_authority() + .try_acquire("alice", Some(1)) + .is_some() + ); +} + #[test] fn test_telemetry_policy_disables_core_and_user_counters() { let stats = Stats::new(); @@ -289,6 +316,20 @@ fn test_quota_used_is_authoritative_and_independent_from_octets_telemetry() { assert_eq!(stats.get_user_quota_used(user), 7); } +#[test] +fn old_quota_reservation_refund_does_not_reduce_post_reset_usage() { + let stats = Stats::new(); + let user = "quota-reset-generation-user"; + let user_stats = stats.get_or_create_user_stats_handle(user); + let reservation = user_stats.quota_reserve(80, 100).unwrap(); + + stats.reset_user_quota(user); + stats.quota_charge_post_write(user_stats.as_ref(), 40); + drop(reservation); + + assert_eq!(stats.get_user_quota_used(user), 40); +} + #[test] fn test_cached_handle_survives_map_cleanup_until_last_drop() { let stats = Stats::new(); diff --git a/src/stats/tests/replay_checker_security_tests.rs b/src/stats/tests/replay_checker_security_tests.rs index 8e73204..e20f39f 100644 --- a/src/stats/tests/replay_checker_security_tests.rs +++ b/src/stats/tests/replay_checker_security_tests.rs @@ -1,4 +1,5 @@ use super::*; +use std::sync::{Arc, Barrier}; use std::time::Duration; #[test] @@ -78,3 +79,32 @@ fn replay_checker_stats_reflect_dual_shard_domains() { "stats should expose both shard domains (handshake + TLS)" ); } + +#[test] +fn concurrent_tls_claim_allows_exactly_one_pending_handshake() { + const WORKERS: usize = 16; + + let checker = Arc::new(ReplayChecker::new(128, Duration::from_secs(1))); + let start = Arc::new(Barrier::new(WORKERS)); + let finish = Arc::new(Barrier::new(WORKERS)); + let mut handles = Vec::with_capacity(WORKERS); + + for _ in 0..WORKERS { + let checker = Arc::clone(&checker); + let start = Arc::clone(&start); + let finish = Arc::clone(&finish); + handles.push(std::thread::spawn(move || { + start.wait(); + let claim = checker.claim_tls_digest(b"parallel-client-hello"); + finish.wait(); + claim.map(|claim| claim.commit()).unwrap_or(false) + })); + } + + let accepted = handles + .into_iter() + .map(|handle| handle.join().expect("TLS claim worker must not panic")) + .filter(|accepted| *accepted) + .count(); + assert_eq!(accepted, 1, "only one concurrent TLS claim may commit"); +} diff --git a/src/stats/tls_fingerprints.rs b/src/stats/tls_fingerprints.rs index 424aa4e..755f170 100644 --- a/src/stats/tls_fingerprints.rs +++ b/src/stats/tls_fingerprints.rs @@ -10,6 +10,7 @@ use dashmap::DashMap; use dashmap::mapref::entry::Entry; use crate::protocol::tls_fingerprint::TlsClientFingerprint; +use crate::slot_budget::SlotBudget; use super::Stats; @@ -68,15 +69,33 @@ struct TlsFingerprintEntry { bad_or_probe: AtomicU64, } -#[derive(Default)] pub struct TlsFingerprintCollector { entries: DashMap, + slots: SlotBudget, + capacity: usize, dropped_total: AtomicU64, parse_error_total: AtomicU64, last_cleanup_epoch_secs: AtomicU64, } +impl Default for TlsFingerprintCollector { + fn default() -> Self { + Self::with_capacity(MAX_TLS_FINGERPRINT_BUCKETS) + } +} + impl TlsFingerprintCollector { + fn with_capacity(capacity: usize) -> Self { + Self { + entries: DashMap::new(), + slots: SlotBudget::new(capacity), + capacity, + dropped_total: AtomicU64::new(0), + parse_error_total: AtomicU64::new(0), + last_cleanup_epoch_secs: AtomicU64::new(0), + } + } + pub fn record_observed( &self, fingerprint: &TlsClientFingerprint, @@ -228,7 +247,7 @@ impl TlsFingerprintCollector { TlsFingerprintSnapshot { retention_secs: ttl.as_secs(), - capacity: MAX_TLS_FINGERPRINT_BUCKETS, + capacity: self.capacity, dropped_total: self.dropped_total.load(Ordering::Relaxed), parse_error_total: self.parse_error_total.load(Ordering::Relaxed), by_fingerprint, @@ -286,22 +305,6 @@ impl TlsFingerprintCollector { ja4_raw: fingerprint.ja4_raw.clone(), }; - if let Some(entry) = self.entries.get(&key) { - update_entry( - entry.value(), - now_epoch_secs, - count_total, - count_auth_success, - count_bad_or_probe, - ); - return; - } - - if self.entries.len() >= MAX_TLS_FINGERPRINT_BUCKETS { - self.dropped_total.fetch_add(1, Ordering::Relaxed); - return; - } - match self.entries.entry(key) { Entry::Occupied(entry) => { update_entry( @@ -313,12 +316,17 @@ impl TlsFingerprintCollector { ); } Entry::Vacant(entry) => { + let Some(slot) = self.slots.try_acquire() else { + self.dropped_total.fetch_add(1, Ordering::Relaxed); + return; + }; entry.insert(TlsFingerprintEntry::new( now_epoch_secs, if count_total { 1 } else { 0 }, if count_auth_success { 1 } else { 0 }, if count_bad_or_probe { 1 } else { 0 }, )); + slot.commit(); } } } @@ -339,14 +347,16 @@ impl TlsFingerprintCollector { } fn cleanup(&self, now_epoch_secs: u64, ttl_secs: u64) { - if ttl_secs == 0 { - self.entries.clear(); - return; - } + let mut removed = 0usize; self.entries.retain(|_, entry| { let last_seen = entry.last_seen_epoch_secs.load(Ordering::Relaxed); - now_epoch_secs.saturating_sub(last_seen) <= ttl_secs + let retained = ttl_secs != 0 && now_epoch_secs.saturating_sub(last_seen) <= ttl_secs; + if !retained { + removed += 1; + } + retained }); + self.slots.release_many(removed); } } @@ -526,31 +536,4 @@ impl Stats { } #[cfg(test)] -mod tests { - use super::*; - - fn fp() -> TlsClientFingerprint { - TlsClientFingerprint { - ja3: "ja3".to_string(), - ja3_raw: "771,4865,,,0".to_string(), - ja4: "t13d010100_hash_hash".to_string(), - ja4_raw: "raw".to_string(), - } - } - - #[test] - fn aggregates_ip_cidr_and_user_scopes() { - let collector = TlsFingerprintCollector::default(); - let ip: IpAddr = "192.0.2.15".parse().expect("test IP parses"); - collector.record_observed(&fp(), ip, Duration::from_secs(60)); - collector.record_auth_success(&fp(), ip, "alice", Duration::from_secs(60)); - let snapshot = collector.snapshot(Duration::from_secs(60), 10); - - assert_eq!(snapshot.by_fingerprint[0].total, 1); - assert_eq!(snapshot.by_fingerprint[0].auth_success, 1); - assert_eq!(snapshot.by_ip[0].scope_key, "192.0.2.15"); - assert_eq!(snapshot.by_cidr[0].scope_key, "192.0.2.0/24"); - assert_eq!(snapshot.by_user[0].scope_key, "alice"); - assert_eq!(snapshot.by_user[0].total, 1); - } -} +mod tests; diff --git a/src/stats/tls_fingerprints/tests.rs b/src/stats/tls_fingerprints/tests.rs new file mode 100644 index 0000000..c0b1edb --- /dev/null +++ b/src/stats/tls_fingerprints/tests.rs @@ -0,0 +1,58 @@ +use super::*; + +fn fp() -> TlsClientFingerprint { + TlsClientFingerprint { + ja3: "ja3".to_string(), + ja3_raw: "771,4865,,,0".to_string(), + ja4: "t13d010100_hash_hash".to_string(), + ja4_raw: "raw".to_string(), + } +} + +#[test] +fn aggregates_ip_cidr_and_user_scopes() { + let collector = TlsFingerprintCollector::default(); + let ip: IpAddr = "192.0.2.15".parse().expect("test IP parses"); + collector.record_observed(&fp(), ip, Duration::from_secs(60)); + collector.record_auth_success(&fp(), ip, "alice", Duration::from_secs(60)); + let snapshot = collector.snapshot(Duration::from_secs(60), 10); + + assert_eq!(snapshot.by_fingerprint[0].total, 1); + assert_eq!(snapshot.by_fingerprint[0].auth_success, 1); + assert_eq!(snapshot.by_ip[0].scope_key, "192.0.2.15"); + assert_eq!(snapshot.by_cidr[0].scope_key, "192.0.2.0/24"); + assert_eq!(snapshot.by_user[0].scope_key, "alice"); + assert_eq!(snapshot.by_user[0].total, 1); +} + +#[test] +fn parallel_distinct_insertions_respect_exact_capacity() { + const CAPACITY: usize = 127; + const ATTEMPTS: usize = 10_000; + + let collector = std::sync::Arc::new(TlsFingerprintCollector::with_capacity(CAPACITY)); + std::thread::scope(|scope| { + for worker in 0..16 { + let collector = std::sync::Arc::clone(&collector); + scope.spawn(move || { + for index in (worker..ATTEMPTS).step_by(16) { + collector.record_scoped( + (TlsFingerprintScopeKind::Ip, format!("192.0.2.{index}")), + &fp(), + 1, + true, + false, + false, + ); + } + }); + } + }); + + assert_eq!(collector.entries.len(), CAPACITY); + assert_eq!(collector.slots.used(), CAPACITY); + assert_eq!( + collector.dropped_total.load(Ordering::Relaxed), + (ATTEMPTS - CAPACITY) as u64 + ); +} diff --git a/src/stats/users.rs b/src/stats/users.rs index 2c9d4af..e092b7b 100644 --- a/src/stats/users.rs +++ b/src/stats/users.rs @@ -1,5 +1,32 @@ use super::*; +/// Mirrors one live connection into optional per-user telemetry. +#[must_use = "the observation must be retained for the connection lifetime"] +pub(crate) struct UserConnectionObservation { + stats: Arc, +} + +impl Drop for UserConnectionObservation { + fn drop(&mut self) { + decrement_current_connections(&self.stats.curr_connects); + } +} + +fn decrement_current_connections(counter: &AtomicU64) { + let mut current = counter.load(Ordering::Relaxed); + while current != 0 { + match counter.compare_exchange_weak( + current, + current - 1, + Ordering::Relaxed, + Ordering::Relaxed, + ) { + Ok(_) => return, + Err(actual) => current = actual, + } + } +} + impl Stats { pub fn increment_user_connects(&self, user: &str) { if !self.telemetry_user_enabled() { @@ -19,6 +46,25 @@ impl Stats { stats.curr_connects.fetch_add(1, Ordering::Relaxed); } + /// Starts an optional telemetry observation without owning admission policy. + pub(crate) fn observe_user_current_connection( + &self, + user: &str, + ) -> Option { + if !self.telemetry_user_enabled() { + return None; + } + let stats = self.get_or_create_user_stats_handle(user); + self.touch_user_stats(stats.as_ref()); + stats.curr_connects.fetch_add(1, Ordering::Relaxed); + Some(UserConnectionObservation { stats }) + } + + /// Returns the process-scoped count used by connection admission. + pub(crate) fn get_process_user_curr_connects(&self, user: &str) -> u64 { + self.connection_authority.active(user) + } + pub fn try_acquire_user_curr_connects(&self, user: &str, limit: Option) -> bool { if !self.telemetry_user_enabled() { return true; @@ -50,22 +96,7 @@ impl Stats { pub fn decrement_user_curr_connects(&self, user: &str) { if let Some(stats) = self.user_stats.get(user) { self.touch_user_stats(stats.value().as_ref()); - let counter = &stats.curr_connects; - let mut current = counter.load(Ordering::Relaxed); - loop { - if current == 0 { - break; - } - match counter.compare_exchange_weak( - current, - current - 1, - Ordering::Relaxed, - Ordering::Relaxed, - ) { - Ok(_) => break, - Err(actual) => current = actual, - } - } + decrement_current_connections(&stats.curr_connects); } } @@ -122,6 +153,27 @@ impl Stats { self.quota_store.used(user) } + /// Returns quota ownership for the exact authenticated user incarnation. + pub(crate) fn quota_handle_for_incarnation( + &self, + user: &str, + incarnation: crate::proxy::user_admission::UserIncarnation, + ) -> Option { + if let Some(handle) = self.quota_store.handle_exact(user, incarnation) { + return Some(handle); + } + if incarnation != 0 { + return None; + } + let handle = self.quota_store.current_or_legacy_handle(user); + (handle.incarnation() == incarnation).then_some(handle) + } + + /// Returns the currently published quota owner for compatibility relay entrypoints. + pub(crate) fn current_user_quota_handle(&self, user: &str) -> UserQuotaHandle { + self.quota_store.current_or_legacy_handle(user) + } + pub fn load_user_quota_state(&self, user: &str, used_bytes: u64, last_reset_epoch_secs: u64) { self.quota_store .load(user, used_bytes, last_reset_epoch_secs); diff --git a/src/synlimit_control/command.rs b/src/synlimit_control/command.rs index 19355c0..d631d92 100644 --- a/src/synlimit_control/command.rs +++ b/src/synlimit_control/command.rs @@ -1,14 +1,16 @@ -use std::path::PathBuf; - use tokio::io::AsyncWriteExt; use tokio::process::Command; +use crate::util::trusted_command::resolve_trusted_helper; + +const COMMAND_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30); + pub(super) async fn run_command( binary: &str, args: &[&str], stdin: Option, ) -> Result<(), String> { - let Some(command_path) = resolve_command(binary) else { + let Some(command_path) = resolve_trusted_helper(binary) else { return Err(format!("{binary} is not available")); }; let mut command = Command::new(command_path); @@ -18,21 +20,26 @@ pub(super) async fn run_command( } command.stdout(std::process::Stdio::null()); command.stderr(std::process::Stdio::piped()); + command.kill_on_drop(true); let mut child = command .spawn() .map_err(|e| format!("spawn {binary} failed: {e}"))?; - if let Some(blob) = stdin - && let Some(mut writer) = child.stdin.take() - { - writer - .write_all(blob.as_bytes()) + let output = tokio::time::timeout(COMMAND_TIMEOUT, async move { + if let Some(blob) = stdin + && let Some(mut writer) = child.stdin.take() + { + writer + .write_all(blob.as_bytes()) + .await + .map_err(|e| format!("stdin write {binary} failed: {e}"))?; + } + child + .wait_with_output() .await - .map_err(|e| format!("stdin write {binary} failed: {e}"))?; - } - let output = child - .wait_with_output() - .await - .map_err(|e| format!("wait {binary} failed: {e}"))?; + .map_err(|e| format!("wait {binary} failed: {e}")) + }) + .await + .map_err(|_| format!("{binary} timed out after {}s", COMMAND_TIMEOUT.as_secs()))??; if output.status.success() { return Ok(()); } @@ -45,13 +52,14 @@ pub(super) async fn run_command( } pub(super) async fn run_command_stdout(binary: &str, args: &[&str]) -> Result { - let Some(command_path) = resolve_command(binary) else { + let Some(command_path) = resolve_trusted_helper(binary) else { return Err(format!("{binary} is not available")); }; - let output = Command::new(command_path) - .args(args) - .output() + let mut command = Command::new(command_path); + command.args(args).kill_on_drop(true); + let output = tokio::time::timeout(COMMAND_TIMEOUT, command.output()) .await + .map_err(|_| format!("{binary} timed out after {}s", COMMAND_TIMEOUT.as_secs()))? .map_err(|e| format!("wait {binary} failed: {e}"))?; if output.status.success() { return Ok(String::from_utf8_lossy(&output.stdout).to_string()); @@ -64,16 +72,6 @@ pub(super) async fn run_command_stdout(binary: &str, args: &[&str]) -> Result Option { - let mut dirs = std::env::var_os("PATH") - .map(|path| std::env::split_paths(&path).collect::>()) - .unwrap_or_default(); - dirs.extend(["/usr/sbin", "/sbin", "/usr/bin", "/bin"].map(PathBuf::from)); - dirs.into_iter() - .map(|dir| dir.join(binary)) - .find(|candidate| candidate.exists() && candidate.is_file()) -} - pub(super) fn has_firewall_privileges() -> bool { #[cfg(target_os = "linux")] { diff --git a/src/tls_front/cache/disk.rs b/src/tls_front/cache/disk.rs index 3a73263..c516ee0 100644 --- a/src/tls_front/cache/disk.rs +++ b/src/tls_front/cache/disk.rs @@ -1,27 +1,26 @@ use std::path::Path; -use tokio::io::AsyncReadExt; - use super::*; pub(super) async fn read_disk_entry_bounded(path: &Path) -> std::io::Result> { - let file = tokio::fs::File::open(path).await?; - if file.metadata().await?.len() > TLS_FRONT_DISK_ENTRY_MAX_BYTES { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidData, - "TLS cache entry exceeds the 1 MiB limit", - )); + #[cfg(unix)] + { + crate::util::secure_fs::read_regular_limited_async( + path.to_path_buf(), + TLS_FRONT_DISK_ENTRY_MAX_BYTES as usize, + ) + .await } - let mut bytes = Vec::new(); - file.take(TLS_FRONT_DISK_ENTRY_MAX_BYTES.saturating_add(1)) - .read_to_end(&mut bytes) - .await?; - if bytes.len() as u64 > TLS_FRONT_DISK_ENTRY_MAX_BYTES { - return Err(std::io::Error::new( - std::io::ErrorKind::InvalidData, - "TLS cache entry grew beyond the 1 MiB limit while reading", - )); + #[cfg(not(unix))] + { + let bytes = tokio::fs::read(path).await?; + if bytes.len() as u64 > TLS_FRONT_DISK_ENTRY_MAX_BYTES { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "TLS cache entry exceeds the 1 MiB limit", + )); + } + Ok(bytes) } - Ok(bytes) } pub(super) fn cert_info_matches_domain(cached: &CachedTlsData) -> bool { diff --git a/src/tls_front/cache/runtime.rs b/src/tls_front/cache/runtime.rs index 91132e4..fca4039 100644 --- a/src/tls_front/cache/runtime.rs +++ b/src/tls_front/cache/runtime.rs @@ -231,9 +231,6 @@ impl TlsFrontCache { pub async fn load_from_disk(&self) { let path = self.disk_path.clone(); - if tokio::fs::create_dir_all(&path).await.is_err() { - return; - } let mut loaded = 0usize; for name in &self.disk_entry_names { let entry_path = path.join(name); @@ -297,9 +294,6 @@ impl TlsFrontCache { } async fn persist(&self, domain: &str, data: &CachedTlsData) { - if tokio::fs::create_dir_all(&self.disk_path).await.is_err() { - return; - } let fname = format!("{}.json", domain.replace(['/', '\\'], "_")); let path = self.disk_path.join(fname); if let Ok(json) = serde_json::to_vec_pretty(data) { @@ -311,7 +305,9 @@ impl TlsFrontCache { ); return; } - // best-effort write + #[cfg(unix)] + let _ = crate::util::secure_fs::atomic_replace_async(path, json, 0o600).await; + #[cfg(not(unix))] let _ = tokio::fs::write(path, json).await; } } diff --git a/src/tls_front/fetcher.rs b/src/tls_front/fetcher.rs index e49ae86..a831f1f 100644 --- a/src/tls_front/fetcher.rs +++ b/src/tls_front/fetcher.rs @@ -172,6 +172,12 @@ fn sweep_expired_profile_cache(ttl: Duration, now: Instant) { profile_cache().retain(|_, value| now.saturating_duration_since(value.updated_at) <= ttl); } +fn remove_profile_if_unchanged(key: &ProfileCacheKey, observed: ProfileCacheValue) { + profile_cache().remove_if(key, |_, current| { + current.profile == observed.profile && current.updated_at == observed.updated_at + }); +} + /// Current number of adaptive TLS fetch profile-cache entries. pub(crate) fn profile_cache_entries_for_metrics() -> usize { profile_cache().len() @@ -270,8 +276,9 @@ fn order_profiles( if let Some(cached) = profile_cache().get(key) { let age = now.saturating_duration_since(cached.updated_at); if age > strategy.profile_cache_ttl { + let observed = *cached; drop(cached); - profile_cache().remove(key); + remove_profile_if_unchanged(key, observed); return ordered; } diff --git a/src/tls_front/fetcher/tests.rs b/src/tls_front/fetcher/tests.rs index 2309754..63556f3 100644 --- a/src/tls_front/fetcher/tests.rs +++ b/src/tls_front/fetcher/tests.rs @@ -6,6 +6,7 @@ use super::{ TLS_NAMED_GROUP_X25519MLKEM768, TlsFetchStrategy, X25519_KEY_SHARE_LEN, build_client_hello, build_tls_fetch_proxy_header, derive_behavior_profile, encode_tls13_certificate_message, fetch_via_rustls_stream, order_profiles, profile_alpn, profile_cache, profile_cache_key, + remove_profile_if_unchanged, }; use crate::config::TlsFetchProfile; use crate::crypto::SecureRandom; @@ -225,6 +226,31 @@ fn test_order_profiles_drops_expired_cached_winner() { assert!(profile_cache().get(&cache_key).is_none()); } +#[test] +fn expired_profile_removal_preserves_concurrent_refresh() { + let cache_key = profile_cache_key("mask3.example", 443, "tls3.example", None, None, 0, None); + let observed = ProfileCacheValue { + profile: TlsFetchProfile::CompatTls12, + updated_at: Instant::now() - Duration::from_secs(60), + }; + let refreshed = ProfileCacheValue { + profile: TlsFetchProfile::ModernChromeLike, + updated_at: Instant::now(), + }; + profile_cache().insert(cache_key.clone(), observed); + profile_cache().insert(cache_key.clone(), refreshed); + + remove_profile_if_unchanged(&cache_key, observed); + + let current = profile_cache() + .get(&cache_key) + .expect("concurrent refresh must remain cached"); + assert_eq!(current.profile, refreshed.profile); + assert_eq!(current.updated_at, refreshed.updated_at); + drop(current); + profile_cache().remove(&cache_key); +} + #[test] fn test_deterministic_client_hello_is_stable() { let rng = SecureRandom::new(); diff --git a/src/transport/middle_proxy/config_updater.rs b/src/transport/middle_proxy/config_updater.rs index 50278b5..0c1bd6e 100644 --- a/src/transport/middle_proxy/config_updater.rs +++ b/src/transport/middle_proxy/config_updater.rs @@ -12,6 +12,8 @@ use tracing::{debug, info, warn}; use crate::config::ProxyConfig; use crate::error::Result; use crate::transport::UpstreamManager; +#[cfg(unix)] +use crate::util::secure_fs::{atomic_replace_async, read_regular_limited_async}; use super::MePool; use super::http_fetch::{HTTPS_RESPONSE_BODY_MAX_BYTES, https_get}; @@ -73,25 +75,34 @@ pub fn parse_proxy_config_text(text: &str, http_status: u16) -> ProxyConfigData } pub async fn load_proxy_config_cache(path: &str) -> Result { - let text = tokio::fs::read_to_string(path).await.map_err(|e| { + #[cfg(unix)] + let bytes = + read_regular_limited_async(Path::new(path).to_path_buf(), HTTPS_RESPONSE_BODY_MAX_BYTES) + .await; + #[cfg(not(unix))] + let bytes = tokio::fs::read(path).await; + let bytes = bytes.map_err(|e| { crate::error::ProxyError::Proxy(format!("read proxy-config cache '{path}' failed: {e}")) })?; + let text = String::from_utf8(bytes).map_err(|e| { + crate::error::ProxyError::Proxy(format!( + "proxy-config cache '{path}' is not valid UTF-8: {e}" + )) + })?; Ok(parse_proxy_config_text(&text, 200)) } pub async fn save_proxy_config_cache(path: &str, raw_text: &str) -> Result<()> { - if let Some(parent) = Path::new(path).parent() - && !parent.as_os_str().is_empty() - { - tokio::fs::create_dir_all(parent).await.map_err(|e| { - crate::error::ProxyError::Proxy(format!( - "create proxy-config cache dir '{}' failed: {e}", - parent.display() - )) - })?; - } - - tokio::fs::write(path, raw_text).await.map_err(|e| { + #[cfg(unix)] + let write = atomic_replace_async( + Path::new(path).to_path_buf(), + raw_text.as_bytes().to_vec(), + 0o640, + ) + .await; + #[cfg(not(unix))] + let write = tokio::fs::write(path, raw_text).await; + write.map_err(|e| { crate::error::ProxyError::Proxy(format!("write proxy-config cache '{path}' failed: {e}")) })?; Ok(()) diff --git a/src/transport/middle_proxy/health.rs b/src/transport/middle_proxy/health.rs index f51b7f5..40c8743 100644 --- a/src/transport/middle_proxy/health.rs +++ b/src/transport/middle_proxy/health.rs @@ -16,9 +16,13 @@ use crate::crypto::SecureRandom; use crate::network::IpFamily; use super::MePool; -use super::pool::MeFamilyRuntimeState; +use super::pool::{ + MeFamilyRuntimeState, RefillTargetKey, WriterContour, WriterOpenIntent, WriterRole, +}; +use super::pool_writer::WriterReplacementPurpose; -const JITTER_FRAC_NUM: u64 = 2; // jitter up to 50% of backoff +// Bounds reconnect jitter to at most half of the current backoff. +const JITTER_FRAC_NUM: u64 = 2; #[allow(dead_code)] const MAX_CONCURRENT_PER_DC_DEFAULT: usize = 1; const SHADOW_ROTATE_RETRY_SECS: u64 = 30; diff --git a/src/transport/middle_proxy/health/drain.rs b/src/transport/middle_proxy/health/drain.rs index 336517b..5c9e629 100644 --- a/src/transport/middle_proxy/health/drain.rs +++ b/src/transport/middle_proxy/health/drain.rs @@ -31,7 +31,7 @@ pub(in crate::transport::middle_proxy) async fn reap_draining_writers( let mut force_close_writer_ids = Vec::::new(); let writers = pool.writers.read().await; for writer in writers.iter() { - if !writer.draining.load(std::sync::atomic::Ordering::Relaxed) { + if !writer.draining.load(std::sync::atomic::Ordering::Acquire) { continue; } if activity diff --git a/src/transport/middle_proxy/health/family.rs b/src/transport/middle_proxy/health/family.rs index c454675..a23a02d 100644 --- a/src/transport/middle_proxy/health/family.rs +++ b/src/transport/middle_proxy/health/family.rs @@ -25,9 +25,11 @@ pub(super) async fn check_family( let mut family_degraded = false; let mut dc_endpoints = HashMap::>::new(); + let endpoint_snapshot = pool.endpoint_snapshot.load(); + let endpoint_revision = endpoint_snapshot.revision; let map_guard = match family { - IpFamily::V4 => pool.proxy_map_v4.read().await, - IpFamily::V6 => pool.proxy_map_v6.read().await, + IpFamily::V4 => &endpoint_snapshot.map_v4, + IpFamily::V6 => &endpoint_snapshot.map_v6, }; for (dc, addrs) in map_guard.iter() { let entry = dc_endpoints.entry(*dc).or_default(); @@ -35,7 +37,7 @@ pub(super) async fn check_family( entry.push(SocketAddr::new(ip, port)); } } - drop(map_guard); + drop(endpoint_snapshot); for endpoints in dc_endpoints.values_mut() { endpoints.sort_unstable(); endpoints.dedup(); @@ -47,6 +49,7 @@ pub(super) async fn check_family( let mut live_addr_counts = HashMap::<(i32, SocketAddr), usize>::new(); let mut live_writer_ids_by_addr = HashMap::<(i32, SocketAddr), Vec>::new(); + let active_generation = pool.current_generation(); for writer in pool .writers .read() @@ -62,6 +65,15 @@ pub(super) async fn check_family( ) { continue; } + if writer.generation != active_generation { + continue; + } + if !dc_endpoints + .get(&writer.writer_dc) + .is_some_and(|endpoints| endpoints.contains(&writer.addr)) + { + continue; + } let key = (writer.writer_dc, writer.addr); *live_addr_counts.entry(key).or_insert(0) += 1; live_writer_ids_by_addr @@ -93,6 +105,13 @@ pub(super) async fn check_family( floor_plan.active_writers_current, floor_plan.warm_writers_current, ); + let floor_targets_by_dc = Arc::new( + floor_plan + .by_dc + .iter() + .map(|(dc, entry)| (*dc, entry.target_required)) + .collect::>(), + ); let live_writer_ids_by_addr = Arc::new(live_writer_ids_by_addr); let writer_idle_since = Arc::new(writer_idle_since); let bound_clients_by_writer = Arc::new(bound_clients_by_writer); @@ -162,22 +181,23 @@ pub(super) async fn check_family( ); } + maybe_refresh_idle_writer_for_dc( + pool, + rng, + key, + dc, + family, + &endpoints, + alive, + required, + live_writer_ids_by_addr.as_ref(), + writer_idle_since.as_ref(), + bound_clients_by_writer.as_ref(), + idle_refresh_next_attempt, + ) + .await; + if alive >= required { - maybe_refresh_idle_writer_for_dc( - pool, - rng, - key, - dc, - family, - &endpoints, - alive, - required, - live_writer_ids_by_addr.as_ref(), - writer_idle_since.as_ref(), - bound_clients_by_writer.as_ref(), - idle_refresh_next_attempt, - ) - .await; maybe_rotate_single_endpoint_shadow( pool, rng, @@ -230,9 +250,12 @@ pub(super) async fn check_family( continue; } if pool - .has_refill_inflight_for_dc_key(crate::transport::middle_proxy::pool::RefillDcKey { + .has_refill_inflight_for_target(RefillTargetKey { dc, family, + generation: pool.current_generation(), + endpoint_revision, + contour: WriterContour::Active, }) .await { @@ -254,6 +277,7 @@ pub(super) async fn check_family( let live_writer_ids_by_addr_for_dc = live_writer_ids_by_addr.clone(); let writer_idle_since_for_dc = writer_idle_since.clone(); let bound_clients_by_writer_for_dc = bound_clients_by_writer.clone(); + let floor_targets_by_dc_for_reconnect = floor_targets_by_dc.clone(); let active_cap_effective_total = floor_plan.active_cap_effective_total; reconnect_set.spawn(async move { let mut restored = 0usize; @@ -261,6 +285,8 @@ pub(super) async fn check_family( let Ok(reconnect_permit) = reconnect_sem_for_dc.clone().try_acquire_owned() else { break; }; + let base_req = pool_for_reconnect + .required_writers_for_dc_with_floor_mode(endpoints_for_dc.len(), false); if pool_for_reconnect.active_contour_writer_count_total().await >= active_cap_effective_total { @@ -270,6 +296,8 @@ pub(super) async fn check_family( dc, family, &endpoints_for_dc, + required, + floor_targets_by_dc_for_reconnect.as_ref(), live_writer_ids_by_addr_for_dc.as_ref(), writer_idle_since_for_dc.as_ref(), bound_clients_by_writer_for_dc.as_ref(), @@ -283,8 +311,6 @@ pub(super) async fn check_family( continue; } - let base_req = pool_for_reconnect - .required_writers_for_dc_with_floor_mode(endpoints_for_dc.len(), false); if alive + restored >= base_req { pool_for_reconnect .stats @@ -304,12 +330,20 @@ pub(super) async fn check_family( } } pool_for_reconnect.stats.increment_me_reconnect_attempt(); + let intent = if alive + restored < base_req { + WriterOpenIntent::Coverage + } else { + WriterOpenIntent::Normal + }; let res = tokio::time::timeout( pool_for_reconnect.reconnect_runtime.me_one_timeout, - pool_for_reconnect.connect_endpoints_round_robin( + pool_for_reconnect.connect_endpoints_round_robin_with_generation_contour( dc, &endpoints_for_dc, rng_for_reconnect.as_ref(), + pool_for_reconnect.current_generation(), + WriterContour::Active, + intent, ), ) .await; diff --git a/src/transport/middle_proxy/health/floor_plan.rs b/src/transport/middle_proxy/health/floor_plan.rs index 0ad16eb..01b41b3 100644 --- a/src/transport/middle_proxy/health/floor_plan.rs +++ b/src/transport/middle_proxy/health/floor_plan.rs @@ -22,6 +22,7 @@ pub(super) async fn live_active_writers_for_dc_family( dc: i32, family: IpFamily, ) -> usize { + let active_generation = pool.current_generation(); let writers = pool.writers.read().await; writers .iter() @@ -32,6 +33,9 @@ pub(super) async fn live_active_writers_for_dc_family( if writer.writer_dc != dc { return false; } + if writer.generation != active_generation { + return false; + } if !matches!( crate::transport::middle_proxy::pool::WriterContour::from_u8( writer.contour.load(std::sync::atomic::Ordering::Relaxed), diff --git a/src/transport/middle_proxy/health/idle_refresh.rs b/src/transport/middle_proxy/health/idle_refresh.rs index 3cc26dc..53054f1 100644 --- a/src/transport/middle_proxy/health/idle_refresh.rs +++ b/src/transport/middle_proxy/health/idle_refresh.rs @@ -6,16 +6,43 @@ pub(super) async fn maybe_swap_idle_writer_for_cap( dc: i32, family: IpFamily, endpoints: &[SocketAddr], + required: usize, + floor_targets_by_dc: &HashMap, live_writer_ids_by_addr: &HashMap<(i32, SocketAddr), Vec>, writer_idle_since: &HashMap, bound_clients_by_writer: &HashMap, ) -> bool { + let Some(replacement_endpoint) = endpoints + .iter() + .min_by_key(|endpoint| { + live_writer_ids_by_addr + .get(&(dc, **endpoint)) + .map_or(0, Vec::len) + }) + .copied() + else { + return false; + }; + + let mut alive_by_dc = HashMap::::new(); + for ((writer_dc, endpoint), writer_ids) in live_writer_ids_by_addr { + if endpoint.is_ipv4() == matches!(family, IpFamily::V4) { + *alive_by_dc.entry(*writer_dc).or_insert(0) += writer_ids.len(); + } + } + let now_epoch_secs = MePool::now_epoch_secs(); - let mut candidate: Option<(u64, SocketAddr, u64)> = None; - for endpoint in endpoints { - let Some(writer_ids) = live_writer_ids_by_addr.get(&(dc, *endpoint)) else { + let mut candidates = Vec::<(u64, SocketAddr, i32, u64, usize)>::new(); + for ((writer_dc, endpoint), writer_ids) in live_writer_ids_by_addr { + if endpoint.is_ipv4() != matches!(family, IpFamily::V4) { + continue; + } + let Some(donor_floor) = floor_targets_by_dc.get(writer_dc).copied() else { continue; }; + if alive_by_dc.get(writer_dc).copied().unwrap_or(0) <= donor_floor { + continue; + } for writer_id in writer_ids { if bound_clients_by_writer.get(writer_id).copied().unwrap_or(0) > 0 { continue; @@ -23,67 +50,96 @@ pub(super) async fn maybe_swap_idle_writer_for_cap( let Some(idle_since_epoch_secs) = writer_idle_since.get(writer_id).copied() else { continue; }; - let idle_age_secs = now_epoch_secs.saturating_sub(idle_since_epoch_secs); - if candidate - .as_ref() - .map(|(_, _, age)| idle_age_secs > *age) - .unwrap_or(true) - { - candidate = Some((*writer_id, *endpoint, idle_age_secs)); + candidates.push(( + *writer_id, + *endpoint, + *writer_dc, + now_epoch_secs.saturating_sub(idle_since_epoch_secs), + donor_floor, + )); + } + } + candidates.sort_unstable_by(|left, right| right.3.cmp(&left.3)); + + for (old_writer_id, donor_endpoint, donor_dc, idle_age_secs, donor_floor) in candidates { + let expected_role = { + let writers = pool.writers.read().await; + writers + .iter() + .find(|writer| writer.id == old_writer_id) + .map(WriterRole::from_writer) + }; + let Some(expected_role) = expected_role else { + continue; + }; + if expected_role.dc != donor_dc + || expected_role.family != family + || expected_role.contour != WriterContour::Active + { + continue; + } + let Some(mut reservation) = pool + .registry + .try_reserve_writer_replacement(old_writer_id) + .await + else { + continue; + }; + let replace = pool.replace_writer_with_generation_contour_for_dc( + replacement_endpoint, + rng.as_ref(), + pool.current_generation(), + WriterContour::Active, + dc, + expected_role, + WriterReplacementPurpose::FloorRebalance { + donor_floor, + receiver_floor: required, + }, + &mut reservation, + ); + match tokio::time::timeout(pool.reconnect_runtime.me_one_timeout, replace).await { + Ok(Ok(())) => { + info!( + dc = %dc, + ?family, + %replacement_endpoint, + donor_dc = %donor_dc, + %donor_endpoint, + old_writer_id, + idle_age_secs, + "Adaptive floor cap rebalance committed" + ); + return true; + } + Ok(Err(error)) => { + debug!( + dc = %dc, + ?family, + %replacement_endpoint, + donor_dc = %donor_dc, + old_writer_id, + idle_age_secs, + %error, + "Adaptive floor cap rebalance failed" + ); + return false; + } + Err(_) => { + debug!( + dc = %dc, + ?family, + %replacement_endpoint, + donor_dc = %donor_dc, + old_writer_id, + idle_age_secs, + "Adaptive floor cap rebalance timed out" + ); + return false; } } } - - let Some((old_writer_id, endpoint, idle_age_secs)) = candidate else { - return false; - }; - - let connected = match tokio::time::timeout( - pool.reconnect_runtime.me_one_timeout, - pool.connect_one_for_dc(endpoint, dc, rng.as_ref()), - ) - .await - { - Ok(Ok(())) => true, - Ok(Err(error)) => { - debug!( - dc = %dc, - ?family, - %endpoint, - old_writer_id, - idle_age_secs, - %error, - "Adaptive floor cap swap connect failed" - ); - false - } - Err(_) => { - debug!( - dc = %dc, - ?family, - %endpoint, - old_writer_id, - idle_age_secs, - "Adaptive floor cap swap connect timed out" - ); - false - } - }; - if !connected { - return false; - } - - pool.mark_writer_draining_with_timeout(old_writer_id, pool.force_close_timeout(), false) - .await; - info!( - dc = %dc, - ?family, - %endpoint, - old_writer_id, - idle_age_secs, - "Adaptive floor cap swap: idle writer rotated" - ); - true + false } pub(super) async fn maybe_refresh_idle_writer_for_dc( @@ -100,10 +156,6 @@ pub(super) async fn maybe_refresh_idle_writer_for_dc( bound_clients_by_writer: &HashMap, idle_refresh_next_attempt: &mut HashMap<(i32, IpFamily), Instant>, ) { - if alive < required { - return; - } - let now = Instant::now(); if let Some(next) = idle_refresh_next_attempt.get(&key) && now < *next @@ -143,12 +195,46 @@ pub(super) async fn maybe_refresh_idle_writer_for_dc( let Some((old_writer_id, endpoint, idle_age_secs, threshold_secs)) = candidate else { return; }; - - let rotate_ok = match tokio::time::timeout( - pool.reconnect_runtime.me_one_timeout, - pool.connect_one_for_dc(endpoint, dc, rng.as_ref()), - ) - .await + let expected_role = { + let writers = pool.writers.read().await; + writers + .iter() + .find(|writer| writer.id == old_writer_id) + .map(WriterRole::from_writer) + }; + let Some(expected_role) = expected_role else { + return; + }; + if expected_role.dc != dc + || expected_role.family != family + || expected_role.contour != WriterContour::Active + { + return; + } + let Some(mut reservation) = pool + .registry + .try_reserve_writer_replacement(old_writer_id) + .await + else { + return; + }; + let generation = pool.current_generation(); + let purpose = if expected_role.generation == generation { + WriterReplacementPurpose::IdleRefresh + } else { + WriterReplacementPurpose::GenerationConvergence + }; + let replace = pool.replace_writer_with_generation_contour_for_dc( + endpoint, + rng.as_ref(), + generation, + WriterContour::Active, + dc, + expected_role, + purpose, + &mut reservation, + ); + let rotate_ok = match tokio::time::timeout(pool.reconnect_runtime.me_one_timeout, replace).await { Ok(Ok(())) => true, Ok(Err(error)) => { @@ -160,7 +246,7 @@ pub(super) async fn maybe_refresh_idle_writer_for_dc( idle_age_secs, threshold_secs, %error, - "Idle writer pre-refresh connect failed" + "Idle writer pre-refresh replacement failed" ); false } @@ -172,7 +258,7 @@ pub(super) async fn maybe_refresh_idle_writer_for_dc( old_writer_id, idle_age_secs, threshold_secs, - "Idle writer pre-refresh connect timed out" + "Idle writer pre-refresh replacement timed out" ); false } @@ -183,8 +269,6 @@ pub(super) async fn maybe_refresh_idle_writer_for_dc( return; } - pool.mark_writer_draining_with_timeout(old_writer_id, pool.force_close_timeout(), false) - .await; idle_refresh_next_attempt.insert( key, now + Duration::from_secs(IDLE_REFRESH_SUCCESS_GUARD_SECS), @@ -198,6 +282,7 @@ pub(super) async fn maybe_refresh_idle_writer_for_dc( threshold_secs, alive, required, + generation_convergence = expected_role.generation != generation, "Idle writer refreshed before upstream idle timeout" ); } diff --git a/src/transport/middle_proxy/health/recovery.rs b/src/transport/middle_proxy/health/recovery.rs index cfeee67..709ca0b 100644 --- a/src/transport/middle_proxy/health/recovery.rs +++ b/src/transport/middle_proxy/health/recovery.rs @@ -209,10 +209,38 @@ pub(super) async fn maybe_rotate_single_endpoint_shadow( ); return; }; + let expected_role = { + let writers = pool.writers.read().await; + writers + .iter() + .find(|writer| writer.id == old_writer_id) + .map(WriterRole::from_writer) + }; + let Some(expected_role) = expected_role else { + shadow_rotate_deadline.insert(key, now + Duration::from_secs(SHADOW_ROTATE_RETRY_SECS)); + return; + }; + let Some(mut reservation) = pool + .registry + .try_reserve_writer_replacement(old_writer_id) + .await + else { + shadow_rotate_deadline.insert(key, now + Duration::from_secs(SHADOW_ROTATE_RETRY_SECS)); + return; + }; let rotate_ok = match tokio::time::timeout( pool.reconnect_runtime.me_one_timeout, - pool.connect_one_for_dc(endpoint, dc, rng.as_ref()), + pool.replace_writer_with_generation_contour_for_dc( + endpoint, + rng.as_ref(), + pool.current_generation(), + WriterContour::Active, + dc, + expected_role, + WriterReplacementPurpose::ShadowRotation, + &mut reservation, + ), ) .await { @@ -246,8 +274,6 @@ pub(super) async fn maybe_rotate_single_endpoint_shadow( return; } - pool.mark_writer_draining_with_timeout(old_writer_id, pool.force_close_timeout(), false) - .await; pool.stats .increment_me_single_endpoint_shadow_rotate_total(); shadow_rotate_deadline.insert(key, now + interval); diff --git a/src/transport/middle_proxy/health/tests.rs b/src/transport/middle_proxy/health/tests.rs index 0e612ef..2b223d1 100644 --- a/src/transport/middle_proxy/health/tests.rs +++ b/src/transport/middle_proxy/health/tests.rs @@ -7,7 +7,7 @@ use std::time::{Duration, Instant}; use tokio::sync::mpsc; use tokio_util::sync::CancellationToken; -use super::{ScheduledReconnects, reap_draining_writers}; +use super::{ScheduledReconnects, maybe_refresh_idle_writer_for_dc, reap_draining_writers}; use crate::config::{GeneralConfig, MeRouteNoWriterMode, MeSocksKdfPolicy, MeWriterPickMode}; use crate::crypto::SecureRandom; use crate::network::IpFamily; @@ -226,6 +226,77 @@ async fn insert_live_writer(pool: &Arc, writer_id: u64, writer_dc: i32) pool.conn_count.fetch_add(1, Ordering::Relaxed); } +async fn insert_active_writer_at( + pool: &Arc, + writer_id: u64, + writer_dc: i32, + addr: SocketAddr, +) -> MeWriter { + let (tx, _writer_rx) = mpsc::channel::(8); + let byte_budget = pool.new_writer_byte_budget(); + let writer = MeWriter { + id: writer_id, + addr, + source_ip: addr.ip(), + writer_dc, + generation: pool.current_generation(), + contour: Arc::new(AtomicU8::new(WriterContour::Active.as_u8())), + created_at: Instant::now(), + tx: tx.clone(), + byte_budget: byte_budget.clone(), + cancel: CancellationToken::new(), + degraded: Arc::new(AtomicBool::new(false)), + rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)), + draining: Arc::new(AtomicBool::new(false)), + draining_started_at_epoch_secs: Arc::new(AtomicU64::new(0)), + drain_deadline_epoch_secs: Arc::new(AtomicU64::new(0)), + allow_drain_fallback: Arc::new(AtomicBool::new(false)), + }; + pool.writers.write().await.push(writer.clone()); + pool.registry + .register_writer(writer_id, tx, byte_budget) + .await; + pool.conn_count.fetch_add(1, Ordering::Relaxed); + writer +} + +#[tokio::test] +async fn under_floor_idle_writer_still_enters_transactional_refresh() { + let pool = make_pool(128).await; + let listener = std::net::TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).unwrap(); + let endpoint = listener.local_addr().unwrap(); + drop(listener); + let writer_id = 7001; + let writer = insert_active_writer_at(&pool, writer_id, 2, endpoint).await; + let key = (2, IpFamily::V4); + let live_writer_ids_by_addr = HashMap::from([((2, endpoint), vec![writer_id])]); + let writer_idle_since = + HashMap::from([(writer_id, MePool::now_epoch_secs().saturating_sub(60))]); + let bound_clients_by_writer = HashMap::from([(writer_id, 0)]); + let mut next_attempt = HashMap::new(); + let rng = Arc::new(SecureRandom::new()); + + maybe_refresh_idle_writer_for_dc( + &pool, + &rng, + key, + 2, + IpFamily::V4, + &[endpoint], + 1, + 10, + &live_writer_ids_by_addr, + &writer_idle_since, + &bound_clients_by_writer, + &mut next_attempt, + ) + .await; + + assert!(next_attempt.contains_key(&key)); + assert!(!writer.draining.load(Ordering::Acquire)); + assert_eq!(pool.registry.writer_replacement_counts(), (0, 0)); +} + #[tokio::test] async fn reap_draining_writers_force_closes_oldest_over_threshold() { let pool = make_pool(2).await; diff --git a/src/transport/middle_proxy/health/zombie_watchdog.rs b/src/transport/middle_proxy/health/zombie_watchdog.rs index 074cad4..62ccd6c 100644 --- a/src/transport/middle_proxy/health/zombie_watchdog.rs +++ b/src/transport/middle_proxy/health/zombie_watchdog.rs @@ -35,7 +35,7 @@ pub async fn me_zombie_writer_watchdog(pool: Arc) { continue; }; ws.iter() - .filter(|w| w.draining.load(std::sync::atomic::Ordering::Relaxed)) + .filter(|w| w.draining.load(std::sync::atomic::Ordering::Acquire)) .filter_map(|w| { let deadline = w .drain_deadline_epoch_secs diff --git a/src/transport/middle_proxy/mod.rs b/src/transport/middle_proxy/mod.rs index 2c38fb5..ead2fa2 100644 --- a/src/transport/middle_proxy/mod.rs +++ b/src/transport/middle_proxy/mod.rs @@ -33,6 +33,9 @@ mod pool_runtime_api; mod pool_status; mod pool_writer; #[cfg(test)] +#[path = "tests/pool_writer_publication_tests.rs"] +mod pool_writer_publication_tests; +#[cfg(test)] #[path = "tests/pool_writer_security_tests.rs"] mod pool_writer_security_tests; mod reader; @@ -63,6 +66,7 @@ pub use ping::{ pub use pool::MePool; #[allow(unused_imports)] pub use pool_nat::{detect_public_ip, stun_probe}; +pub(crate) use pool_status::MeApiHardswapSnapshot; pub(crate) use registry::ConnLease; pub use registry::ConnRegistry; pub use rotation::{MeReinitTrigger, me_reinit_scheduler, me_rotation_task}; diff --git a/src/transport/middle_proxy/ping.rs b/src/transport/middle_proxy/ping.rs index 85888bd..8a86204 100644 --- a/src/transport/middle_proxy/ping.rs +++ b/src/transport/middle_proxy/ping.rs @@ -329,13 +329,14 @@ mod tests { pub async fn run_me_ping(pool: &Arc, rng: &SecureRandom) -> Vec { let mut reports = Vec::new(); + let endpoint_snapshot = pool.endpoint_snapshot.load_full(); let v4_map = if pool.decision.ipv4_me { - pool.proxy_map_v4.read().await.clone() + endpoint_snapshot.map_v4.clone() } else { HashMap::new() }; let v6_map = if pool.decision.ipv6_me { - pool.proxy_map_v6.read().await.clone() + endpoint_snapshot.map_v6.clone() } else { HashMap::new() }; diff --git a/src/transport/middle_proxy/pool.rs b/src/transport/middle_proxy/pool.rs index 232b556..6d58df2 100644 --- a/src/transport/middle_proxy/pool.rs +++ b/src/transport/middle_proxy/pool.rs @@ -29,9 +29,27 @@ use super::pool_lifecycle::MePoolLifecycle; const ME_FORCE_CLOSE_SAFETY_FALLBACK_SECS: u64 = 300; #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] -pub(super) struct RefillDcKey { +/// Exact lifecycle role used to coalesce refill work without cross-generation drift. +pub(super) struct RefillTargetKey { + /// Telegram DC owning the writer. pub dc: i32, + /// Address family of the writer endpoint. pub family: IpFamily, + /// Generation that retains publication authority. + pub generation: u64, + /// Endpoint snapshot revision targeted by this refill producer. + pub endpoint_revision: u64, + /// Lifecycle contour that the replacement must preserve. + pub contour: WriterContour, +} + +#[derive(Debug, Clone, Copy, Default)] +/// Bounded queued-loss state for one exact refill target. +pub(super) struct RefillTargetState { + /// Additional lost writers waiting behind the active refill producer. + pub(super) pending_count: usize, + /// Most recently lost endpoint, used as the next same-endpoint preference. + pub(super) next_addr: Option, } #[derive(Clone)] @@ -88,14 +106,6 @@ impl WritersState { } } - pub(super) async fn update(&self, f: F) -> R - where - F: FnOnce(&mut Vec) -> R, - { - let mut guard = self.write().await; - f(&mut guard) - } - fn debug_assert_store_guarded(&self) { debug_assert!( self.writers_write_guard.try_lock().is_err(), @@ -129,6 +139,13 @@ impl DerefMut for WritersWriteGuard<'_> { } } +impl WritersWriteGuard<'_> { + /// Publishes the current vector while retaining exclusive mutation ownership. + pub(super) fn publish_current(&self) { + self.state.store_guarded(self.writers.clone()); + } +} + impl Drop for WritersWriteGuard<'_> { fn drop(&mut self) { let writers = std::mem::take(&mut self.writers); @@ -136,7 +153,7 @@ impl Drop for WritersWriteGuard<'_> { } } -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] #[repr(u8)] pub(super) enum WriterContour { Warm = 0, @@ -144,18 +161,6 @@ pub(super) enum WriterContour { Draining = 2, } -pub(super) struct WriterOpenReservation<'a> { - counter: Option<&'a AtomicUsize>, -} - -impl Drop for WriterOpenReservation<'_> { - fn drop(&mut self) { - if let Some(counter) = self.counter { - counter.fetch_sub(1, Ordering::AcqRel); - } - } -} - impl WriterContour { pub(super) fn as_u8(self) -> u8 { self as u8 @@ -263,7 +268,23 @@ pub struct RoutingCore { pub(super) writers: Arc, pub(super) rr: AtomicU64, pub(super) writer_epoch: watch::Sender, - pub(super) preferred_endpoints_by_dc: ArcSwap>>, + /// Coherent immutable authority for endpoint maps and reverse indexes. + pub(super) endpoint_snapshot: ArcSwap, +} + +/// Immutable endpoint routing authority published as one coherent revision. +#[derive(Clone, Debug)] +pub(super) struct EndpointSnapshot { + /// Monotonic revision covering every endpoint-derived index in this snapshot. + pub(super) revision: u64, + /// IPv4 endpoint map by Telegram DC. + pub(super) map_v4: HashMap>, + /// IPv6 endpoint map by Telegram DC. + pub(super) map_v6: HashMap>, + /// Reverse lookup from an endpoint to its optional Telegram DC. + pub(super) endpoint_dc_map: HashMap>, + /// Ordered endpoint candidates used for per-DC writer selection. + pub(super) preferred_endpoints_by_dc: HashMap>, } pub(super) struct ReinitCore { @@ -291,6 +312,7 @@ pub(super) struct ReinitStatusSnapshot { pub(super) pending_hardswap_generation: u64, pub(super) pending_hardswap_started_at_epoch_secs: u64, pub(super) pending_hardswap_map_hash: u64, + pub(super) pending_hardswap_endpoint_revision: u64, pub(super) inflight: usize, } @@ -299,12 +321,16 @@ pub(super) struct ReinitPendingState { pub(super) generation: u64, pub(super) started_at_epoch_secs: u64, pub(super) map_hash: u64, + /// Endpoint authority revision targeted by the pending generation. + pub(super) endpoint_revision: u64, } #[derive(Clone, Copy)] pub(super) struct ReinitAttemptState { pub(super) generation: u64, pub(super) map_hash: u64, + /// Endpoint authority revision captured by this attempt. + pub(super) endpoint_revision: u64, pub(super) hardswap: bool, pub(super) committed: bool, } @@ -313,6 +339,8 @@ pub(super) struct ReinitCoordinatorState { pub(super) next_attempt_id: u64, pub(super) active_generation: u64, pub(super) desired_map_hash: u64, + /// Latest endpoint authority revision accepted by the coordinator. + pub(super) endpoint_revision: u64, pub(super) pending: Option, pub(super) attempts: HashMap, } @@ -472,15 +500,15 @@ pub struct MePool { pub(super) rng: Arc, pub(super) proxy_tag: Option>, pub(super) proxy_secret: Arc>, - pub(super) proxy_map_v4: Arc>>>, - pub(super) proxy_map_v6: Arc>>>, - pub(super) endpoint_dc_map: Arc>>>, pub(super) default_dc: AtomicI32, pub(super) next_writer_id: AtomicU64, pub(super) writer_connect_active_reserved: AtomicUsize, pub(super) writer_connect_warm_reserved: AtomicUsize, + /// Replacement connections opened but not yet committed to writer visibility. + pub(super) writer_replacement_open_reserved: AtomicUsize, pub(super) rtt_stats: Arc>>, - pub(super) refill_states: Arc>>>, + /// Coalesced refill state keyed by exact generation and contour ownership. + pub(super) refill_states: Arc>>, pub(super) refill_running: AtomicUsize, pub(super) refill_pending: AtomicUsize, pub(super) conn_count: AtomicUsize, @@ -516,5 +544,6 @@ mod transport_policy; mod selection_policy; // Bounded writer-open admission and coverage accounting. mod writer_admission; +pub(super) use writer_admission::{WriterOpenIntent, WriterOpenReservation, WriterRole}; // Endpoint-to-DC routing and health timing policy. mod routing; diff --git a/src/transport/middle_proxy/pool/construction.rs b/src/transport/middle_proxy/pool/construction.rs index dd76e56..5eede32 100644 --- a/src/transport/middle_proxy/pool/construction.rs +++ b/src/transport/middle_proxy/pool/construction.rs @@ -120,9 +120,8 @@ impl MePool { me_route_inline_recovery_wait_ms: u64, me_connection_cleanup_capacity: usize, ) -> Arc { - let endpoint_dc_map = Self::build_endpoint_dc_map_from_maps(&proxy_map_v4, &proxy_map_v6); - let preferred_endpoints_by_dc = - Self::build_preferred_endpoints_by_dc(&decision, &proxy_map_v4, &proxy_map_v6); + let endpoint_snapshot = + Self::build_endpoint_snapshot(&decision, proxy_map_v4, proxy_map_v6, 1); let registry = Arc::new(ConnRegistry::with_route_and_cleanup_capacity( me_route_channel_capacity, me_connection_cleanup_capacity, @@ -140,6 +139,7 @@ impl MePool { pending_hardswap_generation: 0, pending_hardswap_started_at_epoch_secs: 0, pending_hardswap_map_hash: 0, + pending_hardswap_endpoint_revision: 0, inflight: 0, }; stats.set_me_writer_byte_budget_limit_bytes(me_writer_byte_budget_bytes); @@ -149,7 +149,7 @@ impl MePool { writers: Arc::new(WritersState::new()), rr: AtomicU64::new(0), writer_epoch, - preferred_endpoints_by_dc: ArcSwap::from_pointee(preferred_endpoints_by_dc), + endpoint_snapshot: ArcSwap::from_pointee(endpoint_snapshot), }), reinit: Arc::new(ReinitCore { generation: AtomicU64::new(1), @@ -164,6 +164,7 @@ impl MePool { next_attempt_id: 1, active_generation: 1, desired_map_hash: 0, + endpoint_revision: 1, pending: None, attempts: HashMap::new(), }), @@ -391,13 +392,11 @@ impl MePool { })), stats, pool_size: 2, - proxy_map_v4: Arc::new(RwLock::new(proxy_map_v4)), - proxy_map_v6: Arc::new(RwLock::new(proxy_map_v6)), - endpoint_dc_map: Arc::new(RwLock::new(endpoint_dc_map)), default_dc: AtomicI32::new(default_dc.unwrap_or(2)), next_writer_id: AtomicU64::new(1), writer_connect_active_reserved: AtomicUsize::new(0), writer_connect_warm_reserved: AtomicUsize::new(0), + writer_replacement_open_reserved: AtomicUsize::new(0), rtt_stats: Arc::new(Mutex::new(HashMap::new())), refill_states: Arc::new(ParkingMutex::new(HashMap::new())), refill_running: AtomicUsize::new(0), diff --git a/src/transport/middle_proxy/pool/routing.rs b/src/transport/middle_proxy/pool/routing.rs index c999730..187a50c 100644 --- a/src/transport/middle_proxy/pool/routing.rs +++ b/src/transport/middle_proxy/pool/routing.rs @@ -1,5 +1,37 @@ use super::*; +impl EndpointSnapshot { + /// Returns authoritative configured endpoints for one exact DC and address family. + pub(in crate::transport::middle_proxy) fn endpoints_for_dc_family( + &self, + dc: i32, + family: IpFamily, + ) -> &[(IpAddr, u16)] { + match family { + IpFamily::V4 => self.map_v4.get(&dc), + IpFamily::V6 => self.map_v6.get(&dc), + } + .map(Vec::as_slice) + .unwrap_or_default() + } + + /// Checks control-plane membership without applying data-plane family preference. + pub(in crate::transport::middle_proxy) fn contains_dc_endpoint( + &self, + dc: i32, + endpoint: SocketAddr, + ) -> bool { + let family = if endpoint.is_ipv4() { + IpFamily::V4 + } else { + IpFamily::V6 + }; + self.endpoints_for_dc_family(dc, family) + .iter() + .any(|(ip, port)| *ip == endpoint.ip() && *port == endpoint.port()) + } +} + impl MePool { pub(in crate::transport::middle_proxy) fn single_endpoint_outage_mode_enabled(&self) -> bool { self.single_endpoint_runtime @@ -76,16 +108,23 @@ impl MePool { &self, dc: i32, ) -> bool { + let snapshot = self.endpoint_snapshot.load(); if self.decision.ipv4_me { - let map = self.proxy_map_v4.read().await; - if map.get(&dc).is_some_and(|endpoints| !endpoints.is_empty()) { + if snapshot + .map_v4 + .get(&dc) + .is_some_and(|endpoints| !endpoints.is_empty()) + { return true; } } if self.decision.ipv6_me { - let map = self.proxy_map_v6.read().await; - if map.get(&dc).is_some_and(|endpoints| !endpoints.is_empty()) { + if snapshot + .map_v6 + .get(&dc) + .is_some_and(|endpoints| !endpoints.is_empty()) + { return true; } } @@ -112,7 +151,12 @@ impl MePool { &self, addr: SocketAddr, ) -> i32 { - if let Some(cached) = self.endpoint_dc_map.read().await.get(&addr).copied() + if let Some(cached) = self + .endpoint_snapshot + .load() + .endpoint_dc_map + .get(&addr) + .copied() && let Some(dc) = cached { return dc; @@ -125,9 +169,42 @@ impl MePool { &self, family: IpFamily, ) -> HashMap> { + let snapshot = self.endpoint_snapshot.load(); match family { - IpFamily::V4 => self.proxy_map_v4.read().await.clone(), - IpFamily::V6 => self.proxy_map_v6.read().await.clone(), + IpFamily::V4 => snapshot.map_v4.clone(), + IpFamily::V6 => snapshot.map_v6.clone(), + } + } + + /// Builds all endpoint-derived indexes under one immutable revision. + pub(in crate::transport::middle_proxy) fn build_endpoint_snapshot( + decision: &NetworkDecision, + mut map_v4: HashMap>, + mut map_v6: HashMap>, + revision: u64, + ) -> EndpointSnapshot { + Self::mirror_negative_dcs(&mut map_v4); + Self::mirror_negative_dcs(&mut map_v6); + let endpoint_dc_map = Self::build_endpoint_dc_map_from_maps(&map_v4, &map_v6); + let preferred_endpoints_by_dc = + Self::build_preferred_endpoints_by_dc(decision, &map_v4, &map_v6); + EndpointSnapshot { + revision, + map_v4, + map_v6, + endpoint_dc_map, + preferred_endpoints_by_dc, + } + } + + fn mirror_negative_dcs(map: &mut HashMap>) { + let positive_dcs = map.keys().copied().filter(|dc| *dc > 0).collect::>(); + for dc in positive_dcs { + if !map.contains_key(&-dc) + && let Some(endpoints) = map.get(&dc).cloned() + { + map.insert(-dc, endpoints); + } } } @@ -224,17 +301,12 @@ impl MePool { endpoint_dc_map } - pub(in crate::transport::middle_proxy) async fn rebuild_endpoint_dc_map(&self) { - let map_v4 = self.proxy_map_v4.read().await.clone(); - let map_v6 = self.proxy_map_v6.read().await.clone(); - let rebuilt = Self::build_endpoint_dc_map_from_maps(&map_v4, &map_v6); - let preferred = Self::build_preferred_endpoints_by_dc(&self.decision, &map_v4, &map_v6); - *self.endpoint_dc_map.write().await = rebuilt; - self.preferred_endpoints_by_dc.store(Arc::new(preferred)); + /// Removes runtime endpoint state absent from the current coherent snapshot. + pub(in crate::transport::middle_proxy) async fn prune_endpoint_runtime_state(&self) { let configured_endpoints = self + .endpoint_snapshot + .load() .endpoint_dc_map - .read() - .await .keys() .copied() .collect::>(); @@ -253,8 +325,12 @@ impl MePool { &self, dc: i32, ) -> Vec { - let guard = self.preferred_endpoints_by_dc.load(); - guard.get(&dc).cloned().unwrap_or_default() + self.endpoint_snapshot + .load() + .preferred_endpoints_by_dc + .get(&dc) + .cloned() + .unwrap_or_default() } pub(in crate::transport::middle_proxy) fn health_interval_unhealthy(&self) -> Duration { diff --git a/src/transport/middle_proxy/pool/writer_admission.rs b/src/transport/middle_proxy/pool/writer_admission.rs index 49b8c01..722eed3 100644 --- a/src/transport/middle_proxy/pool/writer_admission.rs +++ b/src/transport/middle_proxy/pool/writer_admission.rs @@ -1,102 +1,177 @@ use super::*; +const WRITER_REPLACEMENT_OPEN_LIMIT_MAX: usize = 128; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +/// Immutable identity of a writer lifecycle role at replacement reservation time. +pub(in crate::transport::middle_proxy) struct WriterRole { + /// Telegram DC owning the writer. + pub(in crate::transport::middle_proxy) dc: i32, + /// Address family of the writer endpoint. + pub(in crate::transport::middle_proxy) family: IpFamily, + /// Pool generation owning the writer. + pub(in crate::transport::middle_proxy) generation: u64, + /// Lifecycle contour assigned to the writer. + pub(in crate::transport::middle_proxy) contour: WriterContour, +} + +impl WriterRole { + /// Captures the current role of an installed writer. + pub(in crate::transport::middle_proxy) fn from_writer(writer: &MeWriter) -> Self { + Self { + dc: writer.writer_dc, + family: if writer.addr.is_ipv4() { + IpFamily::V4 + } else { + IpFamily::V6 + }, + generation: writer.generation, + contour: WriterContour::from_u8(writer.contour.load(Ordering::Acquire)), + } + } + + /// Revalidates that an installed writer still has this exact role. + pub(in crate::transport::middle_proxy) fn matches(self, writer: &MeWriter) -> bool { + self == Self::from_writer(writer) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +/// Capacity policy applied while opening a writer. +pub(in crate::transport::middle_proxy) enum WriterOpenIntent { + /// Ordinary pool growth constrained by the configured contour cap. + Normal, + /// Required active coverage allowed to exceed an undersized configured cap temporarily. + Coverage, + /// Replacement-before-drain capacity owned by an existing victim reservation. + Replacement, +} + +/// RAII ownership of one bounded in-flight writer open. +pub(in crate::transport::middle_proxy) struct WriterOpenReservation<'a> { + counter: Option<&'a AtomicUsize>, +} + +impl Drop for WriterOpenReservation<'_> { + fn drop(&mut self) { + if let Some(counter) = self.counter { + counter.fetch_sub(1, Ordering::AcqRel); + } + } +} + impl MePool { + /// Computes the authoritative active-writer floor across enabled families and DCs. pub(in crate::transport::middle_proxy) async fn active_coverage_required_total(&self) -> usize { let now_epoch_secs = Self::now_epoch_secs(); - let mut endpoints_by_dc = HashMap::>::new(); + let mut required_total = 0usize; + let endpoint_snapshot = self.endpoint_snapshot.load_full(); if self.family_enabled_for_drain_coverage(IpFamily::V4, now_epoch_secs) { - let map = self.proxy_map_v4.read().await; - for (dc, addrs) in map.iter() { - let entry = endpoints_by_dc.entry(*dc).or_default(); + for addrs in endpoint_snapshot.map_v4.values() { + let mut endpoints = HashSet::::new(); for (ip, port) in addrs.iter().copied() { - entry.insert(SocketAddr::new(ip, port)); + endpoints.insert(SocketAddr::new(ip, port)); } + required_total = required_total.saturating_add( + self.required_writers_for_dc_with_floor_mode(endpoints.len(), false), + ); } } if self.family_enabled_for_drain_coverage(IpFamily::V6, now_epoch_secs) { - let map = self.proxy_map_v6.read().await; - for (dc, addrs) in map.iter() { - let entry = endpoints_by_dc.entry(*dc).or_default(); + for addrs in endpoint_snapshot.map_v6.values() { + let mut endpoints = HashSet::::new(); for (ip, port) in addrs.iter().copied() { - entry.insert(SocketAddr::new(ip, port)); + endpoints.insert(SocketAddr::new(ip, port)); } + required_total = required_total.saturating_add( + self.required_writers_for_dc_with_floor_mode(endpoints.len(), false), + ); } } - endpoints_by_dc - .values() - .map(|endpoints| self.required_writers_for_dc_with_floor_mode(endpoints.len(), false)) - .sum() + required_total } + /// Reports whether one writer may be opened under the selected contour policy. pub(in crate::transport::middle_proxy) async fn can_open_writer_for_contour( &self, contour: WriterContour, - allow_coverage_override: bool, + intent: WriterOpenIntent, writer_dc: i32, + target_addr: SocketAddr, ) -> bool { - let (active_writers, warm_writers, _) = self.non_draining_writer_counts_by_contour().await; - match contour { - WriterContour::Active => { - let active_cap = self.adaptive_floor_active_cap_configured_total(); - if active_writers < active_cap { - return true; - } - if !allow_coverage_override { - return false; - } - - let mut endpoints_len = 0; - let now_epoch = Self::now_epoch_secs(); - if self.family_enabled_for_drain_coverage(IpFamily::V4, now_epoch) { - if let Some(addrs) = self.proxy_map_v4.read().await.get(&writer_dc) { - endpoints_len += addrs.len(); - } - } - if self.family_enabled_for_drain_coverage(IpFamily::V6, now_epoch) { - if let Some(addrs) = self.proxy_map_v6.read().await.get(&writer_dc) { - endpoints_len += addrs.len(); - } - } - - if endpoints_len > 0 { - let base_req = - self.required_writers_for_dc_with_floor_mode(endpoints_len, false); - let active_for_dc = { - let ws = self.writers.read().await; - ws.iter() - .filter(|w| { - !w.draining.load(std::sync::atomic::Ordering::Relaxed) - && w.writer_dc == writer_dc - && matches!( - WriterContour::from_u8( - w.contour.load(std::sync::atomic::Ordering::Relaxed), - ), - WriterContour::Active - ) - }) - .count() - }; - if active_for_dc < base_req { - return true; - } - } - - let coverage_required = self.active_coverage_required_total().await; - active_writers < coverage_required - } - WriterContour::Warm => warm_writers < self.adaptive_floor_warm_cap_configured_total(), - WriterContour::Draining => true, + if intent == WriterOpenIntent::Replacement { + return true; } + let family = if target_addr.is_ipv4() { + IpFamily::V4 + } else { + IpFamily::V6 + }; + let endpoint_snapshot = self.endpoint_snapshot.load_full(); + let endpoints = endpoint_snapshot.endpoints_for_dc_family(writer_dc, family); + if !endpoint_snapshot.contains_dc_endpoint(writer_dc, target_addr) { + return false; + } + let (active_writers, warm_writers, _) = self.non_draining_writer_counts_by_contour().await; + let live = match contour { + WriterContour::Active => active_writers, + WriterContour::Warm => warm_writers, + WriterContour::Draining => return true, + }; + let configured_cap = match contour { + WriterContour::Active => self.adaptive_floor_active_cap_configured_total(), + WriterContour::Warm => self.adaptive_floor_warm_cap_configured_total(), + WriterContour::Draining => usize::MAX, + }; + if live < configured_cap { + return true; + } + if intent != WriterOpenIntent::Coverage { + return false; + } + + let endpoint_count = endpoints.len(); + if endpoint_count == 0 { + return false; + } + let required = self.required_writers_for_dc_with_floor_mode(endpoint_count, false); + let status = self.reinit.status.load(); + let generation = match contour { + WriterContour::Active => status.active_generation, + WriterContour::Warm => status.pending_hardswap_generation, + WriterContour::Draining => 0, + }; + let family_count = { + let writers = self.writers.read().await; + writers + .iter() + .filter(|writer| { + !writer.draining.load(Ordering::Relaxed) + && writer.writer_dc == writer_dc + && writer.generation == generation + && WriterContour::from_u8(writer.contour.load(Ordering::Relaxed)) == contour + && writer.addr.is_ipv4() == (family == IpFamily::V4) + && endpoint_snapshot.contains_dc_endpoint(writer_dc, writer.addr) + }) + .count() + }; + if family_count < required { + return true; + } + + false } + /// Reserves bounded transient capacity for a writer open attempt. pub(in crate::transport::middle_proxy) async fn reserve_writer_open( &self, contour: WriterContour, - allow_coverage_override: bool, + intent: WriterOpenIntent, writer_dc: i32, + target_addr: SocketAddr, ) -> Option> { let counter = match contour { WriterContour::Active => &self.writer_connect_active_reserved, @@ -106,9 +181,54 @@ impl MePool { } }; + if intent == WriterOpenIntent::Replacement { + let configured_cap = match contour { + WriterContour::Active => self.adaptive_floor_active_cap_configured_total(), + WriterContour::Warm => self.adaptive_floor_warm_cap_configured_total(), + WriterContour::Draining => usize::MAX, + }; + let effective_cap = match contour { + WriterContour::Active => self + .floor_runtime + .me_adaptive_floor_active_cap_effective + .load(Ordering::Acquire) as usize, + WriterContour::Warm => self + .floor_runtime + .me_adaptive_floor_warm_cap_effective + .load(Ordering::Acquire) as usize, + WriterContour::Draining => usize::MAX, + }; + let replacement_limit = configured_cap + .max(effective_cap) + .max(1) + .min(WRITER_REPLACEMENT_OPEN_LIMIT_MAX); + loop { + let reserved = self + .writer_replacement_open_reserved + .load(Ordering::Acquire); + if reserved >= replacement_limit { + return None; + } + if self + .writer_replacement_open_reserved + .compare_exchange_weak( + reserved, + reserved + 1, + Ordering::AcqRel, + Ordering::Acquire, + ) + .is_ok() + { + return Some(WriterOpenReservation { + counter: Some(&self.writer_replacement_open_reserved), + }); + } + } + } + loop { if !self - .can_open_writer_for_contour(contour, allow_coverage_override, writer_dc) + .can_open_writer_for_contour(contour, intent, writer_dc, target_addr) .await { return None; @@ -125,7 +245,9 @@ impl MePool { WriterContour::Warm => self.adaptive_floor_warm_cap_configured_total(), WriterContour::Draining => usize::MAX, }; - if contour == WriterContour::Active && allow_coverage_override { + if intent == WriterOpenIntent::Coverage + && matches!(contour, WriterContour::Active | WriterContour::Warm) + { limit = limit .max(self.active_coverage_required_total().await) .saturating_add( @@ -150,6 +272,7 @@ impl MePool { } } + /// Resolves a DC writer floor for static or adaptive-idle operation. pub(in crate::transport::middle_proxy) fn required_writers_for_dc_with_floor_mode( &self, endpoint_count: usize, diff --git a/src/transport/middle_proxy/pool_config.rs b/src/transport/middle_proxy/pool_config.rs index 6e29918..280f45b 100644 --- a/src/transport/middle_proxy/pool_config.rs +++ b/src/transport/middle_proxy/pool_config.rs @@ -5,7 +5,8 @@ use std::time::Duration; use tracing::warn; -use super::pool::MePool; +use super::pool::{MePool, WriterContour, WriterRole}; +use super::pool_writer::WriterReplacementPurpose; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum SnapshotApplyOutcome { @@ -30,48 +31,35 @@ impl MePool { return SnapshotApplyOutcome::RejectedEmpty; } - let mut changed = false; - { - let mut guard = self.proxy_map_v4.write().await; - if !new_v4.is_empty() && *guard != new_v4 { - *guard = new_v4; - changed = true; + let changed = { + // Endpoint publication and reinit commit share this barrier. + let mut coordinator = self.reinit.coordinator.lock(); + let current = self.endpoint_snapshot.load_full(); + let map_v4 = if new_v4.is_empty() { + current.map_v4.clone() + } else { + new_v4 + }; + let map_v6 = match new_v6 { + Some(map) if !map.is_empty() => map, + _ => current.map_v6.clone(), + }; + let candidate = Self::build_endpoint_snapshot( + &self.decision, + map_v4, + map_v6, + current.revision.saturating_add(1), + ); + if candidate.map_v4 == current.map_v4 && candidate.map_v6 == current.map_v6 { + false + } else { + coordinator.endpoint_revision = candidate.revision; + self.endpoint_snapshot.store(Arc::new(candidate)); + true } - } - if let Some(v6) = new_v6 { - let mut guard = self.proxy_map_v6.write().await; - if !v6.is_empty() && *guard != v6 { - *guard = v6; - changed = true; - } - } - // Ensure negative DC entries mirror positives when absent (Telegram convention). - { - let mut guard = self.proxy_map_v4.write().await; - let keys: Vec = guard.keys().cloned().collect(); - for k in keys.iter().cloned().filter(|k| *k > 0) { - if !guard.contains_key(&-k) - && let Some(addrs) = guard.get(&k).cloned() - { - guard.insert(-k, addrs); - changed = true; - } - } - } - { - let mut guard = self.proxy_map_v6.write().await; - let keys: Vec = guard.keys().cloned().collect(); - for k in keys.iter().cloned().filter(|k| *k > 0) { - if !guard.contains_key(&-k) - && let Some(addrs) = guard.get(&k).cloned() - { - guard.insert(-k, addrs); - changed = true; - } - } - } + }; if changed { - self.rebuild_endpoint_dc_map().await; + self.prune_endpoint_runtime_state().await; self.notify_writer_epoch(); } if changed { @@ -113,11 +101,33 @@ impl MePool { pub async fn reconnect_all(self: &Arc) { let ws = self.writers.read().await.clone(); for w in ws.iter() { - if let Ok(()) = self - .connect_one_for_dc(w.addr, w.writer_dc, self.rng.as_ref()) - .await + let role = WriterRole::from_writer(w); + if w.draining.load(std::sync::atomic::Ordering::Acquire) + || role.contour == WriterContour::Draining + { + continue; + } + let Some(mut reservation) = self + .registry + .try_reserve_writer_replacement_preserving_clients(w.id) + .await + else { + continue; + }; + if self + .replace_writer_with_generation_contour_for_dc( + w.addr, + self.rng.as_ref(), + role.generation, + role.contour, + role.dc, + role, + WriterReplacementPurpose::SecretRotation, + &mut reservation, + ) + .await + .is_ok() { - self.mark_writer_draining(w.id).await; tokio::time::sleep(Duration::from_secs(2)).await; } } diff --git a/src/transport/middle_proxy/pool_init.rs b/src/transport/middle_proxy/pool_init.rs index 0936246..4a17298 100644 --- a/src/transport/middle_proxy/pool_init.rs +++ b/src/transport/middle_proxy/pool_init.rs @@ -9,7 +9,7 @@ use tracing::{debug, info, warn}; use crate::crypto::SecureRandom; use crate::error::{ProxyError, Result}; -use super::pool::MePool; +use super::pool::{MePool, WriterOpenIntent}; impl MePool { pub async fn init(self: &Arc, pool_size: usize, rng: &Arc) -> Result<()> { @@ -19,7 +19,7 @@ impl MePool { .me_reconnect_max_concurrent_per_dc .max(1) as usize; let ks = self.key_selector().await; - let me_servers = self.proxy_map_v4.read().await.len(); + let me_servers = self.endpoint_snapshot.load().map_v4.len(); let secret_len = self.proxy_secret.read().await.secret.len(); info!( me_servers, @@ -76,7 +76,7 @@ impl MePool { target_writers, rng_clone, connect_concurrency, - true, + WriterOpenIntent::Coverage, ) .await }); @@ -125,7 +125,7 @@ impl MePool { target_writers, rng_clone_local, connect_concurrency, - false, + WriterOpenIntent::Normal, ) .await }); @@ -162,7 +162,7 @@ impl MePool { target_writers: usize, rng: Arc, connect_concurrency: usize, - allow_coverage_override: bool, + intent: WriterOpenIntent, ) -> bool { if addrs.is_empty() { return false; @@ -204,7 +204,7 @@ impl MePool { rng_clone.as_ref(), generation, super::pool::WriterContour::Active, - allow_coverage_override, + intent, ) .await }); @@ -238,7 +238,9 @@ impl MePool { if !progress { let active_writers_current = self.active_contour_writer_count_total().await; let active_cap_configured = self.adaptive_floor_active_cap_configured_total(); - if !allow_coverage_override && active_writers_current >= active_cap_configured { + if intent != WriterOpenIntent::Coverage + && active_writers_current >= active_cap_configured + { info!( dc = %dc, alive = alive_after, diff --git a/src/transport/middle_proxy/pool_refill.rs b/src/transport/middle_proxy/pool_refill.rs index 30a7809..2f9e0f5 100644 --- a/src/transport/middle_proxy/pool_refill.rs +++ b/src/transport/middle_proxy/pool_refill.rs @@ -9,26 +9,49 @@ use tracing::{debug, info, warn}; use crate::crypto::SecureRandom; use crate::network::IpFamily; -use super::pool::{MePool, RefillDcKey, WriterContour}; +use super::pool::{ + MePool, RefillTargetKey, RefillTargetState, WriterContour, WriterOpenIntent, WriterRole, +}; const ME_FLAP_UPTIME_THRESHOLD_SECS: u64 = 20; const ME_FLAP_QUARANTINE_SECS: u64 = 25; const ME_FLAP_MIN_UPTIME_MILLIS: u64 = 500; const ME_REFILL_TOTAL_ATTEMPT_CAP: u32 = 20; +const ME_REFILL_PENDING_PER_TARGET_MAX: usize = 128; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum RefillOutcome { + Restored, + Failed, + Obsolete, +} + +fn refill_open_intent(contour: WriterContour) -> WriterOpenIntent { + match contour { + WriterContour::Active | WriterContour::Warm => WriterOpenIntent::Coverage, + WriterContour::Draining => WriterOpenIntent::Normal, + } +} struct RefillRunGuard { pool: Arc, - key: RefillDcKey, + key: RefillTargetKey, active: bool, } impl RefillRunGuard { fn next_or_finish(&mut self) -> Option { let mut states = self.pool.refill_states.lock(); - let next = states.get_mut(&self.key).and_then(Option::take); - if next.is_some() { + let next = states.get_mut(&self.key).and_then(|state| { + if state.pending_count == 0 { + return None; + } + state.pending_count -= 1; + state.next_addr + }); + if let Some(next) = next { self.pool.refill_pending.fetch_sub(1, Ordering::AcqRel); - return next; + return Some(next); } states.remove(&self.key); self.pool.refill_running.fetch_sub(1, Ordering::AcqRel); @@ -42,10 +65,10 @@ impl Drop for RefillRunGuard { if !self.active { return; } - if let Some(pending) = self.pool.refill_states.lock().remove(&self.key) - && pending.is_some() - { - self.pool.refill_pending.fetch_sub(1, Ordering::AcqRel); + if let Some(state) = self.pool.refill_states.lock().remove(&self.key) { + self.pool + .refill_pending + .fetch_sub(state.pending_count, Ordering::AcqRel); } self.pool.refill_running.fetch_sub(1, Ordering::AcqRel); } @@ -54,9 +77,9 @@ impl Drop for RefillRunGuard { impl MePool { pub(super) async fn sweep_endpoint_quarantine(&self) { let configured = self + .endpoint_snapshot + .load() .endpoint_dc_map - .read() - .await .keys() .copied() .collect::>(); @@ -165,7 +188,8 @@ impl MePool { self.connectable_endpoints(endpoints).await } - pub(super) async fn has_refill_inflight_for_dc_key(&self, key: RefillDcKey) -> bool { + /// Reports whether the exact generation and contour already has a refill producer. + pub(super) async fn has_refill_inflight_for_target(&self, key: RefillTargetKey) -> bool { self.refill_states.lock().contains_key(&key) } @@ -181,7 +205,7 @@ impl MePool { rng, self.current_generation(), WriterContour::Active, - false, + WriterOpenIntent::Normal, ) .await } @@ -193,14 +217,14 @@ impl MePool { rng: &SecureRandom, generation: u64, contour: WriterContour, - allow_coverage_override: bool, + intent: WriterOpenIntent, ) -> bool { let mut candidates = self.connectable_endpoints(endpoints).await; if candidates.is_empty() { return false; } if candidates.len() > 1 { - let mut active_by_endpoint = HashMap::::new(); + let mut matching_by_endpoint = HashMap::::new(); let ws = self.writers.read().await; for writer in ws.iter() { if writer.draining.load(Ordering::Relaxed) { @@ -209,32 +233,26 @@ impl MePool { if writer.writer_dc != dc { continue; } - if !matches!( - super::pool::WriterContour::from_u8(writer.contour.load(Ordering::Relaxed),), - super::pool::WriterContour::Active - ) { + if writer.generation != generation + || WriterContour::from_u8(writer.contour.load(Ordering::Acquire)) != contour + { continue; } if candidates.contains(&writer.addr) { - *active_by_endpoint.entry(writer.addr).or_insert(0) += 1; + *matching_by_endpoint.entry(writer.addr).or_insert(0) += 1; } } drop(ws); candidates - .sort_by_key(|addr| (active_by_endpoint.get(addr).copied().unwrap_or(0), *addr)); + .sort_by_key(|addr| (matching_by_endpoint.get(addr).copied().unwrap_or(0), *addr)); } let start = (self.rr.fetch_add(1, Ordering::Relaxed) as usize) % candidates.len(); for offset in 0..candidates.len() { let idx = (start + offset) % candidates.len(); let addr = candidates[idx]; match self - .connect_one_with_generation_contour_for_dc_with_cap_policy( - addr, - rng, - generation, - contour, - dc, - allow_coverage_override, + .connect_one_with_generation_contour_for_dc_with_intent( + addr, rng, generation, contour, dc, intent, ) .await { @@ -245,43 +263,71 @@ impl MePool { false } - async fn endpoints_for_dc(&self, target_dc: i32) -> Vec { + async fn endpoints_for_refill_target(&self, target: RefillTargetKey) -> Vec { let now_epoch_secs = Self::now_epoch_secs(); - let mut endpoints = HashSet::::new(); - - if self.family_enabled_for_drain_coverage(IpFamily::V4, now_epoch_secs) { - let map = self.proxy_map_v4.read().await; - if let Some(addrs) = map.get(&target_dc) { - for (ip, port) in addrs { - endpoints.insert(SocketAddr::new(*ip, *port)); - } - } + if !self.family_enabled_for_drain_coverage(target.family, now_epoch_secs) { + return Vec::new(); } - - if self.family_enabled_for_drain_coverage(IpFamily::V6, now_epoch_secs) { - let map = self.proxy_map_v6.read().await; - if let Some(addrs) = map.get(&target_dc) { - for (ip, port) in addrs { - endpoints.insert(SocketAddr::new(*ip, *port)); - } - } - } - - let mut sorted: Vec = endpoints.into_iter().collect(); - sorted.sort_unstable(); - sorted + let snapshot = self.endpoint_snapshot.load(); + let map = match target.family { + IpFamily::V4 => &snapshot.map_v4, + IpFamily::V6 => &snapshot.map_v6, + }; + let mut endpoints = map + .get(&target.dc) + .into_iter() + .flatten() + .map(|(ip, port)| SocketAddr::new(*ip, *port)) + .collect::>(); + endpoints.sort_unstable(); + endpoints.dedup(); + endpoints } - async fn refill_writer_after_loss(self: &Arc, addr: SocketAddr, writer_dc: i32) -> bool { + fn refill_target_is_authoritative(&self, target: RefillTargetKey) -> bool { + if !self.family_enabled_for_drain_coverage(target.family, Self::now_epoch_secs()) { + return false; + } + let status = self.reinit.status.load(); + let role_is_authoritative = match target.contour { + WriterContour::Active => target.generation == status.active_generation, + WriterContour::Warm => { + status.pending_hardswap_generation != 0 + && target.generation == status.pending_hardswap_generation + } + WriterContour::Draining => false, + }; + let pending_revision_matches = target.contour != WriterContour::Warm + || target.endpoint_revision == status.pending_hardswap_endpoint_revision; + let endpoint_snapshot = self.endpoint_snapshot.load(); + role_is_authoritative + && pending_revision_matches + && target.endpoint_revision == endpoint_snapshot.revision + && !endpoint_snapshot + .endpoints_for_dc_family(target.dc, target.family) + .is_empty() + } + + async fn refill_writer_after_loss( + self: &Arc, + addr: SocketAddr, + target: RefillTargetKey, + ) -> RefillOutcome { + if !self.refill_target_is_authoritative(target) { + return RefillOutcome::Obsolete; + } + let open_intent = refill_open_intent(target.contour); let fast_retries = self.reconnect_runtime.me_reconnect_fast_retry_count.max(1); let mut total_attempts = 0u32; let same_endpoint_quarantined = self.is_endpoint_quarantined(addr).await; - let dc_endpoints = self.endpoints_for_dc(writer_dc).await; + let dc_endpoints = self.endpoints_for_refill_target(target).await; let single_endpoint_dc = dc_endpoints.len() == 1 && dc_endpoints[0] == addr; let bypass_quarantine_for_single_endpoint = single_endpoint_dc && self.single_endpoint_outage_disable_quarantine(); - if !same_endpoint_quarantined || bypass_quarantine_for_single_endpoint { + if dc_endpoints.contains(&addr) + && (!same_endpoint_quarantined || bypass_quarantine_for_single_endpoint) + { if same_endpoint_quarantined && bypass_quarantine_for_single_endpoint { debug!( %addr, @@ -289,13 +335,23 @@ impl MePool { ); } for attempt in 0..fast_retries { + if !self.refill_target_is_authoritative(target) { + return RefillOutcome::Obsolete; + } if total_attempts >= ME_REFILL_TOTAL_ATTEMPT_CAP { break; } total_attempts = total_attempts.saturating_add(1); self.stats.increment_me_reconnect_attempt(); match self - .connect_one_for_dc(addr, writer_dc, self.rng.as_ref()) + .connect_one_with_generation_contour_for_dc_with_intent( + addr, + self.rng.as_ref(), + target.generation, + target.contour, + target.dc, + open_intent, + ) .await { Ok(()) => { @@ -307,7 +363,7 @@ impl MePool { attempt = attempt + 1, "ME writer restored on the same endpoint" ); - return true; + return RefillOutcome::Restored; } Err(e) => { debug!( @@ -328,17 +384,27 @@ impl MePool { if dc_endpoints.is_empty() { self.stats.increment_me_refill_failed_total(); - return false; + return RefillOutcome::Failed; } for attempt in 0..fast_retries { + if !self.refill_target_is_authoritative(target) { + return RefillOutcome::Obsolete; + } if total_attempts >= ME_REFILL_TOTAL_ATTEMPT_CAP { break; } total_attempts = total_attempts.saturating_add(1); self.stats.increment_me_reconnect_attempt(); if self - .connect_endpoints_round_robin(writer_dc, &dc_endpoints, self.rng.as_ref()) + .connect_endpoints_round_robin_with_generation_contour( + target.dc, + &dc_endpoints, + self.rng.as_ref(), + target.generation, + target.contour, + open_intent, + ) .await { self.stats.increment_me_reconnect_success(); @@ -348,12 +414,12 @@ impl MePool { attempt = attempt + 1, "ME writer restored via DC fallback endpoint" ); - return true; + return RefillOutcome::Restored; } } self.stats.increment_me_refill_failed_total(); - false + RefillOutcome::Failed } pub(crate) fn trigger_immediate_refill_for_dc( @@ -361,35 +427,59 @@ impl MePool { addr: SocketAddr, writer_dc: i32, ) { + self.trigger_immediate_refill_for_role( + addr, + WriterRole { + dc: writer_dc, + family: if addr.is_ipv4() { + IpFamily::V4 + } else { + IpFamily::V6 + }, + generation: self.current_generation(), + contour: WriterContour::Active, + }, + ); + } + + /// Coalesces unexpected writer loss without changing its generation or contour role. + pub(super) fn trigger_immediate_refill_for_role( + self: &Arc, + addr: SocketAddr, + role: WriterRole, + ) { + let target = RefillTargetKey { + dc: role.dc, + family: role.family, + generation: role.generation, + endpoint_revision: self.endpoint_snapshot.load().revision, + contour: role.contour, + }; + if !self.refill_target_is_authoritative(target) { + return; + } let Some(registration) = self.lifecycle.try_register() else { return; }; - let dc_key = RefillDcKey { - dc: writer_dc, - family: if addr.is_ipv4() { - IpFamily::V4 - } else { - IpFamily::V6 - }, - }; { let mut states = self.refill_states.lock(); - if let Some(pending) = states.get_mut(&dc_key) { - if pending.is_none() { + if let Some(state) = states.get_mut(&target) { + if state.pending_count < ME_REFILL_PENDING_PER_TARGET_MAX { + state.pending_count += 1; + state.next_addr = Some(addr); self.refill_pending.fetch_add(1, Ordering::AcqRel); } - *pending = Some(addr); self.stats.increment_me_refill_skipped_inflight_total(); return; } - states.insert(dc_key, None); + states.insert(target, RefillTargetState::default()); self.refill_running.fetch_add(1, Ordering::AcqRel); } let pool = Arc::clone(self); let mut run_guard = RefillRunGuard { pool: Arc::clone(&pool), - key: dc_key, + key: target, active: true, }; self.lifecycle @@ -397,9 +487,24 @@ impl MePool { let mut current_addr = addr; loop { pool.stats.increment_me_refill_triggered_total(); - let restored = pool.refill_writer_after_loss(current_addr, writer_dc).await; - if !restored { - warn!(%current_addr, dc = writer_dc, "ME immediate refill failed"); + let outcome = pool.refill_writer_after_loss(current_addr, target).await; + if outcome == RefillOutcome::Failed { + warn!( + %current_addr, + dc = target.dc, + generation = target.generation, + contour = ?target.contour, + "ME immediate refill failed" + ); + } else if outcome == RefillOutcome::Obsolete { + debug!( + %current_addr, + dc = target.dc, + generation = target.generation, + contour = ?target.contour, + "ME immediate refill target is no longer authoritative" + ); + return; } let Some(next_addr) = run_guard.next_or_finish() else { @@ -410,3 +515,16 @@ impl MePool { }); } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn warm_refill_preserves_required_floor_coverage() { + assert_eq!( + refill_open_intent(WriterContour::Warm), + WriterOpenIntent::Coverage + ); + } +} diff --git a/src/transport/middle_proxy/pool_reinit.rs b/src/transport/middle_proxy/pool_reinit.rs index 46174cd..7af140a 100644 --- a/src/transport/middle_proxy/pool_reinit.rs +++ b/src/transport/middle_proxy/pool_reinit.rs @@ -10,12 +10,13 @@ use rand::seq::SliceRandom; use std::collections::hash_map::DefaultHasher; use tracing::{debug, info, warn}; +use crate::config::MeBindStaleMode; use crate::crypto::SecureRandom; use crate::network::IpFamily; use super::pool::{ - MeDrainGateReason, MePool, ReinitAttemptState, ReinitCoordinatorState, ReinitCore, - ReinitPendingState, ReinitStatusSnapshot, WriterContour, + EndpointSnapshot, MeDrainGateReason, MePool, ReinitAttemptState, ReinitCoordinatorState, + ReinitCore, ReinitPendingState, ReinitStatusSnapshot, WriterContour, WriterOpenIntent, }; // Reinitialization admission, generation state, and coverage checks. @@ -23,6 +24,8 @@ mod coordination; // Generation warmup and stale-writer reconciliation. mod reconcile; +#[cfg(test)] +mod dual_family_tests; #[cfg(test)] mod tests; const ME_HARDSWAP_PENDING_TTL_SECS: u64 = 1800; @@ -33,6 +36,7 @@ struct ReinitAttemptGuard { generation: u64, previous_generation: u64, map_hash: u64, + endpoint_revision: u64, hardswap: bool, } @@ -51,6 +55,48 @@ struct ReinitReservation { pending_age_secs: u64, } +struct ReinitCommitOutcome { + coverage_ratio: f32, + missing_dc: Vec, + missing_groups: Vec, + stale_writer_ids: Vec, + force_close_writer_ids: Vec, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +/// One independently enforced hardswap writer-floor group. +pub(in crate::transport::middle_proxy) struct DcFamilyGroup { + /// Telegram DC owning the group. + pub(in crate::transport::middle_proxy) dc: i32, + /// Address family whose floor is evaluated independently. + pub(in crate::transport::middle_proxy) family: IpFamily, +} + +/// Complete floor-coverage result for one candidate hardswap generation. +pub(in crate::transport::middle_proxy) struct HardswapCoverage { + /// Fraction of configured DC-family groups that reached their full floor. + pub(in crate::transport::middle_proxy) ratio: f32, + /// Stable list of DC-family groups that remain below their floor. + pub(in crate::transport::middle_proxy) missing_groups: Vec, + /// Total writer count still required across all missing groups. + pub(in crate::transport::middle_proxy) writer_deficit: usize, +} + +#[derive(Debug)] +enum ReinitCommitFailure { + Superseded, + Coverage { + coverage_ratio: f32, + missing_dc: Vec, + missing_groups: Vec, + }, + Redundancy { + coverage_ratio: f32, + missing_dc: Vec, + missing_groups: Vec, + }, +} + fn publish_reinit_state(reinit: &ReinitCore, state: &ReinitCoordinatorState) { let mut warm_generations = state .attempts @@ -58,6 +104,9 @@ fn publish_reinit_state(reinit: &ReinitCore, state: &ReinitCoordinatorState) { .filter(|attempt| attempt.hardswap && !attempt.committed) .map(|attempt| attempt.generation) .collect::>(); + if let Some(pending) = state.pending { + warm_generations.push(pending.generation); + } warm_generations.sort_unstable(); warm_generations.dedup(); let pending = state.pending; @@ -68,6 +117,7 @@ fn publish_reinit_state(reinit: &ReinitCore, state: &ReinitCoordinatorState) { pending_hardswap_started_at_epoch_secs: pending .map_or(0, |value| value.started_at_epoch_secs), pending_hardswap_map_hash: pending.map_or(0, |value| value.map_hash), + pending_hardswap_endpoint_revision: pending.map_or(0, |value| value.endpoint_revision), inflight: state.attempts.len(), }; reinit @@ -95,17 +145,24 @@ fn commit_reinit_state( attempt_id: u64, generation: u64, map_hash: u64, + endpoint_revision: u64, hardswap: bool, ) -> bool { let Some(record) = state.attempts.get(&attempt_id).copied() else { return false; }; - if record.map_hash != state.desired_map_hash || record.map_hash != map_hash { + if record.map_hash != state.desired_map_hash + || record.map_hash != map_hash + || record.endpoint_revision != endpoint_revision + || record.endpoint_revision != state.endpoint_revision + { return false; } if hardswap { let pending_matches = state.pending.is_some_and(|pending| { - pending.generation == generation && pending.map_hash == map_hash + pending.generation == generation + && pending.map_hash == map_hash + && pending.endpoint_revision == endpoint_revision }); if !pending_matches || generation < state.active_generation { return false; diff --git a/src/transport/middle_proxy/pool_reinit/coordination.rs b/src/transport/middle_proxy/pool_reinit/coordination.rs index 64e79d3..d928fde 100644 --- a/src/transport/middle_proxy/pool_reinit/coordination.rs +++ b/src/transport/middle_proxy/pool_reinit/coordination.rs @@ -1,7 +1,10 @@ use super::*; impl MePool { - pub(super) fn desired_map_hash(desired_by_dc: &HashMap>) -> u64 { + /// Hashes the sorted desired endpoint map for generation authority checks. + pub(in crate::transport::middle_proxy) fn desired_map_hash( + desired_by_dc: &HashMap>, + ) -> u64 { let mut hasher = DefaultHasher::new(); let mut dcs: Vec = desired_by_dc.keys().copied().collect(); dcs.sort_unstable(); @@ -19,13 +22,18 @@ impl MePool { hasher.finish() } + /// Reserves one generation attempt and publishes its pending ownership snapshot. pub(super) fn reserve_reinit_attempt( self: &Arc, hardswap: bool, map_hash: u64, + endpoint_revision: u64, now_epoch_secs: u64, - ) -> ReinitReservation { + ) -> Option { let mut state = self.reinit.coordinator.lock(); + if state.endpoint_revision != endpoint_revision { + return None; + } state.desired_map_hash = map_hash; let previous_generation = state.active_generation; let mut pending_reused = false; @@ -39,6 +47,7 @@ impl MePool { && pending_age_secs > ME_HARDSWAP_PENDING_TTL_SECS; pending.generation >= previous_generation && pending.map_hash == map_hash + && pending.endpoint_revision == endpoint_revision && !pending_expired }); if let Some(pending) = reusable { @@ -50,6 +59,7 @@ impl MePool { generation, started_at_epoch_secs: now_epoch_secs, map_hash, + endpoint_revision, }); generation } @@ -65,41 +75,114 @@ impl MePool { ReinitAttemptState { generation, map_hash, + endpoint_revision, hardswap, committed: false, }, ); publish_reinit_state(self.reinit.as_ref(), &state); - ReinitReservation { + Some(ReinitReservation { attempt: ReinitAttemptGuard { reinit: Arc::clone(&self.reinit), attempt_id, generation, previous_generation, map_hash, + endpoint_revision, hardswap, }, pending_reused, pending_expired, pending_age_secs, - } + }) } - pub(super) fn commit_reinit_attempt(&self, attempt: &ReinitAttemptGuard) -> bool { + /// Revalidates coverage and commits generation ownership under the publication barrier. + pub(super) async fn commit_reinit_attempt( + &self, + attempt: &ReinitAttemptGuard, + desired_by_dc: &HashMap>, + min_ratio: f32, + ) -> std::result::Result { + let writers = self.writers.write().await; + let mut registry_registration = self.registry.prepare_writer_registration().await; let mut state = self.reinit.coordinator.lock(); + let Some(record) = state.attempts.get(&attempt.attempt_id).copied() else { + return Err(ReinitCommitFailure::Superseded); + }; + if record.generation != attempt.generation + || record.map_hash != state.desired_map_hash + || record.map_hash != attempt.map_hash + || record.endpoint_revision != attempt.endpoint_revision + || record.endpoint_revision != state.endpoint_revision + || (attempt.hardswap + && !state.pending.is_some_and(|pending| { + pending.generation == attempt.generation + && pending.map_hash == attempt.map_hash + && pending.endpoint_revision == attempt.endpoint_revision + })) + { + return Err(ReinitCommitFailure::Superseded); + } + + let authoritative_writer_addrs = writers + .iter() + .filter(|writer| !writer.draining.load(Ordering::Acquire)) + .filter(|writer| { + if attempt.hardswap { + writer.generation == attempt.generation + } else { + writer.generation == state.active_generation + && WriterContour::from_u8(writer.contour.load(Ordering::Acquire)) + == WriterContour::Active + } + }) + .map(|writer| (writer.writer_dc, writer.addr)) + .collect::>(); + let (coverage_ratio, missing_dc, missing_groups) = if attempt.hardswap { + let coverage = self.hardswap_coverage(desired_by_dc, &authoritative_writer_addrs); + let missing_dc = Self::missing_group_dcs(&coverage.missing_groups); + (coverage.ratio, missing_dc, coverage.missing_groups) + } else { + let authoritative_writer_addrs = authoritative_writer_addrs + .iter() + .copied() + .collect::>(); + let (coverage_ratio, missing_dc) = + Self::coverage_ratio(desired_by_dc, &authoritative_writer_addrs); + (coverage_ratio, missing_dc, Vec::new()) + }; + if coverage_ratio < min_ratio { + return Err(ReinitCommitFailure::Coverage { + coverage_ratio, + missing_dc, + missing_groups, + }); + } + if attempt.hardswap + && !missing_groups.is_empty() + && self.bind_stale_mode() == MeBindStaleMode::Never + { + return Err(ReinitCommitFailure::Redundancy { + coverage_ratio, + missing_dc, + missing_groups, + }); + } if !commit_reinit_state( &mut state, attempt.attempt_id, attempt.generation, attempt.map_hash, + attempt.endpoint_revision, attempt.hardswap, ) { - return false; + return Err(ReinitCommitFailure::Superseded); } + if attempt.hardswap { - let writers = self.writers.snapshot(); for writer in writers.iter() { - if !writer.draining.load(Ordering::Relaxed) + if !writer.draining.load(Ordering::Acquire) && writer.generation == attempt.generation { writer @@ -108,10 +191,67 @@ impl MePool { } } } + + let desired_addrs = desired_by_dc + .iter() + .flat_map(|(dc, endpoints)| endpoints.iter().copied().map(|addr| (*dc, addr))) + .collect::>(); + let missing_group_set = missing_groups.iter().copied().collect::>(); + let mut stale_writer_ids = Vec::::new(); + let mut force_close_writer_ids = Vec::::new(); + for writer in writers.iter() { + if writer.draining.load(Ordering::Acquire) { + continue; + } + let stale = if attempt.hardswap { + writer.generation < attempt.generation + } else { + !desired_addrs.contains(&(writer.writer_dc, writer.addr)) + }; + if !stale { + continue; + } + + let writer_group = DcFamilyGroup { + dc: writer.writer_dc, + family: if writer.addr.is_ipv4() { + IpFamily::V4 + } else { + IpFamily::V6 + }, + }; + let preserve_fallback = attempt.hardswap && missing_group_set.contains(&writer_group); + if !preserve_fallback && attempt.hardswap { + registry_registration.retire(writer.id); + } + self.apply_writer_draining_state( + writer, + self.force_close_timeout(), + preserve_fallback || !attempt.hardswap, + ); + stale_writer_ids.push(writer.id); + if (attempt.hardswap && !preserve_fallback) + || (!attempt.hardswap && missing_dc.is_empty()) + { + force_close_writer_ids.push(writer.id); + } + } publish_reinit_state(self.reinit.as_ref(), &state); - true + drop(state); + drop(registry_registration); + drop(writers); + self.notify_writer_epoch(); + + Ok(ReinitCommitOutcome { + coverage_ratio, + missing_dc, + missing_groups, + stale_writer_ids, + force_close_writer_ids, + }) } + /// Computes desired DC-group coverage and returns missing groups in stable order. pub(super) fn coverage_ratio( desired_by_dc: &HashMap>, active_writer_addrs: &HashSet<(i32, SocketAddr)>, @@ -146,10 +286,72 @@ impl MePool { (ratio, missing_dc) } + /// Evaluates full writer-floor coverage independently for every DC and address family. + pub(in crate::transport::middle_proxy) fn hardswap_coverage( + &self, + desired_by_dc: &HashMap>, + writer_addrs: &[(i32, SocketAddr)], + ) -> HardswapCoverage { + let mut covered = 0usize; + let mut total = 0usize; + let mut writer_deficit = 0usize; + let mut missing_groups = Vec::new(); + for (dc, endpoints) in desired_by_dc { + for family in [IpFamily::V4, IpFamily::V6] { + let endpoint_count = endpoints + .iter() + .filter(|endpoint| endpoint.is_ipv4() == (family == IpFamily::V4)) + .count(); + if endpoint_count == 0 { + continue; + } + total = total.saturating_add(1); + let required = self.required_writers_for_dc(endpoint_count); + let alive = writer_addrs + .iter() + .filter(|(writer_dc, endpoint)| { + *writer_dc == *dc + && endpoint.is_ipv4() == (family == IpFamily::V4) + && endpoints.contains(endpoint) + }) + .count(); + if alive >= required { + covered = covered.saturating_add(1); + } else { + writer_deficit = writer_deficit.saturating_add(required.saturating_sub(alive)); + missing_groups.push(DcFamilyGroup { dc: *dc, family }); + } + } + } + missing_groups + .sort_unstable_by_key(|group| (group.dc, matches!(group.family, IpFamily::V6))); + HardswapCoverage { + ratio: if total == 0 { + 1.0 + } else { + (covered as f32) / (total as f32) + }, + missing_groups, + writer_deficit, + } + } + + fn missing_group_dcs(groups: &[DcFamilyGroup]) -> Vec { + let mut dcs = groups.iter().map(|group| group.dc).collect::>(); + dcs.sort_unstable(); + dcs.dedup(); + dcs + } + + /// Restores at least one active writer for every enabled desired DC group. pub async fn reconcile_connections(self: &Arc, rng: &SecureRandom) { + let endpoint_snapshot = self.endpoint_snapshot.load_full(); for family in self.family_order() { - let map = self.proxy_map_for_family(family).await; - for (dc, addrs) in &map { + let map = match family { + IpFamily::V4 => &endpoint_snapshot.map_v4, + IpFamily::V6 => &endpoint_snapshot.map_v6, + }; + for (dc, addrs) in map { let dc_addrs: Vec = addrs .iter() .map(|(ip, port)| SocketAddr::new(*ip, *port)) @@ -175,26 +377,36 @@ impl MePool { } } - pub(super) async fn desired_dc_endpoints(&self) -> HashMap> { + /// Returns the currently authoritative endpoint set for drain and coverage decisions. + pub(in crate::transport::middle_proxy) async fn desired_dc_endpoints( + &self, + ) -> HashMap> { + let endpoint_snapshot = self.endpoint_snapshot.load_full(); + self.desired_dc_endpoints_from_snapshot(&endpoint_snapshot) + } + + /// Projects desired per-DC endpoint sets from one immutable endpoint revision. + pub(in crate::transport::middle_proxy) fn desired_dc_endpoints_from_snapshot( + &self, + endpoint_snapshot: &EndpointSnapshot, + ) -> HashMap> { let now_epoch_secs = Self::now_epoch_secs(); let mut out: HashMap> = HashMap::new(); if self.family_enabled_for_drain_coverage(IpFamily::V4, now_epoch_secs) { - let map_v4 = self.proxy_map_v4.read().await.clone(); - for (dc, addrs) in map_v4 { - let entry = out.entry(dc).or_default(); + for (dc, addrs) in &endpoint_snapshot.map_v4 { + let entry = out.entry(*dc).or_default(); for (ip, port) in addrs { - entry.insert(SocketAddr::new(ip, port)); + entry.insert(SocketAddr::new(*ip, *port)); } } } if self.family_enabled_for_drain_coverage(IpFamily::V6, now_epoch_secs) { - let map_v6 = self.proxy_map_v6.read().await.clone(); - for (dc, addrs) in map_v6 { - let entry = out.entry(dc).or_default(); + for (dc, addrs) in &endpoint_snapshot.map_v6 { + let entry = out.entry(*dc).or_default(); for (ip, port) in addrs { - entry.insert(SocketAddr::new(ip, port)); + entry.insert(SocketAddr::new(*ip, *port)); } } } @@ -202,38 +414,61 @@ impl MePool { out } - pub(in crate::transport::middle_proxy) async fn has_non_draining_writer_per_desired_dc_group( - &self, - ) -> bool { - let desired_by_dc = self.desired_dc_endpoints().await; - let required_dcs: HashSet = desired_by_dc - .iter() - .filter_map(|(dc, endpoints)| { - if endpoints.is_empty() { - None - } else { - Some(*dc) - } - }) - .collect(); - if required_dcs.is_empty() { - return true; - } + /// Promotes authoritative warm writers and drains warm or active generation orphans. + pub(super) async fn reconcile_writer_generation_roles(&self) -> usize { + let writers = self.writers.write().await; + let mut registry_registration = self.registry.prepare_writer_registration().await; + let state = self.reinit.coordinator.lock(); + let active_generation = state.active_generation; + let pending_generation = state.pending.map(|pending| pending.generation); + let endpoint_snapshot = self.endpoint_snapshot.load(); + let now_epoch_secs = Self::now_epoch_secs(); + let mut changed = 0usize; - let ws = self.writers.read().await; - let mut covered_dcs = HashSet::::with_capacity(required_dcs.len()); - for writer in ws.iter() { - if writer.draining.load(Ordering::Relaxed) { + for writer in writers.iter() { + if writer.draining.load(Ordering::Acquire) { continue; } - if required_dcs.contains(&writer.writer_dc) { - covered_dcs.insert(writer.writer_dc); - if covered_dcs.len() == required_dcs.len() { - return true; - } + let contour = WriterContour::from_u8(writer.contour.load(Ordering::Acquire)); + let family = if writer.addr.is_ipv4() { + IpFamily::V4 + } else { + IpFamily::V6 + }; + let endpoint_is_current = self + .family_enabled_for_drain_coverage(family, now_epoch_secs) + && endpoint_snapshot.contains_dc_endpoint(writer.writer_dc, writer.addr); + if contour == WriterContour::Warm + && writer.generation == active_generation + && endpoint_is_current + { + writer + .contour + .store(WriterContour::Active.as_u8(), Ordering::Release); + changed = changed.saturating_add(1); + continue; } + let authoritative_warm = contour == WriterContour::Warm + && pending_generation == Some(writer.generation) + && endpoint_is_current; + let stale_active = + contour == WriterContour::Active && writer.generation != active_generation; + if authoritative_warm || (contour == WriterContour::Active && !stale_active) { + continue; + } + + registry_registration.retire(writer.id); + self.apply_writer_draining_state(writer, self.force_close_timeout(), false); + changed = changed.saturating_add(1); } - false + drop(endpoint_snapshot); + drop(state); + drop(registry_registration); + drop(writers); + if changed > 0 { + self.notify_writer_epoch(); + } + changed } pub(super) fn hardswap_warmup_connect_delay_ms(&self) -> u64 { @@ -270,6 +505,7 @@ impl MePool { core.saturating_add(rand::rng().random_range(0..=jitter)) } + /// Counts non-draining writers owned by one generation and desired DC endpoint set. pub(super) async fn fresh_writer_count_for_dc_endpoints( &self, generation: u64, @@ -285,14 +521,20 @@ impl MePool { .count() } + /// Counts authoritative active writers for one desired DC endpoint set. pub(in crate::transport::middle_proxy) async fn active_writer_count_for_dc_endpoints( &self, dc: i32, endpoints: &HashSet, ) -> usize { + let generation = self.current_generation(); let ws = self.writers.read().await; ws.iter() .filter(|w| !w.draining.load(Ordering::Relaxed)) + .filter(|w| w.generation == generation) + .filter(|w| { + WriterContour::from_u8(w.contour.load(Ordering::Acquire)) == WriterContour::Active + }) .filter(|w| w.writer_dc == dc) .filter(|w| endpoints.contains(&w.addr)) .count() diff --git a/src/transport/middle_proxy/pool_reinit/dual_family_tests.rs b/src/transport/middle_proxy/pool_reinit/dual_family_tests.rs new file mode 100644 index 0000000..1e4cd0a --- /dev/null +++ b/src/transport/middle_proxy/pool_reinit/dual_family_tests.rs @@ -0,0 +1,129 @@ +use std::collections::HashMap; +use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr}; +use std::sync::atomic::Ordering; + +use super::tests::{insert_writer, insert_writer_floor}; +use crate::network::probe::NetworkDecision; +use crate::transport::middle_proxy::pool::WriterContour; +use crate::transport::middle_proxy::pool_writer_security_tests::make_pool_with_decision; + +#[tokio::test] +async fn nonpreferred_family_warm_generation_retains_authority_and_commits() { + let pool = make_pool_with_decision(NetworkDecision { + ipv4_me: true, + ipv6_me: true, + effective_prefer: 4, + effective_multipath: false, + ..NetworkDecision::default() + }) + .await; + let v4 = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443); + let v6 = SocketAddr::new(IpAddr::V6(Ipv6Addr::LOCALHOST), 443); + pool.update_proxy_maps( + HashMap::from([(2, vec![(v4.ip(), v4.port())])]), + Some(HashMap::from([(2, vec![(v6.ip(), v6.port())])])), + ) + .await; + let desired = pool.desired_dc_endpoints().await; + let map_hash = super::MePool::desired_map_hash(&desired); + let endpoint_revision = pool.endpoint_snapshot.load().revision; + let reservation = pool + .reserve_reinit_attempt(true, map_hash, endpoint_revision, 100) + .expect("current endpoint revision must admit the hardswap"); + assert!(pool.hardswap_warmup_is_authoritative( + reservation.attempt.generation, + map_hash, + endpoint_revision, + )); + let generation = reservation.attempt.generation; + let fresh_v4 = insert_writer_floor(&pool, 10, 2, v4, generation, WriterContour::Warm).await; + let fresh_v6 = insert_writer_floor(&pool, 20, 2, v6, generation, WriterContour::Warm).await; + let fresh_v4_media = + insert_writer_floor(&pool, 30, -2, v4, generation, WriterContour::Warm).await; + let fresh_v6_media = + insert_writer_floor(&pool, 40, -2, v6, generation, WriterContour::Warm).await; + + assert_eq!(pool.reconcile_writer_generation_roles().await, 0); + for writer in fresh_v4 + .iter() + .chain(&fresh_v6) + .chain(&fresh_v4_media) + .chain(&fresh_v6_media) + { + assert!(!writer.draining.load(Ordering::Acquire)); + } + + let outcome = pool + .commit_reinit_attempt(&reservation.attempt, &desired, 1.0) + .await + .expect("full dual-family floor must commit without multipath selection"); + + assert!(outcome.missing_groups.is_empty()); + assert_eq!(pool.current_generation(), generation); +} + +#[tokio::test] +async fn endpoint_revision_fences_pending_hardswap_after_map_aba() { + let pool = make_pool_with_decision(NetworkDecision { + ipv4_me: true, + effective_prefer: 4, + ..NetworkDecision::default() + }) + .await; + let endpoint_a = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 10)), 443); + let endpoint_b = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 11)), 443); + pool.update_proxy_maps( + HashMap::from([(2, vec![(endpoint_a.ip(), endpoint_a.port())])]), + None, + ) + .await; + let desired_a = pool.desired_dc_endpoints().await; + let map_hash = super::MePool::desired_map_hash(&desired_a); + let endpoint_revision = pool.endpoint_snapshot.load().revision; + let reservation = pool + .reserve_reinit_attempt(true, map_hash, endpoint_revision, 100) + .expect("current endpoint revision must admit the hardswap"); + assert!(pool.hardswap_warmup_is_authoritative( + reservation.attempt.generation, + map_hash, + endpoint_revision, + )); + let warm = insert_writer( + &pool, + 100, + 2, + endpoint_a, + reservation.attempt.generation, + WriterContour::Warm, + ) + .await; + + pool.update_proxy_maps( + HashMap::from([(2, vec![(endpoint_b.ip(), endpoint_b.port())])]), + None, + ) + .await; + assert!(!pool.hardswap_warmup_is_authoritative( + reservation.attempt.generation, + map_hash, + endpoint_revision, + )); + pool.update_proxy_maps( + HashMap::from([(2, vec![(endpoint_a.ip(), endpoint_a.port())])]), + None, + ) + .await; + assert!(!pool.hardswap_warmup_is_authoritative( + reservation.attempt.generation, + map_hash, + endpoint_revision, + )); + + let snapshot = pool.api_hardswap_snapshot().await; + + assert!(snapshot.pending); + assert_eq!(snapshot.pending_map_current, Some(false)); + assert_eq!(snapshot.pending_writers_current, 0); + assert_eq!(snapshot.orphan_warm_writers_current, 1); + assert!(!warm.draining.load(Ordering::Acquire)); +} diff --git a/src/transport/middle_proxy/pool_reinit/reconcile.rs b/src/transport/middle_proxy/pool_reinit/reconcile.rs index 5c900fe..be133da 100644 --- a/src/transport/middle_proxy/pool_reinit/reconcile.rs +++ b/src/transport/middle_proxy/pool_reinit/reconcile.rs @@ -1,10 +1,29 @@ use super::*; impl MePool { + /// Checks the exact pending-generation tuple before starting more warmup work. + pub(super) fn hardswap_warmup_is_authoritative( + &self, + generation: u64, + map_hash: u64, + endpoint_revision: u64, + ) -> bool { + let state = self.reinit.coordinator.lock(); + state.desired_map_hash == map_hash + && state.endpoint_revision == endpoint_revision + && state.pending.is_some_and(|pending| { + pending.generation == generation + && pending.map_hash == map_hash + && pending.endpoint_revision == endpoint_revision + }) + } + async fn warmup_generation_for_all_dcs( self: &Arc, rng: &SecureRandom, generation: u64, + map_hash: u64, + endpoint_revision: u64, desired_by_dc: &HashMap>, ) { let extra_passes = self @@ -15,101 +34,137 @@ impl MePool { let total_passes = 1 + extra_passes; for (dc, endpoints) in desired_by_dc { - if endpoints.is_empty() { - continue; + if !self.hardswap_warmup_is_authoritative(generation, map_hash, endpoint_revision) { + return; } - - let mut endpoint_list: Vec = endpoints.iter().copied().collect(); - endpoint_list.sort_unstable(); - let required = self.required_writers_for_dc(endpoint_list.len()); - let mut completed = false; - let mut last_fresh_count = self - .fresh_writer_count_for_dc_endpoints(generation, *dc, endpoints) - .await; - - for pass_idx in 0..total_passes { - if last_fresh_count >= required { - completed = true; - break; + for family in [IpFamily::V4, IpFamily::V6] { + let family_endpoints = endpoints + .iter() + .copied() + .filter(|endpoint| endpoint.is_ipv4() == (family == IpFamily::V4)) + .collect::>(); + if family_endpoints.is_empty() { + continue; } - let missing = required.saturating_sub(last_fresh_count); - debug!( - dc = *dc, - pass = pass_idx + 1, - total_passes, - fresh_count = last_fresh_count, - required, - missing, - endpoint_count = endpoint_list.len(), - "ME hardswap warmup pass started" - ); - - for attempt_idx in 0..missing { - let delay_ms = self.hardswap_warmup_connect_delay_ms(); - tokio::time::sleep(Duration::from_millis(delay_ms)).await; - - let connected = self - .connect_endpoints_round_robin_with_generation_contour( - *dc, - &endpoint_list, - rng, - generation, - WriterContour::Warm, - false, - ) - .await; - debug!( - dc = *dc, - pass = pass_idx + 1, - total_passes, - attempt = attempt_idx + 1, - delay_ms, - connected, - "ME hardswap warmup connect attempt finished" - ); - } - - last_fresh_count = self - .fresh_writer_count_for_dc_endpoints(generation, *dc, endpoints) + let mut endpoint_list = family_endpoints.iter().copied().collect::>(); + endpoint_list.sort_unstable(); + let required = self.required_writers_for_dc(endpoint_list.len()); + let mut completed = false; + let mut last_fresh_count = self + .fresh_writer_count_for_dc_endpoints(generation, *dc, &family_endpoints) .await; - if last_fresh_count >= required { - completed = true; - info!( - dc = *dc, - pass = pass_idx + 1, - total_passes, - fresh_count = last_fresh_count, - required, - "ME hardswap warmup floor reached for DC" - ); - break; - } - if pass_idx + 1 < total_passes { - let backoff_ms = self.hardswap_warmup_backoff_ms(pass_idx); + for pass_idx in 0..total_passes { + if last_fresh_count >= required { + completed = true; + break; + } + + let missing = required.saturating_sub(last_fresh_count); debug!( dc = *dc, + family = ?family, pass = pass_idx + 1, total_passes, fresh_count = last_fresh_count, required, - backoff_ms, - "ME hardswap warmup pass incomplete, delaying next pass" + missing, + endpoint_count = endpoint_list.len(), + "ME hardswap family warmup pass started" ); - tokio::time::sleep(Duration::from_millis(backoff_ms)).await; - } - } - if !completed { - warn!( - dc = *dc, - fresh_count = last_fresh_count, - required, - endpoint_count = endpoint_list.len(), - total_passes, - "ME warmup stopped: unable to reach required writer floor for DC" - ); + for attempt_idx in 0..missing { + if !self.hardswap_warmup_is_authoritative( + generation, + map_hash, + endpoint_revision, + ) { + return; + } + let delay_ms = self.hardswap_warmup_connect_delay_ms(); + tokio::time::sleep(Duration::from_millis(delay_ms)).await; + if !self.hardswap_warmup_is_authoritative( + generation, + map_hash, + endpoint_revision, + ) { + return; + } + + let connected = self + .connect_endpoints_round_robin_with_generation_contour( + *dc, + &endpoint_list, + rng, + generation, + WriterContour::Warm, + WriterOpenIntent::Coverage, + ) + .await; + debug!( + dc = *dc, + family = ?family, + pass = pass_idx + 1, + total_passes, + attempt = attempt_idx + 1, + delay_ms, + connected, + "ME hardswap family warmup connect attempt finished" + ); + } + + last_fresh_count = self + .fresh_writer_count_for_dc_endpoints(generation, *dc, &family_endpoints) + .await; + if last_fresh_count >= required { + completed = true; + info!( + dc = *dc, + family = ?family, + pass = pass_idx + 1, + total_passes, + fresh_count = last_fresh_count, + required, + "ME hardswap writer floor reached for DC family" + ); + break; + } + + if pass_idx + 1 < total_passes { + if !self.hardswap_warmup_is_authoritative( + generation, + map_hash, + endpoint_revision, + ) { + return; + } + let backoff_ms = self.hardswap_warmup_backoff_ms(pass_idx); + debug!( + dc = *dc, + family = ?family, + pass = pass_idx + 1, + total_passes, + fresh_count = last_fresh_count, + required, + backoff_ms, + "ME hardswap family warmup incomplete, delaying next pass" + ); + tokio::time::sleep(Duration::from_millis(backoff_ms)).await; + } + } + + if !completed { + warn!( + dc = *dc, + family = ?family, + fresh_count = last_fresh_count, + required, + endpoint_count = endpoint_list.len(), + total_passes, + "ME warmup stopped below the required DC-family writer floor" + ); + } } } } @@ -118,7 +173,8 @@ impl MePool { self: &Arc, rng: &SecureRandom, ) -> bool { - let desired_by_dc = self.desired_dc_endpoints().await; + let endpoint_snapshot = self.endpoint_snapshot.load_full(); + let desired_by_dc = self.desired_dc_endpoints_from_snapshot(&endpoint_snapshot); let now_epoch_secs = Self::now_epoch_secs(); let v4_suppressed = self.is_family_temporarily_suppressed(IpFamily::V4, now_epoch_secs); let v6_suppressed = self.is_family_temporarily_suppressed(IpFamily::V6, now_epoch_secs); @@ -137,10 +193,30 @@ impl MePool { let desired_map_hash = Self::desired_map_hash(&desired_by_dc); let hardswap = self.reinit.hardswap.load(Ordering::Relaxed); - let reservation = self.reserve_reinit_attempt(hardswap, desired_map_hash, now_epoch_secs); + let Some(reservation) = self.reserve_reinit_attempt( + hardswap, + desired_map_hash, + endpoint_snapshot.revision, + now_epoch_secs, + ) else { + debug!( + endpoint_revision = endpoint_snapshot.revision, + "ME reinit snapshot superseded before reservation" + ); + return false; + }; let attempt = reservation.attempt; let previous_generation = attempt.previous_generation; let generation = attempt.generation; + let reconciled_roles = self.reconcile_writer_generation_roles().await; + if reconciled_roles > 0 { + info!( + reconciled_roles, + active_generation = previous_generation, + pending_generation = generation, + "ME writer generation roles reconciled" + ); + } if reservation.pending_reused { self.stats.increment_me_hardswap_pending_reuse_total(); debug!( @@ -161,8 +237,14 @@ impl MePool { } if hardswap { - self.warmup_generation_for_all_dcs(rng, generation, &desired_by_dc) - .await; + self.warmup_generation_for_all_dcs( + rng, + generation, + desired_map_hash, + endpoint_snapshot.revision, + &desired_by_dc, + ) + .await; } else { self.reconcile_connections(rng).await; } @@ -180,14 +262,10 @@ impl MePool { ); let (coverage_ratio, missing_dc) = Self::coverage_ratio(&desired_by_dc, &active_writer_addrs); - let mut route_quorum_ok = coverage_ratio >= min_ratio; - let mut redundancy_ok = missing_dc.is_empty(); - let mut redundancy_missing_dc = missing_dc.clone(); - let mut gate_coverage_ratio = coverage_ratio; if !hardswap && coverage_ratio < min_ratio { self.set_last_drain_gate( false, - redundancy_ok, + missing_dc.is_empty(), MeDrainGateReason::CoverageQuorum, now_epoch_secs, ); @@ -203,79 +281,107 @@ impl MePool { } if hardswap { - let fresh_writer_addrs: HashSet<(i32, SocketAddr)> = writers + let fresh_writer_addrs: Vec<(i32, SocketAddr)> = writers .iter() .filter(|w| !w.draining.load(Ordering::Relaxed)) .filter(|w| w.generation == generation) .map(|w| (w.writer_dc, w.addr)) .collect(); - let (fresh_coverage_ratio, fresh_missing_dc) = - Self::coverage_ratio(&desired_by_dc, &fresh_writer_addrs); - route_quorum_ok = fresh_coverage_ratio >= min_ratio; - redundancy_ok = fresh_missing_dc.is_empty(); - redundancy_missing_dc = fresh_missing_dc.clone(); - gate_coverage_ratio = fresh_coverage_ratio; - if fresh_coverage_ratio < min_ratio { + let fresh_coverage = self.hardswap_coverage(&desired_by_dc, &fresh_writer_addrs); + if fresh_coverage.ratio < min_ratio { self.set_last_drain_gate( false, - redundancy_ok, + fresh_coverage.missing_groups.is_empty(), MeDrainGateReason::CoverageQuorum, now_epoch_secs, ); warn!( previous_generation, generation, - fresh_coverage_ratio = format_args!("{fresh_coverage_ratio:.3}"), - missing_dc = ?fresh_missing_dc, - "ME hardswap pending: fresh generation DC coverage incomplete" + fresh_coverage_ratio = format_args!("{:.3}", fresh_coverage.ratio), + writer_deficit = fresh_coverage.writer_deficit, + missing_groups = ?fresh_coverage.missing_groups, + "ME hardswap pending: fresh generation DC-family floors incomplete" ); return false; } } + drop(writers); + let commit = self + .commit_reinit_attempt(&attempt, &desired_by_dc, min_ratio) + .await; + let outcome = match commit { + Ok(outcome) => outcome, + Err(ReinitCommitFailure::Superseded) => { + debug!( + previous_generation, + generation, "ME reinit result discarded after a newer desired-map attempt" + ); + return false; + } + Err(ReinitCommitFailure::Coverage { + coverage_ratio, + missing_dc, + missing_groups, + }) => { + self.set_last_drain_gate( + false, + missing_dc.is_empty(), + MeDrainGateReason::CoverageQuorum, + now_epoch_secs, + ); + warn!( + previous_generation, + generation, + coverage_ratio = format_args!("{coverage_ratio:.3}"), + min_ratio = format_args!("{min_ratio:.3}"), + missing_dc = ?missing_dc, + missing_groups = ?missing_groups, + "ME reinit coverage changed before commit; keeping current generation" + ); + return false; + } + Err(ReinitCommitFailure::Redundancy { + coverage_ratio, + missing_dc, + missing_groups, + }) => { + self.set_last_drain_gate( + true, + false, + MeDrainGateReason::Redundancy, + now_epoch_secs, + ); + warn!( + previous_generation, + generation, + coverage_ratio = format_args!("{coverage_ratio:.3}"), + min_ratio = format_args!("{min_ratio:.3}"), + missing_dc = ?missing_dc, + missing_groups = ?missing_groups, + "ME hardswap weighted quorum requires stale-binding fallback" + ); + return false; + } + }; self.set_last_drain_gate( - route_quorum_ok, - redundancy_ok, + true, + outcome.missing_dc.is_empty(), MeDrainGateReason::Open, now_epoch_secs, ); - if !redundancy_ok { + if !outcome.missing_dc.is_empty() { warn!( - missing_dc = ?redundancy_missing_dc, - coverage_ratio = format_args!("{gate_coverage_ratio:.3}"), + missing_dc = ?outcome.missing_dc, + missing_groups = ?outcome.missing_groups, + coverage_ratio = format_args!("{:.3}", outcome.coverage_ratio), min_ratio = format_args!("{min_ratio:.3}"), - "ME reinit proceeds with weighted quorum while some DC groups remain uncovered" + "ME reinit committed with bounded stale fallback for uncovered DC-family groups" ); } - if !self.commit_reinit_attempt(&attempt) { - debug!( - previous_generation, - generation, "ME reinit result discarded after a newer desired-map attempt" - ); - return false; - } - - let desired_addrs: HashSet<(i32, SocketAddr)> = desired_by_dc - .iter() - .flat_map(|(dc, set)| set.iter().copied().map(|addr| (*dc, addr))) - .collect(); - - let stale_writer_ids: Vec = writers - .iter() - .filter(|w| !w.draining.load(Ordering::Relaxed)) - .filter(|w| { - if hardswap { - w.generation < generation - } else { - !desired_addrs.contains(&(w.writer_dc, w.addr)) - } - }) - .map(|w| w.id) - .collect(); - drop(writers); - - if stale_writer_ids.is_empty() { + if outcome.stale_writer_ids.is_empty() { debug!("ME reinit cycle completed with no stale writers"); return true; } @@ -283,35 +389,20 @@ impl MePool { let drain_timeout = self.force_close_timeout(); let drain_timeout_secs = drain_timeout.map(|d| d.as_secs()).unwrap_or(0); info!( - stale_writers = stale_writer_ids.len(), + stale_writers = outcome.stale_writer_ids.len(), + force_close_writers = outcome.force_close_writer_ids.len(), previous_generation, generation, hardswap, - coverage_ratio = format_args!("{coverage_ratio:.3}"), + coverage_ratio = format_args!("{:.3}", outcome.coverage_ratio), min_ratio = format_args!("{min_ratio:.3}"), drain_timeout_secs, "ME reinit cycle covered; processing stale writers" ); self.stats.increment_pool_swap_total(); - let can_drop_with_replacement = self.has_non_draining_writer_per_desired_dc_group().await; - if can_drop_with_replacement { - info!( - stale_writers = stale_writer_ids.len(), - "ME reinit stale writers: replacement coverage ready, force-closing clients for fast rebind" - ); - } else { - warn!( - stale_writers = stale_writer_ids.len(), - "ME reinit stale writers: replacement coverage incomplete, keeping draining fallback" - ); - } - for writer_id in stale_writer_ids { - self.mark_writer_draining_with_timeout(writer_id, drain_timeout, !hardswap) - .await; - if can_drop_with_replacement { - self.stats.increment_pool_force_close_total(); - self.remove_writer_and_close_clients(writer_id).await; - } + for writer_id in outcome.force_close_writer_ids { + self.stats.increment_pool_force_close_total(); + self.remove_writer_and_close_clients(writer_id).await; } true } diff --git a/src/transport/middle_proxy/pool_reinit/tests.rs b/src/transport/middle_proxy/pool_reinit/tests.rs index cb85871..bb274be 100644 --- a/src/transport/middle_proxy/pool_reinit/tests.rs +++ b/src/transport/middle_proxy/pool_reinit/tests.rs @@ -1,15 +1,101 @@ use std::collections::{HashMap, HashSet}; -use std::net::{IpAddr, Ipv4Addr, SocketAddr}; +use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr}; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU32, AtomicU64, Ordering}; +use std::time::Instant; -use super::{MePool, commit_reinit_state}; +use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; + +use super::{DcFamilyGroup, MePool, ReinitCommitFailure, commit_reinit_state}; +use crate::config::MeBindStaleMode; +use crate::transport::middle_proxy::codec::WriterCommand; use crate::transport::middle_proxy::pool::{ - ReinitAttemptState, ReinitCoordinatorState, ReinitPendingState, + MeWriter, ReinitAttemptState, ReinitCoordinatorState, ReinitPendingState, WriterContour, }; +use crate::transport::middle_proxy::pool_writer_security_tests::make_pool; fn addr(octet: u8, port: u16) -> SocketAddr { SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, octet)), port) } +fn addr_v6(segment: u16, port: u16) -> SocketAddr { + SocketAddr::new( + IpAddr::V6(Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, segment)), + port, + ) +} + +pub(super) async fn insert_writer( + pool: &Arc, + writer_id: u64, + writer_dc: i32, + endpoint: SocketAddr, + generation: u64, + contour: WriterContour, +) -> MeWriter { + let (tx, _rx) = mpsc::channel::(8); + let byte_budget = pool.new_writer_byte_budget(); + let writer = MeWriter { + id: writer_id, + addr: endpoint, + source_ip: endpoint.ip(), + writer_dc, + generation, + contour: Arc::new(AtomicU8::new(contour.as_u8())), + created_at: Instant::now(), + tx: tx.clone(), + byte_budget: byte_budget.clone(), + cancel: CancellationToken::new(), + degraded: Arc::new(AtomicBool::new(false)), + rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)), + draining: Arc::new(AtomicBool::new(false)), + draining_started_at_epoch_secs: Arc::new(AtomicU64::new(0)), + drain_deadline_epoch_secs: Arc::new(AtomicU64::new(0)), + allow_drain_fallback: Arc::new(AtomicBool::new(false)), + }; + + pool.registry + .register_writer(writer_id, tx, byte_budget) + .await; + pool.writers.write().await.push(writer.clone()); + pool.conn_count.fetch_add(1, Ordering::Relaxed); + writer +} + +pub(super) async fn insert_writer_floor( + pool: &Arc, + first_writer_id: u64, + writer_dc: i32, + endpoint: SocketAddr, + generation: u64, + contour: WriterContour, +) -> Vec { + let required = pool.required_writers_for_dc(1); + let mut writers = Vec::with_capacity(required); + for offset in 0..required { + writers.push( + insert_writer( + pool, + first_writer_id + offset as u64, + writer_dc, + endpoint, + generation, + contour, + ) + .await, + ); + } + writers +} + +fn desired_two_dcs() -> HashMap> { + HashMap::from([ + (1, HashSet::from([addr(1, 2001)])), + (2, HashSet::from([addr(2, 2002)])), + ]) +} + #[test] fn coverage_ratio_counts_dc_coverage_not_floor() { let dc1 = addr(1, 2001); @@ -62,10 +148,12 @@ fn stale_concurrent_attempt_cannot_regress_active_generation() { next_attempt_id: 3, active_generation: 1, desired_map_hash: 22, + endpoint_revision: 7, pending: Some(ReinitPendingState { generation: 3, started_at_epoch_secs: 1, map_hash: 22, + endpoint_revision: 7, }), attempts: HashMap::from([ ( @@ -73,6 +161,7 @@ fn stale_concurrent_attempt_cannot_regress_active_generation() { ReinitAttemptState { generation: 2, map_hash: 11, + endpoint_revision: 6, hardswap: true, committed: false, }, @@ -82,6 +171,7 @@ fn stale_concurrent_attempt_cannot_regress_active_generation() { ReinitAttemptState { generation: 3, map_hash: 22, + endpoint_revision: 7, hardswap: true, committed: false, }, @@ -89,9 +179,357 @@ fn stale_concurrent_attempt_cannot_regress_active_generation() { ]), }; - assert!(commit_reinit_state(&mut state, 2, 3, 22, true)); + assert!(commit_reinit_state(&mut state, 2, 3, 22, 7, true)); assert_eq!(state.active_generation, 3); - assert!(!commit_reinit_state(&mut state, 1, 2, 11, true)); + assert!(!commit_reinit_state(&mut state, 1, 2, 11, 6, true)); assert_eq!(state.active_generation, 3); assert!(state.pending.is_none()); } + +#[tokio::test] +async fn partial_hardswap_is_rejected_when_stale_binding_is_disabled() { + let pool = make_pool().await; + let desired_by_dc = desired_two_dcs(); + let active_generation = pool.current_generation(); + let old_dc1 = insert_writer( + &pool, + 101, + 1, + addr(1, 2001), + active_generation, + WriterContour::Active, + ) + .await; + let old_dc2 = insert_writer( + &pool, + 102, + 2, + addr(2, 2002), + active_generation, + WriterContour::Active, + ) + .await; + let map_hash = MePool::desired_map_hash(&desired_by_dc); + let endpoint_revision = pool.endpoint_snapshot.load().revision; + let reservation = pool + .reserve_reinit_attempt(true, map_hash, endpoint_revision, 100) + .expect("endpoint revision must remain current"); + insert_writer_floor( + &pool, + 201, + 1, + addr(1, 2001), + reservation.attempt.generation, + WriterContour::Warm, + ) + .await; + + let result = pool + .commit_reinit_attempt(&reservation.attempt, &desired_by_dc, 0.5) + .await; + + assert!(matches!( + result, + Err(ReinitCommitFailure::Redundancy { .. }) + )); + assert_eq!(pool.current_generation(), active_generation); + assert!(!old_dc1.draining.load(Ordering::Acquire)); + assert!(!old_dc2.draining.load(Ordering::Acquire)); +} + +#[tokio::test] +async fn partial_hardswap_preserves_fallback_only_for_missing_dc() { + let pool = make_pool().await; + pool.binding_policy + .me_bind_stale_mode + .store(MeBindStaleMode::Ttl.as_u8(), Ordering::Release); + let desired_by_dc = desired_two_dcs(); + let active_generation = pool.current_generation(); + let old_dc1 = insert_writer( + &pool, + 301, + 1, + addr(1, 2001), + active_generation, + WriterContour::Active, + ) + .await; + let old_dc2 = insert_writer( + &pool, + 302, + 2, + addr(9, 2999), + active_generation, + WriterContour::Active, + ) + .await; + let map_hash = MePool::desired_map_hash(&desired_by_dc); + let endpoint_revision = pool.endpoint_snapshot.load().revision; + let reservation = pool + .reserve_reinit_attempt(true, map_hash, endpoint_revision, 100) + .expect("endpoint revision must remain current"); + let fresh_dc1 = insert_writer_floor( + &pool, + 401, + 1, + addr(1, 2001), + reservation.attempt.generation, + WriterContour::Warm, + ) + .await; + + let outcome = pool + .commit_reinit_attempt(&reservation.attempt, &desired_by_dc, 0.5) + .await + .expect("partial hardswap must commit when bounded stale fallback is enabled"); + + assert_eq!(pool.current_generation(), reservation.attempt.generation); + assert_eq!(outcome.missing_dc, vec![2]); + assert_eq!(outcome.force_close_writer_ids, vec![301]); + assert!(old_dc1.draining.load(Ordering::Acquire)); + assert!(!old_dc1.allow_drain_fallback.load(Ordering::Acquire)); + assert!(old_dc2.draining.load(Ordering::Acquire)); + assert!(old_dc2.allow_drain_fallback.load(Ordering::Acquire)); + assert_eq!( + WriterContour::from_u8(fresh_dc1[0].contour.load(Ordering::Acquire)), + WriterContour::Active + ); +} + +#[tokio::test] +async fn partial_hardswap_preserves_fallback_only_for_underfloor_family() { + let pool = make_pool().await; + pool.binding_policy + .me_bind_stale_mode + .store(MeBindStaleMode::Ttl.as_u8(), Ordering::Release); + let v4 = addr(1, 2001); + let v6 = addr_v6(1, 2001); + let desired_by_dc = HashMap::from([(1, HashSet::from([v4, v6]))]); + let active_generation = pool.current_generation(); + let old_v4 = insert_writer(&pool, 451, 1, v4, active_generation, WriterContour::Active).await; + let old_v6 = insert_writer(&pool, 452, 1, v6, active_generation, WriterContour::Active).await; + let map_hash = MePool::desired_map_hash(&desired_by_dc); + let endpoint_revision = pool.endpoint_snapshot.load().revision; + let reservation = pool + .reserve_reinit_attempt(true, map_hash, endpoint_revision, 100) + .expect("endpoint revision must remain current"); + insert_writer_floor( + &pool, + 461, + 1, + v4, + reservation.attempt.generation, + WriterContour::Warm, + ) + .await; + + let outcome = pool + .commit_reinit_attempt(&reservation.attempt, &desired_by_dc, 0.5) + .await + .expect("one covered family satisfies the configured weighted quorum"); + + assert_eq!( + outcome.missing_groups, + vec![DcFamilyGroup { + dc: 1, + family: crate::network::IpFamily::V6, + }] + ); + assert!(outcome.force_close_writer_ids.contains(&old_v4.id)); + assert!(!outcome.force_close_writer_ids.contains(&old_v6.id)); + assert!(!old_v4.allow_drain_fallback.load(Ordering::Acquire)); + assert!(old_v6.allow_drain_fallback.load(Ordering::Acquire)); +} + +#[tokio::test] +async fn complete_hardswap_promotes_fresh_generation_and_retires_old_writers() { + let pool = make_pool().await; + let desired_by_dc = desired_two_dcs(); + let active_generation = pool.current_generation(); + let old_dc1 = insert_writer( + &pool, + 501, + 1, + addr(1, 2001), + active_generation, + WriterContour::Active, + ) + .await; + let old_dc2 = insert_writer( + &pool, + 502, + 2, + addr(2, 2002), + active_generation, + WriterContour::Active, + ) + .await; + let map_hash = MePool::desired_map_hash(&desired_by_dc); + let endpoint_revision = pool.endpoint_snapshot.load().revision; + let reservation = pool + .reserve_reinit_attempt(true, map_hash, endpoint_revision, 100) + .expect("endpoint revision must remain current"); + let fresh_dc1 = insert_writer_floor( + &pool, + 601, + 1, + addr(1, 2001), + reservation.attempt.generation, + WriterContour::Warm, + ) + .await; + let fresh_dc2 = insert_writer_floor( + &pool, + 611, + 2, + addr(2, 2002), + reservation.attempt.generation, + WriterContour::Warm, + ) + .await; + + let outcome = pool + .commit_reinit_attempt(&reservation.attempt, &desired_by_dc, 1.0) + .await + .expect("fully covered hardswap must commit"); + + assert_eq!(pool.current_generation(), reservation.attempt.generation); + assert!(outcome.missing_dc.is_empty()); + assert_eq!(outcome.force_close_writer_ids, vec![501, 502]); + assert!(old_dc1.draining.load(Ordering::Acquire)); + assert!(old_dc2.draining.load(Ordering::Acquire)); + assert_eq!( + WriterContour::from_u8(fresh_dc1[0].contour.load(Ordering::Acquire)), + WriterContour::Active + ); + assert_eq!( + WriterContour::from_u8(fresh_dc2[0].contour.load(Ordering::Acquire)), + WriterContour::Active + ); +} + +#[tokio::test] +async fn endpoint_revision_change_supersedes_hardswap_before_writer_drain() { + let pool = make_pool().await; + let old_endpoint = addr(1, 2001); + pool.update_proxy_maps( + HashMap::from([(1, vec![(old_endpoint.ip(), old_endpoint.port())])]), + None, + ) + .await; + let active_generation = pool.current_generation(); + let old_writer = insert_writer( + &pool, + 651, + 1, + old_endpoint, + active_generation, + WriterContour::Active, + ) + .await; + let desired_by_dc = HashMap::from([(1, HashSet::from([old_endpoint]))]); + let map_hash = MePool::desired_map_hash(&desired_by_dc); + let endpoint_revision = pool.endpoint_snapshot.load().revision; + let reservation = pool + .reserve_reinit_attempt(true, map_hash, endpoint_revision, 100) + .expect("endpoint revision must remain current"); + let fresh_writer = insert_writer( + &pool, + 652, + 1, + old_endpoint, + reservation.attempt.generation, + WriterContour::Warm, + ) + .await; + + let replacement_endpoint = addr(3, 2003); + pool.update_proxy_maps( + HashMap::from([( + 1, + vec![(replacement_endpoint.ip(), replacement_endpoint.port())], + )]), + None, + ) + .await; + + let result = pool + .commit_reinit_attempt(&reservation.attempt, &desired_by_dc, 1.0) + .await; + + assert!(matches!(result, Err(ReinitCommitFailure::Superseded))); + assert_eq!(pool.current_generation(), active_generation); + assert!(!old_writer.draining.load(Ordering::Acquire)); + assert!(!fresh_writer.draining.load(Ordering::Acquire)); + assert_eq!( + WriterContour::from_u8(fresh_writer.contour.load(Ordering::Acquire)), + WriterContour::Warm + ); +} + +#[tokio::test] +async fn generation_role_reconciliation_promotes_active_warm_and_drains_orphans() { + let pool = make_pool().await; + let endpoint = addr(1, 2001); + pool.update_proxy_maps( + HashMap::from([(1, vec![(endpoint.ip(), endpoint.port())])]), + None, + ) + .await; + let desired_by_dc = pool.desired_dc_endpoints().await; + let map_hash = MePool::desired_map_hash(&desired_by_dc); + let endpoint_revision = pool.endpoint_snapshot.load().revision; + let reservation = pool + .reserve_reinit_attempt(true, map_hash, endpoint_revision, 100) + .expect("endpoint revision must remain current"); + let active_warm = insert_writer( + &pool, + 701, + 1, + endpoint, + pool.current_generation(), + WriterContour::Warm, + ) + .await; + let pending_warm = insert_writer( + &pool, + 702, + 1, + endpoint, + reservation.attempt.generation, + WriterContour::Warm, + ) + .await; + let orphan_warm = insert_writer( + &pool, + 703, + 1, + endpoint, + reservation.attempt.generation + 10, + WriterContour::Warm, + ) + .await; + + let changed = pool.reconcile_writer_generation_roles().await; + + assert_eq!(changed, 2); + assert_eq!( + WriterContour::from_u8(active_warm.contour.load(Ordering::Acquire)), + WriterContour::Active + ); + assert_eq!( + WriterContour::from_u8(pending_warm.contour.load(Ordering::Acquire)), + WriterContour::Warm + ); + assert!(orphan_warm.draining.load(Ordering::Acquire)); + assert!(!orphan_warm.allow_drain_fallback.load(Ordering::Acquire)); + let snapshot = pool.api_hardswap_snapshot().await; + let desired = pool.desired_dc_endpoints().await; + let coverage = pool.hardswap_coverage(&desired, &[(1, endpoint)]); + assert_eq!(snapshot.orphan_warm_writers_current, 0); + assert_eq!(snapshot.pending_writer_deficit, coverage.writer_deficit); + assert_eq!( + snapshot.pending_missing_dc_groups, + coverage.missing_groups.len() + ); +} diff --git a/src/transport/middle_proxy/pool_runtime_api.rs b/src/transport/middle_proxy/pool_runtime_api.rs index 78c50e7..5ac9b83 100644 --- a/src/transport/middle_proxy/pool_runtime_api.rs +++ b/src/transport/middle_proxy/pool_runtime_api.rs @@ -1,7 +1,7 @@ use std::collections::HashMap; use std::time::Instant; -use super::pool::{MeDrainGateReason, MePool, RefillDcKey}; +use super::pool::{MeDrainGateReason, MePool, RefillTargetKey}; use crate::network::IpFamily; #[derive(Clone, Debug)] @@ -61,14 +61,17 @@ impl MePool { let refill_states = self.refill_states.lock(); let inflight_endpoints_total = refill_states .values() - .map(|pending| 1usize + usize::from(pending.is_some())) + .map(|state| 1usize.saturating_add(state.pending_count)) .sum(); let running_dc_total = refill_states.len(); let pending_dc_total = refill_states .values() - .filter(|pending| pending.is_some()) + .filter(|state| state.pending_count > 0) .count(); - let inflight_dc_keys = refill_states.keys().copied().collect::>(); + let inflight_dc_keys = refill_states + .keys() + .copied() + .collect::>(); drop(refill_states); let mut by_dc_map = HashMap::<(i16, &'static str), usize>::new(); diff --git a/src/transport/middle_proxy/pool_status.rs b/src/transport/middle_proxy/pool_status.rs index dcf7fc3..42bf2e3 100644 --- a/src/transport/middle_proxy/pool_status.rs +++ b/src/transport/middle_proxy/pool_status.rs @@ -12,6 +12,12 @@ use crate::transport::upstream::IpPreference; mod status_snapshot; // ME runtime policy and coherent snapshot assembly. mod runtime_snapshot; +// Hardswap ownership, coverage, and writer-replacement lifecycle state. +mod hardswap_snapshot; +pub(crate) use hardswap_snapshot::MeApiHardswapSnapshot; +#[cfg(test)] +#[path = "pool_status/tests.rs"] +mod status_invariant_tests; #[derive(Clone, Debug)] pub(crate) struct MeApiWriterStatusSnapshot { pub writer_id: u64, @@ -98,6 +104,20 @@ pub(crate) struct MeApiRuntimeSnapshot { pub reinit_inflight: usize, pub reinit_max_concurrency_effective: usize, pub hardswap_enabled: bool, + /// Number of authoritative warm writers in the pending generation. + pub pending_writers_current: usize, + /// Number of writers still required to reach the pending generation floor. + pub pending_writer_deficit: usize, + /// Number of desired DC-family groups below their pending-generation floor. + pub pending_missing_dc_groups: usize, + /// Whether the pending generation targets the current desired endpoint map. + pub pending_map_current: Option, + /// Number of warm writers not owned by the current pending generation. + pub orphan_warm_writers_current: usize, + /// Number of writer replacements in the preparatory phase. + pub replacement_preparing_current: usize, + /// Number of replacement victims closed to new client binds. + pub replacement_retiring_current: usize, pub floor_mode: &'static str, pub adaptive_floor_idle_secs: u64, pub adaptive_floor_min_writers_single_endpoint: u8, diff --git a/src/transport/middle_proxy/pool_status/hardswap_snapshot.rs b/src/transport/middle_proxy/pool_status/hardswap_snapshot.rs new file mode 100644 index 0000000..bdfbcbf --- /dev/null +++ b/src/transport/middle_proxy/pool_status/hardswap_snapshot.rs @@ -0,0 +1,99 @@ +use super::*; + +#[derive(Clone, Debug)] +/// Bounded control-plane snapshot of hardswap and replacement progress. +pub(crate) struct MeApiHardswapSnapshot { + /// Whether a hardswap generation is pending. + pub pending: bool, + /// Age of the pending generation when one exists. + pub pending_age_secs: Option, + /// Number of authoritative warm writers in the pending generation. + pub pending_writers_current: usize, + /// Number of writers still required to reach the pending generation floor. + pub pending_writer_deficit: usize, + /// Number of desired DC-family groups below their pending-generation writer floor. + pub pending_missing_dc_groups: usize, + /// Whether the pending generation targets the current desired endpoint map. + pub pending_map_current: Option, + /// Number of warm writers not owned by the current pending generation. + pub orphan_warm_writers_current: usize, + /// Number of writer replacements still in the preparatory phase. + pub replacement_preparing_current: usize, + /// Number of replacement victims already closed to new bindings. + pub replacement_retiring_current: usize, +} + +impl MePool { + /// Returns bounded hardswap progress without exposing generation or endpoint labels. + pub(crate) async fn api_hardswap_snapshot(&self) -> MeApiHardswapSnapshot { + let reinit = self.reinit.status.load_full(); + self.api_hardswap_snapshot_for_reinit(reinit.as_ref()).await + } + + /// Builds the bounded projection from one coherent reinitialization snapshot. + pub(super) async fn api_hardswap_snapshot_for_reinit( + &self, + reinit: &ReinitStatusSnapshot, + ) -> MeApiHardswapSnapshot { + let endpoint_snapshot = self.endpoint_snapshot.load_full(); + let desired_by_dc = self.desired_dc_endpoints_from_snapshot(&endpoint_snapshot); + let desired_hash = Self::desired_map_hash(&desired_by_dc); + let writers = self.writers.read().await; + let pending_generation = reinit.pending_hardswap_generation; + let pending = pending_generation != 0; + let pending_map_current = pending + && reinit.pending_hardswap_map_hash == desired_hash + && reinit.pending_hardswap_endpoint_revision == endpoint_snapshot.revision; + let mut pending_writers_current = 0usize; + let mut pending_writer_addrs = Vec::<(i32, SocketAddr)>::new(); + let mut orphan_warm_writers_current = 0usize; + + for writer in writers.iter() { + if writer.draining.load(Ordering::Acquire) { + continue; + } + let contour = WriterContour::from_u8(writer.contour.load(Ordering::Acquire)); + if contour == WriterContour::Warm + && (!pending_map_current || writer.generation != pending_generation) + { + orphan_warm_writers_current = orphan_warm_writers_current.saturating_add(1); + } + if pending_map_current + && writer.generation == pending_generation + && contour == WriterContour::Warm + && desired_by_dc + .get(&writer.writer_dc) + .is_some_and(|endpoints| endpoints.contains(&writer.addr)) + { + pending_writers_current = pending_writers_current.saturating_add(1); + pending_writer_addrs.push((writer.writer_dc, writer.addr)); + } + } + + let pending_coverage = pending_map_current + .then(|| self.hardswap_coverage(&desired_by_dc, &pending_writer_addrs)); + let pending_writer_deficit = pending_coverage + .as_ref() + .map_or(0, |coverage| coverage.writer_deficit); + let pending_missing_dc_groups = pending_coverage + .as_ref() + .map_or(0, |coverage| coverage.missing_groups.len()); + let (replacement_preparing_current, replacement_retiring_current) = + self.registry.writer_replacement_counts(); + let pending_age_secs = pending.then(|| { + Self::now_epoch_secs().saturating_sub(reinit.pending_hardswap_started_at_epoch_secs) + }); + + MeApiHardswapSnapshot { + pending, + pending_age_secs, + pending_writers_current, + pending_writer_deficit, + pending_missing_dc_groups, + pending_map_current: pending.then_some(pending_map_current), + orphan_warm_writers_current, + replacement_preparing_current, + replacement_retiring_current, + } + } +} diff --git a/src/transport/middle_proxy/pool_status/runtime_snapshot.rs b/src/transport/middle_proxy/pool_status/runtime_snapshot.rs index 9291d85..f466977 100644 --- a/src/transport/middle_proxy/pool_status/runtime_snapshot.rs +++ b/src/transport/middle_proxy/pool_status/runtime_snapshot.rs @@ -16,6 +16,7 @@ impl MePool { let pending_started_at = reinit.pending_hardswap_started_at_epoch_secs; let pending_hardswap_age_secs = (pending_started_at > 0).then_some(now_epoch_secs.saturating_sub(pending_started_at)); + let hardswap = self.api_hardswap_snapshot_for_reinit(reinit).await; let mut quarantined_endpoints = Vec::::new(); { @@ -64,6 +65,13 @@ impl MePool { .max_concurrency_effective .load(Ordering::Acquire), hardswap_enabled: self.reinit.hardswap.load(Ordering::Relaxed), + pending_writers_current: hardswap.pending_writers_current, + pending_writer_deficit: hardswap.pending_writer_deficit, + pending_missing_dc_groups: hardswap.pending_missing_dc_groups, + pending_map_current: hardswap.pending_map_current, + orphan_warm_writers_current: hardswap.orphan_warm_writers_current, + replacement_preparing_current: hardswap.replacement_preparing_current, + replacement_retiring_current: hardswap.replacement_retiring_current, floor_mode: floor_mode_label(self.floor_mode()), adaptive_floor_idle_secs: self .floor_runtime diff --git a/src/transport/middle_proxy/pool_status/status_snapshot.rs b/src/transport/middle_proxy/pool_status/status_snapshot.rs index 660db57..ffecb6f 100644 --- a/src/transport/middle_proxy/pool_status/status_snapshot.rs +++ b/src/transport/middle_proxy/pool_status/status_snapshot.rs @@ -3,12 +3,13 @@ use super::*; impl MePool { pub(crate) async fn admission_ready_conditional_cast(&self) -> bool { let mut endpoints_by_dc = BTreeMap::>::new(); + let endpoint_snapshot = self.endpoint_snapshot.load_full(); if self.decision.ipv4_me { - let map = self.proxy_map_v4.read().await.clone(); + let map = endpoint_snapshot.map_v4.clone(); extend_signed_endpoints(&mut endpoints_by_dc, map); } if self.decision.ipv6_me { - let map = self.proxy_map_v6.read().await.clone(); + let map = endpoint_snapshot.map_v6.clone(); extend_signed_endpoints(&mut endpoints_by_dc, map); } @@ -16,10 +17,15 @@ impl MePool { return false; } + let active_generation = self.reinit.status.load().active_generation; let writers = self.writers.read().await.clone(); let mut live_writers_by_dc = HashMap::::new(); for writer in writers.iter() { - if writer.draining.load(Ordering::Relaxed) { + if writer.draining.load(Ordering::Relaxed) + || writer.generation != active_generation + || WriterContour::from_u8(writer.contour.load(Ordering::Relaxed)) + != WriterContour::Active + { continue; } if let Ok(dc) = i16::try_from(writer.writer_dc) { @@ -40,12 +46,13 @@ impl MePool { #[allow(dead_code)] pub(crate) async fn admission_ready_full_floor(&self) -> bool { let mut endpoints_by_dc = BTreeMap::>::new(); + let endpoint_snapshot = self.endpoint_snapshot.load_full(); if self.decision.ipv4_me { - let map = self.proxy_map_v4.read().await.clone(); + let map = endpoint_snapshot.map_v4.clone(); extend_signed_endpoints(&mut endpoints_by_dc, map); } if self.decision.ipv6_me { - let map = self.proxy_map_v6.read().await.clone(); + let map = endpoint_snapshot.map_v6.clone(); extend_signed_endpoints(&mut endpoints_by_dc, map); } @@ -53,26 +60,34 @@ impl MePool { return false; } + let active_generation = self.reinit.status.load().active_generation; let writers = self.writers.read().await.clone(); - let mut live_writers_by_dc = HashMap::::new(); + let mut live_writers_by_group = HashMap::<(i16, bool), usize>::new(); for writer in writers.iter() { - if writer.draining.load(Ordering::Relaxed) { + if writer.draining.load(Ordering::Relaxed) + || writer.generation != active_generation + || WriterContour::from_u8(writer.contour.load(Ordering::Relaxed)) + != WriterContour::Active + { continue; } if let Ok(dc) = i16::try_from(writer.writer_dc) { - *live_writers_by_dc.entry(dc).or_insert(0) += 1; + *live_writers_by_group + .entry((dc, writer.addr.is_ipv4())) + .or_insert(0) += 1; } } for (dc, endpoints) in endpoints_by_dc { - let endpoint_count = endpoints.len(); - if endpoint_count == 0 { - return false; - } - let required = self.required_writers_for_dc_with_floor_mode(endpoint_count, false); - let alive = live_writers_by_dc.get(&dc).copied().unwrap_or(0); - if alive < required { - return false; + for (ipv4, endpoint_count) in endpoint_family_counts(&endpoints) { + if endpoint_count == 0 { + continue; + } + let required = self.required_writers_for_dc_with_floor_mode(endpoint_count, false); + let alive = live_writers_by_group.get(&(dc, ipv4)).copied().unwrap_or(0); + if alive < required { + return false; + } } } @@ -96,12 +111,13 @@ impl MePool { .load(Ordering::Relaxed); let mut endpoints_by_dc = BTreeMap::>::new(); + let endpoint_snapshot = self.endpoint_snapshot.load_full(); if self.decision.ipv4_me { - let map = self.proxy_map_v4.read().await.clone(); + let map = endpoint_snapshot.map_v4.clone(); extend_signed_endpoints(&mut endpoints_by_dc, map); } if self.decision.ipv6_me { - let map = self.proxy_map_v6.read().await.clone(); + let map = endpoint_snapshot.map_v6.clone(); extend_signed_endpoints(&mut endpoints_by_dc, map); } @@ -110,7 +126,12 @@ impl MePool { let required_writers = endpoints_by_dc .values() - .map(|endpoints| self.required_writers_for_dc_with_floor_mode(endpoints.len(), false)) + .map(|endpoints| { + endpoint_family_counts(endpoints) + .into_iter() + .map(|(_, count)| self.required_writers_for_dc_with_floor_mode(count, false)) + .sum::() + }) .sum(); let idle_since = self.registry.writer_idle_since_snapshot().await; @@ -127,7 +148,7 @@ impl MePool { for writer in writers.iter() { let endpoint = writer.addr; let dc = i16::try_from(writer.writer_dc).ok(); - let draining = writer.draining.load(Ordering::Relaxed); + let draining = writer.draining.load(Ordering::Acquire); let degraded = writer.degraded.load(Ordering::Relaxed); let matches_active_generation = writer.generation == active_generation; let in_desired_map = dc @@ -156,13 +177,18 @@ impl MePool { && drain_ttl_secs > 0 && drain_started_at_epoch_secs .is_some_and(|started| now_epoch_secs.saturating_sub(started) > drain_ttl_secs); - let state = match WriterContour::from_u8(writer.contour.load(Ordering::Relaxed)) { + let contour = WriterContour::from_u8(writer.contour.load(Ordering::Relaxed)); + let state = match contour { WriterContour::Warm => "warm", WriterContour::Active => "active", WriterContour::Draining => "draining", }; - if !draining && let Some(dc_idx) = dc { + let authoritative_active = !draining + && contour == WriterContour::Active + && matches_active_generation + && in_desired_map; + if authoritative_active && let Some(dc_idx) = dc { *live_writers_by_dc_endpoint .entry((dc_idx, endpoint)) .or_insert(0) += 1; @@ -172,9 +198,7 @@ impl MePool { entry.0 += ema_ms; entry.1 += 1; } - if matches_active_generation && in_desired_map { - *fresh_writers_by_dc.entry(dc_idx).or_insert(0) += 1; - } + *fresh_writers_by_dc.entry(dc_idx).or_insert(0) += 1; } writer_rows.push(MeApiWriterStatusSnapshot { @@ -211,39 +235,55 @@ impl MePool { .max(1); for (dc, endpoints) in endpoints_by_dc { let endpoint_count = endpoints.len(); + let family_counts = endpoint_family_counts(&endpoints); let dc_available_endpoints = endpoints .iter() .filter(|endpoint| live_writers_by_dc_endpoint.contains_key(&(dc, **endpoint))) .count(); - let base_required = self.required_writers_for_dc(endpoint_count); - let dc_required_writers = - self.required_writers_for_dc_with_floor_mode(endpoint_count, false); - let floor_min = if endpoint_count <= 1 { - (self - .floor_runtime - .me_adaptive_floor_min_writers_single_endpoint - .load(Ordering::Relaxed) as usize) - .max(1) - .min(base_required.max(1)) - } else { - (self - .floor_runtime - .me_adaptive_floor_min_writers_multi_endpoint - .load(Ordering::Relaxed) as usize) - .max(1) - .min(base_required.max(1)) - }; - let extra_per_core = if endpoint_count <= 1 { - self.floor_runtime - .me_adaptive_floor_max_extra_writers_single_per_core - .load(Ordering::Relaxed) as usize - } else { - self.floor_runtime - .me_adaptive_floor_max_extra_writers_multi_per_core - .load(Ordering::Relaxed) as usize - }; - let floor_max = - base_required.saturating_add(adaptive_cpu_cores.saturating_mul(extra_per_core)); + let base_required = family_counts + .iter() + .filter(|(_, count)| *count > 0) + .map(|(_, count)| self.required_writers_for_dc(*count)) + .sum::(); + let dc_required_writers = family_counts + .iter() + .filter(|(_, count)| *count > 0) + .map(|(_, count)| self.required_writers_for_dc_with_floor_mode(*count, false)) + .sum::(); + let floor_min = family_counts + .iter() + .filter(|(_, count)| *count > 0) + .map(|(_, count)| { + let family_base = self.required_writers_for_dc(*count); + let configured = if *count <= 1 { + self.floor_runtime + .me_adaptive_floor_min_writers_single_endpoint + .load(Ordering::Relaxed) as usize + } else { + self.floor_runtime + .me_adaptive_floor_min_writers_multi_endpoint + .load(Ordering::Relaxed) as usize + }; + configured.max(1).min(family_base.max(1)) + }) + .sum::(); + let floor_max = family_counts + .iter() + .filter(|(_, count)| *count > 0) + .map(|(_, count)| { + let family_base = self.required_writers_for_dc(*count); + let extra_per_core = if *count <= 1 { + self.floor_runtime + .me_adaptive_floor_max_extra_writers_single_per_core + .load(Ordering::Relaxed) as usize + } else { + self.floor_runtime + .me_adaptive_floor_max_extra_writers_multi_per_core + .load(Ordering::Relaxed) as usize + }; + family_base.saturating_add(adaptive_cpu_cores.saturating_mul(extra_per_core)) + }) + .sum::(); let floor_capped = matches!(floor_mode, MeFloorMode::Adaptive) && dc_required_writers < base_required; let dc_alive_writers = live_writers_by_dc.get(&dc).copied().unwrap_or(0); @@ -306,3 +346,11 @@ impl MePool { } } } + +fn endpoint_family_counts(endpoints: &BTreeSet) -> [(bool, usize); 2] { + let ipv4 = endpoints + .iter() + .filter(|endpoint| endpoint.is_ipv4()) + .count(); + [(true, ipv4), (false, endpoints.len().saturating_sub(ipv4))] +} diff --git a/src/transport/middle_proxy/pool_status/tests.rs b/src/transport/middle_proxy/pool_status/tests.rs new file mode 100644 index 0000000..ef479ce --- /dev/null +++ b/src/transport/middle_proxy/pool_status/tests.rs @@ -0,0 +1,74 @@ +use std::collections::HashMap; +use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr}; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU32, AtomicU64}; +use std::time::Instant; + +use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; + +use crate::network::probe::NetworkDecision; +use crate::transport::middle_proxy::codec::WriterCommand; +use crate::transport::middle_proxy::pool::{MePool, MeWriter, WriterContour}; +use crate::transport::middle_proxy::pool_writer_security_tests::make_pool_with_decision; + +fn writer(pool: &Arc, id: u64, dc: i32, addr: SocketAddr) -> MeWriter { + let (tx, _rx) = mpsc::channel::(8); + MeWriter { + id, + addr, + source_ip: addr.ip(), + writer_dc: dc, + generation: pool.current_generation(), + contour: Arc::new(AtomicU8::new(WriterContour::Active.as_u8())), + created_at: Instant::now(), + tx, + byte_budget: pool.new_writer_byte_budget(), + cancel: CancellationToken::new(), + degraded: Arc::new(AtomicBool::new(false)), + rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)), + draining: Arc::new(AtomicBool::new(false)), + draining_started_at_epoch_secs: Arc::new(AtomicU64::new(0)), + drain_deadline_epoch_secs: Arc::new(AtomicU64::new(0)), + allow_drain_fallback: Arc::new(AtomicBool::new(false)), + } +} + +#[tokio::test] +async fn dual_family_status_reports_each_family_floor() { + let pool = make_pool_with_decision(NetworkDecision { + ipv4_me: true, + ipv6_me: true, + effective_prefer: 4, + effective_multipath: false, + ..NetworkDecision::default() + }) + .await; + let v4 = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443); + let v6 = SocketAddr::new(IpAddr::V6(Ipv6Addr::LOCALHOST), 443); + pool.update_proxy_maps( + HashMap::from([(2, vec![(v4.ip(), v4.port())])]), + Some(HashMap::from([(2, vec![(v6.ip(), v6.port())])])), + ) + .await; + let required_per_family = pool.required_writers_for_dc(1); + let mut writers = pool.writers.write().await; + for (group, dc) in [2, -2].into_iter().enumerate() { + for offset in 0..required_per_family { + writers.push(writer(&pool, (group as u64 * 100) + offset as u64, dc, v4)); + } + } + drop(writers); + + let snapshot = pool.api_status_snapshot().await; + + assert_eq!(snapshot.required_writers, required_per_family * 4); + assert_eq!(snapshot.alive_writers, required_per_family * 2); + assert_eq!(snapshot.coverage_pct, 50.0); + for dc in snapshot.dcs { + assert_eq!(dc.required_writers, required_per_family * 2); + assert_eq!(dc.alive_writers, required_per_family); + assert_eq!(dc.coverage_pct, 50.0); + } + assert!(!pool.admission_ready_full_floor().await); +} diff --git a/src/transport/middle_proxy/pool_writer.rs b/src/transport/middle_proxy/pool_writer.rs index 6841a29..a02c868 100644 --- a/src/transport/middle_proxy/pool_writer.rs +++ b/src/transport/middle_proxy/pool_writer.rs @@ -1,6 +1,8 @@ use std::collections::HashMap; +use std::future::Future; use std::io::ErrorKind; use std::net::SocketAddr; +use std::pin::Pin; use std::sync::Arc; use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU32, AtomicU64, Ordering}; use std::time::{Duration, Instant}; @@ -18,13 +20,61 @@ use crate::error::{ProxyError, Result}; use crate::protocol::constants::{RPC_CLOSE_EXT_U32, RPC_PING_U32}; use super::codec::{RpcWriter, WriterCommand, build_control_payload}; -use super::pool::{MePool, MeWriter, WriterContour}; +use super::pool::{ + MePool, MeWriter, WriterContour, WriterOpenIntent, WriterOpenReservation, WriterRole, +}; +use super::pool_lifecycle::MeTaskRegistration; use super::reader::reader_loop; +use super::registry::WriterReplacementReservation; use super::wire::build_proxy_req_payload; // Writer admission, teardown, and drain-state transitions. +mod publication; +mod replacement; mod runtime; +struct PreparedWriter<'a> { + writer: MeWriter, + tx: mpsc::Sender, + byte_budget: Arc, + task_registration: MeTaskRegistration<'a>, + writer_task: Pin + Send + 'static>>, + intent: WriterOpenIntent, + _open_reservation: WriterOpenReservation<'a>, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +/// Commit-time invariant applied to a transactional writer replacement. +pub(super) enum WriterReplacementPurpose { + /// Replaces an idle writer before the upstream idle lifetime expires. + IdleRefresh, + /// Moves one idle slot from a donor above floor to a receiver below floor. + FloorRebalance { + /// Minimum writer count that must remain in the donor group. + donor_floor: usize, + /// Target count above which the receiver no longer needs the slot. + receiver_floor: usize, + }, + /// Replaces a stale-generation writer with an active-generation writer. + GenerationConvergence, + /// Rotates a single-endpoint shadow while retaining coverage. + ShadowRotation, + /// Replaces a writer after cryptographic secret rotation. + SecretRotation, +} + +impl WriterReplacementPurpose { + fn as_str(self) -> &'static str { + match self { + Self::IdleRefresh => "idle_refresh", + Self::FloorRebalance { .. } => "floor_rebalance", + Self::GenerationConvergence => "generation_convergence", + Self::ShadowRotation => "shadow_rotation", + Self::SecretRotation => "secret_rotation", + } + } +} + const ME_ACTIVE_PING_SECS: u64 = 25; const ME_ACTIVE_PING_JITTER_SECS: i64 = 5; const ME_IDLE_KEEPALIVE_MAX_SECS: u64 = 5; diff --git a/src/transport/middle_proxy/pool_writer/publication.rs b/src/transport/middle_proxy/pool_writer/publication.rs new file mode 100644 index 0000000..5647f95 --- /dev/null +++ b/src/transport/middle_proxy/pool_writer/publication.rs @@ -0,0 +1,162 @@ +use super::*; + +impl MePool { + /// Publishes a connected writer only while its generation owns the requested role. + pub(super) async fn publish_connected_writer( + self: &Arc, + prepared: PreparedWriter<'_>, + ) -> Result<()> { + let PreparedWriter { + writer, + tx, + byte_budget, + task_registration, + writer_task, + intent, + _open_reservation, + } = prepared; + let writer_id = writer.id; + let mut writers = self.writers.write().await; + let mut registry_registration = self.registry.prepare_writer_registration().await; + let coordinator = self.reinit.coordinator.lock(); + let contour = self.authorize_writer_publication(&writer, &coordinator)?; + self.authorize_writer_publication_capacity(&writer, contour, intent, writers.as_slice())?; + writer.contour.store(contour.as_u8(), Ordering::Release); + registry_registration.install(writer_id, tx, byte_budget); + writers.push(writer); + self.conn_count.fetch_add(1, Ordering::Relaxed); + self.lifecycle + .spawn_registered_writer(task_registration, writer_task); + drop(coordinator); + drop(registry_registration); + drop(writers); + self.notify_writer_epoch(); + Ok(()) + } + + /// Resolves the writer role against the linearized reinitialization authority. + pub(super) fn authorize_writer_publication( + &self, + writer: &MeWriter, + coordinator: &crate::transport::middle_proxy::pool::ReinitCoordinatorState, + ) -> Result { + let endpoint_is_current = self + .endpoint_snapshot + .load() + .contains_dc_endpoint(writer.writer_dc, writer.addr); + if !endpoint_is_current { + return Err(ProxyError::Proxy( + "ME writer target changed before publication".into(), + )); + } + + let requested = WriterContour::from_u8(writer.contour.load(Ordering::Acquire)); + if writer.generation == coordinator.active_generation + && matches!(requested, WriterContour::Active | WriterContour::Warm) + { + return Ok(WriterContour::Active); + } + if requested == WriterContour::Warm + && coordinator.pending.is_some_and(|pending| { + pending.generation == writer.generation + && pending.map_hash == coordinator.desired_map_hash + && pending.endpoint_revision == coordinator.endpoint_revision + }) + { + return Ok(WriterContour::Warm); + } + Err(ProxyError::Proxy( + "ME writer generation lost publication authority".into(), + )) + } + + /// Revalidates role-local capacity at the serialized publication boundary. + pub(in crate::transport::middle_proxy) fn authorize_writer_publication_capacity( + &self, + writer: &MeWriter, + contour: WriterContour, + intent: WriterOpenIntent, + writers: &[MeWriter], + ) -> Result<()> { + let family = if writer.addr.is_ipv4() { + crate::network::IpFamily::V4 + } else { + crate::network::IpFamily::V6 + }; + let now_epoch_secs = Self::now_epoch_secs(); + if !self.family_enabled_for_drain_coverage(family, now_epoch_secs) { + return Err(ProxyError::Proxy( + "ME writer family lost publication authority".into(), + )); + } + if intent == WriterOpenIntent::Replacement || contour == WriterContour::Draining { + return Ok(()); + } + let endpoint_snapshot = self.endpoint_snapshot.load(); + if !endpoint_snapshot.contains_dc_endpoint(writer.writer_dc, writer.addr) { + return Err(ProxyError::Proxy( + "ME writer target changed before publication".into(), + )); + } + if contour == WriterContour::Active && intent == WriterOpenIntent::Normal { + let current = writers + .iter() + .filter(|candidate| { + !candidate.draining.load(Ordering::Acquire) + && WriterContour::from_u8(candidate.contour.load(Ordering::Acquire)) + == WriterContour::Active + }) + .count(); + if current >= self.adaptive_floor_active_cap_configured_total() { + return Err(ProxyError::Proxy( + "ME active writer cap was reached before publication".into(), + )); + } + return Ok(()); + } + let endpoints = endpoint_snapshot.endpoints_for_dc_family(writer.writer_dc, family); + let required = self.required_writers_for_dc(endpoints.len()); + let current = writers + .iter() + .filter(|candidate| { + !candidate.draining.load(Ordering::Acquire) + && candidate.writer_dc == writer.writer_dc + && candidate.generation == writer.generation + && WriterContour::from_u8(candidate.contour.load(Ordering::Acquire)) == contour + && candidate.addr.is_ipv4() == writer.addr.is_ipv4() + && endpoint_snapshot.contains_dc_endpoint(candidate.writer_dc, candidate.addr) + }) + .count(); + if current >= required { + return Err(ProxyError::Proxy( + "ME writer floor was restored before publication".into(), + )); + } + Ok(()) + } + + /// Commits writer visibility and lifecycle ownership after all cancellation points. + #[allow(clippy::too_many_arguments)] + pub(in crate::transport::middle_proxy) async fn publish_prepared_writer( + self: &Arc, + writer: MeWriter, + tx: mpsc::Sender, + byte_budget: Arc, + task_registration: MeTaskRegistration<'_>, + writer_task: F, + ) where + F: Future + Send + 'static, + { + // Writer publication follows the global writers -> registry binding lock order. + let mut writers = self.writers.write().await; + let mut registry_registration = self.registry.prepare_writer_registration().await; + registry_registration.install(writer.id, tx, byte_budget); + writers.push(writer); + self.conn_count.fetch_add(1, Ordering::Relaxed); + self.lifecycle + .spawn_registered_writer(task_registration, writer_task); + drop(writers); + drop(registry_registration); + self.notify_writer_epoch(); + } +} diff --git a/src/transport/middle_proxy/pool_writer/replacement.rs b/src/transport/middle_proxy/pool_writer/replacement.rs new file mode 100644 index 0000000..b195096 --- /dev/null +++ b/src/transport/middle_proxy/pool_writer/replacement.rs @@ -0,0 +1,393 @@ +use super::*; + +impl MePool { + /// Opens a replacement and atomically publishes it before retiring its reserved victim. + pub(in crate::transport::middle_proxy) async fn replace_writer_with_generation_contour_for_dc( + self: &Arc, + addr: SocketAddr, + rng: &SecureRandom, + generation: u64, + contour: WriterContour, + writer_dc: i32, + expected_victim_role: WriterRole, + purpose: WriterReplacementPurpose, + reservation: &mut WriterReplacementReservation<'_>, + ) -> Result<()> { + let prepared = self + .prepare_writer_with_intent( + addr, + rng, + generation, + contour, + writer_dc, + WriterOpenIntent::Replacement, + ) + .await?; + self.publish_prepared_replacement_writer( + prepared, + expected_victim_role, + purpose, + reservation, + ) + .await + } + + async fn publish_prepared_replacement_writer( + self: &Arc, + prepared: PreparedWriter<'_>, + expected_victim_role: WriterRole, + purpose: WriterReplacementPurpose, + reservation: &mut WriterReplacementReservation<'_>, + ) -> Result<()> { + let PreparedWriter { + writer, + tx, + byte_budget, + task_registration, + writer_task, + intent, + _open_reservation, + } = prepared; + let replacement_writer_id = writer.id; + let victim_writer_id = reservation.writer_id(); + + // Lock order is writers -> registry binding. No cancellation point follows acquisition of + // the registry guard, so publication and victim retirement commit as one state change. + let mut writers = self.writers.write().await; + let mut registry_registration = self.registry.prepare_writer_registration().await; + let coordinator = self.reinit.coordinator.lock(); + let contour = self.authorize_writer_publication(&writer, &coordinator)?; + self.authorize_writer_publication_capacity(&writer, contour, intent, writers.as_slice())?; + writer.contour.store(contour.as_u8(), Ordering::Release); + let Some(victim_pos) = writers + .iter() + .position(|candidate| candidate.id == victim_writer_id) + else { + return Err(ProxyError::Proxy( + "ME replacement victim disappeared before commit".into(), + )); + }; + let victim = &writers[victim_pos]; + if victim.draining.load(Ordering::Acquire) || !expected_victim_role.matches(victim) { + return Err(ProxyError::Proxy( + "ME replacement victim changed role before commit".into(), + )); + } + if let WriterReplacementPurpose::FloorRebalance { + donor_floor, + receiver_floor, + } = purpose + { + if expected_victim_role.generation != coordinator.active_generation + || writer.generation != coordinator.active_generation + { + return Err(ProxyError::Proxy( + "ME floor rebalance lost active-generation authority".into(), + )); + } + let endpoint_snapshot = self.endpoint_snapshot.load(); + let donor_count = writers + .iter() + .filter(|candidate| { + !candidate.draining.load(Ordering::Acquire) + && candidate.writer_dc == expected_victim_role.dc + && candidate.generation == expected_victim_role.generation + && WriterContour::from_u8(candidate.contour.load(Ordering::Acquire)) + == WriterContour::Active + && endpoint_snapshot + .contains_dc_endpoint(candidate.writer_dc, candidate.addr) + && (candidate.addr.is_ipv4() + == matches!(expected_victim_role.family, crate::network::IpFamily::V4)) + }) + .count(); + let receiver_family = if writer.addr.is_ipv4() { + crate::network::IpFamily::V4 + } else { + crate::network::IpFamily::V6 + }; + let receiver_count = writers + .iter() + .filter(|candidate| { + !candidate.draining.load(Ordering::Acquire) + && candidate.writer_dc == writer.writer_dc + && candidate.generation == writer.generation + && WriterContour::from_u8(candidate.contour.load(Ordering::Acquire)) + == WriterContour::Active + && endpoint_snapshot + .contains_dc_endpoint(candidate.writer_dc, candidate.addr) + && (if candidate.addr.is_ipv4() { + crate::network::IpFamily::V4 + } else { + crate::network::IpFamily::V6 + }) == receiver_family + }) + .count(); + if donor_count <= donor_floor || receiver_count >= receiver_floor { + return Err(ProxyError::Proxy( + "ME floor rebalance became unnecessary before commit".into(), + )); + } + } + if !registry_registration.prepare_replacement_commit(reservation) { + return Err(ProxyError::Proxy( + "ME replacement victim became active before commit".into(), + )); + } + + registry_registration.install(replacement_writer_id, tx, byte_budget); + writers.push(writer); + self.conn_count.fetch_add(1, Ordering::Relaxed); + writers.publish_current(); + self.apply_writer_draining_state(&writers[victim_pos], self.force_close_timeout(), false); + self.lifecycle + .spawn_registered_writer(task_registration, writer_task); + reservation.mark_committed(); + drop(coordinator); + drop(registry_registration); + drop(writers); + self.notify_writer_epoch(); + info!( + victim_writer_id, + replacement_writer_id, + purpose = purpose.as_str(), + "ME writer replacement committed" + ); + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::net::{IpAddr, Ipv4Addr}; + + use crate::transport::middle_proxy::pool_writer_security_tests::make_pool; + + use super::*; + + fn endpoint(octet: u8) -> SocketAddr { + SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, octet)), 443) + } + + async fn install_writer( + pool: &Arc, + writer_id: u64, + writer_dc: i32, + addr: SocketAddr, + ) -> MeWriter { + let (tx, _rx) = mpsc::channel::(8); + let byte_budget = pool.new_writer_byte_budget(); + let writer = MeWriter { + id: writer_id, + addr, + source_ip: addr.ip(), + writer_dc, + generation: pool.current_generation(), + contour: Arc::new(AtomicU8::new(WriterContour::Active.as_u8())), + created_at: Instant::now(), + tx: tx.clone(), + byte_budget: byte_budget.clone(), + cancel: CancellationToken::new(), + degraded: Arc::new(AtomicBool::new(false)), + rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)), + draining: Arc::new(AtomicBool::new(false)), + draining_started_at_epoch_secs: Arc::new(AtomicU64::new(0)), + drain_deadline_epoch_secs: Arc::new(AtomicU64::new(0)), + allow_drain_fallback: Arc::new(AtomicBool::new(false)), + }; + let mut writers = pool.writers.write().await; + let mut registry_registration = pool.registry.prepare_writer_registration().await; + registry_registration.install(writer_id, tx, byte_budget); + writers.push(writer.clone()); + pool.conn_count.fetch_add(1, Ordering::Relaxed); + drop(registry_registration); + drop(writers); + writer + } + + async fn prepared_writer<'a>( + pool: &'a Arc, + writer_id: u64, + writer_dc: i32, + addr: SocketAddr, + ) -> PreparedWriter<'a> { + let open_reservation = pool + .reserve_writer_open( + WriterContour::Active, + WriterOpenIntent::Replacement, + writer_dc, + addr, + ) + .await + .expect("replacement open must be admitted"); + let task_registration = pool + .lifecycle + .try_register() + .expect("test pool lifecycle must be open"); + let (tx, rx) = mpsc::channel::(8); + let byte_budget = pool.new_writer_byte_budget(); + PreparedWriter { + writer: MeWriter { + id: writer_id, + addr, + source_ip: addr.ip(), + writer_dc, + generation: pool.current_generation(), + contour: Arc::new(AtomicU8::new(WriterContour::Active.as_u8())), + created_at: Instant::now(), + tx: tx.clone(), + byte_budget: byte_budget.clone(), + cancel: CancellationToken::new(), + degraded: Arc::new(AtomicBool::new(false)), + rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)), + draining: Arc::new(AtomicBool::new(false)), + draining_started_at_epoch_secs: Arc::new(AtomicU64::new(0)), + drain_deadline_epoch_secs: Arc::new(AtomicU64::new(0)), + allow_drain_fallback: Arc::new(AtomicBool::new(false)), + }, + tx, + byte_budget, + task_registration, + writer_task: Box::pin(async move { + drop(rx); + }), + intent: WriterOpenIntent::Replacement, + _open_reservation: open_reservation, + } + } + + #[tokio::test] + async fn replacement_commit_publishes_successor_before_draining_victim() { + let pool = make_pool().await; + let addr = endpoint(1); + pool.update_proxy_maps(HashMap::from([(2, vec![(addr.ip(), addr.port())])]), None) + .await; + let victim = install_writer(&pool, 1001, 2, addr).await; + let expected_role = WriterRole::from_writer(&victim); + let mut reservation = pool + .registry + .try_reserve_writer_replacement(victim.id) + .await + .expect("idle victim must be reservable"); + let prepared = prepared_writer(&pool, 1002, 2, addr).await; + + pool.publish_prepared_replacement_writer( + prepared, + expected_role, + WriterReplacementPurpose::IdleRefresh, + &mut reservation, + ) + .await + .expect("replacement commit must succeed"); + + let writers = pool.writers.read().await; + assert_eq!(writers.len(), 2); + assert_eq!( + writers + .iter() + .filter(|writer| !writer.draining.load(Ordering::Acquire)) + .count(), + 1 + ); + assert!(victim.draining.load(Ordering::Acquire)); + assert!(writers.iter().any(|writer| writer.id == 1002)); + drop(writers); + assert_eq!(pool.conn_count.load(Ordering::Acquire), 2); + assert_eq!(pool.registry.writer_replacement_counts(), (0, 1)); + } + + #[tokio::test] + async fn cancelled_replacement_waiting_for_publication_restores_all_reservations() { + let pool = make_pool().await; + let addr = endpoint(2); + pool.update_proxy_maps(HashMap::from([(2, vec![(addr.ip(), addr.port())])]), None) + .await; + let victim = install_writer(&pool, 2001, 2, addr).await; + let expected_role = WriterRole::from_writer(&victim); + let mut reservation = pool + .registry + .try_reserve_writer_replacement(victim.id) + .await + .expect("idle victim must be reservable"); + let prepared = prepared_writer(&pool, 2002, 2, addr).await; + let writers_guard = pool.writers.write().await; + + let result = tokio::time::timeout( + Duration::from_millis(10), + pool.publish_prepared_replacement_writer( + prepared, + expected_role, + WriterReplacementPurpose::IdleRefresh, + &mut reservation, + ), + ) + .await; + + assert!(result.is_err()); + drop(writers_guard); + drop(reservation); + assert_eq!( + pool.writer_replacement_open_reserved + .load(Ordering::Acquire), + 0 + ); + assert_eq!(pool.registry.writer_replacement_counts(), (0, 0)); + assert!(!victim.draining.load(Ordering::Acquire)); + assert!( + !pool + .writers + .read() + .await + .iter() + .any(|writer| writer.id == 2002) + ); + } + + #[tokio::test] + async fn floor_rebalance_commit_rejects_a_donor_without_surplus() { + let pool = make_pool().await; + let donor_addr = endpoint(3); + let receiver_addr = endpoint(4); + pool.update_proxy_maps( + HashMap::from([ + (1, vec![(donor_addr.ip(), donor_addr.port())]), + (2, vec![(receiver_addr.ip(), receiver_addr.port())]), + ]), + None, + ) + .await; + let victim = install_writer(&pool, 3001, 1, donor_addr).await; + let expected_role = WriterRole::from_writer(&victim); + let mut reservation = pool + .registry + .try_reserve_writer_replacement(victim.id) + .await + .expect("idle victim must be reservable"); + let prepared = prepared_writer(&pool, 3002, 2, receiver_addr).await; + + let result = pool + .publish_prepared_replacement_writer( + prepared, + expected_role, + WriterReplacementPurpose::FloorRebalance { + donor_floor: 1, + receiver_floor: 1, + }, + &mut reservation, + ) + .await; + + assert!(result.is_err()); + drop(reservation); + assert!(!victim.draining.load(Ordering::Acquire)); + assert!( + !pool + .writers + .read() + .await + .iter() + .any(|writer| writer.id == 3002) + ); + assert_eq!(pool.registry.writer_replacement_counts(), (0, 0)); + } +} diff --git a/src/transport/middle_proxy/pool_writer/runtime.rs b/src/transport/middle_proxy/pool_writer/runtime.rs index d59218d..8ec9cdf 100644 --- a/src/transport/middle_proxy/pool_writer/runtime.rs +++ b/src/transport/middle_proxy/pool_writer/runtime.rs @@ -54,23 +54,45 @@ impl MePool { contour: WriterContour, writer_dc: i32, ) -> Result<()> { - self.connect_one_with_generation_contour_for_dc_with_cap_policy( - addr, rng, generation, contour, writer_dc, false, + self.connect_one_with_generation_contour_for_dc_with_intent( + addr, + rng, + generation, + contour, + writer_dc, + WriterOpenIntent::Normal, ) .await } - pub(in crate::transport::middle_proxy) async fn connect_one_with_generation_contour_for_dc_with_cap_policy( + /// Connects and publishes one writer under an explicit capacity intent. + pub(in crate::transport::middle_proxy) async fn connect_one_with_generation_contour_for_dc_with_intent( self: &Arc, addr: SocketAddr, rng: &SecureRandom, generation: u64, contour: WriterContour, writer_dc: i32, - allow_coverage_override: bool, + intent: WriterOpenIntent, ) -> Result<()> { - let Some(_writer_open_reservation) = self - .reserve_writer_open(contour, allow_coverage_override, writer_dc) + let prepared = self + .prepare_writer_with_intent(addr, rng, generation, contour, writer_dc, intent) + .await?; + self.publish_connected_writer(prepared).await + } + + /// Completes all cancellable connection work without publishing pool visibility. + pub(super) async fn prepare_writer_with_intent<'a>( + self: &'a Arc, + addr: SocketAddr, + rng: &SecureRandom, + generation: u64, + contour: WriterContour, + writer_dc: i32, + intent: WriterOpenIntent, + ) -> Result> { + let Some(writer_open_reservation) = self + .reserve_writer_open(contour, intent, writer_dc, addr) .await else { return Err(ProxyError::Proxy(format!( @@ -132,16 +154,6 @@ impl MePool { drain_deadline_epoch_secs: drain_deadline_epoch_secs.clone(), allow_drain_fallback: allow_drain_fallback.clone(), }; - self.writers - .update(|writers| writers.push(writer.clone())) - .await; - self.registry - .register_writer(writer_id, tx.clone(), byte_budget) - .await; - self.registry.mark_writer_idle(writer_id).await; - self.conn_count.fetch_add(1, Ordering::Relaxed); - self.notify_writer_epoch(); - let reg = self.registry.clone(); let writers_arc = self.writers_arc(); let ping_tracker = Arc::new(tokio::sync::Mutex::new(HashMap::::new())); @@ -177,8 +189,10 @@ impl MePool { let route_fairshare_enabled = self.transport_policy.me_route_fairshare_enabled.clone(); let reader_route_data_wait_ms = self.transport_policy.me_reader_route_data_wait_ms.clone(); - self.lifecycle - .spawn_registered_writer(task_registration, async move { + let writer_task = { + // Keep transport ownership behind a stable-size pointer while publication waits for + // locks. + Box::pin(async move { // Reader MUST be the first branch in biased select! to avoid read starvation. let exit = tokio::select! { biased; @@ -269,9 +283,18 @@ impl MePool { let remaining = writers_arc.read().await.len(); debug!(writer_id, remaining, "ME writer lifecycle task finished"); - }); + }) + }; - Ok(()) + Ok(PreparedWriter { + writer, + tx, + byte_budget, + task_registration, + writer_task, + intent, + _open_reservation: writer_open_reservation, + }) } pub(crate) async fn remove_writer_and_close_clients(self: &Arc, writer_id: u64) { @@ -309,7 +332,7 @@ impl MePool { ) -> bool { let mut close_tx: Option> = None; let mut removed_addr: Option = None; - let mut removed_dc: Option = None; + let mut removed_role: Option = None; let mut removed_uptime: Option = None; let mut trigger_refill = false; let mut removed = false; @@ -330,14 +353,19 @@ impl MePool { self.stats.increment_me_writer_removed_total(); w.cancel.cancel(); removed_addr = Some(w.addr); - removed_dc = Some(w.writer_dc); + removed_role = Some(WriterRole::from_writer(&w)); removed_uptime = Some(w.created_at.elapsed()); trigger_refill = !was_draining; if trigger_refill { self.stats.increment_me_writer_removed_unexpected_total(); } close_tx = Some(w.tx.clone()); - self.conn_count.fetch_sub(1, Ordering::Relaxed); + // Teardown remains idempotent if the advisory count was already reconciled. + let _ = + self.conn_count + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |count| { + count.checked_sub(1) + }); removed = true; } } @@ -370,8 +398,8 @@ impl MePool { ); } } - if trigger_refill && let Some(writer_dc) = removed_dc { - self.trigger_immediate_refill_for_dc(addr, writer_dc); + if trigger_refill && let Some(role) = removed_role { + self.trigger_immediate_refill_for_role(addr, role); } } if removed { @@ -388,26 +416,13 @@ impl MePool { ) { let timeout = timeout.filter(|d| !d.is_zero()); let found = { - let mut ws = self.writers.write().await; - if let Some(w) = ws.iter_mut().find(|w| w.id == writer_id) { - let already_draining = w.draining.swap(true, Ordering::Relaxed); - w.allow_drain_fallback - .store(allow_drain_fallback, Ordering::Relaxed); - let now_epoch_secs = Self::now_epoch_secs(); - w.draining_started_at_epoch_secs - .store(now_epoch_secs, Ordering::Relaxed); - let drain_deadline_epoch_secs = timeout - .map(|duration| now_epoch_secs.saturating_add(duration.as_secs())) - .unwrap_or(0); - w.drain_deadline_epoch_secs - .store(drain_deadline_epoch_secs, Ordering::Relaxed); - if !already_draining { - self.stats.increment_pool_drain_active(); - self.increment_draining_active_runtime(); + let ws = self.writers.write().await; + let mut registry_registration = self.registry.prepare_writer_registration().await; + if let Some(w) = ws.iter().find(|w| w.id == writer_id) { + if !allow_drain_fallback { + registry_registration.retire(writer_id); } - w.contour - .store(WriterContour::Draining.as_u8(), Ordering::Relaxed); - w.draining.store(true, Ordering::Relaxed); + self.apply_writer_draining_state(w, timeout, allow_drain_fallback); true } else { false @@ -425,6 +440,39 @@ impl MePool { ); } + /// Publishes drain metadata before the release-store that makes draining observable. + pub(in crate::transport::middle_proxy) fn apply_writer_draining_state( + &self, + writer: &MeWriter, + timeout: Option, + allow_drain_fallback: bool, + ) -> bool { + let timeout = timeout.filter(|duration| !duration.is_zero()); + let already_draining = writer.draining.load(Ordering::Acquire); + writer + .allow_drain_fallback + .store(allow_drain_fallback, Ordering::Release); + let now_epoch_secs = Self::now_epoch_secs(); + writer + .draining_started_at_epoch_secs + .store(now_epoch_secs, Ordering::Release); + let drain_deadline_epoch_secs = timeout + .map(|duration| now_epoch_secs.saturating_add(duration.as_secs())) + .unwrap_or(0); + writer + .drain_deadline_epoch_secs + .store(drain_deadline_epoch_secs, Ordering::Release); + writer + .contour + .store(WriterContour::Draining.as_u8(), Ordering::Release); + writer.draining.store(true, Ordering::Release); + if !already_draining { + self.stats.increment_pool_drain_active(); + self.increment_draining_active_runtime(); + } + !already_draining + } + pub(crate) async fn mark_writer_draining(self: &Arc, writer_id: u64) { self.mark_writer_draining_with_timeout(writer_id, Some(Duration::from_secs(300)), false) .await; @@ -434,7 +482,7 @@ impl MePool { &self, writer: &MeWriter, ) -> bool { - if !writer.draining.load(Ordering::Relaxed) { + if !writer.draining.load(Ordering::Acquire) { return true; } if !writer.allow_drain_fallback.load(Ordering::Relaxed) { diff --git a/src/transport/middle_proxy/registry.rs b/src/transport/middle_proxy/registry.rs index e97a233..29e4315 100644 --- a/src/transport/middle_proxy/registry.rs +++ b/src/transport/middle_proxy/registry.rs @@ -18,8 +18,16 @@ const ROUTE_QUEUED_BYTE_PERMIT_UNIT: usize = 16 * 1024; const ROUTE_QUEUED_PERMITS_PER_SLOT: usize = 4; const ROUTE_QUEUED_MAX_FRAME_PERMITS: usize = 1024; +// Transactional writer registry publication. +mod publication; +// Cancellation-safe idle-writer replacement reservations. +mod replacement; mod writer; +pub(in crate::transport::middle_proxy) use replacement::{ + WriterBindOutcome, WriterReplacementReservation, +}; + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum RouteResult { Routed, @@ -72,6 +80,7 @@ struct WriterTable { struct WriterRoute { tx: mpsc::Sender, byte_budget: Arc, + replacement_state: Arc, } #[derive(Clone)] @@ -92,6 +101,15 @@ struct BindingState { last_meta_for_writer: DashMap, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[repr(u8)] +enum WriterReplacementState { + Open = 0, + Preparing, + Retiring, + Draining, +} + struct BindingInner { writer_for_conn: HashMap, conns_for_writer: HashMap>, diff --git a/src/transport/middle_proxy/registry/publication.rs b/src/transport/middle_proxy/registry/publication.rs new file mode 100644 index 0000000..d5870cf --- /dev/null +++ b/src/transport/middle_proxy/registry/publication.rs @@ -0,0 +1,123 @@ +use std::sync::Arc; +use std::sync::atomic::{AtomicU8, Ordering}; + +use tokio::sync::{MutexGuard, Semaphore, mpsc}; + +use super::super::codec::WriterCommand; +use super::replacement::WriterReplacementReservation; +use super::{BindingInner, ConnRegistry, WriterReplacementState, WriterRoute}; + +/// Holds registry binding ownership until pool writer visibility is published. +pub(in crate::transport::middle_proxy) struct WriterRegistrationGuard<'a> { + registry: &'a ConnRegistry, + binding: MutexGuard<'a, BindingInner>, +} + +impl ConnRegistry { + /// Acquires the only cancellation point required for writer registration. + pub(in crate::transport::middle_proxy) async fn prepare_writer_registration( + &self, + ) -> WriterRegistrationGuard<'_> { + WriterRegistrationGuard { + registry: self, + binding: self.binding.inner.lock().await, + } + } +} + +impl WriterRegistrationGuard<'_> { + /// Prevents new bindings while preserving all existing writer associations. + pub(in crate::transport::middle_proxy) fn retire(&mut self, writer_id: u64) -> bool { + let Some(state) = self + .registry + .writers + .map + .get(&writer_id) + .map(|route| Arc::clone(&route.replacement_state)) + else { + return false; + }; + loop { + let current = state.load(Ordering::Acquire); + if current == WriterReplacementState::Retiring as u8 + || current == WriterReplacementState::Draining as u8 + { + return true; + } + if state + .compare_exchange_weak( + current, + WriterReplacementState::Draining as u8, + Ordering::AcqRel, + Ordering::Acquire, + ) + .is_ok() + { + return true; + } + } + } + + /// Revalidates an idle replacement victim and prevents subsequent client bindings. + pub(in crate::transport::middle_proxy) fn prepare_replacement_commit( + &mut self, + reservation: &WriterReplacementReservation<'_>, + ) -> bool { + let Some(route_state) = self + .registry + .writers + .map + .get(&reservation.writer_id()) + .map(|route| Arc::clone(&route.replacement_state)) + else { + return false; + }; + if !std::ptr::eq(self.registry, reservation.registry()) + || !Arc::ptr_eq(&route_state, reservation.state()) + || reservation.requires_idle() + && self + .binding + .conns_for_writer + .get(&reservation.writer_id()) + .is_none_or(|conn_ids| !conn_ids.is_empty()) + { + return false; + } + route_state + .compare_exchange( + WriterReplacementState::Preparing as u8, + WriterReplacementState::Retiring as u8, + Ordering::AcqRel, + Ordering::Acquire, + ) + .is_ok() + } + + /// Installs registry state while retaining binding ownership for pool publication. + pub(in crate::transport::middle_proxy) fn install( + &mut self, + writer_id: u64, + tx: mpsc::Sender, + byte_budget: Arc, + ) { + self.binding.conns_for_writer.entry(writer_id).or_default(); + self.registry + .binding + .bound_clients_by_writer + .entry(writer_id) + .or_insert(0); + self.registry + .binding + .writer_idle_since_epoch_secs + .entry(writer_id) + .or_insert_with(ConnRegistry::now_epoch_secs); + self.registry.writers.map.insert( + writer_id, + WriterRoute { + tx, + byte_budget, + replacement_state: Arc::new(AtomicU8::new(WriterReplacementState::Open as u8)), + }, + ); + } +} diff --git a/src/transport/middle_proxy/registry/replacement.rs b/src/transport/middle_proxy/registry/replacement.rs new file mode 100644 index 0000000..c4b32f0 --- /dev/null +++ b/src/transport/middle_proxy/registry/replacement.rs @@ -0,0 +1,142 @@ +use std::sync::Arc; +use std::sync::atomic::{AtomicU8, Ordering}; + +use super::{ConnRegistry, WriterReplacementState}; + +/// Result of atomically committing a client-to-writer binding. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(in crate::transport::middle_proxy) enum WriterBindOutcome { + /// The client route was bound to the selected writer. + Bound, + /// The client route disappeared before the binding commit. + RouteMissing, + /// The selected writer disappeared before the binding commit. + WriterMissing, + /// The selected writer is retiring and no longer accepts new clients. + WriterRetiring, +} + +/// Cancellation-safe ownership of one prospective writer replacement. +pub(in crate::transport::middle_proxy) struct WriterReplacementReservation<'a> { + registry: &'a ConnRegistry, + writer_id: u64, + state: Arc, + require_idle: bool, + committed: bool, +} + +impl WriterReplacementReservation<'_> { + /// Returns the stable identifier of the prospective victim. + pub(in crate::transport::middle_proxy) fn writer_id(&self) -> u64 { + self.writer_id + } + + /// Returns the registry whose binding lock linearizes this reservation. + pub(super) fn registry(&self) -> &ConnRegistry { + self.registry + } + + /// Returns the writer-local replacement state identity captured at reservation time. + pub(super) fn state(&self) -> &Arc { + &self.state + } + + /// Reports whether commit must revalidate that the victim remains unbound. + pub(super) fn requires_idle(&self) -> bool { + self.require_idle + } + + /// Transfers retirement ownership to the published replacement. + pub(in crate::transport::middle_proxy) fn mark_committed(&mut self) { + self.committed = true; + } +} + +impl Drop for WriterReplacementReservation<'_> { + fn drop(&mut self) { + if self.committed { + return; + } + let preparing = WriterReplacementState::Preparing as u8; + let open = WriterReplacementState::Open as u8; + let _ = self + .state + .compare_exchange(preparing, open, Ordering::AcqRel, Ordering::Acquire); + } +} + +impl ConnRegistry { + /// Claims an idle writer as a prospective replacement victim. + /// + /// Preparing prevents duplicate replacement work but intentionally permits new client binds. + /// Commit revalidates idleness under the existing binding lock before blocking future binds. + pub(in crate::transport::middle_proxy) async fn try_reserve_writer_replacement( + &self, + writer_id: u64, + ) -> Option> { + self.reserve_writer_replacement(writer_id, true).await + } + + /// Claims a writer while retaining its existing client associations through replacement. + pub(in crate::transport::middle_proxy) async fn try_reserve_writer_replacement_preserving_clients( + &self, + writer_id: u64, + ) -> Option> { + self.reserve_writer_replacement(writer_id, false).await + } + + async fn reserve_writer_replacement( + &self, + writer_id: u64, + require_idle: bool, + ) -> Option> { + let binding = self.binding.inner.lock().await; + let state = self + .writers + .map + .get(&writer_id) + .map(|route| Arc::clone(&route.replacement_state))?; + if require_idle + && binding + .conns_for_writer + .get(&writer_id) + .is_none_or(|conn_ids| !conn_ids.is_empty()) + { + return None; + } + state + .compare_exchange( + WriterReplacementState::Open as u8, + WriterReplacementState::Preparing as u8, + Ordering::AcqRel, + Ordering::Acquire, + ) + .ok()?; + + Some(WriterReplacementReservation { + registry: self, + writer_id, + state, + require_idle, + committed: false, + }) + } + + /// Returns fixed-cardinality gauges for preparing and retiring replacements. + pub(in crate::transport::middle_proxy) fn writer_replacement_counts(&self) -> (usize, usize) { + let mut preparing = 0usize; + let mut retiring = 0usize; + for route in &self.writers.map { + match route.replacement_state.load(Ordering::Acquire) { + state if state == WriterReplacementState::Preparing as u8 => { + preparing = preparing.saturating_add(1); + } + state if state == WriterReplacementState::Retiring as u8 => { + retiring = retiring.saturating_add(1); + } + _ => {} + } + } + (preparing, retiring) + } +} diff --git a/src/transport/middle_proxy/registry/tests.rs b/src/transport/middle_proxy/registry/tests.rs index c90daf0..99e442f 100644 --- a/src/transport/middle_proxy/registry/tests.rs +++ b/src/transport/middle_proxy/registry/tests.rs @@ -4,13 +4,90 @@ use std::sync::Arc; use bytes::Bytes; use tokio::sync::Semaphore; -use super::{ConnMeta, ConnRegistry, RouteResult}; +use super::{ConnMeta, ConnRegistry, RouteResult, WriterBindOutcome}; use crate::transport::middle_proxy::MeResponse; fn writer_byte_budget() -> Arc { Arc::new(Semaphore::new(2049)) } +fn conn_meta() -> ConnMeta { + let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443); + ConnMeta { + target_dc: 2, + client_addr: addr, + our_addr: addr, + proto_flags: 0, + } +} + +#[tokio::test] +async fn preparing_replacement_allows_bind_and_idle_commit_revalidates() { + let registry = ConnRegistry::new(); + let (writer_tx, _writer_rx) = tokio::sync::mpsc::channel(8); + registry + .register_writer(10, writer_tx, writer_byte_budget()) + .await; + let reservation = registry + .try_reserve_writer_replacement(10) + .await + .expect("idle writer must be reservable"); + let (conn_id, _rx) = registry.register().await; + + assert_eq!( + registry + .bind_writer_with_outcome(conn_id, 10, conn_meta()) + .await, + WriterBindOutcome::Bound + ); + let mut registration = registry.prepare_writer_registration().await; + assert!(!registration.prepare_replacement_commit(&reservation)); + drop(registration); +} + +#[tokio::test] +async fn retiring_replacement_rejects_new_bind() { + let registry = ConnRegistry::new(); + let (writer_tx, _writer_rx) = tokio::sync::mpsc::channel(8); + registry + .register_writer(10, writer_tx, writer_byte_budget()) + .await; + let mut reservation = registry + .try_reserve_writer_replacement(10) + .await + .expect("idle writer must be reservable"); + let mut registration = registry.prepare_writer_registration().await; + assert!(registration.prepare_replacement_commit(&reservation)); + drop(registration); + reservation.mark_committed(); + let (conn_id, _rx) = registry.register().await; + + assert_eq!( + registry + .bind_writer_with_outcome(conn_id, 10, conn_meta()) + .await, + WriterBindOutcome::WriterRetiring + ); +} + +#[tokio::test] +async fn dropped_preparing_reservation_restores_writer_admission() { + let registry = ConnRegistry::new(); + let (writer_tx, _writer_rx) = tokio::sync::mpsc::channel(8); + registry + .register_writer(10, writer_tx, writer_byte_budget()) + .await; + let reservation = registry + .try_reserve_writer_replacement(10) + .await + .expect("idle writer must be reservable"); + assert!(registry.try_reserve_writer_replacement(10).await.is_none()); + + drop(reservation); + + assert!(registry.try_reserve_writer_replacement(10).await.is_some()); +} + #[tokio::test] async fn writer_activity_snapshot_tracks_writer_and_dc_load() { let registry = ConnRegistry::new(); diff --git a/src/transport/middle_proxy/registry/writer.rs b/src/transport/middle_proxy/registry/writer.rs index a3d375f..f242a68 100644 --- a/src/transport/middle_proxy/registry/writer.rs +++ b/src/transport/middle_proxy/registry/writer.rs @@ -8,6 +8,7 @@ use tokio::sync::mpsc::error::TrySendError; use super::super::codec::WriterCommand; use super::super::{MeResponse, RouteBytePermit}; +use super::replacement::WriterBindOutcome; use super::{ BoundConn, ConnMeta, ConnRegistry, ConnWriter, HotConnBinding, RouteResult, WriterActivitySnapshot, @@ -58,28 +59,16 @@ impl ConnRegistry { } /// Registers one writer command route and its matching memory budget atomically. + #[allow(dead_code)] pub async fn register_writer( &self, writer_id: u64, tx: mpsc::Sender, byte_budget: Arc, ) { - let mut binding = self.binding.inner.lock().await; - binding - .conns_for_writer - .entry(writer_id) - .or_insert_with(HashSet::new); - self.binding - .bound_clients_by_writer - .entry(writer_id) - .or_insert(0); - self.binding - .writer_idle_since_epoch_secs - .entry(writer_id) - .or_insert_with(Self::now_epoch_secs); - self.writers - .map - .insert(writer_id, super::WriterRoute { tx, byte_budget }); + self.prepare_writer_registration() + .await + .install(writer_id, tx, byte_budget); } /// Unregister connection, returning associated writer_id if any. @@ -298,17 +287,30 @@ impl ConnRegistry { } } - pub async fn bind_writer(&self, conn_id: u64, writer_id: u64, meta: ConnMeta) -> bool { + /// Atomically binds one client route while rejecting retiring writer generations. + pub(in crate::transport::middle_proxy) async fn bind_writer_with_outcome( + &self, + conn_id: u64, + writer_id: u64, + meta: ConnMeta, + ) -> WriterBindOutcome { let mut binding = self.binding.inner.lock().await; // ROUTING IS THE SOURCE OF TRUTH: // never keep/attach writer binding for a connection that is already // absent from the routing table. if !self.routing.map.contains_key(&conn_id) { - return false; + return WriterBindOutcome::RouteMissing; } - if !self.writers.map.contains_key(&writer_id) { - return false; + let Some(writer_route) = self.writers.map.get(&writer_id) else { + return WriterBindOutcome::WriterMissing; + }; + let writer_state = writer_route.replacement_state.load(Ordering::Acquire); + if writer_state == super::WriterReplacementState::Retiring as u8 + || writer_state == super::WriterReplacementState::Draining as u8 + { + return WriterBindOutcome::WriterRetiring; } + drop(writer_route); let previous_writer_id = binding.writer_for_conn.insert(conn_id, writer_id); if let Some(previous_writer_id) = previous_writer_id @@ -343,21 +345,15 @@ impl ConnRegistry { self.hot_binding .map .insert(conn_id, HotConnBinding { writer_id, meta }); - true + WriterBindOutcome::Bound } - pub async fn mark_writer_idle(&self, writer_id: u64) { - let mut binding = self.binding.inner.lock().await; - binding - .conns_for_writer - .entry(writer_id) - .or_insert_with(HashSet::new); - let count = binding - .conns_for_writer - .get(&writer_id) - .map(|set| set.len()) - .unwrap_or(0); - self.set_writer_bound_count(writer_id, count); + pub async fn bind_writer(&self, conn_id: u64, writer_id: u64, meta: ConnMeta) -> bool { + matches!( + self.bind_writer_with_outcome(conn_id, writer_id, meta) + .await, + WriterBindOutcome::Bound + ) } pub async fn get_last_writer_meta(&self, writer_id: u64) -> Option { diff --git a/src/transport/middle_proxy/secret.rs b/src/transport/middle_proxy/secret.rs index 6bae755..442ff50 100644 --- a/src/transport/middle_proxy/secret.rs +++ b/src/transport/middle_proxy/secret.rs @@ -1,4 +1,5 @@ use httpdate; +use std::path::PathBuf; use std::sync::Arc; use std::time::SystemTime; use tracing::{debug, info, warn}; @@ -7,6 +8,8 @@ use super::http_fetch::https_get; use super::selftest::record_timeskew_sample; use crate::error::{ProxyError, Result}; use crate::transport::UpstreamManager; +#[cfg(unix)] +use crate::util::secure_fs::{atomic_replace_async, read_regular_limited_async}; pub const PROXY_SECRET_MIN_LEN: usize = 32; @@ -58,7 +61,12 @@ pub async fn fetch_proxy_secret_with_upstream( match download_proxy_secret_with_max_len_via_upstream(max_len, upstream, proxy_secret_url).await { Ok(data) => { - if let Err(e) = tokio::fs::write(cache, &data).await { + #[cfg(unix)] + let cache_result = + atomic_replace_async(PathBuf::from(cache), data.clone(), 0o600).await; + #[cfg(not(unix))] + let cache_result = tokio::fs::write(cache, &data).await; + if let Err(e) = cache_result { warn!(error = %e, "Failed to cache proxy-secret (non-fatal)"); } else { debug!(path = cache, len = data.len(), "Cached proxy-secret"); @@ -72,7 +80,11 @@ pub async fn fetch_proxy_secret_with_upstream( } // 2) Fallback to cache/file regardless of age; require len in bounds. - match tokio::fs::read(cache).await { + #[cfg(unix)] + let cached = read_regular_limited_async(PathBuf::from(cache), max_len).await; + #[cfg(not(unix))] + let cached = tokio::fs::read(cache).await; + match cached { Ok(data) if validate_proxy_secret_len(data.len(), max_len).is_ok() => { let age_hours = tokio::fs::metadata(cache) .await diff --git a/src/transport/middle_proxy/send.rs b/src/transport/middle_proxy/send.rs index c74aa74..df5aeb4 100644 --- a/src/transport/middle_proxy/send.rs +++ b/src/transport/middle_proxy/send.rs @@ -1,26 +1,29 @@ #![allow(clippy::too_many_arguments)] -use std::cmp::Reverse; use std::net::SocketAddr; use std::sync::Arc; use std::sync::atomic::Ordering; use std::time::{Duration, Instant}; use tokio::sync::mpsc::error::TrySendError; -use tokio::sync::{OwnedSemaphorePermit, Semaphore, TryAcquireError, mpsc}; +use tokio::sync::{OwnedSemaphorePermit, TryAcquireError}; use tracing::{debug, warn}; use super::MePool; -use super::codec::{ProxyReqCommand, WriterBytePermit, WriterCommand}; -use super::registry::ConnMeta; -use super::wire::{build_proxy_req_payload, proxy_req_payload_len}; -use crate::config::defaults::ME_WRITER_BYTE_PERMIT_UNIT_BYTES; -use crate::config::{MeRouteNoWriterMode, MeWriterPickMode}; +use super::codec::WriterCommand; +use super::registry::{ConnMeta, WriterBindOutcome}; +use super::wire::build_proxy_req_payload; +use crate::config::MeRouteNoWriterMode; use crate::error::{ProxyError, Result}; -use crate::stats::Stats; -use crate::stream::PooledBuffer; use rand::seq::SliceRandom; +use self::bound::BoundWriterSendOutcome; +use self::reservation::{ + LEGACY_PROXY_REQ_SOURCE_CAPACITY_OVERHEAD_BYTES, WriterByteReserveError, + WriterCommandReserveError, proxy_req_resident_permits, reserve_writer_bytes, + reserve_writer_command_slot, try_reserve_writer_bytes, writer_send_deadline, +}; + const IDLE_WRITER_PENALTY_MID_SECS: u64 = 45; const IDLE_WRITER_PENALTY_HIGH_SECS: u64 = 55; const HYBRID_GLOBAL_BURST_PERIOD_ROUNDS: u32 = 4; @@ -31,142 +34,15 @@ const PICK_PENALTY_WARM: u64 = 200; const PICK_PENALTY_DRAINING: u64 = 600; const PICK_PENALTY_STALE: u64 = 300; const PICK_PENALTY_DEGRADED: u64 = 250; -const RPC_WRITER_FRAME_CAPACITY_OVERHEAD_BYTES: usize = 27; -const LEGACY_PROXY_REQ_SOURCE_CAPACITY_OVERHEAD_BYTES: usize = 128; +// Send-path submodules isolate delivery, close handling, recovery, reservations, and selection. +mod bound; mod close; +mod pooled; mod recovery; +mod reservation; mod selection; -enum WriterCommandReserveError { - Closed, - TimedOut, -} - -enum WriterByteReserveError { - Closed, - TimedOut, -} - -fn proxy_tag_array(tag: Option<&[u8]>) -> Option<[u8; 16]> { - tag.and_then(|tag| <[u8; 16]>::try_from(tag).ok()) -} - -fn proxy_req_payload_from_command( - cmd: WriterCommand, -) -> Option<(PooledBuffer, OwnedSemaphorePermit)> { - match cmd { - WriterCommand::ProxyReq(command) => Some((command.payload, command._permit)), - _ => None, - } -} - -fn payload_permit_from_data_command(cmd: WriterCommand) -> Option { - match cmd { - WriterCommand::Data { _permit, .. } => _permit, - _ => None, - } -} - -async fn reserve_writer_command_slot( - tx: &mpsc::Sender, - deadline: Option, -) -> std::result::Result, WriterCommandReserveError> { - let reserve = tx.clone().reserve_owned(); - match deadline { - Some(deadline) => { - match tokio::time::timeout(deadline.saturating_duration_since(Instant::now()), reserve) - .await - { - Ok(Ok(permit)) => Ok(permit), - Ok(Err(_)) => Err(WriterCommandReserveError::Closed), - Err(_) => Err(WriterCommandReserveError::TimedOut), - } - } - None => reserve.await.map_err(|_| WriterCommandReserveError::Closed), - } -} - -fn writer_send_deadline(wait: Option) -> Option { - wait.map(|wait| Instant::now() + wait) -} - -fn writer_resident_permits( - source_capacity: usize, - encoded_payload_len: usize, -) -> Option<(u32, usize)> { - let resident_bytes = source_capacity - .checked_add(encoded_payload_len)? - .checked_add(RPC_WRITER_FRAME_CAPACITY_OVERHEAD_BYTES)?; - let permits = resident_bytes.div_ceil(ME_WRITER_BYTE_PERMIT_UNIT_BYTES); - let permits = u32::try_from(permits).ok()?; - let reserved_bytes = (permits as usize).checked_mul(ME_WRITER_BYTE_PERMIT_UNIT_BYTES)?; - Some(( - permits.max(1), - reserved_bytes.max(ME_WRITER_BYTE_PERMIT_UNIT_BYTES), - )) -} - -fn proxy_req_resident_permits( - source_capacity: usize, - data_len: usize, - proxy_tag: Option<&[u8]>, - proto_flags: u32, -) -> Option<(u32, usize)> { - writer_resident_permits( - source_capacity, - proxy_req_payload_len(data_len, proxy_tag, proto_flags), - ) -} - -fn try_reserve_writer_bytes( - byte_budget: &Arc, - permits: u32, - reserved_bytes: usize, - stats: &Arc, -) -> std::result::Result { - byte_budget - .clone() - .try_acquire_many_owned(permits) - .map(|permit| WriterBytePermit::new(permit, reserved_bytes, stats.clone())) -} - -async fn reserve_writer_bytes( - byte_budget: &Arc, - permits: u32, - reserved_bytes: usize, - deadline: Option, - stats: &Arc, -) -> std::result::Result { - match try_reserve_writer_bytes(byte_budget, permits, reserved_bytes, stats) { - Ok(permit) => return Ok(permit), - Err(TryAcquireError::Closed) => return Err(WriterByteReserveError::Closed), - Err(TryAcquireError::NoPermits) => { - stats.increment_me_writer_byte_budget_wait_total(); - } - } - - let acquire = byte_budget.clone().acquire_many_owned(permits); - match deadline { - Some(deadline) => { - match tokio::time::timeout(deadline.saturating_duration_since(Instant::now()), acquire) - .await - { - Ok(Ok(permit)) => Ok(WriterBytePermit::new(permit, reserved_bytes, stats.clone())), - Ok(Err(_)) => Err(WriterByteReserveError::Closed), - Err(_) => { - stats.increment_me_writer_byte_budget_timeout_total(); - Err(WriterByteReserveError::TimedOut) - } - } - } - None => acquire - .await - .map(|permit| WriterBytePermit::new(permit, reserved_bytes, stats.clone())) - .map_err(|_| WriterByteReserveError::Closed), - } -} - impl MePool { /// Send RPC_PROXY_REQ. `tag_override`: per-user ad_tag (from access.user_ad_tags); if None, uses pool default. /// `payload_permit` keeps optional client byte accounting alive until the writer consumes the command. @@ -243,78 +119,22 @@ impl MePool { let mut hybrid_wait_current = hybrid_wait_step; loop { - if let Some((current, current_meta)) = self.registry.get_writer_with_meta(conn_id).await - { - let deadline = - writer_send_deadline(self.route_runtime.me_route_blocking_send_timeout); - let writer_permit = match reserve_writer_bytes( - ¤t.byte_budget, + match self + .try_send_bound_writer( + conn_id, + client_addr, + data, + proto_flags, + tag, writer_byte_permits, writer_reserved_bytes, - deadline, - &self.stats, + payload_permit, ) - .await - { - Ok(permit) => permit, - Err(WriterByteReserveError::TimedOut) => { - self.stats - .increment_me_writer_pick_full_total(self.writer_pick_mode()); - return Err(ProxyError::Proxy( - "ME writer byte budget full within blocking send timeout".into(), - )); - } - Err(WriterByteReserveError::Closed) => { - warn!( - writer_id = current.writer_id, - "ME writer byte budget closed" - ); - self.remove_writer_and_close_clients(current.writer_id) - .await; - continue; - } - }; - let (current_payload, _) = build_routed_payload(current_meta.our_addr); - let command = WriterCommand::Data { - payload: current_payload, - _permit: payload_permit.take(), - writer_permit, - }; - match current.tx.try_send(command) { - Ok(()) => { - self.note_hybrid_route_success(); - return Ok(()); - } - Err(TrySendError::Full(cmd)) => { - match reserve_writer_command_slot(¤t.tx, deadline).await { - Ok(permit) => { - permit.send(cmd); - self.note_hybrid_route_success(); - return Ok(()); - } - Err(WriterCommandReserveError::TimedOut) => { - self.stats - .increment_me_writer_pick_full_total(self.writer_pick_mode()); - return Err(ProxyError::Proxy( - "ME writer channel full within blocking send timeout".into(), - )); - } - Err(WriterCommandReserveError::Closed) => { - payload_permit = payload_permit_from_data_command(cmd); - } - } - warn!(writer_id = current.writer_id, "ME writer channel closed"); - self.remove_writer_and_close_clients(current.writer_id) - .await; - continue; - } - Err(TrySendError::Closed(cmd)) => { - payload_permit = payload_permit_from_data_command(cmd); - warn!(writer_id = current.writer_id, "ME writer channel closed"); - self.remove_writer_and_close_clients(current.writer_id) - .await; - continue; - } + .await? + { + BoundWriterSendOutcome::Sent => return Ok(()), + BoundWriterSendOutcome::Retry(retry_permit) => { + payload_permit = retry_permit; } } @@ -348,8 +168,9 @@ impl MePool { for _ in 0..self.route_runtime.me_route_inline_recovery_attempts.max(1) { - let preferred = self.preferred_endpoints_by_dc.load_full(); - for (dc, addrs) in preferred.iter() { + let endpoint_snapshot = self.endpoint_snapshot.load_full(); + for (dc, addrs) in &endpoint_snapshot.preferred_endpoints_by_dc + { for addr in addrs { let _ = self .connect_one_for_dc(*addr, *dc, self.rng.as_ref()) @@ -470,8 +291,9 @@ impl MePool { } emergency_attempts += 1; let mut endpoints = self - .preferred_endpoints_by_dc + .endpoint_snapshot .load() + .preferred_endpoints_by_dc .get(&routed_dc) .cloned() .unwrap_or_default(); @@ -534,88 +356,9 @@ impl MePool { } hybrid_wait_current = hybrid_wait_step; let pick_mode = self.writer_pick_mode(); - let pick_sample_size = self.writer_pick_sample_size(); - let writer_ids: Vec = candidate_indices - .iter() - .map(|idx| writers_snapshot[*idx].id) - .collect(); - let writer_idle_since = self - .registry - .writer_idle_since_for_writer_ids(&writer_ids) + let ordered_candidate_indices = self + .ordered_candidate_indices(candidate_indices, &writers_snapshot, pick_mode) .await; - let now_epoch_secs = Self::now_epoch_secs(); - let start = self.rr.fetch_add(1, Ordering::Relaxed) as usize % candidate_indices.len(); - let ordered_candidate_indices = if pick_mode == MeWriterPickMode::P2c { - self.p2c_ordered_candidate_indices( - &candidate_indices, - &writers_snapshot, - &writer_idle_since, - now_epoch_secs, - start, - pick_sample_size, - ) - } else { - if self - .writer_selection_policy - .me_deterministic_writer_sort - .load(Ordering::Relaxed) - { - candidate_indices.sort_by(|lhs, rhs| { - let left = &writers_snapshot[*lhs]; - let right = &writers_snapshot[*rhs]; - let left_key = ( - self.writer_contour_rank_for_selection(left), - (left.generation < self.current_generation()) as usize, - left.degraded.load(Ordering::Relaxed) as usize, - self.writer_idle_rank_for_selection( - left, - &writer_idle_since, - now_epoch_secs, - ), - Reverse(left.tx.capacity()), - left.addr, - left.id, - ); - let right_key = ( - self.writer_contour_rank_for_selection(right), - (right.generation < self.current_generation()) as usize, - right.degraded.load(Ordering::Relaxed) as usize, - self.writer_idle_rank_for_selection( - right, - &writer_idle_since, - now_epoch_secs, - ), - Reverse(right.tx.capacity()), - right.addr, - right.id, - ); - left_key.cmp(&right_key) - }); - } else { - candidate_indices.sort_by_key(|idx| { - let w = &writers_snapshot[*idx]; - let degraded = w.degraded.load(Ordering::Relaxed); - let stale = (w.generation < self.current_generation()) as usize; - ( - self.writer_contour_rank_for_selection(w), - stale, - degraded as usize, - self.writer_idle_rank_for_selection( - w, - &writer_idle_since, - now_epoch_secs, - ), - Reverse(w.tx.capacity()), - ) - }); - } - - let mut ordered = Vec::::with_capacity(candidate_indices.len()); - for offset in 0..candidate_indices.len() { - ordered.push(candidate_indices[(start + offset) % candidate_indices.len()]); - } - ordered - }; let mut fallback_blocking_idx: Option = None; for idx in ordered_candidate_indices { @@ -648,13 +391,30 @@ impl MePool { // Keep the advertised proxy IP aligned with the selected ME writer source. let effective_our_addr = SocketAddr::new(w.source_ip, our_addr.port()); let (payload, meta) = build_routed_payload(effective_our_addr); - if !self.registry.bind_writer(conn_id, w.id, meta).await { + let bind_outcome = self + .registry + .bind_writer_with_outcome(conn_id, w.id, meta) + .await; + if bind_outcome != WriterBindOutcome::Bound { + drop(permit); + if bind_outcome == WriterBindOutcome::WriterRetiring { + debug!( + conn_id, + writer_id = w.id, + "ME writer entered replacement retirement before bind commit" + ); + continue; + } + if bind_outcome == WriterBindOutcome::RouteMissing { + return Err(ProxyError::Proxy( + "ME client route disappeared before writer bind".into(), + )); + } debug!( conn_id, writer_id = w.id, "ME writer disappeared before bind commit, pruning stale writer" ); - drop(permit); self.remove_writer_and_close_clients(w.id).await; continue; } @@ -742,13 +502,30 @@ impl MePool { // Keep the advertised proxy IP aligned with the selected ME writer source. let effective_our_addr = SocketAddr::new(w.source_ip, our_addr.port()); let (payload, meta) = build_routed_payload(effective_our_addr); - if !self.registry.bind_writer(conn_id, w.id, meta).await { + let bind_outcome = self + .registry + .bind_writer_with_outcome(conn_id, w.id, meta) + .await; + if bind_outcome != WriterBindOutcome::Bound { + drop(permit); + if bind_outcome == WriterBindOutcome::WriterRetiring { + debug!( + conn_id, + writer_id = w.id, + "ME writer entered replacement retirement before fallback bind commit" + ); + continue; + } + if bind_outcome == WriterBindOutcome::RouteMissing { + return Err(ProxyError::Proxy( + "ME client route disappeared before writer bind".into(), + )); + } debug!( conn_id, writer_id = w.id, "ME writer disappeared before fallback bind commit, pruning stale writer" ); - drop(permit); self.remove_writer_and_close_clients(w.id).await; continue; } @@ -766,168 +543,4 @@ impl MePool { return Ok(()); } } - - /// Send RPC_PROXY_REQ while keeping the first bound-writer path allocation-light. - /// The client byte permit follows the payload until writer completion or command drop. - pub async fn send_proxy_req_pooled( - self: &Arc, - conn_id: u64, - target_dc: i16, - client_addr: SocketAddr, - our_addr: SocketAddr, - payload: PooledBuffer, - _permit: OwnedSemaphorePermit, - proto_flags: u32, - tag_override: Option<[u8; 16]>, - ) -> Result<()> { - let tag = tag_override.or_else(|| proxy_tag_array(self.proxy_tag.as_deref())); - let Some((writer_byte_permits, writer_reserved_bytes)) = proxy_req_resident_permits( - payload.capacity(), - payload.len(), - tag.as_ref().map(|tag| tag.as_slice()), - proto_flags, - ) else { - self.stats.increment_me_writer_byte_budget_oversize_total(); - return Err(ProxyError::Proxy( - "ME writer payload residency calculation overflow".into(), - )); - }; - if writer_byte_permits as usize > self.writer_lifecycle.writer_byte_budget_permits { - self.stats.increment_me_writer_byte_budget_oversize_total(); - return Err(ProxyError::Proxy( - "ME writer payload exceeds configured byte budget".into(), - )); - } - - if let Some((current, current_meta)) = self.registry.get_writer_with_meta(conn_id).await { - let deadline = writer_send_deadline(self.route_runtime.me_route_blocking_send_timeout); - let writer_permit = match reserve_writer_bytes( - ¤t.byte_budget, - writer_byte_permits, - writer_reserved_bytes, - deadline, - &self.stats, - ) - .await - { - Ok(permit) => permit, - Err(WriterByteReserveError::TimedOut) => { - self.stats - .increment_me_writer_pick_full_total(self.writer_pick_mode()); - return Err(ProxyError::Proxy( - "ME writer byte budget full within blocking send timeout".into(), - )); - } - Err(WriterByteReserveError::Closed) => { - warn!( - writer_id = current.writer_id, - "ME writer byte budget closed" - ); - self.remove_writer_and_close_clients(current.writer_id) - .await; - return self - .send_proxy_req( - conn_id, - target_dc, - client_addr, - our_addr, - payload.as_ref(), - proto_flags, - tag.as_ref().map(|tag| tag.as_slice()), - Some(_permit), - ) - .await; - } - }; - let command = WriterCommand::ProxyReq(ProxyReqCommand { - conn_id, - client_addr, - our_addr: current_meta.our_addr, - proto_flags, - proxy_tag: tag, - payload, - _permit, - writer_permit, - }); - match current.tx.try_send(command) { - Ok(()) => { - self.note_hybrid_route_success(); - return Ok(()); - } - Err(TrySendError::Full(cmd)) => { - match reserve_writer_command_slot(¤t.tx, deadline).await { - Ok(permit) => { - permit.send(cmd); - self.note_hybrid_route_success(); - return Ok(()); - } - Err(WriterCommandReserveError::TimedOut) => { - self.stats - .increment_me_writer_pick_full_total(self.writer_pick_mode()); - return Err(ProxyError::Proxy( - "ME writer channel full within blocking send timeout".into(), - )); - } - Err(WriterCommandReserveError::Closed) => { - let Some((payload, _permit)) = proxy_req_payload_from_command(cmd) - else { - return Err(ProxyError::Proxy( - "ME writer rejected unexpected command type".into(), - )); - }; - warn!(writer_id = current.writer_id, "ME writer channel closed"); - self.remove_writer_and_close_clients(current.writer_id) - .await; - return self - .send_proxy_req( - conn_id, - target_dc, - client_addr, - our_addr, - payload.as_ref(), - proto_flags, - tag.as_ref().map(|tag| tag.as_slice()), - Some(_permit), - ) - .await; - } - } - } - Err(TrySendError::Closed(cmd)) => { - let Some((payload, _permit)) = proxy_req_payload_from_command(cmd) else { - return Err(ProxyError::Proxy( - "ME writer rejected unexpected command type".into(), - )); - }; - warn!(writer_id = current.writer_id, "ME writer channel closed"); - self.remove_writer_and_close_clients(current.writer_id) - .await; - return self - .send_proxy_req( - conn_id, - target_dc, - client_addr, - our_addr, - payload.as_ref(), - proto_flags, - tag.as_ref().map(|tag| tag.as_slice()), - Some(_permit), - ) - .await; - } - } - } - - self.send_proxy_req( - conn_id, - target_dc, - client_addr, - our_addr, - payload.as_ref(), - proto_flags, - tag.as_ref().map(|tag| tag.as_slice()), - Some(_permit), - ) - .await - } } diff --git a/src/transport/middle_proxy/send/bound.rs b/src/transport/middle_proxy/send/bound.rs new file mode 100644 index 0000000..39fd105 --- /dev/null +++ b/src/transport/middle_proxy/send/bound.rs @@ -0,0 +1,115 @@ +use std::net::SocketAddr; +use std::sync::Arc; + +use tokio::sync::OwnedSemaphorePermit; +use tokio::sync::mpsc::error::TrySendError; +use tracing::warn; + +use super::super::MePool; +use super::super::codec::WriterCommand; +use super::super::wire::build_proxy_req_payload; +use super::reservation::{ + WriterByteReserveError, WriterCommandReserveError, payload_permit_from_data_command, + reserve_writer_bytes, reserve_writer_command_slot, writer_send_deadline, +}; +use crate::error::{ProxyError, Result}; + +pub(super) enum BoundWriterSendOutcome { + Sent, + Retry(Option), +} + +impl MePool { + pub(super) async fn try_send_bound_writer( + self: &Arc, + conn_id: u64, + client_addr: SocketAddr, + data: &[u8], + proto_flags: u32, + tag: Option<&[u8]>, + writer_byte_permits: u32, + writer_reserved_bytes: usize, + payload_permit: Option, + ) -> Result { + let Some((current, current_meta)) = self.registry.get_writer_with_meta(conn_id).await + else { + return Ok(BoundWriterSendOutcome::Retry(payload_permit)); + }; + let deadline = writer_send_deadline(self.route_runtime.me_route_blocking_send_timeout); + let writer_permit = match reserve_writer_bytes( + ¤t.byte_budget, + writer_byte_permits, + writer_reserved_bytes, + deadline, + &self.stats, + ) + .await + { + Ok(permit) => permit, + Err(WriterByteReserveError::TimedOut) => { + self.stats + .increment_me_writer_pick_full_total(self.writer_pick_mode()); + return Err(ProxyError::Proxy( + "ME writer byte budget full within blocking send timeout".into(), + )); + } + Err(WriterByteReserveError::Closed) => { + warn!( + writer_id = current.writer_id, + "ME writer byte budget closed" + ); + self.remove_writer_and_close_clients(current.writer_id) + .await; + return Ok(BoundWriterSendOutcome::Retry(payload_permit)); + } + }; + let payload = build_proxy_req_payload( + conn_id, + client_addr, + current_meta.our_addr, + data, + tag, + proto_flags, + ); + let command = WriterCommand::Data { + payload, + _permit: payload_permit, + writer_permit, + }; + match current.tx.try_send(command) { + Ok(()) => { + self.note_hybrid_route_success(); + Ok(BoundWriterSendOutcome::Sent) + } + Err(TrySendError::Full(cmd)) => { + match reserve_writer_command_slot(¤t.tx, deadline).await { + Ok(permit) => { + permit.send(cmd); + self.note_hybrid_route_success(); + return Ok(BoundWriterSendOutcome::Sent); + } + Err(WriterCommandReserveError::TimedOut) => { + self.stats + .increment_me_writer_pick_full_total(self.writer_pick_mode()); + return Err(ProxyError::Proxy( + "ME writer channel full within blocking send timeout".into(), + )); + } + Err(WriterCommandReserveError::Closed) => {} + } + let payload_permit = payload_permit_from_data_command(cmd); + warn!(writer_id = current.writer_id, "ME writer channel closed"); + self.remove_writer_and_close_clients(current.writer_id) + .await; + Ok(BoundWriterSendOutcome::Retry(payload_permit)) + } + Err(TrySendError::Closed(cmd)) => { + let payload_permit = payload_permit_from_data_command(cmd); + warn!(writer_id = current.writer_id, "ME writer channel closed"); + self.remove_writer_and_close_clients(current.writer_id) + .await; + Ok(BoundWriterSendOutcome::Retry(payload_permit)) + } + } + } +} diff --git a/src/transport/middle_proxy/send/pooled.rs b/src/transport/middle_proxy/send/pooled.rs new file mode 100644 index 0000000..87a84b0 --- /dev/null +++ b/src/transport/middle_proxy/send/pooled.rs @@ -0,0 +1,182 @@ +use std::net::SocketAddr; +use std::sync::Arc; + +use tokio::sync::OwnedSemaphorePermit; +use tokio::sync::mpsc::error::TrySendError; +use tracing::warn; + +use super::super::MePool; +use super::super::codec::{ProxyReqCommand, WriterCommand}; +use super::reservation::{ + WriterByteReserveError, WriterCommandReserveError, proxy_req_payload_from_command, + proxy_req_resident_permits, proxy_tag_array, reserve_writer_bytes, reserve_writer_command_slot, + writer_send_deadline, +}; +use crate::error::{ProxyError, Result}; +use crate::stream::PooledBuffer; + +impl MePool { + /// Send RPC_PROXY_REQ while keeping the first bound-writer path allocation-light. + /// The client byte permit follows the payload until writer completion or command drop. + pub async fn send_proxy_req_pooled( + self: &Arc, + conn_id: u64, + target_dc: i16, + client_addr: SocketAddr, + our_addr: SocketAddr, + payload: PooledBuffer, + _permit: OwnedSemaphorePermit, + proto_flags: u32, + tag_override: Option<[u8; 16]>, + ) -> Result<()> { + let tag = tag_override.or_else(|| proxy_tag_array(self.proxy_tag.as_deref())); + let Some((writer_byte_permits, writer_reserved_bytes)) = proxy_req_resident_permits( + payload.capacity(), + payload.len(), + tag.as_ref().map(|tag| tag.as_slice()), + proto_flags, + ) else { + self.stats.increment_me_writer_byte_budget_oversize_total(); + return Err(ProxyError::Proxy( + "ME writer payload residency calculation overflow".into(), + )); + }; + if writer_byte_permits as usize > self.writer_lifecycle.writer_byte_budget_permits { + self.stats.increment_me_writer_byte_budget_oversize_total(); + return Err(ProxyError::Proxy( + "ME writer payload exceeds configured byte budget".into(), + )); + } + + if let Some((current, current_meta)) = self.registry.get_writer_with_meta(conn_id).await { + let deadline = writer_send_deadline(self.route_runtime.me_route_blocking_send_timeout); + let writer_permit = match reserve_writer_bytes( + ¤t.byte_budget, + writer_byte_permits, + writer_reserved_bytes, + deadline, + &self.stats, + ) + .await + { + Ok(permit) => permit, + Err(WriterByteReserveError::TimedOut) => { + self.stats + .increment_me_writer_pick_full_total(self.writer_pick_mode()); + return Err(ProxyError::Proxy( + "ME writer byte budget full within blocking send timeout".into(), + )); + } + Err(WriterByteReserveError::Closed) => { + warn!( + writer_id = current.writer_id, + "ME writer byte budget closed" + ); + self.remove_writer_and_close_clients(current.writer_id) + .await; + return self + .send_proxy_req( + conn_id, + target_dc, + client_addr, + our_addr, + payload.as_ref(), + proto_flags, + tag.as_ref().map(|tag| tag.as_slice()), + Some(_permit), + ) + .await; + } + }; + let command = WriterCommand::ProxyReq(ProxyReqCommand { + conn_id, + client_addr, + our_addr: current_meta.our_addr, + proto_flags, + proxy_tag: tag, + payload, + _permit, + writer_permit, + }); + match current.tx.try_send(command) { + Ok(()) => { + self.note_hybrid_route_success(); + return Ok(()); + } + Err(TrySendError::Full(cmd)) => { + match reserve_writer_command_slot(¤t.tx, deadline).await { + Ok(permit) => { + permit.send(cmd); + self.note_hybrid_route_success(); + return Ok(()); + } + Err(WriterCommandReserveError::TimedOut) => { + self.stats + .increment_me_writer_pick_full_total(self.writer_pick_mode()); + return Err(ProxyError::Proxy( + "ME writer channel full within blocking send timeout".into(), + )); + } + Err(WriterCommandReserveError::Closed) => { + let Some((payload, _permit)) = proxy_req_payload_from_command(cmd) + else { + return Err(ProxyError::Proxy( + "ME writer rejected unexpected command type".into(), + )); + }; + warn!(writer_id = current.writer_id, "ME writer channel closed"); + self.remove_writer_and_close_clients(current.writer_id) + .await; + return self + .send_proxy_req( + conn_id, + target_dc, + client_addr, + our_addr, + payload.as_ref(), + proto_flags, + tag.as_ref().map(|tag| tag.as_slice()), + Some(_permit), + ) + .await; + } + } + } + Err(TrySendError::Closed(cmd)) => { + let Some((payload, _permit)) = proxy_req_payload_from_command(cmd) else { + return Err(ProxyError::Proxy( + "ME writer rejected unexpected command type".into(), + )); + }; + warn!(writer_id = current.writer_id, "ME writer channel closed"); + self.remove_writer_and_close_clients(current.writer_id) + .await; + return self + .send_proxy_req( + conn_id, + target_dc, + client_addr, + our_addr, + payload.as_ref(), + proto_flags, + tag.as_ref().map(|tag| tag.as_slice()), + Some(_permit), + ) + .await; + } + } + } + + self.send_proxy_req( + conn_id, + target_dc, + client_addr, + our_addr, + payload.as_ref(), + proto_flags, + tag.as_ref().map(|tag| tag.as_slice()), + Some(_permit), + ) + .await + } +} diff --git a/src/transport/middle_proxy/send/recovery.rs b/src/transport/middle_proxy/send/recovery.rs index ab38b5e..da95a3c 100644 --- a/src/transport/middle_proxy/send/recovery.rs +++ b/src/transport/middle_proxy/send/recovery.rs @@ -88,7 +88,8 @@ impl MePool { pub(super) async fn trigger_async_recovery_global(self: &Arc) { self.stats.increment_me_async_recovery_trigger_total(); - let preferred = self.preferred_endpoints_by_dc.load(); + let endpoint_snapshot = self.endpoint_snapshot.load(); + let preferred = &endpoint_snapshot.preferred_endpoints_by_dc; let mut triggered = 0usize; for (dc, addrs) in preferred.iter() { for addr in addrs { diff --git a/src/transport/middle_proxy/send/reservation.rs b/src/transport/middle_proxy/send/reservation.rs new file mode 100644 index 0000000..6b9c6ff --- /dev/null +++ b/src/transport/middle_proxy/send/reservation.rs @@ -0,0 +1,141 @@ +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use tokio::sync::{OwnedSemaphorePermit, Semaphore, TryAcquireError, mpsc}; + +use super::super::codec::{WriterBytePermit, WriterCommand}; +use crate::config::defaults::ME_WRITER_BYTE_PERMIT_UNIT_BYTES; +use crate::stats::Stats; +use crate::stream::PooledBuffer; + +const RPC_WRITER_FRAME_CAPACITY_OVERHEAD_BYTES: usize = 27; +pub(super) const LEGACY_PROXY_REQ_SOURCE_CAPACITY_OVERHEAD_BYTES: usize = 128; + +pub(super) enum WriterCommandReserveError { + Closed, + TimedOut, +} + +pub(super) enum WriterByteReserveError { + Closed, + TimedOut, +} + +pub(super) fn proxy_tag_array(tag: Option<&[u8]>) -> Option<[u8; 16]> { + tag.and_then(|tag| <[u8; 16]>::try_from(tag).ok()) +} + +pub(super) fn proxy_req_payload_from_command( + cmd: WriterCommand, +) -> Option<(PooledBuffer, OwnedSemaphorePermit)> { + match cmd { + WriterCommand::ProxyReq(command) => Some((command.payload, command._permit)), + _ => None, + } +} + +pub(super) fn payload_permit_from_data_command(cmd: WriterCommand) -> Option { + match cmd { + WriterCommand::Data { _permit, .. } => _permit, + _ => None, + } +} + +pub(super) async fn reserve_writer_command_slot( + tx: &mpsc::Sender, + deadline: Option, +) -> std::result::Result, WriterCommandReserveError> { + let reserve = tx.clone().reserve_owned(); + match deadline { + Some(deadline) => { + match tokio::time::timeout(deadline.saturating_duration_since(Instant::now()), reserve) + .await + { + Ok(Ok(permit)) => Ok(permit), + Ok(Err(_)) => Err(WriterCommandReserveError::Closed), + Err(_) => Err(WriterCommandReserveError::TimedOut), + } + } + None => reserve.await.map_err(|_| WriterCommandReserveError::Closed), + } +} + +pub(super) fn writer_send_deadline(wait: Option) -> Option { + wait.map(|wait| Instant::now() + wait) +} + +fn writer_resident_permits( + source_capacity: usize, + encoded_payload_len: usize, +) -> Option<(u32, usize)> { + let resident_bytes = source_capacity + .checked_add(encoded_payload_len)? + .checked_add(RPC_WRITER_FRAME_CAPACITY_OVERHEAD_BYTES)?; + let permits = resident_bytes.div_ceil(ME_WRITER_BYTE_PERMIT_UNIT_BYTES); + let permits = u32::try_from(permits).ok()?; + let reserved_bytes = (permits as usize).checked_mul(ME_WRITER_BYTE_PERMIT_UNIT_BYTES)?; + Some(( + permits.max(1), + reserved_bytes.max(ME_WRITER_BYTE_PERMIT_UNIT_BYTES), + )) +} + +pub(super) fn proxy_req_resident_permits( + source_capacity: usize, + data_len: usize, + proxy_tag: Option<&[u8]>, + proto_flags: u32, +) -> Option<(u32, usize)> { + writer_resident_permits( + source_capacity, + super::super::wire::proxy_req_payload_len(data_len, proxy_tag, proto_flags), + ) +} + +pub(super) fn try_reserve_writer_bytes( + byte_budget: &Arc, + permits: u32, + reserved_bytes: usize, + stats: &Arc, +) -> std::result::Result { + byte_budget + .clone() + .try_acquire_many_owned(permits) + .map(|permit| WriterBytePermit::new(permit, reserved_bytes, stats.clone())) +} + +pub(super) async fn reserve_writer_bytes( + byte_budget: &Arc, + permits: u32, + reserved_bytes: usize, + deadline: Option, + stats: &Arc, +) -> std::result::Result { + match try_reserve_writer_bytes(byte_budget, permits, reserved_bytes, stats) { + Ok(permit) => return Ok(permit), + Err(TryAcquireError::Closed) => return Err(WriterByteReserveError::Closed), + Err(TryAcquireError::NoPermits) => { + stats.increment_me_writer_byte_budget_wait_total(); + } + } + + let acquire = byte_budget.clone().acquire_many_owned(permits); + match deadline { + Some(deadline) => { + match tokio::time::timeout(deadline.saturating_duration_since(Instant::now()), acquire) + .await + { + Ok(Ok(permit)) => Ok(WriterBytePermit::new(permit, reserved_bytes, stats.clone())), + Ok(Err(_)) => Err(WriterByteReserveError::Closed), + Err(_) => { + stats.increment_me_writer_byte_budget_timeout_total(); + Err(WriterByteReserveError::TimedOut) + } + } + } + None => acquire + .await + .map(|permit| WriterBytePermit::new(permit, reserved_bytes, stats.clone())) + .map_err(|_| WriterByteReserveError::Closed), + } +} diff --git a/src/transport/middle_proxy/send/selection.rs b/src/transport/middle_proxy/send/selection.rs index 837062c..b57bbfc 100644 --- a/src/transport/middle_proxy/send/selection.rs +++ b/src/transport/middle_proxy/send/selection.rs @@ -1,4 +1,6 @@ -use std::collections::{HashMap, HashSet}; +use std::cmp::Reverse; +use std::collections::HashMap; +use std::net::SocketAddr; use std::sync::atomic::Ordering; use super::super::MePool; @@ -7,6 +9,19 @@ use super::{ IDLE_WRITER_PENALTY_HIGH_SECS, IDLE_WRITER_PENALTY_MID_SECS, PICK_PENALTY_DEGRADED, PICK_PENALTY_DRAINING, PICK_PENALTY_STALE, PICK_PENALTY_WARM, }; +use crate::config::MeWriterPickMode; + +struct WriterSelectionKey { + index: usize, + contour_rank: usize, + stale: usize, + degraded: usize, + idle_rank: usize, + queue_remaining: usize, + addr: SocketAddr, + id: u64, + pick_score: u64, +} impl MePool { pub(super) async fn candidate_indices_for_dc( @@ -15,7 +30,8 @@ impl MePool { routed_dc: i32, include_warm: bool, ) -> Vec { - let preferred_snapshot = self.preferred_endpoints_by_dc.load(); + let endpoint_snapshot = self.endpoint_snapshot.load(); + let preferred_snapshot = &endpoint_snapshot.preferred_endpoints_by_dc; let mut out = Vec::new(); if let Some(preferred) = preferred_snapshot .get(&routed_dc) @@ -69,24 +85,20 @@ impl MePool { } } - pub(super) fn writer_contour_rank_for_selection( - &self, - writer: &super::super::pool::MeWriter, - ) -> usize { - match WriterContour::from_u8(writer.contour.load(Ordering::Relaxed)) { + fn writer_contour_rank_for_selection(contour: WriterContour) -> usize { + match contour { WriterContour::Active => 0, WriterContour::Warm => 1, WriterContour::Draining => 2, } } - pub(super) fn writer_idle_rank_for_selection( - &self, - writer: &super::super::pool::MeWriter, + fn writer_idle_rank_for_selection( + writer_id: u64, idle_since_by_writer: &HashMap, now_epoch_secs: u64, ) -> usize { - let Some(idle_since) = idle_since_by_writer.get(&writer.id).copied() else { + let Some(idle_since) = idle_since_by_writer.get(&writer_id).copied() else { return 0; }; let idle_age_secs = now_epoch_secs.saturating_sub(idle_since); @@ -99,51 +111,63 @@ impl MePool { } } - pub(super) fn writer_pick_score( + fn capture_writer_selection_key( &self, + index: usize, writer: &super::super::pool::MeWriter, idle_since_by_writer: &HashMap, now_epoch_secs: u64, - ) -> u64 { - let contour_penalty = match WriterContour::from_u8(writer.contour.load(Ordering::Relaxed)) { + current_generation: u64, + ) -> WriterSelectionKey { + let contour = WriterContour::from_u8(writer.contour.load(Ordering::Relaxed)); + let contour_rank = Self::writer_contour_rank_for_selection(contour); + let contour_penalty = match contour { WriterContour::Active => 0, WriterContour::Warm => PICK_PENALTY_WARM, WriterContour::Draining => PICK_PENALTY_DRAINING, }; - let stale_penalty = if writer.generation < self.current_generation() { - PICK_PENALTY_STALE - } else { - 0 - }; - let degraded_penalty = if writer.degraded.load(Ordering::Relaxed) { + let stale = (writer.generation < current_generation) as usize; + let stale_penalty = if stale != 0 { PICK_PENALTY_STALE } else { 0 }; + let degraded = writer.degraded.load(Ordering::Relaxed) as usize; + let degraded_penalty = if degraded != 0 { PICK_PENALTY_DEGRADED } else { 0 }; - let idle_penalty = - (self.writer_idle_rank_for_selection(writer, idle_since_by_writer, now_epoch_secs) - as u64) - * 100; + let idle_rank = + Self::writer_idle_rank_for_selection(writer.id, idle_since_by_writer, now_epoch_secs); + let idle_penalty = (idle_rank as u64) * 100; let queue_cap = self.writer_lifecycle.writer_cmd_channel_capacity.max(1) as u64; - let queue_remaining = writer.tx.capacity() as u64; - let queue_used = queue_cap.saturating_sub(queue_remaining.min(queue_cap)); + let queue_remaining = writer.tx.capacity(); + let queue_used = queue_cap.saturating_sub((queue_remaining as u64).min(queue_cap)); let queue_util_pct = queue_used.saturating_mul(100) / queue_cap; let queue_penalty = queue_util_pct.saturating_mul(4); let rtt_penalty = ((writer.rtt_ema_ms_x10.load(Ordering::Relaxed) as u64).saturating_add(5) / 10) .min(400); - contour_penalty + let pick_score = contour_penalty .saturating_add(stale_penalty) .saturating_add(degraded_penalty) .saturating_add(idle_penalty) .saturating_add(queue_penalty) - .saturating_add(rtt_penalty) + .saturating_add(rtt_penalty); + WriterSelectionKey { + index, + contour_rank, + stale, + degraded, + idle_rank, + queue_remaining, + addr: writer.addr, + id: writer.id, + pick_score, + } } pub(super) fn p2c_ordered_candidate_indices( &self, - candidate_indices: &[usize], + mut candidate_indices: Vec, writers_snapshot: &[super::super::pool::MeWriter], idle_since_by_writer: &HashMap, now_epoch_secs: u64, @@ -155,32 +179,120 @@ impl MePool { return Vec::new(); } - let mut sampled = Vec::::with_capacity(sample_size.min(total)); - let mut seen = HashSet::::with_capacity(total); - for offset in 0..sample_size.min(total) { - let idx = candidate_indices[(start + offset) % total]; - if seen.insert(idx) { - sampled.push(idx); + candidate_indices.rotate_left(start % total); + let current_generation = self.current_generation(); + let sample_size = sample_size.min(total); + let mut sampled = candidate_indices[..sample_size] + .iter() + .map(|idx| { + self.capture_writer_selection_key( + *idx, + &writers_snapshot[*idx], + idle_since_by_writer, + now_epoch_secs, + current_generation, + ) + }) + .collect::>(); + sampled.sort_by_key(|candidate| (candidate.pick_score, candidate.addr, candidate.id)); + for (target, candidate) in candidate_indices[..sample_size].iter_mut().zip(sampled) { + *target = candidate.index; + } + candidate_indices + } + + pub(super) async fn ordered_candidate_indices( + &self, + mut candidate_indices: Vec, + writers_snapshot: &[super::super::pool::MeWriter], + pick_mode: MeWriterPickMode, + ) -> Vec { + let pick_sample_size = self.writer_pick_sample_size(); + let writer_ids: Vec = candidate_indices + .iter() + .map(|idx| writers_snapshot[*idx].id) + .collect(); + let writer_idle_since = self + .registry + .writer_idle_since_for_writer_ids(&writer_ids) + .await; + let now_epoch_secs = Self::now_epoch_secs(); + let start = self.rr.fetch_add(1, Ordering::Relaxed) as usize % candidate_indices.len(); + if pick_mode == MeWriterPickMode::P2c { + return self.p2c_ordered_candidate_indices( + candidate_indices, + writers_snapshot, + &writer_idle_since, + now_epoch_secs, + start, + pick_sample_size, + ); + } + + if self + .writer_selection_policy + .me_deterministic_writer_sort + .load(Ordering::Relaxed) + { + let current_generation = self.current_generation(); + let mut captured = candidate_indices + .iter() + .map(|idx| { + self.capture_writer_selection_key( + *idx, + &writers_snapshot[*idx], + &writer_idle_since, + now_epoch_secs, + current_generation, + ) + }) + .collect::>(); + captured.sort_by_key(|candidate| { + ( + candidate.contour_rank, + candidate.stale, + candidate.degraded, + candidate.idle_rank, + Reverse(candidate.queue_remaining), + candidate.addr, + candidate.id, + ) + }); + for (target, candidate) in candidate_indices.iter_mut().zip(captured) { + *target = candidate.index; + } + } else { + let current_generation = self.current_generation(); + let mut captured = candidate_indices + .iter() + .map(|idx| { + self.capture_writer_selection_key( + *idx, + &writers_snapshot[*idx], + &writer_idle_since, + now_epoch_secs, + current_generation, + ) + }) + .collect::>(); + captured.sort_by_key(|candidate| { + ( + candidate.contour_rank, + candidate.stale, + candidate.degraded, + candidate.idle_rank, + Reverse(candidate.queue_remaining), + ) + }); + for (target, candidate) in candidate_indices.iter_mut().zip(captured) { + *target = candidate.index; } } - sampled.sort_by_key(|idx| { - let writer = &writers_snapshot[*idx]; - ( - self.writer_pick_score(writer, idle_since_by_writer, now_epoch_secs), - writer.addr, - writer.id, - ) - }); - - let mut ordered = Vec::::with_capacity(total); - ordered.extend(sampled.iter().copied()); - for offset in 0..total { - let idx = candidate_indices[(start + offset) % total]; - if seen.insert(idx) { - ordered.push(idx); - } + if !candidate_indices.is_empty() { + let len = candidate_indices.len(); + candidate_indices.rotate_left(start % len); } - ordered + candidate_indices } } diff --git a/src/transport/middle_proxy/tests/pool_refill_security_tests.rs b/src/transport/middle_proxy/tests/pool_refill_security_tests.rs index a513d1a..b591214 100644 --- a/src/transport/middle_proxy/tests/pool_refill_security_tests.rs +++ b/src/transport/middle_proxy/tests/pool_refill_security_tests.rs @@ -1,15 +1,17 @@ use std::collections::HashMap; -use std::net::{IpAddr, Ipv4Addr, SocketAddr}; +use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr}; use std::sync::Arc; use std::sync::atomic::Ordering; use std::time::{Duration, Instant}; use crate::config::{GeneralConfig, MeRouteNoWriterMode, MeSocksKdfPolicy, MeWriterPickMode}; use crate::crypto::SecureRandom; +use crate::network::IpFamily; use crate::network::probe::NetworkDecision; use crate::stats::Stats; -use super::pool::MePool; +use super::pool::{MePool, ReinitStatusSnapshot, WriterContour, WriterRole}; +use super::pool_writer_security_tests::make_pool_with_decision; async fn make_pool() -> Arc { let general = GeneralConfig::default(); @@ -30,7 +32,10 @@ async fn make_pool() -> Arc { HashMap::new(), HashMap::new(), None, - NetworkDecision::default(), + NetworkDecision { + ipv4_me: true, + ..NetworkDecision::default() + }, None, Arc::new(SecureRandom::new()), Arc::new(Stats::default()), @@ -162,12 +167,118 @@ async fn connectable_endpoints_releases_quarantine_lock_before_sleep() { assert_eq!(endpoints, vec![addr]); } +#[tokio::test] +async fn refill_does_not_queue_a_removed_dc_target() { + let pool = make_pool().await; + let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 31, 0, 20)), 443); + + pool.trigger_immediate_refill_for_dc(addr, 2); + + assert!(pool.refill_states.lock().is_empty()); + assert_eq!(pool.refill_running.load(Ordering::Acquire), 0); + assert_eq!(pool.refill_pending.load(Ordering::Acquire), 0); +} + #[tokio::test(flavor = "current_thread")] -async fn refill_coalesces_one_pending_endpoint_and_cleans_up_before_first_poll() { +async fn refill_accepts_enabled_nonpreferred_family_without_multipath() { + let pool = make_pool_with_decision(NetworkDecision { + ipv4_me: true, + ipv6_me: true, + effective_prefer: 4, + effective_multipath: false, + ..NetworkDecision::default() + }) + .await; + let v4_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443); + let v6_addr = SocketAddr::new(IpAddr::V6(Ipv6Addr::LOCALHOST), 443); + pool.update_proxy_maps( + HashMap::from([(2, vec![(v4_addr.ip(), v4_addr.port())])]), + Some(HashMap::from([(2, vec![(v6_addr.ip(), v6_addr.port())])])), + ) + .await; + + pool.trigger_immediate_refill_for_dc(v6_addr, 2); + + assert_eq!(pool.refill_states.lock().len(), 1); + assert_eq!(pool.refill_running.load(Ordering::Acquire), 1); + pool.begin_shutdown(); + tokio::time::timeout(Duration::from_secs(1), async { + while pool.refill_running.load(Ordering::Acquire) != 0 { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + assert!(pool.refill_states.lock().is_empty()); +} + +#[tokio::test(flavor = "current_thread")] +async fn stale_endpoint_revision_cancels_queued_warm_refill_before_connect() { + let pool = make_pool().await; + let old_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 31, 0, 30)), 443); + let new_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 31, 0, 31)), 443); + pool.update_proxy_maps( + HashMap::from([(2, vec![(old_addr.ip(), old_addr.port())])]), + None, + ) + .await; + let old_revision = pool.endpoint_snapshot.load().revision; + let pending_generation = pool.current_generation().saturating_add(1); + pool.reinit.status.store(Arc::new(ReinitStatusSnapshot { + active_generation: pool.current_generation(), + warm_generations: vec![pending_generation], + pending_hardswap_generation: pending_generation, + pending_hardswap_started_at_epoch_secs: 1, + pending_hardswap_map_hash: 1, + pending_hardswap_endpoint_revision: old_revision, + inflight: 1, + })); + pool.trigger_immediate_refill_for_role( + old_addr, + WriterRole { + dc: 2, + family: IpFamily::V4, + generation: pending_generation, + contour: WriterContour::Warm, + }, + ); + assert_eq!(pool.refill_states.lock().len(), 1); + + pool.update_proxy_maps( + HashMap::from([(2, vec![(new_addr.ip(), new_addr.port())])]), + None, + ) + .await; + tokio::time::timeout(Duration::from_secs(1), async { + while pool.refill_running.load(Ordering::Acquire) != 0 { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + + assert!(pool.refill_states.lock().is_empty()); + assert_eq!(pool.stats.get_me_reconnect_attempts(), 0); +} + +#[tokio::test(flavor = "current_thread")] +async fn refill_preserves_bounded_pending_cardinality_and_cleans_up_before_first_poll() { let pool = make_pool().await; let first = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 31, 0, 21)), 443); let second = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 31, 0, 22)), 443); let latest = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 31, 0, 23)), 443); + pool.update_proxy_maps( + HashMap::from([( + 2, + vec![ + (first.ip(), first.port()), + (second.ip(), second.port()), + (latest.ip(), latest.port()), + ], + )]), + None, + ) + .await; pool.trigger_immediate_refill_for_dc(first, 2); pool.trigger_immediate_refill_for_dc(second, 2); @@ -175,7 +286,16 @@ async fn refill_coalesces_one_pending_endpoint_and_cleans_up_before_first_poll() assert_eq!(pool.refill_states.lock().len(), 1); assert_eq!(pool.refill_running.load(Ordering::Acquire), 1); - assert_eq!(pool.refill_pending.load(Ordering::Acquire), 1); + assert_eq!(pool.refill_pending.load(Ordering::Acquire), 2); + let state = pool + .refill_states + .lock() + .values() + .copied() + .next() + .expect("refill state"); + assert_eq!(state.pending_count, 2); + assert_eq!(state.next_addr, Some(latest)); pool.begin_shutdown(); tokio::time::timeout(Duration::from_secs(1), async { diff --git a/src/transport/middle_proxy/tests/pool_writer_publication_tests.rs b/src/transport/middle_proxy/tests/pool_writer_publication_tests.rs new file mode 100644 index 0000000..814ba21 --- /dev/null +++ b/src/transport/middle_proxy/tests/pool_writer_publication_tests.rs @@ -0,0 +1,463 @@ +use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr}; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU32, AtomicU64, Ordering}; +use std::time::{Duration, Instant}; + +use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; + +use super::codec::WriterCommand; +use super::pool::{MeWriter, WriterContour, WriterOpenIntent}; +use super::pool_writer_security_tests::{make_pool, make_pool_with_decision}; +use super::registry::ConnMeta; +use crate::network::probe::NetworkDecision; + +fn unregistered_writer( + pool: &Arc, + writer_id: u64, + addr: SocketAddr, + generation: u64, + contour: WriterContour, +) -> MeWriter { + let (tx, _rx) = mpsc::channel::(8); + MeWriter { + id: writer_id, + addr, + source_ip: addr.ip(), + writer_dc: 2, + generation, + contour: Arc::new(AtomicU8::new(contour.as_u8())), + created_at: Instant::now(), + tx, + byte_budget: pool.new_writer_byte_budget(), + cancel: CancellationToken::new(), + degraded: Arc::new(AtomicBool::new(false)), + rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)), + draining: Arc::new(AtomicBool::new(false)), + draining_started_at_epoch_secs: Arc::new(AtomicU64::new(0)), + drain_deadline_epoch_secs: Arc::new(AtomicU64::new(0)), + allow_drain_fallback: Arc::new(AtomicBool::new(false)), + } +} + +#[tokio::test] +async fn normal_warm_publication_cannot_race_past_the_dc_floor() { + let pool = make_pool().await; + let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443); + pool.update_proxy_maps( + std::collections::HashMap::from([(2, vec![(addr.ip(), addr.port())])]), + None, + ) + .await; + let generation = 2; + let writers = (1..=3) + .map(|writer_id| { + unregistered_writer(&pool, writer_id, addr, generation, WriterContour::Warm) + }) + .collect::>(); + let candidate = unregistered_writer(&pool, 4, addr, generation, WriterContour::Warm); + + assert!( + pool.authorize_writer_publication_capacity( + &candidate, + WriterContour::Warm, + WriterOpenIntent::Normal, + &writers[..2], + ) + .is_ok() + ); + assert!( + pool.authorize_writer_publication_capacity( + &candidate, + WriterContour::Warm, + WriterOpenIntent::Normal, + &writers, + ) + .is_err() + ); + assert!( + pool.authorize_writer_publication_capacity( + &candidate, + WriterContour::Warm, + WriterOpenIntent::Replacement, + &writers, + ) + .is_ok() + ); +} + +#[tokio::test] +async fn normal_active_publication_cannot_race_past_the_family_floor() { + let pool = make_pool().await; + let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443); + pool.update_proxy_maps( + std::collections::HashMap::from([(2, vec![(addr.ip(), addr.port())])]), + None, + ) + .await; + let generation = pool.current_generation(); + let writers = (1..=3) + .map(|writer_id| { + unregistered_writer(&pool, writer_id, addr, generation, WriterContour::Active) + }) + .collect::>(); + let candidate = unregistered_writer(&pool, 4, addr, generation, WriterContour::Active); + + assert!( + pool.authorize_writer_publication_capacity( + &candidate, + WriterContour::Active, + WriterOpenIntent::Coverage, + &writers[..2], + ) + .is_ok() + ); + assert!( + pool.authorize_writer_publication_capacity( + &candidate, + WriterContour::Active, + WriterOpenIntent::Coverage, + &writers, + ) + .is_err() + ); +} + +#[tokio::test] +async fn stale_same_family_writers_do_not_satisfy_current_endpoint_coverage() { + let pool = make_pool().await; + let current_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443); + let stale_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 2)), 443); + pool.update_proxy_maps( + std::collections::HashMap::from([(2, vec![(current_addr.ip(), current_addr.port())])]), + None, + ) + .await; + pool.floor_runtime + .me_adaptive_floor_cpu_cores_override + .store(1, Ordering::Relaxed); + pool.floor_runtime + .me_adaptive_floor_max_active_writers_per_core + .store(1, Ordering::Relaxed); + pool.floor_runtime + .me_adaptive_floor_max_active_writers_global + .store(1, Ordering::Relaxed); + let generation = pool.current_generation(); + let required = pool.required_writers_for_dc_with_floor_mode(1, false); + pool.writers + .write() + .await + .extend((1..=required.saturating_mul(2)).map(|writer_id| { + unregistered_writer( + &pool, + writer_id as u64, + stale_addr, + generation, + WriterContour::Active, + ) + })); + + assert!( + pool.can_open_writer_for_contour( + WriterContour::Active, + WriterOpenIntent::Coverage, + 2, + current_addr, + ) + .await + ); + assert!( + !pool + .can_open_writer_for_contour( + WriterContour::Active, + WriterOpenIntent::Coverage, + 2, + stale_addr, + ) + .await + ); +} + +#[tokio::test] +async fn covered_group_does_not_consume_another_group_coverage_slot() { + let pool = make_pool().await; + let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443); + pool.update_proxy_maps( + std::collections::HashMap::from([(2, vec![(addr.ip(), addr.port())])]), + None, + ) + .await; + pool.floor_runtime + .me_adaptive_floor_cpu_cores_override + .store(1, Ordering::Relaxed); + pool.floor_runtime + .me_adaptive_floor_max_active_writers_per_core + .store(1, Ordering::Relaxed); + pool.floor_runtime + .me_adaptive_floor_max_active_writers_global + .store(1, Ordering::Relaxed); + let generation = pool.current_generation(); + let required = pool.required_writers_for_dc_with_floor_mode(1, false); + pool.writers + .write() + .await + .extend((1..=required).map(|writer_id| { + unregistered_writer( + &pool, + writer_id as u64, + addr, + generation, + WriterContour::Active, + ) + })); + + assert!( + !pool + .can_open_writer_for_contour( + WriterContour::Active, + WriterOpenIntent::Coverage, + 2, + addr, + ) + .await + ); +} + +#[tokio::test] +async fn normal_active_publication_allows_adaptive_growth_above_family_floor() { + let pool = make_pool().await; + let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443); + pool.update_proxy_maps( + std::collections::HashMap::from([(2, vec![(addr.ip(), addr.port())])]), + None, + ) + .await; + let generation = pool.current_generation(); + let writers = (1..=3) + .map(|writer_id| { + unregistered_writer(&pool, writer_id, addr, generation, WriterContour::Active) + }) + .collect::>(); + let candidate = unregistered_writer(&pool, 4, addr, generation, WriterContour::Active); + + assert!( + pool.authorize_writer_publication_capacity( + &candidate, + WriterContour::Active, + WriterOpenIntent::Normal, + &writers, + ) + .is_ok() + ); +} + +#[tokio::test] +async fn normal_active_publication_rejects_configured_contour_cap() { + let pool = make_pool().await; + let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443); + pool.update_proxy_maps( + std::collections::HashMap::from([(2, vec![(addr.ip(), addr.port())])]), + None, + ) + .await; + pool.floor_runtime + .me_adaptive_floor_cpu_cores_override + .store(1, Ordering::Relaxed); + pool.floor_runtime + .me_adaptive_floor_max_active_writers_per_core + .store(3, Ordering::Relaxed); + pool.floor_runtime + .me_adaptive_floor_max_active_writers_global + .store(3, Ordering::Relaxed); + let generation = pool.current_generation(); + let writers = (1..=3) + .map(|writer_id| { + unregistered_writer(&pool, writer_id, addr, generation, WriterContour::Active) + }) + .collect::>(); + let candidate = unregistered_writer(&pool, 4, addr, generation, WriterContour::Active); + + assert!( + pool.authorize_writer_publication_capacity( + &candidate, + WriterContour::Active, + WriterOpenIntent::Normal, + &writers, + ) + .is_err() + ); +} + +#[tokio::test] +async fn nonpreferred_enabled_family_retains_writer_publication_authority() { + let pool = make_pool_with_decision(NetworkDecision { + ipv4_me: true, + ipv6_me: true, + effective_prefer: 4, + effective_multipath: false, + ..NetworkDecision::default() + }) + .await; + let v4_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443); + let v6_addr = SocketAddr::new(IpAddr::V6(Ipv6Addr::LOCALHOST), 443); + pool.update_proxy_maps( + std::collections::HashMap::from([(2, vec![(v4_addr.ip(), v4_addr.port())])]), + Some(std::collections::HashMap::from([( + 2, + vec![(v6_addr.ip(), v6_addr.port())], + )])), + ) + .await; + let candidate = unregistered_writer( + &pool, + 1, + v6_addr, + pool.current_generation(), + WriterContour::Active, + ); + + assert!( + pool.authorize_writer_publication_capacity( + &candidate, + WriterContour::Active, + WriterOpenIntent::Coverage, + &[], + ) + .is_ok() + ); +} + +#[tokio::test] +async fn successful_writer_publication_is_fully_visible_and_removable() { + let pool = make_pool().await; + let writer_id = 76_002; + let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443); + let (tx, _rx) = mpsc::channel::(8); + let byte_budget = pool.new_writer_byte_budget(); + let cancel = CancellationToken::new(); + let writer = MeWriter { + id: writer_id, + addr, + source_ip: addr.ip(), + writer_dc: 2, + generation: pool.current_generation(), + contour: Arc::new(AtomicU8::new(WriterContour::Active.as_u8())), + created_at: Instant::now(), + tx: tx.clone(), + byte_budget: byte_budget.clone(), + cancel: cancel.clone(), + degraded: Arc::new(AtomicBool::new(false)), + rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)), + draining: Arc::new(AtomicBool::new(false)), + draining_started_at_epoch_secs: Arc::new(AtomicU64::new(0)), + drain_deadline_epoch_secs: Arc::new(AtomicU64::new(0)), + allow_drain_fallback: Arc::new(AtomicBool::new(false)), + }; + let task_started = Arc::new(AtomicBool::new(false)); + let task_started_writer = Arc::clone(&task_started); + let writer_task = async move { + task_started_writer.store(true, Ordering::Release); + cancel.cancelled().await; + }; + let task_registration = pool.lifecycle.try_register().unwrap(); + + pool.publish_prepared_writer(writer, tx, byte_budget, task_registration, writer_task) + .await; + + assert_eq!(pool.writers.read().await.len(), 1); + assert_eq!(pool.conn_count.load(Ordering::Relaxed), 1); + tokio::time::timeout(Duration::from_secs(1), async { + while !task_started.load(Ordering::Acquire) { + tokio::task::yield_now().await; + } + }) + .await + .expect("published writer task must start"); + + let (conn_id, _response_rx) = pool.registry.register().await; + assert!( + pool.registry + .bind_writer( + conn_id, + writer_id, + ConnMeta { + target_dc: 2, + client_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 7300), + our_addr: addr, + proto_flags: 0, + }, + ) + .await + ); + assert_eq!( + pool.registry.get_writer(conn_id).await.unwrap().writer_id, + writer_id + ); + + pool.remove_writer_and_close_clients(writer_id).await; + + assert!(pool.writers.read().await.is_empty()); + assert_eq!(pool.conn_count.load(Ordering::Relaxed), 0); + assert!(pool.registry.get_writer(conn_id).await.is_none()); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn published_writer_removal_race_preserves_count() { + let pool = make_pool().await; + + for writer_id in 80_000..90_000 { + let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443); + let (tx, _rx) = mpsc::channel::(1); + let byte_budget = pool.new_writer_byte_budget(); + let cancel = CancellationToken::new(); + let writer = MeWriter { + id: writer_id, + addr, + source_ip: addr.ip(), + writer_dc: 2, + generation: pool.current_generation(), + contour: Arc::new(AtomicU8::new(WriterContour::Active.as_u8())), + created_at: Instant::now(), + tx: tx.clone(), + byte_budget: byte_budget.clone(), + cancel: cancel.clone(), + degraded: Arc::new(AtomicBool::new(false)), + rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)), + draining: Arc::new(AtomicBool::new(false)), + draining_started_at_epoch_secs: Arc::new(AtomicU64::new(0)), + drain_deadline_epoch_secs: Arc::new(AtomicU64::new(0)), + allow_drain_fallback: Arc::new(AtomicBool::new(false)), + }; + let writer_task = async move { + cancel.cancelled().await; + }; + let task_registration = pool.lifecycle.try_register().unwrap(); + let remover_pool = Arc::clone(&pool); + let remover = tokio::spawn(async move { + loop { + if remover_pool + .writers + .snapshot() + .iter() + .any(|writer| writer.id == writer_id) + { + remover_pool + .remove_writer_and_close_clients(writer_id) + .await; + return; + } + tokio::task::yield_now().await; + } + }); + + pool.publish_prepared_writer(writer, tx, byte_budget, task_registration, writer_task) + .await; + tokio::time::timeout(Duration::from_secs(1), remover) + .await + .expect("published writer must become removable") + .expect("writer remover task must not panic"); + + assert!(pool.writers.read().await.is_empty()); + assert_eq!(pool.conn_count.load(Ordering::Relaxed), 0); + } +} diff --git a/src/transport/middle_proxy/tests/pool_writer_security_tests.rs b/src/transport/middle_proxy/tests/pool_writer_security_tests.rs index 1d57257..5b16b79 100644 --- a/src/transport/middle_proxy/tests/pool_writer_security_tests.rs +++ b/src/transport/middle_proxy/tests/pool_writer_security_tests.rs @@ -15,7 +15,17 @@ use crate::crypto::SecureRandom; use crate::network::probe::NetworkDecision; use crate::stats::Stats; -async fn make_pool() -> Arc { +/// Builds an isolated ME pool for writer-state tests. +pub(super) async fn make_pool() -> Arc { + make_pool_with_decision(NetworkDecision { + ipv4_me: true, + ..NetworkDecision::default() + }) + .await +} + +/// Builds an isolated ME pool with an explicit network-family policy. +pub(super) async fn make_pool_with_decision(decision: NetworkDecision) -> Arc { let general = GeneralConfig::default(); MePool::new( @@ -34,7 +44,7 @@ async fn make_pool() -> Arc { HashMap::new(), HashMap::new(), None, - NetworkDecision::default(), + decision, None, Arc::new(SecureRandom::new()), Arc::new(Stats::new()), @@ -154,6 +164,70 @@ async fn insert_writer( pool.conn_count.fetch_add(1, Ordering::Relaxed); } +#[tokio::test] +async fn cancelled_writer_publication_leaves_no_ghost_state() { + let pool = make_pool().await; + let writer_id = 76_001; + let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443); + let (tx, mut rx) = mpsc::channel::(8); + let byte_budget = pool.new_writer_byte_budget(); + let writer = MeWriter { + id: writer_id, + addr, + source_ip: addr.ip(), + writer_dc: 2, + generation: pool.current_generation(), + contour: Arc::new(AtomicU8::new(WriterContour::Active.as_u8())), + created_at: Instant::now(), + tx: tx.clone(), + byte_budget: byte_budget.clone(), + cancel: CancellationToken::new(), + degraded: Arc::new(AtomicBool::new(false)), + rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)), + draining: Arc::new(AtomicBool::new(false)), + draining_started_at_epoch_secs: Arc::new(AtomicU64::new(0)), + drain_deadline_epoch_secs: Arc::new(AtomicU64::new(0)), + allow_drain_fallback: Arc::new(AtomicBool::new(false)), + }; + let task_started = Arc::new(AtomicBool::new(false)); + let task_started_writer = Arc::clone(&task_started); + let writer_task = async move { + task_started_writer.store(true, Ordering::Release); + let _ = rx.recv().await; + }; + let task_registration = pool.lifecycle.try_register().unwrap(); + let held_registration = pool.registry.prepare_writer_registration().await; + + let publication = + pool.publish_prepared_writer(writer, tx, byte_budget, task_registration, writer_task); + assert!( + tokio::time::timeout(Duration::from_millis(10), publication) + .await + .is_err() + ); + drop(held_registration); + + assert!(pool.writers.read().await.is_empty()); + assert_eq!(pool.conn_count.load(Ordering::Relaxed), 0); + assert!(!task_started.load(Ordering::Acquire)); + let (conn_id, _response_rx) = pool.registry.register().await; + assert!( + !pool + .registry + .bind_writer( + conn_id, + writer_id, + ConnMeta { + target_dc: 2, + client_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 7300), + our_addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443), + proto_flags: 0, + }, + ) + .await + ); +} + async fn current_writer_ids(pool: &Arc) -> HashSet { pool.writers .read() diff --git a/src/transport/middle_proxy/tests/send_adversarial_tests.rs b/src/transport/middle_proxy/tests/send_adversarial_tests.rs index 667eb33..c98e7f8 100644 --- a/src/transport/middle_proxy/tests/send_adversarial_tests.rs +++ b/src/transport/middle_proxy/tests/send_adversarial_tests.rs @@ -154,13 +154,11 @@ async fn insert_writer( }; pool.writers.write().await.push(writer); - { - let mut map = pool.proxy_map_v4.write().await; - map.entry(writer_dc) - .or_insert_with(Vec::new) - .push((addr.ip(), addr.port())); - } - pool.rebuild_endpoint_dc_map().await; + let mut map = pool.endpoint_snapshot.load().map_v4.clone(); + map.entry(writer_dc) + .or_insert_with(Vec::new) + .push((addr.ip(), addr.port())); + pool.update_proxy_maps(map, None).await; if register_in_registry { pool.registry .register_writer(writer_id, tx, byte_budget) @@ -240,14 +238,7 @@ async fn send_proxy_req_uses_live_same_dc_writer_while_preferred_endpoint_refill .await; assert!(pool.admission_ready_conditional_cast().await); - assert_eq!( - pool.preferred_endpoints_by_dc - .load() - .get(&2) - .cloned() - .unwrap_or_default(), - vec![new_addr] - ); + assert_eq!(pool.preferred_endpoints_for_dc(2).await, vec![new_addr]); let (conn_id, _rx) = pool.registry.register().await; let result = pool @@ -324,6 +315,36 @@ async fn send_proxy_req_does_not_replay_when_first_bind_commit_fails() { assert_eq!(bound.expect("writer should be bound").writer_id, 11); } +#[tokio::test] +async fn missing_client_route_does_not_prune_a_healthy_writer() { + let (pool, _rng) = make_pool().await; + let writer_id = 12; + let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 12)), 443); + let _writer_rx = insert_writer(&pool, writer_id, 2, addr, true).await; + + let result = pool + .send_proxy_req( + 999_999, + 2, + SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 30005), + SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443), + b"cancelled-route", + 0, + None, + None, + ) + .await; + + assert!(result.is_err()); + assert!( + pool.writers + .read() + .await + .iter() + .any(|writer| writer.id == writer_id) + ); +} + #[tokio::test] async fn send_proxy_req_prunes_iterative_stale_bind_failures_without_data_replay() { let (pool, _rng) = make_pool().await; diff --git a/src/util/mod.rs b/src/util/mod.rs index a81ef10..cf6c160 100644 --- a/src/util/mod.rs +++ b/src/util/mod.rs @@ -1,7 +1,11 @@ //! Utils pub mod ip; +#[cfg(unix)] +pub mod secure_fs; pub mod time; +#[cfg(unix)] +pub mod trusted_command; #[allow(unused_imports)] pub use ip::*; diff --git a/src/util/secure_fs/mod.rs b/src/util/secure_fs/mod.rs new file mode 100644 index 0000000..f124002 --- /dev/null +++ b/src/util/secure_fs/mod.rs @@ -0,0 +1,19 @@ +//! Descriptor-anchored filesystem operations for privileged runtime paths. +//! +//! Submodules: +//! - `path`: symlink-free directory traversal and anchored path ownership +//! - `write`: regular-file opening and durable atomic replacement + +mod path; +mod write; + +pub(crate) use path::{ + AnchoredPath, chdir_nofollow_or_create, open_dir_nofollow, open_trusted_dir_nofollow_or_create, +}; +pub(crate) use write::{ + atomic_replace, atomic_replace_async, open_append_regular, open_append_regular_at, + read_regular_limited, read_regular_limited_async, +}; + +#[cfg(test)] +mod tests; diff --git a/src/util/secure_fs/path.rs b/src/util/secure_fs/path.rs new file mode 100644 index 0000000..124efc8 --- /dev/null +++ b/src/util/secure_fs/path.rs @@ -0,0 +1,184 @@ +use std::ffi::{OsStr, OsString}; +use std::io; +use std::os::fd::OwnedFd; +use std::os::unix::fs::{MetadataExt, PermissionsExt}; +use std::path::{Component, Path}; + +use nix::fcntl::{OFlag, open, openat}; +use nix::sys::stat::{Mode, mkdirat}; + +const DIRECTORY_FLAGS: OFlag = OFlag::O_RDONLY + .union(OFlag::O_DIRECTORY) + .union(OFlag::O_NOFOLLOW) + .union(OFlag::O_CLOEXEC); + +/// An immutable parent-directory descriptor paired with one final path component. +pub(crate) struct AnchoredPath { + parent: OwnedFd, + name: OsString, +} + +impl AnchoredPath { + /// Opens every parent component without following symbolic links. + pub(crate) fn open(path: &Path) -> io::Result { + Self::open_with_parent_creation(path, None) + } + + /// Creates missing parent directories while traversing without symbolic links. + pub(crate) fn open_creating_parents(path: &Path, mode: u32) -> io::Result { + Self::open_with_parent_creation(path, Some(mode)) + } + + /// Opens a parent chain that cannot be renamed by group or world users. + pub(crate) fn open_trusted_parent(path: &Path) -> io::Result { + let name = path + .file_name() + .filter(|name| !name.is_empty()) + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "path has no file name"))? + .to_os_string(); + let parent_path = path.parent().unwrap_or_else(|| Path::new(".")); + let parent = open_trusted_dir_nofollow(parent_path)?; + Ok(Self { parent, name }) + } + + /// Creates missing parents and opens a chain protected from untrusted renames. + pub(crate) fn open_trusted_parent_or_create(path: &Path, mode: u32) -> io::Result { + let name = path + .file_name() + .filter(|name| !name.is_empty()) + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "path has no file name"))? + .to_os_string(); + let parent_path = path.parent().unwrap_or_else(|| Path::new(".")); + let parent = open_trusted_dir_nofollow_or_create(parent_path, mode)?; + Ok(Self { parent, name }) + } + + fn open_with_parent_creation(path: &Path, create_mode: Option) -> io::Result { + let name = path + .file_name() + .filter(|name| !name.is_empty()) + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidInput, "path has no file name"))? + .to_os_string(); + let parent = path.parent().unwrap_or_else(|| Path::new(".")); + let parent = match create_mode { + Some(mode) => open_or_create_dir_nofollow(parent, mode)?, + None => open_dir_nofollow(parent)?, + }; + Ok(Self { parent, name }) + } + + /// Returns the anchored parent descriptor. + pub(crate) fn parent(&self) -> &OwnedFd { + &self.parent + } + + /// Returns the single final component resolved relative to the parent descriptor. + pub(crate) fn name(&self) -> &OsStr { + &self.name + } +} + +/// Opens a directory by walking every component with `O_NOFOLLOW`. +pub(crate) fn open_dir_nofollow(path: &Path) -> io::Result { + open_dir_components(path, None, false) +} + +fn open_or_create_dir_nofollow(path: &Path, mode: u32) -> io::Result { + open_dir_components(path, Some(mode), false) +} + +/// Opens a directory after securely creating any missing components. +pub(crate) fn open_dir_nofollow_or_create(path: &Path, mode: u32) -> io::Result { + open_or_create_dir_nofollow(path, mode) +} + +/// Opens or creates a directory chain protected from untrusted entry replacement. +pub(crate) fn open_trusted_dir_nofollow_or_create(path: &Path, mode: u32) -> io::Result { + open_dir_components(path, Some(mode), true) +} + +/// Opens a directory only when its entire path is owned by root or the effective user. +fn open_trusted_dir_nofollow(path: &Path) -> io::Result { + open_dir_components(path, None, true) +} + +fn open_dir_components( + path: &Path, + create_mode: Option, + require_trusted: bool, +) -> io::Result { + let mut names = Vec::new(); + for component in path.components() { + match component { + Component::RootDir | Component::CurDir => continue, + Component::Normal(name) => names.push(name.to_os_string()), + Component::ParentDir => names.push(OsString::from("..")), + Component::Prefix(_) => { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "unsupported path prefix", + )); + } + } + } + let start = if path.is_absolute() { + Path::new("/") + } else { + Path::new(".") + }; + let mut current = open(start, DIRECTORY_FLAGS, Mode::empty()).map_err(errno_to_io)?; + if require_trusted { + validate_trusted_directory(¤t, !names.is_empty())?; + } + let component_count = names.len(); + for (index, name) in names.into_iter().enumerate() { + let next = match openat(¤t, name.as_os_str(), DIRECTORY_FLAGS, Mode::empty()) { + Ok(descriptor) => descriptor, + Err(nix::errno::Errno::ENOENT) if create_mode.is_some() => { + let mode = Mode::from_bits_truncate(create_mode.unwrap_or(0o750)); + match mkdirat(¤t, name.as_os_str(), mode) { + Ok(()) | Err(nix::errno::Errno::EEXIST) => {} + Err(error) => return Err(errno_to_io(error)), + } + openat(¤t, name.as_os_str(), DIRECTORY_FLAGS, Mode::empty()) + .map_err(errno_to_io)? + } + Err(error) => return Err(errno_to_io(error)), + }; + if require_trusted { + validate_trusted_directory(&next, index + 1 < component_count)?; + } + current = next; + } + Ok(current) +} + +fn validate_trusted_directory(descriptor: &OwnedFd, allow_sticky_parent: bool) -> io::Result<()> { + let file = std::fs::File::from(descriptor.try_clone()?); + let metadata = file.metadata()?; + let effective_uid = nix::unistd::Uid::effective().as_raw(); + let mode = metadata.permissions().mode(); + let writable_by_others = mode & 0o022 != 0; + let protected_sticky_parent = allow_sticky_parent && mode & 0o1000 != 0; + if !metadata.is_dir() + || (metadata.uid() != 0 && metadata.uid() != effective_uid) + || (writable_by_others && !protected_sticky_parent) + { + return Err(io::Error::new( + io::ErrorKind::PermissionDenied, + "directory path is not owned and protected from group/world writes", + )); + } + Ok(()) +} + +/// Creates missing components and changes cwd to the exact opened directory inode. +pub(crate) fn chdir_nofollow_or_create(path: &Path, mode: u32) -> io::Result<()> { + let descriptor = open_dir_nofollow_or_create(path, mode)?; + nix::unistd::fchdir(&descriptor).map_err(errno_to_io) +} + +/// Converts one `nix` errno without discarding its platform error code. +pub(super) fn errno_to_io(error: nix::errno::Errno) -> io::Error { + io::Error::from_raw_os_error(error as i32) +} diff --git a/src/util/secure_fs/tests.rs b/src/util/secure_fs/tests.rs new file mode 100644 index 0000000..c0b80fc --- /dev/null +++ b/src/util/secure_fs/tests.rs @@ -0,0 +1,78 @@ +use std::os::unix::fs::{PermissionsExt, symlink}; + +use super::path::AnchoredPath; +use super::write::atomic_replace_after_anchor; +use super::*; + +#[test] +fn directory_walk_rejects_intermediate_symlink() { + let directory = tempfile::tempdir().unwrap(); + let real = directory.path().join("real"); + let link = directory.path().join("link"); + std::fs::create_dir(&real).unwrap(); + symlink(&real, &link).unwrap(); + + assert!(open_dir_nofollow(&link).is_err()); +} + +#[test] +fn atomic_replace_does_not_follow_final_symlink() { + let directory = tempfile::tempdir().unwrap(); + let sentinel = directory.path().join("sentinel"); + let target = directory.path().join("target"); + std::fs::write(&sentinel, b"preserve").unwrap(); + symlink(&sentinel, &target).unwrap(); + + atomic_replace(&target, b"replacement", 0o600).unwrap(); + + assert_eq!(std::fs::read(&sentinel).unwrap(), b"preserve"); + assert_eq!(std::fs::read(&target).unwrap(), b"replacement"); +} + +#[test] +fn append_open_rejects_final_symlink() { + let directory = tempfile::tempdir().unwrap(); + let sentinel = directory.path().join("sentinel"); + let target = directory.path().join("target"); + std::fs::write(&sentinel, b"preserve").unwrap(); + symlink(&sentinel, &target).unwrap(); + + assert!(open_append_regular(&target, 0o640).is_err()); + assert_eq!(std::fs::read(&sentinel).unwrap(), b"preserve"); +} + +#[test] +fn trusted_parent_rejects_group_writable_directory() { + let current = std::env::current_dir().unwrap(); + let directory = tempfile::Builder::new() + .prefix("telemt-insecure-parent-") + .tempdir_in(current) + .unwrap(); + std::fs::set_permissions(directory.path(), std::fs::Permissions::from_mode(0o770)).unwrap(); + + assert!(AnchoredPath::open_trusted_parent(&directory.path().join("listener.sock")).is_err()); +} + +#[test] +fn anchored_replace_survives_parent_path_substitution() { + let directory = tempfile::tempdir().unwrap(); + let original = directory.path().join("original"); + let moved = directory.path().join("moved"); + let redirect = directory.path().join("redirect"); + std::fs::create_dir(&original).unwrap(); + std::fs::create_dir(&redirect).unwrap(); + let target = original.join("state"); + let anchored = AnchoredPath::open(&target).unwrap(); + std::fs::rename(&original, &moved).unwrap(); + symlink(&redirect, &original).unwrap(); + + atomic_replace_after_anchor(&anchored, b"anchored", 0o600).unwrap(); + + assert_eq!(std::fs::read(moved.join("state")).unwrap(), b"anchored"); + assert!(!redirect.join("state").exists()); + let mode = std::fs::metadata(moved.join("state")) + .unwrap() + .permissions() + .mode(); + assert_eq!(mode & 0o777, 0o600); +} diff --git a/src/util/secure_fs/write.rs b/src/util/secure_fs/write.rs new file mode 100644 index 0000000..72e0a39 --- /dev/null +++ b/src/util/secure_fs/write.rs @@ -0,0 +1,203 @@ +use std::ffi::OsStr; +use std::io::{self, Read, Write}; +use std::os::fd::{AsFd, OwnedFd}; +use std::path::Path; + +use nix::fcntl::{OFlag, openat, renameat}; +use nix::sys::stat::Mode; +use nix::unistd::{UnlinkatFlags, fsync, unlinkat}; + +use super::path::{AnchoredPath, errno_to_io}; + +fn open_regular_at(anchored: &AnchoredPath, flags: OFlag, mode: u32) -> io::Result { + let descriptor = openat( + anchored.parent(), + anchored.name(), + flags | OFlag::O_NOFOLLOW | OFlag::O_CLOEXEC, + Mode::from_bits_truncate(mode), + ) + .map_err(errno_to_io)?; + let file = std::fs::File::from(descriptor); + let metadata = file.metadata()?; + if !metadata.is_file() { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "target must be a regular file", + )); + } + use std::os::unix::fs::MetadataExt; + if metadata.nlink() != 1 { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "target must have exactly one directory entry", + )); + } + Ok(file) +} + +/// Reads a regular file through an anchored parent with an allocation bound. +pub(crate) fn read_regular_limited(path: &Path, max_bytes: usize) -> io::Result> { + let anchored = AnchoredPath::open(path)?; + let mut file = open_regular_at( + &anchored, + OFlag::O_RDONLY | OFlag::O_NONBLOCK, + Mode::empty().bits(), + )?; + let before = file.metadata()?; + if before.len() > max_bytes as u64 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "file exceeds configured size limit", + )); + } + let mut bytes = Vec::with_capacity(before.len() as usize); + Read::take(&mut file, max_bytes.saturating_add(1) as u64).read_to_end(&mut bytes)?; + if bytes.len() > max_bytes { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "file exceeds configured size limit", + )); + } + let after = file.metadata()?; + if !same_file_version(&before, &after) { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "file changed while it was read", + )); + } + Ok(bytes) +} + +/// Reads a bounded regular file on the blocking pool. +pub(crate) async fn read_regular_limited_async( + path: std::path::PathBuf, + max_bytes: usize, +) -> io::Result> { + tokio::task::spawn_blocking(move || read_regular_limited(&path, max_bytes)) + .await + .map_err(|error| io::Error::other(format!("secure reader task failed: {error}")))? +} + +fn same_file_version(left: &std::fs::Metadata, right: &std::fs::Metadata) -> bool { + use std::os::unix::fs::MetadataExt; + + left.dev() == right.dev() + && left.ino() == right.ino() + && left.len() == right.len() + && left.mtime() == right.mtime() + && left.mtime_nsec() == right.mtime_nsec() + && left.ctime() == right.ctime() + && left.ctime_nsec() == right.ctime_nsec() +} + +/// Opens an append-only regular file without following path components or hard links. +pub(crate) fn open_append_regular(path: &Path, mode: u32) -> io::Result { + let anchored = AnchoredPath::open_creating_parents(path, 0o750)?; + open_regular_at( + &anchored, + OFlag::O_WRONLY | OFlag::O_APPEND | OFlag::O_CREAT, + mode, + ) +} + +/// Opens one append-only file relative to an already anchored directory. +pub(crate) fn open_append_regular_at( + parent: Fd, + name: &OsStr, + mode: u32, +) -> io::Result { + if Path::new(name).components().count() != 1 { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "anchored file name must contain one component", + )); + } + let descriptor = openat( + parent, + name, + OFlag::O_WRONLY | OFlag::O_APPEND | OFlag::O_CREAT | OFlag::O_NOFOLLOW | OFlag::O_CLOEXEC, + Mode::from_bits_truncate(mode), + ) + .map_err(errno_to_io)?; + let file = std::fs::File::from(descriptor); + let metadata = file.metadata()?; + if !metadata.is_file() { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "target must be a regular file", + )); + } + use std::os::unix::fs::MetadataExt; + if metadata.nlink() != 1 { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "target must have exactly one directory entry", + )); + } + Ok(file) +} + +/// Durably replaces a file through a same-directory descriptor-anchored rename. +pub(crate) fn atomic_replace(path: &Path, contents: &[u8], mode: u32) -> io::Result<()> { + let anchored = AnchoredPath::open_creating_parents(path, 0o750)?; + atomic_replace_anchored(&anchored, contents, mode) +} + +/// Durably replaces a file on the blocking pool. +pub(crate) async fn atomic_replace_async( + path: std::path::PathBuf, + contents: Vec, + mode: u32, +) -> io::Result<()> { + tokio::task::spawn_blocking(move || atomic_replace(&path, &contents, mode)) + .await + .map_err(|error| io::Error::other(format!("secure writer task failed: {error}")))? +} + +fn atomic_replace_anchored(anchored: &AnchoredPath, contents: &[u8], mode: u32) -> io::Result<()> { + let temp_name = format!(".telemt.tmp-{}", rand::random::()); + let descriptor = openat( + anchored.parent(), + temp_name.as_str(), + OFlag::O_WRONLY | OFlag::O_CREAT | OFlag::O_EXCL | OFlag::O_NOFOLLOW | OFlag::O_CLOEXEC, + Mode::from_bits_truncate(mode), + ) + .map_err(errno_to_io)?; + let result = write_and_publish(descriptor, anchored, &temp_name, contents); + if result.is_err() { + let _ = unlinkat( + anchored.parent(), + temp_name.as_str(), + UnlinkatFlags::NoRemoveDir, + ); + } + result +} + +fn write_and_publish( + descriptor: OwnedFd, + anchored: &AnchoredPath, + temp_name: &str, + contents: &[u8], +) -> io::Result<()> { + let mut file = std::fs::File::from(descriptor); + file.write_all(contents)?; + file.sync_all()?; + renameat( + anchored.parent(), + temp_name, + anchored.parent(), + anchored.name(), + ) + .map_err(errno_to_io)?; + fsync(anchored.parent()).map_err(errno_to_io) +} + +#[cfg(test)] +pub(super) fn atomic_replace_after_anchor( + anchored: &AnchoredPath, + contents: &[u8], + mode: u32, +) -> io::Result<()> { + atomic_replace_anchored(anchored, contents, mode) +} diff --git a/src/util/trusted_command.rs b/src/util/trusted_command.rs new file mode 100644 index 0000000..cb40cdd --- /dev/null +++ b/src/util/trusted_command.rs @@ -0,0 +1,90 @@ +use std::os::unix::fs::{MetadataExt, PermissionsExt}; +use std::path::{Path, PathBuf}; + +const TRUSTED_HELPER_DIRS: [&str; 4] = ["/usr/sbin", "/usr/bin", "/sbin", "/bin"]; +const TRUSTED_HELPERS: [&str; 12] = [ + "nft", + "iptables", + "iptables-restore", + "ip6tables", + "ip6tables-restore", + "conntrack", + "pfctl", + "systemctl", + "rc-update", + "rc-service", + "sysrc", + "service", +]; + +/// Resolves a privileged helper only through the fixed system allowlist. +pub(crate) fn resolve_trusted_helper(binary: &str) -> Option { + if !TRUSTED_HELPERS.contains(&binary) { + return None; + } + + TRUSTED_HELPER_DIRS + .iter() + .map(|directory| Path::new(directory).join(binary)) + .find_map(|candidate| trusted_executable(&candidate)) +} + +fn trusted_executable(candidate: &Path) -> Option { + let canonical = std::fs::canonicalize(candidate).ok()?; + let metadata = std::fs::metadata(&canonical).ok()?; + if !metadata.is_file() + || metadata.uid() != 0 + || metadata.permissions().mode() & 0o022 != 0 + || metadata.permissions().mode() & 0o111 == 0 + || !trusted_parent_chain(canonical.parent()?) + { + return None; + } + Some(canonical) +} + +fn trusted_parent_chain(path: &Path) -> bool { + for ancestor in path.ancestors() { + let Ok(metadata) = std::fs::symlink_metadata(ancestor) else { + return false; + }; + if !metadata.is_dir() + || metadata.file_type().is_symlink() + || metadata.uid() != 0 + || metadata.permissions().mode() & 0o022 != 0 + { + return false; + } + } + true +} + +#[cfg(test)] +mod tests { + use std::os::unix::fs::PermissionsExt; + + use super::*; + + #[test] + fn privileged_helper_allowlist_rejects_arbitrary_binary() { + assert!(resolve_trusted_helper("sh").is_none()); + assert!(resolve_trusted_helper("../bin/nft").is_none()); + } + + #[test] + fn privileged_helper_allowlist_accepts_exact_restore_names() { + assert!(TRUSTED_HELPERS.contains(&"iptables-restore")); + assert!(TRUSTED_HELPERS.contains(&"ip6tables-restore")); + assert!(!TRUSTED_HELPERS.contains(&"iptables-restore-wrapper")); + } + + #[test] + fn writable_executable_is_not_trusted() { + let directory = tempfile::tempdir().unwrap(); + let executable = directory.path().join("nft"); + std::fs::write(&executable, b"#!/bin/sh\nexit 0\n").unwrap(); + std::fs::set_permissions(&executable, std::fs::Permissions::from_mode(0o777)).unwrap(); + + assert!(trusted_executable(&executable).is_none()); + } +} diff --git a/src/web/bridge.rs b/src/web/bridge.rs index 68f1b74..1281b00 100644 --- a/src/web/bridge.rs +++ b/src/web/bridge.rs @@ -17,6 +17,7 @@ pub(crate) struct BridgePage { #[allow(clippy::too_many_arguments)] pub(crate) fn render( host: &str, + base: &str, bootstrap: &str, batch_limit: usize, queue_limit: usize, @@ -32,19 +33,96 @@ pub(crate) fn render( websocket_open_secs: u64, reconnect_grace_secs: u64, carrier_probe_coalesce_ms: u64, + bridge_diagnostics_enabled: bool, rng: &SecureRandom, ) -> BridgePage { let mut nonce = [0u8; 18]; rng.fill(&mut nonce); let nonce = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(nonce); + let diagnostic_script = if bridge_diagnostics_enabled { + format!("\n") + } else { + String::new() + }; + let diagnostic_hook = |method: &str| { + bridge_diagnostics_enabled + .then(|| { + format!("if(clientDiagnostics)try{{clientDiagnostics.{method}()}}catch(error){{}}") + }) + .unwrap_or_default() + }; + let diagnostic_runtime_started = if bridge_diagnostics_enabled { + format!("{}\n", diagnostic_hook("runtimeStarted")) + } else { + String::new() + }; + let base_prefix = base.strip_suffix('/').unwrap_or(base); let body = DOCUMENT + .replace("__DIAGNOSTIC_RUNTIME__\n", &diagnostic_script) .replace("__RESPONSE_RUNTIME__", RESPONSE_RUNTIME) .replace("__REQUEST_RUNTIME__", REQUEST_RUNTIME) .replace("__BUFFER_RUNTIME__", BUFFER_RUNTIME) .replace("__RECOVERY_RUNTIME__", RECOVERY_RUNTIME) .replace("__RUNTIME__", RUNTIME) + .replace( + "__DIAGNOSTIC_BINDING__;\n", + if bridge_diagnostics_enabled { + "const clientDiagnostics=globalThis.TelemtBridgeDiagnostics;\n" + } else { + "" + }, + ) + .replace( + "__DIAGNOSTIC_RUNTIME_STARTED__;\n", + &diagnostic_runtime_started, + ) + .replace( + "__DIAGNOSTIC_BOUNDARY_ACTIVATED__;", + &diagnostic_hook("boundaryActivated"), + ) + .replace( + "__STATUS_FUNCTION__", + if bridge_diagnostics_enabled { + "state=>{if(port&&!closed){port.postMessage({t:'status',state});if(clientDiagnostics)try{clientDiagnostics.statusPosted()}catch(error){}}}" + } else { + "state=>{if(port&&!closed)port.postMessage({t:'status',state})}" + }, + ) + .replace( + "__DIAGNOSTIC_HELLO_RECEIVED__;", + &diagnostic_hook("helloReceived"), + ) + .replace( + "__HELLO_TIMEOUT_CALLBACK__", + if bridge_diagnostics_enabled { + "()=>{if(clientDiagnostics)try{clientDiagnostics.helloTimeout()}catch(error){}fail('timeout')}" + } else { + "()=>fail('timeout')" + }, + ) + .replace( + "__DIAGNOSTIC_CLIENT_CLOSE__;", + &diagnostic_hook("clientCloseBeforeHello"), + ) + .replace( + "__PAGEHIDE_CALLBACK__", + if bridge_diagnostics_enabled { + "()=>{if(clientDiagnostics)try{clientDiagnostics.documentUnloadedBeforeHello()}catch(error){}fail('navigation')}" + } else { + "()=>fail('navigation')" + }, + ) + .replace( + "__DIAGNOSTIC_BOOTSTRAP_REPLACED__;", + if bridge_diagnostics_enabled { + "if(clientDiagnostics)try{clientDiagnostics.setBootstrap(bootstrap)}catch(error){}" + } else { + "" + }, + ) .replace("__NONCE__", &nonce) .replace("__HOST__", host) + .replace("__BASE_PREFIX__", base_prefix) .replace("__BOOTSTRAP__", bootstrap) .replace("__BATCH_LIMIT__", &batch_limit.to_string()) .replace("__QUEUE_LIMIT__", &queue_limit.to_string()) @@ -88,6 +166,7 @@ pub(crate) fn render( } const DOCUMENT: &str = include_str!("bridge/document.html"); +const DIAGNOSTIC_RUNTIME: &str = include_str!("bridge/diagnostic.js"); const RESPONSE_RUNTIME: &str = include_str!("bridge/response.js"); const REQUEST_RUNTIME: &str = include_str!("bridge/request.js"); const BUFFER_RUNTIME: &str = include_str!("bridge/buffers.js"); diff --git a/src/web/bridge/diagnostic.js b/src/web/bridge/diagnostic.js new file mode 100644 index 0000000..781e1a1 --- /dev/null +++ b/src/web/bridge/diagnostic.js @@ -0,0 +1,36 @@ +(()=>{'use strict'; +let bootstrap="__BOOTSTRAP__",hello=false,emitted=0,boundaryTimer=null; +const relayOrigin='https://__HOST__',relayBase=relayOrigin+'__BASE_PREFIX__',requestMs=__BRIDGE_REQUEST_SECS__*1000; +function eventBit(event){ + if(event==='runtime_started')return 1;if(event==='status_posted')return 2;if(event==='hello_received')return 4;if(event==='boundary_timeout')return 8; + if(event==='hello_timeout')return 16;if(event==='client_close_before_hello')return 32;if(event==='document_unloaded_before_hello')return 64; + return event==='runtime_error_before_hello'?128:0; +} +function discard(response){try{if(response.body)response.body.cancel().catch(()=>{})}catch(error){}} +function report(event){ + if(hello&&event!=='hello_received')return;const bit=eventBit(event);if(!bit||(emitted&bit))return;emitted|=bit; + let timer=null; + try{ + const controller=new AbortController(),body=JSON.stringify({v:1,event});timer=setTimeout(()=>controller.abort(),requestMs); + fetch(relayBase+'/api/v1/diagnostic',{method:'POST',body,signal:controller.signal,keepalive:true,mode:'same-origin',credentials:'omit',cache:'no-store',redirect:'error',referrerPolicy:'no-referrer',headers:{Authorization:'Bearer '+bootstrap,'Content-Type':'application/json'}}) + .then(discard,()=>{}).then(()=>clearTimeout(timer),()=>clearTimeout(timer)); + }catch(error){if(timer)clearTimeout(timer)} +} +const boundaryWall=Date.now()+requestMs,boundaryMonotonic=performance.now()+requestMs; +function waitBoundary(){ + boundaryTimer=null;if(hello||(emitted&2))return;const remaining=Math.min(boundaryWall-Date.now(),boundaryMonotonic-performance.now()); + if(remaining>0){boundaryTimer=setTimeout(waitBoundary,remaining);return}report('boundary_timeout'); +} +function clearBoundary(){try{if(boundaryTimer)clearTimeout(boundaryTimer)}catch(error){}boundaryTimer=null} +function runtimeStarted(){report('runtime_started');try{boundaryTimer=setTimeout(waitBoundary,Math.max(1,Math.min(boundaryWall-Date.now(),boundaryMonotonic-performance.now())))}catch(error){}} +function boundaryActivated(){clearBoundary()} +function statusPosted(){clearBoundary();report('status_posted')} +function helloReceived(){clearBoundary();hello=true;report('hello_received')} +function helloTimeout(){report('hello_timeout')} +function clientCloseBeforeHello(){report('client_close_before_hello')} +function documentUnloadedBeforeHello(){report('document_unloaded_before_hello')} +function setBootstrap(value){if(/^[A-Za-z0-9_-]{43}$/.test(value))bootstrap=value} +addEventListener('error',()=>report('runtime_error_before_hello')); +addEventListener('unhandledrejection',()=>report('runtime_error_before_hello')); +globalThis.TelemtBridgeDiagnostics=Object.freeze({runtimeStarted,boundaryActivated,statusPosted,helloReceived,helloTimeout,clientCloseBeforeHello,documentUnloadedBeforeHello,setBootstrap}); +})(); diff --git a/src/web/bridge/document.html b/src/web/bridge/document.html index ea68ec7..7c98b5d 100644 --- a/src/web/bridge/document.html +++ b/src/web/bridge/document.html @@ -6,6 +6,7 @@ Connection +__DIAGNOSTIC_RUNTIME__ diff --git a/src/web/bridge/request.js b/src/web/bridge/request.js index 26e52e9..27bd638 100644 --- a/src/web/bridge/request.js +++ b/src/web/bridge/request.js @@ -40,7 +40,7 @@ function create(settings){ const timer=setTimeout(()=>{timedOut=true;controller.abort()},Math.max(1,Math.min(attemptLimit,remaining))); let response=null,wait=0; try{ - const fetched=await fetch(settings.origin()+path,requestOptions); + const fetched=await fetch(settings.base()+path,requestOptions); if(retryableStatus(fetched.status)){ lastReason='http';wait=retryAfterMs(fetched);settings.cancel(fetched); }else{ diff --git a/src/web/bridge/runtime.js b/src/web/bridge/runtime.js index 8933e1e..78753ca 100644 --- a/src/web/bridge/runtime.js +++ b/src/web/bridge/runtime.js @@ -1,6 +1,8 @@ (()=>{'use strict'; let bootstrap="__BOOTSTRAP__"; -const relayOrigin='https://__HOST__',carrierCapabilities='https,https-lanes,websocket,websocket-lanes'; +const relayOrigin='https://__HOST__',relayBase=relayOrigin+'__BASE_PREFIX__',carrierCapabilities='https,https-lanes,websocket,websocket-lanes'; +__DIAGNOSTIC_BINDING__; +__DIAGNOSTIC_RUNTIME_STARTED__; const responseBody=globalThis.TelemtBridgeResponse;if(!responseBody)throw new Error('missing response runtime'); const requestSupport=globalThis.TelemtBridgeRequest;if(!requestSupport)throw new Error('missing request runtime'); const bufferSupport=globalThis.TelemtBridgeBuffers;if(!bufferSupport)throw new Error('missing buffer runtime'); @@ -23,10 +25,10 @@ const pending=[],upPending=[],recoveryPending=[],lanes=new Map(),closedLanes=new const canonicalFailures=['timeout','network','upgrade','http','protocol']; const failure=(reason,message)=>Object.assign(new Error(message||reason),{telemtReason:reason}); const failureReason=(error,fallback)=>error&&canonicalFailures.includes(error.telemtReason)?error.telemtReason:fallback; -const status=state=>{if(port&&!closed)port.postMessage({t:'status',state})}; -const socketURL=()=>relayOrigin.replace(/^https:/,'wss:')+'/api/v1/ws'; +const status=__STATUS_FUNCTION__; +const socketURL=()=>relayBase.replace(/^https:/,'wss:')+'/api/v1/ws'; const requestClient=requestSupport.create({ - origin:()=>relayOrigin,closed:()=>closed,retryMs:()=>bridgeRetryMs,longPollMs:()=>longPollMs,requestMs:()=>bridgeRequestMs, + base:()=>relayBase,closed:()=>closed,retryMs:()=>bridgeRetryMs,longPollMs:()=>longPollMs,requestMs:()=>bridgeRequestMs, batchLimit:()=>batchLimit,read:(response,limit,exact,signal)=>responseBody.read(response,limit,exact,signal),cancel:responseBody.cancel, failure,reason:failureReason,retrying:()=>status('reconnecting') }); @@ -62,7 +64,7 @@ function retireCarrier(policy){ } lanes.clear();closedLanes.clear();closedLaneOrder.length=0;releasePending(pending,null);releasePending(recoveryPending,null); for(const id of retireAllStreams())if(port){const frame=closeFrame(id);port.postMessage(frame,[frame])} - bootstrap=policy.bootstrap;batchLimit=policy.limits.carrier_batch_bytes;queueLimit=policy.limits.pending_bytes_per_session; + bootstrap=policy.bootstrap;__DIAGNOSTIC_BOOTSTRAP_REPLACED__;batchLimit=policy.limits.carrier_batch_bytes;queueLimit=policy.limits.pending_bytes_per_session; queueItemLimit=policy.limits.pending_items_per_session;maxStreams=policy.limits.max_streams_per_session; laneQueueLimit=Math.min(queueLimit,8388608);laneItemLimit=Math.min(queueItemLimit,1024); longPollMs=policy.timeouts.long_poll_secs*1000;bridgeRequestMs=policy.timeouts.bridge_request_secs*1000; @@ -389,16 +391,17 @@ function queueLane(value){ } function openLaneSocket(lane){ if(lane.socket||closed)return;lane.socket=new WebSocket(socketURL(),'tproxy-lane-v1.'+sessionToken+'.'+String(lane.id));lane.socket.binaryType='arraybuffer'; - const opened=lane.socket,openTimer=setTimeout(()=>{if(!closed&&lanes.get(lane.id)===lane&&lane.socket===opened)finishLane(lane,true)},websocketOpenMs); - lane.socket.onopen=()=>{if(closed||lanes.get(lane.id)!==lane)return;lane.ready=true;status('connected');runLaneSocketUp(lane)}; + const opened=lane.socket;let upgraded=false,settled=false,openTimer=null;const finishSocket=reason=>{if(settled)return;settled=true;if(openTimer)clearTimeout(openTimer);openTimer=null;if(closed||lanes.get(lane.id)!==lane||lane.socket!==opened)return;lane.ready=false;if(!upgraded){lane.socket=null;opened.close();recoveryController.recover(reason,null);return}finishLane(lane,true)}; + openTimer=setTimeout(()=>finishSocket('timeout'),websocketOpenMs); + lane.socket.onopen=()=>{if(closed||lanes.get(lane.id)!==lane||lane.socket!==opened){opened.close();return}upgraded=true;lane.ready=true;status('connected');runLaneSocketUp(lane)}; lane.socket.onmessage=event=>{ - clearTimeout(openTimer); - if(closed||lanes.get(lane.id)!==lane||!(event.data instanceof ArrayBuffer)){finishLane(lane,true);return} + if(openTimer)clearTimeout(openTimer);openTimer=null; + if(closed||lanes.get(lane.id)!==lane||lane.socket!==opened||!(event.data instanceof ArrayBuffer)){finishLane(lane,true);return} let values;try{values=splitFrames(event.data);for(const value of values)if(value.id!==lane.id)throw new Error('cross-lane frame')}catch(error){finishLane(lane,true);return} if(values.some(value=>value.type===3))lane.remoteClosed=true; observeServerFrames(event.data);port.postMessage({t:'traffic',up:0,down:event.data.byteLength});port.postMessage(event.data,[event.data]);status('connected'); }; - lane.socket.onerror=()=>{};lane.socket.onclose=()=>{clearTimeout(openTimer);lane.ready=false;lane.socket=null;if(!closed)finishLane(lane,true)}; + lane.socket.onerror=()=>{};lane.socket.onclose=()=>finishSocket(upgraded?'network':'upgrade'); } async function runLaneSocketUp(lane){ if(lane.running||!lane.ready)return;lane.running=true;let lease=null; @@ -475,7 +478,7 @@ async function pollLane(lane){ } function deleteSession(){ const token=cleanupToken||sessionToken,headers=canonicalFailures.includes(terminalFailure)?{'X-Carrier-Failure':terminalFailure}:null; - if(token)fetch(relayOrigin+'/api/v1/session',options('DELETE',token,null,headers,undefined,true)).catch(()=>{}); + if(token)fetch(relayBase+'/api/v1/session',options('DELETE',token,null,headers,undefined,true)).catch(()=>{}); } function close(notifyServer){ if(closed)return;closed=true;if(recoveryController)recoveryController.cancel();rejectRecoveryCommit(failure('network','bridge closed'));if(helloTimer)clearTimeout(helloTimer);helloTimer=null;if(carrierTimer)clearTimeout(carrierTimer);clearProbeTimer();if(schedulerTimer)clearTimeout(schedulerTimer);schedulerTimer=null;if(attemptController)attemptController.abort();if(pollController)pollController.abort(); @@ -487,20 +490,20 @@ function close(notifyServer){ buffers.assertEmpty(); } function activatePort(nextPort){ - initialized=true;port=nextPort; + initialized=true;port=nextPort;__DIAGNOSTIC_BOUNDARY_ACTIVATED__; port.onmessage=message=>{ observeResumeTrigger(); if(message.data instanceof ArrayBuffer){ - if(!createStarted){createStarted=true;if(helloTimer)clearTimeout(helloTimer);helloTimer=null;helloFrame=message.data;if(negotiationEnabled){negotiationStartedAt=Date.now();armCarrierDeadline(attemptEpoch)}createSession(attemptEpoch)} + if(!createStarted){__DIAGNOSTIC_HELLO_RECEIVED__;createStarted=true;if(helloTimer)clearTimeout(helloTimer);helloTimer=null;helloFrame=message.data;if(negotiationEnabled){negotiationStartedAt=Date.now();armCarrierDeadline(attemptEpoch)}createSession(attemptEpoch)} else{ let data;try{data=acceptNativeFrames(message.data)}catch(error){fail(error&&error.telemtReason==='capacity'?'capacity':'protocol');return}if(!data)return; if(recoveryController.active()&&!recoveryReplaced){if(!reserve(data,null)){fail('capacity');return}recoveryPending.push(data)} else if(!carrierCommitted){if(!reserve(data,null)){fail('capacity');return}pending.push(data);maybeStartCandidate()} else queueCarrier(data); } - }else if(message.data&&message.data.t==='close'){status('failed');close(true)} + }else if(message.data&&message.data.t==='close'){__DIAGNOSTIC_CLIENT_CLOSE__;status('failed');close(true)} }; - port.start();status('connecting');helloTimer=setTimeout(()=>fail('timeout'),bridgeRequestMs); + port.start();status('connecting');helloTimer=setTimeout(__HELLO_TIMEOUT_CALLBACK__,bridgeRequestMs); } recoveryController=recoverySupport.create({ budgetMs:()=>bridgeRecoveryMs,requestMs:()=>bridgeRequestMs,url:()=>relayOrigin+recoveryPath,token:()=>cleanupToken||sessionToken, @@ -539,5 +542,5 @@ function discoverAndroid(){ discoverAndroid(); addEventListener('online',observeResumeTrigger); if(globalThis.document&&typeof globalThis.document.addEventListener==='function')globalThis.document.addEventListener('visibilitychange',()=>{if(globalThis.document.visibilityState==='visible')observeResumeTrigger()}); -addEventListener('pagehide',()=>fail('navigation'),{once:true}); +addEventListener('pagehide',__PAGEHIDE_CALLBACK__,{once:true}); })(); diff --git a/src/web/bridge/tests.rs b/src/web/bridge/tests.rs index 4ebab32..7775d0d 100644 --- a/src/web/bridge/tests.rs +++ b/src/web/bridge/tests.rs @@ -3,6 +3,7 @@ use super::*; fn render_page(bootstrap: &str, candidate_count: usize) -> BridgePage { render( "proxy.example.com", + "/", bootstrap, 2 * 1024 * 1024, 32 * 1024 * 1024, @@ -18,6 +19,31 @@ fn render_page(bootstrap: &str, candidate_count: usize) -> BridgePage { 15, 120, 0, + false, + &SecureRandom::new(), + ) +} + +fn render_diagnostic_page(bootstrap: &str) -> BridgePage { + render( + "proxy.example.com", + "/", + bootstrap, + 2 * 1024 * 1024, + 32 * 1024 * 1024, + 16 * 1024, + 1024, + true, + 4, + [3, 5, 8, 12], + 25, + 10, + 90, + 15, + 15, + 120, + 0, + true, &SecureRandom::new(), ) } @@ -48,6 +74,49 @@ fn rendered_page_contains_bounded_negotiation_contract() { ); } +#[test] +fn rendered_page_resolves_carriers_against_the_exact_base_path() { + let page = render( + "proxy.example.com", + "/Dobry-Cola/super_app/", + "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA", + 2 * 1024 * 1024, + 32 * 1024 * 1024, + 16 * 1024, + 1024, + true, + 4, + [3, 5, 8, 12], + 25, + 10, + 90, + 15, + 15, + 120, + 0, + true, + &SecureRandom::new(), + ); + + assert!( + page.body + .contains("relayBase=relayOrigin+'/Dobry-Cola/super_app'") + ); + assert!(page.body.contains("fetch(settings.base()+path")); + assert!( + page.body + .contains("relayBase.replace(/^https:/,'wss:')+'/api/v1/ws'") + ); + assert!(page.body.contains("fetch(relayBase+'/api/v1/diagnostic'")); + assert!(page.body.contains("url:()=>relayOrigin+recoveryPath")); + assert!( + !page + .body + .contains("/Dobry-Cola/super_app/Dobry-Cola/super_app") + ); + assert!(!page.body.contains("__BASE_PREFIX__")); +} + #[test] fn rendered_page_preserves_the_ios_bootstrap_literal() { let bootstrap = "BBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBB"; @@ -62,6 +131,7 @@ fn rendered_page_preserves_the_ios_bootstrap_literal() { fn rendered_page_embeds_the_configured_bridge_timing_policy() { let page = render( "proxy.example.com", + "/", "GGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGGG", 2 * 1024 * 1024, 32 * 1024 * 1024, @@ -77,6 +147,7 @@ fn rendered_page_embeds_the_configured_bridge_timing_policy() { 11, 119, 4, + false, &SecureRandom::new(), ); @@ -113,6 +184,7 @@ fn effective_deadline_formula_uses_the_final_checkpoint() { fn disabled_negotiation_does_not_arm_a_carrier_deadline() { let page = render( "proxy.example.com", + "/", "DDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDD", 2 * 1024 * 1024, 32 * 1024 * 1024, @@ -128,6 +200,7 @@ fn disabled_negotiation_does_not_arm_a_carrier_deadline() { 15, 120, 0, + false, &SecureRandom::new(), ); assert!(page.body.contains( @@ -190,3 +263,125 @@ fn rendered_page_preserves_exact_v1_status_control_envelope() { ); assert!(!page.body.contains("port.postMessage({t:'status',state,")); } + +#[test] +fn committed_websocket_lane_escalates_only_pre_upgrade_failure() { + let page = render_page("LLLLLLLLLLLLLLLLLLLLLLLLLLLLLLLLLLLLLLLLLLL", 4); + + assert!(page.body.contains("let upgraded=false,settled=false")); + assert!(page.body.contains( + "if(settled)return;settled=true;if(openTimer)clearTimeout(openTimer);openTimer=null" + )); + assert!( + page.body + .contains("if(closed||lanes.get(lane.id)!==lane||lane.socket!==opened)return") + ); + assert!(page.body.contains( + "if(!upgraded){lane.socket=null;opened.close();recoveryController.recover(reason,null);return}" + )); + assert!( + page.body + .contains("recoveryController.recover(reason,null);return}finishLane(lane,true)") + ); + assert!( + page.body + .contains("openTimer=setTimeout(()=>finishSocket('timeout'),websocketOpenMs)") + ); + assert!(page.body.contains( + "if(closed||lanes.get(lane.id)!==lane||lane.socket!==opened){opened.close();return}" + )); + assert!(page.body.contains("upgraded=true;lane.ready=true")); + assert!(page.body.contains( + "lane.socket.onmessage=event=>{\n if(openTimer)clearTimeout(openTimer);openTimer=null;\n if(closed||lanes.get(lane.id)!==lane||lane.socket!==opened||!(event.data instanceof ArrayBuffer))" + )); + assert!( + page.body + .contains("lane.socket.onclose=()=>finishSocket(upgraded?'network':'upgrade')") + ); + assert!(page.body.contains("port.postMessage({t:'status',state})")); + assert!(!page.body.contains("port.postMessage({t:'status',state,")); +} + +#[test] +fn bridge_diagnostic_sideband_is_absent_by_default() { + let page = render_page("IIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIII", 4); + + assert!(page.body.contains("\n
sidebandyes