From f56089cede704dd83f1610f68014081e0f564296 Mon Sep 17 00:00:00 2001
From: Iain Patterson <me@iain.cx>
Date: Thu, 21 Jul 2016 15:18:53 +0100
Subject: [PATCH] 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.
---
 env.cpp | 89 ++++++++++++++++++++++++++++++++++++++++++++++++---------
 env.h   |  2 ++
 2 files changed, 78 insertions(+), 13 deletions(-)

diff --git a/env.cpp b/env.cpp
index 4427d91..3ecf207 100644
--- a/env.cpp
+++ b/env.cpp
@@ -1,20 +1,27 @@
 #include "nssm.h"
 
-/* Copy an environment block. */
-TCHAR *copy_environment_block(TCHAR *env) {
-  unsigned long len;
-
-  if (! env) return 0;
-  for (len = 0; env[len]; len++) while (env[len]) len++;
-  if (! len++) return 0;
-
-  TCHAR *newenv = (TCHAR *) HeapAlloc(GetProcessHeap(), 0, len * sizeof(TCHAR));
-  if (! newenv) {
-    log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_OUT_OF_MEMORY, _T("environment"), _T("copy_environment_block()"), 0);
-    return 0;
+/* Find the length in characters of an environment block. */
+size_t environment_length(TCHAR *env) {
+  size_t len = 0;
+
+  TCHAR *s;
+  for (s = env; ; s++) {
+    len++;
+    if (*s == _T('\0')) {
+      if (*(s + 1) == _T('\0')) {
+        len++;
+        break;
+      }
+    }
   }
 
-  memmove(newenv, env, len * sizeof(TCHAR));
+  return len;
+}
+
+/* Copy an environment block. */
+TCHAR *copy_environment_block(TCHAR *env) {
+  TCHAR *newenv;
+  if (copy_double_null(env, (unsigned long) environment_length(env), &newenv)) return 0;
   return newenv;
 }
 
@@ -179,3 +186,59 @@ TCHAR *copy_environment() {
   FreeEnvironmentStrings(rawenv);
   return env;
 }
+
+/*
+  Create a new block with all the strings of the first block plus a new string.
+  If the key is already present its value will be overwritten in place.
+  If the key is blank or empty the new block will still be allocated and have
+  non-zero length.
+*/
+int append_to_environment_block(TCHAR *env, unsigned long envlen, TCHAR *string, TCHAR **newenv, unsigned long *newlen) {
+  size_t keylen = 0;
+  if (string && string[0]) {
+    for (; string[keylen]; keylen++) {
+      if (string[keylen] == _T('=')) {
+        keylen++;
+        break;
+      }
+    }
+  }
+  return append_to_double_null(env, envlen, newenv, newlen, string, keylen, false);
+}
+
+/*
+  Create a new block with all the strings of the first block minus the given
+  string.
+  If the key is not present the new block will be a copy of the original.
+  If the string is KEY=VALUE the key will only be removed if its value is
+  VALUE.
+  If the string is just KEY the key will unconditionally be removed.
+  If removing the string results in an empty list the new block will still be
+  allocated and have non-zero length.
+*/
+int remove_from_environment_block(TCHAR *env, unsigned long envlen, TCHAR *string, TCHAR **newenv, unsigned long *newlen) {
+  if (! string || ! string[0] || string[0] == _T('=')) return 1;
+
+  TCHAR *key = 0;
+  size_t len = _tcslen(string);
+  size_t i;
+  for (i = 0; i < len; i++) if (string[i] == _T('=')) break;
+
+  /* Rewrite KEY to KEY= but leave KEY=VALUE alone. */
+  size_t keylen = len;
+  if (i == len) keylen++;
+
+  key = (TCHAR *) HeapAlloc(GetProcessHeap(), 0, (keylen + 1) * sizeof(TCHAR));
+  if (! key) {
+    log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_OUT_OF_MEMORY, _T("key"), _T("remove_from_environment_block()"), 0);
+    return 2;
+  }
+  memmove(key, string, len * sizeof(TCHAR));
+  if (keylen > len) key[keylen - 1] = _T('=');
+  key[keylen] = _T('\0');
+
+  int ret = remove_from_double_null(env, envlen, newenv, newlen, key, keylen, false);
+  HeapFree(GetProcessHeap(), 0, key);
+
+  return ret;
+}
diff --git a/env.h b/env.h
index 73bcedd..a5202d3 100644
--- a/env.h
+++ b/env.h
@@ -11,5 +11,7 @@ int duplicate_environment(TCHAR *);
 int test_environment(TCHAR *);
 void duplicate_environment_strings(TCHAR *);
 TCHAR *copy_environment();
+int append_to_environment_block(TCHAR *, unsigned long, TCHAR *, TCHAR **, unsigned long *);
+int remove_from_environment_block(TCHAR *, unsigned long, TCHAR *, TCHAR **, unsigned long *);
 
 #endif
-- 
2.20.1