eggress-protocol-socks 1.0.4

use std::net::{IpAddr, Ipv4Addr, SocketAddr};

use tokio::io::{AsyncReadExt, AsyncWriteExt};

use super::error::Socks4Error;
use super::{MAX_DOMAIN_LEN, MAX_USER_ID_LEN};

/// SOCKS4 reply status codes.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u8)]
pub enum Socks4Status {
    Granted = 90,
    Failed = 91,
    FailedNoIdent = 92,
    FailedDifferentUser = 93,
}

impl Socks4Status {
    /// Convert from raw status byte.
    pub fn from_u8(val: u8) -> Option<Self> {
        match val {
            90 => Some(Self::Granted),
            91 => Some(Self::Failed),
            92 => Some(Self::FailedNoIdent),
            93 => Some(Self::FailedDifferentUser),
            _ => None,
        }
    }
}

/// Parsed SOCKS4/4a CONNECT request.
#[derive(Debug, Clone)]
pub struct Socks4Request {
    pub command: u8,
    pub port: u16,
    pub addr: SocketAddr,
    pub user_id: String,
    /// For SOCKS4a: the domain name when IP is 0.0.0.x (x != 0).
    pub domain: Option<String>,
}

/// Read a SOCKS4/4a request from the stream.
///
/// Format:
///   +----+----+----+----+----+----+----+----+----+----+....+----+
///   | VN | CD | DSTPORT |      DSTIP        | USERID       |0x00|
///   +----+----+----+----+----+----+----+----+----+----+....+----+
///     1    1      2            4              variable       1
///
/// For SOCKS4a, when DSTIP is 0.0.0.x (x != 0), the domain follows
/// after the NUL-terminated USERID.
pub async fn read_socks4_request<S: tokio::io::AsyncRead + Unpin>(
    stream: &mut S,
) -> Result<Socks4Request, Socks4Error> {
    let mut header = [0u8; 8];
    stream.read_exact(&mut header).await?;

    let version = header[0];
    if version != 0x04 {
        return Err(Socks4Error::InvalidVersion(version));
    }

    let command = header[1];
    if command != 0x01 {
        return Err(Socks4Error::UnsupportedCommand(command));
    }

    let port = u16::from_be_bytes([header[2], header[3]]);
    let ip = Ipv4Addr::new(header[4], header[5], header[6], header[7]);

    // Read NUL-terminated user ID (bounded at MAX_USER_ID_LEN + 1 for NUL).
    let mut user_id_bytes = Vec::with_capacity(64);
    let mut buf = [0u8; 1];
    loop {
        if user_id_bytes.len() > MAX_USER_ID_LEN {
            return Err(Socks4Error::UserIdTooLong);
        }
        let n = stream.read(&mut buf).await?;
        if n == 0 {
            return Err(Socks4Error::MalformedRequest(
                "unexpected EOF reading user ID".into(),
            ));
        }
        if buf[0] == 0x00 {
            break;
        }
        user_id_bytes.push(buf[0]);
    }

    let user_id = String::from_utf8(user_id_bytes)
        .map_err(|_| Socks4Error::MalformedRequest("invalid UTF-8 in user ID".into()))?;

    // SOCKS4a: if IP is 0.0.0.x (x != 0), read domain after user ID.
    let domain =
        if ip.octets()[0] == 0 && ip.octets()[1] == 0 && ip.octets()[2] == 0 && ip.octets()[3] != 0
        {
            let mut domain_bytes = Vec::with_capacity(256);
            loop {
                if domain_bytes.len() > MAX_DOMAIN_LEN {
                    return Err(Socks4Error::DomainTooLong);
                }
                let n = stream.read(&mut buf).await?;
                if n == 0 {
                    return Err(Socks4Error::MalformedRequest(
                        "unexpected EOF reading domain".into(),
                    ));
                }
                if buf[0] == 0x00 {
                    break;
                }
                domain_bytes.push(buf[0]);
            }
            let domain = String::from_utf8(domain_bytes)
                .map_err(|_| Socks4Error::MalformedRequest("invalid UTF-8 in domain".into()))?;
            if domain.is_empty() {
                return Err(Socks4Error::MalformedRequest(
                    "empty domain in SOCKS4a request".into(),
                ));
            }
            Some(domain)
        } else {
            None
        };

    let addr = SocketAddr::new(IpAddr::V4(ip), port);

    Ok(Socks4Request {
        command,
        port,
        addr,
        user_id,
        domain,
    })
}

