mirror of https://github.com/huashengdun/webssh
Merge pull request #306 from klarose/handle-closed-worker
do not process message for closed workerspull/296/merge
commit
a3cb94b45e
|
@ -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)
|
||||
|
|
|
@ -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:
|
||||
|
|
Loading…
Reference in New Issue