Merge pull request #306 from klarose/handle-closed-worker

do not process message for closed workers

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)

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: