env.cpp 4.9 KB

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