server: Fix connection double-polling

src/server/mod.rs

@@ -4,7 +4,6 @@ use std::task::{self, Poll};
 
 use futures::io::{AsyncRead, AsyncWrite};
 use imap_codec::types::response::Capability;
-use tower::Service;
 
 use crate::proto::{Request, Response};
 use accept::Accept;
@@ -51,6 +50,7 @@ where
     S::MakeError: Into<Box<dyn std::error::Error + Send + Sync + 'static>> + std::fmt::Display,
     S::Error: std::fmt::Display,
     S::Future: Send + 'static,
+    S::Service: Send + 'static,
 {
     type Output = Result<(), Error<I::Error>>;
 
@@ -67,9 +67,9 @@ where
 
                 let service_fut = this.make_service.make_service_ref(&conn);
 
-                tokio::task::spawn(Connecting {
+                tokio::task::spawn(conn::Connecting {
                     conn,
-                    service_fut,
+                    state: conn::ConnectingState::Waiting { service_fut },
                     protocol: this.protocol.clone(),
                 });
             } else {
@@ -114,38 +114,67 @@ impl<I: Accept> Builder<I> {
     }
 }
 
-#[pin_project::pin_project]
+mod conn {
-struct Connecting<C, F> {
+    use std::future::Future;
-    conn: C,
+    use std::pin::Pin;
+    use std::task::{self, Poll};
+
+    use futures::io::{AsyncRead, AsyncWrite};
+    use tower::Service;
+
+    use super::Imap;
+    use crate::proto::{Request, Response};
+
+    #[pin_project::pin_project]
+    pub struct Connecting<C, F, S> {
+        pub conn: C,
+        #[pin]
+        pub state: ConnectingState<F, S>,
+        pub protocol: Imap,
+    }
+
+    #[pin_project::pin_project(project = ConnectingStateProj)]
+    pub enum ConnectingState<F, S> {
+        Waiting {
             #[pin]
             service_fut: F,
-    protocol: Imap,
+        },
-}
+        Ready {
+            service: S,
+        },
+    }
 
-impl<C, F, ME, S> Future for Connecting<C, F>
+    impl<C, F, ME, S> Future for Connecting<C, F, S>
-where
+    where
         C: AsyncRead + AsyncWrite + Unpin,
         F: Future<Output = std::result::Result<S, ME>>,
         ME: std::fmt::Display,
         S: Service<Request, Response = Response>,
         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;
 
-        use pipeline::Connection;
+            use super::pipeline::Connection;
 
-        let this = self.project();
+            let mut this = self.project();
 
-        let service = match futures::ready!(this.service_fut.poll(cx)) {
+            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::debug!("Connection error: {}", err);
                                 return Poll::Ready(());
                             }
                         };
+
+                        ConnectingState::Ready { service }
+                    }
+                    ConnectingStateProj::Ready { service } => {
                         let mut conn = Connection::new(this.conn);
 
                         // TODO: Properly handle server greeting
@@ -158,11 +187,19 @@ where
                             conn.send(res).unwrap();
                         }
 
-        let mut server = PipelineServer::new(conn, service);
+                        let server = PipelineServer::new(conn, service);
-        if let Err(err) = futures::ready!(Future::poll(Pin::new(&mut server), cx)) {
+                        futures::pin_mut!(server);
+
+                        if let Err(err) = futures::ready!(server.poll(cx)) {
                             tracing::debug!("Connection error: {}", err);
                         }
 
-        Poll::Ready(())
+                        return Poll::Ready(());
+                    }
+                };
+
+                this.state.set(next);
+            }
+        }
     }
 }