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 /*
 108   Verify an environment block.
 109   Returns:  1 if environment is invalid.
 110             0 if environment is OK.
 111            -1 on error.
 112 */
 113 int test_environment(TCHAR *env) {
 114   TCHAR path[PATH_LENGTH];
 115   GetModuleFileName(0, path, _countof(path));
 116   STARTUPINFO si;
 117   ZeroMemory(&si, sizeof(si));
 118   si.cb = sizeof(si);
 119   PROCESS_INFORMATION pi;
 120   ZeroMemory(&pi, sizeof(pi));
 121   unsigned long flags = CREATE_SUSPENDED;
 122 #ifdef UNICODE
 123   flags |= CREATE_UNICODE_ENVIRONMENT;
 124 #endif
 125
 126   /*
 127     Try to relaunch ourselves but with the candidate environment set.
 128     Assuming no solar flare activity, the only reason this would fail is if
 129     the environment were invalid.
 130   */
 131   if (CreateProcess(0, path, 0, 0, 0, flags, env, 0, &si, &pi)) {
 132     TerminateProcess(pi.hProcess, 0);
 133   }
 134   else {
 135     unsigned long error = GetLastError();
 136     if (error == ERROR_INVALID_PARAMETER) return 1;
 137     else return -1;
 138   }
 139
 140   return 0;
 141 }