diff --git a/websockify/websocket.py b/websockify/websocket.py index 49e60b6..e37af0f 100644 --- a/websockify/websocket.py +++ b/websockify/websocket.py @@ -344,10 +344,11 @@ class WebSocket(object): def recv(self): """Read data from the WebSocket. - This will return any available data on the socket. If the - socket is closed then an empty buffer will be returned. The - reason for the close is found in the 'close_code' and - 'close_reason' properties. + This will return any available data on the socket (which may + be the empty string if the peer sent an empty message or + messages). If the socket is closed then None will be + returned. The reason for the close is found in the + 'close_code' and 'close_reason' properties. Unlike recvmsg() this method may return data from more than one WebSocket message. It is however not guaranteed to return all @@ -362,10 +363,11 @@ class WebSocket(object): def recvmsg(self): """Read a single message from the WebSocket. - This will return a single WebSocket message from the socket. - If the socket is closed then an empty buffer will be returned. - The reason for the close is found in the 'close_code' and - 'close_reason' properties. + This will return a single WebSocket message from the socket + (which will be the empty string if the peer sent an empty + message). If the socket is closed then None will be + returned. The reason for the close is found in the + 'close_code' and 'close_reason' properties. Unlike recv() this method will not return data from more than one WebSocket message. Callers should continue calling @@ -377,30 +379,22 @@ class WebSocket(object): # May have been called to flush out a close if self._received_close: self._flush() - return ''.encode("ascii") + return None # Anything already queued? if self.pending(): - msg = self._recvmsg() - if msg is not None: - return msg - - # Note: We cannot proceed to self._recv() here as we may + return self._recvmsg() + # Note: If self._recvmsg() raised WebSocketWantReadError, + # we cannot proceed to self._recv() here as we may # have already called it once as part of the caller's # "while websock.pending():" loop - raise WebSocketWantReadError # Nope, let's try to read a bit if not self._recv_frames(): - return ''.encode("ascii") + return None # Anything queued now? - msg = self._recvmsg() - if msg is not None: - return msg - - # Still nope - raise WebSocketWantReadError + return self._recvmsg() def pending(self): """Check if any WebSocket data is pending. @@ -594,7 +588,7 @@ class WebSocket(object): if self._sent_close: self._close() - return ''.encode("ascii") + return None if not frame["fin"]: self.shutdown(socket.SHUT_RDWR, 1003, "Unsupported: Fragmented close") @@ -621,7 +615,7 @@ class WebSocket(object): self.close_reason = reason self.shutdown(code, reason) - return ''.encode("ascii") + return None elif frame["opcode"] == 0x9: if not frame["fin"]: self.shutdown(socket.SHUT_RDWR, 1003, "Unsupported: Fragmented ping") @@ -637,7 +631,7 @@ class WebSocket(object): else: self.shutdown(socket.SHUT_RDWR, 1003, "Unsupported: Unknown opcode 0x%02x" % frame["opcode"]) - return None + raise WebSocketWantReadError def _flush(self): # Writes pending data to the socket diff --git a/websockify/websocketproxy.py b/websockify/websocketproxy.py index 6c17c10..6aa76e6 100644 --- a/websockify/websocketproxy.py +++ b/websockify/websocketproxy.py @@ -231,7 +231,7 @@ 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: + if buf is None: if self.verbose: self.log_message("%s:%s: Target closed connection", self.server.target_host, self.server.target_port) diff --git a/websockify/websockifyserver.py b/websockify/websockifyserver.py index a531cb4..5bb755e 100644 --- a/websockify/websockifyserver.py +++ b/websockify/websockifyserver.py @@ -181,7 +181,7 @@ class WebSockifyRequestHandler(WebSocketRequestHandler, SimpleHTTPRequestHandler self.print_traffic("}.") break - if len(buf) == 0: + if buf is None: closed = {'code': self.request.close_code, 'reason': self.request.close_reason} return bufs, closed