Introduce global InterceptorEnabled for interceptor control

This adds a globally accessible `InterceptorEnabled` variable to control interceptor behavior dynamically. The changes ensure interceptors can be selectively disabled during specific instructions, with proper restoration of the state afterward. This improves flexibility in handling memory intrinsics and thread sanitization instrumentation.
focs-lab · MorthimerMcMare · Oct 14, 2024 · Oct 15, 2024 · Oct 18, 2024 · Oct 18, 2024
commit c94629915ebd75b3a6e7063bcfa01329b494b717
diff --git a/compiler-rt/lib/sanitizer_common/sanitizer_common.cpp b/compiler-rt/lib/sanitizer_common/sanitizer_common.cpp
@@ -20,6 +20,8 @@
 #include "sanitizer_libc.h"
 #include "sanitizer_placement_new.h"
 
+extern "C" bool InterceptorEnabled = true;
+
 namespace __sanitizer {
 
 const char *SanitizerToolName = "SanitizerTool";

diff --git a/compiler-rt/lib/sanitizer_common/sanitizer_common_interceptors_memintrinsics.inc b/compiler-rt/lib/sanitizer_common/sanitizer_common_interceptors_memintrinsics.inc
@@ -31,6 +31,8 @@
 #include "interception/interception.h"
 #include "sanitizer_platform_interceptors.h"
 
+extern "C" bool InterceptorEnabled;
+
 // Platform-specific options.
 #if SANITIZER_APPLE
 #define PLATFORM_HAS_DIFFERENT_MEMCPY_AND_MEMMOVE 0
@@ -40,13 +42,10 @@
 #define PLATFORM_HAS_DIFFERENT_MEMCPY_AND_MEMMOVE 1
 #endif  // SANITIZER_APPLE
 
-bool InterceptorEnabled = 1;
-
 #ifndef COMMON_INTERCEPTOR_MEMSET_IMPL
 #define COMMON_INTERCEPTOR_MEMSET_IMPL(ctx, dst, v, size)   \
   {                                                         \
     if (InterceptorEnabled) {                               \
-      Printf("INTERCEPTOR_ENABLED\n");                      \
       if (COMMON_INTERCEPTOR_NOTHING_IS_INITIALIZED)        \
         return internal_memset(dst, v, size);               \
       COMMON_INTERCEPTOR_ENTER(ctx, memset, dst, v, size);  \
@@ -58,31 +57,35 @@ bool InterceptorEnabled = 1;
 #endif
 
 #ifndef COMMON_INTERCEPTOR_MEMMOVE_IMPL
-#define COMMON_INTERCEPTOR_MEMMOVE_IMPL(ctx, dst, src, size) \
-  {                                                          \
-    if (COMMON_INTERCEPTOR_NOTHING_IS_INITIALIZED)           \
-      return internal_memmove(dst, src, size);               \
-    COMMON_INTERCEPTOR_ENTER(ctx, memmove, dst, src, size);  \
-    if (common_flags()->intercept_intrin) {                  \
-      COMMON_INTERCEPTOR_WRITE_RANGE(ctx, dst, size);        \
-      COMMON_INTERCEPTOR_READ_RANGE(ctx, src, size);         \
-    }                                                        \
-    return REAL(memmove)(dst, src, size);                    \
+#define COMMON_
8000
INTERCEPTOR_MEMMOVE_IMPL(ctx, dst, src, size)   \
+  {                                                            \
+    if (InterceptorEnabled) {                                  \
+      if (COMMON_INTERCEPTOR_NOTHING_IS_INITIALIZED)           \
+        return internal_memmove(dst, src, size);               \
+      COMMON_INTERCEPTOR_ENTER(ctx, memmove, dst, src, size);  \
+      if (common_flags()->intercept_intrin) {                  \
+        COMMON_INTERCEPTOR_WRITE_RANGE(ctx, dst, size);        \
+        COMMON_INTERCEPTOR_READ_RANGE(ctx, src, size);         \
+      }                                                        \
+    }                                                          \
+    return REAL(memmove)(dst, src, size);                      \
   }
 #endif
 
 #ifndef COMMON_INTERCEPTOR_MEMCPY_IMPL
-#define COMMON_INTERCEPTOR_MEMCPY_IMPL(ctx, dst, src, size) \
-  {                                                         \
-    if (COMMON_INTERCEPTOR_NOTHING_IS_INITIALIZED) {        \
-      return internal_memmove(dst, src, size);              \
-    }                                                       \
-    COMMON_INTERCEPTOR_ENTER(ctx, memcpy, dst, src, size);  \
-    if (common_flags()->intercept_intrin) {                 \
-      COMMON_INTERCEPTOR_WRITE_RANGE(ctx, dst, size);       \
-      COMMON_INTERCEPTOR_READ_RANGE(ctx, src, size);        \
-    }                                                       \
-    return REAL(memcpy)(dst, src, size);                    \
+#define COMMON_INTERCEPTOR_MEMCPY_IMPL(ctx, dst, src, size)   \
+  {                                                           \
+    if (InterceptorEnabled) {                                 \
+      if (COMMON_INTERCEPTOR_NOTHING_IS_INITIALIZED) {        \
+        return internal_memmove(dst, src, size);              \
+      }                                                       \
+      COMMON_INTERCEPTOR_ENTER(ctx, memcpy, dst, src, size);  \
+      if (common_flags()->intercept_intrin) {                 \
+        COMMON_INTERCEPTOR_WRITE_RANGE(ctx, dst, size);       \
+        COMMON_INTERCEPTOR_READ_RANGE(ctx, src, size);        \
+      }                                                       \
+    }                                                         \
+    return REAL(memcpy)(dst, src, size);                      \
   }
 #endif
 

diff --git a/compiler-rt/lib/tsan/rtl/tsan_rtl.h b/compiler-rt/lib/tsan/rtl/tsan_rtl.h
@@ -499,7 +499,7 @@ bool OutputReport(ThreadState *thr, const ScopedReport &srep);
 bool IsFiredSuppression(Context *ctx, ReportType type, StackTrace trace);
 bool IsExpectedReport(uptr addr, uptr size);
 
-#define TSAN_DEBUG_OUTPUT 2
+//#define TSAN_DEBUG_OUTPUT 2
 
 #if defined(TSAN_DEBUG_OUTPUT) && TSAN_DEBUG_OUTPUT >= 1
 # define DPrintf Printf

diff --git a/llvm/lib/Transforms/Instrumentation/ThreadSanitizer.cpp b/llvm/lib/Transforms/Instrumentation/ThreadSanitizer.cpp
@@ -123,6 +123,8 @@ const char kTsanInitName[] = "__tsan_init";
 
 namespace {
 
+GlobalVariable *InterceptorEnabled;
+
 /// ThreadSanitizer: instrument the code in module to find races.
 ///
 /// Instantiating ThreadSanitizer inserts the tsan runtime library API function
@@ -167,6 +169,7 @@ struct ThreadSanitizer {
   void initialize(Module &M, const TargetLibraryInfo &TLI);
   bool instrumentLoadOrStore(const InstructionInfo &II, const DataLayout &DL);
   bool instrumentAtomic(Instruction *I, const DataLayout &DL);
+  void disableInterceptorForInstr(Instruction *I, InstrumentationIRBuilder &IRB);
   bool
   instrumentMemIntrinsic(Instruction *I,
                          std::optional<EscapeAnalysisGlobalInfo *> EAIGlobal);
@@ -206,8 +209,6 @@ struct ThreadSanitizer {
   FunctionCallee TsanVptrUpdate;
   FunctionCallee TsanVptrLoad;
   FunctionCallee MemmoveFn, MemcpyFn, MemsetFn;
-
-  GlobalVariable *InterceptorEnabled;
 };
 
 void insertModuleCtor(Module &M) {
@@ -256,6 +257,12 @@ PreservedAnalyses ModuleThreadSanitizerPass::run(Module &M,
     MAM.getResult<EscapeAnalysisGlobal>(M);
 
   insertModuleCtor(M);
+
+  // Declare an external global variable InterceptorEnabled in the module
+  InterceptorEnabled = new GlobalVariable(
+      M, Type::getInt1Ty(M.getContext()), /*isConstant=*/false,
+      GlobalValue::ExternalLinkage, nullptr, "InterceptorEnabled");
+
   return PreservedAnalyses::none();
 }
 
@@ -405,12 +412,6 @@ void ThreadSanitizer::initialize(Module &M, const TargetLibraryInfo &TLI) {
       "__tsan_memset",
       TLI.getAttrList(&Ctx, {1}, /*Signed=*/true, /*Ret=*/false, Attr),
       IRB.getPtrTy(), IRB.getPtrTy(), IRB.getInt32Ty(), IntptrTy);
-
-  // Declare a global variable InterceptorEnabled in the LLVM module
-  InterceptorEnabled =
-      new GlobalVariable(M, Type::getInt1Ty(Ctx), /*isConstant=*/false,
-                         GlobalValue::ExternalLinkage,
-                         ConstantInt::getTrue(Ctx), "InterceptorEnabled");
 }
 
 static bool isVtableAccess(Instruction *I) {
@@ -809,9 +810,15 @@ static ConstantInt *createOrdering(IRBuilder<> *IRB, AtomicOrdering ord) {
   return IRB->getInt32(v);
 }
 
-/**
- *
- */
+void ThreadSanitizer::disableInterceptorForInstr(
+    Instruction *I, InstrumentationIRBuilder &IRB) {
+  // Disable interceptors for this call
+  IRB.CreateStore(IRB.getInt1(false), InterceptorEnabled);
+  // After MemsetFn, set the InterceptorEnabled back to true
+  IRB.SetInsertPoint(++BasicBlock::iterator(I));
+  IRB.CreateStore(IRB.getInt1(true), InterceptorEnabled);
+}
+
 // So, we either need to ensure the intrinsic is not inlined, or instrument it.
 // We do not instrument memset/memmove/memcpy intrinsics (too complicated),
 // instead we simply replace them with regular function calls, which are then
@@ -839,6 +846,7 @@ bool ThreadSanitizer::instrumentMemIntrinsic(Instruction *I,
 
     // Check if pointer is not escape
     if (!isPointerEscaped(M->getArgOperand(0))) {
+      disableInterceptorForInstr(I, IRB);
       return false;
     }
 
@@ -850,8 +858,11 @@ bool ThreadSanitizer::instrumentMemIntrinsic(Instruction *I,
     I->eraseFromParent();
   } else if (MemTransferInst *M = dyn_cast<MemTransferInst>(I)) {
     // Check if pointers are not escape
-    if (!isPointerEscaped(M->getArgOperand(0)) && !isPointerEscaped(M->getArgOperand(1)))
+    if (!isPointerEscaped(M->getArgOperand(0)) &&
+        !isPointerEscaped(M->getArgOperand(1))) {
+      disableInterceptorForInstr(I, IRB);
       return false;
+    }
 
     IRB.CreateCall(
         isa<MemCpyInst>(M) ? MemcpyFn : MemmoveFn,