socks5-impl 0.3.3

use crate::protocol::{handshake::password_method, AuthMethod, UserKey};
use async_trait::async_trait;
use tokio::net::TcpStream;

/// This trait is for defining the socks5 authentication method.
///
/// Pre-defined authentication methods can be found in the [`auth`](https://docs.rs/socks5-impl/latest/socks5_impl/server/auth/index.html) module.
///
/// You can create your own authentication method by implementing this trait. Since GAT is not stabled yet,
/// [async_trait](https://docs.rs/async-trait/latest/async_trait/index.html) needs to be used.
///
/// # Example
/// ```rust
/// use async_trait::async_trait;
/// use socks5_impl::protocol::AuthMethod;
/// use socks5_impl::server::AuthExecutor;
/// use tokio::net::TcpStream;
///
/// pub struct MyAuth;
///
/// #[async_trait]
/// impl AuthExecutor for MyAuth {
///     fn auth_method(&self) -> AuthMethod {
///         AuthMethod::from(0x80)
///     }
///
///     async fn execute(&self, stream: &mut TcpStream) -> std::io::Result<()> {
///         // do something
///         Ok(())
///     }
/// }
/// ```
#[async_trait]
pub trait AuthExecutor {
    fn auth_method(&self) -> AuthMethod;
    async fn execute(&self, stream: &mut TcpStream) -> std::io::Result<()>;
}

/// No authentication as the socks5 handshake method.
#[derive(Debug, Default)]
pub struct NoAuth;

#[async_trait]
impl AuthExecutor for NoAuth {
    fn auth_method(&self) -> AuthMethod {
        AuthMethod::NoAuth
    }

    async fn execute(&self, _: &mut TcpStream) -> std::io::Result<()> {
        Ok(())
    }
}

/// Username and password as the socks5 handshake method.
pub struct PasswordAuth {
    user_key: UserKey,
}

impl PasswordAuth {
    pub fn new(username: &str, password: &str) -> Self {
        let user_key = UserKey::new(username, password);
        Self { user_key }
    }
}

#[async_trait]
impl AuthExecutor for PasswordAuth {
    fn auth_method(&self) -> AuthMethod {
        AuthMethod::UserPass
    }

    async fn execute(&self, stream: &mut TcpStream) -> std::io::Result<()> {
        let req = password_method::Request::rebuild_from_stream(stream).await?;

        if req.user_key == self.user_key {
            let resp = password_method::Response::new(true);
            resp.write_to_stream(stream).await?;
            Ok(())
        } else {
            let resp = password_method::Response::new(false);
            resp.write_to_stream(stream).await?;
            let err = "SOCKS5 username / password authentication failed";
            Err(std::io::Error::new(std::io::ErrorKind::InvalidData, err))
        }
    }
}