alksocks 0.3.2

//! Forked from fast-socks5 v1.0.0 (upstream `src/util/target_addr.rs`;
//! extraction; ownership per ADR-013, see the crate root's
//! `mod.rs`).
//!
//! Upstream carries `TargetAddr`, the address codec (`read_address`),
//! and the `resolve_dns` helper; only the codec rode the extraction
//! (DNS resolution is the cut net region, since
//! `tokio::net::lookup_host`
//! is not in the wasm-clean graph). The `resolve_dns` method is cut
//! with its two `AddrError` variants (`DNSResolutionFailed`,
//! `NoDNSRecords`).

use crate::rfc1928::consts;
use crate::rfc1928::consts::SOCKS5_ADDR_TYPE_IPV4;
use crate::rfc1928::util::stream::read_exact;
use crate::rfc1928::ReplyError;
use std::fmt;
use std::io;
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6};
use std::vec::IntoIter;
use tokio::io::{AsyncRead, AsyncReadExt};

/// SOCKS5 reply code
#[derive(thiserror::Error, Debug)]
pub enum AddrError {
    #[error("Domain length {0} exceeded maximum")]
    DomainLenTooLong(usize),
    #[error("Can't read IPv4: {0}")]
    IPv4Unreadable(#[source] io::Error),
    #[error("Can't read IPv6: {0}")]
    IPv6Unreadable(#[source] io::Error),
    #[error("Can't read port number: {0}")]
    PortNumberUnreadable(#[source] io::Error),
    #[error("Can't read domain len: {0}")]
    DomainLenUnreadable(#[source] io::Error),
    #[error("Can't read domain content: {0}")]
    DomainContentUnreadable(#[source] io::Error),
    #[error("Can't convert address: {0}")]
    AddrConversionFailed(#[source] io::Error),
    #[error("Malformed UTF-8")]
    Utf8(#[source] std::string::FromUtf8Error),
    #[error("Unknown address type")]
    IncorrectAddressType,
}

impl AddrError {
    pub fn to_reply_error(&self) -> ReplyError {
        match self {
            AddrError::IncorrectAddressType => ReplyError::AddressTypeNotSupported,
            _ => ReplyError::ConnectionRefused,
        }
    }
}

/// A description of a connection target.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum TargetAddr {
    /// Connect to an IP address.
    Ip(SocketAddr),
    /// Connect to a fully qualified domain name.
    ///
    /// The domain name will be passed along to the proxy server and DNS lookup
    /// will happen there.
    Domain(String, u16),
}

impl TargetAddr {
    pub fn is_ip(&self) -> bool {
        matches!(self, TargetAddr::Ip(_))
    }

    pub fn is_domain(&self) -> bool {
        !self.is_ip()
    }

    pub fn to_be_bytes(&self) -> Result<Vec<u8>, AddrError> {
        let mut buf = vec![];
        match self {
            TargetAddr::Ip(SocketAddr::V4(addr)) => {
                buf.extend_from_slice(&[SOCKS5_ADDR_TYPE_IPV4]);

                buf.extend_from_slice(&(addr.ip()).octets()); // ip
                buf.extend_from_slice(&addr.port().to_be_bytes()); // port
            }
            TargetAddr::Ip(SocketAddr::V6(addr)) => {
                buf.extend_from_slice(&[consts::SOCKS5_ADDR_TYPE_IPV6]);

                buf.extend_from_slice(&(addr.ip()).octets()); // ip
                buf.extend_from_slice(&addr.port().to_be_bytes()); // port
            }
            TargetAddr::Domain(ref domain, port) => {
                if domain.len() > u8::MAX as usize {
                    return Err(AddrError::DomainLenTooLong(domain.len()));
                }
                buf.extend_from_slice(&[consts::SOCKS5_ADDR_TYPE_DOMAIN_NAME, domain.len() as u8]);
                buf.extend_from_slice(domain.as_bytes()); // domain content
                buf.extend_from_slice(&port.to_be_bytes());
                // port content (.to_be_bytes() convert from u16 to u8 type)
            }
        }
        Ok(buf)
    }

    pub fn into_string_and_port(self) -> (String, u16) {
        match self {
            TargetAddr::Ip(socket_addr) => (socket_addr.ip().to_string(), socket_addr.port()),
            TargetAddr::Domain(domain, port) => (domain, port),
        }
    }
}

/// The blocking `ToSocketAddrs` form: IP targets resolve directly;
/// domain targets cannot be resolved here (domain resolution is the
/// egress's job, per ADR-012) and error.
impl std::net::ToSocketAddrs for TargetAddr {
    type Iter = IntoIter<SocketAddr>;

    fn to_socket_addrs(&self) -> io::Result<IntoIter<SocketAddr>> {
        match *self {
            TargetAddr::Ip(addr) => Ok(vec![addr].into_iter()),
            TargetAddr::Domain(_, _) => Err(io::Error::other(
                "Domain name cannot be resolved by the blocking ToSocketAddrs interface; resolve it before constructing the target.",
            )),
        }
    }
}

impl fmt::Display for TargetAddr {
    #[inline]
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            TargetAddr::Ip(ref addr) => write!(f, "{}", addr),
            TargetAddr::Domain(ref addr, ref port) => write!(f, "{}:{}", addr, port),
        }
    }
}

/// A trait for objects that can be converted to `TargetAddr`.
pub trait ToTargetAddr {
    /// Converts the value of `self` to a `TargetAddr`.
    fn to_target_addr(&self) -> io::Result<TargetAddr>;
}

impl ToTargetAddr for (&str, u16) {
    fn to_target_addr(&self) -> io::Result<TargetAddr> {
        // try to parse as an IP first
        if let Ok(addr) = self.0.parse::<Ipv4Addr>() {
            return (addr, self.1).to_target_addr();
        }

        if let Ok(addr) = self.0.parse::<Ipv6Addr>() {
            return (addr, self.1).to_target_addr();
        }

        Ok(TargetAddr::Domain(self.0.to_owned(), self.1))
    }
}

impl ToTargetAddr for SocketAddr {
    fn to_target_addr(&self) -> io::Result<TargetAddr> {
        Ok(TargetAddr::Ip(*self))
    }
}

impl ToTargetAddr for SocketAddrV4 {
    fn to_target_addr(&self) -> io::Result<TargetAddr> {
        SocketAddr::V4(*self).to_target_addr()
    }
}

impl ToTargetAddr for SocketAddrV6 {
    fn to_target_addr(&self) -> io::Result<TargetAddr> {
        SocketAddr::V6(*self).to_target_addr()
    }
}

impl ToTargetAddr for (IpAddr, u16) {
    fn to_target_addr(&self) -> io::Result<TargetAddr> {
        match self.0 {
            IpAddr::V4(ipv4_addr) => (ipv4_addr, self.1).to_target_addr(),
            IpAddr::V6(ipv6_addr) => (ipv6_addr, self.1).to_target_addr(),
        }
    }
}

impl ToTargetAddr for (Ipv4Addr, u16) {
    fn to_target_addr(&self) -> io::Result<TargetAddr> {
        SocketAddrV4::new(self.0, self.1).to_target_addr()
    }
}

impl ToTargetAddr for (Ipv6Addr, u16) {
    fn to_target_addr(&self) -> io::Result<TargetAddr> {
        SocketAddrV6::new(self.0, self.1, 0, 0).to_target_addr()
    }
}

impl ToTargetAddr for TargetAddr {
    fn to_target_addr(&self) -> io::Result<TargetAddr> {
        Ok(self.clone())
    }
}

impl ToTargetAddr for &TargetAddr {
    fn to_target_addr(&self) -> io::Result<TargetAddr> {
        Ok((*self).clone())
    }
}

#[derive(Debug)]
pub enum Addr {
    V4([u8; 4]),
    V6([u8; 16]),
    // `String`, validated UTF-8 at the codec boundary (`read_address`
    // converts before constructing).
    Domain(String),
}

/// This function is used by the client & the server.
///
/// The domain arm's allocation is bounded by construction: the wire
/// length field is one byte, so the claimed length is at most 255;
/// the attacker-controlled byte cannot ask for more. Any change that
/// widens that field (an FRAG-style extension) requires an explicit
/// cap here before the allocation.
pub async fn read_address<T: AsyncRead + Unpin>(
    stream: &mut T,
    atyp: u8,
) -> Result<TargetAddr, AddrError> {
    let addr = match atyp {
        consts::SOCKS5_ADDR_TYPE_IPV4 => {
            Addr::V4(read_exact!(stream, [0u8; 4]).map_err(AddrError::IPv4Unreadable)?)
        }
        consts::SOCKS5_ADDR_TYPE_IPV6 => {
            Addr::V6(read_exact!(stream, [0u8; 16]).map_err(AddrError::IPv6Unreadable)?)
        }
        consts::SOCKS5_ADDR_TYPE_DOMAIN_NAME => {
            let len = read_exact!(stream, [0]).map_err(AddrError::DomainLenUnreadable)?[0];
            let domain = read_exact!(stream, vec![0u8; len as usize])
                .map_err(AddrError::DomainContentUnreadable)?;
            // make sure the bytes are correct utf8 string
            let domain = String::from_utf8(domain).map_err(AddrError::Utf8)?;

            Addr::Domain(domain)
        }
        _ => return Err(AddrError::IncorrectAddressType),
    };

    // Find port number
    let port = read_exact!(stream, [0u8; 2]).map_err(AddrError::PortNumberUnreadable)?;
    // Convert (u8 * 2) into u16
    let port = (port[0] as u16) << 8 | port[1] as u16;

    // Merge ADDRESS + PORT into a TargetAddr
    let addr: TargetAddr = match addr {
        Addr::V4([a, b, c, d]) => (Ipv4Addr::new(a, b, c, d), port)
            .to_target_addr()
            .map_err(AddrError::AddrConversionFailed)?,
        Addr::V6(x) => (Ipv6Addr::from(x), port)
            .to_target_addr()
            .map_err(AddrError::AddrConversionFailed)?,
        Addr::Domain(domain) => TargetAddr::Domain(domain, port),
    };

    Ok(addr)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::rfc1928::consts;
    use tokio::io::AsyncWriteExt;

    async fn read_target_addr(bytes: &[u8]) -> Result<TargetAddr, AddrError> {
        let (mut client, mut server) = tokio::io::duplex(64);
        client.write_all(&bytes[1..]).await.unwrap();
        drop(client);
        read_address(&mut server, bytes[0]).await
    }

    #[tokio::test]
    async fn read_address_ipv4_round_trip() {
        let addr = read_target_addr(&[consts::SOCKS5_ADDR_TYPE_IPV4, 192, 168, 0, 1, 0, 80])
            .await
            .unwrap();
        assert_eq!(addr, TargetAddr::Ip("192.168.0.1:80".parse().unwrap()));
        assert!(addr.is_ip());
        let encoded = addr.to_be_bytes().unwrap();
        assert_eq!(
            encoded,
            vec![consts::SOCKS5_ADDR_TYPE_IPV4, 192, 168, 0, 1, 0, 80,]
        );
    }

    #[tokio::test]
    async fn read_address_ipv6_round_trip() {
        let addr = read_target_addr(&[
            consts::SOCKS5_ADDR_TYPE_IPV6,
            0x20,
            0x01,
            0x0d,
            0xb8,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0,
            0x01,
            0x01,
            0xbb,
        ])
        .await
        .unwrap();
        assert_eq!(addr, TargetAddr::Ip("[2001:db8::1]:443".parse().unwrap()));
        let encoded = addr.to_be_bytes().unwrap();
        assert_eq!(
            encoded,
            vec![
                consts::SOCKS5_ADDR_TYPE_IPV6,
                0x20,
                0x01,
                0x0d,
                0xb8,
                0,
                0,
                0,
                0,
                0,
                0,
                0,
                0,
                0,
                0,
                0,
                0x01,
                0x01,
                0xbb,
            ]
        );
    }

    #[tokio::test]
    async fn read_address_domain_round_trip() {
        let addr = read_target_addr(&[
            consts::SOCKS5_ADDR_TYPE_DOMAIN_NAME,
            11,
            b'e',
            b'x',
            b'a',
            b'm',
            b'p',
            b'l',
            b'e',
            b'.',
            b'c',
            b'o',
            b'm',
            0,
            80,
        ])
        .await
        .unwrap();
        assert_eq!(addr, TargetAddr::Domain("example.com".to_owned(), 80));
        assert!(addr.is_domain());
        let encoded = addr.to_be_bytes().unwrap();
        assert_eq!(
            encoded,
            vec![
                consts::SOCKS5_ADDR_TYPE_DOMAIN_NAME,
                11,
                b'e',
                b'x',
                b'a',
                b'm',
                b'p',
                b'l',
                b'e',
                b'.',
                b'c',
                b'o',
                b'm',
                0,
                80
            ]
        );
        assert_eq!(addr.into_string_and_port(), ("example.com".to_owned(), 80));
    }

    #[tokio::test]
    async fn read_address_unknown_atyp_is_incorrect() {
        let err = read_target_addr(&[0x02, 0, 0]).await.unwrap_err();
        assert!(matches!(err, AddrError::IncorrectAddressType));
        assert!(matches!(
            err.to_reply_error(),
            ReplyError::AddressTypeNotSupported
        ));
    }

    #[test]
    fn to_be_bytes_domain_too_long() {
        let long_domain = "a".repeat(256);
        let addr = TargetAddr::Domain(long_domain, 80);
        assert!(matches!(
            addr.to_be_bytes(),
            Err(AddrError::DomainLenTooLong(256))
        ));
    }

    #[test]
    fn to_socket_addrs_domain_is_explicit_resolution_only() {
        let addr = TargetAddr::Domain("example.com".to_owned(), 80);
        assert!(std::net::ToSocketAddrs::to_socket_addrs(&addr).is_err());
        let ip: SocketAddr = "10.0.0.1:80".parse().unwrap();
        let addr = TargetAddr::Ip(ip);
        assert_eq!(
            std::net::ToSocketAddrs::to_socket_addrs(&addr)
                .unwrap()
                .next(),
            Some(ip)
        );
    }

    #[test]
    fn to_target_addr_conversions() {
        assert_eq!(
            ("example.com", 80u16).to_target_addr().unwrap(),
            TargetAddr::Domain("example.com".to_owned(), 80)
        );
        assert_eq!(
            ("1.2.3.4", 80u16).to_target_addr().unwrap(),
            TargetAddr::Ip("1.2.3.4:80".parse().unwrap())
        );
        assert_eq!(
            (IpAddr::V4(Ipv4Addr::new(1, 2, 3, 4)), 80u16)
                .to_target_addr()
                .unwrap(),
            TargetAddr::Ip("1.2.3.4:80".parse().unwrap())
        );
        assert_eq!(
            (Ipv6Addr::LOCALHOST, 80u16).to_target_addr().unwrap(),
            TargetAddr::Ip("[::1]:80".parse().unwrap())
        );
        let target = TargetAddr::Domain("example.com".to_owned(), 80);
        assert_eq!(target.to_target_addr().unwrap(), target);
    }
}