#include "nssm.h"
-/* 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;
- *newlen = envlen;
+/* Copy an environment block. */
+TCHAR *copy_environment_block(TCHAR *env) {
+ unsigned long len;
- if (! *newlen) {
- *formatted = 0;
+ 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;
}
- for (i = 0; i < envlen; i++) if (! env[i] && env[i + 1]) ++*newlen;
+ memmove(newenv, env, len * sizeof(TCHAR));
+ return newenv;
+}
- *formatted = (TCHAR *) HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, *newlen * sizeof(TCHAR));
- if (! *formatted) {
- *newlen = 0;
- return 1;
- }
+/*
+ The environment block starts with variables of the form
+ =C:=C:\Windows\System32 which we ignore.
+*/
+TCHAR *useful_environment(TCHAR *rawenv) {
+ TCHAR *env = rawenv;
- for (i = 0, j = 0; i < envlen; i++) {
- (*formatted)[j] = env[i];
- if (! env[i]) {
- if (env[i + 1]) {
- (*formatted)[j] = _T('\r');
- (*formatted)[++j] = _T('\n');
- }
+ if (env) {
+ while (*env == _T('=')) {
+ for ( ; *env; env++);
+ env++;
}
- j++;
}
- return 0;
+ return env;
}
-/* Strip CR and replace LF with NULL. */
-int unformat_environment(TCHAR *env, unsigned long envlen, TCHAR **unformatted, unsigned long *newlen) {
- unsigned long i, j;
- *newlen = 0;
+/* 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 (! envlen) {
- *unformatted = 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;
}
- for (i = 0; i < envlen; i++) if (env[i] != _T('\r')) ++*newlen;
- /* Must end with two NULLs. */
- *newlen += 2;
+ return ret;
+}
- *unformatted = (TCHAR *) HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, *newlen * sizeof(TCHAR));
- if (! *unformatted) return 1;
+/*
+ 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;
- for (i = 0, j = 0; i < envlen; i++) {
- if (env[i] == _T('\r')) continue;
- if (env[i] == _T('\n')) (*unformatted)[j] = _T('\0');
- else (*unformatted)[j] = env[i];
- j++;
+ 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 0;
+ 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;
}
/*
return 0;
}
+
+/*
+ Duplicate an environment block returned by GetEnvironmentStrings().
+ Since such a block is by definition readonly, and duplicate_environment()
+ modifies its inputs, this function takes a copy of the input and operates
+ on that.
+*/
+void duplicate_environment_strings(TCHAR *env) {
+ TCHAR *newenv = copy_environment_block(env);
+ if (! newenv) return;
+
+ duplicate_environment(newenv);
+ HeapFree(GetProcessHeap(), 0, newenv);
+}