From 158663b84f4e98ba19a81845ba0398e3d2a87c2b Mon Sep 17 00:00:00 2001 From: liris Date: Wed, 15 Oct 2014 14:53:29 +0900 Subject: [PATCH] refs #117 - improve close frame. --- compliance/test_fuzzingclient.py | 8 ++++--- websocket/_abnf.py | 40 +++++++++++++++++++++++++++++++ websocket/_core.py | 19 ++++----------- websocket/tests/test_websocket.py | 2 ++ 4 files changed, 51 insertions(+), 18 deletions(-) diff --git a/compliance/test_fuzzingclient.py b/compliance/test_fuzzingclient.py index 63fc326..3dfd10c 100644 --- a/compliance/test_fuzzingclient.py +++ b/compliance/test_fuzzingclient.py @@ -22,6 +22,7 @@ failed = 0 for case in range(1, count+1): url = SERVER + '/runCase?case={}&agent={}'.format(case, AGENT) + status = websocket.STATUS_NORMAL try: ws = websocket.create_connection(url) opcode, msg = ws.recv_data() @@ -33,12 +34,13 @@ for case in range(1, count+1): except UnicodeDecodeError: # this case is ok. success += 1 + status = websocket.STATUS_PROTOCOL_ERROR except Exception as e: failed += 1 - # print("[Faield] Test Case: " + str(case)) - #print(traceback.format_exc()) + status = websocket.STATUS_PROTOCOL_ERROR + print(traceback.format_exc()) finally: - ws.close() + ws.close(status) print("Ran {} test cases. success: {}, faield: {}".format(case, success, failed)) url = SERVER + '/updateReports?agent={}'.format(AGENT) diff --git a/websocket/_abnf.py b/websocket/_abnf.py index 24cf28b..c5ef295 100644 --- a/websocket/_abnf.py +++ b/websocket/_abnf.py @@ -23,8 +23,33 @@ import array import struct import os from ._exceptions import * +from ._utils import validate_utf8 +# closing frame status codes. +STATUS_NORMAL = 1000 +STATUS_GOING_AWAY = 1001 +STATUS_PROTOCOL_ERROR = 1002 +STATUS_UNSUPPORTED_DATA_TYPE = 1003 +STATUS_STATUS_NOT_AVAILABLE = 1005 +STATUS_ABNORMAL_CLOSED = 1006 +STATUS_INVALID_PAYLOAD = 1007 +STATUS_POLICY_VIOLATION = 1008 +STATUS_MESSAGE_TOO_BIG = 1009 +STATUS_INVALID_EXTENSION = 1010 +STATUS_UNEXPECTED_CONDITION = 1011 +STATUS_TLS_HANDSHAKE_ERROR = 1015 +VALID_CLOSE_STATUS = ( + STATUS_NORMAL, + STATUS_GOING_AWAY, + STATUS_PROTOCOL_ERROR, + STATUS_UNSUPPORTED_DATA_TYPE, + STATUS_INVALID_PAYLOAD, + STATUS_POLICY_VIOLATION, + STATUS_MESSAGE_TOO_BIG, + STATUS_INVALID_EXTENSION, + STATUS_UNEXPECTED_CONDITION, + ) class ABNF(object): """ @@ -85,6 +110,21 @@ class ABNF(object): if self.opcode == ABNF.OPCODE_PING and not self.fin: raise WebSocketException("Invalid ping frame.") + if self.opcode == ABNF.OPCODE_CLOSE: + l = len(self.data) + if not l: + return + if l == 1 or l >= 126: + raise WebSocketException("Invalid close frame.") + if l > 2 and not validate_utf8(self.data[2:]): + raise WebSocketException("Invalid close frame.") + code = 256*six.byte2int(self.data[0]) + six.byte2int(self.data[1]) + if not self._is_valid_close_status(code): + raise WebSocketException("Invalid close opcode.") + + def _is_valid_close_status(self, code): + return code in VALID_CLOSE_STATUS or (3000 <= code <5000) + def __str__(self): return "fin=" + str(self.fin) \ + " opcode=" + str(self.opcode) \ diff --git a/websocket/_core.py b/websocket/_core.py index a0ffc37..9c2abf2 100644 --- a/websocket/_core.py +++ b/websocket/_core.py @@ -56,7 +56,7 @@ import logging # websocket modules from ._exceptions import * -from ._abnf import ABNF +from ._abnf import * from ._utils import NoLock, validate_utf8 """ @@ -71,19 +71,6 @@ Please see http://tools.ietf.org/html/rfc6455 for protocol. # websocket supported version. VERSION = 13 -# closing frame status codes. -STATUS_NORMAL = 1000 -STATUS_GOING_AWAY = 1001 -STATUS_PROTOCOL_ERROR = 1002 -STATUS_UNSUPPORTED_DATA_TYPE = 1003 -STATUS_STATUS_NOT_AVAILABLE = 1005 -STATUS_ABNORMAL_CLOSED = 1006 -STATUS_INVALID_PAYLOAD = 1007 -STATUS_POLICY_VIOLATION = 1008 -STATUS_MESSAGE_TOO_BIG = 1009 -STATUS_INVALID_EXTENSION = 1010 -STATUS_UNEXPECTED_CONDITION = 1011 -STATUS_TLS_HANDSHAKE_ERROR = 1015 DEFAULT_SOCKET_OPTION = [(socket.SOL_TCP, socket.TCP_NODELAY, 1),] if hasattr(socket, "SO_KEEPALIVE"): @@ -733,7 +720,7 @@ class WebSocket(object): self._cont_data = None frame.data = data[1] if not self.fire_cont_frame and data[0] == ABNF.OPCODE_TEXT and not validate_utf8(frame.data): - raise UnicodeDecodeError("cannot decode: " + repr(frame.data)) + raise WebSocketException("cannot decode: " + repr(frame.data)) return [data[0], frame] elif frame.opcode == ABNF.OPCODE_CLOSE: @@ -742,6 +729,8 @@ class WebSocket(object): elif frame.opcode == ABNF.OPCODE_PING: if len(frame.data) < 126: self.pong(frame.data) + else: + raise WebSocketException("Protocol Error") if control_frame: return (frame.opcode, frame) elif frame.opcode == ABNF.OPCODE_PONG: diff --git a/websocket/tests/test_websocket.py b/websocket/tests/test_websocket.py index eab6c47..ca901d0 100644 --- a/websocket/tests/test_websocket.py +++ b/websocket/tests/test_websocket.py @@ -530,6 +530,8 @@ class UtilsTest(unittest.TestCase): self.assertEqual(state, True) state = validate_utf8(six.b('\xce\xba\xe1\xbd\xb9\xcf\x83\xce\xbc\xce\xb5\xed\xa0\x80edited')) self.assertEqual(state, False) + state = validate_utf8(six.b('')) + self.assertEqual(state, True) if __name__ == "__main__": unittest.main()