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

Middle Proxy is so real

Browse Source
This commit is contained in:
Alexey
2026-02-14 01:36:14 +03:00
parent 9b850b0bfb
commit 70859aa5cf
14 changed files with 2028 additions and 785 deletions
Download Patch File Download Diff File Expand all files Collapse all files

931
src/proxy/client.rs
View File

@@ -1,661 +1,354 @@
//! Client Handler
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Duration;
use tokio::net::TcpStream;
use tokio::io::{AsyncRead, AsyncWrite, AsyncReadExt, AsyncWriteExt};
use tokio::time::timeout;
use tracing::{debug, info, warn, error, trace};
use crate::config::ProxyConfig;
use crate::error::{ProxyError, Result, HandshakeResult};
use crate::protocol::constants::*;
use crate::protocol::tls;
use crate::stats::{Stats, ReplayChecker};
use crate::transport::{configure_client_socket, UpstreamManager};
use crate::transport::middle_proxy::{MePool, MeResponse, proto_flags_for_tag};
use crate::stream::{CryptoReader, CryptoWriter, FakeTlsReader, FakeTlsWriter, BufferPool};
use crate::crypto::{AesCtr, SecureRandom};
use crate::proxy::handshake::{
handle_tls_handshake, handle_mtproto_handshake,
HandshakeSuccess, generate_tg_nonce, encrypt_tg_nonce_with_ciphers,
};
use crate::proxy::relay::relay_bidirectional;
use crate::proxy::masking::handle_bad_client;
pub struct ClientHandler;
pub struct RunningClientHandler {
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Duration;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite};
use tokio::net::TcpStream;
use tokio::time::timeout;
use tracing::{debug, warn};
use crate::config::ProxyConfig;
use crate::crypto::SecureRandom;
use crate::error::{HandshakeResult, ProxyError, Result};
use crate::protocol::constants::*;
use crate::protocol::tls;
use crate::stats::{ReplayChecker, Stats};
use crate::stream::{BufferPool, CryptoReader, CryptoWriter};
use crate::transport::middle_proxy::MePool;
use crate::transport::{UpstreamManager, configure_client_socket};
use crate::proxy::direct_relay::handle_via_direct;
use crate::proxy::handshake::{HandshakeSuccess, handle_mtproto_handshake, handle_tls_handshake};
use crate::proxy::masking::handle_bad_client;
use crate::proxy::middle_relay::handle_via_middle_proxy;
pub struct ClientHandler;
pub struct RunningClientHandler {
stream: TcpStream,
peer: SocketAddr,
config: Arc<ProxyConfig>,
stats: Arc<Stats>,
replay_checker: Arc<ReplayChecker>,
upstream_manager: Arc<UpstreamManager>,
buffer_pool: Arc<BufferPool>,
rng: Arc<SecureRandom>,
me_pool: Option<Arc<MePool>>,
}
impl ClientHandler {
pub fn new(
stream: TcpStream,
peer: SocketAddr,
config: Arc<ProxyConfig>,
stats: Arc<Stats>,
replay_checker: Arc<ReplayChecker>,
upstream_manager: Arc<UpstreamManager>,
replay_checker: Arc<ReplayChecker>,
buffer_pool: Arc<BufferPool>,
rng: Arc<SecureRandom>,
me_pool: Option<Arc<MePool>>,
) -> RunningClientHandler {
RunningClientHandler {
stream,
peer,
config,
stats,
replay_checker,
upstream_manager,
buffer_pool,
rng,
me_pool,
}
}
impl ClientHandler {
pub fn new(
stream: TcpStream,
peer: SocketAddr,
config: Arc<ProxyConfig>,
stats: Arc<Stats>,
upstream_manager: Arc<UpstreamManager>,
replay_checker: Arc<ReplayChecker>,
buffer_pool: Arc<BufferPool>,
rng: Arc<SecureRandom>,
me_pool: Option<Arc<MePool>>,
) -> RunningClientHandler {
RunningClientHandler {
stream, peer, config, stats, replay_checker,
upstream_manager, buffer_pool, rng, me_pool,
}
impl RunningClientHandler {
pub async fn run(mut self) -> Result<()> {
self.stats.increment_connects_all();
let peer = self.peer;
debug!(peer = %peer, "New connection");
if let Err(e) = configure_client_socket(
&self.stream,
self.config.timeouts.client_keepalive,
self.config.timeouts.client_ack,
) {
debug!(peer = %peer, error = %e, "Failed to configure client socket");
}
let handshake_timeout = Duration::from_secs(self.config.timeouts.client_handshake);
let stats = self.stats.clone();
let result = timeout(handshake_timeout, self.do_handshake()).await;
match result {
Ok(Ok(())) => {
debug!(peer = %peer, "Connection handled successfully");
Ok(())
}
Ok(Err(e)) => {
debug!(peer = %peer, error = %e, "Handshake failed");
Err(e)
}
Err(_) => {
stats.increment_handshake_timeouts();
debug!(peer = %peer, "Handshake timeout");
Err(ProxyError::TgHandshakeTimeout)
}
}
}
impl RunningClientHandler {
pub async fn run(mut self) -> Result<()> {
self.stats.increment_connects_all();
let peer = self.peer;
debug!(peer = %peer, "New connection");
if let Err(e) = configure_client_socket(
&self.stream,
self.config.timeouts.client_keepalive,
self.config.timeouts.client_ack,
) {
debug!(peer = %peer, error = %e, "Failed to configure client socket");
}
let handshake_timeout = Duration::from_secs(self.config.timeouts.client_handshake);
let stats = self.stats.clone();
let result = timeout(handshake_timeout, self.do_handshake()).await;
match result {
Ok(Ok(())) => {
debug!(peer = %peer, "Connection handled successfully");
Ok(())
}
Ok(Err(e)) => {
debug!(peer = %peer, error = %e, "Handshake failed");
Err(e)
}
Err(_) => {
stats.increment_handshake_timeouts();
debug!(peer = %peer, "Handshake timeout");
Err(ProxyError::TgHandshakeTimeout)
}
}
async fn do_handshake(mut self) -> Result<()> {
let mut first_bytes = [0u8; 5];
self.stream.read_exact(&mut first_bytes).await?;
let is_tls = tls::is_tls_handshake(&first_bytes[..3]);
let peer = self.peer;
debug!(peer = %peer, is_tls = is_tls, "Handshake type detected");
if is_tls {
self.handle_tls_client(first_bytes).await
} else {
self.handle_direct_client(first_bytes).await
}
async fn do_handshake(mut self) -> Result<()> {
let mut first_bytes = [0u8; 5];
self.stream.read_exact(&mut first_bytes).await?;
let is_tls = tls::is_tls_handshake(&first_bytes[..3]);
let peer = self.peer;
debug!(peer = %peer, is_tls = is_tls, "Handshake type detected");
if is_tls {
self.handle_tls_client(first_bytes).await
} else {
self.handle_direct_client(first_bytes).await
}
}
async fn handle_tls_client(mut self, first_bytes: [u8; 5]) -> Result<()> {
let peer = self.peer;
let tls_len = u16::from_be_bytes([first_bytes[3], first_bytes[4]]) as usize;
debug!(peer = %peer, tls_len = tls_len, "Reading TLS handshake");
if tls_len < 512 {
debug!(peer = %peer, tls_len = tls_len, "TLS handshake too short");
self.stats.increment_connects_bad();
let (reader, writer) = self.stream.into_split();
handle_bad_client(reader, writer, &first_bytes, &self.config).await;
return Ok(());
}
async fn handle_tls_client(mut self, first_bytes: [u8; 5]) -> Result<()> {
let peer = self.peer;
let tls_len = u16::from_be_bytes([first_bytes[3], first_bytes[4]]) as usize;
debug!(peer = %peer, tls_len = tls_len, "Reading TLS handshake");
if tls_len < 512 {
debug!(peer = %peer, tls_len = tls_len, "TLS handshake too short");
self.stats.increment_connects_bad();
let (reader, writer) = self.stream.into_split();
handle_bad_client(reader, writer, &first_bytes, &self.config).await;
let mut handshake = vec![0u8; 5 + tls_len];
handshake[..5].copy_from_slice(&first_bytes);
self.stream.read_exact(&mut handshake[5..]).await?;
let config = self.config.clone();
let replay_checker = self.replay_checker.clone();
let stats = self.stats.clone();
let buffer_pool = self.buffer_pool.clone();
let local_addr = self.stream.local_addr().map_err(ProxyError::Io)?;
let (read_half, write_half) = self.stream.into_split();
let (mut tls_reader, tls_writer, _tls_user) = match handle_tls_handshake(
&handshake,
read_half,
write_half,
peer,
&config,
&replay_checker,
&self.rng,
)
.await
{
HandshakeResult::Success(result) => result,
HandshakeResult::BadClient { reader, writer } => {
stats.increment_connects_bad();
handle_bad_client(reader, writer, &handshake, &config).await;
return Ok(());
}
let mut handshake = vec![0u8; 5 + tls_len];
handshake[..5].copy_from_slice(&first_bytes);
self.stream.read_exact(&mut handshake[5..]).await?;
let config = self.config.clone();
let replay_checker = self.replay_checker.clone();
let stats = self.stats.clone();
let buffer_pool = self.buffer_pool.clone();
let (read_half, write_half) = self.stream.into_split();
let (mut tls_reader, tls_writer, _tls_user) = match handle_tls_handshake(
&handshake, read_half, write_half, peer,
&config, &replay_checker, &self.rng,
).await {
HandshakeResult::Success(result) => result,
HandshakeResult::BadClient { reader, writer } => {
stats.increment_connects_bad();
handle_bad_client(reader, writer, &handshake, &config).await;
return Ok(());
}
HandshakeResult::Error(e) => return Err(e),
};
debug!(peer = %peer, "Reading MTProto handshake through TLS");
let mtproto_data = tls_reader.read_exact(HANDSHAKE_LEN).await?;
let mtproto_handshake: [u8; HANDSHAKE_LEN] = mtproto_data[..].try_into()
.map_err(|_| ProxyError::InvalidHandshake("Short MTProto handshake".into()))?;
let (crypto_reader, crypto_writer, success) = match handle_mtproto_handshake(
&mtproto_handshake, tls_reader, tls_writer, peer,
&config, &replay_checker, true,
).await {
HandshakeResult::Success(result) => result,
HandshakeResult::BadClient { reader: _, writer: _ } => {
stats.increment_connects_bad();
debug!(peer = %peer, "Valid TLS but invalid MTProto handshake");
return Ok(());
}
HandshakeResult::Error(e) => return Err(e),
};
Self::handle_authenticated_static(
crypto_reader, crypto_writer, success,
self.upstream_manager, self.stats, self.config,
buffer_pool, self.rng, self.me_pool,
).await
}
async fn handle_direct_client(mut self, first_bytes: [u8; 5]) -> Result<()> {
let peer = self.peer;
if !self.config.general.modes.classic && !self.config.general.modes.secure {
debug!(peer = %peer, "Non-TLS modes disabled");
self.stats.increment_connects_bad();
let (reader, writer) = self.stream.into_split();
handle_bad_client(reader, writer, &first_bytes, &self.config).await;
HandshakeResult::Error(e) => return Err(e),
};
debug!(peer = %peer, "Reading MTProto handshake through TLS");
let mtproto_data = tls_reader.read_exact(HANDSHAKE_LEN).await?;
let mtproto_handshake: [u8; HANDSHAKE_LEN] = mtproto_data[..]
.try_into()
.map_err(|_| ProxyError::InvalidHandshake("Short MTProto handshake".into()))?;
let (crypto_reader, crypto_writer, success) = match handle_mtproto_handshake(
&mtproto_handshake,
tls_reader,
tls_writer,
peer,
&config,
&replay_checker,
true,
)
.await
{
HandshakeResult::Success(result) => result,
HandshakeResult::BadClient {
reader: _,
writer: _,
} => {
stats.increment_connects_bad();
debug!(peer = %peer, "Valid TLS but invalid MTProto handshake");
return Ok(());
}
let mut handshake = [0u8; HANDSHAKE_LEN];
handshake[..5].copy_from_slice(&first_bytes);
self.stream.read_exact(&mut handshake[5..]).await?;
let config = self.config.clone();
let replay_checker = self.replay_checker.clone();
let stats = self.stats.clone();
let buffer_pool = self.buffer_pool.clone();
let (read_half, write_half) = self.stream.into_split();
let (crypto_reader, crypto_writer, success) = match handle_mtproto_handshake(
&handshake, read_half, write_half, peer,
&config, &replay_checker, false,
).await {
HandshakeResult::Success(result) => result,
HandshakeResult::BadClient { reader, writer } => {
stats.increment_connects_bad();
handle_bad_client(reader, writer, &handshake, &config).await;
return Ok(());
}
HandshakeResult::Error(e) => return Err(e),
};
Self::handle_authenticated_static(
crypto_reader, crypto_writer, success,
self.upstream_manager, self.stats, self.config,
buffer_pool, self.rng, self.me_pool,
).await
HandshakeResult::Error(e) => return Err(e),
};
Self::handle_authenticated_static(
crypto_reader,
crypto_writer,
success,
self.upstream_manager,
self.stats,
self.config,
buffer_pool,
self.rng,
self.me_pool,
local_addr,
)
.await
}
async fn handle_direct_client(mut self, first_bytes: [u8; 5]) -> Result<()> {
let peer = self.peer;
if !self.config.general.modes.classic && !self.config.general.modes.secure {
debug!(peer = %peer, "Non-TLS modes disabled");
self.stats.increment_connects_bad();
let (reader, writer) = self.stream.into_split();
handle_bad_client(reader, writer, &first_bytes, &self.config).await;
return Ok(());
}
/// Main dispatch after successful handshake.
/// Two modes:
/// - Direct: TCP relay to TG DC (existing behavior)
/// - Middle Proxy: RPC multiplex through ME pool (new — supports CDN DCs)
async fn handle_authenticated_static<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>,
me_pool: Option<Arc<MePool>>,
) -> Result<()>
where
R: AsyncRead + Unpin + Send + 'static,
W: AsyncWrite + Unpin + Send + 'static,
let mut handshake = [0u8; HANDSHAKE_LEN];
handshake[..5].copy_from_slice(&first_bytes);
self.stream.read_exact(&mut handshake[5..]).await?;
let config = self.config.clone();
let replay_checker = self.replay_checker.clone();
let stats = self.stats.clone();
let buffer_pool = self.buffer_pool.clone();
let local_addr = self.stream.local_addr().map_err(ProxyError::Io)?;
let (read_half, write_half) = self.stream.into_split();
let (crypto_reader, crypto_writer, success) = match handle_mtproto_handshake(
&handshake,
read_half,
write_half,
peer,
&config,
&replay_checker,
false,
)
.await
{
let user = &success.user;
if let Err(e) = Self::check_user_limits_static(user, &config, &stats) {
warn!(user = %user, error = %e, "User limit exceeded");
return Err(e);
HandshakeResult::Success(result) => result,
HandshakeResult::BadClient { reader, writer } => {
stats.increment_connects_bad();
handle_bad_client(reader, writer, &handshake, &config).await;
return Ok(());
}
// Decide: middle proxy or direct
if config.general.use_middle_proxy {
if let Some(ref pool) = me_pool {
return Self::handle_via_middle_proxy(
client_reader, client_writer, success,
pool.clone(), stats, config, buffer_pool,
).await;
}
warn!("use_middle_proxy=true but MePool not initialized, falling back to direct");
}
// Direct mode (original behavior)
Self::handle_via_direct(
client_reader, client_writer, success,
upstream_manager, stats, config, buffer_pool, rng,
).await
}
// =====================================================================
// Direct mode — TCP relay to Telegram DC
// =====================================================================
async fn handle_via_direct<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>,
) -> Result<()>
where
R: AsyncRead + Unpin + Send + 'static,
W: AsyncWrite + Unpin + Send + 'static,
{
let user = &success.user;
let dc_addr = Self::get_dc_addr_static(success.dc_idx, &config)?;
info!(
user = %user,
peer = %success.peer,
dc = success.dc_idx,
dc_addr = %dc_addr,
proto = ?success.proto_tag,
mode = "direct",
"Connecting to Telegram DC"
);
let tg_stream = upstream_manager.connect(dc_addr, Some(success.dc_idx)).await?;
debug!(peer = %success.peer, dc_addr = %dc_addr, "Connected, performing TG handshake");
let (tg_reader, tg_writer) = Self::do_tg_handshake_static(
tg_stream, &success, &config, rng.as_ref(),
).await?;
debug!(peer = %success.peer, "TG handshake complete, starting relay");
stats.increment_user_connects(user);
stats.increment_user_curr_connects(user);
let relay_result = relay_bidirectional(
client_reader, client_writer,
tg_reader, tg_writer,
user, Arc::clone(&stats), buffer_pool,
).await;
stats.decrement_user_curr_connects(user);
match &relay_result {
Ok(()) => debug!(user = %user, "Direct relay completed"),
Err(e) => debug!(user = %user, error = %e, "Direct relay ended with error"),
}
relay_result
}
// =====================================================================
// Middle Proxy mode — RPC multiplex through ME pool
// =====================================================================
/// Middle Proxy RPC relay
///
/// Architecture (matches C MTProxy):
/// ```text
/// Client ←AES-CTR→ [telemt] ←RPC/AES-CBC→ ME ←internal→ DC (any, incl CDN 203)
/// ```
///
/// Key difference from direct mode:
/// - No per-client TCP to DC; all clients share ME pool connections
/// - ME internally routes to correct DC based on client's encrypted auth_key_id
/// - CDN DCs (203+) work because ME knows their internal addresses
/// - We pass raw client MTProto bytes in RPC_PROXY_REQ envelope
/// - ME returns responses in RPC_PROXY_ANS envelope
async fn handle_via_middle_proxy<R, W>(
crypto_reader: CryptoReader<R>,
crypto_writer: CryptoWriter<W>,
HandshakeResult::Error(e) => return Err(e),
};
Self::handle_authenticated_static(
crypto_reader,
crypto_writer,
success,
self.upstream_manager,
self.stats,
self.config,
buffer_pool,
self.rng,
self.me_pool,
local_addr,
)
.await
}
/// Main dispatch after successful handshake.
/// Two modes:
/// - Direct: TCP relay to TG DC (existing behavior)
/// - Middle Proxy: RPC multiplex through ME pool (new — supports CDN DCs)
async fn handle_authenticated_static<R, W>(
client_reader: CryptoReader<R>,
client_writer: CryptoWriter<W>,
success: HandshakeSuccess,
me_pool: Arc<MePool>,
upstream_manager: Arc<UpstreamManager>,
stats: Arc<Stats>,
config: Arc<ProxyConfig>,
_buffer_pool: Arc<BufferPool>,
buffer_pool: Arc<BufferPool>,
rng: Arc<SecureRandom>,
me_pool: Option<Arc<MePool>>,
local_addr: SocketAddr,
) -> Result<()>
where
R: AsyncRead + Unpin + Send + 'static,
W: AsyncWrite + Unpin + Send + 'static,
{
let mut client_reader = crypto_reader;
let mut client_writer = crypto_writer;
let user = success.user.clone();
let peer = success.peer;
let user = &success.user;
info!(
user = %user,
peer = %peer,
dc = success.dc_idx,
proto = ?success.proto_tag,
mode = "middle_proxy",
"Routing via Middle-End"
);
if let Err(e) = Self::check_user_limits_static(user, &config, &stats) {
warn!(user = %user, error = %e, "User limit exceeded");
return Err(e);
}
let (conn_id, mut me_rx) = me_pool.registry().register().await;
let our_addr: SocketAddr = format!("0.0.0.0:{}", config.server.port)
.parse().unwrap_or_else(|_| "0.0.0.0:443".parse().unwrap());
stats.increment_user_connects(&user);
stats.increment_user_curr_connects(&user);
let proto_flags = proto_flags_for_tag(success.proto_tag);
debug!(user = %user, conn_id, proto_flags = format_args!("0x{:08x}", proto_flags), "ME relay started");
// We need to handle framing here.
// Client sends: [Len:4][Payload...] (Intermediate/Secure)
// We must strip Len and send Payload to ME.
// ME sends: [Payload...]
// We must add [Len:4] and send to Client.
// For Secure mode, Len has padding bit (MSB).
let is_secure = success.proto_tag == crate::protocol::constants::ProtoTag::Secure;
let mut client_closed = false;
let mut server_closed = false;
// Split client_reader/writer to use in select!
// CryptoReader/Writer don't support splitting easily without Arc/Mutex or unsafe,
// but here we are in a loop.
// We can't easily split them because they wrap the underlying stream.
// However, we can use a loop with select! on read and rx.
let mut len_buf = [0u8; 4];
let mut reading_len = true;
let mut current_payload_len = 0;
let mut payload_buf = Vec::new();
let result: Result<()> = loop {
tokio::select! {
// C->S: Read length, then payload
res = async {
if reading_len {
client_reader.read_exact(&mut len_buf).await.map(|_| true)
} else {
// Read payload
// We need to read exactly current_payload_len
if payload_buf.len() < current_payload_len {
let needed = current_payload_len - payload_buf.len();
let mut chunk = vec![0u8; needed];
let n = client_reader.read(&mut chunk).await?;
if n == 0 { return Ok(false); } // EOF
payload_buf.extend_from_slice(&chunk[..n]);
Ok(true)
} else {
Ok(true) // Should not happen
}
}
}, if !client_closed => {
match res {
Ok(true) => {
if reading_len {
// Got length
let raw_len = u32::from_le_bytes(len_buf);
// In secure mode, MSB is padding flag. In intermediate, it's just len.
// But wait, standard intermediate doesn't use MSB for padding.
// Secure mode DOES.
// Let's trust the protocol tag.
let len = if is_secure {
raw_len & 0x7FFFFFFF
} else {
raw_len
};
current_payload_len = len as usize;
// Sanity check
if current_payload_len > 16 * 1024 * 1024 {
debug!(conn_id, len=current_payload_len, "Client sent huge frame");
break Err(ProxyError::Proxy("Frame too large".into()));
}
payload_buf.clear();
payload_buf.reserve(current_payload_len);
reading_len = false;
} else {
// Got some payload data
if payload_buf.len() == current_payload_len {
// Full frame received
trace!(conn_id, bytes = current_payload_len, "C->ME (Frame complete)");
stats.add_user_octets_from(&user, current_payload_len as u64);
// Send to ME
// Note: In secure mode, we send the PADDING bytes too?
// Erlang mtp_intermediate: strips 4 bytes len.
// Erlang mtp_secure: strips 4 bytes len.
// The payload includes the padding if it was added?
// Actually, secure layer (mtp_secure.erl) handles padding removal?
// No, mtp_secure just sets padding=>true for intermediate codec.
// The intermediate codec (mtp_intermediate.erl) just extracts the packet.
// The packet passed to RPC is the payload.
// If secure mode adds random padding at the end, it is part of the payload
// that ME receives?
// Let's look at C code.
// ext-server.c: reads packet_len.
// if (packet_len & 0x80000000) -> has padding.
// It reads the full packet.
// Then it passes it to forward_tcp_query.
// So YES, we send the full payload including padding to ME.
if let Err(e) = me_pool.send_proxy_req(
conn_id, peer, our_addr, &payload_buf, proto_flags
).await {
break Err(e);
}
// Reset for next frame
reading_len = true;
}
}
}
Ok(false) => {
// EOF
debug!(conn_id, "Client EOF");
client_closed = true;
let _ = me_pool.send_close(conn_id).await;
if server_closed { break Ok(()); }
}
Err(e) => {
debug!(conn_id, error = %e, "Client read error");
break Err(ProxyError::Io(e));
}
}
}
// S->C: ME sends data, we wrap and send to client
me_msg = me_rx.recv(), if !server_closed => {
match me_msg {
Some(MeResponse::Data(data)) => {
trace!(conn_id, bytes = data.len(), "ME->C");
stats.add_user_octets_to(&user, data.len() as u64);
// Wrap in intermediate frame
let len = data.len() as u32;
// For secure mode, we might need to add padding?
// C code: forward_mtproto_packet -> just sends data.
// But wait, C code adds framing in net-tcp-rpc-ext-server.c?
// No, forward_tcp_query sends RPC_PROXY_REQ.
// ME sends RPC_PROXY_ANS.
// The data in ANS is the MTProto packet.
// We need to send it to client.
// If client is Intermediate/Secure, we MUST add the 4-byte length prefix.
// Secure mode: usually we don't ADD padding on response, we just send valid packets.
// But we MUST send the length.
if let Err(e) = client_writer.write_all(&len.to_le_bytes()).await {
break Err(ProxyError::Io(e));
}
if let Err(e) = client_writer.write_all(&data).await {
break Err(ProxyError::Io(e));
}
if let Err(e) = client_writer.flush().await {
break Err(ProxyError::Io(e));
}
}
Some(MeResponse::Ack(_)) => {
trace!(conn_id, "ME ACK");
}
Some(MeResponse::Close) => {
debug!(conn_id, "ME sent CLOSE");
server_closed = true;
if client_closed { break Ok(()); }
// We should probably close client connection too
break Ok(());
}
None => {
debug!(conn_id, "ME channel closed");
server_closed = true;
if client_closed { break Ok(()); }
break Err(ProxyError::Proxy("ME connection lost".into()));
}
}
}
// Decide: middle proxy or direct
if config.general.use_middle_proxy {
if let Some(ref pool) = me_pool {
return handle_via_middle_proxy(
client_reader,
client_writer,
success,
pool.clone(),
stats,
config,
buffer_pool,
local_addr,
)
.await;
}
};
warn!("use_middle_proxy=true but MePool not initialized, falling back to direct");
}
// Cleanup
debug!(user = %user, conn_id, "ME relay cleanup");
me_pool.registry().unregister(conn_id).await;
stats.decrement_user_curr_connects(&user);
result
// Direct mode (original behavior)
handle_via_direct(
client_reader,
client_writer,
success,
upstream_manager,
stats,
config,
buffer_pool,
rng,
)
.await
}
fn check_user_limits_static(user: &str, config: &ProxyConfig, stats: &Stats) -> Result<()> {
if let Some(expiration) = config.access.user_expirations.get(user) {
if chrono::Utc::now() > *expiration {
return Err(ProxyError::UserExpired { user: user.to_string() });
}
if let Some(expiration) = config.access.user_expirations.get(user) {
if chrono::Utc::now() > *expiration {
return Err(ProxyError::UserExpired {
user: user.to_string(),
});
}
if let Some(limit) = config.access.user_max_tcp_conns.get(user) {
if stats.get_user_curr_connects(user) >= *limit as u64 {
return Err(ProxyError::ConnectionLimitExceeded { user: user.to_string() });
}
}
if let Some(quota) = config.access.user_data_quota.get(user) {
if stats.get_user_total_octets(user) >= *quota {
return Err(ProxyError::DataQuotaExceeded { user: user.to_string() });
}
}
Ok(())
}
/// Resolve DC index to target address (used only in direct mode)
fn get_dc_addr_static(dc_idx: i16, config: &ProxyConfig) -> Result<SocketAddr> {
let datacenters = if config.general.prefer_ipv6 {
&*TG_DATACENTERS_V6
} else {
&*TG_DATACENTERS_V4
};
let num_dcs = datacenters.len();
let dc_key = dc_idx.to_string();
if let Some(addr_str) = config.dc_overrides.get(&dc_key) {
match addr_str.parse::<SocketAddr>() {
Ok(addr) => {
debug!(dc_idx = dc_idx, addr = %addr, "Using DC override from config");
return Ok(addr);
}
Err(_) => {
warn!(dc_idx = dc_idx, addr_str = %addr_str,
"Invalid DC override address in config, ignoring");
}
}
}
let abs_dc = dc_idx.unsigned_abs() as usize;
if abs_dc >= 1 && abs_dc <= num_dcs {
return Ok(SocketAddr::new(datacenters[abs_dc - 1], TG_DATACENTER_PORT));
}
let default_dc = config.default_dc.unwrap_or(2) as usize;
let fallback_idx = if default_dc >= 1 && default_dc <= num_dcs {
default_dc - 1
} else {
1
};
info!(
original_dc = dc_idx,
fallback_dc = (fallback_idx + 1) as u16,
fallback_addr = %datacenters[fallback_idx],
"Special DC ---> default_cluster"
);
Ok(SocketAddr::new(datacenters[fallback_idx], TG_DATACENTER_PORT))
}
/// Perform obfuscated handshake with Telegram DC (direct mode only)
async fn do_tg_handshake_static(
mut stream: TcpStream,
success: &HandshakeSuccess,
config: &ProxyConfig,
rng: &SecureRandom,
) -> Result<(CryptoReader<tokio::net::tcp::OwnedReadHalf>, CryptoWriter<tokio::net::tcp::OwnedWriteHalf>)> {
let (nonce, tg_enc_key, tg_enc_iv, tg_dec_key, tg_dec_iv) = generate_tg_nonce(
success.proto_tag,
success.dc_idx,
&success.dec_key,
success.dec_iv,
rng,
config.general.fast_mode,
);
let (encrypted_nonce, mut tg_encryptor, tg_decryptor) = encrypt_tg_nonce_with_ciphers(&nonce);
debug!(
peer = %success.peer,
nonce_head = %hex::encode(&nonce[..16]),
"Sending nonce to Telegram"
);
stream.write_all(&encrypted_nonce).await?;
stream.flush().await?;
let (read_half, write_half) = stream.into_split();
Ok((
CryptoReader::new(read_half, tg_decryptor),
CryptoWriter::new(write_half, tg_encryptor),
))
if let Some(limit) = config.access.user_max_tcp_conns.get(user) {
if stats.get_user_curr_connects(user) >= *limit as u64 {
return Err(ProxyError::ConnectionLimitExceeded {
user: user.to_string(),
});
}
}
if let Some(quota) = config.access.user_data_quota.get(user) {
if stats.get_user_total_octets(user) >= *quota {
return Err(ProxyError::DataQuotaExceeded {
user: user.to_string(),
});
}
}
Ok(())
}
}