proxychains-masq 0.1.5

mod helpers;

use std::{
    net::{IpAddr, Ipv4Addr},
    sync::{
        atomic::{AtomicU32, Ordering},
        Arc,
    },
};

use helpers::{echo_server, socks5_server};
use proxychains_masq::{
    chain::{ChainConfig, ChainEngine, ChainType, DnatRule, LocalNet, ProxyEntry, ProxyType},
    proxy::Target,
};
use tokio::io::{AsyncReadExt, AsyncWriteExt};

fn socks5_proxy(addr: Ipv4Addr, port: u16) -> ProxyEntry {
    ProxyEntry {
        proxy_type: ProxyType::Socks5,
        addr: IpAddr::V4(addr),
        port,
        username: None,
        password: None,
    }
}

async fn make_proxy(hits: Arc<AtomicU32>) -> (u16, ProxyEntry) {
    let port = socks5_server(None, hits).await;
    let entry = socks5_proxy(Ipv4Addr::LOCALHOST, port);
    (port, entry)
}

// T-17
#[tokio::test]
async fn test_strict_chain_two_hops() {
    let (echo_port, echo_hits) = echo_server().await;
    let a_hits = Arc::new(AtomicU32::new(0));
    let b_hits = Arc::new(AtomicU32::new(0));
    let (_, proxy_a) = make_proxy(a_hits.clone()).await;
    let (_, proxy_b) = make_proxy(b_hits.clone()).await;

    let engine = ChainEngine::new(ChainConfig {
        proxies: vec![proxy_a, proxy_b],
        chain_type: ChainType::Strict,
        ..Default::default()
    });

    let mut stream = engine
        .connect(Target::Ip(IpAddr::V4(Ipv4Addr::LOCALHOST), echo_port))
        .await
        .unwrap();

    stream.write_all(b"strict").await.unwrap();
    let mut buf = [0u8; 6];
    stream.read_exact(&mut buf).await.unwrap();
    assert_eq!(&buf, b"strict");
    assert_eq!(a_hits.load(Ordering::SeqCst), 1);
    assert_eq!(b_hits.load(Ordering::SeqCst), 1);
    let _ = echo_hits;
}

// T-18
#[tokio::test]
async fn test_strict_chain_dead_proxy_fails() {
    let engine = ChainEngine::new(ChainConfig {
        proxies: vec![socks5_proxy(Ipv4Addr::LOCALHOST, 19999)],
        chain_type: ChainType::Strict,
        ..Default::default()
    });
    let result = engine
        .connect(Target::Ip(IpAddr::V4(Ipv4Addr::LOCALHOST), 80))
        .await;
    assert!(result.is_err());
}

// T-19
#[tokio::test]
async fn test_dynamic_chain_skips_dead() {
    let (echo_port, _) = echo_server().await;
    let b_hits = Arc::new(AtomicU32::new(0));
    let (_, proxy_b) = make_proxy(b_hits.clone()).await;

    let dead_proxy = socks5_proxy(Ipv4Addr::LOCALHOST, 19998);

    let engine = ChainEngine::new(ChainConfig {
        proxies: vec![dead_proxy, proxy_b],
        chain_type: ChainType::Dynamic,
        ..Default::default()
    });

    let mut stream = engine
        .connect(Target::Ip(IpAddr::V4(Ipv4Addr::LOCALHOST), echo_port))
        .await
        .unwrap();

    stream.write_all(b"dyn").await.unwrap();
    let mut buf = [0u8; 3];
    stream.read_exact(&mut buf).await.unwrap();
    assert_eq!(&buf, b"dyn");
    assert_eq!(b_hits.load(Ordering::SeqCst), 1);
}

// T-19b: dead first-hop causes anchor to advance; remaining proxies form the chain.
//
// Layout: [dead_x, proxy_a(alive), proxy_c(alive)]
//
// Anchor 0: [dead_x, proxy_a, proxy_c] — TCP connect to dead_x fails → advance anchor.
// Anchor 1: [proxy_a, proxy_c] — proxy_a reachable, proxy_a CONNECTs to proxy_c,
// proxy_c CONNECTs to echo → success.  Both proxy_a and proxy_c are used.
#[tokio::test]
async fn test_dynamic_chain_slides_past_dead_anchor() {
    let (echo_port, _) = echo_server().await;

    let a_hits = Arc::new(AtomicU32::new(0));
    let c_hits = Arc::new(AtomicU32::new(0));
    let (_, proxy_a) = make_proxy(a_hits.clone()).await;
    let (_, proxy_c) = make_proxy(c_hits.clone()).await;

    let dead_x = socks5_proxy(Ipv4Addr::LOCALHOST, 19995);

    let engine = ChainEngine::new(ChainConfig {
        proxies: vec![dead_x, proxy_a, proxy_c],
        chain_type: ChainType::Dynamic,
        chain_len: 2,
        ..Default::default()
    });

    let mut stream = engine
        .connect(Target::Ip(IpAddr::V4(Ipv4Addr::LOCALHOST), echo_port))
        .await
        .unwrap();

    stream.write_all(b"mid").await.unwrap();
    let mut buf = [0u8; 3];
    stream.read_exact(&mut buf).await.unwrap();
    assert_eq!(&buf, b"mid");
    // Window [dead, proxy_a] failed at TCP → slid to [proxy_a, proxy_c]
    assert_eq!(a_hits.load(Ordering::SeqCst), 1, "proxy_a used as first hop of winning window");
    assert_eq!(c_hits.load(Ordering::SeqCst), 1, "proxy_c used as second hop of winning window");
}

