From 95da414511808faab0b962905bb47191a05240fe Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Sun, 27 Feb 2022 20:45:19 +0530
Subject: [PATCH] Avoid one record read

---
 kittens/ssh/main.py                |  4 ++--
 shell-integration/ssh/bootstrap.sh | 34 +++++++++++++++++-------------
 2 files changed, 21 insertions(+), 17 deletions(-)

diff --git a/kittens/ssh/main.py b/kittens/ssh/main.py
index 2061ddb89..4fa83a300 100644
--- a/kittens/ssh/main.py
+++ b/kittens/ssh/main.py
@@ -88,11 +88,12 @@ def make_tarfile(ssh_opts: SSHOptions, base_env: Dict[str, str]) -> bytes:
             env[q] = val
     env.update(ssh_opts.env)
     env['KITTY_SHELL_INTEGRATION'] = ksi or DELETE_ENV_VAR
+    env['KITTY_SSH_KITTEN_DATA_DIR'] = ssh_opts.remote_dir
     env_script = serialize_env(env, base_env)
     buf = io.BytesIO()
     with tarfile.open(mode='w:bz2', fileobj=buf, encoding='utf-8') as tf:
         rd = ssh_opts.remote_dir.rstrip('/')
-        add_data_as_file(tf, rd + '/settings/env-vars.sh', env_script)
+        add_data_as_file(tf, 'kitty-ssh-kitten-data.sh', env_script)
         if ksi:
             tf.add(shell_integration_dir, arcname=rd + '/shell-integration', filter=filter_files)
         tf.add(terminfo_dir, arcname='.terminfo', filter=filter_files)
@@ -145,7 +146,6 @@ def get_ssh_data(msg: str, ssh_opts: Optional[Dict[str, SSHOptions]] = None) ->
                 from base64 import standard_b64encode
                 encoded_data = standard_b64encode(data)
                 yield fmt_prefix(len(encoded_data))
-                yield fmt_prefix(resolved_ssh_opts.remote_dir)
                 yield encoded_data
 
 
diff --git a/shell-integration/ssh/bootstrap.sh b/shell-integration/ssh/bootstrap.sh
index 7f29d1101..e8e3d2acc 100644
--- a/shell-integration/ssh/bootstrap.sh
+++ b/shell-integration/ssh/bootstrap.sh
@@ -38,18 +38,28 @@ dsc_to_kitty "ssh" "hostname=$hostname:pwfile=$password_filename:pw=$data_passwo
 size=""
 record_separator=$(printf "\036")
 
-untar() {
+untar_and_read_env() {
     # extract the tar file atomically, in the sense that any file from the 
     # tarfile is only put into place after it has been fully written to disk
-    tdir=$(mktemp -d "$HOME/.kitty-ssh-kitten-untar-XXXXXXXXXXXX")
-    [ $? = 0 ] || die "Creating temp directory failed"
-    command base64 -d | command tar xjf - --no-same-owner -C "$tdir" 
+
+    tdir=$(mktemp -d "$HOME/.kitty-ssh-kitten-untar-XXXXXXXXXXXX");
+    [ $? = 0 ] || die "Creating temp directory failed";
+    # using dd with bs=1 is very slow, so use head. On non GNU coreutils head
+    # does not limit itself to reading -c bytes only from the pipe so we can potentially lose
+    # some trailing data, for instance if the user starts typing. Cant be helped.
+    command head -c "$size" < /dev/tty | command base64 -d | command tar xjf - --no-same-owner -C "$tdir";
+    data_file="$tdir/kitty-ssh-kitten-data.sh";
+    [ -f "$data_file" ] && . "$data_file";
+    data_dir="$HOME/$KITTY_SSH_KITTEN_DATA_DIR"
+    command rm -f "$data_file";
     cwd="$PWD";
     cd "$tdir";
-    find . -type d -exec mkdir -p "${HOME}/{}" ";"
-    find . -type f -exec mv "{}" "${HOME}/{}" ";"
+    command find . -type d -exec mkdir -p "${HOME}/{}" ";"
+    command find . -type f -exec mv "{}" "${HOME}/{}" ";"
     cd "$cwd";
-    rm -rf "$tdir";
+    command rm -rf "$tdir";
+    [ -z "KITTY_SSH_KITTEN_DATA_DIR" ] && die "Failed to read SSH data from tty";
+    unset KITTY_SSH_KITTEN_DATA_DIR;
 }
 
 read_record() {
@@ -80,9 +90,7 @@ get_data() {
             die "$size"
             ;;
     esac
-    data_dir="$HOME/$(read_record)"
-    # using dd with bs=1 is very slow on Linux, so use head 
-    command head -c "$size" < /dev/tty | untar
+    untar_and_read_env "$size"
 }
 
 get_data
@@ -96,8 +104,7 @@ if [ -n "$leading_data" ]; then
 fi
 shell_integration_dir="$data_dir/shell-integration"
 settings_dir="$data_dir/settings"
-env_var_file="$settings_dir/env-vars.sh"
-[ -f "$HOME/.terminfo/kitty.terminfo" -a -f "$env_var_file"  ] || die "Incomplete extraction of ssh data";
+[ -f "$HOME/.terminfo/kitty.terminfo" ] || die "Incomplete extraction of ssh data";
 
 # export TERMINFO
 tname=".terminfo"
@@ -107,9 +114,6 @@ if [ -e "/usr/share/misc/terminfo.cdb" ]; then
 fi
 export TERMINFO="$HOME/$tname"
 
-# setup env vars
-. "$env_var_file"
-
 # compile terminfo for this system
 if [ -x "$(command -v tic)" ]; then
     tic_out=$(command tic -x -o "$HOME/$tname" "$HOME/.terminfo/kitty.terminfo" 2>&1)