diff options
| -rw-r--r-- | netlib/websockets.py | 17 | ||||
| -rw-r--r-- | test/test_websockets.py | 13 |
2 files changed, 5 insertions, 25 deletions
diff --git a/netlib/websockets.py b/netlib/websockets.py index a03185fa..0cd4dba1 100644 --- a/netlib/websockets.py +++ b/netlib/websockets.py @@ -29,10 +29,6 @@ class CONST(object): MAX_64_BIT_INT = (1 << 64) -class WebSocketFrameValidationException(Exception): - pass - - class Frame(object): """ Represents one websockets frame. @@ -95,7 +91,8 @@ class Frame(object): if from_client: mask_bit = 1 - masking_key = random_masking_key() + # Random masking key + masking_key = os.urandom(4) payload = apply_mask(message, masking_key) else: mask_bit = 0 @@ -164,12 +161,6 @@ class Frame(object): """ return cls.from_file(io.BytesIO(bytestring)) - def safe_to_bytes(self): - if self.is_valid(): - return self.to_bytes() - else: - raise WebSocketFrameValidationException() - def to_bytes(self): """ Serialize the frame back into the wire format, returns a bytestring @@ -308,10 +299,6 @@ def apply_mask(message, masking_key): return result -def random_masking_key(): - return os.urandom(4) - - def client_handshake_headers(key=None, version=VERSION): """ Create the headers for a valid HTTP upgrade request. If Key is not diff --git a/test/test_websockets.py b/test/test_websockets.py index 9b27e810..3fc67dfe 100644 --- a/test/test_websockets.py +++ b/test/test_websockets.py @@ -1,5 +1,4 @@ -from netlib import tcp, test, websockets, http, odict -import io +from netlib import tcp, test, websockets, http import os from nose.tools import raises @@ -120,17 +119,11 @@ class TestWebSockets(test.ServerTestBase): self.random_bytes(num_bytes), is_client ) assert frame == websockets.Frame.from_bytes( - frame.safe_to_bytes() + frame.to_bytes() ) bytes = b'\x81\x03cba' - assert websockets.Frame.from_bytes(bytes).safe_to_bytes() == bytes - - @raises(websockets.WebSocketFrameValidationException) - def test_safe_to_bytes(self): - frame = websockets.Frame.default(self.random_bytes(8)) - frame.actual_payload_length = 1 # corrupt the frame - frame.safe_to_bytes() + assert websockets.Frame.from_bytes(bytes).to_bytes() == bytes def test_check_server_handshake(self): resp = http.Response( |