From 5e457da30bf02a995974a92f6ef579e68786e386 Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Mon, 7 Mar 2022 20:44:26 +0530
Subject: [PATCH] Only use bz2 compression if bzip2 is available

---
 kittens/ssh/main.py                | 7 ++++---
 shell-integration/ssh/bootstrap.py | 2 +-
 shell-integration/ssh/bootstrap.sh | 8 ++++++--
 3 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/kittens/ssh/main.py b/kittens/ssh/main.py
index 7da9e6d22..2458c98c8 100644
--- a/kittens/ssh/main.py
+++ b/kittens/ssh/main.py
@@ -62,7 +62,7 @@ def serialize_env(env: Dict[str, str], base_env: Dict[str, str]) -> bytes:
     return '\n'.join(lines).encode('utf-8')
 
 
-def make_tarfile(ssh_opts: SSHOptions, base_env: Dict[str, str]) -> bytes:
+def make_tarfile(ssh_opts: SSHOptions, base_env: Dict[str, str], compression: str = 'gz') -> bytes:
 
     def normalize_tarinfo(tarinfo: tarfile.TarInfo) -> tarfile.TarInfo:
         tarinfo.uname = tarinfo.gname = ''
@@ -117,7 +117,7 @@ def make_tarfile(ssh_opts: SSHOptions, base_env: Dict[str, str]) -> bytes:
         env['KITTY_LOGIN_CWD'] = ssh_opts.cwd
     env_script = serialize_env(env, base_env)
     buf = io.BytesIO()
-    with tarfile.open(mode='w:bz2', fileobj=buf, encoding='utf-8') as tf:
+    with tarfile.open(mode=f'w:{compression}', fileobj=buf, encoding='utf-8') as tf:
         rd = ssh_opts.remote_dir.rstrip('/')
         for ci in ssh_opts.copy.values():
             tf.add(ci.local_path, arcname=ci.arcname, filter=filter_from_globs(*ci.exclude_patterns))
@@ -147,6 +147,7 @@ def get_ssh_data(msg: str, request_id: str) -> Iterator[bytes]:
         pwfilename = md['pwfile']
         username = md['user']
         rq_id = md['id']
+        compression = md['compression']
     except Exception:
         traceback.print_exc()
         yield fmt_prefix('!invalid ssh data request message')
@@ -169,7 +170,7 @@ def get_ssh_data(msg: str, request_id: str) -> Iterator[bytes]:
             resolved_ssh_opts = options_for_host(hostname, username, ssh_opts, cli_hostname, cli_uname)
             resolved_ssh_opts.copy = {k: CopyInstruction(*v) for k, v in resolved_ssh_opts.copy.items()}
             try:
-                data = make_tarfile(resolved_ssh_opts, env_data['env'])
+                data = make_tarfile(resolved_ssh_opts, env_data['env'], compression)
             except Exception:
                 traceback.print_exc()
                 yield fmt_prefix('!error while gathering ssh data')
diff --git a/shell-integration/ssh/bootstrap.py b/shell-integration/ssh/bootstrap.py
index f84b6ca84..50f5800f3 100644
--- a/shell-integration/ssh/bootstrap.py
+++ b/shell-integration/ssh/bootstrap.py
@@ -60,7 +60,7 @@ def dcs_to_kitty(type, payload):
 def send_data_request():
     hostname = os.environ.get('HOSTNAME') or os.uname().nodename
     write_all(tty_fd, dcs_to_kitty(
-        'ssh', 'id=REQUEST_ID:hostname={}:pwfile=PASSWORD_FILENAME:user={}:pw=DATA_PASSWORD'.format(hostname, getpass.getuser())))
+        'ssh', 'id=REQUEST_ID:hostname={}:pwfile=PASSWORD_FILENAME:user={}:compression=bz2:pw=DATA_PASSWORD'.format(hostname, getpass.getuser())))
 
 
 def debug(msg):
diff --git a/shell-integration/ssh/bootstrap.sh b/shell-integration/ssh/bootstrap.sh
index f9c26485f..5e5c99029 100644
--- a/shell-integration/ssh/bootstrap.sh
+++ b/shell-integration/ssh/bootstrap.sh
@@ -120,7 +120,9 @@ login_cwd=""
 
 init_tty && trap "cleanup_on_bootstrap_exit" EXIT
 if [ "$tty_ok" = "y" ]; then
-    dcs_to_kitty "ssh" "id="REQUEST_ID":hostname="$hostname":pwfile="PASSWORD_FILENAME":user="$USER":pw="DATA_PASSWORD""
+    compression="gz"
+    command -v "bzip2" > /dev/null 2> /dev/null && compression="bz2"
+    dcs_to_kitty "ssh" "id="REQUEST_ID":hostname="$hostname":pwfile="PASSWORD_FILENAME":user="$USER":compression="$compression":pw="DATA_PASSWORD""
 record_separator=$(printf "\036")
 fi
 
@@ -157,7 +159,9 @@ untar_and_read_env() {
 
     tdir=$(command mktemp -d "$HOME/.kitty-ssh-kitten-untar-XXXXXXXXXXXX")
     [ $? = 0 ] || die "Creating temp directory failed"
-    read_n_bytes_from_tty "$1" | base64_decode | command tar xpjf - -C "$tdir"
+    cflag="j"
+    [ "$compression" = "gz" ] && cflag="z"
+    read_n_bytes_from_tty "$1" | base64_decode | command tar "x${cflag}pf" - -C "$tdir"
     data_file="$tdir/data.sh"
     [ -f "$data_file" ] && . "$data_file"
     [ -z "$KITTY_SSH_KITTEN_DATA_DIR" ] && die "Failed to read SSH data from tty"