More flexible responses

Cargo.toml

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

examples/simple.rs

@@ -1,21 +1,22 @@
 use miette::{IntoDiagnostic, Result};
 
-use boitalettres::proto::{Request, Response};
+use boitalettres::proto::{Request, Response as BalResponse};
 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 as ImapData};
+use imap_codec::types::response::{Response, Data, Status, Capability};
 
-    use boitalettres::proto::res::{body::Data, Status};
+async fn handle_req(req: Request) -> Result<BalResponse> {
+
+    tracing::debug!("Got request: {:#?}", req);
 
     let capabilities = vec![Capability::Imap4Rev1, Capability::Idle];
-    let body: Vec<Data> = vec![
-        Status::ok("Yeah")?.into(),
-        ImapData::Capability(capabilities).into(),
+    let res = vec![
+        Response::Data(Data::Capability(capabilities)),
+        Response::Status(Status::ok(Some(req.tag), None, "Done").unwrap()),
     ];
 
-    Ok(Response::ok("Done")?.with_body(body))
+    Ok(res)
 }
 
 #[tokio::main]

src/lib.rs

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

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

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

src/proto/req/mod.rs

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

src/proto/res.rs

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

src/proto/res/body.rs

@@ -1,102 +0,0 @@
-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

@@ -1,136 +0,0 @@
-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>,
-    pub(crate) close: bool,
-}
-
-impl Response {
-    pub fn status(code: StatusCode, msg: &str) -> Result<Response> {
-        Ok(Response {
-            status: Status::new(code, msg)?,
-            body: None,
-            close: false,
-        })
-    }
-
-    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
-    }
-
-    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)]
-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

@@ -1,28 +0,0 @@
-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/conn.rs

@@ -11,6 +11,8 @@ use super::pipeline::Connection;
 use super::Imap;
 use crate::proto::{Request, Response};
 
+use imap_codec::types::response::{Response as ImapResponse, Status};
+
 pub struct Connecting<C, F, S>
 where
     C: AsyncRead + AsyncWrite + Unpin,
@@ -79,8 +81,7 @@ where
                 {
                     use futures::SinkExt;
 
-                    let greeting = Response::ok("Hello").unwrap(); // "Hello" is a valid
-                                                                   // greeting
+                    let greeting = vec![ImapResponse::Status(Status::greeting(None, "Hello").unwrap())]; // "Hello" is a valid greeting
                     conn.start_send_unpin((None, greeting)).unwrap();
                 }
 
@@ -157,7 +158,7 @@ where
     fn call(&mut self, req: Request) -> Self::Future {
         use futures::{FutureExt, TryFutureExt};
 
-        let tag = req.command.tag.clone();
+        let tag = req.tag.clone();
 
         self.inner.call(req).map_ok(|res| (Some(tag), res)).boxed()
     }

src/server/pipeline.rs

@@ -2,13 +2,13 @@ use std::pin::Pin;
 use std::task::{self, Poll};
 
 use bytes::BytesMut;
-use futures::prelude::*;
-use futures::stream::BoxStream;
+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::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>;
@@ -16,28 +16,17 @@ type Result<T, E = Error> = std::result::Result<T, E>;
 #[pin_project::pin_project]
 pub struct Connection<C> {
     #[pin]
-    pub conn: C,
-
+    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::with_capacity(1024),
-
-            outbox: ConcatAll::new(),
+            read_buf: BytesMut::new(),
             write_buf: BytesMut::new(),
-
-            close: false,
         }
     }
 }
@@ -59,15 +48,7 @@ where
                 Err(e) if e.is_incomplete() => {
                     let mut buf = [0u8; 256];
 
-                    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;
-                        }
-                    };
+                    let read = futures::ready!(this.conn.as_mut().poll_read(cx, &mut buf))?;
                     tracing::trace!(read = read, "transport.poll_next");
 
                     if read == 0 {
@@ -87,7 +68,7 @@ where
 
             *this.read_buf = input.into();
 
-            let req = Request { command };
+            let req = command;
             return Poll::Ready(Some(Ok(req)));
         }
     }
@@ -98,37 +79,15 @@ where
     C: AsyncWrite,
 {
     fn poll_flush_buffer(self: Pin<&mut Self>, cx: &mut task::Context) -> Poll<Result<()>> {
-        use bytes::{Buf, BufMut};
-        use imap_codec::codec::Encode;
+        use bytes::Buf;
 
         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(()))
     }
@@ -149,13 +108,24 @@ where
     }
 
     fn start_send(
-        mut self: Pin<&mut Self>,
+        self: Pin<&mut Self>,
         (tag, res): (Option<Tag>, Response),
     ) -> Result<(), Self::Error> {
-        use crate::proto::res::stream::response_stream;
+        use bytes::BufMut;
+        use imap_codec::codec::Encode;
 
-        tracing::debug!(?tag, ?res, "transport.start_send");
-        self.outbox.push(Box::pin(response_stream(res, tag)));
+        debug_assert!(self.write_buf.is_empty());
+
+        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(())
     }
@@ -178,35 +148,3 @@ 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

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

src/util/stream.rs

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