3 files changed, 35 insertions, 29 deletions

diff --git a/compiler-rt/lib/interception/interception_linux.cc b/compiler-rt/lib/interception/interception_linux.cc
index d019d482d6d..d07f060b5b6 100644
--- a/compiler-rt/lib/interception/interception_linux.cc
+++ b/compiler-rt/lib/interception/interception_linux.cc
@@ -33,13 +33,7 @@ static int StrCmp(const char *s1, const char *s2) {
 }
 #endif
 
-bool GetRealFunctionAddress(const char *func_name, uptr *func_addr,
-    uptr real, uptr wrapper) {
-  *func_addr = (uptr)GetFuncAddr(func_name);
-  return real == wrapper;
-}
-
-void *GetFuncAddr(const char *name) {
+static void *GetFuncAddr(const char *name) {
 #if SANITIZER_NETBSD
   // FIXME: Find a better way to handle renames
   if (StrCmp(name, "sigaction"))
@@ -57,11 +51,25 @@ void *GetFuncAddr(const char *name) {
   return addr;
 }
 
+bool InterceptFunction(const char *name, uptr *ptr_to_real, uptr func,
+                       uptr wrapper) {
+  void *addr = GetFuncAddr(name);
+  *ptr_to_real = (uptr)addr;
+  return addr && (func == wrapper);
+}
+
 // Android and Solaris do not have dlvsym
 #if !SANITIZER_ANDROID && !SANITIZER_SOLARIS && !SANITIZER_OPENBSD
-void *GetFuncAddrVer(const char *name, const char *ver) {
+static void *GetFuncAddr(const char *name, const char *ver) {
   return dlvsym(RTLD_NEXT, name, ver);
 }
+
+bool InterceptFunction(const char *name, const char *ver, uptr *ptr_to_real,
+                       uptr func, uptr wrapper) {
+  void *addr = GetFuncAddr(name, ver);
+  *ptr_to_real = (uptr)addr;
+  return addr && (func == wrapper);
+}
 #endif  // !SANITIZER_ANDROID
 
 }  // namespace __interception
diff --git a/compiler-rt/lib/interception/interception_linux.h b/compiler-rt/lib/interception/interception_linux.h
index 4c60bd88090..e578da0cf64 100644
--- a/compiler-rt/lib/interception/interception_linux.h
+++ b/compiler-rt/lib/interception/interception_linux.h
@@ -22,24 +22,27 @@
 #define INTERCEPTION_LINUX_H
 
 namespace __interception {
-// returns true if a function with the given name was found.
-bool GetRealFunctionAddress(const char *func_name, uptr *func_addr,
-    uptr real, uptr wrapper);
-void *GetFuncAddr(const char *name);
-void *GetFuncAddrVer(const char *name, const char *ver);
+bool InterceptFunction(const char *name, uptr *ptr_to_real, uptr func,
+                       uptr wrapper);
+bool InterceptFunction(const char *name, const char *ver, uptr *ptr_to_real,
+                       uptr func, uptr wrapper);
 }  // namespace __interception
 
-#define INTERCEPT_FUNCTION_LINUX_OR_FREEBSD(func)                          \
-  ::__interception::GetRealFunctionAddress(                                \
-      #func, (::__interception::uptr *)&__interception::PTR_TO_REAL(func), \
-      (::__interception::uptr) & (func),                                   \
+#define INTERCEPT_FUNCTION_LINUX_OR_FREEBSD(func) \
+  ::__interception::InterceptFunction(            \
+      #func,                                      \
+      (::__interception::uptr *) & REAL(func),    \
+      (::__interception::uptr) & (func),          \
       (::__interception::uptr) & WRAP(func))
 
 // Android,  Solaris and OpenBSD do not have dlvsym
 #if !SANITIZER_ANDROID && !SANITIZER_SOLARIS && !SANITIZER_OPENBSD
 #define INTERCEPT_FUNCTION_VER_LINUX_OR_FREEBSD(func, symver) \
-  (::__interception::real_##func = (func##_type)(                \
-       unsigned long)::__interception::GetFuncAddrVer(#func, symver))
+  ::__interception::InterceptFunction(                        \
+      #func, symver,                                          \
+      (::__interception::uptr *) & REAL(func),                \
+      (::__interception::uptr) & (func),                      \
+      (::__interception::uptr) & WRAP(func))
 #else
 #define INTERCEPT_FUNCTION_VER_LINUX_OR_FREEBSD(func, symver) \
   INTERCEPT_FUNCTION_LINUX_OR_FREEBSD(func)
diff --git a/compiler-rt/lib/interception/tests/interception_linux_test.cc b/compiler-rt/lib/interception/tests/interception_linux_test.cc
index 5a1ca6ecbf5..3e859cb1a55 100644
--- a/compiler-rt/lib/interception/tests/interception_linux_test.cc
+++ b/compiler-rt/lib/interception/tests/interception_linux_test.cc
@@ -33,24 +33,19 @@ INTERCEPTOR(int, isdigit, int d) {
 
 namespace __interception {
 
-TEST(Interception, GetRealFunctionAddress) {
+TEST(Interception, InterceptFunction) {
   uptr malloc_address = 0;
-  EXPECT_TRUE(GetRealFunctionAddress("malloc", &malloc_address, 0, 0));
+  EXPECT_TRUE(InterceptFunction("malloc", &malloc_address, 0, 0));
   EXPECT_NE(0U, malloc_address);
+  EXPECT_FALSE(InterceptFunction("malloc", &malloc_address, 0, 1));
 
   uptr dummy_address = 0;
-  EXPECT_TRUE(
-      GetRealFunctionAddress("dummy_doesnt_exist__", &dummy_address, 0, 0));
+  EXPECT_FALSE(InterceptFunction("dummy_doesnt_exist__", &dummy_address, 0, 0));
   EXPECT_EQ(0U, dummy_address);
 }
 
-TEST(Interception, GetFuncAddr) {
-  EXPECT_NE(GetFuncAddr("malloc"), nullptr);
-  EXPECT_EQ(GetFuncAddr("does_not_exist"), nullptr);
-}
-
 TEST(Interception, Basic) {
-  INTERCEPT_FUNCTION(isdigit);
+  EXPECT_TRUE(INTERCEPT_FUNCTION(isdigit));
 
   // After interception, the counter should be incremented.
   InterceptorFunctionCalled = 0;