astelm/telemt
1
0
Fork 0
mirror of https://github.com/telemt/telemt.git synced 2026-09-13 22:14:08 +03:00
Code Issues Packages Projects Releases Wiki Activity

WEB

Browse Source 
Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
Co-Authored-By: John Preston <17900494+john-preston@users.noreply.github.com>
This commit is contained in:
Alexey
2026-08-23 03:12:11 +03:00
parent 8dbd24b11b
commit 1029703c2c
58 changed files with 7460 additions and 2646 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/proxy/authenticated.rs
+325
View File
@@ -0,0 +1,325 @@
use std::net::{IpAddr, SocketAddr};
use std::sync::Arc;
use tokio::io::{AsyncRead, AsyncWrite};
use tokio::sync::RwLock;
use tracing::warn;
use crate::config::ProxyConfig;
use crate::crypto::SecureRandom;
use crate::error::{ProxyError, Result};
use crate::ip_tracker::UserIpTracker;
use crate::proxy::direct_relay::handle_via_direct_with_shared_and_conntrack;
use crate::proxy::handshake::HandshakeSuccess;
use crate::proxy::middle_relay::{
handle_via_middle_proxy, handle_via_middle_proxy_with_conntrack,
};
use crate::proxy::route_mode::{RelayRouteMode, RouteRuntimeController};
use crate::proxy::shared_state::{ConntrackClosePolicy, ProxySharedState};
use crate::stats::Stats;
use crate::stream::{BufferPool, CryptoReader, CryptoWriter};
use crate::transport::UpstreamManager;
use crate::transport::middle_proxy::MePool;
/// Immutable dependency snapshot pinned by one authenticated client stream.
#[derive(Clone)]
pub(crate) struct ClientRuntimeDeps {
/// Immutable effective configuration pinned for this stream.
pub(crate) config: Arc<ProxyConfig>,
/// Process statistics registry.
pub(crate) stats: Arc<Stats>,
/// Direct Telegram upstream connector.
pub(crate) upstream_manager: Arc<UpstreamManager>,
/// Shared relay buffer pool.
pub(crate) buffer_pool: Arc<BufferPool>,
/// Process cryptographic random source.
pub(crate) rng: Arc<SecureRandom>,
/// Startup Middle-End pool, when immediately available.
pub(crate) me_pool: Option<Arc<MePool>>,
/// Hot-swappable Middle-End pool holder.
pub(crate) me_pool_runtime: Option<Arc<RwLock<Option<Arc<MePool>>>>>,
/// Route-mode controller shared by active generations.
pub(crate) route_runtime: Arc<RouteRuntimeController>,
/// Per-user source-IP admission tracker.
pub(crate) ip_tracker: Arc<UserIpTracker>,
/// Process-shared admission and relay coordination state.
pub(crate) shared: Arc<ProxySharedState>,
}
/// Runs admission and relay after a successful MTProxy handshake.
pub(crate) async fn run_authenticated<R, W>(
client_reader: CryptoReader<R>,
client_writer: CryptoWriter<W>,
success: HandshakeSuccess,
deps: ClientRuntimeDeps,
local_addr: SocketAddr,
peer_addr: SocketAddr,
conntrack_close_policy: ConntrackClosePolicy,
) -> Result<()>
where
R: AsyncRead + Unpin + Send + 'static,
W: AsyncWrite + Unpin + Send + 'static,
{
let user = success.user.clone();
if !deps.shared.is_user_enabled(&user) {
warn!(user = %user, "Disabled user rejected");
return Err(ProxyError::UserDisabled { user });
}
let user_reservation = acquire_user_connection_reservation(
&user,
&deps.config,
Arc::clone(&deps.stats),
peer_addr,
Arc::clone(&deps.ip_tracker),
)
.await
.map_err(|error| {
warn!(user = %user, error = %error, "User admission check failed");
error
})?;
let route_snapshot = deps.route_runtime.snapshot();
let session_id = deps.rng.u64();
let user_session = deps.shared.register_user_session(&user, session_id);
let session_cancel = user_session.token();
let selected_me_pool = if deps.config.general.use_middle_proxy
&& matches!(route_snapshot.mode, RelayRouteMode::Middle)
{
if let Some(pool) = &deps.me_pool {
Some(Arc::clone(pool))
} else if let Some(pool_runtime) = &deps.me_pool_runtime {
pool_runtime.read().await.clone()
} else {
None
}
} else {
None
};
let relay_result = if deps.config.general.use_middle_proxy
&& matches!(route_snapshot.mode, RelayRouteMode::Middle)
{
if let Some(pool) = selected_me_pool {
if conntrack_close_policy == ConntrackClosePolicy::Publish {
handle_via_middle_proxy(
client_reader,
client_writer,
success,
pool,
Arc::clone(&deps.stats),
Arc::clone(&deps.config),
Arc::clone(&deps.buffer_pool),
local_addr,
Arc::clone(&deps.rng),
deps.route_runtime.subscribe(),
route_snapshot,
session_id,
session_cancel.clone(),
Arc::clone(&deps.shared),
)
.await
} else {
handle_via_middle_proxy_with_conntrack(
client_reader,
client_writer,
success,
pool,
Arc::clone(&deps.stats),
Arc::clone(&deps.config),
Arc::clone(&deps.buffer_pool),
local_addr,
Arc::clone(&deps.rng),
deps.route_runtime.subscribe(),
route_snapshot,
session_id,
session_cancel.clone(),
Arc::clone(&deps.shared),
ConntrackClosePolicy::Suppress,
)
.await
}
} else {
warn!("use_middle_proxy=true but MePool not initialized, falling back to direct");
run_direct(
client_reader,
client_writer,
success,
&deps,
route_snapshot,
session_id,
local_addr,
session_cancel.clone(),
conntrack_close_policy,
)
.await
}
} else {
run_direct(
client_reader,
client_writer,
success,
&deps,
route_snapshot,
session_id,
local_addr,
session_cancel,
conntrack_close_policy,
)
.await
};
user_reservation.release().await;
relay_result
}
async fn run_direct<R, W>(
client_reader: CryptoReader<R>,
client_writer: CryptoWriter<W>,
success: HandshakeSuccess,
deps: &ClientRuntimeDeps,
route_snapshot: crate::proxy::route_mode::RouteCutoverState,
session_id: u64,
local_addr: SocketAddr,
session_cancel: tokio_util::sync::CancellationToken,
conntrack_close_policy: ConntrackClosePolicy,
) -> Result<()>
where
R: AsyncRead + Unpin + Send + 'static,
W: AsyncWrite + Unpin + Send + 'static,
{
handle_via_direct_with_shared_and_conntrack(
client_reader,
client_writer,
success,
Arc::clone(&deps.upstream_manager),
Arc::clone(&deps.stats),
Arc::clone(&deps.config),
Arc::clone(&deps.buffer_pool),
Arc::clone(&deps.rng),
deps.route_runtime.subscribe(),
route_snapshot,
session_id,
local_addr,
session_cancel,
Arc::clone(&deps.shared),
conntrack_close_policy,
)
.await
}
#[must_use = "the reservation owns user and IP admission until release or drop"]
/// Owns one authenticated user's connection and source-IP admission slots.
pub(crate) struct UserConnectionReservation {
stats: Arc<Stats>,
ip_tracker: Arc<UserIpTracker>,
user: String,
ip: IpAddr,
tracks_ip: bool,
active: bool,
}
impl UserConnectionReservation {
/// Creates an active reservation after both admission counters were acquired.
pub(crate) fn new(
stats: Arc<Stats>,
ip_tracker: Arc<UserIpTracker>,
user: String,
ip: IpAddr,
tracks_ip: bool,
) -> Self {
Self {
stats,
ip_tracker,
user,
ip,
tracks_ip,
active: true,
}
}
/// Releases both admission counters through the asynchronous cleanup path.
pub(crate) async fn release(mut self) {
if !self.active {
return;
}
self.active = false;
if self.tracks_ip {
self.ip_tracker.remove_ip(&self.user, self.ip).await;
}
self.stats.decrement_user_curr_connects(&self.user);
}
}
impl Drop for UserConnectionReservation {
fn drop(&mut self) {
if !self.active {
return;
}
self.active = false;
self.stats.increment_session_drop_fallback_total();
self.stats.decrement_user_curr_connects(&self.user);
if self.tracks_ip {
self.ip_tracker.enqueue_cleanup(self.user.clone(), self.ip);
}
}
}
/// Applies user quota, connection, and source-IP admission atomically.
pub(crate) async fn acquire_user_connection_reservation(
user: &str,
config: &ProxyConfig,
stats: Arc<Stats>,
peer_addr: SocketAddr,
ip_tracker: Arc<UserIpTracker>,
) -> Result<UserConnectionReservation> {
if let Some(expiration) = config.access.user_expirations.get(user)
&& chrono::Utc::now() > *expiration
{
return Err(ProxyError::UserExpired {
user: user.to_string(),
});
}
if let Some(quota) = config.access.user_data_quota.get(user)
&& stats.get_user_quota_used(user) >= *quota
{
return Err(ProxyError::DataQuotaExceeded {
user: user.to_string(),
});
}
let limit = config
.access
.user_max_tcp_conns
.get(user)
.copied()
.filter(|limit| *limit > 0)
.or((config.access.user_max_tcp_conns_global_each > 0)
.then_some(config.access.user_max_tcp_conns_global_each))
.map(|value| value as u64);
if !stats.try_acquire_user_curr_connects(user, limit) {
return Err(ProxyError::ConnectionLimitExceeded {
user: user.to_string(),
});
}
if let Err(reason) = ip_tracker.check_and_add(user, peer_addr.ip()).await {
stats.decrement_user_curr_connects(user);
warn!(
user = %user,
ip = %peer_addr.ip(),
reason = %reason,
"IP limit exceeded"
);
return Err(ProxyError::ConnectionLimitExceeded {
user: user.to_string(),
});
}
Ok(UserConnectionReservation::new(
stats,
ip_tracker,
user.to_string(),
peer_addr.ip(),
true,
))
}
src/proxy/client.rs
+29 -225
View File
@@ -26,72 +26,6 @@ enum HandshakeOutcome {
NeedsMasking(PostHandshakeFuture),
}
#[must_use = "UserConnectionReservation must be kept alive to retain user/IP reservation until release or drop"]
struct UserConnectionReservation {
stats: Arc<Stats>,
ip_tracker: Arc<UserIpTracker>,
user: String,
ip: IpAddr,
tracks_ip: bool,
state: SessionReservationState,
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum SessionReservationState {
Active,
Released,
}
impl UserConnectionReservation {
fn new(
stats: Arc<Stats>,
ip_tracker: Arc<UserIpTracker>,
user: String,
ip: IpAddr,
tracks_ip: bool,
) -> Self {
Self {
stats,
ip_tracker,
user,
ip,
tracks_ip,
state: SessionReservationState::Active,
}
}
fn mark_released(&mut self) -> bool {
if self.state != SessionReservationState::Active {
return false;
}
self.state = SessionReservationState::Released;
true
}
async fn release(mut self) {
if !self.mark_released() {
return;
}
if self.tracks_ip {
self.ip_tracker.remove_ip(&self.user, self.ip).await;
}
self.stats.decrement_user_curr_connects(&self.user);
}
}
impl Drop for UserConnectionReservation {
fn drop(&mut self) {
if !self.mark_released() {
return;
}
self.stats.increment_session_drop_fallback_total();
self.stats.decrement_user_curr_connects(&self.user);
if self.tracks_ip {
self.ip_tracker.enqueue_cleanup(self.user.clone(), self.ip);
}
}
}
use crate::config::ProxyConfig;
use crate::crypto::SecureRandom;
use crate::error::{HandshakeResult, ProxyError, Result, StreamError};
@@ -107,7 +41,11 @@ use crate::transport::middle_proxy::MePool;
use crate::transport::socket::normalize_ip;
use crate::transport::{UpstreamManager, configure_client_socket, parse_proxy_protocol};
use crate::proxy::direct_relay::handle_via_direct_with_shared;
use crate::proxy::authenticated::{ClientRuntimeDeps, run_authenticated};
#[cfg(test)]
use crate::proxy::authenticated::{
UserConnectionReservation, acquire_user_connection_reservation,
};
use crate::proxy::handshake::{
HandshakeSuccess, TlsResponseWriteOptions, handle_mtproto_handshake_with_shared,
handle_tls_handshake_with_shared, handle_tls_handshake_with_shared_and_options,
@@ -115,9 +53,10 @@ use crate::proxy::handshake::{
#[cfg(test)]
use crate::proxy::handshake::{handle_mtproto_handshake, handle_tls_handshake};
use crate::proxy::masking::handle_bad_client_with_shared;
use crate::proxy::middle_relay::handle_via_middle_proxy;
use crate::proxy::route_mode::{RelayRouteMode, RouteRuntimeController};
use crate::proxy::shared_state::ProxySharedState;
use crate::proxy::route_mode::RouteRuntimeController;
#[cfg(test)]
use crate::proxy::route_mode::RelayRouteMode;
use crate::proxy::shared_state::{ConntrackClosePolicy, ProxySharedState};
fn beobachten_ttl(config: &ProxyConfig) -> Duration {
const BEOBACHTEN_TTL_MAX_MINUTES: u64 = 24 * 60;
@@ -1688,112 +1627,30 @@ impl RunningClientHandler {
R: AsyncRead + Unpin + Send + 'static,
W: AsyncWrite + Unpin + Send + 'static,
{
let user = success.user.clone();
if !shared.is_user_enabled(&user) {
warn!(user = %user, "Disabled user rejected");
return Err(ProxyError::UserDisabled { user });
}
let user_limit_reservation = match Self::acquire_user_connection_reservation_static(
&user,
&config,
stats.clone(),
peer_addr,
ip_tracker,
)
.await
{
Ok(reservation) => reservation,
Err(e) => {
warn!(user = %user, error = %e, "User admission check failed");
return Err(e);
}
};
let route_snapshot = route_runtime.snapshot();
let session_id = rng.u64();
let _user_session = shared.register_user_session(&user, session_id);
let session_cancel = _user_session.token();
let selected_me_pool = if config.general.use_middle_proxy
&& matches!(route_snapshot.mode, RelayRouteMode::Middle)
{
if let Some(ref pool) = me_pool {
Some(pool.clone())
} else if let Some(pool_runtime) = me_pool_runtime.as_ref() {
pool_runtime.read().await.clone()
} else {
None
}
} else {
None
};
let relay_result = if config.general.use_middle_proxy
&& matches!(route_snapshot.mode, RelayRouteMode::Middle)
{
if let Some(pool) = selected_me_pool {
handle_via_middle_proxy(
client_reader,
client_writer,
success,
pool,
stats.clone(),
config,
buffer_pool,
local_addr,
rng,
route_runtime.subscribe(),
route_snapshot,
session_id,
session_cancel.clone(),
shared.clone(),
)
.await
} else {
warn!("use_middle_proxy=true but MePool not initialized, falling back to direct");
handle_via_direct_with_shared(
client_reader,
client_writer,
success,
upstream_manager,
stats.clone(),
config,
buffer_pool,
rng,
route_runtime.subscribe(),
route_snapshot,
session_id,
local_addr,
session_cancel.clone(),
shared.clone(),
)
.await
}
} else {
// Direct mode (original behavior)
handle_via_direct_with_shared(
client_reader,
client_writer,
success,
upstream_manager,
stats.clone(),
run_authenticated(
client_reader,
client_writer,
success,
ClientRuntimeDeps {
config,
stats,
upstream_manager,
buffer_pool,
rng,
route_runtime.subscribe(),
route_snapshot,
session_id,
local_addr,
session_cancel,
shared.clone(),
)
.await
};
user_limit_reservation.release().await;
relay_result
me_pool,
me_pool_runtime,
route_runtime,
ip_tracker,
shared,
},
local_addr,
peer_addr,
ConntrackClosePolicy::Publish,
)
.await
}
#[cfg(test)]
async fn acquire_user_connection_reservation_static(
user: &str,
config: &ProxyConfig,
@@ -1801,60 +1658,7 @@ impl RunningClientHandler {
peer_addr: SocketAddr,
ip_tracker: Arc<UserIpTracker>,
) -> Result<UserConnectionReservation> {
if let Some(expiration) = config.access.user_expirations.get(user)
&& chrono::Utc::now() > *expiration
{
return Err(ProxyError::UserExpired {
user: user.to_string(),
});
}
if let Some(quota) = config.access.user_data_quota.get(user)
&& stats.get_user_quota_used(user) >= *quota
{
return Err(ProxyError::DataQuotaExceeded {
user: user.to_string(),
});
}
let limit = config
.access
.user_max_tcp_conns
.get(user)
.copied()
.filter(|limit| *limit > 0)
.or((config.access.user_max_tcp_conns_global_each > 0)
.then_some(config.access.user_max_tcp_conns_global_each))
.map(|v| v as u64);
if !stats.try_acquire_user_curr_connects(user, limit) {
return Err(ProxyError::ConnectionLimitExceeded {
user: user.to_string(),
});
}
match ip_tracker.check_and_add(user, peer_addr.ip()).await {
Ok(()) => {}
Err(reason) => {
stats.decrement_user_curr_connects(user);
warn!(
user = %user,
ip = %peer_addr.ip(),
reason = %reason,
"IP limit exceeded"
);
return Err(ProxyError::ConnectionLimitExceeded {
user: user.to_string(),
});
}
}
Ok(UserConnectionReservation::new(
stats,
ip_tracker,
user.to_string(),
peer_addr.ip(),
true,
))
acquire_user_connection_reservation(user, config, stats, peer_addr, ip_tracker).await
}
#[cfg(test)]
src/proxy/direct_relay.rs
+59 -12
View File
@@ -22,7 +22,8 @@ use crate::proxy::route_mode::{
RelayRouteMode, RouteCutoverState, affected_cutover_state, cutover_stagger_delay,
};
use crate::proxy::shared_state::{
ConntrackCloseEvent, ConntrackClosePublishResult, ConntrackCloseReason, ProxySharedState,
ConntrackCloseEvent, ConntrackClosePolicy, ConntrackClosePublishResult, ConntrackCloseReason,
ProxySharedState,
};
use crate::stats::Stats;
use crate::stream::{BufferPool, CryptoReader, CryptoWriter};
@@ -229,6 +230,7 @@ fn unknown_dc_test_lock() -> &'static Mutex<()> {
}
#[allow(dead_code)]
/// Runs Direct relay with standalone cancellation and shared-state defaults.
pub(crate) async fn handle_via_direct<R, W>(
client_reader: CryptoReader<R>,
client_writer: CryptoWriter<W>,
@@ -265,7 +267,49 @@ where
.await
}
/// Runs Direct relay for a kernel-backed TCP client tuple.
pub(crate) async fn handle_via_direct_with_shared<R, W>(
client_reader: CryptoReader<R>,
client_writer: CryptoWriter<W>,
success: HandshakeSuccess,
upstream_manager: Arc<UpstreamManager>,
stats: Arc<Stats>,
config: Arc<ProxyConfig>,
buffer_pool: Arc<BufferPool>,
rng: Arc<SecureRandom>,
route_rx: watch::Receiver<RouteCutoverState>,
route_snapshot: RouteCutoverState,
session_id: u64,
local_addr: SocketAddr,
session_cancel: CancellationToken,
shared: Arc<ProxySharedState>,
) -> Result<()>
where
R: AsyncRead + Unpin + Send + 'static,
W: AsyncWrite + Unpin + Send + 'static,
{
handle_via_direct_with_shared_and_conntrack(
client_reader,
client_writer,
success,
upstream_manager,
stats,
config,
buffer_pool,
rng,
route_rx,
route_snapshot,
session_id,
local_addr,
session_cancel,
shared,
ConntrackClosePolicy::Publish,
)
.await
}
/// Runs Direct relay with explicit kernel-conntrack close publication policy.
pub(crate) async fn handle_via_direct_with_shared_and_conntrack<R, W>(
client_reader: CryptoReader<R>,
client_writer: CryptoWriter<W>,
success: HandshakeSuccess,
@@ -280,6 +324,7 @@ pub(crate) async fn handle_via_direct_with_shared<R, W>(
local_addr: SocketAddr,
session_cancel: CancellationToken,
shared: Arc<ProxySharedState>,
conntrack_close_policy: ConntrackClosePolicy,
) -> Result<()>
where
R: AsyncRead + Unpin + Send + 'static,
@@ -407,17 +452,19 @@ where
pool_snapshot.allocated.saturating_sub(pool_snapshot.pooled),
);
let close_reason = classify_conntrack_close_reason(&relay_result);
let publish_result = shared.publish_conntrack_close_event(ConntrackCloseEvent {
src: success.peer,
dst: local_addr,
reason: close_reason,
});
if !matches!(
publish_result,
ConntrackClosePublishResult::Sent | ConntrackClosePublishResult::Disabled
) {
stats.increment_conntrack_close_event_drop_total();
if conntrack_close_policy == ConntrackClosePolicy::Publish {
let close_reason = classify_conntrack_close_reason(&relay_result);
let publish_result = shared.publish_conntrack_close_event(ConntrackCloseEvent {
src: success.peer,
dst: local_addr,
reason: close_reason,
});
if !matches!(
publish_result,
ConntrackClosePublishResult::Sent | ConntrackClosePublishResult::Disabled
) {
stats.increment_conntrack_close_event_drop_total();
}
}
relay_result
src/proxy/handshake.rs
+2 -1
View File
@@ -20,7 +20,7 @@ use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
use tracing::{debug, info, trace, warn};
use zeroize::{Zeroize, Zeroizing};
use crate::config::{ProxyConfig, UnknownSniAction};
use crate::config::{ProxyConfig, UnknownSniAction, WebSecretMode};
use crate::crypto::{AesCtr, SecureRandom, sha256};
use crate::error::{HandshakeResult, ProxyError};
use crate::protocol::constants::*;
@@ -58,6 +58,7 @@ pub(crate) use self::auth_probe::{AuthProbeSaturationState, AuthProbeState};
#[cfg(test)]
pub use self::mtproto::handle_mtproto_handshake;
pub use self::mtproto::handle_mtproto_handshake_with_shared;
pub(crate) use self::mtproto::handle_mtproto_handshake_for_web_user;
#[allow(unused_imports)]
pub use self::nonce::{encrypt_tg_nonce, encrypt_tg_nonce_with_ciphers, generate_tg_nonce};
pub use self::session::HandshakeSuccess;
src/proxy/handshake/auth_candidates.rs
+70 -1
View File
@@ -11,6 +11,12 @@ pub(super) struct MtprotoCandidateValidation {
pub(super) encryptor: AesCtr,
}
#[derive(Clone, Copy)]
pub(super) enum MtprotoModePolicy {
Configured,
Web(WebSecretMode),
}
pub(super) fn sni_hint_hash(sni: &str) -> u64 {
let mut hasher = DefaultHasher::new();
for byte in sni.bytes() {
@@ -146,6 +152,7 @@ pub(super) fn validate_mtproto_secret_candidate(
secret: &[u8; ACCESS_SECRET_BYTES],
config: &ProxyConfig,
is_tls: bool,
mode_policy: MtprotoModePolicy,
) -> Option<MtprotoCandidateValidation> {
let mut dec_key_input = Zeroizing::new(Vec::with_capacity(PREKEY_LEN + secret.len()));
dec_key_input.extend_from_slice(dec_prekey);
@@ -163,7 +170,7 @@ pub(super) fn validate_mtproto_secret_candidate(
decrypted[PROTO_TAG_POS + 3],
];
let proto_tag = ProtoTag::from_bytes(tag_bytes)?;
if !mode_enabled_for_proto(config, proto_tag, is_tls) {
if !mode_enabled_for_proto_with_policy(config, proto_tag, is_tls, mode_policy) {
return None;
}
@@ -267,6 +274,28 @@ pub(super) fn mode_enabled_for_proto(
proto_tag: ProtoTag,
is_tls: bool,
) -> bool {
mode_enabled_for_proto_with_policy(
config,
proto_tag,
is_tls,
MtprotoModePolicy::Configured,
)
}
fn mode_enabled_for_proto_with_policy(
config: &ProxyConfig,
proto_tag: ProtoTag,
is_tls: bool,
policy: MtprotoModePolicy,
) -> bool {
if let MtprotoModePolicy::Web(secret_mode) = policy {
return match secret_mode {
WebSecretMode::Plain => {
matches!(proto_tag, ProtoTag::Intermediate | ProtoTag::Abridged)
}
WebSecretMode::Dd => matches!(proto_tag, ProtoTag::Secure),
};
}
match proto_tag {
ProtoTag::Secure => {
if is_tls {
@@ -279,6 +308,46 @@ pub(super) fn mode_enabled_for_proto(
}
}
#[cfg(test)]
mod web_mode_tests {
use super::*;
#[test]
fn web_secret_mode_isolates_inner_protocol_tags() {
let config = ProxyConfig::default();
assert!(mode_enabled_for_proto_with_policy(
&config,
ProtoTag::Abridged,
false,
MtprotoModePolicy::Web(WebSecretMode::Plain),
));
assert!(mode_enabled_for_proto_with_policy(
&config,
ProtoTag::Intermediate,
false,
MtprotoModePolicy::Web(WebSecretMode::Plain),
));
assert!(!mode_enabled_for_proto_with_policy(
&config,
ProtoTag::Secure,
false,
MtprotoModePolicy::Web(WebSecretMode::Plain),
));
assert!(mode_enabled_for_proto_with_policy(
&config,
ProtoTag::Secure,
false,
MtprotoModePolicy::Web(WebSecretMode::Dd),
));
assert!(!mode_enabled_for_proto_with_policy(
&config,
ProtoTag::Intermediate,
false,
MtprotoModePolicy::Web(WebSecretMode::Dd),
));
}
}
pub(super) fn decode_user_secrets_in(
shared: &ProxySharedState,
config: &ProxyConfig,
src/proxy/handshake/mtproto.rs
+63 -10
View File
@@ -1,6 +1,6 @@
use super::*;
/// Handle MTProto obfuscation handshake
/// Handles an MTProto obfuscation handshake with isolated test state.
#[cfg(test)]
pub async fn handle_mtproto_handshake<R, W>(
handshake: &[u8; HANDSHAKE_LEN],
@@ -26,11 +26,14 @@ where
replay_checker,
is_tls,
preferred_user,
None,
MtprotoModePolicy::Configured,
shared.as_ref(),
)
.await
}
/// Handles an MTProto obfuscation handshake with process-shared defenses.
pub async fn handle_mtproto_handshake_with_shared<R, W>(
handshake: &[u8; HANDSHAKE_LEN],
reader: R,
@@ -55,6 +58,40 @@ where
replay_checker,
is_tls,
preferred_user,
None,
MtprotoModePolicy::Configured,
shared,
)
.await
}
/// Authenticates one WEB logical stream against exactly one user and secret mode.
pub(crate) async fn handle_mtproto_handshake_for_web_user<R, W>(
handshake: &[u8; HANDSHAKE_LEN],
reader: R,
writer: W,
peer: SocketAddr,
config: &ProxyConfig,
replay_checker: &ReplayChecker,
exact_user: &str,
secret_mode: WebSecretMode,
shared: &ProxySharedState,
) -> HandshakeResult<(CryptoReader<R>, CryptoWriter<W>, HandshakeSuccess), R, W>
where
R: AsyncRead + Unpin + Send,
W: AsyncWrite + Unpin + Send,
{
handle_mtproto_handshake_impl(
handshake,
reader,
writer,
peer,
config,
replay_checker,
false,
None,
Some(exact_user),
MtprotoModePolicy::Web(secret_mode),
shared,
)
.await
@@ -69,6 +106,8 @@ async fn handle_mtproto_handshake_impl<R, W>(
replay_checker: &ReplayChecker,
is_tls: bool,
preferred_user: Option<&str>,
exact_user: Option<&str>,
mode_policy: MtprotoModePolicy,
shared: &ProxySharedState,
) -> HandshakeResult<(CryptoReader<R>, CryptoWriter<W>, HandshakeSuccess), R, W>
where
@@ -113,8 +152,11 @@ where
let sticky_ip_hint = sticky_hint_get_by_ip(shared, peer.ip());
let sticky_prefix_hint = sticky_hint_get_by_ip_prefix(shared, peer.ip());
let preferred_user_id = preferred_user.and_then(|user| snapshot.user_id_by_name(user));
let has_hint =
sticky_ip_hint.is_some() || sticky_prefix_hint.is_some() || preferred_user_id.is_some();
let exact_user_id = exact_user.and_then(|user| snapshot.user_id_by_name(user));
let has_hint = sticky_ip_hint.is_some()
|| sticky_prefix_hint.is_some()
|| preferred_user_id.is_some()
|| exact_user_id.is_some();
let overload = auth_probe_saturation_is_throttled_in(shared, Instant::now());
let candidate_budget = budget_for_validation(snapshot.entries().len(), overload, has_hint);
@@ -145,6 +187,7 @@ where
&entry.secret,
config,
is_tls,
mode_policy,
) {
matched_user = entry.user.clone();
matched_user_id = Some($user_id);
@@ -159,20 +202,20 @@ where
}};
}
let mut matched = false;
if let Some(user_id) = sticky_ip_hint {
let mut matched = exact_user_id.is_some_and(|user_id| try_user_id!(user_id));
if exact_user.is_none() && let Some(user_id) = sticky_ip_hint {
matched = try_user_id!(user_id);
}
if !matched && let Some(user_id) = preferred_user_id {
if exact_user.is_none() && !matched && let Some(user_id) = preferred_user_id {
matched = try_user_id!(user_id);
}
if !matched && let Some(user_id) = sticky_prefix_hint {
if exact_user.is_none() && !matched && let Some(user_id) = sticky_prefix_hint {
matched = try_user_id!(user_id);
}
if !matched && !budget_exhausted {
if exact_user.is_none() && !matched && !budget_exhausted {
let ring = &shared.handshake.recent_user_ring;
if !ring.is_empty() {
let next_seq = shared
@@ -197,7 +240,7 @@ where
}
}
if !matched && !budget_exhausted {
if exact_user.is_none() && !matched && !budget_exhausted {
for idx in 0..snapshot.entries().len() {
let Some(user_id) = u32::try_from(idx).ok() else {
break;
@@ -317,7 +360,16 @@ where
success,
));
} else {
let decoded_users = decode_user_secrets_in(shared, config, preferred_user);
let decoded_users = match exact_user {
Some(user) => config
.access
.users
.get(user)
.and_then(|secret| decode_user_secret(shared, user, secret))
.map(|secret| vec![(user.to_string(), secret)])
.unwrap_or_default(),
None => decode_user_secrets_in(shared, config, preferred_user),
};
let mut validation_checks = 0usize;
for (user, secret) in decoded_users {
@@ -337,6 +389,7 @@ where
&secret_arr,
config,
is_tls,
mode_policy,
) else {
continue;
};
src/proxy/middle_relay.rs
+45 -2
View File
@@ -26,7 +26,8 @@ use crate::proxy::route_mode::{
RelayRouteMode, RouteCutoverState, affected_cutover_state, cutover_stagger_delay,
};
use crate::proxy::shared_state::{
ConntrackCloseEvent, ConntrackClosePublishResult, ConntrackCloseReason, ProxySharedState,
ConntrackCloseEvent, ConntrackClosePolicy, ConntrackClosePublishResult, ConntrackCloseReason,
ProxySharedState,
};
use crate::proxy::traffic_limiter::{RateDirection, TrafficLease, next_refill_delay};
use crate::stats::{
@@ -44,7 +45,7 @@ mod session;
pub(crate) use self::desync::DesyncDedupRotationState;
pub(crate) use self::idle::{RelayIdleCandidateRegistry, note_global_relay_pressure};
pub(crate) use self::session::handle_via_middle_proxy;
pub(crate) use self::session::handle_via_middle_proxy_with_conntrack;
use self::c2me::{
C2MeCommand, acquire_c2me_payload_permit, c2me_queued_permit_budget, enqueue_c2me_command_in,
@@ -91,6 +92,47 @@ pub(crate) use self::idle::{
set_relay_pressure_state_for_testing,
};
/// Runs Middle-End relay for a kernel-backed TCP client tuple.
pub(crate) async fn handle_via_middle_proxy<R, W>(
crypto_reader: CryptoReader<R>,
crypto_writer: CryptoWriter<W>,
success: HandshakeSuccess,
me_pool: Arc<MePool>,
stats: Arc<Stats>,
config: Arc<ProxyConfig>,
buffer_pool: Arc<BufferPool>,
local_addr: SocketAddr,
rng: Arc<SecureRandom>,
route_rx: watch::Receiver<RouteCutoverState>,
route_snapshot: RouteCutoverState,
session_id: u64,
session_cancel: CancellationToken,
shared: Arc<ProxySharedState>,
) -> Result<()>
where
R: AsyncRead + Unpin + Send + 'static,
W: AsyncWrite + Unpin + Send + 'static,
{
handle_via_middle_proxy_with_conntrack(
crypto_reader,
crypto_writer,
success,
me_pool,
stats,
config,
buffer_pool,
local_addr,
rng,
route_rx,
route_snapshot,
session_id,
session_cancel,
shared,
ConntrackClosePolicy::Publish,
)
.await
}
const DESYNC_DEDUP_WINDOW: Duration = Duration::from_secs(60);
const DESYNC_DEDUP_MAX_ENTRIES: usize = 65_536;
const DESYNC_FULL_CACHE_EMIT_MIN_INTERVAL: Duration = Duration::from_millis(1000);
@@ -98,6 +140,7 @@ const DESYNC_ERROR_CLASS: &str = "frame_too_large_crypto_desync";
const C2ME_CHANNEL_CAPACITY_FALLBACK: usize = 128;
const C2ME_SOFT_PRESSURE_MIN_FREE_SLOTS: usize = 64;
const C2ME_SENDER_FAIRNESS_BUDGET: usize = 32;
const C2ME_QUEUED_BYTE_PERMIT_UNIT: usize = 16 * 1024;
const C2ME_QUEUED_PERMITS_PER_SLOT: usize = 4;
const RELAY_IDLE_IO_POLL_MAX: Duration = Duration::from_secs(1);
src/proxy/middle_relay/session.rs
+18 -14
View File
@@ -1,6 +1,7 @@
use super::*;
pub(crate) async fn handle_via_middle_proxy<R, W>(
/// Runs Middle-End relay with explicit kernel-conntrack close publication policy.
pub(crate) async fn handle_via_middle_proxy_with_conntrack<R, W>(
mut crypto_reader: CryptoReader<R>,
crypto_writer: CryptoWriter<W>,
success: HandshakeSuccess,
@@ -15,6 +16,7 @@ pub(crate) async fn handle_via_middle_proxy<R, W>(
session_id: u64,
session_cancel: CancellationToken,
shared: Arc<ProxySharedState>,
conntrack_close_policy: ConntrackClosePolicy,
) -> Result<()>
where
R: AsyncRead + Unpin + Send + 'static,
@@ -78,7 +80,7 @@ where
return Err(ProxyError::RouteSwitched);
}
// Per-user ad_tag from access.user_ad_tags; fallback to general.ad_tag (hot-reloadable)
// Prefer the hot-reloadable per-user ad tag over the global fallback.
let user_tag: Option<Vec<u8>> = config
.access
.user_ad_tags
@@ -785,7 +787,7 @@ where
}
};
// When client closes, but ME channel stopped as unregistered - it isnt error
// A client-initiated close can unregister the ME channel before its writer exits.
if client_closed && matches!(writer_result, Err(ProxyError::MiddleConnectionLost)) {
writer_result = Ok(());
}
@@ -808,17 +810,19 @@ where
"ME relay cleanup"
);
let close_reason = classify_conntrack_close_reason(&result);
let publish_result = shared.publish_conntrack_close_event(ConntrackCloseEvent {
src: peer,
dst: local_addr,
reason: close_reason,
});
if !matches!(
publish_result,
ConntrackClosePublishResult::Sent | ConntrackClosePublishResult::Disabled
) {
stats.increment_conntrack_close_event_drop_total();
if conntrack_close_policy == ConntrackClosePolicy::Publish {
let close_reason = classify_conntrack_close_reason(&result);
let publish_result = shared.publish_conntrack_close_event(ConntrackCloseEvent {
src: peer,
dst: local_addr,
reason: close_reason,
});
if !matches!(
publish_result,
ConntrackClosePublishResult::Sent | ConntrackClosePublishResult::Disabled
) {
stats.increment_conntrack_close_event_drop_total();
}
}
clear_relay_idle_candidate_in(shared.as_ref(), conn_id);
src/proxy/mod.rs
+2
View File
@@ -59,6 +59,8 @@
)]
pub mod adaptive_buffers;
// Shared authenticated admission and relay orchestration for TCP and WEB streams.
pub(crate) mod authenticated;
pub mod client;
// Process-wide Direct relay copy-buffer ownership and pressure policy.
pub(crate) mod direct_buffer_budget;
src/proxy/shared_state.rs
+9
View File
@@ -41,6 +41,15 @@ pub(crate) enum ConntrackClosePublishResult {
QueueClosed,
}
/// Controls whether a relay tuple maps to a real kernel conntrack entry.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum ConntrackClosePolicy {
/// Publish closure for a tuple backed by an accepted kernel TCP flow.
Publish,
/// Suppress closure for a virtual transport tuple with no kernel flow.
Suppress,
}
pub(crate) struct HandshakeSharedState {
pub(crate) auth_probe: DashMap<IpAddr, AuthProbeState>,
pub(crate) auth_probe_saturation: Mutex<Option<AuthProbeSaturationState>>,