From db6b51fb99c108fb66cc4db9da1bdc58a1e409c9 Mon Sep 17 00:00:00 2001
From: Aditya Mandaleeka <adityam@microsoft.com>
Date: Thu, 3 Aug 2023 20:19:47 -0700
Subject: [PATCH] Make stdout/stderr redirection thread timeout configurable.

---
 .../CommonLib/StandardStreamRedirection.cpp   | 32 +++++++++++++++----
 .../CommonLib/StandardStreamRedirection.h     | 13 ++++++--
 2 files changed, 35 insertions(+), 10 deletions(-)

diff --git a/src/Servers/IIS/AspNetCoreModuleV2/CommonLib/StandardStreamRedirection.cpp b/src/Servers/IIS/AspNetCoreModuleV2/CommonLib/StandardStreamRedirection.cpp
index 6629ea0d6d45..4abb04f85b93 100644
--- a/src/Servers/IIS/AspNetCoreModuleV2/CommonLib/StandardStreamRedirection.cpp
+++ b/src/Servers/IIS/AspNetCoreModuleV2/CommonLib/StandardStreamRedirection.cpp
@@ -9,10 +9,7 @@
 #include "StdWrapper.h"
 #include "ntassert.h"
 #include "StringHelpers.h"
-
-#define LOG_IF_DUPFAIL(err) do { if (err == -1) { LOG_IF_FAILED(HRESULT_FROM_WIN32(_doserrno)); } } while (0, 0);
-#define LOG_IF_ERRNO(err) do { if (err != 0) { LOG_IF_FAILED(HRESULT_FROM_WIN32(_doserrno)); } } while (0, 0);
-
+#include "Environment.h"
 
 StandardStreamRedirection::StandardStreamRedirection(RedirectionOutput& output, bool commandLineLaunch) :
     m_output(output),
@@ -23,6 +20,28 @@ StandardStreamRedirection::StandardStreamRedirection(RedirectionOutput& output,
     m_commandLineLaunch(commandLineLaunch)
 {
     TryStartRedirection();
+
+    // Allow users to override the default termination timeout for the redirection thread.
+    auto timeoutMsStr = Environment::GetEnvironmentVariableValue(L"ASPNETCORE_OUTPUT_REDIRECTION_TERMINATION_TIMEOUT_MS");
+    if (timeoutMsStr.has_value())
+    {
+        try
+        {
+            int timeoutMs = std::stoi(timeoutMsStr.value());
+            if (timeoutMs > 0 && timeoutMs <= PIPE_OUTPUT_THREAD_TIMEOUT_MS_MAX)
+            {
+                m_terminationTimeoutMs = timeoutMs;
+            }
+            else
+            {
+                LOG_WARN(L"ASPNETCORE_OUTPUT_REDIRECTION_TERMINATION_TIMEOUT_MS must be an integer between 0 and 1800000. Ignoring.");
+            }
+        }
+        catch (...)
+        {
+            LOG_WARN(L"ASPNETCORE_OUTPUT_REDIRECTION_TERMINATION_TIMEOUT_MS must be an integer between 0 and 1800000. Ignoring.");
+        }
+    }
 }
 
 StandardStreamRedirection::~StandardStreamRedirection() noexcept(false)
@@ -80,8 +99,6 @@ void StandardStreamRedirection::Start()
 // be thrown away.
 void StandardStreamRedirection::Stop()
 {
-    DWORD    dwThreadStatus = 0;
-
     if (m_disposed)
     {
         return;
@@ -126,12 +143,13 @@ void StandardStreamRedirection::Stop()
     }
 
     // GetExitCodeThread returns 0 on failure; thread status code is invalid.
+    DWORD dwThreadStatus = 0;
     if (m_hErrThread != nullptr &&
         !LOG_LAST_ERROR_IF(!GetExitCodeThread(m_hErrThread, &dwThreadStatus)) &&
         dwThreadStatus == STILL_ACTIVE)
     {
         // Wait for graceful shutdown, i.e., the exit of the background thread or timeout
-        if (WaitForSingleObject(m_hErrThread, PIPE_OUTPUT_THREAD_TIMEOUT) != WAIT_OBJECT_0)
+        if (WaitForSingleObject(m_hErrThread, m_terminationTimeoutMs) != WAIT_OBJECT_0)
         {
             // If the thread is still running, we need kill it first before exit to avoid AV
             if (!LOG_LAST_ERROR_IF(GetExitCodeThread(m_hErrThread, &dwThreadStatus) == 0) &&
diff --git a/src/Servers/IIS/AspNetCoreModuleV2/CommonLib/StandardStreamRedirection.h b/src/Servers/IIS/AspNetCoreModuleV2/CommonLib/StandardStreamRedirection.h
index 40250933f611..d1cfe8556111 100644
--- a/src/Servers/IIS/AspNetCoreModuleV2/CommonLib/StandardStreamRedirection.h
+++ b/src/Servers/IIS/AspNetCoreModuleV2/CommonLib/StandardStreamRedirection.h
@@ -9,9 +9,15 @@
 
 class StandardStreamRedirection : NonCopyable
 {
-    // Timeout to be used if a thread never exits
-    #define PIPE_OUTPUT_THREAD_TIMEOUT 2000
-    #define PIPE_READ_SIZE 4096
+    // Default timeout for the redirection thread to exit before it is forcefully terminated
+    // This can be overridden with ASPNETCORE_OUTPUT_REDIRECTION_TERMINATION_TIMEOUT_MS
+    static constexpr int PIPE_OUTPUT_THREAD_TIMEOUT_MS_DEFAULT = 2000;
+
+    // Maximum allowed timeout value
+    static constexpr int PIPE_OUTPUT_THREAD_TIMEOUT_MS_MAX = 1800000; // 30 minutes
+
+    // Size of the buffer used to read from the pipe
+    static constexpr int PIPE_READ_SIZE = 4096;
 
 public:
     StandardStreamRedirection(RedirectionOutput& output, bool commandLineLaunch);
@@ -63,4 +69,5 @@ class StandardStreamRedirection : NonCopyable
     std::unique_ptr<StdWrapper> stdoutWrapper;
     std::unique_ptr<StdWrapper> stderrWrapper;
     RedirectionOutput& m_output;
+    int m_terminationTimeoutMs = PIPE_OUTPUT_THREAD_TIMEOUT_MS_DEFAULT;
 };