socks_lib/v5/
server.rs

1use std::io;
2
3use bytes::BytesMut;
4use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
5
6use super::{Method, Request, Response, Stream};
7
8impl<T> Stream<T>
9where
10    T: AsyncRead + AsyncWrite + Unpin,
11{
12    /// # Methods
13    ///
14    /// ```text
15    ///  +----+----------+----------+
16    ///  |VER | NMETHODS | METHODS  |
17    ///  +----+----------+----------+
18    ///  | 1  |    1     | 1 to 255 |
19    ///  +----+----------+----------+
20    /// ```
21    #[inline]
22    pub async fn read_methods(&mut self) -> io::Result<Vec<Method>> {
23        let mut buffer = [0u8; 2];
24        self.0.read_exact(&mut buffer).await?;
25
26        let method_num = buffer[1];
27        if method_num == 1 {
28            let method = self.0.read_u8().await?;
29            return Ok(vec![Method::from_u8(method)]);
30        }
31
32        let mut methods = vec![0u8; method_num as usize];
33        self.0.read_exact(&mut methods).await?;
34
35        let result = methods.into_iter().map(|e| Method::from_u8(e)).collect();
36
37        Ok(result)
38    }
39
40    ///
41    /// ```text
42    ///  +----+--------+
43    ///  |VER | METHOD |
44    ///  +----+--------+
45    ///  | 1  |   1    |
46    ///  +----+--------+
47    ///  ```
48    #[inline]
49    pub async fn write_auth_method(&mut self, method: Method) -> io::Result<usize> {
50        let bytes = [self.version().into(), method.as_u8()];
51        self.0.write(&bytes).await
52    }
53
54    ///
55    /// ```text
56    ///  +----+-----+-------+------+----------+----------+
57    ///  |VER | CMD |  RSV  | ATYP | DST.ADDR | DST.PORT |
58    ///  +----+-----+-------+------+----------+----------+
59    ///  | 1  |  1  | X'00' |  1   | Variable |    2     |
60    ///  +----+-----+-------+------+----------+----------+
61    /// ```
62    ///
63    #[inline]
64    pub async fn read_request(&mut self) -> io::Result<Request> {
65        let _version = self.0.read_u8().await?;
66        Request::from_async_read(&mut self.0).await
67    }
68
69    ///
70    /// ```text
71    ///  +----+-----+-------+------+----------+----------+
72    ///  |VER | REP |  RSV  | ATYP | BND.ADDR | BND.PORT |
73    ///  +----+-----+-------+------+----------+----------+
74    ///  | 1  |  1  | X'00' |  1   | Variable |    2     |
75    ///  +----+-----+-------+------+----------+----------+
76    /// ```
77    ///
78    #[inline]
79    pub async fn write_response<'a>(&mut self, resp: &Response<'a>) -> io::Result<usize> {
80        let bytes = prepend_u8(resp.to_bytes(), self.version().into());
81        self.0.write(&bytes).await
82    }
83}
84
85fn prepend_u8(mut bytes: BytesMut, value: u8) -> BytesMut {
86    bytes.reserve(1);
87
88    unsafe {
89        let ptr = bytes.as_mut_ptr();
90        std::ptr::copy(ptr, ptr.add(1), bytes.len());
91        std::ptr::write(ptr, value);
92        let new_len = bytes.len() + 1;
93        bytes.set_len(new_len);
94    }
95
96    bytes
97}
socks_lib/v5/server.rs

socks_lib/v5/
server.rs
