9 files changed, 117 insertions, 2 deletions
diff --git a/compiler-rt/lib/esan/esan.cpp b/compiler-rt/lib/esan/esan.cpp
index e7399e219e6..e6d6aff3184 100644
--- a/compiler-rt/lib/esan/esan.cpp
+++ b/compiler-rt/lib/esan/esan.cpp
@@ -197,6 +197,9 @@ void initializeLibrary(ToolType Tool) {
   }
 
   initializeShadow();
+  if (__esan_which_tool == ESAN_WorkingSet)
+    initializeShadowWorkingSet();
+
   initializeInterceptors();
 
   if (__esan_which_tool == ESAN_CacheFrag) {
diff --git a/compiler-rt/lib/esan/esan_interceptors.cpp b/compiler-rt/lib/esan/esan_interceptors.cpp
index 5d4edb51d99..7aefeb6646e 100644
--- a/compiler-rt/lib/esan/esan_interceptors.cpp
+++ b/compiler-rt/lib/esan/esan_interceptors.cpp
@@ -17,6 +17,7 @@
 #include "interception/interception.h"
 #include "sanitizer_common/sanitizer_common.h"
 #include "sanitizer_common/sanitizer_libc.h"
+#include "sanitizer_common/sanitizer_linux.h"
 #include "sanitizer_common/sanitizer_stacktrace.h"
 
 using namespace __esan; // NOLINT
@@ -397,6 +398,11 @@ INTERCEPTOR(int, sigaction, int signum, const struct sigaction *act,
 // This is required to properly use internal_sigaction.
 namespace __sanitizer {
 int real_sigaction(int signum, const void *act, void *oldact) {
+  if (REAL(sigaction) == nullptr) {
+    // With an instrumented allocator, this is called during interceptor init
+    // and we need a raw syscall solution.
+    return internal_sigaction_syscall(signum, act, oldact);
+  }
   return REAL(sigaction)(signum, (const struct sigaction *)act,
                          (struct sigaction *)oldact);
 }
diff --git a/compiler-rt/lib/esan/working_set.cpp b/compiler-rt/lib/esan/working_set.cpp
index 9e7520f492e..622fd29e846 100644
--- a/compiler-rt/lib/esan/working_set.cpp
+++ b/compiler-rt/lib/esan/working_set.cpp
@@ -182,10 +182,13 @@ static void takeSample(void *Arg) {
   }
 }
 
-void initializeWorkingSet() {
+// Initialization that must be done before any instrumented code is executed.
+void initializeShadowWorkingSet() {
   CHECK(getFlags()->cache_line_size == CacheLineSize);
   registerMemoryFaultHandler();
+}
 
+void initializeWorkingSet() {
   if (getFlags()->record_snapshots) {
     for (u32 i = 0; i < NumFreq; ++i)
       SizePerFreq[i].initialize(CircularBufferSizes[i]);
diff --git a/compiler-rt/lib/esan/working_set.h b/compiler-rt/lib/esan/working_set.h
index 3750a480a11..034dfe6d991 100644
--- a/compiler-rt/lib/esan/working_set.h
+++ b/compiler-rt/lib/esan/working_set.h
@@ -21,6 +21,7 @@
 namespace __esan {
 
 void initializeWorkingSet();
+void initializeShadowWorkingSet();
 int finalizeWorkingSet();
 void processRangeAccessWorkingSet(uptr PC, uptr Addr, SIZE_T Size,
                                   bool IsWrite);
diff --git a/compiler-rt/lib/sanitizer_common/CMakeLists.txt b/compiler-rt/lib/sanitizer_common/CMakeLists.txt
index c59ca591bb4..9b9ba4b56c4 100644
--- a/compiler-rt/lib/sanitizer_common/CMakeLists.txt
+++ b/compiler-rt/lib/sanitizer_common/CMakeLists.txt
@@ -34,6 +34,11 @@ set(SANITIZER_SOURCES_NOTERMINATION
   sanitizer_thread_registry.cc
   sanitizer_win.cc)
 
+if(NOT APPLE)
+  list(APPEND SANITIZER_SOURCES_NOTERMINATION
+    sanitizer_linux_x86_64.S)
+endif()
+
 set(SANITIZER_SOURCES
   ${SANITIZER_SOURCES_NOTERMINATION} sanitizer_termination.cc)
 
@@ -133,6 +138,17 @@ append_list_if(SANITIZER_LIMIT_FRAME_SIZE -Wframe-larger-than=570
 append_list_if(COMPILER_RT_HAS_WGLOBAL_CONSTRUCTORS_FLAG -Wglobal-constructors
                SANITIZER_CFLAGS)
 
+if (LLVM_ENABLE_PEDANTIC AND NOT APPLE)
+  # With -pedantic, our .S files raise warnings about empty macro arguments
+  # from __USER_LABEL_PREFIX__ being an empty arg to GLUE().  Unfortunately,
+  # there is no simple way to test for an empty define, nor to disable just
+  # that warning or to disable -pedantic.  There is also no simple way to
+  # remove -pedantic from just this file (we'd have to remove from
+  # CMAKE_C*_FLAGS and re-add as a source property to all the non-.S files).
+  set_source_files_properties(sanitizer_linux_x86_64.S
+    PROPERTIES COMPILE_FLAGS "-w")
+endif ()
+
 if(APPLE)
   set(OS_OPTION OS ${SANITIZER_COMMON_SUPPORTED_OS})
 endif()
diff --git a/compiler-rt/lib/sanitizer_common/sanitizer_linux.cc b/compiler-rt/lib/sanitizer_common/sanitizer_linux.cc
index 75e21a6078b..90c3fe8d3d5 100644
--- a/compiler-rt/lib/sanitizer_common/sanitizer_linux.cc
+++ b/compiler-rt/lib/sanitizer_common/sanitizer_linux.cc
@@ -99,6 +99,12 @@ const int FUTEX_WAKE = 1;
 # define SANITIZER_LINUX_USES_64BIT_SYSCALLS 0
 #endif
 
+#if defined(__x86_64__)
+extern "C" {
+extern void internal_sigreturn();
+}
+#endif
+
 namespace __sanitizer {
 
 #if SANITIZER_LINUX && defined(__x86_64__)
@@ -616,7 +622,8 @@ int internal_fork() {
 
 #if SANITIZER_LINUX
 #define SA_RESTORER 0x04000000
-// Doesn't set sa_restorer, use with caution (see below).
+// Doesn't set sa_restorer if the caller did not set it, so use with caution
+//(see below).
 int internal_sigaction_norestorer(int signum, const void *act, void *oldact) {
   __sanitizer_kernel_sigaction_t k_act, k_oldact;
   internal_memset(&k_act, 0, sizeof(__sanitizer_kernel_sigaction_t));
@@ -660,6 +667,23 @@ int internal_sigaction_norestorer(int signum, const void *act, void *oldact) {
   }
   return result;
 }
+
+// Invokes sigaction via a raw syscall with a restorer, but does not support
+// all platforms yet.
+// We disable for Go simply because we have not yet added to buildgo.sh.
+#if defined(__x86_64__) && !SANITIZER_GO
+int internal_sigaction_syscall(int signum, const void *act, void *oldact) {
+  __sanitizer_sigaction u_adjust;
+  internal_memcpy(&u_adjust, act, sizeof(u_adjust));
+#if !SANITIZER_ANDROID || !SANITIZER_MIPS32
+    if (u_adjust.sa_restorer == nullptr) {
+      u_adjust.sa_restorer = internal_sigreturn;
+    }
+#endif
+    return internal_sigaction_norestorer(signum, (const void *)&u_adjust,
+                                         oldact);
+}
+#endif // defined(__x86_64__) && !SANITIZER_GO
 #endif  // SANITIZER_LINUX
 
 uptr internal_sigprocmask(int how, __sanitizer_sigset_t *set,
diff --git a/compiler-rt/lib/sanitizer_common/sanitizer_linux.h b/compiler-rt/lib/sanitizer_common/sanitizer_linux.h
index beb43568ad0..526fa4426e3 100644
--- a/compiler-rt/lib/sanitizer_common/sanitizer_linux.h
+++ b/compiler-rt/lib/sanitizer_common/sanitizer_linux.h
@@ -42,6 +42,10 @@ uptr internal_prctl(int option, uptr arg2, uptr arg3, uptr arg4, uptr arg5);
 // (like the process-wide error reporting SEGV handler) must use
 // internal_sigaction instead.
 int internal_sigaction_norestorer(int signum, const void *act, void *oldact);
+#if defined(__x86_64__) && !SANITIZER_GO
+// Uses a raw system call to avoid interceptors.
+int internal_sigaction_syscall(int signum, const void *act, void *oldact);
+#endif
 void internal_sigdelset(__sanitizer_sigset_t *set, int signum);
 #if defined(__x86_64__) || defined(__mips__) || defined(__aarch64__) \
   || defined(__powerpc64__) || defined(__s390__)
diff --git a/compiler-rt/lib/sanitizer_common/sanitizer_linux_x86_64.S b/compiler-rt/lib/sanitizer_common/sanitizer_linux_x86_64.S
new file mode 100644
index 00000000000..8ff909542b6
--- /dev/null
+++ b/compiler-rt/lib/sanitizer_common/sanitizer_linux_x86_64.S
@@ -0,0 +1,25 @@
+// This file is dual licensed under the MIT and the University of Illinois Open
+// Source Licenses. See LICENSE.TXT for details.
+
+// Avoid being marked as needing an executable stack:
+#if defined(__linux__) && defined(__ELF__)
+.section .note.GNU-stack,"",%progbits
+#endif
+
+// Further contents are x86_64-only:
+#if defined(__linux__) && defined(__x86_64__)
+
+#include "../builtins/assembly.h"
+
+// If the "naked" function attribute were supported for x86 we could
+// do this via inline asm.
+.text
+.balign 4
+DEFINE_COMPILERRT_FUNCTION(internal_sigreturn)
+        mov           $0xf,             %eax    // 0xf == SYS_rt_sigreturn
+        mov           %rcx,             %r10
+        syscall
+        ret                                     // Won't normally reach here.
+END_COMPILERRT_FUNCTION(internal_sigreturn)
+
+#endif // defined(__linux__) && defined(__x86_64__)
diff --git a/compiler-rt/test/esan/TestCases/workingset-early-fault.c b/compiler-rt/test/esan/TestCases/workingset-early-fault.c
new file mode 100644
index 00000000000..1c420c368ca
--- /dev/null
+++ b/compiler-rt/test/esan/TestCases/workingset-early-fault.c
@@ -0,0 +1,33 @@
+// Test shadow faults during esan initialization as well as
+// faults during dlsym's calloc during interceptor init.
+//
+// RUN: %clang_esan_wset %s -o %t
+// RUN: %run %t 2>&1 | FileCheck %s
+
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+// Our goal is to emulate an instrumented allocator, whose calloc
+// invoked from dlsym will trigger shadow faults, to test an
+// early shadow fault during esan interceptor init.
+// We do this by replacing calloc:
+void *calloc(size_t size, size_t n) {
+  // Unfortunately we can't print anything to make the test
+  // ensure we got here b/c the sanitizer interceptors can't
+  // handle that during interceptor init.
+
+  // Ensure we trigger a shadow write fault:
+  int x[16];
+  x[0] = size;
+  // Now just emulate calloc.
+  void *res = malloc(size*n);
+  memset(res, 0, size*n);
+  return res;
+}
+
+int main(int argc, char **argv) {
+  printf("all done\n");
+  return 0;
+}
+// CHECK: all done