cgi_service/

lib.rs

//! A Tower service that implements the CGI protocol (RFC 3875).
//!
//! # Example
//!
//! ```no_run
//! use axum::{Router, routing::any_service};
//! use cgi_service::CgiService;
//!
//! let app: Router = Router::new().route(
//!     "/",
//!     any_service(CgiService::new("/usr/lib/cgi-bin/script")),
//! );
//! ```

use std::collections::HashMap;
use std::convert::Infallible;
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};

use bytes::Bytes;
use http::{Request, Response, StatusCode};
use http_body::{Body, Frame};
use http_body_util::{BodyExt, Full, combinators::BoxBody};
use pin_project_lite::pin_project;
use thiserror::Error;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::process::{Child, ChildStdin, ChildStdout, Command};
use tokio::sync::mpsc;
use tower_service::Service;

#[derive(Error, Debug)]
pub enum CgiError {
    #[error("Failed to spawn CGI process: {0}")]
    SpawnError(#[from] std::io::Error),

    #[error("CGI process returned non-zero exit code: {0}")]
    ProcessError(i32),

    #[error("Failed to parse CGI response: {0}")]
    ParseError(String),

    #[error("Invalid CGI response header: {0}")]
    InvalidHeader(String),

    #[error("CGI process was killed by signal")]
    ProcessKilled,

    #[error("CGI process timed out after {0}ms")]
    Timeout(u64),
}

impl CgiError {
    fn into_response(self) -> Response<BoxBody<Bytes, Infallible>> {
        Response::builder()
            .status(StatusCode::INTERNAL_SERVER_ERROR)
            .body(BoxBody::new(
                Full::new(Bytes::from(self.to_string())).map_err(|_: Infallible| unreachable!()),
            ))
            .unwrap()
    }
}

/// Maximum size of CGI response headers before we give up
const MAX_HEADER_SIZE: usize = 64 * 1024;
/// Chunk size for reading response body
const BODY_CHUNK_SIZE: usize = 8 * 1024;
/// Buffer size for the response body channel
const CHANNEL_BUFFER_SIZE: usize = 16;

pin_project! {
    /// A streaming body that reads from a CGI process stdout via a channel.
    struct CgiStreamBody {
        rx: mpsc::Receiver<Result<Bytes, std::io::Error>>,
        done: bool,
    }
}

impl CgiStreamBody {
    fn new(rx: mpsc::Receiver<Result<Bytes, std::io::Error>>) -> Self {
        Self { rx, done: false }
    }
}

impl Body for CgiStreamBody {
    type Data = Bytes;
    type Error = Infallible;

    fn poll_frame(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
        let this = self.project();

        if *this.done {
            return Poll::Ready(None);
        }

        match this.rx.poll_recv(cx) {
            Poll::Ready(Some(Ok(bytes))) if !bytes.is_empty() => {
                Poll::Ready(Some(Ok(Frame::data(bytes))))
            }
            Poll::Pending => Poll::Pending,
            _ => {
                *this.done = true;
                Poll::Ready(None)
            }
        }
    }

    fn is_end_stream(&self) -> bool {
        self.done
    }
}

/// Configuration for a CGI service.
///
/// Use the builder methods to customize the CGI execution environment.
#[derive(Debug, Clone)]
pub struct CgiConfig {
    command: String,
    args: Vec<String>,
    working_dir: Option<String>,
    extra_env: HashMap<String, String>,
    server_software: String,
    server_name: String,
    server_port: u16,
    timeout_ms: Option<u64>,
    inherit_env: Vec<String>,
    script_name: Option<String>,
    document_root: Option<String>,
    pass_through_stderr: bool,
    https: bool,
}

impl CgiConfig {
    /// Creates a new configuration with the given command path.
    pub fn new(command: impl Into<String>) -> Self {
        Self {
            command: command.into(),
            args: Vec::new(),
            working_dir: None,
            extra_env: HashMap::new(),
            server_software: "rust-cgi-server/0.1".to_string(),
            server_name: "localhost".to_string(),
            server_port: 80,
            timeout_ms: None,
            inherit_env: Vec::new(),
            script_name: None,
            document_root: None,
            pass_through_stderr: true,
            https: false,
        }
    }

    /// Sets command-line arguments to pass to the CGI script.
    pub fn args(mut self, args: Vec<String>) -> Self {
        self.args = args;
        self
    }

