Fix poll_ready call for WebSockets upgrade (#1219)

* Fix poll_ready call for WebSockets upgrade

* Poll upgrade service from H1ServiceHandler too

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());
 }