Switch to my fork

Fixes #2

Cargo.toml

@@ -1,26 +1,33 @@
 [package]
 name = "boitalettres"
-version = "0.0.1"
+version = "0.1.0"
 license = "BSD-3-Clause"
 edition = "2021"
 
 [dependencies]
 bytes = "1.1"
+
+miette = "5.1"
 thiserror = "1.0"
+
 tracing = "0.1"
+tracing-futures = "0.2"
 
 # IMAP
-imap-codec = "0.5"
+imap-codec = { git = "https://github.com/superboum/imap-codec.git", branch = "v0.5.x" }
 
 # Async
 async-compat = "0.2"
+async-stream = "0.3"
 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 = "5.1", 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,62 @@
+use miette::{IntoDiagnostic, Result};
+
 use boitalettres::proto::{Request, Response};
 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;
+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};
 
-    Ok(Response::Status(
-        Status::ok(Some(req.tag), None, "Ok").map_err(|e| eyre::eyre!(e))?,
-    ))
+    let capabilities = vec![Capability::Imap4Rev1, Capability::Idle];
+    let body: Vec<Data> = vec![
+        Status::ok("Yeah")?.into(),
+        ImapData::Capability(capabilities).into(),
+    ];
+
+    Ok(Response::ok("Done")?.with_body(body))
 }
 
 #[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,4 @@
+pub mod errors;
 pub mod proto;
 pub mod server;
+mod util;

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.rs

@@ -1 +0,0 @@
-pub type Response = imap_codec::types::response::Response;

src/proto/res/body.rs

@@ -0,0 +1,100 @@
+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),
+}
+
+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),
+        }
+    }
+}
+
+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

@@ -0,0 +1,129 @@
+use imap_codec::types::{
+    core::{Tag, Text},
+    response::{Code as ImapCode, Status as ImapStatus},
+};
+
+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>,
+}
+
+impl Response {
+    pub fn status(code: StatusCode, msg: &str) -> Result<Response> {
+        Ok(Response {
+            status: Status::new(code, msg)?,
+            body: None,
+        })
+    }
+
+    pub fn ok(msg: &str) -> Result<Response> {
+        Self::status(StatusCode::Ok, msg)
+    }
+
+    pub fn no(msg: &str) -> Result<Response> {
+        Self::status(StatusCode::No, msg)
+    }
+
+    pub fn bad(msg: &str) -> Result<Response> {
+        Self::status(StatusCode::Bad, msg)
+    }
+
+    pub fn bye(msg: &str) -> Result<Response> {
+        Self::status(StatusCode::Bye, msg)
+    }
+}
+
+impl Response {
+    pub fn with_extra_code(mut self, extra: ImapCode) -> Self {
+        self.status.extra = Some(extra);
+        self
+    }
+
+    pub fn with_body(mut self, body: impl Into<Body>) -> Self {
+        self.body = Some(body.into());
+        self
+    }
+}
+
+impl Response {
+    pub fn split(self) -> (Option<Body>, Status) {
+        (self.body, self.status)
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct Status {
+    pub(crate) code: StatusCode,
+    pub(crate) extra: Option<ImapCode>,
+    pub(crate) text: Text,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum StatusCode {
+    Ok,
+    No,
+    Bad,
+    PreAuth,
+    Bye,
+}
+
+impl Status {
+    fn new(code: StatusCode, msg: &str) -> Result<Status> {
+        Ok(Status {
+            code,
+            extra: None,
+            text: msg.try_into().map_err(Error::text)?,
+        })
+    }
+
+    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 {
+                tag,
+                code: self.extra,
+                text: self.text,
+            },
+            StatusCode::No => ImapStatus::No {
+                tag,
+                code: self.extra,
+                text: self.text,
+            },
+            StatusCode::Bad => ImapStatus::Bad {
+                tag,
+                code: self.extra,
+                text: self.text,
+            },
+            StatusCode::PreAuth => ImapStatus::PreAuth {
+                code: self.extra,
+                text: self.text,
+            },
+            StatusCode::Bye => ImapStatus::Bye {
+                code: self.extra,
+                text: self.text,
+            },
+        }
+    }
+}

src/proto/res/stream.rs

@@ -0,0 +1,23 @@
+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 (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();
+    }
+}

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,162 @@ 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>,
+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,
     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.command.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

@@ -2,11 +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>;
@@ -14,8 +16,12 @@ 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,
 }
 
@@ -23,7 +29,10 @@ 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(),
         }
     }
@@ -46,7 +55,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 {
@@ -66,7 +83,7 @@ where
 
             *this.read_buf = input.into();
 
-            let req = command;
+            let req = Request { command };
             return Poll::Ready(Some(Ok(req)));
         }
     }
@@ -77,33 +94,35 @@ 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 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();
 
         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 +136,14 @@ where
         Poll::Ready(Ok(()))
     }
 
-    fn start_send(self: Pin<&mut Self>, item: Response) -> Result<(), Self::Error> {
-        debug_assert!(self.write_buf.is_empty());
-        self.get_mut().send(item)?;
+    fn start_send(
+        mut self: Pin<&mut Self>,
+        (tag, res): (Option<Tag>, Response),
+    ) -> Result<(), Self::Error> {
+        use crate::proto::res::stream::response_stream;
+
+        tracing::debug!(?tag, ?res, "transport.start_send");
+        self.outbox.push(Box::pin(response_stream(res, tag)));
 
         Ok(())
     }
@@ -142,3 +166,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),
+            }
+        }
+    }
+}