socks5_proto/
udp.rs

1use crate::{address::AddressError, Address, Error, ProtocolError};
2use bytes::{BufMut, BytesMut};
3use std::io::Error as IoError;
4use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
5
6/// SOCKS5 UDP packet header
7///
8/// ```plain
9/// +-----+------+------+----------+----------+----------+
10/// | RSV | FRAG | ATYP | DST.ADDR | DST.PORT |   DATA   |
11/// +-----+------+------+----------+----------+----------+
12/// |  2  |  1   |  1   | Variable |    2     | Variable |
13/// +-----+------+------+----------+----------+----------+
14/// ```
15#[derive(Clone, Debug)]
16pub struct UdpHeader {
17    pub frag: u8,
18    pub address: Address,
19}
20
21impl UdpHeader {
22    pub const fn new(frag: u8, address: Address) -> Self {
23        Self { frag, address }
24    }
25
26    pub async fn read_from<R>(r: &mut R) -> Result<Self, Error>
27    where
28        R: AsyncRead + Unpin,
29    {
30        r.read_exact(&mut [0; 2]).await?;
31
32        let frag = r.read_u8().await?;
33
34        let addr = Address::read_from(r).await.map_err(|err| match err {
35            AddressError::Io(err) => Error::Io(err),
36            AddressError::InvalidType(code) => {
37                Error::Protocol(ProtocolError::InvalidAddressTypeInUdpHeader {
38                    frag,
39                    address_type: code,
40                })
41            }
42        })?;
43
44        Ok(Self::new(frag, addr))
45    }
46
47    pub async fn write_to<W>(&self, w: &mut W) -> Result<(), IoError>
48    where
49        W: AsyncWrite + Unpin,
50    {
51        let mut buf = BytesMut::with_capacity(self.serialized_len());
52        self.write_to_buf(&mut buf);
53        w.write_all(&buf).await?;
54
55        Ok(())
56    }
57
58    pub fn write_to_buf<B: BufMut>(&self, buf: &mut B) {
59        buf.put_bytes(0x00, 2);
60        buf.put_u8(self.frag);
61        self.address.write_to_buf(buf);
62    }
63
64    pub fn serialized_len(&self) -> usize {
65        2 + 1 + self.address.serialized_len()
66    }
67}
socks5_proto/udp.rs

socks5_proto/
udp.rs
