From ced282c06c925d666f5143f15bf517b1696d1689 Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Sat, 28 Sep 2024 12:31:13 +0530
Subject: [PATCH] Infrastructure to pass arbitrary fds to spawned child

---
 kitty/child.c             | 13 +++++++++++--
 kitty/child.py            | 19 +++++++++++++++++--
 kitty/fast_data_types.pyi |  1 +
 3 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/kitty/child.c b/kitty/child.c
index aa328ee42..492d935fb 100644
--- a/kitty/child.c
+++ b/kitty/child.c
@@ -79,11 +79,11 @@ wait_for_terminal_ready(int fd) {
 
 static PyObject*
 spawn(PyObject *self UNUSED, PyObject *args) {
-    PyObject *argv_p, *env_p, *handled_signals_p;
+    PyObject *argv_p, *env_p, *handled_signals_p, *pass_fds;
     int master, slave, stdin_read_fd, stdin_write_fd, ready_read_fd, ready_write_fd, forward_stdio;
     const char *kitten_exe;
     char *cwd, *exe;
-    if (!PyArg_ParseTuple(args, "ssO!O!iiiiiiO!sp", &exe, &cwd, &PyTuple_Type, &argv_p, &PyTuple_Type, &env_p, &master, &slave, &stdin_read_fd, &stdin_write_fd, &ready_read_fd, &ready_write_fd, &PyTuple_Type, &handled_signals_p, &kitten_exe, &forward_stdio)) return NULL;
+    if (!PyArg_ParseTuple(args, "ssO!O!iiiiiiO!spO!", &exe, &cwd, &PyTuple_Type, &argv_p, &PyTuple_Type, &env_p, &master, &slave, &stdin_read_fd, &stdin_write_fd, &ready_read_fd, &ready_write_fd, &PyTuple_Type, &handled_signals_p, &kitten_exe, &forward_stdio, &PyTuple_Type, &pass_fds)) return NULL;
     char name[2048] = {0};
     if (ttyname_r(slave, name, sizeof(name) - 1) != 0) { PyErr_SetFromErrno(PyExc_OSError); return NULL; }
     char **argv = serialize_string_tuple(argv_p);
@@ -130,6 +130,15 @@ spawn(PyObject *self UNUSED, PyObject *args) {
             safe_close(tfd, __FILE__, __LINE__);
 
             int min_closed_fd = 3;
+            for (Py_ssize_t i = 0; i < PyTuple_GET_SIZE(pass_fds); i++) {
+                PyObject *pfd = PyTuple_GET_ITEM(pass_fds, i);
+                if (!PyLong_Check(pfd)) exit_on_err("pass_fds must contain only integers");
+                int fd = PyLong_AsLong(pfd);
+                if (fd > -1) {
+                    if (fd == min_closed_fd) min_closed_fd++;
+                    else if (safe_dup2(fd, min_closed_fd++) == -1) exit_on_err("dup2() failed for forwarded fd 1");
+                }
+            }
             if (forward_stdio) {
                 if (safe_dup2(STDOUT_FILENO, min_closed_fd++) == -1) exit_on_err("dup2() failed for forwarded fd 1");
                 if (safe_dup2(STDERR_FILENO, min_closed_fd++) == -1) exit_on_err("dup2() failed for forwarded fd 2");
diff --git a/kitty/child.py b/kitty/child.py
index 55dbd5f19..d1894178f 100644
--- a/kitty/child.py
+++ b/kitty/child.py
@@ -7,7 +7,7 @@ from collections import defaultdict
 from collections.abc import Generator, Sequence
 from contextlib import contextmanager, suppress
 from itertools import count
-from typing import TYPE_CHECKING, DefaultDict, Optional
+from typing import TYPE_CHECKING, DefaultDict, Optional, Protocol, Union
 
 import kitty.fast_data_types as fast_data_types
 
@@ -23,6 +23,12 @@ if TYPE_CHECKING:
     from .window import CwdRequest
 
 
+class InheritableFile(Protocol):
+
+    def close(self) -> None: ...
+    def fileno(self) -> int: ...
+
+
 if is_macos:
     from kitty.fast_data_types import cmdline_of_process as cmdline_
     from kitty.fast_data_types import cwd_of_process as _cwd
@@ -211,11 +217,13 @@ class Child:
         is_clone_launch: str = '',
         add_listen_on_env_var: bool = True,
         hold: bool = False,
+        pass_fds: tuple[Union[int, InheritableFile], ...] = (),
     ):
         self.is_clone_launch = is_clone_launch
         self.id = next(child_counter)
         self.add_listen_on_env_var = add_listen_on_env_var
         self.argv = list(argv)
+        self.pass_fds = pass_fds
         if cwd_from:
             try:
                 cwd = cwd_from.modify_argv_for_launch_with_cwd(self.argv, env) or cwd
@@ -331,10 +339,17 @@ class Child:
             argv = cmdline_for_hold(argv)
             final_exe = argv[0]
         env = tuple(f'{k}={v}' for k, v in self.final_env.items())
+        pass_fds = tuple(sorted(x if isinstance(x, int) else x.fileno() for x in self.pass_fds))
         pid = fast_data_types.spawn(
             final_exe, cwd, tuple(argv), env, master, slave, stdin_read_fd, stdin_write_fd,
-            ready_read_fd, ready_write_fd, tuple(handled_signals), kitten_exe(), opts.forward_stdio)
+            ready_read_fd, ready_write_fd, tuple(handled_signals), kitten_exe(), opts.forward_stdio, pass_fds)
         os.close(slave)
+        for x in self.pass_fds:
+            if isinstance(x, int):
+                os.close(x)
+            else:
+                x.close()
+        self.pass_fds = ()
         self.pid = pid
         self.child_fd = master
         if stdin is not None:
diff --git a/kitty/fast_data_types.pyi b/kitty/fast_data_types.pyi
index 0f703b615..6e255118f 100644
--- a/kitty/fast_data_types.pyi
+++ b/kitty/fast_data_types.pyi
@@ -1460,6 +1460,7 @@ def spawn(
     handled_signals: Tuple[int, ...],
     kitten_exe: str,
     forward_stdio: bool,
+    pass_fds: tuple[int, ...],
 ) -> int:
     pass