socks5x 0.1.3

use bytes::{Buf, BytesMut};
use futures_util::SinkExt;
use std::io;
use tokio::io::{AsyncRead, AsyncWrite};
use tokio::net::TcpStream;
use tokio_stream::StreamExt;
use tokio_util::codec::{Decoder, Encoder, Framed};

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

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

/// Messages that can be sent from server to client during SOCKS5 handshake
enum Socks5ServerMessage {
    /// Authentication method selection response
    AuthMethodSelection { method: AuthMethod },
    /// Username/password authentication response
    AuthResponse { status: u8 },
    /// SOCKS5 response (connection result)
    Response(Socks5Response),
}

/// Messages that can be received by server from client during SOCKS5 handshake
enum Socks5ClientMessage {
    /// Client greeting with supported auth methods
    Greeting { methods: Vec<AuthMethod> },
    /// Username/password authentication attempt
    AuthRequest { username: String, password: String },
    /// SOCKS5 request (connect, bind, etc.)
    Request(Socks5Request),
}

/// Tokio codec for encoding/decoding SOCKS5 messages on the server side.
struct ServerCodec {
    /// Current state of the SOCKS5 handshake
    state: HandshakeState,
}

#[derive(Clone, Copy)]
enum HandshakeState {
    /// Expecting client greeting
    WaitingGreeting,
    /// Expecting authentication request (username/password)
    WaitingAuth,
    /// Skip auth, go directly to request
    WaitingRequestNoAuth,
    /// Expecting SOCKS5 request
    WaitingRequest,
    /// Handshake complete, raw data mode
    Connected,
}

impl ServerCodec {
    fn new() -> Self {
        Self {
            state: HandshakeState::WaitingGreeting,
        }
    }
}

impl Decoder for ServerCodec {
    type Item = Socks5ClientMessage;
    type Error = io::Error;

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

                let version = src[0];
                let nmethods = src[1] as usize;

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

                if src.len() < 2 + nmethods {
                    return Ok(None);
                }

                let methods = src[2..2 + nmethods]
                    .iter()
                    .filter_map(|&method| AuthMethod::try_from(method).ok())
                    .collect();

                src.advance(2 + nmethods);
                self.state = HandshakeState::WaitingAuth;

                Ok(Some(Socks5ClientMessage::Greeting { methods }))
            }
            HandshakeState::WaitingAuth => self.decode_auth_request(src),
            HandshakeState::WaitingRequestNoAuth => {
                self.state = HandshakeState::WaitingRequest;
                // For no-auth, we transition directly to request parsing
                self.decode(src)
            }
            HandshakeState::WaitingRequest => {
                if let Some((command, address, port)) = decode_res(src)? {
                    self.state = HandshakeState::Connected;
                    Ok(Some(Socks5ClientMessage::Request(Socks5Request {
                        address,
                        port,
                        command,
                    })))
                } else {
                    Ok(None)
                }
            }
            HandshakeState::Connected => {
                // In connected state, we don't decode messages anymore
                Ok(None)
            }
        }
    }
}

impl ServerCodec {
    fn decode_auth_request(
        &mut self,
        src: &mut BytesMut,
    ) -> Result<Option<Socks5ClientMessage>, io::Error> {
        if src.len() < 2 {
            return Ok(None);
        }

        let auth_ver = src[0];
        if auth_ver != SOCKS5_USERNAME_AUTH_VER {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!("Unsupported auth version: {auth_ver}"),
            ));
        }

        let username_len = src[1] as usize;
        if src.len() < 2 + username_len + 1 {
            return Ok(None);
        }

        let username = String::from_utf8_lossy(&src[2..2 + username_len]).to_string();
        let password_len = src[2 + username_len] as usize;

        if src.len() < 2 + username_len + 1 + password_len {
            return Ok(None);
        }

        let password =
            String::from_utf8_lossy(&src[3 + username_len..3 + username_len + password_len])
                .to_string();

        src.advance(3 + username_len + password_len);
        self.state = HandshakeState::WaitingRequest;

        Ok(Some(Socks5ClientMessage::AuthRequest {
            username,
            password,
        }))
    }
}

impl Encoder<Socks5ServerMessage> for ServerCodec {
    type Error = io::Error;

    fn encode(&mut self, item: Socks5ServerMessage, dst: &mut BytesMut) -> Result<(), Self::Error> {
        match item {
            Socks5ServerMessage::AuthMethodSelection { method } => {
                dst.reserve(2);
                dst.extend_from_slice(&[SOCKS5_VERSION, method as u8]);
            }
            Socks5ServerMessage::AuthResponse { status } => {
                dst.reserve(2);
                dst.extend_from_slice(&[SOCKS5_USERNAME_AUTH_VER, status]);
            }
            Socks5ServerMessage::Response(response) => {
                dst.reserve(22); // Max size for IPv6 response
                dst.extend_from_slice(&[SOCKS5_VERSION, response.response as u8, SOCKS5_RESERVED]);
                response.address.to_bytes(dst);
                dst.extend_from_slice(&response.port.to_be_bytes());
            }
        }
        Ok(())
    }
}

