socks5x 0.1.3

//!
//! This crate provides both client and server implementations of the SOCKS5 protocol
//! (RFC 1928) with support for various authentication methods.
//!
//! # Quick Start
//!
//! ## Client Usage
//!
//! ```rust,no_run
//! use socks5x::Socks5Address;
//! use socks5x::client::Socks5Client;
//! use tokio::io::{AsyncReadExt, AsyncWriteExt};
//!
//! #[tokio::main]
//! async fn main() -> Result<(), Box<dyn std::error::Error>> {
//!     let client = Socks5Client::connect(
//!     "127.0.0.1:1080",
//!     None // Some(("username".to_string(), "password".to_string())),
//!     )
//!     .await?;
//!
//!     // Request connection to destination
//!     let mut stream = client
//!         .request_connect(Socks5Address::from("httpbin.org"), 80)
//!         .await?;
//!
//!     // Use the stream normally - data is transparently proxied
//!     let request = "GET /ip HTTP/1.1\r\nHost: httpbin.org\r\nConnection: close\r\n\r\n";
//!     stream.write_all(request.as_bytes()).await?;
//!
//!     let mut response = Vec::new();
//!     stream.read_to_end(&mut response).await?;
//!     println!("Response: {}", String::from_utf8_lossy(&response));
//!
//!     Ok(())
//! }
//! ```
//!
//! ## Server Usage
//!
//! ```rust,no_run
//! use std::sync::Arc;
//!
//! use socks5x::server::{ClientHandler, DefaultConnectionCreator};
//! use tokio::net::TcpListener;
//!
//! #[tokio::main]
//! async fn main() -> Result<(), Box<dyn std::error::Error>> {
//!     let listener = TcpListener::bind("127.0.0.1:1080").await?;
//!     println!("SOCKS5 proxy server listening on 127.0.0.1:1080");
//!
//!     let client_handler = Arc::new(ClientHandler::no_auth(DefaultConnectionCreator));
//!
//!     loop {
//!         let (socket, addr) = listener.accept().await?;
//!         println!("New client connection from: {}", addr);
//!         let client_handler = client_handler.clone();
//!
//!         tokio::spawn(async move {
//!             if let Err(e) = client_handler.handle(socket).await {
//!                 eprintln!("Error handling client {}: {}", addr, e);
//!             } else {
//!                 println!("Client {} disconnected", addr);
//!             }
//!         });
//!     }
//! }
//! ```

use bytes::{Buf, BufMut, BytesMut};
use std::fmt::Display;
use std::io;
use std::net::{Ipv4Addr, Ipv6Addr};

#[cfg(feature = "client")]
pub mod client;
#[cfg(feature = "server")]
pub mod server;

const SOCKS5_VERSION: u8 = 5;
const SOCKS5_RESERVED: u8 = 0x00;
const SOCKS5_USERNAME_AUTH_VER: u8 = 0x01;

/// SOCKS5 command types as defined in RFC 1928.
enum Socks5Command {
    Connect = 0x01,
    Bind = 0x02,
    UdpAssociate = 0x03,
}

impl TryFrom<u8> for Socks5Command {
    type Error = io::Error;

    fn try_from(value: u8) -> Result<Self, Self::Error> {
        match value {
            v if v == Socks5Command::Connect as u8 => Ok(Socks5Command::Connect),
            v if v == Socks5Command::Bind as u8 => Ok(Socks5Command::Bind),
            v if v == Socks5Command::UdpAssociate as u8 => Ok(Socks5Command::UdpAssociate),
            _ => Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!("Unsupported SOCKS5 command: {}", value),
            )),
        }
    }
}

