webssh

commit b88a159fbb3bdebf087f732672962b41430d2610
parent e8924f1a5a3101d09e29379ab487cb3885e5c96c
Author: Sheng <webmaster0115@gmail.com>
Date:   Sat, 14 Apr 2018 16:54:20 +0800

Added a parameter for worker to close with a reason

Diffstat:
M
handler.py
 | 
8
+++++---
M
worker.py
 | 
20
+++++++++++---------

2 files changed, 16 insertions(+), 12 deletions(-)
diff --git a/handler.py b/handler.py
@@ -126,7 +126,7 @@ class IndexHandler(MixinHandler, tornado.web.RequestHandler):
         except socket.error:
             raise ValueError('Unable to connect to {}:{}'.format(*dst_addr))
         except paramiko.BadAuthenticationType:
-            raise ValueError('Authentication failed.')
+            raise ValueError('SSH authentication failed.')
         except paramiko.BadHostKeyException:
             raise ValueError('Bad host key.')
 
@@ -190,7 +190,7 @@ class WsockHandler(MixinHandler, tornado.websocket.WebSocketHandler):
             self.worker_ref = weakref.ref(worker)
             self.loop.add_handler(worker.fd, worker, IOLoop.READ)
         else:
-            self.close()
+            self.close(reason='Websocket authentication failed.')
 
     def on_message(self, message):
         logging.debug('{!r} from {}:{}'.format(message, *self.src_addr))
@@ -202,4 +202,6 @@ class WsockHandler(MixinHandler, tornado.websocket.WebSocketHandler):
         logging.info('Disconnected from {}:{}'.format(*self.src_addr))
         worker = self.worker_ref() if self.worker_ref else None
         if worker:
-            worker.close()
+            if self.close_reason is None:
+                self.close_reason = 'client disconnected'
+            worker.close(reason=self.close_reason)
diff --git a/worker.py b/worker.py
@@ -13,9 +13,9 @@ workers = {}
 def recycle_worker(worker):
     if worker.handler:
         return
-    logging.debug('Recycling worker {}'.format(worker.id))
+    logging.warn('Recycling worker {}'.format(worker.id))
     workers.pop(worker.id, None)
-    worker.close()
+    worker.close(reason='worker recycled')
 
 
 class Worker(object):
@@ -36,7 +36,7 @@ class Worker(object):
         if events & IOLoop.WRITE:
             self.on_write()
         if events & IOLoop.ERROR:
-            self.close()
+            self.close(reason='error event occurred')
 
     def set_handler(self, handler):
         if not self.handler:
@@ -54,18 +54,18 @@ class Worker(object):
         except (OSError, IOError) as e:
             logging.error(e)
             if errno_from_exception(e) in _ERRNO_CONNRESET:
-                self.close()
+                self.close(reason='chan error on reading')
         else:
             logging.debug('{!r} from {}:{}'.format(data, *self.dst_addr))
             if not data:
-                self.close()
+                self.close(reason='chan closed')
                 return
 
             logging.debug('{!r} to {}:{}'.format(data, *self.handler.src_addr))
             try:
                 self.handler.write_message(data)
             except tornado.websocket.WebSocketClosedError:
-                self.close()
+                self.close(reason='websocket closed')
 
     def on_write(self):
         logging.debug('worker {} on write'.format(self.id))
@@ -80,7 +80,7 @@ class Worker(object):
         except (OSError, IOError) as e:
             logging.error(e)
             if errno_from_exception(e) in _ERRNO_CONNRESET:
-                self.close()
+                self.close(reason='chan error on writing')
             else:
                 self.update_handler(IOLoop.WRITE)
         else:
@@ -92,8 +92,10 @@ class Worker(object):
             else:
                 self.update_handler(IOLoop.READ)
 
-    def close(self):
-        logging.debug('Closing worker {}'.format(self.id))
+    def close(self, reason=None):
+        logging.info(
+            'Closing worker {} with reason {}'.format(self.id, reason)
+        )
         if self.handler:
             self.loop.remove_handler(self.fd)
             self.handler.close()