Compare commits

...

32 Commits

Author SHA1 Message Date
Alexey 73afeccae1 Rustfmt 2026-07-13 12:20:24 +03:00
Alexey feb51cbf57 Fix RustCrypto zeroize feature wiring 2026-07-12 09:27:15 +03:00
Alexey 8c65cd868c Enable expanded AES key schedule zeroization
Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
2026-07-11 21:43:01 +03:00
Alexey ea296bbdc8 Replace per-session pool trimming with pressure hysteresis
Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
2026-07-11 21:36:01 +03:00
Alexey fb042f826e Optimize crypto and Fake-TLS buffer residency
Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
2026-07-11 21:22:17 +03:00
Alexey 96425f15c8 Bound Direct relay buffers with an adaptive global memory envelope
Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
2026-07-11 20:56:54 +03:00
Alexey d4c4980e5a Bound ME writer queues by resident payload bytes
Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
2026-07-11 18:43:42 +03:00
Alexey 893ce0cf36 Hold C2ME byte permits through ME writer completion
Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
2026-07-10 16:35:39 +03:00
Alexey 2ac93c6d49 Update Cargo.lock 2026-07-10 11:47:09 +03:00
Alexey a51e58009b Update Cargo.toml 2026-07-10 11:46:33 +03:00
Alexey d523406c0a Merge pull request #869 from telemt/flow
Regression coverage + Synlimit netfilter rules per-target + CidrRateLimitKey with IpNetwork + Telemt-Referenz der Konfigurationsparameter
2026-07-10 11:39:22 +03:00
Alexey 5b3ad0096b Create CONFIG_PARAMS.de.md 2026-07-09 22:42:07 +03:00
Alexey b587fdbf94 Update CONFIG_PARAMS.ru.md 2026-07-08 20:52:26 +03:00
Alexey 77a45e509a Update CONFIG_PARAMS.en.md 2026-07-07 22:46:03 +03:00
Alexey b8be805aed Rustfmt 2026-07-06 20:12:36 +03:00
Alexey a1ebd44cee Update load_basic_tests.rs 2026-07-05 19:58:37 +03:00
Alexey 25d02a8e0e Update traffic_limiter.rs 2026-07-04 15:45:51 +03:00
Alexey 3375017460 Add the new key to the hot-reload snapshot type 2026-07-04 13:11:08 +03:00
Alexey 25e0abae8a Validate duplicate normalized auto-templates 2026-07-04 12:35:22 +03:00
Alexey 50538d234e CidrRateLimitKey with IpNetwork parsing and serialization added 2026-07-04 10:54:55 +03:00
Alexey e3a7be6786 Update CONTRIBUTING.md 2026-07-03 23:39:46 +03:00
Alexey 3fc2877205 Reset ports configuration in merge docker-compose file: merge pull request #866 from dvrass/patch-1
Reset ports configuration in merge docker-compose file
2026-07-03 23:34:04 +03:00
Alexey cd1dc2f4c9 Update docker-compose.host-netfilter.yml 2026-07-03 23:32:30 +03:00
Alexey 451227da60 Namespace synlimit netfilter rules per target set
Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
2026-07-01 01:47:14 +03:00
dvrass f55d8479e3 Reset ports configuration in merge docker-compose file
Reset ports configuration in docker-compose file for network_mode:host
2026-06-30 13:23:56 +03:00
Alexey 81ae483201 Add regression coverage for ME routing, D2C padding, synlimit, and MSS bulk validation
Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
2026-06-30 13:13:11 +03:00
Alexey ed1895d6df Merge pull request #865 from telemt/flow
Secure + VersionD Outbound Paddings Fix
2026-06-29 17:37:13 +03:00
Alexey 88d161a5e9 Bump -> 3.4.22 2026-06-29 16:38:09 +03:00
Alexey a0ac108807 Secure + VersionD Outbound Paddings Fix 2026-06-29 13:56:16 +03:00
Alexey 809352fac5 Merge pull request #864 from telemt/flow
Restore ME writer source IP for initial proxy request binding
2026-06-29 12:51:04 +03:00
Alexey 22627b498d Bump -> 3.4.21 2026-06-29 12:44:03 +03:00
Alexey b9c5c71dbc Restore ME writer source IP for initial proxy request binding 2026-06-29 12:37:31 +03:00
62 changed files with 7524 additions and 410 deletions
+13 -11
View File
@@ -1,18 +1,20 @@
# Issues
## Warnung
### Warnung
Before opening Issue, if it is more question than problem or bug - ask about that [in our chat](https://t.me/telemtrs)
## What it is not
- NOT Question and Answer
- NOT Helpdesk
***Each of your Issues triggers attempts to reproduce problems and analyze them, which are done manually by people***
Issues is **NOT** about:
- Question and Answer
- Helpdesk
- Configuration or Intergraion Support
---
# Pull Requests
## General
### General
- ONLY signed and verified commits
- ONLY from your name
- DO NOT commit with `codex`, `claude`, or other AI tools as author/committer
@@ -20,7 +22,7 @@ Before opening Issue, if it is more question than problem or bug - ask about tha
---
## Definition of Ready (MANDATORY)
### Definition of Ready (MANDATORY)
A Pull Request WILL be ignored or closed if:
@@ -32,14 +34,14 @@ A Pull Request WILL be ignored or closed if:
---
## Blessed Principles
### Blessed Principles
- PR must build
- PR must pass tests
- PR must be understood by author
---
## AI Usage Policy
### AI Usage Policy
AI tools (Claude, ChatGPT, Codex, DeepSeek, etc.) are allowed as **assistants**, NOT as decision-makers.
@@ -60,7 +62,7 @@ PRs that look like unverified AI dumps WILL be closed
---
## Maintainer Policy
### Maintainer Policy
Maintainers reserve the right to:
@@ -72,7 +74,7 @@ Respect the reviewers time
---
## Enforcement
### Enforcement
Pull Requests that violate project standards may be closed without review.
Generated
+2 -1
View File
@@ -21,6 +21,7 @@ dependencies = [
"cfg-if",
"cipher",
"cpufeatures 0.2.17",
"zeroize",
]
[[package]]
@@ -2899,7 +2900,7 @@ checksum = "7b2093cf4c8eb1e67749a6762251bc9cd836b6fc171623bd0a9d324d37af2417"
[[package]]
name = "telemt"
version = "3.4.19"
version = "3.4.24"
dependencies = [
"aes",
"anyhow",
+3 -3
View File
@@ -1,6 +1,6 @@
[package]
name = "telemt"
version = "3.4.19"
version = "3.4.24"
edition = "2024"
[features]
@@ -15,8 +15,8 @@ tokio = { version = "1.52.3", features = ["full", "tracing"] }
tokio-util = { version = "0.7.18", features = ["full"] }
# Crypto
aes = "0.8.4"
ctr = "0.9.2"
aes = { version = "0.8.4", features = ["zeroize"] }
ctr = { version = "0.9.2", features = ["zeroize"] }
cbc = "0.1.2"
sha2 = "0.10.9"
sha1 = "0.10.6"
+2 -3
View File
@@ -4,7 +4,6 @@ services:
context: .
target: prod-netfilter
network_mode: host
ports: []
ports: !reset []
cap_add:
- NET_BIND_SERVICE
- NET_ADMIN
- NET_ADMIN
File diff suppressed because it is too large Load Diff
+5 -3
View File
@@ -3186,7 +3186,7 @@ If your backend or network is very bandwidth-constrained, reduce cap first. If p
| [`replay_window_secs`](#replay_window_secs) | `u64` | `120` | `` |
| [`ignore_time_skew`](#ignore_time_skew) | `bool` | `false` | `` |
| [`user_rate_limits`](#user_rate_limits) | `Map<String, RateLimitBps>` | `{}` | `` |
| [`cidr_rate_limits`](#cidr_rate_limits) | `Map<IpNetwork, RateLimitBps>` | `{}` | `` |
| [`cidr_rate_limits`](#cidr_rate_limits) | `Map<CidrRateLimitKey, RateLimitBps>` | `{}` | `` |
## users
- **Constraints / validation**: Must not be empty (at least one user must exist). Each value must be **exactly 32 hex characters**.
@@ -3349,13 +3349,15 @@ If your backend or network is very bandwidth-constrained, reduce cap first. If p
alice = { up_bps = 1048576, down_bps = 2097152 }
```
## cidr_rate_limits
- **Constraints / validation**: Table `CIDR -> { up_bps, down_bps }`. CIDR must parse as `IpNetwork`; at least one direction must be non-zero.
- **Description**: Source-subnet bandwidth caps applied alongside per-user limits.
- **Constraints / validation**: Table `CIDR or auto-template -> { up_bps, down_bps }`. Explicit CIDR keys must parse as `IpNetwork`; auto-template keys must be `*4/N` (`N=0..32`), `*6/N` (`N=0..128`), or `*/N` (`N=0..32`). At least one direction must be non-zero. Duplicate normalized auto-templates are rejected.
- **Description**: Source-subnet bandwidth caps applied alongside per-user limits. Explicit CIDR rules use longest-prefix-wins and take priority over auto-templates. Auto-templates create buckets lazily per matched source subnet: `*4/N` for IPv4, `*6/N` for IPv6, and `*/N` as a dual-stack shorthand where IPv4 uses `/N` and IPv6 uses `/(N * 4)`.
- **Example**:
```toml
[access.cidr_rate_limits]
"203.0.113.0/24" = { up_bps = 0, down_bps = 1048576 }
"*4/32" = { up_bps = 262144, down_bps = 1048576 }
"*6/64" = { up_bps = 262144, down_bps = 1048576 }
```
# [[upstreams]]
+5 -3
View File
@@ -3112,7 +3112,7 @@
| [`replay_window_secs`](#replay_window_secs) | `u64` | `120` | `✘` |
| [`ignore_time_skew`](#ignore_time_skew) | `bool` | `false` | `✘` |
| [`user_rate_limits`](#user_rate_limits) | `Map<String, RateLimitBps>` | `{}` | `✔` |
| [`cidr_rate_limits`](#cidr_rate_limits) | `Map<IpNetwork, RateLimitBps>` | `{}` | `✔` |
| [`cidr_rate_limits`](#cidr_rate_limits) | `Map<CidrRateLimitKey, RateLimitBps>` | `{}` | `✔` |
## users
- **Ограничения / валидация**: Не должно быть пустым (должен существовать хотя бы один пользователь). Каждое значение должно состоять **ровно из 32 шестнадцатеричных символов**.
@@ -3265,13 +3265,15 @@
alice = { up_bps = 1048576, down_bps = 2097152 }
```
## cidr_rate_limits
- **Ограничения / валидация**: Таблица `CIDR -> { up_bps, down_bps }`. CIDR должен корректно разбираться как `IpNetwork`; хотя бы одно направление должно быть ненулевым.
- **Описание**: Лимиты скорости для подсетей источников, применяются поверх пользовательских ограничений.
- **Ограничения / валидация**: Таблица `CIDR или auto-template -> { up_bps, down_bps }`. Explicit CIDR-ключи должны корректно разбираться как `IpNetwork`; auto-template ключи должны иметь вид `*4/N` (`N=0..32`), `*6/N` (`N=0..128`) или `*/N` (`N=0..32`). Хотя бы одно направление должно быть ненулевым. Дублирующиеся нормализованные auto-template отклоняются.
- **Описание**: Лимиты скорости для подсетей источников, применяются поверх пользовательских ограничений. Explicit CIDR-правила используют longest-prefix-wins и имеют приоритет над auto-template. Auto-template создают bucket’ы лениво по matched source subnet: `*4/N` для IPv4, `*6/N` для IPv6, а `*/N` является dual-stack shorthand, где IPv4 использует `/N`, а IPv6 — `/(N * 4)`.
- **Example**:
```toml
[access.cidr_rate_limits]
"203.0.113.0/24" = { up_bps = 0, down_bps = 1048576 }
"*4/32" = { up_bps = 262144, down_bps = 1048576 }
"*6/64" = { up_bps = 262144, down_bps = 1048576 }
```
# [[upstreams]]
+22
View File
@@ -24,6 +24,10 @@ const DEFAULT_ME_ADAPTIVE_FLOOR_MAX_WARM_WRITERS_GLOBAL: u32 = 256;
const DEFAULT_ME_ROUTE_BACKPRESSURE_ENABLED: bool = false;
const DEFAULT_ME_ROUTE_FAIRSHARE_ENABLED: bool = false;
const DEFAULT_ME_WRITER_CMD_CHANNEL_CAPACITY: usize = 4096;
pub(crate) const ME_WRITER_BYTE_PERMIT_UNIT_BYTES: usize = 16 * 1024;
pub(crate) const ME_WRITER_FRAME_OVERHEAD_RESERVE_BYTES: usize = 256;
const DEFAULT_ME_WRITER_BYTE_BUDGET_BYTES: usize =
32 * 1024 * 1024 + ME_WRITER_BYTE_PERMIT_UNIT_BYTES;
const DEFAULT_ME_ROUTE_CHANNEL_CAPACITY: usize = 768;
const DEFAULT_ME_C2ME_CHANNEL_CAPACITY: usize = 1024;
const DEFAULT_ME_READER_ROUTE_DATA_WAIT_MS: u64 = 2;
@@ -35,6 +39,8 @@ const DEFAULT_ME_QUOTA_SOFT_OVERSHOOT_BYTES: u64 = 64 * 1024;
const DEFAULT_ME_D2C_FRAME_BUF_SHRINK_THRESHOLD_BYTES: usize = 256 * 1024;
const DEFAULT_DIRECT_RELAY_COPY_BUF_C2S_BYTES: usize = 64 * 1024;
const DEFAULT_DIRECT_RELAY_COPY_BUF_S2C_BYTES: usize = 256 * 1024;
pub(crate) const DIRECT_RELAY_BUFFER_BUDGET_UNIT_BYTES: usize = 4 * 1024;
const DEFAULT_DIRECT_RELAY_BUFFER_BUDGET_MAX_BYTES: usize = 0;
const DEFAULT_ME_WRITER_PICK_SAMPLE_SIZE: u8 = 3;
const DEFAULT_ME_HEALTH_INTERVAL_MS_UNHEALTHY: u64 = 1000;
const DEFAULT_ME_HEALTH_INTERVAL_MS_HEALTHY: u64 = 3000;
@@ -455,6 +461,18 @@ pub(crate) fn default_me_writer_cmd_channel_capacity() -> usize {
DEFAULT_ME_WRITER_CMD_CHANNEL_CAPACITY
}
pub(crate) fn default_me_writer_byte_budget_bytes() -> usize {
DEFAULT_ME_WRITER_BYTE_BUDGET_BYTES
}
pub(crate) fn minimum_me_writer_byte_budget_bytes(max_client_frame: usize) -> usize {
max_client_frame
.saturating_mul(2)
.saturating_add(ME_WRITER_FRAME_OVERHEAD_RESERVE_BYTES)
.div_ceil(ME_WRITER_BYTE_PERMIT_UNIT_BYTES)
.saturating_mul(ME_WRITER_BYTE_PERMIT_UNIT_BYTES)
}
pub(crate) fn default_me_route_channel_capacity() -> usize {
DEFAULT_ME_ROUTE_CHANNEL_CAPACITY
}
@@ -499,6 +517,10 @@ pub(crate) fn default_direct_relay_copy_buf_s2c_bytes() -> usize {
DEFAULT_DIRECT_RELAY_COPY_BUF_S2C_BYTES
}
pub(crate) fn default_direct_relay_buffer_budget_max_bytes() -> usize {
DEFAULT_DIRECT_RELAY_BUFFER_BUDGET_MAX_BYTES
}
pub(crate) fn default_me_writer_pick_sample_size() -> u8 {
DEFAULT_ME_WRITER_PICK_SAMPLE_SIZE
}
+3 -4
View File
@@ -36,8 +36,8 @@ use tracing::{error, info, warn};
use super::load::{LoadedConfig, ProxyConfig};
use crate::config::{
ListenerConfig, LogLevel, MeBindStaleMode, MeFloorMode, MeSocksKdfPolicy, MeTelemetryLevel,
MeWriterPickMode, SynLimitMode,
CidrRateLimitKey, ListenerConfig, LogLevel, MeBindStaleMode, MeFloorMode, MeSocksKdfPolicy,
MeTelemetryLevel, MeWriterPickMode, SynLimitMode,
};
const HOT_RELOAD_DEBOUNCE: Duration = Duration::from_millis(50);
@@ -128,8 +128,7 @@ pub struct HotFields {
pub user_expirations: std::collections::HashMap<String, chrono::DateTime<chrono::Utc>>,
pub user_data_quota: std::collections::HashMap<String, u64>,
pub user_rate_limits: std::collections::HashMap<String, crate::config::RateLimitBps>,
pub cidr_rate_limits:
std::collections::HashMap<ipnetwork::IpNetwork, crate::config::RateLimitBps>,
pub cidr_rate_limits: std::collections::HashMap<CidrRateLimitKey, crate::config::RateLimitBps>,
pub user_max_unique_ips: std::collections::HashMap<String, usize>,
pub user_max_unique_ips_global_each: usize,
pub user_max_unique_ips_mode: crate::config::UserMaxUniqueIpsMode,
+51
View File
@@ -39,8 +39,11 @@ use self::validation::{
};
const MAX_ME_WRITER_CMD_CHANNEL_CAPACITY: usize = 16_384;
const MAX_ME_WRITER_BYTE_BUDGET_BYTES: usize = 256 * 1024 * 1024;
const MAX_ME_ROUTE_CHANNEL_CAPACITY: usize = 8_192;
const MAX_ME_C2ME_CHANNEL_CAPACITY: usize = 8_192;
const MIN_DIRECT_RELAY_BUFFER_BUDGET_BYTES: usize = 16 * 1024 * 1024;
const MAX_DIRECT_RELAY_BUFFER_BUDGET_BYTES: usize = 2 * 1024 * 1024 * 1024;
const MIN_MAX_CLIENT_FRAME_BYTES: usize = 4 * 1024;
const MAX_MAX_CLIENT_FRAME_BYTES: usize = 16 * 1024 * 1024;
const MAX_API_REQUEST_BODY_LIMIT_BYTES: usize = 1024 * 1024;
@@ -540,6 +543,22 @@ impl ProxyConfig {
)));
}
let min_writer_byte_budget =
minimum_me_writer_byte_budget_bytes(config.general.max_client_frame);
if config.general.me_writer_byte_budget_bytes % ME_WRITER_BYTE_PERMIT_UNIT_BYTES != 0 {
return Err(ProxyError::Config(format!(
"general.me_writer_byte_budget_bytes must be a multiple of {ME_WRITER_BYTE_PERMIT_UNIT_BYTES}"
)));
}
if !(min_writer_byte_budget..=MAX_ME_WRITER_BYTE_BUDGET_BYTES)
.contains(&config.general.me_writer_byte_budget_bytes)
{
return Err(ProxyError::Config(format!(
"general.me_writer_byte_budget_bytes must be within [{min_writer_byte_budget}, {MAX_ME_WRITER_BYTE_BUDGET_BYTES}] for general.max_client_frame={}",
config.general.max_client_frame
)));
}
if config.general.me_c2me_send_timeout_ms > 60_000 {
return Err(ProxyError::Config(
"general.me_c2me_send_timeout_ms must be within [0, 60000]".to_string(),
@@ -599,6 +618,24 @@ impl ProxyConfig {
));
}
if config.general.direct_relay_buffer_budget_max_bytes != 0 {
if config.general.direct_relay_buffer_budget_max_bytes
% DIRECT_RELAY_BUFFER_BUDGET_UNIT_BYTES
!= 0
{
return Err(ProxyError::Config(format!(
"general.direct_relay_buffer_budget_max_bytes must be 0 or a multiple of {DIRECT_RELAY_BUFFER_BUDGET_UNIT_BYTES}"
)));
}
if !(MIN_DIRECT_RELAY_BUFFER_BUDGET_BYTES..=MAX_DIRECT_RELAY_BUFFER_BUDGET_BYTES)
.contains(&config.general.direct_relay_buffer_budget_max_bytes)
{
return Err(ProxyError::Config(format!(
"general.direct_relay_buffer_budget_max_bytes must be 0 or within [{MIN_DIRECT_RELAY_BUFFER_BUDGET_BYTES}, {MAX_DIRECT_RELAY_BUFFER_BUDGET_BYTES}]"
)));
}
}
if config.general.me_health_interval_ms_unhealthy == 0 {
return Err(ProxyError::Config(
"general.me_health_interval_ms_unhealthy must be > 0".to_string(),
@@ -673,6 +710,16 @@ impl ProxyConfig {
)));
}
}
let mut cidr_auto_templates = HashSet::new();
for cidr in config.access.cidr_rate_limits.keys() {
for template in cidr.auto_templates().into_iter().flatten() {
if !cidr_auto_templates.insert(template) {
return Err(ProxyError::Config(format!(
"access.cidr_rate_limits.{cidr} duplicates normalized auto-template {template}"
)));
}
}
}
if config.general.me_reinit_every_secs == 0 {
return Err(ProxyError::Config(
@@ -958,6 +1005,10 @@ impl ProxyConfig {
.server
.client_mss_value()
.map_err(|error| ProxyError::Config(format!("server.client_mss {error}")))?;
config
.server
.client_mss_bulk_value()
.map_err(|error| ProxyError::Config(format!("server.client_mss_bulk {error}")))?;
for (idx, listener) in config.server.listeners.iter().enumerate() {
if listener.client_mss.is_some() {
listener
+2
View File
@@ -52,6 +52,7 @@ const GENERAL_CONFIG_KEYS: &[&str] = &[
"me_keepalive_payload_random",
"rpc_proxy_req_every",
"me_writer_cmd_channel_capacity",
"me_writer_byte_budget_bytes",
"me_route_channel_capacity",
"me_c2me_channel_capacity",
"me_c2me_send_timeout_ms",
@@ -64,6 +65,7 @@ const GENERAL_CONFIG_KEYS: &[&str] = &[
"me_d2c_frame_buf_shrink_threshold_bytes",
"direct_relay_copy_buf_c2s_bytes",
"direct_relay_copy_buf_s2c_bytes",
"direct_relay_buffer_budget_max_bytes",
"crypto_pending_buffer",
"max_client_frame",
"desync_all_full",
+104
View File
@@ -1,4 +1,5 @@
use super::*;
use crate::config::CidrRateLimitKey;
const TEST_SHADOWSOCKS_URL: &str =
"ss://2022-blake3-aes-256-gcm:MDEyMzQ1Njc4OTAxMjM0NTY3ODkwMTIzNDU2Nzg5MDE=@127.0.0.1:8388";
@@ -319,6 +320,80 @@ fn logging_config_is_loaded_from_strict_config() {
assert_eq!(cfg.logging.max_age_secs, 60);
}
#[test]
fn cidr_rate_limits_accept_auto_templates_in_strict_config() {
let cfg = load_config_from_temp_toml(
r#"
[general]
config_strict = true
[censorship]
tls_domain = "example.com"
[access.users]
user = "00000000000000000000000000000000"
[access.cidr_rate_limits]
"*/24" = { up_bps = 1024, down_bps = 0 }
"*4/30" = { up_bps = 0, down_bps = 2048 }
"*6/64" = { up_bps = 4096, down_bps = 0 }
"#,
);
assert!(
cfg.access
.cidr_rate_limits
.contains_key(&CidrRateLimitKey::AutoDual(24))
);
assert!(
cfg.access
.cidr_rate_limits
.contains_key(&CidrRateLimitKey::AutoV4(30))
);
assert!(
cfg.access
.cidr_rate_limits
.contains_key(&CidrRateLimitKey::AutoV6(64))
);
}
#[test]
fn cidr_rate_limits_reject_invalid_auto_template_prefix() {
let error = load_config_error_from_temp_toml(
r#"
[censorship]
tls_domain = "example.com"
[access.users]
user = "00000000000000000000000000000000"
[access.cidr_rate_limits]
"*4/33" = { up_bps = 1024, down_bps = 0 }
"#,
);
assert!(error.contains("prefix must be within 0..=32"));
}
#[test]
fn cidr_rate_limits_reject_duplicate_normalized_auto_templates() {
let error = load_config_error_from_temp_toml(
r#"
[censorship]
tls_domain = "example.com"
[access.users]
user = "00000000000000000000000000000000"
[access.cidr_rate_limits]
"*/32" = { up_bps = 1024, down_bps = 0 }
"*6/128" = { up_bps = 2048, down_bps = 0 }
"#,
);
assert!(error.contains("duplicates normalized auto-template *6/128"));
}
#[test]
fn file_logging_requires_path() {
let error = load_config_error_from_temp_toml(
@@ -1652,6 +1727,7 @@ fn client_mss_custom_value_is_accepted() {
let toml = r#"
[server]
client_mss = "4096"
client_mss_bulk = "1400"
[censorship]
tls_domain = "example.com"
@@ -1665,6 +1741,7 @@ fn client_mss_custom_value_is_accepted() {
let cfg = ProxyConfig::load(&path).unwrap();
assert_eq!(cfg.server.client_mss_value(), Ok(Some(4096)));
assert_eq!(cfg.server.client_mss_bulk_value(), Ok(Some(1400)));
let _ = std::fs::remove_file(path);
}
@@ -1693,6 +1770,33 @@ fn client_mss_out_of_range_is_rejected() {
}
}
#[test]
fn client_mss_bulk_out_of_range_is_rejected() {
for value in ["87", "4097"] {
let toml = format!(
r#"
[server]
client_mss_bulk = "{value}"
[censorship]
tls_domain = "example.com"
[access.users]
user = "00000000000000000000000000000000"
"#
);
let dir = std::env::temp_dir();
let path = dir.join(format!(
"telemt_client_mss_bulk_out_of_range_{value}_test.toml"
));
std::fs::write(&path, toml).unwrap();
let err = ProxyConfig::load(&path).unwrap_err().to_string();
assert!(err.contains("server.client_mss_bulk custom value must be within [88, 4096]"));
let _ = std::fs::remove_file(path);
}
}
#[test]
fn client_mss_unquoted_number_is_rejected() {
let toml = r#"
@@ -95,6 +95,101 @@ max_client_frame = 16777217
remove_temp_config(&path);
}
#[test]
fn load_rejects_writer_byte_budget_below_frame_residency_minimum() {
let path = write_temp_config(
r#"
[general]
max_client_frame = 16777216
me_writer_byte_budget_bytes = 33554432
"#,
);
let err = ProxyConfig::load(&path)
.expect_err("writer byte budget below frame residency minimum must fail");
let msg = err.to_string();
assert!(
msg.contains("general.me_writer_byte_budget_bytes must be within [33570816, 268435456]"),
"error must explain writer byte budget minimum, got: {msg}"
);
remove_temp_config(&path);
}
#[test]
fn load_rejects_unaligned_writer_byte_budget() {
let path = write_temp_config(
r#"
[general]
me_writer_byte_budget_bytes = 33570817
"#,
);
let err = ProxyConfig::load(&path)
.expect_err("writer byte budget outside permit granularity must fail");
let msg = err.to_string();
assert!(
msg.contains("general.me_writer_byte_budget_bytes must be a multiple of 16384"),
"error must explain writer byte budget alignment, got: {msg}"
);
remove_temp_config(&path);
}
#[test]
fn load_rejects_writer_byte_budget_above_hard_cap() {
let path = write_temp_config(
r#"
[general]
me_writer_byte_budget_bytes = 268451840
"#,
);
let err = ProxyConfig::load(&path).expect_err("writer byte budget above hard cap must fail");
let msg = err.to_string();
assert!(
msg.contains("general.me_writer_byte_budget_bytes must be within [33570816, 268435456]"),
"error must explain writer byte budget hard cap, got: {msg}"
);
remove_temp_config(&path);
}
#[test]
fn load_rejects_unaligned_direct_relay_buffer_budget() {
let path = write_temp_config(
r#"
[general]
direct_relay_buffer_budget_max_bytes = 16777217
"#,
);
let err = ProxyConfig::load(&path).expect_err("unaligned direct relay buffer budget must fail");
assert!(
err.to_string().contains(
"general.direct_relay_buffer_budget_max_bytes must be 0 or a multiple of 4096"
)
);
remove_temp_config(&path);
}
#[test]
fn load_rejects_direct_relay_buffer_budget_above_hard_cap() {
let path = write_temp_config(
r#"
[general]
direct_relay_buffer_budget_max_bytes = 2147487744
"#,
);
let err =
ProxyConfig::load(&path).expect_err("direct relay buffer budget above hard cap must fail");
assert!(err.to_string().contains(
"general.direct_relay_buffer_budget_max_bytes must be 0 or within [16777216, 2147483648]"
));
remove_temp_config(&path);
}
#[test]
fn load_rejects_listen_backlog_above_i32_upper_bound() {
let path = write_temp_config(
@@ -139,16 +234,23 @@ fn load_accepts_memory_limits_at_hard_upper_bounds() {
r#"
[general]
me_writer_cmd_channel_capacity = 16384
me_writer_byte_budget_bytes = 268435456
me_route_channel_capacity = 8192
me_c2me_channel_capacity = 8192
direct_relay_buffer_budget_max_bytes = 2147483648
max_client_frame = 16777216
"#,
);
let cfg = ProxyConfig::load(&path).expect("hard upper bound values must be accepted");
assert_eq!(cfg.general.me_writer_cmd_channel_capacity, 16384);
assert_eq!(cfg.general.me_writer_byte_budget_bytes, 256 * 1024 * 1024);
assert_eq!(cfg.general.me_route_channel_capacity, 8192);
assert_eq!(cfg.general.me_c2me_channel_capacity, 8192);
assert_eq!(
cfg.general.direct_relay_buffer_budget_max_bytes,
2 * 1024 * 1024 * 1024
);
assert_eq!(cfg.general.max_client_frame, 16 * 1024 * 1024);
remove_temp_config(&path);
+166 -6
View File
@@ -2,6 +2,7 @@ use chrono::{DateTime, Utc};
use ipnetwork::IpNetwork;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fmt;
use std::net::IpAddr;
use std::path::PathBuf;
@@ -578,6 +579,10 @@ pub struct GeneralConfig {
#[serde(default = "default_me_writer_cmd_channel_capacity")]
pub me_writer_cmd_channel_capacity: usize,
/// Resident-memory budget in bytes for each ME writer data queue.
#[serde(default = "default_me_writer_byte_budget_bytes")]
pub me_writer_byte_budget_bytes: usize,
/// Capacity of per-connection ME response route channel.
#[serde(default = "default_me_route_channel_capacity")]
pub me_route_channel_capacity: usize,
@@ -621,7 +626,7 @@ pub struct GeneralConfig {
#[serde(default = "default_me_d2c_frame_buf_shrink_threshold_bytes")]
pub me_d2c_frame_buf_shrink_threshold_bytes: usize,
/// Copy buffer size for client->DC direction in direct relay.
/// Copy buffer ceiling for client->DC direction in direct relay.
///
/// This is also the upper bound for one amortized upload rate-limit burst:
/// upload debt is settled before the next relay read instead of blocking
@@ -629,13 +634,18 @@ pub struct GeneralConfig {
#[serde(default = "default_direct_relay_copy_buf_c2s_bytes")]
pub direct_relay_copy_buf_c2s_bytes: usize,
/// Copy buffer size for DC->client direction in direct relay.
/// Copy buffer ceiling for DC->client direction in direct relay.
///
/// This bounds one direct download rate-limit grant because writes are
/// clipped to the currently available shaper budget.
#[serde(default = "default_direct_relay_copy_buf_s2c_bytes")]
pub direct_relay_copy_buf_s2c_bytes: usize,
/// Process-wide hard ceiling for Direct relay copy buffers.
/// `0` derives the ceiling from host and cgroup memory limits.
#[serde(default = "default_direct_relay_buffer_budget_max_bytes")]
pub direct_relay_buffer_budget_max_bytes: usize,
/// Max pending ciphertext buffer per client writer (bytes).
/// Controls FakeTLS backpressure vs throughput.
#[serde(default = "default_crypto_pending_buffer")]
@@ -1102,6 +1112,7 @@ impl Default for GeneralConfig {
me_keepalive_payload_random: default_true(),
rpc_proxy_req_every: default_rpc_proxy_req_every(),
me_writer_cmd_channel_capacity: default_me_writer_cmd_channel_capacity(),
me_writer_byte_budget_bytes: default_me_writer_byte_budget_bytes(),
me_route_channel_capacity: default_me_route_channel_capacity(),
me_c2me_channel_capacity: default_me_c2me_channel_capacity(),
me_c2me_send_timeout_ms: default_me_c2me_send_timeout_ms(),
@@ -1115,6 +1126,7 @@ impl Default for GeneralConfig {
default_me_d2c_frame_buf_shrink_threshold_bytes(),
direct_relay_copy_buf_c2s_bytes: default_direct_relay_copy_buf_c2s_bytes(),
direct_relay_copy_buf_s2c_bytes: default_direct_relay_copy_buf_s2c_bytes(),
direct_relay_buffer_budget_max_bytes: default_direct_relay_buffer_budget_max_bytes(),
me_warmup_stagger_enabled: default_true(),
me_warmup_step_delay_ms: default_warmup_step_delay_ms(),
me_warmup_step_jitter_ms: default_warmup_step_jitter_ms(),
@@ -2098,11 +2110,13 @@ pub struct AccessConfig {
/// Per-CIDR aggregate transport rate limits in bits-per-second.
///
/// Matching uses longest-prefix-wins semantics. A value of `0` in one
/// direction means "unlimited" for that direction. Limits are amortized
/// with the same bounded-burst contract as per-user rate limits.
/// Explicit CIDR keys use longest-prefix-wins semantics. Auto-template
/// keys (`*4/N`, `*6/N`, `*/N`) lazily create per-source-subnet buckets
/// after explicit CIDR matching misses. A value of `0` in one direction
/// means "unlimited" for that direction. Limits are amortized with the
/// same bounded-burst contract as per-user rate limits.
#[serde(default)]
pub cidr_rate_limits: HashMap<IpNetwork, RateLimitBps>,
pub cidr_rate_limits: HashMap<CidrRateLimitKey, RateLimitBps>,
/// Per-username client source IP/CIDR deny list. Checked after successful
/// authentication; matching IPs get the same rejection path as invalid auth
@@ -2171,10 +2185,156 @@ impl AccessConfig {
}
}
/// Key used by `access.cidr_rate_limits`.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum CidrRateLimitKey {
/// Explicit source CIDR rule.
Network(IpNetwork),
/// IPv4 auto-template that creates one bucket for each matching `/N`.
AutoV4(u8),
/// IPv6 auto-template that creates one bucket for each matching `/N`.
AutoV6(u8),
/// Dual-stack auto-template; IPv4 uses `/N`, IPv6 uses `/(N * 4)`.
AutoDual(u8),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub(crate) enum CidrAutoTemplateFamily {
V4,
V6,
}
impl CidrAutoTemplateFamily {
pub(crate) fn marker(self) -> &'static str {
match self {
Self::V4 => "*4",
Self::V6 => "*6",
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub(crate) struct CidrAutoTemplate {
pub(crate) family: CidrAutoTemplateFamily,
pub(crate) prefix_len: u8,
}
impl fmt::Display for CidrAutoTemplate {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(formatter, "{}/{}", self.family.marker(), self.prefix_len)
}
}
impl CidrRateLimitKey {
pub(crate) fn auto_templates(&self) -> [Option<CidrAutoTemplate>; 2] {
match *self {
Self::Network(_) => [None, None],
Self::AutoV4(prefix_len) => [
Some(CidrAutoTemplate {
family: CidrAutoTemplateFamily::V4,
prefix_len,
}),
None,
],
Self::AutoV6(prefix_len) => [
Some(CidrAutoTemplate {
family: CidrAutoTemplateFamily::V6,
prefix_len,
}),
None,
],
Self::AutoDual(prefix_len) => [
Some(CidrAutoTemplate {
family: CidrAutoTemplateFamily::V4,
prefix_len,
}),
Some(CidrAutoTemplate {
family: CidrAutoTemplateFamily::V6,
prefix_len: prefix_len.saturating_mul(4),
}),
],
}
}
}
impl fmt::Display for CidrRateLimitKey {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Network(cidr) => write!(formatter, "{cidr}"),
Self::AutoV4(prefix_len) => write!(formatter, "*4/{prefix_len}"),
Self::AutoV6(prefix_len) => write!(formatter, "*6/{prefix_len}"),
Self::AutoDual(prefix_len) => write!(formatter, "*/{prefix_len}"),
}
}
}
impl Serialize for CidrRateLimitKey {
fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
serializer.collect_str(self)
}
}
impl<'de> Deserialize<'de> for CidrRateLimitKey {
fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
let value = String::deserialize(deserializer)?;
parse_cidr_rate_limit_key(&value).map_err(serde::de::Error::custom)
}
}
fn parse_cidr_rate_limit_key(value: &str) -> std::result::Result<CidrRateLimitKey, String> {
if let Some(prefix) = value.strip_prefix("*4/") {
return parse_cidr_auto_prefix(value, prefix, 32).map(CidrRateLimitKey::AutoV4);
}
if let Some(prefix) = value.strip_prefix("*6/") {
return parse_cidr_auto_prefix(value, prefix, 128).map(CidrRateLimitKey::AutoV6);
}
if let Some(prefix) = value.strip_prefix("*/") {
return parse_cidr_auto_prefix(value, prefix, 32).map(CidrRateLimitKey::AutoDual);
}
if value.starts_with('*') {
return Err(format!(
"invalid CIDR rate limit key {value:?}; expected CIDR, *4/N, *6/N, or */N"
));
}
value
.parse::<IpNetwork>()
.map(CidrRateLimitKey::Network)
.map_err(|error| {
format!(
"invalid CIDR rate limit key {value:?}: {error}; expected CIDR, *4/N, *6/N, or */N"
)
})
}
fn parse_cidr_auto_prefix(
key: &str,
prefix: &str,
max_prefix: u8,
) -> std::result::Result<u8, String> {
let prefix = prefix.parse::<u8>().map_err(|_| {
format!("invalid CIDR auto-template key {key:?}; prefix must be within 0..={max_prefix}")
})?;
if prefix > max_prefix {
return Err(format!(
"invalid CIDR auto-template key {key:?}; prefix must be within 0..={max_prefix}"
));
}
Ok(prefix)
}
/// Transport rate limit in bits-per-second.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct RateLimitBps {
/// Upload direction limit in bits-per-second; `0` means unlimited.
#[serde(default)]
pub up_bps: u64,
/// Download direction limit in bits-per-second; `0` means unlimited.
#[serde(default)]
pub down_bps: u64,
}
+18 -12
View File
@@ -4,12 +4,12 @@
//!
//! ## Zeroize policy
//!
//! - `AesCbc` stores raw key/IV bytes and zeroizes them on drop.
//! - `AesCtr` wraps an opaque `Aes256Ctr` cipher from the `ctr` crate.
//! The expanded key schedule lives inside that type and cannot be
//! zeroized from outside. Callers that hold raw key material (e.g.
//! `HandshakeSuccess`, `ObfuscationParams`) are responsible for
//! zeroizing their own copies.
//! - `AesCbc` stores raw key/IV bytes and zeroizes them on drop. Temporary
//! expanded key schedules are also zeroized by the RustCrypto backend.
//! - `AesCtr` uses the RustCrypto `zeroize` contract to clear its expanded
//! key schedule, counter, and buffered keystream on drop.
//! - Callers that hold raw key material (e.g. `HandshakeSuccess`,
//! `ObfuscationParams`) remain responsible for zeroizing their own copies.
#![allow(dead_code)]
@@ -19,24 +19,28 @@ use ctr::{
Ctr128BE,
cipher::{KeyIvInit, StreamCipher},
};
use zeroize::Zeroize;
use zeroize::{Zeroize, ZeroizeOnDrop};
type Aes256Ctr = Ctr128BE<Aes256>;
static_assertions::assert_impl_all!(Aes256: ZeroizeOnDrop);
static_assertions::assert_impl_all!(Aes256Ctr: ZeroizeOnDrop);
// ============= AES-256-CTR =============
/// AES-256-CTR encryptor/decryptor
///
/// CTR mode is symmetric — encryption and decryption are the same operation.
///
/// **Zeroize note:** The inner `Aes256Ctr` cipher state (expanded key schedule
/// + counter) is opaque and cannot be zeroized. If you need to protect key
/// material, zeroize the `[u8; 32]` key and `u128` IV at the call site
/// before dropping them.
/// **Zeroize note:** The inner `Aes256Ctr` zeroizes its expanded key schedule,
/// counter, and buffered keystream on drop. Callers remain responsible for
/// zeroizing their own raw key and IV copies.
pub struct AesCtr {
cipher: Aes256Ctr,
}
impl ZeroizeOnDrop for AesCtr {}
impl AesCtr {
/// Create new AES-CTR cipher with key and IV
pub fn new(key: &[u8; 32], iv: u128) -> Self {
@@ -92,7 +96,7 @@ impl AesCtr {
/// are different operations. This implementation handles CBC chaining
/// correctly across multiple blocks.
///
/// Key and IV are zeroized on drop.
/// Key, IV, and temporary expanded key schedules are zeroized on drop.
pub struct AesCbc {
key: [u8; 32],
iv: [u8; 16],
@@ -105,6 +109,8 @@ impl Drop for AesCbc {
}
}
impl ZeroizeOnDrop for AesCbc {}
impl AesCbc {
/// AES block size
const BLOCK_SIZE: usize = 16;
+1
View File
@@ -279,6 +279,7 @@ pub(crate) async fn initialize_me_pool(
config.general.me_writer_pick_sample_size,
config.general.me_socks_kdf_policy,
config.general.me_writer_cmd_channel_capacity,
config.general.me_writer_byte_budget_bytes,
config.general.me_route_channel_capacity,
config.general.me_route_backpressure_enabled,
config.general.me_route_fairshare_enabled,
+20 -1
View File
@@ -33,6 +33,9 @@ use crate::conntrack_control;
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,
};
use crate::proxy::route_mode::{RelayRouteMode, RouteRuntimeController};
use crate::proxy::shared_state::ProxySharedState;
use crate::startup::{
@@ -473,7 +476,16 @@ async fn run_telemt_core(
config.network.dns_overrides.len()
);
}
let shared_state = ProxySharedState::new();
let direct_buffer_hard_limit =
resolve_direct_buffer_hard_limit(config.general.direct_relay_buffer_budget_max_bytes).await;
let direct_buffer_budget = DirectBufferBudget::new(direct_buffer_hard_limit);
info!(
hard_limit_bytes = direct_buffer_hard_limit,
configured_override_bytes = config.general.direct_relay_buffer_budget_max_bytes,
"Direct relay buffer budget initialized"
);
let shared_state =
ProxySharedState::new_with_direct_buffer_budget(direct_buffer_budget.clone());
shared_state.apply_user_enabled_config(&config.access.user_enabled);
shared_state.traffic_limiter.apply_policy(
config.access.user_rate_limits.clone(),
@@ -882,6 +894,13 @@ async fn run_telemt_core(
stats.clone(),
shared_state.clone(),
);
spawn_direct_buffer_budget_controller(
direct_buffer_budget,
buffer_pool.clone(),
stats.clone(),
shared_state.clone(),
config.server.max_connections,
);
let bound = listeners::bind_listeners(
&config,
+151
View File
@@ -595,6 +595,81 @@ async fn render_metrics(
"telemt_buffer_pool_buffers_total{{kind=\"in_use\"}} {}",
stats.get_buffer_pool_in_use_gauge()
);
let _ = writeln!(
out,
"# HELP telemt_buffer_pool_events_total Buffer-pool allocation lifecycle events"
);
let _ = writeln!(out, "# TYPE telemt_buffer_pool_events_total counter");
let _ = writeln!(
out,
"telemt_buffer_pool_events_total{{event=\"replaced_nonstandard\"}} {}",
stats.get_buffer_pool_replaced_nonstandard_total()
);
let direct_budget = shared_state.direct_buffer_budget.snapshot();
let _ = writeln!(
out,
"# HELP telemt_direct_relay_buffer_budget_bytes Direct relay copy-buffer budget and memory inputs"
);
let _ = writeln!(out, "# TYPE telemt_direct_relay_buffer_budget_bytes gauge");
for (kind, value) in [
("hard_limit", direct_budget.hard_limit_bytes),
("target", direct_budget.target_bytes),
("reserved", direct_budget.reserved_bytes),
("memory_total", direct_budget.memory_total_bytes),
("memory_available", direct_budget.memory_available_bytes),
("process_rss", direct_budget.process_rss_bytes),
] {
let _ = writeln!(
out,
"telemt_direct_relay_buffer_budget_bytes{{kind=\"{}\"}} {}",
kind, value
);
}
let _ = writeln!(
out,
"# HELP telemt_direct_relay_buffer_budget_events_total Direct relay buffer-budget lifecycle events"
);
let _ = writeln!(
out,
"# TYPE telemt_direct_relay_buffer_budget_events_total counter"
);
for (result, value) in [
("promotion", direct_budget.promotion_total),
("promotion_denied", direct_budget.promotion_denied_total),
("minimum_fallback", direct_budget.minimum_fallback_total),
("admission_rejected", direct_budget.admission_rejected_total),
("quiet_demotion", direct_budget.quiet_demotion_total),
(
"write_pressure_demotion",
direct_budget.write_pressure_demotion_total,
),
(
"global_pressure_demotion",
direct_budget.global_pressure_demotion_total,
),
] {
let _ = writeln!(
out,
"telemt_direct_relay_buffer_budget_events_total{{result=\"{}\"}} {}",
result, value
);
}
let _ = writeln!(
out,
"# HELP telemt_direct_relay_buffer_sessions Current Direct relay sessions by adaptive tier"
);
let _ = writeln!(out, "# TYPE telemt_direct_relay_buffer_sessions gauge");
for (tier, value) in ["base", "tier1", "tier2", "tier3"]
.into_iter()
.zip(direct_budget.tier_sessions)
{
let _ = writeln!(
out,
"telemt_direct_relay_buffer_sessions{{tier=\"{}\"}} {}",
tier, value
);
}
let _ = writeln!(
out,
@@ -2435,6 +2510,82 @@ async fn render_metrics(
}
);
let _ = writeln!(
out,
"# HELP telemt_me_writer_byte_budget_limit_bytes Configured resident-memory budget per ME writer"
);
let _ = writeln!(out, "# TYPE telemt_me_writer_byte_budget_limit_bytes gauge");
let _ = writeln!(
out,
"telemt_me_writer_byte_budget_limit_bytes {}",
if me_allows_normal {
stats.get_me_writer_byte_budget_limit_bytes_gauge()
} else {
0
}
);
let _ = writeln!(
out,
"# HELP telemt_me_writer_byte_budget_reserved_bytes Aggregate ME writer memory reservations by lifecycle state"
);
let _ = writeln!(
out,
"# TYPE telemt_me_writer_byte_budget_reserved_bytes gauge"
);
let _ = writeln!(
out,
"telemt_me_writer_byte_budget_reserved_bytes{{state=\"queued\"}} {}",
if me_allows_normal {
stats.get_me_writer_byte_budget_queued_bytes_gauge()
} else {
0
}
);
let _ = writeln!(
out,
"telemt_me_writer_byte_budget_reserved_bytes{{state=\"inflight\"}} {}",
if me_allows_normal {
stats.get_me_writer_byte_budget_inflight_bytes_gauge()
} else {
0
}
);
let _ = writeln!(
out,
"# HELP telemt_me_writer_byte_budget_events_total ME writer byte-budget outcomes"
);
let _ = writeln!(
out,
"# TYPE telemt_me_writer_byte_budget_events_total counter"
);
let _ = writeln!(
out,
"telemt_me_writer_byte_budget_events_total{{result=\"wait\"}} {}",
if me_allows_normal {
stats.get_me_writer_byte_budget_wait_total()
} else {
0
}
);
let _ = writeln!(
out,
"telemt_me_writer_byte_budget_events_total{{result=\"timeout\"}} {}",
if me_allows_normal {
stats.get_me_writer_byte_budget_timeout_total()
} else {
0
}
);
let _ = writeln!(
out,
"telemt_me_writer_byte_budget_events_total{{result=\"oversize\"}} {}",
if me_allows_normal {
stats.get_me_writer_byte_budget_oversize_total()
} else {
0
}
);
let _ = writeln!(
out,
"# HELP telemt_me_writer_pick_total ME writer-pick outcomes by mode and result"
+9 -3
View File
@@ -246,7 +246,7 @@ pub fn secure_payload_len_from_wire_len(wire_len: usize) -> Option<usize> {
}
/// Generate padding length for Secure Intermediate protocol.
/// Telegram Desktop uses a 4-bit random padding length for VersionD packets.
/// Outbound padding is 1..=3 so a receiver can strip it by 4-byte alignment.
pub fn secure_padding_len(data_len: usize, rng: &SecureRandom) -> usize {
debug_assert!(
is_valid_secure_payload_len(data_len),
@@ -425,15 +425,21 @@ mod tests {
}
#[test]
fn secure_padding_matches_tdesktop_range() {
fn secure_padding_never_produces_aligned_total() {
let rng = SecureRandom::new();
for data_len in (0..1000).step_by(4) {
for _ in 0..100 {
let padding = secure_padding_len(data_len, &rng);
assert!(
padding <= 15,
(1..=3).contains(&padding),
"padding out of range: data_len={data_len}, padding={padding}"
);
assert_ne!(
(data_len + padding) % 4,
0,
"invariant violated: data_len={data_len}, padding={padding}, total={}",
data_len + padding
);
}
}
}
+4 -4
View File
@@ -8,8 +8,8 @@ pub(crate) const INTERMEDIATE_QUICKACK_FLAG: u32 = 0x8000_0000;
/// Payload length mask used by Intermediate and Secure Intermediate headers.
pub(crate) const INTERMEDIATE_WIRE_LEN_MASK: u32 = 0x7fff_ffff;
/// Maximum random tail length used by Telegram Desktop VersionD packets.
pub(crate) const SECURE_VERSION_D_PADDING_MAX: usize = 15;
/// Maximum outbound Secure tail length that keeps wire lengths non-aligned.
pub(crate) const SECURE_VERSION_D_PADDING_MAX: usize = 3;
/// Parsed Intermediate/Secure Intermediate length header.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
@@ -51,9 +51,9 @@ pub(crate) fn secure_version_d_body_len_from_wire_len(wire_len: usize) -> Option
Some(wire_len - (wire_len % 4))
}
/// Generate Telegram Desktop-compatible VersionD random tail length.
/// Generate outbound Secure tail length without ambiguous full-word padding.
pub(crate) fn secure_version_d_padding_len(rng: &SecureRandom) -> usize {
rng.range(SECURE_VERSION_D_PADDING_MAX + 1)
rng.range(SECURE_VERSION_D_PADDING_MAX) + 1
}
#[cfg(test)]
+137 -41
View File
@@ -1,7 +1,4 @@
#![allow(dead_code)]
// Adaptive buffer policy is staged and retained for deterministic rollout.
// Keep definitions compiled for compatibility and security test scaffolding.
// Adaptive buffer policy shared by active Direct relay sessions.
use dashmap::DashMap;
use std::cmp::max;
@@ -13,29 +10,42 @@ const PROFILE_TTL: Duration = Duration::from_secs(300);
const THROUGHPUT_UP_BPS: f64 = 8_000_000.0;
const THROUGHPUT_DOWN_BPS: f64 = 2_000_000.0;
const RATIO_CONFIRM_THRESHOLD: f64 = 1.12;
const TIER1_HOLD_TICKS: u32 = 8;
const TIER2_HOLD_TICKS: u32 = 4;
const QUIET_DEMOTE_TICKS: u32 = 480;
const HARD_COOLDOWN_TICKS: u32 = 20;
const TIER1_HOLD: Duration = Duration::from_secs(2);
const TIER2_HOLD: Duration = Duration::from_secs(1);
const QUIET_DEMOTE: Duration = Duration::from_secs(120);
const HARD_COOLDOWN: Duration = Duration::from_secs(5);
const SUSTAINED_PRESSURE_DEMOTE: Duration = Duration::from_secs(30);
const PRESSURE_DEMOTE_COOLDOWN: Duration = Duration::from_secs(60);
const HARD_PENDING_THRESHOLD: u32 = 3;
const HARD_PARTIAL_RATIO_THRESHOLD: f64 = 0.25;
#[cfg(test)]
const DIRECT_C2S_CAP_BYTES: usize = 128 * 1024;
#[cfg(test)]
const DIRECT_S2C_CAP_BYTES: usize = 512 * 1024;
#[cfg(test)]
const ME_FRAMES_CAP: usize = 96;
#[cfg(test)]
const ME_BYTES_CAP: usize = 384 * 1024;
#[cfg(test)]
const ME_DELAY_MIN_US: u64 = 150;
const MAX_USER_PROFILES_ENTRIES: usize = 50_000;
const MAX_USER_KEY_BYTES: usize = 512;
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
/// Per-session Direct copy-buffer capacity tier.
pub enum AdaptiveTier {
/// Conservative baseline capacity.
Base = 0,
/// First throughput promotion.
Tier1 = 1,
/// Sustained bidirectional pressure promotion.
Tier2 = 2,
/// Configured per-direction ceilings.
Tier3 = 3,
}
impl AdaptiveTier {
/// Returns the next larger tier, saturating at `Tier3`.
pub fn promote(self) -> Self {
match self {
Self::Base => Self::Tier1,
@@ -45,6 +55,7 @@ impl AdaptiveTier {
}
}
/// Returns the next smaller tier, saturating at `Base`.
pub fn demote(self) -> Self {
match self {
Self::Base => Self::Base,
@@ -54,6 +65,7 @@ impl AdaptiveTier {
}
}
#[cfg(test)]
fn ratio(self) -> (usize, usize) {
match self {
Self::Base => (1, 1),
@@ -63,49 +75,70 @@ impl AdaptiveTier {
}
}
/// Returns the stable numeric tier used by bounded metrics.
pub fn as_u8(self) -> u8 {
self as u8
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
/// Signal that caused an accepted adaptive tier transition.
pub enum TierTransitionReason {
/// Sustained throughput and directional ratio confirmation.
SoftConfirmed,
/// Short pending or partial-write pressure burst.
HardPressure,
/// Sustained low-throughput period.
QuietDemotion,
/// Sustained pending or partial-write pressure.
SustainedWritePressure,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
/// Proposed transition emitted by the per-session controller.
pub struct TierTransition {
/// Tier active before the observation.
pub from: AdaptiveTier,
/// Tier requested after the observation.
pub to: AdaptiveTier,
/// Pressure or throughput condition that requested the transition.
pub reason: TierTransitionReason,
}
#[derive(Debug, Clone, Copy, Default)]
/// Directional byte and write-pressure deltas for one observation period.
pub struct RelaySignalSample {
/// Client-to-DC bytes copied during the period.
pub c2s_bytes: u64,
/// Bytes offered to DC-to-client writes during the period.
pub s2c_requested_bytes: u64,
/// Bytes accepted by DC-to-client writes during the period.
pub s2c_written_bytes: u64,
/// Successful DC-to-client write operations during the period.
pub s2c_write_ops: u64,
/// Partial DC-to-client write operations during the period.
pub s2c_partial_writes: u64,
/// Consecutive pending DC-to-client writes at sample time.
pub s2c_consecutive_pending_writes: u32,
}
#[derive(Debug, Clone, Copy)]
/// Stateful hysteresis controller for one active Direct session.
pub struct SessionAdaptiveController {
tier: AdaptiveTier,
max_tier_seen: AdaptiveTier,
throughput_ema_bps: f64,
incoming_ema_bps: f64,
outgoing_ema_bps: f64,
tier1_hold_ticks: u32,
tier2_hold_ticks: u32,
quiet_ticks: u32,
hard_cooldown_ticks: u32,
tier1_hold: Duration,
tier2_hold: Duration,
quiet: Duration,
hard_cooldown: Duration,
sustained_pressure: Duration,
}
impl SessionAdaptiveController {
/// Creates a controller at the tier whose memory reservation was accepted.
pub fn new(initial_tier: AdaptiveTier) -> Self {
Self {
tier: initial_tier,
@@ -113,25 +146,33 @@ impl SessionAdaptiveController {
throughput_ema_bps: 0.0,
incoming_ema_bps: 0.0,
outgoing_ema_bps: 0.0,
tier1_hold_ticks: 0,
tier2_hold_ticks: 0,
quiet_ticks: 0,
hard_cooldown_ticks: 0,
tier1_hold: Duration::ZERO,
tier2_hold: Duration::ZERO,
quiet: Duration::ZERO,
hard_cooldown: Duration::ZERO,
sustained_pressure: Duration::ZERO,
}
}
/// Returns the highest tier proposed by controller observations.
#[allow(dead_code)]
pub fn max_tier_seen(&self) -> AdaptiveTier {
self.max_tier_seen
}
/// Returns the controller's current logical tier.
pub fn tier(&self) -> AdaptiveTier {
self.tier
}
/// Observes one period and returns at most one hysteresis-controlled transition.
pub fn observe(&mut self, sample: RelaySignalSample, tick_secs: f64) -> Option<TierTransition> {
if tick_secs <= f64::EPSILON {
return None;
}
if self.hard_cooldown_ticks > 0 {
self.hard_cooldown_ticks -= 1;
}
let tick = Duration::from_secs_f64(tick_secs);
self.hard_cooldown = self.hard_cooldown.saturating_sub(tick);
let c2s_bps = (sample.c2s_bytes as f64 * 8.0) / tick_secs;
let incoming_bps = (sample.s2c_requested_bytes as f64 * 8.0) / tick_secs;
@@ -144,9 +185,9 @@ impl SessionAdaptiveController {
let tier1_now = self.throughput_ema_bps >= THROUGHPUT_UP_BPS;
if tier1_now {
self.tier1_hold_ticks = self.tier1_hold_ticks.saturating_add(1);
self.tier1_hold = self.tier1_hold.saturating_add(tick);
} else {
self.tier1_hold_ticks = 0;
self.tier1_hold = Duration::ZERO;
}
let ratio = if self.outgoing_ema_bps <= f64::EPSILON {
@@ -156,9 +197,9 @@ impl SessionAdaptiveController {
};
let tier2_now = ratio >= RATIO_CONFIRM_THRESHOLD;
if tier2_now {
self.tier2_hold_ticks = self.tier2_hold_ticks.saturating_add(1);
self.tier2_hold = self.tier2_hold.saturating_add(tick);
} else {
self.tier2_hold_ticks = 0;
self.tier2_hold = Duration::ZERO;
}
let partial_ratio = if sample.s2c_write_ops == 0 {
@@ -169,24 +210,37 @@ impl SessionAdaptiveController {
let hard_now = sample.s2c_consecutive_pending_writes >= HARD_PENDING_THRESHOLD
|| partial_ratio >= HARD_PARTIAL_RATIO_THRESHOLD;
if hard_now && self.hard_cooldown_ticks == 0 {
return self.promote(TierTransitionReason::HardPressure, HARD_COOLDOWN_TICKS);
if hard_now {
self.sustained_pressure = self.sustained_pressure.saturating_add(tick);
if self.sustained_pressure >= SUSTAINED_PRESSURE_DEMOTE {
self.sustained_pressure = Duration::ZERO;
return self.demote(
TierTransitionReason::SustainedWritePressure,
PRESSURE_DEMOTE_COOLDOWN,
);
}
} else {
self.sustained_pressure = Duration::ZERO;
}
if self.tier1_hold_ticks >= TIER1_HOLD_TICKS && self.tier2_hold_ticks >= TIER2_HOLD_TICKS {
return self.promote(TierTransitionReason::SoftConfirmed, 0);
if hard_now && self.hard_cooldown.is_zero() {
return self.promote(TierTransitionReason::HardPressure, HARD_COOLDOWN);
}
if self.tier1_hold >= TIER1_HOLD && self.tier2_hold >= TIER2_HOLD {
return self.promote(TierTransitionReason::SoftConfirmed, Duration::ZERO);
}
let demote_candidate =
self.throughput_ema_bps < THROUGHPUT_DOWN_BPS && !tier2_now && !hard_now;
if demote_candidate {
self.quiet_ticks = self.quiet_ticks.saturating_add(1);
if self.quiet_ticks >= QUIET_DEMOTE_TICKS {
self.quiet_ticks = 0;
return self.demote(TierTransitionReason::QuietDemotion);
self.quiet = self.quiet.saturating_add(tick);
if self.quiet >= QUIET_DEMOTE {
self.quiet = Duration::ZERO;
return self.demote(TierTransitionReason::QuietDemotion, Duration::ZERO);
}
} else {
self.quiet_ticks = 0;
self.quiet = Duration::ZERO;
}
None
@@ -195,7 +249,7 @@ impl SessionAdaptiveController {
fn promote(
&mut self,
reason: TierTransitionReason,
hard_cooldown_ticks: u32,
hard_cooldown: Duration,
) -> Option<TierTransition> {
let from = self.tier;
let to = from.promote();
@@ -204,22 +258,27 @@ impl SessionAdaptiveController {
}
self.tier = to;
self.max_tier_seen = max(self.max_tier_seen, to);
self.hard_cooldown_ticks = hard_cooldown_ticks;
self.tier1_hold_ticks = 0;
self.tier2_hold_ticks = 0;
self.quiet_ticks = 0;
self.hard_cooldown = hard_cooldown;
self.tier1_hold = Duration::ZERO;
self.tier2_hold = Duration::ZERO;
self.quiet = Duration::ZERO;
Some(TierTransition { from, to, reason })
}
fn demote(&mut self, reason: TierTransitionReason) -> Option<TierTransition> {
fn demote(
&mut self,
reason: TierTransitionReason,
hard_cooldown: Duration,
) -> Option<TierTransition> {
let from = self.tier;
let to = from.demote();
if from == to {
return None;
}
self.tier = to;
self.tier1_hold_ticks = 0;
self.tier2_hold_ticks = 0;
self.hard_cooldown = hard_cooldown;
self.tier1_hold = Duration::ZERO;
self.tier2_hold = Duration::ZERO;
Some(TierTransition { from, to, reason })
}
}
@@ -235,6 +294,8 @@ fn profiles() -> &'static DashMap<String, UserAdaptiveProfile> {
USER_PROFILES.get_or_init(DashMap::new)
}
/// Returns a fresh user's recent successful Direct tier, or `Base` when stale.
#[allow(dead_code)]
pub fn seed_tier_for_user(user: &str) -> AdaptiveTier {
if user.len() > MAX_USER_KEY_BYTES {
return AdaptiveTier::Base;
@@ -253,6 +314,8 @@ pub fn seed_tier_for_user(user: &str) -> AdaptiveTier {
AdaptiveTier::Base
}
/// Records the highest successfully allocated tier for bounded session seeding.
#[allow(dead_code)]
pub fn record_user_tier(user: &str, tier: AdaptiveTier) {
if user.len() > MAX_USER_KEY_BYTES {
return;
@@ -282,6 +345,8 @@ pub fn record_user_tier(user: &str, tier: AdaptiveTier) {
}
}
#[cfg(test)]
/// Returns the legacy staged scaling policy retained by security fixtures.
pub fn direct_copy_buffers_for_tier(
tier: AdaptiveTier,
base_c2s: usize,
@@ -294,6 +359,32 @@ pub fn direct_copy_buffers_for_tier(
)
}
/// Maps an adaptive tier to independent capacities within configured ceilings.
pub(crate) fn direct_copy_buffers_for_tier_with_ceilings(
tier: AdaptiveTier,
base_c2s: usize,
base_s2c: usize,
ceiling_c2s: usize,
ceiling_s2c: usize,
) -> (usize, usize) {
(
direct_direction_size(tier, base_c2s, ceiling_c2s),
direct_direction_size(tier, base_s2c, ceiling_s2c),
)
}
fn direct_direction_size(tier: AdaptiveTier, base: usize, ceiling: usize) -> usize {
let target = match tier {
AdaptiveTier::Base => base,
AdaptiveTier::Tier1 => ceiling / 4,
AdaptiveTier::Tier2 => ceiling / 2,
AdaptiveTier::Tier3 => ceiling,
};
target.max(base).min(ceiling.max(base)).max(1)
}
#[cfg(test)]
/// Returns the staged Middle-End flush policy retained by security fixtures.
pub fn me_flush_policy_for_tier(
tier: AdaptiveTier,
base_frames: usize,
@@ -323,6 +414,7 @@ fn ema(prev: f64, value: f64) -> f64 {
}
}
#[cfg(test)]
fn scale(base: usize, numerator: usize, denominator: usize, cap: usize) -> usize {
let scaled = base
.saturating_mul(numerator)
@@ -338,6 +430,10 @@ mod adaptive_buffers_security_tests;
#[path = "tests/adaptive_buffers_record_race_security_tests.rs"]
mod adaptive_buffers_record_race_security_tests;
#[cfg(test)]
#[path = "tests/adaptive_direct_budget_policy_tests.rs"]
mod adaptive_direct_budget_policy_tests;
#[cfg(test)]
mod tests {
use super::*;
@@ -396,7 +492,7 @@ mod tests {
fn test_quiet_demotion_is_slow_and_stepwise() {
let mut ctrl = SessionAdaptiveController::new(AdaptiveTier::Tier2);
let mut demotion = None;
for _ in 0..QUIET_DEMOTE_TICKS {
for _ in 0..480 {
demotion = ctrl.observe(sample(1, 1, 1, 1, 0, 0), 0.25);
}
+570
View File
@@ -0,0 +1,570 @@
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
use tokio::sync::watch;
use crate::stats::Stats;
use crate::stream::BufferPool;
use super::shared_state::ProxySharedState;
/// Accounting granularity for process-wide Direct copy-buffer reservations.
pub(crate) const DIRECT_BUFFER_UNIT_BYTES: usize = 4 * 1024;
/// Minimum client-to-DC copy-buffer capacity for one Direct session.
pub(crate) const DIRECT_BASE_C2S_BYTES: usize = 4 * 1024;
/// Minimum DC-to-client copy-buffer capacity for one Direct session.
pub(crate) const DIRECT_BASE_S2C_BYTES: usize = 8 * 1024;
const AUTO_HARD_MIN_BYTES: usize = 64 * 1024 * 1024;
const AUTO_HARD_MAX_BYTES: usize = 2 * 1024 * 1024 * 1024;
const AUTO_HARD_FALLBACK_BYTES: usize = 512 * 1024 * 1024;
const TARGET_FLOOR_MIN_BYTES: usize = 16 * 1024 * 1024;
const CONTROL_INTERVAL: Duration = Duration::from_secs(1);
const HEALTHY_RECOVERY_SAMPLES: u8 = 30;
const BUFFER_POOL_TRIM_LOW_WATERMARK: usize = 64;
const BUFFER_POOL_TRIM_HIGH_WATERMARK: usize = 128;
#[derive(Debug, Clone, Copy, Default)]
/// Lock-free observability snapshot of the Direct copy-buffer envelope.
pub(crate) struct DirectBufferBudgetSnapshot {
/// Absolute process-wide copy-buffer ceiling.
pub(crate) hard_limit_bytes: u64,
/// Current pressure-adjusted promotion target.
pub(crate) target_bytes: u64,
/// Bytes currently covered by active session leases.
pub(crate) reserved_bytes: u64,
/// Effective host or cgroup memory limit.
pub(crate) memory_total_bytes: u64,
/// Effective host or cgroup memory headroom.
pub(crate) memory_available_bytes: u64,
/// Current process resident set size.
pub(crate) process_rss_bytes: u64,
/// Successful tier growth reservations.
pub(crate) promotion_total: u64,
/// Tier growth attempts rejected by the adaptive target.
pub(crate) promotion_denied_total: u64,
/// Sessions admitted at minimum size above the adaptive target.
pub(crate) minimum_fallback_total: u64,
/// Sessions rejected by the absolute ceiling.
pub(crate) admission_rejected_total: u64,
/// Quiet-period tier reductions.
pub(crate) quiet_demotion_total: u64,
/// Sustained write-pressure tier reductions.
pub(crate) write_pressure_demotion_total: u64,
/// Process-wide pressure tier reductions.
pub(crate) global_pressure_demotion_total: u64,
/// Current sessions for Base through Tier3.
pub(crate) tier_sessions: [u64; 4],
}
#[derive(Debug, Clone, Copy, Default)]
struct SystemMemorySample {
total_bytes: u64,
available_bytes: u64,
process_rss_bytes: u64,
}
/// Process-wide hard envelope and adaptive target for Direct copy buffers.
pub(crate) struct DirectBufferBudget {
hard_limit_bytes: u64,
target_bytes: AtomicU64,
reserved_bytes: AtomicU64,
pressure_generation: AtomicU64,
pressure_tx: watch::Sender<u64>,
memory_total_bytes: AtomicU64,
memory_available_bytes: AtomicU64,
process_rss_bytes: AtomicU64,
promotion_total: AtomicU64,
promotion_denied_total: AtomicU64,
minimum_fallback_total: AtomicU64,
admission_rejected_total: AtomicU64,
quiet_demotion_total: AtomicU64,
write_pressure_demotion_total: AtomicU64,
global_pressure_demotion_total: AtomicU64,
tier_sessions: [AtomicU64; 4],
}
impl DirectBufferBudget {
/// Creates an envelope with a fixed absolute ceiling.
pub(crate) fn new(hard_limit_bytes: usize) -> Arc<Self> {
let hard_limit_bytes = align_down(hard_limit_bytes.max(DIRECT_BUFFER_UNIT_BYTES)) as u64;
let (pressure_tx, _) = watch::channel(0);
Arc::new(Self {
hard_limit_bytes,
target_bytes: AtomicU64::new(hard_limit_bytes),
reserved_bytes: AtomicU64::new(0),
pressure_generation: AtomicU64::new(0),
pressure_tx,
memory_total_bytes: AtomicU64::new(0),
memory_available_bytes: AtomicU64::new(0),
process_rss_bytes: AtomicU64::new(0),
promotion_total: AtomicU64::new(0),
promotion_denied_total: AtomicU64::new(0),
minimum_fallback_total: AtomicU64::new(0),
admission_rejected_total: AtomicU64::new(0),
quiet_demotion_total: AtomicU64::new(0),
write_pressure_demotion_total: AtomicU64::new(0),
global_pressure_demotion_total: AtomicU64::new(0),
tier_sessions: std::array::from_fn(|_| AtomicU64::new(0)),
})
}
/// Returns the current pressure-adjusted reservation target.
pub(crate) fn target_bytes(&self) -> usize {
self.target_bytes.load(Ordering::Relaxed) as usize
}
/// Subscribes to target reductions that require prompt session demotion.
pub(crate) fn subscribe_pressure(&self) -> watch::Receiver<u64> {
self.pressure_tx.subscribe()
}
/// Reserves bytes against either the adaptive target or the absolute ceiling.
pub(crate) fn try_reserve(
self: &Arc<Self>,
bytes: usize,
allow_above_target: bool,
) -> Option<DirectBufferLease> {
let bytes = align_up(bytes) as u64;
let limit = if allow_above_target {
self.hard_limit_bytes
} else {
self.target_bytes
.load(Ordering::Relaxed)
.min(self.hard_limit_bytes)
};
if !self.try_add_reserved(bytes, limit) {
return None;
}
self.tier_sessions[0].fetch_add(1, Ordering::Relaxed);
Some(DirectBufferLease {
budget: Arc::clone(self),
reserved_bytes: bytes,
tier: 0,
})
}
fn try_add_reserved(&self, bytes: u64, limit: u64) -> bool {
let mut current = self.reserved_bytes.load(Ordering::Acquire);
loop {
if bytes > limit.saturating_sub(current) {
return false;
}
match self.reserved_bytes.compare_exchange_weak(
current,
current + bytes,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => return true,
Err(observed) => current = observed,
}
}
}
fn target_floor_bytes(&self) -> u64 {
(self.hard_limit_bytes / 8)
.max(TARGET_FLOOR_MIN_BYTES as u64)
.min(self.hard_limit_bytes)
}
fn set_target_bytes(&self, target: u64) {
let target =
align_down(target.clamp(self.target_floor_bytes(), self.hard_limit_bytes) as usize)
as u64;
let previous = self.target_bytes.swap(target, Ordering::AcqRel);
if target < previous {
let generation = self
.pressure_generation
.fetch_add(1, Ordering::AcqRel)
.wrapping_add(1);
self.pressure_tx.send_replace(generation);
}
}
fn update_system_sample(&self, sample: SystemMemorySample) {
self.memory_total_bytes
.store(sample.total_bytes, Ordering::Relaxed);
self.memory_available_bytes
.store(sample.available_bytes, Ordering::Relaxed);
self.process_rss_bytes
.store(sample.process_rss_bytes, Ordering::Relaxed);
}
/// Records a session that had to bypass the adaptive target at minimum size.
pub(crate) fn increment_minimum_fallback(&self) {
self.minimum_fallback_total.fetch_add(1, Ordering::Relaxed);
}
/// Records a session rejected because the absolute ceiling was exhausted.
pub(crate) fn increment_admission_rejected(&self) {
self.admission_rejected_total
.fetch_add(1, Ordering::Relaxed);
}
/// Records a tier reduction after sustained low throughput.
pub(crate) fn increment_quiet_demotion(&self) {
self.quiet_demotion_total.fetch_add(1, Ordering::Relaxed);
}
/// Records a tier reduction after sustained partial or pending writes.
pub(crate) fn increment_write_pressure_demotion(&self) {
self.write_pressure_demotion_total
.fetch_add(1, Ordering::Relaxed);
}
/// Records a tier reduction requested by the process-wide controller.
pub(crate) fn increment_global_pressure_demotion(&self) {
self.global_pressure_demotion_total
.fetch_add(1, Ordering::Relaxed);
}
/// Captures all bounded metrics without allocating or locking.
pub(crate) fn snapshot(&self) -> DirectBufferBudgetSnapshot {
DirectBufferBudgetSnapshot {
hard_limit_bytes: self.hard_limit_bytes,
target_bytes: self.target_bytes.load(Ordering::Relaxed),
reserved_bytes: self.reserved_bytes.load(Ordering::Relaxed),
memory_total_bytes: self.memory_total_bytes.load(Ordering::Relaxed),
memory_available_bytes: self.memory_available_bytes.load(Ordering::Relaxed),
process_rss_bytes: self.process_rss_bytes.load(Ordering::Relaxed),
promotion_total: self.promotion_total.load(Ordering::Relaxed),
promotion_denied_total: self.promotion_denied_total.load(Ordering::Relaxed),
minimum_fallback_total: self.minimum_fallback_total.load(Ordering::Relaxed),
admission_rejected_total: self.admission_rejected_total.load(Ordering::Relaxed),
quiet_demotion_total: self.quiet_demotion_total.load(Ordering::Relaxed),
write_pressure_demotion_total: self
.write_pressure_demotion_total
.load(Ordering::Relaxed),
global_pressure_demotion_total: self
.global_pressure_demotion_total
.load(Ordering::Relaxed),
tier_sessions: std::array::from_fn(|index| {
self.tier_sessions[index].load(Ordering::Relaxed)
}),
}
}
}
/// Returns the conservative ceiling used when memory discovery is unavailable.
pub(crate) fn fallback_direct_buffer_hard_limit() -> usize {
AUTO_HARD_FALLBACK_BYTES
}
/// RAII ownership of all copy-buffer bytes retained by one Direct session.
pub(crate) struct DirectBufferLease {
budget: Arc<DirectBufferBudget>,
reserved_bytes: u64,
tier: usize,
}
impl DirectBufferLease {
/// Returns the currently covered allocation rounded to accounting units.
pub(crate) fn reserved_bytes(&self) -> usize {
self.reserved_bytes as usize
}
/// Attempts to cover a larger tier before its buffers are resized.
pub(crate) fn try_grow_to(&mut self, bytes: usize) -> bool {
let bytes = align_up(bytes) as u64;
if bytes <= self.reserved_bytes {
return true;
}
let delta = bytes - self.reserved_bytes;
let limit = self
.budget
.target_bytes
.load(Ordering::Relaxed)
.min(self.budget.hard_limit_bytes);
if !self.budget.try_add_reserved(delta, limit) {
self.budget
.promotion_denied_total
.fetch_add(1, Ordering::Relaxed);
return false;
}
self.reserved_bytes = bytes;
self.budget.promotion_total.fetch_add(1, Ordering::Relaxed);
true
}
/// Releases bytes only after both directional buffers report smaller coverage.
pub(crate) fn shrink_to(&mut self, bytes: usize) {
let bytes = align_up(bytes) as u64;
if bytes >= self.reserved_bytes {
return;
}
let released = self.reserved_bytes - bytes;
self.reserved_bytes = bytes;
self.budget
.reserved_bytes
.fetch_sub(released, Ordering::AcqRel);
}
/// Updates bounded per-tier session gauges for an accepted transition.
pub(crate) fn set_tier(&mut self, tier: usize) {
let tier = tier.min(self.budget.tier_sessions.len() - 1);
if tier == self.tier {
return;
}
decrement_saturating(&self.budget.tier_sessions[self.tier]);
self.budget.tier_sessions[tier].fetch_add(1, Ordering::Relaxed);
self.tier = tier;
}
}
impl Drop for DirectBufferLease {
fn drop(&mut self) {
self.budget
.reserved_bytes
.fetch_sub(self.reserved_bytes, Ordering::AcqRel);
decrement_saturating(&self.budget.tier_sessions[self.tier]);
}
}
/// Resolves the startup hard ceiling from config, cgroup, and host memory.
pub(crate) async fn resolve_direct_buffer_hard_limit(configured: usize) -> usize {
if configured != 0 {
return align_down(configured);
}
let sample = read_system_memory_sample().await;
if sample.total_bytes == 0 {
return AUTO_HARD_FALLBACK_BYTES;
}
let derived = (sample.total_bytes / 4)
.clamp(AUTO_HARD_MIN_BYTES as u64, AUTO_HARD_MAX_BYTES as u64)
.min(sample.total_bytes);
align_down(derived as usize).max(DIRECT_BUFFER_UNIT_BYTES)
}
/// Starts the single control-plane task for Direct budget and shared pool pressure.
pub(crate) fn spawn_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;
loop {
interval.tick().await;
let sample = read_system_memory_sample().await;
budget.update_system_sample(sample);
let snapshot = budget.snapshot();
let denied_delta = snapshot
.promotion_denied_total
.saturating_sub(previous_denied);
previous_denied = snapshot.promotion_denied_total;
let fallback_delta = snapshot
.minimum_fallback_total
.saturating_sub(previous_fallback);
previous_fallback = snapshot.minimum_fallback_total;
let rejected_delta = snapshot
.admission_rejected_total
.saturating_sub(previous_rejected);
previous_rejected = snapshot.admission_rejected_total;
let connection_pct = connection_fill_pct(stats.as_ref(), max_connections);
let memory_available_pct = percentage(sample.available_bytes, sample.total_bytes);
let target_utilization_pct = percentage(snapshot.reserved_bytes, snapshot.target_bytes);
let pressure = shared.conntrack_pressure_active()
|| connection_pct.is_some_and(|value| value >= 85)
|| memory_available_pct.is_some_and(|value| value <= 15)
|| target_utilization_pct.is_some_and(|value| value >= 90)
|| denied_delta > 0
|| fallback_delta > 0
|| rejected_delta > 0;
if !pressure {
pool_trim_armed = true;
} else if pool_trim_armed && buffer_pool.pooled() > pool_trim_high {
buffer_pool.trim_to(pool_trim_low);
pool_trim_armed = false;
}
let pool_snapshot = buffer_pool.stats();
stats.set_buffer_pool_gauges(
pool_snapshot.pooled,
pool_snapshot.allocated,
pool_snapshot.allocated.saturating_sub(pool_snapshot.pooled),
);
stats.set_buffer_pool_replaced_nonstandard_total(pool_snapshot.replaced_nonstandard);
let headroom_target = if sample.total_bytes == 0 {
snapshot.hard_limit_bytes
} else {
snapshot
.reserved_bytes
.saturating_add(sample.available_bytes / 4)
.min(snapshot.hard_limit_bytes)
};
if pressure {
healthy_streak = 0;
let reduced = snapshot.target_bytes.saturating_mul(3) / 4;
budget.set_target_bytes(reduced.min(headroom_target));
continue;
}
let healthy = memory_available_pct.is_none_or(|value| value >= 30)
&& connection_pct.is_none_or(|value| value <= 70);
if !healthy {
healthy_streak = 0;
if headroom_target < snapshot.target_bytes {
budget.set_target_bytes(headroom_target);
}
continue;
}
healthy_streak = healthy_streak.saturating_add(1);
if healthy_streak >= HEALTHY_RECOVERY_SAMPLES {
healthy_streak = 0;
let increment = (snapshot.target_bytes / 16).max(4 * 1024 * 1024);
budget.set_target_bytes(
snapshot
.target_bytes
.saturating_add(increment)
.min(headroom_target),
);
}
}
});
}
fn connection_fill_pct(stats: &Stats, max_connections: u32) -> Option<u8> {
if max_connections == 0 {
return None;
}
Some(
((stats.get_current_connections_total().saturating_mul(100)) / u64::from(max_connections))
.min(100) as u8,
)
}
fn percentage(value: u64, total: u64) -> Option<u8> {
if total == 0 {
return None;
}
Some(((value.saturating_mul(100)) / total).min(100) as u8)
}
async fn read_system_memory_sample() -> SystemMemorySample {
#[cfg(target_os = "linux")]
{
let meminfo = tokio::fs::read_to_string("/proc/meminfo")
.await
.unwrap_or_default();
let status = tokio::fs::read_to_string("/proc/self/status")
.await
.unwrap_or_default();
let host_total = parse_kib_field(&meminfo, "MemTotal:");
let host_available = parse_kib_field(&meminfo, "MemAvailable:");
let process_rss = parse_kib_field(&status, "VmRSS:");
let cgroup_v2_max = read_cgroup_limit("/sys/fs/cgroup/memory.max").await;
let cgroup_v2_current = read_u64_file("/sys/fs/cgroup/memory.current").await;
let cgroup_v1_max = read_cgroup_limit("/sys/fs/cgroup/memory/memory.limit_in_bytes").await;
let cgroup_v1_current = read_u64_file("/sys/fs/cgroup/memory/memory.usage_in_bytes").await;
let cgroup_max = cgroup_v2_max.or(cgroup_v1_max);
let cgroup_current = cgroup_v2_current.or(cgroup_v1_current);
let total = match (host_total, cgroup_max) {
(0, Some(limit)) => limit,
(host, Some(limit)) => host.min(limit),
(host, None) => host,
};
let cgroup_available = cgroup_max
.zip(cgroup_current)
.map(|(limit, current)| limit.saturating_sub(current));
let available = match (host_available, cgroup_available) {
(0, Some(value)) => value,
(host, Some(value)) => host.min(value),
(host, None) => host,
};
return SystemMemorySample {
total_bytes: total,
available_bytes: available,
process_rss_bytes: process_rss,
};
}
#[cfg(not(target_os = "linux"))]
{
SystemMemorySample::default()
}
}
#[cfg(target_os = "linux")]
async fn read_cgroup_limit(path: &str) -> Option<u64> {
let raw = tokio::fs::read_to_string(path).await.ok()?;
let raw = raw.trim();
if raw == "max" {
return None;
}
let value = raw.parse::<u64>().ok()?;
(value < (1u64 << 60)).then_some(value)
}
#[cfg(target_os = "linux")]
async fn read_u64_file(path: &str) -> Option<u64> {
tokio::fs::read_to_string(path)
.await
.ok()?
.trim()
.parse()
.ok()
}
#[cfg(target_os = "linux")]
fn parse_kib_field(raw: &str, key: &str) -> u64 {
raw.lines()
.find_map(|line| {
let value = line.strip_prefix(key)?.split_whitespace().next()?;
value.parse::<u64>().ok()
})
.unwrap_or(0)
.saturating_mul(1024)
}
fn align_up(bytes: usize) -> usize {
bytes
.div_ceil(DIRECT_BUFFER_UNIT_BYTES)
.saturating_mul(DIRECT_BUFFER_UNIT_BYTES)
}
fn align_down(bytes: usize) -> usize {
bytes / DIRECT_BUFFER_UNIT_BYTES * DIRECT_BUFFER_UNIT_BYTES
}
fn decrement_saturating(value: &AtomicU64) {
let mut current = value.load(Ordering::Relaxed);
while current != 0 {
match value.compare_exchange_weak(
current,
current - 1,
Ordering::Relaxed,
Ordering::Relaxed,
) {
Ok(_) => return,
Err(observed) => current = observed,
}
}
}
#[cfg(test)]
#[path = "tests/direct_buffer_budget_tests.rs"]
mod tests;
+16 -17
View File
@@ -345,22 +345,22 @@ where
} else {
Duration::from_secs(1800)
};
let relay_result =
crate::proxy::relay::relay_bidirectional_with_activity_timeout_lease_and_cancel(
client_reader,
client_writer,
tg_reader,
tg_writer,
config.general.direct_relay_copy_buf_c2s_bytes,
config.general.direct_relay_copy_buf_s2c_bytes,
user,
Arc::clone(&stats),
config.access.user_data_quota.get(user).copied(),
buffer_pool,
traffic_lease,
relay_activity_timeout,
session_cancel.clone(),
);
let relay_result = crate::proxy::relay::relay_direct_adaptive(
client_reader,
client_writer,
tg_reader,
tg_writer,
config.general.direct_relay_copy_buf_c2s_bytes,
config.general.direct_relay_copy_buf_s2c_bytes,
config.server.max_connections,
user,
Arc::clone(&stats),
config.access.user_data_quota.get(user).copied(),
traffic_lease,
relay_activity_timeout,
session_cancel.clone(),
Arc::clone(&shared.direct_buffer_budget),
);
tokio::pin!(relay_result);
let relay_result = loop {
if let Some(cutover) =
@@ -400,7 +400,6 @@ where
Err(e) => debug!(user = %user, error = %e, "Direct relay ended with error"),
}
buffer_pool_trim.trim_to(buffer_pool_trim.max_buffers().min(64));
let pool_snapshot = buffer_pool_trim.stats();
stats.set_buffer_pool_gauges(
pool_snapshot.pooled,
+7 -1
View File
@@ -69,7 +69,9 @@ use self::quota::{
#[cfg(test)]
use self::c2me::enqueue_c2me_command;
#[cfg(test)]
use self::d2c::{compute_intermediate_secure_wire_len, process_me_writer_response};
use self::d2c::{
compute_intermediate_secure_wire_len, process_me_writer_response, write_client_payload,
};
#[cfg(test)]
pub(crate) use self::desync::{
clear_desync_dedup_for_testing_in_shared, desync_dedup_get_for_testing,
@@ -166,3 +168,7 @@ mod middle_relay_atomic_quota_invariant_tests;
#[cfg(test)]
#[path = "tests/middle_relay_baseline_invariant_tests.rs"]
mod middle_relay_baseline_invariant_tests;
#[cfg(test)]
#[path = "tests/middle_relay_d2c_flush_padding_security_tests.rs"]
mod middle_relay_d2c_flush_padding_security_tests;
+1 -1
View File
@@ -149,6 +149,7 @@ where
peer,
translated_local_addr,
payload,
_permit,
flags,
effective_tag_array,
)
@@ -822,7 +823,6 @@ where
clear_relay_idle_candidate_in(shared.as_ref(), conn_id);
me_pool.registry().unregister(conn_id).await;
buffer_pool.trim_to(buffer_pool.max_buffers().min(64));
let pool_snapshot = buffer_pool.stats();
stats.set_buffer_pool_gauges(
pool_snapshot.pooled,
+2
View File
@@ -60,6 +60,8 @@
pub mod adaptive_buffers;
pub mod client;
// Process-wide Direct relay copy-buffer ownership and pressure policy.
pub(crate) mod direct_buffer_budget;
pub mod direct_relay;
pub mod handshake;
pub mod masking;
+4
View File
@@ -84,8 +84,11 @@ fn watchdog_delta(current: u64, previous: u64) -> u64 {
current.saturating_sub(previous)
}
mod adaptive_copy;
mod io;
pub(crate) use self::adaptive_copy::relay_direct_adaptive;
use self::io::{CombinedStream, SharedCounters, StatsIo, is_quota_io_error};
#[cfg(test)]
use self::io::{quota_adaptive_interval_bytes, should_immediate_quota_check};
@@ -217,6 +220,7 @@ where
.await
}
#[allow(dead_code)]
pub async fn relay_bidirectional_with_activity_timeout_lease_and_cancel<CR, CW, SR, SW>(
client_reader: CR,
client_writer: CW,
+527
View File
@@ -0,0 +1,527 @@
use std::io;
use std::pin::Pin;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::task::{Context, Poll};
use std::time::Duration;
use tokio::io::{AsyncBufRead, AsyncRead, AsyncWrite, AsyncWriteExt, ReadBuf, copy_buf};
use tokio::time::Instant;
use tokio_util::sync::CancellationToken;
use tracing::{debug, warn};
use crate::error::{ProxyError, Result};
use crate::proxy::adaptive_buffers::{
AdaptiveTier, RelaySignalSample, SessionAdaptiveController, TierTransitionReason,
direct_copy_buffers_for_tier_with_ceilings,
};
use crate::proxy::direct_buffer_budget::{
DIRECT_BASE_C2S_BYTES, DIRECT_BASE_S2C_BYTES, DirectBufferBudget, DirectBufferLease,
};
use crate::proxy::traffic_limiter::TrafficLease;
use crate::stats::Stats;
use super::WATCHDOG_INTERVAL;
use super::io::{SharedCounters, StatsIo, is_quota_io_error};
use super::watchdog_delta;
mod write_pressure;
use self::write_pressure::WritePressureIo;
struct AdaptiveBufferState {
desired_bytes: AtomicUsize,
actual_bytes: AtomicUsize,
}
impl AdaptiveBufferState {
fn new(bytes: usize) -> Arc<Self> {
Arc::new(Self {
desired_bytes: AtomicUsize::new(bytes.max(1)),
actual_bytes: AtomicUsize::new(bytes.max(1)),
})
}
}
struct AdaptiveBufReader<R> {
inner: R,
buffer: Box<[u8]>,
pos: usize,
cap: usize,
state: Arc<AdaptiveBufferState>,
}
impl<R> AdaptiveBufReader<R> {
fn new(inner: R, state: Arc<AdaptiveBufferState>) -> Self {
let bytes = state.actual_bytes.load(Ordering::Relaxed).max(1);
Self {
inner,
buffer: vec![0; bytes].into_boxed_slice(),
pos: 0,
cap: 0,
state,
}
}
fn resize_if_drained(&mut self) {
if self.pos != self.cap {
return;
}
let desired = self.state.desired_bytes.load(Ordering::Acquire).max(1);
if desired == self.buffer.len() {
return;
}
self.buffer = vec![0; desired].into_boxed_slice();
self.pos = 0;
self.cap = 0;
self.state.actual_bytes.store(desired, Ordering::Release);
}
}
impl<R: AsyncRead + Unpin> AsyncRead for AdaptiveBufReader<R> {
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
output: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> {
let this = self.get_mut();
if this.pos < this.cap {
let available = &this.buffer[this.pos..this.cap];
let copied = available.len().min(output.remaining());
output.put_slice(&available[..copied]);
this.pos += copied;
return Poll::Ready(Ok(()));
}
this.resize_if_drained();
Pin::new(&mut this.inner).poll_read(cx, output)
}
}
impl<R: AsyncRead + Unpin> AsyncBufRead for AdaptiveBufReader<R> {
fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
let this = self.get_mut();
if this.pos < this.cap {
return Poll::Ready(Ok(&this.buffer[this.pos..this.cap]));
}
this.resize_if_drained();
let mut read_buf = ReadBuf::new(&mut this.buffer);
match Pin::new(&mut this.inner).poll_read(cx, &mut read_buf) {
Poll::Ready(Ok(())) => {
this.pos = 0;
this.cap = read_buf.filled().len();
Poll::Ready(Ok(&this.buffer[..this.cap]))
}
Poll::Ready(Err(error)) => Poll::Ready(Err(error)),
Poll::Pending => Poll::Pending,
}
}
fn consume(self: Pin<&mut Self>, amount: usize) {
let this = self.get_mut();
this.pos = this.pos.saturating_add(amount).min(this.cap);
}
}
enum AdaptiveRelayOutcome {
Copy(io::Result<(u64, u64)>),
ActivityTimeout,
UserDisabled,
}
#[allow(clippy::too_many_arguments)]
/// Relays one Direct session with independently resizable directional buffers.
pub(crate) async fn relay_direct_adaptive<CR, CW, SR, SW>(
client_reader: CR,
client_writer: CW,
server_reader: SR,
server_writer: SW,
ceiling_c2s_bytes: usize,
ceiling_s2c_bytes: usize,
max_connections: u32,
user: &str,
stats: Arc<Stats>,
quota_limit: Option<u64>,
traffic_lease: Option<Arc<TrafficLease>>,
activity_timeout: Duration,
session_cancel: CancellationToken,
budget: Arc<DirectBufferBudget>,
) -> Result<()>
where
CR: AsyncRead + Unpin + Send + 'static,
CW: AsyncWrite + Unpin + Send + 'static,
SR: AsyncRead + Unpin + Send + 'static,
SW: AsyncWrite + Unpin + Send + 'static,
{
let activity_timeout = activity_timeout.max(Duration::from_secs(1));
let epoch = Instant::now();
let counters = Arc::new(SharedCounters::new());
let quota_exceeded = Arc::new(AtomicBool::new(false));
let user_owned = user.to_string();
let (base_c2s, base_s2c) = initial_base_sizes(
ceiling_c2s_bytes,
ceiling_s2c_bytes,
max_connections,
budget.target_bytes(),
);
let base_total = base_c2s.saturating_add(base_s2c);
let mut lease = match budget.try_reserve(base_total, false) {
Some(lease) => lease,
None => {
let minimum_total = DIRECT_BASE_C2S_BYTES + DIRECT_BASE_S2C_BYTES;
match budget.try_reserve(minimum_total, true) {
Some(lease) => {
budget.increment_minimum_fallback();
lease
}
None => {
budget.increment_admission_rejected();
return Err(ProxyError::Proxy(
"Direct relay buffer pressure: budget exhausted".to_string(),
));
}
}
}
};
let effective_base = if lease.reserved_bytes() < base_total {
(DIRECT_BASE_C2S_BYTES, DIRECT_BASE_S2C_BYTES)
} else {
(base_c2s, base_s2c)
};
let c2s_state = AdaptiveBufferState::new(effective_base.0);
let s2c_state = AdaptiveBufferState::new(effective_base.1);
let mut controller = SessionAdaptiveController::new(AdaptiveTier::Base);
let c2s_client = StatsIo::new_with_traffic_lease(
client_reader,
Arc::clone(&counters),
Arc::clone(&stats),
user_owned.clone(),
traffic_lease.clone(),
quota_limit,
Arc::clone(&quota_exceeded),
epoch,
);
let client_writer = StatsIo::new_with_traffic_lease(
client_writer,
Arc::clone(&counters),
Arc::clone(&stats),
user_owned.clone(),
traffic_lease,
quota_limit,
Arc::clone(&quota_exceeded),
epoch,
);
let mut client_writer = WritePressureIo::new(client_writer, Arc::clone(&counters));
let mut c2s_reader = AdaptiveBufReader::new(c2s_client, Arc::clone(&c2s_state));
let mut s2c_reader = AdaptiveBufReader::new(server_reader, Arc::clone(&s2c_state));
let mut server_writer = server_writer;
let mut pressure_rx = budget.subscribe_pressure();
let relay_outcome = {
let copy = async {
let c2s = async {
let copied = copy_buf(&mut c2s_reader, &mut server_writer).await?;
server_writer.shutdown().await?;
Ok::<u64, io::Error>(copied)
};
let s2c = async {
let copied = copy_buf(&mut s2c_reader, &mut client_writer).await?;
client_writer.shutdown().await?;
Ok::<u64, io::Error>(copied)
};
tokio::try_join!(c2s, s2c)
};
tokio::pin!(copy);
let mut interval = tokio::time::interval(WATCHDOG_INTERVAL);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
interval.tick().await;
let mut previous = RelaySignalSample::default();
let mut previous_log_c2s = 0u64;
let mut previous_log_s2c = 0u64;
let mut previous_sample_at = epoch;
loop {
tokio::select! {
result = &mut copy => break AdaptiveRelayOutcome::Copy(result),
_ = session_cancel.cancelled() => break AdaptiveRelayOutcome::UserDisabled,
changed = pressure_rx.changed() => {
if changed.is_ok() {
apply_global_pressure_demotion(
&mut controller,
&mut lease,
&c2s_state,
&s2c_state,
effective_base,
(ceiling_c2s_bytes, ceiling_s2c_bytes),
budget.as_ref(),
);
reconcile_reservation(&mut lease, &c2s_state, &s2c_state);
}
}
_ = interval.tick() => {
let now = Instant::now();
let idle = counters.idle_duration(now, epoch);
if quota_exceeded.load(Ordering::Acquire) {
warn!(user = %user_owned, "User data quota reached, closing relay");
break AdaptiveRelayOutcome::ActivityTimeout;
}
if idle >= activity_timeout {
warn!(
user = %user_owned,
c2s_bytes = counters.c2s_bytes.load(Ordering::Relaxed),
s2c_bytes = counters.s2c_bytes.load(Ordering::Relaxed),
idle_secs = idle.as_secs(),
"Activity timeout"
);
break AdaptiveRelayOutcome::ActivityTimeout;
}
let sample = current_sample(counters.as_ref());
let c2s_delta = watchdog_delta(sample.c2s_bytes, previous_log_c2s);
let s2c_delta = watchdog_delta(sample.s2c_written_bytes, previous_log_s2c);
if c2s_delta > 0 || s2c_delta > 0 {
let secs = now.saturating_duration_since(previous_sample_at).as_secs_f64();
debug!(
user = %user_owned,
c2s_kbps = (c2s_delta as f64 / secs / 1024.0) as u64,
s2c_kbps = (s2c_delta as f64 / secs / 1024.0) as u64,
c2s_total = sample.c2s_bytes,
s2c_total = sample.s2c_written_bytes,
"Relay active"
);
}
let delta = sample_delta(sample, previous);
let tick_secs = now.saturating_duration_since(previous_sample_at).as_secs_f64();
if let Some(transition) = controller.observe(delta, tick_secs) {
apply_controller_transition(
transition,
&mut controller,
&mut lease,
&c2s_state,
&s2c_state,
effective_base,
(ceiling_c2s_bytes, ceiling_s2c_bytes),
budget.as_ref(),
);
}
reconcile_reservation(&mut lease, &c2s_state, &s2c_state);
previous = sample;
previous_log_c2s = sample.c2s_bytes;
previous_log_s2c = sample.s2c_written_bytes;
previous_sample_at = now;
}
}
}
};
let _ = client_writer.shutdown().await;
let _ = server_writer.shutdown().await;
let c2s_ops = counters.c2s_ops.load(Ordering::Relaxed);
let s2c_ops = counters.s2c_ops.load(Ordering::Relaxed);
let duration = epoch.elapsed();
match relay_outcome {
AdaptiveRelayOutcome::Copy(Ok((c2s, s2c))) => {
debug!(
user = %user_owned,
c2s_bytes = c2s,
s2c_bytes = s2c,
c2s_msgs = c2s_ops,
s2c_msgs = s2c_ops,
duration_secs = duration.as_secs(),
"Relay finished"
);
Ok(())
}
AdaptiveRelayOutcome::Copy(Err(error)) if is_quota_io_error(&error) => {
warn!(
user = %user_owned,
c2s_bytes = counters.c2s_bytes.load(Ordering::Relaxed),
s2c_bytes = counters.s2c_bytes.load(Ordering::Relaxed),
c2s_msgs = c2s_ops,
s2c_msgs = s2c_ops,
duration_secs = duration.as_secs(),
"Data quota reached, closing relay"
);
Err(ProxyError::DataQuotaExceeded { user: user_owned })
}
AdaptiveRelayOutcome::Copy(Err(error)) => {
debug!(
user = %user_owned,
c2s_bytes = counters.c2s_bytes.load(Ordering::Relaxed),
s2c_bytes = counters.s2c_bytes.load(Ordering::Relaxed),
c2s_msgs = c2s_ops,
s2c_msgs = s2c_ops,
duration_secs = duration.as_secs(),
error = %error,
"Relay error"
);
Err(error.into())
}
AdaptiveRelayOutcome::ActivityTimeout => {
debug!(
user = %user_owned,
c2s_bytes = counters.c2s_bytes.load(Ordering::Relaxed),
s2c_bytes = counters.s2c_bytes.load(Ordering::Relaxed),
c2s_msgs = c2s_ops,
s2c_msgs = s2c_ops,
duration_secs = duration.as_secs(),
"Relay finished (activity timeout)"
);
Ok(())
}
AdaptiveRelayOutcome::UserDisabled => {
debug!(
user = %user_owned,
c2s_bytes = counters.c2s_bytes.load(Ordering::Relaxed),
s2c_bytes = counters.s2c_bytes.load(Ordering::Relaxed),
c2s_msgs = c2s_ops,
s2c_msgs = s2c_ops,
duration_secs = duration.as_secs(),
"Relay finished (user disabled)"
);
Err(ProxyError::UserDisabled { user: user_owned })
}
}
}
fn initial_base_sizes(
ceiling_c2s: usize,
ceiling_s2c: usize,
max_connections: u32,
target_bytes: usize,
) -> (usize, usize) {
let configured_total = ceiling_c2s.saturating_add(ceiling_s2c);
let configured_worst_case = configured_total.saturating_mul(max_connections as usize);
if max_connections != 0 && configured_worst_case <= target_bytes {
return (ceiling_c2s, ceiling_s2c);
}
(
DIRECT_BASE_C2S_BYTES.min(ceiling_c2s),
DIRECT_BASE_S2C_BYTES.min(ceiling_s2c),
)
}
fn current_sample(counters: &SharedCounters) -> RelaySignalSample {
RelaySignalSample {
c2s_bytes: counters.c2s_bytes.load(Ordering::Relaxed),
s2c_requested_bytes: counters.s2c_requested_bytes.load(Ordering::Relaxed),
s2c_written_bytes: counters.s2c_bytes.load(Ordering::Relaxed),
s2c_write_ops: counters.s2c_ops.load(Ordering::Relaxed),
s2c_partial_writes: counters.s2c_partial_writes.load(Ordering::Relaxed),
s2c_consecutive_pending_writes: counters
.s2c_consecutive_pending_writes
.load(Ordering::Relaxed),
}
}
fn sample_delta(current: RelaySignalSample, previous: RelaySignalSample) -> RelaySignalSample {
RelaySignalSample {
c2s_bytes: current.c2s_bytes.saturating_sub(previous.c2s_bytes),
s2c_requested_bytes: current
.s2c_requested_bytes
.saturating_sub(previous.s2c_requested_bytes),
s2c_written_bytes: current
.s2c_written_bytes
.saturating_sub(previous.s2c_written_bytes),
s2c_write_ops: current.s2c_write_ops.saturating_sub(previous.s2c_write_ops),
s2c_partial_writes: current
.s2c_partial_writes
.saturating_sub(previous.s2c_partial_writes),
s2c_consecutive_pending_writes: current.s2c_consecutive_pending_writes,
}
}
#[allow(clippy::too_many_arguments)]
fn apply_controller_transition(
transition: crate::proxy::adaptive_buffers::TierTransition,
controller: &mut SessionAdaptiveController,
lease: &mut DirectBufferLease,
c2s_state: &AdaptiveBufferState,
s2c_state: &AdaptiveBufferState,
base: (usize, usize),
ceilings: (usize, usize),
budget: &DirectBufferBudget,
) {
let sizes = direct_copy_buffers_for_tier_with_ceilings(
transition.to,
base.0,
base.1,
ceilings.0,
ceilings.1,
);
if transition.to > transition.from {
if !lease.try_grow_to(sizes.0.saturating_add(sizes.1)) {
*controller = SessionAdaptiveController::new(transition.from);
return;
}
} else {
match transition.reason {
TierTransitionReason::QuietDemotion => budget.increment_quiet_demotion(),
TierTransitionReason::SustainedWritePressure => {
budget.increment_write_pressure_demotion();
}
TierTransitionReason::SoftConfirmed | TierTransitionReason::HardPressure => {}
}
}
set_desired_sizes(c2s_state, s2c_state, sizes);
lease.set_tier(transition.to.as_u8() as usize);
}
fn apply_global_pressure_demotion(
controller: &mut SessionAdaptiveController,
lease: &mut DirectBufferLease,
c2s_state: &AdaptiveBufferState,
s2c_state: &AdaptiveBufferState,
base: (usize, usize),
ceilings: (usize, usize),
budget: &DirectBufferBudget,
) {
let current = controller.tier();
let target = current.demote();
if target == current {
return;
}
*controller = SessionAdaptiveController::new(target);
let sizes =
direct_copy_buffers_for_tier_with_ceilings(target, base.0, base.1, ceilings.0, ceilings.1);
set_desired_sizes(c2s_state, s2c_state, sizes);
lease.set_tier(target.as_u8() as usize);
budget.increment_global_pressure_demotion();
}
fn set_desired_sizes(
c2s_state: &AdaptiveBufferState,
s2c_state: &AdaptiveBufferState,
sizes: (usize, usize),
) {
c2s_state
.desired_bytes
.store(sizes.0.max(1), Ordering::Release);
s2c_state
.desired_bytes
.store(sizes.1.max(1), Ordering::Release);
}
fn reconcile_reservation(
lease: &mut DirectBufferLease,
c2s_state: &AdaptiveBufferState,
s2c_state: &AdaptiveBufferState,
) {
// Promotion reserves the desired allocation before either reader grows.
// Demotion keeps the actual allocation covered until its buffered bytes drain.
let covered_c2s = c2s_state
.actual_bytes
.load(Ordering::Acquire)
.max(c2s_state.desired_bytes.load(Ordering::Acquire));
let covered_s2c = s2c_state
.actual_bytes
.load(Ordering::Acquire)
.max(s2c_state.desired_bytes.load(Ordering::Acquire));
lease.shrink_to(covered_c2s.saturating_add(covered_s2c));
}
@@ -0,0 +1,72 @@
use std::io;
use std::pin::Pin;
use std::sync::Arc;
use std::sync::atomic::Ordering;
use std::task::{Context, Poll};
use tokio::io::AsyncWrite;
use super::super::io::SharedCounters;
/// Direct-only writer wrapper that exposes bounded backpressure signals.
pub(super) struct WritePressureIo<W> {
inner: W,
counters: Arc<SharedCounters>,
}
impl<W> WritePressureIo<W> {
/// Wraps the client writer without changing its I/O or error contract.
pub(super) fn new(inner: W, counters: Arc<SharedCounters>) -> Self {
Self { inner, counters }
}
}
impl<W: AsyncWrite + Unpin> AsyncWrite for WritePressureIo<W> {
fn poll_write(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buffer: &[u8],
) -> Poll<io::Result<usize>> {
let this = self.get_mut();
if !buffer.is_empty() {
this.counters
.s2c_requested_bytes
.fetch_add(buffer.len() as u64, Ordering::Relaxed);
}
match Pin::new(&mut this.inner).poll_write(cx, buffer) {
Poll::Ready(Ok(written)) => {
this.counters
.s2c_consecutive_pending_writes
.store(0, Ordering::Relaxed);
if written < buffer.len() {
this.counters
.s2c_partial_writes
.fetch_add(1, Ordering::Relaxed);
}
Poll::Ready(Ok(written))
}
Poll::Ready(Err(error)) => {
this.counters
.s2c_consecutive_pending_writes
.store(0, Ordering::Relaxed);
Poll::Ready(Err(error))
}
Poll::Pending => {
let _ = this.counters.s2c_consecutive_pending_writes.fetch_update(
Ordering::Relaxed,
Ordering::Relaxed,
|current| Some(current.saturating_add(1)),
);
Poll::Pending
}
}
}
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
Pin::new(&mut self.get_mut().inner).poll_flush(cx)
}
fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
Pin::new(&mut self.get_mut().inner).poll_shutdown(cx)
}
}
+10 -1
View File
@@ -1,4 +1,4 @@
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
use std::time::Duration;
use tokio::time::Instant;
@@ -20,6 +20,12 @@ pub(in crate::proxy::relay) struct SharedCounters {
pub(in crate::proxy::relay) c2s_ops: AtomicU64,
/// Number of poll_write completions (≈ S→C chunks)
pub(in crate::proxy::relay) s2c_ops: AtomicU64,
/// Bytes presented to client writes, including retried pending writes.
pub(in crate::proxy::relay) s2c_requested_bytes: AtomicU64,
/// Successful client writes that consumed only part of the offered slice.
pub(in crate::proxy::relay) s2c_partial_writes: AtomicU64,
/// Consecutive pending client writes observed by the active copy loop.
pub(in crate::proxy::relay) s2c_consecutive_pending_writes: AtomicU32,
/// Milliseconds since relay epoch of last I/O activity
last_activity_ms: AtomicU64,
}
@@ -31,6 +37,9 @@ impl SharedCounters {
s2c_bytes: AtomicU64::new(0),
c2s_ops: AtomicU64::new(0),
s2c_ops: AtomicU64::new(0),
s2c_requested_bytes: AtomicU64::new(0),
s2c_partial_writes: AtomicU64::new(0),
s2c_consecutive_pending_writes: AtomicU32::new(0),
last_activity_ms: AtomicU64::new(0),
}
}
+12
View File
@@ -9,6 +9,7 @@ use dashmap::DashMap;
use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc};
use tokio_util::sync::CancellationToken;
use crate::proxy::direct_buffer_budget::{DirectBufferBudget, fallback_direct_buffer_hard_limit};
use crate::proxy::handshake::{AuthProbeSaturationState, AuthProbeState};
use crate::proxy::middle_relay::{DesyncDedupRotationState, RelayIdleCandidateRegistry};
use crate::proxy::traffic_limiter::TrafficLimiter;
@@ -69,6 +70,7 @@ pub(crate) struct ProxySharedState {
pub(crate) handshake: HandshakeSharedState,
pub(crate) middle_relay: MiddleRelaySharedState,
pub(crate) traffic_limiter: Arc<TrafficLimiter>,
pub(crate) direct_buffer_budget: Arc<DirectBufferBudget>,
disabled_users: DashMap<String, ()>,
active_user_sessions: DashMap<(String, u64), CancellationToken>,
pub(crate) conntrack_pressure_active: AtomicBool,
@@ -101,6 +103,15 @@ impl Drop for UserSessionGuard {
impl ProxySharedState {
pub(crate) fn new() -> Arc<Self> {
Self::new_with_direct_buffer_budget(DirectBufferBudget::new(
fallback_direct_buffer_hard_limit(),
))
}
/// Creates process state with the startup-resolved Direct buffer envelope.
pub(crate) fn new_with_direct_buffer_budget(
direct_buffer_budget: Arc<DirectBufferBudget>,
) -> Arc<Self> {
Arc::new(Self {
handshake: HandshakeSharedState {
auth_probe: DashMap::new(),
@@ -129,6 +140,7 @@ impl ProxySharedState {
relay_idle_mark_seq: AtomicU64::new(0),
},
traffic_limiter: TrafficLimiter::new(),
direct_buffer_budget,
disabled_users: DashMap::new(),
active_user_sessions: DashMap::new(),
conntrack_pressure_active: AtomicBool::new(false),
@@ -0,0 +1,58 @@
use super::*;
#[test]
fn configured_direct_sizes_are_strict_tier_ceilings() {
let base = (4 * 1024, 8 * 1024);
let ceilings = (64 * 1024, 256 * 1024);
assert_eq!(
direct_copy_buffers_for_tier_with_ceilings(
AdaptiveTier::Base,
base.0,
base.1,
ceilings.0,
ceilings.1,
),
base
);
assert_eq!(
direct_copy_buffers_for_tier_with_ceilings(
AdaptiveTier::Tier3,
base.0,
base.1,
ceilings.0,
ceilings.1,
),
ceilings
);
}
#[test]
fn sustained_pending_pressure_demotes_after_transient_promotion() {
let mut controller = SessionAdaptiveController::new(AdaptiveTier::Tier1);
let pressure = RelaySignalSample {
c2s_bytes: 0,
s2c_requested_bytes: 1024,
s2c_written_bytes: 0,
s2c_write_ops: 0,
s2c_partial_writes: 0,
s2c_consecutive_pending_writes: 3,
};
let first = controller
.observe(pressure, 10.0)
.expect("transient pressure must retain the staged promotion");
assert_eq!(first.reason, TierTransitionReason::HardPressure);
let second = controller
.observe(pressure, 10.0)
.expect("bounded transient pressure may promote one additional tier");
assert_eq!(second.reason, TierTransitionReason::HardPressure);
let sustained = controller
.observe(pressure, 10.0)
.expect("sustained pressure must release one tier");
assert_eq!(
sustained.reason,
TierTransitionReason::SustainedWritePressure
);
assert_eq!(sustained.to, AdaptiveTier::Tier2);
}
@@ -0,0 +1,38 @@
use super::*;
#[test]
fn lease_drop_releases_the_complete_reservation() {
let budget = DirectBufferBudget::new(16 * 1024);
{
let lease = budget
.try_reserve(12 * 1024, false)
.expect("minimum reservation must fit");
assert_eq!(lease.reserved_bytes(), 12 * 1024);
assert_eq!(budget.snapshot().reserved_bytes, 12 * 1024);
}
assert_eq!(budget.snapshot().reserved_bytes, 0);
}
#[test]
fn absolute_ceiling_rejects_excess_minimum_reservations() {
let budget = DirectBufferBudget::new(16 * 1024);
let _lease = budget
.try_reserve(12 * 1024, true)
.expect("first minimum reservation must fit");
assert!(budget.try_reserve(8 * 1024, true).is_none());
assert_eq!(budget.snapshot().reserved_bytes, 12 * 1024);
}
#[test]
fn growth_and_shrink_keep_accounting_balanced() {
let budget = DirectBufferBudget::new(32 * 1024);
let mut lease = budget
.try_reserve(12 * 1024, false)
.expect("base reservation must fit");
assert!(lease.try_grow_to(24 * 1024));
assert_eq!(budget.snapshot().reserved_bytes, 24 * 1024);
lease.shrink_to(16 * 1024);
assert_eq!(budget.snapshot().reserved_bytes, 16 * 1024);
drop(lease);
assert_eq!(budget.snapshot().reserved_bytes, 0);
}
@@ -0,0 +1,150 @@
use std::io;
use std::pin::Pin;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};
use tokio::io::AsyncWrite;
use super::*;
use crate::crypto::AesCtr;
use crate::protocol::framing::INTERMEDIATE_WIRE_LEN_MASK;
#[derive(Clone, Default)]
struct RecordingWriter {
writes: Arc<Mutex<Vec<u8>>>,
flushes: Arc<AtomicUsize>,
}
impl RecordingWriter {
fn captured(&self) -> Vec<u8> {
self.writes
.lock()
.expect("test writer capture lock must not be poisoned")
.clone()
}
}
impl AsyncWrite for RecordingWriter {
fn poll_write(
self: Pin<&mut Self>,
_cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
self.writes
.lock()
.expect("test writer capture lock must not be poisoned")
.extend_from_slice(buf);
Poll::Ready(Ok(buf.len()))
}
fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
self.flushes.fetch_add(1, Ordering::Relaxed);
Poll::Ready(Ok(()))
}
fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
Poll::Ready(Ok(()))
}
}
fn crypto_writer(inner: RecordingWriter) -> CryptoWriter<RecordingWriter> {
let key = [0u8; 32];
CryptoWriter::new(inner, AesCtr::new(&key, 0), 8 * 1024 * 1024)
}
fn decrypt_capture(mut encrypted: Vec<u8>) -> Vec<u8> {
let key = [0u8; 32];
let mut cipher = AesCtr::new(&key, 0);
cipher.apply(&mut encrypted);
encrypted
}
fn secure_wire_len(cleartext: &[u8]) -> usize {
let header = cleartext
.get(..4)
.expect("secure frame must include an intermediate header");
(u32::from_le_bytes(
header
.try_into()
.expect("secure frame header must be four bytes"),
) & INTERMEDIATE_WIRE_LEN_MASK) as usize
}
async fn write_secure_payload(payload_len: usize) -> (MeD2cWriteMode, Vec<u8>) {
let inner = RecordingWriter::default();
let capture = inner.clone();
let mut writer = crypto_writer(inner);
let payload = vec![0xa5; payload_len];
let mut frame_buf = Vec::new();
let cancel = CancellationToken::new();
let rng = SecureRandom::new();
let mode = write_client_payload(
&mut writer,
ProtoTag::Secure,
0,
&payload,
&rng,
&mut frame_buf,
&cancel,
)
.await
.expect("secure payload write must succeed");
flush_client_or_cancel(&mut writer, &cancel)
.await
.expect("secure payload flush must succeed");
(mode, decrypt_capture(capture.captured()))
}
fn assert_secure_payload_with_tail_padding(cleartext: &[u8], payload_len: usize) {
let wire_len = secure_wire_len(cleartext);
assert_eq!(cleartext.len(), 4 + wire_len);
assert!(
cleartext[4..4 + payload_len]
.iter()
.all(|byte| *byte == 0xa5)
);
let padding_len = wire_len
.checked_sub(payload_len)
.expect("secure wire length must include payload bytes");
assert!((1..=3).contains(&padding_len));
assert_ne!(wire_len % 4, 0);
}
#[tokio::test]
async fn queue_drain_flush_reason_performs_physical_client_flush() {
let inner = RecordingWriter::default();
let flushes = inner.flushes.clone();
let mut writer = crypto_writer(inner);
let cancel = CancellationToken::new();
assert!(me_d2c_flush_reason_requires_client_flush(
MeD2cFlushReason::QueueDrain
));
flush_client_or_cancel(&mut writer, &cancel)
.await
.expect("client flush must succeed");
assert_eq!(flushes.load(Ordering::Relaxed), 1);
}
#[tokio::test]
async fn secure_payload_coalesced_path_keeps_tail_padding() {
let payload_len = 8;
let (mode, cleartext) = write_secure_payload(payload_len).await;
assert!(matches!(mode, MeD2cWriteMode::Coalesced));
assert_secure_payload_with_tail_padding(&cleartext, payload_len);
}
#[tokio::test]
async fn secure_payload_split_path_keeps_tail_padding() {
let payload_len = ME_D2C_SINGLE_WRITE_COALESCE_MAX_BYTES;
let (mode, cleartext) = write_secure_payload(payload_len).await;
assert!(matches!(mode, MeD2cWriteMode::Split));
assert_secure_payload_with_tail_padding(&cleartext, payload_len);
}
+164 -19
View File
@@ -10,7 +10,7 @@ use arc_swap::ArcSwap;
use dashmap::DashMap;
use ipnetwork::IpNetwork;
use crate::config::RateLimitBps;
use crate::config::{CidrRateLimitKey, RateLimitBps};
const REGISTRY_SHARDS: usize = 64;
const FAIR_EPOCH_MS: u64 = 20;
@@ -413,16 +413,29 @@ struct CidrRule {
prefix_len: u8,
}
#[derive(Clone, Copy)]
struct CidrAutoRule {
prefix_len: u8,
limits: RateLimitBps,
}
enum CidrPolicyMatch<'a> {
Explicit(&'a CidrRule),
Auto { key: String, limits: RateLimitBps },
}
#[derive(Default)]
struct PolicySnapshot {
user_limits: HashMap<String, RateLimitBps>,
cidr_rules_v4: Vec<CidrRule>,
cidr_rules_v6: Vec<CidrRule>,
cidr_auto_rules_v4: Vec<CidrAutoRule>,
cidr_auto_rules_v6: Vec<CidrAutoRule>,
cidr_rule_keys: HashSet<String>,
}
impl PolicySnapshot {
fn match_cidr(&self, ip: IpAddr) -> Option<&CidrRule> {
fn match_cidr(&self, ip: IpAddr) -> Option<CidrPolicyMatch<'_>> {
match ip {
IpAddr::V4(_) => self
.cidr_rules_v4
@@ -433,6 +446,20 @@ impl PolicySnapshot {
.iter()
.find(|rule| rule.cidr.contains(ip)),
}
.map(CidrPolicyMatch::Explicit)
.or_else(|| self.match_auto_cidr(ip))
}
fn match_auto_cidr(&self, ip: IpAddr) -> Option<CidrPolicyMatch<'_>> {
let rule = match ip {
IpAddr::V4(_) => self.cidr_auto_rules_v4.first()?,
IpAddr::V6(_) => self.cidr_auto_rules_v6.first()?,
};
let key = auto_cidr_bucket_key(ip, rule.prefix_len)?;
Some(CidrPolicyMatch::Auto {
key,
limits: rule.limits,
})
}
}
@@ -633,7 +660,7 @@ impl TrafficLimiter {
pub fn apply_policy(
&self,
user_limits: HashMap<String, RateLimitBps>,
cidr_limits: HashMap<IpNetwork, RateLimitBps>,
cidr_limits: HashMap<CidrRateLimitKey, RateLimitBps>,
) {
let filtered_users = user_limits
.into_iter()
@@ -642,39 +669,64 @@ impl TrafficLimiter {
let mut cidr_rules_v4 = Vec::new();
let mut cidr_rules_v6 = Vec::new();
let mut cidr_auto_rules_v4 = Vec::new();
let mut cidr_auto_rules_v6 = Vec::new();
let mut cidr_rule_keys = HashSet::new();
for (cidr, limits) in cidr_limits {
for (key, limits) in cidr_limits {
if limits.up_bps == 0 && limits.down_bps == 0 {
continue;
}
let key = cidr.to_string();
let rule = CidrRule {
key: key.clone(),
cidr,
limits,
prefix_len: cidr.prefix(),
};
cidr_rule_keys.insert(key);
match rule.cidr {
IpNetwork::V4(_) => cidr_rules_v4.push(rule),
IpNetwork::V6(_) => cidr_rules_v6.push(rule),
match key {
CidrRateLimitKey::Network(cidr) => {
let key = cidr.to_string();
let rule = CidrRule {
key: key.clone(),
cidr,
limits,
prefix_len: cidr.prefix(),
};
cidr_rule_keys.insert(key);
match rule.cidr {
IpNetwork::V4(_) => cidr_rules_v4.push(rule),
IpNetwork::V6(_) => cidr_rules_v6.push(rule),
}
}
CidrRateLimitKey::AutoV4(prefix_len) => {
cidr_auto_rules_v4.push(CidrAutoRule { prefix_len, limits });
}
CidrRateLimitKey::AutoV6(prefix_len) => {
cidr_auto_rules_v6.push(CidrAutoRule { prefix_len, limits });
}
CidrRateLimitKey::AutoDual(prefix_len) => {
cidr_auto_rules_v4.push(CidrAutoRule { prefix_len, limits });
cidr_auto_rules_v6.push(CidrAutoRule {
prefix_len: prefix_len.saturating_mul(4),
limits,
});
}
}
}
cidr_rules_v4.sort_by(|a, b| b.prefix_len.cmp(&a.prefix_len));
cidr_rules_v6.sort_by(|a, b| b.prefix_len.cmp(&a.prefix_len));
cidr_auto_rules_v4.sort_by(|a, b| b.prefix_len.cmp(&a.prefix_len));
cidr_auto_rules_v6.sort_by(|a, b| b.prefix_len.cmp(&a.prefix_len));
let cidr_policy_entries =
cidr_rule_keys.len() + cidr_auto_rules_v4.len() + cidr_auto_rules_v6.len();
self.user_scope
.policy_entries
.store(filtered_users.len() as u64, Ordering::Relaxed);
self.cidr_scope
.policy_entries
.store(cidr_rule_keys.len() as u64, Ordering::Relaxed);
.store(cidr_policy_entries as u64, Ordering::Relaxed);
self.policy.store(Arc::new(PolicySnapshot {
user_limits: filtered_users,
cidr_rules_v4,
cidr_rules_v6,
cidr_auto_rules_v4,
cidr_auto_rules_v6,
cidr_rule_keys,
}));
@@ -703,11 +755,15 @@ impl TrafficLimiter {
let mut cidr_bucket = None;
let mut cidr_user_key = None;
let mut cidr_user_share = None;
if let Some(rule) = policy.match_cidr(client_ip) {
if let Some(rule_match) = policy.match_cidr(client_ip) {
let (key, limits) = match &rule_match {
CidrPolicyMatch::Explicit(rule) => (rule.key.as_str(), rule.limits),
CidrPolicyMatch::Auto { key, limits } => (key.as_str(), *limits),
};
let bucket = self
.cidr_buckets
.get_or_insert_with(rule.key.as_str(), || CidrBucket::new(rule.limits));
bucket.set_rates(rule.limits);
.get_or_insert_with(key, || CidrBucket::new(limits));
bucket.set_rates(limits);
bucket.active_leases.fetch_add(1, Ordering::Relaxed);
self.cidr_scope
.active_leases
@@ -841,6 +897,16 @@ fn bytes_per_epoch(bps: u64) -> u64 {
bytes.max(1)
}
fn auto_cidr_bucket_key(ip: IpAddr, prefix_len: u8) -> Option<String> {
let cidr = IpNetwork::new(ip, prefix_len).ok()?;
let network = IpNetwork::new(cidr.network(), prefix_len).ok()?;
let family = match network {
IpNetwork::V4(_) => "4",
IpNetwork::V6(_) => "6",
};
Some(format!("auto:{family}:{network}"))
}
fn current_epoch() -> u64 {
let start = limiter_epoch_start();
let elapsed_ms = start.elapsed().as_millis() as u64;
@@ -851,3 +917,82 @@ fn limiter_epoch_start() -> &'static Instant {
static START: OnceLock<Instant> = OnceLock::new();
START.get_or_init(Instant::now)
}
#[cfg(test)]
mod tests {
use super::*;
fn rate(up_bps: u64, down_bps: u64) -> RateLimitBps {
RateLimitBps { up_bps, down_bps }
}
#[test]
fn explicit_cidr_rule_wins_over_auto_template() {
let limiter = TrafficLimiter::new();
let mut cidr_limits = HashMap::new();
cidr_limits.insert(CidrRateLimitKey::AutoV4(24), rate(1_000, 0));
cidr_limits.insert(
CidrRateLimitKey::Network("203.0.113.7/32".parse().unwrap()),
rate(2_000, 0),
);
limiter.apply_policy(HashMap::new(), cidr_limits);
let policy = limiter.policy.load_full();
let matched = policy.match_cidr("203.0.113.7".parse().unwrap()).unwrap();
match matched {
CidrPolicyMatch::Explicit(rule) => assert_eq!(rule.key.as_str(), "203.0.113.7/32"),
CidrPolicyMatch::Auto { .. } => panic!("explicit CIDR must have priority"),
}
}
#[test]
fn auto_template_uses_longest_prefix() {
let limiter = TrafficLimiter::new();
let mut cidr_limits = HashMap::new();
cidr_limits.insert(CidrRateLimitKey::AutoV4(24), rate(1_000, 0));
cidr_limits.insert(CidrRateLimitKey::AutoV4(32), rate(2_000, 0));
limiter.apply_policy(HashMap::new(), cidr_limits);
let policy = limiter.policy.load_full();
let matched = policy.match_cidr("203.0.113.129".parse().unwrap()).unwrap();
match matched {
CidrPolicyMatch::Auto { key, limits } => {
assert_eq!(key, "auto:4:203.0.113.129/32");
assert_eq!(limits.up_bps, 2_000);
}
CidrPolicyMatch::Explicit(_) => panic!("auto-template match expected"),
}
}
#[test]
fn dual_auto_template_maps_v6_prefix_by_four() {
let limiter = TrafficLimiter::new();
let mut cidr_limits = HashMap::new();
cidr_limits.insert(CidrRateLimitKey::AutoDual(32), rate(1_000, 0));
limiter.apply_policy(HashMap::new(), cidr_limits);
let policy = limiter.policy.load_full();
let matched = policy.match_cidr("2001:db8::1".parse().unwrap()).unwrap();
match matched {
CidrPolicyMatch::Auto { key, .. } => {
assert_eq!(key, "auto:6:2001:db8::1/128");
}
CidrPolicyMatch::Explicit(_) => panic!("auto-template match expected"),
}
}
#[test]
fn auto_cidr_bucket_key_canonicalizes_network_address() {
assert_eq!(
auto_cidr_bucket_key("203.0.113.129".parse().unwrap(), 24).unwrap(),
"auto:4:203.0.113.0/24"
);
assert_eq!(
auto_cidr_bucket_key("2001:db8::abcd".parse().unwrap(), 64).unwrap(),
"auto:6:2001:db8::/64"
);
}
}
+36
View File
@@ -204,6 +204,12 @@ impl Stats {
self.buffer_pool_in_use_gauge.load(Ordering::Relaxed)
}
/// Returns the count of non-standard buffers replaced before pooling.
pub fn get_buffer_pool_replaced_nonstandard_total(&self) -> u64 {
self.buffer_pool_replaced_nonstandard_total
.load(Ordering::Relaxed)
}
pub fn get_me_c2me_send_full_total(&self) -> u64 {
self.me_c2me_send_full_total.load(Ordering::Relaxed)
}
@@ -269,6 +275,36 @@ impl Stats {
self.me_writer_pick_mode_switch_total
.load(Ordering::Relaxed)
}
/// Returns the configured resident-memory limit per ME writer.
pub fn get_me_writer_byte_budget_limit_bytes_gauge(&self) -> u64 {
self.me_writer_byte_budget_limit_bytes_gauge
.load(Ordering::Relaxed)
}
/// Returns aggregate queued or enqueueing writer memory reservations.
pub fn get_me_writer_byte_budget_queued_bytes_gauge(&self) -> u64 {
self.me_writer_byte_budget_queued_bytes_gauge
.load(Ordering::Relaxed)
}
/// Returns aggregate writer reservations currently owned by socket writes.
pub fn get_me_writer_byte_budget_inflight_bytes_gauge(&self) -> u64 {
self.me_writer_byte_budget_inflight_bytes_gauge
.load(Ordering::Relaxed)
}
/// Returns the count of blocking writer byte-budget waits.
pub fn get_me_writer_byte_budget_wait_total(&self) -> u64 {
self.me_writer_byte_budget_wait_total
.load(Ordering::Relaxed)
}
/// Returns the count of writer byte-budget wait timeouts.
pub fn get_me_writer_byte_budget_timeout_total(&self) -> u64 {
self.me_writer_byte_budget_timeout_total
.load(Ordering::Relaxed)
}
/// Returns the count of payloads that cannot fit the configured writer budget.
pub fn get_me_writer_byte_budget_oversize_total(&self) -> u64 {
self.me_writer_byte_budget_oversize_total
.load(Ordering::Relaxed)
}
pub fn get_me_socks_kdf_strict_reject(&self) -> u64 {
self.me_socks_kdf_strict_reject.load(Ordering::Relaxed)
}
+7
View File
@@ -274,6 +274,7 @@ pub struct Stats {
buffer_pool_pooled_gauge: AtomicU64,
buffer_pool_allocated_gauge: AtomicU64,
buffer_pool_in_use_gauge: AtomicU64,
buffer_pool_replaced_nonstandard_total: AtomicU64,
// C2ME enqueue observability
me_c2me_send_full_total: AtomicU64,
me_c2me_send_high_water_total: AtomicU64,
@@ -292,6 +293,12 @@ pub struct Stats {
me_writer_pick_p2c_no_candidate_total: AtomicU64,
me_writer_pick_blocking_fallback_total: AtomicU64,
me_writer_pick_mode_switch_total: AtomicU64,
me_writer_byte_budget_limit_bytes_gauge: AtomicU64,
me_writer_byte_budget_queued_bytes_gauge: AtomicU64,
me_writer_byte_budget_inflight_bytes_gauge: AtomicU64,
me_writer_byte_budget_wait_total: AtomicU64,
me_writer_byte_budget_timeout_total: AtomicU64,
me_writer_byte_budget_oversize_total: AtomicU64,
me_socks_kdf_strict_reject: AtomicU64,
me_socks_kdf_compat_fallback: AtomicU64,
secure_padding_invalid: AtomicU64,
+58
View File
@@ -88,6 +88,56 @@ impl Stats {
.fetch_add(1, Ordering::Relaxed);
}
}
/// Publishes the configured resident-memory limit applied to each ME writer.
pub fn set_me_writer_byte_budget_limit_bytes(&self, bytes: usize) {
self.me_writer_byte_budget_limit_bytes_gauge
.store(bytes as u64, Ordering::Relaxed);
}
pub(crate) fn add_me_writer_byte_budget_queued_bytes(&self, bytes: u64) {
self.me_writer_byte_budget_queued_bytes_gauge
.fetch_add(bytes, Ordering::Relaxed);
}
pub(crate) fn move_me_writer_byte_budget_to_inflight(&self, bytes: u64) {
let _ = self.me_writer_byte_budget_queued_bytes_gauge.fetch_update(
Ordering::Relaxed,
Ordering::Relaxed,
|current| Some(current.saturating_sub(bytes)),
);
self.me_writer_byte_budget_inflight_bytes_gauge
.fetch_add(bytes, Ordering::Relaxed);
}
pub(crate) fn release_me_writer_byte_budget_queued_bytes(&self, bytes: u64) {
let _ = self.me_writer_byte_budget_queued_bytes_gauge.fetch_update(
Ordering::Relaxed,
Ordering::Relaxed,
|current| Some(current.saturating_sub(bytes)),
);
}
pub(crate) fn release_me_writer_byte_budget_inflight_bytes(&self, bytes: u64) {
let _ = self
.me_writer_byte_budget_inflight_bytes_gauge
.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
Some(current.saturating_sub(bytes))
});
}
pub(crate) fn increment_me_writer_byte_budget_wait_total(&self) {
if self.telemetry_me_allows_normal() {
self.me_writer_byte_budget_wait_total
.fetch_add(1, Ordering::Relaxed);
}
}
pub(crate) fn increment_me_writer_byte_budget_timeout_total(&self) {
if self.telemetry_me_allows_normal() {
self.me_writer_byte_budget_timeout_total
.fetch_add(1, Ordering::Relaxed);
}
}
pub(crate) fn increment_me_writer_byte_budget_oversize_total(&self) {
if self.telemetry_me_allows_normal() {
self.me_writer_byte_budget_oversize_total
.fetch_add(1, Ordering::Relaxed);
}
}
pub fn increment_me_socks_kdf_strict_reject(&self) {
if self.telemetry_me_allows_normal() {
self.me_socks_kdf_strict_reject
@@ -502,6 +552,14 @@ impl Stats {
}
}
/// Publishes the cumulative count of non-standard pool buffer replacements.
pub fn set_buffer_pool_replaced_nonstandard_total(&self, value: usize) {
if self.telemetry_me_allows_normal() {
self.buffer_pool_replaced_nonstandard_total
.store(value as u64, Ordering::Relaxed);
}
}
pub fn increment_me_c2me_send_full_total(&self) {
if self.telemetry_me_allows_normal() {
self.me_c2me_send_full_total.fetch_add(1, Ordering::Relaxed);
+20 -18
View File
@@ -68,8 +68,8 @@ use crate::crypto::AesCtr;
/// Actual limit is supplied at runtime from configuration.
const DEFAULT_MAX_PENDING_WRITE: usize = 64 * 1024;
/// Default read buffer capacity (reader mostly decrypts in-place into caller buffer).
const DEFAULT_READ_CAPACITY: usize = 16 * 1024;
/// Maximum scratch capacity retained after a completed write.
const MAX_RETAINED_SCRATCH_CAPACITY: usize = 32 * 1024;
// ============= CryptoReader State =============
@@ -110,10 +110,6 @@ pub struct CryptoReader<R> {
upstream: R,
decryptor: AesCtr,
state: CryptoReaderState,
/// Reserved for future coalescing optimizations.
#[allow(dead_code)]
read_buf: BytesMut,
}
impl<R> CryptoReader<R> {
@@ -122,7 +118,6 @@ impl<R> CryptoReader<R> {
upstream,
decryptor,
state: CryptoReaderState::Idle,
read_buf: BytesMut::with_capacity(DEFAULT_READ_CAPACITY),
}
}
@@ -321,7 +316,7 @@ struct PendingCiphertext {
impl PendingCiphertext {
fn new(max_len: usize) -> Self {
Self {
buf: BytesMut::with_capacity(16 * 1024),
buf: BytesMut::new(),
pos: 0,
max_len,
}
@@ -372,15 +367,12 @@ impl PendingCiphertext {
}
}
/// Replace the entire pending ciphertext by moving `src` in (swap, no copy).
fn replace_with(&mut self, mut src: BytesMut) {
/// Replace the entire pending ciphertext by moving `src` in without copying.
fn replace_with(&mut self, src: BytesMut) {
debug_assert!(src.len() <= self.max_len);
self.buf.clear();
self.buf = src;
self.pos = 0;
// Swap: keep allocations hot and avoid copying bytes.
std::mem::swap(&mut self.buf, &mut src);
}
/// Append plaintext and encrypt appended range in-place.
@@ -465,7 +457,7 @@ impl<W> CryptoWriter<W> {
upstream,
encryptor,
state: CryptoWriterState::Idle,
scratch: BytesMut::with_capacity(16 * 1024),
scratch: BytesMut::new(),
max_pending_write: max_pending.max(4 * 1024),
}
}
@@ -502,6 +494,7 @@ impl<W> CryptoWriter<W> {
}
fn poison(&mut self, error: io::Error) {
self.scratch = BytesMut::new();
self.state = CryptoWriterState::Poisoned { error: Some(error) };
}
@@ -552,6 +545,15 @@ impl<W> CryptoWriter<W> {
scratch.extend_from_slice(plaintext);
encryptor.apply(&mut scratch[..]);
}
/// Clear reusable scratch while releasing allocations inflated by large writes.
fn recycle_scratch(scratch: &mut BytesMut) {
if scratch.capacity() > MAX_RETAINED_SCRATCH_CAPACITY {
*scratch = BytesMut::new();
} else {
scratch.clear();
}
}
}
impl<W: AsyncWrite + Unpin> CryptoWriter<W> {
@@ -698,13 +700,13 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for CryptoWriter<W> {
Poll::Ready(Ok(n)) => {
if n == this.scratch.len() {
this.scratch.clear();
Self::recycle_scratch(&mut this.scratch);
return Poll::Ready(Ok(to_accept));
}
// Partial upstream write of ciphertext
let remainder = this.scratch.split_off(n);
this.scratch.clear();
let mut remainder = std::mem::take(&mut this.scratch);
let _ = remainder.split_to(n);
let pending = Self::ensure_pending(&mut this.state, this.max_pending_write);
pending.replace_with(remainder);
+6 -11
View File
@@ -312,7 +312,7 @@ fn encode_secure(frame: &Frame, dst: &mut BytesMut, rng: &SecureRandom) -> io::R
));
}
// Telegram Desktop VersionD uses a 4-bit random padding length.
// Outbound Secure padding avoids full-word tails that readers cannot strip.
let padding_len = secure_padding_len(data.len(), rng);
let total_len = data.len() + padding_len;
@@ -521,13 +521,7 @@ mod tests {
use tokio_util::codec::{FramedRead, FramedWrite};
fn assert_secure_decoded_payload(decoded: &[u8], original: &[u8]) {
assert!(decoded.starts_with(original));
assert!(
(original.len()..=original.len() + 12).contains(&decoded.len()),
"Secure decoded payload may retain up to 12 bytes of full-word padding, got {}",
decoded.len()
);
assert_eq!(decoded.len() % 4, 0);
assert_eq!(decoded, original);
}
#[tokio::test]
@@ -653,7 +647,7 @@ mod tests {
}
#[test]
fn secure_codec_uses_tdesktop_padding_range_and_jitters_wire_length() {
fn secure_codec_uses_non_aligned_padding_and_jitters_wire_length() {
let codec = SecureCodec::new(Arc::new(SecureRandom::new()));
let payload = Bytes::from_static(&[1, 2, 3, 4, 5, 6, 7, 8]);
let mut wire_lens = HashSet::new();
@@ -666,9 +660,10 @@ mod tests {
let wire_len = u32::from_le_bytes([out[0], out[1], out[2], out[3]]) as usize;
assert_eq!(out.len(), 4 + wire_len);
assert!(
(payload.len()..=payload.len() + 15).contains(&wire_len),
"Secure wire length must be payload+0..15, got {wire_len}"
(payload.len() + 1..=payload.len() + 3).contains(&wire_len),
"Secure wire length must be payload+1..3, got {wire_len}"
);
assert_ne!(wire_len % 4, 0);
wire_lens.insert(wire_len);
}
+2 -8
View File
@@ -367,7 +367,7 @@ impl<W: AsyncWrite + Unpin> SecureIntermediateFrameWriter<W> {
));
}
// Telegram Desktop VersionD uses a 4-bit random padding length.
// Outbound Secure padding avoids full-word tails that readers cannot strip.
let padding_len = secure_padding_len(data.len(), &self.rng);
let padding = self.rng.bytes(padding_len);
@@ -633,13 +633,7 @@ mod tests {
use tokio::time::{Duration, timeout};
fn assert_secure_decoded_payload(decoded: &[u8], original: &[u8]) {
assert!(decoded.starts_with(original));
assert!(
(original.len()..=original.len() + 12).contains(&decoded.len()),
"Secure decoded payload may retain up to 12 bytes of full-word padding, got {}",
decoded.len()
);
assert_eq!(decoded.len() % 4, 0);
assert_eq!(decoded, original);
}
#[tokio::test]
+66 -34
View File
@@ -40,7 +40,7 @@ use std::pin::Pin;
use std::task::{Context, Poll};
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, ReadBuf};
use super::state::{HeaderBuffer, StreamState, WriteBuffer, YieldBuffer};
use super::state::{HeaderBuffer, StreamState, YieldBuffer};
use crate::protocol::constants::{
MAX_TLS_CIPHERTEXT_SIZE, MAX_TLS_PLAINTEXT_SIZE, TLS_RECORD_ALERT, TLS_RECORD_APPLICATION,
TLS_RECORD_CHANGE_CIPHER, TLS_RECORD_HANDSHAKE, TLS_VERSION,
@@ -59,6 +59,9 @@ const MAX_TLS_PAYLOAD: usize = MAX_TLS_CIPHERTEXT_SIZE;
/// Note: we never queue unlimited amount of data here; state holds at most one record.
const MAX_PENDING_WRITE: usize = 64 * 1024;
/// Maximum record buffer capacity retained between writes.
const MAX_RETAINED_RECORD_CAPACITY: usize = 2 * (TLS_HEADER_SIZE + MAX_TLS_PAYLOAD);
// ============= TLS Record Types =============
/// Parsed TLS record header (5 bytes)
@@ -639,7 +642,7 @@ enum TlsWriterState {
/// Writing a complete TLS record (header + body), possibly partially
WritingRecord {
record: WriteBuffer,
position: usize,
payload_size: usize,
},
@@ -676,6 +679,7 @@ impl StreamState for TlsWriterState {
pub struct FakeTlsWriter<W> {
upstream: W,
state: TlsWriterState,
record_buffer: BytesMut,
}
impl<W> FakeTlsWriter<W> {
@@ -683,6 +687,7 @@ impl<W> FakeTlsWriter<W> {
Self {
upstream,
state: TlsWriterState::Idle,
record_buffer: BytesMut::new(),
}
}
@@ -704,10 +709,15 @@ impl<W> FakeTlsWriter<W> {
}
pub fn has_pending(&self) -> bool {
matches!(&self.state, TlsWriterState::WritingRecord { record, .. } if !record.is_empty())
matches!(
&self.state,
TlsWriterState::WritingRecord { position, .. }
if *position < self.record_buffer.len()
)
}
fn poison(&mut self, error: io::Error) {
self.record_buffer = BytesMut::new();
self.state = TlsWriterState::Poisoned { error: Some(error) };
}
@@ -720,17 +730,26 @@ impl<W> FakeTlsWriter<W> {
}
}
fn build_record(data: &[u8]) -> BytesMut {
fn build_record(&mut self, data: &[u8]) {
let header = TlsRecordHeader {
record_type: TLS_RECORD_APPLICATION,
version: TLS_VERSION,
length: data.len() as u16,
};
let mut record = BytesMut::with_capacity(TLS_HEADER_SIZE + data.len());
record.extend_from_slice(&header.to_bytes());
record.extend_from_slice(data);
record
self.record_buffer.clear();
self.record_buffer.reserve(TLS_HEADER_SIZE + data.len());
self.record_buffer.extend_from_slice(&header.to_bytes());
self.record_buffer.extend_from_slice(data);
debug_assert!(self.record_buffer.len() <= MAX_PENDING_WRITE);
}
fn recycle_record_buffer(&mut self) {
if self.record_buffer.capacity() > MAX_RETAINED_RECORD_CAPACITY {
self.record_buffer = BytesMut::new();
} else {
self.record_buffer.clear();
}
}
}
@@ -744,10 +763,11 @@ impl<W: AsyncWrite + Unpin> FakeTlsWriter<W> {
fn poll_flush_record_inner(
upstream: &mut W,
cx: &mut Context<'_>,
record: &mut WriteBuffer,
record: &[u8],
position: &mut usize,
) -> FlushResult {
while !record.is_empty() {
let data = record.pending();
while *position < record.len() {
let data = &record[*position..];
match Pin::new(&mut *upstream).poll_write(cx, data) {
Poll::Pending => return FlushResult::Pending,
Poll::Ready(Err(e)) => return FlushResult::Error(e),
@@ -757,7 +777,7 @@ impl<W: AsyncWrite + Unpin> FakeTlsWriter<W> {
"upstream returned 0 bytes written",
));
}
Poll::Ready(Ok(n)) => record.advance(n),
Poll::Ready(Ok(n)) => *position += n,
}
}
@@ -780,14 +800,19 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for FakeTlsWriter<W> {
}
TlsWriterState::WritingRecord {
mut record,
mut position,
payload_size,
} => {
// Finish writing previous record before accepting new bytes.
match Self::poll_flush_record_inner(&mut this.upstream, cx, &mut record) {
match Self::poll_flush_record_inner(
&mut this.upstream,
cx,
&this.record_buffer,
&mut position,
) {
FlushResult::Pending => {
this.state = TlsWriterState::WritingRecord {
record,
position,
payload_size,
};
return Poll::Pending;
@@ -797,6 +822,7 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for FakeTlsWriter<W> {
return Poll::Ready(Err(e));
}
FlushResult::Complete(_) => {
this.recycle_record_buffer();
this.state = TlsWriterState::Idle;
// continue to accept new buf below
}
@@ -818,19 +844,17 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for FakeTlsWriter<W> {
let chunk = &buf[..chunk_size];
// Build the complete record (header + payload)
let record_data = Self::build_record(chunk);
this.build_record(chunk);
match Pin::new(&mut this.upstream).poll_write(cx, &record_data) {
Poll::Ready(Ok(n)) if n == record_data.len() => Poll::Ready(Ok(chunk_size)),
match Pin::new(&mut this.upstream).poll_write(cx, &this.record_buffer) {
Poll::Ready(Ok(n)) if n == this.record_buffer.len() => {
this.recycle_record_buffer();
Poll::Ready(Ok(chunk_size))
}
Poll::Ready(Ok(n)) => {
// Partial write of the record: store remainder.
let mut write_buffer = WriteBuffer::with_max_size(MAX_PENDING_WRITE);
// record_data length is <= 16389, fits MAX_PENDING_WRITE
let _ = write_buffer.extend(&record_data[n..]);
this.state = TlsWriterState::WritingRecord {
record: write_buffer,
position: n,
payload_size: chunk_size,
};
@@ -844,12 +868,8 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for FakeTlsWriter<W> {
}
Poll::Pending => {
// Buffer entire record and report success for this chunk.
let mut write_buffer = WriteBuffer::with_max_size(MAX_PENDING_WRITE);
let _ = write_buffer.extend(&record_data);
this.state = TlsWriterState::WritingRecord {
record: write_buffer,
position: 0,
payload_size: chunk_size,
};
@@ -871,12 +891,17 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for FakeTlsWriter<W> {
}
TlsWriterState::WritingRecord {
mut record,
mut position,
payload_size,
} => match Self::poll_flush_record_inner(&mut this.upstream, cx, &mut record) {
} => match Self::poll_flush_record_inner(
&mut this.upstream,
cx,
&this.record_buffer,
&mut position,
) {
FlushResult::Pending => {
this.state = TlsWriterState::WritingRecord {
record,
position,
payload_size,
};
return Poll::Pending;
@@ -886,6 +911,7 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for FakeTlsWriter<W> {
return Poll::Ready(Err(e));
}
FlushResult::Complete(_) => {
this.recycle_record_buffer();
this.state = TlsWriterState::Idle;
}
},
@@ -905,11 +931,17 @@ impl<W: AsyncWrite + Unpin> AsyncWrite for FakeTlsWriter<W> {
match state {
TlsWriterState::WritingRecord {
mut record,
mut position,
payload_size: _,
} => {
// Best-effort flush (do not block shutdown forever).
let _ = Self::poll_flush_record_inner(&mut this.upstream, cx, &mut record);
let _ = Self::poll_flush_record_inner(
&mut this.upstream,
cx,
&this.record_buffer,
&mut position,
);
this.recycle_record_buffer();
this.state = TlsWriterState::Idle;
}
_ => {
+135 -57
View File
@@ -1,10 +1,8 @@
use std::net::IpAddr;
use super::command::run_command;
use super::model::{SynLimitRule, SynLimitTargets, synlimit_rate_arg};
use super::model::{SynLimitNamespace, SynLimitRule, SynLimitTargets, synlimit_rate_arg};
const IPTABLES_CHAIN: &str = "TELEMT_SYNLIMIT";
const IPTABLES_HASHLIMIT_PREFIX: &str = "TMT-SYN";
const IPV4_IOS_PACKET_LENGTH: u16 = 64;
const IPV6_IOS_PACKET_LENGTH: u16 = 84;
const IOS_TTL_LIMIT: u8 = 65;
@@ -38,49 +36,57 @@ impl IpTablesFamily {
}
}
pub(super) async fn apply_synlimit_rules(targets: &SynLimitTargets) -> Result<(), String> {
apply_rules_for_binary("iptables", &targets.iptables_v4, IpTablesFamily::V4).await?;
apply_rules_for_binary("ip6tables", &targets.iptables_v6, IpTablesFamily::V6).await
pub(super) async fn apply_synlimit_rules(
targets: &SynLimitTargets,
namespace: &SynLimitNamespace,
) -> Result<(), String> {
apply_rules_for_binary(
"iptables",
&targets.iptables_v4,
IpTablesFamily::V4,
namespace,
)
.await?;
apply_rules_for_binary(
"ip6tables",
&targets.iptables_v6,
IpTablesFamily::V6,
namespace,
)
.await
}
async fn apply_rules_for_binary(
binary: &str,
targets: &[SynLimitRule],
family: IpTablesFamily,
namespace: &SynLimitNamespace,
) -> Result<(), String> {
if targets.is_empty() {
return Ok(());
}
let _ = run_command(binary, &["-t", "filter", "-N", IPTABLES_CHAIN], None).await;
run_command(binary, &["-t", "filter", "-F", IPTABLES_CHAIN], None).await?;
if run_command(
binary,
&["-t", "filter", "-C", "INPUT", "-j", IPTABLES_CHAIN],
None,
)
.await
.is_err()
let chain = namespace.iptables_chain.as_str();
let _ = run_command(binary, &["-t", "filter", "-N", chain], None).await;
run_command(binary, &["-t", "filter", "-F", chain], None).await?;
if run_command(binary, &["-t", "filter", "-C", "INPUT", "-j", chain], None)
.await
.is_err()
{
run_command(
binary,
&["-t", "filter", "-I", "INPUT", "1", "-j", IPTABLES_CHAIN],
&["-t", "filter", "-I", "INPUT", "1", "-j", chain],
None,
)
.await?;
}
for (idx, target) in targets.iter().enumerate() {
for rule in iptables_synfix_rule_args(target, idx, family) {
for rule in iptables_synfix_rule_args(target, idx, family, namespace) {
let refs: Vec<&str> = rule.iter().map(String::as_str).collect();
run_command(binary, &refs, None).await?;
}
}
run_command(
binary,
&["-t", "filter", "-A", IPTABLES_CHAIN, "-j", "RETURN"],
None,
)
.await?;
run_command(binary, &["-t", "filter", "-A", chain, "-j", "RETURN"], None).await?;
Ok(())
}
@@ -89,12 +95,13 @@ fn iptables_synfix_rule_args(
target: &SynLimitRule,
idx: usize,
family: IpTablesFamily,
namespace: &SynLimitNamespace,
) -> Vec<Vec<String>> {
vec![
iptables_ios_accept_rule_args(target, idx, family),
iptables_ios_reject_rule_args(target, family),
iptables_generic_accept_rule_args(target, idx, family),
iptables_generic_reject_rule_args(target),
iptables_ios_accept_rule_args(target, idx, family, namespace),
iptables_ios_reject_rule_args(target, family, namespace),
iptables_generic_accept_rule_args(target, idx, family, namespace),
iptables_generic_reject_rule_args(target, namespace),
]
}
@@ -102,12 +109,15 @@ fn iptables_ios_accept_rule_args(
target: &SynLimitRule,
idx: usize,
family: IpTablesFamily,
namespace: &SynLimitNamespace,
) -> Vec<String> {
let hashlimit_name = format!(
"{IPTABLES_HASHLIMIT_PREFIX}-I{}-{idx}",
"{}-I{}-{idx}",
namespace.iptables_hashlimit_prefix,
family.hashlimit_tag()
);
let mut args = iptables_base_rule_args(target.ip, target.port);
let mut args =
iptables_base_rule_args(namespace.iptables_chain.as_str(), target.ip, target.port);
args.extend(iptables_ios_match_args(family));
args.extend(iptables_hashlimit_args(
&hashlimit_name,
@@ -121,8 +131,13 @@ fn iptables_ios_accept_rule_args(
args
}
fn iptables_ios_reject_rule_args(target: &SynLimitRule, family: IpTablesFamily) -> Vec<String> {
let mut args = iptables_base_rule_args(target.ip, target.port);
fn iptables_ios_reject_rule_args(
target: &SynLimitRule,
family: IpTablesFamily,
namespace: &SynLimitNamespace,
) -> Vec<String> {
let mut args =
iptables_base_rule_args(namespace.iptables_chain.as_str(), target.ip, target.port);
args.extend(iptables_ios_match_args(family));
args.extend(iptables_reject_args());
args
@@ -132,12 +147,15 @@ fn iptables_generic_accept_rule_args(
target: &SynLimitRule,
idx: usize,
family: IpTablesFamily,
namespace: &SynLimitNamespace,
) -> Vec<String> {
let hashlimit_name = format!(
"{IPTABLES_HASHLIMIT_PREFIX}-G{}-{idx}",
"{}-G{}-{idx}",
namespace.iptables_hashlimit_prefix,
family.hashlimit_tag()
);
let mut args = iptables_base_rule_args(target.ip, target.port);
let mut args =
iptables_base_rule_args(namespace.iptables_chain.as_str(), target.ip, target.port);
args.extend(iptables_hashlimit_args(
&hashlimit_name,
target.generic_seconds,
@@ -150,18 +168,22 @@ fn iptables_generic_accept_rule_args(
args
}
fn iptables_generic_reject_rule_args(target: &SynLimitRule) -> Vec<String> {
let mut args = iptables_base_rule_args(target.ip, target.port);
fn iptables_generic_reject_rule_args(
target: &SynLimitRule,
namespace: &SynLimitNamespace,
) -> Vec<String> {
let mut args =
iptables_base_rule_args(namespace.iptables_chain.as_str(), target.ip, target.port);
args.extend(iptables_reject_args());
args
}
fn iptables_base_rule_args(ip: Option<IpAddr>, port: u16) -> Vec<String> {
fn iptables_base_rule_args(chain: &str, ip: Option<IpAddr>, port: u16) -> Vec<String> {
let mut args = vec![
"-t".to_string(),
"filter".to_string(),
"-A".to_string(),
IPTABLES_CHAIN.to_string(),
chain.to_string(),
"-p".to_string(),
"tcp".to_string(),
"--syn".to_string(),
@@ -226,17 +248,18 @@ fn iptables_reject_args() -> Vec<String> {
]
}
pub(super) async fn clear_rules_for_binary(binary: &str) -> Result<(), String> {
pub(super) async fn clear_rules_for_binary(
binary: &str,
namespace: &SynLimitNamespace,
) -> Result<bool, String> {
let mut errors = Vec::new();
let mut removed = false;
let chain = namespace.iptables_chain.as_str();
for _ in 0..8 {
match run_command(
binary,
&["-t", "filter", "-D", "INPUT", "-j", IPTABLES_CHAIN],
None,
)
.await
{
Ok(()) => {}
match run_command(binary, &["-t", "filter", "-D", "INPUT", "-j", chain], None).await {
Ok(()) => {
removed = true;
}
Err(error) if is_missing_command_or_iptables_rule(&error) => break,
Err(error) => {
errors.push(format!("{binary} delete INPUT jump failed: {error}"));
@@ -244,19 +267,27 @@ pub(super) async fn clear_rules_for_binary(binary: &str) -> Result<(), String> {
}
}
}
if let Err(error) = run_command(binary, &["-t", "filter", "-F", IPTABLES_CHAIN], None).await
&& !is_missing_command_or_iptables_rule(&error)
{
errors.push(format!("{binary} flush chain failed: {error}"));
match run_command(binary, &["-t", "filter", "-F", chain], None).await {
Ok(()) => {
removed = true;
}
Err(error) if is_missing_command_or_iptables_rule(&error) => {}
Err(error) => {
errors.push(format!("{binary} flush chain failed: {error}"));
}
}
if let Err(error) = run_command(binary, &["-t", "filter", "-X", IPTABLES_CHAIN], None).await
&& !is_missing_command_or_iptables_rule(&error)
{
errors.push(format!("{binary} delete chain failed: {error}"));
match run_command(binary, &["-t", "filter", "-X", chain], None).await {
Ok(()) => {
removed = true;
}
Err(error) if is_missing_command_or_iptables_rule(&error) => {}
Err(error) => {
errors.push(format!("{binary} delete chain failed: {error}"));
}
}
if errors.is_empty() {
Ok(())
Ok(removed)
} else {
Err(errors.join(", "))
}
@@ -266,6 +297,7 @@ fn is_missing_command_or_iptables_rule(error: &str) -> bool {
error.contains("is not available")
|| error.contains("No chain/target/match by that name")
|| error.contains("does not exist")
|| error.contains("Couldn't load target")
}
#[cfg(test)]
@@ -280,12 +312,26 @@ mod tests {
.any(|pair| pair[0].as_str() == key && pair[1].as_str() == value)
}
fn has_key(args: &[String], key: &str) -> bool {
args.iter().any(|arg| arg == key)
}
fn test_namespace() -> SynLimitNamespace {
SynLimitNamespace {
nft_table: "telemt_synlimit_test".to_string(),
iptables_chain: "TMT_SYN_TEST".to_string(),
iptables_hashlimit_prefix: "TMTTEST".to_string(),
}
}
#[test]
fn iptables_rules_use_synfix_order_and_rejects() {
let target = test_rule(Some(IpAddr::V4(Ipv4Addr::new(203, 0, 113, 7))), 443);
let rules = iptables_synfix_rule_args(&target, 0, IpTablesFamily::V4);
let namespace = test_namespace();
let rules = iptables_synfix_rule_args(&target, 0, IpTablesFamily::V4, &namespace);
assert_eq!(rules.len(), 4);
assert!(has_pair(&rules[0], "-A", "TMT_SYN_TEST"));
assert!(has_pair(&rules[0], "--length", "64"));
assert!(has_pair(&rules[0], "--ttl-lt", "65"));
assert!(has_pair(&rules[0], "--hashlimit-upto", "12/second"));
@@ -302,10 +348,42 @@ mod tests {
#[test]
fn ip6tables_rules_use_ipv6_hoplimit_classifier() {
let target = test_rule(Some(IpAddr::V6(Ipv6Addr::LOCALHOST)), 443);
let rules = iptables_synfix_rule_args(&target, 0, IpTablesFamily::V6);
let namespace = test_namespace();
let rules = iptables_synfix_rule_args(&target, 0, IpTablesFamily::V6, &namespace);
assert!(has_pair(&rules[0], "--length", "84"));
assert!(has_pair(&rules[0], "--hl-lt", "65"));
assert!(has_pair(&rules[0], "-d", "::1"));
}
#[test]
fn iptables_missing_rule_errors_are_cleanup_benign() {
assert!(is_missing_command_or_iptables_rule(
"iptables is not available"
));
assert!(is_missing_command_or_iptables_rule(
"iptables: No chain/target/match by that name."
));
assert!(is_missing_command_or_iptables_rule(
"iptables: Chain TELEMT_SYNLIMIT does not exist."
));
assert!(is_missing_command_or_iptables_rule(
"Couldn't load target `TELEMT_SYNLIMIT': No such file or directory"
));
assert!(!is_missing_command_or_iptables_rule(
"iptables: Permission denied"
));
}
#[test]
fn iptables_wildcard_rule_omits_destination_match() {
let target = test_rule(None, 443);
let namespace = test_namespace();
let rules = iptables_synfix_rule_args(&target, 0, IpTablesFamily::V4, &namespace);
for rule in rules {
assert!(!has_key(&rule, "-d"));
assert!(has_pair(&rule, "--dport", "443"));
}
}
}
+89 -16
View File
@@ -1,4 +1,4 @@
use std::sync::Arc;
use std::sync::{Arc, Mutex};
use tokio::sync::watch;
use tracing::warn;
@@ -11,7 +11,9 @@ mod model;
mod nftables;
use self::command::has_cap_net_admin;
use self::model::synlimit_targets;
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") {
@@ -39,14 +41,26 @@ async fn wait_for_config_channel_close_and_reconcile(
}
pub(crate) async fn reconcile_synlimit_rules(cfg: &ProxyConfig) {
if let Err(error) = clear_synlimit_rules_all_backends().await {
warn!(error = %error, "Failed to clear existing SYN limiter rules before reconcile");
let targets = synlimit_targets(cfg);
let namespace = synlimit_namespace(&targets);
if let Some(previous_namespace) = set_active_synlimit_namespace(namespace.clone()) {
match clear_synlimit_rules_for_namespace(&previous_namespace).await {
Ok(true) => {
warn!("Removed previous SYN limiter namespace before reconcile");
}
Ok(false) => {}
Err(error) => {
warn!(error = %error, "Failed to clear previous SYN limiter namespace before reconcile");
}
}
}
let targets = synlimit_targets(cfg);
if targets.is_empty() {
return;
}
let Some(namespace) = namespace else {
return;
};
if !has_cap_net_admin() {
warn!(
"SYN limiter configured but CAP_NET_ADMIN is not available; netfilter rules not applied"
@@ -54,41 +68,100 @@ pub(crate) async fn reconcile_synlimit_rules(cfg: &ProxyConfig) {
return;
}
match clear_synlimit_rules_for_namespace(&namespace).await {
Ok(true) => {
warn!("Removed stale SYN limiter rules left by a previous run before reconcile");
}
Ok(false) => {}
Err(error) => {
warn!(error = %error, "Failed to clear stale SYN limiter rules before reconcile");
}
}
if targets.has_iptables_targets()
&& let Err(error) = iptables::apply_synlimit_rules(&targets).await
&& let Err(error) = iptables::apply_synlimit_rules(&targets, &namespace).await
{
warn!(error = %error, "Failed to apply iptables SYN limiter rules");
}
if targets.has_nft_targets()
&& let Err(error) = nftables::apply_synlimit_rules(&targets).await
&& let Err(error) = nftables::apply_synlimit_rules(&targets, &namespace).await
{
warn!(error = %error, "Failed to apply nftables SYN limiter rules");
}
}
pub(crate) async fn clear_synlimit_rules_all_backends() -> Result<(), String> {
pub(crate) async fn clear_synlimit_rules_all_backends() -> Result<bool, String> {
let Some(namespace) = take_active_synlimit_namespace() else {
return Ok(false);
};
clear_synlimit_rules_for_namespace(&namespace).await
}
async fn clear_synlimit_rules_for_namespace(namespace: &SynLimitNamespace) -> Result<bool, String> {
if !has_cap_net_admin() {
return Ok(());
return Ok(false);
}
let mut errors = Vec::new();
if let Err(error) = nftables::clear_rules_all_families().await {
errors.push(error);
let mut removed = false;
match nftables::clear_rules_all_families(namespace).await {
Ok(value) => {
removed |= value;
}
Err(error) => {
errors.push(error);
}
}
if let Err(error) = iptables::clear_rules_for_binary("iptables").await {
errors.push(error);
match iptables::clear_rules_for_binary("iptables", namespace).await {
Ok(value) => {
removed |= value;
}
Err(error) => {
errors.push(error);
}
}
if let Err(error) = iptables::clear_rules_for_binary("ip6tables").await {
errors.push(error);
match iptables::clear_rules_for_binary("ip6tables", namespace).await {
Ok(value) => {
removed |= value;
}
Err(error) => {
errors.push(error);
}
}
if errors.is_empty() {
Ok(())
Ok(removed)
} else {
Err(errors.join("; "))
}
}
fn set_active_synlimit_namespace(next: Option<SynLimitNamespace>) -> Option<SynLimitNamespace> {
match ACTIVE_SYNLIMIT_NAMESPACE.lock() {
Ok(mut active) => {
if *active == next {
None
} else {
std::mem::replace(&mut *active, next)
}
}
Err(error) => {
warn!(error = %error, "Failed to update active SYN limiter namespace");
None
}
}
}
fn take_active_synlimit_namespace() -> Option<SynLimitNamespace> {
match ACTIVE_SYNLIMIT_NAMESPACE.lock() {
Ok(mut active) => active.take(),
Err(error) => {
warn!(error = %error, "Failed to read active SYN limiter namespace");
None
}
}
}
fn has_synlimit_config(cfg: &ProxyConfig) -> bool {
cfg.server
.listeners
+239
View File
@@ -17,6 +17,13 @@ pub(super) struct SynLimitRule {
pub(super) hashlimit_size: u32,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub(super) struct SynLimitNamespace {
pub(super) nft_table: String,
pub(super) iptables_chain: String,
pub(super) iptables_hashlimit_prefix: String,
}
#[derive(Default)]
pub(super) struct SynLimitTargets {
pub(super) iptables_v4: Vec<SynLimitRule>,
@@ -42,6 +49,44 @@ impl SynLimitTargets {
}
}
struct SynLimitNamespaceHasher {
value: u64,
}
impl SynLimitNamespaceHasher {
const OFFSET: u64 = 0xcbf2_9ce4_8422_2325;
const PRIME: u64 = 0x0000_0100_0000_01b3;
fn new() -> Self {
Self {
value: Self::OFFSET,
}
}
fn write(&mut self, bytes: &[u8]) {
for byte in bytes {
self.value ^= u64::from(*byte);
self.value = self.value.wrapping_mul(Self::PRIME);
}
}
fn write_u8(&mut self, value: u8) {
self.write(&[value]);
}
fn write_u16(&mut self, value: u16) {
self.write(&value.to_le_bytes());
}
fn write_u32(&mut self, value: u32) {
self.write(&value.to_le_bytes());
}
fn finish(self) -> u64 {
self.value
}
}
pub(super) fn synlimit_targets(cfg: &ProxyConfig) -> SynLimitTargets {
let mut iptables_v4 = BTreeSet::new();
let mut iptables_v6 = BTreeSet::new();
@@ -91,6 +136,64 @@ pub(super) fn synlimit_targets(cfg: &ProxyConfig) -> SynLimitTargets {
}
}
pub(super) fn synlimit_namespace(targets: &SynLimitTargets) -> Option<SynLimitNamespace> {
if targets.is_empty() {
return None;
}
let mut hasher = SynLimitNamespaceHasher::new();
write_namespace_rule_group(&mut hasher, b"iptables-v4", &targets.iptables_v4);
write_namespace_rule_group(&mut hasher, b"iptables-v6", &targets.iptables_v6);
write_namespace_rule_group(&mut hasher, b"nft-v4", &targets.nft_v4);
write_namespace_rule_group(&mut hasher, b"nft-v6", &targets.nft_v6);
let suffix = format!("{:016x}", hasher.finish());
let iptables_suffix = &suffix[..12];
let hashlimit_suffix = &suffix[..10];
Some(SynLimitNamespace {
nft_table: format!("telemt_synlimit_{suffix}"),
iptables_chain: format!("TMT_SYN_{iptables_suffix}"),
iptables_hashlimit_prefix: format!("TMT{hashlimit_suffix}"),
})
}
fn write_namespace_rule_group(
hasher: &mut SynLimitNamespaceHasher,
group: &[u8],
rules: &[SynLimitRule],
) {
hasher.write(group);
hasher.write_u32(rules.len() as u32);
for rule in rules {
write_namespace_rule(hasher, rule);
}
}
fn write_namespace_rule(hasher: &mut SynLimitNamespaceHasher, rule: &SynLimitRule) {
match rule.ip {
Some(IpAddr::V4(ip)) => {
hasher.write_u8(4);
hasher.write(&ip.octets());
}
Some(IpAddr::V6(ip)) => {
hasher.write_u8(6);
hasher.write(&ip.octets());
}
None => {
hasher.write_u8(0);
}
}
hasher.write_u16(rule.port);
hasher.write_u32(rule.generic_seconds);
hasher.write_u32(rule.generic_hitcount);
hasher.write_u32(rule.generic_burst);
hasher.write_u32(rule.ios_seconds);
hasher.write_u32(rule.ios_hitcount);
hasher.write_u32(rule.ios_burst);
hasher.write_u32(rule.hashlimit_expire_ms);
hasher.write_u32(rule.hashlimit_size);
}
pub(super) fn synlimit_rate_arg(seconds: u32, hitcount: u32) -> String {
let seconds = u64::from(seconds.max(1));
let hitcount = u64::from(hitcount.max(1));
@@ -124,3 +227,139 @@ pub(super) fn test_rule(ip: Option<IpAddr>, port: u16) -> SynLimitRule {
hashlimit_size: 32_768,
}
}
#[cfg(test)]
mod tests {
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
use super::*;
use crate::config::ListenerConfig;
fn listener(ip: IpAddr, port: Option<u16>, synlimit: SynLimitMode) -> ListenerConfig {
ListenerConfig {
ip,
port,
client_mss: None,
synlimit,
synlimit_seconds: 60,
synlimit_hitcount: 48,
synlimit_burst: 1,
synlimit_ios_seconds: 1,
synlimit_ios_hitcount: 12,
synlimit_ios_burst: 24,
synlimit_hashlimit_expire_ms: 60_000,
synlimit_hashlimit_size: 32_768,
announce: None,
announce_ip: None,
proxy_protocol: None,
reuse_allow: false,
}
}
#[test]
fn synlimit_targets_deduplicate_and_use_legacy_port_fallback() {
let mut cfg = ProxyConfig::default();
cfg.server.port = 9443;
cfg.server.listeners = vec![
listener(
IpAddr::V4(Ipv4Addr::UNSPECIFIED),
None,
SynLimitMode::Iptables,
),
listener(
IpAddr::V4(Ipv4Addr::UNSPECIFIED),
None,
SynLimitMode::Iptables,
),
];
let targets = synlimit_targets(&cfg);
assert_eq!(targets.iptables_v4.len(), 1);
assert_eq!(targets.iptables_v4[0].ip, None);
assert_eq!(targets.iptables_v4[0].port, 9443);
assert!(targets.iptables_v6.is_empty());
assert!(targets.nft_v4.is_empty());
assert!(targets.nft_v6.is_empty());
}
#[test]
fn synlimit_targets_separate_backends_and_ip_families() {
let mut cfg = ProxyConfig::default();
cfg.server.listeners = vec![
listener(
IpAddr::V4(Ipv4Addr::new(203, 0, 113, 1)),
Some(443),
SynLimitMode::Iptables,
),
listener(
IpAddr::V6(Ipv6Addr::LOCALHOST),
Some(443),
SynLimitMode::Iptables,
),
listener(
IpAddr::V4(Ipv4Addr::new(203, 0, 113, 2)),
Some(444),
SynLimitMode::Nftables,
),
listener(
IpAddr::V6(Ipv6Addr::UNSPECIFIED),
Some(444),
SynLimitMode::Nftables,
),
];
let targets = synlimit_targets(&cfg);
assert_eq!(targets.iptables_v4.len(), 1);
assert_eq!(targets.iptables_v6.len(), 1);
assert_eq!(targets.nft_v4.len(), 1);
assert_eq!(targets.nft_v6.len(), 1);
assert_eq!(
targets.iptables_v4[0].ip,
Some(IpAddr::V4(Ipv4Addr::new(203, 0, 113, 1)))
);
assert_eq!(
targets.iptables_v6[0].ip,
Some(IpAddr::V6(Ipv6Addr::LOCALHOST))
);
assert_eq!(
targets.nft_v4[0].ip,
Some(IpAddr::V4(Ipv4Addr::new(203, 0, 113, 2)))
);
assert_eq!(targets.nft_v6[0].ip, None);
}
#[test]
fn synlimit_namespace_is_stable_and_changes_by_targets() {
let mut cfg = ProxyConfig::default();
cfg.server.listeners = vec![listener(
IpAddr::V4(Ipv4Addr::new(203, 0, 113, 1)),
Some(443),
SynLimitMode::Nftables,
)];
let first = synlimit_namespace(&synlimit_targets(&cfg))
.expect("configured targets must have a namespace");
let second = synlimit_namespace(&synlimit_targets(&cfg))
.expect("configured targets must have a namespace");
cfg.server.listeners[0].port = Some(444);
let changed = synlimit_namespace(&synlimit_targets(&cfg))
.expect("configured targets must have a namespace");
assert_eq!(first, second);
assert_ne!(first, changed);
assert!(first.nft_table.starts_with("telemt_synlimit_"));
assert!(first.iptables_chain.starts_with("TMT_SYN_"));
assert!(first.iptables_chain.len() <= 28);
assert!(first.iptables_hashlimit_prefix.starts_with("TMT"));
}
#[test]
fn synlimit_rate_arg_uses_native_units_without_fractional_rates() {
assert_eq!(synlimit_rate_arg(1, 12), "12/second");
assert_eq!(synlimit_rate_arg(60, 48), "48/minute");
assert_eq!(synlimit_rate_arg(3600, 121), "121/hour");
assert_eq!(synlimit_rate_arg(86400, 241), "241/day");
}
}
+90 -30
View File
@@ -1,7 +1,6 @@
use super::command::{run_command, run_command_stdout};
use super::model::{SynLimitRule, SynLimitTargets, synlimit_rate_arg};
use super::model::{SynLimitNamespace, SynLimitRule, SynLimitTargets, synlimit_rate_arg};
const NFT_TABLE: &str = "telemt_synlimit";
const NFT_CHAIN: &str = "input";
const NFT_INPUT_PRIORITY: i16 = -5;
const IPV4_IOS_PACKET_LENGTH: u16 = 64;
@@ -38,10 +37,13 @@ impl NftFamily {
}
}
pub(super) async fn apply_synlimit_rules(targets: &SynLimitTargets) -> Result<(), String> {
pub(super) async fn apply_synlimit_rules(
targets: &SynLimitTargets,
namespace: &SynLimitNamespace,
) -> Result<(), String> {
let families = detect_nft_table_families().await;
for plan in nft_apply_plan(families, &targets.nft_v4, &targets.nft_v6) {
let script = nft_synlimit_script(plan);
let script = nft_synlimit_script(plan, namespace);
run_command("nft", &["-f", "-"], Some(script)).await?;
}
@@ -114,9 +116,13 @@ fn nft_apply_plan<'a>(
Vec::new()
}
fn nft_synlimit_script(plan: NftApplyPlan<'_>) -> String {
fn nft_synlimit_script(plan: NftApplyPlan<'_>, namespace: &SynLimitNamespace) -> String {
let mut script = String::new();
script.push_str(&format!("table {} {NFT_TABLE} {{\n", plan.family.as_str()));
script.push_str(&format!(
"table {} {} {{\n",
plan.family.as_str(),
namespace.nft_table
));
script.push_str(&format!(" chain {NFT_CHAIN} {{\n"));
script.push_str(&format!(
" type filter hook input priority {NFT_INPUT_PRIORITY}; policy accept;\n"
@@ -186,26 +192,29 @@ fn push_nft_v6_rules(script: &mut String, target: &SynLimitRule, idx: usize) {
));
}
pub(super) async fn clear_rules_all_families() -> Result<(), String> {
pub(super) async fn clear_rules_all_families(
namespace: &SynLimitNamespace,
) -> Result<bool, String> {
let mut errors = Vec::new();
let mut removed = false;
let table = namespace.nft_table.as_str();
for family in [NftFamily::Inet, NftFamily::Ip, NftFamily::Ip6] {
if let Err(error) = run_command(
"nft",
&["delete", "table", family.as_str(), NFT_TABLE],
None,
)
.await
&& !is_missing_command_or_nft_table(&error)
{
errors.push(format!(
"nft delete table {} {NFT_TABLE} failed: {error}",
family.as_str()
));
match run_command("nft", &["delete", "table", family.as_str(), table], None).await {
Ok(()) => {
removed = true;
}
Err(error) if is_missing_command_or_nft_table(&error) => {}
Err(error) => {
errors.push(format!(
"nft delete table {} {table} failed: {error}",
family.as_str(),
));
}
}
}
if errors.is_empty() {
Ok(())
Ok(removed)
} else {
Err(errors.join(", "))
}
@@ -222,15 +231,28 @@ mod tests {
use super::*;
use crate::synlimit_control::model::test_rule;
fn test_namespace(table: &str) -> SynLimitNamespace {
SynLimitNamespace {
nft_table: table.to_string(),
iptables_chain: "TMT_SYN_TEST".to_string(),
iptables_hashlimit_prefix: "TMTTEST".to_string(),
}
}
#[test]
fn nft_script_uses_synfix_v4_rules_and_early_priority() {
let rule = test_rule(Some(IpAddr::V4(Ipv4Addr::new(203, 0, 113, 7))), 443);
let script = nft_synlimit_script(NftApplyPlan {
family: NftFamily::Inet,
v4_targets: &[rule],
v6_targets: &[],
});
let namespace = test_namespace("telemt_synlimit_test_a");
let script = nft_synlimit_script(
NftApplyPlan {
family: NftFamily::Inet,
v4_targets: &[rule],
v6_targets: &[],
},
&namespace,
);
assert!(script.contains("table inet telemt_synlimit_test_a"));
assert!(script.contains("type filter hook input priority -5; policy accept;"));
assert!(script.contains("ip daddr 203.0.113.7"));
assert!(script.contains("meta length 64 ip ttl < 65"));
@@ -242,15 +264,53 @@ mod tests {
#[test]
fn nft_script_uses_ipv6_hoplimit_classifier() {
let rule = test_rule(Some(IpAddr::V6(Ipv6Addr::LOCALHOST)), 443);
let script = nft_synlimit_script(NftApplyPlan {
family: NftFamily::Inet,
v4_targets: &[],
v6_targets: &[rule],
});
let namespace = test_namespace("telemt_synlimit_test_b");
let script = nft_synlimit_script(
NftApplyPlan {
family: NftFamily::Inet,
v4_targets: &[],
v6_targets: &[rule],
},
&namespace,
);
assert!(script.contains("table inet telemt_synlimit_test_b"));
assert!(script.contains("ip6 daddr ::1"));
assert!(script.contains("meta length 84 ip6 hoplimit < 65"));
assert!(script.contains("ip6 saddr limit rate over 12/second burst 24 packets"));
assert!(script.contains("ip6 saddr limit rate over 48/minute burst 1 packets"));
}
#[test]
fn nft_missing_table_errors_are_cleanup_benign() {
assert!(is_missing_command_or_nft_table("nft is not available"));
assert!(is_missing_command_or_nft_table(
"Error: No such file or directory"
));
assert!(!is_missing_command_or_nft_table(
"Error: Operation not permitted"
));
}
#[test]
fn nft_apply_plan_keeps_dual_stack_rules_in_inet_table() {
let v4_rule = test_rule(Some(IpAddr::V4(Ipv4Addr::new(203, 0, 113, 7))), 443);
let v6_rule = test_rule(Some(IpAddr::V6(Ipv6Addr::LOCALHOST)), 443);
let v4_rules = [v4_rule];
let v6_rules = [v6_rule];
let plans = nft_apply_plan(
NftTableFamilies {
inet: false,
ip: false,
ip6: false,
},
&v4_rules,
&v6_rules,
);
assert_eq!(plans.len(), 1);
assert_eq!(plans[0].family.as_str(), "inet");
assert_eq!(plans[0].v4_targets, v4_rules.as_slice());
assert_eq!(plans[0].v6_targets, v6_rules.as_slice());
}
}
+100 -3
View File
@@ -1,9 +1,12 @@
use bytes::Bytes;
use std::sync::Arc;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::sync::OwnedSemaphorePermit;
use crate::crypto::{AesCbc, crc32, crc32c};
use crate::error::{ProxyError, Result};
use crate::protocol::constants::*;
use crate::stats::Stats;
use crate::stream::PooledBuffer;
use super::wire::{append_proxy_req_payload_into, proxy_req_payload_len};
@@ -11,9 +14,66 @@ use super::wire::{append_proxy_req_payload_into, proxy_req_payload_len};
const RPC_WRITER_FRAME_BUF_SHRINK_THRESHOLD: usize = 256 * 1024;
const RPC_WRITER_FRAME_BUF_RETAIN: usize = 64 * 1024;
enum WriterBytePermitState {
Queued,
Inflight,
}
/// Holds one writer memory reservation through queueing and socket write completion.
pub(crate) struct WriterBytePermit {
_permit: OwnedSemaphorePermit,
reserved_bytes: u64,
stats: Arc<Stats>,
state: WriterBytePermitState,
}
impl WriterBytePermit {
/// Creates a queued reservation and publishes its rounded resident-byte cost.
pub(crate) fn new(
permit: OwnedSemaphorePermit,
reserved_bytes: usize,
stats: Arc<Stats>,
) -> Self {
let reserved_bytes = reserved_bytes as u64;
stats.add_me_writer_byte_budget_queued_bytes(reserved_bytes);
Self {
_permit: permit,
reserved_bytes,
stats,
state: WriterBytePermitState::Queued,
}
}
/// Moves the reservation from queued to in-flight accounting exactly once.
pub(crate) fn mark_inflight(&mut self) {
if matches!(self.state, WriterBytePermitState::Queued) {
self.stats
.move_me_writer_byte_budget_to_inflight(self.reserved_bytes);
self.state = WriterBytePermitState::Inflight;
}
}
}
impl Drop for WriterBytePermit {
fn drop(&mut self) {
match self.state {
WriterBytePermitState::Queued => self
.stats
.release_me_writer_byte_budget_queued_bytes(self.reserved_bytes),
WriterBytePermitState::Inflight => self
.stats
.release_me_writer_byte_budget_inflight_bytes(self.reserved_bytes),
}
}
}
/// Commands sent to dedicated writer tasks to avoid mutex contention on TCP writes.
pub(crate) enum WriterCommand {
Data(Bytes),
Data {
payload: Bytes,
_permit: Option<OwnedSemaphorePermit>,
writer_permit: WriterBytePermit,
},
DataAndFlush(Bytes),
ProxyReq(ProxyReqCommand),
ControlAndFlush([u8; 12]),
@@ -28,6 +88,8 @@ pub(crate) struct ProxyReqCommand {
pub(crate) proto_flags: u32,
pub(crate) proxy_tag: Option<[u8; 16]>,
pub(crate) payload: PooledBuffer,
pub(crate) _permit: OwnedSemaphorePermit,
pub(crate) writer_permit: WriterBytePermit,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -82,7 +144,7 @@ fn build_rpc_frame_into(
) {
let total_len = 4 + 4 + payload.len() + 4;
frame.clear();
frame.reserve(total_len + 15);
frame.reserve_exact(total_len + 15);
let total_len = total_len as u32;
frame.extend_from_slice(&total_len.to_le_bytes());
frame.extend_from_slice(&seq_no.to_le_bytes());
@@ -274,7 +336,7 @@ impl RpcWriter {
);
let total_len = 4 + 4 + payload_len + 4;
self.frame_buf.clear();
self.frame_buf.reserve(total_len + 15);
self.frame_buf.reserve_exact(total_len + 15);
self.frame_buf
.extend_from_slice(&(total_len as u32).to_le_bytes());
self.frame_buf.extend_from_slice(&self.seq_no.to_le_bytes());
@@ -322,3 +384,38 @@ impl RpcWriter {
self.writer.flush().await.map_err(ProxyError::Io)
}
}
#[cfg(test)]
mod tests {
use super::*;
use tokio::sync::Semaphore;
#[test]
fn writer_byte_permit_tracks_queued_and_inflight_lifecycle() {
let stats = Arc::new(Stats::default());
let semaphore = Arc::new(Semaphore::new(2));
let permit = semaphore
.clone()
.try_acquire_many_owned(2)
.expect("writer byte permits must be available");
let mut writer_permit = WriterBytePermit::new(permit, 32 * 1024, stats.clone());
assert_eq!(
stats.get_me_writer_byte_budget_queued_bytes_gauge(),
32 * 1024
);
assert_eq!(stats.get_me_writer_byte_budget_inflight_bytes_gauge(), 0);
writer_permit.mark_inflight();
assert_eq!(stats.get_me_writer_byte_budget_queued_bytes_gauge(), 0);
assert_eq!(
stats.get_me_writer_byte_budget_inflight_bytes_gauge(),
32 * 1024
);
drop(writer_permit);
assert_eq!(stats.get_me_writer_byte_budget_queued_bytes_gauge(), 0);
assert_eq!(stats.get_me_writer_byte_budget_inflight_bytes_gauge(), 0);
assert_eq!(semaphore.available_permits(), 2);
}
}
+114
View File
@@ -675,8 +675,122 @@ fn hex_dump(data: &[u8]) -> String {
mod tests {
use super::*;
use std::io::ErrorKind;
use std::net::{Ipv4Addr, Ipv6Addr};
use tokio::net::{TcpListener, TcpStream};
fn upstream_egress(
route_kind: UpstreamRouteKind,
socks_bound_addr: Option<SocketAddr>,
) -> UpstreamEgressInfo {
UpstreamEgressInfo {
upstream_id: 7,
route_kind,
local_addr: None,
direct_bind_ip: None,
socks_bound_addr,
socks_proxy_addr: None,
}
}
#[test]
fn socks_bound_addr_is_used_only_for_public_same_family_tuple() {
let v4_bound = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(93, 184, 216, 34)), 443);
let v6_bound = SocketAddr::new(
IpAddr::V6(
"2606:4700:4700::1111"
.parse::<Ipv6Addr>()
.expect("test IPv6 address must parse"),
),
443,
);
assert_eq!(
MePool::select_socks_bound_addr(
IpFamily::V4,
Some(upstream_egress(UpstreamRouteKind::Socks5, Some(v4_bound)))
),
Some(v4_bound)
);
assert_eq!(
MePool::select_socks_bound_addr(
IpFamily::V6,
Some(upstream_egress(UpstreamRouteKind::Socks5, Some(v6_bound)))
),
Some(v6_bound)
);
}
#[test]
fn socks_bound_addr_rejects_bogon_unspecified_wrong_family_and_non_socks_routes() {
let bogon_bound = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1)), 443);
let unspecified_bound = SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), 443);
let public_v4_bound = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(93, 184, 216, 34)), 443);
assert_eq!(
MePool::select_socks_bound_addr(
IpFamily::V4,
Some(upstream_egress(
UpstreamRouteKind::Socks5,
Some(bogon_bound)
))
),
None
);
assert_eq!(
MePool::select_socks_bound_addr(
IpFamily::V4,
Some(upstream_egress(
UpstreamRouteKind::Socks5,
Some(unspecified_bound)
))
),
None
);
assert_eq!(
MePool::select_socks_bound_addr(
IpFamily::V6,
Some(upstream_egress(
UpstreamRouteKind::Socks5,
Some(public_v4_bound)
))
),
None
);
assert_eq!(
MePool::select_socks_bound_addr(
IpFamily::V4,
Some(upstream_egress(
UpstreamRouteKind::Direct,
Some(public_v4_bound)
))
),
None
);
}
#[test]
fn kdf_client_port_source_tracks_only_valid_socks_bound_port() {
let bound = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(93, 184, 216, 34)), 443);
let zero_port_bound = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(93, 184, 216, 34)), 0);
assert_eq!(
KdfClientPortSource::from_socks_bound_port(Some(bound.port())),
KdfClientPortSource::SocksBound
);
assert_eq!(
KdfClientPortSource::from_socks_bound_port(
Some(zero_port_bound)
.filter(|addr| addr.port() != 0)
.map(|addr| addr.port())
),
KdfClientPortSource::LocalSocket
);
assert_eq!(
KdfClientPortSource::from_socks_bound_port(None),
KdfClientPortSource::LocalSocket
);
}
#[tokio::test]
async fn test_configure_keepalive_loopback() {
let listener = match TcpListener::bind("127.0.0.1:0").await {
+11 -2
View File
@@ -1795,6 +1795,7 @@ mod tests {
general.me_writer_pick_sample_size,
MeSocksKdfPolicy::default(),
general.me_writer_cmd_channel_capacity,
general.me_writer_byte_budget_bytes,
general.me_route_channel_capacity,
general.me_route_backpressure_enabled,
general.me_route_fairshare_enabled,
@@ -1821,6 +1822,7 @@ mod tests {
) -> u64 {
let (conn_id, _rx) = pool.registry.register().await;
let (tx, _writer_rx) = mpsc::channel::<WriterCommand>(8);
let byte_budget = pool.new_writer_byte_budget();
let writer = MeWriter {
id: writer_id,
addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 4000 + writer_id as u16),
@@ -1830,6 +1832,7 @@ mod tests {
contour: Arc::new(AtomicU8::new(WriterContour::Draining.as_u8())),
created_at: Instant::now() - Duration::from_secs(writer_id),
tx: tx.clone(),
byte_budget: byte_budget.clone(),
cancel: CancellationToken::new(),
degraded: Arc::new(AtomicBool::new(false)),
rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)),
@@ -1839,7 +1842,9 @@ mod tests {
allow_drain_fallback: Arc::new(AtomicBool::new(false)),
};
pool.writers.write().await.push(writer);
pool.registry.register_writer(writer_id, tx).await;
pool.registry
.register_writer(writer_id, tx, byte_budget)
.await;
pool.conn_count.fetch_add(1, Ordering::Relaxed);
assert!(
pool.registry
@@ -1860,6 +1865,7 @@ mod tests {
async fn insert_live_writer(pool: &Arc<MePool>, writer_id: u64, writer_dc: i32) {
let (tx, _writer_rx) = mpsc::channel::<WriterCommand>(8);
let byte_budget = pool.new_writer_byte_budget();
let writer = MeWriter {
id: writer_id,
addr: SocketAddr::new(
@@ -1877,6 +1883,7 @@ mod tests {
contour: Arc::new(AtomicU8::new(WriterContour::Active.as_u8())),
created_at: Instant::now(),
tx: tx.clone(),
byte_budget: byte_budget.clone(),
cancel: CancellationToken::new(),
degraded: Arc::new(AtomicBool::new(false)),
rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)),
@@ -1886,7 +1893,9 @@ mod tests {
allow_drain_fallback: Arc::new(AtomicBool::new(false)),
};
pool.writers.write().await.push(writer);
pool.registry.register_writer(writer_id, tx).await;
pool.registry
.register_writer(writer_id, tx, byte_budget)
.await;
pool.conn_count.fetch_add(1, Ordering::Relaxed);
}
+16 -1
View File
@@ -10,7 +10,7 @@ use std::sync::atomic::{
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use arc_swap::ArcSwap;
use tokio::sync::{Mutex, RwLock, mpsc, watch};
use tokio::sync::{Mutex, RwLock, Semaphore, mpsc, watch};
use tokio_util::sync::CancellationToken;
use crate::config::{
@@ -48,6 +48,8 @@ pub struct MeWriter {
pub contour: Arc<AtomicU8>,
pub created_at: Instant,
pub tx: mpsc::Sender<WriterCommand>,
/// Aggregate resident-memory budget shared by all data commands for this writer.
pub byte_budget: Arc<Semaphore>,
pub cancel: CancellationToken,
pub degraded: Arc<AtomicBool>,
pub rtt_ema_ms_x10: Arc<AtomicU32>,
@@ -277,6 +279,7 @@ pub(super) struct WriterLifecycleCore {
pub(super) me_keepalive_payload_random: bool,
pub(super) rpc_proxy_req_every_secs: AtomicU64,
pub(super) writer_cmd_channel_capacity: usize,
pub(super) writer_byte_budget_permits: usize,
}
pub(super) struct RouteRuntimeCore {
@@ -553,6 +556,7 @@ impl MePool {
me_writer_pick_sample_size: u8,
me_socks_kdf_policy: MeSocksKdfPolicy,
me_writer_cmd_channel_capacity: usize,
me_writer_byte_budget_bytes: usize,
me_route_channel_capacity: usize,
me_route_backpressure_enabled: bool,
me_route_fairshare_enabled: bool,
@@ -583,6 +587,7 @@ impl MePool {
);
let (writer_epoch, _) = watch::channel(0u64);
let now_epoch_secs = Self::now_epoch_secs();
stats.set_me_writer_byte_budget_limit_bytes(me_writer_byte_budget_bytes);
Arc::new(Self {
routing: Arc::new(RoutingCore {
registry,
@@ -615,6 +620,9 @@ impl MePool {
me_keepalive_payload_random,
rpc_proxy_req_every_secs: AtomicU64::new(rpc_proxy_req_every_secs),
writer_cmd_channel_capacity: me_writer_cmd_channel_capacity.max(1),
writer_byte_budget_permits: me_writer_byte_budget_bytes
.div_ceil(crate::config::defaults::ME_WRITER_BYTE_PERMIT_UNIT_BYTES)
.max(1),
}),
route_runtime: Arc::new(RouteRuntimeCore {
me_route_no_writer_mode: AtomicU8::new(me_route_no_writer_mode.as_u8()),
@@ -833,6 +841,13 @@ impl MePool {
})
}
/// Creates the immutable byte semaphore assigned to one ME writer generation.
pub(crate) fn new_writer_byte_budget(&self) -> Arc<Semaphore> {
Arc::new(Semaphore::new(
self.writer_lifecycle.writer_byte_budget_permits,
))
}
pub fn current_generation(&self) -> u64 {
self.reinit.active_generation.load(Ordering::Relaxed)
}
+13 -3
View File
@@ -63,13 +63,19 @@ async fn writer_command_loop(
_ = cancel.cancelled() => return Ok(()),
cmd = rx.recv() => {
match cmd {
Some(WriterCommand::Data(payload)) => {
Some(WriterCommand::Data {
payload,
_permit,
mut writer_permit,
}) => {
writer_permit.mark_inflight();
rpc_writer.send(&payload).await?;
}
Some(WriterCommand::DataAndFlush(payload)) => {
rpc_writer.send_and_flush(&payload).await?;
}
Some(WriterCommand::ProxyReq(command)) => {
Some(WriterCommand::ProxyReq(mut command)) => {
command.writer_permit.mark_inflight();
rpc_writer.send_proxy_req(&command).await?;
}
Some(WriterCommand::ControlAndFlush(payload)) => {
@@ -419,6 +425,7 @@ impl MePool {
let draining_started_at_epoch_secs = Arc::new(AtomicU64::new(0));
let drain_deadline_epoch_secs = Arc::new(AtomicU64::new(0));
let allow_drain_fallback = Arc::new(AtomicBool::new(false));
let byte_budget = self.new_writer_byte_budget();
let (tx, rx) =
mpsc::channel::<WriterCommand>(self.writer_lifecycle.writer_cmd_channel_capacity);
let rpc_writer = RpcWriter {
@@ -438,6 +445,7 @@ impl MePool {
contour: contour.clone(),
created_at: Instant::now(),
tx: tx.clone(),
byte_budget: byte_budget.clone(),
cancel: cancel.clone(),
degraded: degraded.clone(),
rtt_ema_ms_x10: rtt_ema_ms_x10.clone(),
@@ -449,7 +457,9 @@ impl MePool {
self.writers
.update(|writers| writers.push(writer.clone()))
.await;
self.registry.register_writer(writer_id, tx.clone()).await;
self.registry
.register_writer(writer_id, tx.clone(), byte_budget)
.await;
self.registry.mark_writer_idle(writer_id).await;
self.conn_count.fetch_add(1, Ordering::Relaxed);
self.notify_writer_epoch();
+9 -1
View File
@@ -48,6 +48,8 @@ pub struct BoundConn {
pub struct ConnWriter {
pub writer_id: u64,
pub tx: mpsc::Sender<WriterCommand>,
/// Writer-local memory budget used by the hot bound-client route.
pub byte_budget: Arc<Semaphore>,
}
#[derive(Clone, Debug, Default)]
@@ -62,7 +64,13 @@ struct RoutingTable {
}
struct WriterTable {
map: DashMap<u64, mpsc::Sender<WriterCommand>>,
map: DashMap<u64, WriterRoute>,
}
#[derive(Clone)]
struct WriterRoute {
tx: mpsc::Sender<WriterCommand>,
byte_budget: Arc<Semaphore>,
}
#[derive(Clone)]
+30 -8
View File
@@ -1,10 +1,16 @@
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use std::sync::Arc;
use bytes::Bytes;
use tokio::sync::Semaphore;
use super::{ConnMeta, ConnRegistry, RouteResult};
use crate::transport::middle_proxy::MeResponse;
fn writer_byte_budget() -> Arc<Semaphore> {
Arc::new(Semaphore::new(2049))
}
#[tokio::test]
async fn writer_activity_snapshot_tracks_writer_and_dc_load() {
let registry = ConnRegistry::new();
@@ -14,8 +20,12 @@ async fn writer_activity_snapshot_tracks_writer_and_dc_load() {
let (conn_c, _rx_c) = registry.register().await;
let (writer_tx_a, _writer_rx_a) = tokio::sync::mpsc::channel(8);
let (writer_tx_b, _writer_rx_b) = tokio::sync::mpsc::channel(8);
registry.register_writer(10, writer_tx_a.clone()).await;
registry.register_writer(20, writer_tx_b.clone()).await;
registry
.register_writer(10, writer_tx_a.clone(), writer_byte_budget())
.await;
registry
.register_writer(20, writer_tx_b.clone(), writer_byte_budget())
.await;
let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443);
assert!(
@@ -124,8 +134,12 @@ async fn bind_writer_rebinds_conn_atomically() {
let (conn_id, _rx) = registry.register().await;
let (writer_tx_a, _writer_rx_a) = tokio::sync::mpsc::channel(8);
let (writer_tx_b, _writer_rx_b) = tokio::sync::mpsc::channel(8);
registry.register_writer(10, writer_tx_a).await;
registry.register_writer(20, writer_tx_b).await;
registry
.register_writer(10, writer_tx_a, writer_byte_budget())
.await;
registry
.register_writer(20, writer_tx_b, writer_byte_budget())
.await;
let client_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443);
let first_our_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(1, 1, 1, 1)), 443);
@@ -184,8 +198,12 @@ async fn writer_lost_does_not_drop_rebound_conn() {
let (conn_id, _rx) = registry.register().await;
let (writer_tx_a, _writer_rx_a) = tokio::sync::mpsc::channel(8);
let (writer_tx_b, _writer_rx_b) = tokio::sync::mpsc::channel(8);
registry.register_writer(10, writer_tx_a).await;
registry.register_writer(20, writer_tx_b).await;
registry
.register_writer(10, writer_tx_a, writer_byte_budget())
.await;
registry
.register_writer(20, writer_tx_b, writer_byte_budget())
.await;
let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443);
assert!(
@@ -262,8 +280,12 @@ async fn non_empty_writer_ids_returns_only_writers_with_bound_clients() {
let (conn_id, _rx) = registry.register().await;
let (writer_tx_a, _writer_rx_a) = tokio::sync::mpsc::channel(8);
let (writer_tx_b, _writer_rx_b) = tokio::sync::mpsc::channel(8);
registry.register_writer(10, writer_tx_a).await;
registry.register_writer(20, writer_tx_b).await;
registry
.register_writer(10, writer_tx_a, writer_byte_budget())
.await;
registry
.register_writer(20, writer_tx_b, writer_byte_budget())
.await;
let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443);
assert!(
+15 -4
View File
@@ -1,4 +1,5 @@
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use std::sync::atomic::Ordering;
use std::time::Duration;
@@ -56,7 +57,13 @@ impl ConnRegistry {
}
}
pub async fn register_writer(&self, writer_id: u64, tx: mpsc::Sender<WriterCommand>) {
/// Registers one writer command route and its matching memory budget atomically.
pub async fn register_writer(
&self,
writer_id: u64,
tx: mpsc::Sender<WriterCommand>,
byte_budget: Arc<tokio::sync::Semaphore>,
) {
let mut binding = self.binding.inner.lock().await;
binding
.conns_for_writer
@@ -70,7 +77,9 @@ impl ConnRegistry {
.writer_idle_since_epoch_secs
.entry(writer_id)
.or_insert_with(Self::now_epoch_secs);
self.writers.map.insert(writer_id, tx);
self.writers
.map
.insert(writer_id, super::WriterRoute { tx, byte_budget });
}
/// Unregister connection, returning associated writer_id if any.
@@ -417,7 +426,8 @@ impl ConnRegistry {
.map(|entry| entry.value().clone())?;
Some(ConnWriter {
writer_id,
tx: writer,
tx: writer.tx,
byte_budget: writer.byte_budget,
})
}
@@ -438,7 +448,8 @@ impl ConnRegistry {
Some((
ConnWriter {
writer_id,
tx: writer,
tx: writer.tx,
byte_budget: writer.byte_budget,
},
meta,
))
+329 -65
View File
@@ -6,16 +6,18 @@ use std::sync::Arc;
use std::sync::atomic::Ordering;
use std::time::{Duration, Instant};
use tokio::sync::mpsc;
use tokio::sync::mpsc::error::TrySendError;
use tokio::sync::{OwnedSemaphorePermit, Semaphore, TryAcquireError, mpsc};
use tracing::{debug, warn};
use super::MePool;
use super::codec::{ProxyReqCommand, WriterCommand};
use super::codec::{ProxyReqCommand, WriterBytePermit, WriterCommand};
use super::registry::ConnMeta;
use super::wire::build_proxy_req_payload;
use super::wire::{build_proxy_req_payload, proxy_req_payload_len};
use crate::config::defaults::ME_WRITER_BYTE_PERMIT_UNIT_BYTES;
use crate::config::{MeRouteNoWriterMode, MeWriterPickMode};
use crate::error::{ProxyError, Result};
use crate::stats::Stats;
use crate::stream::PooledBuffer;
use rand::seq::SliceRandom;
@@ -29,6 +31,8 @@ const PICK_PENALTY_WARM: u64 = 200;
const PICK_PENALTY_DRAINING: u64 = 600;
const PICK_PENALTY_STALE: u64 = 300;
const PICK_PENALTY_DEGRADED: u64 = 250;
const RPC_WRITER_FRAME_CAPACITY_OVERHEAD_BYTES: usize = 27;
const LEGACY_PROXY_REQ_SOURCE_CAPACITY_OVERHEAD_BYTES: usize = 128;
mod close;
mod recovery;
@@ -39,34 +43,133 @@ enum WriterCommandReserveError {
TimedOut,
}
enum WriterByteReserveError {
Closed,
TimedOut,
}
fn proxy_tag_array(tag: Option<&[u8]>) -> Option<[u8; 16]> {
tag.and_then(|tag| <[u8; 16]>::try_from(tag).ok())
}
fn proxy_req_payload_from_command(cmd: WriterCommand) -> Option<PooledBuffer> {
fn proxy_req_payload_from_command(
cmd: WriterCommand,
) -> Option<(PooledBuffer, OwnedSemaphorePermit)> {
match cmd {
WriterCommand::ProxyReq(command) => Some(command.payload),
WriterCommand::ProxyReq(command) => Some((command.payload, command._permit)),
_ => None,
}
}
fn payload_permit_from_data_command(cmd: WriterCommand) -> Option<OwnedSemaphorePermit> {
match cmd {
WriterCommand::Data { _permit, .. } => _permit,
_ => None,
}
}
async fn reserve_writer_command_slot(
tx: &mpsc::Sender<WriterCommand>,
wait: Option<Duration>,
deadline: Option<Instant>,
) -> std::result::Result<mpsc::OwnedPermit<WriterCommand>, WriterCommandReserveError> {
let reserve = tx.clone().reserve_owned();
match wait {
Some(wait) => match tokio::time::timeout(wait, reserve).await {
Ok(Ok(permit)) => Ok(permit),
Ok(Err(_)) => Err(WriterCommandReserveError::Closed),
Err(_) => Err(WriterCommandReserveError::TimedOut),
},
match deadline {
Some(deadline) => {
match tokio::time::timeout(deadline.saturating_duration_since(Instant::now()), reserve)
.await
{
Ok(Ok(permit)) => Ok(permit),
Ok(Err(_)) => Err(WriterCommandReserveError::Closed),
Err(_) => Err(WriterCommandReserveError::TimedOut),
}
}
None => reserve.await.map_err(|_| WriterCommandReserveError::Closed),
}
}
fn writer_send_deadline(wait: Option<Duration>) -> Option<Instant> {
wait.map(|wait| Instant::now() + wait)
}
fn writer_resident_permits(
source_capacity: usize,
encoded_payload_len: usize,
) -> Option<(u32, usize)> {
let resident_bytes = source_capacity
.checked_add(encoded_payload_len)?
.checked_add(RPC_WRITER_FRAME_CAPACITY_OVERHEAD_BYTES)?;
let permits = resident_bytes.div_ceil(ME_WRITER_BYTE_PERMIT_UNIT_BYTES);
let permits = u32::try_from(permits).ok()?;
let reserved_bytes = (permits as usize).checked_mul(ME_WRITER_BYTE_PERMIT_UNIT_BYTES)?;
Some((
permits.max(1),
reserved_bytes.max(ME_WRITER_BYTE_PERMIT_UNIT_BYTES),
))
}
fn proxy_req_resident_permits(
source_capacity: usize,
data_len: usize,
proxy_tag: Option<&[u8]>,
proto_flags: u32,
) -> Option<(u32, usize)> {
writer_resident_permits(
source_capacity,
proxy_req_payload_len(data_len, proxy_tag, proto_flags),
)
}
fn try_reserve_writer_bytes(
byte_budget: &Arc<Semaphore>,
permits: u32,
reserved_bytes: usize,
stats: &Arc<Stats>,
) -> std::result::Result<WriterBytePermit, TryAcquireError> {
byte_budget
.clone()
.try_acquire_many_owned(permits)
.map(|permit| WriterBytePermit::new(permit, reserved_bytes, stats.clone()))
}
async fn reserve_writer_bytes(
byte_budget: &Arc<Semaphore>,
permits: u32,
reserved_bytes: usize,
deadline: Option<Instant>,
stats: &Arc<Stats>,
) -> std::result::Result<WriterBytePermit, WriterByteReserveError> {
match try_reserve_writer_bytes(byte_budget, permits, reserved_bytes, stats) {
Ok(permit) => return Ok(permit),
Err(TryAcquireError::Closed) => return Err(WriterByteReserveError::Closed),
Err(TryAcquireError::NoPermits) => {
stats.increment_me_writer_byte_budget_wait_total();
}
}
let acquire = byte_budget.clone().acquire_many_owned(permits);
match deadline {
Some(deadline) => {
match tokio::time::timeout(deadline.saturating_duration_since(Instant::now()), acquire)
.await
{
Ok(Ok(permit)) => Ok(WriterBytePermit::new(permit, reserved_bytes, stats.clone())),
Ok(Err(_)) => Err(WriterByteReserveError::Closed),
Err(_) => {
stats.increment_me_writer_byte_budget_timeout_total();
Err(WriterByteReserveError::TimedOut)
}
}
}
None => acquire
.await
.map(|permit| WriterBytePermit::new(permit, reserved_bytes, stats.clone()))
.map_err(|_| WriterByteReserveError::Closed),
}
}
impl MePool {
/// Send RPC_PROXY_REQ. `tag_override`: per-user ad_tag (from access.user_ad_tags); if None, uses pool default.
/// `payload_permit` keeps optional client byte accounting alive until the writer consumes the command.
pub async fn send_proxy_req(
self: &Arc<Self>,
conn_id: u64,
@@ -76,8 +179,32 @@ impl MePool {
data: &[u8],
proto_flags: u32,
tag_override: Option<&[u8]>,
mut payload_permit: Option<OwnedSemaphorePermit>,
) -> Result<()> {
let tag = tag_override.or(self.proxy_tag.as_deref());
let Some(source_capacity) = data
.len()
.checked_add(LEGACY_PROXY_REQ_SOURCE_CAPACITY_OVERHEAD_BYTES)
else {
self.stats.increment_me_writer_byte_budget_oversize_total();
return Err(ProxyError::Proxy(
"ME writer payload residency calculation overflow".into(),
));
};
let Some((writer_byte_permits, writer_reserved_bytes)) =
proxy_req_resident_permits(source_capacity, data.len(), tag, proto_flags)
else {
self.stats.increment_me_writer_byte_budget_oversize_total();
return Err(ProxyError::Proxy(
"ME writer payload residency calculation overflow".into(),
));
};
if writer_byte_permits as usize > self.writer_lifecycle.writer_byte_budget_permits {
self.stats.increment_me_writer_byte_budget_oversize_total();
return Err(ProxyError::Proxy(
"ME writer payload exceeds configured byte budget".into(),
));
}
let build_routed_payload = |effective_our_addr: SocketAddr| {
(
build_proxy_req_payload(
@@ -118,19 +245,48 @@ impl MePool {
loop {
if let Some((current, current_meta)) = self.registry.get_writer_with_meta(conn_id).await
{
let deadline =
writer_send_deadline(self.route_runtime.me_route_blocking_send_timeout);
let writer_permit = match reserve_writer_bytes(
&current.byte_budget,
writer_byte_permits,
writer_reserved_bytes,
deadline,
&self.stats,
)
.await
{
Ok(permit) => permit,
Err(WriterByteReserveError::TimedOut) => {
self.stats
.increment_me_writer_pick_full_total(self.writer_pick_mode());
return Err(ProxyError::Proxy(
"ME writer byte budget full within blocking send timeout".into(),
));
}
Err(WriterByteReserveError::Closed) => {
warn!(
writer_id = current.writer_id,
"ME writer byte budget closed"
);
self.remove_writer_and_close_clients(current.writer_id)
.await;
continue;
}
};
let (current_payload, _) = build_routed_payload(current_meta.our_addr);
match current.tx.try_send(WriterCommand::Data(current_payload)) {
let command = WriterCommand::Data {
payload: current_payload,
_permit: payload_permit.take(),
writer_permit,
};
match current.tx.try_send(command) {
Ok(()) => {
self.note_hybrid_route_success();
return Ok(());
}
Err(TrySendError::Full(cmd)) => {
match reserve_writer_command_slot(
&current.tx,
self.route_runtime.me_route_blocking_send_timeout,
)
.await
{
match reserve_writer_command_slot(&current.tx, deadline).await {
Ok(permit) => {
permit.send(cmd);
self.note_hybrid_route_success();
@@ -143,14 +299,17 @@ impl MePool {
"ME writer channel full within blocking send timeout".into(),
));
}
Err(WriterCommandReserveError::Closed) => {}
Err(WriterCommandReserveError::Closed) => {
payload_permit = payload_permit_from_data_command(cmd);
}
}
warn!(writer_id = current.writer_id, "ME writer channel closed");
self.remove_writer_and_close_clients(current.writer_id)
.await;
continue;
}
Err(TrySendError::Closed(_)) => {
Err(TrySendError::Closed(cmd)) => {
payload_permit = payload_permit_from_data_command(cmd);
warn!(writer_id = current.writer_id, "ME writer channel closed");
self.remove_writer_and_close_clients(current.writer_id)
.await;
@@ -464,9 +623,31 @@ impl MePool {
if !self.writer_accepts_new_binding(w) {
continue;
}
let (payload, meta) = build_routed_payload(our_addr);
let writer_permit = match try_reserve_writer_bytes(
&w.byte_budget,
writer_byte_permits,
writer_reserved_bytes,
&self.stats,
) {
Ok(permit) => permit,
Err(TryAcquireError::NoPermits) => {
if fallback_blocking_idx.is_none() {
fallback_blocking_idx = Some(idx);
}
continue;
}
Err(TryAcquireError::Closed) => {
self.stats.increment_me_writer_pick_closed_total(pick_mode);
warn!(writer_id = w.id, "ME writer byte budget closed");
self.remove_writer_and_close_clients(w.id).await;
continue;
}
};
match w.tx.clone().try_reserve_owned() {
Ok(permit) => {
// Keep the advertised proxy IP aligned with the selected ME writer source.
let effective_our_addr = SocketAddr::new(w.source_ip, our_addr.port());
let (payload, meta) = build_routed_payload(effective_our_addr);
if !self.registry.bind_writer(conn_id, w.id, meta).await {
debug!(
conn_id,
@@ -477,7 +658,11 @@ impl MePool {
self.remove_writer_and_close_clients(w.id).await;
continue;
}
permit.send(WriterCommand::Data(payload));
permit.send(WriterCommand::Data {
payload,
_permit: payload_permit.take(),
writer_permit,
});
self.stats
.increment_me_writer_pick_success_try_total(pick_mode);
if w.generation < self.current_generation() {
@@ -519,50 +704,71 @@ impl MePool {
}
self.stats
.increment_me_writer_pick_blocking_fallback_total();
let (payload, meta) = build_routed_payload(our_addr);
let reserve_result =
if let Some(timeout) = self.route_runtime.me_route_blocking_send_timeout {
match tokio::time::timeout(timeout, w.tx.clone().reserve_owned()).await {
Ok(result) => result,
Err(_) => {
self.stats.increment_me_writer_pick_full_total(pick_mode);
continue;
}
}
} else {
w.tx.clone().reserve_owned().await
};
match reserve_result {
Ok(permit) => {
if !self.registry.bind_writer(conn_id, w.id, meta).await {
debug!(
conn_id,
writer_id = w.id,
"ME writer disappeared before fallback bind commit, pruning stale writer"
);
drop(permit);
self.remove_writer_and_close_clients(w.id).await;
continue;
}
permit.send(WriterCommand::Data(payload));
self.stats
.increment_me_writer_pick_success_fallback_total(pick_mode);
if w.generation < self.current_generation() {
self.stats.increment_pool_stale_pick_total();
}
self.note_hybrid_route_success();
return Ok(());
let deadline = writer_send_deadline(self.route_runtime.me_route_blocking_send_timeout);
let writer_permit = match reserve_writer_bytes(
&w.byte_budget,
writer_byte_permits,
writer_reserved_bytes,
deadline,
&self.stats,
)
.await
{
Ok(permit) => permit,
Err(WriterByteReserveError::TimedOut) => {
self.stats.increment_me_writer_pick_full_total(pick_mode);
continue;
}
Err(_) => {
Err(WriterByteReserveError::Closed) => {
self.stats.increment_me_writer_pick_closed_total(pick_mode);
warn!(writer_id = w.id, "ME writer byte budget closed (blocking)");
self.remove_writer_and_close_clients(w.id).await;
continue;
}
};
let permit = match reserve_writer_command_slot(&w.tx, deadline).await {
Ok(permit) => permit,
Err(WriterCommandReserveError::TimedOut) => {
self.stats.increment_me_writer_pick_full_total(pick_mode);
continue;
}
Err(WriterCommandReserveError::Closed) => {
self.stats.increment_me_writer_pick_closed_total(pick_mode);
warn!(writer_id = w.id, "ME writer channel closed (blocking)");
self.remove_writer_and_close_clients(w.id).await;
continue;
}
};
// Keep the advertised proxy IP aligned with the selected ME writer source.
let effective_our_addr = SocketAddr::new(w.source_ip, our_addr.port());
let (payload, meta) = build_routed_payload(effective_our_addr);
if !self.registry.bind_writer(conn_id, w.id, meta).await {
debug!(
conn_id,
writer_id = w.id,
"ME writer disappeared before fallback bind commit, pruning stale writer"
);
drop(permit);
self.remove_writer_and_close_clients(w.id).await;
continue;
}
permit.send(WriterCommand::Data {
payload,
_permit: payload_permit.take(),
writer_permit,
});
self.stats
.increment_me_writer_pick_success_fallback_total(pick_mode);
if w.generation < self.current_generation() {
self.stats.increment_pool_stale_pick_total();
}
self.note_hybrid_route_success();
return Ok(());
}
}
/// Send RPC_PROXY_REQ while keeping the first bound-writer path allocation-light.
/// The client byte permit follows the payload until writer completion or command drop.
pub async fn send_proxy_req_pooled(
self: &Arc<Self>,
conn_id: u64,
@@ -570,12 +776,69 @@ impl MePool {
client_addr: SocketAddr,
our_addr: SocketAddr,
payload: PooledBuffer,
_permit: OwnedSemaphorePermit,
proto_flags: u32,
tag_override: Option<[u8; 16]>,
) -> Result<()> {
let tag = tag_override.or_else(|| proxy_tag_array(self.proxy_tag.as_deref()));
let Some((writer_byte_permits, writer_reserved_bytes)) = proxy_req_resident_permits(
payload.capacity(),
payload.len(),
tag.as_ref().map(|tag| tag.as_slice()),
proto_flags,
) else {
self.stats.increment_me_writer_byte_budget_oversize_total();
return Err(ProxyError::Proxy(
"ME writer payload residency calculation overflow".into(),
));
};
if writer_byte_permits as usize > self.writer_lifecycle.writer_byte_budget_permits {
self.stats.increment_me_writer_byte_budget_oversize_total();
return Err(ProxyError::Proxy(
"ME writer payload exceeds configured byte budget".into(),
));
}
if let Some((current, current_meta)) = self.registry.get_writer_with_meta(conn_id).await {
let deadline = writer_send_deadline(self.route_runtime.me_route_blocking_send_timeout);
let writer_permit = match reserve_writer_bytes(
&current.byte_budget,
writer_byte_permits,
writer_reserved_bytes,
deadline,
&self.stats,
)
.await
{
Ok(permit) => permit,
Err(WriterByteReserveError::TimedOut) => {
self.stats
.increment_me_writer_pick_full_total(self.writer_pick_mode());
return Err(ProxyError::Proxy(
"ME writer byte budget full within blocking send timeout".into(),
));
}
Err(WriterByteReserveError::Closed) => {
warn!(
writer_id = current.writer_id,
"ME writer byte budget closed"
);
self.remove_writer_and_close_clients(current.writer_id)
.await;
return self
.send_proxy_req(
conn_id,
target_dc,
client_addr,
our_addr,
payload.as_ref(),
proto_flags,
tag.as_ref().map(|tag| tag.as_slice()),
Some(_permit),
)
.await;
}
};
let command = WriterCommand::ProxyReq(ProxyReqCommand {
conn_id,
client_addr,
@@ -583,6 +846,8 @@ impl MePool {
proto_flags,
proxy_tag: tag,
payload,
_permit,
writer_permit,
});
match current.tx.try_send(command) {
Ok(()) => {
@@ -590,12 +855,7 @@ impl MePool {
return Ok(());
}
Err(TrySendError::Full(cmd)) => {
match reserve_writer_command_slot(
&current.tx,
self.route_runtime.me_route_blocking_send_timeout,
)
.await
{
match reserve_writer_command_slot(&current.tx, deadline).await {
Ok(permit) => {
permit.send(cmd);
self.note_hybrid_route_success();
@@ -609,7 +869,8 @@ impl MePool {
));
}
Err(WriterCommandReserveError::Closed) => {
let Some(payload) = proxy_req_payload_from_command(cmd) else {
let Some((payload, _permit)) = proxy_req_payload_from_command(cmd)
else {
return Err(ProxyError::Proxy(
"ME writer rejected unexpected command type".into(),
));
@@ -626,13 +887,14 @@ impl MePool {
payload.as_ref(),
proto_flags,
tag.as_ref().map(|tag| tag.as_slice()),
Some(_permit),
)
.await;
}
}
}
Err(TrySendError::Closed(cmd)) => {
let Some(payload) = proxy_req_payload_from_command(cmd) else {
let Some((payload, _permit)) = proxy_req_payload_from_command(cmd) else {
return Err(ProxyError::Proxy(
"ME writer rejected unexpected command type".into(),
));
@@ -649,6 +911,7 @@ impl MePool {
payload.as_ref(),
proto_flags,
tag.as_ref().map(|tag| tag.as_slice()),
Some(_permit),
)
.await;
}
@@ -663,6 +926,7 @@ impl MePool {
payload.as_ref(),
proto_flags,
tag.as_ref().map(|tag| tag.as_slice()),
Some(_permit),
)
.await
}
+12 -6
View File
@@ -10,7 +10,7 @@ use crate::protocol::constants::{RPC_CLOSE_CONN_U32, RPC_CLOSE_EXT_U32};
use super::super::MePool;
use super::super::codec::{WriterCommand, build_control_payload};
use super::{WriterCommandReserveError, reserve_writer_command_slot};
use super::{WriterCommandReserveError, reserve_writer_command_slot, writer_send_deadline};
const ME_CLOSE_SIGNAL_SEND_TIMEOUT: Duration = Duration::from_millis(50);
@@ -22,8 +22,11 @@ impl MePool {
match w.tx.try_send(WriterCommand::ControlAndFlush(payload)) {
Ok(()) => {}
Err(TrySendError::Full(cmd)) => {
match reserve_writer_command_slot(&w.tx, Some(ME_CLOSE_SIGNAL_SEND_TIMEOUT))
.await
match reserve_writer_command_slot(
&w.tx,
writer_send_deadline(Some(ME_CLOSE_SIGNAL_SEND_TIMEOUT)),
)
.await
{
Ok(permit) => {
permit.send(cmd);
@@ -63,9 +66,12 @@ impl MePool {
match w.tx.try_send(WriterCommand::ControlAndFlush(payload)) {
Ok(()) => {}
Err(TrySendError::Full(cmd)) => {
let _ = reserve_writer_command_slot(&w.tx, Some(ME_CLOSE_SIGNAL_SEND_TIMEOUT))
.await
.map(|permit| permit.send(cmd));
let _ = reserve_writer_command_slot(
&w.tx,
writer_send_deadline(Some(ME_CLOSE_SIGNAL_SEND_TIMEOUT)),
)
.await
.map(|permit| permit.send(cmd));
}
Err(TrySendError::Closed(_)) => {
debug!(conn_id, "ME close_conn skipped: writer channel closed");
@@ -105,6 +105,7 @@ async fn make_pool(
general.me_writer_pick_sample_size,
MeSocksKdfPolicy::default(),
general.me_writer_cmd_channel_capacity,
general.me_writer_byte_budget_bytes,
general.me_route_channel_capacity,
general.me_route_backpressure_enabled,
general.me_route_fairshare_enabled,
@@ -134,6 +135,7 @@ async fn insert_draining_writer(
drain_deadline_epoch_secs: u64,
) {
let (tx, _writer_rx) = mpsc::channel::<WriterCommand>(8);
let byte_budget = pool.new_writer_byte_budget();
let writer = MeWriter {
id: writer_id,
addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 6000 + writer_id as u16),
@@ -143,6 +145,7 @@ async fn insert_draining_writer(
contour: Arc::new(AtomicU8::new(WriterContour::Draining.as_u8())),
created_at: Instant::now() - Duration::from_secs(writer_id),
tx: tx.clone(),
byte_budget: byte_budget.clone(),
cancel: CancellationToken::new(),
degraded: Arc::new(AtomicBool::new(false)),
rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)),
@@ -153,7 +156,9 @@ async fn insert_draining_writer(
};
pool.writers.write().await.push(writer);
pool.registry.register_writer(writer_id, tx).await;
pool.registry
.register_writer(writer_id, tx, byte_budget)
.await;
pool.conn_count.fetch_add(1, Ordering::Relaxed);
for idx in 0..bound_clients {
@@ -103,6 +103,7 @@ async fn make_pool(
general.me_writer_pick_sample_size,
MeSocksKdfPolicy::default(),
general.me_writer_cmd_channel_capacity,
general.me_writer_byte_budget_bytes,
general.me_route_channel_capacity,
general.me_route_backpressure_enabled,
general.me_route_fairshare_enabled,
@@ -131,6 +132,7 @@ async fn insert_draining_writer(
drain_deadline_epoch_secs: u64,
) {
let (tx, _writer_rx) = mpsc::channel::<WriterCommand>(8);
let byte_budget = pool.new_writer_byte_budget();
let writer = MeWriter {
id: writer_id,
addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 5500 + writer_id as u16),
@@ -140,6 +142,7 @@ async fn insert_draining_writer(
contour: Arc::new(AtomicU8::new(WriterContour::Draining.as_u8())),
created_at: Instant::now() - Duration::from_secs(writer_id),
tx: tx.clone(),
byte_budget: byte_budget.clone(),
cancel: CancellationToken::new(),
degraded: Arc::new(AtomicBool::new(false)),
rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)),
@@ -149,7 +152,9 @@ async fn insert_draining_writer(
allow_drain_fallback: Arc::new(AtomicBool::new(false)),
};
pool.writers.write().await.push(writer);
pool.registry.register_writer(writer_id, tx).await;
pool.registry
.register_writer(writer_id, tx, byte_budget)
.await;
pool.conn_count.fetch_add(1, Ordering::Relaxed);
for idx in 0..bound_clients {
let (conn_id, _rx) = pool.registry.register().await;
@@ -98,6 +98,7 @@ async fn make_pool(me_pool_drain_threshold: u64) -> Arc<MePool> {
general.me_writer_pick_sample_size,
MeSocksKdfPolicy::default(),
general.me_writer_cmd_channel_capacity,
general.me_writer_byte_budget_bytes,
general.me_route_channel_capacity,
general.me_route_backpressure_enabled,
general.me_route_fairshare_enabled,
@@ -126,6 +127,7 @@ async fn insert_draining_writer(
) -> Vec<u64> {
let mut conn_ids = Vec::with_capacity(bound_clients);
let (tx, _writer_rx) = mpsc::channel::<WriterCommand>(8);
let byte_budget = pool.new_writer_byte_budget();
let writer = MeWriter {
id: writer_id,
addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 4500 + writer_id as u16),
@@ -135,6 +137,7 @@ async fn insert_draining_writer(
contour: Arc::new(AtomicU8::new(WriterContour::Draining.as_u8())),
created_at: Instant::now() - Duration::from_secs(writer_id),
tx: tx.clone(),
byte_budget: byte_budget.clone(),
cancel: CancellationToken::new(),
degraded: Arc::new(AtomicBool::new(false)),
rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)),
@@ -144,7 +147,9 @@ async fn insert_draining_writer(
allow_drain_fallback: Arc::new(AtomicBool::new(false)),
};
pool.writers.write().await.push(writer);
pool.registry.register_writer(writer_id, tx).await;
pool.registry
.register_writer(writer_id, tx, byte_budget)
.await;
pool.conn_count.fetch_add(1, Ordering::Relaxed);
for idx in 0..bound_clients {
let (conn_id, _rx) = pool.registry.register().await;
@@ -87,6 +87,7 @@ async fn make_pool() -> Arc<MePool> {
general.me_writer_pick_sample_size,
MeSocksKdfPolicy::default(),
general.me_writer_cmd_channel_capacity,
general.me_writer_byte_budget_bytes,
general.me_route_channel_capacity,
general.me_route_backpressure_enabled,
general.me_route_fairshare_enabled,
@@ -92,6 +92,7 @@ async fn make_pool() -> Arc<MePool> {
general.me_writer_pick_sample_size,
MeSocksKdfPolicy::default(),
general.me_writer_cmd_channel_capacity,
general.me_writer_byte_budget_bytes,
general.me_route_channel_capacity,
general.me_route_backpressure_enabled,
general.me_route_fairshare_enabled,
@@ -120,6 +121,7 @@ async fn insert_writer(
created_at: Instant,
) {
let (tx, _rx) = mpsc::channel::<WriterCommand>(8);
let byte_budget = pool.new_writer_byte_budget();
let contour = if draining {
WriterContour::Draining
} else {
@@ -134,6 +136,7 @@ async fn insert_writer(
contour: Arc::new(AtomicU8::new(contour.as_u8())),
created_at,
tx: tx.clone(),
byte_budget: byte_budget.clone(),
cancel: CancellationToken::new(),
degraded: Arc::new(AtomicBool::new(false)),
rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)),
@@ -144,7 +147,9 @@ async fn insert_writer(
};
pool.writers.write().await.push(writer);
pool.registry.register_writer(writer_id, tx).await;
pool.registry
.register_writer(writer_id, tx, byte_budget)
.await;
pool.conn_count.fetch_add(1, Ordering::Relaxed);
}
@@ -98,6 +98,7 @@ async fn make_pool() -> (Arc<MePool>, Arc<SecureRandom>) {
general.me_writer_pick_sample_size,
MeSocksKdfPolicy::default(),
general.me_writer_cmd_channel_capacity,
general.me_writer_byte_budget_bytes,
general.me_route_channel_capacity,
general.me_route_backpressure_enabled,
general.me_route_fairshare_enabled,
@@ -127,6 +128,7 @@ async fn insert_writer(
register_in_registry: bool,
) -> mpsc::Receiver<WriterCommand> {
let (tx, rx) = mpsc::channel::<WriterCommand>(8);
let byte_budget = pool.new_writer_byte_budget();
let writer = MeWriter {
id: writer_id,
addr,
@@ -136,6 +138,7 @@ async fn insert_writer(
contour: Arc::new(AtomicU8::new(WriterContour::Active.as_u8())),
created_at: Instant::now(),
tx: tx.clone(),
byte_budget: byte_budget.clone(),
cancel: CancellationToken::new(),
degraded: Arc::new(AtomicBool::new(false)),
rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)),
@@ -154,7 +157,9 @@ async fn insert_writer(
}
pool.rebuild_endpoint_dc_map().await;
if register_in_registry {
pool.registry.register_writer(writer_id, tx).await;
pool.registry
.register_writer(writer_id, tx, byte_budget)
.await;
}
rx
}
@@ -165,7 +170,7 @@ async fn recv_data_count(rx: &mut mpsc::Receiver<WriterCommand>, budget: Duratio
while Instant::now().duration_since(start) < budget {
let remaining = budget.saturating_sub(Instant::now().duration_since(start));
match tokio::time::timeout(remaining.min(Duration::from_millis(10)), rx.recv()).await {
Ok(Some(WriterCommand::Data(_))) => data_count += 1,
Ok(Some(WriterCommand::Data { .. })) => data_count += 1,
Ok(Some(WriterCommand::DataAndFlush(_))) => data_count += 1,
Ok(Some(WriterCommand::ProxyReq(_))) => data_count += 1,
Ok(Some(WriterCommand::ControlAndFlush(_))) => data_count += 1,
@@ -185,7 +190,7 @@ async fn recv_first_data_payload(
while Instant::now().duration_since(start) < budget {
let remaining = budget.saturating_sub(Instant::now().duration_since(start));
match tokio::time::timeout(remaining.min(Duration::from_millis(10)), rx.recv()).await {
Ok(Some(WriterCommand::Data(payload))) => return Some(payload.to_vec()),
Ok(Some(WriterCommand::Data { payload, .. })) => return Some(payload.to_vec()),
Ok(Some(WriterCommand::DataAndFlush(payload))) => return Some(payload.to_vec()),
Ok(Some(_)) => {}
Ok(None) => break,
@@ -240,6 +245,7 @@ async fn send_proxy_req_does_not_replay_when_first_bind_commit_fails() {
b"hello",
0,
None,
None,
)
.await;
@@ -299,6 +305,7 @@ async fn send_proxy_req_prunes_iterative_stale_bind_failures_without_data_replay
b"storm",
0,
None,
None,
)
.await;
@@ -323,7 +330,7 @@ async fn send_proxy_req_prunes_iterative_stale_bind_failures_without_data_replay
}
#[tokio::test]
async fn send_proxy_req_preserves_client_facing_our_addr_when_writer_source_ip_differs() {
async fn send_proxy_req_uses_writer_source_ip_when_advertised_our_addr_differs() {
let (pool, _rng) = make_pool().await;
pool.rr.store(0, Ordering::Relaxed);
@@ -356,6 +363,7 @@ async fn send_proxy_req_preserves_client_facing_our_addr_when_writer_source_ip_d
b"route",
0,
None,
None,
)
.await;
@@ -363,5 +371,69 @@ async fn send_proxy_req_preserves_client_facing_our_addr_when_writer_source_ip_d
let payload = recv_first_data_payload(&mut live_rx, Duration::from_millis(50))
.await
.expect("writer must receive routed payload");
assert_eq!(proxy_req_our_addr_from_payload(&payload), our_addr);
assert_eq!(
proxy_req_our_addr_from_payload(&payload),
SocketAddr::new(IpAddr::V4(Ipv4Addr::new(203, 0, 113, 31)), our_addr.port())
);
}
#[tokio::test]
async fn send_proxy_req_blocking_fallback_uses_writer_source_ip() {
let (pool, _rng) = make_pool().await;
pool.rr.store(0, Ordering::Relaxed);
let (conn_id, _rx) = pool.registry.register().await;
let mut live_rx = insert_writer(
&pool,
32,
2,
SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 2, 32)), 443),
true,
)
.await;
let source_ip = IpAddr::V4(Ipv4Addr::new(203, 0, 113, 32));
let tx = {
let mut writers = pool.writers.write().await;
let writer = writers
.iter_mut()
.find(|writer| writer.id == 32)
.expect("test writer must exist");
writer.source_ip = source_ip;
writer.tx.clone()
};
for _ in 0..8 {
tx.try_send(WriterCommand::Close)
.expect("test writer channel must accept preload");
}
let our_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(198, 51, 100, 8)), 9443);
let pool_for_send = pool.clone();
let send_task = tokio::spawn(async move {
pool_for_send
.send_proxy_req(
conn_id,
2,
SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 0, 0, 8)), 30003),
our_addr,
b"blocking",
0,
None,
None,
)
.await
});
tokio::time::sleep(Duration::from_millis(10)).await;
assert!(matches!(live_rx.recv().await, Some(WriterCommand::Close)));
let result = send_task.await.expect("send task must not panic");
assert!(result.is_ok());
let payload = recv_first_data_payload(&mut live_rx, Duration::from_millis(50))
.await
.expect("writer must receive blocking fallback payload");
assert_eq!(
proxy_req_our_addr_from_payload(&payload),
SocketAddr::new(source_ip, our_addr.port())
);
}