diff --git a/src/filters/ws.rs b/src/filters/ws.rs index 2bf79a20a..98961bad2 100644 --- a/src/filters/ws.rs +++ b/src/filters/ws.rs @@ -203,7 +203,7 @@ impl Stream for WebSocket { fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll> { match ready!(Pin::new(&mut self.inner).poll_next(cx)) { - Some(Ok(item)) => Poll::Ready(Some(Ok(Message { inner: item }))), + Some(Ok(item)) => Poll::Ready(Some(Ok(Message::from_protocol(item)))), Some(Err(e)) => { tracing::debug!("websocket poll error: {}", e); Poll::Ready(Some(Err(crate::Error::new(e)))) @@ -227,7 +227,7 @@ impl Sink for WebSocket { } fn start_send(mut self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> { - match Pin::new(&mut self.inner).start_send(item.inner) { + match Pin::new(&mut self.inner).start_send(item.into_protocol()) { Ok(()) => Ok(()), Err(e) => { tracing::debug!("websocket start_send error: {}", e); @@ -260,35 +260,63 @@ impl fmt::Debug for WebSocket { } } -/// A WebSocket message. -/// -/// This will likely become a `non-exhaustive` enum in the future, once that -/// language feature has stabilized. -#[derive(Eq, PartialEq, Clone)] -pub struct Message { - inner: protocol::Message, +/// Represents a WebSocket message. +#[non_exhaustive] +#[derive(Debug, Eq, PartialEq, Clone)] +pub enum Message { + /// Regular string message. + Text(String), + /// Regular binary message. + Binary(Vec), + /// Ping message, used to check latency, typically responded to with a pong message. + Ping(Vec), + /// Pong message, used to check latency, typically sent in response to a ping message. + Pong(Vec), + /// Regular closure of the websocket, with a closure code and reason. + Close(u16, Cow<'static, str>), + /// An abrupt unexpected closure of the websocket. + AbruptClose, } impl Message { + fn from_protocol(m: protocol::Message) -> Message { + match m { + protocol::Message::Text(s) => Message::Text(s), + protocol::Message::Binary(v) => Message::Binary(v), + protocol::Message::Ping(v) => Message::Ping(v), + protocol::Message::Pong(v) => Message::Pong(v), + protocol::Message::Close(Some(c)) => Message::Close(c.code.into(), c.reason), + protocol::Message::Close(None) => Message::AbruptClose, + } + } + + fn into_protocol(self) -> protocol::Message { + match self { + Message::Text(s) => protocol::Message::Text(s), + Message::Binary(v) => protocol::Message::Binary(v), + Message::Ping(v) => protocol::Message::Ping(v), + Message::Pong(v) => protocol::Message::Pong(v), + Message::Close(code, reason) => protocol::Message::Close(Some(protocol::CloseFrame { + code: code.into(), + reason, + })), + Message::AbruptClose => protocol::Message::Close(None), + } + } + /// Construct a new Text `Message`. pub fn text>(s: S) -> Message { - Message { - inner: protocol::Message::text(s), - } + Message::Text(s.into()) } /// Construct a new Binary `Message`. pub fn binary>>(v: V) -> Message { - Message { - inner: protocol::Message::binary(v), - } + Message::Binary(v.into()) } /// Construct a new Ping `Message`. pub fn ping>>(v: V) -> Message { - Message { - inner: protocol::Message::Ping(v.into()), - } + Message::Ping(v.into()) } /// Construct a new Pong `Message`. @@ -297,90 +325,78 @@ impl Message { /// automatically responds to the Ping messages it receives. Manual construction might still be useful in some cases /// like in tests or to send unidirectional heartbeats. pub fn pong>>(v: V) -> Message { - Message { - inner: protocol::Message::Pong(v.into()), - } + Message::Pong(v.into()) } /// Construct the default Close `Message`. pub fn close() -> Message { - Message { - inner: protocol::Message::Close(None), - } + Message::AbruptClose } /// Construct a Close `Message` with a code and reason. pub fn close_with(code: impl Into, reason: impl Into>) -> Message { - Message { - inner: protocol::Message::Close(Some(protocol::frame::CloseFrame { - code: protocol::frame::coding::CloseCode::from(code.into()), - reason: reason.into(), - })), - } + Message::Close(code.into(), reason.into()) } /// Returns true if this message is a Text message. pub fn is_text(&self) -> bool { - self.inner.is_text() + matches!(self, Message::Text(_)) } /// Returns true if this message is a Binary message. pub fn is_binary(&self) -> bool { - self.inner.is_binary() + matches!(self, Message::Binary(_)) } /// Returns true if this message a is a Close message. pub fn is_close(&self) -> bool { - self.inner.is_close() + matches!(self, Message::Close(_, _)) || matches!(self, Message::AbruptClose) } /// Returns true if this message is a Ping message. pub fn is_ping(&self) -> bool { - self.inner.is_ping() + matches!(self, Message::Ping(_)) } /// Returns true if this message is a Pong message. pub fn is_pong(&self) -> bool { - self.inner.is_pong() + matches!(self, Message::Pong(_)) } /// Try to get the close frame (close code and reason) pub fn close_frame(&self) -> Option<(u16, &str)> { - if let protocol::Message::Close(Some(ref close_frame)) = self.inner { - Some((close_frame.code.into(), close_frame.reason.as_ref())) - } else { - None + match self { + Message::Close(code, reason) => Some((*code, reason.as_ref())), + Message::AbruptClose => None, + _ => None, } } /// Try to get a reference to the string text, if this is a Text message. pub fn to_str(&self) -> Result<&str, ()> { - match self.inner { - protocol::Message::Text(ref s) => Ok(s), + match self { + Message::Text(s) => Ok(s), _ => Err(()), } } /// Return the bytes of this message, if the message can contain data. pub fn as_bytes(&self) -> &[u8] { - match self.inner { - protocol::Message::Text(ref s) => s.as_bytes(), - protocol::Message::Binary(ref v) => v, - protocol::Message::Ping(ref v) => v, - protocol::Message::Pong(ref v) => v, - protocol::Message::Close(_) => &[], + match self { + Message::Text(s) => s.as_bytes(), + Message::Binary(v) | Message::Ping(v) | Message::Pong(v) => v, + _ => &[], } } /// Destructure this message into binary data. pub fn into_bytes(self) -> Vec { - self.inner.into_data() - } -} - -impl fmt::Debug for Message { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - fmt::Debug::fmt(&self.inner, f) + match self { + Message::Text(s) => s.into_bytes(), + Message::Binary(v) | Message::Ping(v) | Message::Pong(v) => v, + Message::Close(_, s) => s.into_owned().into_bytes(), + Message::AbruptClose => Vec::new(), + } } }