proxychains_masq/

chain.rs

use std::{
    net::{IpAddr, Ipv4Addr, Ipv6Addr},
    sync::atomic::{AtomicUsize, Ordering},
    time::Duration,
};

use anyhow::{bail, Context, Result};
use rand::{rngs::OsRng, seq::SliceRandom};
use tokio::{net::TcpStream, time::timeout};
use tracing::info;

use crate::proxy::{http, https, raw, socks4, socks5, BoxStream, Target};

// ─── Public types re-exported from config ─────────────────────────────────────

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ChainType {
    Strict,
    Dynamic,
    Random,
    RoundRobin,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ProxyType {
    Socks4,
    Socks5,
    Http,
    /// HTTP CONNECT over TLS; certificate validation is skipped so that
    /// self-signed corporate proxy certs are accepted.
    Https,
    Raw,
}

/// A single proxy in the chain.
#[derive(Debug, Clone)]
pub struct ProxyEntry {
    pub proxy_type: ProxyType,
    pub addr: IpAddr,
    pub port: u16,
    pub username: Option<String>,
    pub password: Option<String>,
}

/// A localnet exclusion rule.
#[derive(Debug, Clone)]
pub struct LocalNet {
    pub addr: IpAddr,
    pub mask_v4: Option<Ipv4Addr>,
    pub prefix_v6: Option<u8>,
    pub port: Option<u16>,
}

/// A DNAT rewrite rule.
#[derive(Debug, Clone)]
pub struct DnatRule {
    pub orig_addr: Ipv4Addr,
    pub orig_port: Option<u16>,
    pub new_addr: Ipv4Addr,
    pub new_port: Option<u16>,
}

/// Configuration for [`ChainEngine`].
#[derive(Debug, Clone)]
pub struct ChainConfig {
    pub proxies: Vec<ProxyEntry>,
    pub chain_type: ChainType,
    /// Number of proxies to use per connection (for Random / RoundRobin).
    pub chain_len: usize,
    pub connect_timeout: Duration,
    pub localnets: Vec<LocalNet>,
    pub dnats: Vec<DnatRule>,
}

impl Default for ChainConfig {
    fn default() -> Self {
        ChainConfig {
            proxies: Vec::new(),
            chain_type: ChainType::Dynamic,
            chain_len: 1,
            connect_timeout: Duration::from_secs(10),
            localnets: Vec::new(),
            dnats: Vec::new(),
        }
    }
}

// ─── ChainEngine ─────────────────────────────────────────────────────────────

/// Routes outbound TCP connections through a configurable chain of proxies.
pub struct ChainEngine {
    config: ChainConfig,
    /// Shared offset for round-robin mode.
    rr_offset: AtomicUsize,
}

impl ChainEngine {
    /// Create a new engine from `config`.
    pub fn new(config: ChainConfig) -> Self {
        ChainEngine {
            config,
            rr_offset: AtomicUsize::new(0),
        }
    }

    /// Open a connection to `target` through the proxy chain.
    ///
    /// If the target matches a `localnet` rule, a direct connection is made.
    /// DNAT rewrites are applied before localnet and proxy selection.
    ///
    /// Returns a [`BoxStream`] so that TLS-wrapped hops (HTTPS proxies) and
    /// plain TCP hops share the same return type.
    pub async fn connect(&self, target: Target) -> Result<BoxStream> {
        let target = self.apply_dnat(target);

        if self.is_localnet(&target) {
            return self.direct_connect(&target).await;
        }

        match self.config.chain_type {
            ChainType::Strict => self.connect_strict(target).await,
            ChainType::Dynamic => self.connect_dynamic(target).await,
            ChainType::Random => self.connect_random(target).await,
            ChainType::RoundRobin => self.connect_round_robin(target).await,
        }
    }

    // ── Chain modes ───────────────────────────────────────────────────────────
    //
    // All four modes use the same anchor-advance strategy via `chain_proxies`:
    // each proxy in the (mode-specific) selection is tried as the chain entry
    // point in order; if the TCP connect or any subsequent handshake fails, the
    // anchor advances to the next candidate.  This skips broken proxies in all
    // modes without risking infinite loops.

    async fn connect_strict(&self, target: Target) -> Result<BoxStream> {
        let refs: Vec<&ProxyEntry> = self.config.proxies.iter().collect();
        if refs.is_empty() {
            bail!("strict_chain: no proxies configured");
        }
        self.chain_proxies(&refs, target, "strict")
            .await
            .context("strict_chain")
    }

    async fn connect_dynamic(&self, target: Target) -> Result<BoxStream> {
        let refs: Vec<&ProxyEntry> = self.config.proxies.iter().collect();
        if refs.is_empty() {
            bail!("dynamic_chain: no proxies configured");
        }
        self.chain_proxies(&refs, target, "dynamic")
            .await
            .context("dynamic_chain")
    }

    async fn connect_random(&self, target: Target) -> Result<BoxStream> {
        let chain_len = self.config.chain_len.max(1);
        let proxies = &self.config.proxies;
        if proxies.len() < chain_len {
            bail!(
                "random_chain: need {chain_len} proxies, only {} available",
                proxies.len()
            );
        }
        // OsRng reads directly from the OS entropy source (getrandom) on every
        // call, guaranteeing different selections across runs and across
        // connections within a run.  It is also Send, so the future remains
        // Send without any drop-before-await gymnastics.
        let mut selected: Vec<&ProxyEntry> = proxies.iter().collect();
        selected.shuffle(&mut OsRng);
        selected.truncate(chain_len);
        self.chain_proxies(&selected, target, "random")
            .await
            .context("random_chain")
    }

    async fn connect_round_robin(&self, target: Target) -> Result<BoxStream> {
        let chain_len = self.config.chain_len.max(1);
        let proxies = &self.config.proxies;
        if proxies.is_empty() {
            bail!("round_robin_chain: no proxies");
        }
        let n = proxies.len();
        let offset = self.rr_offset.fetch_add(chain_len, Ordering::SeqCst) % n;
        let selected: Vec<&ProxyEntry> =
            (0..chain_len).map(|i| &proxies[(offset + i) % n]).collect();
        self.chain_proxies(&selected, target, "round-robin")
            .await
            .context("round_robin_chain")
    }

    // ── Helpers ───────────────────────────────────────────────────────────────

    /// Build a proxy chain from a pre-selected proxy slice, skipping broken hops.
    ///
    /// Tries each proxy in `proxies` as the chain anchor in order.  If the TCP
    /// connect succeeds but a downstream handshake fails, the anchor advances so
    /// the same dead proxy is never retried as an entry point (guaranteeing
    /// termination in O(n) attempts).  Succeeds as long as at least one proxy
    /// in the slice can reach the target.
    ///
    /// On success logs the chain path at `info` level:
    /// `|<label>-chain| proxy1:port → proxy2:port → target OK`
    async fn chain_proxies(
        &self,
        proxies: &[&ProxyEntry],
        target: Target,
        label: &str,
    ) -> Result<BoxStream> {
        if proxies.is_empty() {
            bail!("chain_proxies: empty proxy list");
        }
        for (anchor, proxy) in proxies.iter().enumerate() {
            let stream = match self.tcp_connect(proxy.addr, proxy.port).await {
                Ok(s) => s,
                Err(_) => continue, // anchor unreachable; try the next one
            };
            match chain_from(stream, anchor, proxies, target.clone()).await {
                Ok(s) => {
                    // Log the successful chain path so the user can see which
                    // proxies were actually used (matching proxychains-ng style).
                    let hops = proxies[anchor..]
                        .iter()
                        .map(|p| format!("{}:{}", p.addr, p.port))
                        .collect::<Vec<_>>()
                        .join(" → ");
                    info!("|{label}-chain| {hops} → {target}");
                    return Ok(s);
                }
                Err(_) => continue, // dead intermediate hop; advance anchor
            }
        }
        bail!("chain_proxies: all proxies are unreachable")
    }

    async fn tcp_connect(&self, addr: IpAddr, port: u16) -> Result<BoxStream> {
        let stream = timeout(
            self.config.connect_timeout,
            TcpStream::connect((addr, port)),
        )
        .await
        .context("tcp connect timed out")?
        .context("tcp connect failed")?;
        Ok(Box::new(stream))
    }

    async fn direct_connect(&self, target: &Target) -> Result<BoxStream> {
        let stream = match target {
            Target::Ip(ip, port) => timeout(
                self.config.connect_timeout,
                TcpStream::connect((*ip, *port)),
            )
            .await
            .context("direct connect timed out")?
            .context("direct connect failed")?,
            Target::Host(h, p) => timeout(
                self.config.connect_timeout,
                TcpStream::connect(format!("{h}:{p}").as_str()),
            )
            .await
            .context("direct connect timed out")?
            .context("direct connect failed")?,
        };
        Ok(Box::new(stream))
    }

    fn apply_dnat(&self, target: Target) -> Target {
        if let Target::Ip(IpAddr::V4(ip), port) = &target {
            for rule in &self.config.dnats {
                if rule.orig_addr == *ip {
                    if let Some(orig_port) = rule.orig_port {
                        if orig_port != *port {
                            continue;
                        }
                    }
                    let new_port = rule.new_port.unwrap_or(*port);
                    return Target::Ip(IpAddr::V4(rule.new_addr), new_port);
                }
            }
        }
        target
    }

    fn is_localnet(&self, target: &Target) -> bool {
        let (ip, port) = match target {
            Target::Ip(ip, p) => (Some(*ip), *p),
            Target::Host(_, p) => (None, *p),
        };
        let Some(ip) = ip else { return false };

        for ln in &self.config.localnets {
            if let Some(p) = ln.port {
                if p != port {
                    continue;
                }
            }
            match (ip, ln.addr) {
                (IpAddr::V4(tip), IpAddr::V4(laddr)) => {
                    if let Some(mask) = ln.mask_v4 {
                        let t = u32::from(tip);
                        let l = u32::from(laddr);
                        let m = u32::from(mask);
                        if (t & m) == (l & m) {
                            return true;
                        }
                    }
                }
                (IpAddr::V6(tip), IpAddr::V6(laddr)) => {
                    if let Some(prefix) = ln.prefix_v6 {
                        if ipv6_match(tip, laddr, prefix) {
                            return true;
                        }
                    }
                }
                _ => {}
            }
        }
        false
    }
}

fn ipv6_match(a: Ipv6Addr, b: Ipv6Addr, prefix: u8) -> bool {
    let a = a.octets();
    let b = b.octets();
    let full = (prefix / 8) as usize;
    let rem = prefix % 8;
    if a[..full] != b[..full] {
        return false;
    }
    if rem > 0 {
        let mask = 0xFFu8 << (8 - rem);
        if (a[full] & mask) != (b[full] & mask) {
            return false;
        }
    }
    true
}

/// Build a [`Target`] pointing to `proxy`'s address (used as intermediate hop).
fn hop_target(proxy: &ProxyEntry) -> Target {
    Target::Ip(proxy.addr, proxy.port)
}

/// Perform the appropriate protocol handshake for `prev_proxy` to connect to `next_target`.
async fn handshake(
    stream: BoxStream,
    prev_proxy: &ProxyEntry,
    next_target: Target,
) -> Result<BoxStream> {
    let user = prev_proxy.username.as_deref();
    let pass = prev_proxy.password.as_deref();
    match prev_proxy.proxy_type {
        ProxyType::Socks4 => socks4::connect(stream, &next_target, user).await,
        ProxyType::Socks5 => socks5::connect(stream, &next_target, user, pass).await,
        ProxyType::Http => http::connect(stream, &next_target, user, pass).await,
        ProxyType::Https => https::connect(stream, &next_target, user, pass, prev_proxy.addr).await,
        ProxyType::Raw => raw::connect(stream, &next_target).await,
    }
}

/// Drive the proxy chain from `proxies[anchor]` outward to `target`.
///
/// The caller has already opened a TCP connection to `proxies[anchor]`.  This
/// function issues CONNECT handshakes for each subsequent hop and finally for
/// `target`, returning the fully-tunneled stream on success.
///
/// Returns `Err` on the first failed handshake so that `chain_proxies` can
/// advance the anchor and retry from the next proxy.
async fn chain_from(
    mut stream: BoxStream,
    anchor: usize,
    proxies: &[&ProxyEntry],
    target: Target,
) -> Result<BoxStream> {
    // Issue CONNECT for each intermediate hop: anchor → anchor+1 → … → last.
    for i in anchor..proxies.len() - 1 {
        stream = handshake(stream, proxies[i], hop_target(proxies[i + 1])).await?;
    }
    // Final handshake: last proxy in the slice → target.
    handshake(stream, proxies[proxies.len() - 1], target).await
}