env.cpp

   1 #include "nssm.h"\r
   2 \r
   3 /*\r
   4   Environment block is of the form:\r
   5 \r
   6     KEY1=VALUE1 NULL\r
   7     KEY2=VALUE2 NULL\r
   8     NULL\r
   9 \r
  10   A single variable KEY=VALUE has length 15:\r
  11 \r
  12     KEY=VALUE (13) NULL (1)\r
  13     NULL (1)\r
  14 \r
  15   Environment variable names are case-insensitive!\r
  16 */\r
  17 \r
  18 /* Find the length in characters of an environment block. */\r
  19 size_t environment_length(TCHAR *env) {\r
  20   size_t len = 0;\r
  21 \r
  22   TCHAR *s;\r
  23   for (s = env; ; s++) {\r
  24     len++;\r
  25     if (*s == _T('\0')) {\r
  26       if (*(s + 1) == _T('\0')) {\r
  27         len++;\r
  28         break;\r
  29       }\r
  30     }\r
  31   }\r
  32 \r
  33   return len;\r
  34 }\r
  35 \r
  36 /* Copy an environment block. */\r
  37 TCHAR *copy_environment_block(TCHAR *env) {\r
  38   TCHAR *newenv;\r
  39   if (copy_double_null(env, (unsigned long) environment_length(env), &newenv)) return 0;\r
  40   return newenv;\r
  41 }\r
  42 \r
  43 /*\r
  44   The environment block starts with variables of the form\r
  45   =C:=C:\Windows\System32 which we ignore.\r
  46 */\r
  47 TCHAR *useful_environment(TCHAR *rawenv) {\r
  48   TCHAR *env = rawenv;\r
  49 \r
  50   if (env) {\r
  51     while (*env == _T('=')) {\r
  52       for ( ; *env; env++);\r
  53       env++;\r
  54     }\r
  55   }\r
  56 \r
  57   return env;\r
  58 }\r
  59 \r
  60 /* Expand an environment variable.  Must call HeapFree() on the result. */\r
  61 TCHAR *expand_environment_string(TCHAR *string) {\r
  62   unsigned long len;\r
  63 \r
  64   len = ExpandEnvironmentStrings(string, 0, 0);\r
  65   if (! len) {\r
  66     log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_EXPANDENVIRONMENTSTRINGS_FAILED, string, error_string(GetLastError()), 0);\r
  67     return 0;\r
  68   }\r
  69 \r
  70   TCHAR *ret = (TCHAR *) HeapAlloc(GetProcessHeap(), 0, len * sizeof(TCHAR));\r
  71   if (! ret) {\r
  72     log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_OUT_OF_MEMORY, _T("ExpandEnvironmentStrings()"), _T("expand_environment_string"), 0);\r
  73     return 0;\r
  74   }\r
  75 \r
  76   if (! ExpandEnvironmentStrings(string, ret, len)) {\r
  77     log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_EXPANDENVIRONMENTSTRINGS_FAILED, string, error_string(GetLastError()), 0);\r
  78     HeapFree(GetProcessHeap(), 0, ret);\r
  79     return 0;\r
  80   }\r
  81 \r
  82   return ret;\r
  83 }\r
  84 \r
  85 /*\r
  86   Set all the environment variables from an environment block in the current\r
  87   environment or remove all the variables in the block from the current\r
  88   environment.\r
  89 */\r
  90 static int set_environment_block(TCHAR *env, bool set) {\r
  91   int ret = 0;\r
  92 \r
  93   TCHAR *s, *t;\r
  94   for (s = env; *s; s++) {\r
  95     for (t = s; *t && *t != _T('='); t++);\r
  96     if (*t == _T('=')) {\r
  97       *t = _T('\0');\r
  98       if (set) {\r
  99         TCHAR *expanded = expand_environment_string(++t);\r
 100         if (expanded) {\r
 101           if (! SetEnvironmentVariable(s, expanded)) ret++;\r
 102           HeapFree(GetProcessHeap(), 0, expanded);\r
 103         }\r
 104         else {\r
 105           if (! SetEnvironmentVariable(s, t)) ret++;\r
 106         }\r
 107       }\r
 108       else {\r
 109         if (! SetEnvironmentVariable(s, NULL)) ret++;\r
 110       }\r
 111       for (t++; *t; t++);\r
 112     }\r
 113     s = t;\r
 114   }\r
 115 \r
 116   return ret;\r
 117 }\r
 118 \r
 119 int set_environment_block(TCHAR *env) {\r
 120   return set_environment_block(env, true);\r
 121 }\r
 122 \r
 123 static int unset_environment_block(TCHAR *env) {\r
 124   return set_environment_block(env, false);\r
 125 }\r
 126 \r
 127 /* Remove all variables from the process environment. */\r
 128 int clear_environment() {\r
 129   TCHAR *rawenv = GetEnvironmentStrings();\r
 130   TCHAR *env = useful_environment(rawenv);\r
 131 \r
 132   int ret = unset_environment_block(env);\r
 133 \r
 134   if (rawenv) FreeEnvironmentStrings(rawenv);\r
 135 \r
 136   return ret;\r
 137 }\r
 138 \r
 139 /* Set the current environment to exactly duplicate an environment block. */\r
 140 int duplicate_environment(TCHAR *rawenv) {\r
 141   int ret = clear_environment();\r
 142   TCHAR *env = useful_environment(rawenv);\r
 143   ret += set_environment_block(env);\r
 144   return ret;\r
 145 }\r
 146 \r
 147 /*\r
 148   Verify an environment block.\r
 149   Returns:  1 if environment is invalid.\r
 150             0 if environment is OK.\r
 151            -1 on error.\r
 152 */\r
 153 int test_environment(TCHAR *env) {\r
 154   TCHAR *path = (TCHAR *) nssm_imagepath();\r
 155   STARTUPINFO si;\r
 156   ZeroMemory(&si, sizeof(si));\r
 157   si.cb = sizeof(si);\r
 158   PROCESS_INFORMATION pi;\r
 159   ZeroMemory(&pi, sizeof(pi));\r
 160   unsigned long flags = CREATE_SUSPENDED;\r
 161 #ifdef UNICODE\r
 162   flags |= CREATE_UNICODE_ENVIRONMENT;\r
 163 #endif\r
 164 \r
 165   /*\r
 166     Try to relaunch ourselves but with the candidate environment set.\r
 167     Assuming no solar flare activity, the only reason this would fail is if\r
 168     the environment were invalid.\r
 169   */\r
 170   if (CreateProcess(0, path, 0, 0, 0, flags, env, 0, &si, &pi)) {\r
 171     TerminateProcess(pi.hProcess, 0);\r
 172   }\r
 173   else {\r
 174     unsigned long error = GetLastError();\r
 175     if (error == ERROR_INVALID_PARAMETER) return 1;\r
 176     else return -1;\r
 177   }\r
 178 \r
 179   return 0;\r
 180 }\r
 181 \r
 182 /*\r
 183   Duplicate an environment block returned by GetEnvironmentStrings().\r
 184   Since such a block is by definition readonly, and duplicate_environment()\r
 185   modifies its inputs, this function takes a copy of the input and operates\r
 186   on that.\r
 187 */\r
 188 void duplicate_environment_strings(TCHAR *env) {\r
 189   TCHAR *newenv = copy_environment_block(env);\r
 190   if (! newenv) return;\r
 191 \r
 192   duplicate_environment(newenv);\r
 193   HeapFree(GetProcessHeap(), 0, newenv);\r
 194 }\r
 195 \r
 196 /* Safely get a copy of the current environment. */\r
 197 TCHAR *copy_environment() {\r
 198   TCHAR *rawenv = GetEnvironmentStrings();\r
 199   if (! rawenv) return NULL;\r
 200   TCHAR *env = copy_environment_block(rawenv);\r
 201   FreeEnvironmentStrings(rawenv);\r
 202   return env;\r
 203 }\r
 204 \r
 205 /*\r
 206   Create a new block with all the strings of the first block plus a new string.\r
 207   If the key is already present its value will be overwritten in place.\r
 208   If the key is blank or empty the new block will still be allocated and have\r
 209   non-zero length.\r
 210 */\r
 211 int append_to_environment_block(TCHAR *env, unsigned long envlen, TCHAR *string, TCHAR **newenv, unsigned long *newlen) {\r
 212   size_t keylen = 0;\r
 213   if (string && string[0]) {\r
 214     for (; string[keylen]; keylen++) {\r
 215       if (string[keylen] == _T('=')) {\r
 216         keylen++;\r
 217         break;\r
 218       }\r
 219     }\r
 220   }\r
 221   return append_to_double_null(env, envlen, newenv, newlen, string, keylen, false);\r
 222 }\r
 223 \r
 224 /*\r
 225   Create a new block with all the strings of the first block minus the given\r
 226   string.\r
 227   If the key is not present the new block will be a copy of the original.\r
 228   If the string is KEY=VALUE the key will only be removed if its value is\r
 229   VALUE.\r
 230   If the string is just KEY the key will unconditionally be removed.\r
 231   If removing the string results in an empty list the new block will still be\r
 232   allocated and have non-zero length.\r
 233 */\r
 234 int remove_from_environment_block(TCHAR *env, unsigned long envlen, TCHAR *string, TCHAR **newenv, unsigned long *newlen) {\r
 235   if (! string || ! string[0] || string[0] == _T('=')) return 1;\r
 236 \r
 237   TCHAR *key = 0;\r
 238   size_t len = _tcslen(string);\r
 239   size_t i;\r
 240   for (i = 0; i < len; i++) if (string[i] == _T('=')) break;\r
 241 \r
 242   /* Rewrite KEY to KEY= but leave KEY=VALUE alone. */\r
 243   size_t keylen = len;\r
 244   if (i == len) keylen++;\r
 245 \r
 246   key = (TCHAR *) HeapAlloc(GetProcessHeap(), 0, (keylen + 1) * sizeof(TCHAR));\r
 247   if (! key) {\r
 248     log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_OUT_OF_MEMORY, _T("key"), _T("remove_from_environment_block()"), 0);\r
 249     return 2;\r
 250   }\r
 251   memmove(key, string, len * sizeof(TCHAR));\r
 252   if (keylen > len) key[keylen - 1] = _T('=');\r
 253   key[keylen] = _T('\0');\r
 254 \r
 255   int ret = remove_from_double_null(env, envlen, newenv, newlen, key, keylen, false);\r
 256   HeapFree(GetProcessHeap(), 0, key);\r
 257 \r
 258   return ret;\r
 259 }\r