split header and data functions so we can send them without concatenating them together - V2

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'

websockify/websocket.py

@@ -154,8 +154,8 @@ class WebSocketRequestHandler(SimpleHTTPRequestHandler):
             return data.tostring()
 
     @staticmethod
-    def encode_hybi(buf, opcode, base64=False):
+    def encode_hybi_header(payload_len, opcode):
-        """ Encode a HyBi style WebSocket frame.
+        """ 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
+    def send_hybi(self, buf, opcode, base64=False, record=None):
-        payload_len = len(buf)
+        """ Send a HyBi style WebSocket frame. """
-        if payload_len <= 125:
+        header, buf = self.encode_hybi(buf, opcode, base64)
-            header = pack('>BB', b1, payload_len)
+        if record:
-        elif payload_len > 125 and payload_len < 65536:
+            record.write("%s,\n" % repr("{%s{" % tdelta + encbufs[1]))
-            header = pack('>BBH', b1, 126, payload_len)
+        if len(buf)<=4096:
-        elif payload_len >= 65536:
+            self.request.send(header+buf)
-            header = pack('>BBQ', b1, 127, payload_len)
+        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):
@@ -311,16 +324,10 @@ class WebSocketRequestHandler(SimpleHTTPRequestHandler):
         if bufs:
             for buf in bufs:
                 if self.base64:
-                    encbuf, lenhead, lentail = self.encode_hybi(buf, opcode=1, base64=True)
+                    opcode = 1
                 else:
-                    encbuf, lenhead, lentail = self.encode_hybi(buf, opcode=2, base64=False)
+                    opcode = 2
-
+                encbufs = self.send_hybi(buf, opcode, base64=self.base64, record=self.rec)
-                if self.rec:
-                    self.rec.write("%s,\n" %
-                            repr("{%s{" % tdelta
-                                + encbuf[lenhead:len(encbuf)-lentail]))
-
-                self.send_parts.append(encbuf)
 
         while self.send_parts:
             # Send pending frames
@@ -412,18 +419,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.send_hybi(msg, opcode=0x08)
-        self.request.send(buf)
 
     def send_pong(self, data=''):
         """ Send a WebSocket pong frame. """
-        buf, h, t = self.encode_hybi(s2b(data), opcode=0x0A, base64=False)
+        self.send_hybi(s2b(data), opcode=0x0A)
-        self.request.send(buf)
 
     def send_ping(self, data=''):
         """ Send a WebSocket ping frame. """
-        buf, h, t = self.encode_hybi(s2b(data), opcode=0x09, base64=False)
+        self.send_hybi(s2b(data), opcode=0x09)
-        self.request.send(buf)
 
     def do_websocket_handshake(self):
         h = self.headers