From f0e2ddb821accc7493e276e0e448f199ac7743eb Mon Sep 17 00:00:00 2001
From: Kyle Larose <kyle@agilicus.com>
Date: Fri, 18 Nov 2022 08:38:51 -0500
Subject: [PATCH] do not process message for closed workers

WsockHandler stores a weak reference to the ssh backend worker. The
worker closes itself if the backend connection closes (e.g. the user
exists the ssh session). That happens in parallel to the websocket
handler processing messages, so it is possible for a message to arrive
when the worker no longer has any strong references, leading to an
exception being thrown.

Handle this case by treating the None worker the same way we do invalid
messages: by simply returning.
---
 tests/test_handler.py | 17 +++++++++++++++++
 webssh/handler.py     |  8 ++++++++
 2 files changed, 25 insertions(+)

diff --git a/tests/test_handler.py b/tests/test_handler.py
index 020847d..6a01cfb 100644
--- a/tests/test_handler.py
+++ b/tests/test_handler.py
@@ -5,6 +5,7 @@ from tornado.httputil import HTTPServerRequest
 from tornado.options import options
 from tests.utils import read_file, make_tests_data_path
 from webssh import handler
+from webssh import worker
 from webssh.handler import (
     MixinHandler, WsockHandler, PrivateKey, InvalidValueError
 )
@@ -277,3 +278,19 @@ class TestWsockHandler(unittest.TestCase):
         obj.origin_policy = '*'
         origin = 'https://blog.example.org'
         self.assertTrue(WsockHandler.check_origin(obj, origin))
+
+    def test_failed_weak_ref(self):
+        request = HTTPServerRequest(uri='/')
+        obj = Mock(spec=WsockHandler, request=request)
+        obj.src_addr = ("127.0.0.1", 8888)
+        class FakeWeakRef:
+            def __init__(self):
+                self.count = 0
+            def __call__(self):
+                self.count += 1
+                return None
+
+        ref = FakeWeakRef()
+        obj.worker_ref = ref
+        WsockHandler.on_message(obj, b'{"data": "somestuff"}')
+        self.assertGreaterEqual(ref.count, 1)
diff --git a/webssh/handler.py b/webssh/handler.py
index ced7819..0fe6afc 100644
--- a/webssh/handler.py
+++ b/webssh/handler.py
@@ -559,6 +559,14 @@ class WsockHandler(MixinHandler, tornado.websocket.WebSocketHandler):
     def on_message(self, message):
         logging.debug('{!r} from {}:{}'.format(message, *self.src_addr))
         worker = self.worker_ref()
+        if not worker:
+            # The worker has likely been closed. Do not process.
+            logging.debug(
+                "received message to closed worker from {}:{}".format(
+                    *self.src_addr
+                )
+            )
+            return
         try:
             msg = json.loads(message)
         except JSONDecodeError: