From 3df88ec335da7c3ed1625014d6f9b05fd5b54f31 Mon Sep 17 00:00:00 2001
From: Andy Yankovsky <weratt@gmail.com>
Date: Mon, 7 Feb 2022 20:37:38 +0000
Subject: [PATCH] [Support] Don't print stacktrace if DbgHelp.dll hasn't been
 loaded yet

On Windows certain function from `Signals.h` require that `DbgHelp.dll` is loaded. This typically happens when the main program calls `llvm::InitLLVM`, however in some cases main program doesn't do that (e.g. when the application is using LLDB via `liblldb.dll`). This patch adds a safe guard to prevent crashes. More discussion in
https://reviews.llvm.org/D119009.

Reviewed By: aganea

Differential Revision: https://reviews.llvm.org/D119181
---
 llvm/lib/Support/Windows/Signals.inc          | 12 +++++++++++-
 llvm/unittests/Support/ProgramTest.cpp        | 25 +++++++++++++++++++++++++
 llvm/utils/unittest/UnitTestMain/TestMain.cpp |  9 +++++++--
 3 files changed, 43 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Support/Windows/Signals.inc b/llvm/lib/Support/Windows/Signals.inc
index 32186bb..fab3886 100644
--- a/llvm/lib/Support/Windows/Signals.inc
+++ b/llvm/lib/Support/Windows/Signals.inc
@@ -159,6 +159,10 @@ static fpSymInitialize fSymInitialize;
 typedef BOOL (WINAPI *fpEnumerateLoadedModules)(HANDLE,PENUMLOADED_MODULES_CALLBACK64,PVOID);
 static fpEnumerateLoadedModules fEnumerateLoadedModules;
 
+static bool isDebugHelpInitialized() {
+  return fStackWalk64 && fSymInitialize && fSymSetOptions && fMiniDumpWriteDump;
+}
+
 static bool load64BitDebugHelp(void) {
   HMODULE hLib = ::LoadLibraryW(L"Dbghelp.dll");
   if (hLib) {
@@ -181,7 +185,7 @@ static bool load64BitDebugHelp(void) {
     fEnumerateLoadedModules = (fpEnumerateLoadedModules)
       ::GetProcAddress(hLib, "EnumerateLoadedModules64");
   }
-  return fStackWalk64 && fSymInitialize && fSymSetOptions && fMiniDumpWriteDump;
+  return isDebugHelpInitialized();
 }
 
 using namespace llvm;
@@ -296,6 +300,12 @@ static bool findModulesAndOffsets(void **StackTrace, int Depth,
 static void PrintStackTraceForThread(llvm::raw_ostream &OS, HANDLE hProcess,
                                      HANDLE hThread, STACKFRAME64 &StackFrame,
                                      CONTEXT *Context) {
+  // It's possible that DbgHelp.dll hasn't been loaded yet (e.g. if this
+  // function is called before the main program called `llvm::InitLLVM`).
+  // In this case just return, not stacktrace will be printed.
+  if (!isDebugHelpInitialized())
+    return;
+
   // Initialize the symbol handler.
   fSymSetOptions(SYMOPT_DEFERRED_LOADS | SYMOPT_LOAD_LINES);
   fSymInitialize(hProcess, NULL, TRUE);
diff --git a/llvm/unittests/Support/ProgramTest.cpp b/llvm/unittests/Support/ProgramTest.cpp
index fbbcd84..fdd0478 100644
--- a/llvm/unittests/Support/ProgramTest.cpp
+++ b/llvm/unittests/Support/ProgramTest.cpp
@@ -12,6 +12,7 @@
 #include "llvm/Support/ConvertUTF.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/Path.h"
+#include "llvm/Support/Signals.h"
 #include "gtest/gtest.h"
 #include <stdlib.h>
 #include <thread>
@@ -426,4 +427,28 @@ TEST_F(ProgramEnvTest, TestLockFile) {
   sys::fs::remove(LockedFile);
 }
 
+TEST_F(ProgramEnvTest, TestExecuteWithNoStacktraceHandler) {
+  using namespace llvm::sys;
+
+  if (getenv("LLVM_PROGRAM_TEST_NO_STACKTRACE_HANDLER")) {
+    sys::PrintStackTrace(errs());
+    exit(0);
+  }
+
+  std::string Executable =
+      sys::fs::getMainExecutable(TestMainArgv0, &ProgramTestStringArg1);
+  StringRef argv[] = {
+      Executable,
+      "--gtest_filter=ProgramEnvTest.TestExecuteWithNoStacktraceHandler"};
+
+  addEnvVar("LLVM_PROGRAM_TEST_NO_STACKTRACE_HANDLER=1");
+
+  std::string Error;
+  bool ExecutionFailed;
+  int RetCode = ExecuteAndWait(Executable, argv, getEnviron(), {}, 0, 0, &Error,
+                               &ExecutionFailed);
+  EXPECT_FALSE(ExecutionFailed) << Error;
+  ASSERT_EQ(0, RetCode);
+}
+
 } // end anonymous namespace
diff --git a/llvm/utils/unittest/UnitTestMain/TestMain.cpp b/llvm/utils/unittest/UnitTestMain/TestMain.cpp
index 1be819f..35ba72b 100644
--- a/llvm/utils/unittest/UnitTestMain/TestMain.cpp
+++ b/llvm/utils/unittest/UnitTestMain/TestMain.cpp
@@ -10,6 +10,7 @@
 #include "llvm/Support/Signals.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
+#include <stdlib.h>
 
 #if defined(_WIN32)
 # include <windows.h>
@@ -21,8 +22,12 @@
 const char *TestMainArgv0;
 
 int main(int argc, char **argv) {
-  llvm::sys::PrintStackTraceOnErrorSignal(argv[0],
-                                          true /* Disable crash reporting */);
+  // Skip setting up signal handlers for tests that need to test things without
+  // them configured.
+  if (!getenv("LLVM_PROGRAM_TEST_NO_STACKTRACE_HANDLER")) {
+    llvm::sys::PrintStackTraceOnErrorSignal(argv[0],
+                                            true /* Disable crash reporting */);
+  }
 
   // Initialize both gmock and gtest.
   testing::InitGoogleMock(&argc, argv);
-- 
2.7.4