Support closing connection after a Response

Closes #11

Cargo.toml

@@ -1,13 +1,13 @@
 [package]
 name = "boitalettres"
-version = "0.0.1"
+version = "0.1.0"
 license = "BSD-3-Clause"
 edition = "2021"
 
 [dependencies]
 bytes = "1.1"
 
-miette = "4.7"
+miette = "5.1"
 thiserror = "1.0"
 
 tracing = "0.1"
@@ -18,6 +18,7 @@ imap-codec = "0.5"
 
 # Async
 async-compat = "0.2"
+async-stream = "0.3"
 futures = "0.3"
 pin-project = "1.0"
 
@@ -26,7 +27,7 @@ tokio-tower = "0.6"
 tower = { version = "0.4", features = ["full"] }
 
 [dev-dependencies]
-miette = { version = "4.7", features = ["fancy"] }
+miette = { version = "5.1", features = ["fancy"] }
 
 tracing-subscriber = "0.3"
 console-subscriber = "0.1"

examples/simple.rs

@@ -4,13 +4,16 @@ use boitalettres::proto::{Request, Response};
 use boitalettres::server::accept::addr::{AddrIncoming, AddrStream};
 use boitalettres::server::Server;
 
-async fn handle_req(req: Request) -> Result<Response> {
-    use imap_codec::types::response::{Capability, Data};
+async fn handle_req(_req: Request) -> Result<Response> {
+    use imap_codec::types::response::{Capability, Data as ImapData};
 
-    tracing::debug!("Got request: {:#?}", req);
+    use boitalettres::proto::res::{body::Data, Status};
 
     let capabilities = vec![Capability::Imap4Rev1, Capability::Idle];
-    let body = vec![Data::Capability(capabilities)];
+    let body: Vec<Data> = vec![
+        Status::ok("Yeah")?.into(),
+        ImapData::Capability(capabilities).into(),
+    ];
 
     Ok(Response::ok("Done")?.with_body(body))
 }

src/lib.rs

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

src/proto/body.rs

@@ -1,32 +0,0 @@
-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,6 +1,5 @@
 pub use self::req::Request;
 pub use self::res::Response;
 
-pub mod body;
 pub mod req;
 pub mod res;

src/proto/req.rs

@@ -1,3 +0,0 @@
-use imap_codec::types::command::Command;
-
-pub type Request = Command;

src/proto/req/mod.rs

@@ -0,0 +1,6 @@
+use imap_codec::types::command::Command;
+
+#[derive(Debug)]
+pub struct Request {
+    pub command: Command,
+}

src/proto/res/body.rs

@@ -0,0 +1,102 @@
+use std::{fmt, io};
+
+use futures::prelude::*;
+use futures::stream::BoxStream;
+use imap_codec::codec::Encode;
+use imap_codec::types::response::{Data as ImapData, Status as ImapStatus};
+
+use super::Status;
+
+pub enum Body {
+    Once(Vec<Data>),
+    Stream(BoxStream<'static, Data>),
+}
+
+impl Body {
+    pub fn from_stream<St: Stream<Item = Data> + Send + 'static>(stream: St) -> Self {
+        Body::Stream(stream.boxed())
+    }
+}
+
+impl Body {
+    pub(crate) fn into_stream(self) -> BoxStream<'static, Data> {
+        match self {
+            Body::Once(data) => futures::stream::iter(data).boxed(),
+            Body::Stream(stream) => stream,
+        }
+    }
+}
+
+impl FromIterator<Data> for Body {
+    fn from_iter<T: IntoIterator<Item = Data>>(iter: T) -> Self {
+        Body::Once(Vec::from_iter(iter))
+    }
+}
+
+impl FromIterator<ImapData> for Body {
+    fn from_iter<T: IntoIterator<Item = ImapData>>(iter: T) -> Self {
+        Body::from_iter(iter.into_iter().map(Data::Data))
+    }
+}
+
+impl From<Vec<Data>> for Body {
+    fn from(body: Vec<Data>) -> Self {
+        Body::from_iter(body)
+    }
+}
+
+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])
+    }
+}
+
+impl fmt::Debug for Body {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Body::Once(ref data) => f.debug_struct("Body::Once").field("data", data).finish(),
+            Body::Stream(_) => f.debug_struct("Body::Stream").finish_non_exhaustive(),
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+pub enum Data {
+    Data(ImapData),
+    Status(ImapStatus),
+    Close,
+}
+
+impl Encode for Data {
+    fn encode(&self, writer: &mut impl io::Write) -> std::io::Result<()> {
+        match self {
+            Data::Data(ref data) => data.encode(writer),
+            Data::Status(ref status) => status.encode(writer),
+            Data::Close => Ok(()),
+        }
+    }
+}
+
+impl From<ImapData> for Data {
+    fn from(data: ImapData) -> Self {
+        Data::Data(data)
+    }
+}
+
+impl From<ImapStatus> for Data {
+    fn from(status: ImapStatus) -> Self {
+        Data::Status(status)
+    }
+}
+
+impl From<Status> for Data {
+    fn from(status: Status) -> Self {
+        status.into_imap(None).into()
+    }
+}

src/proto/res/mod.rs

@@ -3,13 +3,17 @@ use imap_codec::types::{
     response::{Code as ImapCode, Status as ImapStatus},
 };
 
-use super::body::Body;
+use self::body::Body;
 use crate::errors::{Error, Result};
 
+pub mod body;
+pub(crate) mod stream;
+
 #[derive(Debug)]
 pub struct Response {
     pub(crate) status: Status,
     pub(crate) body: Option<Body>,
+    pub(crate) close: bool,
 }
 
 impl Response {
@@ -17,6 +21,7 @@ impl Response {
         Ok(Response {
             status: Status::new(code, msg)?,
             body: None,
+            close: false,
         })
     }
 
@@ -47,6 +52,17 @@ impl Response {
         self.body = Some(body.into());
         self
     }
+
+    pub fn close(mut self, close: bool) -> Self {
+        self.close = close;
+        self
+    }
+}
+
+impl Response {
+    pub fn split(self) -> (Option<Body>, Status) {
+        (self.body, self.status)
+    }
 }
 
 #[derive(Debug, Clone)]
@@ -74,6 +90,22 @@ impl Status {
         })
     }
 
