[ASan][Win][compiler-rt] Fixes stack overflow when ntdll has mem* calls during exception handling

compiler-rt/lib/asan/asan_interceptors_memintrinsics.cpp

@@ -20,43 +20,99 @@
 #include "asan_stack.h"
 #include "asan_suppressions.h"
 
+#if SANITIZER_WINDOWS64
+#  include "asan_poisoning.h"
+#  include "sanitizer_common/sanitizer_win.h"
+
+extern "C" void WINAPI RestoreLastError(DWORD);
+extern "C" DWORD WINAPI GetLastError();
+#endif
+
+namespace __asan {
+// On x64, the ShadowExceptionHandler is expected to handle all AVs that happen
+// as a result of uncommitted shadow memory pages. However, in programs that use
+// ntdll (a Windows-specific library that contains some memory intrinsics as
+// well as Windows-specific exception handling mechanisms) as their C Runtime,
+// or in cases where ntdll uses mem* functions inside
+// its exception handling infrastructure, ASAN can end up rethrowing a shadow
+// memory AV until a stack overflow occurs. In other words, ntdll can call back
+// into ASAN for a poisoning check, which creates infinite recursion. To remedy
+// this, we precommit the shadow memory of the address being accessed on x64 for
+// ntdll callees.
+bool ShouldReplaceIntrinsic(bool isNtdllCallee, void *addr, uptr size,
+                            const void *from = nullptr) {
+#if SANITIZER_WINDOWS64
+  if (isNtdllCallee) {
+    const auto lastError = GetLastError();  // Grab last error here to maintain
+                                            // status for after internal calls
+    if (addr) {
+      CommitShadowMemory(reinterpret_cast<uptr>(addr), size);
+    }
+    if (from) {
+      CommitShadowMemory(reinterpret_cast<uptr>(from), size);
+    }
+    RestoreLastError(lastError);
+  }
+#endif
+  return replace_intrin_cached;
+}
+}  // namespace __asan
+
 using namespace __asan;
 
+#if SANITIZER_WINDOWS64
+#  define IS_NTDLL_CALLEE __sanitizer::IsNtdllCallee(_ReturnAddress())
+#else
+#  define IS_NTDLL_CALLEE false
+#endif
+
 // memcpy is called during __asan_init() from the internals of printf(...).
 // We do not treat memcpy with to==from as a bug.
 // See http://llvm.org/bugs/show_bug.cgi?id=11763.
