Merge pull request #196 from fuchsnj/websocket_close_reason

Websocket close reason

src/ws/client.rs

@@ -5,7 +5,6 @@ use std::time::Duration;
 use std::{fmt, io, str};
 
 use base64;
-use byteorder::{ByteOrder, NetworkEndian};
 use bytes::Bytes;
 use cookie::Cookie;
 use futures::unsync::mpsc::{unbounded, UnboundedSender};
@@ -27,7 +26,7 @@ use client::{ClientConnector, ClientRequest, ClientRequestBuilder, ClientRespons
              HttpResponseParserError, SendRequest, SendRequestError};
 
 use super::frame::Frame;
-use super::proto::{CloseCode, OpCode};
+use super::proto::{CloseReason, OpCode};
 use super::{Message, ProtocolError};
 
 /// Websocket client error
@@ -468,10 +467,8 @@ impl Stream for ClientReader {
                     }
                     OpCode::Close => {
                         inner.closed = true;
-                        let code = NetworkEndian::read_uint(payload.as_ref(), 2) as u16;
+                        let close_reason = Frame::parse_close_payload(&payload);
-                        Ok(Async::Ready(Some(Message::Close(CloseCode::from(
+                        Ok(Async::Ready(Some(Message::Close(close_reason))))
-                            code,
-                        )))))
                     }
                     OpCode::Ping => Ok(Async::Ready(Some(Message::Ping(
                         String::from_utf8_lossy(payload.as_ref()).into(),
@@ -560,7 +557,7 @@ impl ClientWriter {
 
     /// Send close frame
     #[inline]
-    pub fn close(&mut self, code: CloseCode, reason: &str) {
+    pub fn close(&mut self, reason: Option<CloseReason>) {
-        self.write(Frame::close(code, reason, true));
+        self.write(Frame::close(reason, true));
     }
 }

src/ws/context.rs

@@ -15,7 +15,7 @@ use error::{Error, ErrorInternalServerError};
 use httprequest::HttpRequest;
 
 use ws::frame::Frame;
-use ws::proto::{CloseCode, OpCode};
+use ws::proto::{CloseReason, OpCode};
 
 /// Execution context for `WebSockets` actors
 pub struct WebsocketContext<A, S = ()>
@@ -177,8 +177,8 @@ where
 
     /// Send close frame
     #[inline]
-    pub fn close(&mut self, code: CloseCode, reason: &str) {
+    pub fn close(&mut self, reason: Option<CloseReason>) {
-        self.write(Frame::close(code, reason, false));
+        self.write(Frame::close(reason, false));
     }
 
     /// Returns drain future

src/ws/frame.rs

@@ -2,7 +2,6 @@ use byteorder::{BigEndian, ByteOrder, NetworkEndian};
 use bytes::{BufMut, Bytes, BytesMut};
 use futures::{Async, Poll, Stream};
 use rand;
-use std::iter::FromIterator;
 use std::{fmt, mem, ptr};
 
 use body::Binary;
@@ -11,7 +10,7 @@ use payload::PayloadHelper;
 
 use ws::ProtocolError;
 use ws::mask::apply_mask;
-use ws::proto::{CloseCode, OpCode};
+use ws::proto::{CloseCode, CloseReason, OpCode};
 
 /// A struct representing a `WebSocket` frame.
 #[derive(Debug)]
@@ -29,21 +28,19 @@ impl Frame {
 
     /// Create a new Close control frame.
     #[inline]
-    pub fn close(code: CloseCode, reason: &str, genmask: bool) -> Binary {
+    pub fn close(reason: Option<CloseReason>, genmask: bool) -> Binary {
-        let raw: [u8; 2] = unsafe {
+        let payload = match reason {
-            let u: u16 = code.into();
+            None => Vec::new(),
-            mem::transmute(u.to_be())
+            Some(reason) => {
-        };
+                let mut code_bytes = [0; 2];
+                NetworkEndian::write_u16(&mut code_bytes, reason.code.into());
 
-        let payload = if let CloseCode::Empty = code {
+                let mut payload = Vec::from(&code_bytes[..]);
-            Vec::new()
+                if let Some(description) = reason.description{
-        } else {
+                    payload.extend(description.as_bytes());
-            Vec::from_iter(
+                }
-                raw[..]
+                payload
-                    .iter()
+            }
-                    .chain(reason.as_bytes().iter())
-                    .cloned(),
-            )
         };
 
         Frame::message(payload, OpCode::Close, true, genmask)
@@ -281,6 +278,22 @@ impl Frame {
         })))
     }
 
+    /// Parse the payload of a close frame.
+    pub fn parse_close_payload(payload: &Binary) -> Option<CloseReason> {
+        if payload.len() >= 2 {
+            let raw_code = NetworkEndian::read_uint(payload.as_ref(), 2) as u16;
+            let code = CloseCode::from(raw_code);
+            let description = if payload.len() > 2 {
+                Some(String::from_utf8_lossy(&payload.as_ref()[2..]).into())
+            } else {
+                None
+            };
+            Some(CloseReason { code, description })
+        } else {
+            None
+        }
+    }
+
     /// Generate binary representation
     pub fn message<B: Into<Binary>>(
         data: B, code: OpCode, finished: bool, genmask: bool
@@ -518,10 +531,17 @@ mod tests {
 
     #[test]
     fn test_close_frame() {
-        let frame = Frame::close(CloseCode::Normal, "data", false);
+        let reason = (CloseCode::Normal, "data");
+        let frame = Frame::close(Some(reason.into()), false);
 
         let mut v = vec![136u8, 6u8, 3u8, 232u8];
         v.extend(b"data");
         assert_eq!(frame, v.into());
     }
+
+    #[test]
+    fn test_empty_close_frame() {
+        let frame = Frame::close(None, false);
+        assert_eq!(frame, vec![0x88, 0x00].into());
+    }
 }

src/ws/mod.rs

@@ -43,7 +43,6 @@
 //! #      .finish();
 //! # }
 //! ```
-use byteorder::{ByteOrder, NetworkEndian};
 use bytes::Bytes;
 use futures::{Async, Poll, Stream};
 use http::{header, Method, StatusCode};
@@ -66,8 +65,7 @@ mod proto;
 pub use self::client::{Client, ClientError, ClientHandshake, ClientReader, ClientWriter};
 pub use self::context::WebsocketContext;
 pub use self::frame::Frame;
-pub use self::proto::CloseCode;
+pub use self::proto::{CloseCode, CloseReason, OpCode};
-pub use self::proto::OpCode;
 
 /// Websocket protocol errors
 #[derive(Fail, Debug)]
@@ -164,7 +162,7 @@ pub enum Message {
     Binary(Binary),
     Ping(String),
     Pong(String),
-    Close(CloseCode),
+    Close(Option<CloseReason>),
 }
 
 /// Do websocket handshake and start actor
@@ -310,15 +308,8 @@ where
                     }
                     OpCode::Close => {
                         self.closed = true;
-                        let close_code = if payload.len() >= 2 {
+                        let close_reason = Frame::parse_close_payload(&payload);
-                            let raw_code =
+                        Ok(Async::Ready(Some(Message::Close(close_reason))))
-                                NetworkEndian::read_uint(payload.as_ref(), 2) as u16;
-                            CloseCode::from(raw_code)
-                        } else {
-                            CloseCode::Status
-                        };
-
-                        Ok(Async::Ready(Some(Message::Close(close_code))))
                     }
                     OpCode::Ping => Ok(Async::Ready(Some(Message::Ping(
                         String::from_utf8_lossy(payload.as_ref()).into(),

src/ws/proto.rs

@@ -90,10 +90,6 @@ pub enum CloseCode {
     /// endpoint that understands only text data MAY send this if it
     /// receives a binary message).
     Unsupported,
-    /// Indicates that no status code was included in a closing frame. This
-    /// close code makes it possible to use a single method, `on_close` to
-    /// handle even cases where no close code was provided.
-    Status,
     /// Indicates an abnormal closure. If the abnormal closure was due to an
     /// error, this close code will not be used. Instead, the `on_error` method
     /// of the handler will be called with the error. However, if the connection
@@ -138,8 +134,6 @@ pub enum CloseCode {
     #[doc(hidden)]
     Tls,
     #[doc(hidden)]
-    Empty,
-    #[doc(hidden)]
     Other(u16),
 }
 
@@ -150,7 +144,6 @@ impl Into<u16> for CloseCode {
             Away => 1001,
             Protocol => 1002,
             Unsupported => 1003,
-            Status => 1005,
             Abnormal => 1006,
             Invalid => 1007,
             Policy => 1008,
@@ -160,7 +153,6 @@ impl Into<u16> for CloseCode {
             Restart => 1012,
             Again => 1013,
             Tls => 1015,
-            Empty => 0,
             Other(code) => code,
         }
     }
@@ -173,7 +165,6 @@ impl From<u16> for CloseCode {
             1001 => Away,
             1002 => Protocol,
             1003 => Unsupported,
-            1005 => Status,
             1006 => Abnormal,
             1007 => Invalid,
             1008 => Policy,
@@ -183,12 +174,35 @@ impl From<u16> for CloseCode {
             1012 => Restart,
             1013 => Again,
             1015 => Tls,
-            0 => Empty,
             _ => Other(code),
         }
     }
 }
 
+#[derive(Debug, Eq, PartialEq, Clone)]
+pub struct CloseReason {
+    pub code: CloseCode,
+    pub description: Option<String>,
+}
+
+impl From<CloseCode> for CloseReason {
+    fn from(code: CloseCode) -> Self {
+        CloseReason {
+            code,
+            description: None,
+        }
+    }
+}
+
+impl <T: Into<String>> From<(CloseCode, T)> for CloseReason {
+    fn from(info: (CloseCode, T)) -> Self {
+        CloseReason{
+            code: info.0,
+            description: Some(info.1.into())
+        }
+    }
+}
+
 static WS_GUID: &'static str = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
 
 // TODO: hash is always same size, we dont need String
@@ -269,7 +283,6 @@ mod test {
         assert_eq!(CloseCode::from(1001u16), CloseCode::Away);
         assert_eq!(CloseCode::from(1002u16), CloseCode::Protocol);
         assert_eq!(CloseCode::from(1003u16), CloseCode::Unsupported);
-        assert_eq!(CloseCode::from(1005u16), CloseCode::Status);
         assert_eq!(CloseCode::from(1006u16), CloseCode::Abnormal);
         assert_eq!(CloseCode::from(1007u16), CloseCode::Invalid);
         assert_eq!(CloseCode::from(1008u16), CloseCode::Policy);
@@ -279,7 +292,6 @@ mod test {
         assert_eq!(CloseCode::from(1012u16), CloseCode::Restart);
         assert_eq!(CloseCode::from(1013u16), CloseCode::Again);
         assert_eq!(CloseCode::from(1015u16), CloseCode::Tls);
-        assert_eq!(CloseCode::from(0u16), CloseCode::Empty);
         assert_eq!(CloseCode::from(2000u16), CloseCode::Other(2000));
     }
 
@@ -289,7 +301,6 @@ mod test {
         assert_eq!(1001u16, Into::<u16>::into(CloseCode::Away));
         assert_eq!(1002u16, Into::<u16>::into(CloseCode::Protocol));
         assert_eq!(1003u16, Into::<u16>::into(CloseCode::Unsupported));
-        assert_eq!(1005u16, Into::<u16>::into(CloseCode::Status));
         assert_eq!(1006u16, Into::<u16>::into(CloseCode::Abnormal));
         assert_eq!(1007u16, Into::<u16>::into(CloseCode::Invalid));
         assert_eq!(1008u16, Into::<u16>::into(CloseCode::Policy));
@@ -299,7 +310,6 @@ mod test {
         assert_eq!(1012u16, Into::<u16>::into(CloseCode::Restart));
         assert_eq!(1013u16, Into::<u16>::into(CloseCode::Again));
         assert_eq!(1015u16, Into::<u16>::into(CloseCode::Tls));
-        assert_eq!(0u16, Into::<u16>::into(CloseCode::Empty));
         assert_eq!(2000u16, Into::<u16>::into(CloseCode::Other(2000)));
     }
 }

tests/test_ws.rs

@@ -27,7 +27,7 @@ impl StreamHandler<ws::Message, ws::ProtocolError> for Ws {
             ws::Message::Ping(msg) => ctx.pong(&msg),
             ws::Message::Text(text) => ctx.text(text),
             ws::Message::Binary(bin) => ctx.binary(bin),
-            ws::Message::Close(reason) => ctx.close(reason, ""),
+            ws::Message::Close(reason) => ctx.close(reason),
             _ => (),
         }
     }
@@ -55,9 +55,9 @@ fn test_simple() {
     let (item, reader) = srv.execute(reader.into_future()).unwrap();
     assert_eq!(item, Some(ws::Message::Pong("ping".to_owned())));
 
-    writer.close(ws::CloseCode::Normal, "");
+    writer.close(Some(ws::CloseCode::Normal.into()));
     let (item, _) = srv.execute(reader.into_future()).unwrap();
-    assert_eq!(item, Some(ws::Message::Close(ws::CloseCode::Normal)));
+    assert_eq!(item, Some(ws::Message::Close(Some(ws::CloseCode::Normal.into()))));
 }
 
 #[test]
@@ -65,9 +65,20 @@ fn test_empty_close_code() {
     let mut srv = test::TestServer::new(|app| app.handler(|req| ws::start(req, Ws)));
     let (reader, mut writer) = srv.ws().unwrap();
 
-    writer.close(ws::CloseCode::Empty, "");
+    writer.close(None);
     let (item, _) = srv.execute(reader.into_future()).unwrap();
-    assert_eq!(item, Some(ws::Message::Close(ws::CloseCode::Status)));
+    assert_eq!(item, Some(ws::Message::Close(None)));
+}
+
+#[test]
+fn test_close_description() {
+    let mut srv = test::TestServer::new(|app| app.handler(|req| ws::start(req, Ws)));
+    let (reader, mut writer) = srv.ws().unwrap();
+
+    let close_reason:ws::CloseReason = (ws::CloseCode::Normal, "close description").into();
+    writer.close(Some(close_reason.clone()));
+    let (item, _) = srv.execute(reader.into_future()).unwrap();
+    assert_eq!(item, Some(ws::Message::Close(Some(close_reason))));
 }
 
 #[test]
@@ -147,7 +158,7 @@ impl StreamHandler<ws::Message, ws::ProtocolError> for Ws2 {
             ws::Message::Ping(msg) => ctx.pong(&msg),
             ws::Message::Text(text) => ctx.text(text),
             ws::Message::Binary(bin) => ctx.binary(bin),
-            ws::Message::Close(reason) => ctx.close(reason, ""),
+            ws::Message::Close(reason) => ctx.close(reason),
             _ => (),
         }
     }