+    pub fn ok(msg: &str) -> Result<Status> {
+        Self::new(StatusCode::Ok, msg)
+    }
+
+    pub fn no(msg: &str) -> Result<Status> {
+        Self::new(StatusCode::No, msg)
+    }
+
+    pub fn bad(msg: &str) -> Result<Status> {
+        Self::new(StatusCode::Bad, msg)
+    }
+
+    pub fn bye(msg: &str) -> Result<Status> {
+        Self::new(StatusCode::Bye, msg)
+    }
+
     pub(crate) fn into_imap(self, tag: Option<Tag>) -> ImapStatus {
         match self.code {
             StatusCode::Ok => ImapStatus::Ok {

src/proto/res/stream.rs

@@ -0,0 +1,28 @@
+use futures::prelude::*;
+use imap_codec::types::core::Tag;
+
+use super::body::Data;
+use super::Response;
+
+pub fn response_stream(res: Response, tag: Option<Tag>) -> impl Stream<Item = Data> {
+    let close = res.close;
+    let (body, status) = res.split();
+    let body = body.map(|body| body.into_stream());
+
+    async_stream::stream! {
+        if let Some(body) = body {
+            for await item in body {
+                tracing::trace!(?item, "response_stream.yield");
+                yield item;
+            }
+        }
+
+        let item = status.into_imap(tag);
+        tracing::trace!(?item, "response_stream.yield");
+        yield item.into();
+
+        if close {
+            yield Data::Close;
+        }
+    }
+}

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

@@ -4,35 +4,111 @@ 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>,
+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 = Response::ok("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,
@@ -40,51 +116,28 @@ where
 {
     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 futures::SinkExt;
-
-                        let greeting = Response::ok("Hello").unwrap(); // "Hello" is a valid
-                                                                       // greeting
-                        this.conn.start_send_unpin((None, greeting)).unwrap();
-                    }
-
-                    ConnectingState::Ready { service }
-                }
-                ConnectingStateProj::Ready { service } => {
-                    let server = PipelineServer::new(this.conn, PipelineService { inner: 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;
+            }
         }
     }
 }
 
-struct PipelineService<S> {
+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,
 }
 
@@ -104,7 +157,7 @@ where
     fn call(&mut self, req: Request) -> Self::Future {
         use futures::{FutureExt, TryFutureExt};
 
-        let tag = req.tag.clone();
+        let tag = req.command.tag.clone();
 
         self.inner.call(req).map_ok(|res| (Some(tag), res)).boxed()
     }

src/server/mod.rs

@@ -46,8 +46,8 @@ where
     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,
 {
@@ -66,8 +66,10 @@ where
                 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

@@ -2,12 +2,13 @@ use std::pin::Pin;
 use std::task::{self, Poll};
 
 use bytes::BytesMut;
-use futures::io::{AsyncRead, AsyncWrite};
-use futures::sink::Sink;
-use futures::stream::Stream;
+use futures::prelude::*;
+use futures::stream::BoxStream;
 use imap_codec::types::core::Tag;
 
+use crate::proto::res::body::Data;
 use crate::proto::{Request, Response};
+use crate::util::stream::ConcatAll;
 
 type Error = tower::BoxError;
 type Result<T, E = Error> = std::result::Result<T, E>;
@@ -15,17 +16,28 @@ type Result<T, E = Error> = std::result::Result<T, E>;
 #[pin_project::pin_project]
 pub struct Connection<C> {
     #[pin]
-    conn: C,
+    pub conn: C,
+
     read_buf: BytesMut,
+
+    #[pin]
+    outbox: ConcatAll<BoxStream<'static, Data>>,
     write_buf: BytesMut,
+
+    close: bool,
 }
 
 impl<C> Connection<C> {
     pub fn new(conn: C) -> Self {
         Self {
             conn,
-            read_buf: BytesMut::new(),
+
+            read_buf: BytesMut::with_capacity(1024),
+
+            outbox: ConcatAll::new(),
             write_buf: BytesMut::new(),
+
+            close: false,
         }
     }
 }
@@ -47,7 +59,15 @@ where
                 Err(e) if e.is_incomplete() => {
                     let mut buf = [0u8; 256];
 
-                    let read = futures::ready!(this.conn.as_mut().poll_read(cx, &mut buf))?;
+                    tracing::trace!("transport.poll_read");
+                    // let read = futures::ready!(this.conn.as_mut().poll_read(cx, &mut buf))?;
+                    let read = match this.conn.as_mut().poll_read(cx, &mut buf) {
+                        Poll::Ready(res) => res?,
+                        Poll::Pending => {
+                            tracing::trace!("transport.pending");
+                            return Poll::Pending;
+                        }
+                    };
                     tracing::trace!(read = read, "transport.poll_next");
 
                     if read == 0 {
@@ -67,7 +87,7 @@ where
 
             *this.read_buf = input.into();
 
-            let req = command;
+            let req = Request { command };
             return Poll::Ready(Some(Ok(req)));
         }
     }
@@ -78,15 +98,37 @@ where
     C: AsyncWrite,
 {
     fn poll_flush_buffer(self: Pin<&mut Self>, cx: &mut task::Context) -> Poll<Result<()>> {
-        use bytes::Buf;
+        use bytes::{Buf, BufMut};
+        use imap_codec::codec::Encode;
 
         let mut this = self.project();
 
+        tracing::debug!(size = this.outbox.len(), "transport.flush_outbox");
+        let mut writer = this.write_buf.writer();
+
+        while let Poll::Ready(Some(data)) = this.outbox.as_mut().poll_next(cx) {
+            tracing::trace!(?data, "transport.write_buf");
+
+            if let Data::Close = data {
+                *this.close = true;
+            }
+
+            if let Err(err) = data.encode(&mut writer) {
+                tracing::error!(?err, "transport.encode_error");
+                return Poll::Ready(Err(Box::new(err)));
+            }
+        }
+
         tracing::debug!(size = this.write_buf.len(), "transport.flush_buffer");
         while !this.write_buf.is_empty() {
             let written = futures::ready!(this.conn.as_mut().poll_write(cx, this.write_buf))?;
             this.write_buf.advance(written);
         }
+        this.write_buf.clear();
+
+        if *this.close {
+            futures::ready!(this.conn.as_mut().poll_close(cx))?;
+        }
 
         Poll::Ready(Ok(()))
     }
@@ -107,23 +149,13 @@ where
     }
 
     fn start_send(
-        self: Pin<&mut Self>,
+        mut self: Pin<&mut Self>,
         (tag, res): (Option<Tag>, Response),
     ) -> Result<(), Self::Error> {
-        use bytes::BufMut;
-        use imap_codec::codec::Encode;
+        use crate::proto::res::stream::response_stream;
 
-        debug_assert!(self.write_buf.is_empty());
-
-        let write_buf = &mut self.get_mut().write_buf;
-        let mut writer = write_buf.writer();
-
-        let body = res.body.into_iter().flat_map(|body| body.into_data());
-        for data in body {
-            data.encode(&mut writer)?;
-        }
-
-        res.status.into_imap(tag).encode(&mut writer)?;
+        tracing::debug!(?tag, ?res, "transport.start_send");
+        self.outbox.push(Box::pin(response_stream(res, tag)));
 
         Ok(())
     }
@@ -146,3 +178,35 @@ where
         Poll::Ready(Ok(()))
     }
 }
+
+impl<C> Sink<Response> for Connection<C>
+where
+    Self: Sink<(Option<Tag>, Response)>,
+{
+    type Error = <Connection<C> as Sink<(Option<Tag>, Response)>>::Error;
+
+    fn poll_ready(
+        self: Pin<&mut Self>,
+        cx: &mut task::Context<'_>,
+    ) -> Poll<Result<(), Self::Error>> {
+        <Self as Sink<(Option<Tag>, Response)>>::poll_ready(self, cx)
+    }
+
+    fn start_send(self: Pin<&mut Self>, item: Response) -> Result<(), Self::Error> {
+        <Self as Sink<(Option<Tag>, Response)>>::start_send(self, (None, item))
+    }
+
+    fn poll_flush(
+        self: Pin<&mut Self>,
+        cx: &mut task::Context<'_>,
+    ) -> Poll<Result<(), Self::Error>> {
+        <Self as Sink<(Option<Tag>, Response)>>::poll_flush(self, cx)
+    }
+
+    fn poll_close(
+        self: Pin<&mut Self>,
+        cx: &mut task::Context<'_>,
+    ) -> Poll<Result<(), Self::Error>> {
+        <Self as Sink<(Option<Tag>, Response)>>::poll_close(self, cx)
+    }
+}

src/util/mod.rs

@@ -0,0 +1,2 @@
+pub mod buf;
+pub mod stream;

src/util/stream.rs

@@ -0,0 +1,54 @@
+use std::fmt;
+use std::pin::Pin;
+use std::task::{self, Poll};
+
+use futures::prelude::*;
+use futures::stream::{FuturesOrdered, Stream, StreamFuture};
+
+/// [`SelectAll`](futures::stream::SelectAll) but ordered
+#[pin_project::pin_project]
+pub struct ConcatAll<St: Stream + Unpin> {
+    #[pin]
+    inner: FuturesOrdered<StreamFuture<St>>,
+}
+
+impl<St: Stream + Unpin> ConcatAll<St> {
+    pub fn new() -> Self {
+        Self {
+            inner: FuturesOrdered::new(),
+        }
+    }
+
+    pub fn len(&self) -> usize {
+        self.inner.len()
+    }
+
+    pub fn push(&mut self, stream: St) {
+        use futures::StreamExt;
+
+        self.inner.push(stream.into_future());
+    }
+}
+
+impl<St: Stream + Unpin> fmt::Debug for ConcatAll<St> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("ConcatAll").finish()
+    }
+}
+
+impl<St: Stream + Unpin> Stream for ConcatAll<St> {
+    type Item = St::Item;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Option<Self::Item>> {
+        loop {
+            match futures::ready!(self.inner.poll_next_unpin(cx)) {
+                Some((Some(item), remaining)) => {
+                    self.push(remaining);
+                    return Poll::Ready(Some(item));
+                }
+                Some((None, _)) => {}
+                _ => return Poll::Ready(None),
+            }
+        }
+    }
+}