crabapples
/
websockify
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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

This commit is contained in:
Antoine Martin 2016-08-06 22:31:55 +07:00
parent 72ce5c1ae8
commit 212e7d0ddd
2 changed files with 36 additions and 30 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
tests/test_websocket.py
View File

@ -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'

60
websockify/websocket.py
View File

@ -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(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):
@ -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)
if self.rec:
self.rec.write("%s,\n" %
repr("{%s{" % tdelta
+ encbuf[lenhead:len(encbuf)-lentail]))
self.send_parts.append(encbuf)
opcode = 2
encbufs = self.send_hybi(buf, opcode, base64=self.base64, record=self.rec)
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.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