Fix poll_ready call for WebSockets upgrade (#1219)

* Fix poll_ready call for WebSockets upgrade

* Poll upgrade service from H1ServiceHandler too

actix-http/CHANGES.md

@@ -1,5 +1,11 @@
 # Changes
 
+## [1.0.xx] - 2019-12-xx
+
+### Fixed
+
+* Poll upgrade service's readiness from HTTP service handlers
+
 ## [1.0.0] - 2019-12-13
 
 ### Added

actix-http/src/h1/dispatcher.rs

@@ -66,7 +66,6 @@ where
     U::Error: fmt::Display,
 {
     Normal(InnerDispatcher<T, S, B, X, U>),
-    UpgradeReadiness(InnerDispatcher<T, S, B, X, U>, Request),
     Upgrade(U::Future),
     None,
 }
@@ -764,8 +763,16 @@ where
                             if let DispatcherState::Normal(inner) =
                                 std::mem::replace(&mut self.inner, DispatcherState::None)
                             {
-                                self.inner =
-                                    DispatcherState::UpgradeReadiness(inner, req);
+                                let mut parts = FramedParts::with_read_buf(
+                                    inner.io,
+                                    inner.codec,
+                                    inner.read_buf,
+                                );
+                                parts.write_buf = inner.write_buf;
+                                let framed = Framed::from_parts(parts);
+                                self.inner = DispatcherState::Upgrade(
+                                    inner.upgrade.unwrap().call((req, framed)),
+                                );
                                 return self.poll(cx);
                             } else {
                                 panic!()
@@ -815,35 +822,6 @@ where
                     }
                 }
             }
