Tighten up handling of close frames

src/calibre/srv/tests/web_sockets.py

@@ -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 = []
+                frags = []
-            for i, b in enumerate(byte_data):
+                for i, b in enumerate(q):
-                frags.append({'opcode':(TEXT if i == 0 else CONTINUATION), 'fin':1 if i == len(byte_data)-1 else 0, 'payload':b})
+                    frags.append({'opcode':(TEXT if i == 0 else CONTINUATION), 'fin':1 if i == len(q)-1 else 0, 'payload':b})
-            simple_test(frags, [byte_data.decode('utf-8')])
+                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)

src/calibre/srv/web_socket.py

@@ -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:
+        if self.opcode in CONTROL_CODES and self.payload_length > 125:
-            conn.log.error('Too large ping packet from client')
+            conn.log.error('Too large control frame from client')
-            conn.websocket_close(PROTOCOL_ERROR, 'Ping packet too large')
+            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''):