From 30e635a934f9eeb5fa72bc1df592a7d4d78c39e1 Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Sat, 5 Mar 2022 10:31:00 +0530
Subject: [PATCH] Use zsh's builtin functions to do I/O with TTY

Far more robust than the system tools
---
 shell-integration/ssh/bootstrap.sh | 68 ++++++++++++++++++++++--------
 1 file changed, 51 insertions(+), 17 deletions(-)

diff --git a/shell-integration/ssh/bootstrap.sh b/shell-integration/ssh/bootstrap.sh
index 670a1af4f..affce018e 100644
--- a/shell-integration/ssh/bootstrap.sh
+++ b/shell-integration/ssh/bootstrap.sh
@@ -2,14 +2,61 @@
 # Copyright (C) 2022 Kovid Goyal <kovid at kovidgoyal.net>
 # Distributed under terms of the GPLv3 license.
 
-# read the transmitted data from STDIN
 cleanup_on_bootstrap_exit() {
     [ -n "$saved_tty_settings" ] && command stty "$saved_tty_settings" 2> /dev/null < /dev/tty
     saved_tty_settings=""
 }
 
+# try to use zsh's builtin functions for reading/writing to TTY fd as they are superior to the POSIX variants
+tty_fd=-1
+if builtin zmodload zsh/system 2> /dev/null; then
+    builtin sysopen -o cloexec -rwu tty_fd -- $TTY 2> /dev/null
+    [ $tty_fd = -1 ] && builtin sysopen -o cloexec -rwu tty_fd -- /dev/tty 2> /dev/null
+fi
+if [ $tty_fd -gt -1 ]; then
+    dcs_to_kitty() {
+        builtin local b64data
+        b64data=$(builtin printf "%s" "$2" | builtin command base64) 
+        builtin print -nu "$tty_fd" '\eP@kitty-'"${1}|${b64data//[[:space:]]}"'\e\\'
+    }
+    read_one_byte_from_tty() {
+        builtin sysread -s "1" -i "$tty_fd" n 2> /dev/null
+        return $?
+    }
+    read_n_bytes_from_tty() {
+        builtin let num_left=$1
+        while [ $num_left -gt 0 ]; do
+            builtin sysread -c num_read -s "$num_left" -i "$tty_fd" -o "1" 2> /dev/null || die "Failed to read $num_left bytes from TTY using sysread"
+            builtin let num_left=$num_left-$num_read
+        done
+    }
+else
+    dcs_to_kitty() { printf "\033P@kitty-$1|%s\033\\" "$(printf "%s" "$2" | command base64 | command tr -d \\n)" > /dev/tty; }
+
+    read_one_byte_from_tty() {
+        # We need a way to read a single byte at a time and to read a specified number of bytes in one invocation.
+        # The options are head -c, read -N and dd
+        #
+        # read -N is not in POSIX and dash/posh dont implement it. Also bash seems to read beyond
+        # the specified number of bytes into an internal buffer.
+        #
+        # head -c reads beyond the specified number of bytes into an internal buffer on macOS
+        #
+        # POSIX dd works for one byte at a time but for reading X bytes it needs the GNU iflag=count_bytes
+        # extension, and is anyway unsafe as it can lead to corrupt output when the read syscall is interrupted.
+        n=$(command dd bs=1 count=1 2> /dev/null < /dev/tty) 
+        return $?
+    }
+
+    read_n_bytes_from_tty() {
+        # 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 "$1" < /dev/tty 
+    }
+fi
+
 die() { printf "\033[31m%s\033[m\n\r" "$*" > /dev/stderr; cleanup_on_bootstrap_exit; exit 1; }
-dcs_to_kitty() { printf "\033P@kitty-$1|%s\033\\" "$(printf "%s" "$2" | base64 | tr -d \\n)" > /dev/tty; }
 debug() { dcs_to_kitty "print" "debug $1"; }
 echo_via_kitty() { dcs_to_kitty "echo" "$1"; }
 saved_tty_settings=$(command stty -g 2> /dev/null < /dev/tty)
@@ -63,10 +110,7 @@ untar_and_read_env() {
 
     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 "$1" < /dev/tty | command base64 -d | command tar xjf - --no-same-owner -C "$tdir";
+    read_n_bytes_from_tty "$1" | command base64 -d | command tar xjf - --no-same-owner -C "$tdir";
     data_file="$tdir/data.sh";
     [ -f "$data_file" ] && . "$data_file";
     data_dir="$HOME/$KITTY_SSH_KITTEN_DATA_DIR"
@@ -78,19 +122,9 @@ untar_and_read_env() {
 }
 
 read_record() {
-    # We need a way to read a single byte at a time and to read a specified number of bytes in one invocation.
-    # The options are head -c, read -N and dd
-    #
-    # read -N is not in POSIX and dash/posh dont implement it. Also bash seems to read beyond
-    # the specified number of bytes into an internal buffer.
-    #
-    # head -c reads beyond the specified number of bytes into an internal buffer on macOS
-    #
-    # POSIX dd works for one byte at a time but for reading X bytes it needs the GNU iflag=count_bytes
-    # extension, and is anyway unsafe as it can lead to corrupt output when the read syscall is interrupted.
     record=""
     while :; do
-        n=$(command dd bs=1 count=1 2> /dev/null < /dev/tty) 
+        read_one_byte_from_tty || die "Reading a byte from the TTY failed"
         [ "$n" = "$record_separator" ] && break
         record="$record$n"
     done