From 06d88d07d290debef9cfbfee83d482b8c7c8628d Mon Sep 17 00:00:00 2001
From: Antoine Martin <antoine@nagafix.co.uk>
Date: Sat, 6 Aug 2016 22:31:55 +0700
Subject: [PATCH] split header and data functions so we can send them without
 concatenating them together

---
 tests/test_websocket.py |  6 ++--
 websockify/websocket.py | 61 +++++++++++++++++++++--------------------
 2 files changed, 35 insertions(+), 32 deletions(-)

diff --git a/tests/test_websocket.py b/tests/test_websocket.py
index d8acb4f..8f4694b 100644
--- a/tests/test_websocket.py
+++ b/tests/test_websocket.py
@@ -409,9 +409,11 @@ class HyBiEncodeDecodeTestCase(unittest.TestCase):
 
     def test_encode_hybi_basic(self):
         res = websocket.WebSocketRequestHandler.encode_hybi(b'Hello', 0x1)
-        expected = (b'\x81\x05\x48\x65\x6c\x6c\x6f', 2, 0)
+        expected_header = (b'\x81\x05')
+        expected_data = (b'\x48\x65\x6c\x6c\x6f')
 
-        self.assertEqual(res, expected)
+        self.assertEqual(res[0], expected_header)
+        self.assertEqual(res[1], expected_data)
 
     def test_strict_mode_refuses_unmasked_client_frames(self):
         buf = b'\x81\x05\x48\x65\x6c\x6c\x6f'
diff --git a/websockify/websocket.py b/websockify/websocket.py
index 16e1dc4..770faa3 100644
--- a/websockify/websocket.py
+++ b/websockify/websocket.py
@@ -154,8 +154,8 @@ class WebSocketRequestHandler(SimpleHTTPRequestHandler):
             return data.tostring()
 
     @staticmethod
-    def encode_hybi(buf, opcode, base64=False):
-        """ Encode a HyBi style WebSocket frame.
+    def encode_hybi_header(payload_len, opcode):
+        """ Encode a HyBi style WebSocket frame header.
         Optional opcode:
             0x0 - continuation
             0x1 - text frame (base64 encode buf)
@@ -164,21 +164,34 @@ class WebSocketRequestHandler(SimpleHTTPRequestHandler):
             0x9 - ping
             0xA - pong
         """
+        b1 = 0x80 | (opcode & 0x0f) # FIN + opcode
+        if payload_len <= 125:
+            return pack('>BB', b1, payload_len)
+        elif payload_len > 125 and payload_len < 65536:
+            return pack('>BBH', b1, 126, payload_len)
+        else:
+            assert payload_len >= 65536
+            return pack('>BBQ', b1, 127, payload_len)
+
+    @staticmethod
+    def encode_hybi(buf, opcode, base64=False):
+        """ Encode a HyBi style WebSocket frame. """
         if base64:
             buf = b64encode(buf)
+        header = WebSocketRequestHandler.encode_hybi_header(len(buf), opcode)
+        return (header, buf)
 
-        b1 = 0x80 | (opcode & 0x0f) # FIN + opcode
-        payload_len = len(buf)
-        if payload_len <= 125:
-            header = pack('>BB', b1, payload_len)
-        elif payload_len > 125 and payload_len < 65536:
-            header = pack('>BBH', b1, 126, payload_len)
-        elif payload_len >= 65536:
-            header = pack('>BBQ', b1, 127, payload_len)
+    def send_hybi(self, buf, opcode, base64=False, record=None):
+        """ Send a HyBi style WebSocket frame. """
+        header, buf = self.encode_hybi(buf, opcode, base64)
+        if record:
+            record.write("%s,\n" % repr("{%s{" % tdelta + encbufs[1]))
+        if len(buf)<=4096:
+            self.request.send("".join((header, buf)))
+        else:
+            self.request.send(header)
+            self.request.send(buf)
 
-        #self.msg("Encoded: %s", repr(header + buf))
-
-        return header + buf, len(header), 0
 
     @staticmethod
     def decode_hybi(buf, base64=False, logger=None, strict=True):
@@ -310,17 +323,8 @@ class WebSocketRequestHandler(SimpleHTTPRequestHandler):
 
         if bufs:
             for buf in bufs:
-                if self.base64:
-                    encbuf, lenhead, lentail = self.encode_hybi(buf, opcode=1, base64=True)
-                else:
-                    encbuf, lenhead, lentail = self.encode_hybi(buf, opcode=2, base64=False)
-
-                if self.rec:
-                    self.rec.write("%s,\n" %
-                            repr("{%s{" % tdelta
-                                + encbuf[lenhead:len(encbuf)-lentail]))
-
-                self.send_parts.append(encbuf)
+                opcode = 2-int(self.base64)     #based64: opcode=1, binary: opcode=2
+                encbufs = self.send_hybi(buf, opcode, base64=self.base64, record=self.rec)
 
         while self.send_parts:
             # Send pending frames
@@ -412,18 +416,15 @@ class WebSocketRequestHandler(SimpleHTTPRequestHandler):
         """ Send a WebSocket orderly close frame. """
 
         msg = pack(">H%ds" % len(reason), code, s2b(reason))
-        buf, h, t = self.encode_hybi(msg, opcode=0x08, base64=False)
-        self.request.send(buf)
+        self.send_hybi(msg, opcode=0x08)
 
     def send_pong(self, data=''):
         """ Send a WebSocket pong frame. """
-        buf, h, t = self.encode_hybi(s2b(data), opcode=0x0A, base64=False)
-        self.request.send(buf)
+        self.send_hybi(s2b(data), opcode=0x0A)
 
     def send_ping(self, data=''):
         """ Send a WebSocket ping frame. """
-        buf, h, t = self.encode_hybi(s2b(data), opcode=0x09, base64=False)
-        self.request.send(buf)
+        self.send_hybi(s2b(data), opcode=0x09)
 
     def do_websocket_handshake(self):
         h = self.headers