/// SOCKS5 response status codes as defined in RFC 1928.
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
enum Socks5Status {
    /// Request granted.
    RequestGranted = 0x00,
    /// General SOCKS server failure.
    GeneralFailure = 0x01,
    /// Connection not allowed by ruleset.
    ConnectionNotAllowed = 0x02,
    /// Network unreachable.
    NetworkUnreachable = 0x03,
    /// Host unreachable.
    HostUnreachable = 0x04,
    /// Connection refused.
    ConnectionRefused = 0x05,
    /// TTL expired.
    TtlExpired = 0x06,
    /// Command not supported.
    CommandNotSupported = 0x07,
    /// Address type not supported.
    AddressTypeNotSupported = 0x08,
}

impl TryFrom<u8> for Socks5Status {
    type Error = io::Error;

    fn try_from(value: u8) -> Result<Self, Self::Error> {
        match value {
            v if v == Socks5Status::RequestGranted as u8 => Ok(Socks5Status::RequestGranted),
            v if v == Socks5Status::GeneralFailure as u8 => Ok(Socks5Status::GeneralFailure),
            v if v == Socks5Status::ConnectionNotAllowed as u8 => {
                Ok(Socks5Status::ConnectionNotAllowed)
            }
            v if v == Socks5Status::NetworkUnreachable as u8 => {
                Ok(Socks5Status::NetworkUnreachable)
            }
            v if v == Socks5Status::HostUnreachable as u8 => Ok(Socks5Status::HostUnreachable),
            v if v == Socks5Status::ConnectionRefused as u8 => Ok(Socks5Status::ConnectionRefused),
            v if v == Socks5Status::TtlExpired as u8 => Ok(Socks5Status::TtlExpired),
            v if v == Socks5Status::CommandNotSupported as u8 => {
                Ok(Socks5Status::CommandNotSupported)
            }
            v if v == Socks5Status::AddressTypeNotSupported as u8 => {
                Ok(Socks5Status::AddressTypeNotSupported)
            }
            _ => Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!("Unsupported SOCKS5 command: {}", value),
            )),
        }
    }
}

enum Socks5AddressType {
    IPv4 = 0x01,
    Domain = 0x03,
    IPv6 = 0x04,
}

impl TryFrom<u8> for Socks5AddressType {
    type Error = io::Error;

    fn try_from(value: u8) -> Result<Self, Self::Error> {
        match value {
            v if v == Socks5AddressType::IPv4 as u8 => Ok(Socks5AddressType::IPv4),
            v if v == Socks5AddressType::Domain as u8 => Ok(Socks5AddressType::Domain),
            v if v == Socks5AddressType::IPv6 as u8 => Ok(Socks5AddressType::IPv6),
            _ => Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!("Unsupported SOCKS5 address type: {}", value),
            )),
        }
    }
}

/// Address types supported by the SOCKS5 protocol.
///
/// SOCKS5 supports three types of destination addresses: IPv4 addresses,
/// IPv6 addresses, and domain names. Domain names are resolved by the
/// SOCKS5 server, allowing clients to connect to hosts without knowing
/// their IP addresses.
///
/// # Examples
///
/// ```rust
/// use socks5x::Socks5Address;
/// use std::net::{Ipv4Addr, Ipv6Addr};
///
/// // From string literals
/// let addr1 = Socks5Address::from("example.com");
/// let addr2 = Socks5Address::from("192.168.1.1");
/// let addr3 = Socks5Address::from("2001:db8::1");
///
/// // From IP addresses
/// let addr4 = Socks5Address::IPv4(Ipv4Addr::new(127, 0, 0, 1));
/// let addr5 = Socks5Address::IPv6(Ipv6Addr::LOCALHOST);
/// let addr6 = Socks5Address::Domain("github.com".to_string());
/// ```
/// Address types supported by SOCKS5 protocol.
#[derive(Clone)]
pub enum Socks5Address {
    /// IPv4 address.
    IPv4(Ipv4Addr),
    /// Domain name to be resolved by the SOCKS5 server.
    Domain(String),
    /// IPv6 address.
    IPv6(Ipv6Addr),
}

