proxychains-masq 0.1.5

use anyhow::{bail, Context, Result};
use tokio::io::{AsyncReadExt, AsyncWriteExt};

use super::{BoxStream, Target};

/// Perform a SOCKS4 / SOCKS4a CONNECT handshake on an already-connected `stream`.
///
/// If `target` is [`Target::Host`], SOCKS4a extension is used (fake IP `0.0.0.1`,
/// hostname appended after the username field).
///
/// # Errors
///
/// Returns an error if the handshake fails or the proxy reports a non-success status.
pub async fn connect(
    mut stream: BoxStream,
    target: &Target,
    username: Option<&str>,
) -> Result<BoxStream> {
    let port = target.port();
    let user_bytes = username.unwrap_or("").as_bytes();

    // ── Build request ─────────────────────────────────────────────────────────
    let mut req: Vec<u8> = Vec::with_capacity(32);
    req.push(4); // version
    req.push(1); // CONNECT

    // Destination port (big-endian)
    req.extend_from_slice(&port.to_be_bytes());

    match target {
        Target::Ip(addr, _) => {
            match addr {
                std::net::IpAddr::V4(v4) => req.extend_from_slice(&v4.octets()),
                std::net::IpAddr::V6(_) => bail!("SOCKS4 does not support IPv6 addresses"),
            }
            req.extend_from_slice(user_bytes);
            req.push(0); // null-terminate username
        }
        Target::Host(hostname, _) => {
            // SOCKS4a: fake IP 0.0.0.1
            req.extend_from_slice(&[0, 0, 0, 1]);
            req.extend_from_slice(user_bytes);
            req.push(0); // null-terminate username
            req.extend_from_slice(hostname.as_bytes());
            req.push(0); // null-terminate hostname
        }
    }

    stream
        .write_all(&req)
        .await
        .context("socks4: write request")?;

    // ── Read response (8 bytes) ────────────────────────────────────────────────
    let mut resp = [0u8; 8];
    stream
        .read_exact(&mut resp)
        .await
        .context("socks4: read response")?;

    // Byte 0 must be 0 (null), byte 1 must be 90 (0x5A = granted)
    if resp[0] != 0 {
        bail!("socks4: unexpected version byte in response: {}", resp[0]);
    }
    if resp[1] != 90 {
        bail!("socks4: connection denied (status {})", resp[1]);
    }

    Ok(stream)
}