Add close_frame function to ws::Message (#775)

Author: aumetra

src/filters/ws.rs

@@ -324,6 +324,15 @@ impl Message {
         self.inner.is_pong()
     }
 
+    /// Try to get the close frame (close code and reason)
+    pub fn close_frame(&self) -> Option<(u16, &str)> {
+        if let protocol::Message::Close(Some(ref close_frame)) = self.inner {
+            Some((close_frame.code.into(), close_frame.reason.as_ref()))
+        } else {
+            None
+        }
+    }
+
     /// 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 {

tests/ws.rs

@@ -81,6 +81,21 @@ async fn binary() {
     assert_eq!(msg.as_bytes(), &b"bonk"[..]);
 }
 
+#[tokio::test]
+async fn close_frame() {
+    let _ = pretty_env_logger::try_init();
+
+    let route = warp::ws().map(|ws: warp::ws::Ws| {
+        ws.on_upgrade(|mut websocket| async move {
+            let msg = websocket.next().await.expect("item").expect("ok");
+            let _ = msg.close_frame().expect("close frame");
+        })
+    });
+
+    let client = warp::test::ws().handshake(route).await.expect("handshake");
+    drop(client);
+}
+
 #[tokio::test]
 async fn send_ping() {
     let _ = pretty_env_logger::try_init();