athene 2.0.4

use crate::{
    context::Context,
    error::Error,
    middleware::{ImplNext, MiddlewareBuilder, Next},
    request::Request,
    response::Response,
    router::{HandlerChain, Route, Router, RouterBuilder},
    tokio_io::TokioIO,
};
use hyper::{server::conn::http1, service::Service};
use std::{future::Future, net::SocketAddr, pin::Pin, sync::Arc};
use tokio::net::TcpListener;

#[doc(hidden)]
pub struct Builder<C, M>
where
    C: 'static + Route + Unpin + Send + Sync,
    M: 'static + Next + Unpin + Send + Sync,
{
    router: RouterBuilder<C>,
    middlewares: MiddlewareBuilder<M>,
}

impl<C, M> Builder<C, M>
where
    C: 'static + Route + Unpin + Send + Sync,
    M: 'static + Next + Unpin + Send + Sync,
{
    #[inline]
    pub fn new() -> Builder<HandlerChain, ImplNext> {
        Builder {
            router: RouterBuilder::default(),
            middlewares: MiddlewareBuilder::default(),
        }
    }

    #[inline]
    pub fn router<R, F>(self, f: F) -> Builder<R, M>
    where
        R: Route + Send + Sync + Unpin,
        F: FnOnce(RouterBuilder<C>) -> RouterBuilder<R>,
    {
        Builder {
            router: f(self.router),
            middlewares: self.middlewares,
        }
    }

    #[inline]
    pub fn middleware<N, F>(self, f: F) -> Builder<C, N>
    where
        N: Next + Unpin + Send + Sync,
        F: FnOnce(MiddlewareBuilder<M>) -> MiddlewareBuilder<N>,
    {
        Builder {
            router: self.router,
            middlewares: f(self.middlewares),
        }
    }

    #[inline]
    pub async fn listen(self, addr: &str) -> Result<(), Error> {
        let mut server = Server {
            router: self.router.build(),
            middlewares: self.middlewares.build(),
            addr: None,
        };

        let listener = TcpListener::bind(addr).await?;
        println!("{}", addr);
        loop {
            let (stream, addr) = listener.accept().await?;
            server.addr = Some(addr);
            let server = server.clone();
            let io = TokioIO::new(stream);
            tokio::task::spawn(async move {
                #[cfg(feature = "websocket")]
                {
                    if let Err(error) = http1::Builder::new()
                        .serve_connection(io, server)
                        .with_upgrades()
                        .await
                    {
                        eprintln!("Error while serving HTTP connection: {error}");
                    }
                }
                #[cfg(not(feature = "websocket"))]
                {
                    if let Err(error) = http1::Builder::new()
                        .serve_connection(io, server)
                        .await
                    {
                        eprintln!("Error while serving HTTP connection: {error}");
                    }
                }
            });
        }
    }

    #[cfg(feature = "https")]
    #[inline]
    pub async fn listen_tls(self, addr: &str, tls: crate::rustls::TlsConfig) -> Result<(), Error> {
        let mut server = Server {
            router: self.router.build(),
            middlewares: self.middlewares.build(),
            addr: None,
        };

        let tls: tokio_rustls::TlsAcceptor = tls.build().map(std::sync::Arc::new)?.into();
        println!("{}{}", "\x1b[31m", addr);
        loop {
            let (stream, addr) = listener.accept().await?;
            server.addr = Some(addr);
            let server = server.clone();
            let tls_stream = tls.accept(stream).await?;
            let io = TokioIO::new(tls_stream);
            tokio::task::spawn(async move {
                #[cfg(feature = "websocket")]
                {
                    if let Err(error) = http1::Builder::new()
                        .serve_connection(io, server)
                        .with_upgrades()
                        .await
                    {
                        eprintln!("Error while serving HTTP connection: {error}");
                    }
                }
                #[cfg(not(feature = "websocket"))]
                {
                    if let Err(error) = http1::Builder::new()
                        .serve_connection(io, server)
                        .await
                    {
                        eprintln!("Error while serving HTTP connection: {error}");
                    }
                }
            });
        }
    }
}

#[doc(hidden)]
#[derive(Clone)]
struct Server {
    router: Router,
    middlewares: Arc<dyn Next>,
    addr: Option<SocketAddr>,
}

impl Server {
    #[allow(unused_mut)]
    #[inline]
    async fn serve(
        request: hyper::Request<hyper::body::Incoming>,
        router: Router,
        middlewares: Arc<dyn Next>,
        addr: SocketAddr,
    ) -> Result<Response, Error> {
        let mut req = Request::new(request.map(crate::body::HttpBody::Incoming), addr);
        let meta = router.resolve_metadata(&mut req);
        let ctx = Context::new(req, router.clone(), meta);
        let res = middlewares
            .next(ctx)
            .await
            .and_then(|mut ctx| ctx.state.take_response())
            .or_else(|e| {
                let builder = crate::response::Builder::new();
                let res = e.response_builder(builder).build()?;
                Ok(res)
            });
        res
    }
}

impl Service<hyper::Request<hyper::body::Incoming>> for Server {
    type Error = Error;
    type Response = Response;
    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;

    #[inline]
    fn call(&self, req: hyper::Request<hyper::body::Incoming>) -> Self::Future {
        Box::pin(Self::serve(
            req,
            self.router.clone(),
            self.middlewares.clone(),
            self.addr.unwrap(),
        ))
    }
}