123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259 |
- #include "nssm.h"
- /*
- Environment block is of the form:
- KEY1=VALUE1 NULL
- KEY2=VALUE2 NULL
- NULL
- A single variable KEY=VALUE has length 15:
- KEY=VALUE (13) NULL (1)
- NULL (1)
- Environment variable names are case-insensitive!
- */
- /* 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;
- }
- }
- }
- 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;
- }
- /*
- 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;
- }
- /*
- Verify an environment block.
- Returns: 1 if environment is invalid.
- 0 if environment is OK.
- -1 on error.
- */
- int test_environment(TCHAR *env) {
- TCHAR *path = (TCHAR *) nssm_imagepath();
- STARTUPINFO si;
- ZeroMemory(&si, sizeof(si));
- si.cb = sizeof(si);
- PROCESS_INFORMATION pi;
- ZeroMemory(&pi, sizeof(pi));
- unsigned long flags = CREATE_SUSPENDED;
- #ifdef UNICODE
- flags |= CREATE_UNICODE_ENVIRONMENT;
- #endif
- /*
- Try to relaunch ourselves but with the candidate environment set.
- Assuming no solar flare activity, the only reason this would fail is if
- the environment were invalid.
- */
- if (CreateProcess(0, path, 0, 0, 0, flags, env, 0, &si, &pi)) {
- TerminateProcess(pi.hProcess, 0);
- }
- else {
- unsigned long error = GetLastError();
- if (error == ERROR_INVALID_PARAMETER) return 1;
- else return -1;
- }
- 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);
- }
- /* Safely get a copy of the current environment. */
- TCHAR *copy_environment() {
- TCHAR *rawenv = GetEnvironmentStrings();
- if (! rawenv) return NULL;
- TCHAR *env = copy_environment_block(rawenv);
- 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;
- }
|