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