wcgi-host 0.2.0

use std::collections::HashMap;

use http::{header::HeaderName, HeaderMap, HeaderValue, StatusCode};
use tokio::io::{AsyncBufRead, AsyncBufReadExt};

use crate::CgiError;

const SERVER_SOFTWARE: &str = concat!(env!("CARGO_PKG_NAME"), " ", env!("CARGO_PKG_VERSION"));

/// The RFC says certain headers that should be mapped to environment variables,
/// if present.
const KNOWN_META_VARIABLES: [(&str, &str); 3] = [
    ("Content-Length", "CONTENT_LENGTH"),
    ("Content-Type", "CONTENT_TYPE"),
    ("Authorization", "AUTH_TYPE"),
];

pub(crate) fn prepare_environment_variables(
    parts: http::request::Parts,
    env: &mut HashMap<String, String>,
) {
    env.insert(
        "REQUEST_METHOD".to_string(),
        parts.method.as_str().to_string(),
    );
    env.insert(
        "QUERY_STRING".to_string(),
        parts.uri.query().unwrap_or("").to_string(),
    );
    env.insert("PATH_INFO".to_string(), parts.uri.path().to_string());
    env.insert("SERVER_SOFTWARE".to_string(), SERVER_SOFTWARE.to_string());

    // FIXME(Michael-F-Bryan): we hard-code the assumption that our CGI files
    // were mounted under /app/. This should be configurable.
    // https://github.com/wasmerio/wcgi/issues/21
    env.insert("DOCUMENT_ROOT".to_string(), "/app/".to_string());
    env.insert(
        "SCRIPT_FILENAME".to_string(),
        format!("/app{}", parts.uri.path()),
    );

    if let Some(protocol) = server_protocol(parts.version) {
        env.insert("SERVER_PROTOCOL".to_string(), protocol.to_string());
    }

    if let Some(content_length) = parts
        .headers
        .get("Content-Length")
        .and_then(|v| v.to_str().ok())
    {
        env.insert("CONTENT_LENGTH".to_string(), content_length.to_string());
    }
    for (header_name, env_variable) in KNOWN_META_VARIABLES {
        if let Some(value) = parts.headers.get(header_name).and_then(|v| v.to_str().ok()) {
            env.insert(env_variable.to_string(), value.to_string());
        }
    }

    // All "protocol-specific" HTTP sure headers are passed in as $HTTP_xxx
    for (name, value) in &parts.headers {
        if let Ok(value) = value.to_str() {
            let name = format!("HTTP_{}", name.to_string().to_uppercase().replace('-', "_"));
            env.insert(name, value.to_string());
        }
    }
}

pub(crate) async fn extract_response_header(
    stdout: &mut (impl AsyncBufRead + Unpin),
) -> Result<http::response::Parts, CgiError> {
    let mut headers = parse_cgi_headers(stdout).await?;

    let mut builder = http::response::Builder::new();

    let status = headers.remove("Status").and_then(|status| {
        let status = status.to_str().ok()?;
        parse_status_header(status)
    });

    if let Some(status) = status {
        builder = builder.status(status);
    }

    // Note: This can only panic when the TryInto calls used in the builder's
    // various builder methods fail. However, this should never happen because
    // we parse the headers and status code ourselves.
    //
    // Don't look too closely or you'll notice that we call into_parts() just so
    // we can get a Parts object that the caller can pass to
    // Response::from_parts() later on.
    let (mut parts, _) = builder
        .body(())
        .expect("All builder inputs should already be validated")
        .into_parts();

    parts.headers.extend(headers);

    Ok(parts)
}

/// Parse the header section from the response.
///
/// # Implementation Notes
///
/// Any invalid headers will be silently discarded. This might happen if lines
/// aren't in the form `name: value` or when header name/value contains invalid
/// characters.
async fn parse_cgi_headers(
    stdout_receiver: &mut (impl AsyncBufRead + Unpin),
) -> Result<HeaderMap, CgiError> {
    let mut headers = HeaderMap::new();
    let mut buffer = String::new();

    loop {
        buffer.clear();
        stdout_receiver
            .read_line(&mut buffer)
            .await
            .map_err(CgiError::StdoutRead)?;

        let line = buffer.trim();

        if line.is_empty() {
            // We found the CRLF CRLF that indicates the end of the header
            // section.
            break;
        }

        let (key, value) = match line.split_once(':') {
            Some((k, v)) => (k, v),
            None => {
                // Let's be lenient and ignore the invalid header line.
                continue;
            }
        };

        if let Some((key, value)) = parse_header_pair(key, value) {
            headers.append(key, value);
        }
    }

    Ok(headers)
}

