borer-core 0.5.9

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

use anyhow::{Context, anyhow};
use bytes::{BufMut, BytesMut};
use socks5_proto::Address;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};

use crate::{CRLF, stream::peekable::AsyncPeek};

use super::padding::Padding;

/// Trojan request
///
/// ```plain
/// +-----------------------+---------+----------------+---------+----------+
/// | hex(SHA224(password)) |  CRLF   | Trojan Request |  CRLF   | Payload  |
/// +-----------------------+---------+----------------+---------+----------+
/// |          56           | X'0D0A' |    Variable    | X'0D0A' | Variable |
/// +-----------------------+---------+----------------+---------+----------+
///
/// where Trojan Request is a SOCKS5-like request:
///
/// +-----+------+----------+----------+
/// | CMD | ATYP | DST.ADDR | DST.PORT |
/// +-----+------+----------+----------+
/// |  1  |  1   | Variable |    2     |
/// +-----+------+----------+----------+
///
/// ```
#[derive(Clone, Debug)]
/// Parsed Trojan request header and destination metadata.
pub struct Request {
    pub hash: String,
    pub command: Command,
    pub address: Address,
}

impl Request {
    const ATYP_IPV4: u8 = 0x01;
    const ATYP_FQDN: u8 = 0x03;
    const ATYP_IPV6: u8 = 0x04;

    /// Create a Trojan request value from its already-decoded parts.
    pub const fn new(hash: String, command: Command, address: Address) -> Self {
        Self {
            hash,
            command,
            address,
        }
    }

    /// Peek the Trojan hash prefix without consuming the underlying stream.
    pub async fn peek_head<T>(r: &mut T) -> anyhow::Result<Vec<u8>>
    where
        T: AsyncRead + AsyncPeek + Unpin,
    {
        let mut buf = Vec::new();
        for _i in 0..56 {
            let b1 = r.peek_u8().await.context("trojan peek u8 failed")?;
            if b1 == b'\r' {
                let b2 = r.peek_u8().await.context("trojan peek u8 failed")?;
                if b2 == b'\n' {
                    buf.push(b1);
                    buf.push(b2);
                    break;
                }
            } else {
                buf.push(b1);
            }
        }

        Ok(buf)
    }

    /// Read a full Trojan request from the stream.
    pub async fn read_from<R>(r: &mut R) -> anyhow::Result<Self>
    where
        R: AsyncRead + Unpin,
    {
        let mut buf: [u8; 56] = [0; 56];
        let len = r
            .read(&mut buf[..])
            .await
            .context("trojan read hash failed")?;
        if len != 56 {
            return Err(anyhow!("the Request not Trojan"));
        }

        let hash = String::from_utf8_lossy(&buf[..]).to_string();

        let _crlf = r.read_u16().await?;

        let (cmd, addr) = Self::read_address_from(r)
            .await
            .context("trojan read Address failed")?;

        let _crlf = r.read_u16().await?;

        if let Command::Padding = cmd {
            let _padding = Padding::read_from(r)
                .await
                .context("trojan read padding failed")?;
        }

        Ok(Self::new(hash, cmd, addr))
    }

