yggdrasil/userspace/netutils/src/http.rs

use std::{
    fs::File,
    io::{self, stdout, Stdout, Write},
    path::{Path, PathBuf},
    process::ExitCode,
    str::FromStr,
};

use clap::{Parser, Subcommand};
use hclient::{
    connection::tls, AutoConnector, Bytes, HttpBody, HttpClient, HttpConnection, HttpConnector,
    HttpError, HttpResponse,
};
use http::{
    header::{self, ToStrError},
    Method, Uri,
};

#[derive(Debug, thiserror::Error)]
enum Error {
    #[error("I/O error: {0}")]
    IoError(#[from] io::Error),
    #[error("HTTP error: {0}")]
    HttpError(#[from] HttpError<tls::Error>),
    #[error("Too many redirects")]
    TooManyRedirects,
    #[error("Malformed response: {0}")]
    MalformedResponse(&'static str),
    #[error("Invalid header '{0}' value: {1}")]
    InvalidHeaderValue(header::HeaderName, ToStrError),
    #[error("Invalid URL ('{0}'): {1}")]
    InvalidUrl(String, http::uri::InvalidUri),
}

#[derive(Debug, Parser)]
struct Arguments {
    #[clap(short, long)]
    output: Option<PathBuf>,
    #[clap(short, long)]
    follow: bool,
    #[clap(subcommand)]
    method: RequestMethod,
}

#[derive(Debug, Subcommand)]
enum RequestMethod {
    #[clap(arg_required_else_help = true)]
    Get {
        #[clap(help = "URL to GET")]
        url: Uri,
    },
}

enum Output {
    File(File),
    Stdout(Stdout),
}

impl Output {
    fn open<P: AsRef<Path>>(path: Option<P>) -> Result<Self, io::Error> {
        match path {
            Some(path) => File::create(path).map(Self::File),
            None => Ok(Self::Stdout(stdout())),
        }
    }
}

impl Write for Output {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        match self {
            Self::File(file) => file.write(buf),
            Self::Stdout(file) => file.write(buf),
        }
    }

    fn flush(&mut self) -> io::Result<()> {
        match self {
            Self::File(file) => file.flush(),
            Self::Stdout(file) => file.flush(),
        }
    }
}

fn call<C: HttpConnector, B: HttpBody>(
    client: &mut HttpClient<C>,
    method: Method,
    mut url: Uri,
) -> Result<HttpResponse<C::Connection, B>, Error>
where
    Error: From<HttpError<C::Error>>,
{
    const MAX_REDIRECTS: usize = 5;

    let mut redirect_count = 0;

    while redirect_count < MAX_REDIRECTS {
        log::info!("Try URL: {url:?}");
        let response = client
            .request(method.clone(), url)
            .header(header::CONNECTION, "close")?
            .call()?;

        if response.status().is_redirection() {
            let location =
                response
                    .headers()
                    .get(header::LOCATION)
                    .ok_or(Error::MalformedResponse(
                        "No 'Location' header in a redirect response",
                    ))?;
            let redirect_url = location
                .to_str()
                .map_err(|e| Error::InvalidHeaderValue(header::LOCATION, e))?;
            let redirect_url = Uri::from_str(redirect_url)
                .map_err(|e| Error::InvalidUrl(redirect_url.into(), e))?;

            log::info!("Redirect to {redirect_url:?}");

            url = redirect_url;
            redirect_count += 1;
        } else {
            return Ok(response);
        }
    }

    Err(Error::TooManyRedirects)
}

fn write_response<C: HttpConnection>(
    output: &mut Output,
    mut response: HttpResponse<C, Bytes>,
) -> Result<(), Error> {
    let mut buffer = [0; 4096];
    loop {
        let len = response.read(&mut buffer).unwrap();
        if len == 0 {
            break;
        }
        output.write_all(&buffer[..len])?;
    }
    Ok(())
}

fn get(url: Uri, output: Option<PathBuf>) -> Result<(), Error> {
    let mut output = Output::open(output)?;
    // TODO proper root certificate store
    let mut client = HttpClient::from(AutoConnector::with_insecure_tls());
    let response = call(&mut client, Method::GET, url)?;
    write_response(&mut output, response)?;
    Ok(())
}

fn main() -> ExitCode {
    logsink::setup_logging(false);
    let args = Arguments::parse();

    let result = match args.method {
        RequestMethod::Get { url } => get(url, args.output),
    };

    match result {
        Ok(()) => ExitCode::SUCCESS,
        Err(error) => {
            eprintln!("Error: {}", error);
            ExitCode::FAILURE
        }
    }
}