Merge pull request #159 from jumpserver/dev

Dev

coco/models.py

@@ -33,7 +33,7 @@ class Connection(object):
     def __str__(self):
         return '<{} from {}>'.format(self.user, self.addr)
 
-    def new_client(self, tid):
+    def new_client(self, tid, chan=None):
         client = Client(
             tid=tid, user=self.user, addr=self.addr,
             login_from=self.login_from
@@ -107,13 +107,13 @@ class Client(object):
     ```
     """
 
-    def __init__(self, tid=None, user=None, addr=None, login_from=None):
+    def __init__(self, tid=None, user=None, addr=None, login_from=None, chan=None):
         if tid is None:
             tid = str(uuid.uuid4())
         self.id = tid
         self.user = user
         self.addr = addr
-        self.chan = None
+        self.chan = chan
         self.request = Request()
         self.connection_id = None
         self.login_from = login_from
@@ -123,7 +123,11 @@ class Client(object):
         return self.chan.fileno()
 
     def send(self, b):
+        try:
             return self.chan.send(b)
+        except (OSError, EOFError, socket.error) as e:
+            logger.error('Send to client {} error: {}'.format(self, e))
+            return 0
 
     def send_unicode(self, s):
         b = s.encode()
@@ -138,6 +142,7 @@ class Client(object):
 
     def close(self):
         logger.info("Client {} close".format(self))
+        if self.chan:
             self.chan.close()
         return
 

coco/session.py

@@ -166,10 +166,12 @@ class Session:
         self.sel.register(self.server, selectors.EVENT_READ)
         self.sel.register(self.stop_evt, selectors.EVENT_READ)
         self.sel.register(self.client.change_size_evt, selectors.EVENT_READ)
-        if self.client.closed:
-            return
         while not self.is_finished:
             events = self.sel.select(timeout=60)
+            if self.client.closed:
+                break
+            if self.server.closed:
+                break
             for sock in [key.fileobj for key, _ in events]:
                 data = sock.recv(BUF_SIZE)
                 if sock == self.server: