[llvm][rtsan] Add transform pass for sanitize_realtime_unsafe (#109543)

Author: davidtrevelyan

llvm/lib/Transforms/Instrumentation/RealtimeSanitizer.cpp

@@ -15,49 +15,77 @@
 
 #include "llvm/IR/Analysis.h"
 #include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/InstIterator.h"
 #include "llvm/IR/Module.h"
 
+#include "llvm/Demangle/Demangle.h"
 #include "llvm/Transforms/Instrumentation/RealtimeSanitizer.h"
 
 using namespace llvm;
 
+static SmallVector<Type *> getArgTypes(ArrayRef<Value *> FunctionArgs) {
+  SmallVector<Type *> Types;
+  for (Value *Arg : FunctionArgs)
+    Types.push_back(Arg->getType());
+  return Types;
+}
+
 static void insertCallBeforeInstruction(Function &Fn, Instruction &Instruction,
-                                        const char *FunctionName) {
+                                        const char *FunctionName,
+                                        ArrayRef<Value *> FunctionArgs) {
   LLVMContext &Context = Fn.getContext();
-  FunctionType *FuncType = FunctionType::get(Type::getVoidTy(Context), false);
+  FunctionType *FuncType = FunctionType::get(Type::getVoidTy(Context),
+                                             getArgTypes(FunctionArgs), false);
   FunctionCallee Func =
       Fn.getParent()->getOrInsertFunction(FunctionName, FuncType);
   IRBuilder<> Builder{&Instruction};
-  Builder.CreateCall(Func, {});
+  Builder.CreateCall(Func, FunctionArgs);
 }
 
 static void insertCallAtFunctionEntryPoint(Function &Fn,
-                                           const char *InsertFnName) {
-
-  insertCallBeforeInstruction(Fn, Fn.front().front(), InsertFnName);
+                                           const char *InsertFnName,
+                                           ArrayRef<Value *> FunctionArgs) {
+  insertCallBeforeInstruction(Fn, Fn.front().front(), InsertFnName,
+                              FunctionArgs);
 }
 
 static void insertCallAtAllFunctionExitPoints(Function &Fn,
-                                              const char *InsertFnName) {
-  for (auto &BB : Fn)
-    for (auto &I : BB)
-      if (isa<ReturnInst>(&I))
-        insertCallBeforeInstruction(Fn, I, InsertFnName);
+                                              const char *InsertFnName,
+                                              ArrayRef<Value *> FunctionArgs) {
+  for (auto &I : instructions(Fn))
+    if (isa<ReturnInst>(&I))
+      insertCallBeforeInstruction(Fn, I, InsertFnName, FunctionArgs);
+}
+
+static PreservedAnalyses rtsanPreservedCFGAnalyses() {
+  PreservedAnalyses PA;
+  PA.preserveSet<CFGAnalyses>();
+  return PA;
+}
+
+static PreservedAnalyses runSanitizeRealtime(Function &Fn) {
+  insertCallAtFunctionEntryPoint(Fn, "__rtsan_realtime_enter", {});
+  insertCallAtAllFunctionExitPoints(Fn, "__rtsan_realtime_exit", {});
+  return rtsanPreservedCFGAnalyses();
+}
+
+static PreservedAnalyses runSanitizeRealtimeUnsafe(Function &Fn) {
+  IRBuilder<> Builder(&Fn.front().front());
+  Value *Name = Builder.CreateGlobalString(demangle(Fn.getName()));
+  insertCallAtFunctionEntryPoint(Fn, "__rtsan_notify_blocking_call", {Name});
+  return rtsanPreservedCFGAnalyses();
 }
 
 RealtimeSanitizerPass::RealtimeSanitizerPass(
     const RealtimeSanitizerOptions &Options) {}
 
-PreservedAnalyses RealtimeSanitizerPass::run(Function &F,
+PreservedAnalyses RealtimeSanitizerPass::run(Function &Fn,
                                              AnalysisManager<Function> &AM) {
-  if (F.hasFnAttribute(Attribute::SanitizeRealtime)) {
-    insertCallAtFunctionEntryPoint(F, "__rtsan_realtime_enter");
-    insertCallAtAllFunctionExitPoints(F, "__rtsan_realtime_exit");
-
-    PreservedAnalyses PA;
-    PA.preserveSet<CFGAnalyses>();
-    return PA;
-  }
+  if (Fn.hasFnAttribute(Attribute::SanitizeRealtime))
+    return runSanitizeRealtime(Fn);
+
+  if (Fn.hasFnAttribute(Attribute::SanitizeRealtimeUnsafe))
+    return runSanitizeRealtimeUnsafe(Fn);
 
   return PreservedAnalyses::all();
 }

llvm/test/Instrumentation/RealtimeSanitizer/rtsan_unsafe.ll

@@ -0,0 +1,31 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-globals all --version 5
+; RUN: opt < %s -passes=rtsan -S | FileCheck %s
+
+; RealtimeSanitizer pass should create the demangled function name as a global string and,
+; at the function entrypoint, pass it as an argument to the rtsan notify method
+
+;.
+; CHECK: @[[GLOB0:[0-9]+]] = private unnamed_addr constant [20 x i8] c"blocking_function()\00", align 1
+;.
+define void @_Z17blocking_functionv() #0 {
+; CHECK-LABEL: define void @_Z17blocking_functionv(
+; CHECK-SAME: ) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:    call void @__rtsan_notify_blocking_call(ptr @[[GLOB0]])
+; CHECK-NEXT:    ret void
+;
+  ret void
+}
+
+define noundef i32 @main() #2 {
+; CHECK-LABEL: define noundef i32 @main() {
+; CHECK-NEXT:    call void @_Z17blocking_functionv()
+; CHECK-NEXT:    ret i32 0
+;
+  call void @_Z17blocking_functionv() #4
+  ret i32 0
+}
+
+attributes #0 = { mustprogress noinline sanitize_realtime_unsafe optnone ssp uwtable(sync) }
+;.
+; CHECK: attributes #[[ATTR0]] = { mustprogress noinline optnone sanitize_realtime_unsafe ssp uwtable(sync) }
+;.