Bug 951827 - Part 5: Tests for shared hooks. r=ehsan

toolkit/xre/test/win/TestDllInterceptor.cpp

@@ -38,6 +38,31 @@ patched_rotatePayload(payload p)
   return orig_rotatePayload(p);
 }
 
+__declspec(noinline) bool AlwaysTrue(int, int, int, int, int, int) {
+  // Dummy function that makes the caller recognizable by the detour patcher
+  return true;
+}
+
+extern "C" __declspec(dllexport) __declspec(noinline) uint32_t SetBits(uint32_t x)
+{
+  if (AlwaysTrue(1, 2, 3, 4, 5, 6)) {
+    return x | 0x11;
+  }
+  return 0;
+}
+
+static uint32_t (*orig_SetBits_early)(uint32_t);
+static uint32_t patched_SetBits_early(uint32_t x)
+{
+  return orig_SetBits_early(x) | 0x2200;
+}
+
+static uint32_t (*orig_SetBits_late)(uint32_t);
+static uint32_t patched_SetBits_late(uint32_t x)
+{
+  return orig_SetBits_late(x) | 0x330000;
+}
+
 bool TestHook(const char *dll, const char *func)
 {
   void *orig_func;
@@ -57,6 +82,25 @@ bool TestHook(const char *dll, const char *func)
   }
 }
 
+bool TestSharedHook(const char *dll, const char *func)
+{
+  void *orig_func;
+  bool successful = false;
+  {
+    WindowsDllInterceptor TestIntercept;
+    TestIntercept.Init(dll);
+    successful = TestIntercept.AddSharedHook(func, 0, &orig_func);
+  }
+
+  if (successful) {
+    printf("TEST-PASS | WindowsDllInterceptor | Could hook (shared) %s from %s\n", func, dll);
+    return true;
+  } else {
+    printf("TEST-UNEXPECTED-FAIL | WindowsDllInterceptor | Failed to hook (shared) %s from %s\n", func, dll);
+    return false;
+  }
+}
+
 int main()
 {
   payload initial = { 0x12345678, 0xfc4e9d31, 0x87654321 };
@@ -112,6 +156,66 @@ int main()
     return 1;
   }
 
+#ifdef _M_IX86
+  // The x64 detour patcher does understand the assembly code of SetBits.
+  // We only need these shared hook tests on x86 anyway, because shared hooks
+  // are the same as regular hooks on x64.
+
+  // Despite the noinline annotation, the compiler may try to re-use the
+  // return value of SetBits(0). Force it to call the function every time.
+  uint32_t (*volatile SetBitsVolatile)(uint32_t) = SetBits;
+
+  {
+    WindowsDllInterceptor ExeInterceptEarly;
+    ExeInterceptEarly.Init("TestDllInterceptor.exe");
+    if (ExeInterceptEarly.AddSharedHook("SetBits", reinterpret_cast<intptr_t>(patched_SetBits_early), (void**) &orig_SetBits_early)) {
+      printf("TEST-PASS | WindowsDllInterceptor | Early hook added\n");
+    } else {
+      printf("TEST-UNEXPECTED-FAIL | WindowsDllInterceptor | Failed to add early hook\n");
+      return 1;
+    }
+
+    if (SetBitsVolatile(0) == 0x2211) {
+      printf("TEST-PASS | WindowsDllInterceptor | Early hook was called\n");
+    } else {
+      printf("TEST-UNEXPECTED-FAIL | WindowsDllInterceptor | Early hook was not called\n");
+      return 1;
+    }
+
+    {
+      WindowsDllInterceptor ExeInterceptLate;
+      ExeInterceptLate.Init("TestDllInterceptor.exe");
+      if (ExeInterceptLate.AddHook("SetBits", reinterpret_cast<intptr_t>(patched_SetBits_late), (void**) &orig_SetBits_late)) {
+        printf("TEST-PASS | WindowsDllInterceptor | Late hook added\n");
+      } else {
+        printf("TEST-UNEXPECTED-FAIL | WindowsDllInterceptor | Failed to add late hook\n");
+        return 1;
+      }
+
+      if (SetBitsVolatile(0) == 0x332211) {
+        printf("TEST-PASS | WindowsDllInterceptor | Late hook was called\n");
+      } else {
+        printf("TEST-UNEXPECTED-FAIL | WindowsDllInterceptor | Late hook was not called\n");
+        return 1;
+      }
+    }
+
+    if (SetBitsVolatile(0) == 0x2211) {
+      printf("TEST-PASS | WindowsDllInterceptor | Late hook was unregistered\n");
+    } else {
+      printf("TEST-UNEXPECTED-FAIL | WindowsDllInterceptor | Late hook was not unregistered\n");
+      return 1;
+    }
+  }
+
+  if (SetBitsVolatile(0) == 0x11) {
+    printf("TEST-PASS | WindowsDllInterceptor | Early hook was unregistered\n");
+  } else {
+    printf("TEST-UNEXPECTED-FAIL | WindowsDllInterceptor | Early hook was not unregistered\n");
+    return 1;
+  }
+#endif
+
   if (TestHook("user32.dll", "GetWindowInfo") &&
 #ifdef _WIN64
       TestHook("user32.dll", "SetWindowLongPtrA") &&
@@ -138,6 +242,9 @@ int main()
       TestHook("kernel32.dll", "VirtualAlloc") &&
       TestHook("kernel32.dll", "MapViewOfFile") &&
       TestHook("gdi32.dll", "CreateDIBSection") &&
+#endif
+#ifdef _M_IX86 // Shared hooks are the same as regular hooks on x64
+      TestSharedHook("ntdll.dll", "LdrLoadDll") &&
 #endif
       TestHook("ntdll.dll", "LdrLoadDll")) {
     printf("TEST-PASS | WindowsDllInterceptor | all checks passed\n");