impl From<&str> for Socks5Address {
    fn from(value: &str) -> Self {
        if let Ok(ipv4) = value.parse::<Ipv4Addr>() {
            Self::IPv4(ipv4)
        } else if let Ok(ipv6) = value.parse::<Ipv6Addr>() {
            Self::IPv6(ipv6)
        } else if value.contains('.')
            && value
                .chars()
                .all(|c| c.is_alphanumeric() || c == '.' || c == '-')
        {
            Self::Domain(value.to_string())
        } else {
            // Fallback to a safe default if invalid
            eprintln!("Invalid address format: '{}', treating as domain", value);
            Self::Domain(value.to_string())
        }
    }
}

impl Socks5Address {
    /// Encodes address to SOCKS5 wire format.
    fn to_bytes(&self, buf: &mut BytesMut) {
        match self {
            Self::IPv4(ip) => {
                buf.put_u8(Socks5AddressType::IPv4 as u8);
                buf.extend_from_slice(&ip.octets());
            }
            Self::Domain(domain) => {
                buf.put_u8(Socks5AddressType::Domain as u8);
                let len = domain.len() as u8;
                buf.put_u8(len);
                buf.extend_from_slice(domain.as_bytes());
            }
            Self::IPv6(ip) => {
                buf.put_u8(Socks5AddressType::IPv6 as u8);
                buf.extend_from_slice(&ip.octets());
            }
        }
    }
    /// Parses address from SOCKS5 wire format.
    fn parse(atyp: Socks5AddressType, src: &mut BytesMut) -> Result<Self, io::Error> {
        match atyp {
            Socks5AddressType::IPv4 => {
                if src.len() < 4 {
                    return Err(io::Error::new(
                        io::ErrorKind::WouldBlock,
                        "Not enough data for IPv4",
                    ));
                }
                let ip = Ipv4Addr::new(src[0], src[1], src[2], src[3]);
                src.advance(4);
                Ok(Socks5Address::IPv4(ip))
            }
            Socks5AddressType::Domain => {
                if src.is_empty() {
                    return Err(io::Error::new(
                        io::ErrorKind::WouldBlock,
                        "Not enough data for domain length",
                    ));
                }
                let len = src[0] as usize;
                src.advance(1);
                if src.len() < len {
                    return Err(io::Error::new(
                        io::ErrorKind::WouldBlock,
                        "Not enough data for domain",
                    ));
                }
                let domain = String::from_utf8_lossy(&src[..len]).to_string();
                src.advance(len);
                Ok(Socks5Address::Domain(domain))
            }
            Socks5AddressType::IPv6 => {
                if src.len() < 16 {
                    return Err(io::Error::new(
                        io::ErrorKind::WouldBlock,
                        "Not enough data for IPv6",
                    ));
                }
                let ip = Ipv6Addr::new(
                    u16::from_be_bytes([src[0], src[1]]),
                    u16::from_be_bytes([src[2], src[3]]),
                    u16::from_be_bytes([src[4], src[5]]),
                    u16::from_be_bytes([src[6], src[7]]),
                    u16::from_be_bytes([src[8], src[9]]),
                    u16::from_be_bytes([src[10], src[11]]),
                    u16::from_be_bytes([src[12], src[13]]),
                    u16::from_be_bytes([src[14], src[15]]),
                );
                src.advance(16);
                Ok(Socks5Address::IPv6(ip))
            }
        }
    }
}

impl Display for Socks5Address {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Socks5Address::IPv4(ip) => write!(f, "{}", ip),
            Socks5Address::Domain(domain) => write!(f, "{}", domain),
            Socks5Address::IPv6(ip) => write!(f, "{}", ip),
        }
    }
}
struct Socks5Response {
    response: Socks5Status,
    address: Socks5Address,
    port: u16,
}
pub struct Socks5Request {
    command: Socks5Command,
    address: Socks5Address,
    port: u16,
}

