Use UTF-8 functions when working with explicit encodings.
[nssm.git] / env.cpp
diff --git a/env.cpp b/env.cpp
index 76c187e..3366fc6 100644 (file)
--- a/env.cpp
+++ b/env.cpp
-#include "nssm.h"
-
-/*
-  The environment block starts with variables of the form
-  =C:=C:\Windows\System32 which we ignore.
-*/
-TCHAR *useful_environment(TCHAR *rawenv) {
-  TCHAR *env = rawenv;
-
-  if (env) {
-    while (*env == _T('=')) {
-      for ( ; *env; env++);
-      env++;
-    }
-  }
-
-  return env;
-}
-
-/* Expand an environment variable.  Must call HeapFree() on the result. */
-TCHAR *expand_environment_string(TCHAR *string) {
-  unsigned long len;
-
-  len = ExpandEnvironmentStrings(string, 0, 0);
-  if (! len) {
-    log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_EXPANDENVIRONMENTSTRINGS_FAILED, string, error_string(GetLastError()), 0);
-    return 0;
-  }
-
-  TCHAR *ret = (TCHAR *) HeapAlloc(GetProcessHeap(), 0, len * sizeof(TCHAR));
-  if (! ret) {
-    log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_OUT_OF_MEMORY, _T("ExpandEnvironmentStrings()"), _T("expand_environment_string"), 0);
-    return 0;
-  }
-
-  if (! ExpandEnvironmentStrings(string, ret, len)) {
-    log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_EXPANDENVIRONMENTSTRINGS_FAILED, string, error_string(GetLastError()), 0);
-    HeapFree(GetProcessHeap(), 0, ret);
-    return 0;
-  }
-
-  return ret;
-}
-
-/*
-  Set all the environment variables from an environment block in the current
-  environment or remove all the variables in the block from the current
-  environment.
-*/
-static int set_environment_block(TCHAR *env, bool set) {
-  int ret = 0;
-
-  TCHAR *s, *t;
-  for (s = env; *s; s++) {
-    for (t = s; *t && *t != _T('='); t++);
-    if (*t == _T('=')) {
-      *t = _T('\0');
-      if (set) {
-        TCHAR *expanded = expand_environment_string(++t);
-        if (expanded) {
-          if (! SetEnvironmentVariable(s, expanded)) ret++;
-          HeapFree(GetProcessHeap(), 0, expanded);
-        }
-        else {
-          if (! SetEnvironmentVariable(s, t)) ret++;
-        }
-      }
-      else {
-        if (! SetEnvironmentVariable(s, NULL)) ret++;
-      }
-      for (t++ ; *t; t++);
-    }
-    s = t;
-  }
-
-  return ret;
-}
-
-int set_environment_block(TCHAR *env) {
-  return set_environment_block(env, true);
-}
-
-static int unset_environment_block(TCHAR *env) {
-  return set_environment_block(env, false);
-}
-
-/* Remove all variables from the process environment. */
-int clear_environment() {
-  TCHAR *rawenv = GetEnvironmentStrings();
-  TCHAR *env = useful_environment(rawenv);
-
-  int ret = unset_environment_block(env);
-
-  if (rawenv) FreeEnvironmentStrings(rawenv);
-
-  return ret;
-}
-
-/* Set the current environment to exactly duplicate an environment block. */
-int duplicate_environment(TCHAR *rawenv) {
-  int ret = clear_environment();
-  TCHAR *env = useful_environment(rawenv);
-  ret += set_environment_block(env);
-  return ret;
-}
-
-/*
-  Verify an environment block.
-  Returns:  1 if environment is invalid.
-            0 if environment is OK.
-           -1 on error.
-*/
-int test_environment(TCHAR *env) {
-  TCHAR path[PATH_LENGTH];
-  GetModuleFileName(0, path, _countof(path));
-  STARTUPINFO si;
-  ZeroMemory(&si, sizeof(si));
-  si.cb = sizeof(si);
-  PROCESS_INFORMATION pi;
-  ZeroMemory(&pi, sizeof(pi));
-  unsigned long flags = CREATE_SUSPENDED;
-#ifdef UNICODE
-  flags |= CREATE_UNICODE_ENVIRONMENT;
-#endif
-
-  /*
-    Try to relaunch ourselves but with the candidate environment set.
-    Assuming no solar flare activity, the only reason this would fail is if
-    the environment were invalid.
-  */
-  if (CreateProcess(0, path, 0, 0, 0, flags, env, 0, &si, &pi)) {
-    TerminateProcess(pi.hProcess, 0);
-  }
-  else {
-    unsigned long error = GetLastError();
-    if (error == ERROR_INVALID_PARAMETER) return 1;
-    else return -1;
-  }
-
-  return 0;
-}
+#include "nssm.h"\r
+\r
+/*\r
+  Environment block is of the form:\r
+\r
+    KEY1=VALUE1 NULL\r
+    KEY2=VALUE2 NULL\r
+    NULL\r
+\r
+  A single variable KEY=VALUE has length 15:\r
+\r
+    KEY=VALUE (13) NULL (1)\r
+    NULL (1)\r
+\r
+  Environment variable names are case-insensitive!\r
+*/\r
+\r
+/* Find the length in characters of an environment block. */\r
+size_t environment_length(TCHAR *env) {\r
+  size_t len = 0;\r
+\r
+  TCHAR *s;\r
+  for (s = env; ; s++) {\r
+    len++;\r
+    if (*s == _T('\0')) {\r
+      if (*(s + 1) == _T('\0')) {\r
+        len++;\r
+        break;\r
+      }\r
+    }\r
+  }\r
+\r
+  return len;\r
+}\r
+\r
+/* Copy an environment block. */\r
+TCHAR *copy_environment_block(TCHAR *env) {\r
+  TCHAR *newenv;\r
+  if (copy_double_null(env, (unsigned long) environment_length(env), &newenv)) return 0;\r
+  return newenv;\r
+}\r
+\r
+/*\r
+  The environment block starts with variables of the form\r
+  =C:=C:\Windows\System32 which we ignore.\r
+*/\r
+TCHAR *useful_environment(TCHAR *rawenv) {\r
+  TCHAR *env = rawenv;\r
+\r
+  if (env) {\r
+    while (*env == _T('=')) {\r
+      for ( ; *env; env++);\r
+      env++;\r
+    }\r
+  }\r
+\r
+  return env;\r
+}\r
+\r
+/* Expand an environment variable.  Must call HeapFree() on the result. */\r
+TCHAR *expand_environment_string(TCHAR *string) {\r
+  unsigned long len;\r
+\r
+  len = ExpandEnvironmentStrings(string, 0, 0);\r
+  if (! len) {\r
+    log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_EXPANDENVIRONMENTSTRINGS_FAILED, string, error_string(GetLastError()), 0);\r
+    return 0;\r
+  }\r
+\r
+  TCHAR *ret = (TCHAR *) HeapAlloc(GetProcessHeap(), 0, len * sizeof(TCHAR));\r
+  if (! ret) {\r
+    log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_OUT_OF_MEMORY, _T("ExpandEnvironmentStrings()"), _T("expand_environment_string"), 0);\r
+    return 0;\r
+  }\r
+\r
+  if (! ExpandEnvironmentStrings(string, ret, len)) {\r
+    log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_EXPANDENVIRONMENTSTRINGS_FAILED, string, error_string(GetLastError()), 0);\r
+    HeapFree(GetProcessHeap(), 0, ret);\r
+    return 0;\r
+  }\r
+\r
+  return ret;\r
+}\r
+\r
+/*\r
+  Set all the environment variables from an environment block in the current\r
+  environment or remove all the variables in the block from the current\r
+  environment.\r
+*/\r
+static int set_environment_block(TCHAR *env, bool set) {\r
+  int ret = 0;\r
+\r
+  TCHAR *s, *t;\r
+  for (s = env; *s; s++) {\r
+    for (t = s; *t && *t != _T('='); t++);\r
+    if (*t == _T('=')) {\r
+      *t = _T('\0');\r
+      if (set) {\r
+        TCHAR *expanded = expand_environment_string(++t);\r
+        if (expanded) {\r
+          if (! SetEnvironmentVariable(s, expanded)) ret++;\r
+          HeapFree(GetProcessHeap(), 0, expanded);\r
+        }\r
+        else {\r
+          if (! SetEnvironmentVariable(s, t)) ret++;\r
+        }\r
+      }\r
+      else {\r
+        if (! SetEnvironmentVariable(s, NULL)) ret++;\r
+      }\r
+      for (t++; *t; t++);\r
+    }\r
+    s = t;\r
+  }\r
+\r
+  return ret;\r
+}\r
+\r
+int set_environment_block(TCHAR *env) {\r
+  return set_environment_block(env, true);\r
+}\r
+\r
+static int unset_environment_block(TCHAR *env) {\r
+  return set_environment_block(env, false);\r
+}\r
+\r
+/* Remove all variables from the process environment. */\r
+int clear_environment() {\r
+  TCHAR *rawenv = GetEnvironmentStrings();\r
+  TCHAR *env = useful_environment(rawenv);\r
+\r
+  int ret = unset_environment_block(env);\r
+\r
+  if (rawenv) FreeEnvironmentStrings(rawenv);\r
+\r
+  return ret;\r
+}\r
+\r
+/* Set the current environment to exactly duplicate an environment block. */\r
+int duplicate_environment(TCHAR *rawenv) {\r
+  int ret = clear_environment();\r
+  TCHAR *env = useful_environment(rawenv);\r
+  ret += set_environment_block(env);\r
+  return ret;\r
+}\r
+\r
+/*\r
+  Verify an environment block.\r
+  Returns:  1 if environment is invalid.\r
+            0 if environment is OK.\r
+           -1 on error.\r
+*/\r
+int test_environment(TCHAR *env) {\r
+  TCHAR *path = (TCHAR *) nssm_imagepath();\r
+  STARTUPINFO si;\r
+  ZeroMemory(&si, sizeof(si));\r
+  si.cb = sizeof(si);\r
+  PROCESS_INFORMATION pi;\r
+  ZeroMemory(&pi, sizeof(pi));\r
+  unsigned long flags = CREATE_SUSPENDED;\r
+#ifdef UNICODE\r
+  flags |= CREATE_UNICODE_ENVIRONMENT;\r
+#endif\r
+\r
+  /*\r
+    Try to relaunch ourselves but with the candidate environment set.\r
+    Assuming no solar flare activity, the only reason this would fail is if\r
+    the environment were invalid.\r
+  */\r
+  if (CreateProcess(0, path, 0, 0, 0, flags, env, 0, &si, &pi)) {\r
+    TerminateProcess(pi.hProcess, 0);\r
+  }\r
+  else {\r
+    unsigned long error = GetLastError();\r
+    if (error == ERROR_INVALID_PARAMETER) return 1;\r
+    else return -1;\r
+  }\r
+\r
+  return 0;\r
+}\r
+\r
+/*\r
+  Duplicate an environment block returned by GetEnvironmentStrings().\r
+  Since such a block is by definition readonly, and duplicate_environment()\r
+  modifies its inputs, this function takes a copy of the input and operates\r
+  on that.\r
+*/\r
+void duplicate_environment_strings(TCHAR *env) {\r
+  TCHAR *newenv = copy_environment_block(env);\r
+  if (! newenv) return;\r
+\r
+  duplicate_environment(newenv);\r
+  HeapFree(GetProcessHeap(), 0, newenv);\r
+}\r
+\r
+/* Safely get a copy of the current environment. */\r
+TCHAR *copy_environment() {\r
+  TCHAR *rawenv = GetEnvironmentStrings();\r
+  if (! rawenv) return NULL;\r
+  TCHAR *env = copy_environment_block(rawenv);\r
+  FreeEnvironmentStrings(rawenv);\r
+  return env;\r
+}\r
+\r
+/*\r
+  Create a new block with all the strings of the first block plus a new string.\r
+  If the key is already present its value will be overwritten in place.\r
+  If the key is blank or empty the new block will still be allocated and have\r
+  non-zero length.\r
+*/\r
+int append_to_environment_block(TCHAR *env, unsigned long envlen, TCHAR *string, TCHAR **newenv, unsigned long *newlen) {\r
+  size_t keylen = 0;\r
+  if (string && string[0]) {\r
+    for (; string[keylen]; keylen++) {\r
+      if (string[keylen] == _T('=')) {\r
+        keylen++;\r
+        break;\r
+      }\r
+    }\r
+  }\r
+  return append_to_double_null(env, envlen, newenv, newlen, string, keylen, false);\r
+}\r
+\r
+/*\r
+  Create a new block with all the strings of the first block minus the given\r
+  string.\r
+  If the key is not present the new block will be a copy of the original.\r
+  If the string is KEY=VALUE the key will only be removed if its value is\r
+  VALUE.\r
+  If the string is just KEY the key will unconditionally be removed.\r
+  If removing the string results in an empty list the new block will still be\r
+  allocated and have non-zero length.\r
+*/\r
+int remove_from_environment_block(TCHAR *env, unsigned long envlen, TCHAR *string, TCHAR **newenv, unsigned long *newlen) {\r
+  if (! string || ! string[0] || string[0] == _T('=')) return 1;\r
+\r
+  TCHAR *key = 0;\r
+  size_t len = _tcslen(string);\r
+  size_t i;\r
+  for (i = 0; i < len; i++) if (string[i] == _T('=')) break;\r
+\r
+  /* Rewrite KEY to KEY= but leave KEY=VALUE alone. */\r
+  size_t keylen = len;\r
+  if (i == len) keylen++;\r
+\r
+  key = (TCHAR *) HeapAlloc(GetProcessHeap(), 0, (keylen + 1) * sizeof(TCHAR));\r
+  if (! key) {\r
+    log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_OUT_OF_MEMORY, _T("key"), _T("remove_from_environment_block()"), 0);\r
+    return 2;\r
+  }\r
+  memmove(key, string, len * sizeof(TCHAR));\r
+  if (keylen > len) key[keylen - 1] = _T('=');\r
+  key[keylen] = _T('\0');\r
+\r
+  int ret = remove_from_double_null(env, envlen, newenv, newlen, key, keylen, false);\r
+  HeapFree(GetProcessHeap(), 0, key);\r
+\r
+  return ret;\r
+}\r