From b0343eb22d8a371fb84cdf304e4eb58f94dad700 Mon Sep 17 00:00:00 2001
From: Nikolay Kim <fafhrd91@gmail.com>
Date: Sun, 17 Mar 2019 22:31:10 -0700
Subject: [PATCH] simplify ws stream interface

---
 actix-web-actors/src/ws.rs        | 35 ++++++++++++++++++++++++-------
 actix-web-actors/tests/test_ws.rs | 14 ++++++-------
 2 files changed, 34 insertions(+), 15 deletions(-)

diff --git a/actix-web-actors/src/ws.rs b/actix-web-actors/src/ws.rs
index b5f5c08c2..546326272 100644
--- a/actix-web-actors/src/ws.rs
+++ b/actix-web-actors/src/ws.rs
@@ -12,9 +12,9 @@ use actix::{
     Message as ActixMessage, SpawnHandle,
 };
 use actix_codec::{Decoder, Encoder};
-use actix_http::ws::hash_key;
+use actix_http::ws::{hash_key, Codec};
 pub use actix_http::ws::{
-    CloseCode, CloseReason, Codec, Frame, HandshakeError, Message, ProtocolError,
+    CloseCode, CloseReason, Frame, HandshakeError, Message, ProtocolError,
 };
 
 use actix_web::dev::{Head, HttpResponseBuilder};
@@ -28,7 +28,7 @@ use futures::{Async, Future, Poll, Stream};
 /// Do websocket handshake and start ws actor.
 pub fn start<A, T>(actor: A, req: &HttpRequest, stream: T) -> Result<HttpResponse, Error>
 where
-    A: Actor<Context = WebsocketContext<A>> + StreamHandler<Frame, ProtocolError>,
+    A: Actor<Context = WebsocketContext<A>> + StreamHandler<Message, ProtocolError>,
     T: Stream<Item = Bytes, Error = PayloadError> + 'static,
 {
     let mut res = handshake(req)?;
@@ -170,7 +170,7 @@ where
     /// Create a new Websocket context from a request and an actor
     pub fn create<S>(actor: A, stream: S) -> impl Stream<Item = Bytes, Error = Error>
     where
-        A: StreamHandler<Frame, ProtocolError>,
+        A: StreamHandler<Message, ProtocolError>,
         S: Stream<Item = Bytes, Error = PayloadError> + 'static,
     {
         let mb = Mailbox::default();
@@ -190,7 +190,7 @@ where
     ) -> impl Stream<Item = Bytes, Error = Error>
     where
         F: FnOnce(&mut Self) -> A + 'static,
-        A: StreamHandler<Frame, ProtocolError>,
+        A: StreamHandler<Message, ProtocolError>,
         S: Stream<Item = Bytes, Error = PayloadError> + 'static,
     {
         let mb = Mailbox::default();
@@ -368,7 +368,7 @@ impl<S> Stream for WsStream<S>
 where
     S: Stream<Item = Bytes, Error = PayloadError>,
 {
-    type Item = Frame;
+    type Item = Message;
     type Error = ProtocolError;
 
     fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
@@ -401,7 +401,28 @@ where
                     Ok(Async::NotReady)
                 }
             }
-            Some(frm) => Ok(Async::Ready(Some(frm))),
+            Some(frm) => {
+                let msg = match frm {
+                    Frame::Text(data) => {
+                        if let Some(data) = data {
+                            Message::Text(
+                                std::str::from_utf8(&data)
+                                    .map_err(|_| ProtocolError::BadEncoding)?
+                                    .to_string(),
+                            )
+                        } else {
+                            Message::Text(String::new())
+                        }
+                    }
+                    Frame::Binary(data) => Message::Binary(
+                        data.map(|b| b.freeze()).unwrap_or_else(|| Bytes::new()),
+                    ),
+                    Frame::Ping(s) => Message::Ping(s),
+                    Frame::Pong(s) => Message::Pong(s),
+                    Frame::Close(reason) => Message::Close(reason),
+                };
+                Ok(Async::Ready(Some(msg)))
+            }
         }
     }
 }
diff --git a/actix-web-actors/tests/test_ws.rs b/actix-web-actors/tests/test_ws.rs
index ea9c8d8fe..202d562ca 100644
--- a/actix-web-actors/tests/test_ws.rs
+++ b/actix-web-actors/tests/test_ws.rs
@@ -12,15 +12,13 @@ impl Actor for Ws {
     type Context = ws::WebsocketContext<Self>;
 }
 
-impl StreamHandler<ws::Frame, ws::ProtocolError> for Ws {
-    fn handle(&mut self, msg: ws::Frame, ctx: &mut Self::Context) {
+impl StreamHandler<ws::Message, ws::ProtocolError> for Ws {
+    fn handle(&mut self, msg: ws::Message, ctx: &mut Self::Context) {
         match msg {
-            ws::Frame::Ping(msg) => ctx.pong(&msg),
-            ws::Frame::Text(text) => {
-                ctx.text(String::from_utf8_lossy(&text.unwrap())).to_owned()
-            }
-            ws::Frame::Binary(bin) => ctx.binary(bin.unwrap()),
-            ws::Frame::Close(reason) => ctx.close(reason),
+            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),
             _ => (),
         }
     }