cataclysm 0.5.5

use tokio::{
    sync::{Semaphore},
    net::{TcpListener}
    //io::AsyncWriteExt
};
use bytes::Buf;
use crate::metafunctions::callback::PipelineKind;
use crate::{
    Stream,
    Branch, Shared, Additional, Cors, branch::PureBranch, Pipeline, Error, session::SessionCreator,
    http::{Request, RequestHeader, Response, Method}
};
use std::sync::{Arc};

// Default max connections for the server
const MAX_CONNECTIONS: usize = 2_000;
const KEEP_ALIVE_TIMEOUT: usize = 5;
const KEEP_ALIVE_MAX: usize = 100;
const RESPONSE_CHUNK_SIZE: usize = 4_096;
const READ_CHUNK_SIZE: usize = 8_192;

/// Builder pattern for the server structure
///
/// It is the main method for building a server and configuring certain behaviour
pub struct ServerBuilder<T> {
    branch: Branch<T>,
    shared: Option<Shared<T>>,
    session_creator: Option<Arc<Box<dyn SessionCreator>>>,
    log_string: Option<String>,
    cors: Option<Cors>,
    max_connections: usize,
    keep_alive_timeout: Option<usize>,
    keep_alive_max: Option<usize>,
    timeout: std::time::Duration
}

impl<T: Sync + Send> ServerBuilder<T> {
    /// Creates a new server from a given branch
    ///
    /// ```rust,no_run
    /// # use cataclysm::{ServerBuilder, Branch, http::{Method, Response}};
    /// let branch: Branch<()> = Branch::new("/").with(Method::Get.to(|| async {Response::ok().body("Ok!")}));
    /// let mut server_builder = ServerBuilder::new(branch);
    /// // ...
    /// ```
    pub fn new(branch: Branch<T>) -> ServerBuilder<T> {
        ServerBuilder {
            branch,
            shared: None,
            session_creator: None,
            log_string: None,
            cors: None,
            max_connections: MAX_CONNECTIONS,
            keep_alive_timeout: None,
            keep_alive_max: None,
            timeout: std::time::Duration::from_millis(15_000)
        }
    }

    /// Declare some information to be shared with the [Shared](crate::Shared) extractor
    ///
    /// ```rust,no_run
    /// use cataclysm::{Server, Branch, Shared, http::{Response, Method, Path}};
    /// 
    /// // Receives a string, and concatenates the shared suffix
    /// async fn index(path: Path<(String,)>, shared: Shared<String>) -> Response {
    ///     let (prefix,) = path.into_inner();
    ///     Response::ok().body(format!("{}{}", prefix, *shared))
    /// }
    /// 
    /// #[tokio::main]
    /// async fn main() {
    ///     // We create our tree structure
    ///     let branch = Branch::new("/{:prefix}").with(Method::Get.to(index));
    ///     // We create a server with the given tree structure
    ///     let server = Server::builder(branch).share("!!!".into()).build().unwrap();
    ///     // And we launch it on the following address
    ///     server.run("127.0.0.1:8000").await.unwrap();
    /// }
    /// ```
    ///
    /// If you intend to share a mutable variable, consider using rust's [Mutex](https://doc.rust-lang.org/std/sync/struct.Mutex.html), as the shared value is already inside an [Arc](https://doc.rust-lang.org/std/sync/struct.Arc.html).
    pub fn share(mut self, shared: T) -> ServerBuilder<T> {
        self.shared = Some(Shared::new(shared));
        self
    }

