Better close code/reason handling.

This commit is contained in:
Joel Martin 2012-04-25 13:44:01 -05:00
parent 37c0c80a93
commit 9a88f1800c
2 changed files with 24 additions and 14 deletions

View File

@ -92,9 +92,14 @@ Sec-WebSocket-Accept: %s\r
policy_response = """<cross-domain-policy><allow-access-from domain="*" to-ports="*" /></cross-domain-policy>\n"""
# An exception before the WebSocket connection was established
class EClose(Exception):
pass
# An exception while the WebSocket client was connected
class CClose(Exception):
pass
def __init__(self, listen_host='', listen_port=None, source_is_ipv6=False,
verbose=False, cert='', key='', ssl_only=None,
daemon=False, record='', web='',
@ -303,8 +308,8 @@ Sec-WebSocket-Accept: %s\r
'length' : 0,
'payload' : None,
'left' : 0,
'close_code' : None,
'close_reason' : None}
'close_code' : 1000,
'close_reason' : ''}
blen = len(buf)
f['left'] = blen
@ -470,7 +475,7 @@ Sec-WebSocket-Accept: %s\r
buf = self.client.recv(self.buffer_size)
if len(buf) == 0:
closed = "Client closed abruptly"
closed = {'code': 1000, 'reason': "Client closed abruptly"}
return bufs, closed
if self.recv_part:
@ -492,14 +497,14 @@ Sec-WebSocket-Accept: %s\r
break
else:
if frame['opcode'] == 0x8: # connection close
closed = "Client closed, reason: %s - %s" % (
frame['close_code'],
frame['close_reason'])
closed = {'code': frame['close_code'],
'reason': frame['close_reason']}
break
else:
if buf[0:2] == s2b('\xff\x00'):
closed = "Client sent orderly close frame"
closed = {'code': 1000,
'reason': "Client sent orderly close frame"}
break
elif buf[0:2] == s2b('\x00\xff'):
@ -532,13 +537,11 @@ Sec-WebSocket-Accept: %s\r
return bufs, closed
def send_close(self, code=None, reason=''):
def send_close(self, code=1000, reason=''):
""" Send a WebSocket orderly close frame. """
if self.version.startswith("hybi"):
msg = s2b('')
if code != None:
msg = pack(">H%ds" % (len(reason)), code)
msg = pack(">H%ds" % len(reason), code, reason)
buf, h, t = self.encode_hybi(msg, opcode=0x08, base64=False)
self.client.send(buf)
@ -762,6 +765,11 @@ Sec-WebSocket-Accept: %s\r
self.ws_connection = True
self.new_client()
except self.CClose as e:
# Close the client
_, exc, _ = sys.exc_info()
if self.client:
self.send_close(exc.args[0], exc.args[1])
except self.EClose:
_, exc, _ = sys.exc_info()
# Connection was not a WebSockets connection
@ -778,6 +786,8 @@ Sec-WebSocket-Accept: %s\r
self.rec.close()
if self.client and self.client != startsock:
# Close the SSL wrapped socket
# Original socket closed by caller
self.client.close()
def new_client(self):

View File

@ -190,7 +190,8 @@ Traffic Legend:
if target in ins:
# Receive target data, encode it and queue for client
buf = target.recv(self.buffer_size)
if len(buf) == 0: raise self.EClose("Target closed")
if len(buf) == 0:
raise self.CClose(1000, "Target closed")
cqueue.append(buf)
self.traffic("{")
@ -210,8 +211,7 @@ Traffic Legend:
if closed:
# TODO: What about blocking on client socket?
self.send_close()
raise self.EClose(closed)
raise self.CClose(closed['code'], closed['reason'])
def websockify_init():
usage = "\n %prog [options]"