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