diff --git a/src/config/mod.rs b/src/config/mod.rs index dbf8afa..84a630a 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -154,6 +154,14 @@ pub struct GeneralConfig { #[serde(default)] pub middle_proxy_nat_ip: Option, + /// Enable STUN-based NAT probing to discover public IP:port for ME KDF. + #[serde(default)] + pub middle_proxy_nat_probe: bool, + + /// Optional STUN server address (host:port) for NAT probing. + #[serde(default)] + pub middle_proxy_nat_stun: Option, + #[serde(default)] pub log_level: LogLevel, } @@ -168,6 +176,8 @@ impl Default for GeneralConfig { ad_tag: None, proxy_secret_path: None, middle_proxy_nat_ip: None, + middle_proxy_nat_probe: false, + middle_proxy_nat_stun: None, log_level: LogLevel::Normal, } } diff --git a/src/crypto/hash.rs b/src/crypto/hash.rs index cf3ba0d..6b3c654 100644 --- a/src/crypto/hash.rs +++ b/src/crypto/hash.rs @@ -55,6 +55,49 @@ pub fn crc32(data: &[u8]) -> u32 { crc32fast::hash(data) } +/// Build the exact prekey buffer used by Telegram Middle Proxy KDF. +/// +/// Returned buffer layout (IPv4): +/// nonce_srv | nonce_clt | clt_ts | srv_ip | clt_port | purpose | clt_ip | srv_port | secret | nonce_srv | [clt_v6 | srv_v6] | nonce_clt +pub fn build_middleproxy_prekey( + nonce_srv: &[u8; 16], + nonce_clt: &[u8; 16], + clt_ts: &[u8; 4], + srv_ip: Option<&[u8]>, + clt_port: &[u8; 2], + purpose: &[u8], + clt_ip: Option<&[u8]>, + srv_port: &[u8; 2], + secret: &[u8], + clt_ipv6: Option<&[u8; 16]>, + srv_ipv6: Option<&[u8; 16]>, +) -> Vec { + const EMPTY_IP: [u8; 4] = [0, 0, 0, 0]; + + let srv_ip = srv_ip.unwrap_or(&EMPTY_IP); + let clt_ip = clt_ip.unwrap_or(&EMPTY_IP); + + let mut s = Vec::with_capacity(256); + s.extend_from_slice(nonce_srv); + s.extend_from_slice(nonce_clt); + s.extend_from_slice(clt_ts); + s.extend_from_slice(srv_ip); + s.extend_from_slice(clt_port); + s.extend_from_slice(purpose); + s.extend_from_slice(clt_ip); + s.extend_from_slice(srv_port); + s.extend_from_slice(secret); + s.extend_from_slice(nonce_srv); + + if let (Some(clt_v6), Some(srv_v6)) = (clt_ipv6, srv_ipv6) { + s.extend_from_slice(clt_v6); + s.extend_from_slice(srv_v6); + } + + s.extend_from_slice(nonce_clt); + s +} + /// Middle Proxy key derivation /// /// Uses MD5 + SHA-1 as mandated by the Telegram Middle Proxy protocol. @@ -73,30 +116,20 @@ pub fn derive_middleproxy_keys( clt_ipv6: Option<&[u8; 16]>, srv_ipv6: Option<&[u8; 16]>, ) -> ([u8; 32], [u8; 16]) { - const EMPTY_IP: [u8; 4] = [0, 0, 0, 0]; - - let srv_ip = srv_ip.unwrap_or(&EMPTY_IP); - let clt_ip = clt_ip.unwrap_or(&EMPTY_IP); - - let mut s = Vec::with_capacity(256); - s.extend_from_slice(nonce_srv); - s.extend_from_slice(nonce_clt); - s.extend_from_slice(clt_ts); - s.extend_from_slice(srv_ip); - s.extend_from_slice(clt_port); - s.extend_from_slice(purpose); - s.extend_from_slice(clt_ip); - s.extend_from_slice(srv_port); - s.extend_from_slice(secret); - s.extend_from_slice(nonce_srv); - - if let (Some(clt_v6), Some(srv_v6)) = (clt_ipv6, srv_ipv6) { - s.extend_from_slice(clt_v6); - s.extend_from_slice(srv_v6); - } - - s.extend_from_slice(nonce_clt); - + let s = build_middleproxy_prekey( + nonce_srv, + nonce_clt, + clt_ts, + srv_ip, + clt_port, + purpose, + clt_ip, + srv_port, + secret, + clt_ipv6, + srv_ipv6, + ); + let md5_1 = md5(&s[1..]); let sha1_sum = sha1(&s); let md5_2 = md5(&s[2..]); @@ -106,4 +139,40 @@ pub fn derive_middleproxy_keys( key[12..].copy_from_slice(&sha1_sum); (key, md5_2) -} \ No newline at end of file +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn middleproxy_prekey_sha_is_stable() { + let nonce_srv = [0x11u8; 16]; + let nonce_clt = [0x22u8; 16]; + let clt_ts = 0x44332211u32.to_le_bytes(); + let srv_ip = Some([149u8, 154, 175, 50].as_ref()); + let clt_ip = Some([10u8, 0, 0, 1].as_ref()); + let clt_port = 0x1f90u16.to_le_bytes(); // 8080 + let srv_port = 0x22b8u16.to_le_bytes(); // 8888 + let secret = vec![0x55u8; 128]; + + let prekey = build_middleproxy_prekey( + &nonce_srv, + &nonce_clt, + &clt_ts, + srv_ip, + &clt_port, + b"CLIENT", + clt_ip, + &srv_port, + &secret, + None, + None, + ); + let digest = sha256(&prekey); + assert_eq!( + hex::encode(digest), + "a4595b75f1f610f2575ace802ddc65c91b5acef3b0e0d18189e0c7c9f787d15c" + ); + } +} diff --git a/src/crypto/mod.rs b/src/crypto/mod.rs index ee97b27..40951c6 100644 --- a/src/crypto/mod.rs +++ b/src/crypto/mod.rs @@ -5,5 +5,5 @@ pub mod hash; pub mod random; pub use aes::{AesCtr, AesCbc}; -pub use hash::{sha256, sha256_hmac, sha1, md5, crc32, derive_middleproxy_keys}; -pub use random::SecureRandom; \ No newline at end of file +pub use hash::{sha256, sha256_hmac, sha1, md5, crc32, derive_middleproxy_keys, build_middleproxy_prekey}; +pub use random::SecureRandom; diff --git a/src/main.rs b/src/main.rs index e03f600..719d39f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -229,7 +229,13 @@ async fn main() -> std::result::Result<(), Box> { "Proxy-secret loaded" ); - let pool = MePool::new(proxy_tag, proxy_secret, config.general.middle_proxy_nat_ip); + let pool = MePool::new( + proxy_tag, + proxy_secret, + config.general.middle_proxy_nat_ip, + config.general.middle_proxy_nat_probe, + config.general.middle_proxy_nat_stun.clone(), + ); match pool.init(2, &rng).await { Ok(()) => { diff --git a/src/transport/middle_proxy/pool.rs b/src/transport/middle_proxy/pool.rs index 9a978cc..650a029 100644 --- a/src/transport/middle_proxy/pool.rs +++ b/src/transport/middle_proxy/pool.rs @@ -12,7 +12,7 @@ use tokio::sync::{Mutex, RwLock}; use tokio::time::{Instant, timeout}; use tracing::{debug, info, warn}; -use crate::crypto::{SecureRandom, derive_middleproxy_keys, sha256}; +use crate::crypto::{SecureRandom, build_middleproxy_prekey, derive_middleproxy_keys, sha256}; use crate::error::{ProxyError, Result}; use crate::protocol::constants::*; @@ -35,6 +35,8 @@ pub struct MePool { proxy_secret: Vec, pub(super) nat_ip_cfg: Option, pub(super) nat_ip_detected: OnceLock, + pub(super) nat_probe: bool, + pub(super) nat_stun: Option, pool_size: usize, } @@ -43,6 +45,8 @@ impl MePool { proxy_tag: Option>, proxy_secret: Vec, nat_ip: Option, + nat_probe: bool, + nat_stun: Option, ) -> Arc { Arc::new(Self { registry: Arc::new(ConnRegistry::new()), @@ -52,6 +56,8 @@ impl MePool { proxy_secret, nat_ip_cfg: nat_ip, nat_ip_detected: OnceLock::new(), + nat_probe, + nat_stun, pool_size: 2, }) } @@ -143,7 +149,12 @@ impl MePool { let local_addr = stream.local_addr().map_err(ProxyError::Io)?; let peer_addr = stream.peer_addr().map_err(ProxyError::Io)?; let _ = self.maybe_detect_nat_ip(local_addr.ip()).await; - let local_addr_nat = self.translate_our_addr(local_addr); + let reflected = if self.nat_probe { + self.maybe_reflect_public_addr().await + } else { + None + }; + let local_addr_nat = self.translate_our_addr_with_reflection(local_addr, reflected); let peer_addr_nat = SocketAddr::new(self.translate_ip_for_nat(peer_addr.ip()), peer_addr.port()); let (mut rd, mut wr) = tokio::io::split(stream); @@ -205,6 +216,7 @@ impl MePool { info!( %local_addr, %local_addr_nat, + reflected_ip = reflected.map(|r| r.ip()).as_ref().map(ToString::to_string), %peer_addr, %peer_addr_nat, key_selector = format_args!("0x{ks:08x}"), @@ -237,6 +249,38 @@ impl MePool { } }; + let diag_level: u8 = std::env::var("ME_DIAG") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(0); + + let prekey_client = build_middleproxy_prekey( + &srv_nonce, + &my_nonce, + &ts_bytes, + srv_ip_opt.as_ref().map(|x| &x[..]), + &client_port_bytes, + b"CLIENT", + clt_ip_opt.as_ref().map(|x| &x[..]), + &server_port_bytes, + secret, + clt_v6_opt.as_ref(), + srv_v6_opt.as_ref(), + ); + let prekey_server = build_middleproxy_prekey( + &srv_nonce, + &my_nonce, + &ts_bytes, + srv_ip_opt.as_ref().map(|x| &x[..]), + &client_port_bytes, + b"SERVER", + clt_ip_opt.as_ref().map(|x| &x[..]), + &server_port_bytes, + secret, + clt_v6_opt.as_ref(), + srv_v6_opt.as_ref(), + ); + let (wk, wi) = derive_middleproxy_keys( &srv_nonce, &my_nonce, @@ -264,24 +308,39 @@ impl MePool { srv_v6_opt.as_ref(), ); - let diag = std::env::var("ME_DIAG").map(|v| v == "1").unwrap_or(false); let hs_payload = build_handshake_payload(hs_our_ip, local_addr.port(), hs_peer_ip, peer_addr.port()); let hs_frame = build_rpc_frame(-1, &hs_payload); - if diag { + if diag_level >= 1 { info!( write_key = %hex_dump(&wk), write_iv = %hex_dump(&wi), read_key = %hex_dump(&rk), read_iv = %hex_dump(&ri), + srv_ip = %srv_ip_opt.map(|ip| hex_dump(&ip)).unwrap_or_default(), + clt_ip = %clt_ip_opt.map(|ip| hex_dump(&ip)).unwrap_or_default(), + srv_port = %hex_dump(&server_port_bytes), + clt_port = %hex_dump(&client_port_bytes), + crypto_ts = %hex_dump(&ts_bytes), + nonce_srv = %hex_dump(&srv_nonce), + nonce_clt = %hex_dump(&my_nonce), + prekey_sha256_client = %hex_dump(&sha256(&prekey_client)), + prekey_sha256_server = %hex_dump(&sha256(&prekey_server)), hs_plain = %hex_dump(&hs_frame), proxy_secret_sha256 = %hex_dump(&sha256(secret)), "ME diag: derived keys and handshake plaintext" ); } + if diag_level >= 2 { + info!( + prekey_client = %hex_dump(&prekey_client), + prekey_server = %hex_dump(&prekey_server), + "ME diag: full prekey buffers" + ); + } let (encrypted_hs, write_iv) = cbc_encrypt_padded(&wk, &wi, &hs_frame)?; - if diag { + if diag_level >= 1 { info!( hs_cipher = %hex_dump(&encrypted_hs), "ME diag: handshake ciphertext" diff --git a/src/transport/middle_proxy/pool_nat.rs b/src/transport/middle_proxy/pool_nat.rs index 6891919..2a37ec4 100644 --- a/src/transport/middle_proxy/pool_nat.rs +++ b/src/transport/middle_proxy/pool_nat.rs @@ -31,6 +31,26 @@ impl MePool { } } + pub(super) fn translate_our_addr_with_reflection( + &self, + addr: std::net::SocketAddr, + reflected: Option, + ) -> std::net::SocketAddr { + let ip = if let Some(r) = reflected { + // Use reflected IP (not port) only when local address is non-public. + if is_privateish(addr.ip()) || addr.ip().is_loopback() || addr.ip().is_unspecified() { + r.ip() + } else { + self.translate_ip_for_nat(addr.ip()) + } + } else { + self.translate_ip_for_nat(addr.ip()) + }; + + // Keep the kernel-assigned TCP source port; STUN port can differ. + std::net::SocketAddr::new(ip, addr.port()) + } + pub(super) async fn maybe_detect_nat_ip(&self, local_ip: IpAddr) -> Option { if self.nat_ip_cfg.is_some() { return self.nat_ip_cfg; @@ -57,6 +77,25 @@ impl MePool { } } } + + pub(super) async fn maybe_reflect_public_addr(&self) -> Option { + let stun_addr = self + .nat_stun + .clone() + .unwrap_or_else(|| "stun.l.google.com:19302".to_string()); + match fetch_stun_binding(&stun_addr).await { + Ok(sa) => { + if let Some(sa) = sa { + info!(%sa, "NAT probe: reflected address"); + } + sa + } + Err(e) => { + warn!(error = %e, "NAT probe failed"); + None + } + } + } } async fn fetch_public_ipv4() -> Result> { @@ -72,6 +111,87 @@ async fn fetch_public_ipv4() -> Result> { Ok(ip) } +async fn fetch_stun_binding(stun_addr: &str) -> Result> { + use rand::RngCore; + use tokio::net::UdpSocket; + + let socket = UdpSocket::bind("0.0.0.0:0") + .await + .map_err(|e| ProxyError::Proxy(format!("STUN bind failed: {e}")))?; + socket + .connect(stun_addr) + .await + .map_err(|e| ProxyError::Proxy(format!("STUN connect failed: {e}")))?; + + // Build minimal Binding Request. + let mut req = vec![0u8; 20]; + req[0..2].copy_from_slice(&0x0001u16.to_be_bytes()); // Binding Request + req[2..4].copy_from_slice(&0u16.to_be_bytes()); // length + req[4..8].copy_from_slice(&0x2112A442u32.to_be_bytes()); // magic cookie + rand::thread_rng().fill_bytes(&mut req[8..20]); + + socket + .send(&req) + .await + .map_err(|e| ProxyError::Proxy(format!("STUN send failed: {e}")))?; + + let mut buf = [0u8; 128]; + let n = socket + .recv(&mut buf) + .await + .map_err(|e| ProxyError::Proxy(format!("STUN recv failed: {e}")))?; + if n < 20 { + return Ok(None); + } + + // Parse attributes. + let mut idx = 20; + while idx + 4 <= n { + let atype = u16::from_be_bytes(buf[idx..idx + 2].try_into().unwrap()); + let alen = u16::from_be_bytes(buf[idx + 2..idx + 4].try_into().unwrap()) as usize; + idx += 4; + if idx + alen > n { + break; + } + match atype { + 0x0020 /* XOR-MAPPED-ADDRESS */ | 0x0001 /* MAPPED-ADDRESS */ => { + if alen < 8 { + break; + } + let family = buf[idx + 1]; + if family != 0x01 { + // only IPv4 supported here + break; + } + let port_bytes = [buf[idx + 2], buf[idx + 3]]; + let ip_bytes = [buf[idx + 4], buf[idx + 5], buf[idx + 6], buf[idx + 7]]; + + let (port, ip) = if atype == 0x0020 { + let magic = 0x2112A442u32.to_be_bytes(); + let port = u16::from_be_bytes(port_bytes) ^ ((magic[0] as u16) << 8 | magic[1] as u16); + let ip = [ + ip_bytes[0] ^ magic[0], + ip_bytes[1] ^ magic[1], + ip_bytes[2] ^ magic[2], + ip_bytes[3] ^ magic[3], + ]; + (port, ip) + } else { + (u16::from_be_bytes(port_bytes), ip_bytes) + }; + return Ok(Some(std::net::SocketAddr::new( + IpAddr::V4(Ipv4Addr::new(ip[0], ip[1], ip[2], ip[3])), + port, + ))); + } + _ => {} + } + idx += (alen + 3) & !3; // 4-byte alignment + } + + Ok(None) +} + fn is_privateish(ip: IpAddr) -> bool { match ip { IpAddr::V4(v4) => v4.is_private() || v4.is_link_local(), diff --git a/src/transport/middle_proxy/secret.rs b/src/transport/middle_proxy/secret.rs index 5b201fd..b998411 100644 --- a/src/transport/middle_proxy/secret.rs +++ b/src/transport/middle_proxy/secret.rs @@ -8,42 +8,47 @@ use crate::error::{ProxyError, Result}; pub async fn fetch_proxy_secret(cache_path: Option<&str>) -> Result> { let cache = cache_path.unwrap_or("proxy-secret"); - if let Ok(metadata) = tokio::fs::metadata(cache).await { - if let Ok(modified) = metadata.modified() { - let age = std::time::SystemTime::now() - .duration_since(modified) - .unwrap_or(Duration::from_secs(u64::MAX)); - if age < Duration::from_secs(86_400) { - if let Ok(data) = tokio::fs::read(cache).await { - if data.len() >= 32 { - info!( - path = cache, - len = data.len(), - age_hours = age.as_secs() / 3600, - "Loaded proxy-secret from cache" - ); - return Ok(data); - } - warn!( - path = cache, - len = data.len(), - "Cached proxy-secret too short" - ); - } + // 1) Try fresh download first. + match download_proxy_secret().await { + Ok(data) => { + if let Err(e) = tokio::fs::write(cache, &data).await { + warn!(error = %e, "Failed to cache proxy-secret (non-fatal)"); + } else { + debug!(path = cache, len = data.len(), "Cached proxy-secret"); } + return Ok(data); + } + Err(download_err) => { + warn!(error = %download_err, "Proxy-secret download failed, trying cache/file fallback"); + // Fall through to cache/file. } } - info!("Downloading proxy-secret from core.telegram.org..."); - let data = download_proxy_secret().await?; - - if let Err(e) = tokio::fs::write(cache, &data).await { - warn!(error = %e, "Failed to cache proxy-secret (non-fatal)"); - } else { - debug!(path = cache, len = data.len(), "Cached proxy-secret"); + // 2) Fallback to cache/file regardless of age; require len>=32. + match tokio::fs::read(cache).await { + Ok(data) if data.len() >= 32 => { + let age_hours = tokio::fs::metadata(cache) + .await + .ok() + .and_then(|m| m.modified().ok()) + .and_then(|m| std::time::SystemTime::now().duration_since(m).ok()) + .map(|d| d.as_secs() / 3600); + info!( + path = cache, + len = data.len(), + age_hours, + "Loaded proxy-secret from cache/file after download failure" + ); + Ok(data) + } + Ok(data) => Err(ProxyError::Proxy(format!( + "Cached proxy-secret too short: {} bytes (need >= 32)", + data.len() + ))), + Err(e) => Err(ProxyError::Proxy(format!( + "Failed to read proxy-secret cache after download failure: {e}" + ))), } - - Ok(data) } async fn download_proxy_secret() -> Result> {