Use dict_T to pass env vars to process spawning code

Co-authored-by: Matthieu Coudron <mattator@gmail.com>

7 files changed, 133 insertions, 63 deletions

diff --git a/src/nvim/channel.c b/src/nvim/channel.c
index 37cbfb968b..5628ede2e6 100644
--- a/src/nvim/channel.c
+++ b/src/nvim/channel.c
@@ -304,7 +304,8 @@ Channel *channel_job_start(char **argv, CallbackReader on_stdout,
                            bool pty, bool rpc, bool overlapped, bool detach,
                            const char *cwd,
                            uint16_t pty_width, uint16_t pty_height,
-                           char *term_name, char **env, varnumber_T *status_out)
+                           char *term_name, dict_T *env,
+                           varnumber_T *status_out)
 {
   assert(cwd == NULL || os_isdir_executable(cwd));
 
@@ -358,7 +359,9 @@ Channel *channel_job_start(char **argv, CallbackReader on_stdout,
   if (status) {
     EMSG3(_(e_jobspawn), os_strerror(status), cmd);
     xfree(cmd);
-    os_free_fullenv(proc->env);
+    if (proc->env) {
+      tv_dict_free(proc->env);
+    }
     if (proc->type == kProcessTypePty) {
       xfree(chan->stream.pty.term_name);
     }
@@ -367,8 +370,9 @@ Channel *channel_job_start(char **argv, CallbackReader on_stdout,
     return NULL;
   }
   xfree(cmd);
-  os_free_fullenv(proc->env);
-
+  if (proc->env) {
+    tv_dict_free(proc->env);
+  }
 
   wstream_init(&proc->in, 0);
   if (has_out) {
diff --git a/src/nvim/eval/funcs.c b/src/nvim/eval/funcs.c
index 8235d74cbb..04e539d309 100644
--- a/src/nvim/eval/funcs.c
+++ b/src/nvim/eval/funcs.c
@@ -4887,7 +4887,7 @@ static void f_jobstart(typval_T *argvars, typval_T *rettv, FunPtr fptr)
 
   bool executable = true;
   char **argv = tv_to_argv(&argvars[0], NULL, &executable);
-  char **env = NULL;
+  dict_T *env = NULL;
   if (!argv) {
     rettv->vval.v_number = executable ? 0 : -1;
     return;  // Did error message in tv_to_argv.
@@ -4936,7 +4936,7 @@ static void f_jobstart(typval_T *argvars, typval_T *rettv, FunPtr fptr)
 #endif
 
     char *new_cwd = tv_dict_get_string(job_opts, "cwd", false);
-    if (new_cwd && strlen(new_cwd) > 0) {
+    if (new_cwd && *new_cwd != NUL) {
       cwd = new_cwd;
       // The new cwd must be a directory.
       if (!os_isdir_executable((const char *)cwd)) {
@@ -4945,48 +4945,30 @@ static void f_jobstart(typval_T *argvars, typval_T *rettv, FunPtr fptr)
         return;
       }
     }
-    dictitem_T *job_env = tv_dict_find(job_opts, S_LEN("env"));
-    if (job_env) {
-      if (job_env->di_tv.v_type != VAR_DICT) {
-        EMSG2(_(e_invarg2), "env");
-        shell_free_argv(argv);
-        return;
-      }
-
-      size_t custom_env_size = (size_t)tv_dict_len(job_env->di_tv.vval.v_dict);
-      size_t i = 0;
-      size_t env_size = 0;
-
-      if (clear_env) {
-        // + 1 for last null entry
-        env = xmalloc((custom_env_size + 1) * sizeof(*env));
-        env_size = 0;
-      } else {
-        env_size = os_get_fullenv_size();
 
-        env = xmalloc((custom_env_size + env_size + 1) * sizeof(*env));
-
-        os_copy_fullenv(env, env_size);
-        i = env_size;
-      }
-      assert(env);  // env must be allocated at this point
+    dictitem_T *job_env = tv_dict_find(job_opts, S_LEN("env"));
+    if (job_env && job_env->di_tv.v_type != VAR_DICT) {
+      EMSG2(_(e_invarg2), "env");
+      shell_free_argv(argv);
+      return;
+    }
 
-      TV_DICT_ITER(job_env->di_tv.vval.v_dict, var, {
-        const char *str = tv_get_string(&var->di_tv);
-        assert(str);
-        size_t len = STRLEN(var->di_key) + strlen(str) + strlen("=") + 1;
-        env[i] = xmalloc(len);
-        snprintf(env[i], len, "%s=%s", (char *)var->di_key, str);
-        i++;
-      });
+    env = tv_dict_alloc();
 
-      // must be null terminated
-      env[env_size + custom_env_size] = NULL;
+    if (!clear_env) {
+      typval_T temp_env = TV_INITIAL_VALUE;
+      f_environ(NULL, &temp_env, NULL);
+      tv_dict_extend(env, temp_env.vval.v_dict, "force");
+      tv_dict_free(temp_env.vval.v_dict);
     }
 
+    if (job_env) {
+      tv_dict_extend(env, job_env->di_tv.vval.v_dict, "force");
+    }
 
     if (!common_job_callbacks(job_opts, &on_stdout, &on_stderr, &on_exit)) {
       shell_free_argv(argv);
+      tv_dict_free(env);
       return;
     }
   }
diff --git a/src/nvim/eval/typval.c b/src/nvim/eval/typval.c
index 02d32a4f86..2f8776def4 100644
--- a/src/nvim/eval/typval.c
+++ b/src/nvim/eval/typval.c
@@ -1523,6 +1523,33 @@ varnumber_T tv_dict_get_number(const dict_T *const d, const char *const key)
   return tv_get_number(&di->di_tv);
 }
 
+/// Converts a dict to an environment
+///
+///
+char **tv_dict_to_env(dict_T *denv)
+{
+  size_t env_size = (size_t)tv_dict_len(denv);
+
+  size_t i = 0;
+  char **env = NULL;
+
+  // + 1 for NULL
+  env = xmalloc((env_size + 1) * sizeof(*env));
+
+  TV_DICT_ITER(denv, var, {
+    const char *str = tv_get_string(&var->di_tv);
+    assert(str);
+    size_t len = STRLEN(var->di_key) + strlen(str) + strlen("=") + 1;
+    env[i] = xmalloc(len);
+    snprintf(env[i], len, "%s=%s", (char *)var->di_key, str);
+    i++;
+  });
+
+  // must be null terminated
+  env[env_size] = NULL;
+  return env;
+}
+
 /// Get a string item from a dictionary
 ///
 /// @param[in]  d  Dictionary to get item from.
diff --git a/src/nvim/event/libuv_process.c b/src/nvim/event/libuv_process.c
index 13517d3df1..72e45878f6 100644
--- a/src/nvim/event/libuv_process.c
+++ b/src/nvim/event/libuv_process.c
@@ -41,7 +41,6 @@ int libuv_process_spawn(LibuvProcess *uvproc)
 #endif
   uvproc->uvopts.exit_cb = exit_cb;
   uvproc->uvopts.cwd = proc->cwd;
-  uvproc->uvopts.env = proc->env;
   uvproc->uvopts.stdio = uvproc->uvstdio;
   uvproc->uvopts.stdio_count = 3;
   uvproc->uvstdio[0].flags = UV_IGNORE;
@@ -49,6 +48,12 @@ int libuv_process_spawn(LibuvProcess *uvproc)
   uvproc->uvstdio[2].flags = UV_IGNORE;
   uvproc->uv.data = proc;
 
+  if (proc->env) {
+    uvproc->uvopts.env = tv_dict_to_env(proc->env);
+  } else {
+    uvproc->uvopts.env = NULL;
+  }
+
   if (!proc->in.closed) {
     uvproc->uvstdio[0].flags = UV_CREATE_PIPE | UV_READABLE_PIPE;
 #ifdef WIN32
diff --git a/src/nvim/event/process.h b/src/nvim/event/process.h
index 84e81238e9..24debdb276 100644
--- a/src/nvim/event/process.h
+++ b/src/nvim/event/process.h
@@ -4,6 +4,7 @@
 #include "nvim/event/loop.h"
 #include "nvim/event/rstream.h"
 #include "nvim/event/wstream.h"
+#include "nvim/eval/typval.h"
 
 typedef enum {
   kProcessTypeUv,
@@ -23,7 +24,7 @@ struct process {
   uint64_t stopped_time;  // process_stop() timestamp
   const char *cwd;
   char **argv;
-  char **env;
+  dict_T *env;
   Stream in, out, err;
   process_exit_cb cb;
   internal_process_cb internal_exit_cb, internal_close_cb;
diff --git a/src/nvim/os/pty_process_unix.c b/src/nvim/os/pty_process_unix.c
index 4d7d9a45df..a3c9e726c5 100644
--- a/src/nvim/os/pty_process_unix.c
+++ b/src/nvim/os/pty_process_unix.c
@@ -20,6 +20,10 @@
 # include <pty.h>
 #endif
 
+#ifdef __APPLE__
+# include <crt_externs.h>
+#endif
+
 #include <uv.h>
 
 #include "nvim/lib/klist.h"
@@ -151,31 +155,24 @@ void pty_process_teardown(Loop *loop)
   uv_signal_stop(&loop->children_watcher);
 }
 
+static const char *ignored_env_vars[] = {
+  "COLUMNS",
+  "LINES",
+  "TERMCAP",
+  "COLORFGBG"
+};
+
 static void init_child(PtyProcess *ptyproc)
   FUNC_ATTR_NONNULL_ALL
 {
+#if defined(HAVE__NSGETENVIRON)
+  char **environ = *_NSGetEnviron();
+#else
+  extern char **environ;
+#endif
   // New session/process-group. #6530
   setsid();
 
-  os_unsetenv("COLUMNS");
-  os_unsetenv("LINES");
-  os_unsetenv("TERMCAP");
-  os_unsetenv("COLORFGBG");
-  // setting COLORTERM to "truecolor" if termguicolors is set and 256
-  // otherwise, but only if it was set in the parent terminal at all
-  if (os_env_exists("COLORTERM")) {
-    const char *colorterm = os_getenv("COLORTERM");
-    if (colorterm != NULL) {
-      if (p_tgc) {
-        os_setenv("COLORTERM", "truecolor", 1);
-      } else {
-        os_setenv("COLORTERM", "256", 1);
-      }
-    } else {
-      os_unsetenv("COLORTERM");
-    }
-  }
-
   signal(SIGCHLD, SIG_DFL);
   signal(SIGHUP, SIG_DFL);
   signal(SIGINT, SIG_DFL);
@@ -190,9 +187,49 @@ static void init_child(PtyProcess *ptyproc)
   }
 
   char *prog = ptyproc->process.argv[0];
-  os_setenv("TERM", ptyproc->term_name ? ptyproc->term_name : "ansi", 1);
-  execvp(prog, ptyproc->process.argv);
+  if (proc->env) {
+    for (size_t i = 0; i < ARRAY_SIZE(ignored_env_vars); i++) {
+      dictitem_T *dv = tv_dict_find(proc->env, ignored_env_vars[i], -1);
+      if (dv) {
+        tv_dict_item_remove(proc->env, dv);
+      }
+    }
+    tv_dict_add_str(proc->env, S_LEN("TERM"), ptyproc->term_name ? ptyproc->term_name : "ansi");
+
+    // setting COLORTERM to "truecolor" if termguicolors is set and 256
+    // otherwise, but only if it was set in the parent terminal at all
+    dictitem_T *dv = tv_dict_find(proc->env, S_LEN("COLORTERM"));
+    if (dv) {
+      tv_dict_item_remove(proc->env, dv);
+      tv_dict_add_str(proc->env, S_LEN("COLORTERM"), p_tgc ? "truecolor" : "256");
+    }
+
+    environ = tv_dict_to_env(proc->env);
+  } else {
+    for (size_t i = 0; i < ARRAY_SIZE(ignored_env_vars); i++) {
+      os_unsetenv(ignored_env_vars[i]);
+    }
+
+    // setting COLORTERM to "truecolor" if termguicolors is set and 256
+    // otherwise, but only if it was set in the parent terminal at all
+    if (os_env_exists("COLORTERM")) {
+      const char *colorterm = os_getenv("COLORTERM");
+      if (colorterm != NULL) {
+        if (p_tgc) {
+          os_setenv("COLORTERM", "truecolor", 1);
+        } else {
+          os_setenv("COLORTERM", "256", 1);
+        }
+      } else {
+        os_unsetenv("COLORTERM");
+      }
+    }
+
+    os_setenv("TERM", ptyproc->term_name ? ptyproc->term_name : "ansi", 1);
+  }
+  execvp(prog, proc->argv);
   ELOG("execvp failed: %s: %s", strerror(errno), prog);
+
   _exit(122);  // 122 is EXEC_FAILED in the Vim source.
 }
 
diff --git a/test/functional/core/job_spec.lua b/test/functional/core/job_spec.lua
index 6d1182478a..3f7595c149 100644
--- a/test/functional/core/job_spec.lua
+++ b/test/functional/core/job_spec.lua
@@ -75,6 +75,20 @@ describe('jobs', function()
     })
   end)
 
+  it('append environment with pty #env', function()
+    nvim('command', "let $VAR = 'abc'")
+    nvim('command', "let g:job_opts.pty = v:true")
+    nvim('command', "let g:job_opts.env = {'TOTO': 'hello world'}")
+    if iswin() then
+      nvim('command', [[call jobstart('echo %TOTO% %VAR%', g:job_opts)]])
+    else
+      nvim('command', [[call jobstart('echo $TOTO $VAR', g:job_opts)]])
+    end
+    expect_msg_seq({
+      {'notification', 'stdout', {0, {'hello world abc', ''}}},
+    })
+  end)
+
   it('replace environment #env', function()
     nvim('command', "let $VAR = 'abc'")
     nvim('command', "let g:job_opts.env = {'TOTO': 'hello world'}")