eggress_protocol_socks/socks5/
server.rs

1use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
2
3use crate::error::Socks5Error;
4
5/// SOCKS5 address types.
6pub const ATYP_IPV4: u8 = 0x01;
7pub const ATYP_DOMAIN: u8 = 0x03;
8pub const ATYP_IPV6: u8 = 0x04;
9
10/// SOCKS5 commands.
11pub const CMD_CONNECT: u8 = 0x01;
12pub const CMD_BIND: u8 = 0x02;
13pub const CMD_UDP_ASSOCIATE: u8 = 0x03;
14
15/// Parsed SOCKS5 command.
16#[derive(Debug, Clone, Copy, PartialEq, Eq)]
17pub enum Socks5Command {
18    Connect,
19    Bind,
20    UdpAssociate,
21}
22
23/// Parse a SOCKS5 command byte.
24pub fn parse_command(cmd: u8) -> Result<Socks5Command, Socks5Error> {
25    match cmd {
26        CMD_CONNECT => Ok(Socks5Command::Connect),
27        CMD_BIND => Ok(Socks5Command::Bind),
28        CMD_UDP_ASSOCIATE => Ok(Socks5Command::UdpAssociate),
29        _ => Err(Socks5Error::UnsupportedCommand(cmd)),
30    }
31}
32
33/// SOCKS5 reply codes.
34pub const REP_SUCCESS: u8 = 0x00;
35pub const REP_GENERAL_FAILURE: u8 = 0x01;
36pub const REP_NOT_ALLOWED: u8 = 0x02;
37pub const REP_COMMAND_NOT_SUPPORTED: u8 = 0x07;
38pub const REP_ADDRESS_TYPE_NOT_SUPPORTED: u8 = 0x08;
39
40/// SOCKS5 authentication methods.
41const AUTH_NONE: u8 = 0x00;
42const AUTH_USERNAME_PASSWORD: u8 = 0x02;
43const AUTH_NO_ACCEPTABLE: u8 = 0xFF;
44
45/// Username/password auth version.
46const AUTH_VERSION: u8 = 0x01;
47
48/// Maximum length for username or password in SOCKS5 auth.
49const MAX_CRED_LEN: usize = 255;
50
51/// A parsed SOCKS5 CONNECT request.
52#[derive(Debug, Clone, PartialEq, Eq, Hash)]
53pub enum SocksAddr {
54    IPv4([u8; 4], u16),
55    Domain(String, u16),
56    IPv6([u8; 16], u16),
57}
58
59impl SocksAddr {
60    /// Returns the host as a displayable string.
61    pub fn host_str(&self) -> String {
62        match self {
63            SocksAddr::IPv4(addr, _) => format!("{}.{}.{}.{}", addr[0], addr[1], addr[2], addr[3]),
64            SocksAddr::Domain(domain, _) => domain.clone(),
65            SocksAddr::IPv6(addr, _) => {
66                // Format as [ipv6]:port
67                let segments: Vec<String> = addr
68                    .chunks(2)
69                    .map(|chunk| format!("{:02x}{:02x}", chunk[0], chunk[1]))
70                    .collect();
71                format!("[{}]", segments.join(":"))
72            }
73        }
74    }
75
76    /// Returns the port.
77    pub fn port(&self) -> u16 {
78        match self {
79            SocksAddr::IPv4(_, port) | SocksAddr::Domain(_, port) | SocksAddr::IPv6(_, port) => {
80                *port
81            }
82        }
83    }
84
85    /// Encode this address into bytes for a SOCKS5 reply.
86    pub fn encode_reply(&self) -> Result<Vec<u8>, Socks5Error> {
87        let mut buf = Vec::new();
88        match self {
89            SocksAddr::IPv4(addr, port) => {
90                buf.push(ATYP_IPV4);
91                buf.extend_from_slice(addr);
92                buf.extend_from_slice(&port.to_be_bytes());
93            }
94            SocksAddr::Domain(domain, port) => {
95                let byte_len = domain.len();
96                if byte_len > 255 {
97                    return Err(Socks5Error::DomainTooLong(byte_len));
98                }
99                buf.push(ATYP_DOMAIN);
100                buf.push(byte_len as u8);
101                buf.extend_from_slice(domain.as_bytes());
102                buf.extend_from_slice(&port.to_be_bytes());
103            }
104            SocksAddr::IPv6(addr, port) => {
105                buf.push(ATYP_IPV6);
106                buf.extend_from_slice(addr);
107                buf.extend_from_slice(&port.to_be_bytes());
108            }
109        }
110        Ok(buf)
111    }
112}
113
114/// Parse the SOCKS5 method negotiation bytes synchronously.
115///
116/// Returns `(methods, remaining)` on success, or a [`Socks5Error`] if the
117/// buffer is truncated or the version byte is wrong. Exposed for fuzzing.
118pub fn parse_method_negotiation(buf: &[u8]) -> Result<(Vec<u8>, &[u8]), Socks5Error> {
119    if buf.is_empty() {
120        return Err(Socks5Error::UnexpectedEof);
121    }
122    let version = buf[0];
123    if version != 0x05 {
124        return Err(Socks5Error::UnsupportedVersion(version));
125    }
126    if buf.len() < 2 {
127        return Err(Socks5Error::UnexpectedEof);
128    }
129    let nmethods = buf[1] as usize;
130    if buf.len() < 2 + nmethods {
131        return Err(Socks5Error::UnexpectedEof);
132    }
133    Ok((buf[2..2 + nmethods].to_vec(), &buf[2 + nmethods..]))
134}
135
136/// Parse a SOCKS5 CONNECT request synchronously.
137///
138/// Returns `(target, remaining)` on success. Exposed for fuzzing.
139pub fn parse_connect_request(buf: &[u8]) -> Result<(SocksAddr, &[u8]), Socks5Error> {
140    if buf.len() < 4 {
141        return Err(Socks5Error::UnexpectedEof);
142    }
143    let version = buf[0];
144    if version != 0x05 {
145        return Err(Socks5Error::UnsupportedVersion(version));
146    }
147    let cmd = buf[1];
148    if cmd != CMD_CONNECT {
149        return Err(Socks5Error::UnsupportedCommand(cmd));
150    }
151    let rsv = buf[2];
152    if rsv != 0 {
153        return Err(Socks5Error::InvalidReservedByte(rsv));
154    }
155    let atyp = buf[3];
156
157    let (addr, rest) = match atyp {
158        ATYP_IPV4 => {
159            if buf.len() < 4 + 4 + 2 {
160                return Err(Socks5Error::UnexpectedEof);
161            }
162            let mut octets = [0u8; 4];
163            octets.copy_from_slice(&buf[4..8]);
164            let port = u16::from_be_bytes([buf[8], buf[9]]);
165            (SocksAddr::IPv4(octets, port), &buf[10..])
166        }
167        ATYP_DOMAIN => {
168            if buf.len() < 5 {
169                return Err(Socks5Error::UnexpectedEof);
170            }
171            let len = buf[4] as usize;
172            if buf.len() < 5 + len + 2 {
173                return Err(Socks5Error::UnexpectedEof);
174            }
175            let domain_bytes = &buf[5..5 + len];
176            let domain = String::from_utf8(domain_bytes.to_vec())
177                .map_err(|e| Socks5Error::MalformedMessage(format!("invalid domain: {e}")))?;
178            let port = u16::from_be_bytes([buf[5 + len], buf[5 + len + 1]]);
179            (SocksAddr::Domain(domain, port), &buf[5 + len + 2..])
180        }
181        ATYP_IPV6 => {
182            if buf.len() < 4 + 16 + 2 {
183                return Err(Socks5Error::UnexpectedEof);
184            }
185            let mut octets = [0u8; 16];
186            octets.copy_from_slice(&buf[4..20]);
187            let port = u16::from_be_bytes([buf[20], buf[21]]);
188            (SocksAddr::IPv6(octets, port), &buf[22..])
189        }
190        _ => return Err(Socks5Error::UnsupportedAddressType(atyp)),
191    };
192
193    Ok((addr, rest))
194}
195
196/// Parse a SOCKS5 generic request synchronously (CONNECT, BIND, or UDP_ASSOCIATE).
197///
198/// Returns `(command, target, remaining)`. Exposed for fuzzing.
199pub fn parse_socks5_request(buf: &[u8]) -> Result<(Socks5Command, SocksAddr, &[u8]), Socks5Error> {
200    if buf.len() < 4 {
201        return Err(Socks5Error::UnexpectedEof);
202    }
203    let version = buf[0];
204    if version != 0x05 {
205        return Err(Socks5Error::UnsupportedVersion(version));
206    }
207    let cmd = buf[1];
208    let command = parse_command(cmd)?;
209    let rsv = buf[2];
210    if rsv != 0 {
211        return Err(Socks5Error::InvalidReservedByte(rsv));
212    }
213    let atyp = buf[3];
214
215    let (addr, rest) = match atyp {
216        ATYP_IPV4 => {
217            if buf.len() < 4 + 4 + 2 {
218                return Err(Socks5Error::UnexpectedEof);
219            }
220            let mut octets = [0u8; 4];
221            octets.copy_from_slice(&buf[4..8]);
222            let port = u16::from_be_bytes([buf[8], buf[9]]);
223            (SocksAddr::IPv4(octets, port), &buf[10..])
224        }
225        ATYP_DOMAIN => {
226            if buf.len() < 5 {
227                return Err(Socks5Error::UnexpectedEof);
228            }
229            let len = buf[4] as usize;
230            if buf.len() < 5 + len + 2 {
231                return Err(Socks5Error::UnexpectedEof);
232            }
233            let domain_bytes = &buf[5..5 + len];
234            let domain = String::from_utf8(domain_bytes.to_vec())
235                .map_err(|e| Socks5Error::MalformedMessage(format!("invalid domain: {e}")))?;
236            let port = u16::from_be_bytes([buf[5 + len], buf[5 + len + 1]]);
237            (SocksAddr::Domain(domain, port), &buf[5 + len + 2..])
238        }
239        ATYP_IPV6 => {
240            if buf.len() < 4 + 16 + 2 {
241                return Err(Socks5Error::UnexpectedEof);
242            }
243            let mut octets = [0u8; 16];
244            octets.copy_from_slice(&buf[4..20]);
245            let port = u16::from_be_bytes([buf[20], buf[21]]);
246            (SocksAddr::IPv6(octets, port), &buf[22..])
247        }
248        _ => return Err(Socks5Error::UnsupportedAddressType(atyp)),
249    };
250
251    Ok((command, addr, rest))
252}
253
254/// Read a complete method negotiation message from the client.
255///
256/// Returns the list of methods the client supports.
257pub async fn read_method_negotiation<R: AsyncRead + Unpin>(
258    reader: &mut R,
259) -> Result<Vec<u8>, Socks5Error> {
260    let version = reader.read_u8().await?;
261    if version != 0x05 {
262        return Err(Socks5Error::UnsupportedVersion(version));
263    }
264
265    let nmethods = reader.read_u8().await?;
266    let mut methods = vec![0u8; nmethods as usize];
267    reader.read_exact(&mut methods).await?;
268
269    Ok(methods)
270}
271
272/// Send a method selection response to the client.
273///
274/// If the client supports `AUTH_NONE` and no password is required, selects no auth.
275/// If `password` is Some and the client supports username/password auth, selects that.
276/// Otherwise, sends 0xFF (no acceptable methods).
277pub async fn send_method_selection<W: AsyncWrite + Unpin>(
278    writer: &mut W,
279    methods: &[u8],
280    password: Option<&str>,
281) -> Result<(), Socks5Error> {
282    let method = if password.is_some() && methods.contains(&AUTH_USERNAME_PASSWORD) {
283        AUTH_USERNAME_PASSWORD
284    } else if methods.contains(&AUTH_NONE) {
285        AUTH_NONE
286    } else {
287        AUTH_NO_ACCEPTABLE
288    };
289
290    writer.write_all(&[0x05, method]).await?;
291    writer.flush().await?;
292
293    if method == AUTH_NO_ACCEPTABLE {
294        return Err(Socks5Error::MethodNegotiationFailed);
295    }
296
297    Ok(())
298}
299
300/// Read and validate username/password authentication from the client.
301///
302/// Returns Ok(()) on success, or an error if auth fails.
303pub async fn read_auth_request<R: AsyncRead + Unpin>(
304    reader: &mut R,
305    expected_password: &str,
306) -> Result<String, Socks5Error> {
307    let version = reader.read_u8().await?;
308    if version != AUTH_VERSION {
309        return Err(Socks5Error::UnsupportedVersion(version));
310    }
311
312    let ulen = reader.read_u8().await? as usize;
313    if ulen > MAX_CRED_LEN {
314        return Err(Socks5Error::CredentialsTooLong);
315    }
316    let mut username = vec![0u8; ulen];
317    reader.read_exact(&mut username).await?;
318
319    let plen = reader.read_u8().await? as usize;
320    if plen > MAX_CRED_LEN {
321        return Err(Socks5Error::CredentialsTooLong);
322    }
323    let mut password_bytes = vec![0u8; plen];
324    reader.read_exact(&mut password_bytes).await?;
325
326    let password_str = String::from_utf8_lossy(&password_bytes);
327    use subtle::ConstantTimeEq;
328    let passwords_match: bool = password_str
329        .as_bytes()
330        .ct_eq(expected_password.as_bytes())
331        .into();
332    if !passwords_match {
333        return Err(Socks5Error::AuthFailed);
334    }
335
336    Ok(String::from_utf8_lossy(&username).to_string())
337}
338
339/// Send an authentication response to the client.
340pub async fn send_auth_response<W: AsyncWrite + Unpin>(
341    writer: &mut W,
342    success: bool,
343) -> Result<(), Socks5Error> {
344    let status = if success { 0x00 } else { 0x01 };
345    writer.write_all(&[AUTH_VERSION, status]).await?;
346    writer.flush().await?;
347    Ok(())
348}
349
350/// Read a CONNECT request from the client.
351///
352/// Returns the target address.
353pub async fn read_connect_request<R: AsyncRead + Unpin>(
354    reader: &mut R,
355) -> Result<SocksAddr, Socks5Error> {
356    let version = reader.read_u8().await?;
357    if version != 0x05 {
358        return Err(Socks5Error::UnsupportedVersion(version));
359    }
360
361    let cmd = reader.read_u8().await?;
362    if cmd != CMD_CONNECT {
363        // Send reply for command not supported before returning error
364        return Err(Socks5Error::UnsupportedCommand(cmd));
365    }
366
367    let rsv = reader.read_u8().await?;
368    if rsv != 0 {
369        return Err(Socks5Error::InvalidReservedByte(rsv));
370    }
371
372    let atyp = reader.read_u8().await?;
373
374    let addr = match atyp {
375        ATYP_IPV4 => {
376            let mut buf = [0u8; 4];
377            reader.read_exact(&mut buf).await?;
378            let port = reader.read_u16().await?;
379            SocksAddr::IPv4(buf, port)
380        }
381        ATYP_DOMAIN => {
382            let len = reader.read_u8().await? as usize;
383            let mut domain = vec![0u8; len];
384            reader.read_exact(&mut domain).await?;
385            let domain = String::from_utf8(domain)
386                .map_err(|e| Socks5Error::MalformedMessage(format!("invalid domain: {e}")))?;
387            let port = reader.read_u16().await?;
388            SocksAddr::Domain(domain, port)
389        }
390        ATYP_IPV6 => {
391            let mut buf = [0u8; 16];
392            reader.read_exact(&mut buf).await?;
393            let port = reader.read_u16().await?;
394            SocksAddr::IPv6(buf, port)
395        }
396        _ => return Err(Socks5Error::UnsupportedAddressType(atyp)),
397    };
398
399    Ok(addr)
400}
401
402/// Read a SOCKS5 request from the client.
403///
404/// Returns the command and target address. Does not reject non-CONNECT commands.
405pub async fn read_socks5_request<R: AsyncRead + Unpin>(
406    reader: &mut R,
407) -> Result<(Socks5Command, SocksAddr), Socks5Error> {
408    let version = reader.read_u8().await?;
409    if version != 0x05 {
410        return Err(Socks5Error::UnsupportedVersion(version));
411    }
412
413    let cmd = reader.read_u8().await?;
414    let command = parse_command(cmd)?;
415
416    let rsv = reader.read_u8().await?;
417    if rsv != 0 {
418        return Err(Socks5Error::InvalidReservedByte(rsv));
419    }
420
421    let atyp = reader.read_u8().await?;
422
423    let addr = match atyp {
424        ATYP_IPV4 => {
425            let mut buf = [0u8; 4];
426            reader.read_exact(&mut buf).await?;
427            let port = reader.read_u16().await?;
428            SocksAddr::IPv4(buf, port)
429        }
430        ATYP_DOMAIN => {
431            let len = reader.read_u8().await? as usize;
432            let mut domain = vec![0u8; len];
433            reader.read_exact(&mut domain).await?;
434            let domain = String::from_utf8(domain)
435                .map_err(|e| Socks5Error::MalformedMessage(format!("invalid domain: {e}")))?;
436            let port = reader.read_u16().await?;
437            SocksAddr::Domain(domain, port)
438        }
439        ATYP_IPV6 => {
440            let mut buf = [0u8; 16];
441            reader.read_exact(&mut buf).await?;
442            let port = reader.read_u16().await?;
443            SocksAddr::IPv6(buf, port)
444        }
445        _ => return Err(Socks5Error::UnsupportedAddressType(atyp)),
446    };
447
448    Ok((command, addr))
449}
450
451/// Send a UDP ASSOCIATE reply to the client with the relay bind address.
452pub async fn send_udp_associate_reply<W: AsyncWrite + Unpin>(
453    writer: &mut W,
454    bind_addr: &SocksAddr,
455) -> Result<(), Socks5Error> {
456    send_connect_reply(writer, REP_SUCCESS, bind_addr).await
457}
458
459/// Send a CONNECT reply to the client.
460pub async fn send_connect_reply<W: AsyncWrite + Unpin>(
461    writer: &mut W,
462    rep: u8,
463    bind_addr: &SocksAddr,
464) -> Result<(), Socks5Error> {
465    let mut reply = vec![0x05, rep, 0x00]; // version, reply, reserved
466    reply.extend_from_slice(&bind_addr.encode_reply()?);
467    writer.write_all(&reply).await?;
468    writer.flush().await?;
469    Ok(())
470}
471
472/// Handle a complete SOCKS5 server handshake.
473///
474/// This reads the method negotiation, optionally handles username/password auth,
475/// and reads the CONNECT request. Returns the target address on success.
476///
477/// # Arguments
478/// * `reader` - The stream to read from.
479/// * `writer` - The stream to write to.
480/// * `password` - If Some, require username/password authentication with this password.
481pub async fn handle_socks5_handshake<R: AsyncRead + Unpin, W: AsyncWrite + Unpin>(
482    reader: &mut R,
483    writer: &mut W,
484    password: Option<&str>,
485) -> Result<SocksAddr, Socks5Error> {
486    // Step 1: Method negotiation
487    let methods = read_method_negotiation(reader).await?;
488    send_method_selection(writer, &methods, password).await?;
489
490    // Step 2: Auth (if password required)
491    if let Some(pwd) = password {
492        read_auth_request(reader, pwd).await?;
493        send_auth_response(writer, true).await?;
494    }
495
496    // Step 3: CONNECT request
497    let target = read_connect_request(reader).await?;
498
499    Ok(target)
500}
501
502/// Send a rejection reply for unsupported commands.
503pub async fn reject_command<W: AsyncWrite + Unpin>(
504    writer: &mut W,
505    cmd: u8,
506    target: &SocksAddr,
507) -> Result<(), Socks5Error> {
508    let _ = cmd;
509    send_connect_reply(writer, REP_COMMAND_NOT_SUPPORTED, target).await?;
510    Ok(())
511}
512
513/// Get the success reply code.
514pub const fn success_reply() -> u8 {
515    REP_SUCCESS
516}
517
518/// Get the general failure reply code.
519pub const fn general_failure_reply() -> u8 {
520    REP_GENERAL_FAILURE
521}
522
523/// Get the command not supported reply code.
524pub const fn command_not_supported_reply() -> u8 {
525    REP_COMMAND_NOT_SUPPORTED
526}
527
528/// Get the address type not supported reply code.
529pub const fn address_type_not_supported_reply() -> u8 {
530    REP_ADDRESS_TYPE_NOT_SUPPORTED
531}
532
533#[cfg(test)]
534mod tests {
535    use super::*;
536    use tokio::io::duplex;
537
538    #[tokio::test]
539    async fn test_method_negotiation_no_auth() {
540        let (mut client, mut server) = duplex(1024);
541
542        // Client sends: version=5, nmethods=1, method=0x00 (no auth)
543        client.write_all(&[0x05, 0x01, 0x00]).await.unwrap();
544
545        let methods = read_method_negotiation(&mut server).await.unwrap();
546        assert_eq!(methods, vec![0x00]);
547
548        send_method_selection(&mut server, &methods, None)
549            .await
550            .unwrap();
551
552        let mut response = [0u8; 2];
553        client.read_exact(&mut response).await.unwrap();
554        assert_eq!(response, [0x05, 0x00]); // Selected no auth
555    }
556
557    #[tokio::test]
558    async fn test_method_negotiation_username_password() {
559        let (mut client, mut server) = duplex(1024);
560
561        // Client offers: no auth and username/password
562        client.write_all(&[0x05, 0x02, 0x00, 0x02]).await.unwrap();
563
564        let methods = read_method_negotiation(&mut server).await.unwrap();
565        assert_eq!(methods, vec![0x00, 0x02]);
566
567        // Server requires password
568        send_method_selection(&mut server, &methods, Some("secret"))
569            .await
570            .unwrap();
571
572        let mut response = [0u8; 2];
573        client.read_exact(&mut response).await.unwrap();
574        assert_eq!(response, [0x05, 0x02]); // Selected username/password
575    }
576
577    #[tokio::test]
578    async fn test_method_negotiation_no_acceptable() {
579        let (mut client, mut server) = duplex(1024);
580
581        // Client offers only GSSAPI (0x01) which we don't support
582        client.write_all(&[0x05, 0x01, 0x01]).await.unwrap();
583
584        let methods = read_method_negotiation(&mut server).await.unwrap();
585        let result = send_method_selection(&mut server, &methods, None).await;
586        assert!(matches!(result, Err(Socks5Error::MethodNegotiationFailed)));
587    }
588
589    #[tokio::test]
590    async fn test_auth_success() {
591        let (mut client, mut server) = duplex(1024);
592
593        // Auth request: version=1, ulen=4, username="user", plen=6, password="secret"
594        client
595            .write_all(&[0x01, 0x04, b'u', b's', b'e', b'r', 0x06])
596            .await
597            .unwrap();
598        client.write_all(b"secret").await.unwrap();
599
600        let username = read_auth_request(&mut server, "secret").await.unwrap();
601        assert_eq!(username, "user");
602
603        send_auth_response(&mut server, true).await.unwrap();
604
605        let mut response = [0u8; 2];
606        client.read_exact(&mut response).await.unwrap();
607        assert_eq!(response, [0x01, 0x00]); // Success
608    }
609
610    #[tokio::test]
611    async fn test_auth_failure() {
612        let (mut client, mut server) = duplex(1024);
613
614        // Auth request with wrong password
615        client
616            .write_all(&[0x01, 0x04, b'u', b's', b'e', b'r', 0x05])
617            .await
618            .unwrap();
619        client.write_all(b"wrong").await.unwrap();
620
621        let result = read_auth_request(&mut server, "secret").await;
622        assert!(matches!(result, Err(Socks5Error::AuthFailed)));
623    }
624
625    #[tokio::test]
626    async fn test_connect_ipv4() {
627        let (mut client, mut server) = duplex(1024);
628
629        // CONNECT request: version=5, cmd=1, rsv=0, atyp=1 (IPv4), addr=192.168.1.1, port=8080
630        client
631            .write_all(&[0x05, 0x01, 0x00, 0x01, 192, 168, 1, 1])
632            .await
633            .unwrap();
634        client.write_all(&8080u16.to_be_bytes()).await.unwrap();
635
636        let target = read_connect_request(&mut server).await.unwrap();
637        assert_eq!(target, SocksAddr::IPv4([192, 168, 1, 1], 8080));
638    }
639
640    #[tokio::test]
641    async fn test_connect_domain() {
642        let (mut client, mut server) = duplex(1024);
643
644        let domain = "example.com";
645        // version=5, cmd=1, rsv=0, atyp=3, len, domain, port
646        client
647            .write_all(&[0x05, 0x01, 0x00, 0x03, domain.len() as u8])
648            .await
649            .unwrap();
650        client.write_all(domain.as_bytes()).await.unwrap();
651        client.write_all(&443u16.to_be_bytes()).await.unwrap();
652
653        let target = read_connect_request(&mut server).await.unwrap();
654        assert_eq!(target, SocksAddr::Domain("example.com".to_string(), 443));
655    }
656
657    #[tokio::test]
658    async fn test_connect_ipv6() {
659        let (mut client, mut server) = duplex(1024);
660
661        // CONNECT request: version=5, cmd=1, rsv=0, atyp=4 (IPv6), addr=::1, port=443
662        let ipv6_addr = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1];
663        client.write_all(&[0x05, 0x01, 0x00, 0x04]).await.unwrap();
664        client.write_all(&ipv6_addr).await.unwrap();
665        client.write_all(&443u16.to_be_bytes()).await.unwrap();
666
667        let target = read_connect_request(&mut server).await.unwrap();
668        assert_eq!(target, SocksAddr::IPv6(ipv6_addr, 443));
669    }
670
671    #[tokio::test]
672    async fn test_connect_reply_success() {
673        let (mut client, mut server) = duplex(1024);
674
675        let bind_addr = SocksAddr::IPv4([0, 0, 0, 0], 0);
676        send_connect_reply(&mut server, REP_SUCCESS, &bind_addr)
677            .await
678            .unwrap();
679
680        let mut response = [0u8; 10];
681        client.read_exact(&mut response).await.unwrap();
682        assert_eq!(response[0], 0x05); // version
683        assert_eq!(response[1], 0x00); // success
684        assert_eq!(response[2], 0x00); // reserved
685        assert_eq!(response[3], 0x01); // atyp IPv4
686    }
687
688    #[tokio::test]
689    async fn test_unsupported_version() {
690        let (mut client, mut server) = duplex(1024);
691
692        // Send version 4 instead of 5
693        client.write_all(&[0x04, 0x01, 0x00]).await.unwrap();
694
695        let result = read_method_negotiation(&mut server).await;
696        assert!(matches!(result, Err(Socks5Error::UnsupportedVersion(0x04))));
697    }
698
699    #[tokio::test]
700    async fn test_unsupported_command() {
701        let (mut client, mut server) = duplex(1024);
702
703        // BIND command (0x02)
704        let ipv4_addr = [192, 168, 1, 1];
705        client.write_all(&[0x05, 0x02, 0x00, 0x01]).await.unwrap();
706        client.write_all(&ipv4_addr).await.unwrap();
707        client.write_all(&80u16.to_be_bytes()).await.unwrap();
708
709        let result = read_connect_request(&mut server).await;
710        assert!(matches!(result, Err(Socks5Error::UnsupportedCommand(0x02))));
711    }
712
713    #[tokio::test]
714    async fn test_unsupported_address_type() {
715        let (mut client, mut server) = duplex(1024);
716
717        // atyp=0x05 (unsupported)
718        client.write_all(&[0x05, 0x01, 0x00, 0x05]).await.unwrap();
719
720        let result = read_connect_request(&mut server).await;
721        assert!(matches!(
722            result,
723            Err(Socks5Error::UnsupportedAddressType(0x05))
724        ));
725    }
726
727    #[tokio::test]
728    async fn test_reject_bind_command() {
729        let (mut client, mut server) = duplex(1024);
730
731        let target = SocksAddr::IPv4([192, 168, 1, 1], 80);
732        reject_command(&mut server, 0x02, &target).await.unwrap();
733
734        let mut response = [0u8; 10];
735        client.read_exact(&mut response).await.unwrap();
736        assert_eq!(response[0], 0x05); // version
737        assert_eq!(response[1], 0x07); // command not supported (RFC 1928 §6)
738    }
739
740    #[tokio::test]
741    async fn test_reject_unknown_command() {
742        let (mut client, mut server) = duplex(1024);
743
744        let target = SocksAddr::IPv4([192, 168, 1, 1], 80);
745        reject_command(&mut server, 0x04, &target).await.unwrap();
746
747        let mut response = [0u8; 10];
748        client.read_exact(&mut response).await.unwrap();
749        assert_eq!(response[0], 0x05); // version
750        assert_eq!(response[1], 0x07); // command not supported
751    }
752
753    #[tokio::test]
754    async fn test_creds_too_long() {
755        // ulen is u8 so max 255; CredentialsTooLong can only happen
756        // if we manually construct bad data. Verify boundary with 255-byte credential.
757    }
758
759    #[tokio::test]
760    async fn test_boundary_credentials_length() {
761        let (mut client, mut server) = duplex(2048);
762
763        let username = "a".repeat(255);
764        let password = "b".repeat(255);
765
766        // Auth request: version=1, ulen=255, username, plen=255, password
767        client.write_all(&[0x01, 255]).await.unwrap();
768        client.write_all(username.as_bytes()).await.unwrap();
769        client.write_all(&[255]).await.unwrap();
770        client.write_all(password.as_bytes()).await.unwrap();
771
772        let result = read_auth_request(&mut server, &password).await;
773        assert!(result.is_ok());
774        assert_eq!(result.unwrap(), username);
775    }
776
777    #[tokio::test]
778    async fn test_full_handshake_no_auth() {
779        let (mut client, mut server) = duplex(1024);
780
781        // Method negotiation: client offers no auth
782        client.write_all(&[0x05, 0x01, 0x00]).await.unwrap();
783
784        // Server reads and selects method
785        let methods = read_method_negotiation(&mut server).await.unwrap();
786        send_method_selection(&mut server, &methods, None)
787            .await
788            .unwrap();
789
790        // Read method selection response
791        let mut response = [0u8; 2];
792        client.read_exact(&mut response).await.unwrap();
793        assert_eq!(response, [0x05, 0x00]);
794
795        // CONNECT request
796        client
797            .write_all(&[0x05, 0x01, 0x00, 0x01, 10, 0, 0, 1])
798            .await
799            .unwrap();
800        client.write_all(&443u16.to_be_bytes()).await.unwrap();
801
802        let target = read_connect_request(&mut server).await.unwrap();
803        assert_eq!(target, SocksAddr::IPv4([10, 0, 0, 1], 443));
804
805        // Send success reply
806        let bind_addr = SocksAddr::IPv4([0, 0, 0, 0], 0);
807        send_connect_reply(&mut server, REP_SUCCESS, &bind_addr)
808            .await
809            .unwrap();
810
811        let mut reply = [0u8; 10];
812        client.read_exact(&mut reply).await.unwrap();
813        assert_eq!(reply[0], 0x05);
814        assert_eq!(reply[1], 0x00);
815    }
816
817    #[tokio::test]
818    async fn test_full_handshake_with_auth() {
819        let (mut client, mut server) = duplex(2048);
820
821        // Method negotiation: client offers both methods
822        client.write_all(&[0x05, 0x02, 0x00, 0x02]).await.unwrap();
823
824        // Server requires password
825        let methods = read_method_negotiation(&mut server).await.unwrap();
826        send_method_selection(&mut server, &methods, Some("mypass"))
827            .await
828            .unwrap();
829
830        // Read method selection
831        let mut response = [0u8; 2];
832        client.read_exact(&mut response).await.unwrap();
833        assert_eq!(response, [0x05, 0x02]);
834
835        // Auth request
836        client
837            .write_all(&[0x01, 0x04, b'u', b's', b'e', b'r'])
838            .await
839            .unwrap();
840        client
841            .write_all(&[0x06, b'm', b'y', b'p', b'a', b's', b's'])
842            .await
843            .unwrap();
844
845        let username = read_auth_request(&mut server, "mypass").await.unwrap();
846        assert_eq!(username, "user");
847        send_auth_response(&mut server, true).await.unwrap();
848
849        // Read auth response
850        let mut auth_response = [0u8; 2];
851        client.read_exact(&mut auth_response).await.unwrap();
852        assert_eq!(auth_response, [0x01, 0x00]);
853
854        // CONNECT request
855        let domain = "example.com";
856        client
857            .write_all(&[0x05, 0x01, 0x00, 0x03, domain.len() as u8])
858            .await
859            .unwrap();
860        client.write_all(domain.as_bytes()).await.unwrap();
861        client.write_all(&443u16.to_be_bytes()).await.unwrap();
862
863        let target = read_connect_request(&mut server).await.unwrap();
864        assert_eq!(target, SocksAddr::Domain("example.com".to_string(), 443));
865    }
866
867    #[tokio::test]
868    async fn test_fragged_handshake() {
869        let (mut client, mut server) = duplex(1024);
870
871        // Send method negotiation in fragments
872        client.write_all(&[0x05]).await.unwrap();
873        // Small delay to simulate fragmentation
874        tokio::time::sleep(std::time::Duration::from_millis(10)).await;
875        client.write_all(&[0x01]).await.unwrap();
876        tokio::time::sleep(std::time::Duration::from_millis(10)).await;
877        client.write_all(&[0x00]).await.unwrap();
878
879        let methods = read_method_negotiation(&mut server).await.unwrap();
880        assert_eq!(methods, vec![0x00]);
881    }
882
883    #[tokio::test]
884    async fn test_socks_addr_display() {
885        let ipv4 = SocksAddr::IPv4([192, 168, 1, 1], 8080);
886        assert_eq!(ipv4.host_str(), "192.168.1.1");
887        assert_eq!(ipv4.port(), 8080);
888
889        let domain = SocksAddr::Domain("example.com".to_string(), 443);
890        assert_eq!(domain.host_str(), "example.com");
891        assert_eq!(domain.port(), 443);
892
893        let ipv6 = SocksAddr::IPv6([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1], 443);
894        assert_eq!(ipv6.port(), 443);
895    }
896
897    #[tokio::test]
898    async fn test_socks_addr_encode_reply() {
899        let ipv4 = SocksAddr::IPv4([192, 168, 1, 1], 8080);
900        let encoded = ipv4.encode_reply().unwrap();
901        assert_eq!(encoded[0], ATYP_IPV4);
902        assert_eq!(&encoded[1..5], &[192, 168, 1, 1]);
903        assert_eq!(&encoded[5..7], &8080u16.to_be_bytes());
904
905        let domain = SocksAddr::Domain("example.com".to_string(), 443);
906        let encoded = domain.encode_reply().unwrap();
907        assert_eq!(encoded[0], ATYP_DOMAIN);
908        assert_eq!(encoded[1], 11); // "example.com" length
909        assert_eq!(&encoded[2..13], b"example.com");
910        assert_eq!(&encoded[13..15], &443u16.to_be_bytes());
911    }
912
913    #[test]
914    fn test_parse_command() {
915        assert_eq!(parse_command(0x01).unwrap(), Socks5Command::Connect);
916        assert_eq!(parse_command(0x02).unwrap(), Socks5Command::Bind);
917        assert_eq!(parse_command(0x03).unwrap(), Socks5Command::UdpAssociate);
918        assert!(parse_command(0x04).is_err());
919        assert!(parse_command(0xFF).is_err());
920    }
921
922    #[tokio::test]
923    async fn test_udp_associate_ipv4() {
924        let (mut client, mut server) = duplex(1024);
925
926        // UDP ASSOCIATE request: version=5, cmd=3, rsv=0, atyp=1 (IPv4), addr=0.0.0.0, port=0
927        client
928            .write_all(&[0x05, 0x03, 0x00, 0x01, 0, 0, 0, 0])
929            .await
930            .unwrap();
931        client.write_all(&0u16.to_be_bytes()).await.unwrap();
932
933        let (cmd, target) = read_socks5_request(&mut server).await.unwrap();
934        assert_eq!(cmd, Socks5Command::UdpAssociate);
935        assert_eq!(target, SocksAddr::IPv4([0, 0, 0, 0], 0));
936    }
937
938    #[tokio::test]
939    async fn test_udp_associate_ipv6() {
940        let (mut client, mut server) = duplex(1024);
941
942        let ipv6_addr = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1];
943        client.write_all(&[0x05, 0x03, 0x00, 0x04]).await.unwrap();
944        client.write_all(&ipv6_addr).await.unwrap();
945        client.write_all(&0u16.to_be_bytes()).await.unwrap();
946
947        let (cmd, target) = read_socks5_request(&mut server).await.unwrap();
948        assert_eq!(cmd, Socks5Command::UdpAssociate);
949        assert_eq!(target, SocksAddr::IPv6(ipv6_addr, 0));
950    }
951
952    #[tokio::test]
953    async fn test_udp_associate_domain() {
954        let (mut client, mut server) = duplex(1024);
955
956        let domain = "example.com";
957        client
958            .write_all(&[0x05, 0x03, 0x00, 0x03, domain.len() as u8])
959            .await
960            .unwrap();
961        client.write_all(domain.as_bytes()).await.unwrap();
962        client.write_all(&0u16.to_be_bytes()).await.unwrap();
963
964        let (cmd, target) = read_socks5_request(&mut server).await.unwrap();
965        assert_eq!(cmd, Socks5Command::UdpAssociate);
966        assert_eq!(target, SocksAddr::Domain("example.com".to_string(), 0));
967    }
968
969    #[tokio::test]
970    async fn test_connect_still_works_via_read_socks5_request() {
971        let (mut client, mut server) = duplex(1024);
972
973        client
974            .write_all(&[0x05, 0x01, 0x00, 0x01, 192, 168, 1, 1])
975            .await
976            .unwrap();
977        client.write_all(&8080u16.to_be_bytes()).await.unwrap();
978
979        let (cmd, target) = read_socks5_request(&mut server).await.unwrap();
980        assert_eq!(cmd, Socks5Command::Connect);
981        assert_eq!(target, SocksAddr::IPv4([192, 168, 1, 1], 8080));
982    }
983
984    #[tokio::test]
985    async fn test_reject_bind_only_not_udp_associate() {
986        let (mut client, mut server) = duplex(1024);
987
988        let target = SocksAddr::IPv4([192, 168, 1, 1], 80);
989        reject_command(&mut server, 0x02, &target).await.unwrap();
990
991        let mut response = [0u8; 10];
992        client.read_exact(&mut response).await.unwrap();
993        assert_eq!(response[0], 0x05);
994        assert_eq!(response[1], 0x07); // command not supported (RFC 1928 §6)
995    }
996
997    #[tokio::test]
998    async fn test_send_udp_associate_reply() {
999        let (mut client, mut server) = duplex(1024);
1000
1001        let bind_addr = SocksAddr::IPv4([127, 0, 0, 1], 1080);
1002        send_udp_associate_reply(&mut server, &bind_addr)
1003            .await
1004            .unwrap();
1005
1006        let mut response = [0u8; 10];
1007        client.read_exact(&mut response).await.unwrap();
1008        assert_eq!(response[0], 0x05);
1009        assert_eq!(response[1], 0x00); // success
1010        assert_eq!(response[3], 0x01); // atyp IPv4
1011        assert_eq!(&response[4..8], &[127, 0, 0, 1]);
1012        assert_eq!(&response[8..10], &1080u16.to_be_bytes());
1013    }
1014
1015    #[test]
1016    fn parse_socks5_request_rsv_zero_accepted() {
1017        let buf = [0x05, 0x01, 0x00, 0x01, 127, 0, 0, 1, 0x00, 0x50];
1018        let result = parse_socks5_request(&buf);
1019        assert!(result.is_ok());
1020    }
1021
1022    #[test]
1023    fn parse_socks5_request_rsv_nonzero_rejected() {
1024        let buf = [0x05, 0x01, 0x01, 0x01, 127, 0, 0, 1, 0x00, 0x50];
1025        let result = parse_socks5_request(&buf);
1026        assert!(matches!(
1027            result,
1028            Err(Socks5Error::InvalidReservedByte(0x01))
1029        ));
1030    }
1031
1032    #[test]
1033    fn parse_connect_request_rsv_nonzero_rejected() {
1034        let buf = [0x05, 0x01, 0x42, 0x01, 127, 0, 0, 1, 0x00, 0x50];
1035        let result = parse_connect_request(&buf);
1036        assert!(matches!(
1037            result,
1038            Err(Socks5Error::InvalidReservedByte(0x42))
1039        ));
1040    }
1041
1042    #[tokio::test]
1043    async fn read_socks5_request_rsv_nonzero_rejected() {
1044        let (mut client, mut server) = duplex(1024);
1045
1046        // version=5, cmd=1 (CONNECT), rsv=0x01 (non-zero), atyp=1 (IPv4)
1047        client
1048            .write_all(&[0x05, 0x01, 0x01, 0x01, 127, 0, 0, 1])
1049            .await
1050            .unwrap();
1051        client.write_all(&80u16.to_be_bytes()).await.unwrap();
1052
1053        let result = read_socks5_request(&mut server).await;
1054        assert!(matches!(
1055            result,
1056            Err(Socks5Error::InvalidReservedByte(0x01))
1057        ));
1058    }
1059
1060    #[tokio::test]
1061    async fn read_connect_request_rsv_nonzero_rejected() {
1062        let (mut client, mut server) = duplex(1024);
1063
1064        // version=5, cmd=1 (CONNECT), rsv=0xFF, atyp=1 (IPv4)
1065        client
1066            .write_all(&[0x05, 0x01, 0xFF, 0x01, 127, 0, 0, 1])
1067            .await
1068            .unwrap();
1069        client.write_all(&80u16.to_be_bytes()).await.unwrap();
1070
1071        let result = read_connect_request(&mut server).await;
1072        assert!(matches!(
1073            result,
1074            Err(Socks5Error::InvalidReservedByte(0xFF))
1075        ));
1076    }
1077
1078    #[test]
1079    fn encode_reply_domain_255_bytes_ok() {
1080        let domain = "a".repeat(255);
1081        let addr = SocksAddr::Domain(domain, 80);
1082        let encoded = addr.encode_reply().unwrap();
1083        assert_eq!(encoded[0], ATYP_DOMAIN);
1084        assert_eq!(encoded[1], 255);
1085        assert_eq!(&encoded[2..257], "a".repeat(255).as_bytes());
1086        assert_eq!(&encoded[257..259], &80u16.to_be_bytes());
1087    }
1088
1089    #[test]
1090    fn encode_reply_domain_256_bytes_error() {
1091        let domain = "a".repeat(256);
1092        let addr = SocksAddr::Domain(domain, 80);
1093        let result = addr.encode_reply();
1094        assert!(matches!(result, Err(Socks5Error::DomainTooLong(256))));
1095    }
1096
1097    #[test]
1098    fn encode_reply_multibyte_utf8_domain_exceeding_255_bytes() {
1099        // Each 'é' is 2 bytes in UTF-8. 128 * 2 = 256 bytes > 255
1100        let domain = "é".repeat(128);
1101        let addr = SocksAddr::Domain(domain, 80);
1102        let result = addr.encode_reply();
1103        assert!(matches!(result, Err(Socks5Error::DomainTooLong(256))));
1104    }
1105
1106    #[test]
1107    fn encode_reply_ipv4_unchanged() {
1108        let addr = SocksAddr::IPv4([192, 168, 1, 1], 8080);
1109        let encoded = addr.encode_reply().unwrap();
1110        assert_eq!(encoded, vec![ATYP_IPV4, 192, 168, 1, 1, 0x1F, 0x90]);
1111    }
1112
1113    #[test]
1114    fn encode_reply_ipv6_unchanged() {
1115        let addr = SocksAddr::IPv6([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1], 443);
1116        let encoded = addr.encode_reply().unwrap();
1117        assert_eq!(encoded[0], ATYP_IPV6);
1118        assert_eq!(
1119            &encoded[1..17],
1120            &[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1]
1121        );
1122        assert_eq!(&encoded[17..19], &443u16.to_be_bytes());
1123    }
1124
1125    #[test]
1126    fn encode_reply_domain_unchanged_for_valid() {
1127        let addr = SocksAddr::Domain("example.com".to_string(), 443);
1128        let encoded = addr.encode_reply().unwrap();
1129        assert_eq!(encoded[0], ATYP_DOMAIN);
1130        assert_eq!(encoded[1], 11);
1131        assert_eq!(&encoded[2..13], b"example.com");
1132        assert_eq!(&encoded[13..15], &443u16.to_be_bytes());
1133    }
1134
1135    #[test]
1136    fn udp_rsv_first_byte_nonzero_rejected() {
1137        use super::super::udp_codec::{decode_socks5_udp_datagram, UdpCodecError};
1138        let pkt = vec![0x01, 0x00, 0x00, ATYP_IPV4, 1, 2, 3, 4, 0x00, 0x50];
1139        let result = decode_socks5_udp_datagram(&pkt);
1140        assert!(matches!(result, Err(UdpCodecError::BadReserved)));
1141    }
1142
1143    #[test]
1144    fn udp_rsv_second_byte_nonzero_rejected() {
1145        use super::super::udp_codec::{decode_socks5_udp_datagram, UdpCodecError};
1146        let pkt = vec![0x00, 0x01, 0x00, ATYP_IPV4, 1, 2, 3, 4, 0x00, 0x50];
1147        let result = decode_socks5_udp_datagram(&pkt);
1148        assert!(matches!(result, Err(UdpCodecError::BadReserved)));
1149    }
1150}
eggress_protocol_socks/socks5/server.rs

eggress_protocol_socks/socks5/
server.rs
