proxychains-masq 0.1.5

mod helpers;

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

use helpers::{echo_server, http_connect_server, socks4_server, socks5_server, tls_connect_server};
use proxychains_masq::proxy::{http, https, raw, socks4, socks5, BoxStream, Target};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;

// ─── SOCKS4 / SOCKS4a ─────────────────────────────────────────────────────────

// T-08
#[tokio::test]
async fn test_socks4_connect_ipv4() {
    let (echo_port, _) = echo_server().await;
    let hits = Arc::new(AtomicU32::new(0));
    let proxy_port = socks4_server(hits.clone()).await;

    let stream: BoxStream = Box::new(TcpStream::connect(("127.0.0.1", proxy_port)).await.unwrap());
    let mut stream = socks4::connect(
        stream,
        &Target::Ip(IpAddr::V4(Ipv4Addr::LOCALHOST), echo_port),
        None,
    )
    .await
    .unwrap();

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

#[tokio::test]
async fn test_socks4_connect_with_username() {
    let (echo_port, _) = echo_server().await;
    let hits = Arc::new(AtomicU32::new(0));
    let proxy_port = socks4_server(hits.clone()).await;

    let stream: BoxStream = Box::new(TcpStream::connect(("127.0.0.1", proxy_port)).await.unwrap());
    let mut stream = socks4::connect(
        stream,
        &Target::Ip(IpAddr::V4(Ipv4Addr::LOCALHOST), echo_port),
        Some("testuser"),
    )
    .await
    .unwrap();

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

// T-09
#[tokio::test]
async fn test_socks4a_connect_hostname() {
    let (echo_port, _) = echo_server().await;
    let hits = Arc::new(AtomicU32::new(0));
    let proxy_port = socks4_server(hits.clone()).await;

    let stream: BoxStream = Box::new(TcpStream::connect(("127.0.0.1", proxy_port)).await.unwrap());
    let mut stream = socks4::connect(
        stream,
        &Target::Host("127.0.0.1".to_owned(), echo_port),
        None,
    )
    .await
    .unwrap();

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

// ─── SOCKS5 ──────────────────────────────────────────────────────────────────

// T-10
#[tokio::test]
async fn test_socks5_connect_no_auth() {
    let (echo_port, _) = echo_server().await;
    let hits = Arc::new(AtomicU32::new(0));
    let proxy_port = socks5_server(None, hits.clone()).await;

    let stream: BoxStream = Box::new(TcpStream::connect(("127.0.0.1", proxy_port)).await.unwrap());
    let mut stream = socks5::connect(
        stream,
        &Target::Ip(IpAddr::V4(Ipv4Addr::LOCALHOST), echo_port),
        None,
        None,
    )
    .await
    .unwrap();

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

// T-11
#[tokio::test]
async fn test_socks5_connect_user_pass() {
    let (echo_port, _) = echo_server().await;
    let hits = Arc::new(AtomicU32::new(0));
    let proxy_port = socks5_server(Some(("user", "pass")), hits.clone()).await;

    let stream: BoxStream = Box::new(TcpStream::connect(("127.0.0.1", proxy_port)).await.unwrap());
    let mut stream = socks5::connect(
        stream,
        &Target::Ip(IpAddr::V4(Ipv4Addr::LOCALHOST), echo_port),
        Some("user"),
        Some("pass"),
    )
    .await
    .unwrap();

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

#[tokio::test]
async fn test_socks5_wrong_password_errors() {
    let (echo_port, _) = echo_server().await;
    let hits = Arc::new(AtomicU32::new(0));
    let proxy_port = socks5_server(Some(("user", "correct")), hits.clone()).await;

    let stream: BoxStream = Box::new(TcpStream::connect(("127.0.0.1", proxy_port)).await.unwrap());
    let result = socks5::connect(
        stream,
        &Target::Ip(IpAddr::V4(Ipv4Addr::LOCALHOST), echo_port),
        Some("user"),
        Some("wrong"),
    )
    .await;

    assert!(result.is_err(), "expected error for wrong password");
}

// T-12
#[tokio::test]
async fn test_socks5_connect_hostname() {
    let (echo_port, _) = echo_server().await;
    let hits = Arc::new(AtomicU32::new(0));
    let proxy_port = socks5_server(None, hits.clone()).await;

    let stream: BoxStream = Box::new(TcpStream::connect(("127.0.0.1", proxy_port)).await.unwrap());
    let mut stream = socks5::connect(
        stream,
        &Target::Host("127.0.0.1".to_owned(), echo_port),
        None,
        None,
    )
    .await
    .unwrap();

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

// T-13
#[tokio::test]
async fn test_socks5_connect_ipv6() {
    let listener = tokio::net::TcpListener::bind("[::1]:0").await.unwrap();
    let echo_port = listener.local_addr().unwrap().port();
    tokio::spawn(async move {
        loop {
            let (mut s, _) = listener.accept().await.unwrap();
            tokio::spawn(async move {
                let mut buf = vec![0u8; 4096];
                loop {
                    let n = s.read(&mut buf).await.unwrap_or(0);
                    if n == 0 {
                        break;
                    }
                    if s.write_all(&buf[..n]).await.is_err() {
                        break;
                    }
                }
            });
        }
    });

    let hits = Arc::new(AtomicU32::new(0));
    let proxy_port = socks5_server(None, hits.clone()).await;

    let stream: BoxStream = Box::new(TcpStream::connect(("127.0.0.1", proxy_port)).await.unwrap());
    let mut stream = socks5::connect(
        stream,
        &Target::Ip(IpAddr::V6(Ipv6Addr::LOCALHOST), echo_port),
        None,
        None,
    )
    .await
    .unwrap();

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

// ─── HTTP CONNECT ─────────────────────────────────────────────────────────────

// T-14
#[tokio::test]
async fn test_http_connect_no_auth() {
    let (echo_port, _) = echo_server().await;
    let hits = Arc::new(AtomicU32::new(0));
    let proxy_port = http_connect_server(None, hits.clone()).await;

    let stream: BoxStream = Box::new(TcpStream::connect(("127.0.0.1", proxy_port)).await.unwrap());
    let mut stream = http::connect(
        stream,
        &Target::Host("127.0.0.1".to_owned(), echo_port),
        None,
        None,
    )
    .await
    .unwrap();

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

// T-15
#[tokio::test]
async fn test_http_connect_basic_auth() {
    let (echo_port, _) = echo_server().await;
    let hits = Arc::new(AtomicU32::new(0));
    let proxy_port = http_connect_server(Some(("u", "p")), hits.clone()).await;

    let stream: BoxStream = Box::new(TcpStream::connect(("127.0.0.1", proxy_port)).await.unwrap());
    let mut stream = http::connect(
        stream,
        &Target::Host("127.0.0.1".to_owned(), echo_port),
        Some("u"),
        Some("p"),
    )
    .await
    .unwrap();

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

#[tokio::test]
async fn test_http_connect_407_errors() {
    let (echo_port, _) = echo_server().await;
    let hits = Arc::new(AtomicU32::new(0));
    let proxy_port = http_connect_server(Some(("u", "p")), hits.clone()).await;

    let stream: BoxStream = Box::new(TcpStream::connect(("127.0.0.1", proxy_port)).await.unwrap());
    let result = http::connect(
        stream,
        &Target::Host("127.0.0.1".to_owned(), echo_port),
        None,
        None,
    )
    .await;

    assert!(result.is_err(), "expected 407 error");
}

// ─── RAW ──────────────────────────────────────────────────────────────────────

// ─── HTTPS CONNECT ────────────────────────────────────────────────────────────

// T-08b: HTTPS proxy with insecure (skip-verify) TLS and no auth
#[tokio::test]
async fn test_https_connect_no_auth() {
    let (echo_port, hits) = echo_server().await;
    let proxy_port = tls_connect_server(None, hits.clone()).await;

    let stream: BoxStream = Box::new(TcpStream::connect(("127.0.0.1", proxy_port)).await.unwrap());
    let mut stream = https::connect(
        stream,
        &Target::Ip(IpAddr::V4(Ipv4Addr::LOCALHOST), echo_port),
        None,
        None,
        "127.0.0.1".parse().unwrap(),
        true, // test proxy uses a self-signed certificate
    )
    .await
    .unwrap();

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

// T-08c: HTTPS proxy with basic auth
#[tokio::test]
async fn test_https_connect_basic_auth() {
    let (echo_port, hits) = echo_server().await;
    let proxy_port = tls_connect_server(Some(("user", "pass")), hits.clone()).await;

    let stream: BoxStream = Box::new(TcpStream::connect(("127.0.0.1", proxy_port)).await.unwrap());
    let mut stream = https::connect(
        stream,
        &Target::Ip(IpAddr::V4(Ipv4Addr::LOCALHOST), echo_port),
        Some("user"),
        Some("pass"),
        "127.0.0.1".parse().unwrap(),
        true, // test proxy uses a self-signed certificate
    )
    .await
    .unwrap();

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

// ─── RAW ──────────────────────────────────────────────────────────────────────

// T-16
#[tokio::test]
async fn test_raw_passthrough() {
    let (echo_port, _) = echo_server().await;

    let stream: BoxStream = Box::new(TcpStream::connect(("127.0.0.1", echo_port)).await.unwrap());
    let mut stream = raw::connect(
        stream,
        &Target::Ip(IpAddr::V4(Ipv4Addr::LOCALHOST), echo_port),
    )
    .await
    .unwrap();

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