/// Handles the initial SOCKS5 client greeting and authentication method selection.
async fn handle_client_greeting(
    framed: &mut Framed<TcpStream, ServerCodec>,
    auth_validator: &Auth<impl AuthValidator>,
) -> io::Result<()> {
    if let Some(message) = framed.next().await.transpose()?
        && let Socks5ClientMessage::Greeting { methods } = message
    {
        let expected_method = match auth_validator {
            Auth::NoAuth => AuthMethod::NoAuth,
            Auth::UserPassword(_) => AuthMethod::UsernamePassword,
        };

        let found_method = methods
            .iter()
            .find(|&method| method == &expected_method)
            .cloned()
            .ok_or_else(|| {
                io::Error::new(io::ErrorKind::InvalidData, "Suitable auth method not found")
            })?;

        // Update codec state based on selected auth method
        match found_method {
            AuthMethod::NoAuth => {
                framed.codec_mut().state = HandshakeState::WaitingRequestNoAuth;
            }
            AuthMethod::UsernamePassword => {
                framed.codec_mut().state = HandshakeState::WaitingAuth;
            }
            auth_method => {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    format!("Unsupported auth method: {auth_method}"),
                ));
            }
        }

        framed
            .send(Socks5ServerMessage::AuthMethodSelection {
                method: found_method,
            })
            .await?;

        return Ok(());
    }

    Err(io::Error::new(
        io::ErrorKind::InvalidData,
        "Expected greeting message",
    ))
}

/// Handles username/password authentication if required.
async fn handle_client_authenticate<A: AuthValidator>(
    framed: &mut Framed<TcpStream, ServerCodec>,
    auth_validator: &Auth<A>,
) -> io::Result<()> {
    match auth_validator {
        Auth::NoAuth => Ok(()),
        Auth::UserPassword(validator) => {
            if let Some(message) = framed.next().await.transpose()?
                && let Socks5ClientMessage::AuthRequest { username, password } = message
            {
                let is_valid = validator.validate(&username, &password).await;
                let status = if is_valid { 0 } else { 1 };

                framed
                    .send(Socks5ServerMessage::AuthResponse { status })
                    .await?;

                if status != 0 {
                    return Err(io::Error::new(
                        io::ErrorKind::PermissionDenied,
                        "Authentication failed",
                    ));
                }

                return Ok(());
            }

            Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "Expected auth request message",
            ))
        }
    }
}

/// Validates the client command (currently only CONNECT is supported).
async fn handle_client_command(
    request: &Socks5Request,
    framed: &mut Framed<TcpStream, ServerCodec>,
) -> io::Result<()> {
    if !matches!(request.command, Socks5Command::Connect) {
        framed
            .send(Socks5ServerMessage::Response(Socks5Response {
                response: Socks5Status::CommandNotSupported,
                address: Socks5Address::IPv4(std::net::Ipv4Addr::new(0, 0, 0, 0)),
                port: 0,
            }))
            .await?;
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            "Only CONNECT supported",
        ));
    }
    Ok(())
}

/// Creates a connection to the target destination.
async fn handle_create_stream<C: ConnectionCreator>(
    request: Socks5Request,
    framed: &mut Framed<TcpStream, ServerCodec>,
    connection_creator: &C,
) -> io::Result<C::Stream> {
    let dest_stream = match connection_creator
        .create_stream(request.address, request.port)
        .await
    {
        Ok(stream) => stream,
        Err(_) => {
            framed
                .send(Socks5ServerMessage::Response(Socks5Response {
                    response: Socks5Status::HostUnreachable,
                    address: Socks5Address::IPv4(std::net::Ipv4Addr::new(0, 0, 0, 0)),
                    port: 0,
                }))
                .await?;
            return Err(io::Error::new(
                io::ErrorKind::ConnectionRefused,
                "Connection failed",
            ));
        }
    };

    Ok(dest_stream)
}

