Merge pull request #1 from drunkdream/feature/connect-timeout

Feature/connect timeout

.github/workflows/ssh.yml

@@ -23,4 +23,4 @@ jobs:
           python -m pip install -r requirements.txt
       - name: SSH connection to Actions
         run: |
-            python -m tmate
+            python -m tmate --wait-timeout=600

tmate/__init__.py

@@ -1,3 +1,3 @@
 # -*- coding: utf-8 -*-
 
-VERSION = "0.1.2"
+VERSION = "0.1.3"

tmate/__main__.py

@@ -8,11 +8,12 @@
 from . import client
 from . import utils
 
-async def run(host, port):
+
+async def run(host, port, wait_timeout=None):
     cli = client.TMateClient(host, port)
     await cli.connect()
     await cli.handshake()
-    await cli.serve()
+    await cli.serve(wait_timeout=wait_timeout)
 
 
 def main():
@@ -22,14 +23,17 @@ def main():
     parser.add_argument("-f", help="set the config file path")
     parser.add_argument("--host", help="tmate server host", default="ssh.tmate.io")
     parser.add_argument("--port", help="tmate server port", type=int, default=22)
+    parser.add_argument(
+        "--wait-timeout", help="timeout wait for client", type=float, default=0
+    )
 
     args = parser.parse_args()
 
     if sys.platform == "win32":
         utils.enable_ansi_code()
         loop = asyncio.ProactorEventLoop()
         asyncio.set_event_loop(loop)
-    
+
     handler = logging.StreamHandler()
     formatter = logging.Formatter("%(message)s")
     handler.setFormatter(formatter)
@@ -39,8 +43,9 @@ def main():
     utils.logger.addHandler(handler)
 
     loop = asyncio.get_event_loop()
-    loop.run_until_complete(run(args.host, args.port))
+    loop.run_until_complete(run(args.host, args.port, args.wait_timeout))
     loop.close()
 
+
 if __name__ == "__main__":
     sys.exit(main())

tmate/client.py

@@ -24,7 +24,12 @@
     ("#", "list-buffers"),
     ("$", "command-prompt", "-I#S", "rename-session '%%'"),
     ("%", "split-window", "-h"),
-    ("&", "confirm-before", "-pkill-window #W? (y/n)", "kill-window",),
+    (
+        "&",
+        "confirm-before",
+        "-pkill-window #W? (y/n)",
+        "kill-window",
+    ),
     ("'", "command-prompt", "-pindex", "select-window -t ':%%'"),
     ("(", "switch-client", "-p"),
     (")", "switch-client", "-n"),
@@ -90,7 +95,14 @@
     ("-r", "C-Down", "resize-pane", "-D"),
     ("-r", "C-Left", "resize-pane", "-L"),
     ("-r", "C-Right", "resize-pane", "-R"),
-    ("-n", "MouseDown1Pane", "select-pane", "-t=;", "send-keys", "-M",),
+    (
+        "-n",
+        "MouseDown1Pane",
+        "select-pane",
+        "-t=;",
+        "send-keys",
+        "-M",
+    ),
     ("-n", "MouseDrag1Border", "resize-pane", "-M"),
     ("-n", "MouseDown1Status", "select-window", "-t="),
     ("-n", "WheelDownStatus", "next-window"),
@@ -206,11 +218,20 @@ async def handshake(self):
         self._channel_writer.write(buffer)
         self.update_layout(self._size)
 
-    async def serve(self, timeout=None):
+    async def serve(self, wait_timeout=None, timeout=None):
         time0 = time.time()
         while timeout is None or time.time() - time0 < timeout:
+            if not self._shell and wait_timeout and time.time() - time0 >= wait_timeout:
+                utils.logger.warn(
+                    "[%s] Wait for client timeout" % self.__class__.__name__
+                )
+                return
             try:
-                buffer = await self._channel_reader.read(4096)
+                buffer = await asyncio.wait_for(
+                    self._channel_reader.read(4096), timeout=1
+                )
+            except asyncio.TimeoutError:
+                continue
             except asyncssh.ConnectionLost:
                 break
             unpacker = msgpack.Unpacker()