From c78ac777682cd5a2a9a99f943b42d458e7d25581 Mon Sep 17 00:00:00 2001
From: jw-msft <84477130+jw-msft@users.noreply.github.com>
Date: Thu, 20 Feb 2025 17:55:14 +0000
Subject: [PATCH] add signal handler to adu-shell for restart race

---
 src/adu-shell/inc/adushell.hpp |  7 ++++++-
 src/adu-shell/src/main.cpp     | 29 ++++++++++++++++++++++++++++-
 2 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/src/adu-shell/inc/adushell.hpp b/src/adu-shell/inc/adushell.hpp
index 7b09759f8..e9532ba13 100644
--- a/src/adu-shell/inc/adushell.hpp
+++ b/src/adu-shell/inc/adushell.hpp
@@ -33,11 +33,16 @@ typedef struct tagADUShell_LaunchArguments
 } ADUShell_LaunchArguments;
 
 /**
- * @brief A exit code from an ADUShell child process.
+ * @brief An exit code from an ADUShell child process.
  * We're using EXIT_SUCCESS (0) for success case and EXIT_FAILURE (1) for general errors.
  */
 #define ADUSHELL_EXIT_UNSUPPORTED 3
 
+/**
+ * @brief An exit code for early exit of adu-shell process due to SIGTERM/SIGINT signal handling.
+ */
+#define ADUSHELL_EXIT_SIGNAL_HANDLING 4
+
 /**
  * @brief A result from ADUShell task.
  */
diff --git a/src/adu-shell/src/main.cpp b/src/adu-shell/src/main.cpp
index ea652aca1..cf1f3bd5a 100644
--- a/src/adu-shell/src/main.cpp
+++ b/src/adu-shell/src/main.cpp
@@ -6,7 +6,7 @@
  * Licensed under the MIT License.
  */
 #include <getopt.h>
-
+#include <signal.h> // signal
 #include <string.h>
 
 #include <aducpal/grp.h> // getgrnam
@@ -43,6 +43,11 @@ namespace ScriptTasks = Adu::Shell::Tasks::Script;
 
 namespace adushconst = Adu::Shell::Const;
 
+/**
+ * @brief Indicates whether a task is in progress or not. SIGTERM/SIGINT signal handling should not exit process if true.
+ */
+static bool s_task_in_progress = false;
+
 /**
  * @brief Parse command-line arguments.
  * @param argc arguments count.
@@ -301,6 +306,24 @@ bool ADUShell_PermissionCheck()
     return isTrusted;
 }
 
+/**
+ * @brief Called when a terminate (SIGINT, SIGTERM) signal is detected.
+ *
+ * @param sig Signal value.
+ */
+void OnSignal(int sig)
+{
+    if (s_task_in_progress)
+    {
+        Log_Warn("Task already in progress. Ignoring SIG: %d", sig);
+        return;
+    }
+
+    // Assume handling SIGINT or SIGTERM at this point.
+    Log_Warn("Exit process for SIG: %d", sig);
+    exit(ADUSHELL_EXIT_SIGNAL_HANDLING);
+}
+
 /**
  * @brief Main method.
  *
@@ -333,6 +356,9 @@ int main(int argc, char** argv)
         goto done;
     }
 
+    signal(SIGINT, OnSignal);
+    signal(SIGTERM, OnSignal);
+
     ADUCPAL_setenv(ADUC_CONFIG_FOLDER_ENV, launchArgs.configFolder, 1);
     config = ADUC_ConfigInfo_GetInstance();
     if (config == NULL)
@@ -371,6 +397,7 @@ int main(int argc, char** argv)
             effectiveUserId,
             ADUCPAL_getegid());
 
+        s_task_in_progress = true;
         ret = ADUShell_Dowork(launchArgs);
 
         ADUC_Logging_Uninit();