env.cpp 4.8 KB

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