http-cli/src/main.rs

use clap::{Parser, ValueEnum};
use futures::stream::StreamExt;
use json::JsonValue;
use reqwest::header::HeaderMap;
use reqwest::{
    header::{HeaderName, HeaderValue, ACCEPT, CONTENT_TYPE, LOCATION},
    redirect, Body, RequestBuilder, Response, StatusCode, Url, Version,
};
use std::{env, fmt, os::fd::AsRawFd, process::exit, str::FromStr};
use tokio::io::{stdin, stdout, AsyncReadExt, AsyncWriteExt};

#[derive(Debug, Clone, Parser)]
#[command(name = "http")]
#[command(author, version, about = "HTTP(S) request utility")]
struct Params {
    /// HTTP request method
    method: RequestMethod,

    /// URL (with or without protocol)
    url: String,

    /// Request body (`-` for stdin)
    body: Option<String>,

    /// Be verbose
    #[arg(short = 'v', long)]
    verbose: bool,

    /// Follow redirects
    #[arg(short = 'n', long)]
    nofollow: bool,

    /// Print HTTP status code
    #[arg(short = 's')]
    status: bool,

    /// Additional request headers
    #[arg(short = 'H')]
    header: Option<Vec<String>>,

    /// Show response headers
    #[arg(short = 'i', long, conflicts_with = "headers")]
    include_headers: bool,

    /// Only show response headers
    #[arg(short = 'I', long, conflicts_with = "include_headers")]
    headers: bool,

    /// Never pretty-print supported content types
    #[arg(short = 'r', long)]
    raw: bool,

    /// Shortcut for setting the `Content-Type` header
    #[arg(short = 't', long)]
    content_type: Option<String>,

    /// Shortcut for setting the `Accept` header
    #[arg(short = 'a', long)]
    accept: Option<String>,

    /// Print copyright, warranty, and license information.
    #[arg(long)]
    license: bool,
}

#[derive(Debug, Copy, Clone, PartialEq, ValueEnum)]
enum RequestMethod {
    Get,
    Head,
    Post,
    Put,
    Delete,
    Options,
    Patch,
}

impl From<RequestMethod> for reqwest::Method {
    fn from(value: RequestMethod) -> reqwest::Method {
        match value {
            RequestMethod::Get => reqwest::Method::GET,
            RequestMethod::Head => reqwest::Method::HEAD,
            RequestMethod::Post => reqwest::Method::POST,
            RequestMethod::Put => reqwest::Method::PUT,
            RequestMethod::Delete => reqwest::Method::DELETE,
            RequestMethod::Options => reqwest::Method::OPTIONS,
            RequestMethod::Patch => reqwest::Method::PATCH,
        }
    }
}

#[tokio::main(flavor = "current_thread")]
async fn main() {
    if let Some(arg1) = env::args().nth(1) {
        if arg1.as_str() == "--license" {
            print!("{LICENSE}");
            exit(0);
        }
    }

    let mut params = Params::parse();
    if params.license {
        print!("{LICENSE}");
        exit(0);
    }

    if params.method == RequestMethod::Options {
        params.headers = true;
    }

    while let Some(next_url) = do_request(params.clone()).await {
        if params.verbose {
            eprintln!("Following redirect to {next_url}");
        }
        params.url = next_url;
    }
}

