Skip to main content

socks5_impl/protocol/handshake/
request.rs

1#[cfg(feature = "tokio")]
2use crate::AsyncStreamOperation;
3use crate::{
4    StreamOperation,
5    protocol::{AuthMethod, Version},
6};
7#[cfg(feature = "tokio")]
8use tokio::io::{AsyncRead, AsyncReadExt};
9
10/// SOCKS5 handshake request
11///
12/// ```plain
13/// +-----+----------+----------+
14/// | VER | NMETHODS | METHODS  |
15/// +-----+----------+----------+
16/// |  1  |    1     | 1 to 255 |
17/// +-----+----------+----------|
18/// ```
19#[derive(Clone, Debug)]
20pub struct Request {
21    methods: Vec<AuthMethod>,
22}
23
24impl Request {
25    pub fn new(methods: Vec<AuthMethod>) -> Self {
26        Self { methods }
27    }
28
29    pub fn evaluate_method(&self, server_method: AuthMethod) -> bool {
30        self.methods.contains(&server_method)
31    }
32}
33
34impl StreamOperation for Request {
35    fn retrieve_from_stream<R: std::io::Read>(r: &mut R) -> std::io::Result<Self> {
36        let mut ver = [0; 1];
37        r.read_exact(&mut ver)?;
38        let ver = Version::try_from(ver[0])?;
39
40        if ver != Version::V5 {
41            let err = format!("Unsupported SOCKS version {0:#x}", u8::from(ver));
42            return Err(std::io::Error::new(std::io::ErrorKind::Unsupported, err));
43        }
44
45        let mut mlen = [0; 1];
46        r.read_exact(&mut mlen)?;
47        let mlen = mlen[0];
48
49        let mut methods = vec![0; mlen as usize];
50        r.read_exact(&mut methods)?;
51
52        let methods = methods.into_iter().map(AuthMethod::from).collect();
53
54        Ok(Self { methods })
55    }
56
57    fn write_to_buf<B: bytes::BufMut>(&self, buf: &mut B) {
58        buf.put_u8(Version::V5.into());
59        buf.put_u8(u8::try_from(self.methods.len()).expect("too many handshake methods for SOCKS5"));
60
61        let methods = self.methods.iter().map(u8::from).collect::<Vec<u8>>();
62        buf.put_slice(&methods);
63    }
64
65    fn len(&self) -> usize {
66        2 + self.methods.len()
67    }
68}
69
70#[cfg(feature = "tokio")]
71#[async_trait::async_trait]
72impl AsyncStreamOperation for Request {
73    async fn retrieve_from_async_stream<R>(r: &mut R) -> std::io::Result<Self>
74    where
75        R: AsyncRead + Unpin + Send + ?Sized,
76    {
77        let ver = Version::try_from(r.read_u8().await?)?;
78
79        if ver != Version::V5 {
80            let err = format!("Unsupported SOCKS version {0:#x}", u8::from(ver));
81            return Err(std::io::Error::new(std::io::ErrorKind::Unsupported, err));
82        }
83
84        let mlen = r.read_u8().await?;
85        let mut methods = vec![0; mlen as usize];
86        r.read_exact(&mut methods).await?;
87
88        let methods = methods.into_iter().map(AuthMethod::from).collect();
89
90        Ok(Self { methods })
91    }
92}