Upstreams for ME + Egress-data from UM + ME-over-SOCKS + Bind-aware STUN

src/transport/middle_proxy/handshake.rs

@@ -17,10 +17,12 @@ use tracing::{debug, info, warn};
 use crate::crypto::{SecureRandom, build_middleproxy_prekey, derive_middleproxy_keys, sha256};
 use crate::error::{ProxyError, Result};
 use crate::network::IpFamily;
+use crate::network::probe::is_bogon;
 use crate::protocol::constants::{
     ME_CONNECT_TIMEOUT_SECS, ME_HANDSHAKE_TIMEOUT_SECS, RPC_CRYPTO_AES_U32,
     RPC_HANDSHAKE_ERROR_U32, rpc_crypto_flags,
 };
+use crate::transport::{UpstreamEgressInfo, UpstreamRouteKind};
 
 use super::codec::{
     RpcChecksumMode, build_handshake_payload, build_nonce_payload, build_rpc_frame,
@@ -43,33 +45,125 @@ pub(crate) struct HandshakeOutput {
 }
 
 impl MePool {
-    /// TCP connect with timeout + return RTT in milliseconds.
-    pub(crate) async fn connect_tcp(&self, addr: SocketAddr) -> Result<(TcpStream, f64)> {
-        let start = Instant::now();
-        let connect_fut = async {
-            if addr.is_ipv6()
-                && let Some(v6) = self.detected_ipv6
-            {
-                match TcpSocket::new_v6() {
-                    Ok(sock) => {
-                        if let Err(e) = sock.bind(SocketAddr::new(IpAddr::V6(v6), 0)) {
-                            debug!(error = %e, bind_ip = %v6, "ME IPv6 bind failed, falling back to default bind");
-                        } else {
-                            match sock.connect(addr).await {
-                                Ok(stream) => return Ok(stream),
-                                Err(e) => debug!(error = %e, target = %addr, "ME IPv6 bound connect failed, retrying default connect"),
-                            }
-                        }
+    async fn resolve_dc_idx_for_endpoint(&self, addr: SocketAddr) -> Option<i16> {
+        if addr.is_ipv4() {
+            let map = self.proxy_map_v4.read().await;
+            for (dc, addrs) in map.iter() {
+                if addrs
+                    .iter()
+                    .any(|(ip, port)| SocketAddr::new(*ip, *port) == addr)
+                {
+                    let abs_dc = dc.abs();
+                    if abs_dc > 0
+                        && let Ok(dc_idx) = i16::try_from(abs_dc)
+                    {
+                        return Some(dc_idx);
                     }
-                    Err(e) => debug!(error = %e, "ME IPv6 socket creation failed, falling back to default connect"),
                 }
             }
-            TcpStream::connect(addr).await
+        } else {
+            let map = self.proxy_map_v6.read().await;
+            for (dc, addrs) in map.iter() {
+                if addrs
+                    .iter()
+                    .any(|(ip, port)| SocketAddr::new(*ip, *port) == addr)
+                {
+                    let abs_dc = dc.abs();
+                    if abs_dc > 0
+                        && let Ok(dc_idx) = i16::try_from(abs_dc)
+                    {
+                        return Some(dc_idx);
+                    }
+                }
+            }
+        }
+        None
+    }
+
+    fn direct_bind_ip_for_stun(
+        family: IpFamily,
+        upstream_egress: Option<UpstreamEgressInfo>,
+    ) -> Option<IpAddr> {
+        let info = upstream_egress?;
+        if info.route_kind != UpstreamRouteKind::Direct {
+            return None;
+        }
+        match (family, info.direct_bind_ip) {
+            (IpFamily::V4, Some(IpAddr::V4(ip))) => Some(IpAddr::V4(ip)),
+            (IpFamily::V6, Some(IpAddr::V6(ip))) => Some(IpAddr::V6(ip)),
+            _ => None,
+        }
+    }
+
+    fn select_socks_bound_addr(
+        family: IpFamily,
+        upstream_egress: Option<UpstreamEgressInfo>,
+    ) -> Option<SocketAddr> {
+        let info = upstream_egress?;
+        if !matches!(
+            info.route_kind,
+            UpstreamRouteKind::Socks4 | UpstreamRouteKind::Socks5
+        ) {
+            return None;
+        }
+        let bound = info.socks_bound_addr?;
+        let family_matches = matches!(
+            (family, bound.ip()),
+            (IpFamily::V4, IpAddr::V4(_)) | (IpFamily::V6, IpAddr::V6(_))
+        );
+        if !family_matches || is_bogon(bound.ip()) || bound.ip().is_unspecified() {
+            return None;
+        }
+        Some(bound)
+    }
+
+    /// TCP connect with timeout + return RTT in milliseconds.
+    pub(crate) async fn connect_tcp(
+        &self,
+        addr: SocketAddr,
+    ) -> Result<(TcpStream, f64, Option<UpstreamEgressInfo>)> {
+        let start = Instant::now();
+        let (stream, upstream_egress) = if let Some(upstream) = &self.upstream {
+            let dc_idx = self.resolve_dc_idx_for_endpoint(addr).await;
+            let (stream, egress) = timeout(
+                Duration::from_secs(ME_CONNECT_TIMEOUT_SECS),
+                upstream.connect_with_details(addr, dc_idx, None),
+            )
+            .await
+            .map_err(|_| ProxyError::ConnectionTimeout {
+                addr: addr.to_string(),
+            })??;
+            (stream, Some(egress))
+        } else {
+            let connect_fut = async {
+                if addr.is_ipv6()
+                    && let Some(v6) = self.detected_ipv6
+                {
+                    match TcpSocket::new_v6() {
+                        Ok(sock) => {
+                            if let Err(e) = sock.bind(SocketAddr::new(IpAddr::V6(v6), 0)) {
+                                debug!(error = %e, bind_ip = %v6, "ME IPv6 bind failed, falling back to default bind");
+                            } else {
+                                match sock.connect(addr).await {
+                                    Ok(stream) => return Ok(stream),
+                                    Err(e) => debug!(error = %e, target = %addr, "ME IPv6 bound connect failed, retrying default connect"),
+                                }
+                            }
+                        }
+                        Err(e) => debug!(error = %e, "ME IPv6 socket creation failed, falling back to default connect"),
+                    }
+                }
+                TcpStream::connect(addr).await
+            };
+
+            let stream = timeout(Duration::from_secs(ME_CONNECT_TIMEOUT_SECS), connect_fut)
+                .await
+                .map_err(|_| ProxyError::ConnectionTimeout {
+                    addr: addr.to_string(),
+                })??;
+            (stream, None)
         };
 
-        let stream = timeout(Duration::from_secs(ME_CONNECT_TIMEOUT_SECS), connect_fut)
-            .await
-            .map_err(|_| ProxyError::ConnectionTimeout { addr: addr.to_string() })??;
         let connect_ms = start.elapsed().as_secs_f64() * 1000.0;
         stream.set_nodelay(true).ok();
         if let Err(e) = Self::configure_keepalive(&stream) {
@@ -79,7 +173,7 @@ impl MePool {
         if let Err(e) = Self::configure_user_timeout(stream.as_raw_fd()) {
             warn!(error = %e, "ME TCP_USER_TIMEOUT setup failed");
         }
-        Ok((stream, connect_ms))
+        Ok((stream, connect_ms, upstream_egress))
     }
 
     fn configure_keepalive(stream: &TcpStream) -> std::io::Result<()> {
@@ -117,12 +211,14 @@ impl MePool {
         &self,
         stream: TcpStream,
         addr: SocketAddr,
+        upstream_egress: Option<UpstreamEgressInfo>,
         rng: &SecureRandom,
     ) -> Result<HandshakeOutput> {
         let hs_start = Instant::now();
 
         let local_addr = stream.local_addr().map_err(ProxyError::Io)?;
-        let peer_addr = stream.peer_addr().map_err(ProxyError::Io)?;
+        let transport_peer_addr = stream.peer_addr().map_err(ProxyError::Io)?;
+        let peer_addr = addr;
 
         let _ = self.maybe_detect_nat_ip(local_addr.ip()).await;
         let family = if local_addr.ip().is_ipv4() {
@@ -130,8 +226,12 @@ impl MePool {
         } else {
             IpFamily::V6
         };
-        let reflected = if self.nat_probe {
-            self.maybe_reflect_public_addr(family).await
+        let socks_bound_addr = Self::select_socks_bound_addr(family, upstream_egress);
+        let reflected = if let Some(bound) = socks_bound_addr {
+            Some(bound)
+        } else if self.nat_probe {
+            let bind_ip = Self::direct_bind_ip_for_stun(family, upstream_egress);
+            self.maybe_reflect_public_addr(family, bind_ip).await
         } else {
             None
         };
@@ -197,7 +297,9 @@ impl MePool {
             %local_addr_nat,
             reflected_ip = reflected.map(|r| r.ip()).as_ref().map(ToString::to_string),
             %peer_addr,
+            %transport_peer_addr,
             %peer_addr_nat,
+            socks_bound_addr = socks_bound_addr.map(|v| v.to_string()),
             key_selector = format_args!("0x{ks:08x}"),
             crypto_schema = format_args!("0x{schema:08x}"),
             skew_secs = skew,
@@ -206,7 +308,11 @@ impl MePool {
 
         let ts_bytes = crypto_ts.to_le_bytes();
         let server_port_bytes = peer_addr_nat.port().to_le_bytes();
-        let client_port_bytes = local_addr_nat.port().to_le_bytes();
+        let client_port_for_kdf = socks_bound_addr
+            .map(|bound| bound.port())
+            .filter(|port| *port != 0)
+            .unwrap_or(local_addr_nat.port());
+        let client_port_bytes = client_port_for_kdf.to_le_bytes();
 
         let server_ip = extract_ip_material(peer_addr_nat);
         let client_ip = extract_ip_material(local_addr_nat);