miguelgrinberg / simple-websocket

Simple WebSocket server and client for Python.
MIT License
79 stars 17 forks source link

Improve handling of immediately closed sockets with Ping/Pong enabled #35

Closed lkedziora closed 10 months ago

lkedziora commented 10 months ago

This PR fixes two smaller issues around handling of disconnected sockets when Ping/Pong is enabled.

  1. When Ping/Pong is enabled, the handler thread first registers the socket in the chosen selector. However, it is possible for a different thread to call ws.close() before the handler thread gets a chance to run, which results in the thread dying with an exception when registering the (now closed) socket.
  2. Ping/Pong logic may attempt to send the Ping in an invalid WS state, which results in an exception.

Reproducing

Both of these issues can be reproduced using the following:

server.py

from flask import Flask, Response, request
from simple_websocket import Server, ConnectionClosed
import logging
import traceback

app = Flask(__name__)

def patch_server_for_cleaner_output():
    original = Server._thread

    def patched_wrapper(self, *args, **kwargs):
        try:
            original(self, *args, **kwargs)
        except Exception as e:
            print()
            print("Caught exception in WebSocket handler thread:", e, flush=True)
            print("Socket state:", self.sock, flush=True)
            print("Is connected?", self.connected, flush=True)
            traceback.print_exc()

    Server._thread = patched_wrapper

# Comment this out for cleaner output
# patch_server_for_cleaner_output()

@app.route('/command', websocket=True)
def command():
    # The issue only occurs when ping is enabled
    ws = Server.accept(request.environ, ping_interval=25)

    # Real-life scenario: the server can handle the client, but
    # for whatever reason a process fails and the connection is
    # closed very quickly.
    # The key here is that the connection is closed before the
    # handler thread has had the chance to run.
    try:
        ws.close(4000, "server-side process failed")
    except ConnectionClosed:
        pass

    return Response()

if __name__ == '__main__':
    # Silence request logging
    log = logging.getLogger('werkzeug')
    log.setLevel(logging.ERROR)
    app.run()

client.py

from simple_websocket import Client
import time

def main():
    # What the client does is generally irrelevant to both issues
    # This just helps with quickly getting a reproduction
    while True:
        ws = Client.connect('ws://localhost:5000/command')
        try:
            _ = ws.receive()
            ws.close()
        except:
            pass
        time.sleep(0.1)

if __name__ == "__main__":
    main()

Steps for reproduction:

  1. python server.py
  2. python client.py
  3. Wait for the exceptions to occur on the server.

One of the following can occur:

Exception in thread Thread-120 (simple_websocket.Base._thread):
Traceback (most recent call last):
  File "/usr/lib/python3.11/threading.py", line 1038, in _bootstrap_inner
    self.run()
  File "/usr/lib/python3.11/threading.py", line 975, in run
    self._target(*self._args, **self._kwargs)
  File "/sources/.venv/lib/python3.11/site-packages/simple_websocket/ws.py", line 135, in _thread
    sel.register(self.sock, selectors.EVENT_READ, True)
  File "/usr/lib/python3.11/selectors.py", line 352, in register
    key = super().register(fileobj, events, data)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3.11/selectors.py", line 238, in register
    key = SelectorKey(fileobj, self._fileobj_lookup(fileobj), events, data)
                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3.11/selectors.py", line 225, in _fileobj_lookup
    return _fileobj_to_fd(fileobj)
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3.11/selectors.py", line 42, in _fileobj_to_fd
    raise ValueError("Invalid file descriptor: {}".format(fd))
ValueError: Invalid file descriptor: -1
Exception in thread Thread-350 (simple_websocket.Base._thread):
Traceback (most recent call last):
  File "/usr/lib/python3.11/threading.py", line 1038, in _bootstrap_inner
    self.run()
  File "/usr/lib/python3.11/threading.py", line 975, in run
    self._target(*self._args, **self._kwargs)
  File "/sources/.venv/lib/python3.11/site-packages/simple_websocket/ws.py", line 148, in _thread
    self.sock.send(self.ws.send(Ping()))
                   ^^^^^^^^^^^^^^^^^^^^
  File "/sources/.venv/lib/python3.11/site-packages/wsproto/__init__.py", line 64, in send
    data += self.connection.send(event)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/sources/.venv/lib/python3.11/site-packages/wsproto/connection.py", line 107, in send
    raise LocalProtocolError(
wsproto.utilities.LocalProtocolError: Event Ping(payload=b'') cannot be sent in state ConnectionState.LOCAL_CLOSING.

These are inherently racy, so it might take a bit for both of them to appear.

codecov-commenter commented 10 months ago

Codecov Report

All modified and coverable lines are covered by tests :white_check_mark:

Project coverage is 100.00%. Comparing base (5a5dbec) to head (96f66e6). Report is 9 commits behind head on main.

Additional details and impacted files ```diff @@ Coverage Diff @@ ## main #35 +/- ## ========================================= Coverage 100.00% 100.00% ========================================= Files 4 4 Lines 435 436 +1 Branches 81 81 ========================================= + Hits 435 436 +1 ```

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

miguelgrinberg commented 10 months ago

Looks good. Thanks!