socks5_impl/protocol/handshake/
request.rs

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