Grasscutters/mitmproxy
3
0
Fork 0
mirror of https://github.com/Grasscutters/mitmproxy.git synced 2024-11-23 08:11:00 +00:00
Code Issues Projects Releases Wiki Activity

websockets: more whitespace, WebSocketFrame -> Frame

Browse Source
This commit is contained in:
Aldo Cortesi 2015-04-17 14:29:20 +12:00
parent 488c25d812
commit 7defb5be86
3 changed files with 81 additions and 76 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
netlib/websockets/implementations.py
View File

@ -22,11 +22,11 @@ class WebSocketsEchoHandler(tcp.BaseHandler):
self.read_next_message()
def read_next_message(self):
decoded = ws.WebSocketsFrame.from_byte_stream(self.rfile.read).decoded_payload
decoded = ws.Frame.from_byte_stream(self.rfile.read).decoded_payload
self.on_message(decoded)
def send_message(self, message):
frame = ws.WebSocketsFrame.default(message, from_client = False)
frame = ws.Frame.default(message, from_client = False)
self.wfile.write(frame.safe_to_bytes())
self.wfile.flush()
@ -72,9 +72,9 @@ class WebSocketsClient(tcp.TCPClient):
self.close()
def read_next_message(self):
return ws.WebSocketsFrame.from_byte_stream(self.rfile.read).payload
return ws.Frame.from_byte_stream(self.rfile.read).payload
def send_message(self, message):
frame = ws.WebSocketsFrame.default(message, from_client = True)
frame = ws.Frame.default(message, from_client = True)
self.wfile.write(frame.safe_to_bytes())
self.wfile.flush()

22
netlib/websockets/websockets.py
View File

@ -29,7 +29,7 @@ class WebSocketFrameValidationException(Exception):
pass
class WebSocketsFrame(object):
class Frame(object):
"""
Represents one websockets frame.
Constructor takes human readable forms of the frame components
@ -120,7 +120,7 @@ class WebSocketsFrame(object):
def is_valid(self):
"""
Validate websocket frame invariants, call at anytime to ensure the
WebSocketsFrame has not been corrupted.
Frame has not been corrupted.
"""
try:
assert 0 <= self.fin <= 1
@ -157,7 +157,8 @@ class WebSocketsFrame(object):
("masking_key - " + str(self.masking_key)),
("payload - " + str(self.payload)),
("decoded_payload - " + str(self.decoded_payload)),
("actual_payload_length - " + str(self.actual_payload_length))])
("actual_payload_length - " + str(self.actual_payload_length))
])
def safe_to_bytes(self):
if self.is_valid():
@ -167,11 +168,10 @@ class WebSocketsFrame(object):
def to_bytes(self):
"""
Serialize the frame back into the wire format, returns a bytestring If
you haven't checked is_valid_frame() then there's no guarentees that
the serialized bytes will be correct. see safe_to_bytes()
Serialize the frame back into the wire format, returns a bytestring
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)
@ -199,13 +199,10 @@ class WebSocketsFrame(object):
if self.actual_payload_length < 126:
pass
elif self.actual_payload_length < 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:
# '!Q' = pack as 64 bit unsigned long long
# add 8 bytes extended payload length
@ -215,7 +212,6 @@ class WebSocketsFrame(object):
bytes += self.masking_key
bytes += self.payload # already will be encoded if neccessary
return bytes
@classmethod
@ -286,7 +282,9 @@ class WebSocketsFrame(object):
self.masking_key == other.masking_key and
self.payload == other.payload and
self.decoded_payload == other.decoded_payload and
self.actual_payload_length == other.actual_payload_length)
self.actual_payload_length == other.actual_payload_length
)
def apply_mask(message, masking_key):
"""

33
test/test_websockets.py
View File

@ -5,6 +5,7 @@ from netlib.websockets import websockets as ws
import os
from nose.tools import raises
class TestWebSockets(test.ServerTestBase):
handler = impl.WebSocketsEchoHandler
@ -22,9 +23,12 @@ class TestWebSockets(test.ServerTestBase):
self.echo("hello I'm the client")
def test_frame_sizes(self):
small_msg = self.random_bytes(100) # length can fit in the the 7 bit payload length
medium_msg = self.random_bytes(50000) # 50kb, sligthly larger than can fit in a 7 bit int
large_msg = self.random_bytes(150000) # 150kb, slightly larger than can fit in a 16 bit int
# length can fit in the the 7 bit payload length
small_msg = self.random_bytes(100)
# 50kb, sligthly larger than can fit in a 7 bit int
medium_msg = self.random_bytes(50000)
# 150kb, slightly larger than can fit in a 16 bit int
large_msg = self.random_bytes(150000)
self.echo(small_msg)
self.echo(medium_msg)
@ -35,41 +39,44 @@ class TestWebSockets(test.ServerTestBase):
default builder should always generate valid frames
"""
msg = self.random_bytes()
client_frame = ws.WebSocketsFrame.default(msg, from_client = True)
client_frame = ws.Frame.default(msg, from_client = True)
assert client_frame.is_valid()
server_frame = ws.WebSocketsFrame.default(msg, from_client = False)
server_frame = ws.Frame.default(msg, from_client = False)
assert server_frame.is_valid()
def test_serialization_bijection(self):
"""
Ensure that various frame types can be serialized/deserialized back and forth
between to_bytes() and from_bytes()
Ensure that various frame types can be serialized/deserialized back
and forth between to_bytes() and from_bytes()
"""
for is_client in [True, False]:
for num_bytes in [100, 50000, 150000]:
frame = ws.WebSocketsFrame.default(self.random_bytes(num_bytes), is_client)
assert frame == ws.WebSocketsFrame.from_bytes(frame.to_bytes())
frame = ws.Frame.default(
self.random_bytes(num_bytes), is_client
)
assert frame == ws.Frame.from_bytes(frame.to_bytes())
bytes = b'\x81\x11cba'
assert ws.WebSocketsFrame.from_bytes(bytes).to_bytes() == bytes
assert ws.Frame.from_bytes(bytes).to_bytes() == bytes
@raises(ws.WebSocketFrameValidationException)
def test_safe_to_bytes(self):
frame = ws.WebSocketsFrame.default(self.random_bytes(8))
frame.actual_payload_length = 1 #corrupt the frame
frame = ws.Frame.default(self.random_bytes(8))
frame.actual_payload_length = 1 # corrupt the frame
frame.safe_to_bytes()
class BadHandshakeHandler(impl.WebSocketsEchoHandler):
def handshake(self):
client_hs = ws.read_handshake(self.rfile.read, 1)
key = ws.process_handshake_from_client(client_hs)
ws.process_handshake_from_client(client_hs)
response = ws.create_server_handshake("malformed_key")
self.wfile.write(response)
self.wfile.flush()
self.handshake_done = True
class TestBadHandshake(test.ServerTestBase):
"""
Ensure that the client disconnects if the server handshake is malformed