From 80f385e703d1c8f8d927dfc7ce6a8fe7c7a2629d Mon Sep 17 00:00:00 2001
From: Denis Kolodin <deniskolodin@gmail.com>
Date: Wed, 2 May 2018 08:25:31 +0300
Subject: [PATCH] Add WsWriter trait

`WsWriter` trait is a common interface for writing to a websocket and
it's implemented for both: `WebScoketContext` and `ClientWriter`.
---
 src/lib.rs        |  1 +
 src/ws/client.rs  | 26 +++++++-------
 src/ws/context.rs | 87 +++++++++++++++++++++++++----------------------
 src/ws/mod.rs     | 14 ++++++++
 4 files changed, 75 insertions(+), 53 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 14e2cfc4e..5bef60c4b 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -175,6 +175,7 @@ pub use httprequest::HttpRequest;
 pub use httpresponse::HttpResponse;
 pub use json::Json;
 pub use scope::Scope;
+pub use ws::WsWriter;
 
 #[cfg(feature = "openssl")]
 pub(crate) const HAS_OPENSSL: bool = true;
diff --git a/src/ws/client.rs b/src/ws/client.rs
index 8a4abcae5..07b44b4d8 100644
--- a/src/ws/client.rs
+++ b/src/ws/client.rs
@@ -27,7 +27,7 @@ use client::{ClientConnector, ClientRequest, ClientRequestBuilder, ClientRespons
 
 use super::frame::Frame;
 use super::proto::{CloseReason, OpCode};
-use super::{Message, ProtocolError};
+use super::{Message, ProtocolError, WsWriter};
 
 /// Websocket client error
 #[derive(Fail, Debug)]
@@ -503,13 +503,6 @@ pub struct ClientWriter {
     inner: Rc<UnsafeCell<Inner>>,
 }
 
-impl ClientWriter {
-    #[inline]
-    fn as_mut(&mut self) -> &mut Inner {
-        unsafe { &mut *self.inner.get() }
-    }
-}
-
 impl ClientWriter {
     /// Write payload
     #[inline]
@@ -521,21 +514,28 @@ impl ClientWriter {
         }
     }
 
+    #[inline]
+    fn as_mut(&mut self) -> &mut Inner {
+        unsafe { &mut *self.inner.get() }
+    }
+}
+
+impl WsWriter for ClientWriter {
     /// Send text frame
     #[inline]
-    pub fn text<T: Into<Binary>>(&mut self, text: T) {
+    fn text<T: Into<Binary>>(&mut self, text: T) {
         self.write(Frame::message(text.into(), OpCode::Text, true, true));
     }
 
     /// Send binary frame
     #[inline]
-    pub fn binary<B: Into<Binary>>(&mut self, data: B) {
+    fn binary<B: Into<Binary>>(&mut self, data: B) {
         self.write(Frame::message(data, OpCode::Binary, true, true));
     }
 
     /// Send ping frame
     #[inline]
-    pub fn ping(&mut self, message: &str) {
+    fn ping(&mut self, message: &str) {
         self.write(Frame::message(
             Vec::from(message),
             OpCode::Ping,
@@ -546,7 +546,7 @@ impl ClientWriter {
 
     /// Send pong frame
     #[inline]
-    pub fn pong(&mut self, message: &str) {
+    fn pong(&mut self, message: &str) {
         self.write(Frame::message(
             Vec::from(message),
             OpCode::Pong,
@@ -557,7 +557,7 @@ impl ClientWriter {
 
     /// Send close frame
     #[inline]
-    pub fn close(&mut self, reason: Option<CloseReason>) {
+    fn close(&mut self, reason: Option<CloseReason>) {
         self.write(Frame::close(reason, true));
     }
 }
diff --git a/src/ws/context.rs b/src/ws/context.rs
index b5a2456c1..723b215ad 100644
--- a/src/ws/context.rs
+++ b/src/ws/context.rs
@@ -13,6 +13,7 @@ use context::{ActorHttpContext, Drain, Frame as ContextFrame};
 use error::{Error, ErrorInternalServerError};
 use httprequest::HttpRequest;
 
+use ws::WsWriter;
 use ws::frame::Frame;
 use ws::proto::{CloseReason, OpCode};
 
@@ -140,46 +141,6 @@ where
         &mut self.request
     }
 
-    /// Send text frame
-    #[inline]
-    pub fn text<T: Into<Binary>>(&mut self, text: T) {
-        self.write(Frame::message(text.into(), OpCode::Text, true, false));
-    }
-
-    /// Send binary frame
-    #[inline]
-    pub fn binary<B: Into<Binary>>(&mut self, data: B) {
-        self.write(Frame::message(data, OpCode::Binary, true, false));
-    }
-
-    /// Send ping frame
-    #[inline]
-    pub fn ping(&mut self, message: &str) {
-        self.write(Frame::message(
-            Vec::from(message),
-            OpCode::Ping,
-            true,
-            false,
-        ));
-    }
-
-    /// Send pong frame
-    #[inline]
-    pub fn pong(&mut self, message: &str) {
-        self.write(Frame::message(
-            Vec::from(message),
-            OpCode::Pong,
-            true,
-            false,
-        ));
-    }
-
-    /// Send close frame
-    #[inline]
-    pub fn close(&mut self, reason: Option<CloseReason>) {
-        self.write(Frame::close(reason, false));
-    }
-
     /// Returns drain future
     pub fn drain(&mut self) -> Drain<A> {
         let (tx, rx) = oneshot::channel();
@@ -213,6 +174,52 @@ where
     }
 }
 
+impl<A, S> WsWriter for WebsocketContext<A, S>
+where
+    A: Actor<Context = Self>,
+    S: 'static,
+{
+    /// Send text frame
+    #[inline]
+    fn text<T: Into<Binary>>(&mut self, text: T) {
+        self.write(Frame::message(text.into(), OpCode::Text, true, false));
+    }
+
+    /// Send binary frame
+    #[inline]
+    fn binary<B: Into<Binary>>(&mut self, data: B) {
+        self.write(Frame::message(data, OpCode::Binary, true, false));
+    }
+
+    /// Send ping frame
+    #[inline]
+    fn ping(&mut self, message: &str) {
+        self.write(Frame::message(
+            Vec::from(message),
+            OpCode::Ping,
+            true,
+            false,
+        ));
+    }
+
+    /// Send pong frame
+    #[inline]
+    fn pong(&mut self, message: &str) {
+        self.write(Frame::message(
+            Vec::from(message),
+            OpCode::Pong,
+            true,
+            false,
+        ));
+    }
+
+    /// Send close frame
+    #[inline]
+    fn close(&mut self, reason: Option<CloseReason>) {
+        self.write(Frame::close(reason, false));
+    }
+}
+
 impl<A, S> ActorHttpContext for WebsocketContext<A, S>
 where
     A: Actor<Context = Self>,
diff --git a/src/ws/mod.rs b/src/ws/mod.rs
index 402f2bdf6..a39b95b1d 100644
--- a/src/ws/mod.rs
+++ b/src/ws/mod.rs
@@ -340,6 +340,20 @@ where
     }
 }
 
+/// Common writing methods for a websocket.
+pub trait WsWriter {
+    /// Send a text
+    fn text<T: Into<Binary>>(&mut self, text: T);
+    /// Send a binary
+    fn binary<B: Into<Binary>>(&mut self, data: B);
+    /// Send a ping message
+    fn ping(&mut self, message: &str);
+    /// Send a pong message
+    fn pong(&mut self, message: &str);
+    /// Close the connection
+    fn close(&mut self, reason: Option<CloseReason>);
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;