socks5_impl/server/
auth.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
use crate::protocol::{handshake::password_method, AsyncStreamOperation, AuthMethod, UserKey};
use as_any::AsAny;
use async_trait::async_trait;
use std::sync::Arc;
use tokio::net::TcpStream;

/// This trait is for defining the socks5 authentication method.
///
/// Pre-defined authentication methods can be found in the [`auth`](crate::server::auth) 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 {
///     type Output = std::io::Result<usize>;
///
///     fn auth_method(&self) -> AuthMethod {
///         AuthMethod::from(0x80)
///     }
///
///     async fn execute(&self, stream: &mut TcpStream) -> Self::Output {
///         // do something
///         Ok(1145141919810)
///     }
/// }
/// ```
#[async_trait]
pub trait AuthExecutor {
    type Output: AsAny;
    fn auth_method(&self) -> AuthMethod;
    async fn execute(&self, stream: &mut TcpStream) -> Self::Output;
}

pub type AuthAdaptor<O> = Arc<dyn AuthExecutor<Output = O> + Send + Sync>;

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

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

    async fn execute(&self, _: &mut TcpStream) -> Self::Output {}
}

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

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

#[async_trait]
impl AuthExecutor for UserKeyAuth {
    type Output = std::io::Result<bool>;

    fn auth_method(&self) -> AuthMethod {
        AuthMethod::UserPass
    }

    async fn execute(&self, stream: &mut TcpStream) -> Self::Output {
        use password_method::{Request, Response, Status::*};
        let req = Request::retrieve_from_async_stream(stream).await?;

        let is_equal = req.user_key == self.user_key;
        let resp = Response::new(if is_equal { Succeeded } else { Failed });
        resp.write_to_async_stream(stream).await?;
        if is_equal {
            Ok(true)
        } else {
            Err(std::io::Error::new(std::io::ErrorKind::Other, "username or password is incorrect"))
        }
    }
}