From 768859513a5b3eeb1741ab7d5180c0c93eabfafe Mon Sep 17 00:00:00 2001
From: anthonyjchriste <anthony.christe@gmail.com>
Date: Thu, 27 Jun 2019 18:49:03 -1000
Subject: [PATCH] Expose the max limit for payload sizes in Websocket Actors.
 #925 (#933)

* Expose the max limit for payload sizes in Websocket Actors.

* Revert to previous not-formatted code.

* Implement WebsocketContext::with_codec and make Codec Copy and Clone.

* Fix formatting.

* Fix formatting.
---
 actix-http/src/ws/codec.rs |  2 +-
 actix-web-actors/src/ws.rs | 33 +++++++++++++++++++++++++--------
 2 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/actix-http/src/ws/codec.rs b/actix-http/src/ws/codec.rs
index ad599ffa7..9891bfa6e 100644
--- a/actix-http/src/ws/codec.rs
+++ b/actix-http/src/ws/codec.rs
@@ -37,7 +37,7 @@ pub enum Frame {
     Close(Option<CloseReason>),
 }
 
-#[derive(Debug)]
+#[derive(Debug, Copy, Clone)]
 /// WebSockets protocol codec
 pub struct Codec {
     max_size: usize,
diff --git a/actix-web-actors/src/ws.rs b/actix-web-actors/src/ws.rs
index 0ef3c9169..16f475a7f 100644
--- a/actix-web-actors/src/ws.rs
+++ b/actix-web-actors/src/ws.rs
@@ -177,9 +177,26 @@ where
             inner: ContextParts::new(mb.sender_producer()),
             messages: VecDeque::new(),
         };
-        ctx.add_stream(WsStream::new(stream));
+        ctx.add_stream(WsStream::new(stream, Codec::new()));
 
-        WebsocketContextFut::new(ctx, actor, mb)
+        WebsocketContextFut::new(ctx, actor, mb, Codec::new())
+    }
+
+    #[inline]
+    /// Create a new Websocket context from a request, an actor, and a codec
+    pub fn with_codec<S>(actor: A, stream: S, codec: Codec) -> impl Stream<Item = Bytes, Error = Error>
+        where
+            A: StreamHandler<Message, ProtocolError>,
+            S: Stream<Item = Bytes, Error = PayloadError> + 'static,
+    {
+        let mb = Mailbox::default();
+        let mut ctx = WebsocketContext {
+            inner: ContextParts::new(mb.sender_producer()),
+            messages: VecDeque::new(),
+        };
+        ctx.add_stream(WsStream::new(stream, codec));
+
+        WebsocketContextFut::new(ctx, actor, mb, codec)
     }
 
     /// Create a new Websocket context
@@ -197,11 +214,11 @@ where
             inner: ContextParts::new(mb.sender_producer()),
             messages: VecDeque::new(),
         };
-        ctx.add_stream(WsStream::new(stream));
+        ctx.add_stream(WsStream::new(stream, Codec::new()));
 
         let act = f(&mut ctx);
 
-        WebsocketContextFut::new(ctx, act, mb)
+        WebsocketContextFut::new(ctx, act, mb, Codec::new())
     }
 }
 
@@ -288,11 +305,11 @@ impl<A> WebsocketContextFut<A>
 where
     A: Actor<Context = WebsocketContext<A>>,
 {
-    fn new(ctx: WebsocketContext<A>, act: A, mailbox: Mailbox<A>) -> Self {
+    fn new(ctx: WebsocketContext<A>, act: A, mailbox: Mailbox<A>, codec: Codec) -> Self {
         let fut = ContextFut::new(ctx, act, mailbox);
         WebsocketContextFut {
             fut,
-            encoder: Codec::new(),
+            encoder: codec,
             buf: BytesMut::new(),
             closed: false,
         }
@@ -353,10 +370,10 @@ impl<S> WsStream<S>
 where
     S: Stream<Item = Bytes, Error = PayloadError>,
 {
-    fn new(stream: S) -> Self {
+    fn new(stream: S, codec: Codec) -> Self {
         Self {
             stream,
-            decoder: Codec::new(),
+            decoder: codec,
             buf: BytesMut::new(),
             closed: false,
         }