async fn do_request(params: Params) -> Option<String> {
    let client = ok_or_exit(
        reqwest::ClientBuilder::new()
            .redirect(redirect::Policy::none())
            .build(),
    );

    let mut builder = client.request(params.method.into(), get_url(params.url));
    if let Some(accept) = params.accept {
        builder = builder.header(ACCEPT, accept);
    }
    if let Some(content_type) = params.content_type {
        builder = builder.header(CONTENT_TYPE, content_type);
    }
    if let Some(headers) = params.header {
        builder = append_headers(builder, &headers);
    }
    if let Some(body) = params.body {
        builder = append_body(builder, body).await;
    }
    let response = ok_or_exit(builder.send().await);

    let status = response.status();
    if status.is_client_error() || status.is_server_error() {
        eprintln!("HTTP/{} {}", http_version(&response), response.status());
        exit(2);
    }
    let redirect_url = if status.is_redirection() && status != StatusCode::SEE_OTHER {
        response
            .headers()
            .get(LOCATION)
            .and_then(|location| String::from_utf8(Vec::from(location.as_bytes())).ok())
    } else {
        None
    };

    if redirect_url.is_none() {
        if params.verbose {
            println!("HTTP/{} {}", http_version(&response), response.status());
        }

        if params.include_headers || params.headers {
            print_headers(&response);
            println!();
        }

        if !params.headers {
            let pretty = !params.raw && isatty(1);
            let len = print_body(response, pretty).await;
            if len == 0 && params.verbose {
                eprintln!("Response body is empty");
            }
        }
    }

    redirect_url
}

async fn print_body(response: Response, pretty: bool) -> usize {
    if pretty && content_type_is_json(response.headers()) {
        return print_body_json(response).await;
    }

    let mut stdout = stdout();
    let stream = response.bytes_stream();
    tokio::pin!(stream);

    let mut total_bytes = 0;
    let mut ends_with_newline = true;
    while let Some(chunk) = stream.next().await {
        let chunk = ok_or_exit(chunk);
        ends_with_newline = chunk.as_ref()[chunk.len() - 1] == b'\n';
        total_bytes += chunk.len();
        ok_or_exit(stdout.write(chunk.as_ref()).await);
    }
    if pretty && !ends_with_newline {
        ok_or_exit(stdout.write(b"\x1b[1;97;101m%\x1b[0m\n").await);
    }
    total_bytes
}

fn content_type_is_json(headers: &HeaderMap) -> bool {
    if let Some(content_type) = headers.get(CONTENT_TYPE) {
        let content_type = content_type.as_bytes();
        if !content_type.starts_with(b"application/") {
            return false;
        }

        let after_plus_pos = content_type
            .iter()
            .enumerate()
            .skip(12)
            .find_map(|(i, &c)| (c == b'+').then_some(i + 1));
        let after_plus = if let Some(plus_pos) = after_plus_pos {
            &content_type[plus_pos..]
        } else {
            &content_type[12..]
        };
        after_plus.starts_with(b"json")
    } else {
        false
    }
}

async fn print_body_json(response: Response) -> usize {
    let bytes = ok_or_exit(response.bytes().await);
    let nbyte = bytes.len();
    let s = ok_or_exit(String::from_utf8(bytes.to_vec()));
    if let Ok(json) = json::parse(&s) {
        let mut s = String::with_capacity(nbyte);
        format_json(&json, &mut s, 0);
        if !s.is_empty() {
            s.push('\n');
        }
        ok_or_exit(stdout().write(s.as_bytes()).await);
    } else {
        ok_or_exit(stdout().write(s.as_bytes()).await);
    }

    nbyte
}

fn format_json(json: &JsonValue, s: &mut String, indent: usize) {
    match json {
        JsonValue::Null => *s += "\x1b[0;96mnull\x1b[0m",
        JsonValue::Boolean(b) => s.push_str(&format!("\x1b[0;96m{b}\x1b[0m")),
        JsonValue::Number(n) => s.push_str(&format!("\x1b[0;33m{n}\x1b[0m")),
        JsonValue::String(val) => s.push_str(&format!("\x1b[0;32m\"{val}\"\x1b[0m")),
        JsonValue::Short(val) => s.push_str(&format!("\x1b[0;32m\"{val}\"\x1b[0m")),
        JsonValue::Array(a) => {
            if a.is_empty() {
                s.push_str("[]");
            } else {
                s.push('[');
                for (i, e) in a.iter().enumerate() {
                    let indent = indent + 4;
                    if i > 0 {
                        s.push(',');
                    }
                    s.push('\n');
                    s.extend((0..indent).into_iter().map(|_| ' '));
                    format_json(e, s, indent);
                }
                s.push('\n');
                s.extend((0..indent).into_iter().map(|_| ' '));
                s.push(']');
            }
        }
        JsonValue::Object(o) => {
            if o.is_empty() {
                s.push_str("{}");
            } else {
                s.push('{');
                for (i, (k, v)) in o.iter().enumerate() {
                    let indent = indent + 4;
                    if i > 0 {
                        s.push(',');
                    }
                    s.push('\n');
                    s.extend((0..indent).into_iter().map(|_| ' '));
                    s.push_str(&format!("\x1b[0;36m\"{k}\"\x1b[0m: "));
                    format_json(v, s, indent);
                }
                s.push('\n');
                s.extend((0..indent).into_iter().map(|_| ' '));
                s.push('}');
            }
        }
    }
}

