From f7a7e39a36a7daddde87eb09ae1d38a8819e0d24 Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Wed, 20 Jul 2022 21:15:51 +0530
Subject: [PATCH] prewarm wrapper should indicate it is active via the process
 environment

---
 kitty_tests/prewarm.py | 14 +++++++++++---
 launcher.c             |  6 +++---
 prewarm-launcher.c     |  8 +++++++-
 3 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/kitty_tests/prewarm.py b/kitty_tests/prewarm.py
index dc211920f..416336b26 100644
--- a/kitty_tests/prewarm.py
+++ b/kitty_tests/prewarm.py
@@ -107,12 +107,20 @@ def socket_child_main(exit_code=0, initial_print=''):
                 pty.wait_till(lambda: 'SIGTSTP received' in pty.screen_contents())
                 wait_for_death(19)
 
-        with self.subTest(msg='test SIGWINCH handling'):
+            with self.subTest(msg='test SIGWINCH handling'):
+                pty = self.create_pty(
+                    argv=[kitty_exe(), '+runpy', src + 'socket_child_main(initial_print="child ready:")'], cols=cols, env=env, cwd=cwd)
+                pty.wait_till(lambda: 'child ready:' in pty.screen_contents())
+                pty.set_window_size(columns=cols + 3)
+                pty.wait_till(lambda: f'Screen size changed: {cols + 3}' in pty.screen_contents())
+                os.close(pty.master_fd)
+
+        with self.subTest(msg='test env rewrite'):
             pty = self.create_pty(
                 argv=[kitty_exe(), '+runpy', src + 'socket_child_main(initial_print="child ready:")'], cols=cols, env=env, cwd=cwd)
             pty.wait_till(lambda: 'child ready:' in pty.screen_contents())
-            pty.set_window_size(columns=cols + 3)
-            pty.wait_till(lambda: f'Screen size changed: {cols + 3}' in pty.screen_contents())
+            from kitty.child import environ_of_process
+            self.assertTrue(environ_of_process(pty.child_pid).get('KITTY_PWPR'))
             os.close(pty.master_fd)
 
         with self.subTest(msg='test passing of data via cwd, env vars and stdin/stdout redirection'):
diff --git a/launcher.c b/launcher.c
index 78df050ea..b24f5a73f 100644
--- a/launcher.c
+++ b/launcher.c
@@ -268,11 +268,11 @@ read_exe_path(char *exe, size_t buf_sz) {
 }
 #endif // }}}
 
-extern void use_prewarmed_process(int argc, char *argv[]);
+extern void use_prewarmed_process(int argc, char *argv[], char *envp[]);
 
-int main(int argc, char *argv[]) {
+int main(int argc, char *argv[], char* envp[]) {
     if (argc < 1 || !argv) { fprintf(stderr, "Invalid argc/argv\n"); return 1; }
-    use_prewarmed_process(argc, argv);
+    use_prewarmed_process(argc, argv, envp);
     char exe[PATH_MAX+1] = {0};
     char exe_dir_buf[PATH_MAX+1] = {0};
     FREE_AFTER_FUNCTION const char *lc_ctype = NULL;
diff --git a/prewarm-launcher.c b/prewarm-launcher.c
index f334d3d70..396977812 100644
--- a/prewarm-launcher.c
+++ b/prewarm-launcher.c
@@ -673,7 +673,7 @@ check_socket_addr(char *addr) {
 }
 
 void
-use_prewarmed_process(int argc, char *argv[]) {
+use_prewarmed_process(int argc, char *argv[], char *envp[]) {
     char *env_addr = getenv("KITTY_PREWARM_SOCKET");
     if (!env_addr || !*env_addr || !is_prewarmable(argc, argv)) return;
     env_addr = check_socket_addr(env_addr);
@@ -698,6 +698,12 @@ use_prewarmed_process(int argc, char *argv[]) {
     to_child_tty.buf = from_child_tty.buf + IO_BUZ_SZ;
 #undef fail
 
+    while (*envp) {
+        if (strncmp(*envp, "KITTY_PREWARM_SOCKET=", 2) == 0) {
+            sprintf(*envp,  "KITTY_PWPR=%s", child_tty_name);
+        }
+        envp++;
+    }
     loop();
     flush_data();
     cleanup();