env.cpp

   1 #include "nssm.h"
   2
   3 /*
   4   The environment block starts with variables of the form
   5   =C:=C:\Windows\System32 which we ignore.
   6 */
   7 TCHAR *useful_environment(TCHAR *rawenv) {
   8   TCHAR *env = rawenv;
   9
  10   if (env) {
  11     while (*env == _T('=')) {
  12       for ( ; *env; env++);
  13       env++;
  14     }
  15   }
  16
  17   return env;
  18 }
  19
  20 /* Expand an environment variable.  Must call HeapFree() on the result. */
  21 TCHAR *expand_environment_string(TCHAR *string) {
  22   unsigned long len;
  23
  24   len = ExpandEnvironmentStrings(string, 0, 0);
  25   if (! len) {
  26     log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_EXPANDENVIRONMENTSTRINGS_FAILED, string, error_string(GetLastError()), 0);
  27     return 0;
  28   }
  29
  30   TCHAR *ret = (TCHAR *) HeapAlloc(GetProcessHeap(), 0, len * sizeof(TCHAR));
  31   if (! ret) {
  32     log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_OUT_OF_MEMORY, _T("ExpandEnvironmentStrings()"), _T("expand_environment_string"), 0);
  33     return 0;
  34   }
  35
  36   if (! ExpandEnvironmentStrings(string, ret, len)) {
  37     log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_EXPANDENVIRONMENTSTRINGS_FAILED, string, error_string(GetLastError()), 0);
  38     HeapFree(GetProcessHeap(), 0, ret);
  39     return 0;
  40   }
  41
  42   return ret;
  43 }
  44
  45 /*
  46   Set all the environment variables from an environment block in the current
  47   environment or remove all the variables in the block from the current
  48   environment.
  49 */
  50 static int set_environment_block(TCHAR *env, bool set) {
  51   int ret = 0;
  52
  53   TCHAR *s, *t;
  54   for (s = env; *s; s++) {
  55     for (t = s; *t && *t != _T('='); t++);
  56     if (*t == _T('=')) {
  57       *t = _T('\0');
  58       if (set) {
  59         TCHAR *expanded = expand_environment_string(++t);
  60         if (expanded) {
  61           if (! SetEnvironmentVariable(s, expanded)) ret++;
  62           HeapFree(GetProcessHeap(), 0, expanded);
  63         }
  64         else {
  65           if (! SetEnvironmentVariable(s, t)) ret++;
  66         }
  67       }
  68       else {
  69         if (! SetEnvironmentVariable(s, NULL)) ret++;
  70       }
  71       for (t++ ; *t; t++);
  72     }
  73     s = t;
  74   }
  75
  76   return ret;
  77 }
  78
  79 int set_environment_block(TCHAR *env) {
  80   return set_environment_block(env, true);
  81 }
  82
  83 static int unset_environment_block(TCHAR *env) {
  84   return set_environment_block(env, false);
  85 }
  86
  87 /* Remove all variables from the process environment. */
  88 int clear_environment() {
  89   TCHAR *rawenv = GetEnvironmentStrings();
  90   TCHAR *env = useful_environment(rawenv);
  91
  92   int ret = unset_environment_block(env);
  93
  94   if (rawenv) FreeEnvironmentStrings(rawenv);
  95
  96   return ret;
  97 }
  98
  99 /* Set the current environment to exactly duplicate an environment block. */
 100 int duplicate_environment(TCHAR *rawenv) {
 101   int ret = clear_environment();
 102   TCHAR *env = useful_environment(rawenv);
 103   ret += set_environment_block(env);
 104   return ret;
 105 }
 106
 107 /* Replace NULL with CRLF. Leave NULL NULL as the end marker. */
 108 int format_environment(TCHAR *env, unsigned long envlen, TCHAR **formatted, unsigned long *newlen) {
 109   unsigned long i, j;
 110   *newlen = envlen;
 111
 112   if (! *newlen) {
 113     *formatted = 0;
 114     return 0;
 115   }
 116
 117   for (i = 0; i < envlen; i++) if (! env[i] && env[i + 1]) ++*newlen;
 118
 119   *formatted = (TCHAR *) HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, *newlen * sizeof(TCHAR));
 120   if (! *formatted) {
 121     *newlen = 0;
 122     return 1;
 123   }
 124
 125   for (i = 0, j = 0; i < envlen; i++) {
 126     (*formatted)[j] = env[i];
 127     if (! env[i]) {
 128       if (env[i + 1]) {
 129         (*formatted)[j] = _T('\r');
 130         (*formatted)[++j] = _T('\n');
 131       }
 132     }
 133     j++;
 134   }
 135
 136   return 0;
 137 }
 138
 139 /* Strip CR and replace LF with NULL. */
 140 int unformat_environment(TCHAR *env, unsigned long envlen, TCHAR **unformatted, unsigned long *newlen) {
 141   unsigned long i, j;
 142   *newlen = 0;
 143
 144   if (! envlen) {
 145     *unformatted = 0;
 146     return 0;
 147   }
 148
 149   for (i = 0; i < envlen; i++) if (env[i] != _T('\r')) ++*newlen;
 150   /* Must end with two NULLs. */
 151   *newlen += 2;
 152
 153   *unformatted = (TCHAR *) HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, *newlen * sizeof(TCHAR));
 154   if (! *unformatted) return 1;
 155
 156   for (i = 0, j = 0; i < envlen; i++) {
 157     if (env[i] == _T('\r')) continue;
 158     if (env[i] == _T('\n')) (*unformatted)[j] = _T('\0');
 159     else (*unformatted)[j] = env[i];
 160     j++;
 161   }
 162
 163   return 0;
 164 }
 165
 166 /*
 167   Verify an environment block.
 168   Returns:  1 if environment is invalid.
 169             0 if environment is OK.
 170            -1 on error.
 171 */
 172 int test_environment(TCHAR *env) {
 173   TCHAR path[PATH_LENGTH];
 174   GetModuleFileName(0, path, _countof(path));
 175   STARTUPINFO si;
 176   ZeroMemory(&si, sizeof(si));
 177   si.cb = sizeof(si);
 178   PROCESS_INFORMATION pi;
 179   ZeroMemory(&pi, sizeof(pi));
 180   unsigned long flags = CREATE_SUSPENDED;
 181 #ifdef UNICODE
 182   flags |= CREATE_UNICODE_ENVIRONMENT;
 183 #endif
 184
 185   /*
 186     Try to relaunch ourselves but with the candidate environment set.
 187     Assuming no solar flare activity, the only reason this would fail is if
 188     the environment were invalid.
 189   */
 190   if (CreateProcess(0, path, 0, 0, 0, flags, env, 0, &si, &pi)) {
 191     TerminateProcess(pi.hProcess, 0);
 192   }
 193   else {
 194     unsigned long error = GetLastError();
 195     if (error == ERROR_INVALID_PARAMETER) return 1;
 196     else return -1;
 197   }
 198
 199   return 0;
 200 }