fn parse_header_pair(key: &str, value: &str) -> Option<(HeaderName, HeaderValue)> {
    let key = HeaderName::from_bytes(key.trim().as_bytes()).ok()?;
    let value = value.trim().parse().ok()?;

    Some((key, value))
}

fn parse_status_header(status_code: &str) -> Option<StatusCode> {
    // Note: the status code header may contain just the number (i.e. "200") or
    // also the reason ("200 OK").
    let src = match status_code.split_once(' ') {
        Some((s, _)) => s,
        None => status_code,
    };

    src.parse().ok()
}

fn server_protocol(version: http::Version) -> Option<&'static str> {
    match version {
        http::Version::HTTP_09 => Some("HTTP/0.9"),
        http::Version::HTTP_10 => Some("HTTP/1.0"),
        http::Version::HTTP_11 => Some("HTTP/1.1"),
        http::Version::HTTP_2 => Some("HTTP/2.0"),
        http::Version::HTTP_3 => Some("HTTP/3.0"),
        _ => None,
    }
}

#[cfg(test)]
mod tests {
    use tokio::io::AsyncReadExt;

    use super::*;

    #[test]
    fn parse_status_codes() {
        let codes = [
            ("200", Some(StatusCode::OK)),
            ("200 OK", Some(StatusCode::OK)),
            ("", None),
        ];

        for (src, expected) in codes {
            let got = parse_status_header(src);
            assert_eq!(got, expected);
        }
    }

    #[tokio::test]
    async fn parse_response_parts() {
        let src = [
            "Status: 503 Database Unavailable",
            "Content-type: text/html",
            "",
            "<HTML>",
            "<HEAD><TITLE>503 Database Unavailable</TITLE></HEAD>",
            "<BODY>",
            "  <H1>Error</H1>",
            "  <P>Sorry, the database is currently not available. Please",
            "    try again later.</P>",
            "</BODY>",
            "</HTML>",
        ]
        .join("\r\n");
        let mut reader = tokio::io::BufReader::new(src.as_bytes());

        let parts = extract_response_header(&mut reader).await.unwrap();

        assert_eq!(parts.status, StatusCode::SERVICE_UNAVAILABLE);
        assert_eq!(parts.headers["Content-type"].to_str().unwrap(), "text/html");
        // Make sure we stopped reading at the \r\n\r\n
        let mut body = String::new();
        reader.read_to_string(&mut body).await.unwrap();
        assert!(body.starts_with("<HTML>\r\n<HEAD><TITLE>503 Database Unavailable</TITLE></HEAD>"));
    }

    #[tokio::test]
    async fn respect_duplicate_headers() {
        let src = [
            "Status: 503 Database Unavailable",
            "Content-type: text/html",
            "Cookie: first=x",
            "Cookie: second=y",
        ]
        .join("\r\n");
        let mut reader = tokio::io::BufReader::new(src.as_bytes());

        let parts = extract_response_header(&mut reader).await.unwrap();

        let cookies: Vec<_> = parts
            .headers
            .get_all("Cookie")
            .iter()
            .map(|v| v.to_str().unwrap())
            .collect();
        assert_eq!(cookies, ["first=x", "second=y"]);
    }

    #[tokio::test]
    async fn parse_response_parts_with_empty_body() {
        let src = "Location: https://google.com/\r\n";
        let mut reader = tokio::io::BufReader::new(src.as_bytes());

        let parts = extract_response_header(&mut reader).await.unwrap();

        assert_eq!(
            parts.headers["Location"].to_str().unwrap(),
            "https://google.com/"
        );
        // Make sure we stopped reading at the \r\n\r\n
        let mut body = String::new();
        reader.read_to_string(&mut body).await.unwrap();
        assert!(body.is_empty());
    }
}