http-server-cc/src/main.rs

// #![feature(if_let_guard)]

use std::{ collections::HashMap, path::PathBuf, sync::Arc };
use anyhow::{bail, Result};

use itertools::Itertools;
use tokio::{
    fs::File,
    io::{ AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader },
    net::{ TcpListener, TcpStream },
};

mod utils;
use utils::*;

#[derive(Debug, Clone)]
struct Args {
    pub directory: Option<PathBuf>,
}

type A = Arc<Args>;


fn parse_args () -> Args {
    let directory = std::env::args()
        .position(|e| e == "--directory")
        .and_then(|d| std::env::args().nth(d + 1))
        .map(|d| PathBuf::from(d));

    Args { directory }
}

#[tokio::main]
async fn main() -> Result<()> {
    let listener = TcpListener::bind("127.0.0.1:4221").await.unwrap();
    let args = Arc::new(parse_args());

    loop {
        let args = args.clone();
        let Ok((socket, info)) = listener.accept().await else { continue; } ;
        println!("Connection from {:?} accepted, processing…", info);
        let _ = process(socket, args).await;
    }

}


async fn process (mut stream: TcpStream, args: A) -> Result<()> {
    let buf_reader = BufReader::new(&mut stream);

    let mut data = buf_reader.lines();

    let (method, target, _version) = 'outer : {
        'inner : {
            let Ok(Some(start_line)) = data.next_line().await else { break 'inner };
            let mut parts = start_line.split_ascii_whitespace();
            let Some(Ok(method))  = parts.next().map(Method::try_from) else { break 'inner };
            let Some(path)    = parts.next() else { break 'inner };
            let Some(version) = parts.next() else { break 'inner };
            break 'outer (method, path.to_owned(), version.to_owned());
        }
        // this is either the best or the worst piece of code i've written

        let _ = stream.write_all(&Response::_400.build()).await;
        let _ = stream.flush();
        bail!(E::InvalidRequest)
    };


    let mut data = data.into_inner();

    let headers = Headers::parse(&mut data).await;
    let encoding = Encoding::parse(headers.get("Accept-Encoding"));

    println!("{:?}", headers);
    println!("{:?}", encoding);

    use Method as M;
    let response = match (method, target.as_str()) {
        (M::GET, "/") => Response::Empty,
        (M::GET, "/user-agent") => Response::TextPlain(headers.get("User-Agent").to_owned(), encoding),
        // p if let Some(echo) = p.strip_prefix("/echo/") => Response::TextPlain(echo), // a nicer way to do that, not available in stable yet
        (M::GET, r) if r.starts_with("/echo/") => Response::TextPlain(r.trim_start_matches("/echo/").to_owned(), encoding),
        (M::GET, r) if r.starts_with("/files/") => 'file : {
            let Some(path) = &args.directory else { break 'file Response::_500; };
            let path = path.join(r.trim_start_matches("/files/"));
            let Ok(mut f) = File::open(path).await else { break 'file Response::_404; };

            let mut buf = vec![];
            let _ = f.read_to_end(&mut buf).await;
            Response::OctetStream(buf)
        },
        (M::POST, r) if r.starts_with("/files") => 'file : {
            let length = headers.get("Content-Length").parse().unwrap();
            let body = parse_req_body(&mut data, length).await;

            let Some(path) = &args.directory else { break 'file Response::_500; };
            let path = path.join(r.trim_start_matches("/files/"));
            let Ok(mut f) = File::create(path).await else { break 'file Response::_500; };
            let Ok(_) = f.write_all(&body).await else { break 'file Response::_500 };

            Response::_201
        },
        _ => Response::_404,
    };

    println!("{:?}", String::from_utf8_lossy(&response.clone().build()));

    let _ = stream.write_all(&response.build()).await;
    let _ = stream.flush().await;

    Ok(())
}

pub async fn parse_req_body (reader: &mut BufReader<&mut TcpStream>, length: usize) -> Vec<u8> {
    let mut v = vec![0; length];
    let _ = reader.read_exact(&mut v).await;
    v
}

#[derive(Debug, Clone)]
pub struct Headers (HashMap<String, String>);

impl Headers {
    pub async fn parse (reader: &mut BufReader<&'_ mut TcpStream>) -> Self {
        let mut map = HashMap::new();
        let mut buf = String::new();
        while let Ok(_) = reader.read_line(&mut buf).await {
            if let Some((k, v)) = buf.split_once(":") {
                map.insert(k.trim().to_lowercase(), v.trim().to_owned());
            } else {
                break;
            }

            buf.clear();
        }
        Self(map)
    }

    pub fn get <'a> (&'a self, key: &str) -> &'a str {
        self.0.get(&key.to_lowercase()).map(|e| e.as_str()).unwrap_or_default()
    }
}


#[derive(Debug, Clone)]
enum Encoding {
    Gzip,
    Invalid,
    None,
}


impl Encoding {
    pub fn header (&self) -> &'static str {
        match self {
            Self::Gzip => "gzip",
            _ => d!()
        }
    }

    pub fn parse (s: &str) -> Vec<Self> {
        s.split(',').map(str::trim).map(From::from).collect()
    }
}

impl From<&str> for Encoding {
    fn from (s: &str) -> Self {
        match s.to_ascii_lowercase().as_str()  {
            "" => Self::None,
            "gzip" => Self::Gzip,
            _ => Self::Invalid
        }
    }
}


#[derive(Debug, Clone, PartialEq, Eq)]
enum Method {
    GET,
    POST,
    PUT,
    DELETE,
    UPDATE,
}

impl TryFrom<&str> for Method {
    type Error = anyhow::Error;
    fn try_from (s: &str) -> Result<Self> {
        let mut s = s.to_string();
        s.make_ascii_lowercase();
        Ok(match s.as_str() {
            "get" => Self::GET,
            "post" => Self::POST,
            "put" => Self::PUT,
            "delete" => Self::DELETE,
            "update" => Self::UPDATE,
            _ => bail!(E::UnknownMethod)
        })
    }
}

#[derive(Debug, Clone)]
enum Response {
    _201,
    _400,
    _404,
    _500,
    Empty,
    TextPlain (String, Vec<Encoding>),
    OctetStream (Vec<u8>)
}

#[allow(non_upper_case_globals)]
impl Response {
    fn build (self) -> Vec<u8> {

        let headers = self.headers().join("\r\n");
        let code = self.code();

        let mut v: Vec<u8> = f!("HTTP/1.1 {code}\r\n{headers}\r\n\r\n").into();

        match self {
            Self::OctetStream(bytes) => {
                v.extend_from_slice(&bytes);
            },
            Self::TextPlain(text, _) => {
                v.extend_from_slice(text.as_bytes());
            },
            _ => ()
        }

        v
    }

    fn code (&self) -> &'static str {
        match self {
            Self::_201 => "201 Created",
            Self::_400 => "400 Bad Request",
            Self::_404 => "404 Not Found",
            Self::_500 => "500 Internal Server Error",
            _ => "200 OK",
        }
    }

    fn headers (&self) -> Vec<String> {
        match self {
            Self::TextPlain(text, enc) => {
                let mut v = vec![
                    f!("Content-Type: text/plain"),
                    format!("Content-Length: {}", text.len()),
                ];
                let enc = enc.into_iter().map(Encoding::header).filter(|e| !e.is_empty()).join(", ");
                if !enc.is_empty() { v.push(f!("Content-Encoding: {enc}")) }
                v
            },
            Self::OctetStream(bytes) => vec![
                     f!("Content-Type: application/octet-stream"),
                format!("Content-Length: {}", bytes.len())
            ],
            _ => d!()
        }
    }
}