Repository: seanmonstar/warp (language: rust)

Issue:
warp::filters::ws::Message::ping() not functioning
I am attempting to use the code I previously contributed to send WebSocket ping messages, but it does not appear to function.

Specifically, constructing a `ping` message and then sending it through an existing WebSocket connection does not result in the client receiving that ping.

I am working around the issue by sending text messages as a less-elegant form of a keep-alive ping, but wanted to raise it here—especially given that I contributed the code that doesn't seem to work correctly.  


Discussion / hints:
It might be that there is [code that specifically ignores sending pings](https://github.com/seanmonstar/warp/blob/4687c03ed9620e1ae8aaa4fd6cbd8e100091e584/src/filters/ws.rs#L329-L337) 😆 

We should probably add a unit test to `tests/ws.rs` when fixing this.

Fix patch:
diff --git a/src/filters/ext.rs b/src/filters/ext.rs
index ad55c44b4..483a6db91 100644
--- a/src/filters/ext.rs
+++ b/src/filters/ext.rs
@@ -27,9 +27,7 @@ pub fn get<T: Clone + Send + Sync + 'static>(
 /// If the extension doesn't exist, it yields `None`.
 pub fn optional<T: Clone + Send + Sync + 'static>(
 ) -> impl Filter<Extract = (Option<T>,), Error = Infallible> + Copy {
-    filter_fn_one(|route| {
-        future::ok(route.extensions().get::<T>().cloned())
-    })
+    filter_fn_one(|route| future::ok(route.extensions().get::<T>().cloned()))
 }
 
 unit_error! {
diff --git a/src/filters/ws.rs b/src/filters/ws.rs
index 7b707e858..d34af0fbe 100644
--- a/src/filters/ws.rs
+++ b/src/filters/ws.rs
@@ -284,34 +284,18 @@ impl Stream for WebSocket {
     type Item = Result<Message, crate::Error>;
 
     fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
-        loop {
-            let msg = match (*self).with_context(Some(cx), |s| s.read_message()) {
-                Ok(item) => item,
-                Err(::tungstenite::Error::Io(ref err))
-                    if err.kind() == io::ErrorKind::WouldBlock =>
-                {
-                    return Poll::Pending;
-                }
-                Err(::tungstenite::Error::ConnectionClosed) => {
-                    log::trace!("websocket closed");
-                    return Poll::Ready(None);
-                }
-                Err(e) => {
-                    log::debug!("websocket poll error: {}", e);
-                    return Poll::Ready(Some(Err(crate::Error::new(e))));
-                }
-            };
-
-            match msg {
-                msg @ protocol::Message::Text(..)
-                | msg @ protocol::Message::Binary(..)
-                | msg @ protocol::Message::Close(..)
-                | msg @ protocol::Message::Ping(..) => {
-                    return Poll::Ready(Some(Ok(Message { inner: msg })));
-                }
-                protocol::Message::Pong(payload) => {
-                    log::trace!("websocket client pong: {:?}", payload);
-                }
+        match (*self).with_context(Some(cx), |s| s.read_message()) {
+            Ok(item) => Poll::Ready(Some(Ok(Message { inner: item }))),
+            Err(::tungstenite::Error::Io(ref err)) if err.kind() == io::ErrorKind::WouldBlock => {
+                Poll::Pending
+            }
+            Err(::tungstenite::Error::ConnectionClosed) => {
+                log::trace!("websocket closed");
+                Poll::Ready(None)
+            }
+            Err(e) => {
+                log::debug!("websocket poll error: {}", e);
+                Poll::Ready(Some(Err(crate::Error::new(e))))
             }
         }
     }
@@ -327,16 +311,6 @@ impl Sink<Message> for WebSocket {
     }
 
     fn start_send(mut self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> {
-        if let protocol::Message::Ping(..) = item.inner {
-            // warp doesn't yet expose a way to construct a `Ping` message,
-            // so the only way this could is if the user is forwarding the
-            // received `Ping`s straight back.
-            //
-            // tungstenite already auto-reponds to `Ping`s with a `Pong`,
-            // so this just prevents accidentally sending extra pings.
-            return Ok(());
-        }
-
         match self.with_context(None, |s| s.write_message(item.inner)) {
             Ok(()) => Ok(()),
             // Err(::tungstenite::Error::SendQueueFull(inner)) => {
@@ -449,6 +423,11 @@ impl Message {
         self.inner.is_ping()
     }
 
+    /// Returns true if this message is a Pong message.
+    pub fn is_pong(&self) -> bool {
+        self.inner.is_pong()
+    }
+
     /// Try to get a reference to the string text, if this is a Text message.
     pub fn to_str(&self) -> Result<&str, ()> {
         match self.inner {


Test patch:
diff --git a/tests/ws.rs b/tests/ws.rs
index f0824f0c2..b53fd2f16 100644
--- a/tests/ws.rs
+++ b/tests/ws.rs
@@ -1,6 +1,7 @@
 #![deny(warnings)]
 
-use futures::{FutureExt, StreamExt};
+use futures::{FutureExt, SinkExt, StreamExt};
+use warp::ws::Message;
 use warp::Filter;
 
 #[tokio::test]
@@ -79,6 +80,59 @@ async fn binary() {
     assert_eq!(msg.as_bytes(), &b"bonk"[..]);
 }
 
+#[tokio::test]
+async fn send_ping() {
+    let _ = pretty_env_logger::try_init();
+
+    let filter = warp::ws().map(|ws: warp::ws::Ws| {
+        ws.on_upgrade(|mut websocket| {
+            async move {
+                websocket.send(Message::ping("srv")).await.unwrap();
+                // assume the client will pong back
+                let msg = websocket.next().await.expect("item").expect("ok");
+                assert!(msg.is_pong());
+                assert_eq!(msg.as_bytes(), &b"srv"[..]);
+            }
+        })
+    });
+
+    let mut client = warp::test::ws().handshake(filter).await.expect("handshake");
+
+    let msg = client.recv().await.expect("recv");
+    assert!(msg.is_ping());
+    assert_eq!(msg.as_bytes(), &b"srv"[..]);
+
+    client.recv_closed().await.expect("closed");
+}
+
+#[tokio::test]
+async fn echo_pings() {
+    let _ = pretty_env_logger::try_init();
+
+    let mut client = warp::test::ws()
+        .handshake(ws_echo())
+        .await
+        .expect("handshake");
+
+    client.send(Message::ping("clt")).await;
+
+    // tungstenite sends the PONG first
+    let msg = client.recv().await.expect("recv");
+    assert!(msg.is_pong());
+    assert_eq!(msg.as_bytes(), &b"clt"[..]);
+
+    // and then `ws_echo` sends us back the same PING
+    let msg = client.recv().await.expect("recv");
+    assert!(msg.is_ping());
+    assert_eq!(msg.as_bytes(), &b"clt"[..]);
+
+    // and then our client would have sent *its* PONG
+    // and `ws_echo` would send *that* back too
+    let msg = client.recv().await.expect("recv");
+    assert!(msg.is_pong());
+    assert_eq!(msg.as_bytes(), &b"clt"[..]);
+}
+
 #[tokio::test]
 async fn closed() {
     let _ = pretty_env_logger::try_init();
@@ -120,7 +174,9 @@ fn ws_echo() -> impl Filter<Extract = impl warp::Reply, Error = warp::Rejection>
         ws.on_upgrade(|websocket| {
             // Just echo all messages back...
             let (tx, rx) = websocket.split();
-            rx.forward(tx).map(|_| ())
+            rx.inspect(|i| log::debug!("ws recv: {:?}", i))
+                .forward(tx)
+                .map(|_| ())
         })
     })
 }