maj/src/response.rs

use crate::{gemini, StatusCode};
use num::FromPrimitive;
use std::fmt;
use std::io::{self, prelude::*, ErrorKind};

/// A Gemini response as specified in [the spec](https://gemini.circumlunar.space/docs/specification.html).
#[derive(Default)]
pub struct Response {
    pub status: StatusCode,
    pub meta: String,
    pub body: Vec<u8>,
}

#[derive(thiserror::Error, Debug)]
pub struct ResponseStatusError {
    status: StatusCode,
    meta: String,
}

impl fmt::Display for ResponseStatusError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "{:?} ({}): {}",
            self.status, self.status as u8, self.meta
        )
    }
}

impl Response {
    pub fn with_body(meta: String, body: Vec<u8>) -> Response {
        Response {
            status: StatusCode::Success,
            meta: meta,
            body: body,
        }
    }

    pub fn gemini(body: Vec<u8>) -> Response {
        Response {
            status: StatusCode::Success,
            meta: "text/gemini".to_string(),
            body: body,
        }
    }

    pub fn render(body: Vec<gemini::Node>) -> Response {
        let mut buf: Vec<u8> = vec![];
        gemini::render(body, &mut buf).unwrap();

        Response {
            status: StatusCode::Success,
            meta: "text/gemini".to_string(),
            body: buf,
        }
    }

    pub fn perm_redirect(to: String) -> Response {
        Response {
            status: StatusCode::PermanentRedirect,
            meta: to,
            body: vec![],
        }
    }

    pub fn no_proxy() -> Response {
        Response {
            status: StatusCode::ProxyRequestRefused,
            meta: "Wrong host".to_string(),
            body: vec![],
        }
    }

    pub fn not_found() -> Response {
        Response {
            status: StatusCode::NotFound,
            meta: "Not found".to_string(),
            body: vec![],
        }
    }

    pub fn input<T: Into<String>>(msg: T) -> Response {
        Response {
            status: StatusCode::Input,
            meta: msg.into(),
            body: vec![],
        }
    }

    pub fn need_cert<T: Into<String>>(msg: T) -> Response {
        Response {
            status: StatusCode::ClientCertificateRequired,
            meta: msg.into(),
            body: vec![],
        }
    }
}

/// The parser state.
#[derive(Debug)]
enum State {
    ReadStatusCode { data: Vec<u8> },
    ReadWhitespace,
    ReadMeta { data: Vec<u8> },
    ReadBody { data: Vec<u8> },
}

/// Response error.
#[derive(thiserror::Error, Debug)]
pub enum Error {
    #[error("unexpected end of file found while parsing response")]
    EOF,

    #[error("I/O error")]
    IO(#[from] std::io::Error),

    #[error("invalid status code character {0}")]
    InvalidStatusCode(u8),

    #[error("UTF-8 error: {0}")]
    Utf8(#[from] std::str::Utf8Error),

    #[error("Number parsing error: {0}")]
    NumParse(#[from] std::num::ParseIntError),

    #[error("None found when none should not be found")]
    NoneFound,

    #[error("Response meta is too long")]
    ResponseMetaTooLong,
}

impl Response {
    pub fn parse(inp: &mut impl Read) -> Result<Response, Error> {
        let mut state = State::ReadStatusCode { data: vec![] };
        let mut buf = [0; 1];
        let mut result = Response::default();

        loop {
            match inp.read(&mut buf) {
                Ok(n) => {
                    if n == 0 {
                        if let State::ReadBody { data } = state {
                            result.body = data;
                            return Ok(result);
                        }
                        panic!("got here: {}, {:?}", n, state);
                    }
                }

                Err(why) => {
                    if why.kind() == ErrorKind::ConnectionAborted {
                        if let State::ReadBody { data } = state {
                            result.body = data;
                            return Ok(result);
                        }
                    }

                    return Err(Error::IO(why));
                }
            }

            match &mut state {
                State::ReadStatusCode { data } => match buf[0] as char {
                    '1' | '2' | '3' | '4' | '5' | '6' | '7' | '8' | '9' | '0' => {
                        data.push(buf[0]);
                    }
                    ' ' | '\t' => {
                        let status_code: &str = std::str::from_utf8(data)?;
                        let status_code: u8 = status_code.parse()?;
                        result.status = StatusCode::from_u8(status_code).ok_or(Error::NoneFound)?;
                        state = State::ReadWhitespace;
                    }
                    foo => return Err(Error::InvalidStatusCode(foo as u8)),
                },

                State::ReadWhitespace => match buf[0] as char {
                    ' ' | '\t' => {}
                    _ => {
                        state = State::ReadMeta { data: vec![buf[0]] };
                    }
                },

                State::ReadMeta { data } => match buf[0] as char {
                    '\r' => {}
                    '\n' => {
                        result.meta = std::str::from_utf8(data)?.to_string();
                        state = State::ReadBody { data: vec![] };
                    }
                    _ => {
                        if data.len() == 1024 {
                            return Err(Error::ResponseMetaTooLong);
                        }
                        data.push(buf[0]);
                    }
                },

                State::ReadBody { data } => data.push(buf[0]),
            }
        }
    }

    pub fn write(self, out: &mut impl Write) -> io::Result<()> {
        write!(out, "{} {}\r\n", self.status.num(), self.meta)?;
        out.write(&self.body)?;

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::*;

    #[test]
    fn success() -> Result<(), Error> {
        let _ = pretty_env_logger::try_init();
        let mut fin = std::fs::File::open("./testdata/simple_response.txt")?;

        let resp = Response::parse(&mut fin)?;
        assert_eq!(resp.meta, "text/gemini".to_string());
        assert_eq!(resp.status, StatusCode::Success);

        Ok(())
    }

    #[test]
    fn error() -> Result<(), Error> {
        let _ = pretty_env_logger::try_init();
        let mut fin = std::fs::File::open("./testdata/error_response.txt")?;

        let resp = Response::parse(&mut fin)?;
        assert_eq!(resp.status, StatusCode::PermanentFailure);

        Ok(())
    }

    #[test]
    fn not_found() -> Result<(), Error> {
        let _ = pretty_env_logger::try_init();
        let mut fin = std::fs::File::open("./testdata/notfound_response.txt")?;

        let resp = Response::parse(&mut fin)?;
        assert_eq!(resp.status, StatusCode::NotFound);

        Ok(())
    }

    #[test]
    fn meta_too_long() {
        let _ = pretty_env_logger::try_init();
        let mut fin = std::fs::File::open("./testdata/meta_too_long.txt").unwrap();

        match Response::parse(&mut fin) {
            Ok(_) => panic!("wanted error but didn't get one"),
            Err(why) => {
                if let ResponseError::ResponseMetaTooLong = why {
                    println!("ok");
                } else {
                    panic!("wanted ResponseError::ResponseMetaTooLong")
                }
            }
        }
    }
}