-#define ASAN_MEMCPY_IMPL(ctx, to, from, size)                 \
-  do {                                                        \
-    if (LIKELY(replace_intrin_cached)) {                      \
-      if (LIKELY(to != from)) {                               \
-        CHECK_RANGES_OVERLAP("memcpy", to, size, from, size); \
-      }                                                       \
-      ASAN_READ_RANGE(ctx, from, size);                       \
-      ASAN_WRITE_RANGE(ctx, to, size);                        \
-    } else if (UNLIKELY(!AsanInited())) {                     \
-      return internal_memcpy(to, from, size);                 \
-    }                                                         \
-    return REAL(memcpy)(to, from, size);                      \
+#define ASAN_MEMCPY_IMPL(ctx, to, from, size)                       \
+  do {                                                              \
+    if (SANITIZER_WINDOWS64 &&                                      \
+        !ShouldReplaceIntrinsic(IS_NTDLL_CALLEE, to, size, from)) { \
+      return REAL(memcpy)(to, from, size);                          \
+    }                                                               \
+    if (LIKELY(replace_intrin_cached)) {                            \
+      if (LIKELY(to != from)) {                                     \
+        CHECK_RANGES_OVERLAP("memcpy", to, size, from, size);       \
+      }                                                             \
+      ASAN_READ_RANGE(ctx, from, size);                             \
+      ASAN_WRITE_RANGE(ctx, to, size);                              \
+    } else if (UNLIKELY(!AsanInited())) {                           \
+      return internal_memcpy(to, from, size);                       \
+    }                                                               \
+    return REAL(memcpy)(to, from, size);                            \
   } while (0)
 
 // memset is called inside Printf.
-#define ASAN_MEMSET_IMPL(ctx, block, c, size) \
-  do {                                        \
-    if (LIKELY(replace_intrin_cached)) {      \
-      ASAN_WRITE_RANGE(ctx, block, size);     \
-    } else if (UNLIKELY(!AsanInited())) {     \
-      return internal_memset(block, c, size); \
-    }                                         \
-    return REAL(memset)(block, c, size);      \
+#define ASAN_MEMSET_IMPL(ctx, block, c, size)                    \
+  do {                                                           \
+    if (SANITIZER_WINDOWS64 &&                                   \
+        !ShouldReplaceIntrinsic(IS_NTDLL_CALLEE, block, size)) { \
+      return REAL(memset)(block, c, size);                       \
+    }                                                            \
+    if (LIKELY(replace_intrin_cached)) {                         \
+      ASAN_WRITE_RANGE(ctx, block, size);                        \
+    } else if (UNLIKELY(!AsanInited())) {                        \
+      return internal_memset(block, c, size);                    \
+    }                                                            \
+    return REAL(memset)(block, c, size);                         \
   } while (0)
 
-#define ASAN_MEMMOVE_IMPL(ctx, to, from, size) \
-  do {                                         \
-    if (LIKELY(replace_intrin_cached)) {       \
-      ASAN_READ_RANGE(ctx, from, size);        \
-      ASAN_WRITE_RANGE(ctx, to, size);         \
-    }                                          \
-    return internal_memmove(to, from, size);   \
+#define ASAN_MEMMOVE_IMPL(ctx, to, from, size)                      \
+  do {                                                              \
+    if (SANITIZER_WINDOWS64 &&                                      \
+        !ShouldReplaceIntrinsic(IS_NTDLL_CALLEE, to, size, from)) { \
+      return internal_memmove(to, from, size);                      \
+    }                                                               \
+    if (LIKELY(replace_intrin_cached)) {                            \
+      ASAN_READ_RANGE(ctx, from, size);                             \
+      ASAN_WRITE_RANGE(ctx, to, size);                              \
+    }                                                               \
+    return internal_memmove(to, from, size);                        \
   } while (0)
 
 void *__asan_memcpy(void *to, const void *from, uptr size) {

compiler-rt/lib/asan/asan_poisoning.h

@@ -17,6 +17,25 @@
 #include "sanitizer_common/sanitizer_flags.h"
 #include "sanitizer_common/sanitizer_platform.h"
 
+#if SANITIZER_WINDOWS64
+#  include "sanitizer_common/sanitizer_win.h"
+#  include "sanitizer_common/sanitizer_win_defs.h"
+
+// These definitions are duplicated from Window.h in order to avoid conflicts
+// with other types in Windows.h.
+// These functions and types are used to manipulate the shadow memory on
+// x64 Windows.
+typedef unsigned long DWORD;
+typedef void *LPVOID;
+typedef int BOOL;
+
+constexpr DWORD MEM_COMMIT = 0x00001000;
+constexpr DWORD MEM_DECOMMIT = 0x00004000;
+constexpr DWORD PAGE_READWRITE = 0x04;
+
+extern "C" LPVOID WINAPI VirtualAlloc(LPVOID, size_t, DWORD, DWORD);
+#endif
+
 namespace __asan {
 
 // Enable/disable memory poisoning.
@@ -95,4 +114,16 @@ ALWAYS_INLINE void FastPoisonShadowPartialRightRedzone(
 // [MemToShadow(p), MemToShadow(p+size)].
 void FlushUnneededASanShadowMemory(uptr p, uptr size);
 
+// Commits the shadow memory for a range of aligned memory. This only matters
+// on 64-bit Windows where relying on pages to get paged in on access
+// violation is inefficient when we know the memory range ahead of time.
+ALWAYS_INLINE void CommitShadowMemory(uptr aligned_beg, uptr aligned_size) {
+#if SANITIZER_WINDOWS64
+  uptr shadow_beg = MEM_TO_SHADOW(aligned_beg);
+  uptr shadow_end =
+      MEM_TO_SHADOW(aligned_beg + aligned_size - ASAN_SHADOW_GRANULARITY) + 1;
+  ::VirtualAlloc((LPVOID)shadow_beg, (size_t)(shadow_end - shadow_beg),
+                 MEM_COMMIT, PAGE_READWRITE);
+#endif
+}
 }  // namespace __asan

compiler-rt/lib/asan/asan_rtl.cpp

@@ -463,6 +463,10 @@ static bool AsanInitInternal() {
   if (SANITIZER_START_BACKGROUND_THREAD_IN_ASAN_INTERNAL)
     MaybeStartBackgroudThread();
 
+#if SANITIZER_WINDOWS64
+  __sanitizer::InitializeNtdllInfo();
+#endif
+
   // On Linux AsanThread::ThreadStart() calls malloc() that's why asan_inited
   // should be set to 1 prior to initializing the threads.
   replace_intrin_cached = flags()->replace_intrin;

compiler-rt/lib/sanitizer_common/sanitizer_win.cpp

@@ -1264,6 +1264,18 @@ void LogFullErrorReport(const char *buffer) {
 
 void InitializePlatformCommonFlags(CommonFlags *cf) {}
 
+static MODULEINFO ntdllInfo;
+void InitializeNtdllInfo() {
+  GetModuleInformation(GetCurrentProcess(), GetModuleHandle(L"ntdll.dll"),
+                       &ntdllInfo, sizeof(ntdllInfo));
+}
+
+bool IsNtdllCallee(void *calleeAddr) {
+  return (uptr)ntdllInfo.lpBaseOfDll <= (uptr)calleeAddr &&
+         ((uptr)ntdllInfo.lpBaseOfDll + (uptr)ntdllInfo.SizeOfImage) >=
+             (uptr)calleeAddr;
+}
+
 }  // namespace __sanitizer
 
 #endif  // _WIN32

compiler-rt/lib/sanitizer_common/sanitizer_win.h

@@ -19,6 +19,12 @@
 namespace __sanitizer {
 // Check based on flags if we should handle the exception.
 bool IsHandledDeadlyException(DWORD exceptionCode);
+
+// Initializes module information of ntdll for referencing callee addresses
+void InitializeNtdllInfo();
+
+// Returns whether or not the callee address lies within ntdll
+bool IsNtdllCallee(void* calleeAddr);
 }  // namespace __sanitizer
 
 #endif  // SANITIZER_WINDOWS

compiler-rt/test/asan/TestCases/Windows/ntdll_regression_tests.cpp

@@ -0,0 +1,26 @@
+// RUN: %clang_cl_asan /Od %s -Fe%t
+// RUN: not %run %t 2>&1 | FileCheck %s
+
+#include <Windows.h>
+#include <iostream>
+
+// Small sanity test to make sure ASAN does not stomp on
+// GetLastError values. This is motivated by __asan::ShouldReplaceIntrinsic,
+// which remedied infinite recursion due to ntdll exception handling paths calling
+// instrumented functions on startup before shadow memory can be committed.
+int TestNTStatusMaintained() {
+  ::SetLastError(ERROR_SUCCESS);
+  constexpr unsigned long c_initialSizeGuess = 1;
+  wchar_t szBuffer[c_initialSizeGuess];
+  unsigned long size = ::GetDllDirectoryW(c_initialSizeGuess, szBuffer);
+  auto le = ::GetLastError();
+  if (size == 0 && le != ERROR_SUCCESS) {
+    std::cerr << "Last error is different.\n";
+    return -1;
+  }
+  std::cerr << "Success.\n";
+  return 0;
+  // CHECK: Success.
+}
+
+int main() { return TestNTStatusMaintained(); }