    /// Sets the working directory for the CGI process.
    pub fn working_dir(mut self, dir: impl Into<String>) -> Self {
        self.working_dir = Some(dir.into());
        self
    }

    /// Adds an environment variable to pass to the CGI process.
    pub fn env(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
        self.extra_env.insert(key.into(), value.into());
        self
    }

    /// Sets the SERVER_SOFTWARE meta-variable.
    pub fn server_software(mut self, software: impl Into<String>) -> Self {
        self.server_software = software.into();
        self
    }

    /// Sets the SERVER_NAME meta-variable.
    pub fn server_name(mut self, name: impl Into<String>) -> Self {
        self.server_name = name.into();
        self
    }

    /// Sets the SERVER_PORT meta-variable.
    pub fn server_port(mut self, port: u16) -> Self {
        self.server_port = port;
        self
    }

    /// Set whether the request was received over HTTPS.
    /// When true, sets the HTTPS environment variable to "on".
    pub fn https(mut self, https: bool) -> Self {
        self.https = https;
        self
    }

    /// Sets a timeout for CGI script execution.
    pub fn timeout(mut self, timeout: std::time::Duration) -> Self {
        self.timeout_ms = Some(timeout.as_millis() as u64);
        self
    }

    /// Sets environment variables to inherit from the parent process.
    pub fn inherit_env(mut self, vars: Vec<String>) -> Self {
        self.inherit_env = vars;
        self
    }

    /// Adds a single environment variable to inherit from the parent process.
    pub fn inherit(mut self, var: impl Into<String>) -> Self {
        self.inherit_env.push(var.into());
        self
    }

    /// When set, PATH_INFO will be computed as the request path minus this script name.
    pub fn script_name(mut self, name: impl Into<String>) -> Self {
        self.script_name = Some(name.into());
        self
    }

    /// Set the document root for computing PATH_TRANSLATED.
    /// When PATH_INFO is non-empty, PATH_TRANSLATED = document_root + PATH_INFO.
    pub fn document_root(mut self, root: impl Into<String>) -> Self {
        self.document_root = Some(root.into());
        self
    }

