From 01ef6a6a554b58bdd9f3c233e02a61fefe6a0829 Mon Sep 17 00:00:00 2001 From: Pierre Ossman Date: Wed, 30 Sep 2020 14:58:19 +0200 Subject: [PATCH] Require same arguments on insufficient socket space This matches the behaviour of SSLSocket, which we are trying to mimic. It also closely matches the behaviour of normal Socket which can be assumed to not have sent anything if an error occurs. We might actually send some data, but the caller cannot really see that and must call us again as if no data was sent. --- websockify/websocket.py | 85 ++++++++++++++++++++++++++++------ websockify/websockifyserver.py | 10 +--- 2 files changed, 73 insertions(+), 22 deletions(-) diff --git a/websockify/websocket.py b/websockify/websocket.py index d15113b..a0dd5b8 100644 --- a/websockify/websocket.py +++ b/websockify/websocket.py @@ -93,6 +93,8 @@ class WebSocket(object): self._recv_queue = [] self._send_buffer = ''.encode("ascii") + self._previous_sendmsg = None + self._sent_close = False self._received_close = False @@ -254,8 +256,8 @@ class WebSocket(object): the value "websocket" in such cases. WebSocketWantWriteError can be raised if the response cannot be - sent right away. Repeated calls to accept() does not need to - retain the arguments. + sent right away. accept() must be called again once more space + is available using the same arguments. """ # This is a state machine in order to handle @@ -419,7 +421,8 @@ class WebSocket(object): data from other calls, or split it over multiple messages. WebSocketWantWriteError can be raised if there is insufficient - space in the underlying socket. + space in the underlying socket. send() must be called again + once more space is available using the same arguments. """ if len(bytes) == 0: return 0 @@ -434,31 +437,81 @@ class WebSocket(object): single WebSocket message. WebSocketWantWriteError can be raised if there is insufficient - space in the underlying socket. + space in the underlying socket. sendmsg() must be called again + once more space is available using the same arguments. """ if not isinstance(msg, bytes): raise TypeError - if not self._sent_close: - # Only called to flush? - self._sendmsg(0x2, msg) + if self._sent_close: + return 0 + + if self._previous_sendmsg is not None: + if self._previous_sendmsg != msg: + raise ValueError + + self._flush() + self._previous_sendmsg = None + + return len(msg) + + try: + self._sendmsg(0x2, msg) + except WebSocketWantWriteError: + self._previous_sendmsg = msg + raise - self._flush() return len(msg) def ping(self, data=''.encode('ascii')): - """Write a ping message to the WebSocket.""" + """Write a ping message to the WebSocket + + WebSocketWantWriteError can be raised if there is insufficient + space in the underlying socket. ping() must be called again once + more space is available using the same arguments. + """ if not isinstance(data, bytes): raise TypeError - self._sendmsg(0x9, data) + if self._previous_sendmsg is not None: + if self._previous_sendmsg != data: + raise ValueError + + self._flush() + self._previous_sendmsg = None + + return + + try: + self._sendmsg(0x9, data) + except WebSocketWantWriteError: + self._previous_sendmsg = data + raise def pong(self, data=''.encode('ascii')): - """Write a pong message to the WebSocket.""" + """Write a pong message to the WebSocket + + WebSocketWantWriteError can be raised if there is insufficient + space in the underlying socket. pong() must be called again once + more space is available using the same arguments. + """ if not isinstance(data, bytes): raise TypeError - self._sendmsg(0xA, data) + if self._previous_sendmsg is not None: + if self._previous_sendmsg != data: + raise ValueError + + self._flush() + self._previous_sendmsg = None + + return + + try: + self._sendmsg(0xA, data) + except WebSocketWantWriteError: + self._previous_sendmsg = data + raise def shutdown(self, how, code=1000, reason=None): """Gracefully terminate the WebSocket connection. @@ -470,7 +523,9 @@ class WebSocket(object): ignored. WebSocketWantWriteError can be raised if there is insufficient - space in the underlying socket for the close message. + space in the underlying socket for the close message. shutdown() + must be called again once more space is available using the same + arguments. The how argument is currently ignored. """ @@ -502,7 +557,9 @@ class WebSocket(object): a close message to the peer. WebSocketWantWriteError can be raised if there is insufficient - space in the underlying socket for the close message. + space in the underlying socket for the close message. close() + must be called again once more space is available using the same + arguments. """ self.shutdown(socket.SHUT_RDWR, code, reason) self._close() diff --git a/websockify/websockifyserver.py b/websockify/websockifyserver.py index c801725..670209e 100644 --- a/websockify/websockifyserver.py +++ b/websockify/websockifyserver.py @@ -146,20 +146,14 @@ class WebSockifyRequestHandler(WebSocketRequestHandlerMixIn, SimpleHTTPRequestHa self.rec.write("'{{{0}{{{1}',\n".format(tdelta, bufstr)) self.send_parts.append(buf) - # Flush any previously queued data - try: - self.request.sendmsg('') - except WebSocketWantWriteError: - return True - while self.send_parts: # Send pending frames - buf = self.send_parts.pop(0) try: - self.request.sendmsg(buf) + self.request.sendmsg(self.send_parts[0]) except WebSocketWantWriteError: self.print_traffic("<.") return True + self.send_parts.pop() self.print_traffic("<") return False