Add os_getfullenv_size/os_copyfullenv

1 files changed, 123 insertions, 29 deletions

diff --git a/src/nvim/os/env.c b/src/nvim/os/env.c
index 15153e9bd7..ac442ee2e8 100644
--- a/src/nvim/os/env.c
+++ b/src/nvim/os/env.c
@@ -184,39 +184,139 @@ int os_unsetenv(const char *name)
   return r == 0 ? 0 : -1;
 }
 
-char **os_getfullenv(void)
+/// Returns number of variables in the current environment variables block
+size_t os_get_fullenv_size(void)
 {
+  size_t len = 0;
 #ifdef _WIN32
-  wchar_t *env = GetEnvironmentStringsW();
-  if (!env) {
+  wchar_t *envstrings = GetEnvironmentStringsW();
+  wchar_t *p = envstrings;
+  size_t l;
+  if (!envstrings) {
+    return len;
+  }
+  // GetEnvironmentStringsW() result has this format:
+  //    var1=value1\0var2=value2\0...varN=valueN\0\0
+  while ((l = wcslen(p)) != 0) {
+    p += l + 1;
+    len++;
+  }
+
+  FreeEnvironmentStringsW(envstrings);
+#else
+# if defined(HAVE__NSGETENVIRON)
+  char **environ = *_NSGetEnviron();
+# else
+  extern char         **environ;
+# endif
+
+  while (environ[len] != NULL) {
+    len++;
+  }
+
+#endif
+  return len;
+}
+
+void os_free_fullenv(char **env)
+{
+  if (!env) { return; }
+  for (char **it = env; *it; it++) {
+    XFREE_CLEAR(*it);
+  }
+  xfree(env);
+}
+
+/// Copies the current environment variables into the given array, `env`.  Each
+/// array element is of the form "NAME=VALUE".
+/// Result must be freed by the caller.
+///
+/// @param[out]  env  array to populate with environment variables
+/// @param  env_size  size of `env`, @see os_fullenv_size
+void os_copy_fullenv(char **env, size_t env_size)
+{
+#ifdef _WIN32
+  wchar_t *envstrings = GetEnvironmentStringsW();
+  if (!envstrings) {
+    return;
+  }
+  wchar_t *p = envstrings;
+  size_t i = 0;
+  size_t l;
+  // GetEnvironmentStringsW() result has this format:
+  //    var1=value1\0var2=value2\0...varN=valueN\0\0
+  while ((l = wcslen(p)) != 0 && i < env_size) {
+    char *utf8_str;
+    int conversion_result = utf16_to_utf8(p, -1, &utf8_str);
+    if (conversion_result != 0) {
+      EMSG2("utf16_to_utf8 failed: %d", conversion_result);
+      break;
+    }
+    p += l + 1;
+
+    env[i] = utf8_str;
+    i++;
+  }
+
+  FreeEnvironmentStringsW(envstrings);
+#else
+# if defined(HAVE__NSGETENVIRON)
+  char **environ = *_NSGetEnviron();
+# else
+  extern char         **environ;
+# endif
+
+  size_t i = 0;
+  while (environ[i] != NULL && i < env_size) {
+    env[i] = xstrdup(environ[i]);
+    i++;
+  }
+#endif
+}
+
+/// Copy value of the environment variable at `index` in the current
+/// environment variables block.
+/// Result must be freed by the caller.
+///
+/// @param index nth item in environment variables block
+/// @return [allocated] environment variable's value, or NULL
+char *os_getenvname_at_index(size_t index)
+{
+#ifdef _WIN32
+  wchar_t *envstrings = GetEnvironmentStringsW();
+  if (!envstrings) {
     return NULL;
   }
+  wchar_t *p = envstrings;
   char *name = NULL;
-  size_t current_index = 0;
+  size_t i = 0;
+  size_t l;
   // GetEnvironmentStringsW() result has this format:
   //    var1=value1\0var2=value2\0...varN=valueN\0\0
-  for (wchar_t *it = env; *it != L'\0' || *(it + 1) != L'\0'; it++) {
-    if (index == current_index) {
+  while ((l = wcslen(p)) != 0 && i <= index) {
+    if (i == index) {
       char *utf8_str;
-      int conversion_result = utf16_to_utf8(it, -1, &utf8_str);
+      int conversion_result = utf16_to_utf8(p, -1, &utf8_str);
       if (conversion_result != 0) {
         EMSG2("utf16_to_utf8 failed: %d", conversion_result);
         break;
       }
-      size_t namesize = 0;
-      while (utf8_str[namesize] != '=' && utf8_str[namesize] != NUL) {
-        namesize++;
-      }
-      name = (char *)vim_strnsave((char_u *)utf8_str, namesize);
+
+      const char * const end = strchr(utf8_str, '=');
+      assert(end != NULL);
+      ptrdiff_t len = end - utf8_str;
+      assert(len > 0);
+      name = xstrndup(utf8_str, (size_t)len);
       xfree(utf8_str);
       break;
     }
-    if (*it == L'\0') {
-      current_index++;
-    }
+
+    // Advance past the name and NUL
+    p += l + 1;
+    i++;
   }
 
-  FreeEnvironmentStringsW(env);
+  FreeEnvironmentStringsW(envstrings);
   return name;
 #else
 # if defined(HAVE__NSGETENVIRON)
@@ -224,26 +324,20 @@ char **os_getfullenv(void)
 # else
   extern char         **environ;
 # endif
-  return environ;
-}
 
-char *os_getenvname_at_index(size_t index)
-{
-  char **env = os_getfullenv();
   // check if index is inside the environ array
   for (size_t i = 0; i <= index; i++) {
-    if (env[i] == NULL) {
+    if (environ[i] == NULL) {
       return NULL;
     }
   }
-  char *str = env[index];
+  char *str = environ[index];
   assert(str != NULL);
-  size_t namesize = 0;
-  while (str[namesize] != '=' && str[namesize] != NUL) {
-    namesize++;
-  }
-  char *name = (char *)vim_strnsave((char_u *)str, namesize);
-  return name;
+  const char * const end = strchr(str, '=');
+  assert(end != NULL);
+  ptrdiff_t len = end - str;
+  assert(len > 0);
+  return xstrndup(str, (size_t)len);
 #endif
 }