crabapples
/
websockify
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

6
websockify
View File

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