borer-core 0.5.9

use std::{str::Split, sync::Arc};

use anyhow::{Context, anyhow};
use log::debug;
use socks5_proto::Address;
use tokio::io::{AsyncRead, AsyncReadExt};

const MAX_HTTP_REQUEST_SIZE: usize = 16384;
const BAD_REQUEST: &str = "BadRequest";

/// Parsed CONNECT-style or proxied HTTP request metadata.
pub struct HttpConnectRequest {
    addr: Address,
    mode: RequestMode,
}

#[derive(Clone, Debug, Eq, PartialEq)]
enum RequestMode {
    Connect,
    Forward(Nugget),
}

#[derive(Eq, PartialEq, Debug, Clone)]
/// Preserved raw HTTP request bytes that should be forwarded upstream.
pub struct Nugget {
    data: Arc<Vec<u8>>,
}

/// Read bytes until the end of the HTTP header block (`\r\n\r\n`).
pub async fn read_http_request_end<T: AsyncRead + Unpin>(r: &mut T) -> anyhow::Result<Vec<u8>> {
    let mut buf = Vec::new();
    for _i in 0..MAX_HTTP_REQUEST_SIZE {
        let u1 = r.read_u8().await?;
        buf.push(u1);
        if u1 == b'\r' {
            let [u2, u3, u4] = {
                let mut x = [0u8; 3];
                r.read_exact(&mut x).await.map(|_| x)
            }?;
            buf.push(u2);
            buf.push(u3);
            buf.push(u4);
            if u2 == b'\n' && u3 == b'\r' && u4 == b'\n' {
                break;
            }
        }
    }
    Ok(buf)
}

impl HttpConnectRequest {
    /// Parse an HTTP request into a target address and optional forwardable payload.
    pub fn parse(http_request: &[u8]) -> anyhow::Result<Self> {
        Self::precondition_size(http_request)?;
        Self::precondition_legal_characters(http_request)?;

        let http_request_as_string =
            String::from_utf8(http_request.to_vec()).context("contains only ASCII")?;

        let mut lines = http_request_as_string.split("\r\n");
        let request_line =
            Self::parse_request_line(lines.next().ok_or_else(|| anyhow!(BAD_REQUEST))?)?;

        let (host, mode) = match request_line.mode {
            ParsedRequestMode::Connect => (request_line.target.to_string(), RequestMode::Connect),
            ParsedRequestMode::Forward => (
                Self::extract_destination_host(&mut lines, request_line.target)
                    .unwrap_or_else(|| request_line.target.to_string()),
                RequestMode::Forward(Nugget::new(http_request)),
            ),
        };

        Ok(Self {
            addr: Self::host_to_address(host)?,
            mode,
        })
    }

    /// Return the upstream address extracted from this request.
    pub fn addr(&self) -> &Address {
        &self.addr
    }

    /// Return the preserved raw request when this is a forward-proxy request.
    pub fn nugget(&self) -> Option<&Nugget> {
        match &self.mode {
            RequestMode::Connect => None,
            RequestMode::Forward(nugget) => Some(nugget),
        }
    }

    fn host_to_address(host: String) -> anyhow::Result<Address> {
        let mut parts = host.rsplitn(2, ':');
        let port = parts
            .next()
            .ok_or_else(|| anyhow!("parse http target port failed"))?;
        let domain = parts
            .next()
            .ok_or_else(|| anyhow!("parse http target host failed: {host}"))?;

        if domain.is_empty() {
            Err(anyhow!("parse http target host failed: {host}"))
        } else {
            Ok(Address::DomainAddress(
                domain.as_bytes().to_vec(),
                port.parse()?,
            ))
        }
    }

    fn extract_destination_host(lines: &mut Split<&str>, endpoint: &str) -> Option<String> {
        const HOST_HEADER: &str = "host:";

        lines
            .find(|line| line.to_ascii_lowercase().starts_with(HOST_HEADER))
            .map(|line| line[HOST_HEADER.len()..].trim())
            .map(|host| {
                let mut host = String::from(host);
                if host.rfind(':').is_none() {
                    let default_port = if endpoint.to_ascii_lowercase().starts_with("https://") {
                        ":443"
                    } else {
                        ":80"
                    };
                    host.push_str(default_port);
                }
                host
            })
    }

    fn parse_request_line(request_line: &str) -> anyhow::Result<ParsedRequestLine<'_>> {
        let request_line_items = request_line.split(' ').collect::<Vec<&str>>();
        Self::precondition_well_formed(request_line, &request_line_items)?;

        let method = request_line_items[0];
        let target = request_line_items[1];
        let version = request_line_items[2];

        let mode = Self::parse_request_mode(method);
        Self::check_version(version)?;

