From de8a09254d913334ab6f59269767df530369f2f7 Mon Sep 17 00:00:00 2001 From: Nathan Fox Date: Sat, 21 Apr 2018 16:50:27 -0400 Subject: [PATCH 1/4] 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 522ae02af..5a8f10e04 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 ef333b411..b38312583 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 d9159e541..07e592868 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 06b771126..7db1cf3be 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 0c2eab0f4..0dbb5fda8 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 61f4af424..624f91593 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), _ => (), } } From f6fd9e70f96039ccc6116414e475a21966765b01 Mon Sep 17 00:00:00 2001 From: Nathan Fox Date: Sat, 21 Apr 2018 16:53:55 -0400 Subject: [PATCH 2/4] code cleanup --- src/ws/frame.rs | 3 +-- src/ws/mod.rs | 1 - 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/src/ws/frame.rs b/src/ws/frame.rs index 07e592868..dfef26b0e 100644 --- a/src/ws/frame.rs +++ b/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; @@ -30,7 +29,7 @@ impl Frame { /// Create a new Close control frame. #[inline] pub fn close(reason: Option, genmask: bool) -> Binary { - let payload:Vec = match reason { + let payload = match reason { None => Vec::new(), Some(reason) => { let mut code_bytes = [0; 2]; diff --git a/src/ws/mod.rs b/src/ws/mod.rs index 7db1cf3be..93fd431ca 100644 --- a/src/ws/mod.rs +++ b/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}; From b7b61afaccf718ecb37cbc67260ad59ef34908c3 Mon Sep 17 00:00:00 2001 From: Nathan Fox Date: Sat, 21 Apr 2018 17:20:23 -0400 Subject: [PATCH 3/4] add ws close description parse test --- src/ws/proto.rs | 2 +- tests/test_ws.rs | 11 +++++++++++ 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/src/ws/proto.rs b/src/ws/proto.rs index 0dbb5fda8..954f3f4a7 100644 --- a/src/ws/proto.rs +++ b/src/ws/proto.rs @@ -179,7 +179,7 @@ impl From for CloseCode { } } -#[derive(Debug, PartialEq)] +#[derive(Debug, Eq, PartialEq, Clone)] pub struct CloseReason { pub code: CloseCode, pub description: Option, diff --git a/tests/test_ws.rs b/tests/test_ws.rs index 624f91593..1283b2e89 100644 --- a/tests/test_ws.rs +++ b/tests/test_ws.rs @@ -70,6 +70,17 @@ fn test_empty_close_code() { 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] fn test_large_text() { let data = rand::thread_rng() From f8b75c157fafeb244f2b1bd741a82508cde18b78 Mon Sep 17 00:00:00 2001 From: Nathan Fox Date: Sun, 22 Apr 2018 11:43:47 -0400 Subject: [PATCH 4/4] fix style --- src/ws/frame.rs | 52 ++++++++++++++++++++++++------------------------- src/ws/proto.rs | 28 +++++++++++++------------- 2 files changed, 40 insertions(+), 40 deletions(-) diff --git a/src/ws/frame.rs b/src/ws/frame.rs index dfef26b0e..de78b31d2 100644 --- a/src/ws/frame.rs +++ b/src/ws/frame.rs @@ -29,21 +29,21 @@ impl Frame { /// Create a new Close control frame. #[inline] pub fn close(reason: Option, genmask: bool) -> Binary { - let payload = match reason { - None => Vec::new(), - Some(reason) => { - let mut code_bytes = [0; 2]; - NetworkEndian::write_u16(&mut code_bytes, reason.code.into()); + let payload = match reason { + None => Vec::new(), + Some(reason) => { + let mut code_bytes = [0; 2]; + NetworkEndian::write_u16(&mut code_bytes, reason.code.into()); - let mut payload = Vec::from(&code_bytes[..]); - if let Some(description) = reason.description{ - payload.extend(description.as_bytes()); - } - payload - } - }; + 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))] @@ -529,19 +529,19 @@ mod tests { assert_eq!(frame, v.into()); } - #[test] - fn test_close_frame() { - let reason = (CloseCode::Normal, "data"); - let frame = Frame::close(Some(reason.into()), 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()); - } + #[test] + fn test_empty_close_frame() { + let frame = Frame::close(None, false); + assert_eq!(frame, vec![0x88, 0x00].into()); + } } diff --git a/src/ws/proto.rs b/src/ws/proto.rs index 954f3f4a7..2851fb9e9 100644 --- a/src/ws/proto.rs +++ b/src/ws/proto.rs @@ -181,26 +181,26 @@ impl From for CloseCode { #[derive(Debug, Eq, PartialEq, Clone)] pub struct CloseReason { - pub code: CloseCode, - pub description: Option, + pub code: CloseCode, + pub description: Option, } impl From for CloseReason { - fn from(code: CloseCode) -> Self { - CloseReason { - code, - description: None, - } - } + 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()) - } - } + 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";