michael@0: /* This Source Code Form is subject to the terms of the Mozilla Public
michael@0:  * License, v. 2.0. If a copy of the MPL was not distributed with this
michael@0:  * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
michael@0: 
michael@0: #include <stdio.h>
michael@0: #include "nsWindowsDllInterceptor.h"
michael@0: 
michael@0: using namespace mozilla;
michael@0: 
michael@0: struct payload {
michael@0:   UINT64 a;
michael@0:   UINT64 b;
michael@0:   UINT64 c;
michael@0: 
michael@0:   bool operator==(const payload &other) const {
michael@0:     return (a == other.a &&
michael@0:             b == other.b &&
michael@0:             c == other.c);
michael@0:   }
michael@0: };
michael@0: 
michael@0: extern "C" __declspec(dllexport) __declspec(noinline) payload rotatePayload(payload p) {
michael@0:   UINT64 tmp = p.a;
michael@0:   p.a = p.b;
michael@0:   p.b = p.c;
michael@0:   p.c = tmp;
michael@0:   return p;
michael@0: }
michael@0: 
michael@0: static bool patched_func_called = false;
michael@0: 
michael@0: static payload (*orig_rotatePayload)(payload);
michael@0: 
michael@0: static payload
michael@0: patched_rotatePayload(payload p)
michael@0: {
michael@0:   patched_func_called = true;
michael@0:   return orig_rotatePayload(p);
michael@0: }
michael@0: 
michael@0: bool TestHook(const char *dll, const char *func)
michael@0: {
michael@0:   void *orig_func;
michael@0:   bool successful = false;
michael@0:   {
michael@0:     WindowsDllInterceptor TestIntercept;
michael@0:     TestIntercept.Init(dll);
michael@0:     successful = TestIntercept.AddHook(func, 0, &orig_func);
michael@0:   }
michael@0: 
michael@0:   if (successful) {
michael@0:     printf("TEST-PASS | WindowsDllInterceptor | Could hook %s from %s\n", func, dll);
michael@0:     return true;
michael@0:   } else {
michael@0:     printf("TEST-UNEXPECTED-FAIL | WindowsDllInterceptor | Failed to hook %s from %s\n", func, dll);
michael@0:     return false;
michael@0:   }
michael@0: }
michael@0: 
michael@0: bool TestDetour(const char *dll, const char *func)
michael@0: {
michael@0:   void *orig_func;
michael@0:   bool successful = false;
michael@0:   {
michael@0:     WindowsDllInterceptor TestIntercept;
michael@0:     TestIntercept.Init(dll);
michael@0:     successful = TestIntercept.AddDetour(func, 0, &orig_func);
michael@0:   }
michael@0: 
michael@0:   if (successful) {
michael@0:     printf("TEST-PASS | WindowsDllInterceptor | Could detour %s from %s\n", func, dll);
michael@0:     return true;
michael@0:   } else {
michael@0:     printf("TEST-UNEXPECTED-FAIL | WindowsDllInterceptor | Failed to detour %s from %s\n", func, dll);
michael@0:     return false;
michael@0:   }
michael@0: }
michael@0: 
michael@0: int main()
michael@0: {
michael@0:   payload initial = { 0x12345678, 0xfc4e9d31, 0x87654321 };
michael@0:   payload p0, p1;
michael@0:   ZeroMemory(&p0, sizeof(p0));
michael@0:   ZeroMemory(&p1, sizeof(p1));
michael@0: 
michael@0:   p0 = rotatePayload(initial);
michael@0: 
michael@0:   {
michael@0:     WindowsDllInterceptor ExeIntercept;
michael@0:     ExeIntercept.Init("TestDllInterceptor.exe");
michael@0:     if (ExeIntercept.AddHook("rotatePayload", reinterpret_cast<intptr_t>(patched_rotatePayload), (void**) &orig_rotatePayload)) {
michael@0:       printf("TEST-PASS | WindowsDllInterceptor | Hook added\n");
michael@0:     } else {
michael@0:       printf("TEST-UNEXPECTED-FAIL | WindowsDllInterceptor | Failed to add hook\n");
michael@0:       return 1;
michael@0:     }
michael@0: 
michael@0:     p1 = rotatePayload(initial);
michael@0: 
michael@0:     if (patched_func_called) {
michael@0:       printf("TEST-PASS | WindowsDllInterceptor | Hook called\n");
michael@0:     } else {
michael@0:       printf("TEST-UNEXPECTED-FAIL | WindowsDllInterceptor | Hook was not called\n");
michael@0:       return 1;
michael@0:     }
michael@0: 
michael@0:     if (p0 == p1) {
michael@0:       printf("TEST-PASS | WindowsDllInterceptor | Hook works properly\n");
michael@0:     } else {
michael@0:       printf("TEST-UNEXPECTED-FAIL | WindowsDllInterceptor | Hook didn't return the right information\n");
michael@0:       return 1;
michael@0:     }
michael@0:   }
michael@0: 
michael@0:   patched_func_called = false;
michael@0:   ZeroMemory(&p1, sizeof(p1));
michael@0: 
michael@0:   p1 = rotatePayload(initial);
michael@0: 
michael@0:   if (!patched_func_called) {
michael@0:     printf("TEST-PASS | WindowsDllInterceptor | Hook was not called after unregistration\n");
michael@0:   } else {
michael@0:     printf("TEST-UNEXPECTED-FAIL | WindowsDllInterceptor | Hook was still called after unregistration\n");
michael@0:     return 1;
michael@0:   }
michael@0: 
michael@0:   if (p0 == p1) {
michael@0:     printf("TEST-PASS | WindowsDllInterceptor | Original function worked properly\n");
michael@0:   } else {
michael@0:     printf("TEST-UNEXPECTED-FAIL | WindowsDllInterceptor | Original function didn't return the right information\n");
michael@0:     return 1;
michael@0:   }
michael@0: 
michael@0:   if (TestHook("user32.dll", "GetWindowInfo") &&
michael@0: #ifdef _WIN64
michael@0:       TestHook("user32.dll", "SetWindowLongPtrA") &&
michael@0:       TestHook("user32.dll", "SetWindowLongPtrW") &&
michael@0: #else
michael@0:       TestHook("user32.dll", "SetWindowLongA") &&
michael@0:       TestHook("user32.dll", "SetWindowLongW") &&
michael@0: #endif
michael@0:       TestHook("user32.dll", "TrackPopupMenu") &&
michael@0: #ifdef _M_IX86
michael@0:       // We keep this test to hook complex code on x86. (Bug 850957)
michael@0:       TestHook("ntdll.dll", "NtFlushBuffersFile") &&
michael@0: #endif
michael@0:       TestHook("ntdll.dll", "NtCreateFile") &&
michael@0:       TestHook("ntdll.dll", "NtReadFile") &&
michael@0:       TestHook("ntdll.dll", "NtReadFileScatter") &&
michael@0:       TestHook("ntdll.dll", "NtWriteFile") &&
michael@0:       TestHook("ntdll.dll", "NtWriteFileGather") &&
michael@0:       TestHook("ntdll.dll", "NtQueryFullAttributesFile") &&
michael@0:       // Bug 733892: toolkit/crashreporter/nsExceptionHandler.cpp
michael@0:       TestHook("kernel32.dll", "SetUnhandledExceptionFilter") &&
michael@0: #ifdef _M_IX86
michael@0:       // Bug 670967: xpcom/base/AvailableMemoryTracker.cpp
michael@0:       TestHook("kernel32.dll", "VirtualAlloc") &&
michael@0:       TestHook("kernel32.dll", "MapViewOfFile") &&
michael@0:       TestHook("gdi32.dll", "CreateDIBSection") &&
michael@0: #endif
michael@0:       TestDetour("ntdll.dll", "LdrLoadDll")) {
michael@0:     printf("TEST-PASS | WindowsDllInterceptor | all checks passed\n");
michael@0:     return 0;
michael@0:   }
michael@0: 
michael@0:   return 1;
michael@0: }