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: