Cutlery/calibre
1
0
Fork 0
mirror of https://github.com/kovidgoyal/calibre.git synced 2025-07-09 03:04:10 -04:00
Code Issues Packages Projects Releases Wiki Activity

Tighten up handling of close frames

Browse Source
This commit is contained in:
Kovid Goyal 2015-10-26 12:58:19 +05:30
parent 2535953fc6
commit 3f1b5782f0
2 changed files with 54 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

34
src/calibre/srv/tests/web_sockets.py
View File

@ -13,7 +13,7 @@ from hashlib import sha1
from calibre.srv.tests.base import BaseTest, TestServer
from calibre.srv.web_socket import (
GUID_STR, BINARY, TEXT, MessageWriter, create_frame, CLOSE, NORMAL_CLOSE,
PING, PONG, PROTOCOL_ERROR, CONTINUATION)
PING, PONG, PROTOCOL_ERROR, CONTINUATION, INCONSISTENT_DATA)
from calibre.utils.monotonic import monotonic
from calibre.utils.socket_inheritance import set_socket_inherit
@ -243,6 +243,7 @@ class WebSocketTest(BaseTest):
simple_test([(PONG, payload)], [])
fragments = 'Hello-µ@ßöä üàá-UTF-8!!'.split()
nc = struct.pack(b'!H', NORMAL_CLOSE)
with server.silence_log:
for rsv in xrange(1, 7):
@ -254,6 +255,7 @@ class WebSocketTest(BaseTest):
simple_test([
{'opcode':opcode, 'payload':'f1', 'fin':0}, {'opcode':opcode, 'payload':'f2'}
], close_code=PROTOCOL_ERROR, send_close=False)
simple_test([(CLOSE, nc + b'x'*124)], send_close=False, close_code=PROTOCOL_ERROR)
for fin in (0, 1):
simple_test([{'opcode':0, 'fin': fin, 'payload':b'non-continuation frame'}, 'some text'], close_code=PROTOCOL_ERROR, send_close=False)
@ -266,6 +268,16 @@ class WebSocketTest(BaseTest):
{'opcode':TEXT, 'payload':fragments[0], 'fin':0}, {'opcode':TEXT, 'payload':fragments[1]},
], close_code=PROTOCOL_ERROR, send_close=False)
frags = []
for payload in (b'\xce\xba\xe1\xbd\xb9\xcf\x83\xce\xbc\xce\xb5', b'\xed\xa0\x80', b'\x80\x65\x64\x69\x74\x65\x64'):
frags.append({'opcode':(CONTINUATION if frags else TEXT), 'fin':1 if len(frags) == 2 else 0, 'payload':payload})
simple_test(frags, close_code=INCONSISTENT_DATA, send_close=False)
frags, q = [], b'\xce\xba\xe1\xbd\xb9\xcf\x83\xce\xbc\xce\xb5\xed\xa0\x80\x80\x65\x64\x69\x74\x65\x64'
for i, b in enumerate(q):
frags.append({'opcode':(TEXT if i == 0 else CONTINUATION), 'fin':1 if i == len(q)-1 else 0, 'payload':b})
simple_test(frags, close_code=INCONSISTENT_DATA, send_close=False)
simple_test([
{'opcode':TEXT, 'payload':fragments[0], 'fin':0}, {'opcode':CONTINUATION, 'payload':fragments[1]}
], [''.join(fragments)])
@ -288,9 +300,21 @@ class WebSocketTest(BaseTest):
simple_test([
{'opcode':TEXT, 'fin':0}, {'opcode':CONTINUATION, 'fin':0, 'payload':'x'}, {'opcode':CONTINUATION},], ['x'])
byte_data = "Hello-µ@ßöäüàá-UTF-8!!".encode('utf-8')
for q in (b'\xce\xba\xe1\xbd\xb9\xcf\x83\xce\xbc\xce\xb5', "Hello-µ@ßöäüàá-UTF-8!!".encode('utf-8')):
frags = []
for i, b in enumerate(byte_data):
frags.append({'opcode':(TEXT if i == 0 else CONTINUATION), 'fin':1 if i == len(byte_data)-1 else 0, 'payload':b})
simple_test(frags, [byte_data.decode('utf-8')])
for i, b in enumerate(q):
frags.append({'opcode':(TEXT if i == 0 else CONTINUATION), 'fin':1 if i == len(q)-1 else 0, 'payload':b})
simple_test(frags, [q.decode('utf-8')])
simple_test([(CLOSE, nc), (CLOSE, b'\x01\x01')], send_close=False)
simple_test([(CLOSE, nc), (PING, b'ping')], send_close=False)
simple_test([(CLOSE, nc), 'xxx'], send_close=False)
simple_test([{'opcode':TEXT, 'payload':'xxx', 'fin':0}, (CLOSE, nc), {'opcode':CONTINUATION, 'payload':'yyy'}], send_close=False)
simple_test([(CLOSE, b'')], send_close=False)
simple_test([(CLOSE, b'\x01')], send_close=False, close_code=PROTOCOL_ERROR)
simple_test([(CLOSE, nc + b'x'*123)], send_close=False)
simple_test([(CLOSE, nc + b'a\x80\x80')], send_close=False, close_code=PROTOCOL_ERROR)
for code in (1000,1001,1002,1003,1007,1008,1009,1010,1011,3000,3999,4000,4999):
simple_test([(CLOSE, struct.pack(b'!H', code))], send_close=False, close_code=code)
for code in (0,999,1004,1005,1006,1012,1013,1014,1015,1016,1100,2000,2999):
simple_test([(CLOSE, struct.pack(b'!H', code))], send_close=False, close_code=PROTOCOL_ERROR)

30
src/calibre/srv/web_socket.py
View File

@ -48,6 +48,8 @@ POLICY_VIOLATION = 1008
MESSAGE_TOO_BIG = 1009
UNEXPECTED_ERROR = 1011
RESERVED_CLOSE_CODES = (1004,1005,1006,)
class ReadFrame(object): # {{{
def __init__(self):
@ -93,9 +95,9 @@ class ReadFrame(object): # {{{
self.reset()
return
self.payload_length = b & 0b01111111
if self.opcode in (PING, PONG) and self.payload_length > 125:
conn.log.error('Too large ping packet from client')
conn.websocket_close(PROTOCOL_ERROR, 'Ping packet too large')
if self.opcode in CONTROL_CODES and self.payload_length > 125:
conn.log.error('Too large control frame from client')
conn.websocket_close(PROTOCOL_ERROR, 'Control frame too large')
self.reset()
return
self.mask_buf = b''
@ -347,13 +349,29 @@ class WebSocketConnection(HTTPConnection):
def ws_control_frame(self, opcode, data):
if opcode in (PING, CLOSE):
rcode = PONG if opcode == PING else CLOSE
if opcode == CLOSE:
self.ws_close_received = True
self.stop_reading = True
if data:
try:
close_code = struct.unpack_from(b'!H', data)[0]
except struct.error:
data = struct.pack(b'!H', PROTOCOL_ERROR) + b'close frame data must be atleast two bytes'
else:
try:
data[2:].decode('utf-8')
except UnicodeDecodeError:
data = struct.pack(b'!H', PROTOCOL_ERROR) + b'close frame data must be valid UTF-8'
else:
if close_code < 1000 or close_code in RESERVED_CLOSE_CODES or (1011 < close_code < 3000):
data = struct.pack(b'!H', PROTOCOL_ERROR) + b'close code reserved'
else:
close_code = NORMAL_CLOSE
data = struct.pack(b'!H', close_code)
f = BytesIO(create_frame(1, rcode, data))
f.is_close_frame = opcode == CLOSE
with self.cf_lock:
self.control_frames.append(f)
if opcode == CLOSE:
self.ws_close_received = True
self.stop_reading = True
self.set_ws_state()
def websocket_close(self, code=NORMAL_CLOSE, reason=b''):