better_socks/
udp.rs

1use crate::{tcp::*, Error, Result, TargetAddr, ToProxyAddrs};
2use bytes::{BufMut, Bytes, BytesMut};
3use std::{
4    io,
5    net::{Ipv4Addr, Ipv6Addr, SocketAddr},
6    pin::Pin,
7    result::Result as StdResult,
8    task::{Context, Poll},
9};
10use tokio::{
11    io::{AsyncRead, ReadBuf},
12    net::{TcpStream, ToSocketAddrs, UdpSocket},
13};
14use tokio_util::{
15    codec::{Decoder, Encoder},
16    udp::UdpFramed,
17};
18
19use futures_core::Stream;
20use futures_sink::Sink;
21use pin_project::pin_project;
22
23#[pin_project]
24pub struct Socks5UdpFramed {
25    #[pin]
26    framed: UdpFramed<Socks5UdpCodec, UdpSocket>,
27    #[pin]
28    stream: Socks5Stream<TcpStream>,
29    socks_addr: SocketAddr,
30}
31
32// +----+------+------+----------+----------+----------+
33// |RSV | FRAG | ATYP | DST.ADDR | DST.PORT |   DATA   |
34// +----+------+------+----------+----------+----------+
35// | 2  |  1   |  1   | Variable |    2     | Variable |
36// +----+------+------+----------+----------+----------+
37#[derive(Debug)]
38pub struct Socks5UdpMessage {
39    pub rsv: [u8; 2],
40    pub frag: u8,
41    pub atyp: u8,
42    pub dst_addr: TargetAddr<'static>,
43    pub data: BytesMut,
44}
45
46#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash, Default)]
47pub struct Socks5UdpCodec;
48
49impl Socks5UdpFramed {
50    pub async fn connect<P, T>(proxy: P, bind_addr: Option<T>) -> Result<Self>
51    where
52        P: ToProxyAddrs,
53        T: ToSocketAddrs,
54    {
55        let socket = match bind_addr {
56            None => UdpSocket::bind(SocketAddr::from(([0, 0, 0, 0], 0))).await?,
57            Some(addr) => UdpSocket::bind(addr).await?,
58        };
59        let local = socket.local_addr()?;
60        let stream = Socks5Stream::associate(proxy, local).await?;
61        let framed = UdpFramed::new(socket, Socks5UdpCodec::new());
62        let socks_addr = relay_socket_addr(stream.target_addr())?;
63        Ok(Self {
64            framed,
65            stream,
66            socks_addr,
67        })
68    }
69
70    pub async fn connect_with_password<'a, P, T>(
71        proxy: P,
72        bind_addr: Option<T>,
73        username: &'a str,
74        password: &'a str,
75    ) -> Result<Self>
76    where
77        P: ToProxyAddrs,
78        T: ToSocketAddrs,
79    {
80        let socket = match bind_addr {
81            None => UdpSocket::bind(SocketAddr::from(([0, 0, 0, 0], 0))).await?,
82            Some(addr) => UdpSocket::bind(addr).await?,
83        };
84        let local = socket.local_addr()?;
85        let stream = Socks5Stream::associate_with_password(proxy, local, username, password).await?;
86        let framed = UdpFramed::new(socket, Socks5UdpCodec::new());
87        let socks_addr = relay_socket_addr(stream.target_addr())?;
88        Ok(Self {
89            framed,
90            stream,
91            socks_addr,
92        })
93    }
94
95    pub fn socks_addr(&self) -> &SocketAddr {
96        &self.socks_addr
97    }
98
99    pub fn local_addr(&self) -> io::Result<SocketAddr> {
100        self.framed.get_ref().local_addr()
101    }
102}
103
104impl Stream for Socks5UdpFramed {
105    type Item = StdResult<(<Socks5UdpCodec as Decoder>::Item, SocketAddr), <Socks5UdpCodec as Decoder>::Error>;
106
107    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
108        let this = self.project();
109        if let Poll::Ready(d) = this.framed.poll_next(cx) {
110            return Poll::Ready(d);
111        }
112
113        let mut buf = [0u8; 512];
114        let mut buf = ReadBuf::new(&mut buf[..]);
115        match this.stream.poll_read(cx, &mut buf) {
116            Poll::Ready(Err(e)) => {
117                return Poll::Ready(Some(Err(Error::Io(e))));
118            },
119            Poll::Ready(Ok(())) => {
120                // EOF on the TCP control connection means the UDP association is gone.
121                if buf.filled().is_empty() {
122                    return Poll::Ready(None);
123                }
124                // Unexpected control-channel data is ignored; keep waiting for UDP.
125            },
126            Poll::Pending => {},
127        }
128
129        Poll::Pending
130    }
131}
132
133impl Sink<(Bytes, TargetAddr<'static>)> for Socks5UdpFramed {
134    type Error = <Socks5UdpCodec as Encoder<(Bytes, TargetAddr<'static>)>>::Error;
135
136    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<StdResult<(), Self::Error>> {
137        self.project().framed.poll_ready(cx)
138    }
139
140    fn start_send(self: Pin<&mut Self>, item: (Bytes, TargetAddr<'static>)) -> StdResult<(), Self::Error> {
141        let send_addr = *self.socks_addr();
142        self.project().framed.start_send((item, send_addr))
143    }
144
145    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<StdResult<(), Self::Error>> {
146        self.project().framed.poll_flush(cx)
147    }
148
149    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<StdResult<(), Self::Error>> {
150        self.project().framed.poll_close(cx)
151    }
152}
153
154fn relay_socket_addr(addr: TargetAddr<'_>) -> Result<SocketAddr> {
155    match addr {
156        TargetAddr::Ip(addr) => Ok(addr),
157        TargetAddr::Domain(_, _) => Err(Error::InvalidTargetAddress(
158            "UDP ASSOCIATE bind address must be an IP address",
159        )),
160    }
161}
162
163impl Socks5UdpCodec {
164    pub fn new() -> Self {
165        Self {}
166    }
167}
168
169impl Default for Socks5UdpMessage {
170    fn default() -> Self {
171        Self::new()
172    }
173}
174
175impl Socks5UdpMessage {
176    pub fn new() -> Self {
177        Self {
178            rsv: [0u8; 2],
179            frag: 0u8,
180            atyp: 0u8,
181            dst_addr: TargetAddr::Ip(SocketAddr::from(([0, 0, 0, 0], 0))),
182            data: BytesMut::new(),
183        }
184    }
185}
186
187// +----+------+------+----------+----------+----------+
188// |RSV | FRAG | ATYP | DST.ADDR | DST.PORT |   DATA   |
189// +----+------+------+----------+----------+----------+
190// | 2  |  1   |  1   | Variable |    2     | Variable |
191// +----+------+------+----------+----------+----------+
192impl Decoder for Socks5UdpCodec {
193    type Error = Error;
194    type Item = Socks5UdpMessage;
195
196    fn decode(&mut self, buf: &mut BytesMut) -> StdResult<Option<Self::Item>, Self::Error> {
197        if buf.is_empty() {
198            return Ok(None);
199        }
200
201        // RSV(2) + FRAG(1) + ATYP(1). UdpFramed delivers a whole datagram, so a
202        // short buffer is a truncated packet rather than a partial frame.
203        if buf.len() < 4 {
204            return Err(Error::InvalidTargetAddress("UDP datagram shorter than SOCKS5 header"));
205        }
206
207        let mut msg = Socks5UdpMessage::new();
208        msg.rsv.copy_from_slice(&buf[0..2]);
209        if msg.rsv != [0u8, 0u8] {
210            return Err(Error::InvalidReservedByte);
211        }
212
213        msg.frag = buf[2];
214        if msg.frag != 0 {
215            return Err(Error::FragmentationNotSupported);
216        }
217        msg.atyp = buf[3];
218
219        let header_len = match msg.atyp {
220            0x01 => {
221                if buf.len() < 10 {
222                    return Err(Error::InvalidTargetAddress("truncated IPv4 UDP datagram"));
223                }
224                let ip = Ipv4Addr::new(buf[4], buf[5], buf[6], buf[7]);
225                let port = u16::from_be_bytes([buf[8], buf[9]]);
226                msg.dst_addr = TargetAddr::Ip(SocketAddr::from((ip, port)));
227                10
228            },
229            0x04 => {
230                if buf.len() < 22 {
231                    return Err(Error::InvalidTargetAddress("truncated IPv6 UDP datagram"));
232                }
233                let mut octets = [0u8; 16];
234                octets.copy_from_slice(&buf[4..20]);
235                let port = u16::from_be_bytes([buf[20], buf[21]]);
236                msg.dst_addr = TargetAddr::Ip(SocketAddr::from((Ipv6Addr::from(octets), port)));
237                22
238            },
239            0x03 => {
240                if buf.len() < 5 {
241                    return Err(Error::InvalidTargetAddress("truncated domain UDP datagram"));
242                }
243                let len = buf[4] as usize;
244                let header_len = 5 + len + 2;
245                if buf.len() < header_len {
246                    return Err(Error::InvalidTargetAddress("truncated domain UDP datagram"));
247                }
248                let domain = String::from_utf8(buf[5..5 + len].to_vec())
249                    .map_err(|_| Error::InvalidTargetAddress("not a valid UTF-8 string"))?;
250                let port = u16::from_be_bytes([buf[5 + len], buf[5 + len + 1]]);
251                msg.dst_addr = TargetAddr::Domain(domain.into(), port);
252                header_len
253            },
254            _ => return Err(Error::UnknownAddressType),
255        };
256
257        msg.data = buf.split_off(header_len);
258        buf.clear();
259        Ok(Some(msg))
260    }
261}
262
263// +----+------+------+----------+----------+----------+
264// |RSV | FRAG | ATYP | DST.ADDR | DST.PORT |   DATA   |
265// +----+------+------+----------+----------+----------+
266// | 2  |  1   |  1   | Variable |    2     | Variable |
267// +----+------+------+----------+----------+----------+
268impl Encoder<(Bytes, TargetAddr<'static>)> for Socks5UdpCodec {
269    type Error = Error;
270
271    // TODO: consider fragment
272    fn encode(&mut self, (data, addr): (Bytes, TargetAddr<'static>), buf: &mut BytesMut) -> StdResult<(), Self::Error> {
273        let mut header = BytesMut::new();
274        header.resize(4, 0u8);
275
276        let mut addr_port = BytesMut::new();
277        match addr {
278            TargetAddr::Ip(SocketAddr::V4(addr)) => {
279                addr_port.reserve(6);
280                header[3] = 0x01;
281                addr_port.put_slice(&addr.ip().octets());
282                addr_port.put_slice(&addr.port().to_be_bytes());
283            },
284            TargetAddr::Ip(SocketAddr::V6(addr)) => {
285                addr_port.reserve(18);
286                header[3] = 0x04;
287                addr_port.put_slice(&addr.ip().octets());
288                addr_port.put_slice(&addr.port().to_be_bytes());
289            },
290            TargetAddr::Domain(domain, port) => {
291                let domain_len = domain.len();
292                if domain_len > 255 {
293                    return Err(Error::InvalidTargetAddress("overlong domain"));
294                }
295                addr_port.reserve(1 + domain_len + 2);
296                header[3] = 0x03;
297                addr_port.put_u8(domain_len as u8);
298                addr_port.put_slice(domain.as_bytes());
299                addr_port.put_slice(&port.to_be_bytes());
300            },
301        }
302        header.extend(addr_port);
303
304        buf.clear();
305        buf.extend(header);
306        buf.extend(data);
307
308        Ok(())
309    }
310}
311
312#[cfg(test)]
313mod tests {
314    use super::*;
315    use std::net::{Ipv4Addr, SocketAddrV6};
316
317    fn roundtrip(addr: TargetAddr<'static>, payload: &[u8]) -> Socks5UdpMessage {
318        let mut codec = Socks5UdpCodec::new();
319        let mut buf = BytesMut::new();
320        codec
321            .encode((Bytes::copy_from_slice(payload), addr), &mut buf)
322            .expect("encode");
323        codec.decode(&mut buf).expect("decode").expect("frame")
324    }
325
326    #[test]
327    fn roundtrip_ipv4() {
328        let addr = TargetAddr::Ip(SocketAddr::from((Ipv4Addr::new(1, 2, 3, 4), 9050)));
329        let msg = roundtrip(addr.clone(), b"hello");
330        assert_eq!(msg.atyp, 0x01);
331        assert_eq!(msg.frag, 0);
332        assert_eq!(msg.dst_addr, addr);
333        assert_eq!(&msg.data[..], b"hello");
334    }
335
336    #[test]
337    fn roundtrip_ipv6() {
338        let addr = TargetAddr::Ip(SocketAddr::V6(SocketAddrV6::new(
339            Ipv6Addr::LOCALHOST,
340            1080,
341            0,
342            0,
343        )));
344        let msg = roundtrip(addr.clone(), b"v6");
345        assert_eq!(msg.atyp, 0x04);
346        assert_eq!(msg.dst_addr, addr);
347        assert_eq!(&msg.data[..], b"v6");
348    }
349
350    #[test]
351    fn roundtrip_domain() {
352        let addr = TargetAddr::Domain("example.com".into(), 80);
353        let msg = roundtrip(addr.clone(), b"hi");
354        assert_eq!(msg.atyp, 0x03);
355        assert_eq!(msg.dst_addr, addr);
356        assert_eq!(&msg.data[..], b"hi");
357    }
358
359    #[test]
360    fn empty_payload_is_valid() {
361        let addr = TargetAddr::Ip(SocketAddr::from((Ipv4Addr::new(8, 8, 8, 8), 53)));
362        let msg = roundtrip(addr.clone(), b"");
363        assert_eq!(msg.dst_addr, addr);
364        assert!(msg.data.is_empty());
365    }
366
367    #[test]
368    fn short_domain_does_not_panic() {
369        // Regression: the old decoder sliced `buf[5..(len - 2)]`, which panics when
370        // `len < 2` and corrupts longer domain names.
371        let addr = TargetAddr::Domain("ab".into(), 443);
372        let msg = roundtrip(addr.clone(), b"x");
373        assert_eq!(msg.dst_addr, addr);
374    }
375
376    #[test]
377    fn truncated_datagram_is_error() {
378        let mut codec = Socks5UdpCodec::new();
379        let mut buf = BytesMut::from(&b"\x00\x00\x00"[..]);
380        assert!(codec.decode(&mut buf).is_err());
381
382        let mut buf = BytesMut::from(&b"\x00\x00\x00\x01\x01\x02\x03\x04\x00"[..]);
383        assert!(codec.decode(&mut buf).is_err());
384    }
385
386    #[test]
387    fn empty_buffer_is_incomplete() {
388        let mut codec = Socks5UdpCodec::new();
389        let mut buf = BytesMut::new();
390        assert!(codec.decode(&mut buf).unwrap().is_none());
391    }
392
393    #[test]
394    fn nonzero_frag_is_error() {
395        let mut codec = Socks5UdpCodec::new();
396        let mut buf = BytesMut::new();
397        buf.extend_from_slice(&[0, 0, 1, 1, 1, 2, 3, 4, 0, 53, b'x']);
398        assert!(matches!(codec.decode(&mut buf), Err(Error::FragmentationNotSupported)));
399    }
400
401    #[test]
402    fn overlong_domain_encode_fails() {
403        let mut codec = Socks5UdpCodec::new();
404        let mut buf = BytesMut::new();
405        let domain = "a".repeat(256);
406        let addr = TargetAddr::Domain(domain.into(), 80);
407        assert!(codec.encode((Bytes::new(), addr), &mut buf).is_err());
408    }
409}
better_socks/udp.rs

better_socks/
udp.rs
