webssh

commit a893507946bf1451ddd4235bd60966e87556c6e8
parent 2c36c386531d43c51210bde96019a82076dd1cb3
Author: Sheng <webmaster0115@gmail.com>
Date:   Sat, 26 May 2018 15:51:53 +0800

Made WsockHandler more robust

Diffstat:
M
tests/sshserver.py
 | 
3
++-
M
tests/test_app.py
 | 
18
++++++++++++++++++
M
webssh/handler.py
 | 
18
+++++++++++++++---

3 files changed, 35 insertions(+), 4 deletions(-)
diff --git a/tests/sshserver.py b/tests/sshserver.py
@@ -114,7 +114,8 @@ def run_ssh_server(port=2200, running=True):
 
         chan.send('\r\n\r\nWelcome!\r\n\r\n')
         if username == 'bar':
-            print(chan.recv(1024))
+            msg = chan.recv(1024)
+            chan.send(msg)
 
         chan.close()
         t.close()
diff --git a/tests/test_app.py b/tests/test_app.py
@@ -136,5 +136,23 @@ class TestApp(AsyncHTTPTestCase):
         ws = yield tornado.websocket.websocket_connect(ws_url)
         msg = yield ws.read_message()
         self.assertIn(b'Welcome!', msg)
+
+        # message will be ignored silently
+        yield ws.write_message('hello')
+        yield ws.write_message('"hello"')
+        yield ws.write_message('[hello]')
+        yield ws.write_message(json.dumps({'resize': []}))
+        yield ws.write_message(json.dumps({'resize': {}}))
+        yield ws.write_message(json.dumps({'resize': [100]}))
+        yield ws.write_message(json.dumps({'resize': [100]*10}))
+        yield ws.write_message(json.dumps({'resize': [-1, -1]}))
+        yield ws.write_message(json.dumps({'data': [1]}))
+        yield ws.write_message(json.dumps({'data': (1,)}))
+        yield ws.write_message(json.dumps({'data': {'a': 2}}))
+        yield ws.write_message(json.dumps({'data': 1}))
+        yield ws.write_message(json.dumps({'data': 2.1}))
+        yield ws.write_message(json.dumps({'key-non-existed': 'hello'}))
         yield ws.write_message(json.dumps({'resize': [79, 23], 'data': 'bye'}))
+        msg = yield ws.read_message()
+        self.assertEqual(b'bye', msg)
         ws.close()
diff --git a/webssh/handler.py b/webssh/handler.py
@@ -2,6 +2,7 @@ import io
 import json
 import logging
 import socket
+import struct
 import threading
 import traceback
 import weakref
@@ -9,6 +10,7 @@ import paramiko
 import tornado.web
 
 from tornado.ioloop import IOLoop
+from tornado.util import basestring_type
 from webssh.worker import Worker, recycle_worker, workers
 
 try:
@@ -200,15 +202,25 @@ class WsockHandler(MixinHandler, tornado.websocket.WebSocketHandler):
     def on_message(self, message):
         logging.debug('{!r} from {}:{}'.format(message, *self.src_addr))
         worker = self.worker_ref()
-        msg = json.loads(message)
+        try:
+            msg = json.loads(message)
+        except ValueError:  # py2
+            return
+        except json.decoder.JSONDecodeError:  # py3
+            return
+
+        if not isinstance(msg, dict):
+            return
+
         resize = msg.get('resize')
         if resize:
             try:
                 worker.chan.resize_pty(*resize)
-            except paramiko.SSHException:
+            except (TypeError, struct.error, paramiko.SSHException):
                 pass
+
         data = msg.get('data')
-        if data:
+        if data and isinstance(data, basestring_type):
             worker.data_to_dst.append(data)
             worker.on_write()