    /// Read the Trojan command and destination address from the stream.
    pub async fn read_address_from<R>(r: &mut R) -> anyhow::Result<(Command, Address)>
    where
        R: AsyncRead + Unpin,
    {
        let cmd = r.read_u8().await.context("address read cmd failed")?;
        let cmd = Command::try_from(cmd).map_err(|cmd| anyhow!("Unknown cmd {cmd}"))?;

        let atyp = r.read_u8().await.context("address read atyp failed")?;

        match atyp {
            Self::ATYP_IPV4 => {
                let mut buf = [0; 6];
                r.read_exact(&mut buf)
                    .await
                    .context("address read ipv4 failed")?;

                let addr = Ipv4Addr::new(buf[0], buf[1], buf[2], buf[3]);

                let port = u16::from_be_bytes([buf[4], buf[5]]);

                let addr = Address::SocketAddress(SocketAddr::from((addr, port)));
                Ok((cmd, addr))
            }
            Self::ATYP_FQDN => {
                let len = r.read_u8().await? as usize;

                let mut buf = vec![0; len + 2];
                r.read_exact(&mut buf)
                    .await
                    .context("address read domain failed")?;

                let port = u16::from_be_bytes([buf[len], buf[len + 1]]);
                buf.truncate(len);

                let addr = Address::DomainAddress(buf, port);
                Ok((cmd, addr))
            }
            Self::ATYP_IPV6 => {
                let mut buf = [0; 18];
                r.read_exact(&mut buf)
                    .await
                    .context("address read ipv6 failed")?;

                let addr = Ipv6Addr::new(
                    u16::from_be_bytes([buf[0], buf[1]]),
                    u16::from_be_bytes([buf[2], buf[3]]),
                    u16::from_be_bytes([buf[4], buf[5]]),
                    u16::from_be_bytes([buf[6], buf[7]]),
                    u16::from_be_bytes([buf[8], buf[9]]),
                    u16::from_be_bytes([buf[10], buf[11]]),
                    u16::from_be_bytes([buf[12], buf[13]]),
                    u16::from_be_bytes([buf[14], buf[15]]),
                );

                let port = u16::from_be_bytes([buf[16], buf[17]]);

                let addr = Address::SocketAddress(SocketAddr::from((addr, port)));
                Ok((cmd, addr))
            }
            atyp => Err(anyhow!("invalid type {atyp}")),
        }
    }

    /// Serialize this request to the stream.
    pub async fn write_to<W>(&self, w: &mut W) -> anyhow::Result<()>
    where
        W: AsyncWrite + Unpin,
    {
        let mut buf = BytesMut::with_capacity(self.serialized_len());
        self.write_to_buf(&mut buf);
        w.write_all(&buf).await.context("trojan Write buf failed")?;

        Ok(())
    }

    /// Serialize this request into an existing byte buffer.
    pub fn write_to_buf<B: BufMut>(&self, buf: &mut B) {
        buf.put_slice(self.hash.as_bytes());
        buf.put_slice(&CRLF);
        buf.put_u8(u8::from(self.command));
        self.write_to_buf_address(buf);
        buf.put_slice(&CRLF);
        if self.is_padding() {
            Padding::default().write_to_buf(buf)
        }
    }
    pub fn write_to_buf_address<B: BufMut>(&self, buf: &mut B) {
        match &self.address {
            Address::SocketAddress(SocketAddr::V4(addr)) => {
                buf.put_u8(Self::ATYP_IPV4);
                buf.put_slice(&addr.ip().octets());
                buf.put_u16(addr.port());
            }
            Address::SocketAddress(SocketAddr::V6(addr)) => {
                buf.put_u8(Self::ATYP_IPV6);
                for seg in addr.ip().segments() {
                    buf.put_u16(seg);
                }
                buf.put_u16(addr.port());
            }
            Address::DomainAddress(addr, port) => {
                buf.put_u8(Self::ATYP_FQDN);
                buf.put_u8(addr.len() as u8);
                buf.put_slice(addr);
                buf.put_u16(*port);
            }
        }
    }

    /// Return the serialized length of the Trojan request header.
    pub fn serialized_len(&self) -> usize {
        56 + 2 + 1 + self.address.serialized_len() + 2
    }

    /// Whether this request uses the Trojan padding command.
    pub fn is_padding(&self) -> bool {
        Command::Padding == self.command
    }
}

/// SOCKS5 command
#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
/// Trojan command code.
pub enum Command {
    Connect,
    Bind,
    Associate,
    Padding,
}

impl Command {
    const CONNECT: u8 = 0x01;
    const BIND: u8 = 0x02;
    const ASSOCIATE: u8 = 0x03;
    const PADDING: u8 = 0x04;
}

impl TryFrom<u8> for Command {
    type Error = u8;

    fn try_from(code: u8) -> Result<Self, Self::Error> {
        match code {
            Self::CONNECT => Ok(Self::Connect),
            Self::BIND => Ok(Self::Bind),
            Self::ASSOCIATE => Ok(Self::Associate),
            Self::PADDING => Ok(Self::Padding),
            code => Err(code),
        }
    }
}