/// Authentication methods supported by SOCKS5.
///
/// SOCKS5 supports various authentication methods as defined in RFC 1928
/// and related RFCs. This library currently implements NoAuth and
/// UsernamePassword methods, with the other variants provided for
/// completeness and future extensibility.
#[derive(Clone, PartialEq)]
enum AuthMethod {
    /// No authentication (0x00)
    NoAuth = 0x00,
    /// GSSAPI (RFC 1961) (0x01)
    Gssapi = 0x01,
    /// Username/password (RFC 1929) (0x02)
    UsernamePassword = 0x02,
    /// Challenge-Handshake Authentication Protocol (0x03)
    Chap = 0x03,
    /// Challenge-Response Authentication Method (0x05)
    Cram = 0x05,
    /// Secure Sockets Layer (0x06)
    Ssl = 0x06,
    /// NDS Authentication (0x07)
    Nds = 0x07,
    /// Multi-Authentication Framework (0x08)
    Maf = 0x08,
    /// JSON Parameter Block (0x09)
    Json = 0x09,
}

impl Display for AuthMethod {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let name = match self {
            AuthMethod::NoAuth => "NoAuth",
            AuthMethod::Gssapi => "Gssapi",
            AuthMethod::UsernamePassword => "UsernamePassword",
            AuthMethod::Chap => "Chap",
            AuthMethod::Cram => "Cram",
            AuthMethod::Ssl => "Ssl",
            AuthMethod::Nds => "Nds",
            AuthMethod::Maf => "Maf",
            AuthMethod::Json => "Json",
        };
        write!(f, "{}", name)
    }
}
impl TryFrom<u8> for AuthMethod {
    type Error = io::Error;
    fn try_from(value: u8) -> Result<Self, Self::Error> {
        match value {
            v if v == Self::NoAuth as u8 => Ok(Self::NoAuth),
            v if v == Self::Gssapi as u8 => Ok(Self::Gssapi),
            v if v == Self::UsernamePassword as u8 => Ok(Self::UsernamePassword),
            v if v == Self::Chap as u8 => Ok(Self::Chap),
            v if v == Self::Cram as u8 => Ok(Self::Cram),
            v if v == Self::Ssl as u8 => Ok(Self::Ssl),
            v if v == Self::Nds as u8 => Ok(Self::Nds),
            v if v == Self::Maf as u8 => Ok(Self::Maf),
            v if v == Self::Json as u8 => Ok(Self::Json),
            v => Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!("Unknown auth method {v}"),
            )),
        }
    }
}

fn decode_res<S>(src: &mut BytesMut) -> io::Result<Option<(S, Socks5Address, u16)>>
where
    S: TryFrom<u8>,
    std::io::Error: From<<S as TryFrom<u8>>::Error>,
{
    //   +----+-----+-------+------+----------+----------+
    //   |VER | CMD |  RSV  | ATYP | DST.ADDR | DST.PORT |
    //   +----+-----+-------+------+----------+----------+
    //   | 1  |  1  | X'00' |  1   | Variable |    2     |
    //   +----+-----+-------+------+----------+----------+

    if src.len() < 4 {
        return Ok(None);
    }

    let version = src[0];
    if version != SOCKS5_VERSION {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            "Invalid SOCKS version",
        ));
    }

    let res = S::try_from(src[1])?;
    let atyp = Socks5AddressType::try_from(src[3])?;
    src.advance(4); // Consume VER, CMD, RSV, ATYP

    let address = match Socks5Address::parse(atyp, src) {
        Ok(addr) => addr,
        Err(e) if e.kind() == io::ErrorKind::WouldBlock => return Ok(None),
        Err(e) => return Err(e),
    };

    if src.len() < 2 {
        return Ok(None);
    }
    let port = u16::from_be_bytes([src[0], src[1]]);
    src.advance(2);

    Ok(Some((res, address, port)))
}