astelm/telemt
1
0
Fork 0
mirror of https://github.com/telemt/telemt.git synced 2026-04-21 12:34:11 +03:00
Code Issues Packages Projects Releases Wiki Activity

Refactor proxy and transport modules for improved safety and performance

Browse Source 
- Enhanced linting rules in `src/proxy/mod.rs` to enforce stricter code quality checks in production.
- Updated hash functions in `src/proxy/middle_relay.rs` for better efficiency.
- Added new security tests in `src/proxy/tests/middle_relay_stub_completion_security_tests.rs` to validate desynchronization behavior.
- Removed ignored test stubs in `src/proxy/tests/middle_relay_security_tests.rs` to clean up the test suite.
- Improved error handling and code readability in various transport modules, including `src/transport/middle_proxy/config_updater.rs` and `src/transport/middle_proxy/pool.rs`.
- Introduced new padding functions in `src/stream/frame_stream_padding_security_tests.rs` to ensure consistent behavior across different implementations.
- Adjusted TLS stream validation in `src/stream/tls_stream.rs` for better boundary checking.
- General code cleanup and dead code elimination across multiple files to enhance maintainability.
This commit is contained in:
David Osipov
2026-03-21 20:05:07 +04:00
parent a6c298b633
commit 4c32370b25
35 changed files with 794 additions and 174 deletions
Download Patch File Download Diff File Expand all files Collapse all files

76
src/transport/middle_proxy/config_updater.rs
View File

