boitalettres/src/server/mod.rs

use std::future::Future;
use std::pin::Pin;
use std::task::{self, Poll};

use futures::io::{AsyncRead, AsyncWrite};
use imap_codec::types::response::Capability;
use tower::Service;

use crate::proto::{Request, Response};
use accept::Accept;
pub use service::MakeServiceRef;

pub mod accept;
mod pipeline;
mod service;

#[derive(Debug, thiserror::Error)]
pub enum Error<A> {
    Accept(#[source] A),
    MakeService(#[source] Box<dyn std::error::Error + Send + Sync + 'static>),
}

#[derive(Debug, Default, Clone)]
pub struct Imap {
    pub capabilities: Vec<Capability>,
}

#[pin_project::pin_project]
pub struct Server<I, S> {
    #[pin]
    incoming: I,
    make_service: S,
    protocol: Imap,
}

impl<I> Server<I, ()>
where
    I: Accept,
{
    #[allow(clippy::new_ret_no_self)]
    pub fn new(incoming: I) -> Builder<I> {
        Builder::new(incoming)
    }
}

impl<I, S> Future for Server<I, S>
where
    I: Accept,
    I::Conn: AsyncRead + AsyncWrite + Unpin + Send + 'static,
    S: MakeServiceRef<I::Conn, Request, Response = Response>,
    S::MakeError: Into<Box<dyn std::error::Error + Send + Sync + 'static>> + std::fmt::Display,
    S::Error: std::fmt::Display,
    S::Future: Send + 'static,
{
    type Output = Result<(), Error<I::Error>>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
        loop {
            let this = self.as_mut().project();

            if let Some(conn) = futures::ready!(this.incoming.poll_accept(cx)) {
                let conn = conn.map_err(Error::Accept)?;

                futures::ready!(this.make_service.poll_ready_ref(cx))
                    .map_err(Into::into)
                    .map_err(Error::MakeService)?;

                let service_fut = this.make_service.make_service_ref(&conn);

                tokio::task::spawn(Connecting {
                    conn,
                    service_fut,
                    protocol: this.protocol.clone(),
                });
            } else {
                return Poll::Ready(Ok(()));
            }
        }
    }
}

pub struct Builder<I: Accept> {
    incoming: I,
    protocol: Imap,
}

#[allow(clippy::needless_update)]
impl<I: Accept> Builder<I> {
    pub fn new(incoming: I) -> Self {
        Self {
            incoming,
            protocol: Default::default(),
        }
    }

    pub fn capabilities(self, capabilities: impl IntoIterator<Item = Capability>) -> Self {
        let protocol = Imap {
            capabilities: capabilities.into_iter().collect(),
            ..self.protocol
        };

        Self { protocol, ..self }
    }

    pub fn serve<S>(self, make_service: S) -> Server<I, S>
    where
        S: MakeServiceRef<I::Conn, Request, Response = Response>,
    {
        Server {
            incoming: self.incoming,
            make_service,
            protocol: self.protocol,
        }
    }
}

#[pin_project::pin_project]
struct Connecting<C, F> {
    conn: C,
    #[pin]
    service_fut: F,
    protocol: Imap,
}

impl<C, F, ME, S> Future for Connecting<C, F>
where
    C: AsyncRead + AsyncWrite + Unpin,
    F: Future<Output = std::result::Result<S, ME>>,
    ME: std::fmt::Display,
    S: Service<Request, Response = Response>,
    S::Error: std::fmt::Display,
{
    type Output = ();

    fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
        use tokio_tower::pipeline::Server as PipelineServer;

        use pipeline::Connection;

        let this = self.project();

        let service = match futures::ready!(this.service_fut.poll(cx)) {
            Ok(service) => service,
            Err(err) => {
                tracing::debug!("Connection error: {}", err);
                return Poll::Ready(());
            }
        };
        let mut conn = Connection::new(this.conn);

        // TODO: Properly handle server greeting
        {
            use imap_codec::types::response::{Response, Status};

            let status = Status::ok(None, None, "Hello").unwrap();
            let res = Response::Status(status);

            conn.send(res).unwrap();
        }

        let mut server = PipelineServer::new(conn, service);
        if let Err(err) = futures::ready!(Future::poll(Pin::new(&mut server), cx)) {
            tracing::debug!("Connection error: {}", err);
        }

        Poll::Ready(())
    }
}