/// Write a SOCKS4 reply to the stream.
///
/// Format:
///   +----+----+----+----+----+----+----+----+
///   | VN | CD | DSTPORT |      DSTIP        |
///   +----+----+----+----+----+----+----+----+
///     1    1      2            4
pub async fn write_socks4_reply<S: tokio::io::AsyncWrite + Unpin>(
    stream: &mut S,
    status: Socks4Status,
    addr: SocketAddr,
) -> Result<(), Socks4Error> {
    let ip = match addr.ip() {
        IpAddr::V4(v4) => v4.octets(),
        IpAddr::V6(_) => Ipv4Addr::UNSPECIFIED.octets(),
    };
    let port = addr.port().to_be_bytes();
    let reply: [u8; 8] = [
        0x00,
        status as u8,
        port[0],
        port[1],
        ip[0],
        ip[1],
        ip[2],
        ip[3],
    ];
    stream.write_all(&reply).await?;
    stream.flush().await?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    fn encode_request(uid: &[u8], port: u16, domain: Option<&[u8]>) -> Vec<u8> {
        let mut buf = vec![0x04, 0x01];
        buf.extend_from_slice(&port.to_be_bytes());
        // SOCKS4a sentinel IP for domain; otherwise use 127.0.0.1
        if domain.is_some() {
            buf.extend_from_slice(&[0, 0, 0, 1]);
        } else {
            buf.extend_from_slice(&[127, 0, 0, 1]);
        }
        buf.extend_from_slice(uid);
        buf.push(0x00);
        if let Some(d) = domain {
            buf.extend_from_slice(d);
            buf.push(0x00);
        }
        buf
    }

    #[tokio::test]
    async fn accepts_max_user_id_length() {
        let uid = vec![b'A'; MAX_USER_ID_LEN];
        let req = encode_request(&uid, 80, None);
        let mut stream: &[u8] = &req;
        let parsed = read_socks4_request(&mut stream)
            .await
            .expect("255-byte UID should parse");
        assert_eq!(parsed.user_id.len(), MAX_USER_ID_LEN);
    }

    #[tokio::test]
    async fn rejects_overlong_user_id() {
        let uid = vec![b'A'; MAX_USER_ID_LEN + 1];
        let req = encode_request(&uid, 80, None);
        let mut stream: &[u8] = &req;
        let result = read_socks4_request(&mut stream).await;
        assert!(matches!(result, Err(Socks4Error::UserIdTooLong)));
    }

    #[tokio::test]
    async fn accepts_max_domain_length() {
        let domain = vec![b'a'; MAX_DOMAIN_LEN];
        let mut req = vec![0x04, 0x01];
        req.extend_from_slice(&80u16.to_be_bytes());
        req.extend_from_slice(&[0, 0, 0, 1]); // SOCKS4a sentinel
        req.push(0x00); // empty user id
        req.extend_from_slice(&domain);
        req.push(0x00);
        let mut stream: &[u8] = &req;
        let parsed = read_socks4_request(&mut stream)
            .await
            .expect("255-byte domain should parse");
        assert_eq!(
            parsed.domain.as_ref().map(String::len),
            Some(MAX_DOMAIN_LEN)
        );
    }

    #[tokio::test]
    async fn rejects_overlong_domain() {
        let domain = vec![b'a'; MAX_DOMAIN_LEN + 1];
        let mut req = vec![0x04, 0x01];
        req.extend_from_slice(&80u16.to_be_bytes());
        req.extend_from_slice(&[0, 0, 0, 1]);
        req.push(0x00);
        req.extend_from_slice(&domain);
        req.push(0x00);
        let mut stream: &[u8] = &req;
        let result = read_socks4_request(&mut stream).await;
        assert!(matches!(result, Err(Socks4Error::DomainTooLong)));
    }
}