[rtsan] Decouple assertions from error actions (llvm#109535)

Decouples sanitizer assertion `ExpectNotRealtime` from the action that should happen if a real-time context is detected.
ROCm · Sep 21, 2024 · a04db2c · a04db2c
1 parent cf57a67
commit a04db2c
Show file tree

Hide file tree

Showing 5 changed files with 43 additions and 54 deletions.
diff --git a/compiler-rt/lib/rtsan/CMakeLists.txt b/compiler-rt/lib/rtsan/CMakeLists.txt
@@ -2,7 +2,6 @@ include_directories(..)
 
 set(RTSAN_CXX_SOURCES
   rtsan.cpp
-  rtsan_assertions.cpp
   rtsan_context.cpp
   rtsan_diagnostics.cpp
   rtsan_flags.cpp

diff --git a/compiler-rt/lib/rtsan/rtsan.cpp b/compiler-rt/lib/rtsan/rtsan.cpp
@@ -10,6 +10,7 @@
 
 #include <rtsan/rtsan.h>
 #include <rtsan/rtsan_assertions.h>
+#include <rtsan/rtsan_diagnostics.h>
 #include <rtsan/rtsan_flags.h>
 #include <rtsan/rtsan_interceptors.h>
 
@@ -28,6 +29,13 @@ static void SetInitialized() {
   atomic_store(&rtsan_initialized, 1, memory_order_release);
 }
 
+static auto PrintDiagnosticsAndDieAction(DiagnosticsInfo info) {
+  return [info]() {
+    __rtsan::PrintDiagnostics(info);
+    Die();
+  };
+}
+
 extern "C" {
 
 SANITIZER_INTERFACE_ATTRIBUTE void __rtsan_init() {
@@ -74,22 +82,21 @@ SANITIZER_INTERFACE_ATTRIBUTE void __rtsan_enable() {
 }
 
 SANITIZER_INTERFACE_ATTRIBUTE void
-__rtsan_notify_intercepted_call(const char *intercepted_function_name) {
+__rtsan_notify_intercepted_call(const char *func_name) {
   __rtsan_ensure_initialized();
-
   GET_CALLER_PC_BP;
-  DiagnosticsInfo info = {InterceptedCallInfo{intercepted_function_name}, pc,
-                          bp};
-  ExpectNotRealtime(GetContextForThisThread(), info);
+  ExpectNotRealtime(
+      GetContextForThisThread(),
+      PrintDiagnosticsAndDieAction({InterceptedCallInfo{func_name}, pc, bp}));
 }
 
 SANITIZER_INTERFACE_ATTRIBUTE void
-__rtsan_notify_blocking_call(const char *blocking_function_name) {
+__rtsan_notify_blocking_call(const char *func_name) {
   __rtsan_ensure_initialized();
-
   GET_CALLER_PC_BP;
-  DiagnosticsInfo info = {BlockingCallInfo{blocking_function_name}, pc, bp};
-  ExpectNotRealtime(GetContextForThisThread(), info);
+  ExpectNotRealtime(
+      GetContextForThisThread(),
+      PrintDiagnosticsAndDieAction({BlockingCallInfo{func_name}, pc, bp}));
 }
 
 } // extern "C"
diff --git a/compiler-rt/lib/rtsan/rtsan_assertions.cpp b/compiler-rt/lib/rtsan/rtsan_assertions.cpp
diff --git a/compiler-rt/lib/rtsan/rtsan_assertions.h b/compiler-rt/lib/rtsan/rtsan_assertions.h
@@ -12,10 +12,19 @@
 
 #pragma once
 
+#include "rtsan/rtsan.h"
 #include "rtsan/rtsan_context.h"
-#include "rtsan/rtsan_diagnostics.h"
 
 namespace __rtsan {
 
-void ExpectNotRealtime(Context &context, const DiagnosticsInfo &info);
+template <typename OnViolationAction>
+void ExpectNotRealtime(Context &context, OnViolationAction &&OnViolation) {
+  CHECK(__rtsan_is_initialized());
+  if (context.InRealtimeContext() && !context.IsBypassed()) {
+    context.BypassPush();
+    OnViolation();
+    context.BypassPop();
+  }
+}
+
 } // namespace __rtsan
diff --git a/compiler-rt/lib/rtsan/tests/rtsan_test_assertions.cpp b/compiler-rt/lib/rtsan/tests/rtsan_test_assertions.cpp
@@ -13,9 +13,8 @@
 #include "rtsan_test_utilities.h"
 
 #include "rtsan/rtsan_assertions.h"
-#include "rtsan/rtsan_diagnostics.h"
 
-#include <gtest/gtest.h>
+#include <gmock/gmock.h>
 
 using namespace __rtsan;
 
@@ -24,30 +23,33 @@ class TestRtsanAssertions : public ::testing::Test {
   void SetUp() override { __rtsan_ensure_initialized(); }
 };
 
-DiagnosticsInfo FakeDiagnosticsInfo() {
-  DiagnosticsInfo info;
-  info.pc = 0;
-  info.bp = 0;
-  info.call_info = InterceptedCallInfo{"fake_function_name"};
-  return info;
+static void ExpectViolationAction(__rtsan::Context &context,
+                                  bool expect_violation_callback) {
+  ::testing::MockFunction<void()> mock_on_violation;
+  EXPECT_CALL(mock_on_violation, Call).Times(expect_violation_callback ? 1 : 0);
+  ExpectNotRealtime(context, mock_on_violation.AsStdFunction());
 }
 
-TEST_F(TestRtsanAssertions, ExpectNotRealtimeDoesNotDieIfNotInRealtimeContext) {
+TEST_F(TestRtsanAssertions,
+       ExpectNotRealtimeDoesNotCallViolationActionIfNotInRealtimeContext) {
   __rtsan::Context context{};
   ASSERT_FALSE(context.InRealtimeContext());
-  ExpectNotRealtime(context, FakeDiagnosticsInfo());
+  ExpectViolationAction(context, false);
 }
 
-TEST_F(TestRtsanAssertions, ExpectNotRealtimeDiesIfInRealtimeContext) {
+TEST_F(TestRtsanAssertions,
+       ExpectNotRealtimeCallsViolationActionIfInRealtimeContext) {
   __rtsan::Context context{};
   context.RealtimePush();
   ASSERT_TRUE(context.InRealtimeContext());
-  EXPECT_DEATH(ExpectNotRealtime(context, FakeDiagnosticsInfo()), "");
+  ExpectViolationAction(context, true);
 }
 
-TEST_F(TestRtsanAssertions, ExpectNotRealtimeDoesNotDieIfRealtimeButBypassed) {
+TEST_F(TestRtsanAssertions,
+       ExpectNotRealtimeDoesNotCallViolationActionIfRealtimeButBypassed) {
   __rtsan::Context context{};
   context.RealtimePush();
   context.BypassPush();
-  ExpectNotRealtime(context, FakeDiagnosticsInfo());
+  ASSERT_TRUE(context.IsBypassed());
+  ExpectViolationAction(context, false);
 }