From: Iain Patterson Date: Thu, 21 Jul 2016 14:18:53 +0000 (+0100) Subject: Added append_to/remove_from_environment_block(). X-Git-Url: http://git.iain.cx/?a=commitdiff_plain;h=f56089cede704dd83f1610f68014081e0f564296;p=nssm.git 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 --- 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