naughtysoul/telemt
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

HOL Minimized + Random conn_id + Target DC Magics

Browse Source 
Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
This commit is contained in:
Alexey
2026-02-14 01:52:49 +03:00
parent a8c3128c50
commit d405756b94
3 changed files with 69 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
src/proxy/middle_relay.rs
View File

@@ -61,7 +61,14 @@ where
Ok(Some(payload)) => {
trace!(conn_id, bytes = payload.len(), "C->ME frame");
stats.add_user_octets_from(&user, payload.len() as u64);
me_pool.send_proxy_req(conn_id, peer, translated_local_addr, &payload, proto_flags).await?;
me_pool.send_proxy_req(
conn_id,
success.dc_idx,
peer,
translated_local_addr,
&payload,
proto_flags,
).await?;
}
Ok(None) => {
debug!(conn_id, "Client EOF");

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

@@ -26,7 +26,7 @@ const ME_ACTIVE_PING_SECS: u64 = 25;
pub struct MePool {
pub(super) registry: Arc<ConnRegistry>,
pub(super) writers: Arc<RwLock<Vec<Arc<Mutex<RpcWriter>>>>>,
pub(super) writers: Arc<RwLock<Vec<(SocketAddr, Arc<Mutex<RpcWriter>>)>>> ,
pub(super) rr: AtomicU64,
pub(super) proxy_tag: Option<Vec<u8>>,
proxy_secret: Vec<u8>,
@@ -82,7 +82,8 @@ impl MePool {
}
}
fn writers_arc(&self) -> Arc<RwLock<Vec<Arc<Mutex<RpcWriter>>>>> {
fn writers_arc(&self) -> Arc<RwLock<Vec<(SocketAddr, Arc<Mutex<RpcWriter>>)>>>
{
self.writers.clone()
}
@@ -348,7 +349,7 @@ impl MePool {
iv: write_iv,
seq_no: 0,
}));
self.writers.write().await.push(rpc_w.clone());
self.writers.write().await.push((addr, rpc_w.clone()));
let reg = self.registry.clone();
let w_pong = rpc_w.clone();
@@ -362,7 +363,7 @@ impl MePool {
warn!(error = %e, "ME reader ended");
}
let mut ws = w_pool.write().await;
ws.retain(|w| !Arc::ptr_eq(w, &w_pong));
ws.retain(|(_, w)| !Arc::ptr_eq(w, &w_pong));
info!(remaining = ws.len(), "Dead ME writer removed from pool");
});
tokio::spawn(async move {
@@ -376,7 +377,7 @@ impl MePool {
if let Err(e) = w_ping.lock().await.send(&p).await {
debug!(error = %e, "Active ME ping failed, removing dead writer");
let mut ws = w_pool_ping.write().await;
ws.retain(|w| !Arc::ptr_eq(w, &w_ping));
ws.retain(|(_, w)| !Arc::ptr_eq(w, &w_ping));
break;
}
}

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

@@ -6,7 +6,7 @@ use tokio::sync::Mutex;
use tracing::{debug, warn};
use crate::error::{ProxyError, Result};
use crate::protocol::constants::RPC_CLOSE_EXT_U32;
use crate::protocol::constants::{RPC_CLOSE_EXT_U32, TG_MIDDLE_PROXIES_V4};
use super::MePool;
use super::codec::RpcWriter;
@@ -16,6 +16,7 @@ impl MePool {
pub async fn send_proxy_req(
&self,
conn_id: u64,
target_dc: i16,
client_addr: SocketAddr,
our_addr: SocketAddr,
data: &[u8],
@@ -35,14 +36,20 @@ impl MePool {
if ws.is_empty() {
return Err(ProxyError::Proxy("All ME connections dead".into()));
}
let writers: Vec<Arc<Mutex<RpcWriter>>> = ws.iter().cloned().collect();
let start = self.rr.fetch_add(1, Ordering::Relaxed) as usize % writers.len();
let writers: Vec<(SocketAddr, Arc<Mutex<RpcWriter>>)> = ws.iter().cloned().collect();
drop(ws);
let candidate_indices = candidate_indices_for_dc(&writers, target_dc);
if candidate_indices.is_empty() {
return Err(ProxyError::Proxy("No ME writers available for target DC".into()));
}
let start = self.rr.fetch_add(1, Ordering::Relaxed) as usize % candidate_indices.len();
// Prefer immediately available writer to avoid waiting on stalled connection.
for offset in 0..writers.len() {
let idx = (start + offset) % writers.len();
let w = writers[idx].clone();
for offset in 0..candidate_indices.len() {
let cidx = (start + offset) % candidate_indices.len();
let idx = candidate_indices[cidx];
let w = writers[idx].1.clone();
if let Ok(mut guard) = w.try_lock() {
let send_res = guard.send(&payload).await;
drop(guard);
@@ -51,7 +58,7 @@ impl MePool {
Err(e) => {
warn!(error = %e, "ME write failed, removing dead conn");
let mut ws = self.writers.write().await;
ws.retain(|o| !Arc::ptr_eq(o, &w));
ws.retain(|(_, o)| !Arc::ptr_eq(o, &w));
if ws.is_empty() {
return Err(ProxyError::Proxy("All ME connections dead".into()));
}
@@ -62,13 +69,13 @@ impl MePool {
}
// All writers are currently busy, wait for the selected one.
let w = writers[start].clone();
let w = writers[candidate_indices[start]].1.clone();
match w.lock().await.send(&payload).await {
Ok(()) => return Ok(()),
Err(e) => {
warn!(error = %e, "ME write failed, removing dead conn");
let mut ws = self.writers.write().await;
ws.retain(|o| !Arc::ptr_eq(o, &w));
ws.retain(|(_, o)| !Arc::ptr_eq(o, &w));
if ws.is_empty() {
return Err(ProxyError::Proxy("All ME connections dead".into()));
}
@@ -80,7 +87,7 @@ impl MePool {
pub async fn send_close(&self, conn_id: u64) -> Result<()> {
let ws = self.writers.read().await;
if !ws.is_empty() {
let w = ws[0].clone();
let w = ws[0].1.clone();
drop(ws);
let mut p = Vec::with_capacity(12);
p.extend_from_slice(&RPC_CLOSE_EXT_U32.to_le_bytes());
@@ -88,7 +95,7 @@ impl MePool {
if let Err(e) = w.lock().await.send(&p).await {
debug!(error = %e, "ME close write failed");
let mut ws = self.writers.write().await;
ws.retain(|o| !Arc::ptr_eq(o, &w));
ws.retain(|(_, o)| !Arc::ptr_eq(o, &w));
}
}
@@ -100,3 +107,40 @@ impl MePool {
self.writers.try_read().map(|w| w.len()).unwrap_or(0)
}
}
fn candidate_indices_for_dc(
writers: &[(SocketAddr, Arc<Mutex<RpcWriter>>)],
target_dc: i16,
) -> Vec<usize> {
let mut preferred = Vec::<SocketAddr>::new();
let key = target_dc as i32;
if let Some(v) = TG_MIDDLE_PROXIES_V4.get(&key) {
preferred.extend(v.iter().map(|(ip, port)| SocketAddr::new(*ip, *port)));
}
if preferred.is_empty() {
let abs = key.abs();
if let Some(v) = TG_MIDDLE_PROXIES_V4.get(&abs) {
preferred.extend(v.iter().map(|(ip, port)| SocketAddr::new(*ip, *port)));
}
}
if preferred.is_empty() {
let abs = key.abs();
if let Some(v) = TG_MIDDLE_PROXIES_V4.get(&-abs) {
preferred.extend(v.iter().map(|(ip, port)| SocketAddr::new(*ip, *port)));
}
}
if preferred.is_empty() {
return (0..writers.len()).collect();
}
let mut out = Vec::new();
for (idx, (addr, _)) in writers.iter().enumerate() {
if preferred.iter().any(|p| p == addr) {
out.push(idx);
}
}
if out.is_empty() {
return (0..writers.len()).collect();
}
out
}