rewriting socket connection, optimizing big POST requests

tinyhttp-codegen/src/lib.rs

@@ -14,8 +14,8 @@ pub fn get(attr: TokenStream, item: TokenStream) -> TokenStream {
     let body = item_fn.block.deref();
     let return_type = sig.output;
 
-    let body_args = sig.inputs;
-    let is_body_args = !body_args.is_empty();
+    let fn_args = sig.inputs;
+    let is_body_args = !fn_args.is_empty();
     //eprintln!("LEN: {}", body_args.len());
 
     let mut path = value.value();
@@ -74,14 +74,25 @@ pub fn get(attr: TokenStream, item: TokenStream) -> TokenStream {
     //            get_route = get_route.set_body(body);
     //        }
     let new_get_body = if is_body_args {
+        let first_arg_name = fn_args.first().unwrap();
+        let arg_type = match first_arg_name {
+            syn::FnArg::Typed(i) => i.to_owned(),
+            _ => todo!(),
+        };
         quote! {
             let mut get_route = GetRouteWithReqAndRes::new()
                 .set_path(#path.into());
 
-            fn body(#body_args) -> Response {
+            fn body<'b>(try_from_req: &'b mut Request) -> Response {
+                let #arg_type = try_from_req.into();
                 #body.into()
             }
 
+            // OG
+            // fn body(#body_args) -> Response {
+            // #body.into()
+            // }
+
             get_route = get_route.set_body(body);
         }
     } else {

tinyhttp-internal/src/codegen/route.rs

@@ -175,7 +175,7 @@ pub struct GetRouteWithReqAndRes {
     path: Option<&'static str>,
     method: Method,
     wildcard: Option<String>,
-    get_body: Option<fn(Request) -> Response>,
+    get_body: Option<fn(&mut Request) -> Response>,
 }
 
 impl Default for GetRouteWithReqAndRes {
@@ -205,19 +205,19 @@ impl GetRouteWithReqAndRes {
         self.wildcard = Some(wildcard);
         self
     }
-    pub fn set_body(mut self, body: fn(Request) -> Response) -> Self {
+    pub fn set_body(mut self, body: fn(&'_ mut Request) -> Response) -> Self {
         self.get_body = Some(body);
         self
     }
 
-    pub fn get_body(&self) -> Option<fn(Request) -> Response> {
+    pub fn get_body(&self) -> Option<fn(&'_ mut Request) -> Response> {
         self.get_body
     }
 }
 
 impl ToResponse for GetRouteWithReqAndRes {
-    fn to_res(&self, req: Request) -> Response {
-        self.get_body().unwrap()(req)
+    fn to_res(&self, mut req: Request) -> Response {
+        self.get_body().unwrap()(&mut req)
     }
 }
 

tinyhttp-internal/src/http.rs

@@ -1,6 +1,6 @@
 use std::{
-    io::{self, BufReader},
-    iter::FromIterator,
+    collections::HashMap,
+    io::{self, BufRead, BufReader},
     path::Path,
     sync::Arc,
 };
@@ -17,6 +17,12 @@ use crate::{
 use flate2::{write::GzEncoder, Compression};
 
 pub fn start_http(http: HttpListener, config: Config) {
+    #[cfg(feature = "log")]
+    log::info!(
+        "Listening on port {}",
+        http.socket.local_addr().unwrap().port()
+    );
+
     let arc_config = Arc::new(config);
     for stream in http.get_stream() {
         let mut conn = stream.unwrap();
@@ -56,93 +62,85 @@ pub fn start_http(http: HttpListener, config: Config) {
     }
 }
 
-fn build_and_parse_req(buf: Vec<u8>) -> Result<Request, RequestError> {
-    #[cfg(feature = "log")]
-    log::trace!("build_and_parse_req ~~> buf: {:?}", buf);
+fn build_and_parse_req<P: Read>(conn: &mut P) -> Result<Request, RequestError> {
+    let mut buf_reader = BufReader::new(conn);
+    let mut status_line_str = String::new();
 
-    let mut safe_http_index = buf.windows(2).enumerate();
-    let status_line_index_opt = safe_http_index
-        .find(|(_, w)| matches!(*w, b"\r\n"))
-        .map(|(i, _)| i);
+    buf_reader.read_line(&mut status_line_str).unwrap();
+    status_line_str.drain(status_line_str.len() - 2..status_line_str.len());
 
-    let status_line_index = if let Some(status) = status_line_index_opt {
-        status
-    } else {
-        #[cfg(feature = "log")]
-        log::error!("failed parsing status line!");
-
-        return Err(RequestError::StatusLineErr);
-    };
-
-    let first_header_index = if let Some(first_header_index) = safe_http_index
-        .find(|(_, w)| matches!(*w, b"\r\n"))
-        .map(|(i, _)| i)
-    {
-        first_header_index
-    } else {
-        #[cfg(feature = "log")]
-        log::warn!("no headers found!");
+    #[cfg(feature = "log")]
+    log::trace!("STATUS LINE: {:#?}", status_line_str);
 
-        0usize
-    };
+    let iter = buf_reader.fill_buf().unwrap();
+    let header_end_idx = iter
+        .windows(4)
+        .position(|w| matches!(w, b"\r\n\r\n"))
+        .unwrap();
 
     #[cfg(feature = "log")]
-    log::trace!(
-        "STATUS LINE: {:#?}",
-        std::str::from_utf8(&buf[..status_line_index])
-    );
+    log::trace!("Body starts at {}", header_end_idx);
+    let headers_buf = iter[..header_end_idx + 2].to_vec();
 
-    #[cfg(feature = "log")]
-    log::trace!(
-        "FIRST HEADER: {:#?}",
-        std::str::from_utf8(&buf[status_line_index + 2..first_header_index])
-    );
+    buf_reader.consume(header_end_idx + 4); // Add 4 bytes since header_end_idx does not count
+                                            // \r\n\r\n
 
-    let mut headers = Vec::<String>::new();
-    let mut headers_index = first_header_index + 2;
+    let mut headers = HashMap::new();
+    let mut headers_index = 0;
 
-    // Sort through all request headers
-    while let Some(header_index) = safe_http_index
+    let mut headers_buf_iter = headers_buf.windows(2).enumerate();
+
+    //Sort through all request headers
+    while let Some(header_index) = headers_buf_iter
         .find(|(_, w)| matches!(*w, b"\r\n"))
-        .map(|(i, _)| i + 2)
+        .map(|(i, _)| i)
     {
         #[cfg(feature = "log")]
         log::trace!("header index: {}", header_index);
 
-        let header = std::str::from_utf8(&buf[headers_index..header_index - 2])
+        let header = std::str::from_utf8(&headers_buf[headers_index..header_index])
             .unwrap()
             .to_lowercase();
+
         if header.is_empty() {
             break;
         }
         #[cfg(feature = "log")]
-        log::trace!("HEADER: {:?}", headers);
+        log::trace!("HEADER: {:?}", header);
 
-        headers_index = header_index;
-        headers.push(header);
-    }
+        headers_index = header_index + 2;
 
-    let iter_status_line = std::str::from_utf8(&buf[..status_line_index]).unwrap();
+        let mut colon_split = header.splitn(2, ':');
+        headers.insert(
+            colon_split.next().unwrap().to_string(),
+            colon_split.next().unwrap().trim().to_string(),
+        );
+    }
 
-    //let headers = parse_headers(http.to_string());
-    let str_status_line = Vec::from_iter(iter_status_line.split_whitespace());
-    let status_line: Vec<String> = str_status_line.into_iter().map(|i| i.to_string()).collect();
-    #[cfg(feature = "log")]
-    log::trace!("{:#?}", status_line);
-    let body_index = buf
-        .windows(4)
-        .enumerate()
-        .find(|(_, w)| matches!(*w, b"\r\n\r\n"))
-        .map(|(i, _)| i)
+    let body_len = headers
+        .get("content-length")
+        .unwrap_or(&String::from("0"))
+        .parse::<usize>()
         .unwrap();
 
-    let raw_body = &buf[body_index + 4..];
-    //    #[cfg(feature = "log")]
-    //    log::debug!(
-    //        "BODY (TOP): {:#?}",
-    //        std::str::from_utf8(&buf[body_index + 4..]).unwrap()
-    //    );
-    Ok(Request::new(raw_body, headers, status_line, None))
+    let mut raw_body = Vec::with_capacity(body_len);
+
+    // Zero-init raw_body
+    (0..body_len).for_each(|_| {
+        raw_body.push(0);
+    });
+
+    buf_reader.read_exact(&mut raw_body).unwrap();
+
+    Ok(Request::new(
+        raw_body,
+        headers,
+        status_line_str
+            .split_whitespace()
+            .map(|s| s.to_string())
+            .collect(),
+        None,
+    ))
 }
 
 fn build_res(mut req: Request, config: &Config) -> Response {
@@ -265,8 +263,7 @@ fn build_res(mut req: Request, config: &Config) -> Response {
 }
 
 pub fn parse_request<P: Read + Write>(conn: &mut P, config: Arc<Config>) {
-    let buf = read_stream(conn);
-    let request = build_and_parse_req(buf);
+    let request = build_and_parse_req(conn);
 
     if let Err(e) = request {
         let specific_err = match e {
@@ -381,29 +378,3 @@ fn read_to_vec<P: AsRef<Path>>(path: P) -> io::Result<Vec<u8>> {
     }
     inner(path.as_ref())
 }
-
-pub(crate) fn read_stream<P: Read>(stream: &mut P) -> Vec<u8> {
-    let buffer_size = 1024;
-    let mut request_buffer = vec![];
-    loop {
-        let mut buffer = vec![0; buffer_size];
-        match stream.read(&mut buffer) {
-            Ok(n) => {
-                if n == 0 {
-                    break;
-                } else if n < buffer_size {
-                    request_buffer.append(&mut buffer[..n].to_vec());
-                    break;
-                } else {
-                    request_buffer.append(&mut buffer);
-                }
-            }
-            Err(e) => {
-                println!("Error: Could not read string!: {}", e);
-                std::process::exit(1);
-            }
-        }
-    }
-
-    request_buffer
-}

tinyhttp-internal/src/request.rs

@@ -1,11 +1,17 @@
-use std::{collections::HashMap, mem};
+use std::{collections::HashMap, fmt::Display, mem};
 
 #[derive(Clone, Debug, Default)]
-pub struct Wildcard<T> {
+pub struct Wildcard<T: Display> {
     wildcard: T,
 }
 
-impl<T> Wildcard<T> {
+impl<T: Display> Display for Wildcard<T> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.get_wildcard())
+    }
+}
+
+impl<T: Display> Wildcard<T> {
     pub fn get_wildcard(&self) -> &T {
         &self.wildcard
     }
@@ -18,7 +24,7 @@ impl<T> Wildcard<T> {
 /// body is used when the body of the request is not a String
 #[derive(Clone, Debug, Default)]
 pub struct Request {
-    raw_headers: Vec<String>,
+    raw_headers: HashMap<String, String>,
     status_line: Vec<String>,
     body: Vec<u8>,
     wildcard: Option<String>,
@@ -33,13 +39,13 @@ pub enum BodyType {
 
 impl Request {
     pub fn new(
-        raw_body: &[u8],
-        raw_headers: Vec<String>,
+        body: Vec<u8>,
+        raw_headers: HashMap<String, String>,
         status_line: Vec<String>,
         wildcard: Option<String>,
     ) -> Request {
         Request {
-            body: raw_body.to_vec(),
+            body,
             raw_headers,
             status_line,
             wildcard,
@@ -67,19 +73,16 @@ impl Request {
     }
 
     /// Get request headers in a HashMap
-    pub fn get_headers(&self) -> HashMap<&str, &str> {
+    pub fn get_headers(&self) -> &HashMap<String, String> {
         #[cfg(feature = "log")]
         log::trace!("Headers: {:#?}", self.raw_headers);
-        self.raw_headers
-            .iter()
-            .map(|i| i.split(": "))
-            .map(|mut i| (i.next().unwrap(), i.next().unwrap()))
-            .collect::<HashMap<&str, &str>>()
+
+        &self.raw_headers
     }
 
     /// Get status line of request
-    pub fn get_status_line<'a>(&'a self) -> &'a [String] {
-        &*self.status_line
+    pub fn get_status_line(&self) -> &[String] {
+        &self.status_line
     }
 
     pub fn get_wildcard(&self) -> Option<&String> {
@@ -104,11 +107,17 @@ impl<'a> From<&'a mut Request> for Wildcard<&'a str> {
     }
 }
 
-impl<'a> From<&'a mut Request> for Wildcard<&'a [u8]> {
+//impl<'a> From<&'a mut Request> for Wildcard<&'a [u8]> {
+//    fn from(value: &'a mut Request) -> Self {
+//        Wildcard {
+//            wildcard: value.wildcard.as_ref().unwrap().as_bytes(),
+//        }
+//    }
+//}
+
+impl<'a> From<&'a mut Request> for &'a HashMap<String, String> {
     fn from(value: &'a mut Request) -> Self {
-        Wildcard {
-            wildcard: value.wildcard.as_ref().unwrap().as_bytes(),
-        }
+        value.get_headers()
     }
 }