diff options
| -rw-r--r-- | netlib/websockets.py | 62 | ||||
| -rw-r--r-- | test/test_websockets.py | 60 |
2 files changed, 77 insertions, 45 deletions
diff --git a/netlib/websockets.py b/netlib/websockets.py index 83e90238..5b9d8fbd 100644 --- a/netlib/websockets.py +++ b/netlib/websockets.py @@ -25,6 +25,11 @@ from . import utils websockets_magic = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11' +class CONST(object): + MAX_16_BIT_INT = (1 << 16) + MAX_64_BIT_INT = (1 << 64) + + class WebSocketFrameValidationException(Exception): pass @@ -82,14 +87,6 @@ class Frame(object): self.actual_payload_length = actual_payload_length @classmethod - def from_bytes(cls, bytestring): - """ - Construct a websocket frame from an in-memory bytestring to construct - a frame from a stream of bytes, use from_byte_stream() directly - """ - return cls.from_byte_stream(io.BytesIO(bytestring).read) - - @classmethod def default(cls, message, from_client = False): """ Construct a basic websocket frame from some default values. @@ -145,7 +142,7 @@ class Frame(object): except AssertionError: return False - def human_readable(self): + def human_readable(self): # pragma: nocover return "\n".join([ ("fin - " + str(self.fin)), ("rsv1 - " + str(self.rsv1)), @@ -160,6 +157,14 @@ class Frame(object): ("actual_payload_length - " + str(self.actual_payload_length)) ]) + @classmethod + def from_bytes(cls, bytestring): + """ + Construct a websocket frame from an in-memory bytestring + to construct a frame from a stream of bytes, use from_file() directly + """ + return cls.from_file(io.BytesIO(bytestring)) + def safe_to_bytes(self): if self.is_valid(): return self.to_bytes() @@ -172,8 +177,6 @@ class Frame(object): If you haven't checked is_valid_frame() then there's no guarentees that the serialized bytes will be correct. see safe_to_bytes() """ - max_16_bit_int = (1 << 16) - max_64_bit_int = (1 << 63) # break down of the bit-math used to construct the first byte from the # frame's integer values first shift the significant bit into the @@ -199,11 +202,11 @@ class Frame(object): if self.actual_payload_length < 126: pass - elif self.actual_payload_length < max_16_bit_int: + elif self.actual_payload_length < CONST.MAX_16_BIT_INT: # '!H' pack as 16 bit unsigned short # add 2 byte extended payload length bytes += struct.pack('!H', self.actual_payload_length) - elif self.actual_payload_length < max_64_bit_int: + elif self.actual_payload_length < CONST.MAX_64_BIT_INT: # '!Q' = pack as 64 bit unsigned long long # add 8 bytes extended payload length bytes += struct.pack('!Q', self.actual_payload_length) @@ -214,17 +217,20 @@ class Frame(object): bytes += self.payload # already will be encoded if neccessary return bytes + def to_file(self, writer): + writer.write(self.to_bytes()) + writer.flush() + @classmethod - def from_byte_stream(cls, read_bytes): + def from_file(cls, reader): """ read a websockets frame sent by a server or client - - read_bytes is a function that can be backed - by sockets or by any byte reader. So this - function may be used to read frames from disk/wire/memory - """ - first_byte = utils.bytes_to_int(read_bytes(1)) - second_byte = utils.bytes_to_int(read_bytes(1)) + + reader is a "file like" object that could be backed by a network stream or a disk + or an in memory stream reader + """ + first_byte = utils.bytes_to_int(reader.read(1)) + second_byte = utils.bytes_to_int(reader.read(1)) # grab the left most bit fin = first_byte >> 7 @@ -241,18 +247,18 @@ class Frame(object): actual_payload_length = payload_length elif payload_length == 126: - actual_payload_length = utils.bytes_to_int(read_bytes(2)) + actual_payload_length = utils.bytes_to_int(reader.read(2)) elif payload_length == 127: - actual_payload_length = utils.bytes_to_int(read_bytes(8)) + actual_payload_length = utils.bytes_to_int(reader.read(8)) # masking key only present if mask bit set if mask_bit == 1: - masking_key = read_bytes(4) + masking_key = reader.read(4) else: masking_key = None - payload = read_bytes(actual_payload_length) + payload = reader.read(actual_payload_length) if mask_bit == 1: decoded_payload = apply_mask(payload, masking_key) @@ -344,7 +350,7 @@ def build_handshake(headers, request): return b'\r\n'.join(handshake) -def read_handshake(read_bytes, num_bytes_per_read): +def read_handshake(reader, num_bytes_per_read): """ From provided function that reads bytes, read in a complete HTTP request, which terminates with a CLRF @@ -352,7 +358,7 @@ def read_handshake(read_bytes, num_bytes_per_read): response = b'' doubleCLRF = b'\r\n\r\n' while True: - bytes = read_bytes(num_bytes_per_read) + bytes = reader.read(num_bytes_per_read) if not bytes: break response += bytes @@ -386,7 +392,7 @@ def process_handshake_from_client(handshake): return key -def process_handshake_from_server(handshake, client_nounce): +def process_handshake_from_server(handshake): headers = headers_from_http_message(handshake) if headers.get("Upgrade", None) != "websocket": return diff --git a/test/test_websockets.py b/test/test_websockets.py index 62268423..035f9e17 100644 --- a/test/test_websockets.py +++ b/test/test_websockets.py @@ -1,6 +1,7 @@ from netlib import tcp from netlib import test from netlib import websockets +import io import os from nose.tools import raises @@ -20,16 +21,15 @@ class WebSocketsEchoHandler(tcp.BaseHandler): self.read_next_message() def read_next_message(self): - decoded = websockets.Frame.from_byte_stream(self.rfile.read).decoded_payload + decoded = websockets.Frame.from_file(self.rfile).decoded_payload self.on_message(decoded) def send_message(self, message): frame = websockets.Frame.default(message, from_client = False) - self.wfile.write(frame.safe_to_bytes()) - self.wfile.flush() - + frame.to_file(self.wfile) + def handshake(self): - client_hs = websockets.read_handshake(self.rfile.read, 1) + client_hs = websockets.read_handshake(self.rfile, 1) key = websockets.process_handshake_from_client(client_hs) response = websockets.create_server_handshake(key) self.wfile.write(response) @@ -62,22 +62,18 @@ class WebSocketsClient(tcp.TCPClient): self.wfile.write(handshake) self.wfile.flush() - server_handshake = websockets.read_handshake(self.rfile.read, 1) - server_nounce = websockets.process_handshake_from_server( - server_handshake, self.client_nounce - ) + server_handshake = websockets.read_handshake(self.rfile, 1) + server_nounce = websockets.process_handshake_from_server(server_handshake) if not server_nounce == websockets.create_server_nounce(self.client_nounce): self.close() def read_next_message(self): - return websockets.Frame.from_byte_stream(self.rfile.read).payload + return websockets.Frame.from_file(self.rfile).payload def send_message(self, message): frame = websockets.Frame.default(message, from_client = True) - self.wfile.write(frame.safe_to_bytes()) - self.wfile.flush() - + frame.to_file(self.wfile) class TestWebSockets(test.ServerTestBase): handler = WebSocketsEchoHandler @@ -128,10 +124,10 @@ class TestWebSockets(test.ServerTestBase): frame = websockets.Frame.default( self.random_bytes(num_bytes), is_client ) - assert frame == websockets.Frame.from_bytes(frame.to_bytes()) + assert frame == websockets.Frame.from_bytes(frame.safe_to_bytes()) - bytes = b'\x81\x11cba' - assert websockets.Frame.from_bytes(bytes).to_bytes() == bytes + bytes = b'\x81\x03cba' + assert websockets.Frame.from_bytes(bytes).safe_to_bytes() == bytes @raises(websockets.WebSocketFrameValidationException) def test_safe_to_bytes(self): @@ -139,10 +135,40 @@ class TestWebSockets(test.ServerTestBase): frame.actual_payload_length = 1 # corrupt the frame frame.safe_to_bytes() + def test_handshake(self): + bad_upgrade = "not_websockets" + bad_header_handshake = websockets.build_handshake([ + ('Host', '%s:%s' % ("a", "b")), + ('Connection', "c"), + ('Upgrade', bad_upgrade), + ('Sec-WebSocket-Key', "d"), + ('Sec-WebSocket-Version', "e") + ], "f") + + # check behavior when required header values are missing + assert None == websockets.process_handshake_from_server(bad_header_handshake) + assert None == websockets.process_handshake_from_client(bad_header_handshake) + + key = "test_key" + + client_handshake = websockets.create_client_handshake("a","b",key,"d","e") + assert key == websockets.process_handshake_from_client(client_handshake) + + server_handshake = websockets.create_server_handshake(key) + assert websockets.create_server_nounce(key) == websockets.process_handshake_from_server(server_handshake) + + handshake = websockets.create_client_handshake("a","b","c","d","e") + stream = io.BytesIO(handshake) + assert handshake == websockets.read_handshake(stream, 1) + + # ensure readhandshake doesn't loop forever on empty stream + empty_stream = io.BytesIO("") + assert "" == websockets.read_handshake(empty_stream, 1) + class BadHandshakeHandler(WebSocketsEchoHandler): def handshake(self): - client_hs = websockets.read_handshake(self.rfile.read, 1) + client_hs = websockets.read_handshake(self.rfile, 1) websockets.process_handshake_from_client(client_hs) response = websockets.create_server_handshake("malformed_key") self.wfile.write(response) |