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)
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

View File

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