socks5x 0.1.3

use bytes::{Buf, BufMut, BytesMut};
use core::fmt;
use futures_util::SinkExt;
use std::io;
use tokio_stream::StreamExt;

use std::pin::Pin;
use std::task::Poll;
use tokio::io::{AsyncRead, AsyncWrite};
use tokio::net::{TcpStream, ToSocketAddrs};
use tokio_util::codec::{Decoder, Encoder, Framed};

use crate::{SOCKS5_RESERVED, SOCKS5_USERNAME_AUTH_VER, SOCKS5_VERSION};

use super::{
    AuthMethod, Socks5Address, Socks5Command, Socks5Request, Socks5Response, Socks5Status,
    decode_res,
};

struct ClientCodec {
    /// Track what type of response we're expecting next
    expecting: ResponseType,
}

#[derive(Clone, Copy)]
enum ResponseType {
    /// Expecting method selection response (2 bytes)
    MethodSelection,
    /// Expecting auth response (2 bytes)
    AuthResponse,
    /// Expecting connect response (variable length)
    ConnectResponse,
}

impl ClientCodec {
    fn new() -> Self {
        Self {
            expecting: ResponseType::MethodSelection,
        }
    }
}

enum ClientRequest {
    /// Initial handshake to negotiate authentication method.
    Handshake(AuthMethod),
    /// Username/password authentication request.
    UserPassAuth((String, String)),
    /// Connection request to destination.
    Connect(Socks5Request),
}

impl fmt::Display for ClientRequest {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ClientRequest::Handshake(auth_method) => {
                write!(f, "ClientRequest::Handshake({auth_method})")
            }
            ClientRequest::UserPassAuth((username, _)) => {
                write!(
                    f,
                    "ClientRequest::UserPassAuth(username: {username}, password: ***)"
                )
            }
            ClientRequest::Connect(_socks5_request) => {
                write!(f, "ClientRequest::Connect")
            }
        }
    }
}

/// Represents responses received from the SOCKS5 server.
enum ClientResponse {
    /// Server selected authentication method.
    MethodSelected(AuthMethod),
    /// Response to connection request.
    ConnectResponse(Socks5Response),
    /// Authentication result (true = success, false = failure).
    AuthResponse(bool),
}

impl fmt::Display for ClientResponse {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ClientResponse::MethodSelected(auth_method) => {
                write!(f, "ClientResponse::MethodSelected({auth_method})")
            }
            ClientResponse::ConnectResponse(_socks5_response) => {
                write!(f, "ClientResponse::ConnectResponse")
            }
            ClientResponse::AuthResponse(success) => {
                write!(f, "ClientResponse::AuthResponse(success: {})", success)
            }
        }
    }
}

impl Decoder for ClientCodec {
    type Item = ClientResponse;
    type Error = io::Error;

    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
        match self.expecting {
            ResponseType::MethodSelection => {
                if src.len() < 2 {
                    return Ok(None);
                }

                let version = src[0];
                let method = src[1];

                if version != SOCKS5_VERSION {
                    return Err(io::Error::new(
                        io::ErrorKind::InvalidData,
                        format!(
                            "Invalid SOCKS version: expected {}, got {}",
                            SOCKS5_VERSION, version
                        ),
                    ));
                }

                src.advance(2);

                let auth_method = AuthMethod::try_from(method)?;

                // Update expectation based on selected method
                match auth_method {
                    AuthMethod::NoAuth => {
                        self.expecting = ResponseType::ConnectResponse;
                    }
                    AuthMethod::UsernamePassword => {
                        self.expecting = ResponseType::AuthResponse;
                    }
                    _ => {
                        return Err(io::Error::new(
                            io::ErrorKind::InvalidData,
                            format!("Unimplemented method requested {auth_method}"),
                        ));
                    }
                }

                Ok(Some(ClientResponse::MethodSelected(auth_method)))
            }
            ResponseType::AuthResponse => {
                if src.len() < 2 {
                    return Ok(None);
                }

                let version = src[0];
                let status = src[1];

                if version != SOCKS5_USERNAME_AUTH_VER {
                    return Err(io::Error::new(
                        io::ErrorKind::InvalidData,
                        format!(
                            "Invalid auth version: expected {}, got {}",
                            SOCKS5_USERNAME_AUTH_VER, version
                        ),
                    ));
                }

                src.advance(2);
                self.expecting = ResponseType::ConnectResponse;

                Ok(Some(ClientResponse::AuthResponse(status == 0)))
            }
            ResponseType::ConnectResponse => {
                match decode_res(src)? {
                    Some((response, address, port)) => {
                        Ok(Some(ClientResponse::ConnectResponse(Socks5Response {
                            address,
                            port,
                            response,
                        })))
                    }
                    None => Ok(None), // Not enough data yet
                }
            }
        }
    }
}

