mirror of
https://github.com/telemt/telemt.git
synced 2026-07-20 20:50:20 +03:00
Compare commits
16 Commits
flow-aescr
...
main
| Author | SHA1 | Date | |
|---|---|---|---|
| d851200e47 | |||
| a5216d77fb | |||
| 5b5cd952c8 | |||
| a40578b278 | |||
| bd99493622 | |||
| d862deebb2 | |||
| 0d869d716c | |||
| 1e85b91ad3 | |||
| 4679bdcfd5 | |||
| 7291c3192c | |||
| 4a5dc0b21b | |||
| fabd98ce89 | |||
| 7df3dab5e8 | |||
| c6f40e3717 | |||
| 91e05265be | |||
| 991d5b2c38 |
@@ -0,0 +1,49 @@
|
||||
name: Coverage
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
branches: [ "*" ]
|
||||
workflow_dispatch:
|
||||
|
||||
env:
|
||||
CARGO_TERM_COLOR: always
|
||||
|
||||
jobs:
|
||||
coverage:
|
||||
name: LLVM coverage report
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- uses: dtolnay/rust-toolchain@stable
|
||||
with:
|
||||
components: llvm-tools-preview
|
||||
|
||||
- name: Cache cargo
|
||||
uses: actions/cache@v4
|
||||
with:
|
||||
path: |
|
||||
~/.cargo/bin
|
||||
~/.cargo/registry
|
||||
~/.cargo/git
|
||||
target
|
||||
key: ${{ runner.os }}-cargo-llvm-cov-${{ hashFiles('**/Cargo.lock') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-cargo-llvm-cov-
|
||||
${{ runner.os }}-cargo-
|
||||
|
||||
- name: Install cargo-llvm-cov
|
||||
run: cargo install --locked cargo-llvm-cov || true
|
||||
|
||||
- name: Generate LCOV report
|
||||
run: cargo llvm-cov --locked --lcov --output-path lcov.info
|
||||
|
||||
- name: Upload LCOV report
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: telemt-lcov
|
||||
path: lcov.info
|
||||
Generated
+1
-1
@@ -2900,7 +2900,7 @@ checksum = "7b2093cf4c8eb1e67749a6762251bc9cd836b6fc171623bd0a9d324d37af2417"
|
||||
|
||||
[[package]]
|
||||
name = "telemt"
|
||||
version = "3.4.24"
|
||||
version = "3.4.25"
|
||||
dependencies = [
|
||||
"aes",
|
||||
"anyhow",
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "telemt"
|
||||
version = "3.4.24"
|
||||
version = "3.4.25"
|
||||
edition = "2024"
|
||||
|
||||
[features]
|
||||
|
||||
@@ -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 |
|
||||
@@ -324,6 +334,7 @@ Returned by `PATCH /v1/config` on success (`200`).
|
||||
| `connections_bad_total` | `u64` | Failed/invalid client connections. |
|
||||
| `connections_bad_by_class` | `ClassCount[]` | Failed/invalid connections grouped by class. |
|
||||
| `handshake_failures_by_class` | `ClassCount[]` | Handshake failures grouped by class. |
|
||||
| `handshake_failures_by_stage` | `StageCount[]` | Handshake failures grouped by state-machine stage. |
|
||||
| `handshake_timeouts_total` | `u64` | Handshake timeout count. |
|
||||
| `configured_users` | `usize` | Number of configured users in config. |
|
||||
|
||||
@@ -333,6 +344,38 @@ Returned by `PATCH /v1/config` on success (`200`).
|
||||
| `class` | `string` | Failure class label. |
|
||||
| `total` | `u64` | Counter value for this class. |
|
||||
|
||||
#### `StageCount`
|
||||
| Field | Type | Description |
|
||||
| --- | --- | --- |
|
||||
| `stage` | `string` | State-machine stage label. |
|
||||
| `total` | `u64` | Counter value for this stage. |
|
||||
|
||||
#### Handshake failure stage diagnostics
|
||||
|
||||
`handshake_failures_by_class` and `telemt_handshake_failures_by_class_total` describe the error kind. `handshake_failures_by_stage` and `telemt_handshake_failures_by_stage_total` describe where the same failure happened in the handshake state machine.
|
||||
|
||||
This does not add a DPI verdict or any protocol decision. The stage is derived from the existing Telemt handshake control flow and is counted only when the existing handshake failure or timeout accounting path is reached.
|
||||
|
||||
Fixed stage labels:
|
||||
|
||||
| Stage | Meaning |
|
||||
| --- | --- |
|
||||
| `first_packet_prelude` | Reading the first 5 bytes before selecting the TLS or direct branch. |
|
||||
| `tls_clienthello_body` | Reading the TLS ClientHello body after the TLS record header. |
|
||||
| `tls_core` | Running the TLS-F handshake/auth flow. |
|
||||
| `tls_post_serverhello_mtproto` | Waiting for the 64-byte MTProto handshake after TLS ServerHello. |
|
||||
| `direct_mtproto` | Reading the direct classic/secure 64-byte MTProto handshake. |
|
||||
|
||||
Example:
|
||||
|
||||
```text
|
||||
telemt_handshake_failures_by_class_total{class="expected_64_got_0_unexpected_eof"} 3
|
||||
telemt_handshake_failures_by_stage_total{stage="direct_mtproto"} 1
|
||||
telemt_handshake_failures_by_stage_total{stage="tls_post_serverhello_mtproto"} 2
|
||||
```
|
||||
|
||||
This means the same EOF-while-reading-64-bytes failure happened once in the direct MTProto path and twice after TLS ServerHello.
|
||||
|
||||
### `SystemInfoData`
|
||||
| Field | Type | Description |
|
||||
| --- | --- | --- |
|
||||
@@ -917,6 +960,7 @@ JA3 follows the Salesforce ClientHello field order. JA4 follows the FoxIO TLS-cl
|
||||
| `connections_bad_total` | `u64` | Failed/invalid connections. |
|
||||
| `connections_bad_by_class` | `ClassCount[]` | Failed/invalid connections grouped by class. |
|
||||
| `handshake_failures_by_class` | `ClassCount[]` | Handshake failures grouped by class. |
|
||||
| `handshake_failures_by_stage` | `StageCount[]` | Handshake failures grouped by state-machine stage. |
|
||||
| `handshake_timeouts_total` | `u64` | Handshake timeouts. |
|
||||
| `accept_permit_timeout_total` | `u64` | Listener admission permit acquisition timeouts. |
|
||||
| `configured_users` | `usize` | Configured user count. |
|
||||
@@ -1376,24 +1420,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 +1471,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 +1483,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`. |
|
||||
|
||||
+51
-2
@@ -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"));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
+170
-13
@@ -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::{ReloadAccepted, ReloadControl, ReloadRequest, ReloadSubmitError};
|
||||
use crate::proxy::route_mode::RouteRuntimeController;
|
||||
use crate::proxy::shared_state::ProxySharedState;
|
||||
use crate::startup::StartupTracker;
|
||||
@@ -29,11 +32,13 @@ 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;
|
||||
mod patch;
|
||||
#[cfg(test)]
|
||||
mod reload_tests;
|
||||
mod runtime_edge;
|
||||
mod runtime_init;
|
||||
mod runtime_min;
|
||||
@@ -44,7 +49,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 +113,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 +132,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 +178,41 @@ 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())
|
||||
}
|
||||
|
||||
async fn submit_reload_from_disk(
|
||||
config_path: &std::path::Path,
|
||||
mutation_lock: &Mutex<()>,
|
||||
reload_control: &ReloadControl,
|
||||
expected_revision: Option<&str>,
|
||||
request: ReloadRequest,
|
||||
) -> Result<(ReloadAccepted, String), ApiFailure> {
|
||||
let _guard = mutation_lock.lock().await;
|
||||
ensure_expected_revision(config_path, expected_revision).await?;
|
||||
let revision = current_revision(config_path).await?;
|
||||
let config = Arc::new(load_config_for_reload(config_path).await?);
|
||||
let accepted = 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((accepted, revision))
|
||||
}
|
||||
|
||||
fn allowed_methods_for_path(path: &str) -> Option<&'static str> {
|
||||
match path {
|
||||
"/v1/health"
|
||||
@@ -175,12 +244,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 +271,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 +333,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 +342,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 +376,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 +410,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 +754,33 @@ 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 (accepted, revision) = submit_reload_from_disk(
|
||||
&shared.config_path,
|
||||
shared.mutation_lock.as_ref(),
|
||||
&shared.reload_control,
|
||||
expected_revision.as_deref(),
|
||||
request,
|
||||
)
|
||||
.await?;
|
||||
Ok(success_response(StatusCode::ACCEPTED, accepted, revision))
|
||||
}
|
||||
("PATCH", "/v1/config") => {
|
||||
if api_cfg.read_only {
|
||||
return Ok(error_response(
|
||||
@@ -663,11 +793,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 +817,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/")
|
||||
|
||||
@@ -0,0 +1,119 @@
|
||||
use super::*;
|
||||
use crate::config::ProxyConfig;
|
||||
|
||||
async fn config_file() -> (tempfile::TempDir, PathBuf, String) {
|
||||
let directory = tempfile::tempdir().unwrap();
|
||||
let path = directory.path().join("config.toml");
|
||||
let mut config = ProxyConfig::default();
|
||||
config.server.max_connections = 4_242;
|
||||
let body = toml::to_string_pretty(&config).unwrap();
|
||||
tokio::fs::write(&path, &body).await.unwrap();
|
||||
let revision = config_store::compute_revision(&body);
|
||||
(directory, path, revision)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reload_submission_uses_matching_disk_revision_and_snapshot() {
|
||||
let (_directory, path, revision) = config_file().await;
|
||||
let mutation_lock = Mutex::new(());
|
||||
let (control, mut commands) = ReloadControl::channel(1);
|
||||
let request = ReloadRequest::default();
|
||||
|
||||
let (accepted, response_revision) = submit_reload_from_disk(
|
||||
&path,
|
||||
&mutation_lock,
|
||||
&control,
|
||||
Some(&revision),
|
||||
request.clone(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let command = commands.recv().await.unwrap();
|
||||
|
||||
assert_eq!(response_revision, revision);
|
||||
assert_eq!(accepted.config_revision, revision);
|
||||
assert_eq!(command.config_revision, revision);
|
||||
assert_eq!(command.request, request);
|
||||
assert_eq!(command.config.server.max_connections, 4_242);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn revision_conflict_rejects_without_enqueuing_reload() {
|
||||
let (_directory, path, _revision) = config_file().await;
|
||||
let mutation_lock = Mutex::new(());
|
||||
let (control, _commands) = ReloadControl::channel(1);
|
||||
|
||||
let error = submit_reload_from_disk(
|
||||
&path,
|
||||
&mutation_lock,
|
||||
&control,
|
||||
Some("stale-revision"),
|
||||
ReloadRequest::default(),
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
|
||||
assert_eq!(error.status, StatusCode::CONFLICT);
|
||||
assert_eq!(error.code, "revision_conflict");
|
||||
assert_eq!(control.in_progress().await, None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reload_conflict_and_closed_coordinator_map_to_http_contract() {
|
||||
let (_directory, path, _revision) = config_file().await;
|
||||
let mutation_lock = Mutex::new(());
|
||||
let (control, mut commands) = ReloadControl::channel(1);
|
||||
let _accepted = submit_reload_from_disk(
|
||||
&path,
|
||||
&mutation_lock,
|
||||
&control,
|
||||
None,
|
||||
ReloadRequest::default(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let _command = commands.recv().await.unwrap();
|
||||
|
||||
let conflict = submit_reload_from_disk(
|
||||
&path,
|
||||
&mutation_lock,
|
||||
&control,
|
||||
None,
|
||||
ReloadRequest::default(),
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert_eq!(conflict.status, StatusCode::CONFLICT);
|
||||
assert_eq!(conflict.code, "reload_in_progress");
|
||||
|
||||
control.fail(1, "test cleanup").await;
|
||||
drop(commands);
|
||||
let unavailable = submit_reload_from_disk(
|
||||
&path,
|
||||
&mutation_lock,
|
||||
&control,
|
||||
None,
|
||||
ReloadRequest::default(),
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert_eq!(unavailable.status, StatusCode::SERVICE_UNAVAILABLE);
|
||||
assert_eq!(unavailable.code, "maestro_unavailable");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reload_routes_expose_only_documented_methods_and_ids() {
|
||||
assert_eq!(
|
||||
allowed_methods_for_path("/v1/system/reload"),
|
||||
Some(ALLOW_POST)
|
||||
);
|
||||
assert_eq!(
|
||||
allowed_methods_for_path("/v1/system/reload/42"),
|
||||
Some(ALLOW_GET)
|
||||
);
|
||||
assert_eq!(reload_status_route_id("/v1/system/reload/42"), Some(42));
|
||||
assert_eq!(
|
||||
reload_status_route_id("/v1/system/reload/not-a-number"),
|
||||
None
|
||||
);
|
||||
}
|
||||
+292
-39
@@ -4,58 +4,184 @@ 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");
|
||||
}
|
||||
});
|
||||
let _config_watcher = spawn_config_watcher(
|
||||
runtime_watch_rx.clone(),
|
||||
runtime_state.clone(),
|
||||
runtime_events.clone(),
|
||||
);
|
||||
let _admission_watcher =
|
||||
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::task::JoinHandle<()> {
|
||||
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::task::JoinHandle<()> {
|
||||
tokio::spawn(async move {
|
||||
let Some(mut current) = runtime_watch_rx.borrow().clone() else {
|
||||
return;
|
||||
};
|
||||
record_admission_state(&runtime_state, &runtime_events, ¤t);
|
||||
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, ¤t);
|
||||
}
|
||||
}
|
||||
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, ¤t);
|
||||
continue;
|
||||
}
|
||||
if active_generation_id(&runtime_watch_rx) == Some(current.generation_id) {
|
||||
record_admission_state(&runtime_state, &runtime_events, ¤t);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
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 {
|
||||
@@ -64,3 +190,130 @@ fn now_epoch_secs() -> u64 {
|
||||
.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
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn watcher_recovers_from_closed_generation_and_exits_with_process_channel() {
|
||||
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));
|
||||
let watcher = spawn_config_watcher(runtime_watch_rx, runtime_state.clone(), events.clone());
|
||||
drop(initial_config_tx);
|
||||
tokio::task::yield_now().await;
|
||||
|
||||
let (next, next_config_tx, _next_admission_tx) = state(2);
|
||||
runtime_watch_tx.send_replace(Some(next));
|
||||
wait_for_count(&runtime_state, 1).await;
|
||||
next_config_tx.send_replace(Arc::new(ProxyConfig::default()));
|
||||
wait_for_count(&runtime_state, 2).await;
|
||||
|
||||
drop(runtime_watch_tx);
|
||||
tokio::time::timeout(Duration::from_secs(1), watcher)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
events
|
||||
.snapshot(16)
|
||||
.events
|
||||
.iter()
|
||||
.filter(|event| event.event_type == "config.reload.applied")
|
||||
.count(),
|
||||
2
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
+25
-1
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -0,0 +1,435 @@
|
||||
use std::future::Future;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicUsize, 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;
|
||||
#[cfg(test)]
|
||||
use crate::proxy::route_mode::RelayRouteMode;
|
||||
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);
|
||||
const SESSION_ADMISSION_CLOSED: usize = 1 << (usize::BITS - 1);
|
||||
const SESSION_REGISTRATION_COUNT: usize = SESSION_ADMISSION_CLOSED - 1;
|
||||
|
||||
struct SessionAdmission {
|
||||
state: AtomicUsize,
|
||||
}
|
||||
|
||||
struct SessionRegistration<'a> {
|
||||
admission: &'a SessionAdmission,
|
||||
}
|
||||
|
||||
impl SessionAdmission {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
state: AtomicUsize::new(0),
|
||||
}
|
||||
}
|
||||
|
||||
fn try_register(&self) -> Option<SessionRegistration<'_>> {
|
||||
let mut state = self.state.load(Ordering::Acquire);
|
||||
loop {
|
||||
if state & SESSION_ADMISSION_CLOSED != 0
|
||||
|| state & SESSION_REGISTRATION_COUNT == SESSION_REGISTRATION_COUNT
|
||||
{
|
||||
return None;
|
||||
}
|
||||
match self.state.compare_exchange_weak(
|
||||
state,
|
||||
state + 1,
|
||||
Ordering::AcqRel,
|
||||
Ordering::Acquire,
|
||||
) {
|
||||
Ok(_) => return Some(SessionRegistration { admission: self }),
|
||||
Err(observed) => state = observed,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn close(&self) {
|
||||
self.state
|
||||
.fetch_or(SESSION_ADMISSION_CLOSED, Ordering::AcqRel);
|
||||
}
|
||||
|
||||
fn reopen(&self) {
|
||||
self.state
|
||||
.fetch_and(!SESSION_ADMISSION_CLOSED, Ordering::AcqRel);
|
||||
}
|
||||
|
||||
async fn wait_for_registrations(&self) {
|
||||
while self.state.load(Ordering::Acquire) & SESSION_REGISTRATION_COUNT != 0 {
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for SessionRegistration<'_> {
|
||||
fn drop(&mut self) {
|
||||
self.admission.state.fetch_sub(1, Ordering::Release);
|
||||
}
|
||||
}
|
||||
|
||||
/// 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 {
|
||||
/// Creates an open generation-owned task scope.
|
||||
pub(crate) fn new() -> Self {
|
||||
Self {
|
||||
tracker: TaskTracker::new(),
|
||||
cancel: CancellationToken::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Spawns one task that is cancelled when the generation stops.
|
||||
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 => {}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/// Returns the cancellation signal shared by generation-owned controllers.
|
||||
pub(crate) fn cancellation_token(&self) -> CancellationToken {
|
||||
self.cancel.clone()
|
||||
}
|
||||
|
||||
/// Cancels the scope and waits within the bounded background-task budget.
|
||||
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,
|
||||
session_admission: SessionAdmission,
|
||||
}
|
||||
|
||||
impl RuntimeGeneration {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
/// Builds one fully owned runtime generation.
|
||||
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(),
|
||||
session_admission: SessionAdmission::new(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Returns the latest hot-reloaded configuration for this generation.
|
||||
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(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the initial or asynchronously published Middle-End pool.
|
||||
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()
|
||||
}
|
||||
|
||||
/// Registers a session only while admission remains open.
|
||||
pub(crate) fn spawn_session<F>(&self, future: F) -> bool
|
||||
where
|
||||
F: Future<Output = ()> + Send + 'static,
|
||||
{
|
||||
let Some(_registration) = self.session_admission.try_register() else {
|
||||
return false;
|
||||
};
|
||||
let cancel = self.session_cancel.clone();
|
||||
self.sessions.spawn(async move {
|
||||
tokio::select! {
|
||||
_ = cancel.cancelled() => {}
|
||||
_ = future => {}
|
||||
}
|
||||
});
|
||||
true
|
||||
}
|
||||
|
||||
/// Closes admission while preserving already registered sessions.
|
||||
pub(crate) fn stop_accepting_sessions(&self) {
|
||||
self.session_admission.close();
|
||||
}
|
||||
|
||||
/// Reopens admission after a candidate activation rolls back.
|
||||
pub(crate) fn resume_accepting_sessions(&self) {
|
||||
self.session_admission.reopen();
|
||||
}
|
||||
|
||||
/// Waits for registered sessions and cancels them when the deadline expires.
|
||||
pub(crate) async fn drain_sessions(&self, timeout: Duration) -> bool {
|
||||
self.stop_accepting_sessions();
|
||||
self.session_admission.wait_for_registrations().await;
|
||||
self.sessions.close();
|
||||
if tokio::time::timeout(timeout, self.sessions.wait())
|
||||
.await
|
||||
.is_ok()
|
||||
{
|
||||
return true;
|
||||
}
|
||||
self.stop_sessions().await;
|
||||
false
|
||||
}
|
||||
|
||||
/// Cancels all sessions and waits within the bounded session-stop budget.
|
||||
pub(crate) async fn stop_sessions(&self) {
|
||||
self.stop_accepting_sessions();
|
||||
self.session_admission.wait_for_registrations().await;
|
||||
self.session_cancel.cancel();
|
||||
self.sessions.close();
|
||||
let _ = tokio::time::timeout(SESSION_STOP_TIMEOUT, self.sessions.wait()).await;
|
||||
}
|
||||
|
||||
/// Stops all background tasks owned by this generation.
|
||||
pub(crate) async fn stop_background_tasks(&self) {
|
||||
self.background_tasks.stop().await;
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
/// Builds a lightweight runtime generation without network startup tasks.
|
||||
pub(super) fn test_runtime_generation(id: u64, config: ProxyConfig) -> Arc<RuntimeGeneration> {
|
||||
let (config_tx, config_rx) = watch::channel(Arc::new(config.clone()));
|
||||
let (_admission_tx, admission_rx) = watch::channel(true);
|
||||
let stats = Arc::new(Stats::new());
|
||||
let upstream_manager = Arc::new(UpstreamManager::new(
|
||||
config.upstreams,
|
||||
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 _config_tx = config_tx;
|
||||
RuntimeGeneration::new(
|
||||
id,
|
||||
config_rx,
|
||||
admission_rx,
|
||||
stats,
|
||||
upstream_manager,
|
||||
Arc::new(ReplayChecker::new(128, Duration::from_secs(60))),
|
||||
Arc::new(BufferPool::with_config(4096, 16)),
|
||||
Arc::new(SecureRandom::new()),
|
||||
None,
|
||||
Arc::new(RwLock::new(None)),
|
||||
Arc::new(RouteRuntimeController::new(RelayRouteMode::Direct)),
|
||||
None,
|
||||
Arc::new(UserIpTracker::new()),
|
||||
Arc::new(BeobachtenStore::new()),
|
||||
ProxySharedState::new(),
|
||||
Arc::new(Semaphore::new(64)),
|
||||
RuntimeTaskScope::new(),
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tokio::sync::Barrier;
|
||||
|
||||
#[tokio::test]
|
||||
async fn stop_sessions_cancels_tracked_future() {
|
||||
let generation = test_runtime_generation(1, ProxyConfig::default());
|
||||
let started = Arc::new(tokio::sync::Notify::new());
|
||||
let dropped = Arc::new(tokio::sync::Notify::new());
|
||||
let started_task = started.clone();
|
||||
let dropped_task = dropped.clone();
|
||||
assert!(generation.spawn_session(async move {
|
||||
struct DropSignal(Arc<tokio::sync::Notify>);
|
||||
impl Drop for DropSignal {
|
||||
fn drop(&mut self) {
|
||||
self.0.notify_one();
|
||||
}
|
||||
}
|
||||
let _drop_signal = DropSignal(dropped_task);
|
||||
started_task.notify_one();
|
||||
std::future::pending::<()>().await;
|
||||
}));
|
||||
started.notified().await;
|
||||
|
||||
generation.stop_sessions().await;
|
||||
|
||||
tokio::time::timeout(Duration::from_secs(1), dropped.notified())
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(!generation.spawn_session(async {}));
|
||||
}
|
||||
|
||||
#[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();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn session_admission_waits_for_registration_started_before_cutover() {
|
||||
let admission = Arc::new(SessionAdmission::new());
|
||||
let registration = admission.try_register().unwrap();
|
||||
admission.close();
|
||||
assert!(admission.try_register().is_none());
|
||||
|
||||
let wait_admission = admission.clone();
|
||||
let waiter = tokio::spawn(async move {
|
||||
wait_admission.wait_for_registrations().await;
|
||||
});
|
||||
tokio::task::yield_now().await;
|
||||
assert!(!waiter.is_finished());
|
||||
|
||||
drop(registration);
|
||||
tokio::time::timeout(Duration::from_secs(1), waiter)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
|
||||
admission.reopen();
|
||||
assert!(admission.try_register().is_some());
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
||||
async fn cutover_never_leaves_late_session_registrations() {
|
||||
const ATTEMPTS: usize = 10_000;
|
||||
|
||||
let admission = Arc::new(SessionAdmission::new());
|
||||
let tracker = TaskTracker::new();
|
||||
let cancel = CancellationToken::new();
|
||||
let start = Arc::new(Barrier::new(ATTEMPTS + 1));
|
||||
let live = Arc::new(AtomicUsize::new(0));
|
||||
let mut attempts = tokio::task::JoinSet::new();
|
||||
|
||||
for _ in 0..ATTEMPTS {
|
||||
let admission = admission.clone();
|
||||
let tracker = tracker.clone();
|
||||
let cancel = cancel.clone();
|
||||
let start = start.clone();
|
||||
let live = live.clone();
|
||||
attempts.spawn(async move {
|
||||
start.wait().await;
|
||||
let Some(_registration) = admission.try_register() else {
|
||||
return;
|
||||
};
|
||||
tracker.spawn(async move {
|
||||
live.fetch_add(1, Ordering::AcqRel);
|
||||
cancel.cancelled().await;
|
||||
live.fetch_sub(1, Ordering::AcqRel);
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
start.wait().await;
|
||||
admission.close();
|
||||
admission.wait_for_registrations().await;
|
||||
tracker.close();
|
||||
cancel.cancel();
|
||||
while attempts.join_next().await.is_some() {}
|
||||
tokio::time::timeout(Duration::from_secs(1), tracker.wait())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(live.load(Ordering::Acquire), 0);
|
||||
assert!(admission.try_register().is_none());
|
||||
}
|
||||
}
|
||||
+43
-196
@@ -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,
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
+184
-200
@@ -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,61 @@ 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 +85,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 +108,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 +368,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 +466,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 +584,69 @@ pub(crate) async fn initialize_me_pool(
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
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();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn supervisor_restarts_exited_child_and_stops_with_runtime_scope() {
|
||||
let scope = RuntimeTaskScope::new();
|
||||
let starts = Arc::new(AtomicUsize::new(0));
|
||||
let restarted = Arc::new(Notify::new());
|
||||
let starts_task = starts.clone();
|
||||
let restarted_task = restarted.clone();
|
||||
scope.spawn(supervise_me_task("restart_test", move || {
|
||||
let starts = starts_task.clone();
|
||||
let restarted = restarted_task.clone();
|
||||
async move {
|
||||
if starts.fetch_add(1, Ordering::AcqRel) + 1 >= 3 {
|
||||
restarted.notify_one();
|
||||
}
|
||||
}
|
||||
}));
|
||||
|
||||
tokio::time::timeout(Duration::from_secs(1), restarted.notified())
|
||||
.await
|
||||
.unwrap();
|
||||
scope.stop().await;
|
||||
let stopped_at = starts.load(Ordering::Acquire);
|
||||
for _ in 0..100 {
|
||||
tokio::task::yield_now().await;
|
||||
}
|
||||
|
||||
assert!(stopped_at >= 3);
|
||||
assert_eq!(starts.load(Ordering::Acquire), stopped_at);
|
||||
}
|
||||
}
|
||||
|
||||
+113
-84
@@ -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("a_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));
|
||||
let reload_supervisor = 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,30 @@ 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,
|
||||
reload_supervisor,
|
||||
)
|
||||
.await;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -0,0 +1,449 @@
|
||||
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 {
|
||||
/// Validates mode-specific request parameters.
|
||||
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"),
|
||||
}
|
||||
}
|
||||
|
||||
/// Parses optional PATCH query parameters into a reload request.
|
||||
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>,
|
||||
}
|
||||
|
||||
struct ReloadStatusState {
|
||||
next_reload_id: u64,
|
||||
active_reload_id: Option<u64>,
|
||||
statuses: VecDeque<ReloadStatus>,
|
||||
accepting_commands: bool,
|
||||
}
|
||||
|
||||
impl Default for ReloadStatusState {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
next_reload_id: 0,
|
||||
active_reload_id: None,
|
||||
statuses: VecDeque::new(),
|
||||
accepting_commands: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct ReloadStatusStore {
|
||||
state: Mutex<ReloadStatusState>,
|
||||
}
|
||||
|
||||
impl ReloadControl {
|
||||
/// Creates the process-scoped coordinator channel and status store.
|
||||
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 },
|
||||
)
|
||||
}
|
||||
|
||||
/// Atomically reserves and enqueues one reload operation.
|
||||
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.try_send(command).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,
|
||||
})
|
||||
}
|
||||
|
||||
/// Returns a retained reload status by identifier.
|
||||
pub(crate) async fn status(&self, reload_id: u64) -> Option<ReloadStatus> {
|
||||
self.status_store.get(reload_id).await
|
||||
}
|
||||
|
||||
/// Returns the identifier of the currently active reload.
|
||||
pub(crate) async fn in_progress(&self) -> Option<u64> {
|
||||
self.status_store.state.lock().await.active_reload_id
|
||||
}
|
||||
|
||||
/// Rejects new commands while preserving an already accepted operation.
|
||||
pub(crate) async fn begin_shutdown(&self) {
|
||||
self.status_store.state.lock().await.accepting_commands = false;
|
||||
}
|
||||
|
||||
/// Records a non-terminal lifecycle phase.
|
||||
pub(crate) async fn mark_phase(&self, reload_id: u64, phase: ReloadPhase) {
|
||||
self.status_store.mark_phase(reload_id, phase).await;
|
||||
}
|
||||
|
||||
/// Records process-owned fields deferred until the next process restart.
|
||||
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;
|
||||
}
|
||||
|
||||
/// Commits the active generation and completes the matching reload.
|
||||
pub(crate) async fn succeed(&self, reload_id: u64, generation: u64) {
|
||||
self.status_store
|
||||
.finish_success(reload_id, generation, &self.active_generation)
|
||||
.await;
|
||||
}
|
||||
|
||||
/// Marks the matching reload as failed.
|
||||
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;
|
||||
}
|
||||
|
||||
/// Marks the matching reload as rolled back.
|
||||
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;
|
||||
}
|
||||
|
||||
/// Appends a non-fatal warning to the matching reload status.
|
||||
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 {
|
||||
/// Receives the next accepted reload command.
|
||||
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 !state.accepting_commands {
|
||||
return Err(ReloadSubmitError::MaestroUnavailable);
|
||||
}
|
||||
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 finish_success(&self, reload_id: u64, generation: u64, active_generation: &AtomicU64) {
|
||||
let mut state = self.state.lock().await;
|
||||
if state.active_reload_id != Some(reload_id) {
|
||||
return;
|
||||
}
|
||||
let Some(status) = state
|
||||
.statuses
|
||||
.iter_mut()
|
||||
.find(|status| status.reload_id == reload_id)
|
||||
else {
|
||||
return;
|
||||
};
|
||||
status.state = ReloadPhase::Succeeded;
|
||||
status.error = None;
|
||||
status.finished_at_epoch_secs = Some(now_epoch_secs());
|
||||
active_generation.store(generation, Ordering::Release);
|
||||
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)]
|
||||
#[path = "reload_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,280 @@
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use arc_swap::ArcSwap;
|
||||
use tokio::sync::watch;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
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::{PreparedRuntime, 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>>,
|
||||
}
|
||||
|
||||
/// Process-owned handle that quiesces reloads before shutdown snapshots the runtime.
|
||||
pub(crate) struct ReloadSupervisorHandle {
|
||||
control: ReloadControl,
|
||||
shutdown: CancellationToken,
|
||||
join: tokio::task::JoinHandle<()>,
|
||||
}
|
||||
|
||||
impl ReloadSupervisorHandle {
|
||||
/// Stops new submissions and waits for the accepted reload to finish.
|
||||
pub(crate) async fn quiesce(self) {
|
||||
self.control.begin_shutdown().await;
|
||||
self.shutdown.cancel();
|
||||
if let Err(error) = self.join.await {
|
||||
warn!(error = %error, "Reload supervisor failed while quiescing");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[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)]
|
||||
/// Starts the process-scoped reload supervisor and returns its shutdown owner.
|
||||
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>>,
|
||||
) -> ReloadSupervisorHandle {
|
||||
let supervisor = Self {
|
||||
active_runtime,
|
||||
control,
|
||||
commands,
|
||||
config_path,
|
||||
quota_store,
|
||||
detected_ips_tx,
|
||||
runtime_log_filter,
|
||||
runtime_watch_tx,
|
||||
};
|
||||
let control = supervisor.control.clone();
|
||||
let shutdown = CancellationToken::new();
|
||||
let join = tokio::spawn(supervisor.run(shutdown.clone()));
|
||||
ReloadSupervisorHandle {
|
||||
control,
|
||||
shutdown,
|
||||
join,
|
||||
}
|
||||
}
|
||||
|
||||
async fn run(mut self, shutdown: CancellationToken) {
|
||||
loop {
|
||||
tokio::select! {
|
||||
biased;
|
||||
_ = shutdown.cancelled() => {
|
||||
if self.control.in_progress().await.is_some()
|
||||
&& let Some(command) = self.commands.recv().await
|
||||
{
|
||||
self.reload(command).await;
|
||||
}
|
||||
break;
|
||||
}
|
||||
command = self.commands.recv() => {
|
||||
let Some(command) = command else {
|
||||
break;
|
||||
};
|
||||
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,
|
||||
);
|
||||
self.activate_prepared(command, old_runtime, prepared, revision_action, |entries| {
|
||||
crate::network::dns_overrides::install_entries(entries)
|
||||
.map_err(|error| error.to_string())
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
async fn activate_prepared<InstallDns>(
|
||||
&self,
|
||||
command: ReloadCommand,
|
||||
old_runtime: Arc<RuntimeGeneration>,
|
||||
prepared: PreparedRuntime,
|
||||
revision_action: RevisionGateAction,
|
||||
install_dns: InstallDns,
|
||||
) where
|
||||
InstallDns: FnOnce(&[String]) -> Result<(), String>,
|
||||
{
|
||||
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) = install_dns(&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)]
|
||||
#[path = "reload_supervisor_tests.rs"]
|
||||
mod tests;
|
||||
@@ -0,0 +1,320 @@
|
||||
use super::*;
|
||||
use crate::config::ProxyConfig;
|
||||
use crate::maestro::generation::test_runtime_generation;
|
||||
use crate::maestro::reload::{ReloadRequest, ReloadSubmitError};
|
||||
use crate::stats::QuotaStore;
|
||||
use tokio::sync::Notify;
|
||||
use tracing_subscriber::{EnvFilter, Registry};
|
||||
|
||||
struct ReloadFixture {
|
||||
supervisor: Arc<ReloadSupervisor>,
|
||||
control: ReloadControl,
|
||||
command: ReloadCommand,
|
||||
old_runtime: Arc<RuntimeGeneration>,
|
||||
new_runtime: Arc<RuntimeGeneration>,
|
||||
runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
|
||||
}
|
||||
|
||||
fn runtime_log_filter() -> RuntimeLogFilter {
|
||||
let (_layer, handle) =
|
||||
tracing_subscriber::reload::Layer::<EnvFilter, Registry>::new(EnvFilter::new("info"));
|
||||
RuntimeLogFilter::new(handle)
|
||||
}
|
||||
|
||||
async fn fixture(request: ReloadRequest) -> ReloadFixture {
|
||||
let old_runtime = test_runtime_generation(1, ProxyConfig::default());
|
||||
let new_config = Arc::new(ProxyConfig::default());
|
||||
let new_runtime = test_runtime_generation(2, new_config.as_ref().clone());
|
||||
let active_runtime = Arc::new(ArcSwap::from(old_runtime.clone()));
|
||||
let (control, commands) = ReloadControl::channel(old_runtime.id);
|
||||
let accepted = control
|
||||
.submit(new_config.clone(), "revision".to_string(), request.clone())
|
||||
.await
|
||||
.unwrap();
|
||||
let (detected_ips_tx, _detected_ips_rx) = watch::channel((None, None));
|
||||
let (runtime_watch_tx, runtime_watch_rx) = watch::channel(Some(old_runtime.watch_state()));
|
||||
let supervisor = Arc::new(ReloadSupervisor {
|
||||
active_runtime,
|
||||
control: control.clone(),
|
||||
commands,
|
||||
config_path: PathBuf::new(),
|
||||
quota_store: Arc::new(QuotaStore::default()),
|
||||
detected_ips_tx,
|
||||
runtime_log_filter: runtime_log_filter(),
|
||||
runtime_watch_tx,
|
||||
});
|
||||
let command = ReloadCommand {
|
||||
reload_id: accepted.reload_id,
|
||||
target_generation: accepted.target_generation,
|
||||
config: new_config,
|
||||
config_revision: accepted.config_revision,
|
||||
request,
|
||||
};
|
||||
ReloadFixture {
|
||||
supervisor,
|
||||
control,
|
||||
command,
|
||||
old_runtime,
|
||||
new_runtime,
|
||||
runtime_watch_rx,
|
||||
}
|
||||
}
|
||||
|
||||
struct DropSignal(Arc<Notify>);
|
||||
|
||||
impl Drop for DropSignal {
|
||||
fn drop(&mut self) {
|
||||
self.0.notify_one();
|
||||
}
|
||||
}
|
||||
|
||||
#[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(_)
|
||||
));
|
||||
}
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn revision_rollback_keeps_old_generation_and_cleans_candidate() {
|
||||
let fixture = fixture(ReloadRequest {
|
||||
failure_policy: ReloadFailurePolicy::Rollback,
|
||||
..ReloadRequest::default()
|
||||
})
|
||||
.await;
|
||||
let candidate_dropped = Arc::new(Notify::new());
|
||||
let candidate_drop = candidate_dropped.clone();
|
||||
assert!(fixture.new_runtime.spawn_session(async move {
|
||||
let _drop_signal = DropSignal(candidate_drop);
|
||||
std::future::pending::<()>().await;
|
||||
}));
|
||||
tokio::task::yield_now().await;
|
||||
|
||||
fixture
|
||||
.supervisor
|
||||
.activate_prepared(
|
||||
fixture.command,
|
||||
fixture.old_runtime.clone(),
|
||||
PreparedRuntime {
|
||||
generation: fixture.new_runtime,
|
||||
detected_ips: (None, None),
|
||||
},
|
||||
RevisionGateAction::Rollback("revision changed".to_string()),
|
||||
|_| -> Result<(), String> { panic!("DNS activation must not run on rollback") },
|
||||
)
|
||||
.await;
|
||||
|
||||
tokio::time::timeout(Duration::from_secs(1), candidate_dropped.notified())
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(fixture.supervisor.active_runtime.load().id, 1);
|
||||
assert_eq!(
|
||||
fixture
|
||||
.runtime_watch_rx
|
||||
.borrow()
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.generation_id,
|
||||
1
|
||||
);
|
||||
assert!(fixture.old_runtime.spawn_session(async {}));
|
||||
let status = fixture.control.status(1).await.unwrap();
|
||||
assert_eq!(status.state, ReloadPhase::RolledBack);
|
||||
fixture.old_runtime.stop_sessions().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dns_failure_policy_controls_rollback_or_keep_new() {
|
||||
for policy in [ReloadFailurePolicy::Rollback, ReloadFailurePolicy::KeepNew] {
|
||||
let fixture = fixture(ReloadRequest {
|
||||
failure_policy: policy,
|
||||
..ReloadRequest::default()
|
||||
})
|
||||
.await;
|
||||
fixture
|
||||
.supervisor
|
||||
.activate_prepared(
|
||||
fixture.command,
|
||||
fixture.old_runtime.clone(),
|
||||
PreparedRuntime {
|
||||
generation: fixture.new_runtime.clone(),
|
||||
detected_ips: (None, None),
|
||||
},
|
||||
RevisionGateAction::Proceed,
|
||||
|_| Err("invalid DNS entry".to_string()),
|
||||
)
|
||||
.await;
|
||||
|
||||
let status = fixture.control.status(1).await.unwrap();
|
||||
match policy {
|
||||
ReloadFailurePolicy::Rollback => {
|
||||
assert_eq!(fixture.supervisor.active_runtime.load().id, 1);
|
||||
assert_eq!(status.state, ReloadPhase::RolledBack);
|
||||
assert!(fixture.old_runtime.spawn_session(async {}));
|
||||
fixture.old_runtime.stop_sessions().await;
|
||||
}
|
||||
ReloadFailurePolicy::KeepNew => {
|
||||
assert_eq!(fixture.supervisor.active_runtime.load().id, 2);
|
||||
assert_eq!(status.state, ReloadPhase::Succeeded);
|
||||
assert_eq!(status.warnings.len(), 1);
|
||||
assert!(!fixture.old_runtime.spawn_session(async {}));
|
||||
fixture.new_runtime.stop_sessions().await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn drain_publishes_new_generation_before_old_sessions_finish() {
|
||||
let mut fixture = fixture(ReloadRequest {
|
||||
mode: ReloadMode::Drain,
|
||||
timeout_secs: Some(30),
|
||||
..ReloadRequest::default()
|
||||
})
|
||||
.await;
|
||||
let old_started = Arc::new(Notify::new());
|
||||
let old_release = Arc::new(Notify::new());
|
||||
let started = old_started.clone();
|
||||
let release = old_release.clone();
|
||||
assert!(fixture.old_runtime.spawn_session(async move {
|
||||
started.notify_one();
|
||||
release.notified().await;
|
||||
}));
|
||||
old_started.notified().await;
|
||||
|
||||
let supervisor = fixture.supervisor.clone();
|
||||
let old_runtime = fixture.old_runtime.clone();
|
||||
let new_runtime = fixture.new_runtime.clone();
|
||||
let activation = tokio::spawn(async move {
|
||||
supervisor
|
||||
.activate_prepared(
|
||||
fixture.command,
|
||||
old_runtime,
|
||||
PreparedRuntime {
|
||||
generation: new_runtime,
|
||||
detected_ips: (None, None),
|
||||
},
|
||||
RevisionGateAction::Proceed,
|
||||
|_| Ok(()),
|
||||
)
|
||||
.await;
|
||||
});
|
||||
|
||||
fixture.runtime_watch_rx.changed().await.unwrap();
|
||||
assert_eq!(
|
||||
fixture
|
||||
.runtime_watch_rx
|
||||
.borrow()
|
||||
.as_ref()
|
||||
.unwrap()
|
||||
.generation_id,
|
||||
2
|
||||
);
|
||||
assert!(!activation.is_finished());
|
||||
assert!(!fixture.old_runtime.spawn_session(async {}));
|
||||
|
||||
old_release.notify_one();
|
||||
activation.await.unwrap();
|
||||
assert_eq!(
|
||||
fixture.control.status(1).await.unwrap().state,
|
||||
ReloadPhase::Succeeded
|
||||
);
|
||||
fixture.new_runtime.stop_sessions().await;
|
||||
}
|
||||
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn drain_timeout_cancels_old_sessions_and_records_one_warning() {
|
||||
let mut fixture = fixture(ReloadRequest {
|
||||
mode: ReloadMode::Drain,
|
||||
timeout_secs: Some(1),
|
||||
..ReloadRequest::default()
|
||||
})
|
||||
.await;
|
||||
let dropped = Arc::new(Notify::new());
|
||||
let drop_signal = dropped.clone();
|
||||
assert!(fixture.old_runtime.spawn_session(async move {
|
||||
let _drop_signal = DropSignal(drop_signal);
|
||||
std::future::pending::<()>().await;
|
||||
}));
|
||||
tokio::task::yield_now().await;
|
||||
|
||||
let supervisor = fixture.supervisor.clone();
|
||||
let old_runtime = fixture.old_runtime.clone();
|
||||
let new_runtime = fixture.new_runtime.clone();
|
||||
let activation = tokio::spawn(async move {
|
||||
supervisor
|
||||
.activate_prepared(
|
||||
fixture.command,
|
||||
old_runtime,
|
||||
PreparedRuntime {
|
||||
generation: new_runtime,
|
||||
detected_ips: (None, None),
|
||||
},
|
||||
RevisionGateAction::Proceed,
|
||||
|_| Ok(()),
|
||||
)
|
||||
.await;
|
||||
});
|
||||
fixture.runtime_watch_rx.changed().await.unwrap();
|
||||
tokio::task::yield_now().await;
|
||||
tokio::time::advance(Duration::from_secs(1)).await;
|
||||
activation.await.unwrap();
|
||||
|
||||
dropped.notified().await;
|
||||
let status = fixture.control.status(1).await.unwrap();
|
||||
assert_eq!(status.state, ReloadPhase::Succeeded);
|
||||
assert_eq!(status.warnings.len(), 1);
|
||||
assert!(status.warnings[0].contains("exceeded drain timeout"));
|
||||
fixture.new_runtime.stop_sessions().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn quiesce_joins_idle_supervisor_and_rejects_later_submissions() {
|
||||
let runtime = test_runtime_generation(1, ProxyConfig::default());
|
||||
let active_runtime = Arc::new(ArcSwap::from(runtime.clone()));
|
||||
let (control, commands) = ReloadControl::channel(runtime.id);
|
||||
let (detected_ips_tx, _detected_ips_rx) = watch::channel((None, None));
|
||||
let (runtime_watch_tx, _runtime_watch_rx) = watch::channel(Some(runtime.watch_state()));
|
||||
let handle = ReloadSupervisor::spawn(
|
||||
active_runtime,
|
||||
control.clone(),
|
||||
commands,
|
||||
PathBuf::new(),
|
||||
Arc::new(QuotaStore::default()),
|
||||
detected_ips_tx,
|
||||
runtime_log_filter(),
|
||||
runtime_watch_tx,
|
||||
);
|
||||
|
||||
tokio::time::timeout(Duration::from_secs(1), handle.quiesce())
|
||||
.await
|
||||
.unwrap();
|
||||
let result = control
|
||||
.submit(
|
||||
Arc::new(ProxyConfig::default()),
|
||||
"revision".to_string(),
|
||||
ReloadRequest::default(),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(result, Err(ReloadSubmitError::MaestroUnavailable));
|
||||
runtime.stop_sessions().await;
|
||||
}
|
||||
@@ -0,0 +1,256 @@
|
||||
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);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn terminal_outcomes_release_slot_and_only_success_advances_generation() {
|
||||
let (control, mut receiver) = ReloadControl::channel(7);
|
||||
|
||||
let failed = control
|
||||
.submit(
|
||||
Arc::new(ProxyConfig::default()),
|
||||
"rev-failed".to_string(),
|
||||
ReloadRequest::default(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let _command = receiver.recv().await.unwrap();
|
||||
control
|
||||
.mark_phase(failed.reload_id, ReloadPhase::Preparing)
|
||||
.await;
|
||||
control.fail(failed.reload_id, "prepare failed").await;
|
||||
let failed_status = control.status(failed.reload_id).await.unwrap();
|
||||
assert_eq!(failed_status.state, ReloadPhase::Failed);
|
||||
assert_eq!(failed_status.error.as_deref(), Some("prepare failed"));
|
||||
assert!(failed_status.started_at_epoch_secs.is_some());
|
||||
assert!(failed_status.finished_at_epoch_secs.is_some());
|
||||
|
||||
let rolled_back = control
|
||||
.submit(
|
||||
Arc::new(ProxyConfig::default()),
|
||||
"rev-rollback".to_string(),
|
||||
ReloadRequest::default(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let _command = receiver.recv().await.unwrap();
|
||||
assert_eq!(rolled_back.target_generation, 8);
|
||||
control
|
||||
.rolled_back(rolled_back.reload_id, "revision changed")
|
||||
.await;
|
||||
|
||||
let succeeded = control
|
||||
.submit(
|
||||
Arc::new(ProxyConfig::default()),
|
||||
"rev-success".to_string(),
|
||||
ReloadRequest::default(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let _command = receiver.recv().await.unwrap();
|
||||
assert_eq!(succeeded.target_generation, 8);
|
||||
control
|
||||
.succeed(succeeded.reload_id, succeeded.target_generation)
|
||||
.await;
|
||||
|
||||
let next = control
|
||||
.submit(
|
||||
Arc::new(ProxyConfig::default()),
|
||||
"rev-next".to_string(),
|
||||
ReloadRequest::default(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(next.target_generation, 9);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn stale_success_cannot_advance_generation_or_release_active_reload() {
|
||||
let (control, mut receiver) = ReloadControl::channel(3);
|
||||
let active = control
|
||||
.submit(
|
||||
Arc::new(ProxyConfig::default()),
|
||||
"rev-active".to_string(),
|
||||
ReloadRequest::default(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let _command = receiver.recv().await.unwrap();
|
||||
|
||||
control.succeed(active.reload_id + 100, 99).await;
|
||||
|
||||
assert_eq!(control.in_progress().await, Some(active.reload_id));
|
||||
control.fail(active.reload_id, "expected failure").await;
|
||||
let next = control
|
||||
.submit(
|
||||
Arc::new(ProxyConfig::default()),
|
||||
"rev-next".to_string(),
|
||||
ReloadRequest::default(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(next.target_generation, 4);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn status_history_retains_only_the_latest_entries() {
|
||||
let (control, mut receiver) = ReloadControl::channel(1);
|
||||
let mut reload_ids = Vec::new();
|
||||
for index in 0..=RELOAD_HISTORY_CAPACITY {
|
||||
let accepted = control
|
||||
.submit(
|
||||
Arc::new(ProxyConfig::default()),
|
||||
format!("rev-{index}"),
|
||||
ReloadRequest::default(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let _command = receiver.recv().await.unwrap();
|
||||
reload_ids.push(accepted.reload_id);
|
||||
control.fail(accepted.reload_id, "expected failure").await;
|
||||
}
|
||||
|
||||
assert!(control.status(reload_ids[0]).await.is_none());
|
||||
assert!(control.status(reload_ids[1]).await.is_some());
|
||||
assert!(control.status(*reload_ids.last().unwrap()).await.is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn closed_command_channel_marks_reload_failed_and_releases_slot() {
|
||||
let (control, receiver) = ReloadControl::channel(1);
|
||||
drop(receiver);
|
||||
|
||||
let result = control
|
||||
.submit(
|
||||
Arc::new(ProxyConfig::default()),
|
||||
"rev-closed".to_string(),
|
||||
ReloadRequest::default(),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(result, Err(ReloadSubmitError::MaestroUnavailable));
|
||||
assert_eq!(control.in_progress().await, None);
|
||||
let status = control.status(1).await.unwrap();
|
||||
assert_eq!(status.state, ReloadPhase::Failed);
|
||||
assert_eq!(
|
||||
status.error.as_deref(),
|
||||
Some("maestro command channel is closed")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn shutdown_gate_rejects_new_commands_without_disturbing_active_status() {
|
||||
let (control, mut receiver) = ReloadControl::channel(4);
|
||||
let active = control
|
||||
.submit(
|
||||
Arc::new(ProxyConfig::default()),
|
||||
"rev-active".to_string(),
|
||||
ReloadRequest::default(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let _command = receiver.recv().await.unwrap();
|
||||
|
||||
control.begin_shutdown().await;
|
||||
let rejected = control
|
||||
.submit(
|
||||
Arc::new(ProxyConfig::default()),
|
||||
"rev-rejected".to_string(),
|
||||
ReloadRequest::default(),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(rejected, Err(ReloadSubmitError::MaestroUnavailable));
|
||||
assert_eq!(control.in_progress().await, Some(active.reload_id));
|
||||
control.fail(active.reload_id, "shutdown test").await;
|
||||
}
|
||||
@@ -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));
|
||||
}
|
||||
}
|
||||
@@ -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(
|
||||
|
||||
+35
-17
@@ -12,17 +12,18 @@ 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 super::reload_supervisor::ReloadSupervisorHandle;
|
||||
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 +49,19 @@ 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: synlimit_control::SynlimitController,
|
||||
reload_supervisor: ReloadSupervisorHandle,
|
||||
) {
|
||||
let signal = wait_for_shutdown_signal().await;
|
||||
perform_shutdown(
|
||||
signal,
|
||||
process_started_at,
|
||||
me_pool,
|
||||
&stats,
|
||||
active_runtime,
|
||||
quota_state_path,
|
||||
synlimit_controller,
|
||||
reload_supervisor,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
@@ -87,13 +90,18 @@ 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: synlimit_control::SynlimitController,
|
||||
reload_supervisor: ReloadSupervisorHandle,
|
||||
) {
|
||||
let shutdown_started_at = Instant::now();
|
||||
info!(signal = %signal, "Received shutdown signal");
|
||||
|
||||
reload_supervisor.quiesce().await;
|
||||
let runtime = active_runtime.load_full();
|
||||
let stats = runtime.stats.as_ref();
|
||||
|
||||
// Dump stats if SIGQUIT
|
||||
if signal == ShutdownSignal::Quit {
|
||||
dump_stats(stats, process_started_at);
|
||||
@@ -103,12 +111,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 +130,11 @@ async fn perform_shutdown(
|
||||
}
|
||||
}
|
||||
|
||||
synlimit_controller.shutdown().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("a_state_path, stats).await {
|
||||
Ok(()) => {
|
||||
info!(
|
||||
@@ -191,7 +202,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 +218,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 +228,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
|
||||
}
|
||||
|
||||
|
||||
+306
-146
@@ -5,11 +5,86 @@ 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 +93,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 +118,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 +157,110 @@ 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::*;
|
||||
use crate::startup::StartupComponentStatus;
|
||||
use crate::stats::Stats;
|
||||
|
||||
fn test_config(cache_dir: &std::path::Path) -> ProxyConfig {
|
||||
let mut config = ProxyConfig::default();
|
||||
config.censorship.tls_emulation = true;
|
||||
config.censorship.tls_domain = "front.example".to_string();
|
||||
config.censorship.mask_host = Some("127.0.0.1".to_string());
|
||||
config.censorship.mask_port = 1;
|
||||
config.censorship.tls_front_dir = cache_dir.display().to_string();
|
||||
config.censorship.tls_fetch.profiles.truncate(1);
|
||||
config.censorship.tls_fetch.attempt_timeout_ms = 10;
|
||||
config.censorship.tls_fetch.total_budget_ms = 20;
|
||||
config
|
||||
}
|
||||
|
||||
fn upstream_manager(config: &ProxyConfig) -> Arc<UpstreamManager> {
|
||||
Arc::new(UpstreamManager::new(
|
||||
Vec::new(),
|
||||
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,
|
||||
Arc::new(Stats::new()),
|
||||
))
|
||||
}
|
||||
|
||||
async fn tls_component_status(tracker: &StartupTracker) -> StartupComponentStatus {
|
||||
tracker
|
||||
.snapshot()
|
||||
.await
|
||||
.components
|
||||
.into_iter()
|
||||
.find(|component| component.id == COMPONENT_TLS_FRONT_BOOTSTRAP)
|
||||
.unwrap()
|
||||
.status
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tls_fetch_host_uses_each_domain_when_mask_host_is_primary_default() {
|
||||
@@ -218,4 +277,105 @@ 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())
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn require_ready_rejects_default_cache_after_bounded_fetch_failure() {
|
||||
let cache_dir = tempfile::tempdir().unwrap();
|
||||
let config = test_config(cache_dir.path());
|
||||
let domains = vec![config.censorship.tls_domain.clone()];
|
||||
let tracker = Arc::new(StartupTracker::new(1));
|
||||
let scope = RuntimeTaskScope::new();
|
||||
|
||||
let result = bootstrap_tls_front(
|
||||
&config,
|
||||
&domains,
|
||||
upstream_manager(&config),
|
||||
&tracker,
|
||||
scope.clone(),
|
||||
TlsBootstrapPolicy::RequireReady,
|
||||
)
|
||||
.await;
|
||||
|
||||
assert!(result.is_err());
|
||||
assert_eq!(
|
||||
tls_component_status(&tracker).await,
|
||||
StartupComponentStatus::Failed
|
||||
);
|
||||
scope.stop().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn require_ready_accepts_non_default_disk_cache_when_refresh_fails() {
|
||||
let cache_dir = tempfile::tempdir().unwrap();
|
||||
let config = test_config(cache_dir.path());
|
||||
let domains = vec![config.censorship.tls_domain.clone()];
|
||||
let seed = TlsFrontCache::new(&domains, config.censorship.fake_cert_len, cache_dir.path());
|
||||
let mut cached = seed.default_entry().as_ref().clone();
|
||||
cached.domain = domains[0].clone();
|
||||
tokio::fs::write(
|
||||
cache_dir.path().join("front.example.json"),
|
||||
serde_json::to_vec(&cached).unwrap(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let tracker = Arc::new(StartupTracker::new(1));
|
||||
let scope = RuntimeTaskScope::new();
|
||||
|
||||
let cache = bootstrap_tls_front(
|
||||
&config,
|
||||
&domains,
|
||||
upstream_manager(&config),
|
||||
&tracker,
|
||||
scope.clone(),
|
||||
TlsBootstrapPolicy::RequireReady,
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
|
||||
assert!(cache.default_profile_domains(&domains).await.is_empty());
|
||||
assert_eq!(
|
||||
tls_component_status(&tracker).await,
|
||||
StartupComponentStatus::Ready
|
||||
);
|
||||
scope.stop().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn best_effort_returns_ready_and_refresh_tasks_are_scope_owned() {
|
||||
let cache_dir = tempfile::tempdir().unwrap();
|
||||
let config = test_config(cache_dir.path());
|
||||
let domains = vec![config.censorship.tls_domain.clone()];
|
||||
let tracker = Arc::new(StartupTracker::new(1));
|
||||
let scope = RuntimeTaskScope::new();
|
||||
|
||||
let cache = bootstrap_tls_front(
|
||||
&config,
|
||||
&domains,
|
||||
upstream_manager(&config),
|
||||
&tracker,
|
||||
scope.clone(),
|
||||
TlsBootstrapPolicy::BestEffort,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(cache.is_some());
|
||||
assert_eq!(
|
||||
tls_component_status(&tracker).await,
|
||||
StartupComponentStatus::Ready
|
||||
);
|
||||
tokio::time::timeout(Duration::from_secs(1), scope.stop())
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
+26
-84
@@ -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,
|
||||
|
||||
@@ -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!(
|
||||
|
||||
@@ -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> {
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
@@ -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) {
|
||||
|
||||
+26
-23
@@ -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());
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
+5
-35
@@ -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 {
|
||||
|
||||
+268
-20
@@ -1,9 +1,11 @@
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use tokio::sync::watch;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::config::{ProxyConfig, SynLimitMode};
|
||||
use crate::maestro::generation::RuntimeWatchState;
|
||||
|
||||
mod command;
|
||||
mod iptables;
|
||||
@@ -15,28 +17,141 @@ 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>>) {
|
||||
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;
|
||||
}
|
||||
|
||||
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");
|
||||
}
|
||||
});
|
||||
/// Process-owned lifecycle handle for the SYN limiter reconciler.
|
||||
pub(crate) struct SynlimitController {
|
||||
shutdown: CancellationToken,
|
||||
join: tokio::task::JoinHandle<()>,
|
||||
}
|
||||
|
||||
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;
|
||||
impl SynlimitController {
|
||||
/// Stops config observation after any in-flight reconcile completes.
|
||||
pub(crate) async fn shutdown(self) {
|
||||
self.shutdown.cancel();
|
||||
let _ = self.join.await;
|
||||
}
|
||||
}
|
||||
|
||||
/// Spawns the process-scoped SYN limiter reconciler for active generations.
|
||||
pub(crate) fn spawn_synlimit_controller(
|
||||
runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
|
||||
) -> SynlimitController {
|
||||
let shutdown = CancellationToken::new();
|
||||
let join = if !cfg!(target_os = "linux") {
|
||||
tokio::spawn(watch_active_runtime_configs(
|
||||
runtime_watch_rx,
|
||||
shutdown.clone(),
|
||||
|_generation_id, cfg| async move {
|
||||
if has_synlimit_config(&cfg) {
|
||||
warn!(
|
||||
"SYN limiter is configured but unsupported on this OS; skipping netfilter rules"
|
||||
);
|
||||
}
|
||||
},
|
||||
))
|
||||
} else {
|
||||
tokio::spawn(watch_active_runtime_configs(
|
||||
runtime_watch_rx,
|
||||
shutdown.clone(),
|
||||
|_generation_id, cfg| async move {
|
||||
reconcile_synlimit_rules(&cfg).await;
|
||||
},
|
||||
))
|
||||
};
|
||||
SynlimitController { shutdown, join }
|
||||
}
|
||||
|
||||
async fn watch_active_runtime_configs<F, Fut>(
|
||||
mut runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
|
||||
shutdown: CancellationToken,
|
||||
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;
|
||||
}
|
||||
tokio::select! {
|
||||
biased;
|
||||
_ = shutdown.cancelled() => return,
|
||||
changed = runtime_watch_rx.changed() => {
|
||||
if changed.is_err() {
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
if shutdown.is_cancelled() {
|
||||
return;
|
||||
}
|
||||
let initial_config = current.config_rx.borrow().clone();
|
||||
on_config(current.generation_id, initial_config).await;
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
biased;
|
||||
_ = shutdown.cancelled() => break,
|
||||
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,
|
||||
&shutdown,
|
||||
).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,
|
||||
shutdown: &CancellationToken,
|
||||
) -> Option<RuntimeWatchState> {
|
||||
loop {
|
||||
if let Some(state) = runtime_watch_rx.borrow().clone()
|
||||
&& state.generation_id != previous_generation_id
|
||||
{
|
||||
return Some(state);
|
||||
}
|
||||
tokio::select! {
|
||||
biased;
|
||||
_ = shutdown.cancelled() => return None,
|
||||
changed = runtime_watch_rx.changed() => {
|
||||
if changed.is_err() {
|
||||
return None;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -168,3 +283,136 @@ fn has_synlimit_config(cfg: &ProxyConfig) -> bool {
|
||||
.iter()
|
||||
.any(|listener| !matches!(listener.synlimit, SynLimitMode::Off))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::time::Duration;
|
||||
use tokio::sync::{Notify, 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,
|
||||
CancellationToken::new(),
|
||||
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();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn shutdown_waits_for_inflight_reconcile_and_stops_future_updates() {
|
||||
let (initial, config_tx, _admission_tx) = runtime_state(1, 10);
|
||||
let (_runtime_tx, runtime_rx) = watch::channel(Some(initial));
|
||||
let shutdown = CancellationToken::new();
|
||||
let started = Arc::new(Notify::new());
|
||||
let release = Arc::new(Notify::new());
|
||||
let calls = Arc::new(AtomicUsize::new(0));
|
||||
let started_callback = started.clone();
|
||||
let release_callback = release.clone();
|
||||
let calls_callback = calls.clone();
|
||||
let watcher_shutdown = shutdown.clone();
|
||||
let watcher = tokio::spawn(watch_active_runtime_configs(
|
||||
runtime_rx,
|
||||
watcher_shutdown,
|
||||
move |_generation_id, _cfg| {
|
||||
let started = started_callback.clone();
|
||||
let release = release_callback.clone();
|
||||
let calls = calls_callback.clone();
|
||||
async move {
|
||||
calls.fetch_add(1, Ordering::AcqRel);
|
||||
started.notify_one();
|
||||
release.notified().await;
|
||||
}
|
||||
},
|
||||
));
|
||||
started.notified().await;
|
||||
|
||||
shutdown.cancel();
|
||||
tokio::task::yield_now().await;
|
||||
assert!(!watcher.is_finished());
|
||||
release.notify_one();
|
||||
tokio::time::timeout(Duration::from_secs(1), watcher)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
|
||||
drop(_runtime_tx);
|
||||
let mut updated = ProxyConfig::default();
|
||||
updated.server.max_connections = 20;
|
||||
assert!(config_tx.send(Arc::new(updated)).is_err());
|
||||
assert_eq!(calls.load(Ordering::Acquire), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn shutdown_before_start_skips_initial_reconcile() {
|
||||
let (initial, _config_tx, _admission_tx) = runtime_state(1, 10);
|
||||
let (_runtime_tx, runtime_rx) = watch::channel(Some(initial));
|
||||
let shutdown = CancellationToken::new();
|
||||
shutdown.cancel();
|
||||
let calls = Arc::new(AtomicUsize::new(0));
|
||||
let calls_callback = calls.clone();
|
||||
|
||||
watch_active_runtime_configs(runtime_rx, shutdown, move |_generation_id, _cfg| {
|
||||
let calls = calls_callback.clone();
|
||||
async move {
|
||||
calls.fetch_add(1, Ordering::AcqRel);
|
||||
}
|
||||
})
|
||||
.await;
|
||||
|
||||
assert_eq!(calls.load(Ordering::Acquire), 0);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -190,6 +190,18 @@ 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 +558,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");
|
||||
|
||||
+18
-19
@@ -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 {
|
||||
|
||||
@@ -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}")))?
|
||||
|
||||
@@ -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");
|
||||
@@ -2019,6 +2052,45 @@ mod tests {
|
||||
const TEST_SHADOWSOCKS_URL: &str =
|
||||
"ss://2022-blake3-aes-256-gcm:MDEyMzQ1Njc4OTAxMjM0NTY3ODkwMTIzNDU2Nzg5MDE=@127.0.0.1:8388";
|
||||
|
||||
fn manager_with_dns(entries: &[String]) -> UpstreamManager {
|
||||
UpstreamManager::new(Vec::new(), 1, 1, 1, 1, 1, false, Arc::new(Stats::new()))
|
||||
.with_dns_overrides(entries)
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn generation_local_dns_overrides_are_isolated_and_case_insensitive() {
|
||||
let active = manager_with_dns(&["Front.Example:443:192.0.2.10".to_string()]);
|
||||
let candidate = manager_with_dns(&["front.example:443:[2001:db8::10]".to_string()]);
|
||||
|
||||
assert_eq!(
|
||||
active.resolve_hostname("front.example", 443).await.unwrap(),
|
||||
"192.0.2.10:443".parse::<SocketAddr>().unwrap()
|
||||
);
|
||||
assert_eq!(
|
||||
candidate
|
||||
.resolve_hostname("FRONT.EXAMPLE", 443)
|
||||
.await
|
||||
.unwrap(),
|
||||
"[2001:db8::10]:443".parse::<SocketAddr>().unwrap()
|
||||
);
|
||||
|
||||
candidate
|
||||
.update_dns_overrides(&["front.example:443:192.0.2.20".to_string()])
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
active.resolve_hostname("FRONT.EXAMPLE", 443).await.unwrap(),
|
||||
"192.0.2.10:443".parse::<SocketAddr>().unwrap()
|
||||
);
|
||||
assert_eq!(
|
||||
candidate
|
||||
.resolve_hostname("front.example", 443)
|
||||
.await
|
||||
.unwrap(),
|
||||
"192.0.2.20:443".parse::<SocketAddr>().unwrap()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn required_healthy_group_count_applies_three_group_threshold() {
|
||||
assert_eq!(UpstreamManager::required_healthy_group_count(0), 0);
|
||||
|
||||
Reference in New Issue
Block a user