From de8a09254d913334ab6f59269767df530369f2f7 Mon Sep 17 00:00:00 2001 From: Nathan Fox Date: Sat, 21 Apr 2018 16:50:27 -0400 Subject: [PATCH] use Optional with websocket close reason --- src/ws/client.rs | 13 ++++----- src/ws/context.rs | 6 ++--- src/ws/frame.rs | 69 ++++++++++++++++++++++++++++++----------------- src/ws/mod.rs | 16 +++-------- src/ws/proto.rs | 38 ++++++++++++++++---------- tests/test_ws.rs | 12 ++++----- 6 files changed, 87 insertions(+), 67 deletions(-) diff --git a/src/ws/client.rs b/src/ws/client.rs index 522ae02a..5a8f10e0 100644 --- a/src/ws/client.rs +++ b/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; - Ok(Async::Ready(Some(Message::Close(CloseCode::from( - code, - ))))) + let close_reason = Frame::parse_close_payload(&payload); + Ok(Async::Ready(Some(Message::Close(close_reason)))) } 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) { - self.write(Frame::close(code, reason, true)); + pub fn close(&mut self, reason: Option) { + self.write(Frame::close(reason, true)); } } diff --git a/src/ws/context.rs b/src/ws/context.rs index ef333b41..b3831258 100644 --- a/src/ws/context.rs +++ b/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 @@ -177,8 +177,8 @@ where /// Send close frame #[inline] - pub fn close(&mut self, code: CloseCode, reason: &str) { - self.write(Frame::close(code, reason, false)); + pub fn close(&mut self, reason: Option) { + self.write(Frame::close(reason, false)); } /// Returns drain future diff --git a/src/ws/frame.rs b/src/ws/frame.rs index d9159e54..07e59286 100644 --- a/src/ws/frame.rs +++ b/src/ws/frame.rs @@ -11,7 +11,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,24 +29,22 @@ impl Frame { /// Create a new Close control frame. #[inline] - pub fn close(code: CloseCode, reason: &str, genmask: bool) -> Binary { - let raw: [u8; 2] = unsafe { - let u: u16 = code.into(); - mem::transmute(u.to_be()) - }; + pub fn close(reason: Option, genmask: bool) -> Binary { + let payload:Vec = match reason { + None => Vec::new(), + Some(reason) => { + let mut code_bytes = [0; 2]; + NetworkEndian::write_u16(&mut code_bytes, reason.code.into()); - let payload = if let CloseCode::Empty = code { - Vec::new() - } else { - Vec::from_iter( - raw[..] - .iter() - .chain(reason.as_bytes().iter()) - .cloned(), - ) - }; + let mut payload = Vec::from(&code_bytes[..]); + if let Some(description) = reason.description{ + payload.extend(description.as_bytes()); + } + payload + } + }; - Frame::message(payload, OpCode::Close, true, genmask) + Frame::message(payload, OpCode::Close, true, genmask) } #[cfg_attr(feature = "cargo-clippy", allow(type_complexity))] @@ -281,6 +279,22 @@ impl Frame { }))) } + /// Parse the payload of a close frame. + pub fn parse_close_payload(payload: &Binary) -> Option { + 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>( data: B, code: OpCode, finished: bool, genmask: bool @@ -516,12 +530,19 @@ mod tests { assert_eq!(frame, v.into()); } - #[test] - fn test_close_frame() { - let frame = Frame::close(CloseCode::Normal, "data", false); + #[test] + fn test_close_frame() { + 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()); - } + 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()); + } } diff --git a/src/ws/mod.rs b/src/ws/mod.rs index 06b77112..7db1cf3b 100644 --- a/src/ws/mod.rs +++ b/src/ws/mod.rs @@ -66,8 +66,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::OpCode; +pub use self::proto::{CloseCode, CloseReason, OpCode}; /// Websocket protocol errors #[derive(Fail, Debug)] @@ -164,7 +163,7 @@ pub enum Message { Binary(Binary), Ping(String), Pong(String), - Close(CloseCode), + Close(Option), } /// Do websocket handshake and start actor @@ -310,15 +309,8 @@ where } OpCode::Close => { self.closed = true; - let close_code = if payload.len() >= 2 { - let raw_code = - NetworkEndian::read_uint(payload.as_ref(), 2) as u16; - CloseCode::from(raw_code) - } else { - CloseCode::Status - }; - - Ok(Async::Ready(Some(Message::Close(close_code)))) + let close_reason = Frame::parse_close_payload(&payload); + Ok(Async::Ready(Some(Message::Close(close_reason)))) } OpCode::Ping => Ok(Async::Ready(Some(Message::Ping( String::from_utf8_lossy(payload.as_ref()).into(), diff --git a/src/ws/proto.rs b/src/ws/proto.rs index 0c2eab0f..0dbb5fda 100644 --- a/src/ws/proto.rs +++ b/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 for CloseCode { Away => 1001, Protocol => 1002, Unsupported => 1003, - Status => 1005, Abnormal => 1006, Invalid => 1007, Policy => 1008, @@ -160,7 +153,6 @@ impl Into for CloseCode { Restart => 1012, Again => 1013, Tls => 1015, - Empty => 0, Other(code) => code, } } @@ -173,7 +165,6 @@ impl From for CloseCode { 1001 => Away, 1002 => Protocol, 1003 => Unsupported, - 1005 => Status, 1006 => Abnormal, 1007 => Invalid, 1008 => Policy, @@ -183,12 +174,35 @@ impl From for CloseCode { 1012 => Restart, 1013 => Again, 1015 => Tls, - 0 => Empty, _ => Other(code), } } } +#[derive(Debug, PartialEq)] +pub struct CloseReason { + pub code: CloseCode, + pub description: Option, +} + +impl From for CloseReason { + fn from(code: CloseCode) -> Self { + CloseReason { + code, + description: None, + } + } +} + +impl > 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::::into(CloseCode::Away)); assert_eq!(1002u16, Into::::into(CloseCode::Protocol)); assert_eq!(1003u16, Into::::into(CloseCode::Unsupported)); - assert_eq!(1005u16, Into::::into(CloseCode::Status)); assert_eq!(1006u16, Into::::into(CloseCode::Abnormal)); assert_eq!(1007u16, Into::::into(CloseCode::Invalid)); assert_eq!(1008u16, Into::::into(CloseCode::Policy)); @@ -299,7 +310,6 @@ mod test { assert_eq!(1012u16, Into::::into(CloseCode::Restart)); assert_eq!(1013u16, Into::::into(CloseCode::Again)); assert_eq!(1015u16, Into::::into(CloseCode::Tls)); - assert_eq!(0u16, Into::::into(CloseCode::Empty)); assert_eq!(2000u16, Into::::into(CloseCode::Other(2000))); } } diff --git a/tests/test_ws.rs b/tests/test_ws.rs index 61f4af42..624f9159 100644 --- a/tests/test_ws.rs +++ b/tests/test_ws.rs @@ -27,7 +27,7 @@ impl StreamHandler 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,9 @@ 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] @@ -147,7 +147,7 @@ impl StreamHandler 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), _ => (), } }