rusty-penguin 0.8.3

//! Client `remote` specification.
//
// SPDX-License-Identifier: Apache-2.0 OR GPL-3.0-or-later

use std::{fmt::Display, str::FromStr};
use thiserror::Error;

#[cfg(not(feature = "default-is-ipv6"))]
/// Default host for local and unspecified addresses.
macro_rules! default_host {
    (local) => {
        String::from("127.0.0.1")
    };
    (unspec) => {
        String::from("0.0.0.0")
    };
}
#[cfg(feature = "default-is-ipv6")]
/// Default host for local and unspecified addresses.
macro_rules! default_host {
    (local) => {
        String::from("::1")
    };
    (unspec) => {
        String::from("::")
    };
}

// Export this macro for use in tests
#[cfg(test)]
pub(crate) use default_host;

/// Default SOCKS port
pub const SOCKS_DEFAULT_PORT: u16 = 1080;
/// Default TPROXY port
pub const TPROXY_DEFAULT_PORT: u16 = 1234;

/// Configuration for one item to forward
#[derive(Debug, Clone, Hash, Eq, PartialEq)]
pub struct Remote {
    /// Local-side forwarding info
    pub local_addr: LocalSpec,
    /// Peer-side forwarding info
    #[expect(clippy::struct_field_names)]
    pub remote_addr: RemoteSpec,
    /// Layer-4 protocol this instance forwards
    pub protocol: Protocol,
}

/// The local side can be either IP+port or "stdio"
#[derive(Debug, Clone, Hash, Eq, PartialEq)]
pub enum LocalSpec {
    /// An IP socket
    Inet((String, u16)),
    /// Standard input/output
    Stdio,
}

/// The remote side can be either IP+port or "socks"
#[derive(Debug, Clone, Hash, Eq, PartialEq)]
pub enum RemoteSpec {
    /// An IP socket
    Inet((String, u16)),
    /// Function as a SOCKS proxy
    Socks,
    /// Configure the listen socket for Transparent Proxy
    Tproxy,
}

/// Protocol can be either "tcp" or "udp".
#[derive(Debug, Copy, Clone, Hash, Eq, PartialEq)]
pub enum Protocol {
    /// Transmission Control Protocol
    Tcp,
    /// User Datagram Protocol
    Udp,
}

/// Errors that can occur when parsing a remote
#[derive(Clone, Error, Debug, PartialEq, Eq)]
pub enum Error {
    /// IPv6 address empty bracket
    #[error("missing address inside brackets")]
    AddressEmpty,
    /// IPv6 address missing closing bracket
    #[error("missing closing `]`")]
    BracketMismatch,
    /// Unexpected character after an IPv6 address
    #[error("found garbage following IPv6 (first offending character `{0}`)")]
    GarbageAfterAddress(char),
    #[error("invalid port or unexpected host `{0}`: {1:?}")]
    Port(String, std::num::IntErrorKind),
    /// Invalid protocol
    #[error("invalid protocol `{0}`")]
    Protocol(String),
    /// `TProxy` needs a socket to listen on
    #[error("stdio cannot accept Transparent Proxy")]
    StdioTproxy,
    /// Too many segments
    #[error("found more than four colon-separated segments")]
    TooManySegments,
    /// SOCKS listen endpoint cannot be UDP
    #[error("SOCKS remote must be TCP")]
    UdpSocks,
}

impl Display for Protocol {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(match self {
            Self::Tcp => "tcp",
            Self::Udp => "udp",
        })
    }
}

impl FromStr for Protocol {
    type Err = Error;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.to_lowercase().as_str() {
            "tcp" => Ok(Self::Tcp),
            "udp" => Ok(Self::Udp),
            other => Err(Error::Protocol(other.to_string())),
        }
    }
}

