From f59de2490dd78ffd9d318928c13ba4c2238d9464 Mon Sep 17 00:00:00 2001
From: Glitch <smallglitch@cryptolab.net>
Date: Thu, 14 Jan 2021 20:51:25 +0100
Subject: [PATCH] Add `close_frame` function to `ws::Message` (#775)

---
 src/filters/ws.rs |  9 +++++++++
 tests/ws.rs       | 15 +++++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/src/filters/ws.rs b/src/filters/ws.rs
index bbd1836d2..8475523af 100644
--- a/src/filters/ws.rs
+++ b/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 {
diff --git a/tests/ws.rs b/tests/ws.rs
index ad7913318..5ff23e482 100644
--- a/tests/ws.rs
+++ b/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();