manserve/src/main.rs

use std::{
    env, io,
    process::{exit, Command},
    string::FromUtf8Error,
    time::Duration,
};

use axum::{
    error_handling::HandleErrorLayer,
    extract::Path,
    http::{header, HeaderValue, Method, StatusCode},
    response::IntoResponse,
    routing::get,
    BoxError, Router,
};

use base64::{engine::general_purpose::STANDARD, Engine};
use tower::{buffer::BufferLayer, limit::rate::RateLimitLayer, ServiceBuilder};
use tower_http::cors::CorsLayer;

#[tokio::main]
async fn main() {
    let args: Vec<String> = env::args().collect();
    if args.len() < 2 {
        println!("Provide a port with manserve $PORT");
        exit(1);
    }
    if args[1].parse::<u16>().is_err() {
        println!("Could not parse port {}", args[1]);
        exit(1);
    }
    let listener = tokio::net::TcpListener::bind(format!("0.0.0.0:{}", args[1]))
        .await
        .unwrap();
    println!("Starting server on port {}", args[1]);
    axum::serve(listener, app()).await.unwrap();
}

fn app() -> Router {
    Router::new()
        .route("/:command", get(command))
        .layer(
            // Throttle incoming traffic
            ServiceBuilder::new()
                .layer(HandleErrorLayer::new(|err: BoxError| async move {
                    (
                        StatusCode::INTERNAL_SERVER_ERROR,
                        format!("Unhandled error: {err}"),
                    )
                }))
                .layer(BufferLayer::new(1024))
                .layer(RateLimitLayer::new(5, Duration::from_secs(1))),
        )
        .layer(
            CorsLayer::new()
                .allow_origin("*".parse::<HeaderValue>().unwrap())
                .allow_methods([Method::GET]),
        )
}

#[derive(Debug)]
enum ManserveErr {
    NotFound,
    IOError,
    InvalidUtf,
}

impl From<io::Error> for ManserveErr {
    fn from(_value: io::Error) -> Self {
        ManserveErr::IOError
    }
}
impl From<FromUtf8Error> for ManserveErr {
    fn from(_value: FromUtf8Error) -> Self {
        ManserveErr::InvalidUtf
    }
}

fn get_man_loc(command: &str) -> Result<String, ManserveErr> {
    let output = match Command::new("man").args(["-w", command]).output() {
        Ok(output) => output,
        Err(error) => {
            println!("{error}");
            return Err(error.into());
        }
    };
    if !output.status.success() {
        return Err(ManserveErr::NotFound);
    }
    let mut loc = String::from_utf8(output.stdout)?;
    loc.pop();
    Ok(loc)
}

fn generate_html(command: &str) -> Result<String, ManserveErr> {
    let man_page = get_man_loc(command)?;
    let output = Command::new("mandoc")
        .arg("-Thtml")
        .arg("-O")
        .arg("fragment")
        .arg(man_page.as_str())
        .output()?;

    let encoded = STANDARD.encode(String::from_utf8(output.stdout)?);
    Ok(encoded)
}

async fn command(Path(command): Path<String>) -> impl IntoResponse {
    let command = command.as_str();
    match generate_html(command) {
        Ok(html) => {
            let json = format!("{{\"command\": \"{command}\", \"html\": \"{html}\"}}");
            (
                StatusCode::OK,
                [(header::CONTENT_TYPE, "application/json")],
                json.into_response(),
            )
        }
        Err(error) => match error {
            ManserveErr::NotFound => {
                let json = format!("{{\"command\": \"{command}\", \"error\": \"Could not find the requested command\"}}");
                (
                    StatusCode::NOT_FOUND,
                    [(header::CONTENT_TYPE, "application/json")],
                    json.into_response(),
                )
            }
            ManserveErr::IOError => {
                let json = format!("{{\"command\": \"{command}\", \"error\": \"The server could not complete the request\"}}");
                (
                    StatusCode::INTERNAL_SERVER_ERROR,
                    [(header::CONTENT_TYPE, "application/json")],
                    json.into_response(),
                )
            }
            ManserveErr::InvalidUtf => {
                let json = format!("{{\"command\": \"{command}\", \"error\": \"There was invalid utf-8 in the request\"}}");
                (
                    StatusCode::BAD_REQUEST,
                    [(header::CONTENT_TYPE, "application/json")],
                    json.into_response(),
                )
            }
        },
    }
}

#[cfg(test)]
mod tests {
    use axum::{body::Body, http::Request};
    use http_body_util::BodyExt;
    use serde::Deserialize;
    use tower::Service;

    use super::*;

    #[tokio::test]
    async fn json() {
        let mut app = app();
        let res = app
            .call(Request::builder().uri("/ls").body(Body::default()).unwrap())
            .await
            .unwrap();
        assert_eq!(res.status(), StatusCode::OK)
    }

    #[tokio::test]
    async fn bad_command() {
        let mut app = app();
        let res = app
            .call(
                Request::builder()
                    .uri("/lollll")
                    .body(Body::default())
                    .unwrap(),
            )
            .await
            .unwrap();
        assert_eq!(res.status(), StatusCode::NOT_FOUND)
    }

    #[tokio::test]
    async fn man_ls() {
        #[derive(Deserialize, Debug)]
        struct Manpage {
            command: String,
            html: String,
        }
        let mut app = app();
        let res = app
            .call(Request::builder().uri("/ls").body(Body::default()).unwrap())
            .await
            .unwrap();

        let expected = generate_html("ls").unwrap();
        let res_str =
            String::from_utf8(res.into_body().collect().await.unwrap().to_bytes().into()).unwrap();
        let manpage: Manpage = serde_json::from_str(&res_str).unwrap();

        assert_eq!(manpage.html, expected)
    }
}