mirror of
https://github.com/telemt/telemt.git
synced 2026-09-26 12:35:59 +03:00
Compare commits
22 Commits
main
..
thundering
| Author | SHA1 | Date | |
|---|---|---|---|
| 326c0ecdb9 | |||
| 2d63fcf376 | |||
| c12c5c73c1 | |||
| 26a574780e | |||
| 08109d53e8 | |||
| f1107c21d9 | |||
| baa9bfbb01 | |||
| 51e706770c | |||
| d706b3f3ba | |||
| 89dacbd17e | |||
| acad414cc7 | |||
| 02f66c542e | |||
| 9a683d8b3d | |||
| 55f3d19ee0 | |||
| 5ec9f85530 | |||
| 0ac236955a | |||
| b37f1ebdeb | |||
| 1ea5f7a1d6 | |||
| 54197e848a | |||
| 479899240a | |||
| 844e41ea34 | |||
| 021ad1fe68 |
@@ -36,6 +36,7 @@ nix = { version = "0.31.3", default-features = false, features = [
|
||||
"net",
|
||||
"user",
|
||||
"process",
|
||||
"dir",
|
||||
"fs",
|
||||
"signal",
|
||||
] }
|
||||
|
||||
+10
-30
@@ -1,32 +1,12 @@
|
||||
# Licensing
|
||||
# LICENSING
|
||||
## Licenses for Versions
|
||||
| Version ≥ | Version ≤ | License |
|
||||
|-----------|-----------|---------------|
|
||||
| 1.0 | 3.3.17 | NO LICNESE |
|
||||
| 3.3.18 | 3.4.0 | TELEMT PL 3 |
|
||||
|
||||
Telemt is currently distributed under the **TELEMT Public License**
|
||||
### License Types
|
||||
- **NO LICENSE** = ***ALL RIGHT RESERVED***
|
||||
- **TELEMT PL** - special Telemt Public License based on Apache License 2 principles
|
||||
|
||||
For the complete and legally binding license terms, see [`LICENSE`](./LICENSE).
|
||||
|
||||
The license file accompanying a particular release is authoritative for that release. A future release may use a different version of the TELEMT Public License or another license without changing the licensing terms of previously released versions.
|
||||
|
||||
## Summary
|
||||
|
||||
Under the TELEMT Public License, you may generally:
|
||||
|
||||
* use and reproduce the Software;
|
||||
* modify it and create derivative works;
|
||||
* publish and redistribute source or binary copies;
|
||||
* sublicense the Software;
|
||||
* sell copies or products containing the Software.
|
||||
|
||||
The license includes several conditions and limitations, including:
|
||||
|
||||
* copyright notices, attribution notices, and the license text must be preserved when redistributing Telemt-derived code;
|
||||
* modified versions must be clearly identified as modified and include a brief description of the changes;
|
||||
* modified or redistributed versions must not be presented as official Telemt releases;
|
||||
* the license does not grant rights to Telemt trademarks, logos, or branding;
|
||||
* contributors provide a patent license for patent claims necessarily infringed by their contributions, subject to defensive termination provisions;
|
||||
* contributions submitted for inclusion in Telemt are licensed under the same license unless explicitly stated otherwise;
|
||||
* operators of publicly accessible services using Telemt are encouraged to provide attribution;
|
||||
* providing corresponding source code and build instructions with binary distributions is encouraged, but not required.
|
||||
|
||||
The Software is provided **WITHOUT ANY KIND OF WARRANTY**, as described in [`LICENSE`](./LICENSE).
|
||||
|
||||
This document is only a summary intended to make Telemt's licensing easier to understand. If anything in this document conflicts with the applicable [`LICENSE`](./LICENSE), the license text controls.
|
||||
## [Telemt Public License 3](https://github.com/telemt/telemt/blob/main/LICENSE)
|
||||
|
||||
+44
-4
@@ -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<ProxyConfig>,
|
||||
response: PatchConfigResponse,
|
||||
@@ -56,6 +62,21 @@ pub(super) async fn patch_config(
|
||||
expected_revision: Option<String>,
|
||||
reload_request: Option<ReloadRequest>,
|
||||
shared: &ApiShared,
|
||||
) -> Result<PatchConfigResponse, ApiFailure> {
|
||||
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<String>,
|
||||
reload_request: Option<ReloadRequest>,
|
||||
shared: &ApiShared,
|
||||
) -> Result<PatchConfigResponse, ApiFailure> {
|
||||
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<String>,
|
||||
) -> Result<PatchConfigResponse, ApiFailure> {
|
||||
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 {
|
||||
|
||||
@@ -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();
|
||||
|
||||
+28
-17
@@ -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<String> {
|
||||
.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<String, ApiFailure> {
|
||||
@@ -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<ProxyConfig, ApiFailure> {
|
||||
|
||||
@@ -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<File>,
|
||||
}
|
||||
|
||||
impl ConfigWriteLock {
|
||||
fn acquire(path: &Path) -> std::io::Result<Self> {
|
||||
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<String, ApiFailure> {
|
||||
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<String> {
|
||||
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<Option<ExistingTarget>> {
|
||||
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<Option<ExistingTarget>> {
|
||||
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<GraphFence<'_>>,
|
||||
) -> std::io::Result<Option<String>> {
|
||||
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::<u64>()
|
||||
);
|
||||
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<GraphFence<'_>>,
|
||||
) -> std::io::Result<Option<String>> {
|
||||
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<GraphFence<'_>>,
|
||||
) -> std::io::Result<Option<String>> {
|
||||
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::<u64>()));
|
||||
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)
|
||||
}
|
||||
@@ -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<String, ApiFailure> {
|
||||
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<String, ApiFailure> {
|
||||
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::<u64>()
|
||||
);
|
||||
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",
|
||||
)
|
||||
}
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
@@ -21,7 +21,6 @@ pub(super) async fn create_user_route(
|
||||
}
|
||||
let expected_revision = parse_if_match(req.headers());
|
||||
let body = read_json::<CreateUserRequest>(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),
|
||||
|
||||
@@ -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::<BTreeSet<_>>();
|
||||
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::<BTreeSet<_>>();
|
||||
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::<PatchUserRequest>(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 {
|
||||
|
||||
+28
-3
@@ -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<ArcSwap<RuntimeGeneration>>,
|
||||
pub(super) web_trace: Arc<WebTraceStore>,
|
||||
pub(super) web_runtime_rx: watch::Receiver<WebRuntimePublication>,
|
||||
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<T, F>(&self, future: F) -> Result<T, ApiFailure>
|
||||
where
|
||||
T: Send + 'static,
|
||||
F: std::future::Future<Output = Result<T, ApiFailure>> + 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(
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
+3
-171
@@ -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<String>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub(super) struct RuntimeMePoolStateGenerationData {
|
||||
pub(super) active_generation: u64,
|
||||
pub(super) warm_generation: u64,
|
||||
pub(super) warm_generations: Vec<u64>,
|
||||
pub(super) pending_hardswap_generation: u64,
|
||||
pub(super) pending_hardswap_age_secs: Option<u64>,
|
||||
pub(super) reinit_inflight: usize,
|
||||
pub(super) reinit_max_concurrency_effective: usize,
|
||||
pub(super) draining_generations: Vec<u64>,
|
||||
}
|
||||
|
||||
#[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<RuntimeMePoolStateRefillDcData>,
|
||||
}
|
||||
|
||||
#[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<RuntimeMePoolStatePayload>,
|
||||
}
|
||||
|
||||
#[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::<u64>::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;
|
||||
|
||||
@@ -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<u64>,
|
||||
pending_hardswap_generation: u64,
|
||||
pending_hardswap_age_secs: Option<u64>,
|
||||
reinit_inflight: usize,
|
||||
reinit_max_concurrency_effective: usize,
|
||||
draining_generations: Vec<u64>,
|
||||
}
|
||||
|
||||
#[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<bool>,
|
||||
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<RuntimeMePoolStateRefillDcData>,
|
||||
}
|
||||
|
||||
#[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<RuntimeMePoolStatePayload>,
|
||||
}
|
||||
|
||||
/// 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::<u64>::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(),
|
||||
},
|
||||
}),
|
||||
}
|
||||
}
|
||||
@@ -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(),
|
||||
|
||||
@@ -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
|
||||
};
|
||||
|
||||
+3
-2
@@ -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,
|
||||
|
||||
+31
-5
@@ -4,6 +4,20 @@ pub(in crate::api) async fn create_user(
|
||||
body: CreateUserRequest,
|
||||
expected_revision: Option<String>,
|
||||
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<String>,
|
||||
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(
|
||||
|
||||
+62
-10
@@ -6,6 +6,22 @@ pub(in crate::api) async fn rotate_secret(
|
||||
body: RotateSecretRequest,
|
||||
expected_revision: Option<String>,
|
||||
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<String>,
|
||||
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<String>,
|
||||
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<String>,
|
||||
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))
|
||||
}
|
||||
|
||||
+80
-9
@@ -5,6 +5,22 @@ pub(in crate::api) async fn patch_user(
|
||||
body: PatchUserRequest,
|
||||
expected_revision: Option<String>,
|
||||
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<String>,
|
||||
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<String>,
|
||||
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<String>,
|
||||
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();
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -67,6 +67,11 @@ pub(super) async fn render(
|
||||
push_filter_form(&mut html, &query);
|
||||
html.push_str("<section><h2>Store</h2><table><tbody>");
|
||||
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(
|
||||
|
||||
@@ -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("<th>sideband</th><td>yes</td>"));
|
||||
assert!(body.contains("0123456789abcdef"));
|
||||
assert!(body.contains("192.0.2.40"));
|
||||
}
|
||||
|
||||
+23
-496
@@ -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<String>,
|
||||
#[cfg(unix)]
|
||||
/// Unix daemon lifecycle options.
|
||||
pub daemon_opts: DaemonOptions,
|
||||
/// Fire-and-forget initialization options.
|
||||
pub init_opts: Option<InitOptions>,
|
||||
}
|
||||
|
||||
@@ -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<i32> {
|
||||
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<i32> {
|
||||
}
|
||||
}
|
||||
|
||||
/// Executes a non-server subcommand on platforms without daemon support.
|
||||
#[cfg(not(unix))]
|
||||
pub fn execute_subcommand(cmd: &ParsedCommand) -> Option<i32> {
|
||||
match cmd.subcommand {
|
||||
@@ -261,487 +272,3 @@ pub fn execute_subcommand(cmd: &ParsedCommand) -> Option<i32> {
|
||||
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<String>,
|
||||
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<InitOptions> {
|
||||
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<dyn std::error::Error>> {
|
||||
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<u8> = (0..16).map(|_| rng.random::<u8>()).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!("===================");
|
||||
}
|
||||
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
+368
@@ -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<String>,
|
||||
/// 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<InitOptions> {
|
||||
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<dyn std::error::Error>> {
|
||||
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<u8> = (0..16).map(|_| rng.random::<u8>()).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!("===================");
|
||||
}
|
||||
@@ -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,
|
||||
|
||||
@@ -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<W: Watcher>(
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
+53
-10
@@ -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<PathBuf>,
|
||||
/// Raw source bytes keyed by canonical source path.
|
||||
/// Raw source bytes keyed by normalized absolute source path.
|
||||
pub(crate) source_contents: BTreeMap<PathBuf, String>,
|
||||
/// 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<PathBuf, String>,
|
||||
/// Include-expanded TOML used for typed deserialization.
|
||||
pub(crate) rendered: String,
|
||||
@@ -174,21 +177,42 @@ impl ProxyConfig {
|
||||
source_overrides: &BTreeMap<PathBuf, String>,
|
||||
) -> Result<ConfigSourceGraph> {
|
||||
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<PathBuf, String>,
|
||||
) -> Result<ConfigSourceGraph> {
|
||||
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() {
|
||||
|
||||
+50
-15
@@ -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,
|
||||
|
||||
@@ -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<UserAuthEntry>,
|
||||
by_name: HashMap<String, u32>,
|
||||
by_hint_key: HashMap<u64, Vec<u32>>,
|
||||
sni_index: HashMap<u64, Vec<u32>>,
|
||||
sni_initial_index: HashMap<u8, Vec<u32>>,
|
||||
}
|
||||
@@ -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<String, String>) -> Result<Self> {
|
||||
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::<Vec<_>>();
|
||||
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)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
+235
-119
@@ -5,13 +5,30 @@ 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;
|
||||
|
||||
// Path-based static snapshot fallback for platforms without directory descriptors.
|
||||
#[cfg(not(unix))]
|
||||
mod static_site_fallback;
|
||||
|
||||
const WEB_CAPABILITY_CONTEXT: &[u8] = b"tdesktop-web-proxy-bridge-v1\n";
|
||||
const WEB_DEBUG_FINGERPRINT_CONTEXT: &[u8] = b"telemt-web-debug-key-fingerprint-v1\0";
|
||||
@@ -63,6 +80,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(),
|
||||
@@ -192,34 +210,31 @@ fn load_static_site(
|
||||
total_files: &mut usize,
|
||||
total_bytes: &mut usize,
|
||||
) -> Result<WebStaticSite> {
|
||||
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 +247,95 @@ fn load_static_site(
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
fn open_static_root(root: &Path) -> Result<Dir> {
|
||||
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<String, WebStaticAsset>,
|
||||
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 +344,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<String, WebStaticAsset>,
|
||||
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<String> {
|
||||
let mut route = String::new();
|
||||
for component in relative.components() {
|
||||
@@ -424,4 +504,40 @@ mod tests {
|
||||
"IpJrt3e7sKtzPyoXy6w-Zj6GGEvsvclN66JzQEfPYLA"
|
||||
);
|
||||
}
|
||||
|
||||
#[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");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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<String, WebStaticAsset>,
|
||||
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<String, WebStaticAsset>,
|
||||
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(())
|
||||
}
|
||||
@@ -326,6 +326,7 @@ const WEB_LIMITS_CONFIG_KEYS: &[&str] = &[
|
||||
|
||||
const WEB_DEBUG_CONFIG_KEYS: &[&str] = &[
|
||||
"enabled",
|
||||
"sideband",
|
||||
"capture_lifecycle",
|
||||
"capture_headers",
|
||||
"capture_timings",
|
||||
|
||||
@@ -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() {
|
||||
|
||||
@@ -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"]
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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());
|
||||
}
|
||||
@@ -301,17 +301,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 +335,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]
|
||||
|
||||
+1
-1
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
|
||||
@@ -35,6 +35,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.
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
+10
-409
@@ -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<NetfilterBackend> {
|
||||
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<u16> {
|
||||
let mut ports: BTreeSet<u16> = 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<IpAddr>, u16)>, Vec<(Option<IpAddr>, u16)>) {
|
||||
let mode = cfg.server.conntrack_control.mode;
|
||||
let mut v4_targets: BTreeSet<(Option<IpAddr>, u16)> = BTreeSet::new();
|
||||
let mut v6_targets: BTreeSet<(Option<IpAddr>, 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::<IpAddr>().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::<IpAddr>().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<String>) -> 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<u8> {
|
||||
let soft_limit = nofile_soft_limit()?;
|
||||
if soft_limit == 0 {
|
||||
@@ -715,29 +339,6 @@ fn nofile_soft_limit() -> Option<u64> {
|
||||
}
|
||||
}
|
||||
|
||||
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::*;
|
||||
|
||||
@@ -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;
|
||||
@@ -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<Option<DesiredState>>,
|
||||
status_rx: watch::Receiver<Option<ReconcileStatus>>,
|
||||
terminal: CancellationToken,
|
||||
closed: Arc<AtomicBool>,
|
||||
completed_flag: Arc<AtomicBool>,
|
||||
cleanup_succeeded: Arc<AtomicBool>,
|
||||
completed: Arc<Notify>,
|
||||
}
|
||||
|
||||
impl FirewallAuthority {
|
||||
/// Starts the single process-owned firewall reconciler.
|
||||
pub(crate) fn spawn(control_plane: &ProcessControlPlane) -> Result<Self, String> {
|
||||
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<ProxyConfig>,
|
||||
stats: Arc<Stats>,
|
||||
) -> 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<ProxyConfig>,
|
||||
stats: Arc<Stats>,
|
||||
) -> 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<AtomicBool>,
|
||||
completed_flag: Arc<AtomicBool>,
|
||||
completed: Arc<Notify>,
|
||||
}
|
||||
|
||||
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<R> {
|
||||
runner: R,
|
||||
desired_rx: watch::Receiver<Option<DesiredState>>,
|
||||
status_tx: watch::Sender<Option<ReconcileStatus>>,
|
||||
terminal: CancellationToken,
|
||||
completion: CompletionGuard,
|
||||
cleanup_succeeded: Arc<AtomicBool>,
|
||||
applied: AppliedState,
|
||||
last_generation: u64,
|
||||
last_policy: Option<DesiredPolicy>,
|
||||
last_stats: Option<Arc<Stats>>,
|
||||
}
|
||||
|
||||
impl<R> FirewallReconciler<R>
|
||||
where
|
||||
R: FirewallCommandRunner + 'static,
|
||||
{
|
||||
pub(super) fn new(
|
||||
runner: R,
|
||||
desired_rx: watch::Receiver<Option<DesiredState>>,
|
||||
status_tx: watch::Sender<Option<ReconcileStatus>>,
|
||||
terminal: CancellationToken,
|
||||
closed: Arc<AtomicBool>,
|
||||
completed_flag: Arc<AtomicBool>,
|
||||
cleanup_succeeded: Arc<AtomicBool>,
|
||||
completed: Arc<Notify>,
|
||||
) -> 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<DesiredState> {
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -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<String>,
|
||||
pub(super) stdin: Option<String>,
|
||||
}
|
||||
|
||||
impl CommandSpec {
|
||||
pub(super) fn new(binary: &'static str, args: impl IntoIterator<Item = &'static str>) -> 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<Item = &'static str>,
|
||||
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<String>) -> 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")
|
||||
}
|
||||
@@ -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<R: FirewallCommandRunner>(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<R: FirewallCommandRunner>(
|
||||
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<R: FirewallCommandRunner>(
|
||||
runner: &R,
|
||||
family: IpFamily,
|
||||
slot: Option<ShadowSlot>,
|
||||
) -> 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<R: FirewallCommandRunner>(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<R: FirewallCommandRunner>(
|
||||
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<R: FirewallCommandRunner>(
|
||||
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<R: FirewallCommandRunner>(
|
||||
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<R: FirewallCommandRunner>(
|
||||
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<ShadowSlot>) -> 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
|
||||
}
|
||||
@@ -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<IpAddr>,
|
||||
pub(super) port: u16,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Eq, PartialEq)]
|
||||
pub(super) enum DesiredPolicy {
|
||||
Empty,
|
||||
Rules {
|
||||
configured_backend: ConntrackBackend,
|
||||
v4: Vec<NotrackTarget>,
|
||||
v6: Vec<NotrackTarget>,
|
||||
},
|
||||
}
|
||||
|
||||
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<Stats>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Eq, PartialEq)]
|
||||
pub(super) enum AppliedPlan {
|
||||
Empty,
|
||||
Iptables {
|
||||
slot: ShadowSlot,
|
||||
v4: Vec<NotrackTarget>,
|
||||
v6: Vec<NotrackTarget>,
|
||||
},
|
||||
Nftables {
|
||||
slot: ShadowSlot,
|
||||
v4: Vec<NotrackTarget>,
|
||||
v6: Vec<NotrackTarget>,
|
||||
},
|
||||
}
|
||||
|
||||
impl AppliedPlan {
|
||||
pub(super) fn slot(&self) -> Option<ShadowSlot> {
|
||||
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<u16> {
|
||||
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<NotrackTarget>, Vec<NotrackTarget>) {
|
||||
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::<IpAddr>() {
|
||||
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(),
|
||||
)
|
||||
}
|
||||
@@ -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<R: FirewallCommandRunner>(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<R: FirewallCommandRunner>(
|
||||
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<R: FirewallCommandRunner>(
|
||||
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<R: FirewallCommandRunner>(
|
||||
runner: &R,
|
||||
slot: ShadowSlot,
|
||||
) -> Result<(), CommandError> {
|
||||
delete_table_if_present(runner, table(slot)).await
|
||||
}
|
||||
|
||||
pub(super) async fn cleanup_all<R: FirewallCommandRunner>(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<R: FirewallCommandRunner>(
|
||||
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<R: FirewallCommandRunner>(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"
|
||||
)
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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<CommandSpec>,
|
||||
binary_calls: BTreeMap<&'static str, usize>,
|
||||
failures: Vec<FailureRule>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct FakeRunner {
|
||||
available: Arc<BTreeSet<&'static str>>,
|
||||
has_cap_net_admin: bool,
|
||||
state: Arc<Mutex<FakeState>>,
|
||||
}
|
||||
|
||||
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<CommandSpec> {
|
||||
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<Notify>,
|
||||
}
|
||||
|
||||
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));
|
||||
}
|
||||
@@ -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"));
|
||||
}
|
||||
@@ -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<R: FirewallCommandRunner> 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<bool>,
|
||||
pub(super) cancelled: bool,
|
||||
}
|
||||
|
||||
pub(super) async fn reconcile_once<I, R>(
|
||||
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<R: FirewallCommandRunner>(
|
||||
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<R: FirewallCommandRunner>(
|
||||
runner: &R,
|
||||
desired: &DesiredPolicy,
|
||||
previous: &AppliedPlan,
|
||||
) -> Result<AppliedPlan, CommandError> {
|
||||
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<R: FirewallCommandRunner>(
|
||||
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<R: FirewallCommandRunner>(
|
||||
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<R: FirewallCommandRunner>(
|
||||
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<R: FirewallCommandRunner>(
|
||||
runner: &R,
|
||||
plan: &AppliedPlan,
|
||||
) -> Result<(), CommandError> {
|
||||
recover_to_empty(runner).await?;
|
||||
install_plan(runner, plan).await
|
||||
}
|
||||
|
||||
pub(super) async fn recover_to_empty<R: FirewallCommandRunner>(
|
||||
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
|
||||
))),
|
||||
}
|
||||
}
|
||||
+28
-279
@@ -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<DaemonizeResult, DaemonError> {
|
||||
// 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<File>,
|
||||
locked: bool,
|
||||
}
|
||||
|
||||
impl PidFile {
|
||||
/// Creates a new PID file manager for the given path.
|
||||
pub fn new<P: AsRef<Path>>(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<Option<i32>, 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<Uid, DaemonError> {
|
||||
// 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<Gid, DaemonError> {
|
||||
}
|
||||
}
|
||||
|
||||
/// Reads PID from a PID file.
|
||||
#[allow(dead_code)]
|
||||
pub fn read_pid_file<P: AsRef<Path>>(path: P) -> Result<i32, DaemonError> {
|
||||
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<P: AsRef<Path>>(
|
||||
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<P: AsRef<Path>>(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());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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<File>,
|
||||
pid_identity: Option<FileIdentity>,
|
||||
lock_file: Option<Flock<File>>,
|
||||
anchor: Option<AnchoredPath>,
|
||||
}
|
||||
|
||||
#[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<P: AsRef<Path>>(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<Option<i32>, 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<File> {
|
||||
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<Option<i32>, 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<Option<i32>, 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<FileIdentity>,
|
||||
) -> 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<fs::Metadata, DaemonError> {
|
||||
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<P: AsRef<Path>>(path: P) -> Result<i32, DaemonError> {
|
||||
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<P: AsRef<Path>>(
|
||||
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<P: AsRef<Path>>(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<bool, DaemonError> {
|
||||
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<OwnedFd, DaemonError> {
|
||||
// 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::<libc::siginfo_t>(),
|
||||
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;
|
||||
@@ -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<std::process::ExitStatus> {
|
||||
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<T: Send + Sync>() {}
|
||||
|
||||
assert_send_sync::<PidFile>();
|
||||
}
|
||||
|
||||
#[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();
|
||||
}
|
||||
+40
-11
@@ -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<String, HashMap<IpAddr, usize>>,
|
||||
recent_ips: HashMap<String, HashMap<IpAddr, Instant>>,
|
||||
incarnations: HashMap<String, UserIncarnation>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct CleanupShard {
|
||||
queue: Mutex<HashMap<String, HashMap<IpAddr, usize>>>,
|
||||
queue: Mutex<CleanupQueue>,
|
||||
}
|
||||
|
||||
type CleanupQueue = HashMap<String, HashMap<UserIncarnation, HashMap<IpAddr, usize>>>;
|
||||
type CleanupBatch = HashMap<(String, UserIncarnation, IpAddr), usize>;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct UserIpLimitPolicy {
|
||||
source_generation: u64,
|
||||
max_ips: Arc<HashMap<String, usize>>,
|
||||
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<AtomicU64>,
|
||||
cleanup_deferred_releases: Arc<AtomicU64>,
|
||||
limit_policy: Arc<ArcSwap<UserIpLimitPolicy>>,
|
||||
policy_update: Arc<Mutex<()>>,
|
||||
last_compact_epoch_secs: Arc<AtomicU64>,
|
||||
cleanup_queue_len: Arc<AtomicU64>,
|
||||
cleanup_shards: Arc<Box<[CleanupShard]>>,
|
||||
@@ -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<String, HashMap<IpAddr, usize>>,
|
||||
) -> 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)
|
||||
|
||||
+112
-30
@@ -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<String, usize>) {
|
||||
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<String, usize>,
|
||||
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) {
|
||||
|
||||
+162
-25
@@ -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<Item = (&(String, UserIncarnation, IpAddr), &usize)> {
|
||||
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::<usize>()
|
||||
};
|
||||
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,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -68,10 +68,13 @@ impl UserIpTracker {
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn run_periodic_maintenance(self: Arc<Self>) {
|
||||
pub async fn run_periodic_maintenance(self: Arc<Self>, 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 {
|
||||
|
||||
+92
-4
@@ -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;
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
+3
-49
@@ -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)]
|
||||
|
||||
+185
-141
@@ -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<Utc>,
|
||||
file: Option<File>,
|
||||
#[cfg(unix)]
|
||||
dir_fd: OwnedFd,
|
||||
now: Box<dyn Fn() -> DateTime<Utc> + 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<Utc>) -> 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<Utc>) {
|
||||
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<Vec<LogFileCandidate>> {
|
||||
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<Vec<LogFileCandidate>> {
|
||||
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<Utc>) -> SystemTime {
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::io::Write;
|
||||
|
||||
use tempfile::tempdir;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn fixed_now() -> DateTime<Utc> {
|
||||
DateTime::<Utc>::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<PathBuf> {
|
||||
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;
|
||||
|
||||
@@ -0,0 +1,143 @@
|
||||
use std::io::Write;
|
||||
|
||||
use tempfile::tempdir;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn fixed_now() -> DateTime<Utc> {
|
||||
DateTime::<Utc>::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<PathBuf> {
|
||||
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()
|
||||
);
|
||||
}
|
||||
+28
-41
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -120,6 +120,34 @@ impl ProcessControlPlane {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Registers work that must finish once accepted, even after shutdown cancellation starts.
|
||||
pub(crate) fn spawn_completion<F>(&self, future: F) -> Result<(), F>
|
||||
where
|
||||
F: Future<Output = ()> + 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<S, F>(&self, spawn: S) -> Result<(), S>
|
||||
where
|
||||
S: FnOnce(CancellationToken) -> F,
|
||||
F: Future<Output = ()> + 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));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
|
||||
@@ -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<Notify>);
|
||||
|
||||
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");
|
||||
}
|
||||
}
|
||||
+33
-599
@@ -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>) -> 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<std::sync::Arc<UpstreamManager>>,
|
||||
) -> Option<ProxyConfigData> {
|
||||
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;
|
||||
|
||||
@@ -0,0 +1,369 @@
|
||||
use std::time::Duration;
|
||||
|
||||
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 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,
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
/// 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>) -> 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<std::sync::Arc<UpstreamManager>>,
|
||||
) -> Option<ProxyConfigData> {
|
||||
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;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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));
|
||||
}
|
||||
}
|
||||
@@ -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<ProxyConfig>,
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
|
||||
+74
-16
@@ -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::<crate::conntrack_control::FirewallAuthority>;
|
||||
|
||||
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;
|
||||
|
||||
@@ -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<Option<RuntimeWatchState>>,
|
||||
listener_manager: Arc<Mutex<ListenerManager>>,
|
||||
web_trace: Arc<WebTraceStore>,
|
||||
conntrack_firewall: Option<FirewallAuthority>,
|
||||
}
|
||||
|
||||
/// Process-owned handle that quiesces reloads before shutdown snapshots the runtime.
|
||||
@@ -106,6 +108,7 @@ impl ReloadSupervisor {
|
||||
runtime_watch_tx: watch::Sender<Option<RuntimeWatchState>>,
|
||||
listener_manager: ListenerManager,
|
||||
web_trace: Arc<WebTraceStore>,
|
||||
conntrack_firewall: Option<FirewallAuthority>,
|
||||
) -> 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,
|
||||
|
||||
@@ -23,10 +23,12 @@ fn runtime_log_filter() -> RuntimeLogFilter {
|
||||
|
||||
fn prepared_runtime(generation: Arc<RuntimeGeneration>) -> 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())
|
||||
|
||||
@@ -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<IpAddr>, Option<IpAddr>),
|
||||
/// Gate opened only after the candidate becomes the active generation.
|
||||
pub(crate) config_watcher_activation: watch::Sender<bool>,
|
||||
/// 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<QuotaStore>,
|
||||
connection_authority: Arc<UserConnectionAuthority>,
|
||||
runtime_log_filter: RuntimeLogFilter,
|
||||
tls_full_cert_budget: Arc<TlsFullCertBudget>,
|
||||
user_admission: Arc<UserAdmissionAuthority>,
|
||||
ip_tracker: Arc<UserIpTracker>,
|
||||
traffic_limiter: Arc<TrafficLimiter>,
|
||||
direct_buffer_budget: Arc<DirectBufferBudget>,
|
||||
max_connections: Arc<Semaphore>,
|
||||
) -> Result<PreparedRuntime, String> {
|
||||
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();
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -56,6 +56,7 @@ pub(super) async fn prepare_runtime(
|
||||
ip_tracker: Arc<UserIpTracker>,
|
||||
shared_state: Arc<ProxySharedState>,
|
||||
direct_buffer_budget: Arc<DirectBufferBudget>,
|
||||
max_connections: Arc<Semaphore>,
|
||||
route_runtime: Arc<RouteRuntimeController>,
|
||||
api_me_pool: Arc<RwLock<Option<Arc<MePool>>>>,
|
||||
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,
|
||||
));
|
||||
|
||||
|
||||
@@ -90,6 +90,7 @@ impl RuntimeLogFilter {
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(crate) async fn spawn_runtime_tasks(
|
||||
generation_id: u64,
|
||||
config: &Arc<ProxyConfig>,
|
||||
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<IpAddr> = 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,
|
||||
|
||||
@@ -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<ArcSwap<RuntimeGeneration>>,
|
||||
quota_state: Arc<QuotaStateOwner>,
|
||||
reload_supervisor: ReloadSupervisorHandle,
|
||||
conntrack_firewall: Option<FirewallAuthority>,
|
||||
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<ArcSwap<RuntimeGeneration>>,
|
||||
quota_state: Arc<QuotaStateOwner>,
|
||||
reload_supervisor: ReloadSupervisorHandle,
|
||||
conntrack_firewall: Option<FirewallAuthority>,
|
||||
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");
|
||||
}
|
||||
|
||||
@@ -27,6 +27,7 @@ mod protocol;
|
||||
mod proxy;
|
||||
mod quota_state;
|
||||
mod service;
|
||||
mod slot_budget;
|
||||
mod startup;
|
||||
mod stats;
|
||||
mod stream;
|
||||
|
||||
@@ -236,6 +236,10 @@ async fn handle<B>(
|
||||
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<B>(
|
||||
tls_cache,
|
||||
tls_full_cert_budget,
|
||||
web_publication,
|
||||
me_hardswap.as_ref(),
|
||||
)
|
||||
.await;
|
||||
let resp = Response::builder()
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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"));
|
||||
}
|
||||
}
|
||||
@@ -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"
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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"));
|
||||
|
||||
+167
-29
@@ -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<R, W>(
|
||||
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<Stats>,
|
||||
ip_tracker: Arc<UserIpTracker>,
|
||||
quota_handle: UserQuotaHandle,
|
||||
_connection_permit: UserConnectionPermit,
|
||||
_stats_observation: Option<UserConnectionObservation>,
|
||||
ip_permit: Option<UserIpPermit>,
|
||||
released: bool,
|
||||
}
|
||||
|
||||
struct UserIpPermit {
|
||||
tracker: Arc<UserIpTracker>,
|
||||
owner: Option<UserIpOwner>,
|
||||
}
|
||||
|
||||
struct UserIpOwner {
|
||||
user: String,
|
||||
incarnation: UserIncarnation,
|
||||
ip: IpAddr,
|
||||
tracks_ip: bool,
|
||||
active: bool,
|
||||
}
|
||||
|
||||
impl UserIpPermit {
|
||||
fn new(
|
||||
tracker: Arc<UserIpTracker>,
|
||||
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<Stats>,
|
||||
ip_tracker: Arc<UserIpTracker>,
|
||||
user: String,
|
||||
ip: IpAddr,
|
||||
incarnation: UserIncarnation,
|
||||
quota_handle: UserQuotaHandle,
|
||||
connection_permit: UserConnectionPermit,
|
||||
stats_observation: Option<UserConnectionObservation>,
|
||||
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<Stats>,
|
||||
peer_addr: SocketAddr,
|
||||
ip_tracker: Arc<UserIpTracker>,
|
||||
) -> Result<UserConnectionReservation> {
|
||||
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<Stats>,
|
||||
peer_addr: SocketAddr,
|
||||
ip_tracker: Arc<UserIpTracker>,
|
||||
) -> Result<UserConnectionReservation> {
|
||||
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,
|
||||
))
|
||||
}
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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<u64>,
|
||||
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<ParkingMutexGuard<'_, ()>> {
|
||||
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<Self>,
|
||||
@@ -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<DirectBufferBudget>,
|
||||
buffer_pool: Arc<BufferPool>,
|
||||
stats: Arc<Stats>,
|
||||
shared: Arc<ProxySharedState>,
|
||||
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<u8> {
|
||||
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<u8> {
|
||||
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<u64> {
|
||||
let raw = tokio::fs::read_to_string(path).await.ok()?;
|
||||
let raw = raw.trim();
|
||||
if raw == "max" {
|
||||
return None;
|
||||
}
|
||||
let value = raw.parse::<u64>().ok()?;
|
||||
(value < (1u64 << 60)).then_some(value)
|
||||
}
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
async fn read_u64_file(path: &str) -> Option<u64> {
|
||||
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::<u64>().ok()
|
||||
})
|
||||
.unwrap_or(0)
|
||||
.saturating_mul(1024)
|
||||
}
|
||||
|
||||
fn align_up(bytes: usize) -> usize {
|
||||
bytes
|
||||
.div_ceil(DIRECT_BUFFER_UNIT_BYTES)
|
||||
|
||||
@@ -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<DirectBufferBudget>,
|
||||
buffer_pool: Arc<BufferPool>,
|
||||
stats: Arc<Stats>,
|
||||
shared: Arc<ProxySharedState>,
|
||||
connection_slots: Arc<Semaphore>,
|
||||
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<u8> {
|
||||
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<u8> {
|
||||
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<u64> {
|
||||
let raw = tokio::fs::read_to_string(path).await.ok()?;
|
||||
let raw = raw.trim();
|
||||
if raw == "max" {
|
||||
return None;
|
||||
}
|
||||
let value = raw.parse::<u64>().ok()?;
|
||||
(value < (1u64 << 60)).then_some(value)
|
||||
}
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
async fn read_u64_file(path: &str) -> Option<u64> {
|
||||
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::<u64>().ok()
|
||||
})
|
||||
.unwrap_or(0)
|
||||
.saturating_mul(1024)
|
||||
}
|
||||
@@ -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<std::fs::File> {
|
||||
#[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
|
||||
|
||||
@@ -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<R, W>(
|
||||
session_cancel: CancellationToken,
|
||||
shared: Arc<ProxySharedState>,
|
||||
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,
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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<u32> {
|
||||
pub(super) fn sticky_hint_get_by_ip(shared: &ProxySharedState, peer_ip: IpAddr) -> Option<u64> {
|
||||
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<u32> {
|
||||
) -> Option<u64> {
|
||||
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<u32> {
|
||||
pub(super) fn sticky_hint_get_by_sni(shared: &ProxySharedState, sni: &str) -> Option<u64> {
|
||||
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<K>(
|
||||
entries: &DashMap<K, u64>,
|
||||
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,
|
||||
|
||||
@@ -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<IpAddr, AuthProbeState>,
|
||||
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<IpAddr, AuthProbeState>,
|
||||
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<u32> {
|
||||
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<IpAddr, AuthProbeState> {
|
||||
&shared.handshake.auth_probe
|
||||
}
|
||||
|
||||
mod testing;
|
||||
#[cfg(test)]
|
||||
pub(crate) fn auth_probe_saturation_state_for_testing_in_shared(
|
||||
shared: &ProxySharedState,
|
||||
) -> &Mutex<Option<AuthProbeSaturationState>> {
|
||||
&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<AuthProbeSaturationState>> {
|
||||
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<HashSet<(String, String)>> {
|
||||
&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> {
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user