diff --git a/kittens/ssh/main.py b/kittens/ssh/main.py index 97019434f..8efb52230 100644 --- a/kittens/ssh/main.py +++ b/kittens/ssh/main.py @@ -17,18 +17,20 @@ import tempfile import time import traceback from base64 import standard_b64decode, standard_b64encode -from contextlib import suppress, contextmanager +from contextlib import contextmanager, suppress from getpass import getuser from typing import ( - Any, Callable, Dict, Iterator, List, NoReturn, Optional, Sequence, Set, - Tuple, Union + Callable, Dict, Iterator, List, NoReturn, Optional, Sequence, Set, Tuple, + Union ) from kitty.constants import ( runtime_dir, shell_integration_dir, ssh_control_master_template, terminfo_dir ) +from kitty.options.types import Options from kitty.shm import SharedMemory +from kitty.types import run_once from kitty.utils import SSHConnectionData from .completion import complete, ssh_options @@ -66,6 +68,12 @@ def serialize_env(env: Dict[str, str], base_env: Dict[str, str]) -> bytes: return '\n'.join(lines).encode('utf-8') +@run_once +def kitty_opts() -> Options: + from kitty.cli import create_default_opts + return create_default_opts() + + def make_tarfile(ssh_opts: SSHOptions, base_env: Dict[str, str], compression: str = 'gz') -> bytes: def normalize_tarinfo(tarinfo: tarfile.TarInfo) -> tarfile.TarInfo: @@ -98,15 +106,13 @@ def make_tarfile(ssh_opts: SSHOptions, base_env: Dict[str, str], compression: st from kitty.shell_integration import get_effective_ksi_env_var if ssh_opts.shell_integration == 'inherited': - from kitty.cli import create_default_opts - ksi = get_effective_ksi_env_var(create_default_opts()) + ksi = get_effective_ksi_env_var(kitty_opts()) else: - from kitty.options.types import Options from kitty.options.utils import shell_integration ksi = get_effective_ksi_env_var(Options({'shell_integration': shell_integration(ssh_opts.shell_integration)})) env = { - 'TERM': os.environ['TERM'], + 'TERM': os.environ.get('TERM') or kitty_opts().term, 'COLORTERM': 'truecolor', } for q in ('KITTY_WINDOW_ID', 'WINDOWID'): @@ -139,12 +145,7 @@ def make_tarfile(ssh_opts: SSHOptions, base_env: Dict[str, str], compression: st def get_ssh_data(msg: str, request_id: str) -> Iterator[bytes]: - record_sep = b'\036' - - def fmt_prefix(msg: Any) -> bytes: - return str(msg).encode('ascii') + record_sep - - yield record_sep # to discard leading data + yield b'\nKITTY_DATA_START\n' # to discard leading data try: msg = standard_b64decode(msg).decode('utf-8') md = dict(x.split('=', 1) for x in msg.split(':')) @@ -153,7 +154,7 @@ def get_ssh_data(msg: str, request_id: str) -> Iterator[bytes]: rq_id = md['id'] except Exception: traceback.print_exc() - yield fmt_prefix('!invalid ssh data request message') + yield b'invalid ssh data request message\n' else: try: with SharedMemory(pwfilename, readonly=True) as shm: @@ -170,13 +171,21 @@ def get_ssh_data(msg: str, request_id: str) -> Iterator[bytes]: raise ValueError('Incorrect request id') except Exception as e: traceback.print_exc() - yield fmt_prefix(f'!{e}') + yield f'{e}\n'.encode('utf-8') else: + yield b'OK\n' ssh_opts = SSHOptions(env_data['opts']) ssh_opts.copy = {k: CopyInstruction(*v) for k, v in ssh_opts.copy.items()} - encoded_data = env_data['tarfile'].encode('ascii') - yield fmt_prefix(len(encoded_data)) - yield encoded_data + encoded_data = memoryview(env_data['tarfile'].encode('ascii')) + # macOS has a 255 byte limit on its input queue as per man stty. + # Not clear if that applies to canonical mode input as well, but + # better to be safe. + line_sz = 254 + while encoded_data: + yield encoded_data[:line_sz] + yield b'\n' + encoded_data = encoded_data[line_sz:] + yield b'KITTY_DATA_END\n' def safe_remove(x: str) -> None: diff --git a/shell-integration/ssh/bootstrap.py b/shell-integration/ssh/bootstrap.py index dc23f15e8..4f728387a 100644 --- a/shell-integration/ssh/bootstrap.py +++ b/shell-integration/ssh/bootstrap.py @@ -8,14 +8,12 @@ import io import os import pwd import re -import select import shutil import subprocess import sys import tarfile import tempfile import termios -import tty tty_fd = -1 original_termios_state = None @@ -126,26 +124,32 @@ def compile_terminfo(base): os.symlink('../x/xterm-kitty', q) +def iter_base64_data(f): + global leading_data + started = 0 + for line in f: + line = line.rstrip() + if started == 0: + if line == b'KITTY_DATA_START': + started = 1 + else: + leading_data += line + elif started == 1: + if line == b'OK': + started = 2 + else: + raise SystemExit(line.decode('utf-8', 'replace').rstrip()) + else: + if line == b'KITTY_DATA_END': + break + yield line + + def get_data(): global data_dir, shell_integration_dir, leading_data - data = b'' - - while data.count(b'\036') < 2: - select.select([tty_fd], [], []) - n = os.read(tty_fd, 64) - if not n: - raise SystemExit('Unexpected EOF while reading data from terminal') - data += n - leading_data, size, data = data.split(b'\036', 2) - if size.startswith(b'!'): - raise SystemExit(size[1:].decode('utf-8', 'replace')) - size = int(size) - while len(data) < size: - select.select([tty_fd], [], []) - n = os.read(tty_fd, size - len(data)) - if not n: - raise SystemExit('Unexpected EOF while reading data from terminal') - data += n + data = [] + with open(tty_fd, 'rb', closefd=False) as f: + data = b''.join(iter_base64_data(f)) cleanup() if leading_data: # clear current line as it might have things echoed on it from leading_data @@ -212,7 +216,7 @@ def exec_with_shell_integration(): def main(): global tty_fd, original_termios_state, login_shell try: - tty_fd = os.open(os.ctermid(), os.O_RDWR | os.O_NONBLOCK | os.O_CLOEXEC) + tty_fd = os.open(os.ctermid(), os.O_RDWR | os.O_CLOEXEC) except OSError: pass else: @@ -222,11 +226,8 @@ def main(): except OSError: pass else: - tty.setraw(tty_fd, termios.TCSANOW) new_state = termios.tcgetattr(tty_fd) new_state[3] &= ~termios.ECHO - new_state[-1][termios.VMIN] = 1 - new_state[-1][termios.VTIME] = 0 termios.tcsetattr(tty_fd, termios.TCSANOW, new_state) try: if original_termios_state is not None: diff --git a/shell-integration/ssh/bootstrap.sh b/shell-integration/ssh/bootstrap.sh index 2ace11432..39f43f75d 100644 --- a/shell-integration/ssh/bootstrap.sh +++ b/shell-integration/ssh/bootstrap.sh @@ -64,103 +64,13 @@ init_tty() { [ -n "$saved_tty_settings" ] && tty_ok="y" if [ "$tty_ok" = "y" ]; then - command stty raw min 1 time 0 -echo 2> /dev/null < /dev/tty || die "stty failed to set raw mode" + command stty -echo 2> /dev/null < /dev/tty || die "stty failed to set raw mode" return 0 fi return 1 } -# try to use zsh's builtin sysread function for reading to TTY -# as it is superior to the POSIX variants. The builtin read function doesn't work -# as it hangs reading N bytes on macOS -tty_fd=-1 -if [ -n "$ZSH_VERSION" ] && 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" | base64_encode) - 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\134" "$(printf "%s" "$2" | base64_encode)" > /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 $? - } - - # 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. - if [ "$(printf "%s" "test" | command ghead -c 3 2> /dev/null)" = "tes" ]; then - # Some BSD based systems use ghead for GNU head which is strictly superior to - # Broken System Design head, so prefer it. - read_n_bytes_from_tty() { - command ghead -c "$1" < /dev/tty - } - elif [ "$(printf "%s" "test" | command head -c 3 2> /dev/null)" = "tes" ]; then - read_n_bytes_from_tty() { - command head -c "$1" < /dev/tty - } - elif detect_python; then - read_n_bytes_from_tty() { - command "$python" "-c" " -import sys, os, errno -def eintr_retry(func, *args): - while True: - try: - return func(*args) - except EnvironmentError as e: - if e.errno != errno.EINTR: - raise -n = $1 -in_fd = sys.stdin.fileno() -out_fd = sys.stdout.fileno() -while n > 0: - d = memoryview(eintr_retry(os.read, in_fd, n)) - n -= len(d) - while d: - nw = eintr_retry(os.write, out_fd, d) - d = d[nw:] -" < /dev/tty - } - elif detect_perl; then - read_n_bytes_from_tty() { - command "$perl" -MList::Util=min -e ' -open(my $fh,"<","/dev/tty");binmode($fh);binmode(STDOUT);my ($n,$buf)=(@ARGV[0],""); -while($n>0){my $rv=sysread($fh,$buf,min(65536,$n));unless($rv){exit(1);};$n-=$rv;print STDOUT $buf;}' "$1" 2> /dev/null - } - else - read_n_bytes_from_tty() { - command dd bs=1 count="$1" 2> /dev/null < /dev/tty - } - fi -fi - +dcs_to_kitty() { printf "\033P@kitty-$1|%s\033\134" "$(printf "%s" "$2" | base64_encode)" > /dev/tty; } debug() { dcs_to_kitty "print" "debug: $1"; } echo_via_kitty() { dcs_to_kitty "echo" "$1"; } @@ -212,13 +122,19 @@ compile_terminfo() { fi } +read_base64_from_tty() { + while IFS= read -r line; do + [ "$line" = "KITTY_DATA_END" ] && return 0 + printf "%s" "$line" + done +} + 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=$(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 "xpzf" "-" "-C" "$tdir" + read_base64_from_tty | base64_decode | command tar "xpzf" "-" "-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" @@ -245,14 +161,20 @@ read_record() { } get_data() { - leading_data=$(read_record) - size=$(read_record) - case "$size" in - ("!"*) - die "$size" - ;; - esac - untar_and_read_env "$size" + started="n" + while IFS= read -r line; do + if [ "$started" = "y" ]; then + [ "$line" = "OK" ] && break + die "$line" + else + if [ "$line" = "KITTY_DATA_START" ]; then + started="y" + else + leading_data="$leading_data$line" + fi + fi + done + untar_and_read_env } if [ "$tty_ok" = "y" ]; then