Win: Use SetConsoleCtrlHandler for SIGINT/SIGTERM

3 files changed, 114 insertions, 7 deletions

diff --git a/src/DownloadEngine.cc b/src/DownloadEngine.cc
index 348efa82..3b70364b 100644
--- a/src/DownloadEngine.cc
+++ b/src/DownloadEngine.cc
@@ -83,6 +83,7 @@ namespace global {
 // 2 ... stop signal processed by DownloadEngine
 // 3 ... 2nd stop signal(force shutdown) detected
 // 4 ... 2nd stop signal processed by DownloadEngine
+// 5 ... main loop exited
 volatile sig_atomic_t globalHaltRequested = 0;
 
 } // namespace global
@@ -141,6 +142,13 @@ void executeCommand(std::deque<std::unique_ptr<Command>>& commands,
 
 int DownloadEngine::run(bool oneshot)
 {
+  class GHR {
+    public:
+      ~GHR() {
+        global::globalHaltRequested = 5;
+      }
+  } ghr;
+
   while(!commands_.empty() || !routineCommands_.empty()) {
     if(!commands_.empty()) {
       waitData();
diff --git a/src/MultiUrlRequestInfo.cc b/src/MultiUrlRequestInfo.cc
index 7eb171bf..8f54c4e3 100644
--- a/src/MultiUrlRequestInfo.cc
+++ b/src/MultiUrlRequestInfo.cc
@@ -86,21 +86,44 @@ extern volatile sig_atomic_t globalHaltRequested;
 } // namespace global
 
 namespace {
-void handler(int signal) {
+
+#ifdef _WIN32
+static const DWORD mainThread = GetCurrentThreadId();
+#endif
+
+static void handler(int signal)
+{
   if(
 #ifdef SIGHUP
      signal == SIGHUP ||
 #endif // SIGHUP
      signal == SIGTERM) {
-    if(global::globalHaltRequested == 0 || global::globalHaltRequested == 2) {
+    if(global::globalHaltRequested <= 2) {
       global::globalHaltRequested = 3;
     }
-  } else {
-    if(global::globalHaltRequested == 0) {
-      global::globalHaltRequested = 1;
-    } else if(global::globalHaltRequested == 2) {
-      global::globalHaltRequested = 3;
+#ifdef _WIN32
+    if (::GetCurrentThreadId() != mainThread) {
+      // SIGTERM may arrive on another thread (via SetConsoleCtrlHandler), and
+      // the process will be forcefully terminated as soon as that thread is
+      // done. So better make sure it isn't done prematurely. ;)
+      while (global::globalHaltRequested != 5) {
+        ::Sleep(100); // Yeah, semi-busy waiting for now.
+      }
     }
+#endif
+    return;
+  }
+
+  // SIGINT
+
+  if (global::globalHaltRequested == 0) {
+    global::globalHaltRequested = 1;
+    return;
+  }
+
+  if (global::globalHaltRequested == 2) {
+    global::globalHaltRequested = 3;
+    return;
   }
 }
 } // namespace
diff --git a/src/util.cc b/src/util.cc
index a671447d..b238b345 100644
--- a/src/util.cc
+++ b/src/util.cc
@@ -84,6 +84,7 @@
 #include "BufferedFile.h"
 #include "SocketCore.h"
 #include "prefs.h"
+#include "Lock.h"
 
 #ifdef ENABLE_MESSAGE_DIGEST
 # include "MessageDigest.h"
@@ -1232,8 +1233,83 @@ bool isNumericHost(const std::string& name)
   return true;
 }
 
+#if _WIN32
+namespace {
+  static Lock win_signal_lock;
+
+  static void(*win_int_handler)(int) = nullptr;
+  static void(*win_term_handler)(int) = nullptr;
+
+  static void win_ign_handler(int) {}
+
+  static BOOL WINAPI HandlerRoutine(DWORD ctrlType)
+  {
+    void(*handler)(int) = nullptr;
+    switch (ctrlType) {
+      case CTRL_C_EVENT:
+      case CTRL_BREAK_EVENT:
+        {
+          // Handler will be called on a new/different thread.
+          LockGuard lg(win_signal_lock);
+          handler = win_int_handler;
+        }
+
+        if (handler) {
+          handler(SIGINT);
+          return TRUE;
+        }
+        return FALSE;
+
+      case CTRL_LOGOFF_EVENT:
+      case CTRL_CLOSE_EVENT:
+      case CTRL_SHUTDOWN_EVENT:
+        {
+          // Handler will be called on a new/different thread.
+          LockGuard lg(win_signal_lock);
+          handler = win_term_handler;;
+        }
+        if (handler) {
+          handler(SIGTERM);
+          return TRUE;
+        }
+        return FALSE;
+    }
+    return FALSE;
+  }
+}
+#endif
+
 void setGlobalSignalHandler(int sig, sigset_t* mask, void (*handler)(int),
                             int flags) {
+#if _WIN32
+  if (sig == SIGINT || sig == SIGTERM) {
+    // Handler will be called on a new/different thread.
+    LockGuard lg(win_signal_lock);
+
+    if (handler == SIG_DFL) {
+      handler = nullptr;
+    }
+    else if (handler == SIG_IGN) {
+      handler = win_ign_handler;
+    }
+    // Not yet in use: add console handler.
+    if (handler && !win_int_handler && !win_term_handler) {
+      ::SetConsoleCtrlHandler(HandlerRoutine, TRUE);
+    }
+    if (sig == SIGINT) {
+      win_int_handler = handler;
+    }
+    else {
+      win_term_handler = handler;
+    }
+    // No handlers set: remove.
+    if (!win_int_handler && !win_term_handler) {
+      ::SetConsoleCtrlHandler(HandlerRoutine, FALSE);
+    }
+    return;
+  }
+#endif
+
 #ifdef HAVE_SIGACTION
   struct sigaction sigact;
   sigact.sa_handler = handler;