-            DispatcherState::UpgradeReadiness(ref mut inner, _) => {
-                let upgrade = inner.upgrade.as_mut().unwrap();
-                match upgrade.poll_ready(cx) {
-                    Poll::Ready(Ok(_)) => {
-                        if let DispatcherState::UpgradeReadiness(inner, req) =
-                            std::mem::replace(&mut self.inner, DispatcherState::None)
-                        {
-                            let mut parts = FramedParts::with_read_buf(
-                                inner.io,
-                                inner.codec,
-                                inner.read_buf,
-                            );
-                            parts.write_buf = inner.write_buf;
-                            let framed = Framed::from_parts(parts);
-                            self.inner = DispatcherState::Upgrade(
-                                inner.upgrade.unwrap().call((req, framed)),
-                            );
-                            self.poll(cx)
-                        } else {
-                            panic!()
-                        }
-                    }
-                    Poll::Pending => Poll::Pending,
-                    Poll::Ready(Err(e)) => {
-                        error!("Upgrade handler readiness check error: {}", e);
-                        Poll::Ready(Err(DispatchError::Upgrade))
-                    }
-                }
-            }
             DispatcherState::Upgrade(ref mut fut) => {
                 unsafe { Pin::new_unchecked(fut) }.poll(cx).map_err(|e| {
                     error!("Upgrade handler error: {}", e);

actix-http/src/h1/service.rs

@@ -72,7 +72,7 @@ where
         Request = (Request, Framed<TcpStream, Codec>),
         Response = (),
     >,
-    U::Error: fmt::Display,
+    U::Error: fmt::Display + Into<Error>,
     U::InitError: fmt::Debug,
 {
     /// Create simple tcp stream service
@@ -115,7 +115,7 @@ mod openssl {
             Request = (Request, Framed<SslStream<TcpStream>, Codec>),
             Response = (),
         >,
-        U::Error: fmt::Display,
+        U::Error: fmt::Display + Into<Error>,
         U::InitError: fmt::Debug,
     {
         /// Create openssl based service
@@ -255,7 +255,7 @@ where
     X::Error: Into<Error>,
     X::InitError: fmt::Debug,
     U: ServiceFactory<Config = (), Request = (Request, Framed<T, Codec>), Response = ()>,
-    U::Error: fmt::Display,
+    U::Error: fmt::Display + Into<Error>,
     U::InitError: fmt::Debug,
 {
     type Config = ();
@@ -412,7 +412,7 @@ where
     X: Service<Request = Request, Response = Request>,
     X::Error: Into<Error>,
     U: Service<Request = (Request, Framed<T, Codec>), Response = ()>,
-    U::Error: fmt::Display,
+    U::Error: fmt::Display + Into<Error>,
 {
     type Request = (T, Option<net::SocketAddr>);
     type Response = ();
@@ -440,6 +440,19 @@ where
             })?
             .is_ready()
             && ready;
+        
+        let ready = if let Some(ref mut upg) = self.upgrade {
+            upg.poll_ready(cx)
+                .map_err(|e| {
+                    let e = e.into();
+                    log::error!("Http service readiness error: {:?}", e);
+                    DispatchError::Service(e)
+                })?
+                .is_ready()
+                && ready
+        } else {
+            ready
+        };
 
         if ready {
             Poll::Ready(Ok(()))

actix-http/src/service.rs

@@ -169,7 +169,7 @@ where
         Request = (Request, Framed<TcpStream, h1::Codec>),
         Response = (),
     >,
-    U::Error: fmt::Display,
+    U::Error: fmt::Display + Into<Error>,
     U::InitError: fmt::Debug,
     <U::Service as Service>::Future: 'static,
 {
@@ -214,7 +214,7 @@ mod openssl {
             Request = (Request, Framed<SslStream<TcpStream>, h1::Codec>),
             Response = (),
         >,
-        U::Error: fmt::Display,
+        U::Error: fmt::Display + Into<Error>,
         U::InitError: fmt::Debug,
         <U::Service as Service>::Future: 'static,
     {
@@ -335,7 +335,7 @@ where
         Request = (Request, Framed<T, h1::Codec>),
         Response = (),
     >,
-    U::Error: fmt::Display,
+    U::Error: fmt::Display + Into<Error>,
     U::InitError: fmt::Debug,
     <U::Service as Service>::Future: 'static,
 {
@@ -493,7 +493,7 @@ where
     X: Service<Request = Request, Response = Request>,
     X::Error: Into<Error>,
     U: Service<Request = (Request, Framed<T, h1::Codec>), Response = ()>,
-    U::Error: fmt::Display,
+    U::Error: fmt::Display + Into<Error>,
 {
     type Request = (T, Protocol, Option<net::SocketAddr>);
     type Response = ();
@@ -522,6 +522,19 @@ where
             .is_ready()
             && ready;
 
+        let ready = if let Some(ref mut upg) = self.upgrade {
+            upg.poll_ready(cx)
+                .map_err(|e| {
+                    let e = e.into();
+                    log::error!("Http service readiness error: {:?}", e);
+                    DispatchError::Service(e)
+                })?
+                .is_ready()
+                && ready
+        } else {
+            ready
+        };
+
         if ready {
             Poll::Ready(Ok(()))
         } else {

actix-http/tests/test_ws.rs

@@ -1,24 +1,70 @@
+use std::cell::Cell;
+use std::marker::PhantomData;
+use std::pin::Pin;
+use std::sync::{Arc, Mutex};
+
 use actix_codec::{AsyncRead, AsyncWrite, Framed};
 use actix_http::{body, h1, ws, Error, HttpService, Request, Response};
 use actix_http_test::test_server;
+use actix_service::{fn_factory, Service};
 use actix_utils::framed::Dispatcher;
 use bytes::Bytes;
 use futures::future;
-use futures::{SinkExt, StreamExt};
+use futures::task::{Context, Poll};
+use futures::{Future, SinkExt, StreamExt};
 
-async fn ws_service<T: AsyncRead + AsyncWrite + Unpin>(
-    (req, mut framed): (Request, Framed<T, h1::Codec>),
-) -> Result<(), Error> {
-    let res = ws::handshake(req.head()).unwrap().message_body(());
+struct WsService<T>(Arc<Mutex<(PhantomData<T>, Cell<bool>)>>);
 
-    framed
-        .send((res, body::BodySize::None).into())
-        .await
-        .unwrap();
+impl<T> WsService<T> {
+    fn new() -> Self {
+        WsService(Arc::new(Mutex::new((PhantomData, Cell::new(false)))))
+    }
 
-    Dispatcher::new(framed.into_framed(ws::Codec::new()), service)
-        .await
-        .map_err(|_| panic!())
+    fn set_polled(&mut self) {
+        *self.0.lock().unwrap().1.get_mut() = true;
+    }
+
+    fn was_polled(&self) -> bool {
+        self.0.lock().unwrap().1.get()
+    }
+}
+
+impl<T> Clone for WsService<T> {
+    fn clone(&self) -> Self {
+        WsService(self.0.clone())
+    }
+}
+
+impl<T> Service for WsService<T>
+where
+    T: AsyncRead + AsyncWrite + Unpin + 'static,
+{
+    type Request = (Request, Framed<T, h1::Codec>);
+    type Response = ();
+    type Error = Error;
+    type Future = Pin<Box<dyn Future<Output = Result<(), Error>>>>;
+
+    fn poll_ready(&mut self, _ctx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.set_polled();
+        Poll::Ready(Ok(()))
+    }
+
+    fn call(&mut self, (req, mut framed): Self::Request) -> Self::Future {
+        let fut = async move {
+            let res = ws::handshake(req.head()).unwrap().message_body(());
+
+            framed
+                .send((res, body::BodySize::None).into())
+                .await
+                .unwrap();
+
+            Dispatcher::new(framed.into_framed(ws::Codec::new()), service)
+                .await
+                .map_err(|_| panic!())
+        };
+
+        Box::pin(fut)
+    }
 }
 
 async fn service(msg: ws::Frame) -> Result<ws::Message, Error> {
@@ -37,11 +83,16 @@ async fn service(msg: ws::Frame) -> Result<ws::Message, Error> {
 
 #[actix_rt::test]
 async fn test_simple() {
-    let mut srv = test_server(|| {
-        HttpService::build()
-            .upgrade(actix_service::fn_service(ws_service))
-            .finish(|_| future::ok::<_, ()>(Response::NotFound()))
-            .tcp()
+    let ws_service = WsService::new();
+    let mut srv = test_server({
+        let ws_service = ws_service.clone();
+        move || {
+            let ws_service = ws_service.clone();
+            HttpService::build()
+                .upgrade(fn_factory(move || future::ok::<_, ()>(ws_service.clone())))
+                .finish(|_| future::ok::<_, ()>(Response::NotFound()))
+                .tcp()
+        }
     });
 
     // client service
@@ -138,4 +189,6 @@ async fn test_simple() {
         item.unwrap().unwrap(),
         ws::Frame::Close(Some(ws::CloseCode::Normal.into()))
     );
+
+    assert!(ws_service.was_polled());
 }