eggress_protocol_socks/socks5/
udp_codec.rs

1use crate::socks5::server::{SocksAddr, ATYP_DOMAIN, ATYP_IPV4, ATYP_IPV6};
2
3pub const MAX_UDP_DATAGRAM_SIZE: usize = 65535;
4
5pub struct Socks5UdpRequest<'a> {
6    pub target: SocksAddr,
7    pub payload: &'a [u8],
8}
9
10#[derive(Debug, thiserror::Error)]
11pub enum UdpCodecError {
12    #[error("packet too short")]
13    TooShort,
14    #[error("non-zero reserved field")]
15    BadReserved,
16    #[error("fragmentation not supported")]
17    FragmentationUnsupported,
18    #[error("unknown address type: {0}")]
19    UnknownAddressType(u8),
20    #[error("zero domain length")]
21    BadDomainLength,
22    #[error("malformed domain name")]
23    MalformedDomain,
24    #[error("missing port")]
25    MissingPort,
26    #[error("packet too large: {0} > {1}")]
27    PacketTooLarge(usize, usize),
28}
29
30pub fn decode_socks5_udp_datagram(packet: &[u8]) -> Result<Socks5UdpRequest<'_>, UdpCodecError> {
31    if packet.len() < 4 {
32        return Err(UdpCodecError::TooShort);
33    }
34
35    // Defense-in-depth: reject oversized attacker-controlled datagrams up
36    // front instead of relying only on kernel receive limits.
37    if packet.len() > MAX_UDP_DATAGRAM_SIZE {
38        return Err(UdpCodecError::PacketTooLarge(
39            packet.len(),
40            MAX_UDP_DATAGRAM_SIZE,
41        ));
42    }
43
44    if packet[0] != 0x00 || packet[1] != 0x00 {
45        return Err(UdpCodecError::BadReserved);
46    }
47
48    if packet[2] != 0x00 {
49        return Err(UdpCodecError::FragmentationUnsupported);
50    }
51
52    let atyp = packet[3];
53    let mut offset = 4;
54
55    let target = match atyp {
56        ATYP_IPV4 => {
57            if packet.len() < offset + 4 + 2 {
58                return Err(UdpCodecError::TooShort);
59            }
60            let mut addr = [0u8; 4];
61            addr.copy_from_slice(&packet[offset..offset + 4]);
62            offset += 4;
63            let port = u16::from_be_bytes([packet[offset], packet[offset + 1]]);
64            offset += 2;
65            SocksAddr::IPv4(addr, port)
66        }
67        ATYP_IPV6 => {
68            if packet.len() < offset + 16 + 2 {
69                return Err(UdpCodecError::TooShort);
70            }
71            let mut addr = [0u8; 16];
72            addr.copy_from_slice(&packet[offset..offset + 16]);
73            offset += 16;
74            let port = u16::from_be_bytes([packet[offset], packet[offset + 1]]);
75            offset += 2;
76            SocksAddr::IPv6(addr, port)
77        }
78        ATYP_DOMAIN => {
79            if packet.len() < offset + 1 {
80                return Err(UdpCodecError::TooShort);
81            }
82            let domain_len = packet[offset] as usize;
83            offset += 1;
84            if domain_len == 0 {
85                return Err(UdpCodecError::BadDomainLength);
86            }
87            if packet.len() < offset + domain_len + 2 {
88                return Err(UdpCodecError::TooShort);
89            }
90            let domain = std::str::from_utf8(&packet[offset..offset + domain_len])
91                .map_err(|_| UdpCodecError::MalformedDomain)?;
92            offset += domain_len;
93            let port = u16::from_be_bytes([packet[offset], packet[offset + 1]]);
94            offset += 2;
95            SocksAddr::Domain(domain.to_string(), port)
96        }
97        _ => return Err(UdpCodecError::UnknownAddressType(atyp)),
98    };
99
100    let payload = &packet[offset..];
101    Ok(Socks5UdpRequest { target, payload })
102}
103
104pub fn encode_socks5_udp_datagram(
105    target: &SocksAddr,
106    payload: &[u8],
107    out: &mut Vec<u8>,
108) -> Result<(), UdpCodecError> {
109    out.clear();
110    out.extend_from_slice(&[0x00, 0x00]);
111    out.push(0x00);
112    let encoded = target
113        .encode_reply()
114        .map_err(|_| UdpCodecError::MalformedDomain)?;
115    out.extend_from_slice(&encoded);
116    out.extend_from_slice(payload);
117    Ok(())
118}
119
120pub fn decode_socks5_udp_request(packet: &[u8]) -> Result<Socks5UdpRequest<'_>, UdpCodecError> {
121    decode_socks5_udp_datagram(packet)
122}
123
124pub fn encode_socks5_udp_response(
125    target: &SocksAddr,
126    payload: &[u8],
127    out: &mut Vec<u8>,
128) -> Result<(), UdpCodecError> {
129    encode_socks5_udp_datagram(target, payload, out)
130}
131
132#[cfg(test)]
133mod tests {
134    use super::*;
135
136    fn ipv4_packet(target: [u8; 4], port: u16, payload: &[u8]) -> Vec<u8> {
137        let mut pkt = vec![0x00, 0x00, 0x00, ATYP_IPV4];
138        pkt.extend_from_slice(&target);
139        pkt.extend_from_slice(&port.to_be_bytes());
140        pkt.extend_from_slice(payload);
141        pkt
142    }
143
144    fn ipv6_packet(target: [u8; 16], port: u16, payload: &[u8]) -> Vec<u8> {
145        let mut pkt = vec![0x00, 0x00, 0x00, ATYP_IPV6];
146        pkt.extend_from_slice(&target);
147        pkt.extend_from_slice(&port.to_be_bytes());
148        pkt.extend_from_slice(payload);
149        pkt
150    }
151
152    fn domain_packet(domain: &str, port: u16, payload: &[u8]) -> Vec<u8> {
153        let mut pkt = vec![0x00, 0x00, 0x00, ATYP_DOMAIN, domain.len() as u8];
154        pkt.extend_from_slice(domain.as_bytes());
155        pkt.extend_from_slice(&port.to_be_bytes());
156        pkt.extend_from_slice(payload);
157        pkt
158    }
159
160    #[test]
161    fn decode_ipv4_target() {
162        let pkt = ipv4_packet([192, 168, 1, 1], 8080, b"hello");
163        let req = decode_socks5_udp_datagram(&pkt).unwrap();
164        assert_eq!(req.target, SocksAddr::IPv4([192, 168, 1, 1], 8080));
165        assert_eq!(req.payload, b"hello");
166    }
167
168    #[test]
169    fn decode_ipv6_target() {
170        let addr = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1];
171        let pkt = ipv6_packet(addr, 443, b"data");
172        let req = decode_socks5_udp_datagram(&pkt).unwrap();
173        assert_eq!(req.target, SocksAddr::IPv6(addr, 443));
174        assert_eq!(req.payload, b"data");
175    }
176
177    #[test]
178    fn decode_domain_target() {
179        let pkt = domain_packet("example.com", 53, b"\x01\x02");
180        let req = decode_socks5_udp_datagram(&pkt).unwrap();
181        assert_eq!(req.target, SocksAddr::Domain("example.com".to_string(), 53));
182        assert_eq!(req.payload, b"\x01\x02");
183    }
184
185    #[test]
186    fn encode_decode_roundtrip_ipv4() {
187        let target = SocksAddr::IPv4([10, 0, 0, 1], 9999);
188        let payload = b"roundtrip";
189        let mut buf = Vec::new();
190        encode_socks5_udp_datagram(&target, payload, &mut buf).unwrap();
191
192        let req = decode_socks5_udp_datagram(&buf).unwrap();
193        assert_eq!(req.target, target);
194        assert_eq!(req.payload, payload);
195    }
196
197    #[test]
198    fn encode_decode_roundtrip_ipv6() {
199        let addr = [0x20, 0x01, 0x0d, 0xb8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1];
200        let target = SocksAddr::IPv6(addr, 80);
201        let payload = b"test payload";
202        let mut buf = Vec::new();
203        encode_socks5_udp_datagram(&target, payload, &mut buf).unwrap();
204
205        let req = decode_socks5_udp_datagram(&buf).unwrap();
206        assert_eq!(req.target, target);
207        assert_eq!(req.payload, payload);
208    }
209
210    #[test]
211    fn encode_decode_roundtrip_domain() {
212        let target = SocksAddr::Domain("localhost".to_string(), 5353);
213        let payload = b"dns query";
214        let mut buf = Vec::new();
215        encode_socks5_udp_datagram(&target, payload, &mut buf).unwrap();
216
217        let req = decode_socks5_udp_datagram(&buf).unwrap();
218        assert_eq!(req.target, target);
219        assert_eq!(req.payload, payload);
220    }
221
222    #[test]
223    fn reject_bad_rsv() {
224        let mut pkt = ipv4_packet([1, 2, 3, 4], 80, b"");
225        pkt[0] = 0x01;
226        assert!(matches!(
227            decode_socks5_udp_datagram(&pkt),
228            Err(UdpCodecError::BadReserved)
229        ));
230    }
231
232    #[test]
233    fn reject_bad_rsv_second_byte() {
234        let mut pkt = ipv4_packet([1, 2, 3, 4], 80, b"");
235        pkt[1] = 0x01;
236        assert!(matches!(
237            decode_socks5_udp_datagram(&pkt),
238            Err(UdpCodecError::BadReserved)
239        ));
240    }
241
242    #[test]
243    fn reject_frag_nonzero() {
244        let mut pkt = ipv4_packet([1, 2, 3, 4], 80, b"");
245        pkt[2] = 0x01;
246        assert!(matches!(
247            decode_socks5_udp_datagram(&pkt),
248            Err(UdpCodecError::FragmentationUnsupported)
249        ));
250    }
251
252    #[test]
253    fn reject_short_packet_too_few_bytes() {
254        assert!(matches!(
255            decode_socks5_udp_datagram(&[0x00, 0x00]),
256            Err(UdpCodecError::TooShort)
257        ));
258    }
259
260    #[test]
261    fn reject_short_packet_ipv4() {
262        // ATYP_IPV4 but missing address bytes
263        let pkt = vec![0x00, 0x00, 0x00, ATYP_IPV4, 1, 2, 3];
264        assert!(matches!(
265            decode_socks5_udp_datagram(&pkt),
266            Err(UdpCodecError::TooShort)
267        ));
268    }
269
270    #[test]
271    fn reject_short_packet_ipv6() {
272        let pkt = vec![0x00, 0x00, 0x00, ATYP_IPV6, 0, 0, 0, 0];
273        assert!(matches!(
274            decode_socks5_udp_datagram(&pkt),
275            Err(UdpCodecError::TooShort)
276        ));
277    }
278
279    #[test]
280    fn reject_short_packet_domain_no_len() {
281        let pkt = vec![0x00, 0x00, 0x00, ATYP_DOMAIN];
282        assert!(matches!(
283            decode_socks5_udp_datagram(&pkt),
284            Err(UdpCodecError::TooShort)
285        ));
286    }
287
288    #[test]
289    fn reject_short_packet_domain_truncated() {
290        let pkt = vec![0x00, 0x00, 0x00, ATYP_DOMAIN, 5, b'a', b'b'];
291        assert!(matches!(
292            decode_socks5_udp_datagram(&pkt),
293            Err(UdpCodecError::TooShort)
294        ));
295    }
296
297    #[test]
298    fn reject_unknown_atyp() {
299        let pkt = vec![0x00, 0x00, 0x00, 0x05, 0x01, 0x02, 0x03, 0x04, 0x00, 0x50];
300        assert!(matches!(
301            decode_socks5_udp_datagram(&pkt),
302            Err(UdpCodecError::UnknownAddressType(0x05))
303        ));
304    }
305
306    #[test]
307    fn reject_zero_domain_length() {
308        let pkt = vec![0x00, 0x00, 0x00, ATYP_DOMAIN, 0x00];
309        assert!(matches!(
310            decode_socks5_udp_datagram(&pkt),
311            Err(UdpCodecError::BadDomainLength)
312        ));
313    }
314
315    #[test]
316    fn preserve_payload_bytes_exactly() {
317        let payload: Vec<u8> = (0..=255).collect();
318        let pkt = ipv4_packet([10, 0, 0, 1], 1234, &payload);
319        let req = decode_socks5_udp_datagram(&pkt).unwrap();
320        assert_eq!(req.payload.len(), 256);
321        for (i, &b) in req.payload.iter().enumerate() {
322            assert_eq!(b, i as u8);
323        }
324    }
325
326    #[test]
327    fn zero_length_payload() {
328        let pkt = ipv4_packet([10, 0, 0, 1], 80, b"");
329        let req = decode_socks5_udp_datagram(&pkt).unwrap();
330        assert_eq!(req.payload.len(), 0);
331    }
332
333    #[test]
334    fn reject_oversized_datagram() {
335        let payload = vec![0u8; MAX_UDP_DATAGRAM_SIZE + 1];
336        let pkt = ipv4_packet([10, 0, 0, 1], 80, &payload);
337        assert!(matches!(
338            decode_socks5_udp_datagram(&pkt),
339            Err(UdpCodecError::PacketTooLarge(_, MAX_UDP_DATAGRAM_SIZE))
340        ));
341    }
342
343    #[test]
344    fn encode_response_format() {
345        let target = SocksAddr::IPv4([192, 168, 1, 1], 80);
346        let payload = b"test";
347        let mut buf = Vec::new();
348        encode_socks5_udp_datagram(&target, payload, &mut buf).unwrap();
349
350        assert_eq!(buf[0], 0x00);
351        assert_eq!(buf[1], 0x00);
352        assert_eq!(buf[2], 0x00);
353        assert_eq!(buf[3], ATYP_IPV4);
354        assert_eq!(&buf[4..8], &[192, 168, 1, 1]);
355        assert_eq!(&buf[8..10], &80u16.to_be_bytes());
356        assert_eq!(&buf[10..], b"test");
357    }
358}
eggress_protocol_socks/socks5/udp_codec.rs

eggress_protocol_socks/socks5/
udp_codec.rs