/// Tokenize a remote string by splitting on `:`.
/// We don't need RE!
fn tokenize_remote(s: &str) -> Result<Vec<&str>, Error> {
    let mut tokens = Vec::new();
    let mut stuff = s;
    macro_rules! check_too_many_and_push {
        ($token:expr) => {
            if tokens.len() >= 4 {
                // Check for this early to reduce work in the error case
                return Err(Error::TooManySegments);
            }
            tokens.push($token);
        };
    }
    loop {
        // IPv6 address in brackets
        if stuff.starts_with('[') {
            // `end` is a byte index on UTF-8 boundaries
            let end = stuff.find(']').ok_or(Error::BracketMismatch)? + 1;
            check_too_many_and_push!(&stuff[..end]);
            // Now stuff[end..] should start with ':', so we assert that and skip it.
            let following = stuff[end..].chars().next();
            if let Some(ch) = following
                && ch != ':'
            {
                return Err(Error::GarbageAfterAddress(ch));
            }
            stuff = stuff.get(end + 1..).ok_or(Error::AddressEmpty)?;
        } else if let Some((token, rest)) = stuff.split_once(':') {
            check_too_many_and_push!(token);
            stuff = rest;
        } else {
            check_too_many_and_push!(stuff);
            return Ok(tokens);
        }
    }
}

impl Display for Remote {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match &self.local_addr {
            LocalSpec::Inet((host, port)) => {
                if host.contains(':') {
                    write!(f, "[{host}]:{port}")?;
                } else {
                    write!(f, "{host}:{port}")?;
                }
            }
            LocalSpec::Stdio => f.write_str("stdio")?,
        }
        match &self.remote_addr {
            RemoteSpec::Inet((host, port)) => {
                if host.contains(':') {
                    write!(f, ":[{host}]:{port}")?;
                } else {
                    write!(f, ":{host}:{port}")?;
                }
            }
            RemoteSpec::Socks => f.write_str(":socks")?,
            RemoteSpec::Tproxy => f.write_str(":tproxy")?,
        }
        write!(f, "/{}", self.protocol)?;
        Ok(())
    }
}

impl FromStr for Remote {
    type Err = Error;

    /// Parse a remote specification.
    #[expect(clippy::too_many_lines)]
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        macro_rules! parse_port_or_bail {
            ($port_str:expr) => {
                $port_str
                    .parse::<u16>()
                    .map_err(|e| Error::Port($port_str.to_string(), *e.kind()))?
            };
        }
        let (rest, proto) = match s.rsplit_once('/') {
            Some((rest, proto)) => (rest, proto.parse()?),
            None => (s, Protocol::Tcp),
        };
        let tokens = tokenize_remote(rest)?;
        let result = match tokens[..] {
            // One element: either "socks", "tproxy" or a port number.
            ["socks"] => Self {
                local_addr: LocalSpec::Inet((default_host!(local), SOCKS_DEFAULT_PORT)),
                remote_addr: RemoteSpec::Socks,
                protocol: proto,
            },
            ["tproxy"] => Self {
                local_addr: LocalSpec::Inet((default_host!(local), TPROXY_DEFAULT_PORT)),
                remote_addr: RemoteSpec::Tproxy,
                protocol: proto,
            },
            [port] => Self {
                local_addr: LocalSpec::Inet((default_host!(unspec), parse_port_or_bail!(port))),
                remote_addr: RemoteSpec::Inet((default_host!(local), parse_port_or_bail!(port))),
                protocol: proto,
            },
            // Two elements: either "socks" or "tproxy" and local port number, or remote host and port number.
            ["stdio", "socks"] => Self {
                local_addr: LocalSpec::Stdio,
                remote_addr: RemoteSpec::Socks,
                protocol: proto,
            },
            ["stdio", "tproxy"] => return Err(Error::StdioTproxy),
            [port, "socks"] => Self {
                local_addr: LocalSpec::Inet((default_host!(local), parse_port_or_bail!(port))),
                remote_addr: RemoteSpec::Socks,
                protocol: proto,
            },
            [port, "tproxy"] => Self {
                local_addr: LocalSpec::Inet((default_host!(local), parse_port_or_bail!(port))),
                remote_addr: RemoteSpec::Tproxy,
                protocol: proto,
            },
            ["stdio", port] => Self {
                local_addr: LocalSpec::Stdio,
                remote_addr: RemoteSpec::Inet((default_host!(local), parse_port_or_bail!(port))),
                protocol: proto,
            },
            [host, port] => Self {
                local_addr: LocalSpec::Inet((default_host!(unspec), parse_port_or_bail!(port))),
                remote_addr: RemoteSpec::Inet((
                    remove_brackets(host).to_string(),
                    parse_port_or_bail!(port),
                )),
                protocol: proto,
            },
            // Three elements:
            // - "stdio", remote host, and port number,
            // - local host, local port number, and "socks", or
            // - local port number, remote host, and port number.
            ["stdio", remote_host, remote_port] => Self {
                local_addr: LocalSpec::Stdio,
                remote_addr: RemoteSpec::Inet((
                    remove_brackets(remote_host).to_string(),
                    parse_port_or_bail!(remote_port),
                )),
                protocol: proto,
            },
            [local_host, local_port, "socks"] => Self {
                local_addr: LocalSpec::Inet((
                    remove_brackets(local_host).to_string(),
                    parse_port_or_bail!(local_port),
                )),
                remote_addr: RemoteSpec::Socks,
                protocol: proto,
            },
            [local_host, local_port, "tproxy"] => Self {
                local_addr: LocalSpec::Inet((
                    remove_brackets(local_host).to_string(),
                    parse_port_or_bail!(local_port),
                )),
                remote_addr: RemoteSpec::Tproxy,
                protocol: proto,
            },
            [local_port, remote_host, remote_port] => Self {
                local_addr: LocalSpec::Inet((
                    default_host!(unspec),
                    parse_port_or_bail!(local_port),
                )),
                remote_addr: RemoteSpec::Inet((
                    remove_brackets(remote_host).to_string(),
                    parse_port_or_bail!(remote_port),
                )),
                protocol: proto,
            },
            [local_host, local_port, remote_host, remote_port] => Self {
                local_addr: LocalSpec::Inet((
                    remove_brackets(local_host).to_string(),
                    parse_port_or_bail!(local_port),
                )),
                remote_addr: RemoteSpec::Inet((
                    remove_brackets(remote_host).to_string(),
                    parse_port_or_bail!(remote_port),
                )),
                protocol: proto,
            },
            _ => {
                // This should be unreachable since we check in `tokenize_remote`
                debug_assert!(
                    false,
                    "`tokenize_remote` did not catch too many segments (this is a bug)"
                );
                return Err(Error::TooManySegments);
            }
        };
        if matches!(
            result,
            Self {
                remote_addr: RemoteSpec::Socks,
                protocol: Protocol::Udp,
                ..
            }
        ) {
            Err(Error::UdpSocks)
        } else {
            Ok(result)
        }
    }
}

