From: Iain Patterson 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 *);