Preserve fd numbers in pass_fds

Also cleanup serialization of argv/env string arrays for exec

Author: kovidgoyal

kitty/child.c

@@ -16,40 +16,25 @@
 #include <sys/ioctl.h>
 #include <termios.h>
 
-static char**
-serialize_string_tuple(PyObject *src) {
-    Py_ssize_t sz = PyTuple_GET_SIZE(src);
+#define EXTRA_ENV_BUFFER_SIZE 64
 
-    char **ans = calloc(sz + 1, sizeof(char*));
-    if (!ans) fatal("Out of memory");
+static char**
+serialize_string_tuple(PyObject *src, Py_ssize_t extra) {
+    const Py_ssize_t sz = PyTuple_GET_SIZE(src);
+    size_t required_size = sizeof(char*) * (1 + sz + extra);
+    required_size += extra * EXTRA_ENV_BUFFER_SIZE;
+    void *block = calloc(required_size, 1);
+    if (!block) { PyErr_NoMemory(); return NULL; }
+    char **ans = block;
     for (Py_ssize_t i = 0; i < sz; i++) {
-        const char *pysrc = PyUnicode_AsUTF8(PyTuple_GET_ITEM(src, i));
-        if (!pysrc) {
-            PyErr_Clear();
-            RAII_PyObject(u8, PyUnicode_AsEncodedString(PyTuple_GET_ITEM(src, i), "UTF-8", "ignore"));
-            if (!u8) { PyErr_Print(); fatal("couldn't parse command line"); }
-            ans[i] = calloc(PyBytes_GET_SIZE(u8) + 1, sizeof(char));
-            if (ans[i] == NULL) fatal("Out of memory");
-            memcpy(ans[i], PyBytes_AS_STRING(u8), PyBytes_GET_SIZE(u8));
-        } else {
-            size_t len = strlen(pysrc);
-            ans[i] = calloc(len + 1, sizeof(char));
-            if (ans[i] == NULL) fatal("Out of memory");
-            memcpy(ans[i], pysrc, len);
-        }
+        PyObject *x = PyTuple_GET_ITEM(src, i);
+        if (!PyUnicode_Check(x)) { free(block); PyErr_SetString(PyExc_TypeError, "string tuple must have only strings"); return NULL; }
+        ans[i] = (char*)PyUnicode_AsUTF8(x);
+        if (!ans[i]) { free(block); return NULL; }
     }
     return ans;
 }
 
-static void
-free_string_tuple(char** data) {
-    size_t i = 0;
-    while(data[i]) free(data[i++]);
-    free(data);
-}
-
-extern char **environ;
-
 static void
 write_to_stderr(const char *text) {
     size_t sz = strlen(text);
@@ -86,8 +71,10 @@ spawn(PyObject *self UNUSED, PyObject *args) {
     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);
-    char **env = serialize_string_tuple(env_p);
+    char **argv = serialize_string_tuple(argv_p, 0);
+    if (!argv) return NULL;
+    char **env = serialize_string_tuple(env_p, 1);
+    if (!env) { free(argv); return NULL; }
     int handled_signals[16] = {0}, num_handled_signals = MIN((int)arraysz(handled_signals), PyTuple_GET_SIZE(handled_signals_p));
     for (Py_ssize_t i = 0; i < num_handled_signals; i++) handled_signals[i] = PyLong_AsLong(PyTuple_GET_ITEM(handled_signals_p, i));
 
@@ -129,23 +116,29 @@ spawn(PyObject *self UNUSED, PyObject *args) {
             if (ioctl(tfd, TIOCSCTTY, 0) == -1) exit_on_err("Failed to set controlling terminal with TIOCSCTTY");
             safe_close(tfd, __FILE__, __LINE__);
 
-            int min_closed_fd = 3;
+            fd_set passed_fds; FD_ZERO(&passed_fds); bool has_preserved_fds = false;
+            if (forward_stdio) {
+                int fd = safe_dup(STDOUT_FILENO);
+                if (fd < 0) exit_on_err("dup() failed for forwarded STDOUT");
+                FD_SET(fd, &passed_fds);
+                size_t s = PyTuple_GET_SIZE(env_p);
+                env[s] = (char*)(env + (s + 2));
+                snprintf(env[s], EXTRA_ENV_BUFFER_SIZE, "KITTY_STDIO_FORWARDED=%d", fd);
+                fd = safe_dup(STDERR_FILENO);
+                if (fd < 0) exit_on_err("dup() failed for forwarded STDERR");
+                FD_SET(fd, &passed_fds);
+                has_preserved_fds = true;
+            }
+
             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");
-                        safe_close(fd, __FILE__, __LINE__);
-                    }
+                if (fd > -1 && fd < FD_SETSIZE) {
+                    FD_SET(fd, &passed_fds);
+                    has_preserved_fds = true;
                 }
             }
-            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");
-            }
             // Redirect stdin/stdout/stderr to the pty
             if (safe_dup2(slave, STDOUT_FILENO) == -1) exit_on_err("dup2() failed for fd number 1");
             if (safe_dup2(slave, STDERR_FILENO) == -1) exit_on_err("dup2() failed for fd number 2");
