diff --git a/src/ws/client.rs b/src/ws/client.rs index 989220474..18789fef8 100644 --- a/src/ws/client.rs +++ b/src/ws/client.rs @@ -27,7 +27,7 @@ use client::{ Pipeline, SendRequest, SendRequestError, }; -use super::frame::Frame; +use super::frame::{Frame, FramedMessage}; use super::proto::{CloseReason, OpCode}; use super::{Message, ProtocolError, WsWriter}; @@ -529,10 +529,10 @@ pub struct ClientWriter { impl ClientWriter { /// Write payload #[inline] - fn write(&mut self, mut data: Binary) { + fn write(&mut self, mut data: FramedMessage) { let inner = self.inner.borrow_mut(); if !inner.closed { - let _ = inner.tx.unbounded_send(data.take()); + let _ = inner.tx.unbounded_send(data.0.take()); } else { warn!("Trying to write to disconnected response"); } diff --git a/src/ws/context.rs b/src/ws/context.rs index ffdd0b559..4db83df5c 100644 --- a/src/ws/context.rs +++ b/src/ws/context.rs @@ -20,7 +20,7 @@ use context::{ActorHttpContext, Drain, Frame as ContextFrame}; use error::{Error, ErrorInternalServerError, PayloadError}; use httprequest::HttpRequest; -use ws::frame::Frame; +use ws::frame::{Frame, FramedMessage}; use ws::proto::{CloseReason, OpCode}; use ws::{Message, ProtocolError, WsStream, WsWriter}; @@ -132,14 +132,19 @@ where A: Actor, { /// Write payload + /// + /// This is a low-level function that accepts framed messages that should + /// be created using `Frame::message()`. If you want to send text or binary + /// data you should prefer the `text()` or `binary()` convenience functions + /// that handle the framing for you. #[inline] - fn write(&mut self, data: Binary) { + pub fn write_raw(&mut self, data: FramedMessage) { if !self.disconnected { if self.stream.is_none() { self.stream = Some(SmallVec::new()); } let stream = self.stream.as_mut().unwrap(); - stream.push(ContextFrame::Chunk(Some(data))); + stream.push(ContextFrame::Chunk(Some(data.0))); } else { warn!("Trying to write to disconnected response"); } @@ -167,19 +172,19 @@ where /// Send text frame #[inline] pub fn text>(&mut self, text: T) { - self.write(Frame::message(text.into(), OpCode::Text, true, false)); + self.write_raw(Frame::message(text.into(), OpCode::Text, true, false)); } /// Send binary frame #[inline] pub fn binary>(&mut self, data: B) { - self.write(Frame::message(data, OpCode::Binary, true, false)); + self.write_raw(Frame::message(data, OpCode::Binary, true, false)); } /// Send ping frame #[inline] pub fn ping(&mut self, message: &str) { - self.write(Frame::message( + self.write_raw(Frame::message( Vec::from(message), OpCode::Ping, true, @@ -190,7 +195,7 @@ where /// Send pong frame #[inline] pub fn pong(&mut self, message: &str) { - self.write(Frame::message( + self.write_raw(Frame::message( Vec::from(message), OpCode::Pong, true, @@ -201,7 +206,7 @@ where /// Send close frame #[inline] pub fn close(&mut self, reason: Option) { - self.write(Frame::close(reason, false)); + self.write_raw(Frame::close(reason, false)); } /// Check if connection still open diff --git a/src/ws/frame.rs b/src/ws/frame.rs index 006d322f6..5e4fd8290 100644 --- a/src/ws/frame.rs +++ b/src/ws/frame.rs @@ -28,7 +28,7 @@ impl Frame { /// Create a new Close control frame. #[inline] - pub fn close(reason: Option, genmask: bool) -> Binary { + pub fn close(reason: Option, genmask: bool) -> FramedMessage { let payload = match reason { None => Vec::new(), Some(reason) => { @@ -295,7 +295,7 @@ impl Frame { /// Generate binary representation pub fn message>( data: B, code: OpCode, finished: bool, genmask: bool, - ) -> Binary { + ) -> FramedMessage { let payload = data.into(); let one: u8 = if finished { 0x80 | Into::::into(code) @@ -325,7 +325,7 @@ impl Frame { buf }; - if genmask { + let binary = if genmask { let mask = rand::random::(); buf.put_u32_le(mask); buf.extend_from_slice(payload.as_ref()); @@ -335,7 +335,9 @@ impl Frame { } else { buf.put_slice(payload.as_ref()); buf.into() - } + }; + + FramedMessage(binary) } } @@ -372,6 +374,10 @@ impl fmt::Display for Frame { } } +/// `WebSocket` message with framing. +#[derive(Debug)] +pub struct FramedMessage(pub(crate) Binary); + #[cfg(test)] mod tests { use super::*; @@ -502,7 +508,7 @@ mod tests { let mut v = vec![137u8, 4u8]; v.extend(b"data"); - assert_eq!(frame, v.into()); + assert_eq!(frame.0, v.into()); } #[test] @@ -511,7 +517,7 @@ mod tests { let mut v = vec![138u8, 4u8]; v.extend(b"data"); - assert_eq!(frame, v.into()); + assert_eq!(frame.0, v.into()); } #[test] @@ -521,12 +527,12 @@ mod tests { let mut v = vec![136u8, 6u8, 3u8, 232u8]; v.extend(b"data"); - assert_eq!(frame, v.into()); + assert_eq!(frame.0, v.into()); } #[test] fn test_empty_close_frame() { let frame = Frame::close(None, false); - assert_eq!(frame, vec![0x88, 0x00].into()); + assert_eq!(frame.0, vec![0x88, 0x00].into()); } } diff --git a/src/ws/mod.rs b/src/ws/mod.rs index ed44e2708..6b37bc7e0 100644 --- a/src/ws/mod.rs +++ b/src/ws/mod.rs @@ -64,7 +64,7 @@ pub use self::client::{ Client, ClientError, ClientHandshake, ClientReader, ClientWriter, }; pub use self::context::WebsocketContext; -pub use self::frame::Frame; +pub use self::frame::{Frame, FramedMessage}; pub use self::proto::{CloseCode, CloseReason, OpCode}; /// Websocket protocol errors