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

Do not confuse an empty message with a closed connection 

Fixes #312.

Signed-off-by: Anders Kaseorg <andersk@mit.edu>
This commit is contained in:
Anders Kaseorg 2017-11-10 06:24:39 -05:00
parent ade9d61c22
commit a29946e978
3 changed files with 16 additions and 24 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

36
websockify/websocket.py
View File

@ -343,9 +343,9 @@ class WebSocket(object):
"""Read data from the WebSocket. """Read data from the WebSocket.
This will return any available data on the socket. If the This will return any available data on the socket. If the
socket is closed then an empty buffer will be returned. The socket is closed then None will be returned. The reason for
reason for the close is found in the 'close_code' and the close is found in the 'close_code' and 'close_reason'
'close_reason' properties. properties.
Unlike recvmsg() this method may return data from more than one Unlike recvmsg() this method may return data from more than one
WebSocket message. It is however not guaranteed to return all WebSocket message. It is however not guaranteed to return all
@ -361,8 +361,8 @@ class WebSocket(object):
"""Read a single message from the WebSocket. """Read a single message from the WebSocket.
This will return a single WebSocket message from the socket. This will return a single WebSocket message from the socket.
If the socket is closed then an empty buffer will be returned. If the socket is closed then None will be returned. The
The reason for the close is found in the 'close_code' and reason for the close is found in the 'close_code' and
'close_reason' properties. 'close_reason' properties.
Unlike recv() this method will not return data from more than Unlike recv() this method will not return data from more than
@ -375,30 +375,22 @@ class WebSocket(object):
# May have been called to flush out a close # May have been called to flush out a close
if self._received_close: if self._received_close:
self._flush() self._flush()
return ''.encode("ascii") return None
# Anything already queued? # Anything already queued?
if self.pending(): if self.pending():
msg = self._recvmsg() return self._recvmsg()
if msg is not None: # Note: If self._recvmsg() raised WebSocketWantReadError,
return msg # we cannot proceed to self._recv() here as we may
# Note: We cannot proceed to self._recv() here as we may
# have already called it once as part of the caller's # have already called it once as part of the caller's
# "while websock.pending():" loop # "while websock.pending():" loop
raise WebSocketWantReadError
# Nope, let's try to read a bit # Nope, let's try to read a bit
if not self._recv_frames(): if not self._recv_frames():
return ''.encode("ascii") return None
# Anything queued now? # Anything queued now?
msg = self._recvmsg() return self._recvmsg()
if msg is not None:
return msg
# Still nope
raise WebSocketWantReadError
def pending(self): def pending(self):
"""Check if any WebSocket data is pending. """Check if any WebSocket data is pending.
@ -592,7 +584,7 @@ class WebSocket(object):
if self._sent_close: if self._sent_close:
self._close() self._close()
return ''.encode("ascii") return None
if not frame["fin"]: if not frame["fin"]:
self.shutdown(socket.SHUT_RDWR, 1003, "Unsupported: Fragmented close") self.shutdown(socket.SHUT_RDWR, 1003, "Unsupported: Fragmented close")
@ -619,7 +611,7 @@ class WebSocket(object):
self.close_reason = reason self.close_reason = reason
self.shutdown(code, reason) self.shutdown(code, reason)
return ''.encode("ascii") return None
elif frame["opcode"] == 0x9: elif frame["opcode"] == 0x9:
if not frame["fin"]: if not frame["fin"]:
self.shutdown(socket.SHUT_RDWR, 1003, "Unsupported: Fragmented ping") self.shutdown(socket.SHUT_RDWR, 1003, "Unsupported: Fragmented ping")
@ -635,7 +627,7 @@ class WebSocket(object):
else: else:
self.shutdown(socket.SHUT_RDWR, 1003, "Unsupported: Unknown opcode 0x%02x" % frame["opcode"]) self.shutdown(socket.SHUT_RDWR, 1003, "Unsupported: Unknown opcode 0x%02x" % frame["opcode"])
return None raise WebSocketWantReadError
def _flush(self): def _flush(self):
# Writes pending data to the socket # Writes pending data to the socket

2
websockify/websocketproxy.py
View File

@ -226,7 +226,7 @@ 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: if buf is None:
if self.verbose: if self.verbose:
self.log_message("%s:%s: Target closed connection", self.log_message("%s:%s: Target closed connection",
self.server.target_host, self.server.target_port) self.server.target_host, self.server.target_port)

2
websockify/websockifyserver.py
View File

@ -178,7 +178,7 @@ class WebSockifyRequestHandler(WebSocketRequestHandler, SimpleHTTPRequestHandler
self.print_traffic("}.") self.print_traffic("}.")
break break
if len(buf) == 0: if buf is None:
closed = {'code': self.request.close_code, closed = {'code': self.request.close_code,
'reason': self.request.close_reason} 'reason': self.request.close_reason}
return bufs, closed return bufs, closed