diff --git a/examples/autobahn-client.rs b/examples/autobahn-client.rs index e883ff4..700e2d7 100644 --- a/examples/autobahn-client.rs +++ b/examples/autobahn-client.rs @@ -5,7 +5,7 @@ extern crate url; use url::Url; -use tungstenite::{connect, Error, Result}; +use tungstenite::{connect, Error, Result, Message}; const AGENT: &'static str = "Tungstenite"; @@ -33,8 +33,14 @@ fn run_test(case: u32) -> Result<()> { ).unwrap(); let mut socket = connect(case_url)?; loop { - let msg = socket.read_message()?; - socket.write_message(msg)?; + match socket.read_message()? { + msg @ Message::Text(_) | + msg @ Message::Binary(_) => { + socket.write_message(msg)?; + } + Message::Ping(_) | + Message::Pong(_) => {} + } } } diff --git a/examples/autobahn-server.rs b/examples/autobahn-server.rs index d0bcbcc..458fa26 100644 --- a/examples/autobahn-server.rs +++ b/examples/autobahn-server.rs @@ -5,7 +5,7 @@ extern crate tungstenite; use std::net::{TcpListener, TcpStream}; use std::thread::spawn; -use tungstenite::{accept, HandshakeError, Error, Result}; +use tungstenite::{accept, HandshakeError, Error, Result, Message}; fn must_not_block(err: HandshakeError) -> Error { match err { @@ -17,8 +17,14 @@ fn must_not_block(err: HandshakeError) -> Error { fn handle_client(stream: TcpStream) -> Result<()> { let mut socket = accept(stream).map_err(must_not_block)?; loop { - let msg = socket.read_message()?; - socket.write_message(msg)?; + match socket.read_message()? { + msg @ Message::Text(_) | + msg @ Message::Binary(_) => { + socket.write_message(msg)?; + } + Message::Ping(_) | + Message::Pong(_) => {} + } } } diff --git a/src/protocol/message.rs b/src/protocol/message.rs index da7daaa..44eda20 100644 --- a/src/protocol/message.rs +++ b/src/protocol/message.rs @@ -141,6 +141,14 @@ pub enum Message { Text(String), /// A binary WebSocket message Binary(Vec), + /// A ping message with the specified payload + /// + /// The payload here must have a length less than 125 bytes + Ping(Vec), + /// A pong message with the specified payload + /// + /// The payload here must have a length less than 125 bytes + Pong(Vec), } impl Message { @@ -163,15 +171,31 @@ impl Message { pub fn is_text(&self) -> bool { match *self { Message::Text(_) => true, - Message::Binary(_) => false, + _ => false, } } /// Indicates whether a message is a binary message. pub fn is_binary(&self) -> bool { match *self { - Message::Text(_) => false, Message::Binary(_) => true, + _ => false, + } + } + + /// Indicates whether a message is a ping message. + pub fn is_ping(&self) -> bool { + match *self { + Message::Ping(_) => true, + _ => false, + } + } + + /// Indicates whether a message is a pong message. + pub fn is_pong(&self) -> bool { + match *self { + Message::Pong(_) => true, + _ => false, } } @@ -179,24 +203,25 @@ impl Message { pub fn len(&self) -> usize { match *self { Message::Text(ref string) => string.len(), - Message::Binary(ref data) => data.len(), + Message::Binary(ref data) | + Message::Ping(ref data) | + Message::Pong(ref data) => data.len(), } } /// Returns true if the WebSocket message has no content. /// For example, if the other side of the connection sent an empty string. pub fn is_empty(&self) -> bool { - match *self { - Message::Text(ref string) => string.is_empty(), - Message::Binary(ref data) => data.is_empty(), - } + self.len() == 0 } /// Consume the WebSocket and return it as binary data. pub fn into_data(self) -> Vec { match self { Message::Text(string) => string.into_bytes(), - Message::Binary(data) => data, + Message::Binary(data) | + Message::Ping(data) | + Message::Pong(data) => data, } } @@ -204,7 +229,9 @@ impl Message { pub fn into_text(self) -> Result { match self { Message::Text(string) => Ok(string), - Message::Binary(data) => Ok(try!( + Message::Binary(data) | + Message::Ping(data) | + Message::Pong(data) => Ok(try!( String::from_utf8(data).map_err(|err| err.utf8_error()))), } } @@ -214,7 +241,9 @@ impl Message { pub fn to_text(&self) -> Result<&str> { match *self { Message::Text(ref string) => Ok(string), - Message::Binary(ref data) => Ok(try!(str::from_utf8(data))), + Message::Binary(ref data) | + Message::Ping(ref data) | + Message::Pong(ref data) => Ok(try!(str::from_utf8(data))), } } diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs index b576fc3..841cd38 100644 --- a/src/protocol/mod.rs +++ b/src/protocol/mod.rs @@ -103,24 +103,27 @@ impl WebSocket { /// This function guarantees that the frame is queued regardless of any errors. /// There is no need to resend the frame. In order to handle WouldBlock or Incomplete, /// call write_pending() afterwards. + /// + /// Note that only the last pong frame is stored to be sent, and only the + /// most recent pong frame is sent if multiple pong frames are queued up. pub fn write_message(&mut self, message: Message) -> Result<()> { - let frame = { - let opcode = match message { - Message::Text(_) => OpData::Text, - Message::Binary(_) => OpData::Binary, - }; - Frame::message(message.into_data(), OpCode::Data(opcode), true) + let frame = match message { + Message::Text(data) => { + Frame::message(data.into(), OpCode::Data(OpData::Text), true) + } + Message::Binary(data) => { + Frame::message(data, OpCode::Data(OpData::Binary), true) + } + Message::Ping(data) => Frame::ping(data), + Message::Pong(data) => { + self.pong = Some(Frame::pong(data)); + return self.write_pending() + } }; self.send_queue.push_back(frame); self.write_pending() } - /// Send ping. - pub fn send_ping(&mut self, payload: Vec) -> Result<()> { - self.send_queue.push_back(Frame::ping(payload)); - self.write_pending() - } - /// Close the connection. /// /// This function guarantees that the close frame will be queued. @@ -212,7 +215,7 @@ impl WebSocket { match frame.opcode() { OpCode::Control(ctl) => { - (match ctl { + match ctl { // All control frames MUST have a payload length of 125 bytes or less // and MUST NOT be fragmented. (RFC 6455) _ if !frame.is_final() => { @@ -222,22 +225,24 @@ impl WebSocket { Err(Error::Protocol("Control frame too big".into())) } OpCtl::Close => { - self.do_close(frame.into_close()?) + self.do_close(frame.into_close()?).map(|_| None) } OpCtl::Reserved(i) => { Err(Error::Protocol(format!("Unknown control frame type {}", i).into())) } OpCtl::Ping | OpCtl::Pong if !self.state.is_active() => { // No ping processing while closing. - Ok(()) + Ok(None) } OpCtl::Ping => { - self.do_ping(frame.into_data()) + let data = frame.into_data(); + self.pong = Some(Frame::pong(data.clone())); + Ok(Some(Message::Ping(data))) } OpCtl::Pong => { - self.do_pong(frame.into_data()) + Ok(Some(Message::Pong(frame.into_data()))) } - }).map(|_| None) + } } OpCode::Data(_) if !self.state.is_active() => { @@ -353,27 +358,6 @@ impl WebSocket { } } - /// Received a ping frame. - fn do_ping(&mut self, ping: Vec) -> Result<()> { - // If an endpoint receives a Ping frame and has not yet sent Pong - // frame(s) in response to previous Ping frame(s), the endpoint MAY - // elect to send a Pong frame for only the most recently processed Ping - // frame. (RFC 6455) - // We do exactly that, keeping a "queue" from one and only Pong frame. - self.pong = Some(Frame::pong(ping)); - Ok(()) - } - - /// Received a pong frame. - fn do_pong(&mut self, _: Vec) -> Result<()> { - // A Pong frame MAY be sent unsolicited. This serves as a - // unidirectional heartbeat. A response to an unsolicited Pong frame is - // not expected. (RFC 6455) - // Due to this, we just don't check pongs right now. - // TODO: check if there was a reply to our ping at all... - Ok(()) - } - /// Send a single pending frame. fn send_one_frame(&mut self, mut frame: Frame) -> Result<()> { match self.role { @@ -463,6 +447,8 @@ mod tests { #[test] fn receive_messages() { let incoming = Cursor::new(vec![ + 0x89, 0x02, 0x01, 0x02, + 0x8a, 0x01, 0x03, 0x01, 0x07, 0x48, 0x65, 0x6c, 0x6c, 0x6f, 0x2c, 0x20, 0x80, 0x06, @@ -471,6 +457,8 @@ mod tests { 0x01, 0x02, 0x03, ]); let mut socket = WebSocket::from_raw_socket(WriteMoc(incoming), Role::Client); + assert_eq!(socket.read_message().unwrap(), Message::Ping(vec![1, 2])); + assert_eq!(socket.read_message().unwrap(), Message::Pong(vec![3])); assert_eq!(socket.read_message().unwrap(), Message::Text("Hello, World!".into())); assert_eq!(socket.read_message().unwrap(), Message::Binary(vec![0x01, 0x02, 0x03])); }