    /// Sets whether to pass through CGI process stderr to the parent process.
    pub fn pass_through_stderr(mut self, pass_through: bool) -> Self {
        self.pass_through_stderr = pass_through;
        self
    }
}

fn build_cgi_env<B>(
    request: &Request<B>,
    config: &CgiConfig,
    remote_addr: Option<&str>,
) -> HashMap<String, String> {
    let mut env = HashMap::new();

    // Required CGI meta-variables (RFC 3875 Section 4.1)
    env.insert("GATEWAY_INTERFACE".to_string(), "CGI/1.1".to_string());
    env.insert("REQUEST_METHOD".to_string(), request.method().to_string());
    env.insert("SERVER_NAME".to_string(), config.server_name.clone());
    env.insert("SERVER_PORT".to_string(), config.server_port.to_string());
    env.insert(
        "SERVER_PROTOCOL".to_string(),
        format!("{:?}", request.version()),
    );
    env.insert(
        "SERVER_SOFTWARE".to_string(),
        config.server_software.clone(),
    );

    let path = request.uri().path();
    let (script_name, path_info) = if let Some(ref sn) = config.script_name {
        let path_info = if path.starts_with(sn) {
            path[sn.len()..].to_string()
        } else {
            String::new()
        };
        (sn.clone(), path_info)
    } else {
        (path.to_string(), String::new())
    };

    env.insert("SCRIPT_NAME".to_string(), script_name);
    env.insert("PATH_INFO".to_string(), path_info.clone());

    // PATH_TRANSLATED: physical path corresponding to PATH_INFO (RFC 3875 Section 4.1.6)
    if !path_info.is_empty()
        && let Some(ref doc_root) = config.document_root
    {
        let path_translated = format!("{}{}", doc_root, path_info);
        env.insert("PATH_TRANSLATED".to_string(), path_translated);
    }

    // Query string (empty string if none, per spec)
    env.insert(
        "QUERY_STRING".to_string(),
        request.uri().query().unwrap_or("").to_string(),
    );

    env.insert(
        "REQUEST_URI".to_string(),
        request
            .uri()
            .path_and_query()
            .map(|pq| pq.to_string())
            .unwrap_or_else(|| path.to_string()),
    );

    // REMOTE_ADDR is required per RFC 3875 Section 4.1.8
    // Default to 127.0.0.1 if not explicitly set
    env.insert(
        "REMOTE_ADDR".to_string(),
        remote_addr.unwrap_or("127.0.0.1").to_string(),
    );

    // HTTPS meta-variable indicates secure connection
    if config.https {
        env.insert("HTTPS".to_string(), "on".to_string());
    }

    if let Some(content_type) = request.headers().get(http::header::CONTENT_TYPE)
        && let Ok(ct) = content_type.to_str()
    {
        env.insert("CONTENT_TYPE".to_string(), ct.to_string());
    }

    if let Some(content_length) = request.headers().get(http::header::CONTENT_LENGTH)
        && let Ok(cl) = content_length.to_str()
    {
        env.insert("CONTENT_LENGTH".to_string(), cl.to_string());
    }

    for var_name in &config.inherit_env {
        if let Ok(value) = std::env::var(var_name) {
            env.insert(var_name.clone(), value);
        }
    }

    // Convert HTTP headers to HTTP_* environment variables
    // RFC 3875 Section 4.1.18: Exclude certain headers for security/protocol reasons
    for (name, value) in request.headers() {
        // Skip Content-Type and Content-Length (handled separately as CGI vars)
        if name == http::header::CONTENT_TYPE || name == http::header::CONTENT_LENGTH {
            continue;
        }

        // Skip Authorization headers for security (RFC 3875 recommendation)
        if name == http::header::AUTHORIZATION || name == http::header::PROXY_AUTHORIZATION {
            continue;
        }

        // Skip hop-by-hop headers that shouldn't be passed to CGI
        if name == http::header::CONNECTION
            || name.as_str().eq_ignore_ascii_case("keep-alive")
            || name == http::header::TRANSFER_ENCODING
            || name == http::header::TE
            || name == http::header::TRAILER
            || name == http::header::UPGRADE
        {
            continue;
        }

        let env_name = format!("HTTP_{}", name.as_str().to_uppercase().replace('-', "_"));

        if let Ok(v) = value.to_str() {
            env.insert(env_name, v.to_string());
        }
    }

    for (key, value) in &config.extra_env {
        env.insert(key.clone(), value.clone());
    }

    env
}

fn find_header_body_separator(data: &[u8]) -> Option<usize> {
    for i in 0..data.len() {
        if i + 3 < data.len() && &data[i..i + 4] == b"\r\n\r\n" {
            return Some(i);
        }
        if i + 1 < data.len() && &data[i..i + 2] == b"\n\n" {
            return Some(i);
        }
    }
    None
}

/// Detect the length of the header/body separator at the given position
fn separator_length_at(data: &[u8], pos: usize) -> usize {
    if data[pos..].starts_with(b"\r\n\r\n") {
        4
    } else {
        2 // \n\n or fallback
    }
}

/// Parsed CGI headers
struct ParsedHeaders {
    status: StatusCode,
    headers: http::HeaderMap,
}

/// Parse CGI headers from raw bytes (without body)
fn parse_cgi_headers(header_bytes: &[u8]) -> Result<ParsedHeaders, CgiError> {
    let header_str = std::str::from_utf8(header_bytes)
        .map_err(|e| CgiError::ParseError(format!("Invalid UTF-8 in headers: {}", e)))?;

    let mut status = StatusCode::OK;
    let mut headers = http::HeaderMap::new();

    for line in header_str.lines() {
        let line = line.trim();
        if line.is_empty() {
            continue;
        }

        if let Some((name, value)) = line.split_once(':') {
            let name = name.trim();
            let value = value.trim();

            if name.eq_ignore_ascii_case("Status") {
                let code_str = value.split_once(' ').map_or(value, |(code, _)| code);
                status = code_str
                    .parse::<u16>()
                    .map_err(|_| {
                        CgiError::InvalidHeader(format!("Invalid status code: {}", code_str))
                    })
                    .and_then(|code| {
                        StatusCode::from_u16(code).map_err(|_| {
                            CgiError::InvalidHeader(format!("Invalid status code: {}", code))
                        })
                    })?;
            } else {
                let header_name = http::header::HeaderName::from_bytes(name.as_bytes())
                    .map_err(|e| CgiError::InvalidHeader(format!("Invalid header name: {}", e)))?;
                let header_value = http::header::HeaderValue::from_str(value)
                    .map_err(|e| CgiError::InvalidHeader(format!("Invalid header value: {}", e)))?;
                headers.append(header_name, header_value);
            }
        }
    }

    Ok(ParsedHeaders { status, headers })
}

/// Write request body chunks to CGI process stdin
async fn write_request_body<B>(mut stdin: ChildStdin, body: B, timeout: Option<std::time::Duration>)
where
    B: Body + Send,
    B::Data: Send,
    B::Error: std::fmt::Debug,
{
    use bytes::Buf;

    let write_future = async {
        let mut body = std::pin::pin!(body);
        loop {
            match std::future::poll_fn(|cx| body.as_mut().poll_frame(cx)).await {
                Some(Ok(frame)) => {
                    if let Ok(data) = frame.into_data()
                        && let Err(e) = stdin.write_all(data.chunk()).await
                    {
                        // Process may have closed stdin early - this is OK
                        // Some CGI scripts don't read the full body
                        if e.kind() != std::io::ErrorKind::BrokenPipe {
                            eprintln!("CGI stdin write error: {}", e);
                        }
                        break;
                    }
                }
                Some(Err(e)) => {
                    eprintln!("Error reading request body: {:?}", e);
                    break;
                }
                None => break,
            }
        }
        // Drop stdin to signal EOF to the CGI process
        drop(stdin);
    };

    if let Some(timeout_duration) = timeout {
        if tokio::time::timeout(timeout_duration, write_future)
            .await
            .is_err()
        {
            eprintln!("Timeout writing request body to CGI stdin");
        }
    } else {
        write_future.await;
    }
}

/// Read stdout until headers are complete, then return parsed headers and any body bytes read.
/// Also takes the child process to check exit status if no output is received.
async fn read_and_parse_headers(
    stdout: &mut ChildStdout,
    child: &mut Child,
    timeout: Option<std::time::Duration>,
) -> Result<(ParsedHeaders, Vec<u8>), CgiError> {
    let read_future = async {
        let mut header_buf = Vec::with_capacity(4096);
        let mut temp_buf = [0u8; 1024];

        loop {
            let n = stdout
                .read(&mut temp_buf)
                .await
                .map_err(CgiError::SpawnError)?;

            if n == 0 {
                // EOF before headers complete
                // If no data was received at all, check if the process failed
                if header_buf.is_empty() {
                    // Try to get process exit status
                    if let Ok(Some(status)) = child.try_wait()
                        && !status.success()
                    {
                        return Err(status
                            .code()
                            .map_or(CgiError::ProcessKilled, CgiError::ProcessError));
                    }
                    // Process exited successfully but produced no output
                    return Err(CgiError::ParseError(
                        "CGI script produced no output".to_string(),
                    ));
                }
                // Some data was received but no header separator - treat as headers only
                let parsed = parse_cgi_headers(&header_buf)?;
                return Ok((parsed, Vec::new()));
            }

            header_buf.extend_from_slice(&temp_buf[..n]);

            if let Some(pos) = find_header_body_separator(&header_buf) {
                let sep_len = separator_length_at(&header_buf, pos);
                let body_start = pos + sep_len;

                let parsed = parse_cgi_headers(&header_buf[..pos])?;
                let initial_body = header_buf[body_start..].to_vec();

                return Ok((parsed, initial_body));
            }

            // Prevent unbounded header growth
            if header_buf.len() > MAX_HEADER_SIZE {
                return Err(CgiError::ParseError(format!(
                    "Headers exceed maximum size of {} bytes",
                    MAX_HEADER_SIZE
                )));
            }
        }
    };

    if let Some(timeout_duration) = timeout {
        tokio::time::timeout(timeout_duration, read_future)
            .await
            .map_err(|_| CgiError::Timeout(timeout_duration.as_millis() as u64))?
    } else {
        read_future.await
    }
}

/// Stream stdout to the response body channel
async fn stream_stdout(
    mut stdout: ChildStdout,
    tx: mpsc::Sender<Result<Bytes, std::io::Error>>,
    initial_body: Vec<u8>,
    mut child: Child,
) {
    if !initial_body.is_empty() && tx.send(Ok(Bytes::from(initial_body))).await.is_err() {
        return;
    }

    let mut buf = vec![0u8; BODY_CHUNK_SIZE];
    loop {
        match stdout.read(&mut buf).await {
            Ok(0) => break,
            Ok(n) => {
                let chunk = Bytes::copy_from_slice(&buf[..n]);
                if tx.send(Ok(chunk)).await.is_err() {
                    break;
                }
            }
            Err(e) => {
                let _ = tx.send(Err(e)).await;
                break;
            }
        }
    }

    drop(stdout);
    drop(tx);

    match child.wait().await {
        Ok(status) if !status.success() => match status.code() {
            Some(code) => eprintln!("CGI process exited with code: {code}"),
            None => eprintln!("CGI process was killed by signal"),
        },
        Err(e) => eprintln!("Error waiting for CGI process: {e}"),
        _ => {}
    }
}

/// Execute CGI script with streaming request and response bodies
async fn execute_cgi<B>(
    config: &CgiConfig,
    env: HashMap<String, String>,
    body: B,
) -> Result<Response<BoxBody<Bytes, Infallible>>, CgiError>
where
    B: Body + Send + 'static,
    B::Data: Send,
    B::Error: std::fmt::Debug + Send,
{
    use std::time::Duration;

    let mut cmd = Command::new(&config.command);

    if !config.args.is_empty() {
        cmd.args(&config.args);
    }

    if let Some(ref dir) = config.working_dir {
        cmd.current_dir(dir);
    }

    cmd.env_clear();
    for (key, value) in &env {
        cmd.env(key, value);
    }

    cmd.stdin(std::process::Stdio::piped());
    cmd.stdout(std::process::Stdio::piped());
    cmd.stderr(std::process::Stdio::inherit()); // Pass through stderr directly

    let mut child = cmd.spawn()?;
    let timeout = config.timeout_ms.map(Duration::from_millis);

    let stdin = child.stdin.take().unwrap();
    tokio::spawn(write_request_body(stdin, body, timeout));

    let mut stdout = child.stdout.take().unwrap();

    // Pass child reference to check exit status if no output is received
    let (parsed_headers, initial_body) =
        read_and_parse_headers(&mut stdout, &mut child, timeout).await?;

    let (tx, rx) = mpsc::channel(CHANNEL_BUFFER_SIZE);
    tokio::spawn(stream_stdout(stdout, tx, initial_body, child));

    let mut response_builder = Response::builder().status(parsed_headers.status);
    for (name, value) in parsed_headers.headers.iter() {
        response_builder = response_builder.header(name, value);
    }

    response_builder
        .body(BoxBody::new(CgiStreamBody::new(rx)))
        .map_err(|e| CgiError::ParseError(format!("Failed to build response: {}", e)))
}

/// A Tower service that executes CGI scripts.
///
/// This service implements [RFC 3875](https://www.rfc-editor.org/rfc/rfc3875) and can be used
/// with any Tower-compatible HTTP server such as Axum or Hyper.
#[derive(Clone)]
pub struct CgiService {
    config: CgiConfig,
    remote_addr: Option<String>,
}

impl CgiService {
    /// Creates a new CGI service with the given command path.
    pub fn new(command: impl Into<String>) -> Self {
        Self {
            config: CgiConfig::new(command),
            remote_addr: None,
        }
    }

    /// Creates a new CGI service with the given configuration.
    pub fn with_config(config: CgiConfig) -> Self {
        Self {
            config,
            remote_addr: None,
        }
    }

    /// Sets the REMOTE_ADDR meta-variable for requests handled by this service.
    pub fn remote_addr(mut self, addr: impl Into<String>) -> Self {
        self.remote_addr = Some(addr.into());
        self
    }
}

impl<B> Service<Request<B>> for CgiService
where
    B: Body + Send + 'static,
    B::Data: Send,
    B::Error: std::fmt::Debug + Send,
{
    type Response = Response<BoxBody<Bytes, Infallible>>;
    type Error = Infallible;
    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;

    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        Poll::Ready(Ok(()))
    }

    fn call(&mut self, request: Request<B>) -> Self::Future {
        let config = self.config.clone();
        let remote_addr = self.remote_addr.clone();

        Box::pin(async move {
            let env = build_cgi_env(&request, &config, remote_addr.as_deref());
            let (_parts, body) = request.into_parts();
            Ok(execute_cgi(&config, env, body)
                .await
                .unwrap_or_else(|e| e.into_response()))
        })
    }
}