From: Iain Patterson <me@iain.cx>
Date: Thu, 23 Jan 2014 21:12:43 +0000 (+0000)
Subject: Added more environment functions.
X-Git-Tag: v2.22~40
X-Git-Url: http://git.iain.cx/?a=commitdiff_plain;h=2e3a41e1d5dea82a7f36b331df59ed7532b30b8c;p=nssm.git

Added more environment functions.

useful_environment() skips the uninteresting variables at the start of
an environment block, returning a pointer to the first useful variable.

expand_environment_string() expands a string with
ExpandEnvironmentStrings() and returns a pointer to the expanded string,
which must be freed after use.

set_environment_block() iterates through all the variables in an
environment block, expands each one with expand_environment_string() and
calls SetEnvironmentVariable() with the expanded value.

clear_environment() removes all variables from the current environment.

duplicate_environment() duplicates an environment block by first calling
clear_environment() then set_environment_block() with the source block.
Thus it ensures that the variables in the block - and only those - are
set in the current environment.
---

diff --git a/env.cpp b/env.cpp
index 922c96f..5b5f34a 100644
--- a/env.cpp
+++ b/env.cpp
@@ -1,5 +1,109 @@
 #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;
+}
+
 /* Replace NULL with CRLF. Leave NULL NULL as the end marker. */
 int format_environment(TCHAR *env, unsigned long envlen, TCHAR **formatted, unsigned long *newlen) {
   unsigned long i, j;
diff --git a/env.h b/env.h
index f23c592..4e0191c 100644
--- a/env.h
+++ b/env.h
@@ -1,6 +1,11 @@
 #ifndef ENV_H
 #define ENV_H
 
+TCHAR *useful_environment(TCHAR *);
+TCHAR *expand_environment_string(TCHAR *);
+int set_environment_block(TCHAR *);
+int clear_environment();
+int duplicate_environment(TCHAR *);
 int format_environment(TCHAR *, unsigned long, TCHAR **, unsigned long *);
 int unformat_environment(TCHAR *, unsigned long, TCHAR **, unsigned long *);
 int test_environment(TCHAR *);