@@ -355,49 +355,49 @@ async fn run_update_cycle(
let mut ready_v4: Option<(ProxyConfigData, u64)> = None;
let cfg_v4 = retry_fetch("https://core.telegram.org/getProxyConfig").await;
if let Some(cfg_v4) = cfg_v4 {
if snapshot_passes_guards(cfg, &cfg_v4, "getProxyConfig") {
let cfg_v4_hash = hash_proxy_config(&cfg_v4);
let stable_hits = state.config_v4.observe(cfg_v4_hash);
if stable_hits < required_cfg_snapshots {
debug!(
stable_hits,
required_cfg_snapshots,
snapshot = format_args!("0x{cfg_v4_hash:016x}"),
"ME config v4 candidate observed"
);
} else if state.config_v4.is_applied(cfg_v4_hash) {
debug!(
snapshot = format_args!("0x{cfg_v4_hash:016x}"),
"ME config v4 stable snapshot already applied"
);
} else {
ready_v4 = Some((cfg_v4, cfg_v4_hash));
}
if let Some(cfg_v4) = cfg_v4
&& snapshot_passes_guards(cfg, &cfg_v4, "getProxyConfig")
{
let cfg_v4_hash = hash_proxy_config(&cfg_v4);
let stable_hits = state.config_v4.observe(cfg_v4_hash);
if stable_hits < required_cfg_snapshots {
debug!(
stable_hits,
required_cfg_snapshots,
snapshot = format_args!("0x{cfg_v4_hash:016x}"),
"ME config v4 candidate observed"
);
} else if state.config_v4.is_applied(cfg_v4_hash) {
debug!(
snapshot = format_args!("0x{cfg_v4_hash:016x}"),
"ME config v4 stable snapshot already applied"
);
} else {
ready_v4 = Some((cfg_v4, cfg_v4_hash));
}
}
let mut ready_v6: Option<(ProxyConfigData, u64)> = None;
let cfg_v6 = retry_fetch("https://core.telegram.org/getProxyConfigV6").await;
if let Some(cfg_v6) = cfg_v6 {
if snapshot_passes_guards(cfg, &cfg_v6, "getProxyConfigV6") {
let cfg_v6_hash = hash_proxy_config(&cfg_v6);
let stable_hits = state.config_v6.observe(cfg_v6_hash);
if stable_hits < required_cfg_snapshots {
debug!(
stable_hits,
required_cfg_snapshots,
snapshot = format_args!("0x{cfg_v6_hash:016x}"),
"ME config v6 candidate observed"
);
} else if state.config_v6.is_applied(cfg_v6_hash) {
debug!(
snapshot = format_args!("0x{cfg_v6_hash:016x}"),
"ME config v6 stable snapshot already applied"
);
} else {
ready_v6 = Some((cfg_v6, cfg_v6_hash));
}
if let Some(cfg_v6) = cfg_v6
&& snapshot_passes_guards(cfg, &cfg_v6, "getProxyConfigV6")
{
let cfg_v6_hash = hash_proxy_config(&cfg_v6);
let stable_hits = state.config_v6.observe(cfg_v6_hash);
if stable_hits < required_cfg_snapshots {
debug!(
stable_hits,
required_cfg_snapshots,
snapshot = format_args!("0x{cfg_v6_hash:016x}"),
"ME config v6 candidate observed"
);
} else if state.config_v6.is_applied(cfg_v6_hash) {
debug!(
snapshot = format_args!("0x{cfg_v6_hash:016x}"),
"ME config v6 stable snapshot already applied"
);
} else {
ready_v6 = Some((cfg_v6, cfg_v6_hash));
}
}

2
src/transport/middle_proxy/health.rs
View File

@@ -1,3 +1,5 @@
#![allow(clippy::too_many_arguments)]
use std::collections::HashMap;
use std::collections::HashSet;
use std::net::SocketAddr;

22
src/transport/middle_proxy/ping.rs
View File

@@ -1,3 +1,5 @@
#![allow(clippy::items_after_test_module)]
use std::collections::HashMap;
use std::net::{IpAddr, SocketAddr};
use std::sync::Arc;
@@ -12,6 +14,8 @@ use crate::transport::{UpstreamEgressInfo, UpstreamRouteKind};
use super::MePool;
type MePingGroup = (MePingFamily, i32, Vec<(IpAddr, u16)>);
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MePingFamily {
V4,
@@ -137,14 +141,14 @@ fn detect_interface_for_ip(ip: IpAddr) -> Option<String> {
if let Ok(addrs) = getifaddrs() {
for iface in addrs {
if let Some(address) = iface.address {
if let Some(v4) = address.as_sockaddr_in() {
if IpAddr::V4(v4.ip()) == ip {
return Some(iface.interface_name);
}
} else if let Some(v6) = address.as_sockaddr_in6() {
if IpAddr::V6(v6.ip()) == ip {
return Some(iface.interface_name);
}
if let Some(v4) = address.as_sockaddr_in()
&& IpAddr::V4(v4.ip()) == ip
{
return Some(iface.interface_name);
} else if let Some(v6) = address.as_sockaddr_in6()
&& IpAddr::V6(v6.ip()) == ip
{
return Some(iface.interface_name);
}
}
}
@@ -329,7 +333,7 @@ pub async fn run_me_ping(pool: &Arc<MePool>, rng: &SecureRandom) -> Vec<MePingRe
HashMap::new()
};
let mut grouped: Vec<(MePingFamily, i32, Vec<(IpAddr, u16)>)> = Vec::new();
let mut grouped: Vec<MePingGroup> = Vec::new();
for (dc, addrs) in v4_map {
grouped.push((MePingFamily::V4, dc, addrs));
}

8
src/transport/middle_proxy/pool.rs
View File

@@ -1,3 +1,5 @@
#![allow(clippy::too_many_arguments, clippy::type_complexity)]
use std::collections::{HashMap, HashSet};
use std::net::{IpAddr, Ipv6Addr, SocketAddr};
use std::sync::Arc;
@@ -619,6 +621,7 @@ impl MePool {
self.runtime_ready.load(Ordering::Relaxed)
}
#[allow(dead_code)]
pub(super) fn set_family_runtime_state(
&self,
family: IpFamily,
@@ -982,28 +985,33 @@ impl MePool {
Some(Duration::from_secs(secs))
}
#[allow(dead_code)]
pub(super) fn drain_soft_evict_enabled(&self) -> bool {
self.me_pool_drain_soft_evict_enabled
.load(Ordering::Relaxed)
}
#[allow(dead_code)]
pub(super) fn drain_soft_evict_grace_secs(&self) -> u64 {
self.me_pool_drain_soft_evict_grace_secs
.load(Ordering::Relaxed)
}
#[allow(dead_code)]
pub(super) fn drain_soft_evict_per_writer(&self) -> usize {
self.me_pool_drain_soft_evict_per_writer
.load(Ordering::Relaxed)
.max(1) as usize
}
#[allow(dead_code)]
pub(super) fn drain_soft_evict_budget_per_core(&self) -> usize {
self.me_pool_drain_soft_evict_budget_per_core
.load(Ordering::Relaxed)
.max(1) as usize
}
#[allow(dead_code)]
pub(super) fn drain_soft_evict_cooldown(&self) -> Duration {
Duration::from_millis(
self.me_pool_drain_soft_evict_cooldown_ms

28
src/transport/middle_proxy/pool_status.rs
View File

@@ -293,20 +293,20 @@ impl MePool {
WriterContour::Draining => "draining",
};
if !draining {
if let Some(dc_idx) = dc {
*live_writers_by_dc_endpoint
.entry((dc_idx, endpoint))
.or_insert(0) += 1;
*live_writers_by_dc.entry(dc_idx).or_insert(0) += 1;
if let Some(ema_ms) = rtt_ema_ms {
let entry = dc_rtt_agg.entry(dc_idx).or_insert((0.0, 0));
entry.0 += ema_ms;
entry.1 += 1;
}
if matches_active_generation && in_desired_map {
*fresh_writers_by_dc.entry(dc_idx).or_insert(0) += 1;
}
if !draining
&& let Some(dc_idx) = dc
{
*live_writers_by_dc_endpoint
.entry((dc_idx, endpoint))
.or_insert(0) += 1;
*live_writers_by_dc.entry(dc_idx).or_insert(0) += 1;
if let Some(ema_ms) = rtt_ema_ms {
let entry = dc_rtt_agg.entry(dc_idx).or_insert((0.0, 0));
entry.0 += ema_ms;
entry.1 += 1;
}
if matches_active_generation && in_desired_map {
*fresh_writers_by_dc.entry(dc_idx).or_insert(0) += 1;
}
}

14
src/transport/middle_proxy/pool_writer.rs
View File

@@ -268,10 +268,10 @@ impl MePool {
cancel_reader_token.cancel();
}
}
if let Err(e) = res {
if !idle_close_by_peer {
warn!(error = %e, "ME reader ended");
}
if let Err(e) = res
&& !idle_close_by_peer
{
warn!(error = %e, "ME reader ended");
}
let remaining = writers_arc.read().await.len();
debug!(writer_id, remaining, "ME reader task finished");
@@ -386,10 +386,9 @@ impl MePool {
if cleanup_for_ping
.compare_exchange(false, true, Ordering::AcqRel, Ordering::Relaxed)
.is_ok()
&& let Some(pool) = pool_ping.upgrade()
{
if let Some(pool) = pool_ping.upgrade() {
pool.remove_writer_and_close_clients(writer_id).await;
}
pool.remove_writer_and_close_clients(writer_id).await;
}
break;
}
@@ -538,6 +537,7 @@ impl MePool {
.await
}
#[allow(dead_code)]
async fn remove_writer_only(self: &Arc<Self>, writer_id: u64) -> bool {
self.remove_writer_with_mode(writer_id, WriterTeardownMode::Any)
.await

2
src/transport/middle_proxy/reader.rs
View File

@@ -1,3 +1,5 @@
#![allow(clippy::too_many_arguments)]
use std::collections::HashMap;
use std::io::ErrorKind;
use std::sync::Arc;

2
src/transport/middle_proxy/registry.rs
View File

@@ -165,6 +165,7 @@ impl ConnRegistry {
None
}
#[allow(dead_code)]
pub async fn route(&self, id: u64, resp: MeResponse) -> RouteResult {
let tx = {
let inner = self.inner.read().await;
@@ -438,6 +439,7 @@ impl ConnRegistry {
.unwrap_or(true)
}
#[allow(dead_code)]
pub async fn unregister_writer_if_empty(&self, writer_id: u64) -> bool {
let mut inner = self.inner.write().await;
let Some(conn_ids) = inner.conns_for_writer.get(&writer_id) else {

6
src/transport/middle_proxy/send.rs
View File

@@ -1,3 +1,5 @@
#![allow(clippy::too_many_arguments)]
use std::cmp::Reverse;
use std::collections::{HashMap, HashSet};
use std::net::SocketAddr;
@@ -593,7 +595,7 @@ impl MePool {
let round = *hybrid_recovery_round;
let target_triggered = self.trigger_async_recovery_for_target_dc(routed_dc).await;
if !target_triggered || round % HYBRID_GLOBAL_BURST_PERIOD_ROUNDS == 0 {
if !target_triggered || round.is_multiple_of(HYBRID_GLOBAL_BURST_PERIOD_ROUNDS) {
self.trigger_async_recovery_global().await;
}
*hybrid_recovery_round = round.saturating_add(1);
@@ -672,7 +674,7 @@ impl MePool {
if !self.writer_eligible_for_selection(w, include_warm) {
continue;
}
if w.writer_dc == routed_dc && preferred.iter().any(|endpoint| *endpoint == w.addr) {
if w.writer_dc == routed_dc && preferred.contains(&w.addr) {
out.push(idx);
}
}