socks_abstract5/

lib.rs

use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
use thiserror::Error;
use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    net::{TcpListener, TcpStream, UdpSocket},
};

#[derive(Debug, Error)]
pub enum Error {
    #[error("IO error: {0}")]
    Io(#[from] std::io::Error),
    #[error("Invalid SOCKS version")]
    InvalidVersion,
    #[error("Authentication failed")]
    AuthFailed,
    #[error("Invalid command: {0}")]
    InvalidCommand(u8),
    #[error("Invalid address type: {0}")]
    InvalidAddrType(u8),
    #[error("UTF-8 decode error: {0}")]
    Utf8Error(#[from] std::string::FromUtf8Error),
    #[error("Connection not allowed by ruleset")]
    ConnectionNotAllowed,
    #[error("Network unreachable")]
    NetworkUnreachable,
    #[error("Host unreachable")]
    HostUnreachable,
    #[error("Connection refused")]
    ConnectionRefused,
    #[error("TTL expired")]
    TtlExpired,
    #[error("Protocol error")]
    ProtocolError,
    #[error("Address type not supported")]
    AddrTypeNotSupported,
}

pub type Result<T> = std::result::Result<T, Error>;

#[derive(Debug)]
pub enum Command {
    Connect = 0x01,
    Bind = 0x02,
    UdpAssociate = 0x03,
}

impl TryFrom<u8> for Command {
    type Error = Error;
    fn try_from(value: u8) -> Result<Self> {
        match value {
            0x01 => Ok(Command::Connect),
            0x02 => Ok(Command::Bind),
            0x03 => Ok(Command::UdpAssociate),
            cmd => Err(Error::InvalidCommand(cmd)),
        }
    }
}

#[derive(Debug)]
pub enum Address {
    Ipv4(Ipv4Addr),
    Ipv6(Ipv6Addr),
    Domain(String),
}

impl Address {
    async fn read_from(stream: &mut TcpStream) -> Result<(Self, u16)> {
        let atyp = stream.read_u8().await?;
        
        let addr = match atyp {
            0x01 => { // IPv4
                let mut buf = [0u8; 4];
                stream.read_exact(&mut buf).await?;
                Address::Ipv4(Ipv4Addr::from(buf))
            }
            0x03 => { // Domain
                let len = stream.read_u8().await? as usize;
                let mut domain = vec![0u8; len];
                stream.read_exact(&mut domain).await?;
                Address::Domain(String::from_utf8(domain)?)
            }
            0x04 => { // IPv6
                let mut buf = [0u8; 16];
                stream.read_exact(&mut buf).await?;
                Address::Ipv6(Ipv6Addr::from(buf))
            }
            _ => return Err(Error::InvalidAddrType(atyp)),
        };

        let mut port_buf = [0u8; 2];
        stream.read_exact(&mut port_buf).await?;
        let port = u16::from_be_bytes(port_buf);

        Ok((addr, port))
    }

    async fn to_socket_addr(&self, port: u16) -> Result<SocketAddr> {
        match self {
            Address::Ipv4(addr) => Ok(SocketAddr::new(IpAddr::V4(*addr), port)),
            Address::Ipv6(addr) => Ok(SocketAddr::new(IpAddr::V6(*addr), port)),
            Address::Domain(domain) => {
                let addrs = tokio::net::lookup_host(format!("{}:{}", domain, port)).await?;
                Ok(addrs.into_iter().next().ok_or(Error::HostUnreachable)?)
            }
        }
    }
}

pub struct Socks5Server {
    listener: TcpListener,
}

impl Socks5Server {
    pub async fn new(addr: &str) -> Result<Self> {
        let listener = TcpListener::bind(addr).await?;
        Ok(Self { listener })
    }

    pub async fn run(&self) -> Result<()> {
        loop {
            let (client, _) = self.listener.accept().await?;
            tokio::spawn(async move {
                if let Err(e) = handle_client(client).await {
                    eprintln!("Client error: {}", e);
                }
            });
        }
    }
}

async fn handle_client(mut client: TcpStream) -> Result<()> {
    // Authentication negotiation
    let ver = client.read_u8().await?;
    if ver != 0x05 {
        return Err(Error::InvalidVersion);
    }

    let nmethods = client.read_u8().await?;
    let mut methods = vec![0u8; nmethods as usize];
    client.read_exact(&mut methods).await?;

    // For now, only support no authentication
    if !methods.contains(&0x00) {
        client.write_all(&[0x05, 0xFF]).await?;
        return Err(Error::AuthFailed);
    }
    client.write_all(&[0x05, 0x00]).await?;

    // Request
    let ver = client.read_u8().await?;
    if ver != 0x05 {
        return Err(Error::InvalidVersion);
    }

    let cmd = Command::try_from(client.read_u8().await?)?;
    let _rsv = client.read_u8().await?; // Reserved byte
    let (target_addr, target_port) = Address::read_from(&mut client).await?;

    match cmd {
        Command::Connect => handle_connect(&mut client, target_addr, target_port).await,
        Command::Bind => handle_bind(&mut client, target_addr, target_port).await,
        Command::UdpAssociate => handle_udp_associate(&mut client).await,
    }
}

async fn handle_connect(client: &mut TcpStream, addr: Address, port: u16) -> Result<()> {
    let target_addr = addr.to_socket_addr(port).await?;
    let mut target = match TcpStream::connect(target_addr).await {
        Ok(stream) => stream,
        Err(e) => {
            let reply = match e.kind() {
                std::io::ErrorKind::ConnectionRefused => 0x05,
                std::io::ErrorKind::TimedOut => 0x06,
                _ => 0x01,
            };
            send_error_reply(client, reply).await?;
            return Err(e.into());
        }
    };

    // Send success response
    let bind_addr = target.local_addr()?;
    send_reply(client, 0x00, &bind_addr).await?;

    // Start proxying
    let (mut cr, mut cw) = client.split();
    let (mut tr, mut tw) = target.split();

    tokio::select! {
        res = tokio::io::copy(&mut cr, &mut tw) => {
            if let Err(e) = res {
                return Err(e.into());
            }
        }
        res = tokio::io::copy(&mut tr, &mut cw) => {
            if let Err(e) = res {
                return Err(e.into());
            }
        }
    }

    Ok(())
}

async fn handle_bind(client: &mut TcpStream, addr: Address, port: u16) -> Result<()> {
    // Create listener for incoming connection
    let listener = TcpListener::bind("0.0.0.0:0").await?;
    let bind_addr = listener.local_addr()?;

    // Send first reply with bind address
    send_reply(client, 0x00, &bind_addr).await?;

    // Wait for incoming connection
    let (mut target, target_addr) = listener.accept().await?;
    
    // Verify connection is from expected address
    let expected_addr = addr.to_socket_addr(port).await?;
    if target_addr.ip() != expected_addr.ip() {
        send_error_reply(client, 0x02).await?;
        return Err(Error::ConnectionNotAllowed);
    }

    // Send second reply with target address
    send_reply(client, 0x00, &target_addr).await?;

    // Start proxying
    let (mut cr, mut cw) = client.split();
    let (mut tr, mut tw) = target.split();

    tokio::select! {
        res = tokio::io::copy(&mut cr, &mut tw) => {
            if let Err(e) = res {
                return Err(e.into());
            }
        }
        res = tokio::io::copy(&mut tr, &mut cw) => {
            if let Err(e) = res {
                return Err(e.into());
            }
        }
    }

    Ok(())
}

async fn handle_udp_associate(client: &mut TcpStream) -> Result<()> {
    // Create UDP socket for relaying
    let relay = UdpSocket::bind("0.0.0.0:0").await?;
    let relay_addr = relay.local_addr()?;

    // Send reply with relay address
    send_reply(client, 0x00, &relay_addr).await?;

    // Keep TCP connection open and monitor it
    let mut keep_alive = [0u8; 1];
    loop {
        match client.read(&mut keep_alive).await {
            Ok(0) | Err(_) => break, // Connection closed
            _ => continue,
        }
    }

    Ok(())
}

async fn send_reply(stream: &mut TcpStream, reply: u8, addr: &SocketAddr) -> Result<()> {
    stream.write_u8(0x05).await?; // VER
    stream.write_u8(reply).await?; // REP
    stream.write_u8(0x00).await?; // RSV

    match addr {
        SocketAddr::V4(addr) => {
            stream.write_u8(0x01).await?; // ATYP
            stream.write_all(&addr.ip().octets()).await?;
            stream.write_all(&addr.port().to_be_bytes()).await?;
        }
        SocketAddr::V6(addr) => {
            stream.write_u8(0x04).await?; // ATYP
            stream.write_all(&addr.ip().octets()).await?;
            stream.write_all(&addr.port().to_be_bytes()).await?;
        }
    }

    Ok(())
}

async fn send_error_reply(stream: &mut TcpStream, reply: u8) -> Result<()> {
    let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)), 0);
    send_reply(stream, reply, &addr).await
}