actix/actix-web
1
0
Fork 0
mirror of https://github.com/fafhrd91/actix-web synced 2024-11-27 17:52:56 +01:00
Code Issues Releases Wiki Activity

use Optional with websocket close reason

Browse Source
This commit is contained in:
Nathan Fox 2018-04-21 16:50:27 -04:00
parent 2adf8a3a48
commit de8a09254d
6 changed files with 87 additions and 67 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
src/ws/client.rs
View File

@ -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<CloseReason>) {
self.write(Frame::close(reason, true));
}
}

6
src/ws/context.rs
View File

@ -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) {
self.write(Frame::close(code, reason, false));
pub fn close(&mut self, reason: Option<CloseReason>) {
self.write(Frame::close(reason, false));
}
/// Returns drain future

53
src/ws/frame.rs
View File

@ -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,21 +29,19 @@ 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<CloseReason>, genmask: bool) -> Binary {
let payload:Vec<u8> = 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)
@ -281,6 +279,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 +532,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());
}
}

16
src/ws/mod.rs
View File

@ -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<CloseReason>),
}
/// 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(),

38
src/ws/proto.rs
View File

@ -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, PartialEq)]
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)));
}
}

12
tests/test_ws.rs
View File

@ -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,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<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),
_ => (),
}
}