env.cpp

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