More flexible responses

Cargo.toml

@@ -6,8 +6,12 @@ edition = "2021"
 
 [dependencies]
 bytes = "1.1"
+
+miette = "4.7"
 thiserror = "1.0"
+
 tracing = "0.1"
+tracing-futures = "0.2"
 
 # IMAP
 imap-codec = "0.5"
@@ -16,11 +20,13 @@ imap-codec = "0.5"
 async-compat = "0.2"
 futures = "0.3"
 pin-project = "1.0"
+
 tokio = { version = "1.18", features = ["full"] }
 tokio-tower = "0.6"
 tower = { version = "0.4", features = ["full"] }
 
 [dev-dependencies]
-eyre = "0.6"
+miette = { version = "4.7", features = ["fancy"] }
+
 tracing-subscriber = "0.3"
 console-subscriber = "0.1"

doc/draft.md

@@ -10,6 +10,12 @@ This example is meant to show basic service-based IMAP server with boitalettres
 $ cargo run --example simple
 ```
 
+If you want to trace this `simple` example with the [`tokio-console`](https://github.com/tokio-rs/console/tree/main/tokio-console) utility, run:
+
+```shell
+$ RUSTFLAGS="tokio_unstable" cargo run --example simple --features "tokio/tracing"
+```
+
 - [Basic python testing script](../scripts/test_imap.py)
 
 ```shell

examples/simple.rs

@@ -1,55 +1,63 @@
-use boitalettres::proto::{Request, Response};
+use miette::{IntoDiagnostic, Result};
+
+use boitalettres::proto::{Request, Response as BalResponse};
 use boitalettres::server::accept::addr::{AddrIncoming, AddrStream};
 use boitalettres::server::Server;
 