/// Mirrors data between client and destination streams.
async fn mirror_stream<S: SocksSplitable>(
    dest_stream: S,
    mut framed: Framed<TcpStream, ServerCodec>,
) -> io::Result<()> {
    framed
        .send(Socks5ServerMessage::Response(Socks5Response {
            response: Socks5Status::RequestGranted,
            address: Socks5Address::IPv4(std::net::Ipv4Addr::new(0, 0, 0, 0)),
            port: 0,
        }))
        .await?;
    framed.flush().await?;

    let mut socket = framed.into_inner();
    let (mut client_read, mut client_write) = socket.split();
    let (mut dest_read, mut dest_write) = dest_stream.split_for_socks()?;

    let client_to_dest = tokio::io::copy(&mut client_read, &mut dest_write);
    let dest_to_client = tokio::io::copy(&mut dest_read, &mut client_write);

    tokio::select! {
        rc2d = client_to_dest => rc2d,
        rd2c = dest_to_client => rd2c,
    }?;
    Ok(())
}

/// Handles SOCKS5 client connections with configurable authentication and connection creation.
pub struct ClientHandler<H, A = NoAuth> {
    auth_validator: Auth<A>,
    connection_handler: H,
}

impl<H> ClientHandler<H>
where
    H: ConnectionCreator,
{
    /// Creates a new client handler with no authentication required.
    pub fn no_auth(connection_creator: H) -> Self {
        ClientHandler {
            auth_validator: Auth::NoAuth,
            connection_handler: connection_creator,
        }
    }
}

impl<A, H> ClientHandler<H, A>
where
    A: AuthValidator,
    H: ConnectionCreator,
{
    /// Creates a new client handler with the specified authentication validator and connection creator.
    pub fn new(auth_validator: Auth<A>, connection_creator: H) -> Self {
        ClientHandler {
            auth_validator,
            connection_handler: connection_creator,
        }
    }

    /// Handles a SOCKS5 client connection through the complete protocol flow.
    ///
    /// This method manages the entire SOCKS5 handshake including greeting,
    /// authentication, command processing, and data proxying.
    pub async fn handle(&self, socket: TcpStream) -> io::Result<()> {
        let mut framed = Framed::new(socket, ServerCodec::new());

        // Handle greeting and method selection
        handle_client_greeting(&mut framed, &self.auth_validator).await?;

        // Handle authentication if required
        handle_client_authenticate(&mut framed, &self.auth_validator).await?;

        // Handle SOCKS5 request
        if let Some(message) = framed.next().await.transpose()?
            && let Socks5ClientMessage::Request(request) = message
        {
            handle_client_command(&request, &mut framed).await?;
            let dest_stream =
                handle_create_stream(request, &mut framed, &self.connection_handler).await?;
            return mirror_stream(dest_stream, framed).await;
        }

        Ok(())
    }
}

/// Placeholder type for no authentication.
pub struct NoAuth;

/// Authentication configuration enum.
pub enum Auth<A = NoAuth> {
    NoAuth,
    UserPassword(A),
}

impl AuthValidator for NoAuth {
    async fn validate(&self, _: &str, _: &str) -> bool {
        unimplemented!()
    }
}

/// Trait for splitting bidirectional streams into separate read/write halves for SOCKS5.
pub trait SocksSplitable {
    /// Splits stream into boxed read and write halves.
    fn split_for_socks(
        self,
    ) -> io::Result<(
        Box<dyn tokio::io::AsyncRead + Unpin + Send>,
        Box<dyn tokio::io::AsyncWrite + Unpin + Send>,
    )>;
}

impl SocksSplitable for TcpStream {
    fn split_for_socks(
        self,
    ) -> io::Result<(
        Box<dyn AsyncRead + Unpin + Send>,
        Box<dyn AsyncWrite + Unpin + Send>,
    )> {
        let (read, write) = self.into_split();
        Ok((Box::new(read), Box::new(write)))
    }
}

/// Trait for validating username/password authentication.
/// 
/// Validates the provided username and password credentials.
///
/// Returns `true` if the credentials are valid, `false` otherwise.
pub trait AuthValidator: Send + Sync {
    fn validate(
        &self,
        username: &str,
        password: &str,
    ) -> impl std::future::Future<Output = bool> + Send;
}


/// Trait for creating connections to destination addresses in a SOCKS5 server.
///
/// Allows customizing how the SOCKS5 server establishes outbound connections,
/// such as using different protocols, connection pooling, or proxying.
pub trait ConnectionCreator: Send + Sync {
    /// Trait for creating connections to destination addresses.
    type Stream: SocksSplitable;

    /// Creates a connection to the specified address and port.
    fn create_stream(
        &self,
        address: Socks5Address,
        port: u16,
    ) -> impl std::future::Future<Output = io::Result<Self::Stream>> + Send;
}

/// Default connection creator that establishes direct TCP connections.
pub struct DefaultConnectionCreator;

impl ConnectionCreator for DefaultConnectionCreator {
    type Stream = TcpStream;
    async fn create_stream(&self, address: Socks5Address, port: u16) -> io::Result<Self::Stream> {
        TcpStream::connect((address.to_string(), port)).await
    }
}