proxychains-masq 0.1.5

use std::net::IpAddr;

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

use super::{BoxStream, Target};

const METHOD_NO_AUTH: u8 = 0x00;
const METHOD_USER_PASS: u8 = 0x02;
const METHOD_NO_ACCEPTABLE: u8 = 0xFF;

/// Perform a SOCKS5 CONNECT handshake on an already-connected `stream`.
///
/// Supports:
/// - No-auth (method 0x00)
/// - Username/password auth (method 0x02, RFC 1929)
/// - IPv4, IPv6, and hostname targets
///
/// # Errors
///
/// Returns an error if the handshake fails, auth is rejected, or the proxy
/// returns a non-success reply.
pub async fn connect(
    mut stream: BoxStream,
    target: &Target,
    username: Option<&str>,
    password: Option<&str>,
) -> Result<BoxStream> {
    // ── Method negotiation ────────────────────────────────────────────────────
    let methods: &[u8] = if username.is_some() {
        &[METHOD_NO_AUTH, METHOD_USER_PASS]
    } else {
        &[METHOD_NO_AUTH]
    };

    let mut greeting = vec![5u8, methods.len() as u8];
    greeting.extend_from_slice(methods);
    stream
        .write_all(&greeting)
        .await
        .context("socks5: write greeting")?;

    let mut sel = [0u8; 2];
    stream
        .read_exact(&mut sel)
        .await
        .context("socks5: read method selection")?;
    if sel[0] != 5 {
        bail!("socks5: unexpected version in method selection: {}", sel[0]);
    }
    if sel[1] == METHOD_NO_ACCEPTABLE {
        bail!("socks5: no acceptable authentication method");
    }

    // ── Auth sub-negotiation (if required) ────────────────────────────────────
    if sel[1] == METHOD_USER_PASS {
        let user = username.unwrap_or("");
        let pass = password.unwrap_or("");

        let mut auth: Vec<u8> = Vec::with_capacity(3 + user.len() + pass.len());
        auth.push(1); // sub-negotiation version
        auth.push(user.len() as u8);
        auth.extend_from_slice(user.as_bytes());
        auth.push(pass.len() as u8);
        auth.extend_from_slice(pass.as_bytes());
        stream
            .write_all(&auth)
            .await
            .context("socks5: write auth")?;

        let mut ar = [0u8; 2];
        stream
            .read_exact(&mut ar)
            .await
            .context("socks5: read auth response")?;
        // Version field can be 1 or 5 (some non-conforming proxies send 5)
        if ar[1] != 0 {
            bail!("socks5: authentication failed (status {})", ar[1]);
        }
    } else if sel[1] != METHOD_NO_AUTH {
        bail!("socks5: server chose unexpected method: {}", sel[1]);
    }

    // ── CONNECT request ───────────────────────────────────────────────────────
    let mut req: Vec<u8> = Vec::with_capacity(32);
    req.extend_from_slice(&[5, 1, 0]); // VER CMD RSV

    match target {
        Target::Ip(IpAddr::V4(v4), _) => {
            req.push(1); // atype IPv4
            req.extend_from_slice(&v4.octets());
        }
        Target::Ip(IpAddr::V6(v6), _) => {
            req.push(4); // atype IPv6
            req.extend_from_slice(&v6.octets());
        }
        Target::Host(hostname, _) => {
            req.push(3); // atype domain name
            req.push(hostname.len() as u8);
            req.extend_from_slice(hostname.as_bytes());
        }
    }
    req.extend_from_slice(&target.port().to_be_bytes());

    stream
        .write_all(&req)
        .await
        .context("socks5: write CONNECT")?;

    // ── Read reply ────────────────────────────────────────────────────────────
    let mut rep = [0u8; 4];
    stream
        .read_exact(&mut rep)
        .await
        .context("socks5: read reply header")?;
    if rep[0] != 5 {
        bail!("socks5: unexpected version in reply: {}", rep[0]);
    }
    if rep[1] != 0 {
        bail!("socks5: CONNECT failed (reply code {})", rep[1]);
    }

    // Drain the bound address from the reply
    let atype = rep[3];
    let addr_len: usize = match atype {
        1 => 4,
        4 => 16,
        3 => {
            let len = stream.read_u8().await.context("socks5: read domain len")? as usize;
            len
        }
        _ => bail!("socks5: unknown address type in reply: {atype}"),
    };
    let mut addr_buf = vec![0u8; addr_len + 2]; // addr + 2-byte port
    stream
        .read_exact(&mut addr_buf)
        .await
        .context("socks5: drain bound addr")?;

    Ok(stream)
}