hook.cpp 16 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410
  1. #include "nssm.h"
  2. typedef struct {
  3. TCHAR *name;
  4. HANDLE process_handle;
  5. unsigned long pid;
  6. unsigned long deadline;
  7. FILETIME creation_time;
  8. kill_t k;
  9. } hook_t;
  10. const TCHAR *hook_event_strings[] = { NSSM_HOOK_EVENT_START, NSSM_HOOK_EVENT_STOP, NSSM_HOOK_EVENT_EXIT, NSSM_HOOK_EVENT_POWER, NSSM_HOOK_EVENT_ROTATE, NULL };
  11. const TCHAR *hook_action_strings[] = { NSSM_HOOK_ACTION_PRE, NSSM_HOOK_ACTION_POST, NSSM_HOOK_ACTION_CHANGE, NSSM_HOOK_ACTION_RESUME, NULL };
  12. static unsigned long WINAPI await_hook(void *arg) {
  13. hook_t *hook = (hook_t *) arg;
  14. if (! hook) return NSSM_HOOK_STATUS_ERROR;
  15. int ret = 0;
  16. if (WaitForSingleObject(hook->process_handle, hook->deadline) == WAIT_TIMEOUT) ret = NSSM_HOOK_STATUS_TIMEOUT;
  17. /* Tidy up hook process tree. */
  18. if (hook->name) hook->k.name = hook->name;
  19. else hook->k.name = _T("hook");
  20. hook->k.process_handle = hook->process_handle;
  21. hook->k.pid = hook->pid;
  22. hook->k.stop_method = ~0;
  23. hook->k.kill_console_delay = NSSM_KILL_CONSOLE_GRACE_PERIOD;
  24. hook->k.kill_window_delay = NSSM_KILL_WINDOW_GRACE_PERIOD;
  25. hook->k.kill_threads_delay = NSSM_KILL_THREADS_GRACE_PERIOD;
  26. hook->k.creation_time = hook->creation_time;
  27. GetSystemTimeAsFileTime(&hook->k.exit_time);
  28. kill_process_tree(&hook->k, hook->pid);
  29. if (ret) {
  30. CloseHandle(hook->process_handle);
  31. if (hook->name) HeapFree(GetProcessHeap(), 0, hook->name);
  32. HeapFree(GetProcessHeap(), 0, hook);
  33. return ret;
  34. }
  35. unsigned long exitcode;
  36. GetExitCodeProcess(hook->process_handle, &exitcode);
  37. CloseHandle(hook->process_handle);
  38. if (hook->name) HeapFree(GetProcessHeap(), 0, hook->name);
  39. HeapFree(GetProcessHeap(), 0, hook);
  40. if (exitcode == NSSM_HOOK_STATUS_ABORT) return NSSM_HOOK_STATUS_ABORT;
  41. if (exitcode) return NSSM_HOOK_STATUS_FAILED;
  42. return NSSM_HOOK_STATUS_SUCCESS;
  43. }
  44. static void set_hook_runtime(TCHAR *v, FILETIME *start, FILETIME *now) {
  45. if (start && now) {
  46. ULARGE_INTEGER s;
  47. s.LowPart = start->dwLowDateTime;
  48. s.HighPart = start->dwHighDateTime;
  49. if (s.QuadPart) {
  50. ULARGE_INTEGER t;
  51. t.LowPart = now->dwLowDateTime;
  52. t.HighPart = now->dwHighDateTime;
  53. if (t.QuadPart && t.QuadPart >= s.QuadPart) {
  54. t.QuadPart -= s.QuadPart;
  55. t.QuadPart /= 10000LL;
  56. TCHAR number[16];
  57. _sntprintf_s(number, _countof(number), _TRUNCATE, _T("%llu"), t.QuadPart);
  58. SetEnvironmentVariable(v, number);
  59. return;
  60. }
  61. }
  62. }
  63. SetEnvironmentVariable(v, _T(""));
  64. }
  65. static void add_thread_handle(hook_thread_t *hook_threads, HANDLE thread_handle, TCHAR *name) {
  66. if (! hook_threads) return;
  67. int num_threads = hook_threads->num_threads + 1;
  68. hook_thread_data_t *data = (hook_thread_data_t *) HeapAlloc(GetProcessHeap(), 0, num_threads * sizeof(hook_thread_data_t));
  69. if (! data) {
  70. log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_OUT_OF_MEMORY, _T("hook_thread_t"), _T("add_thread_handle()"), 0);
  71. return;
  72. }
  73. int i;
  74. for (i = 0; i < hook_threads->num_threads; i++) memmove(&data[i], &hook_threads->data[i], sizeof(data[i]));
  75. memmove(data[i].name, name, sizeof(data[i].name));
  76. data[i].thread_handle = thread_handle;
  77. if (hook_threads->data) HeapFree(GetProcessHeap(), 0, hook_threads->data);
  78. hook_threads->data = data;
  79. hook_threads->num_threads = num_threads;
  80. }
  81. bool valid_hook_name(const TCHAR *hook_event, const TCHAR *hook_action, bool quiet) {
  82. bool valid_event = false;
  83. bool valid_action = false;
  84. /* Exit/Post */
  85. if (str_equiv(hook_event, NSSM_HOOK_EVENT_EXIT)) {
  86. if (str_equiv(hook_action, NSSM_HOOK_ACTION_POST)) return true;
  87. if (quiet) return false;
  88. print_message(stderr, NSSM_MESSAGE_INVALID_HOOK_ACTION, hook_event);
  89. _ftprintf(stderr, _T("%s\n"), NSSM_HOOK_ACTION_POST);
  90. return false;
  91. }
  92. /* Power/{Change,Resume} */
  93. if (str_equiv(hook_event, NSSM_HOOK_EVENT_POWER)) {
  94. if (str_equiv(hook_action, NSSM_HOOK_ACTION_CHANGE)) return true;
  95. if (str_equiv(hook_action, NSSM_HOOK_ACTION_RESUME)) return true;
  96. if (quiet) return false;
  97. print_message(stderr, NSSM_MESSAGE_INVALID_HOOK_ACTION, hook_event);
  98. _ftprintf(stderr, _T("%s\n"), NSSM_HOOK_ACTION_CHANGE);
  99. _ftprintf(stderr, _T("%s\n"), NSSM_HOOK_ACTION_RESUME);
  100. return false;
  101. }
  102. /* Rotate/{Pre,Post} */
  103. if (str_equiv(hook_event, NSSM_HOOK_EVENT_ROTATE)) {
  104. if (str_equiv(hook_action, NSSM_HOOK_ACTION_PRE)) return true;
  105. if (str_equiv(hook_action, NSSM_HOOK_ACTION_POST)) return true;
  106. if (quiet) return false;
  107. print_message(stderr, NSSM_MESSAGE_INVALID_HOOK_ACTION, hook_event);
  108. _ftprintf(stderr, _T("%s\n"), NSSM_HOOK_ACTION_PRE);
  109. _ftprintf(stderr, _T("%s\n"), NSSM_HOOK_ACTION_POST);
  110. return false;
  111. }
  112. /* Start/{Pre,Post} */
  113. if (str_equiv(hook_event, NSSM_HOOK_EVENT_START)) {
  114. if (str_equiv(hook_action, NSSM_HOOK_ACTION_PRE)) return true;
  115. if (str_equiv(hook_action, NSSM_HOOK_ACTION_POST)) return true;
  116. if (quiet) return false;
  117. print_message(stderr, NSSM_MESSAGE_INVALID_HOOK_ACTION, hook_event);
  118. _ftprintf(stderr, _T("%s\n"), NSSM_HOOK_ACTION_PRE);
  119. _ftprintf(stderr, _T("%s\n"), NSSM_HOOK_ACTION_POST);
  120. return false;
  121. }
  122. /* Stop/Pre */
  123. if (str_equiv(hook_event, NSSM_HOOK_EVENT_STOP)) {
  124. if (str_equiv(hook_action, NSSM_HOOK_ACTION_PRE)) return true;
  125. if (quiet) return false;
  126. print_message(stderr, NSSM_MESSAGE_INVALID_HOOK_ACTION, hook_event);
  127. _ftprintf(stderr, _T("%s\n"), NSSM_HOOK_ACTION_PRE);
  128. return false;
  129. }
  130. if (quiet) return false;
  131. print_message(stderr, NSSM_MESSAGE_INVALID_HOOK_EVENT);
  132. _ftprintf(stderr, _T("%s\n"), NSSM_HOOK_EVENT_EXIT);
  133. _ftprintf(stderr, _T("%s\n"), NSSM_HOOK_EVENT_POWER);
  134. _ftprintf(stderr, _T("%s\n"), NSSM_HOOK_EVENT_ROTATE);
  135. _ftprintf(stderr, _T("%s\n"), NSSM_HOOK_EVENT_START);
  136. _ftprintf(stderr, _T("%s\n"), NSSM_HOOK_EVENT_STOP);
  137. return false;
  138. }
  139. void await_hook_threads(hook_thread_t *hook_threads, SERVICE_STATUS_HANDLE status_handle, SERVICE_STATUS *status, unsigned long deadline) {
  140. if (! hook_threads) return;
  141. if (! hook_threads->num_threads) return;
  142. int *retain = (int *) HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, hook_threads->num_threads * sizeof(int));
  143. if (! retain) {
  144. log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_OUT_OF_MEMORY, _T("retain"), _T("await_hook_threads()"), 0);
  145. return;
  146. }
  147. /*
  148. We could use WaitForMultipleObjects() but await_single_object() can update
  149. the service status as well.
  150. */
  151. int num_threads = 0;
  152. int i;
  153. for (i = 0; i < hook_threads->num_threads; i++) {
  154. if (deadline) {
  155. if (await_single_handle(status_handle, status, hook_threads->data[i].thread_handle, hook_threads->data[i].name, _T(__FUNCTION__), deadline) != 1) {
  156. CloseHandle(hook_threads->data[i].thread_handle);
  157. continue;
  158. }
  159. }
  160. else if (WaitForSingleObject(hook_threads->data[i].thread_handle, 0) != WAIT_TIMEOUT) {
  161. CloseHandle(hook_threads->data[i].thread_handle);
  162. continue;
  163. }
  164. retain[num_threads++]= i;
  165. }
  166. if (num_threads) {
  167. hook_thread_data_t *data = (hook_thread_data_t *) HeapAlloc(GetProcessHeap(), 0, num_threads * sizeof(hook_thread_data_t));
  168. if (! data) {
  169. log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_OUT_OF_MEMORY, _T("data"), _T("await_hook_threads()"), 0);
  170. HeapFree(GetProcessHeap(), 0, retain);
  171. return;
  172. }
  173. for (i = 0; i < num_threads; i++) memmove(&data[i], &hook_threads->data[retain[i]], sizeof(data[i]));
  174. HeapFree(GetProcessHeap(), 0, hook_threads->data);
  175. hook_threads->data = data;
  176. hook_threads->num_threads = num_threads;
  177. }
  178. else {
  179. HeapFree(GetProcessHeap(), 0, hook_threads->data);
  180. ZeroMemory(hook_threads, sizeof(*hook_threads));
  181. }
  182. HeapFree(GetProcessHeap(), 0, retain);
  183. }
  184. /*
  185. Returns:
  186. NSSM_HOOK_STATUS_SUCCESS if the hook ran successfully.
  187. NSSM_HOOK_STATUS_NOTFOUND if no hook was found.
  188. NSSM_HOOK_STATUS_ABORT if the hook failed and we should cancel service start.
  189. NSSM_HOOK_STATUS_ERROR on error.
  190. NSSM_HOOK_STATUS_NOTRUN if the hook didn't run.
  191. NSSM_HOOK_STATUS_TIMEOUT if the hook timed out.
  192. NSSM_HOOK_STATUS_FAILED if the hook failed.
  193. */
  194. int nssm_hook(hook_thread_t *hook_threads, nssm_service_t *service, TCHAR *hook_event, TCHAR *hook_action, unsigned long *hook_control, unsigned long deadline, bool async) {
  195. int ret = 0;
  196. hook_t *hook = (hook_t *) HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, sizeof(hook_t));
  197. if (! hook) {
  198. log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_OUT_OF_MEMORY, _T("hook"), _T("nssm_hook()"), 0);
  199. return NSSM_HOOK_STATUS_ERROR;
  200. }
  201. FILETIME now;
  202. GetSystemTimeAsFileTime(&now);
  203. EnterCriticalSection(&service->hook_section);
  204. /* Set the environment. */
  205. set_service_environment(service);
  206. /* ABI version. */
  207. TCHAR number[16];
  208. _sntprintf_s(number, _countof(number), _TRUNCATE, _T("%lu"), NSSM_HOOK_VERSION);
  209. SetEnvironmentVariable(NSSM_HOOK_ENV_VERSION, number);
  210. /* Event triggering this action. */
  211. SetEnvironmentVariable(NSSM_HOOK_ENV_EVENT, hook_event);
  212. /* Hook action. */
  213. SetEnvironmentVariable(NSSM_HOOK_ENV_ACTION, hook_action);
  214. /* Control triggering this action. May be empty. */
  215. if (hook_control) SetEnvironmentVariable(NSSM_HOOK_ENV_TRIGGER, service_control_text(*hook_control));
  216. else SetEnvironmentVariable(NSSM_HOOK_ENV_TRIGGER, _T(""));
  217. /* Last control handled. */
  218. SetEnvironmentVariable(NSSM_HOOK_ENV_LAST_CONTROL, service_control_text(service->last_control));
  219. /* Path to NSSM, unquoted for the environment. */
  220. SetEnvironmentVariable(NSSM_HOOK_ENV_IMAGE_PATH, nssm_unquoted_imagepath());
  221. /* NSSM version. */
  222. SetEnvironmentVariable(NSSM_HOOK_ENV_NSSM_CONFIGURATION, NSSM_CONFIGURATION);
  223. SetEnvironmentVariable(NSSM_HOOK_ENV_NSSM_VERSION, NSSM_VERSION);
  224. SetEnvironmentVariable(NSSM_HOOK_ENV_BUILD_DATE, NSSM_DATE);
  225. /* NSSM PID. */
  226. _sntprintf_s(number, _countof(number), _TRUNCATE, _T("%lu"), GetCurrentProcessId());
  227. SetEnvironmentVariable(NSSM_HOOK_ENV_PID, number);
  228. /* NSSM runtime. */
  229. set_hook_runtime(NSSM_HOOK_ENV_RUNTIME, &service->nssm_creation_time, &now);
  230. /* Application PID. */
  231. if (service->pid) {
  232. _sntprintf_s(number, _countof(number), _TRUNCATE, _T("%lu"), service->pid);
  233. SetEnvironmentVariable(NSSM_HOOK_ENV_APPLICATION_PID, number);
  234. /* Application runtime. */
  235. set_hook_runtime(NSSM_HOOK_ENV_APPLICATION_RUNTIME, &service->creation_time, &now);
  236. /* Exit code. */
  237. SetEnvironmentVariable(NSSM_HOOK_ENV_EXITCODE, _T(""));
  238. }
  239. else {
  240. SetEnvironmentVariable(NSSM_HOOK_ENV_APPLICATION_PID, _T(""));
  241. if (str_equiv(hook_event, NSSM_HOOK_EVENT_START) && str_equiv(hook_action, NSSM_HOOK_ACTION_PRE)) {
  242. SetEnvironmentVariable(NSSM_HOOK_ENV_APPLICATION_RUNTIME, _T(""));
  243. SetEnvironmentVariable(NSSM_HOOK_ENV_EXITCODE, _T(""));
  244. }
  245. else {
  246. set_hook_runtime(NSSM_HOOK_ENV_APPLICATION_RUNTIME, &service->creation_time, &service->exit_time);
  247. /* Exit code. */
  248. _sntprintf_s(number, _countof(number), _TRUNCATE, _T("%lu"), service->exitcode);
  249. SetEnvironmentVariable(NSSM_HOOK_ENV_EXITCODE, number);
  250. }
  251. }
  252. /* Deadline for this script. */
  253. _sntprintf_s(number, _countof(number), _TRUNCATE, _T("%lu"), deadline);
  254. SetEnvironmentVariable(NSSM_HOOK_ENV_DEADLINE, number);
  255. /* Service name. */
  256. SetEnvironmentVariable(NSSM_HOOK_ENV_SERVICE_NAME, service->name);
  257. SetEnvironmentVariable(NSSM_HOOK_ENV_SERVICE_DISPLAYNAME, service->displayname);
  258. /* Times the service was asked to start. */
  259. _sntprintf_s(number, _countof(number), _TRUNCATE, _T("%lu"), service->start_requested_count);
  260. SetEnvironmentVariable(NSSM_HOOK_ENV_START_REQUESTED_COUNT, number);
  261. /* Times the service actually did start. */
  262. _sntprintf_s(number, _countof(number), _TRUNCATE, _T("%lu"), service->start_count);
  263. SetEnvironmentVariable(NSSM_HOOK_ENV_START_COUNT, number);
  264. /* Times the service exited. */
  265. _sntprintf_s(number, _countof(number), _TRUNCATE, _T("%lu"), service->exit_count);
  266. SetEnvironmentVariable(NSSM_HOOK_ENV_EXIT_COUNT, number);
  267. /* Throttled count. */
  268. _sntprintf_s(number, _countof(number), _TRUNCATE, _T("%lu"), service->throttle);
  269. SetEnvironmentVariable(NSSM_HOOK_ENV_THROTTLE_COUNT, number);
  270. /* Command line. */
  271. TCHAR app[CMD_LENGTH];
  272. _sntprintf_s(app, _countof(app), _TRUNCATE, _T("\"%s\" %s"), service->exe, service->flags);
  273. SetEnvironmentVariable(NSSM_HOOK_ENV_COMMAND_LINE, app);
  274. TCHAR cmd[CMD_LENGTH];
  275. if (get_hook(service->name, hook_event, hook_action, cmd, sizeof(cmd))) {
  276. log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_GET_HOOK_FAILED, hook_event, hook_action, service->name, 0);
  277. unset_service_environment(service);
  278. LeaveCriticalSection(&service->hook_section);
  279. HeapFree(GetProcessHeap(), 0, hook);
  280. return NSSM_HOOK_STATUS_ERROR;
  281. }
  282. /* No hook. */
  283. if (! _tcslen(cmd)) {
  284. unset_service_environment(service);
  285. LeaveCriticalSection(&service->hook_section);
  286. HeapFree(GetProcessHeap(), 0, hook);
  287. return NSSM_HOOK_STATUS_NOTFOUND;
  288. }
  289. /* Run the command. */
  290. STARTUPINFO si;
  291. ZeroMemory(&si, sizeof(si));
  292. si.cb = sizeof(si);
  293. PROCESS_INFORMATION pi;
  294. ZeroMemory(&pi, sizeof(pi));
  295. if (service->hook_share_output_handles) (void) use_output_handles(service, &si);
  296. bool inherit_handles = false;
  297. if (si.dwFlags & STARTF_USESTDHANDLES) inherit_handles = true;
  298. unsigned long flags = 0;
  299. #ifdef UNICODE
  300. flags |= CREATE_UNICODE_ENVIRONMENT;
  301. #endif
  302. ret = NSSM_HOOK_STATUS_NOTRUN;
  303. if (CreateProcess(0, cmd, 0, 0, inherit_handles, flags, 0, service->dir, &si, &pi)) {
  304. close_output_handles(&si);
  305. hook->name = (TCHAR *) HeapAlloc(GetProcessHeap(), 0, HOOK_NAME_LENGTH * sizeof(TCHAR));
  306. if (hook->name) _sntprintf_s(hook->name, HOOK_NAME_LENGTH, _TRUNCATE, _T("%s (%s/%s)"), service->name, hook_event, hook_action);
  307. hook->process_handle = pi.hProcess;
  308. hook->pid = pi.dwProcessId;
  309. hook->deadline = deadline;
  310. if (get_process_creation_time(hook->process_handle, &hook->creation_time)) GetSystemTimeAsFileTime(&hook->creation_time);
  311. unsigned long tid;
  312. HANDLE thread_handle = CreateThread(NULL, 0, await_hook, (void *) hook, 0, &tid);
  313. if (thread_handle) {
  314. if (async) {
  315. ret = 0;
  316. await_hook_threads(hook_threads, service->status_handle, &service->status, 0);
  317. add_thread_handle(hook_threads, thread_handle, hook->name);
  318. }
  319. else {
  320. await_single_handle(service->status_handle, &service->status, thread_handle, hook->name, _T(__FUNCTION__), deadline + NSSM_SERVICE_STATUS_DEADLINE);
  321. unsigned long exitcode;
  322. GetExitCodeThread(thread_handle, &exitcode);
  323. ret = (int) exitcode;
  324. CloseHandle(thread_handle);
  325. }
  326. }
  327. else {
  328. log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_CREATETHREAD_FAILED, error_string(GetLastError()), 0);
  329. await_hook(hook);
  330. if (hook->name) HeapFree(GetProcessHeap(), 0, hook->name);
  331. HeapFree(GetProcessHeap(), 0, hook);
  332. }
  333. }
  334. else {
  335. log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_HOOK_CREATEPROCESS_FAILED, hook_event, hook_action, service->name, cmd, error_string(GetLastError()), 0);
  336. HeapFree(GetProcessHeap(), 0, hook);
  337. close_output_handles(&si);
  338. }
  339. /* Restore our environment. */
  340. unset_service_environment(service);
  341. LeaveCriticalSection(&service->hook_section);
  342. return ret;
  343. }
  344. int nssm_hook(hook_thread_t *hook_threads, nssm_service_t *service, TCHAR *hook_event, TCHAR *hook_action, unsigned long *hook_control, unsigned long deadline) {
  345. return nssm_hook(hook_threads, service, hook_event, hook_action, hook_control, deadline, true);
  346. }
  347. int nssm_hook(hook_thread_t *hook_threads, nssm_service_t *service, TCHAR *hook_event, TCHAR *hook_action, unsigned long *hook_control) {
  348. return nssm_hook(hook_threads, service, hook_event, hook_action, hook_control, NSSM_HOOK_DEADLINE);
  349. }