ssh kitten: Start work on connection sharing

Basic sharing works. Now investigate if we can eliminate the round-trip
latency by transmitting the data without waiting for the start message
when using a shared connection

kittens/ssh/main.py

@@ -153,7 +153,7 @@ def get_ssh_data(msg: str, request_id: str) -> Iterator[bytes]:
         yield fmt_prefix('!invalid ssh data request message')
     else:
         try:
-            with open(os.path.join(cache_dir(), pwfilename), 'rb') as f:
+            with open(os.path.join(cache_dir(), 'ssh', pwfilename), 'rb') as f:
                 os.unlink(f.name)
                 env_data = json.load(f)
                 if pw != env_data['pw']:
@@ -216,7 +216,9 @@ def bootstrap_script(
     with open(os.path.join(shell_integration_dir, 'ssh', f'bootstrap.{script_type}')) as f:
         ans = f.read()
     pw = uuid4()
-    with tempfile.NamedTemporaryFile(prefix='ssh-kitten-pw-', suffix='.json', dir=cache_dir(), delete=False) as tf:
+    ddir = os.path.join(cache_dir(), 'ssh')
+    os.makedirs(ddir, exist_ok=True)
+    with tempfile.NamedTemporaryFile(prefix='ssh-kitten-pw-', suffix='.json', dir=ddir, delete=False) as tf:
         data = {'pw': pw, 'env': dict(os.environ), 'opts': ssh_opts_dict, 'cli_hostname': cli_hostname, 'cli_uname': cli_uname}
         tf.write(json.dumps(data).encode('utf-8'))
     atexit.register(safe_remove, tf.name)
@@ -333,7 +335,7 @@ class InvalidSSHArgs(ValueError):
 
 def parse_ssh_args(args: List[str], extra_args: Tuple[str, ...] = ()) -> Tuple[List[str], List[str], bool, Tuple[str, ...]]:
     boolean_ssh_args, other_ssh_args = get_ssh_cli()
-    passthrough_args = {f'-{x}' for x in 'Nnf'}
+    passthrough_args = {f'-{x}' for x in 'NnfG'}
     ssh_args = []
     server_args: List[str] = []
     expecting_option_val = False
@@ -429,6 +431,16 @@ def get_remote_command(
     return wrap_bootstrap_script(sh_script, interpreter)
 
 
+def connection_sharing_args(opts: SSHOptions, kitty_pid: int) -> List[str]:
+    cp = os.path.join(cache_dir(), 'ssh', f'{kitty_pid}-master-%C')
+    ans: List[str] = [
+        '-o', 'ControlMaster=auto',
+        '-o', f'ControlPath={cp}',
+        '-o', 'ControlPersist=yes',
+    ]
+    return ans
+
+
 def main(args: List[str]) -> NoReturn:
     args = args[1:]
     if args and args[0] == 'use-python':
@@ -446,6 +458,7 @@ def main(args: List[str]) -> NoReturn:
         hostname, remote_args = server_args[0], server_args[1:]
         if not remote_args:
             cmd.append('-t')
+        insertion_point = len(cmd)
         cmd.append('--')
         cmd.append(hostname)
         uname = getuser()
@@ -468,7 +481,11 @@ def main(args: List[str]) -> NoReturn:
             overrides.insert(0, f'hostname {uname}@{hostname_for_match}')
         so = init_config(overrides)
         sod = {k: v._asdict() for k, v in so.items()}
-        cmd += get_remote_command(remote_args, hostname, hostname_for_match, uname, options_for_host(hostname_for_match, uname, so).interpreter, sod)
+        host_opts = options_for_host(hostname_for_match, uname, so)
+        use_control_master = 'KITTY_PID' in os.environ and host_opts.share_connections
+        cmd += get_remote_command(remote_args, hostname, hostname_for_match, uname, host_opts.interpreter, sod)
+        if use_control_master:
+            cmd[insertion_point:insertion_point] = connection_sharing_args(host_opts, int(os.environ['KITTY_PID']))
     import subprocess
     with suppress(FileNotFoundError):
         try:

kittens/ssh/options/definition.py

@@ -79,6 +79,12 @@ value are expanded. The default is empty so no changing is done, which
 usually means the home directory is used.
 ''')
 
+opt('share_connections', 'y', option_type='to_bool', long_text='''
+Within a single kitty instance, all connections to a particular server can be
+shared. This reduces startup latency for subsequent connections and means that you have
+to enter the password only once. Under the hood, it uses SSH ControlMasters and
+these are automatically cleaned up by kitty when it quits.
+''')
 
 opt('interpreter', 'sh', long_text='''
 The interpreter to use on the remote host. Must be either a POSIX complaint shell

kittens/ssh/options/parse.py

@@ -2,7 +2,7 @@
 
 import typing
 from kittens.ssh.options.utils import copy, env, hostname, relative_dir
-from kitty.conf.utils import merge_dicts
+from kitty.conf.utils import merge_dicts, to_bool
 
 
 class Parser:
@@ -30,6 +30,9 @@ class Parser:
     def remote_dir(self, val: str, ans: typing.Dict[str, typing.Any]) -> None:
         ans['remote_dir'] = relative_dir(val)
 
+    def share_connections(self, val: str, ans: typing.Dict[str, typing.Any]) -> None:
+        ans['share_connections'] = to_bool(val)
+
     def shell_integration(self, val: str, ans: typing.Dict[str, typing.Any]) -> None:
         ans['shell_integration'] = str(val)
 

kittens/ssh/options/types.py

@@ -12,6 +12,7 @@ option_names = (  # {{{
  'interpreter',
  'login_shell',
  'remote_dir',
+ 'share_connections',
  'shell_integration')  # }}}
 
 
@@ -21,6 +22,7 @@ class Options:
     interpreter: str = 'sh'
     login_shell: str = ''
     remote_dir: str = '.local/share/kitty-ssh-kitten'
+    share_connections: bool = True
     shell_integration: str = 'inherited'
     copy: typing.Dict[str, kittens.ssh.copy.CopyInstruction] = {}
     env: typing.Dict[str, str] = {}

kitty/main.py

@@ -16,7 +16,7 @@ from .cli_stub import CLIOptions
 from .conf.utils import BadLine
 from .config import cached_values_for
 from .constants import (
-    appname, beam_cursor_data_file, config_dir, glfw_path, is_macos,
+    appname, beam_cursor_data_file, cache_dir, config_dir, glfw_path, is_macos,
     is_wayland, kitty_exe, logo_png_file, running_in_kitty
 )
 from .fast_data_types import (
@@ -345,6 +345,19 @@ def set_locale() -> None:
             log_error('Failed to set locale with no LANG')
 
 
+def cleanup_ssh_control_masters() -> None:
+    import glob
+    import subprocess
+    try:
+        files = glob.glob(os.path.join(cache_dir(), 'ssh', f'{os.getpid()}-master-*'))
+    except OSError:
+        return
+    for x in files:
+        subprocess.run(['ssh', '-o', f'ControlPath={x}', '-O', 'exit', 'kitty-unused-host-name'])
+        with suppress(OSError):
+            os.remove(x)
+
+
 def _main() -> None:
     running_in_kitty(True)
     with suppress(AttributeError):  # python compiled without threading
@@ -408,6 +421,7 @@ def _main() -> None:
             run_app(opts, cli_opts, bad_lines)
     finally:
         glfw_terminate()
+        cleanup_ssh_control_masters()
 
 
 def main() -> None: