From 37c185462a2611da89de35486f21fc8f6c9b0385 Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Sat, 26 Feb 2022 14:57:01 +0530
Subject: [PATCH] Handle error messages form terminal during bootstrap

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

diff --git a/kittens/ssh/main.py b/kittens/ssh/main.py
index d84a2c600..f5c66c8f5 100644
--- a/kittens/ssh/main.py
+++ b/kittens/ssh/main.py
@@ -11,7 +11,9 @@ import tarfile
 import tempfile
 import time
 from contextlib import suppress
-from typing import Dict, Iterator, List, NoReturn, Optional, Set, Tuple, Union
+from typing import (
+    Any, Dict, Iterator, List, NoReturn, Optional, Set, Tuple, Union
+)
 
 from kitty.constants import cache_dir, shell_integration_dir, terminfo_dir
 from kitty.shell_integration import get_effective_ksi_env_var
@@ -55,26 +57,30 @@ def make_tarfile(hostname: str = '', shell_integration_dest: str = DEFAULT_SHELL
 
 
 def get_ssh_data(msg: str, shell_integration_dest: str = DEFAULT_SHELL_INTEGRATION_DEST) -> Iterator[bytes]:
+
+    def fmt_prefix(msg: Any) -> bytes:
+        return f'\036{msg}:'.encode('ascii')
+
     try:
         hostname, pwfilename, pw = msg.split(':', 2)
     except Exception:
-        yield b' invalid ssh data request message\n'
+        yield fmt_prefix('!invalid ssh data request message')
     try:
         with open(os.path.join(cache_dir(), pwfilename)) as f:
             os.unlink(f.name)
             if pw != f.read():
                 raise ValueError('Incorrect password')
     except Exception:
-        yield b' incorrect ssh data password\n'
+        yield fmt_prefix('!incorrect ssh data password')
     else:
         try:
             data = make_tarfile(hostname, shell_integration_dest)
         except Exception:
-            yield b' error while gathering ssh data\n'
+            yield fmt_prefix('!error while gathering ssh data')
         else:
             from base64 import standard_b64encode
             encoded_data = standard_b64encode(data)
-            yield f"\036{len(encoded_data)}:".encode('ascii')
+            yield fmt_prefix(len(encoded_data))
             yield encoded_data
 
 
diff --git a/shell-integration/ssh/bootstrap.sh b/shell-integration/ssh/bootstrap.sh
index 6ea32f08c..44890a0d9 100644
--- a/shell-integration/ssh/bootstrap.sh
+++ b/shell-integration/ssh/bootstrap.sh
@@ -7,7 +7,7 @@ cleanup_on_bootstrap_exit() {
     [ ! -z "$saved_tty_settings" ] && command stty "$saved_tty_settings"
 }
 
-die() { echo "$*" > /dev/stderr; cleanup_on_bootstrap_exit; exit 1; }
+die() { printf "\033[31m%s\033[m\n" "$*" > /dev/stderr; cleanup_on_bootstrap_exit; exit 1; }
 dsc_to_kitty() { printf "\033P@kitty-$1|%s\033\\" "$(printf "%s" "$2" | base64)" > /dev/tty; }
 debug() { dsc_to_kitty "print" "debug $1"; }
 echo_via_kitty() { dsc_to_kitty "echo" "$1"; }
@@ -67,6 +67,11 @@ get_data() {
             fi
         fi
     done
+    case "$size" in
+        ("!"*)
+            die "$size"
+            ;;
+    esac
     # using dd with bs=1 is very slow on Linux, so use head 
     command head -c "$size" | untar
     # command dd bs=1 "count=$size" 2> /dev/null | untar