fn http_version(response: &Response) -> &'static str {
    match response.version() {
        Version::HTTP_09 => "0.9",
        Version::HTTP_10 => "1.0",
        Version::HTTP_11 => "1.1",
        Version::HTTP_2 => "2",
        Version::HTTP_3 => "3",
        _ => "?",
    }
}

fn print_headers(response: &Response) {
    for (name, val) in response.headers() {
        let (name, val) = (name.as_str(), val.as_bytes());
        let val = match std::str::from_utf8(val) {
            Ok(val) => val,
            Err(e) => {
                eprintln!("Error parsing header value for \"{name}\": {e}");
                "???"
            }
        };

        print_bold(name);
        println!(": {val}");
    }
}

fn get_url(url: String) -> Url {
    let url = if !url.starts_with("http://") && !url.starts_with("https://") {
        let arg0 = env::args().next();
        if arg0.as_deref() == Some("https") {
            format!("https://{url}")
        } else {
            format!("http://{url}")
        }
    } else {
        url
    };

    ok_or_exit(url.parse())
}

fn append_headers(mut builder: RequestBuilder, headers: &[String]) -> RequestBuilder {
    fn invalid_header(header: &str) -> String {
        format!("Invalid header name: \"{header}\" is not in format \"name: value\"")
    }

    for header in headers {
        let name_end = ok_or_exit(
            header
                .as_bytes()
                .iter()
                .enumerate()
                .find_map(|(index, &c)| (c == b':' || c == b'=').then_some(index))
                .ok_or_else(|| invalid_header(header)),
        );
        let name = &header[..name_end];

        let val_start = ok_or_exit(
            header
                .as_bytes()
                .iter()
                .enumerate()
                .skip(name_end)
                .find_map(|(index, &c)| (!c.is_ascii_whitespace()).then_some(index))
                .ok_or_else(|| invalid_header(header)),
        );
        let val = &header[val_start..];

        let name = ok_or_exit(HeaderName::from_str(name));
        let val = ok_or_exit(HeaderValue::from_str(val));
        builder = builder.header(name, val);
    }

    builder
}

async fn append_body(builder: RequestBuilder, body: String) -> RequestBuilder {
    let body = if body == "-" {
        eprintln!("Reading body from stdin");
        body_from_stdin().await
    } else {
        Body::from(body)
    };

    builder.body(body)
}

async fn body_from_stdin() -> Body {
    let mut buf = Vec::new();
    ok_or_exit(stdin().read_to_end(&mut buf).await);
    Body::from(buf)
}

fn ok_or_exit<T, E>(result: Result<T, E>) -> T
where
    E: fmt::Display,
{
    match result {
        Ok(val) => val,
        Err(e) => {
            eprintln!("{e}");
            exit(1);
        }
    }
}

fn print_bold(s: &str) {
    let stdout = stdout();
    if isatty(stdout.as_raw_fd()) {
        print!("\x1b[1m{s}\x1b[0m");
    } else {
        print!("{s}")
    }
}

#[cfg(unix)]
fn isatty(fd: i32) -> bool {
    let result = unsafe { libc::isatty(fd) };
    result == 1
}

#[cfg(not(unix))]
fn isatty(_: i32) -> bool {
    false
}

const LICENSE: &str = include_str!("../LICENSE");