impl Encoder<ClientRequest> for ClientCodec {
    type Error = io::Error;

    fn encode(&mut self, item: ClientRequest, dst: &mut BytesMut) -> Result<(), Self::Error> {
        match item {
            ClientRequest::UserPassAuth((username, password)) => {
                //    +----+------+----------+------+----------+
                //    |VER | ULEN |  UNAME   | PLEN |  PASSWD  |
                //    +----+------+----------+------+----------+
                //    | 1  |  1   | 1 to 255 |  1   | 1 to 255 |
                //    +----+------+----------+------+----------+
                let uname = username.as_bytes();
                let passwd = password.as_bytes();

                if uname.len() > 255 || passwd.len() > 255 {
                    return Err(io::Error::new(
                        io::ErrorKind::InvalidData,
                        "Username or password too long (max 255 bytes)",
                    ));
                }

                dst.reserve(3 + uname.len() + passwd.len());
                dst.put_u8(SOCKS5_USERNAME_AUTH_VER);
                dst.put_u8(uname.len() as u8);
                dst.extend_from_slice(uname);
                dst.put_u8(passwd.len() as u8);
                dst.extend_from_slice(passwd);
            }
            ClientRequest::Handshake(auth_method) => {
                //    +----+----------+----------+
                //    |VER | NMETHODS | METHODS  |
                //    +----+----------+----------+
                //    | 1  |    1     | 1 to 255 |
                //    +----+----------+----------+
                dst.reserve(3);
                dst.extend_from_slice(&[SOCKS5_VERSION, 0x01, auth_method as u8]);
            }
            ClientRequest::Connect(req) => {
                // +----+-----+-------+------+----------+----------+
                // |VER | CMD |  RSV  | ATYP | DST.ADDR | DST.PORT |
                // +----+-----+-------+------+----------+----------+
                // | 1  |  1  | X'00' |  1   | Variable |    2     |
                // +----+-----+-------+------+----------+----------+
                dst.reserve(22); // Max size for IPv6 request
                dst.extend_from_slice(&[SOCKS5_VERSION, req.command as u8, SOCKS5_RESERVED]);
                req.address.to_bytes(dst);
                dst.extend_from_slice(&req.port.to_be_bytes());
            }
        }
        Ok(())
    }
}

/// An asynchronous SOCKS5 client for connecting through proxy servers.
///
/// # Example
/// ```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",
///         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(())
/// }
/// ```
pub struct Socks5Client {
    framed: Framed<TcpStream, ClientCodec>,
    credentials: Option<(String, String)>,
}