    /// Sets a session creator for the Session extractor to work
    ///
    /// The following is a bare bones example of this function being used, but should not be used for sessions as no mechanism for authentication is provided.
    ///
    /// ```rust,no_run
    /// use cataclysm::{Server, session::{Session, SessionCreator}, Branch, Shared, http::{Response, Request, Method, Path}};
    /// 
    /// async fn index(mut session: Session) -> Response {
    ///     // the session will be empty if the signature was invalid
    ///     // ... do something with the session
    ///     // apply changes to response
    ///     match session.apply(Response::ok()) {
    ///         Ok(response) => response,
    ///         Err(e) => {
    ///             println!("error when applying session to response, {}", e);
    ///             Response::internal_server_error()
    ///         }
    ///     }
    /// }
    ///
    /// struct RecklesslyUnsafeSession{
    /// 
    /// }
    ///
    /// // Exremely bad implementation of a cookie session with no signature, only for demo purposes
    /// impl SessionCreator for RecklesslyUnsafeSession {
    ///     fn parse(&self, req: &Request) -> Result<Option<String>, cataclysm::Error> {
    ///         let empty = vec![];
    ///         let cookie_headers = vec![
    ///             req.headers().get("Cookie").unwrap_or_else(|| &empty),
    ///             req.headers().get("cookie").unwrap_or_else(|| &empty)
    ///         ].into_iter().flatten();
    ///         for cookie_header in cookie_headers {
    ///             for single_cookie in cookie_header.split("; ") {
    ///                 if let Some((name, content)) = single_cookie.split_once("=") {
    ///                     if name == "reckless-and-unsafe" {
    ///                         return Ok(Some(content.to_string()))
    ///                     }
    ///                 }
    ///             }
    ///         }
    ///         Ok(None)
    ///     }
    /// 
    ///     fn apply(&self, content: String, mut res: Response) -> Result<Response, cataclysm::Error> {
    ///         // Absolutely reckless unsafe way to set a cookie, without even percent encoding
    ///         res = res.header("Set-Cookie", format!("reckless-and-unsafe={}", content));
    ///         Ok(res)
    ///     }
    /// }
    /// 
    /// #[tokio::main]
    /// async fn main() {
    ///     // We create our tree structure
    ///     let branch: Branch<()> = Branch::new("/").with(Method::Get.to(index));
    ///     // We create a server with the given tree structure
    ///     let server = Server::builder(branch).session_creator(
    ///         RecklesslyUnsafeSession{} // Default signed cookie session implementation
    ///     ).build().unwrap();
    ///     // And we launch it on the following address
    ///     server.run("127.0.0.1:8000").await.unwrap();
    /// }
    /// ```
    pub fn session_creator<A: 'static + SessionCreator>(mut self, session_creator: A) -> Self {
        self.session_creator = Some(Arc::new(Box::new(session_creator)));
        self
    }

    /// Sets a log string, to log information per call
    ///
    /// ```rust,no_run
    /// # use cataclysm::{Server, Branch, Shared, http::{Response, Method, Path}};
    /// // Tree structure
    /// let branch: Branch<()> = Branch::new("/").with(Method::Get.to(|| async {Response::ok()}));
    /// // Now we configure the server
    /// let server = Server::builder(branch).log_format("[%M %P] %S, from %A").build().unwrap();
    /// ```
    ///
    /// The list of available format elements are the following
    /// 
    /// * `%M`: Method from the request
    /// * `%P`: Path from the request
    /// * `%S`: Status from the response
    /// * `%A`: Socket address and port from the connection
    /// * `%F`: Responder path, where the callback was found (if any). Only available with the `full_log` feature.
    /// * `%f`: Same as previous but skipping file serving.
    /// (more data to be added soon)
    pub fn log_format<A: Into<String>>(mut self, log_string: A) -> Self {
        self.log_string = Some(log_string.into());
        self
    }

    /// Adds the cors "middleware"
    ///
    /// ```rust,no_run
    /// # use cataclysm::{Server, Branch, CorsBuilder, http::{Response, Method}};
    /// // Tree structure
    /// let branch: Branch<()> = Branch::new("/").with(Method::Get.to(|| async {Response::ok()}));
    /// // Now we configure the server
    /// let server = Server::builder(branch)
    ///     .cors(CorsBuilder::new()
    ///         .origin("https://fake.domain")
    ///         .max_age(600)
    ///         .build().unwrap()
    ///     ).build().unwrap();
    /// ```
    pub fn cors(mut self, cors: Cors) -> Self {
        self.cors = Some(cors);
        self
    }

    /// Sets up a maximum number of connections for the server to be dealt with
    ///
    /// ```rust,no_run
    /// # use cataclysm::{Server, Branch, http::{Response, Method}};
    /// // Tree structure
    /// let branch: Branch<()> = Branch::new("/").with(Method::Get.to(|| async {Response::ok()}));
    /// // Now we configure the server
    /// let server = Server::builder(branch).max_connections(10_000).build().unwrap();
    /// ```
    pub fn max_connections(mut self, n: usize) -> Self {
        self.max_connections = n;
        self
    }

    /// Sets up the timeout for keepalive connections.
    ///
    /// If this functions and `ServerBuilder::keep_alive_max` are not called, the keep alive functionality is off. If `ServerBuilder::keep_alive_max` is called only, the default timeout is 5 seconds.
    ///
    /// ```rust,no_run
    /// # use cataclysm::{Server, Branch, http::{Response, Method}};
    /// // Tree structure
    /// let branch: Branch<()> = Branch::new("/").with(Method::Get.to(|| async {Response::ok()}));
    /// // Now we configure the server
    /// let server = Server::builder(branch).keep_alive_timeout(10).build().unwrap();
    /// ```
    pub fn keep_alive_timeout(mut self, seconds: usize) -> Self {
        self.keep_alive_timeout = Some(seconds);
        self
    }

    /// Sets up maximum number of reuses for keepalive connections.
    ///
    /// If this functions and `ServerBuilder::keep_alive_timeout` are not called, the keep alive functionality is off. If `ServerBuilder::keep_alive_timeout` is called only, the default max is 100 times.
    ///
    /// ```rust,no_run
    /// # use cataclysm::{Server, Branch, http::{Response, Method}};
    /// // Tree structure
    /// let branch: Branch<()> = Branch::new("/").with(Method::Get.to(|| async {Response::ok()}));
    /// // Now we configure the server
    /// let server = Server::builder(branch).keep_alive_max(200).build().unwrap();
    /// ```
    pub fn keep_alive_max(mut self, n: usize) -> Self {
        self.keep_alive_max = Some(n);
        self
    }

    /// Sets up a custom timeout for http requests to be finished
    ///
    /// ```rust,no_run
    /// # use cataclysm::{Server, Branch, http::{Response, Method}};
    /// use std::time::Duration;
    /// // Tree structure
    /// let branch: Branch<()> = Branch::new("/").with(Method::Get.to(|| async {Response::ok()}));
    /// // Now we configure the server
    /// let server = Server::builder(branch).timeout(Duration::from_millis(5_000)).build().unwrap();
    /// ```
    pub fn timeout(mut self, duration: std::time::Duration) -> Self {
        self.timeout = duration;
        self
    }

    /// Builds the server
    ///
    /// ```rust,no_run
    /// use cataclysm::{Server, Branch, Shared, http::{Response, Method, Path}};
    /// 
    /// // Receives a string, and concatenates the shared suffix
    /// async fn index() -> Response {
    ///     Response::ok().body("Hello")
    /// }
    /// 
    /// #[tokio::main]
    /// async fn main() {
    ///     // We create our tree structure
    ///     let branch: Branch<()> = Branch::new("/").with(Method::Get.to(index));
    ///     // We create a server with the given tree structure
    ///     let server = Server::builder(branch).build().unwrap();
    ///     // And we launch it on the following address
    ///     server.run("127.0.0.1:8000").await.unwrap();
    /// }
    /// ```
    pub fn build(self) -> Result<Arc<Server<T>>, Error> {
        Ok(Arc::new(Server {
            pure_branch: Arc::new(self.branch.purify()),
            additional: Arc::new(Additional {
                shared: self.shared,
                session_creator: self.session_creator
            }),
            log_string: Arc::new(self.log_string),
            cors: Arc::new(self.cors),
            max_connections: Arc::new(Semaphore::new(self.max_connections)),
            keep_alive_timeout: Arc::new(self.keep_alive_timeout),
            keep_alive_max: Arc::new(self.keep_alive_max),
            timeout: Arc::new(self.timeout)
        }))
    }
}

