server: Fix re-entrant polling

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,87 +4,140 @@ 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 struct Connecting<C, F, S> {
+where
-    pub conn: Connection<C>,
+    C: AsyncRead + AsyncWrite + Unpin,
-    #[pin]
+    S: Service<Request, Response = Response>,
-    pub state: ConnectingState<F, S>,
+    S::Future: Send + 'static,
+{
+    pub state: Option<ConnectingState<C, F, S>>,
     pub protocol: Imap,
 }
 
-#[pin_project::pin_project(project = ConnectingStateProj)]
+pub enum ConnectingState<C, F, S>
-pub enum ConnectingState<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> Future for Connecting<C, F, S>
+impl<C, F, ME, S> ConnectingState<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_new_state(self, cx: &mut task::Context) -> (Self, Option<Poll<()>>) {
-
+        match self {
-    fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
+            ConnectingState::Waiting {
-        use tokio_tower::pipeline::Server as PipelineServer;
+                conn,
-
+                mut service_fut,
-        let mut this = self.project();
+            } => {
-
+                let service = match Pin::new(&mut service_fut).poll(cx) {
-        loop {
+                    Poll::Ready(Ok(service)) => service,
-            let next = match this.state.as_mut().project() {
+                    Poll::Ready(Err(err)) => {
-                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(());
+                        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
-                        this.conn.start_send_unpin((None, greeting)).unwrap();
+                    conn.start_send_unpin((None, greeting)).unwrap();
                 }
 
-                    ConnectingState::Ready { service }
+                (ConnectingState::Ready { conn, service }, None)
             }
-                ConnectingStateProj::Ready { service } => {
+            ConnectingState::Ready { conn, service } => (
-                    let server = PipelineServer::new(this.conn, PipelineService { inner: service });
+                ConnectingState::Serving {
-                    futures::pin_mut!(server);
+                    server: PipelineServer::new(conn, PipelineService { inner: service }),
-
+                },
-                    return server.poll(cx).map(|res| {
+                None,
-                        if let Err(err) = res {
+            ),
-                            tracing::debug!("Connection error: {}", err);
+            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(()))),
-
-            this.state.set(next);
         }
     }
 }
 
-struct PipelineService<S> {
+impl<C, F, ME, S> Future for Connecting<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,
+{
+    type Output = ();
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
+        loop {
+            let state = self.as_mut().state.take().unwrap();
+            let (next, res) = state.poll_new_state(cx);
+
+            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,
 }
 

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::Error: std::fmt::Display + Send + 'static,
-    S::Future: 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 {
+                    state: Some(conn::ConnectingState::Waiting {
                         conn: Connection::new(conn),
-                    state: conn::ConnectingState::Waiting { service_fut },
+                        service_fut,
+                    }),
                     protocol: this.protocol.clone(),
                 });
             } else {