4 files changed, 40 insertions, 19 deletions
diff --git a/lib/interception/tests/CMakeLists.txt b/lib/interception/tests/CMakeLists.txt
index 1da0a455b..96bdda73e 100644
--- a/lib/interception/tests/CMakeLists.txt
+++ b/lib/interception/tests/CMakeLists.txt
@@ -21,6 +21,9 @@ set(INTERCEPTION_TEST_CFLAGS_COMMON
   -Werror=sign-compare
   -Wno-non-virtual-dtor)
 
+set(INTERCEPTION_TEST_LINK_FLAGS_COMMON
+  ${COMPILER_RT_UNITTEST_LINK_FLAGS})
+
 # -gline-tables-only must be enough for these tests, so use it if possible.
 if(COMPILER_RT_TEST_COMPILER_ID MATCHES "Clang")
   list(APPEND INTERCEPTION_TEST_CFLAGS_COMMON -gline-tables-only)
diff --git a/lib/interception/tests/interception_linux_test.cc b/lib/interception/tests/interception_linux_test.cc
index cc09aa09d..3e859cb1a 100644
--- a/lib/interception/tests/interception_linux_test.cc
+++ b/lib/interception/tests/interception_linux_test.cc
@@ -1,9 +1,8 @@
 //===-- interception_linux_test.cc ----------------------------------------===//
 //
-//                     The LLVM Compiler Infrastructure
-//
-// This file is distributed under the University of Illinois Open Source
-// License. See LICENSE.TXT for details.
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
@@ -34,19 +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, Basic) {
-  ASSERT_TRUE(INTERCEPT_FUNCTION(isdigit));
+  EXPECT_TRUE(INTERCEPT_FUNCTION(isdigit));
 
   // After interception, the counter should be incremented.
   InterceptorFunctionCalled = 0;
diff --git a/lib/interception/tests/interception_test_main.cc b/lib/interception/tests/interception_test_main.cc
index 311da51ec..97e5b377a 100644
--- a/lib/interception/tests/interception_test_main.cc
+++ b/lib/interception/tests/interception_test_main.cc
@@ -1,9 +1,8 @@
 //===-- interception_test_main.cc------------------------------------------===//
 //
-//                     The LLVM Compiler Infrastructure
-//
-// This file is distributed under the University of Illinois Open Source
-// License. See LICENSE.TXT for details.
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
diff --git a/lib/interception/tests/interception_win_test.cc b/lib/interception/tests/interception_win_test.cc
index 37ef994f8..c3affc45b 100644
--- a/lib/interception/tests/interception_win_test.cc
+++ b/lib/interception/tests/interception_win_test.cc
@@ -1,9 +1,8 @@
 //===-- interception_win_test.cc ------------------------------------------===//
 //
-//                     The LLVM Compiler Infrastructure
-//
-// This file is distributed under the University of Illinois Open Source
-// License. See LICENSE.TXT for details.
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
@@ -209,6 +208,24 @@ const u8 kUnpatchableCode6[] = {
     0x90, 0x90, 0x90, 0x90,
 };
 
+const u8 kPatchableCode6[] = {
+    0x48, 0x89, 0x54, 0x24, 0xBB, // mov QWORD PTR [rsp + 0xBB], rdx
+    0x33, 0xC9,                   // xor ecx,ecx
+    0xC3,                         // ret
+};
+
+const u8 kPatchableCode7[] = {
+    0x4c, 0x89, 0x4c, 0x24, 0xBB,  // mov QWORD PTR [rsp + 0xBB], r9
+    0x33, 0xC9,                   // xor ecx,ecx
+    0xC3,                         // ret
+};
+
+const u8 kPatchableCode8[] = {
+    0x4c, 0x89, 0x44, 0x24, 0xBB, // mov QWORD PTR [rsp + 0xBB], r8
+    0x33, 0xC9,                   // xor ecx,ecx
+    0xC3,                         // ret
+};
+
 // A buffer holding the dynamically generated code under test.
 u8* ActiveCode;
 const size_t ActiveCodeLength = 4096;
@@ -508,7 +525,6 @@ TEST(Interception, PatchableFunction) {
 #endif
   EXPECT_TRUE(TestFunctionPatching(kPatchableCode4, override));
   EXPECT_TRUE(TestFunctionPatching(kPatchableCode5, override));
-
 #if SANITIZER_WINDOWS64
   EXPECT_TRUE(TestFunctionPatching(kLoadGlobalCode, override));
 #endif
@@ -573,7 +589,11 @@ TEST(Interception, PatchableFunctionWithHotPatch) {
   EXPECT_FALSE(TestFunctionPatching(kPatchableCode2, override, prefix));
   EXPECT_FALSE(TestFunctionPatching(kPatchableCode3, override, prefix));
   EXPECT_FALSE(TestFunctionPatching(kPatchableCode4, override, prefix));
-
+#if SANITIZER_WINDOWS64
+  EXPECT_TRUE(TestFunctionPatching(kPatchableCode6, override, prefix));
+  EXPECT_TRUE(TestFunctionPatching(kPatchableCode7, override, prefix));
+  EXPECT_TRUE(TestFunctionPatching(kPatchableCode8, override, prefix));
+#endif
   EXPECT_FALSE(TestFunctionPatching(kUnpatchableCode1, override, prefix));
   EXPECT_TRUE(TestFunctionPatching(kUnpatchableCode2, override, prefix));
   EXPECT_FALSE(TestFunctionPatching(kUnpatchableCode3, override, prefix));