eggress_protocol_socks/socks4/
server.rs

1use std::net::{IpAddr, Ipv4Addr, SocketAddr};
2
3use tokio::io::{AsyncReadExt, AsyncWriteExt};
4
5use super::error::Socks4Error;
6use super::{MAX_DOMAIN_LEN, MAX_USER_ID_LEN};
7
8/// SOCKS4 reply status codes.
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10#[repr(u8)]
11pub enum Socks4Status {
12    Granted = 90,
13    Failed = 91,
14    FailedNoIdent = 92,
15    FailedDifferentUser = 93,
16}
17
18impl Socks4Status {
19    /// Convert from raw status byte.
20    pub fn from_u8(val: u8) -> Option<Self> {
21        match val {
22            90 => Some(Self::Granted),
23            91 => Some(Self::Failed),
24            92 => Some(Self::FailedNoIdent),
25            93 => Some(Self::FailedDifferentUser),
26            _ => None,
27        }
28    }
29}
30
31/// Parsed SOCKS4/4a CONNECT request.
32#[derive(Debug, Clone)]
33pub struct Socks4Request {
34    pub command: u8,
35    pub port: u16,
36    pub addr: SocketAddr,
37    pub user_id: String,
38    /// For SOCKS4a: the domain name when IP is 0.0.0.x (x != 0).
39    pub domain: Option<String>,
40}
41
42/// Read a SOCKS4/4a request from the stream.
43///
44/// Format:
45///   +----+----+----+----+----+----+----+----+----+----+....+----+
46///   | VN | CD | DSTPORT |      DSTIP        | USERID       |0x00|
47///   +----+----+----+----+----+----+----+----+----+----+....+----+
48///     1    1      2            4              variable       1
49///
50/// For SOCKS4a, when DSTIP is 0.0.0.x (x != 0), the domain follows
51/// after the NUL-terminated USERID.
52pub async fn read_socks4_request<S: tokio::io::AsyncRead + Unpin>(
53    stream: &mut S,
54) -> Result<Socks4Request, Socks4Error> {
55    let mut header = [0u8; 8];
56    stream.read_exact(&mut header).await?;
57
58    let version = header[0];
59    if version != 0x04 {
60        return Err(Socks4Error::InvalidVersion(version));
61    }
62
63    let command = header[1];
64    if command != 0x01 {
65        return Err(Socks4Error::UnsupportedCommand(command));
66    }
67
68    let port = u16::from_be_bytes([header[2], header[3]]);
69    let ip = Ipv4Addr::new(header[4], header[5], header[6], header[7]);
70
71    // Read NUL-terminated user ID (bounded at MAX_USER_ID_LEN + 1 for NUL).
72    let mut user_id_bytes = Vec::with_capacity(64);
73    let mut buf = [0u8; 1];
74    loop {
75        if user_id_bytes.len() > MAX_USER_ID_LEN {
76            return Err(Socks4Error::UserIdTooLong);
77        }
78        let n = stream.read(&mut buf).await?;
79        if n == 0 {
80            return Err(Socks4Error::MalformedRequest(
81                "unexpected EOF reading user ID".into(),
82            ));
83        }
84        if buf[0] == 0x00 {
85            break;
86        }
87        user_id_bytes.push(buf[0]);
88    }
89
90    let user_id = String::from_utf8(user_id_bytes)
91        .map_err(|_| Socks4Error::MalformedRequest("invalid UTF-8 in user ID".into()))?;
92
93    // SOCKS4a: if IP is 0.0.0.x (x != 0), read domain after user ID.
94    let domain =
95        if ip.octets()[0] == 0 && ip.octets()[1] == 0 && ip.octets()[2] == 0 && ip.octets()[3] != 0
96        {
97            let mut domain_bytes = Vec::with_capacity(256);
98            loop {
99                if domain_bytes.len() > MAX_DOMAIN_LEN {
100                    return Err(Socks4Error::DomainTooLong);
101                }
102                let n = stream.read(&mut buf).await?;
103                if n == 0 {
104                    return Err(Socks4Error::MalformedRequest(
105                        "unexpected EOF reading domain".into(),
106                    ));
107                }
108                if buf[0] == 0x00 {
109                    break;
110                }
111                domain_bytes.push(buf[0]);
112            }
113            let domain = String::from_utf8(domain_bytes)
114                .map_err(|_| Socks4Error::MalformedRequest("invalid UTF-8 in domain".into()))?;
115            if domain.is_empty() {
116                return Err(Socks4Error::MalformedRequest(
117                    "empty domain in SOCKS4a request".into(),
118                ));
119            }
120            Some(domain)
121        } else {
122            None
123        };
124
125    let addr = SocketAddr::new(IpAddr::V4(ip), port);
126
127    Ok(Socks4Request {
128        command,
129        port,
130        addr,
131        user_id,
132        domain,
133    })
134}
135
136/// Write a SOCKS4 reply to the stream.
137///
138/// Format:
139///   +----+----+----+----+----+----+----+----+
140///   | VN | CD | DSTPORT |      DSTIP        |
141///   +----+----+----+----+----+----+----+----+
142///     1    1      2            4
143pub async fn write_socks4_reply<S: tokio::io::AsyncWrite + Unpin>(
144    stream: &mut S,
145    status: Socks4Status,
146    addr: SocketAddr,
147) -> Result<(), Socks4Error> {
148    let ip = match addr.ip() {
149        IpAddr::V4(v4) => v4.octets(),
150        IpAddr::V6(_) => Ipv4Addr::UNSPECIFIED.octets(),
151    };
152    let port = addr.port().to_be_bytes();
153    let reply: [u8; 8] = [
154        0x00,
155        status as u8,
156        port[0],
157        port[1],
158        ip[0],
159        ip[1],
160        ip[2],
161        ip[3],
162    ];
163    stream.write_all(&reply).await?;
164    stream.flush().await?;
165    Ok(())
166}
167
168#[cfg(test)]
169mod tests {
170    use super::*;
171
172    fn encode_request(uid: &[u8], port: u16, domain: Option<&[u8]>) -> Vec<u8> {
173        let mut buf = vec![0x04, 0x01];
174        buf.extend_from_slice(&port.to_be_bytes());
175        // SOCKS4a sentinel IP for domain; otherwise use 127.0.0.1
176        if domain.is_some() {
177            buf.extend_from_slice(&[0, 0, 0, 1]);
178        } else {
179            buf.extend_from_slice(&[127, 0, 0, 1]);
180        }
181        buf.extend_from_slice(uid);
182        buf.push(0x00);
183        if let Some(d) = domain {
184            buf.extend_from_slice(d);
185            buf.push(0x00);
186        }
187        buf
188    }
189
190    #[tokio::test]
191    async fn accepts_max_user_id_length() {
192        let uid = vec![b'A'; MAX_USER_ID_LEN];
193        let req = encode_request(&uid, 80, None);
194        let mut stream: &[u8] = &req;
195        let parsed = read_socks4_request(&mut stream)
196            .await
197            .expect("255-byte UID should parse");
198        assert_eq!(parsed.user_id.len(), MAX_USER_ID_LEN);
199    }
200
201    #[tokio::test]
202    async fn rejects_overlong_user_id() {
203        let uid = vec![b'A'; MAX_USER_ID_LEN + 1];
204        let req = encode_request(&uid, 80, None);
205        let mut stream: &[u8] = &req;
206        let result = read_socks4_request(&mut stream).await;
207        assert!(matches!(result, Err(Socks4Error::UserIdTooLong)));
208    }
209
210    #[tokio::test]
211    async fn accepts_max_domain_length() {
212        let domain = vec![b'a'; MAX_DOMAIN_LEN];
213        let mut req = vec![0x04, 0x01];
214        req.extend_from_slice(&80u16.to_be_bytes());
215        req.extend_from_slice(&[0, 0, 0, 1]); // SOCKS4a sentinel
216        req.push(0x00); // empty user id
217        req.extend_from_slice(&domain);
218        req.push(0x00);
219        let mut stream: &[u8] = &req;
220        let parsed = read_socks4_request(&mut stream)
221            .await
222            .expect("255-byte domain should parse");
223        assert_eq!(
224            parsed.domain.as_ref().map(String::len),
225            Some(MAX_DOMAIN_LEN)
226        );
227    }
228
229    #[tokio::test]
230    async fn rejects_overlong_domain() {
231        let domain = vec![b'a'; MAX_DOMAIN_LEN + 1];
232        let mut req = vec![0x04, 0x01];
233        req.extend_from_slice(&80u16.to_be_bytes());
234        req.extend_from_slice(&[0, 0, 0, 1]);
235        req.push(0x00);
236        req.extend_from_slice(&domain);
237        req.push(0x00);
238        let mut stream: &[u8] = &req;
239        let result = read_socks4_request(&mut stream).await;
240        assert!(matches!(result, Err(Socks4Error::DomainTooLong)));
241    }
242}
eggress_protocol_socks/socks4/server.rs

eggress_protocol_socks/socks4/
server.rs
