socks5_proto/handshake/
request.rs

1use super::Method;
2use crate::{Error, ProtocolError};
3use bytes::{BufMut, BytesMut};
4use std::{
5    io::Error as IoError,
6    mem::{self, ManuallyDrop},
7};
8use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
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    pub methods: Vec<Method>,
22}
23
24impl Request {
25    pub const fn new(methods: Vec<Method>) -> Self {
26        Self { methods }
27    }
28
29    pub async fn read_from<R>(r: &mut R) -> Result<Self, Error>
30    where
31        R: AsyncRead + Unpin,
32    {
33        let ver = r.read_u8().await?;
34
35        if ver != crate::SOCKS_VERSION {
36            return Err(Error::Protocol(ProtocolError::ProtocolVersion {
37                version: ver,
38            }));
39        }
40
41        let mlen = r.read_u8().await?;
42        let mut methods = vec![0; mlen as usize];
43        r.read_exact(&mut methods).await?;
44
45        let methods = unsafe {
46            let mut methods = ManuallyDrop::new(methods);
47
48            Vec::from_raw_parts(
49                methods.as_mut_ptr() as *mut Method,
50                methods.len(),
51                methods.capacity(),
52            )
53        };
54
55        Ok(Self::new(methods))
56    }
57
58    pub async fn write_to<W>(&self, w: &mut W) -> Result<(), IoError>
59    where
60        W: AsyncWrite + Unpin,
61    {
62        let mut buf = BytesMut::with_capacity(self.serialized_len());
63        self.write_to_buf(&mut buf);
64        w.write_all(&buf).await?;
65
66        Ok(())
67    }
68
69    pub fn write_to_buf<B: BufMut>(&self, buf: &mut B) {
70        buf.put_u8(crate::SOCKS_VERSION);
71        buf.put_u8(self.methods.len() as u8);
72
73        let methods = unsafe { mem::transmute(self.methods.as_slice()) };
74        buf.put_slice(methods);
75    }
76
77    pub fn serialized_len(&self) -> usize {
78        1 + 1 + self.methods.len()
79    }
80}