// T-20
#[tokio::test]
async fn test_dynamic_chain_all_dead_fails() {
    let engine = ChainEngine::new(ChainConfig {
        proxies: vec![
            socks5_proxy(Ipv4Addr::LOCALHOST, 19997),
            socks5_proxy(Ipv4Addr::LOCALHOST, 19996),
        ],
        chain_type: ChainType::Dynamic,
        ..Default::default()
    });
    let result = engine
        .connect(Target::Ip(IpAddr::V4(Ipv4Addr::LOCALHOST), 80))
        .await;
    assert!(result.is_err());
}

// T-21
#[tokio::test]
async fn test_random_chain_len() {
    let (echo_port, _) = echo_server().await;
    let hits: Vec<Arc<AtomicU32>> = (0..4).map(|_| Arc::new(AtomicU32::new(0))).collect();
    let mut proxies = Vec::new();
    for h in &hits {
        let (_, p) = make_proxy(h.clone()).await;
        proxies.push(p);
    }

    let engine = ChainEngine::new(ChainConfig {
        proxies,
        chain_type: ChainType::Random,
        chain_len: 2,
        ..Default::default()
    });

    // Run 5 connections — each should use exactly 2 proxies
    for _ in 0..5 {
        let mut stream = engine
            .connect(Target::Ip(IpAddr::V4(Ipv4Addr::LOCALHOST), echo_port))
            .await
            .unwrap();
        stream.write_all(b"r").await.unwrap();
        let mut buf = [0u8; 1];
        stream.read_exact(&mut buf).await.unwrap();
    }

    let total_hits: u32 = hits.iter().map(|h| h.load(Ordering::SeqCst)).sum();
    // 5 connections × 2 hops = 10 proxy hits total
    assert_eq!(total_hits, 10);
}

// T-22
#[tokio::test]
async fn test_round_robin_advances() {
    let (echo_port, _) = echo_server().await;
    let hits: Vec<Arc<AtomicU32>> = (0..3).map(|_| Arc::new(AtomicU32::new(0))).collect();
    let mut proxies = Vec::new();
    for h in &hits {
        let (_, p) = make_proxy(h.clone()).await;
        proxies.push(p);
    }

    let engine = ChainEngine::new(ChainConfig {
        proxies,
        chain_type: ChainType::RoundRobin,
        chain_len: 1,
        ..Default::default()
    });

    // 3 connections should each use a different proxy
    for _ in 0..3 {
        let mut stream = engine
            .connect(Target::Ip(IpAddr::V4(Ipv4Addr::LOCALHOST), echo_port))
            .await
            .unwrap();
        stream.write_all(b"rr").await.unwrap();
        let mut buf = [0u8; 2];
        stream.read_exact(&mut buf).await.unwrap();
    }

    for h in &hits {
        assert_eq!(
            h.load(Ordering::SeqCst),
            1,
            "each proxy should be used exactly once"
        );
    }
}

// T-23
#[tokio::test]
async fn test_localnet_bypass_goes_direct() {
    let (echo_port, _) = echo_server().await;
    let proxy_hits = Arc::new(AtomicU32::new(0));
    let (_, proxy) = make_proxy(proxy_hits.clone()).await;

    let engine = ChainEngine::new(ChainConfig {
        proxies: vec![proxy],
        chain_type: ChainType::Strict,
        localnets: vec![LocalNet {
            addr: IpAddr::V4(Ipv4Addr::new(127, 0, 0, 0)),
            mask_v4: Some(Ipv4Addr::new(255, 0, 0, 0)),
            prefix_v6: None,
            port: None,
        }],
        ..Default::default()
    });

    let mut stream = engine
        .connect(Target::Ip(IpAddr::V4(Ipv4Addr::LOCALHOST), echo_port))
        .await
        .unwrap();

    stream.write_all(b"local").await.unwrap();
    let mut buf = [0u8; 5];
    stream.read_exact(&mut buf).await.unwrap();
    assert_eq!(&buf, b"local");
    // Proxy must NOT have been used
    assert_eq!(proxy_hits.load(Ordering::SeqCst), 0);
}

// T-24
#[tokio::test]
async fn test_dnat_rewrites_destination() {
    let (echo_port, _) = echo_server().await;
    let proxy_hits = Arc::new(AtomicU32::new(0));
    let (_, proxy) = make_proxy(proxy_hits.clone()).await;

    // dnat 1.2.3.4:80 -> 127.0.0.1:echo_port
    let engine = ChainEngine::new(ChainConfig {
        proxies: vec![proxy],
        chain_type: ChainType::Strict,
        dnats: vec![DnatRule {
            orig_addr: Ipv4Addr::new(1, 2, 3, 4),
            orig_port: Some(80),
            new_addr: Ipv4Addr::LOCALHOST,
            new_port: Some(echo_port),
        }],
        ..Default::default()
    });

    // Connect to the "fake" destination — should be rewritten to echo server
    let mut stream = engine
        .connect(Target::Ip(IpAddr::V4(Ipv4Addr::new(1, 2, 3, 4)), 80))
        .await
        .unwrap();

    stream.write_all(b"dnat").await.unwrap();
    let mut buf = [0u8; 4];
    stream.read_exact(&mut buf).await.unwrap();
    assert_eq!(&buf, b"dnat");
}