rustzmq2 0.1.0

//! Minimal SOCKS5 client. Supports NO-AUTH and username/password auth.
//!
//! Inline implementation (no external crate) to keep the dependency graph
//! unchanged. Generic over any stream implementing `futures::AsyncRead +
//! AsyncWrite`, so the same code handles both tokio and smol TCP streams
//! (tokio's stream is adapted via `tokio_util::compat` at the call site).

use futures::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use std::io;

/// Perform the SOCKS5 handshake on an already-connected stream. Returns
/// `Ok(())` once the proxy has confirmed the CONNECT request; subsequent
/// I/O on the stream is tunnelled to the target.
pub(crate) async fn handshake<S>(
    stream: &mut S,
    target_host: &str,
    target_port: u16,
    credentials: Option<(&str, &str)>,
) -> io::Result<()>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    // ── Step 1: method negotiation ────────────────────────────────────────
    let methods: &[u8] = if credentials.is_some() {
        &[0x00, 0x02] // no-auth + user/pass
    } else {
        &[0x00] // no-auth only
    };
    let mut greeting = vec![0x05, methods.len() as u8];
    greeting.extend_from_slice(methods);
    stream.write_all(&greeting).await?;

    let mut sel = [0u8; 2];
    stream.read_exact(&mut sel).await?;
    if sel[0] != 0x05 {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            "SOCKS5: bad version from proxy",
        ));
    }
    match sel[1] {
        0x00 => {} // no auth
        0x02 => {
            let (user, pass) = credentials.ok_or_else(|| {
                io::Error::new(
                    io::ErrorKind::PermissionDenied,
                    "SOCKS5: proxy requires auth but none configured",
                )
            })?;
            auth_userpass(stream, user, pass).await?;
        }
        0xFF => {
            return Err(io::Error::new(
                io::ErrorKind::PermissionDenied,
                "SOCKS5: no acceptable auth method",
            ))
        }
        other => {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!("SOCKS5: unsupported auth method {:#x}", other),
            ))
        }
    }

    // ── Step 2: CONNECT request ───────────────────────────────────────────
    let mut req = vec![0x05, 0x01, 0x00]; // VER, CMD=CONNECT, RSV
                                          // Prefer IPv4/IPv6 literal if target_host parses as one; otherwise use
                                          // DOMAINNAME address type and let the proxy resolve.
    if let Ok(ip) = target_host.parse::<std::net::IpAddr>() {
        match ip {
            std::net::IpAddr::V4(v4) => {
                req.push(0x01);
                req.extend_from_slice(&v4.octets());
            }
            std::net::IpAddr::V6(v6) => {
                req.push(0x04);
                req.extend_from_slice(&v6.octets());
            }
        }
    } else {
        if target_host.len() > 255 {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                "SOCKS5: domain name too long",
            ));
        }
        req.push(0x03);
        req.push(target_host.len() as u8);
        req.extend_from_slice(target_host.as_bytes());
    }
    req.extend_from_slice(&target_port.to_be_bytes());
    stream.write_all(&req).await?;

    // ── Step 3: reply ─────────────────────────────────────────────────────
    let mut hdr = [0u8; 4];
    stream.read_exact(&mut hdr).await?;
    if hdr[0] != 0x05 {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            "SOCKS5: bad reply version",
        ));
    }
    if hdr[1] != 0x00 {
        return Err(io::Error::other(format!(
            "SOCKS5: CONNECT failed (reply code {:#x})",
            hdr[1]
        )));
    }
    // Drain the bound-addr portion (ATYP + addr + port).
    match hdr[3] {
        0x01 => {
            let mut buf = [0u8; 4 + 2];
            stream.read_exact(&mut buf).await?;
        }
        0x04 => {
            let mut buf = [0u8; 16 + 2];
            stream.read_exact(&mut buf).await?;
        }
        0x03 => {
            let mut len = [0u8; 1];
            stream.read_exact(&mut len).await?;
            let mut buf = vec![0u8; len[0] as usize + 2];
            stream.read_exact(&mut buf).await?;
        }
        other => {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!("SOCKS5: bad ATYP {:#x}", other),
            ))
        }
    }
    Ok(())
}

async fn auth_userpass<S>(stream: &mut S, user: &str, pass: &str) -> io::Result<()>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    if user.len() > 255 || pass.len() > 255 {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput,
            "SOCKS5: credentials exceed 255 bytes",
        ));
    }
    let mut buf = vec![0x01, user.len() as u8];
    buf.extend_from_slice(user.as_bytes());
    buf.push(pass.len() as u8);
    buf.extend_from_slice(pass.as_bytes());
    stream.write_all(&buf).await?;

    let mut reply = [0u8; 2];
    stream.read_exact(&mut reply).await?;
    if reply[1] != 0x00 {
        return Err(io::Error::new(
            io::ErrorKind::PermissionDenied,
            "SOCKS5: username/password authentication failed",
        ));
    }
    Ok(())
}