/// Http Server instance
///
/// The Server structure hosts all the information to successfully process each call
pub struct Server<T> {
    pure_branch: Arc<PureBranch<T>>,
    additional: Arc<Additional<T>>,
    log_string: Arc<Option<String>>,
    cors: Arc<Option<Cors>>,
    max_connections: Arc<Semaphore>,
    keep_alive_timeout: Arc<Option<usize>>,
    keep_alive_max: Arc<Option<usize>>,
    timeout: Arc<std::time::Duration>
}

impl<T: 'static + Sync + Send> Server<T> {
    // Short for ServerBuilder's `new` function.
    pub fn builder(branch: Branch<T>) -> ServerBuilder<T> {
        ServerBuilder::new(branch)
    }

    pub async fn run<S: AsRef<str>>(self: &Arc<Self>, socket: S) -> Result<(), Error> {
        let listener = TcpListener::bind(socket.as_ref()).await.map_err(|e| Error::Io(e))?;

        log::info!("Cataclysm ongoing \u{26c8}");
        #[cfg(feature = "full_log")]
        log::warn!("using the `full_log` feature might impact performance and leak sensible information. Disable in production.");
        // We need a fused future for the select macro
        tokio::select! {
            _ = async {
                loop {
                    // We lock the loop until one permit becomes available
                    #[cfg(feature = "full_log")]
                    log::trace!("[server] semaphore contains {} available permits", self.max_connections.available_permits());
                    let permit = match self.max_connections.clone().acquire_owned().await {
                        Ok(p) => {
                            #[cfg(feature = "full_log")]
                            log::trace!("[server] permit obtained, {} remaining permits", self.max_connections.available_permits());
                            p
                        },
                        Err(_) => {
                            log::error!("[server] semaphore seems to be closed, terminating al processes");
                            break;
                        }
                    };
                    
                    match listener.accept().await {
                        Ok((socket, addr)) => {
                            #[cfg(feature = "full_log")]
                            log::trace!("[server] socket connection accepted");
                            let server = Arc::clone(self);

                            let stream = Stream::new(socket, Some(permit));
                            
                            tokio::spawn(async move {
                                match server.dispatch(stream, addr, *server.timeout).await {
                                    Ok(_) => {
                                        #[cfg(feature = "full_log")]
                                        log::trace!("[server] connection successfully dispatched");
                                    },
                                    Err(e) => {
                                        if !matches!(e, Error::Timeout) {
                                            log::error!("[server] error on dispatch call, {}", e);
                                        }
                                    }
                                }
                            });

                            #[cfg(feature = "full_log")]
                            log::trace!("[server] waiting for new socket connection...");
                        },
                        Err(e) => {
                            log::error!("[server] error on listening, {}", e);
                        }
                    }
                }
            } => (),
            _ = tokio::signal::ctrl_c() => {
                log::info!("Shutting down server");
            }
        };
        Ok(())
    }

    /// Deals with the read part of the socket stream
    async fn dispatch_read(socket: &Stream, addr: std::net::SocketAddr) -> Result<Option<Request>, Error> {
        let mut request_bytes = Vec::with_capacity(READ_CHUNK_SIZE);
        let mut expected_length = None;
        let mut request_header = None;
        // First we read
        loop {
            socket.readable().await.map_err(|e| Error::Io(e))?;
            
            // being stored in the async task.
            let mut buf = [0; READ_CHUNK_SIZE];

            // Try to read data, this may still fail with `WouldBlock`
            // if the readiness event is a false positive.
            match socket.try_read(&mut buf) {
                Ok(0) => {
                    break
                },
                Ok(n) => {
                    request_bytes.extend_from_slice(&buf[0..n]);

                    if request_header.is_none() {
                        request_header = match RequestHeader::parse(&mut request_bytes, addr) {
                            Ok(r) => {
                                // We check if we need to give a continue 100
                                if r.headers.get("Expect").map(|h| h.get(0).map(|ih| ih == "100-continue")).flatten().unwrap_or(false) {
                                    // We send it
                                    Server::<T>::dispatch_write(&socket, Response::r#continue()).await?;
                                }

                                // We check now if there is a content size hint
                                expected_length = r.headers.get("Content-Length").or_else(|| r.headers.get("content-length")).map(|cl| cl.get(0).map(|v| v.parse::<usize>().ok())).flatten().flatten();
                                #[cfg(feature = "full_log")]
                                log::trace!("expecting to read {:?} bytes in request", expected_length);
                                Some(r)
                            },
                            Err(_e) => {
                                #[cfg(feature = "full_log")]
                                log::debug!("{}", _e);
                                Server::<T>::dispatch_write(&socket, Response::bad_request()).await?;
                                return Ok(None)
                            }
                        };
                    }

                    // And now we check if, given the hint, we need to act upon.
                    if let Some(expected_length) = &expected_length {
                        if *expected_length > request_bytes.len() {
                            continue;
                        } else {
                            break;
                        }
                    } else {
                        break;
                    }
                },
                Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => {
                    continue
                }
                Err(e) => return Err(Error::Io(e))
            }
        }

        Ok(request_header.map(|v| v.content(request_bytes)))
    }

    async fn dispatch_write(socket: &Stream, mut response: Response) -> Result<(), Error> {
        let serialized_response = response.serialize();
        let mut chunks_iter = serialized_response.chunks(RESPONSE_CHUNK_SIZE);
        #[cfg(feature = "full_log")]
        log::trace!("writting {} chunks of maximum {} bytes each", chunks_iter.len(), RESPONSE_CHUNK_SIZE);
        // We check the first chunk
        let mut current_chunk = match chunks_iter.next() {
            Some(v) => v,
            None => return Ok(()) // Zero length response
        };
        loop {
            // Wait for the socket to be writable
            socket.writable().await.map_err(|e| Error::Io(e))?;
    
            // Try to write data, this may still fail with `WouldBlock`
            // if the readiness event is a false positive.        
            match socket.try_write(&current_chunk) {
                Ok(n) => {
                    if n != current_chunk.remaining() {
                        // There are some bytes still to be written in this chunk
                        #[cfg(feature = "full_log")]
                        log::trace!("incomplete chunk, trying to serve remaining bytes ({}/{})", current_chunk.len(), RESPONSE_CHUNK_SIZE);
                        current_chunk.advance(n);
                        continue;
                    } else {
                        current_chunk = match chunks_iter.next() {
                            Some(v) => v,
                            None => return Ok(())
                        }
                    }
                }
                Err(ref e) if e.kind() == tokio::io::ErrorKind::WouldBlock => {
                    continue;
                }
                Err(e) => break Err(Error::Io(e))
            }
        }
    }

    async fn dispatch(self: &Arc<Self>, stream: Stream, addr: std::net::SocketAddr, mut timeout: std::time::Duration) -> Result<(), Error> {
        let mut remaining_per_connection = None;
        let default_max_times = 200;
        #[cfg(feature = "full_log")]
        let mut attended_paths = Vec::new();
        loop {
            if remaining_per_connection == Some(0) {
                break;
            }

            let mut request = tokio::select!{
                res = Server::<T>::dispatch_read(&stream, addr) => match res {
                    Ok(request) => match request {
                        Some(b) => b,
                        None => return Ok(())
                    },
                    Err(e) => return Err(e)
                },
                _ = tokio::time::sleep(timeout) => {
                    #[cfg(feature = "full_log")]
                    log::trace!("[server] timeout for http response, after attending {:?}", attended_paths);
                    return Err(Error::Timeout)
                }
            };

            #[cfg(feature = "full_log")]
            {
                log::trace!("[server] headers: {:?}", request.header.headers);
                attended_paths.push(format!("{}", request.url().path()));
            }
    
            if let Some(cors) = &*self.cors {
                if request.header.method == Method::Options {
                    if let Some(supported_methods) = self.pure_branch.supported_methods(request.url().path()) {
                        #[cfg(feature = "full_log")]
                        log::trace!("[server] processing reply to preflight cors call");
                        stream.response(cors.preflight(&request, &supported_methods)).await?;
                        continue;
                    } // If the method is not options, it will anyways return a not-found
                }
            }
    
            // No clue why this was here
            // request.header.addr = addr;
    
            #[cfg(feature = "full_log")]
            let mut tracker = None;
    
            // The method will take the request, and modify particularly the "variable count" variable
            let mut response = match self.pure_branch.pipeline(&mut request) {
                Some(pipeline_info) => {
                    #[cfg(feature = "full_log")]
                    {
                        tracker = Some(pipeline_info.pipeline_track);
                    }
    
                    match pipeline_info.pipeline_kind {
                        PipelineKind::NormalPipeline{pipeline} => {
                            #[cfg(feature = "full_log")]
                            log::trace!("[server] found normal pipeline for path {} with method {}", request.header.url, request.header.method);
                            match pipeline {
                                Pipeline::Layer(func, pipeline_layer) => func(request.clone(), pipeline_layer, self.additional.clone()),
                                Pipeline::Core(core_fn) => core_fn(request.clone(), self.additional.clone())
                            }.await
                        },
                        #[cfg(feature = "stream")]
                        PipelineKind::StreamPipeline{pipeline} => {
                            #[cfg(feature = "full_log")]
                            log::trace!("[server] found stream pipeline for path {}", request.header.url);
                            pipeline(request.clone(), self.additional.clone(), stream).await;
                            return Ok(())
                        }
                    }
                },
                None => {
                    #[cfg(feature = "full_log")]
                    log::trace!("[server] pipeline for path {} with method {} not found", request.header.url, request.header.method);
                    Response::not_found()
                }
            };
    
            let should_keep_alive = request.header.requests_keep_alive(); 
    
            if let Some(remaining_per_connection) = &mut remaining_per_connection {
                *remaining_per_connection -= 1;
                if *remaining_per_connection == 0 {
                    response = response.header("Connection", "Close");
                }
            } else {
                if should_keep_alive {
                    if self.keep_alive_timeout.is_some() || self.keep_alive_max.is_some() {
                        let keep_alive_timeout = self.keep_alive_timeout.unwrap_or(KEEP_ALIVE_TIMEOUT);
                        let keep_alive_max = self.keep_alive_max.unwrap_or(KEEP_ALIVE_MAX);
                        #[cfg(feature = "full_log")]
                        log::trace!("[server] keep alive request received, setting new timeout to {} seconds, and maximum {} calls", keep_alive_timeout, keep_alive_max);
                        response = response.header("Keep-Alive", format!("timeout={}, max={}", keep_alive_timeout, keep_alive_max));
                        timeout = std::time::Duration::from_secs(5);
                        remaining_per_connection = Some(default_max_times);
                    } else {
                        #[cfg(feature = "full_log")]
                        log::trace!("[server] keep alive request received, but keep alive is not enabled");
                        response = response.header("Connection", "Close");
                        remaining_per_connection = Some(0);
                    }
                } else {
                    remaining_per_connection = Some(0);
                }
            }
    
            // Cors validation, not as an actual pipeline layer
            if let Some(cors) = &*self.cors {
                cors.apply(&request, &mut response);
            }
    
            if let Some(log_string) = &*self.log_string {
                #[allow(unused_mut)]
                let mut final_log_string = log_string.replace("%M", request.header.method.to_str())
                    .replace("%P", &request.url().path())
                    .replace("%A", &format!("{}", addr))
                    .replace("%S", &format!("{}", response.status.0));
                #[cfg(feature = "full_log")]
                {
                    if log_string.contains("%f") {
                        if matches!(tracker, Some(crate::metafunctions::callback::PipelineTrack::File(_))) {
                            final_log_string = "".to_string();
                        }
                    }
                    let replacer = tracker.map(|t| format!("{}", t)).unwrap_or_else(|| "NoTrack".to_string());
                    final_log_string = final_log_string
                        .replace(&"%F", &replacer)
                        .replace(&"%f", &replacer);
                }
    
    
                if !final_log_string.is_empty() {
                    log::info!("{}", final_log_string);
                }
            }
    
            stream.response(response).await?;
        }
        #[cfg(feature = "full_log")]
        log::trace!("[server] leaving dispatch method");
        Ok(())
    }
}