env.cpp

   1 #include "nssm.h"
   2
   3 /* Replace NULL with CRLF. Leave NULL NULL as the end marker. */
   4 int format_environment(TCHAR *env, unsigned long envlen, TCHAR **formatted, unsigned long *newlen) {
   5   unsigned long i, j;
   6   *newlen = envlen;
   7
   8   if (! *newlen) {
   9     *formatted = 0;
  10     return 0;
  11   }
  12
  13   for (i = 0; i < envlen; i++) if (! env[i] && env[i + 1]) ++*newlen;
  14
  15   *formatted = (TCHAR *) HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, *newlen * sizeof(TCHAR));
  16   if (! *formatted) {
  17     *newlen = 0;
  18     return 1;
  19   }
  20
  21   for (i = 0, j = 0; i < envlen; i++) {
  22     (*formatted)[j] = env[i];
  23     if (! env[i]) {
  24       if (env[i + 1]) {
  25         (*formatted)[j] = _T('\r');
  26         (*formatted)[++j] = _T('\n');
  27       }
  28     }
  29     j++;
  30   }
  31
  32   return 0;
  33 }
  34
  35 /* Strip CR and replace LF with NULL. */
  36 int unformat_environment(TCHAR *env, unsigned long envlen, TCHAR **unformatted, unsigned long *newlen) {
  37   unsigned long i, j;
  38   *newlen = 0;
  39
  40   if (! envlen) {
  41     *unformatted = 0;
  42     return 0;
  43   }
  44
  45   for (i = 0; i < envlen; i++) if (env[i] != _T('\r')) ++*newlen;
  46   /* Must end with two NULLs. */
  47   *newlen += 2;
  48
  49   *unformatted = (TCHAR *) HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, *newlen * sizeof(TCHAR));
  50   if (! *unformatted) return 1;
  51
  52   for (i = 0, j = 0; i < envlen; i++) {
  53     if (env[i] == _T('\r')) continue;
  54     if (env[i] == _T('\n')) (*unformatted)[j] = _T('\0');
  55     else (*unformatted)[j] = env[i];
  56     j++;
  57   }
  58
  59   return 0;
  60 }
  61
  62 /*
  63   Verify an environment block.
  64   Returns:  1 if environment is invalid.
  65             0 if environment is OK.
  66            -1 on error.
  67 */
  68 int test_environment(TCHAR *env) {
  69   TCHAR path[PATH_LENGTH];
  70   GetModuleFileName(0, path, _countof(path));
  71   STARTUPINFO si;
  72   ZeroMemory(&si, sizeof(si));
  73   si.cb = sizeof(si);
  74   PROCESS_INFORMATION pi;
  75   ZeroMemory(&pi, sizeof(pi));
  76   unsigned long flags = CREATE_SUSPENDED;
  77 #ifdef UNICODE
  78   flags |= CREATE_UNICODE_ENVIRONMENT;
  79 #endif
  80
  81   /*
  82     Try to relaunch ourselves but with the candidate environment set.
  83     Assuming no solar flare activity, the only reason this would fail is if
  84     the environment were invalid.
  85   */
  86   if (CreateProcess(0, path, 0, 0, 0, flags, env, 0, &si, &pi)) {
  87     TerminateProcess(pi.hProcess, 0);
  88   }
  89   else {
  90     unsigned long error = GetLastError();
  91     if (error == ERROR_INVALID_PARAMETER) return 1;
  92     else return -1;
  93   }
  94
  95   return 0;
  96 }