From f6e0b806dade29a8f10f4887f5a565b393d6c904 Mon Sep 17 00:00:00 2001
From: kazk <kazk.dev@gmail.com>
Date: Tue, 14 Sep 2021 13:27:56 -0700
Subject: [PATCH] Extract compression to a separate method

---
 src/protocol/mod.rs | 36 ++++++++++++++----------------------
 1 file changed, 14 insertions(+), 22 deletions(-)

diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs
index 4e41e67..b240a58 100644
--- a/src/protocol/mod.rs
+++ b/src/protocol/mod.rs
@@ -395,28 +395,8 @@ impl WebSocketContext {
         }
 
         let frame = match message {
-            Message::Text(data) => {
-                if let Some(pmce) = self.pmce.as_mut() {
-                    Frame::compressed_message(
-                        pmce.compress(data.as_bytes())?,
-                        OpCode::Data(OpData::Text),
-                        true,
-                    )
-                } else {
-                    Frame::message(data.into(), OpCode::Data(OpData::Text), true)
-                }
-            }
-            Message::Binary(data) => {
-                if let Some(pmce) = self.pmce.as_mut() {
-                    Frame::compressed_message(
-                        pmce.compress(&data)?,
-                        OpCode::Data(OpData::Binary),
-                        true,
-                    )
-                } else {
-                    Frame::message(data, OpCode::Data(OpData::Binary), true)
-                }
-            }
+            Message::Text(data) => self.prepare_data_frame(data.into(), OpData::Text)?,
+            Message::Binary(data) => self.prepare_data_frame(data, OpData::Binary)?,
             Message::Ping(data) => Frame::ping(data),
             Message::Pong(data) => {
                 self.pong = Some(Frame::pong(data));
@@ -429,6 +409,18 @@ impl WebSocketContext {
         self.write_pending(stream)
     }
 
+    fn prepare_data_frame(&mut self, data: Vec<u8>, opdata: OpData) -> Result<Frame> {
+        debug_assert!(matches!(opdata, OpData::Text | OpData::Binary), "Invalid data frame kind");
+        let opcode = OpCode::Data(opdata);
+        let is_final = true;
+        let frame = if let Some(pmce) = self.pmce.as_mut() {
+            Frame::compressed_message(pmce.compress(&data)?, opcode, is_final)
+        } else {
+            Frame::message(data, opcode, is_final)
+        };
+        Ok(frame)
+    }
+
     /// Flush the pending send queue.
     pub fn write_pending<Stream>(&mut self, stream: &mut Stream) -> Result<()>
     where