impl From<Command> for u8 {
    fn from(cmd: Command) -> Self {
        match cmd {
            Command::Connect => Command::CONNECT,
            Command::Bind => Command::BIND,
            Command::Associate => Command::ASSOCIATE,
            Command::Padding => Command::PADDING,
        }
    }
}

#[cfg(test)]
mod tests {
    use std::{
        io::Cursor,
        net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
    };

    use bytes::BytesMut;
    use socks5_proto::Address;
    use tokio::io::AsyncReadExt;

    use super::{Command, Request};
    use crate::stream::peekable::{AsyncPeek, PeekableStream};

    fn test_hash() -> String {
        "a".repeat(56)
    }

    #[tokio::test]
    async fn peek_head_reads_hash_prefix_without_consuming_stream() {
        let payload = format!("{}\r\nrest", test_hash()).into_bytes();
        let inner = Cursor::new(payload.clone());
        let mut stream = PeekableStream::new(inner);

        let head = Request::peek_head(&mut stream).await.unwrap();
        let drained = stream.drain().unwrap();
        let mut replay = Vec::new();
        stream.read_to_end(&mut replay).await.unwrap();

        assert_eq!(head, test_hash().into_bytes());
        assert_eq!(drained, head);
        assert_eq!(replay, b"\r\nrest");
    }

    #[tokio::test]
    async fn read_address_from_parses_ipv4_domain_and_ipv6() {
        let mut ipv4 = Cursor::new(vec![1, 1, 127, 0, 0, 1, 0x01, 0xbb]);
        let mut domain = Cursor::new(vec![
            4, 3, 11, b'e', b'x', b'a', b'm', b'p', b'l', b'e', b'.', b'c', b'o', b'm', 0, 80,
        ]);
        let mut ipv6 = Cursor::new(vec![
            2, 4, 0x20, 0x01, 0x0d, 0xb8, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0x01, 0xbb,
        ]);

        let (cmd4, addr4) = Request::read_address_from(&mut ipv4).await.unwrap();
        let (cmdd, addrd) = Request::read_address_from(&mut domain).await.unwrap();
        let (cmd6, addr6) = Request::read_address_from(&mut ipv6).await.unwrap();

        assert_eq!(cmd4, Command::Connect);
        assert_eq!(
            addr4,
            Address::SocketAddress(SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443))
        );
        assert_eq!(cmdd, Command::Padding);
        assert_eq!(addrd, Address::DomainAddress(b"example.com".to_vec(), 80));
        assert_eq!(cmd6, Command::Bind);
        assert_eq!(
            addr6,
            Address::SocketAddress(SocketAddr::new(
                IpAddr::V6(Ipv6Addr::new(0x2001, 0xdb8, 0, 0, 0, 0, 0, 1)),
                443
            ))
        );
    }

    #[tokio::test]
    async fn read_from_round_trips_non_padding_request() {
        let request = Request::new(
            test_hash(),
            Command::Associate,
            Address::DomainAddress(b"example.com".to_vec(), 8080),
        );
        let mut buf = BytesMut::new();
        request.write_to_buf(&mut buf);

        let parsed = Request::read_from(&mut Cursor::new(buf.to_vec()))
            .await
            .unwrap();

        assert_eq!(parsed.hash, request.hash);
        assert_eq!(parsed.command, request.command);
        assert_eq!(parsed.address, request.address);
    }

    #[tokio::test]
    async fn read_from_accepts_padding_request() {
        let request = Request::new(
            test_hash(),
            Command::Padding,
            Address::DomainAddress(b"example.com".to_vec(), 443),
        );
        let mut buf = BytesMut::new();
        request.write_to_buf(&mut buf);

        let parsed = Request::read_from(&mut Cursor::new(buf.to_vec()))
            .await
            .unwrap();

        assert_eq!(parsed.hash, request.hash);
        assert_eq!(parsed.command, Command::Padding);
        assert_eq!(parsed.address, request.address);
    }

    #[tokio::test]
    async fn read_from_rejects_short_hash_prefix() {
        let err = Request::read_from(&mut Cursor::new(vec![b'a'; 10]))
            .await
            .unwrap_err();

        assert!(err.to_string().contains("not Trojan"));
    }
}