-async fn handle_req(req: Request) -> eyre::Result<Response> {
-    use imap_codec::types::response::Status;
+use imap_codec::types::response::{Response, Data, Status, Capability};
+
+async fn handle_req(req: Request) -> Result<BalResponse> {
 
     tracing::debug!("Got request: {:#?}", req);
 
-    Ok(Response::Status(
-        Status::ok(Some(req.tag), None, "Ok").map_err(|e| eyre::eyre!(e))?,
-    ))
+    let capabilities = vec![Capability::Imap4Rev1, Capability::Idle];
+    let res = vec![
+        Response::Data(Data::Capability(capabilities)),
+        Response::Status(Status::ok(Some(req.tag), None, "Done").unwrap()),
+    ];
+
+    Ok(res)
 }
 
 #[tokio::main]
-async fn main() -> eyre::Result<()> {
+async fn main() -> Result<()> {
     setup_logging();
 
-    let incoming = AddrIncoming::new("127.0.0.1:4567").await?;
+    let incoming = AddrIncoming::new("127.0.0.1:4567")
+        .await
+        .into_diagnostic()?;
 
     let make_service = tower::service_fn(|addr: &AddrStream| {
         tracing::debug!(remote_addr = %addr.remote_addr, local_addr = %addr.local_addr, "accept");
 
-        let service = tower::ServiceBuilder::new()
-            .buffer(16)
-            .service_fn(handle_req);
+        let service = tower::ServiceBuilder::new().service_fn(handle_req);
 
         futures::future::ok::<_, std::convert::Infallible>(service)
     });
 
     let server = Server::new(incoming).serve(make_service);
-    let _ = server.await?;
 
-    Ok(())
+    server.await.map_err(Into::into)
 }
 
 // Don't mind this, this is just for debugging.
 fn setup_logging() {
     use tracing_subscriber::prelude::*;
 
-    tracing_subscriber::registry()
-        .with(console_subscriber::spawn())
-        .with(
-            tracing_subscriber::fmt::layer().with_filter(
-                tracing_subscriber::filter::Targets::new()
-                    .with_default(tracing::Level::DEBUG)
-                    .with_target("boitalettres", tracing::Level::TRACE)
-                    .with_target("simple", tracing::Level::TRACE)
-                    .with_target("tower", tracing::Level::TRACE)
-                    .with_target("tokio_tower", tracing::Level::TRACE)
-                    .with_target("mio", tracing::Level::TRACE),
-            ),
-        )
-        .init();
+    let registry = tracing_subscriber::registry().with(
+        tracing_subscriber::fmt::layer().with_filter(
+            tracing_subscriber::filter::Targets::new()
+                .with_default(tracing::Level::DEBUG)
+                .with_target("boitalettres", tracing::Level::TRACE)
+                .with_target("simple", tracing::Level::TRACE)
+                .with_target("tower", tracing::Level::TRACE)
+                .with_target("tokio_tower", tracing::Level::TRACE)
+                .with_target("mio", tracing::Level::TRACE),
+        ),
+    );
+
+    #[cfg(tokio_unstable)]
+    let registry = registry.with(console_subscriber::spawn());
+
+    registry.init();
 }

src/errors.rs

@@ -0,0 +1,33 @@
+use miette::Diagnostic;
+
+type BoxedError = tower::BoxError;
+
+#[derive(Debug, thiserror::Error, Diagnostic)]
+pub enum Error {
+    #[error("Error occured when accepting new connections")]
+    #[diagnostic(code(boitalettres::accept))]
+    Accept(#[source] BoxedError),
+
+    #[error("Error occured on service creation")]
+    #[diagnostic(code(boitalettres::make_service))]
+    MakeService(#[source] BoxedError),
+
+    #[error("{0}")]
+    Text(String),
+}
+
+impl Error {
+    pub(crate) fn accept<E: Into<BoxedError>>(err: E) -> Error {
+        Error::Accept(err.into())
+    }
+
+    pub(crate) fn make_service<E: Into<BoxedError>>(err: E) -> Error {
+        Error::MakeService(err.into())
+    }
+
+    pub(crate) fn text<E: Into<String>>(err: E) -> Error {
+        Error::Text(err.into())
+    }
+}
+
+pub type Result<T, E = Error> = std::result::Result<T, E>;

src/lib.rs

@@ -1,2 +1,3 @@
+pub mod errors;
 pub mod proto;
 pub mod server;

src/proto/body.rs

@@ -0,0 +1,32 @@
+use imap_codec::types::response::Data as ImapData;
+
+#[derive(Debug)]
+pub enum Body {
+    Once(Vec<ImapData>),
+}
+
+impl Body {
+    pub(crate) fn into_data(self) -> Vec<ImapData> {
+        match self {
+            Body::Once(data) => data,
+        }
+    }
+}
+
+impl FromIterator<ImapData> for Body {
+    fn from_iter<T: IntoIterator<Item = ImapData>>(iter: T) -> Self {
+        Body::Once(Vec::from_iter(iter))
+    }
+}
+
+impl From<Vec<ImapData>> for Body {
+    fn from(data: Vec<ImapData>) -> Self {
+        Body::from_iter(data)
+    }
+}
+
+impl From<ImapData> for Body {
+    fn from(data: ImapData) -> Self {
+        Body::from_iter([data])
+    }
+}

src/proto/mod.rs

@@ -1,5 +1,6 @@
 pub use self::req::Request;
 pub use self::res::Response;
 
+pub mod body;
 pub mod req;
 pub mod res;

src/proto/res.rs

@@ -1 +1,3 @@
-pub type Response = imap_codec::types::response::Response;
+use imap_codec::types::response::Response as ImapResponse;
+pub type Response = Vec<ImapResponse>;
+

src/server/accept/addr.rs

@@ -62,6 +62,14 @@ impl AsyncRead for AddrStream {
     ) -> Poll<std::io::Result<usize>> {
         self.project().stream.poll_read(cx, buf)
     }
+
+    fn poll_read_vectored(
+        self: Pin<&mut Self>,
+        cx: &mut task::Context<'_>,
+        bufs: &mut [std::io::IoSliceMut<'_>],
+    ) -> Poll<std::io::Result<usize>> {
+        self.project().stream.poll_read_vectored(cx, bufs)
+    }
 }
 
 impl AsyncWrite for AddrStream {

src/server/conn.rs

@@ -3,91 +3,163 @@ use std::pin::Pin;
 use std::task::{self, Poll};
 
 use futures::io::{AsyncRead, AsyncWrite};
+use imap_codec::types::core::Tag;
+use tokio_tower::pipeline::Server as PipelineServer;
 use tower::Service;
 
 use super::pipeline::Connection;
 use super::Imap;
 use crate::proto::{Request, Response};
 
-#[pin_project::pin_project]
-pub struct Connecting<C, F, S> {
-    pub conn: Connection<C>,
-    #[pin]
-    pub state: ConnectingState<F, S>,
+use imap_codec::types::response::{Response as ImapResponse, Status};
+
+pub struct Connecting<C, F, S>
+where
+    C: AsyncRead + AsyncWrite + Unpin,
+    S: Service<Request, Response = Response>,
+    S::Future: Send + 'static,
+{
+    pub state: Option<ConnectingState<C, F, S>>,
     pub protocol: Imap,
 }
 
-#[pin_project::pin_project(project = ConnectingStateProj)]
-pub enum ConnectingState<F, S> {
+pub enum ConnectingState<C, F, S>
+where
+    C: AsyncRead + AsyncWrite + Unpin,
+    S: Service<Request, Response = Response>,
+    S::Future: Send + 'static,
+{
     Waiting {
-        #[pin]
+        conn: Connection<C>,
         service_fut: F,
     },
     Ready {
+        conn: Connection<C>,
         service: S,
     },
+    Serving {
+        server: PipelineServer<Connection<C>, PipelineService<S>>,
+    },
+    Finished,
+}
+
+impl<C, F, ME, S> ConnectingState<C, F, S>
+where
+    C: AsyncRead + AsyncWrite + Unpin,
+    F: Future<Output = std::result::Result<S, ME>> + Unpin,
+    ME: std::fmt::Display,
+    S: Service<Request, Response = Response>,
+    S::Future: Send + 'static,
+    S::Error: std::fmt::Display,
+{
+    fn poll_new_state(self, cx: &mut task::Context) -> (Self, Option<Poll<()>>) {
+        match self {
+            ConnectingState::Waiting {
+                conn,
+                mut service_fut,
+            } => {
+                let service = match Pin::new(&mut service_fut).poll(cx) {
+                    Poll::Ready(Ok(service)) => service,
+                    Poll::Ready(Err(err)) => {
+                        tracing::error!("Connection error: {}", err);
+                        return (
+                            ConnectingState::Waiting { conn, service_fut },
+                            Some(Poll::Ready(())),
+                        );
+                    }
+                    Poll::Pending => {
+                        return (
+                            ConnectingState::Waiting { conn, service_fut },
+                            Some(Poll::Pending),
+                        )
+                    }
+                };
+
+                let mut conn = conn;
+
+                // TODO: Properly handle server greeting
+                {
+                    use futures::SinkExt;
+
+                    let greeting = vec![ImapResponse::Status(Status::greeting(None, "Hello").unwrap())]; // "Hello" is a valid greeting
+                    conn.start_send_unpin((None, greeting)).unwrap();
+                }
+
+                (ConnectingState::Ready { conn, service }, None)
+            }
+            ConnectingState::Ready { conn, service } => (
+                ConnectingState::Serving {
+                    server: PipelineServer::new(conn, PipelineService { inner: service }),
+                },
+                None,
+            ),
+            ConnectingState::Serving { mut server } => match Pin::new(&mut server).poll(cx) {
+                Poll::Ready(Ok(_)) => (ConnectingState::Finished, Some(Poll::Ready(()))),
+                Poll::Ready(Err(err)) => {
+                    tracing::debug!("Connecting error: {}", err);
+                    (ConnectingState::Finished, Some(Poll::Ready(())))
+                }
+                Poll::Pending => (ConnectingState::Serving { server }, Some(Poll::Pending)),
+            },
+            ConnectingState::Finished => (self, Some(Poll::Ready(()))),
+        }
+    }
 }
 
 impl<C, F, ME, S> Future for Connecting<C, F, S>
 where
     C: AsyncRead + AsyncWrite + Unpin,
-    F: Future<Output = std::result::Result<S, ME>>,
+    F: Future<Output = std::result::Result<S, ME>> + Unpin,
     ME: std::fmt::Display,
     S: Service<Request, Response = Response>,
+    S::Future: Send + 'static,
     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;
-
-        let mut this = self.project();
-
+    fn poll(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
         loop {
-            let next = match this.state.as_mut().project() {
-                ConnectingStateProj::Waiting { service_fut } => {
-                    let service = match futures::ready!(service_fut.poll(cx)) {
-                        Ok(service) => service,
-                        Err(err) => {
-                            tracing::error!("Connection error: {}", err);
-                            return Poll::Ready(());
-                        }
-                    };
+            let state = self.as_mut().state.take().unwrap();
+            let (next, res) = state.poll_new_state(cx);
 
-                    // TODO: Properly handle server greeting
-                    {
-                        use imap_codec::types::response::{Response, Status};
-
-                        let status = match Status::ok(None, None, "Hello") {
-                            Ok(status) => status,
-                            Err(err) => {
-                                tracing::error!("Connection error: {}", err);
-                                return Poll::Ready(());
-                            }
-                        };
-                        let res = Response::Status(status);
-
-                        if let Err(err) = this.conn.send(res) {
-                            tracing::error!("Connection error: {}", err);
-                            return Poll::Ready(());
-                        };
-                    }
-
-                    ConnectingState::Ready { service }
-                }
-                ConnectingStateProj::Ready { service } => {
-                    let server = PipelineServer::new(this.conn, service);
-                    futures::pin_mut!(server);
-
-                    return server.poll(cx).map(|res| {
-                        if let Err(err) = res {
-                            tracing::debug!("Connection error: {}", err);
-                        }
-                    });
-                }
-            };
-
-            this.state.set(next);
+            self.state = Some(next);
+            if let Some(res) = res {
+                return res;
+            }
         }
     }
 }
+
+impl<C, F, S> Unpin for Connecting<C, F, S>
+where
+    C: AsyncRead + AsyncWrite + Unpin,
+    S: Service<Request, Response = Response>,
+    S::Future: Send + 'static,
+{
+}
+
+pub struct PipelineService<S> {
+    inner: S,
+}
+
+impl<S> Service<Request> for PipelineService<S>
+where
+    S: Service<Request>,
+    S::Future: Send + 'static,
+{
+    type Response = (Option<Tag>, S::Response);
+    type Error = S::Error;
+    type Future = futures::future::BoxFuture<'static, Result<Self::Response, Self::Error>>;
+
+    fn poll_ready(&mut self, cx: &mut task::Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.inner.poll_ready(cx)
+    }
+
+    fn call(&mut self, req: Request) -> Self::Future {
+        use futures::{FutureExt, TryFutureExt};
+
+        let tag = req.tag.clone();
+
+        self.inner.call(req).map_ok(|res| (Some(tag), res)).boxed()
+    }
+}

src/server/mod.rs

@@ -5,6 +5,7 @@ use std::task::{self, Poll};
 use futures::io::{AsyncRead, AsyncWrite};
 use imap_codec::types::response::Capability;
 
+use crate::errors::{Error, Result};
 use crate::proto::{Request, Response};
 use accept::Accept;
 use pipeline::Connection;
@@ -15,14 +16,6 @@ mod conn;
 mod pipeline;
 mod service;
 
-#[derive(Debug, thiserror::Error)]
-pub enum Error<A> {
-    #[error("Error occured when accepting new connections")]
-    Accept(#[source] A),
-    #[error("Error occured on service creation")]
-    MakeService(#[source] tower::BoxError),
-}
-
 #[derive(Debug, Default, Clone)]
 pub struct Imap {
     pub capabilities: Vec<Capability>,
@@ -50,30 +43,33 @@ impl<I, S> Future for Server<I, S>
 where
     I: Accept,
     I::Conn: AsyncRead + AsyncWrite + Unpin + Send + 'static,
+    I::Error: Send + Sync + 'static,
     S: MakeServiceRef<I::Conn, Request, Response = Response>,
     S::MakeError: Into<tower::BoxError> + std::fmt::Display,
-    S::Error: std::fmt::Display,
-    S::Future: Send + 'static,
+    S::Error: std::fmt::Display + Send + 'static,
+    S::Future: Unpin + Send + 'static,
     S::Service: Send + 'static,
+    <S::Service as tower::Service<Request>>::Future: Send + 'static,
 {
-    type Output = Result<(), Error<I::Error>>;
+    type Output = Result<()>;
 
     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)?;
+                let conn = conn.map_err(Error::accept)?;
 
                 futures::ready!(this.make_service.poll_ready_ref(cx))
-                    .map_err(Into::into)
-                    .map_err(Error::MakeService)?;
+                    .map_err(Error::make_service)?;
 
                 let service_fut = this.make_service.make_service_ref(&conn);
 
                 tokio::task::spawn(conn::Connecting {
-                    conn: Connection::new(conn),
-                    state: conn::ConnectingState::Waiting { service_fut },
+                    state: Some(conn::ConnectingState::Waiting {
+                        conn: Connection::new(conn),
+                        service_fut,
+                    }),
                     protocol: this.protocol.clone(),
                 });
             } else {

src/server/pipeline.rs

@@ -5,6 +5,8 @@ use bytes::BytesMut;
 use futures::io::{AsyncRead, AsyncWrite};
 use futures::sink::Sink;
 use futures::stream::Stream;
+use imap_codec::types::core::Tag;
+use imap_codec::types::response::Response as ImapResponse;
 
 use crate::proto::{Request, Response};
 
@@ -89,21 +91,9 @@ where
 
         Poll::Ready(Ok(()))
     }
-
-    pub(crate) fn send(&mut self, item: Response) -> Result<()> {
-        use bytes::BufMut;
-        use imap_codec::codec::Encode;
-
-        let mut writer = BufMut::writer(&mut self.write_buf);
-
-        tracing::debug!(item = ?item, "transport.send");
-        item.encode(&mut writer)?;
-
-        Ok(())
-    }
 }
 
-impl<C> Sink<Response> for Connection<C>
+impl<C> Sink<(Option<Tag>, Response)> for Connection<C>
 where
     C: AsyncWrite + Unpin,
 {
@@ -117,9 +107,25 @@ where
         Poll::Ready(Ok(()))
     }
 
-    fn start_send(self: Pin<&mut Self>, item: Response) -> Result<(), Self::Error> {
+    fn start_send(
+        self: Pin<&mut Self>,
+        (tag, res): (Option<Tag>, Response),
+    ) -> Result<(), Self::Error> {
+        use bytes::BufMut;
+        use imap_codec::codec::Encode;
+
         debug_assert!(self.write_buf.is_empty());
-        self.get_mut().send(item)?;
+
+        let write_buf = &mut self.get_mut().write_buf;
+        let mut writer = write_buf.writer();
+
+        for elem in res {
+            match elem {
+                ImapResponse::Status(status) => status.encode(&mut writer)?,
+                ImapResponse::Data(data) => data.encode(&mut writer)?,
+                _ => (),
+            }
+        }
 
         Ok(())
     }