From 150bf1a5b0942b40ea3ea66fdeae510a1ea7cb29 Mon Sep 17 00:00:00 2001
From: pagedown <pg.dn@outlook.com>
Date: Wed, 16 Mar 2022 01:03:51 +0800
Subject: [PATCH] Reduce bootstrap script length by removing comments and
 indents

dropbear has 9000 bytes limit on ssh arguments length.
---
 kittens/ssh/main.py                | 8 ++++++--
 kitty_tests/ssh.py                 | 8 +++++++-
 shell-integration/ssh/bootstrap.sh | 9 ++++++---
 3 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/kittens/ssh/main.py b/kittens/ssh/main.py
index 74c31af8a..ec0938d19 100644
--- a/kittens/ssh/main.py
+++ b/kittens/ssh/main.py
@@ -223,13 +223,17 @@ def safe_remove(x: str) -> None:
         os.remove(x)
 
 
-def prepare_script(ans: str, replacements: Dict[str, str]) -> str:
+def prepare_script(ans: str, replacements: Dict[str, str], script_type: str) -> str:
     for k in ('EXEC_CMD',):
         replacements[k] = replacements.get(k, '')
 
     def sub(m: 're.Match[str]') -> str:
         return replacements[m.group()]
 
+    if script_type == 'sh':
+        # Remove comments and indents. The dropbear SSH server has 9000 bytes limit on ssh arguments length.
+        # Needs to be trimmed before replacing EXEC_CMD to avoid affecting the indentation of user commands.
+        ans = re.sub(r'^\s*#.*\n|^\s*', '', ans, flags=re.MULTILINE)
     return re.sub('|'.join(fr'\b{k}\b' for k in replacements), sub, ans)
 
 
@@ -269,7 +273,7 @@ def bootstrap_script(
     if request_data:
         sd.update(sensitive_data)
     replacements.update(sensitive_data)
-    return prepare_script(ans, sd), replacements, shm
+    return prepare_script(ans, sd, script_type), replacements, shm
 
 
 def get_ssh_cli() -> Tuple[Set[str], Set[str]]:
diff --git a/kitty_tests/ssh.py b/kitty_tests/ssh.py
index e0b4d62be..c542b7f52 100644
--- a/kitty_tests/ssh.py
+++ b/kitty_tests/ssh.py
@@ -10,7 +10,8 @@ from functools import lru_cache
 
 from kittens.ssh.config import load_config
 from kittens.ssh.main import (
-    bootstrap_script, get_connection_data, wrap_bootstrap_script
+    bootstrap_script, get_connection_data, get_remote_command,
+    wrap_bootstrap_script
 )
 from kittens.ssh.options.types import Options as SSHOptions
 from kittens.ssh.options.utils import DELETE_ENV_VAR
@@ -78,6 +79,11 @@ print(' '.join(map(str, buf)))'''), lines=13, cols=77)
         self.ae(parse(conf, '1').env, {'a': 'c', 'b': 'b'})
         self.ae(parse(conf, '2').env, {'a': 'c', 'b': 'b'})
 
+    def test_ssh_bootstrap_sh_cmd_limit(self):
+        rcmd, _, _ = get_remote_command([], SSHOptions({'interpreter': 'sh'}))
+        # dropbear has a 9000 bytes maximum command length limit
+        self.assertLessEqual(sum(len(x) for x in rcmd), 9000)
+
     @property
     @lru_cache()
     def all_possible_sh(self):
diff --git a/shell-integration/ssh/bootstrap.sh b/shell-integration/ssh/bootstrap.sh
index fed6a0f1f..617cc324b 100644
--- a/shell-integration/ssh/bootstrap.sh
+++ b/shell-integration/ssh/bootstrap.sh
@@ -274,7 +274,8 @@ exec_zsh_with_integration() {
         export ZDOTDIR="$shell_integration_dir/zsh"
         exec "$login_shell" "-l"
     fi
-    unset KITTY_ORIG_ZDOTDIR  # ensure this is not propagated
+    # ensure this is not propagated
+    unset KITTY_ORIG_ZDOTDIR
 }
 
 exec_fish_with_integration() {
@@ -313,8 +314,10 @@ exec_with_shell_integration() {
 }
 
 execute_sh_with_posix_env() {
-    [ "$shell_name" = "sh" ] || return  # only for sh as that is likely to be POSIX compliant
-    command "$login_shell" -l -c ":" > /dev/null 2> /dev/null && return  # sh supports -l so use that
+    # only for sh as that is likely to be POSIX compliant
+    [ "$shell_name" = "sh" ] || return
+    # sh supports -l so use that
+    command "$login_shell" -l -c ":" > /dev/null 2> /dev/null && return
     [ -z "$shell_integration_dir" ] && die "Could not read data over tty ssh kitten cannot function"
     sh_dir="$shell_integration_dir/sh"
     command mkdir -p "$sh_dir" || die "Creating $sh_dir failed"