        Ok(ParsedRequestLine { target, mode })
    }

    fn precondition_well_formed(
        request_line: &str,
        request_line_items: &[&str],
    ) -> anyhow::Result<()> {
        if request_line_items.len() != 3 {
            debug!("bad request line: `{request_line:?}`");
            Err(anyhow!(BAD_REQUEST))
        } else {
            Ok(())
        }
    }

    fn check_version(version: &str) -> anyhow::Result<()> {
        if version != "HTTP/1.1" && version != "HTTP/1.0" {
            debug!("bad version {}", version);
            Err(anyhow!(BAD_REQUEST))
        } else {
            Ok(())
        }
    }

    fn parse_request_mode(method: &str) -> ParsedRequestMode {
        if method == "CONNECT" {
            ParsedRequestMode::Connect
        } else {
            ParsedRequestMode::Forward
        }
    }

    fn precondition_legal_characters(http_request: &[u8]) -> anyhow::Result<()> {
        for b in http_request {
            match b {
                // non-ascii characters don't make sense in this context
                32..=126 | 9 | 10 | 13 => {}
                _ => {
                    debug!("bad request header. Illegal character: {:#04x}", b);
                    return Err(anyhow!(BAD_REQUEST));
                }
            }
        }
        Ok(())
    }

    fn precondition_size(http_request: &[u8]) -> anyhow::Result<()> {
        if http_request.len() >= MAX_HTTP_REQUEST_SIZE {
            debug!(
                "bad request header. Size {} exceeds limit {}",
                http_request.len(),
                MAX_HTTP_REQUEST_SIZE
            );
            Err(anyhow!(BAD_REQUEST))
        } else {
            Ok(())
        }
    }
}

struct ParsedRequestLine<'a> {
    target: &'a str,
    mode: ParsedRequestMode,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum ParsedRequestMode {
    Connect,
    Forward,
}

impl Nugget {
    /// Store an owned copy of the raw request bytes.
    pub fn new<T: Into<Vec<u8>>>(v: T) -> Self {
        Self {
            data: Arc::new(v.into()),
        }
    }

    /// Access the raw request bytes.
    pub fn data(&self) -> Arc<Vec<u8>> {
        self.data.clone()
    }
}

#[cfg(test)]
mod tests {
    use std::io::Cursor;

    use socks5_proto::Address;

    use super::{HttpConnectRequest, Nugget, read_http_request_end};

    #[tokio::test]
    async fn read_http_request_end_reads_until_double_crlf() {
        let raw = b"CONNECT example.com:443 HTTP/1.1\r\nHost: example.com:443\r\n\r\npayload";
        let mut cursor = Cursor::new(raw.as_slice());

        let actual = read_http_request_end(&mut cursor).await.unwrap();

        assert_eq!(
            actual,
            b"CONNECT example.com:443 HTTP/1.1\r\nHost: example.com:443\r\n\r\n"
        );
    }

    #[test]
    fn parse_connect_request_uses_request_target_and_has_no_nugget() {
        let raw = b"CONNECT example.com:443 HTTP/1.1\r\nHost: ignored.example.com\r\n\r\n";

        let request = HttpConnectRequest::parse(raw).unwrap();

        assert_eq!(
            request.addr().clone(),
            Address::DomainAddress(b"example.com".to_vec(), 443)
        );
        assert!(request.nugget().is_none());
    }

    #[test]
    fn parse_non_connect_request_uses_host_header_and_preserves_request() {
        let raw = b"GET https://upstream.example/path HTTP/1.1\r\nHost: service.internal\r\n\r\n";

        let request = HttpConnectRequest::parse(raw).unwrap();

        assert_eq!(
            request.addr().clone(),
            Address::DomainAddress(b"service.internal".to_vec(), 443)
        );
        assert_eq!(
            request.nugget().cloned().unwrap(),
            Nugget::new(raw.as_slice())
        );
    }

    #[test]
    fn parse_non_connect_without_host_port_adds_http_default_port() {
        let raw = b"GET http://upstream.example/path HTTP/1.1\r\nHost: service.internal\r\n\r\n";

        let request = HttpConnectRequest::parse(raw).unwrap();

        assert_eq!(
            request.addr().clone(),
            Address::DomainAddress(b"service.internal".to_vec(), 80)
        );
    }

    #[test]
    fn parse_rejects_non_ascii_bytes() {
        let raw = b"CONNECT example.com:443 HTTP/1.1\r\nHost: examp\x01e.com\r\n\r\n";

        let err = HttpConnectRequest::parse(raw).err().expect("should reject");

        assert!(err.to_string().contains("BadRequest"));
    }

    #[test]
    fn parse_accepts_http_1_0_connect_request() {
        let raw = b"CONNECT example.com:443 HTTP/1.0\r\nHost: example.com:443\r\n\r\n";

        let request = HttpConnectRequest::parse(raw).unwrap();

        assert_eq!(
            request.addr().clone(),
            Address::DomainAddress(b"example.com".to_vec(), 443)
        );
        assert!(request.nugget().is_none());
    }

    #[test]
    fn parse_rejects_invalid_http_version() {
        let raw = b"CONNECT example.com:443 HTTP/0.9\r\nHost: example.com:443\r\n\r\n";

        let err = HttpConnectRequest::parse(raw).err().expect("should reject");

        assert!(err.to_string().contains("BadRequest"));
    }
}