env.cpp

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