NAT + STUN Probes...

src/config/mod.rs

@@ -154,6 +154,14 @@ pub struct GeneralConfig {
     #[serde(default)]
     pub middle_proxy_nat_ip: Option<IpAddr>,
 
+    /// 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<String>,
+
     #[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,
         }
     }

src/crypto/hash.rs

@@ -55,12 +55,11 @@ pub fn crc32(data: &[u8]) -> u32 {
     crc32fast::hash(data)
 }
 
-/// Middle Proxy key derivation
+/// Build the exact prekey buffer used by Telegram Middle Proxy KDF.
 ///
-/// Uses MD5 + SHA-1 as mandated by the Telegram Middle Proxy protocol.
+/// Returned buffer layout (IPv4):
-/// These algorithms are NOT replaceable here — changing them would break
+/// nonce_srv | nonce_clt | clt_ts | srv_ip | clt_port | purpose | clt_ip | srv_port | secret | nonce_srv | [clt_v6 | srv_v6] | nonce_clt
-/// interoperability with Telegram's middle proxy infrastructure.
+pub fn build_middleproxy_prekey(
-pub fn derive_middleproxy_keys(
     nonce_srv: &[u8; 16],
     nonce_clt: &[u8; 16],
     clt_ts: &[u8; 4],
@@ -72,7 +71,7 @@ pub fn derive_middleproxy_keys(
     secret: &[u8],
     clt_ipv6: Option<&[u8; 16]>,
     srv_ipv6: Option<&[u8; 16]>,
-) -> ([u8; 32], [u8; 16]) {
+) -> Vec<u8> {
     const EMPTY_IP: [u8; 4] = [0, 0, 0, 0];
 
     let srv_ip = srv_ip.unwrap_or(&EMPTY_IP);
@@ -96,6 +95,40 @@ pub fn derive_middleproxy_keys(
     }
 
     s.extend_from_slice(nonce_clt);
+    s
+}
+
+/// Middle Proxy key derivation
+///
+/// Uses MD5 + SHA-1 as mandated by the Telegram Middle Proxy protocol.
+/// These algorithms are NOT replaceable here — changing them would break
+/// interoperability with Telegram's middle proxy infrastructure.
+pub fn derive_middleproxy_keys(
+    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]>,
+) -> ([u8; 32], [u8; 16]) {
+    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);
@@ -107,3 +140,39 @@ pub fn derive_middleproxy_keys(
     
     (key, md5_2)
 }
+
+#[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"
+        );
+    }
+}

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 hash::{sha256, sha256_hmac, sha1, md5, crc32, derive_middleproxy_keys, build_middleproxy_prekey};
 pub use random::SecureRandom;

src/main.rs

@@ -229,7 +229,13 @@ async fn main() -> std::result::Result<(), Box<dyn std::error::Error>> {
                     "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(()) => {

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<u8>,
     pub(super) nat_ip_cfg: Option<IpAddr>,
     pub(super) nat_ip_detected: OnceLock<IpAddr>,
+    pub(super) nat_probe: bool,
+    pub(super) nat_stun: Option<String>,
     pool_size: usize,
 }
 
@@ -43,6 +45,8 @@ impl MePool {
         proxy_tag: Option<Vec<u8>>,
         proxy_secret: Vec<u8>,
         nat_ip: Option<IpAddr>,
+        nat_probe: bool,
+        nat_stun: Option<String>,
     ) -> Arc<Self> {
         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"

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>,
+    ) -> 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<IpAddr> {
         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<std::net::SocketAddr> {
+        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<Option<Ipv4Addr>> {
@@ -72,6 +111,87 @@ async fn fetch_public_ipv4() -> Result<Option<Ipv4Addr>> {
     Ok(ip)
 }
 
+async fn fetch_stun_binding(stun_addr: &str) -> Result<Option<std::net::SocketAddr>> {
+    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(),

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<Vec<u8>> {
     let cache = cache_path.unwrap_or("proxy-secret");
 
-    if let Ok(metadata) = tokio::fs::metadata(cache).await {
+    // 1) Try fresh download first.
-        if let Ok(modified) = metadata.modified() {
+    match download_proxy_secret().await {
-            let age = std::time::SystemTime::now()
+        Ok(data) => {
-                .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"
-                    );
-                }
-            }
-        }
-    }
-
-    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");
             }
+            return Ok(data);
+        }
+        Err(download_err) => {
+            warn!(error = %download_err, "Proxy-secret download failed, trying cache/file fallback");
+            // Fall through to cache/file.
+        }
+    }
 
+    // 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}"
+        ))),
+    }
 }
 
 async fn download_proxy_secret() -> Result<Vec<u8>> {