Added append_to/remove_from_environment_block().

New functions to add an environment variable to an existing environment
block or to remove a variable (specified by KEY=VALUE or just KEY) from
a block.

A new block is always returned.  It will have length 2 (NULL NULL) if
the requested action was a no-op.

diff --git a/env.cpp b/env.cpp
index 4427d91..3ecf207 100644 (file)
--- a/env.cpp
+++ b/env.cpp
@@ -1,20 +1,27 @@
 #include "nssm.h"\r
 \r
-/* Copy an environment block. */\r
-TCHAR *copy_environment_block(TCHAR *env) {\r
-  unsigned long len;\r
-\r
-  if (! env) return 0;\r
-  for (len = 0; env[len]; len++) while (env[len]) len++;\r
-  if (! len++) return 0;\r
-\r
-  TCHAR *newenv = (TCHAR *) HeapAlloc(GetProcessHeap(), 0, len * sizeof(TCHAR));\r
-  if (! newenv) {\r
-    log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_OUT_OF_MEMORY, _T("environment"), _T("copy_environment_block()"), 0);\r
-    return 0;\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
-  memmove(newenv, env, len * sizeof(TCHAR));\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
@@ -179,3 +186,59 @@ TCHAR *copy_environment() {
   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

diff --git a/env.h b/env.h
index 73bcedd..a5202d3 100644 (file)
--- a/env.h
+++ b/env.h
@@ -11,5 +11,7 @@ int duplicate_environment(TCHAR *);
 int test_environment(TCHAR *);\r
 void duplicate_environment_strings(TCHAR *);\r
 TCHAR *copy_environment();\r
+int append_to_environment_block(TCHAR *, unsigned long, TCHAR *, TCHAR **, unsigned long *);\r
+int remove_from_environment_block(TCHAR *, unsigned long, TCHAR *, TCHAR **, unsigned long *);\r
 \r
 #endif\r