env.cpp

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