/// Remove brackets from possbly an IPv6 address
#[must_use]
pub fn remove_brackets(s: &str) -> &str {
    if s.starts_with('[') && s.ends_with(']') {
        &s[1..s.len() - 1]
    } else {
        s
    }
}

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

    // We could have DRYed, but this is to physically show the difference.
    #[cfg(not(feature = "default-is-ipv6"))]
    #[test]
    fn test_default_host() {
        crate::tests::setup_logging();
        assert_eq!(
            std::net::Ipv4Addr::from_str(&default_host!(unspec)).unwrap(),
            std::net::Ipv4Addr::UNSPECIFIED
        );
        assert_eq!(
            std::net::Ipv4Addr::from_str(&default_host!(local)).unwrap(),
            std::net::Ipv4Addr::LOCALHOST
        );
    }
    #[cfg(feature = "default-is-ipv6")]
    #[test]
    fn test_default_host() {
        crate::tests::setup_logging();
        assert_eq!(
            std::net::Ipv6Addr::from_str(&default_host!(unspec)).unwrap(),
            std::net::Ipv6Addr::UNSPECIFIED
        );
        assert_eq!(
            std::net::Ipv6Addr::from_str(&default_host!(local)).unwrap(),
            std::net::Ipv6Addr::LOCALHOST
        );
    }

    #[test]
    #[allow(clippy::too_many_lines)]
    fn test_parse_remote() {
        crate::tests::setup_logging();
        let tests: &[(&str, Remote)] = &[
            // jpillora's tests and an exhausive list of cases
            (
                "3000",
                Remote {
                    local_addr: LocalSpec::Inet((default_host!(unspec), 3000)),
                    remote_addr: RemoteSpec::Inet((default_host!(local), 3000)),
                    protocol: Protocol::Tcp,
                },
            ),
            (
                "4000/udp",
                Remote {
                    local_addr: LocalSpec::Inet((default_host!(unspec), 4000)),
                    remote_addr: RemoteSpec::Inet((default_host!(local), 4000)),
                    protocol: Protocol::Udp,
                },
            ),
            (
                "google.com:80",
                Remote {
                    local_addr: LocalSpec::Inet((default_host!(unspec), 80)),
                    remote_addr: RemoteSpec::Inet((String::from("google.com"), 80)),
                    protocol: Protocol::Tcp,
                },
            ),
            (
                "テスト.net:80",
                Remote {
                    local_addr: LocalSpec::Inet((default_host!(unspec), 80)),
                    remote_addr: RemoteSpec::Inet((String::from("テスト.net"), 80)),
                    protocol: Protocol::Tcp,
                },
            ),
            (
                "8080:example.com:80",
                Remote {
                    local_addr: LocalSpec::Inet((default_host!(unspec), 8080)),
                    remote_addr: RemoteSpec::Inet((String::from("example.com"), 80)),
                    protocol: Protocol::Tcp,
                },
            ),
            (
                "socks",
                Remote {
                    local_addr: LocalSpec::Inet((default_host!(local), SOCKS_DEFAULT_PORT)),
                    remote_addr: RemoteSpec::Socks,
                    protocol: Protocol::Tcp,
                },
            ),
            (
                "9050:socks",
                Remote {
                    local_addr: LocalSpec::Inet((default_host!(local), 9050)),
                    remote_addr: RemoteSpec::Socks,
                    protocol: Protocol::Tcp,
                },
            ),
            (
                "127.0.0.1:1081:socks",
                Remote {
                    local_addr: LocalSpec::Inet((String::from("127.0.0.1"), 1081)),
                    remote_addr: RemoteSpec::Socks,
                    protocol: Protocol::Tcp,
                },
            ),
            (
                "9050:socks",
                Remote {
                    local_addr: LocalSpec::Inet((default_host!(local), 9050)),
                    remote_addr: RemoteSpec::Socks,
                    protocol: Protocol::Tcp,
                },
            ),
            (
                "tproxy",
                Remote {
                    local_addr: LocalSpec::Inet((default_host!(local), TPROXY_DEFAULT_PORT)),
                    remote_addr: RemoteSpec::Tproxy,
                    protocol: Protocol::Tcp,
                },
            ),
            (
                "tproxy/udp",
                Remote {
                    local_addr: LocalSpec::Inet((default_host!(local), TPROXY_DEFAULT_PORT)),
                    remote_addr: RemoteSpec::Tproxy,
                    protocol: Protocol::Udp,
                },
            ),
            (
                "5000:tproxy",
                Remote {
                    local_addr: LocalSpec::Inet((default_host!(local), 5000)),
                    remote_addr: RemoteSpec::Tproxy,
                    protocol: Protocol::Tcp,
                },
            ),
            (
                "4567:tproxy/udp",
                Remote {
                    local_addr: LocalSpec::Inet((default_host!(local), 4567)),
                    remote_addr: RemoteSpec::Tproxy,
                    protocol: Protocol::Udp,
                },
            ),
            (
                "127.0.0.1:1081:tproxy",
                Remote {
                    local_addr: LocalSpec::Inet((String::from("127.0.0.1"), 1081)),
                    remote_addr: RemoteSpec::Tproxy,
                    protocol: Protocol::Tcp,
                },
            ),
            (
                "127.0.0.1:1081:tproxy/udp",
                Remote {
                    local_addr: LocalSpec::Inet((String::from("127.0.0.1"), 1081)),
                    remote_addr: RemoteSpec::Tproxy,
                    protocol: Protocol::Udp,
                },
            ),
            (
                "[::1]:12345:tproxy/udp",
                Remote {
                    local_addr: LocalSpec::Inet((String::from("::1"), 12345)),
                    remote_addr: RemoteSpec::Tproxy,
                    protocol: Protocol::Udp,
                },
            ),
            (
                "1.1.1.1:53/udp",
                Remote {
                    local_addr: LocalSpec::Inet((default_host!(unspec), 53)),
                    remote_addr: RemoteSpec::Inet((String::from("1.1.1.1"), 53)),
                    protocol: Protocol::Udp,
                },
            ),
            (
                "localhost:5353:1.1.1.1:53/udp",
                Remote {
                    local_addr: LocalSpec::Inet((String::from("localhost"), 5353)),
                    remote_addr: RemoteSpec::Inet((String::from("1.1.1.1"), 53)),
                    protocol: Protocol::Udp,
                },
            ),
            (
                "22:example.com:22",
                Remote {
                    local_addr: LocalSpec::Inet((default_host!(unspec), 22)),
                    remote_addr: RemoteSpec::Inet((String::from("example.com"), 22)),
                    protocol: Protocol::Tcp,
                },
            ),
            (
                "[::1]:8080:google.com:80",
                Remote {
                    local_addr: LocalSpec::Inet((String::from("::1"), 8080)),
                    remote_addr: RemoteSpec::Inet((String::from("google.com"), 80)),
                    protocol: Protocol::Tcp,
                },
            ),
            (
                "localhost:5354:[2001:db8:4860:0:0:0:0:8888]:53/udp",
                Remote {
                    local_addr: LocalSpec::Inet((String::from("localhost"), 5354)),
                    remote_addr: RemoteSpec::Inet((String::from("2001:db8:4860:0:0:0:0:8888"), 53)),
                    protocol: Protocol::Udp,
                },
            ),
            (
                // Make sure your editor supports mixed LTR and RTL before editing this line
                "آزمایشی.com:123:δοκιμή.net:9999/tcp",
                Remote {
                    local_addr: LocalSpec::Inet((String::from("آزمایشی.com"), 123)),
                    remote_addr: RemoteSpec::Inet((String::from("δοκιμή.net"), 9999)),
                    protocol: Protocol::Tcp,
                },
            ),
            (
                "stdio:google.com:80",
                Remote {
                    local_addr: LocalSpec::Stdio,
                    remote_addr: RemoteSpec::Inet((String::from("google.com"), 80)),
                    protocol: Protocol::Tcp,
                },
            ),
            (
                "stdio:socks",
                Remote {
                    local_addr: LocalSpec::Stdio,
                    remote_addr: RemoteSpec::Socks,
                    protocol: Protocol::Tcp,
                },
            ),
            (
                "stdio:443",
                Remote {
                    local_addr: LocalSpec::Stdio,
                    remote_addr: RemoteSpec::Inet((default_host!(local), 443)),
                    protocol: Protocol::Tcp,
                },
            ),
            (
                "stdio:5353/udp",
                Remote {
                    local_addr: LocalSpec::Stdio,
                    remote_addr: RemoteSpec::Inet((default_host!(local), 5353)),
                    protocol: Protocol::Udp,
                },
            ),
        ];
        for (s, expected) in tests {
            // Test that the common format is parsed correctly
            let actual = s.parse::<Remote>().unwrap();
            assert_eq!(actual, *expected);
            // Test that the canonical format is made and parsed correctly
            let reparsed = actual.to_string().parse::<Remote>().unwrap();
            assert_eq!(reparsed, *expected);
        }
    }

    #[test]
    #[allow(clippy::too_many_lines)]
    fn test_parse_remote_bad() {
        crate::tests::setup_logging();
        let tests: &[(&str, Error)] = &[
            (
                "just_a_hostname",
                Error::Port(
                    String::from("just_a_hostname"),
                    std::num::IntErrorKind::InvalidDigit,
                ),
            ),
            (
                "99999",
                Error::Port(String::from("99999"), std::num::IntErrorKind::PosOverflow),
            ),
            ("::::", Error::TooManySegments),
            (
                "host:port",
                Error::Port(String::from("port"), std::num::IntErrorKind::InvalidDigit),
            ),
            // More of just abusing Rust's UTF-8 support
            ("[::1]إختبار:80", Error::GarbageAfterAddress('إ')),
            ("[::1:80", Error::BracketMismatch),
            (
                "[::1]:99/nonsense",
                Error::Protocol(String::from("nonsense")),
            ),
            ("socks/udp", Error::UdpSocks),
            ("stdio:tproxy", Error::StdioTproxy),
        ];
        for (s, expected) in tests {
            // Test that an incorrect format is rejected with the correct error
            let actual = s.parse::<Remote>().unwrap_err();
            assert_eq!(actual, *expected);
        }
    }
}