Added get_debug_token().

[nssm.git] / process.cpp

diff --git a/process.cpp b/process.cpp
index d1a9a49..220411b 100644 (file)
--- a/process.cpp
+++ b/process.cpp
@@ -2,6 +2,47 @@
 \r
 extern imports_t imports;\r
 \r
+HANDLE get_debug_token() {\r
+  long error;\r
+  HANDLE token;\r
+  if (! OpenThreadToken(GetCurrentThread(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, false, &token)) {\r
+    error = GetLastError();\r
+    if (error == ERROR_NO_TOKEN) {\r
+      (void) ImpersonateSelf(SecurityImpersonation);\r
+      (void) OpenThreadToken(GetCurrentThread(), TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, false, &token);\r
+    }\r
+  }\r
+  if (! token) return INVALID_HANDLE_VALUE;\r
+\r
+  TOKEN_PRIVILEGES privileges, old;\r
+  unsigned long size = sizeof(TOKEN_PRIVILEGES);\r
+  LUID luid;\r
+  if (! LookupPrivilegeValue(NULL, SE_DEBUG_NAME, &luid)) {\r
+    CloseHandle(token);\r
+    return INVALID_HANDLE_VALUE;\r
+  }\r
+\r
+  privileges.PrivilegeCount = 1;\r
+  privileges.Privileges[0].Luid = luid;\r
+  privileges.Privileges[0].Attributes = 0;\r
+\r
+  if (! AdjustTokenPrivileges(token, false, &privileges, size, &old, &size)) {\r
+    CloseHandle(token);\r
+    return INVALID_HANDLE_VALUE;\r
+  }\r
+\r
+  old.PrivilegeCount = 1;\r
+  old.Privileges[0].Luid = luid;\r
+  old.Privileges[0].Attributes |= SE_PRIVILEGE_ENABLED;\r
+\r
+  if (! AdjustTokenPrivileges(token, false, &old, size, NULL, NULL)) {\r
+    CloseHandle(token);\r
+    return INVALID_HANDLE_VALUE;\r
+  }\r
+\r
+  return token;\r
+}\r
+\r
 void service_kill_t(nssm_service_t *service, kill_t *k) {\r
   if (! service) return;\r
   if (! k) return;\r
@@ -115,7 +156,7 @@ int kill_threads(nssm_service_t *service, kill_t *k) {
 \r
   /* Get a snapshot of all threads in the system. */\r
   HANDLE snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);\r
-  if (! snapshot) {\r
+  if (snapshot == INVALID_HANDLE_VALUE) {\r
     log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_CREATETOOLHELP32SNAPSHOT_THREAD_FAILED, k->name, error_string(GetLastError()), 0);\r
     return 0;\r
   }\r
@@ -241,7 +282,8 @@ int kill_console(nssm_service_t *service, kill_t *k) {
 \r
   /* Ignore the event ourselves. */\r
   ret = 0;\r
-  if (! SetConsoleCtrlHandler(0, TRUE)) {\r
+  BOOL ignored = SetConsoleCtrlHandler(0, TRUE);\r
+  if (! ignored) {\r
     log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_SETCONSOLECTRLHANDLER_FAILED, k->name, error_string(GetLastError()), 0);\r
     ret = 4;\r
   }\r
@@ -262,6 +304,11 @@ int kill_console(nssm_service_t *service, kill_t *k) {
   /* Wait for process to exit. */\r
   if (await_single_handle(k->status_handle, k->status, k->process_handle, k->name, _T(__FUNCTION__), k->kill_console_delay)) ret = 6;\r
 \r
+  /* Remove our handler. */\r
+  if (ignored && ! SetConsoleCtrlHandler(0, FALSE)) {\r
+    log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_SETCONSOLECTRLHANDLER_FAILED, k->name, error_string(GetLastError()), 0);\r
+  }\r
+\r
   return ret;\r
 }\r
 \r
@@ -269,7 +316,7 @@ int kill_console(kill_t *k) {
   return kill_console(NULL, k);\r
 }\r
 \r
-void kill_process_tree(nssm_service_t * service, kill_t *k, unsigned long ppid) {\r
+void walk_process_tree(nssm_service_t *service, walk_function_t fn, kill_t *k, unsigned long ppid) {\r
   if (! k) return;\r
   /* Shouldn't happen unless the service failed to start. */\r
   if (! k->pid) return; /* XXX: needed? */\r
@@ -288,7 +335,7 @@ void kill_process_tree(nssm_service_t * service, kill_t *k, unsigned long ppid)
     _sntprintf_s(ppid_string, _countof(ppid_string), _TRUNCATE, _T("%lu"), ppid);\r
     log_event(EVENTLOG_INFORMATION_TYPE, NSSM_EVENT_KILL_PROCESS_TREE, pid_string, ppid_string, k->name, 0);\r
     k->process_handle = process_handle; /* XXX: open directly? */\r
-    if (! kill_process(k)) {\r
+    if (! fn(service, k)) {\r
       /* Maybe it already died. */\r
       unsigned long ret;\r
       if (! GetExitCodeProcess(process_handle, &ret) || ret == STILL_ACTIVE) {\r
@@ -303,7 +350,7 @@ void kill_process_tree(nssm_service_t * service, kill_t *k, unsigned long ppid)
 \r
   /* Get a snapshot of all processes in the system. */\r
   HANDLE snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);\r
-  if (! snapshot) {\r
+  if (snapshot == INVALID_HANDLE_VALUE) {\r
     log_event(EVENTLOG_ERROR_TYPE, NSSM_EVENT_CREATETOOLHELP32SNAPSHOT_PROCESS_FAILED, k->name, error_string(GetLastError()), 0);\r
     return;\r
   }\r
@@ -321,7 +368,7 @@ void kill_process_tree(nssm_service_t * service, kill_t *k, unsigned long ppid)
   /* This is a child of the doomed process so kill it. */\r
   if (! check_parent(k, &pe, pid)) {\r
     k->pid = pe.th32ProcessID;\r
-    kill_process_tree(k, ppid);\r
+    walk_process_tree(service, fn, k, ppid);\r
   }\r
   k->pid = pid;\r
 \r
@@ -337,7 +384,7 @@ void kill_process_tree(nssm_service_t * service, kill_t *k, unsigned long ppid)
 \r
     if (! check_parent(k, &pe, pid)) {\r
       k->pid = pe.th32ProcessID;\r
-      kill_process_tree(k, ppid);\r
+      walk_process_tree(service, fn, k, ppid);\r
     }\r
     k->pid = pid;\r
   }\r
@@ -346,5 +393,5 @@ void kill_process_tree(nssm_service_t * service, kill_t *k, unsigned long ppid)
 }\r
 \r
 void kill_process_tree(kill_t *k, unsigned long ppid) {\r
-  return kill_process_tree(NULL, k, ppid);\r
+  return walk_process_tree(NULL, kill_process, k, ppid);\r
 }\r