Whitespace, interface simplification

- safe_tobytes doesn't buy us much
- move masking key generation inline
This commit is contained in:
Aldo Cortesi 2015-04-21 22:51:01 +12:00
parent 3e0a71ea34
commit 1b509d5aea
2 changed files with 5 additions and 25 deletions

View File

@ -29,10 +29,6 @@ class CONST(object):
MAX_64_BIT_INT = (1 << 64) MAX_64_BIT_INT = (1 << 64)
class WebSocketFrameValidationException(Exception):
pass
class Frame(object): class Frame(object):
""" """
Represents one websockets frame. Represents one websockets frame.
@ -95,7 +91,8 @@ class Frame(object):
if from_client: if from_client:
mask_bit = 1 mask_bit = 1
masking_key = random_masking_key() # Random masking key
masking_key = os.urandom(4)
payload = apply_mask(message, masking_key) payload = apply_mask(message, masking_key)
else: else:
mask_bit = 0 mask_bit = 0
@ -164,12 +161,6 @@ class Frame(object):
""" """
return cls.from_file(io.BytesIO(bytestring)) 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): def to_bytes(self):
""" """
Serialize the frame back into the wire format, returns a bytestring Serialize the frame back into the wire format, returns a bytestring
@ -308,10 +299,6 @@ def apply_mask(message, masking_key):
return result return result
def random_masking_key():
return os.urandom(4)
def client_handshake_headers(key=None, version=VERSION): def client_handshake_headers(key=None, version=VERSION):
""" """
Create the headers for a valid HTTP upgrade request. If Key is not Create the headers for a valid HTTP upgrade request. If Key is not

View File

@ -1,5 +1,4 @@
from netlib import tcp, test, websockets, http, odict from netlib import tcp, test, websockets, http
import io
import os import os
from nose.tools import raises from nose.tools import raises
@ -120,17 +119,11 @@ class TestWebSockets(test.ServerTestBase):
self.random_bytes(num_bytes), is_client self.random_bytes(num_bytes), is_client
) )
assert frame == websockets.Frame.from_bytes( assert frame == websockets.Frame.from_bytes(
frame.safe_to_bytes() frame.to_bytes()
) )
bytes = b'\x81\x03cba' bytes = b'\x81\x03cba'
assert websockets.Frame.from_bytes(bytes).safe_to_bytes() == bytes assert websockets.Frame.from_bytes(bytes).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()
def test_check_server_handshake(self): def test_check_server_handshake(self):
resp = http.Response( resp = http.Response(