impl Socks5Client {
    /// Creates a new SOCKS5 client and connects to the specified proxy server.
    /// This method establishes a TCP connection to the proxy server and performs
    /// the initial SOCKS5 handshake to negotiate the authentication method.
    pub async fn connect(
        proxy_addr: impl ToSocketAddrs,
        credentials: Option<(String, String)>,
    ) -> io::Result<Self> {
        let stream = TcpStream::connect(proxy_addr).await?;

        if let Some((ref username, ref password)) = credentials {
            //  RFC1929: 1 to 255
            if username.is_empty()
                || username.len() > 255
                || password.is_empty()
                || password.len() > 255
            {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    "Username and password length should be in range [1-255]",
                ));
            }
        }

        let mut client = Socks5Client {
            framed: Framed::new(stream, ClientCodec::new()),
            credentials,
        };

        client.handshake().await?;
        Ok(client)
    }

    /// Performs username/password authentication with the proxy server.
    async fn authenticate(&mut self) -> io::Result<()> {
        let creds = self
            .credentials
            .clone()
            .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "No credentials provided"))?;

        self.framed.send(ClientRequest::UserPassAuth(creds)).await?;
        self.framed.flush().await?;

        match self.framed.next().await.transpose()? {
            Some(ClientResponse::AuthResponse(true)) => Ok(()),
            Some(ClientResponse::AuthResponse(false)) => Err(io::Error::new(
                io::ErrorKind::PermissionDenied,
                "Authentication failed: wrong username/password",
            )),
            Some(response) => Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!("Unexpected response during authentication: {}", response),
            )),
            None => Err(io::Error::new(
                io::ErrorKind::UnexpectedEof,
                "Connection closed during authentication",
            )),
        }
    }

    /// Performs the SOCKS5 authentication handshake with the proxy server.
    async fn handshake(&mut self) -> io::Result<()> {
        let auth_method = match self.credentials {
            Some(_) => AuthMethod::UsernamePassword,
            None => AuthMethod::NoAuth,
        };

        self.framed
            .send(ClientRequest::Handshake(auth_method))
            .await?;
        self.framed.flush().await?;

        match self.framed.next().await.transpose()? {
            Some(ClientResponse::MethodSelected(selected_method)) => match selected_method {
                AuthMethod::NoAuth => {
                    if self.credentials.is_some() {
                        return Err(io::Error::new(
                            io::ErrorKind::InvalidData,
                            "Server selected NoAuth but client provided credentials",
                        ));
                    }
                    Ok(())
                }
                AuthMethod::UsernamePassword => {
                    if self.credentials.is_none() {
                        return Err(io::Error::new(
                            io::ErrorKind::InvalidData,
                            "Server selected UsernamePassword but client provided no credentials",
                        ));
                    }
                    self.authenticate().await
                }
                other => Err(io::Error::new(
                    io::ErrorKind::Unsupported,
                    format!("Server selected unsupported auth method: {}", other),
                )),
            },
            Some(response) => Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!("Unexpected response to handshake: {}", response),
            )),
            None => Err(io::Error::new(
                io::ErrorKind::UnexpectedEof,
                "No handshake response from server",
            )),
        }
    }

    /// Requests a connection to the destination server through the SOCKS5 proxy.
    ///
    /// This method consumes the client and returns a stream that can be used for
    /// direct communication with the destination server. All data sent through
    /// this stream will be tunneled through the SOCKS5 proxy.
    pub async fn request_connect(
        mut self,
        dest_addr: Socks5Address,
        dest_port: u16,
    ) -> io::Result<Socks5ClientStream> {
        let request = Socks5Request {
            address: dest_addr,
            port: dest_port,
            command: Socks5Command::Connect,
        };

        self.framed.send(ClientRequest::Connect(request)).await?;
        self.framed.flush().await?;

        match self.framed.next().await.transpose()? {
            Some(ClientResponse::ConnectResponse(response)) => match response.response {
                Socks5Status::RequestGranted => Ok(Socks5ClientStream(self.framed.into_inner())),
                other => Err(io::Error::new(
                    io::ErrorKind::ConnectionRefused,
                    format!("SOCKS5 request failed with status: {:?}", other),
                )),
            },
            Some(response) => Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!("Unexpected response to connect request: {}", response),
            )),
            None => Err(io::Error::new(
                io::ErrorKind::UnexpectedEof,
                "No response from server to connect request",
            )),
        }
    }
}

pub struct Socks5ClientStream(TcpStream);

impl AsyncRead for Socks5ClientStream {
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut std::task::Context,
        buf: &mut tokio::io::ReadBuf<'_>,
    ) -> Poll<std::io::Result<()>> {
        Pin::new(&mut self.0).poll_read(cx, buf)
    }
}

impl AsyncWrite for Socks5ClientStream {
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut std::task::Context,
        buf: &[u8],
    ) -> Poll<io::Result<usize>> {
        Pin::new(&mut self.0).poll_write(cx, buf)
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut std::task::Context) -> Poll<io::Result<()>> {
        Pin::new(&mut self.0).poll_flush(cx)
    }

    fn poll_shutdown(
        mut self: Pin<&mut Self>,
        cx: &mut std::task::Context,
    ) -> Poll<io::Result<()>> {
        Pin::new(&mut self.0).poll_shutdown(cx)
    }
}