@@ -165,10 +158,10 @@ spawn(PyObject *self UNUSED, PyObject *args) {
             safe_close(ready_read_fd, __FILE__, __LINE__);
 
             // Close any extra fds inherited from parent
-            for (int c = min_closed_fd; c < 201; c++) safe_close(c, __FILE__, __LINE__);
+            if (has_preserved_fds) { for (int c = 3; c < 256; c++) { if (!FD_ISSET(c, &passed_fds)) safe_close(c, __FILE__, __LINE__); } }
+            else for (int c = 3; c < 256; c++) { safe_close(c, __FILE__, __LINE__); }
 
-            environ = env;
-            execvp(exe, argv);
+            execvpe(exe, argv, env);
             // Report the failure and exec kitten instead, so that we are not left
             // with a forked but not exec'ed process
             write_to_stderr("Failed to launch child: ");
@@ -196,32 +189,14 @@ spawn(PyObject *self UNUSED, PyObject *args) {
             break;
     }
 #undef exit_on_err
-    free_string_tuple(argv);
-    free_string_tuple(env);
+    free(argv);
+    free(env);
     if (PyErr_Occurred()) return NULL;
     return PyLong_FromLong(pid);
 }
 
-#ifdef __APPLE__
-#include <crt_externs.h>
-#else
-extern char **environ;
-#endif
-
-static PyObject*
-clearenv_py(PyObject *self UNUSED, PyObject *args UNUSED) {
-#ifdef __APPLE__
-    char **e = *_NSGetEnviron();
-    if (e) *e = NULL;
-#else
-    if (environ) *environ = NULL;
-#endif
-    Py_RETURN_NONE;
-}
-
 static PyMethodDef module_methods[] = {
     METHODB(spawn, METH_VARARGS),
-    {"clearenv", clearenv_py, METH_NOARGS, ""},
     {NULL, NULL, 0, NULL}        /* Sentinel */
 };
 

kitty/child.py

@@ -256,6 +256,7 @@ def get_final_env(self) -> dict[str, str]:
             env['KITTY_LISTEN_ON'] = boss.listening_on
         else:
             env.pop('KITTY_LISTEN_ON', None)
+        env.pop('KITTY_STDIO_FORWARDED', None)
         if self.cwd:
             # needed in case cwd is a symlink, in which case shells
             # can use it to display the current directory name rather
@@ -268,8 +269,6 @@ def get_final_env(self) -> dict[str, str]:
         elif opts.terminfo_type == 'direct':
             env['TERMINFO'] = base64_terminfo_data()
         env['KITTY_INSTALLATION_DIR'] = kitty_base_dir
-        if opts.forward_stdio:
-            env['KITTY_STDIO_FORWARDED'] = '3'
         self.unmodified_argv = list(self.argv)
         if not self.should_run_via_run_shell_kitten and 'disabled' not in opts.shell_integration:
             from .shell_integration import modify_shell_environ
@@ -344,11 +343,6 @@ def fork(self) -> Optional[int]:
             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, 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

kitty/fast_data_types.pyi

@@ -1671,7 +1671,6 @@ class SingleKey:
 
 def set_use_os_log(yes: bool) -> None: ...
 def get_docs_ref_map() -> bytes: ...
-def clearenv() -> None: ...
 def set_clipboard_data_types(ct: int, mime_types: Tuple[str, ...]) -> None: ...
 def get_clipboard_mime(ct: int, mime: Optional[str], callback: Callable[[bytes], None]) -> None: ...
 def run_with_activation_token(func: Callable[[str], None]) -> bool: ...

kitty/options/definition.py

@@ -3311,11 +3311,11 @@
 
 
 opt('forward_stdio', 'no', option_type='to_bool', long_text='''
-Forward STDOUT and STDERR of the kitty process to child processes
-as file descriptors 3 and 4. This is useful for debugging as it
+Forward STDOUT and STDERR of the kitty process to child processes.
+This is useful for debugging as it
 allows child processes to print to kitty's STDOUT directly. For example,
-:code:`echo hello world >&3` in a shell will print to the parent kitty's
-STDOUT. When enabled, this also sets the :code:`KITTY_STDIO_FORWARDED=3`
+:code:`echo hello world >&$KITTY_STDIO_FORWARDED` in a shell will print
+to the parent kitty's STDOUT. Sets the :code:`KITTY_STDIO_FORWARDED=fdnum`
 environment variable so child processes know about the forwarding.
 ''')
 

kitty/safe-wrappers.h

@@ -98,6 +98,13 @@ safe_close(int fd, const char* file UNUSED, const int line UNUSED) {
     while(close(fd) != 0 && errno == EINTR);
 }
 
+static inline int
+safe_dup(int a) {
+    int ret;
+    while((ret = dup(a)) < 0 && errno == EINTR);
+    return ret;
+}
+
 static inline int
 safe_dup2(int a, int b) {
     int ret;