hook.cpp

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