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[PATH_LENGTH];
 133   GetModuleFileName(0, path, _countof(path));
 134   STARTUPINFO si;
 135   ZeroMemory(&si, sizeof(si));
 136   si.cb = sizeof(si);
 137   PROCESS_INFORMATION pi;
 138   ZeroMemory(&pi, sizeof(pi));
 139   unsigned long flags = CREATE_SUSPENDED;
 140 #ifdef UNICODE
 141   flags |= CREATE_UNICODE_ENVIRONMENT;
 142 #endif
 143
 144   /*
 145     Try to relaunch ourselves but with the candidate environment set.
 146     Assuming no solar flare activity, the only reason this would fail is if
 147     the environment were invalid.
 148   */
 149   if (CreateProcess(0, path, 0, 0, 0, flags, env, 0, &si, &pi)) {
 150     TerminateProcess(pi.hProcess, 0);
 151   }
 152   else {
 153     unsigned long error = GetLastError();
 154     if (error == ERROR_INVALID_PARAMETER) return 1;
 155     else return -1;
 156   }
 157
 158   return 0;
 159 }
 160
 161 /*
 162   Duplicate an environment block returned by GetEnvironmentStrings().
 163   Since such a block is by definition readonly, and duplicate_environment()
 164   modifies its inputs, this function takes a copy of the input and operates
 165   on that.
 166 */
 167 void duplicate_environment_strings(TCHAR *env) {
 168   TCHAR *newenv = copy_environment_block(env);
 169   if (! newenv) return;
 170
 171   duplicate_environment(newenv);
 172   HeapFree(GetProcessHeap(), 0, newenv);
 173 }