Compare commits

...

5 Commits

Author SHA1 Message Date
Alexey
fabd98ce89 Cfg Test for handle_bad_client 2026-07-18 14:31:43 +03:00
Alexey
7df3dab5e8 Update API.md
Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
2026-07-18 14:30:03 +03:00
Alexey
c6f40e3717 Harden Maestro reload lifecycle and readiness barriers
Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
2026-07-18 14:27:24 +03:00
Alexey
91e05265be Docs for API In-Proccess Reload
Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
2026-07-17 23:56:42 +03:00
Alexey
991d5b2c38 Maestro: add in-process runtime generation reload
Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
2026-07-17 23:56:00 +03:00
34 changed files with 3276 additions and 1110 deletions

View File

@@ -107,7 +107,9 @@ Notes:
| `GET` | `/v1/stats/users/active-ips` | none | `200` | `UserActiveIps[]` |
| `GET` | `/v1/stats/users` | none | `200` | `UserInfo[]` |
| `GET` | `/v1/config` | none | `200` | `ConfigData` |
| `PATCH` | `/v1/config` | sparse JSON object | `200` | `PatchConfigResponse` |
| `PATCH` | `/v1/config` | sparse JSON object; optional reload query | `200` or `202` | `PatchConfigResponse` |
| `POST` | `/v1/system/reload` | `ReloadRequest` or empty body | `202` | `ReloadAccepted` |
| `GET` | `/v1/system/reload/{id}` | none | `200` | `ReloadStatus` |
| `GET` | `/v1/users` | none | `200` | `UserInfo[]` |
| `POST` | `/v1/users` | `CreateUserRequest` | `201` or `202` | `CreateUserResponse` |
| `GET` | `/v1/users/{username}` | none | `200` | `UserInfo` |
@@ -146,7 +148,9 @@ Notes:
| `GET /v1/stats/users/active-ips` | Returns users that currently have non-empty active source-IP lists. |
| `GET /v1/stats/users` | Alias of `GET /v1/users`; returns disk-first user views with runtime lag flag. |
| `GET /v1/config` | Returns the current editable config sections as JSON (no `access.*`) plus the revision. |
| `PATCH /v1/config` | Applies a sparse patch to editable config sections; validates, writes, and reports restart impact. |
| `PATCH /v1/config` | Applies a sparse patch and optionally submits an in-process runtime reload to Maestro. |
| `POST /v1/system/reload` | Loads and validates the current on-disk config, then asks Maestro to prepare and activate a new runtime generation. |
| `GET /v1/system/reload/{id}` | Returns one retained reload status; the coordinator retains the most recent 32 operations. |
| `GET /v1/users` | Returns disk-first user views sorted by username. |
| `POST /v1/users` | Creates a user and returns the effective user view plus secret. |
| `GET /v1/users/{username}` | Returns one disk-first user view or `404` when absent. |
@@ -170,11 +174,13 @@ Notes:
| `404` | `not_found` | Unknown route, unknown user, or unsupported sub-route. |
| `405` | `method_not_allowed` | Unsupported method for `/v1/users/{username}` route shape. |
| `409` | `revision_conflict` | `If-Match` revision mismatch. |
| `409` | `reload_in_progress` | Another reload operation is non-terminal. |
| `409` | `user_exists` | User already exists on create. |
| `409` | `last_user_forbidden` | Attempt to delete last configured user. |
| `413` | `payload_too_large` | Body exceeds `request_body_limit_bytes`. |
| `500` | `internal_error` | Internal error (I/O, serialization, config load/save). |
| `503` | `api_disabled` | API disabled in config. |
| `503` | `maestro_unavailable` | Maestro's reload command channel is unavailable. |
## Routing and Method Edge Cases
@@ -292,13 +298,17 @@ Sections absent from the config file are absent from the response (not `null`).
### `PatchConfigResponse`
Returned by `PATCH /v1/config` on success (`200`).
Returned by `PATCH /v1/config` on success (`200`, or `202` when a reload was accepted).
| Field | Type | Description |
| --- | --- | --- |
| `revision` | `string` | SHA-256 hex of the config file after the patch was written. |
| `restart_required` | `bool` | `true` when one or more changed fields require a process restart to take effect. Hot-reloadable fields (e.g. `general.log_level`) are applied automatically by the config file watcher; restart-required fields (e.g. any `censorship.*`, `timeouts.*`, `upstreams`, or `general.modes` change) are written to disk but only take effect after the Telemt process is restarted. The caller is responsible for triggering a restart when this flag is `true`. |
| `restart_required` | `bool` | Legacy classifier result: `true` when the old file watcher alone cannot apply every changed field. Use `runtime_reload_required` and `process_restart_required` for new integrations. |
| `runtime_reload_required` | `bool` | `true` when full effect requires a Maestro runtime-generation reload rather than the legacy hot-field overlay. |
| `process_restart_required` | `bool` | `true` when process-owned sockets or paths changed and remain deferred after an in-process reload. |
| `deferred_process_fields` | `string[]` | Process-owned fields that the active process cannot rebind during generation activation. |
| `changed` | `string[]` | Top-level section names that differed between the old and new config (e.g. `["censorship"]`). |
| `reload` | `ReloadAccepted?` | Present only when the patch included a valid reload query and Maestro accepted the operation. |
### `HealthData`
| Field | Type | Description |
@@ -1376,24 +1386,50 @@ Applies a sparse patch to the editable config sections. The merged config is ful
**Read-only mode:** returns `403 read_only` when the API runs with `read_only = true`.
**Success `200` response body** (`data` field of the standard envelope):
**Optional in-process reload query:**
| Query | Required | Description |
| --- | --- | --- |
| `reload=instant` | no | Activates a new generation and cancels sessions owned by the previous generation. |
| `reload=drain` | no | Activates a new generation and lets old sessions finish until `timeout_secs`. |
| `timeout_secs=1..3600` | for `reload=drain` | Bounded old-generation drain interval. Invalid with `reload=instant`. |
| `failure_policy=keep_new\|rollback` | no | Defaults to `keep_new`. `rollback` applies only through the activation barrier, before old-generation teardown. |
Without a `reload` query parameter, the endpoint preserves the legacy behavior: it writes the patch and the file watcher applies only supported hot fields.
**Success `200` or `202` response body** (`data` field of the standard envelope):
```json
{
"revision": "<new-sha256-hex>",
"restart_required": true,
"changed": ["censorship"]
"runtime_reload_required": true,
"process_restart_required": false,
"deferred_process_fields": [],
"changed": ["censorship"],
"reload": {
"reload_id": 7,
"target_generation": 2,
"config_revision": "<new-sha256-hex>",
"state": "accepted",
"mode": "instant",
"failure_policy": "keep_new"
}
}
```
- `revision` — SHA-256 hex of the config file after the write.
- `restart_required``true` when the change affects a field that Telemt cannot hot-reload (e.g. `censorship.*`, `timeouts.*`, `upstreams`, `general.modes`). Hot-reloadable fields (e.g. `general.log_level`) are applied automatically by the config file watcher. Restart-required fields are written to disk but only take effect after the Telemt process is restarted; the caller is responsible for triggering the restart.
- `restart_required`legacy file-watcher classification retained for compatibility.
- `runtime_reload_required` — reports whether a full Maestro generation reload is needed for runtime effect.
- `process_restart_required` and `deferred_process_fields` — report process-owned sockets or paths that remain unchanged by an in-process reload.
- `changed` — list of top-level section names that differed.
- `reload` — accepted operation metadata; omitted when no reload query was supplied.
**Status codes:**
| HTTP | `error.code` | Condition |
| --- | --- | --- |
| `200` | — | Patch applied successfully. |
| `202` | — | Patch applied and runtime reload accepted. |
| `400` | `bad_request` | Invalid JSON, empty patch, or config validation/deserialization failure. |
| `400` | `access_not_editable` | Patch contains an `access` key. |
| `400` | `section_not_editable` | Patch contains `server`, `network`, or an unknown top-level key. |
@@ -1401,6 +1437,7 @@ Applies a sparse patch to the editable config sections. The merged config is ful
| `403` | `read_only` | API is in read-only mode. |
| `405` | `method_not_allowed` | Method other than `GET` or `PATCH` used on `/v1/config`. |
| `409` | `revision_conflict` | `If-Match` header supplied but does not match current revision. |
| `409` | `reload_in_progress` | Another runtime reload is active; the patch is not written. |
| `500` | `internal_error` | I/O or serialization failure. |
**curl example:**
@@ -1412,14 +1449,40 @@ curl -s -H "Authorization: <token>" http://127.0.0.1:<api>/v1/system/info | jq -
curl -s -X PATCH -H "Authorization: <token>" -H "If-Match: <revision>" \
-H "Content-Type: application/json" \
-d '{"censorship":{"tls_domain":"front.example.com"}}' \
http://127.0.0.1:<api>/v1/config
'http://127.0.0.1:<api>/v1/config?reload=instant'
```
## Runtime Reload Endpoints
### `POST /v1/system/reload`
Loads the current on-disk config under the API mutation lock and submits an immutable config snapshot to Maestro. `If-Match` is optional and uses the same revision contract as `PATCH /v1/config`. An empty body defaults to `{"mode":"instant","failure_policy":"keep_new"}`.
```json
{
"mode": "drain",
"timeout_secs": 30,
"failure_policy": "rollback"
}
```
The endpoint returns `202` with `ReloadAccepted`. A concurrent non-terminal reload returns `409 reload_in_progress`. Config parsing or validation failure is reported before a command is submitted.
### `GET /v1/system/reload/{id}`
Returns `ReloadStatus` with `state` equal to `accepted`, `preparing`, `activating`, `draining`, `succeeded`, `rolled_back`, or `failed`. Terminal statuses include `finished_at_epoch_secs`; failures include `error`. Successful activation may include `warnings` for old-generation cleanup failures and `deferred_process_fields` for process-owned settings.
Runtime generation activation rebuilds statistics, upstream routing, replay and buffer state, TLS-front cache, IP tracking, admission/route state, and Middle-End orchestration. Per-user quota accounting is process-scoped and remains continuous across generations. API, metrics, client TCP/Unix listeners, PID ownership, and logging remain process-scoped; changed bind/path fields are reported as deferred and do not cause Maestro to invoke systemd, containerd, or another process supervisor.
Reload preparation requires every configured TLS-front domain to have a non-default cached profile and requires a ready Middle-End pool when direct fallback is disabled. A candidate that does not satisfy either readiness condition fails without replacing the active generation.
The revision is verified again after preparation. With `failure_policy=rollback`, a changed revision or revision read failure rolls the candidate back; with `failure_policy=keep_new`, the condition is reported in `warnings` and activation continues.
## Mutation Semantics
| Endpoint | Notes |
| --- | --- |
| `PATCH /v1/config` | Deep-merges the patch into editable config sections (tables merged per-field; arrays/scalars replaced wholesale). Validates the merged result before writing. Writes only the touched sections via atomic `tmp + rename`. Returns the new revision and which sections changed. |
| `PATCH /v1/config` | Deep-merges and validates the patch, writes touched sections via atomic `tmp + rename`, and optionally submits the exact written revision for an in-process Maestro reload. |
| `POST /v1/users` | Creates user, validates config, then atomically updates only affected `access.*` TOML tables (`access.users` always, plus optional per-user tables present in request). |
| `PATCH /v1/users/{username}` | Partial update of provided fields only. Missing fields remain unchanged; explicit `null` removes optional per-user entries. The write path updates only affected `access.*` TOML tables. |
| `POST /v1/users/{username}/rotate-secret` | Replaces the user's secret with a provided valid 32-hex value or a generated value, then returns the effective secret in `CreateUserResponse`. |

View File

@@ -6,19 +6,28 @@ use toml::Value as Toml;
use super::ApiShared;
use super::config_store::{
EDITABLE_SECTIONS, compute_revision, current_revision, save_sections_to_disk,
EDITABLE_SECTIONS, compute_revision, current_revision, load_config_from_disk,
save_sections_to_disk,
};
use super::model::ApiFailure;
use crate::config::ProxyConfig;
use crate::config::hot_reload::classify_config_changes;
use crate::maestro::reload::{ReloadAccepted, ReloadRequest, ReloadSubmitError};
use crate::maestro::runtime_build::deferred_process_fields;
use serde::Serialize;
use std::path::Path;
use std::sync::Arc;
#[derive(Debug, Serialize)]
pub(super) struct PatchConfigResponse {
pub revision: String,
pub restart_required: bool,
pub runtime_reload_required: bool,
pub process_restart_required: bool,
pub deferred_process_fields: Vec<String>,
pub changed: Vec<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub reload: Option<ReloadAccepted>,
}
/// Shared-state wrapper around [`apply_patch_to_path`]: serializes config
@@ -27,10 +36,40 @@ pub(super) struct PatchConfigResponse {
pub(super) async fn patch_config(
patch_json: Json,
expected_revision: Option<String>,
reload_request: Option<ReloadRequest>,
shared: &ApiShared,
) -> Result<PatchConfigResponse, ApiFailure> {
let _guard = shared.mutation_lock.lock().await;
let resp = apply_patch_to_path(&shared.config_path, &patch_json, expected_revision).await?;
if reload_request.is_some()
&& let Some(reload_id) = shared.reload_control.in_progress().await
{
return Err(ApiFailure::new(
hyper::StatusCode::CONFLICT,
"reload_in_progress",
format!("Reload {} is already in progress", reload_id),
));
}
let mut resp = apply_patch_to_path(&shared.config_path, &patch_json, expected_revision).await?;
if let Some(request) = reload_request {
let config = Arc::new(load_config_from_disk(&shared.config_path).await?);
let accepted = shared
.reload_control
.submit(config, resp.revision.clone(), request)
.await
.map_err(|error| match error {
ReloadSubmitError::InProgress(reload_id) => ApiFailure::new(
hyper::StatusCode::CONFLICT,
"reload_in_progress",
format!("Reload {} is already in progress", reload_id),
),
ReloadSubmitError::MaestroUnavailable => ApiFailure::new(
hyper::StatusCode::SERVICE_UNAVAILABLE,
"maestro_unavailable",
"Maestro reload coordinator is unavailable",
),
})?;
resp.reload = Some(accepted);
}
drop(_guard);
shared
.runtime_events
@@ -114,6 +153,7 @@ pub(super) async fn apply_patch_to_path(
// 4. classify changes (Telemt's own hot/restart rule)
let class = classify_config_changes(&old_cfg, &new_cfg);
let deferred_process_fields = deferred_process_fields(&old_cfg, &new_cfg);
// 5. write only the touched top-level sections
let revision = save_sections_to_disk(config_path, &new_cfg, &touched).await?;
@@ -121,7 +161,11 @@ pub(super) async fn apply_patch_to_path(
Ok(PatchConfigResponse {
revision,
restart_required: class.restart_required,
runtime_reload_required: class.restart_required,
process_restart_required: !deferred_process_fields.is_empty(),
deferred_process_fields,
changed: class.changed,
reload: None,
})
}
@@ -266,6 +310,9 @@ mod tests {
let patch: Json = serde_json::json!({"censorship": {"tls_domain": "b.com"}});
let resp = apply_patch_to_path(&path, &patch, None).await.unwrap();
assert!(resp.restart_required);
assert!(resp.runtime_reload_required);
assert!(!resp.process_restart_required);
assert!(resp.deferred_process_fields.is_empty());
assert!(resp.changed.iter().any(|c| c == "censorship"));
let written = std::fs::read_to_string(&path).unwrap();
assert!(written.contains("tls_domain = \"b.com\""));
@@ -406,6 +453,8 @@ mod tests {
let patch: Json = serde_json::json!({"general": {"log_level": "debug"}});
let resp = apply_patch_to_path(&path, &patch, None).await.unwrap();
assert!(!resp.restart_required);
assert!(!resp.runtime_reload_required);
assert!(!resp.process_restart_required);
assert!(resp.changed.iter().any(|c| c == "general"));
}
}

View File

@@ -72,6 +72,13 @@ pub(super) async fn current_revision(config_path: &Path) -> Result<String, ApiFa
Ok(compute_revision(&content))
}
pub(crate) async fn current_revision_for_maestro(config_path: &Path) -> Result<String, String> {
let content = tokio::fs::read_to_string(config_path)
.await
.map_err(|error| format!("failed to read config: {}", error))?;
Ok(compute_revision(&content))
}
pub(super) fn compute_revision(content: &str) -> String {
let mut hasher = Sha256::new();
hasher.update(content.as_bytes());
@@ -86,6 +93,14 @@ pub(super) async fn load_config_from_disk(config_path: &Path) -> Result<ProxyCon
.map_err(|e| ApiFailure::internal(format!("failed to load config: {}", e)))
}
pub(super) async fn load_config_for_reload(config_path: &Path) -> Result<ProxyConfig, ApiFailure> {
let config_path = config_path.to_path_buf();
tokio::task::spawn_blocking(move || ProxyConfig::load(config_path))
.await
.map_err(|error| ApiFailure::internal(format!("failed to join config loader: {}", error)))?
.map_err(|error| ApiFailure::bad_request(format!("invalid runtime config: {}", error)))
}
#[allow(dead_code)]
pub(super) async fn save_config_to_disk(
config_path: &Path,

View File

@@ -7,6 +7,7 @@ use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::time::Duration;
use arc_swap::ArcSwap;
use http_body_util::Full;
use hyper::body::{Bytes, Incoming};
use hyper::header::AUTHORIZATION;
@@ -19,8 +20,10 @@ use tokio::sync::{Mutex, RwLock, Semaphore, watch};
use tokio::time::timeout;
use tracing::{debug, info, warn};
use crate::config::{ApiGrayAction, ProxyConfig};
use crate::config::ApiGrayAction;
use crate::ip_tracker::UserIpTracker;
use crate::maestro::generation::{RuntimeGeneration, RuntimeWatchState};
use crate::maestro::reload::{ReloadControl, ReloadRequest, ReloadSubmitError};
use crate::proxy::route_mode::RouteRuntimeController;
use crate::proxy::shared_state::ProxySharedState;
use crate::startup::StartupTracker;
@@ -29,7 +32,7 @@ use crate::transport::UpstreamManager;
use crate::transport::middle_proxy::MePool;
mod config_edit;
mod config_store;
pub(crate) mod config_store;
mod events;
mod http_utils;
mod model;
@@ -44,7 +47,8 @@ mod runtime_zero;
mod users;
use config_store::{
current_revision, ensure_expected_revision, load_config_from_disk, parse_if_match,
current_revision, ensure_expected_revision, load_config_for_reload, load_config_from_disk,
parse_if_match,
};
use events::ApiEventStore;
use http_utils::{error_response, read_json, read_optional_json, success_response};
@@ -107,12 +111,15 @@ pub(super) struct ApiShared {
pub(super) minimal_cache: Arc<Mutex<Option<MinimalCacheEntry>>>,
pub(super) runtime_edge_connections_cache: Arc<Mutex<Option<EdgeConnectionsCacheEntry>>>,
pub(super) runtime_edge_recompute_lock: Arc<Mutex<()>>,
pub(super) cache_generation: Arc<AtomicU64>,
pub(super) runtime_events: Arc<ApiEventStore>,
pub(super) request_id: Arc<AtomicU64>,
pub(super) runtime_state: Arc<ApiRuntimeState>,
pub(super) startup_tracker: Arc<StartupTracker>,
pub(super) route_runtime: Arc<RouteRuntimeController>,
pub(super) proxy_shared: Arc<ProxySharedState>,
pub(super) reload_control: ReloadControl,
pub(super) active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
}
impl ApiShared {
@@ -123,6 +130,31 @@ impl ApiShared {
fn detected_link_ips(&self) -> (Option<IpAddr>, Option<IpAddr>) {
*self.detected_ips_rx.borrow()
}
fn for_runtime(&self, runtime: &RuntimeGeneration) -> Self {
Self {
stats: runtime.stats.clone(),
ip_tracker: runtime.ip_tracker.clone(),
me_pool: runtime.me_pool_runtime.clone(),
upstream_manager: runtime.upstream_manager.clone(),
config_path: self.config_path.clone(),
quota_state_path: self.quota_state_path.clone(),
detected_ips_rx: self.detected_ips_rx.clone(),
mutation_lock: self.mutation_lock.clone(),
minimal_cache: self.minimal_cache.clone(),
runtime_edge_connections_cache: self.runtime_edge_connections_cache.clone(),
runtime_edge_recompute_lock: self.runtime_edge_recompute_lock.clone(),
cache_generation: self.cache_generation.clone(),
runtime_events: self.runtime_events.clone(),
request_id: self.request_id.clone(),
runtime_state: self.runtime_state.clone(),
startup_tracker: self.startup_tracker.clone(),
route_runtime: runtime.route_runtime.clone(),
proxy_shared: runtime.proxy_shared.clone(),
reload_control: self.reload_control.clone(),
active_runtime: self.active_runtime.clone(),
}
}
}
fn auth_header_matches(actual: &str, expected: &str) -> bool {
@@ -144,6 +176,12 @@ fn user_action_route_matches(path: &str, suffix: &str) -> bool {
.unwrap_or(false)
}
fn reload_status_route_id(path: &str) -> Option<u64> {
path.strip_prefix("/v1/system/reload/")
.filter(|id| !id.is_empty() && !id.contains('/'))
.and_then(|id| id.parse().ok())
}
fn allowed_methods_for_path(path: &str) -> Option<&'static str> {
match path {
"/v1/health"
@@ -175,12 +213,14 @@ fn allowed_methods_for_path(path: &str) -> Option<&'static str> {
| "/v1/stats/users/active-ips"
| "/v1/stats/users/quota"
| "/v1/stats/users" => Some(ALLOW_GET),
"/v1/system/reload" => Some(ALLOW_POST),
"/v1/users" => Some(ALLOW_GET_POST),
"/v1/config" => Some(ALLOW_GET_PATCH),
_ if user_action_route_matches(path, "/reset-quota") => Some(ALLOW_POST),
_ if user_action_route_matches(path, "/rotate-secret") => Some(ALLOW_POST),
_ if user_action_route_matches(path, "/enable") => Some(ALLOW_POST),
_ if user_action_route_matches(path, "/disable") => Some(ALLOW_POST),
_ if reload_status_route_id(path).is_some() => Some(ALLOW_GET),
_ if path
.strip_prefix("/v1/users/")
.map(|user| !user.is_empty() && !user.contains('/'))
@@ -200,14 +240,35 @@ pub async fn serve(
route_runtime: Arc<RouteRuntimeController>,
proxy_shared: Arc<ProxySharedState>,
upstream_manager: Arc<UpstreamManager>,
config_rx: watch::Receiver<Arc<ProxyConfig>>,
admission_rx: watch::Receiver<bool>,
config_path: PathBuf,
quota_state_path: PathBuf,
detected_ips_rx: watch::Receiver<(Option<IpAddr>, Option<IpAddr>)>,
process_started_at_epoch_secs: u64,
startup_tracker: Arc<StartupTracker>,
reload_control: ReloadControl,
mut active_runtime_rx: watch::Receiver<Option<Arc<ArcSwap<RuntimeGeneration>>>>,
mut runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
) {
let active_runtime = loop {
if let Some(active_runtime) = active_runtime_rx.borrow().clone() {
break active_runtime;
}
if active_runtime_rx.changed().await.is_err() {
warn!("Runtime generation channel closed before API bootstrap");
return;
}
};
let initial_watch_state = loop {
if let Some(watch_state) = runtime_watch_rx.borrow().clone() {
break watch_state;
}
if runtime_watch_rx.changed().await.is_err() {
warn!("Runtime watch channel closed before API bootstrap");
return;
}
};
let config_rx = initial_watch_state.config_rx.clone();
let admission_rx = initial_watch_state.admission_rx.clone();
let listener = match TcpListener::bind(listen).await {
Ok(listener) => listener,
Err(error) => {
@@ -241,6 +302,7 @@ pub async fn serve(
minimal_cache: Arc::new(Mutex::new(None)),
runtime_edge_connections_cache: Arc::new(Mutex::new(None)),
runtime_edge_recompute_lock: Arc::new(Mutex::new(())),
cache_generation: Arc::new(AtomicU64::new(1)),
runtime_events: Arc::new(ApiEventStore::new(
config_rx.borrow().server.api.runtime_edge_events_capacity,
)),
@@ -249,11 +311,12 @@ pub async fn serve(
startup_tracker,
route_runtime,
proxy_shared,
reload_control,
active_runtime,
});
spawn_runtime_watchers(
config_rx.clone(),
admission_rx.clone(),
runtime_watch_rx,
runtime_state.clone(),
shared.runtime_events.clone(),
);
@@ -282,13 +345,11 @@ pub async fn serve(
};
let shared_conn = shared.clone();
let config_rx_conn = config_rx.clone();
tokio::spawn(async move {
let _connection_permit = connection_permit;
let svc = service_fn(move |req: Request<Incoming>| {
let shared_req = shared_conn.clone();
let config_rx_req = config_rx_conn.clone();
async move { handle(req, peer, shared_req, config_rx_req).await }
async move { handle(req, peer, shared_req).await }
});
match timeout(
API_HTTP_CONNECTION_TIMEOUT,
@@ -318,8 +379,19 @@ async fn handle(
req: Request<Incoming>,
peer: SocketAddr,
shared: Arc<ApiShared>,
config_rx: watch::Receiver<Arc<ProxyConfig>>,
) -> Result<Response<Full<Bytes>>, IoError> {
let runtime = shared.active_runtime.load_full();
let previous_cache_generation = shared.cache_generation.swap(runtime.id, Ordering::AcqRel);
if previous_cache_generation != runtime.id {
*shared.minimal_cache.lock().await = None;
*shared.runtime_edge_connections_cache.lock().await = None;
}
let shared = Arc::new(shared.for_runtime(runtime.as_ref()));
let config_rx = runtime.config_rx.clone();
shared
.runtime_state
.admission_open
.store(*runtime.admission_rx.borrow(), Ordering::Relaxed);
let request_id = shared.next_request_id();
let cfg = config_rx.borrow().clone();
let api_cfg = &cfg.server.api;
@@ -651,6 +723,45 @@ async fn handle(
config_edit::read_managed_config(&shared.config_path).await?;
Ok(success_response(StatusCode::OK, value, revision))
}
("POST", "/v1/system/reload") => {
if api_cfg.read_only {
return Ok(error_response(
request_id,
ApiFailure::new(
StatusCode::FORBIDDEN,
"read_only",
"API runs in read-only mode",
),
));
}
let expected_revision = parse_if_match(req.headers());
let request = read_optional_json::<ReloadRequest>(req.into_body(), body_limit)
.await?
.unwrap_or_default();
request.validate().map_err(ApiFailure::bad_request)?;
let _guard = shared.mutation_lock.lock().await;
ensure_expected_revision(&shared.config_path, expected_revision.as_deref()).await?;
let revision = current_revision(&shared.config_path).await?;
let config = Arc::new(load_config_for_reload(&shared.config_path).await?);
let accepted = shared
.reload_control
.submit(config, revision.clone(), request)
.await
.map_err(|error| match error {
ReloadSubmitError::InProgress(reload_id) => ApiFailure::new(
StatusCode::CONFLICT,
"reload_in_progress",
format!("Reload {} is already in progress", reload_id),
),
ReloadSubmitError::MaestroUnavailable => ApiFailure::new(
StatusCode::SERVICE_UNAVAILABLE,
"maestro_unavailable",
"Maestro reload coordinator is unavailable",
),
})?;
Ok(success_response(StatusCode::ACCEPTED, accepted, revision))
}
("PATCH", "/v1/config") => {
if api_cfg.read_only {
return Ok(error_response(
@@ -663,11 +774,20 @@ async fn handle(
));
}
let expected_revision = parse_if_match(req.headers());
let reload_request =
ReloadRequest::from_query(query.as_deref()).map_err(ApiFailure::bad_request)?;
let body = read_json::<serde_json::Value>(req.into_body(), body_limit).await?;
match config_edit::patch_config(body, expected_revision, &shared).await {
match config_edit::patch_config(body, expected_revision, reload_request, &shared)
.await
{
Ok(resp) => {
let revision = resp.revision.clone();
Ok(success_response(StatusCode::OK, resp, revision))
let status = if resp.reload.is_some() {
StatusCode::ACCEPTED
} else {
StatusCode::OK
};
Ok(success_response(status, resp, revision))
}
Err(error) => {
shared
@@ -678,6 +798,24 @@ async fn handle(
}
}
_ => {
if method == Method::GET
&& let Some(reload_id) = reload_status_route_id(normalized_path)
{
let revision = current_revision(&shared.config_path).await?;
let status =
shared
.reload_control
.status(reload_id)
.await
.ok_or_else(|| {
ApiFailure::new(
StatusCode::NOT_FOUND,
"reload_not_found",
format!("Reload {} was not found", reload_id),
)
})?;
return Ok(success_response(StatusCode::OK, status, revision));
}
if method == Method::POST
&& let Some(base_user) = normalized_path
.strip_prefix("/v1/users/")

View File

@@ -4,63 +4,283 @@ use std::time::{SystemTime, UNIX_EPOCH};
use tokio::sync::watch;
use crate::config::ProxyConfig;
use crate::maestro::generation::RuntimeWatchState;
use super::ApiRuntimeState;
use super::events::ApiEventStore;
pub(super) fn spawn_runtime_watchers(
config_rx: watch::Receiver<Arc<ProxyConfig>>,
admission_rx: watch::Receiver<bool>,
runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
runtime_state: Arc<ApiRuntimeState>,
runtime_events: Arc<ApiEventStore>,
) {
let mut config_rx_reload = config_rx;
let runtime_state_reload = runtime_state.clone();
let runtime_events_reload = runtime_events.clone();
tokio::spawn(async move {
loop {
if config_rx_reload.changed().await.is_err() {
break;
}
runtime_state_reload
.config_reload_count
.fetch_add(1, Ordering::Relaxed);
runtime_state_reload
.last_config_reload_epoch_secs
.store(now_epoch_secs(), Ordering::Relaxed);
runtime_events_reload.record("config.reload.applied", "config receiver updated");
}
});
spawn_config_watcher(
runtime_watch_rx.clone(),
runtime_state.clone(),
runtime_events.clone(),
);
spawn_admission_watcher(runtime_watch_rx, runtime_state, runtime_events);
}
let mut admission_rx_watch = admission_rx;
fn spawn_config_watcher(
mut runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
runtime_state: Arc<ApiRuntimeState>,
runtime_events: Arc<ApiEventStore>,
) {
tokio::spawn(async move {
runtime_state
.admission_open
.store(*admission_rx_watch.borrow(), Ordering::Relaxed);
runtime_events.record(
"admission.state",
format!("accepting_new_connections={}", *admission_rx_watch.borrow()),
);
let Some(mut current) = runtime_watch_rx.borrow().clone() else {
return;
};
loop {
if admission_rx_watch.changed().await.is_err() {
break;
tokio::select! {
biased;
changed = runtime_watch_rx.changed() => {
if changed.is_err() {
break;
}
let Some(next) = runtime_watch_rx.borrow().clone() else {
continue;
};
if next.generation_id != current.generation_id {
current = next;
record_config_reload(
&runtime_state,
&runtime_events,
format!("runtime generation {} activated", current.generation_id),
);
}
}
changed = current.config_rx.changed() => {
if changed.is_err() {
let Some(next) = wait_for_new_generation(
&mut runtime_watch_rx,
current.generation_id,
).await else {
break;
};
current = next;
record_config_reload(
&runtime_state,
&runtime_events,
format!("runtime generation {} activated", current.generation_id),
);
continue;
}
if active_generation_id(&runtime_watch_rx) != Some(current.generation_id) {
continue;
}
record_config_reload(
&runtime_state,
&runtime_events,
format!("generation {} config receiver updated", current.generation_id),
);
}
}
let admission_open = *admission_rx_watch.borrow();
runtime_state
.admission_open
.store(admission_open, Ordering::Relaxed);
runtime_events.record(
"admission.state",
format!("accepting_new_connections={}", admission_open),
);
}
});
}
fn spawn_admission_watcher(
mut runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
runtime_state: Arc<ApiRuntimeState>,
runtime_events: Arc<ApiEventStore>,
) {
tokio::spawn(async move {
let Some(mut current) = runtime_watch_rx.borrow().clone() else {
return;
};
record_admission_state(&runtime_state, &runtime_events, &current);
loop {
tokio::select! {
biased;
changed = runtime_watch_rx.changed() => {
if changed.is_err() {
break;
}
let Some(next) = runtime_watch_rx.borrow().clone() else {
continue;
};
if next.generation_id != current.generation_id {
current = next;
record_admission_state(&runtime_state, &runtime_events, &current);
}
}
changed = current.admission_rx.changed() => {
if changed.is_err() {
let Some(next) = wait_for_new_generation(
&mut runtime_watch_rx,
current.generation_id,
).await else {
break;
};
current = next;
record_admission_state(&runtime_state, &runtime_events, &current);
continue;
}
if active_generation_id(&runtime_watch_rx) == Some(current.generation_id) {
record_admission_state(&runtime_state, &runtime_events, &current);
}
}
}
}
});
}
fn active_generation_id(
runtime_watch_rx: &watch::Receiver<Option<RuntimeWatchState>>,
) -> Option<u64> {
runtime_watch_rx
.borrow()
.as_ref()
.map(|state| state.generation_id)
}
async fn wait_for_new_generation(
runtime_watch_rx: &mut watch::Receiver<Option<RuntimeWatchState>>,
previous_generation_id: u64,
) -> Option<RuntimeWatchState> {
loop {
if let Some(state) = runtime_watch_rx.borrow().clone()
&& state.generation_id != previous_generation_id
{
return Some(state);
}
if runtime_watch_rx.changed().await.is_err() {
return None;
}
}
}
fn record_config_reload(
runtime_state: &ApiRuntimeState,
runtime_events: &ApiEventStore,
context: String,
) {
runtime_state
.config_reload_count
.fetch_add(1, Ordering::Relaxed);
runtime_state
.last_config_reload_epoch_secs
.store(now_epoch_secs(), Ordering::Relaxed);
runtime_events.record("config.reload.applied", context);
}
fn record_admission_state(
runtime_state: &ApiRuntimeState,
runtime_events: &ApiEventStore,
current: &RuntimeWatchState,
) {
let admission_open = *current.admission_rx.borrow();
runtime_state
.admission_open
.store(admission_open, Ordering::Relaxed);
runtime_events.record(
"admission.state",
format!(
"generation={} accepting_new_connections={}",
current.generation_id, admission_open
),
);
}
fn now_epoch_secs() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs()
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::ProxyConfig;
use std::sync::atomic::{AtomicBool, AtomicU64};
use std::time::Duration;
fn state(
generation_id: u64,
) -> (
RuntimeWatchState,
watch::Sender<Arc<ProxyConfig>>,
watch::Sender<bool>,
) {
let (config_tx, config_rx) = watch::channel(Arc::new(ProxyConfig::default()));
let (admission_tx, admission_rx) = watch::channel(true);
(
RuntimeWatchState {
generation_id,
config_rx,
admission_rx,
},
config_tx,
admission_tx,
)
}
fn runtime_state() -> Arc<ApiRuntimeState> {
Arc::new(ApiRuntimeState {
process_started_at_epoch_secs: 1,
config_reload_count: AtomicU64::new(0),
last_config_reload_epoch_secs: AtomicU64::new(0),
admission_open: AtomicBool::new(false),
})
}
async fn wait_for_count(runtime_state: &ApiRuntimeState, expected: u64) {
tokio::time::timeout(Duration::from_secs(1), async {
loop {
if runtime_state.config_reload_count.load(Ordering::Relaxed) == expected {
break;
}
tokio::task::yield_now().await;
}
})
.await
.unwrap();
}
#[tokio::test]
async fn watchers_follow_only_the_active_generation() {
let (initial, initial_config_tx, initial_admission_tx) = state(1);
let (runtime_watch_tx, runtime_watch_rx) = watch::channel(Some(initial));
let runtime_state = runtime_state();
let events = Arc::new(ApiEventStore::new(16));
spawn_runtime_watchers(runtime_watch_rx, runtime_state.clone(), events.clone());
tokio::task::yield_now().await;
assert_eq!(runtime_state.config_reload_count.load(Ordering::Relaxed), 0);
initial_config_tx.send_replace(Arc::new(ProxyConfig::default()));
wait_for_count(&runtime_state, 1).await;
let (next, next_config_tx, next_admission_tx) = state(2);
runtime_watch_tx.send_replace(Some(next));
wait_for_count(&runtime_state, 2).await;
initial_config_tx.send_replace(Arc::new(ProxyConfig::default()));
initial_admission_tx.send_replace(false);
tokio::task::yield_now().await;
assert_eq!(runtime_state.config_reload_count.load(Ordering::Relaxed), 2);
assert!(runtime_state.admission_open.load(Ordering::Relaxed));
next_config_tx.send_replace(Arc::new(ProxyConfig::default()));
next_admission_tx.send_replace(false);
wait_for_count(&runtime_state, 3).await;
tokio::time::timeout(Duration::from_secs(1), async {
while runtime_state.admission_open.load(Ordering::Relaxed) {
tokio::task::yield_now().await;
}
})
.await
.unwrap();
let snapshot = events.snapshot(16);
assert_eq!(
snapshot
.events
.iter()
.filter(|event| event.event_type == "config.reload.applied")
.count(),
3
);
}
}

View File

@@ -1,4 +1,5 @@
use std::sync::atomic::Ordering;
use std::time::{SystemTime, UNIX_EPOCH};
use serde::Serialize;
@@ -162,7 +163,7 @@ pub(super) fn build_system_info_data(
build_time_utc,
rustc_version,
process_started_at_epoch_secs: shared.runtime_state.process_started_at_epoch_secs,
uptime_seconds: shared.stats.uptime_secs(),
uptime_seconds: process_uptime_seconds(shared.runtime_state.process_started_at_epoch_secs),
config_path: shared.config_path.display().to_string(),
config_hash: revision.to_string(),
config_reload_count: shared
@@ -173,6 +174,18 @@ pub(super) fn build_system_info_data(
}
}
fn process_uptime_seconds(process_started_at_epoch_secs: u64) -> f64 {
let now_epoch_secs = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
process_uptime_seconds_at(process_started_at_epoch_secs, now_epoch_secs)
}
fn process_uptime_seconds_at(process_started_at_epoch_secs: u64, now_epoch_secs: u64) -> f64 {
now_epoch_secs.saturating_sub(process_started_at_epoch_secs) as f64
}
pub(super) async fn build_runtime_gates_data(
shared: &ApiShared,
cfg: &ProxyConfig,
@@ -339,3 +352,14 @@ fn me_writer_pick_mode_label(mode: MeWriterPickMode) -> &'static str {
MeWriterPickMode::P2c => "p2c",
}
}
#[cfg(test)]
mod tests {
use super::process_uptime_seconds_at;
#[test]
fn process_uptime_is_monotonic_and_saturating() {
assert_eq!(process_uptime_seconds_at(100, 135), 35.0);
assert_eq!(process_uptime_seconds_at(135, 100), 0.0);
}
}

View File

@@ -1403,11 +1403,13 @@ fn reload_config(
/// `detected_ip_v4` / `detected_ip_v6` are the IPs discovered during the
/// startup probe — used when generating proxy links for newly added users,
/// matching the same logic as the startup output.
/// The watcher releases its notify and signal resources when `cancellation` fires.
pub fn spawn_config_watcher(
config_path: PathBuf,
initial: Arc<ProxyConfig>,
detected_ip_v4: Option<IpAddr>,
detected_ip_v6: Option<IpAddr>,
cancellation: tokio_util::sync::CancellationToken,
) -> (watch::Receiver<Arc<ProxyConfig>>, watch::Receiver<LogLevel>) {
let initial_level = initial.general.log_level.clone();
let (config_tx, config_rx) = watch::channel(initial);
@@ -1515,10 +1517,14 @@ pub fn spawn_config_watcher(
_ = sighup.recv() => {
info!("SIGHUP received — reloading {:?}", config_path);
}
_ = cancellation.cancelled() => break,
}
#[cfg(not(unix))]
if notify_rx.recv().await.is_none() {
break;
tokio::select! {
msg = notify_rx.recv() => {
if msg.is_none() { break; }
}
_ = cancellation.cancelled() => break,
}
// Debounce: drain extra events that arrive within a short quiet window.

View File

@@ -7,6 +7,7 @@ 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 crate::config::{ConntrackBackend, ConntrackMode, ProxyConfig};
@@ -57,10 +58,11 @@ impl PressureState {
}
}
pub(crate) fn spawn_conntrack_controller(
pub(crate) async fn run_conntrack_controller(
config_rx: watch::Receiver<Arc<ProxyConfig>>,
stats: Arc<Stats>,
shared: Arc<ProxySharedState>,
cancellation: CancellationToken,
) {
if !cfg!(target_os = "linux") {
let cfg = config_rx.borrow();
@@ -87,16 +89,15 @@ pub(crate) fn spawn_conntrack_controller(
let (tx, rx) = mpsc::channel(CONNTRACK_EVENT_QUEUE_CAPACITY);
shared.set_conntrack_close_sender(tx);
tokio::spawn(async move {
run_conntrack_controller(config_rx, stats, shared, rx).await;
});
run_conntrack_controller_worker(config_rx, stats, shared, rx, cancellation).await;
}
async fn run_conntrack_controller(
async fn run_conntrack_controller_worker(
mut config_rx: watch::Receiver<Arc<ProxyConfig>>,
stats: Arc<Stats>,
shared: Arc<ProxySharedState>,
mut close_rx: mpsc::Receiver<ConntrackCloseEvent>,
cancellation: CancellationToken,
) {
let mut cfg = config_rx.borrow().clone();
let mut pressure_state = PressureState::new(stats.as_ref());
@@ -115,6 +116,7 @@ async fn run_conntrack_controller(
loop {
tokio::select! {
_ = cancellation.cancelled() => break,
changed = config_rx.changed() => {
if changed.is_err() {
break;

View File

@@ -8,6 +8,8 @@ use crate::config::ProxyConfig;
use crate::proxy::route_mode::{RelayRouteMode, RouteRuntimeController};
use crate::transport::middle_proxy::MePool;
use super::generation::RuntimeTaskScope;
const STARTUP_FALLBACK_AFTER: Duration = Duration::from_secs(80);
const RUNTIME_FALLBACK_AFTER: Duration = Duration::from_secs(6);
@@ -19,6 +21,7 @@ pub(crate) async fn configure_admission_gate(
admission_tx: &watch::Sender<bool>,
config_rx: watch::Receiver<Arc<ProxyConfig>>,
me_ready_rx: watch::Receiver<u64>,
task_scope: RuntimeTaskScope,
) {
if config.general.use_middle_proxy {
if me_pool.is_some() || config.general.me2dc_fallback {
@@ -64,7 +67,7 @@ pub(crate) async fn configure_admission_gate(
let mut config_rx_gate = config_rx.clone();
let mut me_ready_rx_gate = me_ready_rx;
let mut admission_poll_ms = config.general.me_admission_poll_ms.max(1);
tokio::spawn(async move {
task_scope.spawn(async move {
let mut gate_open = initial_gate_open;
let mut route_mode = initial_route_mode;
let mut ready_observed = initial_ready;

238
src/maestro/generation.rs Normal file
View File

@@ -0,0 +1,238 @@
use std::future::Future;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use tokio::sync::{RwLock, Semaphore, watch};
use tokio_util::sync::CancellationToken;
use tokio_util::task::TaskTracker;
use crate::config::ProxyConfig;
use crate::crypto::SecureRandom;
use crate::ip_tracker::UserIpTracker;
use crate::proxy::route_mode::RouteRuntimeController;
use crate::proxy::shared_state::ProxySharedState;
use crate::stats::beobachten::BeobachtenStore;
use crate::stats::{ReplayChecker, Stats};
use crate::stream::BufferPool;
use crate::tls_front::TlsFrontCache;
use crate::transport::UpstreamManager;
use crate::transport::middle_proxy::MePool;
const SESSION_STOP_TIMEOUT: Duration = Duration::from_secs(5);
const BACKGROUND_STOP_TIMEOUT: Duration = Duration::from_secs(5);
/// Process-visible control-plane receivers for one active runtime generation.
#[derive(Clone)]
pub(crate) struct RuntimeWatchState {
pub(crate) generation_id: u64,
pub(crate) config_rx: watch::Receiver<Arc<ProxyConfig>>,
pub(crate) admission_rx: watch::Receiver<bool>,
}
/// Cancellation and join ownership for one generation's background tasks.
#[derive(Clone)]
pub(crate) struct RuntimeTaskScope {
tracker: TaskTracker,
cancel: CancellationToken,
}
impl RuntimeTaskScope {
pub(crate) fn new() -> Self {
Self {
tracker: TaskTracker::new(),
cancel: CancellationToken::new(),
}
}
pub(crate) fn spawn<F>(&self, future: F)
where
F: Future<Output = ()> + Send + 'static,
{
let cancel = self.cancel.clone();
self.tracker.spawn(async move {
tokio::select! {
_ = cancel.cancelled() => {}
_ = future => {}
}
});
}
pub(crate) fn cancellation_token(&self) -> CancellationToken {
self.cancel.clone()
}
pub(crate) async fn stop(&self) {
self.cancel.cancel();
self.tracker.close();
let _ = tokio::time::timeout(BACKGROUND_STOP_TIMEOUT, self.tracker.wait()).await;
}
}
/// Runtime-owned data plane and control-plane dependencies for one generation.
pub(crate) struct RuntimeGeneration {
pub(crate) id: u64,
pub(crate) config_rx: watch::Receiver<Arc<ProxyConfig>>,
pub(crate) admission_rx: watch::Receiver<bool>,
pub(crate) stats: Arc<Stats>,
pub(crate) upstream_manager: Arc<UpstreamManager>,
pub(crate) replay_checker: Arc<ReplayChecker>,
pub(crate) buffer_pool: Arc<BufferPool>,
pub(crate) rng: Arc<SecureRandom>,
pub(crate) me_pool: Option<Arc<MePool>>,
pub(crate) me_pool_runtime: Arc<RwLock<Option<Arc<MePool>>>>,
pub(crate) route_runtime: Arc<RouteRuntimeController>,
pub(crate) tls_cache: Option<Arc<TlsFrontCache>>,
pub(crate) ip_tracker: Arc<UserIpTracker>,
pub(crate) beobachten: Arc<BeobachtenStore>,
pub(crate) proxy_shared: Arc<ProxySharedState>,
pub(crate) max_connections: Arc<Semaphore>,
background_tasks: RuntimeTaskScope,
sessions: TaskTracker,
session_cancel: CancellationToken,
accepting_sessions: AtomicBool,
}
impl RuntimeGeneration {
#[allow(clippy::too_many_arguments)]
pub(crate) fn new(
id: u64,
config_rx: watch::Receiver<Arc<ProxyConfig>>,
admission_rx: watch::Receiver<bool>,
stats: Arc<Stats>,
upstream_manager: Arc<UpstreamManager>,
replay_checker: Arc<ReplayChecker>,
buffer_pool: Arc<BufferPool>,
rng: Arc<SecureRandom>,
me_pool: Option<Arc<MePool>>,
me_pool_runtime: Arc<RwLock<Option<Arc<MePool>>>>,
route_runtime: Arc<RouteRuntimeController>,
tls_cache: Option<Arc<TlsFrontCache>>,
ip_tracker: Arc<UserIpTracker>,
beobachten: Arc<BeobachtenStore>,
proxy_shared: Arc<ProxySharedState>,
max_connections: Arc<Semaphore>,
background_tasks: RuntimeTaskScope,
) -> Arc<Self> {
Arc::new(Self {
id,
config_rx,
admission_rx,
stats,
upstream_manager,
replay_checker,
buffer_pool,
rng,
me_pool,
me_pool_runtime,
route_runtime,
tls_cache,
ip_tracker,
beobachten,
proxy_shared,
max_connections,
background_tasks,
sessions: TaskTracker::new(),
session_cancel: CancellationToken::new(),
accepting_sessions: AtomicBool::new(true),
})
}
pub(crate) fn config(&self) -> Arc<ProxyConfig> {
self.config_rx.borrow().clone()
}
/// Returns receivers used by process-scoped observers of this generation.
pub(crate) fn watch_state(&self) -> RuntimeWatchState {
RuntimeWatchState {
generation_id: self.id,
config_rx: self.config_rx.clone(),
admission_rx: self.admission_rx.clone(),
}
}
pub(crate) async fn current_me_pool(&self) -> Option<Arc<MePool>> {
if let Some(pool) = &self.me_pool {
return Some(pool.clone());
}
self.me_pool_runtime.read().await.clone()
}
pub(crate) fn spawn_session<F>(&self, future: F) -> bool
where
F: Future<Output = ()> + Send + 'static,
{
if !self.accepting_sessions.load(Ordering::Acquire) {
return false;
}
let cancel = self.session_cancel.clone();
self.sessions.spawn(async move {
tokio::select! {
_ = cancel.cancelled() => {}
_ = future => {}
}
});
true
}
pub(crate) fn stop_accepting_sessions(&self) {
self.accepting_sessions.store(false, Ordering::Release);
}
pub(crate) fn resume_accepting_sessions(&self) {
self.accepting_sessions.store(true, Ordering::Release);
}
pub(crate) async fn drain_sessions(&self, timeout: Duration) -> bool {
self.stop_accepting_sessions();
self.sessions.close();
if tokio::time::timeout(timeout, self.sessions.wait())
.await
.is_ok()
{
return true;
}
self.stop_sessions().await;
false
}
pub(crate) async fn stop_sessions(&self) {
self.stop_accepting_sessions();
self.session_cancel.cancel();
self.sessions.close();
let _ = tokio::time::timeout(SESSION_STOP_TIMEOUT, self.sessions.wait()).await;
}
pub(crate) async fn stop_background_tasks(&self) {
self.background_tasks.stop().await;
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn stop_sessions_cancels_tracked_future() {
let tracker = TaskTracker::new();
let cancel = CancellationToken::new();
let child_cancel = cancel.clone();
tracker.spawn(async move {
child_cancel.cancelled().await;
});
cancel.cancel();
tracker.close();
tokio::time::timeout(Duration::from_secs(1), tracker.wait())
.await
.unwrap();
}
#[tokio::test]
async fn runtime_task_scope_joins_cancelled_background_task() {
let scope = RuntimeTaskScope::new();
scope.spawn(std::future::pending());
tokio::time::timeout(Duration::from_secs(1), scope.stop())
.await
.unwrap();
}
}

View File

@@ -3,35 +3,33 @@ use std::net::{IpAddr, SocketAddr};
use std::sync::Arc;
use std::time::Duration;
use arc_swap::ArcSwap;
use tokio::net::TcpListener;
#[cfg(unix)]
use tokio::net::UnixListener;
use tokio::sync::{RwLock, Semaphore, watch};
use tracing::{debug, error, info, warn};
use crate::config::{ProxyConfig, RstOnCloseMode};
use crate::crypto::SecureRandom;
use crate::ip_tracker::UserIpTracker;
use crate::proxy::ClientHandler;
use crate::proxy::route_mode::RouteRuntimeController;
use crate::proxy::shared_state::ProxySharedState;
use crate::startup::{COMPONENT_LISTENERS_BIND, StartupTracker};
use crate::stats::beobachten::BeobachtenStore;
use crate::stats::{ReplayChecker, Stats};
use crate::stream::BufferPool;
use crate::tls_front::TlsFrontCache;
use crate::transport::middle_proxy::MePool;
use crate::transport::socket::set_linger_zero;
use crate::transport::{ListenOptions, UpstreamManager, create_listener, find_listener_processes};
use crate::transport::{ListenOptions, create_listener, find_listener_processes};
use super::generation::RuntimeGeneration;
use super::helpers::{
expected_handshake_close_description, is_expected_handshake_eof, peer_close_description,
print_proxy_links,
};
#[cfg(unix)]
mod unix;
#[cfg(unix)]
pub(crate) use unix::spawn_unix_accept_loop;
pub(crate) struct BoundListeners {
pub(crate) listeners: Vec<(TcpListener, bool)>,
pub(crate) has_unix_listener: bool,
#[cfg(unix)]
pub(crate) unix_listener: Option<UnixListener>,
}
fn listener_port_or_legacy(listener: &crate::config::ListenerConfig, config: &ProxyConfig) -> u16 {
@@ -59,21 +57,6 @@ pub(crate) async fn bind_listeners(
detected_ip_v4: Option<IpAddr>,
detected_ip_v6: Option<IpAddr>,
startup_tracker: &Arc<StartupTracker>,
config_rx: watch::Receiver<Arc<ProxyConfig>>,
admission_rx: watch::Receiver<bool>,
stats: Arc<Stats>,
upstream_manager: Arc<UpstreamManager>,
replay_checker: Arc<ReplayChecker>,
buffer_pool: Arc<BufferPool>,
rng: Arc<SecureRandom>,
me_pool: Option<Arc<MePool>>,
me_pool_runtime: Arc<RwLock<Option<Arc<MePool>>>>,
route_runtime: Arc<RouteRuntimeController>,
tls_cache: Option<Arc<TlsFrontCache>>,
ip_tracker: Arc<UserIpTracker>,
beobachten: Arc<BeobachtenStore>,
shared: Arc<ProxySharedState>,
max_connections: Arc<Semaphore>,
) -> Result<BoundListeners, Box<dyn Error>> {
startup_tracker
.start_component(
@@ -218,7 +201,8 @@ pub(crate) async fn bind_listeners(
print_proxy_links(&host, port, config);
}
let mut has_unix_listener = false;
#[cfg(unix)]
let mut unix_listener_out = None;
#[cfg(unix)]
if let Some(ref unix_path) = config.server.listen_unix_sock {
let _ = tokio::fs::remove_file(unix_path).await;
@@ -251,123 +235,14 @@ pub(crate) async fn bind_listeners(
info!("Listening on unix:{}", unix_path);
}
has_unix_listener = true;
let mut config_rx_unix: watch::Receiver<Arc<ProxyConfig>> = config_rx.clone();
let admission_rx_unix = admission_rx.clone();
let stats = stats.clone();
let upstream_manager = upstream_manager.clone();
let replay_checker = replay_checker.clone();
let buffer_pool = buffer_pool.clone();
let rng = rng.clone();
let me_pool = me_pool.clone();
let me_pool_runtime = me_pool_runtime.clone();
let route_runtime = route_runtime.clone();
let tls_cache = tls_cache.clone();
let ip_tracker = ip_tracker.clone();
let beobachten = beobachten.clone();
let shared = shared.clone();
let max_connections_unix = max_connections.clone();
tokio::spawn(async move {
let unix_conn_counter = Arc::new(std::sync::atomic::AtomicU64::new(1));
loop {
match unix_listener.accept().await {
Ok((stream, _)) => {
if !*admission_rx_unix.borrow() {
drop(stream);
continue;
}
let accept_permit_timeout_ms =
config_rx_unix.borrow().server.accept_permit_timeout_ms;
let permit = if accept_permit_timeout_ms == 0 {
match max_connections_unix.clone().acquire_owned().await {
Ok(permit) => permit,
Err(_) => {
error!("Connection limiter is closed");
break;
}
}
} else {
match tokio::time::timeout(
Duration::from_millis(accept_permit_timeout_ms),
max_connections_unix.clone().acquire_owned(),
)
.await
{
Ok(Ok(permit)) => permit,
Ok(Err(_)) => {
error!("Connection limiter is closed");
break;
}
Err(_) => {
stats.increment_accept_permit_timeout_total();
debug!(
timeout_ms = accept_permit_timeout_ms,
"Dropping accepted unix connection: permit wait timeout"
);
drop(stream);
continue;
}
}
};
let conn_id =
unix_conn_counter.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
let fake_peer =
SocketAddr::from(([127, 0, 0, 1], (conn_id % 65535) as u16));
let config = config_rx_unix.borrow_and_update().clone();
let stats = stats.clone();
let upstream_manager = upstream_manager.clone();
let replay_checker = replay_checker.clone();
let buffer_pool = buffer_pool.clone();
let rng = rng.clone();
let me_pool = me_pool.clone();
let me_pool_runtime = me_pool_runtime.clone();
let route_runtime = route_runtime.clone();
let tls_cache = tls_cache.clone();
let ip_tracker = ip_tracker.clone();
let beobachten = beobachten.clone();
let shared = shared.clone();
let proxy_protocol_enabled = config.server.proxy_protocol;
tokio::spawn(async move {
let _permit = permit;
if let Err(e) =
crate::proxy::client::handle_client_stream_with_shared_and_pool_runtime(
stream,
fake_peer,
config,
stats,
upstream_manager,
replay_checker,
buffer_pool,
rng,
me_pool,
Some(me_pool_runtime),
route_runtime,
tls_cache,
ip_tracker,
beobachten,
shared,
proxy_protocol_enabled,
)
.await
{
debug!(error = %e, "Unix socket connection error");
}
});
}
Err(e) => {
error!("Unix socket accept error: {}", e);
tokio::time::sleep(Duration::from_millis(100)).await;
}
}
}
});
unix_listener_out = Some(unix_listener);
}
#[cfg(unix)]
let has_unix_listener = unix_listener_out.is_some();
#[cfg(not(unix))]
let has_unix_listener = false;
startup_tracker
.complete_component(
COMPONENT_LISTENERS_BIND,
@@ -381,51 +256,25 @@ pub(crate) async fn bind_listeners(
Ok(BoundListeners {
listeners,
has_unix_listener,
#[cfg(unix)]
unix_listener: unix_listener_out,
})
}
#[allow(clippy::too_many_arguments)]
pub(crate) fn spawn_tcp_accept_loops(
listeners: Vec<(TcpListener, bool)>,
config_rx: watch::Receiver<Arc<ProxyConfig>>,
admission_rx: watch::Receiver<bool>,
stats: Arc<Stats>,
upstream_manager: Arc<UpstreamManager>,
replay_checker: Arc<ReplayChecker>,
buffer_pool: Arc<BufferPool>,
rng: Arc<SecureRandom>,
me_pool: Option<Arc<MePool>>,
me_pool_runtime: Arc<RwLock<Option<Arc<MePool>>>>,
route_runtime: Arc<RouteRuntimeController>,
tls_cache: Option<Arc<TlsFrontCache>>,
ip_tracker: Arc<UserIpTracker>,
beobachten: Arc<BeobachtenStore>,
shared: Arc<ProxySharedState>,
max_connections: Arc<Semaphore>,
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
) {
for (listener, listener_proxy_protocol) in listeners {
let mut config_rx: watch::Receiver<Arc<ProxyConfig>> = config_rx.clone();
let admission_rx_tcp = admission_rx.clone();
let stats = stats.clone();
let upstream_manager = upstream_manager.clone();
let replay_checker = replay_checker.clone();
let buffer_pool = buffer_pool.clone();
let rng = rng.clone();
let me_pool = me_pool.clone();
let me_pool_runtime = me_pool_runtime.clone();
let route_runtime = route_runtime.clone();
let tls_cache = tls_cache.clone();
let ip_tracker = ip_tracker.clone();
let beobachten = beobachten.clone();
let shared = shared.clone();
let max_connections_tcp = max_connections.clone();
let active_runtime = active_runtime.clone();
tokio::spawn(async move {
loop {
match listener.accept().await {
Ok((stream, peer_addr)) => {
let rst_mode = config_rx.borrow().general.rst_on_close;
let runtime = active_runtime.load_full();
let config = runtime.config();
let rst_mode = config.general.rst_on_close;
#[cfg(unix)]
let raw_fd = {
use std::os::unix::io::AsRawFd;
@@ -434,15 +283,14 @@ pub(crate) fn spawn_tcp_accept_loops(
if matches!(rst_mode, RstOnCloseMode::Errors | RstOnCloseMode::Always) {
let _ = set_linger_zero(&stream);
}
if !*admission_rx_tcp.borrow() {
if !*runtime.admission_rx.borrow() {
debug!(peer = %peer_addr, "Admission gate closed, dropping connection");
drop(stream);
continue;
}
let accept_permit_timeout_ms =
config_rx.borrow().server.accept_permit_timeout_ms;
let accept_permit_timeout_ms = config.server.accept_permit_timeout_ms;
let permit = if accept_permit_timeout_ms == 0 {
match max_connections_tcp.clone().acquire_owned().await {
match runtime.max_connections.clone().acquire_owned().await {
Ok(permit) => permit,
Err(_) => {
error!("Connection limiter is closed");
@@ -452,7 +300,7 @@ pub(crate) fn spawn_tcp_accept_loops(
} else {
match tokio::time::timeout(
Duration::from_millis(accept_permit_timeout_ms),
max_connections_tcp.clone().acquire_owned(),
runtime.max_connections.clone().acquire_owned(),
)
.await
{
@@ -462,7 +310,7 @@ pub(crate) fn spawn_tcp_accept_loops(
break;
}
Err(_) => {
stats.increment_accept_permit_timeout_total();
runtime.stats.increment_accept_permit_timeout_total();
debug!(
peer = %peer_addr,
timeout_ms = accept_permit_timeout_ms,
@@ -473,24 +321,23 @@ pub(crate) fn spawn_tcp_accept_loops(
}
}
};
let config = config_rx.borrow_and_update().clone();
let stats = stats.clone();
let upstream_manager = upstream_manager.clone();
let replay_checker = replay_checker.clone();
let buffer_pool = buffer_pool.clone();
let rng = rng.clone();
let me_pool = me_pool.clone();
let me_pool_runtime = me_pool_runtime.clone();
let route_runtime = route_runtime.clone();
let tls_cache = tls_cache.clone();
let ip_tracker = ip_tracker.clone();
let beobachten = beobachten.clone();
let shared = shared.clone();
let stats = runtime.stats.clone();
let upstream_manager = runtime.upstream_manager.clone();
let replay_checker = runtime.replay_checker.clone();
let buffer_pool = runtime.buffer_pool.clone();
let rng = runtime.rng.clone();
let me_pool = runtime.me_pool.clone();
let me_pool_runtime = runtime.me_pool_runtime.clone();
let route_runtime = runtime.route_runtime.clone();
let tls_cache = runtime.tls_cache.clone();
let ip_tracker = runtime.ip_tracker.clone();
let beobachten = runtime.beobachten.clone();
let shared = runtime.proxy_shared.clone();
let proxy_protocol_enabled = listener_proxy_protocol;
let real_peer_report = Arc::new(std::sync::Mutex::new(None));
let real_peer_report_for_handler = real_peer_report.clone();
tokio::spawn(async move {
let _ = runtime.spawn_session(async move {
let _permit = permit;
if let Err(e) = ClientHandler::new_with_shared(
stream,

View File

@@ -0,0 +1,117 @@
use std::net::SocketAddr;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
use arc_swap::ArcSwap;
use tokio::net::UnixListener;
use tracing::{debug, error};
use super::RuntimeGeneration;
pub(crate) fn spawn_unix_accept_loop(
listener: Option<UnixListener>,
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
) {
let Some(listener) = listener else {
return;
};
tokio::spawn(async move {
let connection_counter = AtomicU64::new(1);
loop {
match listener.accept().await {
Ok((stream, _)) => {
let runtime = active_runtime.load_full();
if !*runtime.admission_rx.borrow() {
drop(stream);
continue;
}
let config = runtime.config();
let timeout_ms = config.server.accept_permit_timeout_ms;
let permit = if timeout_ms == 0 {
match runtime.max_connections.clone().acquire_owned().await {
Ok(permit) => permit,
Err(_) => {
error!("Connection limiter is closed");
break;
}
}
} else {
match tokio::time::timeout(
Duration::from_millis(timeout_ms),
runtime.max_connections.clone().acquire_owned(),
)
.await
{
Ok(Ok(permit)) => permit,
Ok(Err(_)) => {
error!("Connection limiter is closed");
break;
}
Err(_) => {
runtime.stats.increment_accept_permit_timeout_total();
debug!(
timeout_ms,
"Dropping accepted unix connection: permit wait timeout"
);
drop(stream);
continue;
}
}
};
let connection_id = connection_counter.fetch_add(1, Ordering::Relaxed);
let fake_peer =
SocketAddr::from(([127, 0, 0, 1], (connection_id % 65535) as u16));
let stats = runtime.stats.clone();
let upstream_manager = runtime.upstream_manager.clone();
let replay_checker = runtime.replay_checker.clone();
let buffer_pool = runtime.buffer_pool.clone();
let rng = runtime.rng.clone();
let me_pool = runtime.me_pool.clone();
let me_pool_runtime = runtime.me_pool_runtime.clone();
let route_runtime = runtime.route_runtime.clone();
let tls_cache = runtime.tls_cache.clone();
let ip_tracker = runtime.ip_tracker.clone();
let beobachten = runtime.beobachten.clone();
let shared = runtime.proxy_shared.clone();
let proxy_protocol_enabled = config.server.proxy_protocol;
let _ = runtime.spawn_session(async move {
let _permit = permit;
if let Err(error) =
crate::proxy::client::handle_client_stream_with_shared_and_pool_runtime(
stream,
fake_peer,
config,
stats,
upstream_manager,
replay_checker,
buffer_pool,
rng,
me_pool,
Some(me_pool_runtime),
route_runtime,
tls_cache,
ip_tracker,
beobachten,
shared,
proxy_protocol_enabled,
)
.await
{
debug!(error = %error, "Unix socket connection error");
}
});
}
Err(error) => {
error!(error = %error, "Unix socket accept error");
tokio::time::sleep(Duration::from_millis(100)).await;
}
}
}
});
}

View File

@@ -1,9 +1,11 @@
#![allow(clippy::too_many_arguments)]
use std::future::Future;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::{RwLock, watch};
use tokio_util::task::AbortOnDropHandle;
use tracing::{error, info, warn};
use crate::config::ProxyConfig;
@@ -17,8 +19,58 @@ use crate::stats::Stats;
use crate::transport::UpstreamManager;
use crate::transport::middle_proxy::MePool;
use super::generation::RuntimeTaskScope;
use super::helpers::load_startup_proxy_config_snapshot;
async fn supervise_me_task<F, Fut>(task_name: &'static str, mut task: F)
where
F: FnMut() -> Fut,
Fut: Future<Output = ()> + Send + 'static,
{
loop {
let result = AbortOnDropHandle::new(tokio::spawn(task())).await;
match result {
Ok(()) => warn!(task = task_name, "Middle-End supervisor task exited unexpectedly, restarting"),
Err(error) => {
error!(task = task_name, error = %error, "Middle-End supervisor task panicked, restarting in 1s");
tokio::time::sleep(Duration::from_secs(1)).await;
}
}
}
}
fn spawn_me_supervisors(
task_scope: RuntimeTaskScope,
pool: Arc<MePool>,
rng: Arc<SecureRandom>,
min_connections: usize,
) {
let health_pool = pool.clone();
let health_rng = rng;
task_scope.spawn(supervise_me_task("health_monitor", move || {
let pool = health_pool.clone();
let rng = health_rng.clone();
async move {
crate::transport::middle_proxy::me_health_monitor(pool, rng, min_connections).await;
}
}));
let drain_pool = pool.clone();
task_scope.spawn(supervise_me_task("drain_timeout_enforcer", move || {
let pool = drain_pool.clone();
async move {
crate::transport::middle_proxy::me_drain_timeout_enforcer(pool).await;
}
}));
task_scope.spawn(supervise_me_task("zombie_writer_watchdog", move || {
let pool = pool.clone();
async move {
crate::transport::middle_proxy::me_zombie_writer_watchdog(pool).await;
}
}));
}
pub(crate) async fn initialize_me_pool(
use_middle_proxy: bool,
config: &ProxyConfig,
@@ -30,6 +82,7 @@ pub(crate) async fn initialize_me_pool(
stats: Arc<Stats>,
api_me_pool: Arc<RwLock<Option<Arc<MePool>>>>,
me_ready_tx: watch::Sender<u64>,
task_scope: RuntimeTaskScope,
) -> Option<Arc<MePool>> {
if !use_middle_proxy {
return None;
@@ -52,15 +105,8 @@ pub(crate) async fn initialize_me_pool(
.as_ref()
.map(|tag| hex::decode(tag).expect("general.ad_tag must be validated before startup"));
// =============================================================
// CRITICAL: Download Telegram proxy-secret (NOT user secret!)
//
// C MTProxy uses TWO separate secrets:
// -S flag = 16-byte user secret for client obfuscation
// --aes-pwd = 32-512 byte binary file for ME RPC auth
//
// proxy-secret is from: https://core.telegram.org/getProxySecret
// =============================================================
// The Telegram proxy-secret authenticates ME RPC and is distinct from client secrets.
// It corresponds to the C MTProxy --aes-pwd input and may be fetched from Telegram.
let proxy_secret_path = config.general.proxy_secret_path.as_deref();
let pool_size = config.general.middle_proxy_pool_size.max(1);
let proxy_secret = loop {
@@ -319,143 +365,70 @@ pub(crate) async fn initialize_me_pool(
let rng_bg = rng.clone();
let startup_tracker_bg = startup_tracker.clone();
let me_ready_tx_bg = me_ready_tx.clone();
let task_scope_bg = task_scope.clone();
let retry_limit = if me_init_retry_attempts == 0 {
String::from("unlimited")
} else {
me_init_retry_attempts.to_string()
};
std::thread::spawn(move || {
let runtime = match tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
{
Ok(runtime) => runtime,
Err(error) => {
error!(error = %error, "Failed to build background runtime for ME initialization");
return;
}
};
runtime.block_on(async move {
let mut init_attempt: u32 = 0;
loop {
init_attempt = init_attempt.saturating_add(1);
startup_tracker_bg.set_me_init_attempt(init_attempt).await;
match pool_bg.init(pool_size, &rng_bg).await {
Ok(()) => {
startup_tracker_bg.set_me_last_error(None).await;
startup_tracker_bg
.complete_component(
COMPONENT_ME_POOL_INIT_STAGE1,
Some("ME pool initialized".to_string()),
)
.await;
startup_tracker_bg
.set_me_status(StartupMeStatus::Ready, "ready")
.await;
me_ready_tx_bg.send_modify(|version| {
*version = version.saturating_add(1);
});
info!(
task_scope.spawn(async move {
let mut init_attempt: u32 = 0;
loop {
init_attempt = init_attempt.saturating_add(1);
startup_tracker_bg.set_me_init_attempt(init_attempt).await;
match pool_bg.init(pool_size, &rng_bg).await {
Ok(()) => {
startup_tracker_bg.set_me_last_error(None).await;
startup_tracker_bg
.complete_component(
COMPONENT_ME_POOL_INIT_STAGE1,
Some("ME pool initialized".to_string()),
)
.await;
startup_tracker_bg
.set_me_status(StartupMeStatus::Ready, "ready")
.await;
me_ready_tx_bg.send_modify(|version| {
*version = version.saturating_add(1);
});
info!(
attempt = init_attempt,
"Middle-End pool initialized successfully"
);
spawn_me_supervisors(
task_scope_bg,
pool_bg.clone(),
rng_bg.clone(),
pool_size,
);
break;
}
Err(e) => {
startup_tracker_bg
.set_me_last_error(Some(e.to_string()))
.await;
if init_attempt >= me_init_warn_after_attempts {
warn!(
error = %e,
attempt = init_attempt,
"Middle-End pool initialized successfully"
retry_limit = %retry_limit,
retry_in_secs = 2,
"ME pool is not ready yet; retrying background initialization"
);
} else {
info!(
error = %e,
attempt = init_attempt,
retry_limit = %retry_limit,
retry_in_secs = 2,
"ME pool startup warmup: retrying background initialization"
);
// ── Supervised background tasks ──────────────────
// Each task runs inside a nested tokio::spawn so
// that a panic is caught via JoinHandle and the
// outer loop restarts the task automatically.
let pool_health = pool_bg.clone();
let rng_health = rng_bg.clone();
let min_conns = pool_size;
tokio::spawn(async move {
loop {
let p = pool_health.clone();
let r = rng_health.clone();
let res = tokio::spawn(async move {
crate::transport::middle_proxy::me_health_monitor(
p, r, min_conns,
)
.await;
})
.await;
match res {
Ok(()) => warn!("me_health_monitor exited unexpectedly, restarting"),
Err(e) => {
error!(error = %e, "me_health_monitor panicked, restarting in 1s");
tokio::time::sleep(Duration::from_secs(1)).await;
}
}
}
});
let pool_drain_enforcer = pool_bg.clone();
tokio::spawn(async move {
loop {
let p = pool_drain_enforcer.clone();
let res = tokio::spawn(async move {
crate::transport::middle_proxy::me_drain_timeout_enforcer(p).await;
})
.await;
match res {
Ok(()) => warn!("me_drain_timeout_enforcer exited unexpectedly, restarting"),
Err(e) => {
error!(error = %e, "me_drain_timeout_enforcer panicked, restarting in 1s");
tokio::time::sleep(Duration::from_secs(1)).await;
}
}
}
});
let pool_watchdog = pool_bg.clone();
tokio::spawn(async move {
loop {
let p = pool_watchdog.clone();
let res = tokio::spawn(async move {
crate::transport::middle_proxy::me_zombie_writer_watchdog(p).await;
})
.await;
match res {
Ok(()) => warn!("me_zombie_writer_watchdog exited unexpectedly, restarting"),
Err(e) => {
error!(error = %e, "me_zombie_writer_watchdog panicked, restarting in 1s");
tokio::time::sleep(Duration::from_secs(1)).await;
}
}
}
});
// CRITICAL: keep the current-thread runtime
// alive. Without this, block_on() returns,
// the Runtime is dropped, and ALL spawned
// background tasks (health monitor, drain
// enforcer, zombie watchdog) are silently
// cancelled — causing the draining-writer
// leak that brought us here.
std::future::pending::<()>().await;
unreachable!();
}
Err(e) => {
startup_tracker_bg.set_me_last_error(Some(e.to_string())).await;
if init_attempt >= me_init_warn_after_attempts {
warn!(
error = %e,
attempt = init_attempt,
retry_limit = %retry_limit,
retry_in_secs = 2,
"ME pool is not ready yet; retrying background initialization"
);
} else {
info!(
error = %e,
attempt = init_attempt,
retry_limit = %retry_limit,
retry_in_secs = 2,
"ME pool startup warmup: retrying background initialization"
);
}
pool_bg.reset_stun_state();
tokio::time::sleep(Duration::from_secs(2)).await;
}
pool_bg.reset_stun_state();
tokio::time::sleep(Duration::from_secs(2)).await;
}
}
});
}
});
startup_tracker
.set_me_status(StartupMeStatus::Initializing, "background_init")
@@ -490,70 +463,12 @@ pub(crate) async fn initialize_me_pool(
"Middle-End pool initialized successfully"
);
// ── Supervised background tasks ──────────────────
let pool_clone = pool.clone();
let rng_clone = rng.clone();
let min_conns = pool_size;
tokio::spawn(async move {
loop {
let p = pool_clone.clone();
let r = rng_clone.clone();
let res = tokio::spawn(async move {
crate::transport::middle_proxy::me_health_monitor(
p, r, min_conns,
)
.await;
})
.await;
match res {
Ok(()) => warn!(
"me_health_monitor exited unexpectedly, restarting"
),
Err(e) => {
error!(error = %e, "me_health_monitor panicked, restarting in 1s");
tokio::time::sleep(Duration::from_secs(1)).await;
}
}
}
});
let pool_drain_enforcer = pool.clone();
tokio::spawn(async move {
loop {
let p = pool_drain_enforcer.clone();
let res = tokio::spawn(async move {
crate::transport::middle_proxy::me_drain_timeout_enforcer(p).await;
})
.await;
match res {
Ok(()) => warn!(
"me_drain_timeout_enforcer exited unexpectedly, restarting"
),
Err(e) => {
error!(error = %e, "me_drain_timeout_enforcer panicked, restarting in 1s");
tokio::time::sleep(Duration::from_secs(1)).await;
}
}
}
});
let pool_watchdog = pool.clone();
tokio::spawn(async move {
loop {
let p = pool_watchdog.clone();
let res = tokio::spawn(async move {
crate::transport::middle_proxy::me_zombie_writer_watchdog(p).await;
})
.await;
match res {
Ok(()) => warn!(
"me_zombie_writer_watchdog exited unexpectedly, restarting"
),
Err(e) => {
error!(error = %e, "me_zombie_writer_watchdog panicked, restarting in 1s");
tokio::time::sleep(Duration::from_secs(1)).await;
}
}
}
});
spawn_me_supervisors(
task_scope.clone(),
pool.clone(),
rng.clone(),
pool_size,
);
break Some(pool);
}
@@ -666,3 +581,38 @@ pub(crate) async fn initialize_me_pool(
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use tokio::sync::Notify;
struct DropSignal(Arc<Notify>);
impl Drop for DropSignal {
fn drop(&mut self) {
self.0.notify_one();
}
}
#[tokio::test]
async fn scoped_supervisor_aborts_its_current_child() {
let scope = RuntimeTaskScope::new();
let dropped = Arc::new(Notify::new());
let dropped_for_task = dropped.clone();
scope.spawn(supervise_me_task("test", move || {
let dropped = dropped_for_task.clone();
async move {
let _signal = DropSignal(dropped);
std::future::pending::<()>().await;
}
}));
tokio::task::yield_now().await;
scope.stop().await;
tokio::time::timeout(Duration::from_secs(1), dropped.notified())
.await
.unwrap();
}
}

View File

@@ -13,19 +13,24 @@
// - shutdown: graceful shutdown sequence and uptime logging.
mod admission;
mod connectivity;
pub(crate) mod generation;
mod helpers;
mod listeners;
mod me_startup;
pub(crate) mod reload;
mod reload_supervisor;
pub(crate) mod runtime_build;
mod runtime_tasks;
mod shutdown;
mod tls_bootstrap;
use arc_swap::ArcSwap;
use std::net::{IpAddr, SocketAddr};
use std::sync::Arc;
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use tokio::sync::{RwLock, Semaphore, watch};
use tracing::{error, info, warn};
use tracing_subscriber::{EnvFilter, fmt, prelude::*, reload};
use tracing_subscriber::{EnvFilter, fmt, prelude::*, reload as tracing_reload};
use crate::api;
use crate::config::{LogLevel, ProxyConfig};
@@ -34,7 +39,7 @@ use crate::crypto::SecureRandom;
use crate::ip_tracker::UserIpTracker;
use crate::network::probe::{decide_network_capabilities, log_probe_result, run_probe};
use crate::proxy::direct_buffer_budget::{
DirectBufferBudget, resolve_direct_buffer_hard_limit, spawn_direct_buffer_budget_controller,
DirectBufferBudget, resolve_direct_buffer_hard_limit, run_direct_buffer_budget_controller,
};
use crate::proxy::route_mode::{RelayRouteMode, RouteRuntimeController};
use crate::proxy::shared_state::ProxySharedState;
@@ -46,7 +51,7 @@ use crate::startup::{
};
use crate::stats::beobachten::BeobachtenStore;
use crate::stats::telemetry::TelemetryPolicy;
use crate::stats::{ReplayChecker, Stats};
use crate::stats::{QuotaStore, ReplayChecker, Stats};
use crate::stream::BufferPool;
use crate::synlimit_control;
use crate::transport::UpstreamManager;
@@ -343,7 +348,7 @@ async fn run_telemt_core(
}
};
let (filter_layer, filter_handle) =
reload::Layer::new(EnvFilter::new(initial_filter_spec.clone()));
tracing_reload::Layer::new(EnvFilter::new(initial_filter_spec.clone()));
startup_tracker
.start_component(
COMPONENT_TRACING_INIT,
@@ -387,6 +392,7 @@ async fn run_telemt_core(
_logging_guard = Some(guard);
}
}
let runtime_log_filter = runtime_tasks::RuntimeLogFilter::new(filter_handle);
startup_tracker
.complete_component(
@@ -433,21 +439,26 @@ async fn run_telemt_core(
warn!("Using default tls_domain. Consider setting a custom domain.");
}
let stats = Arc::new(Stats::new());
let quota_store = Arc::new(QuotaStore::default());
let stats = Arc::new(Stats::with_quota_store(quota_store.clone()));
let runtime_task_scope = generation::RuntimeTaskScope::new();
stats.apply_telemetry_policy(TelemetryPolicy::from_config(&config.general.telemetry));
let quota_state_path = config.general.quota_state_path.clone();
crate::quota_state::load_quota_state(&quota_state_path, stats.as_ref()).await;
let upstream_manager = Arc::new(UpstreamManager::new(
config.upstreams.clone(),
config.general.upstream_connect_retry_attempts,
config.general.upstream_connect_retry_backoff_ms,
config.general.upstream_connect_budget_ms,
config.general.tg_connect,
config.general.upstream_unhealthy_fail_threshold,
config.general.upstream_connect_failfast_hard_errors,
stats.clone(),
));
let upstream_manager = Arc::new(
UpstreamManager::new(
config.upstreams.clone(),
config.general.upstream_connect_retry_attempts,
config.general.upstream_connect_retry_backoff_ms,
config.general.upstream_connect_budget_ms,
config.general.tg_connect,
config.general.upstream_unhealthy_fail_threshold,
config.general.upstream_connect_failfast_hard_errors,
stats.clone(),
)
.with_dns_overrides(&config.network.dns_overrides)?,
);
let ip_tracker = Arc::new(UserIpTracker::new());
ip_tracker
.load_limits(
@@ -492,11 +503,15 @@ async fn run_telemt_core(
config.access.cidr_rate_limits.clone(),
);
let (api_config_tx, api_config_rx) = watch::channel(Arc::new(config.clone()));
let (detected_ips_tx, detected_ips_rx) = watch::channel((None::<IpAddr>, None::<IpAddr>));
let initial_direct_first = config.general.use_middle_proxy && config.general.me2dc_fallback;
let initial_admission_open = !config.general.use_middle_proxy || initial_direct_first;
let (admission_tx, admission_rx) = watch::channel(initial_admission_open);
let (reload_control, reload_commands) = reload::ReloadControl::channel(1);
let (active_runtime_tx, active_runtime_rx) =
watch::channel(None::<Arc<ArcSwap<generation::RuntimeGeneration>>>);
let (runtime_watch_tx, runtime_watch_rx) =
watch::channel(None::<generation::RuntimeWatchState>);
let initial_route_mode = if !config.general.use_middle_proxy || initial_direct_first {
RelayRouteMode::Direct
} else {
@@ -530,12 +545,13 @@ async fn run_telemt_core(
let upstream_manager_api = upstream_manager.clone();
let route_runtime_api = route_runtime.clone();
let proxy_shared_api = shared_state.clone();
let config_rx_api = api_config_rx.clone();
let admission_rx_api = admission_rx.clone();
let config_path_api = config_path.clone();
let quota_state_path_api = quota_state_path.clone();
let startup_tracker_api = startup_tracker.clone();
let detected_ips_rx_api = detected_ips_rx.clone();
let reload_control_api = reload_control.clone();
let active_runtime_rx_api = active_runtime_rx.clone();
let runtime_watch_rx_api = runtime_watch_rx.clone();
tokio::spawn(async move {
api::serve(
listen,
@@ -545,13 +561,14 @@ async fn run_telemt_core(
route_runtime_api,
proxy_shared_api,
upstream_manager_api,
config_rx_api,
admission_rx_api,
config_path_api,
quota_state_path_api,
detected_ips_rx_api,
process_started_at_epoch_secs,
startup_tracker_api,
reload_control_api,
active_runtime_rx_api,
runtime_watch_rx_api,
)
.await;
});
@@ -591,8 +608,10 @@ async fn run_telemt_core(
&tls_domains,
upstream_manager.clone(),
&startup_tracker,
runtime_task_scope.clone(),
tls_bootstrap::TlsBootstrapPolicy::BestEffort,
)
.await;
.await?;
startup_tracker
.start_component(
@@ -718,6 +737,7 @@ async fn run_telemt_core(
stats.clone(),
api_me_pool.clone(),
me_ready_tx.clone(),
runtime_task_scope.clone(),
)
.await
};
@@ -805,16 +825,22 @@ async fn run_telemt_core(
rng.clone(),
ip_tracker.clone(),
beobachten.clone(),
api_config_tx.clone(),
me_pool.clone(),
shared_state.clone(),
me_ready_tx.clone(),
runtime_task_scope.clone(),
)
.await;
let config_rx = runtime_watches.config_rx;
let log_level_rx = runtime_watches.log_level_rx;
let detected_ip_v4 = runtime_watches.detected_ip_v4;
let detected_ip_v6 = runtime_watches.detected_ip_v6;
runtime_log_filter.start(
has_rust_log,
&effective_log_level,
log_level_rx,
runtime_task_scope.clone(),
);
if direct_first_startup {
let config_bg = config.clone();
@@ -827,7 +853,8 @@ async fn run_telemt_core(
let api_me_pool_bg = api_me_pool.clone();
let me_ready_tx_bg = me_ready_tx.clone();
let config_rx_bg = config_rx.clone();
tokio::spawn(async move {
let task_scope_bg = runtime_task_scope.clone();
runtime_task_scope.spawn(async move {
let mut bootstrap_attempt: u32 = 0;
loop {
bootstrap_attempt = bootstrap_attempt.saturating_add(1);
@@ -842,6 +869,7 @@ async fn run_telemt_core(
stats_bg.clone(),
api_me_pool_bg.clone(),
me_ready_tx_bg.clone(),
task_scope_bg.clone(),
)
.await;
if let Some(pool) = pool {
@@ -851,6 +879,7 @@ async fn run_telemt_core(
pool,
rng_bg,
me_ready_tx_bg,
task_scope_bg,
);
break;
}
@@ -864,7 +893,7 @@ async fn run_telemt_core(
let startup_tracker_ready = startup_tracker.clone();
let api_me_pool_ready = api_me_pool.clone();
let mut me_ready_rx_transport = me_ready_tx.subscribe();
tokio::spawn(async move {
runtime_task_scope.spawn(async move {
if me_ready_rx_transport.changed().await.is_ok() {
if let Some(pool) = api_me_pool_ready.read().await.as_ref() {
pool.set_runtime_ready(true);
@@ -886,20 +915,54 @@ async fn run_telemt_core(
&admission_tx,
config_rx.clone(),
me_ready_rx,
runtime_task_scope.clone(),
)
.await;
let _admission_tx_hold = admission_tx;
conntrack_control::spawn_conntrack_controller(
let conntrack_scope = runtime_task_scope.clone();
runtime_task_scope.spawn(conntrack_control::run_conntrack_controller(
config_rx.clone(),
stats.clone(),
shared_state.clone(),
);
spawn_direct_buffer_budget_controller(
conntrack_scope.cancellation_token(),
));
runtime_task_scope.spawn(run_direct_buffer_budget_controller(
direct_buffer_budget,
buffer_pool.clone(),
stats.clone(),
shared_state.clone(),
config.server.max_connections,
));
let runtime_generation = generation::RuntimeGeneration::new(
1,
config_rx.clone(),
admission_rx.clone(),
stats.clone(),
upstream_manager.clone(),
replay_checker.clone(),
buffer_pool.clone(),
rng.clone(),
me_pool.clone(),
api_me_pool.clone(),
route_runtime.clone(),
tls_cache.clone(),
ip_tracker.clone(),
beobachten.clone(),
shared_state.clone(),
max_connections.clone(),
runtime_task_scope.clone(),
);
let active_runtime = Arc::new(ArcSwap::from(runtime_generation));
reload_supervisor::ReloadSupervisor::spawn(
active_runtime.clone(),
reload_control,
reload_commands,
config_path.clone(),
quota_store,
detected_ips_tx,
runtime_log_filter,
runtime_watch_tx.clone(),
);
let bound = listeners::bind_listeners(
@@ -909,25 +972,15 @@ async fn run_telemt_core(
detected_ip_v4,
detected_ip_v6,
&startup_tracker,
config_rx.clone(),
admission_rx.clone(),
stats.clone(),
upstream_manager.clone(),
replay_checker.clone(),
buffer_pool.clone(),
rng.clone(),
me_pool.clone(),
api_me_pool.clone(),
route_runtime.clone(),
tls_cache.clone(),
ip_tracker.clone(),
beobachten.clone(),
shared_state.clone(),
max_connections.clone(),
)
.await?;
let listeners = bound.listeners;
let has_unix_listener = bound.has_unix_listener;
#[cfg(unix)]
let unix_listener = bound.unix_listener;
#[cfg(unix)]
let has_unix_listener = unix_listener.is_some();
#[cfg(not(unix))]
let has_unix_listener = false;
if listeners.is_empty() && !has_unix_listener {
error!("No listeners. Exiting.");
@@ -937,54 +990,29 @@ async fn run_telemt_core(
// On Unix, caller supplies privilege drop after bind (may require root for port < 1024).
drop_after_bind();
synlimit_control::reconcile_synlimit_rules(&config).await;
synlimit_control::spawn_synlimit_controller(config_rx.clone());
let synlimit_controller = synlimit_control::spawn_synlimit_controller(runtime_watch_rx);
runtime_tasks::apply_runtime_log_filter(
has_rust_log,
&effective_log_level,
filter_handle,
log_level_rx,
)
.await;
runtime_tasks::spawn_metrics_if_configured(
&config,
&startup_tracker,
stats.clone(),
beobachten.clone(),
shared_state.clone(),
ip_tracker.clone(),
tls_cache.clone(),
config_rx.clone(),
)
.await;
runtime_tasks::spawn_metrics_if_configured(&config, &startup_tracker, active_runtime.clone())
.await;
runtime_watch_tx.send_replace(Some(active_runtime.load_full().watch_state()));
active_runtime_tx.send_replace(Some(active_runtime.clone()));
runtime_tasks::mark_runtime_ready(&startup_tracker).await;
// Spawn signal handlers for SIGUSR1/SIGUSR2 (non-shutdown signals)
shutdown::spawn_signal_handlers(stats.clone(), process_started_at);
shutdown::spawn_signal_handlers(active_runtime.clone(), process_started_at);
listeners::spawn_tcp_accept_loops(
listeners,
config_rx.clone(),
admission_rx.clone(),
stats.clone(),
upstream_manager.clone(),
replay_checker.clone(),
buffer_pool.clone(),
rng.clone(),
me_pool.clone(),
api_me_pool.clone(),
route_runtime.clone(),
tls_cache.clone(),
ip_tracker.clone(),
beobachten.clone(),
shared_state,
max_connections.clone(),
);
listeners::spawn_tcp_accept_loops(listeners, active_runtime.clone());
#[cfg(unix)]
listeners::spawn_unix_accept_loop(unix_listener, active_runtime.clone());
shutdown::wait_for_shutdown(process_started_at, me_pool, stats, quota_state_path).await;
shutdown::wait_for_shutdown(
process_started_at,
active_runtime,
quota_state_path,
synlimit_controller,
)
.await;
Ok(())
}

493
src/maestro/reload.rs Normal file
View File

@@ -0,0 +1,493 @@
use std::collections::VecDeque;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};
use serde::{Deserialize, Serialize};
use tokio::sync::{Mutex, mpsc};
use crate::config::ProxyConfig;
const RELOAD_HISTORY_CAPACITY: usize = 32;
const RELOAD_COMMAND_CAPACITY: usize = 1;
const MAX_DRAIN_TIMEOUT_SECS: u64 = 3_600;
/// Session handling policy for an in-process runtime reload.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Deserialize, Serialize)]
#[serde(rename_all = "snake_case")]
pub(crate) enum ReloadMode {
#[default]
Instant,
Drain,
}
/// Failure policy applied during the activation barrier.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Deserialize, Serialize)]
#[serde(rename_all = "snake_case")]
pub(crate) enum ReloadFailurePolicy {
#[default]
KeepNew,
Rollback,
}
/// Request body accepted by the maestro reload endpoint.
#[derive(Debug, Clone, Default, PartialEq, Eq, Deserialize, Serialize)]
#[serde(deny_unknown_fields)]
pub(crate) struct ReloadRequest {
#[serde(default)]
pub(crate) mode: ReloadMode,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) timeout_secs: Option<u64>,
#[serde(default)]
pub(crate) failure_policy: ReloadFailurePolicy,
}
impl ReloadRequest {
pub(crate) fn validate(&self) -> Result<(), &'static str> {
match (self.mode, self.timeout_secs) {
(ReloadMode::Instant, None) => Ok(()),
(ReloadMode::Instant, Some(_)) => Err("timeout_secs is only valid when mode is drain"),
(ReloadMode::Drain, Some(1..=MAX_DRAIN_TIMEOUT_SECS)) => Ok(()),
(ReloadMode::Drain, Some(_)) => Err("timeout_secs must be within 1..=3600"),
(ReloadMode::Drain, None) => Err("timeout_secs is required when mode is drain"),
}
}
pub(crate) fn from_query(query: Option<&str>) -> Result<Option<Self>, String> {
let Some(query) = query.filter(|query| !query.is_empty()) else {
return Ok(None);
};
let mut mode = None;
let mut timeout_secs = None;
let mut failure_policy = None;
for (key, value) in url::form_urlencoded::parse(query.as_bytes()) {
match key.as_ref() {
"reload" if mode.is_none() => {
mode = Some(match value.as_ref() {
"instant" => ReloadMode::Instant,
"drain" => ReloadMode::Drain,
_ => return Err("reload must be instant or drain".to_string()),
});
}
"timeout_secs" if timeout_secs.is_none() => {
timeout_secs = Some(
value
.parse::<u64>()
.map_err(|_| "timeout_secs must be an integer".to_string())?,
);
}
"failure_policy" if failure_policy.is_none() => {
failure_policy = Some(match value.as_ref() {
"keep_new" => ReloadFailurePolicy::KeepNew,
"rollback" => ReloadFailurePolicy::Rollback,
_ => {
return Err("failure_policy must be keep_new or rollback".to_string());
}
});
}
"reload" | "timeout_secs" | "failure_policy" => {
return Err(format!("duplicate query parameter: {}", key));
}
_ => return Err(format!("unknown query parameter: {}", key)),
}
}
let mode = mode.ok_or_else(|| "reload query parameter is required".to_string())?;
let request = Self {
mode,
timeout_secs,
failure_policy: failure_policy.unwrap_or_default(),
};
request.validate().map_err(str::to_string)?;
Ok(Some(request))
}
}
/// Observable phase of one reload operation.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
#[serde(rename_all = "snake_case")]
pub(crate) enum ReloadPhase {
Accepted,
Preparing,
Activating,
Draining,
Succeeded,
RolledBack,
Failed,
}
impl ReloadPhase {
fn is_terminal(self) -> bool {
matches!(
self,
ReloadPhase::Succeeded | ReloadPhase::RolledBack | ReloadPhase::Failed
)
}
}
/// Bounded public status for one reload operation.
#[derive(Debug, Clone, Serialize)]
pub(crate) struct ReloadStatus {
pub(crate) reload_id: u64,
pub(crate) target_generation: u64,
pub(crate) config_revision: String,
pub(crate) state: ReloadPhase,
pub(crate) mode: ReloadMode,
pub(crate) failure_policy: ReloadFailurePolicy,
pub(crate) requested_at_epoch_secs: u64,
#[serde(skip_serializing_if = "Option::is_none")]
pub(crate) started_at_epoch_secs: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub(crate) finished_at_epoch_secs: Option<u64>,
#[serde(
rename = "deferred_process_fields",
default,
skip_serializing_if = "Vec::is_empty"
)]
pub(crate) deferred_fields: Vec<String>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub(crate) warnings: Vec<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub(crate) error: Option<String>,
}
/// Accepted operation metadata returned before asynchronous preparation starts.
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub(crate) struct ReloadAccepted {
pub(crate) reload_id: u64,
pub(crate) target_generation: u64,
pub(crate) config_revision: String,
pub(crate) state: ReloadPhase,
pub(crate) mode: ReloadMode,
pub(crate) failure_policy: ReloadFailurePolicy,
}
pub(crate) struct ReloadCommand {
pub(crate) reload_id: u64,
pub(crate) target_generation: u64,
pub(crate) config: Arc<ProxyConfig>,
pub(crate) config_revision: String,
pub(crate) request: ReloadRequest,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum ReloadSubmitError {
InProgress(u64),
MaestroUnavailable,
}
#[derive(Clone)]
pub(crate) struct ReloadControl {
command_tx: mpsc::Sender<ReloadCommand>,
status_store: Arc<ReloadStatusStore>,
active_generation: Arc<AtomicU64>,
}
pub(crate) struct ReloadCommandReceiver {
command_rx: mpsc::Receiver<ReloadCommand>,
}
#[derive(Default)]
struct ReloadStatusState {
next_reload_id: u64,
active_reload_id: Option<u64>,
statuses: VecDeque<ReloadStatus>,
}
#[derive(Default)]
struct ReloadStatusStore {
state: Mutex<ReloadStatusState>,
}
impl ReloadControl {
pub(crate) fn channel(initial_generation: u64) -> (Self, ReloadCommandReceiver) {
let (command_tx, command_rx) = mpsc::channel(RELOAD_COMMAND_CAPACITY);
(
Self {
command_tx,
status_store: Arc::new(ReloadStatusStore::default()),
active_generation: Arc::new(AtomicU64::new(initial_generation)),
},
ReloadCommandReceiver { command_rx },
)
}
pub(crate) async fn submit(
&self,
config: Arc<ProxyConfig>,
config_revision: String,
request: ReloadRequest,
) -> Result<ReloadAccepted, ReloadSubmitError> {
let target_generation = self
.active_generation
.load(Ordering::Acquire)
.saturating_add(1);
let status = self
.status_store
.reserve(target_generation, config_revision, request.clone())
.await?;
let command = ReloadCommand {
reload_id: status.reload_id,
target_generation,
config,
config_revision: status.config_revision.clone(),
request,
};
if self.command_tx.send(command).await.is_err() {
self.status_store
.finish(
status.reload_id,
ReloadPhase::Failed,
Some("maestro command channel is closed".to_string()),
)
.await;
return Err(ReloadSubmitError::MaestroUnavailable);
}
Ok(ReloadAccepted {
reload_id: status.reload_id,
target_generation,
config_revision: status.config_revision,
state: ReloadPhase::Accepted,
mode: status.mode,
failure_policy: status.failure_policy,
})
}
pub(crate) async fn status(&self, reload_id: u64) -> Option<ReloadStatus> {
self.status_store.get(reload_id).await
}
pub(crate) async fn in_progress(&self) -> Option<u64> {
self.status_store.state.lock().await.active_reload_id
}
pub(crate) async fn mark_phase(&self, reload_id: u64, phase: ReloadPhase) {
self.status_store.mark_phase(reload_id, phase).await;
}
pub(crate) async fn set_deferred_fields(&self, reload_id: u64, fields: Vec<String>) {
self.status_store
.update(reload_id, |status| status.deferred_fields = fields)
.await;
}
pub(crate) async fn succeed(&self, reload_id: u64, generation: u64) {
self.active_generation.store(generation, Ordering::Release);
self.status_store
.finish(reload_id, ReloadPhase::Succeeded, None)
.await;
}
pub(crate) async fn fail(&self, reload_id: u64, error: impl Into<String>) {
self.status_store
.finish(reload_id, ReloadPhase::Failed, Some(error.into()))
.await;
}
pub(crate) async fn rolled_back(&self, reload_id: u64, error: impl Into<String>) {
self.status_store
.finish(reload_id, ReloadPhase::RolledBack, Some(error.into()))
.await;
}
pub(crate) async fn add_warning(&self, reload_id: u64, warning: impl Into<String>) {
let warning = warning.into();
self.status_store
.update(reload_id, |status| status.warnings.push(warning))
.await;
}
}
impl ReloadCommandReceiver {
pub(crate) async fn recv(&mut self) -> Option<ReloadCommand> {
self.command_rx.recv().await
}
}
impl ReloadStatusStore {
async fn reserve(
&self,
target_generation: u64,
config_revision: String,
request: ReloadRequest,
) -> Result<ReloadStatus, ReloadSubmitError> {
let mut state = self.state.lock().await;
if let Some(reload_id) = state.active_reload_id {
return Err(ReloadSubmitError::InProgress(reload_id));
}
state.next_reload_id = state.next_reload_id.saturating_add(1).max(1);
let reload_id = state.next_reload_id;
let status = ReloadStatus {
reload_id,
target_generation,
config_revision,
state: ReloadPhase::Accepted,
mode: request.mode,
failure_policy: request.failure_policy,
requested_at_epoch_secs: now_epoch_secs(),
started_at_epoch_secs: None,
finished_at_epoch_secs: None,
deferred_fields: Vec::new(),
warnings: Vec::new(),
error: None,
};
state.active_reload_id = Some(reload_id);
state.statuses.push_back(status.clone());
while state.statuses.len() > RELOAD_HISTORY_CAPACITY {
state.statuses.pop_front();
}
Ok(status)
}
async fn get(&self, reload_id: u64) -> Option<ReloadStatus> {
self.state
.lock()
.await
.statuses
.iter()
.find(|status| status.reload_id == reload_id)
.cloned()
}
async fn mark_phase(&self, reload_id: u64, phase: ReloadPhase) {
self.update(reload_id, |status| {
status.state = phase;
if status.started_at_epoch_secs.is_none() && phase != ReloadPhase::Accepted {
status.started_at_epoch_secs = Some(now_epoch_secs());
}
})
.await;
}
async fn finish(&self, reload_id: u64, phase: ReloadPhase, error: Option<String>) {
debug_assert!(phase.is_terminal());
let mut state = self.state.lock().await;
if let Some(status) = state
.statuses
.iter_mut()
.find(|status| status.reload_id == reload_id)
{
status.state = phase;
status.error = error;
status.finished_at_epoch_secs = Some(now_epoch_secs());
}
if state.active_reload_id == Some(reload_id) {
state.active_reload_id = None;
}
}
async fn update(&self, reload_id: u64, update: impl FnOnce(&mut ReloadStatus)) {
let mut state = self.state.lock().await;
if let Some(status) = state
.statuses
.iter_mut()
.find(|status| status.reload_id == reload_id)
{
update(status);
}
}
}
fn now_epoch_secs() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn request_defaults_to_instant_keep_new() {
let request: ReloadRequest = serde_json::from_str("{}").unwrap();
assert_eq!(request, ReloadRequest::default());
assert_eq!(request.validate(), Ok(()));
}
#[test]
fn drain_requires_bounded_timeout() {
let missing = ReloadRequest {
mode: ReloadMode::Drain,
..ReloadRequest::default()
};
assert!(missing.validate().is_err());
let valid = ReloadRequest {
mode: ReloadMode::Drain,
timeout_secs: Some(30),
..ReloadRequest::default()
};
assert_eq!(valid.validate(), Ok(()));
}
#[test]
fn patch_query_parses_reload_policy() {
let request =
ReloadRequest::from_query(Some("reload=drain&timeout_secs=30&failure_policy=rollback"))
.unwrap()
.unwrap();
assert_eq!(request.mode, ReloadMode::Drain);
assert_eq!(request.timeout_secs, Some(30));
assert_eq!(request.failure_policy, ReloadFailurePolicy::Rollback);
assert!(ReloadRequest::from_query(Some("timeout_secs=30")).is_err());
}
#[test]
fn status_uses_documented_deferred_process_fields_key() {
let status = ReloadStatus {
reload_id: 1,
target_generation: 2,
config_revision: "revision".to_string(),
state: ReloadPhase::Succeeded,
mode: ReloadMode::Instant,
failure_policy: ReloadFailurePolicy::KeepNew,
requested_at_epoch_secs: 10,
started_at_epoch_secs: Some(11),
finished_at_epoch_secs: Some(12),
deferred_fields: vec!["server.listeners".to_string()],
warnings: Vec::new(),
error: None,
};
let value = serde_json::to_value(status).unwrap();
assert_eq!(
value["deferred_process_fields"],
serde_json::json!(["server.listeners"])
);
assert!(value.get("deferred_fields").is_none());
}
#[tokio::test]
async fn coordinator_rejects_concurrent_reload_and_releases_terminal_slot() {
let (control, mut receiver) = ReloadControl::channel(1);
let first = control
.submit(
Arc::new(ProxyConfig::default()),
"rev-1".to_string(),
ReloadRequest::default(),
)
.await
.unwrap();
let _command = receiver.recv().await.unwrap();
let second = control
.submit(
Arc::new(ProxyConfig::default()),
"rev-2".to_string(),
ReloadRequest::default(),
)
.await;
assert_eq!(second, Err(ReloadSubmitError::InProgress(first.reload_id)));
control
.succeed(first.reload_id, first.target_generation)
.await;
let third = control
.submit(
Arc::new(ProxyConfig::default()),
"rev-3".to_string(),
ReloadRequest::default(),
)
.await
.unwrap();
assert_eq!(third.reload_id, first.reload_id + 1);
}
}

View File

@@ -0,0 +1,256 @@
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
use arc_swap::ArcSwap;
use tokio::sync::watch;
use tracing::{info, warn};
use crate::stats::QuotaStore;
use super::generation::{RuntimeGeneration, RuntimeWatchState};
use super::reload::{
ReloadCommand, ReloadCommandReceiver, ReloadControl, ReloadFailurePolicy, ReloadMode,
ReloadPhase,
};
use super::runtime_build::{deferred_process_fields, prepare_runtime};
use super::runtime_tasks::RuntimeLogFilter;
pub(crate) struct ReloadSupervisor {
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
control: ReloadControl,
commands: ReloadCommandReceiver,
config_path: PathBuf,
quota_store: Arc<QuotaStore>,
detected_ips_tx: watch::Sender<(Option<std::net::IpAddr>, Option<std::net::IpAddr>)>,
runtime_log_filter: RuntimeLogFilter,
runtime_watch_tx: watch::Sender<Option<RuntimeWatchState>>,
}
#[derive(Debug, PartialEq, Eq)]
enum RevisionGateAction {
Proceed,
Warn(String),
Rollback(String),
}
fn revision_gate_action(
accepted_revision: &str,
current_revision: Result<String, String>,
failure_policy: ReloadFailurePolicy,
) -> RevisionGateAction {
let warning = match current_revision {
Ok(current) if current == accepted_revision => return RevisionGateAction::Proceed,
Ok(current) => format!(
"config revision changed during preparation: accepted={} current={}",
accepted_revision, current
),
Err(error) => format!("config revision verification failed: {}", error),
};
match failure_policy {
ReloadFailurePolicy::KeepNew => RevisionGateAction::Warn(warning),
ReloadFailurePolicy::Rollback => RevisionGateAction::Rollback(warning),
}
}
async fn stop_background_and_middle_end(generation: &RuntimeGeneration) -> bool {
generation.stop_background_tasks().await;
let Some(pool) = generation.current_me_pool().await else {
return false;
};
tokio::time::timeout(Duration::from_secs(2), pool.shutdown_send_close_conn_all())
.await
.is_err()
}
async fn cleanup_candidate(generation: &RuntimeGeneration) -> bool {
generation.stop_sessions().await;
stop_background_and_middle_end(generation).await
}
impl ReloadSupervisor {
#[allow(clippy::too_many_arguments)]
pub(crate) fn spawn(
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
control: ReloadControl,
commands: ReloadCommandReceiver,
config_path: PathBuf,
quota_store: Arc<QuotaStore>,
detected_ips_tx: watch::Sender<(Option<std::net::IpAddr>, Option<std::net::IpAddr>)>,
runtime_log_filter: RuntimeLogFilter,
runtime_watch_tx: watch::Sender<Option<RuntimeWatchState>>,
) {
let supervisor = Self {
active_runtime,
control,
commands,
config_path,
quota_store,
detected_ips_tx,
runtime_log_filter,
runtime_watch_tx,
};
tokio::spawn(supervisor.run());
}
async fn run(mut self) {
while let Some(command) = self.commands.recv().await {
self.reload(command).await;
}
}
async fn reload(&self, command: ReloadCommand) {
self.control
.mark_phase(command.reload_id, ReloadPhase::Preparing)
.await;
let old_runtime = self.active_runtime.load_full();
let deferred = deferred_process_fields(&old_runtime.config(), &command.config);
self.control
.set_deferred_fields(command.reload_id, deferred)
.await;
let prepared = match prepare_runtime(
command.target_generation,
command.config.as_ref().clone(),
&self.config_path,
self.quota_store.clone(),
self.runtime_log_filter.clone(),
)
.await
{
Ok(prepared) => prepared,
Err(error) => {
self.control.fail(command.reload_id, error).await;
return;
}
};
let revision_action = revision_gate_action(
&command.config_revision,
crate::api::config_store::current_revision_for_maestro(&self.config_path).await,
command.request.failure_policy,
);
match revision_action {
RevisionGateAction::Proceed => {}
RevisionGateAction::Warn(warning) => {
self.control.add_warning(command.reload_id, warning).await;
}
RevisionGateAction::Rollback(warning) => {
let _ = cleanup_candidate(&prepared.generation).await;
self.runtime_log_filter
.apply_reload(&old_runtime.config().general.log_level);
self.control.rolled_back(command.reload_id, warning).await;
return;
}
}
self.control
.mark_phase(command.reload_id, ReloadPhase::Activating)
.await;
let new_runtime = prepared.generation;
old_runtime.stop_accepting_sessions();
if let Err(error) = crate::network::dns_overrides::install_entries(
&new_runtime.config().network.dns_overrides,
) {
let message = format!("runtime DNS activation failed: {}", error);
if command.request.failure_policy == ReloadFailurePolicy::Rollback {
old_runtime.resume_accepting_sessions();
let _ = cleanup_candidate(&new_runtime).await;
self.runtime_log_filter
.apply_reload(&old_runtime.config().general.log_level);
self.control.rolled_back(command.reload_id, message).await;
return;
}
self.control.add_warning(command.reload_id, message).await;
}
let replaced = self.active_runtime.swap(new_runtime.clone());
self.detected_ips_tx.send_replace(prepared.detected_ips);
self.runtime_log_filter
.apply_reload(&new_runtime.config().general.log_level);
self.runtime_watch_tx
.send_replace(Some(new_runtime.watch_state()));
info!(
reload_id = command.reload_id,
old_generation = replaced.id,
new_generation = new_runtime.id,
config_revision = %command.config_revision,
"Runtime generation activated"
);
match command.request.mode {
ReloadMode::Instant => {
replaced.stop_sessions().await;
}
ReloadMode::Drain => {
self.control
.mark_phase(command.reload_id, ReloadPhase::Draining)
.await;
let timeout = Duration::from_secs(
command
.request
.timeout_secs
.expect("validated drain request must carry timeout_secs"),
);
if !replaced.drain_sessions(timeout).await {
let warning = format!(
"generation {} exceeded drain timeout; remaining sessions were cancelled",
replaced.id
);
warn!(reload_id = command.reload_id, warning = %warning);
self.control.add_warning(command.reload_id, warning).await;
}
}
}
if stop_background_and_middle_end(&replaced).await {
let warning = format!(
"generation {} Middle-End close broadcast timed out",
replaced.id
);
warn!(reload_id = command.reload_id, warning = %warning);
self.control.add_warning(command.reload_id, warning).await;
}
self.control
.succeed(command.reload_id, new_runtime.id)
.await;
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn revision_gate_proceeds_only_on_verified_match() {
assert_eq!(
revision_gate_action(
"accepted",
Ok("accepted".to_string()),
ReloadFailurePolicy::Rollback,
),
RevisionGateAction::Proceed
);
}
#[test]
fn revision_gate_applies_failure_policy_to_mismatch_and_read_error() {
for result in [
Ok("changed".to_string()),
Err("read failed".to_string()),
] {
assert!(matches!(
revision_gate_action(
"accepted",
result.clone(),
ReloadFailurePolicy::KeepNew,
),
RevisionGateAction::Warn(_)
));
assert!(matches!(
revision_gate_action("accepted", result, ReloadFailurePolicy::Rollback),
RevisionGateAction::Rollback(_)
));
}
}
}

View File

@@ -0,0 +1,383 @@
use std::net::IpAddr;
use std::path::Path;
use std::sync::Arc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use tokio::sync::{RwLock, Semaphore, watch};
use crate::config::ProxyConfig;
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::route_mode::{RelayRouteMode, RouteRuntimeController};
use crate::proxy::shared_state::ProxySharedState;
use crate::startup::StartupTracker;
use crate::stats::beobachten::BeobachtenStore;
use crate::stats::telemetry::TelemetryPolicy;
use crate::stats::{QuotaStore, ReplayChecker, Stats};
use crate::stream::BufferPool;
use crate::transport::UpstreamManager;
use crate::transport::middle_proxy::MePool;
use super::admission;
use super::generation::{RuntimeGeneration, RuntimeTaskScope};
use super::runtime_tasks::RuntimeLogFilter;
use super::{me_startup, runtime_tasks, tls_bootstrap};
pub(crate) struct PreparedRuntime {
pub(crate) generation: Arc<RuntimeGeneration>,
pub(crate) detected_ips: (Option<IpAddr>, Option<IpAddr>),
}
pub(crate) async fn prepare_runtime(
generation_id: u64,
config: ProxyConfig,
config_path: &Path,
quota_store: Arc<QuotaStore>,
runtime_log_filter: RuntimeLogFilter,
) -> Result<PreparedRuntime, String> {
let started_at_epoch_secs = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.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));
stats.apply_telemetry_policy(TelemetryPolicy::from_config(&config.general.telemetry));
let upstream_manager = Arc::new(
UpstreamManager::new(
config.upstreams.clone(),
config.general.upstream_connect_retry_attempts,
config.general.upstream_connect_retry_backoff_ms,
config.general.upstream_connect_budget_ms,
config.general.tg_connect,
config.general.upstream_unhealthy_fail_threshold,
config.general.upstream_connect_failfast_hard_errors,
stats.clone(),
)
.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 probe = run_probe(
&config.network,
&config.upstreams,
config.general.middle_proxy_nat_probe,
config.general.stun_nat_probe_concurrency,
)
.await
.map_err(|error| format!("network probe failed: {}", error))?;
let decision =
decide_network_capabilities(&config.network, &probe, config.general.middle_proxy_nat_ip);
let prefer_ipv6 = decision.prefer_ipv6();
let mut tls_domains = Vec::with_capacity(1 + config.censorship.tls_domains.len());
tls_domains.push(config.censorship.tls_domain.clone());
for domain in &config.censorship.tls_domains {
if !tls_domains.contains(domain) {
tls_domains.push(domain.clone());
}
}
let tls_cache = tls_bootstrap::bootstrap_tls_front(
&config,
&tls_domains,
upstream_manager.clone(),
&startup_tracker,
task_scope.clone(),
tls_bootstrap::TlsBootstrapPolicy::RequireReady,
)
.await
.map_err(|error| error.to_string())?;
let beobachten = Arc::new(BeobachtenStore::new());
let rng = Arc::new(SecureRandom::new());
let route_mode = if !config.general.use_middle_proxy || config.general.me2dc_fallback {
RelayRouteMode::Direct
} else {
RelayRouteMode::Middle
};
let route_runtime = Arc::new(RouteRuntimeController::new(route_mode));
let me_pool_runtime = Arc::new(RwLock::new(None::<Arc<MePool>>));
let (me_ready_tx, me_ready_rx) = watch::channel(0_u64);
let direct_first_startup = config.general.use_middle_proxy && config.general.me2dc_fallback;
let me_pool = if direct_first_startup {
None
} else {
me_startup::initialize_me_pool(
config.general.use_middle_proxy,
&config,
&decision,
&probe,
&startup_tracker,
upstream_manager.clone(),
rng.clone(),
stats.clone(),
me_pool_runtime.clone(),
me_ready_tx.clone(),
task_scope.clone(),
)
.await
};
if strict_middle_proxy_unavailable(
config.general.use_middle_proxy,
direct_first_startup,
me_pool.is_some(),
) {
task_scope.stop().await;
return Err(
"Middle-End pool is required but did not become ready during reload preparation"
.to_string(),
);
}
let config = Arc::new(config);
let replay_checker = Arc::new(ReplayChecker::new(
config.access.replay_check_len,
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 watches = runtime_tasks::spawn_runtime_tasks(
&config,
config_path,
&probe,
prefer_ipv6,
decision.ipv4_dc,
decision.ipv6_dc,
&startup_tracker,
stats.clone(),
upstream_manager.clone(),
replay_checker.clone(),
me_pool.clone(),
rng.clone(),
ip_tracker.clone(),
beobachten.clone(),
me_pool.clone(),
proxy_shared.clone(),
me_ready_tx.clone(),
task_scope.clone(),
)
.await;
let config_rx = watches.config_rx;
runtime_log_filter.spawn_watcher(watches.log_level_rx, task_scope.clone());
let initial_admission_open = !config.general.use_middle_proxy || me_pool.is_some();
let (admission_tx, admission_rx) = watch::channel(initial_admission_open);
admission::configure_admission_gate(
&config,
me_pool.clone(),
me_pool_runtime.clone(),
route_runtime.clone(),
&admission_tx,
config_rx.clone(),
me_ready_rx,
task_scope.clone(),
)
.await;
if direct_first_startup {
let config_bg = config.clone();
let decision_bg = decision.clone();
let probe_bg = probe.clone();
let startup_tracker_bg = startup_tracker.clone();
let upstream_manager_bg = upstream_manager.clone();
let rng_bg = rng.clone();
let stats_bg = stats.clone();
let me_pool_runtime_bg = me_pool_runtime.clone();
let me_ready_tx_bg = me_ready_tx.clone();
let config_rx_bg = config_rx.clone();
let task_scope_bg = task_scope.clone();
let retry_limit = config.general.me_init_retry_attempts;
task_scope.spawn(async move {
let mut attempt = 0_u32;
loop {
attempt = attempt.saturating_add(1);
let pool = me_startup::initialize_me_pool(
true,
config_bg.as_ref(),
&decision_bg,
&probe_bg,
&startup_tracker_bg,
upstream_manager_bg.clone(),
rng_bg.clone(),
stats_bg.clone(),
me_pool_runtime_bg.clone(),
me_ready_tx_bg.clone(),
task_scope_bg.clone(),
)
.await;
if let Some(pool) = pool {
runtime_tasks::spawn_middle_proxy_runtime_tasks(
config_bg.as_ref(),
config_rx_bg,
pool,
rng_bg,
me_ready_tx_bg,
task_scope_bg,
);
break;
}
if retry_limit > 0 && attempt >= retry_limit {
break;
}
tokio::time::sleep(Duration::from_secs(2)).await;
}
});
}
let conntrack_scope = task_scope.clone();
task_scope.spawn(crate::conntrack_control::run_conntrack_controller(
config_rx.clone(),
stats.clone(),
proxy_shared.clone(),
conntrack_scope.cancellation_token(),
));
task_scope.spawn(run_direct_buffer_budget_controller(
direct_buffer_budget,
buffer_pool.clone(),
stats.clone(),
proxy_shared.clone(),
config.server.max_connections,
));
let generation = RuntimeGeneration::new(
generation_id,
config_rx,
admission_rx,
stats,
upstream_manager,
replay_checker,
buffer_pool,
rng,
me_pool,
me_pool_runtime,
route_runtime,
tls_cache,
ip_tracker,
beobachten,
proxy_shared,
max_connections,
task_scope,
);
drop(admission_tx);
Ok(PreparedRuntime {
generation,
detected_ips: (
probe.detected_ipv4.map(IpAddr::V4),
probe.detected_ipv6.map(IpAddr::V6),
),
})
}
fn strict_middle_proxy_unavailable(
use_middle_proxy: bool,
direct_first_startup: bool,
pool_available: bool,
) -> bool {
use_middle_proxy && !direct_first_startup && !pool_available
}
pub(crate) fn deferred_process_fields(old: &ProxyConfig, new: &ProxyConfig) -> Vec<String> {
let mut fields = Vec::new();
if old.server.port != new.server.port
|| old.server.proxy_protocol != new.server.proxy_protocol
|| old.server.listen_backlog != new.server.listen_backlog
|| serde_json::to_value(&old.server.listeners).ok()
!= serde_json::to_value(&new.server.listeners).ok()
{
fields.push("server.listeners".to_string());
}
if old.server.listen_unix_sock != new.server.listen_unix_sock
|| old.server.listen_unix_sock_perm != new.server.listen_unix_sock_perm
{
fields.push("server.listen_unix_sock".to_string());
}
if old.server.api.listen != new.server.api.listen
|| old.server.api.enabled != new.server.api.enabled
{
fields.push("server.api.listen".to_string());
}
if old.server.metrics_listen != new.server.metrics_listen
|| old.server.metrics_port != new.server.metrics_port
{
fields.push("server.metrics_listen".to_string());
}
if old.general.quota_state_path != new.general.quota_state_path {
fields.push("general.quota_state_path".to_string());
}
if old.general.disable_colors != new.general.disable_colors {
fields.push("general.disable_colors".to_string());
}
if old.general.data_path != new.general.data_path {
fields.push("general.data_path".to_string());
}
if serde_json::to_value(&old.logging).ok() != serde_json::to_value(&new.logging).ok() {
fields.push("logging".to_string());
}
fields
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn process_socket_and_logging_changes_are_deferred() {
let old = ProxyConfig::default();
let mut new = old.clone();
new.server.listen_backlog = new.server.listen_backlog.saturating_add(1);
new.general.disable_colors = !new.general.disable_colors;
let fields = deferred_process_fields(&old, &new);
assert!(fields.contains(&"server.listeners".to_string()));
assert!(fields.contains(&"general.disable_colors".to_string()));
}
#[test]
fn runtime_only_change_does_not_require_process_rebind() {
let old = ProxyConfig::default();
let mut new = old.clone();
new.censorship.tls_domain = "reload.example".to_string();
assert!(deferred_process_fields(&old, &new).is_empty());
}
#[test]
fn strict_middle_proxy_requires_a_prepared_pool() {
assert!(strict_middle_proxy_unavailable(true, false, false));
assert!(!strict_middle_proxy_unavailable(true, false, true));
assert!(!strict_middle_proxy_unavailable(true, true, false));
assert!(!strict_middle_proxy_unavailable(false, false, false));
}
}

View File

@@ -2,9 +2,11 @@ use std::net::IpAddr;
use std::path::Path;
use std::sync::Arc;
use arc_swap::ArcSwap;
use tokio::sync::{mpsc, watch};
use tracing::{debug, info, warn};
use tracing_subscriber::EnvFilter;
use tracing_subscriber::Registry;
use tracing_subscriber::reload;
use crate::config::hot_reload::spawn_config_watcher;
@@ -21,10 +23,11 @@ use crate::startup::{
use crate::stats::beobachten::BeobachtenStore;
use crate::stats::telemetry::TelemetryPolicy;
use crate::stats::{ReplayChecker, Stats};
use crate::tls_front::TlsFrontCache;
use crate::transport::UpstreamManager;
use crate::transport::middle_proxy::{MePool, MeReinitTrigger};
use super::generation::RuntimeGeneration;
use super::generation::RuntimeTaskScope;
use super::helpers::write_beobachten_snapshot;
pub(crate) struct RuntimeWatches {
@@ -34,6 +37,56 @@ pub(crate) struct RuntimeWatches {
pub(crate) detected_ip_v6: Option<IpAddr>,
}
#[derive(Clone)]
pub(crate) struct RuntimeLogFilter {
handle: reload::Handle<EnvFilter, Registry>,
}
impl RuntimeLogFilter {
pub(crate) fn new(handle: reload::Handle<EnvFilter, Registry>) -> Self {
Self { handle }
}
pub(crate) fn start(
&self,
has_rust_log: bool,
effective_log_level: &LogLevel,
log_level_rx: watch::Receiver<LogLevel>,
task_scope: RuntimeTaskScope,
) {
self.apply(effective_log_level, has_rust_log);
self.spawn_watcher(log_level_rx, task_scope);
}
pub(crate) fn apply_reload(&self, level: &LogLevel) {
self.apply(level, false);
}
pub(crate) fn spawn_watcher(
&self,
mut log_level_rx: watch::Receiver<LogLevel>,
task_scope: RuntimeTaskScope,
) {
let filter = self.clone();
task_scope.spawn(async move {
loop {
if log_level_rx.changed().await.is_err() {
break;
}
let level = log_level_rx.borrow_and_update().clone();
filter.apply_reload(&level);
}
});
}
fn apply(&self, level: &LogLevel, has_rust_log: bool) {
let runtime_filter = EnvFilter::new(log_filter_spec(has_rust_log, level));
if let Err(error) = self.handle.reload(runtime_filter) {
tracing::error!(error = %error, "Failed to update runtime log filter");
}
}
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn spawn_runtime_tasks(
config: &Arc<ProxyConfig>,
@@ -50,14 +103,14 @@ pub(crate) async fn spawn_runtime_tasks(
rng: Arc<SecureRandom>,
ip_tracker: Arc<UserIpTracker>,
beobachten: Arc<BeobachtenStore>,
api_config_tx: watch::Sender<Arc<ProxyConfig>>,
me_pool_for_policy: Option<Arc<MePool>>,
shared_state: Arc<ProxySharedState>,
me_ready_tx: watch::Sender<u64>,
task_scope: RuntimeTaskScope,
) -> RuntimeWatches {
let um_clone = upstream_manager.clone();
let dc_overrides_for_health = config.dc_overrides.clone();
tokio::spawn(async move {
task_scope.spawn(async move {
um_clone
.run_health_checks(
prefer_ipv6,
@@ -69,19 +122,19 @@ pub(crate) async fn spawn_runtime_tasks(
});
let rc_clone = replay_checker.clone();
tokio::spawn(async move {
task_scope.spawn(async move {
rc_clone.run_periodic_cleanup().await;
});
let stats_maintenance = stats.clone();
tokio::spawn(async move {
task_scope.spawn(async move {
stats_maintenance
.run_periodic_user_stats_maintenance()
.await;
});
let ip_tracker_maintenance = ip_tracker.clone();
tokio::spawn(async move {
task_scope.spawn(async move {
ip_tracker_maintenance.run_periodic_maintenance().await;
});
@@ -104,6 +157,7 @@ pub(crate) async fn spawn_runtime_tasks(
config.clone(),
detected_ip_v4,
detected_ip_v6,
task_scope.cancellation_token(),
);
startup_tracker
.complete_component(
@@ -111,21 +165,10 @@ pub(crate) async fn spawn_runtime_tasks(
Some("config hot-reload watcher started".to_string()),
)
.await;
let mut config_rx_api_bridge = config_rx.clone();
let api_config_tx_bridge = api_config_tx.clone();
tokio::spawn(async move {
loop {
if config_rx_api_bridge.changed().await.is_err() {
break;
}
let cfg = config_rx_api_bridge.borrow_and_update().clone();
api_config_tx_bridge.send_replace(cfg);
}
});
let stats_policy = stats.clone();
let upstream_policy = upstream_manager.clone();
let mut config_rx_policy = config_rx.clone();
tokio::spawn(async move {
task_scope.spawn(async move {
loop {
if config_rx_policy.changed().await.is_err() {
break;
@@ -133,6 +176,9 @@ pub(crate) async fn spawn_runtime_tasks(
let cfg = config_rx_policy.borrow_and_update().clone();
stats_policy
.apply_telemetry_policy(TelemetryPolicy::from_config(&cfg.general.telemetry));
if let Err(error) = upstream_policy.update_dns_overrides(&cfg.network.dns_overrides) {
warn!(error = %error, "Failed to update generation DNS overrides");
}
if let Some(pool) = &me_pool_for_policy {
pool.update_runtime_transport_policy(
cfg.general.me_socks_kdf_policy,
@@ -149,7 +195,7 @@ pub(crate) async fn spawn_runtime_tasks(
let ip_tracker_policy = ip_tracker.clone();
let mut config_rx_ip_limits = config_rx.clone();
tokio::spawn(async move {
task_scope.spawn(async move {
let mut prev_limits = config_rx_ip_limits
.borrow()
.access
@@ -205,7 +251,7 @@ pub(crate) async fn spawn_runtime_tasks(
config.access.cidr_rate_limits.clone(),
);
let mut config_rx_rate_limits = config_rx.clone();
tokio::spawn(async move {
task_scope.spawn(async move {
let mut prev_user_limits = config_rx_rate_limits
.borrow()
.access
@@ -236,7 +282,7 @@ pub(crate) async fn spawn_runtime_tasks(
let shared_user_enabled = shared_state.clone();
let mut config_rx_user_enabled = config_rx.clone();
tokio::spawn(async move {
task_scope.spawn(async move {
loop {
if config_rx_user_enabled.changed().await.is_err() {
break;
@@ -257,7 +303,7 @@ pub(crate) async fn spawn_runtime_tasks(
let beobachten_writer = beobachten.clone();
let config_rx_beobachten = config_rx.clone();
tokio::spawn(async move {
task_scope.spawn(async move {
loop {
let cfg = config_rx_beobachten.borrow().clone();
let sleep_secs = cfg.general.beobachten_flush_secs.max(1);
@@ -278,7 +324,14 @@ pub(crate) async fn spawn_runtime_tasks(
});
if let Some(pool) = me_pool {
spawn_middle_proxy_runtime_tasks(config, config_rx.clone(), pool, rng, me_ready_tx);
spawn_middle_proxy_runtime_tasks(
config,
config_rx.clone(),
pool,
rng,
me_ready_tx,
task_scope,
);
}
RuntimeWatches {
@@ -295,6 +348,7 @@ pub(crate) fn spawn_middle_proxy_runtime_tasks(
pool: Arc<MePool>,
rng: Arc<SecureRandom>,
me_ready_tx: watch::Sender<u64>,
task_scope: RuntimeTaskScope,
) {
let reinit_trigger_capacity = config.general.me_reinit_trigger_channel.max(1);
let (reinit_tx, reinit_rx) = mpsc::channel::<MeReinitTrigger>(reinit_trigger_capacity);
@@ -303,7 +357,7 @@ pub(crate) fn spawn_middle_proxy_runtime_tasks(
let rng_clone_sched = rng.clone();
let config_rx_clone_sched = config_rx.clone();
let me_ready_tx_sched = me_ready_tx.clone();
tokio::spawn(async move {
task_scope.spawn(async move {
crate::transport::middle_proxy::me_reinit_scheduler(
pool_clone_sched,
rng_clone_sched,
@@ -317,7 +371,7 @@ pub(crate) fn spawn_middle_proxy_runtime_tasks(
let pool_clone = pool.clone();
let config_rx_clone = config_rx.clone();
let reinit_tx_updater = reinit_tx.clone();
tokio::spawn(async move {
task_scope.spawn(async move {
crate::transport::middle_proxy::me_config_updater(
pool_clone,
config_rx_clone,
@@ -328,37 +382,12 @@ pub(crate) fn spawn_middle_proxy_runtime_tasks(
let config_rx_clone_rot = config_rx.clone();
let reinit_tx_rotation = reinit_tx.clone();
tokio::spawn(async move {
task_scope.spawn(async move {
crate::transport::middle_proxy::me_rotation_task(config_rx_clone_rot, reinit_tx_rotation)
.await;
});
}
pub(crate) async fn apply_runtime_log_filter(
has_rust_log: bool,
effective_log_level: &LogLevel,
filter_handle: reload::Handle<EnvFilter, tracing_subscriber::Registry>,
mut log_level_rx: watch::Receiver<LogLevel>,
) {
let runtime_filter = EnvFilter::new(log_filter_spec(has_rust_log, effective_log_level));
filter_handle
.reload(runtime_filter)
.expect("Failed to switch log filter");
tokio::spawn(async move {
loop {
if log_level_rx.changed().await.is_err() {
break;
}
let level = log_level_rx.borrow_and_update().clone();
let new_filter = tracing_subscriber::EnvFilter::new(log_filter_spec(false, &level));
if let Err(e) = filter_handle.reload(new_filter) {
tracing::error!("config reload: failed to update log filter: {}", e);
}
}
});
}
pub(crate) fn log_filter_spec(has_rust_log: bool, effective_log_level: &LogLevel) -> String {
if has_rust_log {
std::env::var("RUST_LOG")
@@ -373,12 +402,7 @@ pub(crate) fn log_filter_spec(has_rust_log: bool, effective_log_level: &LogLevel
pub(crate) async fn spawn_metrics_if_configured(
config: &Arc<ProxyConfig>,
startup_tracker: &Arc<StartupTracker>,
stats: Arc<Stats>,
beobachten: Arc<BeobachtenStore>,
shared_state: Arc<ProxySharedState>,
ip_tracker: Arc<UserIpTracker>,
tls_cache: Option<Arc<TlsFrontCache>>,
config_rx: watch::Receiver<Arc<ProxyConfig>>,
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
) {
// metrics_listen takes precedence; fall back to metrics_port for backward compat.
let metrics_target: Option<(u16, Option<String>)> =
@@ -408,28 +432,10 @@ pub(crate) async fn spawn_metrics_if_configured(
Some(format!("spawn metrics endpoint on {}", label)),
)
.await;
let stats = stats.clone();
let beobachten = beobachten.clone();
let shared_state = shared_state.clone();
let config_rx_metrics = config_rx.clone();
let ip_tracker_metrics = ip_tracker.clone();
let tls_cache_metrics = tls_cache.clone();
let whitelist = config.server.metrics_whitelist.clone();
let active_runtime = active_runtime.clone();
let listen_backlog = config.server.listen_backlog;
tokio::spawn(async move {
metrics::serve(
port,
listen,
listen_backlog,
stats,
beobachten,
shared_state,
ip_tracker_metrics,
tls_cache_metrics,
config_rx_metrics,
whitelist,
)
.await;
metrics::serve(port, listen, listen_backlog, active_runtime).await;
});
startup_tracker
.complete_component(

View File

@@ -12,17 +12,17 @@ use std::path::PathBuf;
use std::sync::Arc;
use std::time::{Duration, Instant};
use arc_swap::ArcSwap;
#[cfg(not(unix))]
use tokio::signal;
#[cfg(unix)]
use tokio::signal::unix::{SignalKind, signal};
use tracing::{info, warn};
use super::generation::RuntimeGeneration;
use super::helpers::{format_uptime, unit_label};
use crate::stats::Stats;
use crate::synlimit_control;
use crate::transport::middle_proxy::MePool;
use super::helpers::{format_uptime, unit_label};
/// Signal that triggered shutdown.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -48,17 +48,17 @@ impl std::fmt::Display for ShutdownSignal {
/// Waits for a shutdown signal and performs graceful shutdown.
pub(crate) async fn wait_for_shutdown(
process_started_at: Instant,
me_pool: Option<Arc<MePool>>,
stats: Arc<Stats>,
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
quota_state_path: PathBuf,
synlimit_controller: tokio::task::JoinHandle<()>,
) {
let signal = wait_for_shutdown_signal().await;
perform_shutdown(
signal,
process_started_at,
me_pool,
&stats,
active_runtime,
quota_state_path,
synlimit_controller,
)
.await;
}
@@ -87,10 +87,12 @@ async fn wait_for_shutdown_signal() -> ShutdownSignal {
async fn perform_shutdown(
signal: ShutdownSignal,
process_started_at: Instant,
me_pool: Option<Arc<MePool>>,
stats: &Stats,
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
quota_state_path: PathBuf,
synlimit_controller: tokio::task::JoinHandle<()>,
) {
let runtime = active_runtime.load_full();
let stats = runtime.stats.as_ref();
let shutdown_started_at = Instant::now();
info!(signal = %signal, "Received shutdown signal");
@@ -103,12 +105,10 @@ async fn perform_shutdown(
let uptime_secs = process_started_at.elapsed().as_secs();
info!("Uptime: {}", format_uptime(uptime_secs));
if let Err(error) = synlimit_control::clear_synlimit_rules_all_backends().await {
warn!(error = %error, "Failed to clear SYN limiter rules during shutdown");
}
// Graceful ME pool shutdown
if let Some(pool) = &me_pool {
runtime.stop_sessions().await;
runtime.stop_background_tasks().await;
if let Some(pool) = runtime.current_me_pool().await {
match tokio::time::timeout(Duration::from_secs(2), pool.shutdown_send_close_conn_all())
.await
{
@@ -124,6 +124,12 @@ async fn perform_shutdown(
}
}
synlimit_controller.abort();
let _ = synlimit_controller.await;
if let Err(error) = synlimit_control::clear_synlimit_rules_all_backends().await {
warn!(error = %error, "Failed to clear SYN limiter rules during shutdown");
}
match crate::quota_state::save_quota_state(&quota_state_path, stats).await {
Ok(()) => {
info!(
@@ -191,7 +197,10 @@ fn dump_stats(stats: &Stats, process_started_at: Instant) {
/// - SIGUSR1: Log rotation acknowledgment (for external log rotation tools)
/// - SIGUSR2: Dump runtime status to log
#[cfg(unix)]
pub(crate) fn spawn_signal_handlers(stats: Arc<Stats>, process_started_at: Instant) {
pub(crate) fn spawn_signal_handlers(
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
process_started_at: Instant,
) {
tokio::spawn(async move {
let mut sigusr1 =
signal(SignalKind::user_defined1()).expect("Failed to register SIGUSR1 handler");
@@ -204,7 +213,8 @@ pub(crate) fn spawn_signal_handlers(stats: Arc<Stats>, process_started_at: Insta
handle_sigusr1();
}
_ = sigusr2.recv() => {
handle_sigusr2(&stats, process_started_at);
let runtime = active_runtime.load_full();
handle_sigusr2(runtime.stats.as_ref(), process_started_at);
}
}
}
@@ -213,7 +223,10 @@ pub(crate) fn spawn_signal_handlers(stats: Arc<Stats>, process_started_at: Insta
/// No-op on non-Unix platforms.
#[cfg(not(unix))]
pub(crate) fn spawn_signal_handlers(_stats: Arc<Stats>, _process_started_at: Instant) {
pub(crate) fn spawn_signal_handlers(
_active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
_process_started_at: Instant,
) {
// No SIGUSR1/SIGUSR2 on non-Unix
}

View File

@@ -5,11 +5,90 @@ use rand::RngExt;
use tracing::warn;
use crate::config::ProxyConfig;
use crate::error::{ProxyError, Result};
use crate::startup::{COMPONENT_TLS_FRONT_BOOTSTRAP, StartupTracker};
use crate::tls_front::TlsFrontCache;
use crate::tls_front::fetcher::TlsFetchStrategy;
use crate::transport::UpstreamManager;
use super::generation::RuntimeTaskScope;
/// Readiness requirement for TLS-front cache initialization.
#[derive(Clone, Copy)]
pub(crate) enum TlsBootstrapPolicy {
BestEffort,
RequireReady,
}
#[derive(Clone)]
struct TlsFetchContext {
cache: Arc<TlsFrontCache>,
domains: Vec<String>,
mask_host: String,
primary_domain: String,
mask_unix_sock: Option<String>,
tls_fetch_scope: Option<String>,
upstream_manager: Arc<UpstreamManager>,
strategy: TlsFetchStrategy,
port: u16,
proxy_protocol: u8,
}
impl TlsFetchContext {
async fn fetch_all(&self, failure_message: &'static str) {
let mut join = tokio::task::JoinSet::new();
for domain in self.domains.clone() {
let cache = self.cache.clone();
let host = tls_fetch_host_for_domain(
&self.mask_host,
&self.primary_domain,
&domain,
);
let unix_sock = self.mask_unix_sock.clone();
let scope = self.tls_fetch_scope.clone();
let upstream = self.upstream_manager.clone();
let strategy = self.strategy.clone();
let port = self.port;
let proxy_protocol = self.proxy_protocol;
join.spawn(async move {
match crate::tls_front::fetcher::fetch_real_tls_with_strategy(
&host,
port,
&domain,
&strategy,
Some(upstream),
scope.as_deref(),
proxy_protocol,
unix_sock.as_deref(),
)
.await
{
Ok(result) => cache.update_from_fetch(&domain, result).await,
Err(error) => warn!(domain = %domain, error = %error, failure_message),
}
});
}
while let Some(result) = join.join_next().await {
if let Err(error) = result {
warn!(error = %error, "TLS emulation fetch task join failed");
}
}
}
async fn fetch_all_with_budget(&self, phase: &'static str) {
if tokio::time::timeout(self.strategy.total_budget, self.fetch_all(phase))
.await
.is_err()
{
warn!(
phase,
timeout_ms = self.strategy.total_budget.as_millis(),
"TLS emulation fetch budget exhausted"
);
}
}
}
fn tls_fetch_host_for_domain(mask_host: &str, primary_tls_domain: &str, domain: &str) -> String {
if mask_host.eq_ignore_ascii_case(primary_tls_domain) {
domain.to_string()
@@ -18,12 +97,24 @@ fn tls_fetch_host_for_domain(mask_host: &str, primary_tls_domain: &str, domain:
}
}
fn readiness_error(default_domains: &[String]) -> Option<String> {
(!default_domains.is_empty()).then(|| {
format!(
"TLS-front profiles are not ready for domains: {}",
default_domains.join(", ")
)
})
}
/// Initializes the TLS-front cache and generation-owned refresh tasks.
pub(crate) async fn bootstrap_tls_front(
config: &ProxyConfig,
tls_domains: &[String],
upstream_manager: Arc<UpstreamManager>,
startup_tracker: &Arc<StartupTracker>,
) -> Option<Arc<TlsFrontCache>> {
task_scope: RuntimeTaskScope,
policy: TlsBootstrapPolicy,
) -> Result<Option<Arc<TlsFrontCache>>> {
startup_tracker
.start_component(
COMPONENT_TLS_FRONT_BOOTSTRAP,
@@ -31,26 +122,38 @@ pub(crate) async fn bootstrap_tls_front(
)
.await;
let tls_cache: Option<Arc<TlsFrontCache>> = if config.censorship.tls_emulation {
let cache = Arc::new(TlsFrontCache::new(
tls_domains,
config.censorship.fake_cert_len,
&config.censorship.tls_front_dir,
));
cache.load_from_disk().await;
if !config.censorship.tls_emulation {
startup_tracker
.skip_component(
COMPONENT_TLS_FRONT_BOOTSTRAP,
Some("censorship.tls_emulation is false".to_string()),
)
.await;
return Ok(None);
}
let port = config.censorship.mask_port;
let proxy_protocol = config.censorship.mask_proxy_protocol;
let mask_host = config
let cache = Arc::new(TlsFrontCache::new(
tls_domains,
config.censorship.fake_cert_len,
&config.censorship.tls_front_dir,
));
cache.load_from_disk().await;
let tls_fetch = config.censorship.tls_fetch.clone();
let fetch_context = TlsFetchContext {
cache: cache.clone(),
domains: tls_domains.to_vec(),
mask_host: config
.censorship
.mask_host
.clone()
.unwrap_or_else(|| config.censorship.tls_domain.clone());
let mask_unix_sock = config.censorship.mask_unix_sock.clone();
let tls_fetch_scope = (!config.censorship.tls_fetch_scope.is_empty())
.then(|| config.censorship.tls_fetch_scope.clone());
let tls_fetch = config.censorship.tls_fetch.clone();
let fetch_strategy = TlsFetchStrategy {
.unwrap_or_else(|| config.censorship.tls_domain.clone()),
primary_domain: config.censorship.tls_domain.clone(),
mask_unix_sock: config.censorship.mask_unix_sock.clone(),
tls_fetch_scope: (!config.censorship.tls_fetch_scope.is_empty())
.then(|| config.censorship.tls_fetch_scope.clone()),
upstream_manager,
strategy: TlsFetchStrategy {
profiles: tls_fetch.profiles,
strict_route: tls_fetch.strict_route,
attempt_timeout: Duration::from_millis(tls_fetch.attempt_timeout_ms.max(1)),
@@ -58,150 +161,71 @@ pub(crate) async fn bootstrap_tls_front(
grease_enabled: tls_fetch.grease_enabled,
deterministic: tls_fetch.deterministic,
profile_cache_ttl: Duration::from_secs(tls_fetch.profile_cache_ttl_secs),
};
let fetch_timeout = fetch_strategy.total_budget;
},
port: config.censorship.mask_port,
proxy_protocol: config.censorship.mask_proxy_protocol,
};
let cache_initial = cache.clone();
let domains_initial = tls_domains.to_vec();
let host_initial = mask_host.clone();
let primary_initial = config.censorship.tls_domain.clone();
let unix_sock_initial = mask_unix_sock.clone();
let scope_initial = tls_fetch_scope.clone();
let upstream_initial = upstream_manager.clone();
let strategy_initial = fetch_strategy.clone();
tokio::spawn(async move {
let mut join = tokio::task::JoinSet::new();
for domain in domains_initial {
let cache_domain = cache_initial.clone();
let host_domain =
tls_fetch_host_for_domain(&host_initial, &primary_initial, &domain);
let unix_sock_domain = unix_sock_initial.clone();
let scope_domain = scope_initial.clone();
let upstream_domain = upstream_initial.clone();
let strategy_domain = strategy_initial.clone();
join.spawn(async move {
match crate::tls_front::fetcher::fetch_real_tls_with_strategy(
&host_domain,
port,
&domain,
&strategy_domain,
Some(upstream_domain),
scope_domain.as_deref(),
proxy_protocol,
unix_sock_domain.as_deref(),
)
match policy {
TlsBootstrapPolicy::BestEffort => {
let initial_fetch = fetch_context.clone();
let fake_cert_len = config.censorship.fake_cert_len;
task_scope.spawn(async move {
initial_fetch
.fetch_all_with_budget("TLS emulation initial fetch failed")
.await;
for domain in initial_fetch
.cache
.default_profile_domains(&initial_fetch.domains)
.await
{
Ok(res) => cache_domain.update_from_fetch(&domain, res).await,
Err(e) => {
warn!(domain = %domain, error = %e, "TLS emulation initial fetch failed")
}
}
});
}
while let Some(res) = join.join_next().await {
if let Err(e) = res {
warn!(error = %e, "TLS emulation initial fetch task join failed");
}
}
});
let cache_timeout = cache.clone();
let domains_timeout = tls_domains.to_vec();
let fake_cert_len = config.censorship.fake_cert_len;
tokio::spawn(async move {
tokio::time::sleep(fetch_timeout).await;
for domain in domains_timeout {
let cached = cache_timeout.get(&domain).await;
if cached.domain == "default" {
{
warn!(
domain = %domain,
timeout_secs = fetch_timeout.as_secs(),
timeout_ms = initial_fetch.strategy.total_budget.as_millis(),
fake_cert_len,
"TLS-front fetch not ready within timeout; using cache/default fake cert fallback"
);
}
});
}
TlsBootstrapPolicy::RequireReady => {
fetch_context
.fetch_all_with_budget("TLS emulation initial fetch failed")
.await;
let default_domains = cache.default_profile_domains(tls_domains).await;
if let Some(error) = readiness_error(&default_domains) {
startup_tracker
.fail_component(COMPONENT_TLS_FRONT_BOOTSTRAP, Some(error.clone()))
.await;
return Err(ProxyError::Proxy(error));
}
});
let cache_refresh = cache.clone();
let domains_refresh = tls_domains.to_vec();
let host_refresh = mask_host.clone();
let primary_refresh = config.censorship.tls_domain.clone();
let unix_sock_refresh = mask_unix_sock.clone();
let scope_refresh = tls_fetch_scope.clone();
let upstream_refresh = upstream_manager.clone();
let strategy_refresh = fetch_strategy.clone();
tokio::spawn(async move {
loop {
let base_secs = rand::rng().random_range(4 * 3600..=6 * 3600);
let jitter_secs = rand::rng().random_range(0..=7200);
tokio::time::sleep(Duration::from_secs(base_secs + jitter_secs)).await;
let mut join = tokio::task::JoinSet::new();
for domain in domains_refresh.clone() {
let cache_domain = cache_refresh.clone();
let host_domain =
tls_fetch_host_for_domain(&host_refresh, &primary_refresh, &domain);
let unix_sock_domain = unix_sock_refresh.clone();
let scope_domain = scope_refresh.clone();
let upstream_domain = upstream_refresh.clone();
let strategy_domain = strategy_refresh.clone();
join.spawn(async move {
match crate::tls_front::fetcher::fetch_real_tls_with_strategy(
&host_domain,
port,
&domain,
&strategy_domain,
Some(upstream_domain),
scope_domain.as_deref(),
proxy_protocol,
unix_sock_domain.as_deref(),
)
.await
{
Ok(res) => cache_domain.update_from_fetch(&domain, res).await,
Err(e) => {
warn!(domain = %domain, error = %e, "TLS emulation refresh failed")
}
}
});
}
while let Some(res) = join.join_next().await {
if let Err(e) = res {
warn!(error = %e, "TLS emulation refresh task join failed");
}
}
}
});
Some(cache)
} else {
startup_tracker
.skip_component(
COMPONENT_TLS_FRONT_BOOTSTRAP,
Some("censorship.tls_emulation is false".to_string()),
)
.await;
None
};
if tls_cache.is_some() {
startup_tracker
.complete_component(
COMPONENT_TLS_FRONT_BOOTSTRAP,
Some("tls front cache is initialized".to_string()),
)
.await;
}
}
tls_cache
let refresh_context = fetch_context;
task_scope.spawn(async move {
loop {
let base_secs = rand::rng().random_range(4 * 3600..=6 * 3600);
let jitter_secs = rand::rng().random_range(0..=7200);
tokio::time::sleep(Duration::from_secs(base_secs + jitter_secs)).await;
refresh_context
.fetch_all_with_budget("TLS emulation refresh failed")
.await;
}
});
startup_tracker
.complete_component(
COMPONENT_TLS_FRONT_BOOTSTRAP,
Some("tls front cache is initialized".to_string()),
)
.await;
Ok(Some(cache))
}
#[cfg(test)]
mod tests {
use super::tls_fetch_host_for_domain;
use super::{readiness_error, tls_fetch_host_for_domain};
#[test]
fn tls_fetch_host_uses_each_domain_when_mask_host_is_primary_default() {
@@ -218,4 +242,13 @@ mod tests {
"origin.example"
);
}
#[test]
fn readiness_rejects_only_default_profiles() {
assert!(readiness_error(&[]).is_none());
assert_eq!(
readiness_error(&["front.example".to_string()]),
Some("TLS-front profiles are not ready for domains: front.example".to_string())
);
}
}

View File

@@ -4,12 +4,12 @@ use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Duration;
use arc_swap::ArcSwap;
use http_body_util::Full;
use hyper::body::Bytes;
use hyper::server::conn::http1;
use hyper::service::service_fn;
use hyper::{Request, Response, StatusCode};
use ipnetwork::IpNetwork;
use tokio::net::TcpListener;
use tokio::sync::Semaphore;
use tokio::time::timeout;
@@ -17,6 +17,7 @@ use tracing::{debug, info, warn};
use crate::config::ProxyConfig;
use crate::ip_tracker::UserIpTracker;
use crate::maestro::generation::RuntimeGeneration;
use crate::proxy::shared_state::ProxySharedState;
use crate::stats::Stats;
use crate::stats::beobachten::BeobachtenStore;
@@ -36,16 +37,8 @@ pub async fn serve(
port: u16,
listen: Option<String>,
listen_backlog: u32,
stats: Arc<Stats>,
beobachten: Arc<BeobachtenStore>,
shared_state: Arc<ProxySharedState>,
ip_tracker: Arc<UserIpTracker>,
tls_cache: Option<Arc<TlsFrontCache>>,
config_rx: tokio::sync::watch::Receiver<Arc<ProxyConfig>>,
whitelist: Vec<IpNetwork>,
active_runtime: Arc<ArcSwap<RuntimeGeneration>>,
) {
let whitelist = Arc::new(whitelist);
// If `metrics_listen` is set, bind on that single address only.
if let Some(ref listen_addr) = listen {
let addr: SocketAddr = match listen_addr.parse() {
@@ -61,17 +54,7 @@ pub async fn serve(
match bind_metrics_listener(addr, ipv6_only, listen_backlog) {
Ok(listener) => {
info!("Metrics endpoint: http://{}/metrics and /beobachten", addr);
serve_listener(
listener,
stats,
beobachten,
shared_state,
ip_tracker,
tls_cache,
config_rx,
whitelist,
)
.await;
serve_listener(listener, active_runtime).await;
}
Err(e) => {
warn!(error = %e, "Failed to bind metrics on {}", addr);
@@ -117,50 +100,14 @@ pub async fn serve(
warn!("Metrics listener is unavailable on both IPv4 and IPv6");
}
(Some(listener), None) | (None, Some(listener)) => {
serve_listener(
listener,
stats,
beobachten,
shared_state,
ip_tracker,
tls_cache,
config_rx,
whitelist,
)
.await;
serve_listener(listener, active_runtime).await;
}
(Some(listener4), Some(listener6)) => {
let stats_v6 = stats.clone();
let beobachten_v6 = beobachten.clone();
let shared_state_v6 = shared_state.clone();
let ip_tracker_v6 = ip_tracker.clone();
let tls_cache_v6 = tls_cache.clone();
let config_rx_v6 = config_rx.clone();
let whitelist_v6 = whitelist.clone();
let active_runtime_v6 = active_runtime.clone();
tokio::spawn(async move {
serve_listener(
listener6,
stats_v6,
beobachten_v6,
shared_state_v6,
ip_tracker_v6,
tls_cache_v6,
config_rx_v6,
whitelist_v6,
)
.await;
serve_listener(listener6, active_runtime_v6).await;
});
serve_listener(
listener4,
stats,
beobachten,
shared_state,
ip_tracker,
tls_cache,
config_rx,
whitelist,
)
.await;
serve_listener(listener4, active_runtime).await;
}
}
}
@@ -180,16 +127,7 @@ fn bind_metrics_listener(
TcpListener::from_std(socket.into())
}
async fn serve_listener(
listener: TcpListener,
stats: Arc<Stats>,
beobachten: Arc<BeobachtenStore>,
shared_state: Arc<ProxySharedState>,
ip_tracker: Arc<UserIpTracker>,
tls_cache: Option<Arc<TlsFrontCache>>,
config_rx: tokio::sync::watch::Receiver<Arc<ProxyConfig>>,
whitelist: Arc<Vec<IpNetwork>>,
) {
async fn serve_listener(listener: TcpListener, active_runtime: Arc<ArcSwap<RuntimeGeneration>>) {
let connection_permits = Arc::new(Semaphore::new(METRICS_MAX_CONTROL_CONNECTIONS));
loop {
@@ -201,7 +139,15 @@ async fn serve_listener(
}
};
if !whitelist.is_empty() && !whitelist.iter().any(|net| net.contains(peer.ip())) {
let runtime = active_runtime.load_full();
let config = runtime.config();
if !config.server.metrics_whitelist.is_empty()
&& !config
.server
.metrics_whitelist
.iter()
.any(|net| net.contains(peer.ip()))
{
debug!(peer = %peer, "Metrics request denied by whitelist");
continue;
}
@@ -218,21 +164,17 @@ async fn serve_listener(
}
};
let stats = stats.clone();
let beobachten = beobachten.clone();
let shared_state = shared_state.clone();
let ip_tracker = ip_tracker.clone();
let tls_cache = tls_cache.clone();
let config_rx_conn = config_rx.clone();
let active_runtime = active_runtime.clone();
tokio::spawn(async move {
let _connection_permit = connection_permit;
let svc = service_fn(move |req| {
let stats = stats.clone();
let beobachten = beobachten.clone();
let shared_state = shared_state.clone();
let ip_tracker = ip_tracker.clone();
let tls_cache = tls_cache.clone();
let config = config_rx_conn.borrow().clone();
let runtime = active_runtime.load_full();
let stats = runtime.stats.clone();
let beobachten = runtime.beobachten.clone();
let shared_state = runtime.proxy_shared.clone();
let ip_tracker = runtime.ip_tracker.clone();
let tls_cache = runtime.tls_cache.clone();
let config = runtime.config();
async move {
handle(
req,

View File

@@ -8,6 +8,33 @@ use crate::error::{ProxyError, Result};
type OverrideMap = HashMap<(String, u16), IpAddr>;
/// Immutable DNS override snapshot owned by one runtime generation.
#[derive(Debug, Clone, Default)]
pub struct DnsOverrides {
entries: std::sync::Arc<OverrideMap>,
}
impl DnsOverrides {
/// Parses a validated generation-local override snapshot.
pub fn from_entries(entries: &[String]) -> Result<Self> {
Ok(Self {
entries: std::sync::Arc::new(parse_entries(entries)?),
})
}
/// Resolves a generation-local hostname override.
pub fn resolve(&self, host: &str, port: u16) -> Option<IpAddr> {
self.entries
.get(&(host.to_ascii_lowercase(), port))
.copied()
}
/// Resolves a generation-local override as a socket address.
pub fn resolve_socket_addr(&self, host: &str, port: u16) -> Option<SocketAddr> {
self.resolve(host, port).map(|ip| SocketAddr::new(ip, port))
}
}
static DNS_OVERRIDES: OnceLock<RwLock<OverrideMap>> = OnceLock::new();
fn overrides_store() -> &'static RwLock<OverrideMap> {
@@ -180,6 +207,22 @@ mod tests {
assert_eq!(resolved, Some("127.0.0.1".parse().unwrap()));
}
#[test]
fn generation_snapshots_do_not_observe_each_other() {
let first = DnsOverrides::from_entries(&["example.com:443:127.0.0.1".to_string()]).unwrap();
let second =
DnsOverrides::from_entries(&["example.com:443:127.0.0.2".to_string()]).unwrap();
assert_eq!(
first.resolve("example.com", 443),
Some("127.0.0.1".parse().unwrap())
);
assert_eq!(
second.resolve("example.com", 443),
Some("127.0.0.2".parse().unwrap())
);
}
#[test]
fn split_host_port_parses_supported_shapes() {
assert_eq!(

View File

@@ -337,113 +337,111 @@ pub(crate) async fn resolve_direct_buffer_hard_limit(configured: usize) -> usize
align_down(derived as usize).max(DIRECT_BUFFER_UNIT_BYTES)
}
/// Starts the single control-plane task for Direct budget and shared pool pressure.
pub(crate) fn spawn_direct_buffer_budget_controller(
/// 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,
) {
tokio::spawn(async move {
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;
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);
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 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;
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),
);
}
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> {

View File

@@ -959,6 +959,7 @@ fn configure_mask_backend_socket(stream: &TcpStream) {
}
/// Handles a bad client by forwarding it to the configured mask target.
#[cfg(test)]
pub async fn handle_bad_client<R, W>(
reader: R,
writer: W,

View File

@@ -1,6 +1,5 @@
use crate::stats::UserStats;
use std::io;
use std::sync::atomic::Ordering;
#[derive(Debug)]
struct QuotaIoSentinel;
@@ -52,17 +51,5 @@ pub(super) fn refund_reserved_quota_bytes(user_stats: &UserStats, reserved_bytes
if reserved_bytes == 0 {
return;
}
let mut current = user_stats.quota_used.load(Ordering::Relaxed);
loop {
let next = current.saturating_sub(reserved_bytes);
match user_stats.quota_used.compare_exchange_weak(
current,
next,
Ordering::Relaxed,
Ordering::Relaxed,
) {
Ok(_) => return,
Err(observed) => current = observed,
}
}
user_stats.refund_quota(reserved_bytes);
}

View File

@@ -80,7 +80,11 @@ impl Stats {
return handle;
}
let entry = self.user_stats.entry(user.to_string()).or_default();
let quota = self.quota_store.user(user);
let entry = self
.user_stats
.entry(user.to_string())
.or_insert_with(|| Arc::new(UserStats::with_quota(quota)));
if entry.last_seen_epoch_secs.load(Ordering::Relaxed) == 0 {
self.touch_user_stats(entry.value().as_ref());
}
@@ -166,10 +170,7 @@ impl Stats {
#[inline]
pub(crate) fn quota_charge_post_write(&self, user_stats: &UserStats, bytes: u64) -> u64 {
self.touch_user_stats(user_stats);
user_stats
.quota_used
.fetch_add(bytes, Ordering::Relaxed)
.saturating_add(bytes)
user_stats.quota.charge(bytes)
}
pub(super) fn maybe_cleanup_user_stats(&self) {

View File

@@ -8,6 +8,7 @@ mod core_getters;
mod helpers;
mod me_counters;
mod me_getters;
mod quota_store;
mod replay;
pub mod telemetry;
pub mod tls_fingerprints;
@@ -20,6 +21,7 @@ use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU64, Ordering};
use std::time::Instant;
pub(crate) use self::quota_store::QuotaStore;
#[allow(unused_imports)]
pub use self::replay::{ReplayChecker, ReplayStats};
use self::telemetry::TelemetryPolicy;
@@ -344,6 +346,7 @@ pub struct Stats {
cached_epoch_secs: AtomicU64,
tls_fingerprints: tls_fingerprints::TlsFingerprintCollector,
user_stats: DashMap<String, Arc<UserStats>>,
quota_store: Arc<QuotaStore>,
user_stats_last_cleanup_epoch_secs: AtomicU64,
start_time: parking_lot::RwLock<Option<Instant>>,
}
@@ -356,12 +359,7 @@ pub struct UserStats {
pub octets_to_client: AtomicU64,
pub msgs_from_client: AtomicU64,
pub msgs_to_client: AtomicU64,
/// Total bytes charged against per-user quota admission.
///
/// This counter is the single source of truth for quota enforcement and
/// intentionally tracks attempted traffic, not guaranteed delivery.
pub quota_used: AtomicU64,
pub quota_last_reset_epoch_secs: AtomicU64,
quota: Arc<quota_store::UserQuotaCounters>,
pub last_seen_epoch_secs: AtomicU64,
}
@@ -378,9 +376,21 @@ pub enum QuotaReserveError {
}
impl UserStats {
fn with_quota(quota: Arc<quota_store::UserQuotaCounters>) -> Self {
Self {
quota,
..Self::default()
}
}
#[inline]
pub fn quota_used(&self) -> u64 {
self.quota_used.load(Ordering::Relaxed)
self.quota.used()
}
#[inline]
pub(crate) fn refund_quota(&self, bytes: u64) {
self.quota.refund(bytes);
}
/// Attempts one CAS reservation step against the quota counter.
@@ -390,27 +400,20 @@ impl UserStats {
/// with their own contention strategy.
#[inline]
pub fn quota_try_reserve(&self, bytes: u64, limit: u64) -> Result<u64, QuotaReserveError> {
let current = self.quota_used.load(Ordering::Relaxed);
if bytes > limit.saturating_sub(current) {
return Err(QuotaReserveError::LimitExceeded);
}
let next = current.saturating_add(bytes);
match self.quota_used.compare_exchange_weak(
current,
next,
Ordering::Relaxed,
Ordering::Relaxed,
) {
Ok(_) => Ok(next),
Err(_) => Err(QuotaReserveError::Contended),
}
self.quota.try_reserve(bytes, limit)
}
}
impl Stats {
pub fn new() -> Self {
let stats = Self::default();
Self::with_quota_store(Arc::new(QuotaStore::default()))
}
pub(crate) fn with_quota_store(quota_store: Arc<QuotaStore>) -> Self {
let stats = Self {
quota_store,
..Self::default()
};
stats.apply_telemetry_policy(TelemetryPolicy::default());
stats.refresh_cached_epoch_secs();
*stats.start_time.write() = Some(Instant::now());

147
src/stats/quota_store.rs Normal file
View File

@@ -0,0 +1,147 @@
use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use dashmap::DashMap;
use super::{QuotaReserveError, UserQuotaSnapshot};
/// Process-scoped per-user quota accounting shared by runtime generations.
#[derive(Default)]
pub struct QuotaStore {
users: DashMap<String, Arc<UserQuotaCounters>>,
}
/// Atomic quota state for one configured user.
#[derive(Default)]
pub(crate) struct UserQuotaCounters {
used_bytes: AtomicU64,
last_reset_epoch_secs: AtomicU64,
}
impl QuotaStore {
pub(crate) fn user(&self, user: &str) -> Arc<UserQuotaCounters> {
if let Some(existing) = self.users.get(user) {
return Arc::clone(existing.value());
}
Arc::clone(
self.users
.entry(user.to_string())
.or_insert_with(|| Arc::new(UserQuotaCounters::default()))
.value(),
)
}
pub(crate) fn used(&self, user: &str) -> u64 {
self.users.get(user).map(|state| state.used()).unwrap_or(0)
}
pub(crate) fn load(&self, user: &str, used_bytes: u64, last_reset_epoch_secs: u64) {
let state = self.user(user);
state.used_bytes.store(used_bytes, Ordering::Relaxed);
state
.last_reset_epoch_secs
.store(last_reset_epoch_secs, Ordering::Relaxed);
}
pub(crate) fn reset(&self, user: &str, now_epoch_secs: u64) -> UserQuotaSnapshot {
let state = self.user(user);
state.used_bytes.store(0, Ordering::Relaxed);
state
.last_reset_epoch_secs
.store(now_epoch_secs, Ordering::Relaxed);
UserQuotaSnapshot {
used_bytes: 0,
last_reset_epoch_secs: now_epoch_secs,
}
}
pub(crate) fn snapshot(&self) -> HashMap<String, UserQuotaSnapshot> {
let mut out = HashMap::new();
for entry in self.users.iter() {
let state = entry.value();
let used_bytes = state.used();
let last_reset_epoch_secs = state.last_reset_epoch_secs.load(Ordering::Relaxed);
if used_bytes == 0 && last_reset_epoch_secs == 0 {
continue;
}
out.insert(
entry.key().clone(),
UserQuotaSnapshot {
used_bytes,
last_reset_epoch_secs,
},
);
}
out
}
}
impl UserQuotaCounters {
#[inline]
pub(crate) fn used(&self) -> u64 {
self.used_bytes.load(Ordering::Relaxed)
}
#[inline]
pub(crate) fn charge(&self, bytes: u64) -> u64 {
self.used_bytes
.fetch_add(bytes, Ordering::Relaxed)
.saturating_add(bytes)
}
#[inline]
pub(crate) fn refund(&self, bytes: u64) {
let mut current = self.used_bytes.load(Ordering::Relaxed);
loop {
let next = current.saturating_sub(bytes);
match self.used_bytes.compare_exchange_weak(
current,
next,
Ordering::Relaxed,
Ordering::Relaxed,
) {
Ok(_) => return,
Err(observed) => current = observed,
}
}
}
#[inline]
pub(crate) fn try_reserve(&self, bytes: u64, limit: u64) -> Result<u64, QuotaReserveError> {
let current = self.used_bytes.load(Ordering::Relaxed);
if bytes > limit.saturating_sub(current) {
return Err(QuotaReserveError::LimitExceeded);
}
let next = current.saturating_add(bytes);
match self.used_bytes.compare_exchange_weak(
current,
next,
Ordering::Relaxed,
Ordering::Relaxed,
) {
Ok(_) => Ok(next),
Err(_) => Err(QuotaReserveError::Contended),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::stats::Stats;
#[test]
fn quota_counters_are_shared_across_stats_generations() {
let store = Arc::new(QuotaStore::default());
let first = Stats::with_quota_store(store.clone());
store.user("alice").charge(512);
assert_eq!(first.get_user_quota_used("alice"), 512);
let second = Stats::with_quota_store(store);
assert_eq!(second.get_user_quota_used("alice"), 512);
second.reset_user_quota("alice");
assert_eq!(first.get_user_quota_used("alice"), 0);
}
}

View File

@@ -119,51 +119,21 @@ impl Stats {
}
pub fn get_user_quota_used(&self, user: &str) -> u64 {
self.user_stats
.get(user)
.map(|s| s.quota_used.load(Ordering::Relaxed))
.unwrap_or(0)
self.quota_store.used(user)
}
pub fn load_user_quota_state(&self, user: &str, used_bytes: u64, last_reset_epoch_secs: u64) {
let stats = self.get_or_create_user_stats_handle(user);
stats.quota_used.store(used_bytes, Ordering::Relaxed);
stats
.quota_last_reset_epoch_secs
.store(last_reset_epoch_secs, Ordering::Relaxed);
self.quota_store
.load(user, used_bytes, last_reset_epoch_secs);
}
pub fn reset_user_quota(&self, user: &str) -> UserQuotaSnapshot {
let stats = self.get_or_create_user_stats_handle(user);
let last_reset_epoch_secs = Self::now_epoch_secs();
stats.quota_used.store(0, Ordering::Relaxed);
stats
.quota_last_reset_epoch_secs
.store(last_reset_epoch_secs, Ordering::Relaxed);
UserQuotaSnapshot {
used_bytes: 0,
last_reset_epoch_secs,
}
self.quota_store.reset(user, last_reset_epoch_secs)
}
pub fn user_quota_snapshot(&self) -> HashMap<String, UserQuotaSnapshot> {
let mut out = HashMap::new();
for entry in self.user_stats.iter() {
let stats = entry.value();
let used_bytes = stats.quota_used.load(Ordering::Relaxed);
let last_reset_epoch_secs = stats.quota_last_reset_epoch_secs.load(Ordering::Relaxed);
if used_bytes == 0 && last_reset_epoch_secs == 0 {
continue;
}
out.insert(
entry.key().clone(),
UserQuotaSnapshot {
used_bytes,
last_reset_epoch_secs,
},
);
}
out
self.quota_store.snapshot()
}
pub fn get_handshake_timeouts(&self) -> u64 {

View File

@@ -4,6 +4,7 @@ use tokio::sync::watch;
use tracing::warn;
use crate::config::{ProxyConfig, SynLimitMode};
use crate::maestro::generation::RuntimeWatchState;
mod command;
mod iptables;
@@ -15,28 +16,105 @@ use self::model::{SynLimitNamespace, synlimit_namespace, synlimit_targets};
static ACTIVE_SYNLIMIT_NAMESPACE: Mutex<Option<SynLimitNamespace>> = Mutex::new(None);
pub(crate) fn spawn_synlimit_controller(config_rx: watch::Receiver<Arc<ProxyConfig>>) {
/// Spawns the process-scoped SYN limiter reconciler for active generations.
pub(crate) fn spawn_synlimit_controller(
runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
) -> tokio::task::JoinHandle<()> {
if !cfg!(target_os = "linux") {
if has_synlimit_config(&config_rx.borrow()) {
warn!("SYN limiter is configured but unsupported on this OS; skipping netfilter rules");
}
return;
return tokio::spawn(watch_active_runtime_configs(
runtime_watch_rx,
|_generation_id, cfg| async move {
if has_synlimit_config(&cfg) {
warn!(
"SYN limiter is configured but unsupported on this OS; skipping netfilter rules"
);
}
},
));
}
tokio::spawn(async move {
wait_for_config_channel_close_and_reconcile(config_rx).await;
if let Err(error) = clear_synlimit_rules_all_backends().await {
warn!(error = %error, "Failed to clear SYN limiter rules after config channel close");
}
});
tokio::spawn(watch_active_runtime_configs(
runtime_watch_rx,
|_generation_id, cfg| async move {
reconcile_synlimit_rules(&cfg).await;
},
))
}
async fn wait_for_config_channel_close_and_reconcile(
mut config_rx: watch::Receiver<Arc<ProxyConfig>>,
) {
while config_rx.changed().await.is_ok() {
let cfg = config_rx.borrow_and_update().clone();
reconcile_synlimit_rules(&cfg).await;
async fn watch_active_runtime_configs<F, Fut>(
mut runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
mut on_config: F,
)
where
F: FnMut(u64, Arc<ProxyConfig>) -> Fut,
Fut: std::future::Future<Output = ()>,
{
let mut current = loop {
if let Some(state) = runtime_watch_rx.borrow().clone() {
break state;
}
if runtime_watch_rx.changed().await.is_err() {
return;
}
};
let initial_config = current.config_rx.borrow().clone();
on_config(current.generation_id, initial_config).await;
loop {
tokio::select! {
biased;
changed = runtime_watch_rx.changed() => {
if changed.is_err() {
break;
}
let Some(next) = runtime_watch_rx.borrow().clone() else {
continue;
};
if next.generation_id != current.generation_id {
current = next;
let config = current.config_rx.borrow().clone();
on_config(current.generation_id, config).await;
}
}
changed = current.config_rx.changed() => {
if changed.is_err() {
let Some(next) = wait_for_new_runtime(
&mut runtime_watch_rx,
current.generation_id,
).await else {
break;
};
current = next;
let config = current.config_rx.borrow().clone();
on_config(current.generation_id, config).await;
continue;
}
let active_generation_id = runtime_watch_rx
.borrow()
.as_ref()
.map(|state| state.generation_id);
if active_generation_id == Some(current.generation_id) {
let cfg = current.config_rx.borrow_and_update().clone();
on_config(current.generation_id, cfg).await;
}
}
}
}
}
async fn wait_for_new_runtime(
runtime_watch_rx: &mut watch::Receiver<Option<RuntimeWatchState>>,
previous_generation_id: u64,
) -> Option<RuntimeWatchState> {
loop {
if let Some(state) = runtime_watch_rx.borrow().clone()
&& state.generation_id != previous_generation_id
{
return Some(state);
}
if runtime_watch_rx.changed().await.is_err() {
return None;
}
}
}
@@ -168,3 +246,70 @@ fn has_synlimit_config(cfg: &ProxyConfig) -> bool {
.iter()
.any(|listener| !matches!(listener.synlimit, SynLimitMode::Off))
}
#[cfg(test)]
mod tests {
use super::*;
use std::time::Duration;
use tokio::sync::mpsc;
fn runtime_state(
generation_id: u64,
max_connections: u32,
) -> (
RuntimeWatchState,
watch::Sender<Arc<ProxyConfig>>,
watch::Sender<bool>,
) {
let mut config = ProxyConfig::default();
config.server.max_connections = max_connections;
let (config_tx, config_rx) = watch::channel(Arc::new(config));
let (admission_tx, admission_rx) = watch::channel(true);
(
RuntimeWatchState {
generation_id,
config_rx,
admission_rx,
},
config_tx,
admission_tx,
)
}
#[tokio::test]
async fn config_watcher_ignores_retired_generation_updates() {
let (initial, initial_config_tx, _initial_admission_tx) = runtime_state(1, 10);
let (runtime_tx, runtime_rx) = watch::channel(Some(initial));
let (observed_tx, mut observed_rx) = mpsc::unbounded_channel();
let watcher = tokio::spawn(watch_active_runtime_configs(
runtime_rx,
move |generation_id, cfg| {
let observed_tx = observed_tx.clone();
async move {
let _ = observed_tx.send((generation_id, cfg.server.max_connections));
}
},
));
assert_eq!(observed_rx.recv().await, Some((1, 10)));
let (next, next_config_tx, _next_admission_tx) = runtime_state(2, 20);
runtime_tx.send_replace(Some(next));
assert_eq!(observed_rx.recv().await, Some((2, 20)));
let mut stale = ProxyConfig::default();
stale.server.max_connections = 30;
initial_config_tx.send_replace(Arc::new(stale));
assert!(
tokio::time::timeout(Duration::from_millis(50), observed_rx.recv())
.await
.is_err()
);
let mut active = ProxyConfig::default();
active.server.max_connections = 40;
next_config_tx.send_replace(Arc::new(active));
assert_eq!(observed_rx.recv().await, Some((2, 40)));
watcher.abort();
}
}

View File

@@ -190,6 +190,22 @@ impl TlsFrontCache {
(snapshot, suppressed)
}
/// Returns configured domains that still resolve to the synthetic default profile.
pub(crate) async fn default_profile_domains(&self, domains: &[String]) -> Vec<String> {
let guard = self.memory.read().await;
domains
.iter()
.filter(|domain| {
guard
.get(domain.as_str())
.unwrap_or(&self.default)
.domain
== "default"
})
.cloned()
.collect()
}
fn full_cert_sent_shard_index(client_ip: IpAddr) -> usize {
let mut hasher = DefaultHasher::new();
client_ip.hash(&mut hasher);
@@ -546,6 +562,23 @@ mod tests {
assert!(!cert_info_matches_domain(&cached));
}
#[tokio::test]
async fn default_profile_domains_reports_only_unprepared_entries() {
let domains = vec!["ready.example".to_string(), "pending.example".to_string()];
let cache = TlsFrontCache::new(&domains, 1024, "tlsfront-test-cache");
cache
.set(
"ready.example",
cached_with_cert_info("ready.example", None, Vec::new()),
)
.await;
assert_eq!(
cache.default_profile_domains(&domains).await,
vec!["pending.example".to_string()]
);
}
#[tokio::test]
async fn test_take_full_cert_budget_for_ip_uses_ttl() {
let cache = TlsFrontCache::new(&["example.com".to_string()], 1024, "tlsfront-test-cache");

View File

@@ -916,26 +916,22 @@ async fn connect_tcp_with_upstream(
strict_route: bool,
) -> Result<UpstreamStream> {
if let Some(manager) = upstream {
let resolved = if let Some(addr) = resolve_socket_addr(host, port) {
Some(addr)
} else {
match tokio::net::lookup_host((host, port)).await {
Ok(mut addrs) => addrs.find(|a| a.is_ipv4()),
Err(e) => {
if strict_route {
return Err(anyhow!(
"upstream route DNS resolution failed for {host}:{port}: {e}"
));
}
warn!(
host = %host,
port = port,
scope = ?scope,
error = %e,
"Upstream DNS resolution failed, using direct connect"
);
None
let resolved = match manager.resolve_hostname(host, port).await {
Ok(addr) => Some(addr),
Err(e) => {
if strict_route {
return Err(anyhow!(
"upstream route DNS resolution failed for {host}:{port}: {e}"
));
}
warn!(
host = %host,
port = port,
scope = ?scope,
error = %e,
"Upstream DNS resolution failed, using direct connect"
);
None
}
};
@@ -955,6 +951,9 @@ async fn connect_tcp_with_upstream(
error = %e,
"Upstream connect failed, using direct connect"
);
return Ok(UpstreamStream::Tcp(
timeout(connect_timeout, TcpStream::connect(addr)).await??,
));
}
}
} else if strict_route {

View File

@@ -66,33 +66,13 @@ fn extract_host_port_path(url: &str) -> Result<(String, u16, String)> {
Ok((host, port, path))
}
async fn resolve_target_addr(host: &str, port: u16) -> Result<std::net::SocketAddr> {
if let Some(addr) = resolve_socket_addr(host, port) {
return Ok(addr);
}
let addrs: Vec<std::net::SocketAddr> = tokio::net::lookup_host((host, port))
.await
.map_err(|e| ProxyError::Proxy(format!("DNS resolve failed for {host}:{port}: {e}")))?
.collect();
if let Some(addr) = addrs.iter().copied().find(|addr| addr.is_ipv4()) {
return Ok(addr);
}
addrs
.first()
.copied()
.ok_or_else(|| ProxyError::Proxy(format!("DNS returned no addresses for {host}:{port}")))
}
async fn connect_https_transport(
host: &str,
port: u16,
upstream: Option<Arc<UpstreamManager>>,
) -> Result<UpstreamStream> {
if let Some(manager) = upstream {
let target = resolve_target_addr(host, port).await?;
let target = manager.resolve_hostname(host, port).await?;
return timeout(HTTP_CONNECT_TIMEOUT, manager.connect(target, None, None))
.await
.map_err(|_| ProxyError::Proxy(format!("upstream connect timeout for {host}:{port}")))?

View File

@@ -4,6 +4,7 @@
#![allow(deprecated)]
use arc_swap::ArcSwap;
use rand::RngExt;
use std::collections::{BTreeSet, HashMap};
use std::net::{IpAddr, SocketAddr};
@@ -20,7 +21,7 @@ use tracing::{debug, info, trace, warn};
use crate::config::{UpstreamConfig, UpstreamType};
use crate::error::{ProxyError, Result};
use crate::network::dns_overrides::{resolve_socket_addr, split_host_port};
use crate::network::dns_overrides::{DnsOverrides, split_host_port};
use crate::protocol::constants::{TG_DATACENTER_PORT, TG_DATACENTERS_V4, TG_DATACENTERS_V6};
use crate::stats::Stats;
use crate::transport::shadowsocks::{
@@ -333,6 +334,7 @@ pub struct UpstreamManager {
no_upstreams_warn_epoch_ms: Arc<AtomicU64>,
no_healthy_warn_epoch_ms: Arc<AtomicU64>,
stats: Arc<Stats>,
dns_overrides: Arc<ArcSwap<DnsOverrides>>,
}
impl UpstreamManager {
@@ -374,9 +376,37 @@ impl UpstreamManager {
no_upstreams_warn_epoch_ms: Arc::new(AtomicU64::new(0)),
no_healthy_warn_epoch_ms: Arc::new(AtomicU64::new(0)),
stats,
dns_overrides: Arc::new(ArcSwap::from_pointee(DnsOverrides::default())),
}
}
pub(crate) fn with_dns_overrides(mut self, entries: &[String]) -> Result<Self> {
self.dns_overrides = Arc::new(ArcSwap::from_pointee(DnsOverrides::from_entries(entries)?));
Ok(self)
}
pub(crate) fn update_dns_overrides(&self, entries: &[String]) -> Result<()> {
let snapshot = DnsOverrides::from_entries(entries)?;
self.dns_overrides.store(Arc::new(snapshot));
Ok(())
}
pub(crate) async fn resolve_hostname(&self, host: &str, port: u16) -> Result<SocketAddr> {
if let Some(addr) = self.dns_overrides.load().resolve_socket_addr(host, port) {
return Ok(addr);
}
let addrs: Vec<SocketAddr> = tokio::net::lookup_host((host, port))
.await
.map_err(ProxyError::Io)?
.collect();
if let Some(addr) = addrs.iter().copied().find(SocketAddr::is_ipv4) {
return Ok(addr);
}
addrs.first().copied().ok_or_else(|| {
ProxyError::Proxy(format!("DNS returned no addresses for {host}:{port}"))
})
}
fn now_epoch_ms() -> u64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
@@ -709,11 +739,12 @@ impl UpstreamManager {
}
async fn connect_hostname_with_dns_override(
&self,
address: &str,
connect_timeout: Duration,
) -> Result<TcpStream> {
if let Some((host, port)) = split_host_port(address)
&& let Some(addr) = resolve_socket_addr(&host, port)
&& let Some(addr) = self.dns_overrides.load().resolve_socket_addr(&host, port)
{
return match tokio::time::timeout(connect_timeout, TcpStream::connect(addr)).await {
Ok(Ok(stream)) => Ok(stream),
@@ -1203,7 +1234,8 @@ impl UpstreamManager {
"SOCKS4 interface binding is not supported for hostname addresses, ignoring"
);
}
Self::connect_hostname_with_dns_override(address, connect_timeout).await?
self.connect_hostname_with_dns_override(address, connect_timeout)
.await?
};
// replace socks user_id with config.selected_scope, if set
@@ -1291,7 +1323,8 @@ impl UpstreamManager {
"SOCKS5 interface binding is not supported for hostname addresses, ignoring"
);
}
Self::connect_hostname_with_dns_override(address, connect_timeout).await?
self.connect_hostname_with_dns_override(address, connect_timeout)
.await?
};
debug!(config = ?config, "Socks5 connection");