michael@0: /* -*- Mode: C++; tab-width: 40; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
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: #ifndef NS_WINDOWS_DLL_INTERCEPTOR_H_
michael@0: #define NS_WINDOWS_DLL_INTERCEPTOR_H_
michael@0: #include <windows.h>
michael@0: #include <winternl.h>
michael@0: 
michael@0: /*
michael@0:  * Simple function interception.
michael@0:  *
michael@0:  * We have two separate mechanisms for intercepting a function: We can use the
michael@0:  * built-in nop space, if it exists, or we can create a detour.
michael@0:  *
michael@0:  * Using the built-in nop space works as follows: On x86-32, DLL functions
michael@0:  * begin with a two-byte nop (mov edi, edi) and are preceeded by five bytes of
michael@0:  * NOP instructions.
michael@0:  *
michael@0:  * When we detect a function with this prelude, we do the following:
michael@0:  *
michael@0:  * 1. Write a long jump to our interceptor function into the five bytes of NOPs
michael@0:  *    before the function.
michael@0:  *
michael@0:  * 2. Write a short jump -5 into the two-byte nop at the beginning of the function.
michael@0:  *
michael@0:  * This mechanism is nice because it's thread-safe.  It's even safe to do if
michael@0:  * another thread is currently running the function we're modifying!
michael@0:  *
michael@0:  * When the WindowsDllNopSpacePatcher is destroyed, we overwrite the short jump
michael@0:  * but not the long jump, so re-intercepting the same function won't work,
michael@0:  * because its prelude won't match.
michael@0:  *
michael@0:  *
michael@0:  * Unfortunately nop space patching doesn't work on functions which don't have
michael@0:  * this magic prelude (and in particular, x86-64 never has the prelude).  So
michael@0:  * when we can't use the built-in nop space, we fall back to using a detour,
michael@0:  * which works as follows:
michael@0:  *
michael@0:  * 1. Save first N bytes of OrigFunction to trampoline, where N is a
michael@0:  *    number of bytes >= 5 that are instruction aligned.
michael@0:  *
michael@0:  * 2. Replace first 5 bytes of OrigFunction with a jump to the Hook
michael@0:  *    function.
michael@0:  *
michael@0:  * 3. After N bytes of the trampoline, add a jump to OrigFunction+N to
michael@0:  *    continue original program flow.
michael@0:  *
michael@0:  * 4. Hook function needs to call the trampoline during its execution,
michael@0:  *    to invoke the original function (so address of trampoline is
michael@0:  *    returned).
michael@0:  *
michael@0:  * When the WindowsDllDetourPatcher object is destructed, OrigFunction is
michael@0:  * patched again to jump directly to the trampoline instead of going through
michael@0:  * the hook function. As such, re-intercepting the same function won't work, as
michael@0:  * jump instructions are not supported.
michael@0:  *
michael@0:  * Note that this is not thread-safe.  Sad day.
michael@0:  *
michael@0:  */
michael@0: 
michael@0: #include <stdint.h>
michael@0: 
michael@0: namespace mozilla {
michael@0: namespace internal {
michael@0: 
michael@0: class WindowsDllNopSpacePatcher
michael@0: {
michael@0:   typedef unsigned char *byteptr_t;
michael@0:   HMODULE mModule;
michael@0: 
michael@0:   // Dumb array for remembering the addresses of functions we've patched.
michael@0:   // (This should be nsTArray, but non-XPCOM code uses this class.)
michael@0:   static const size_t maxPatchedFns = 128;
michael@0:   byteptr_t mPatchedFns[maxPatchedFns];
michael@0:   int mPatchedFnsLen;
michael@0: 
michael@0: public:
michael@0:   WindowsDllNopSpacePatcher()
michael@0:     : mModule(0)
michael@0:     , mPatchedFnsLen(0)
michael@0:   {}
michael@0: 
michael@0:   ~WindowsDllNopSpacePatcher()
michael@0:   {
michael@0:     // Restore the mov edi, edi to the beginning of each function we patched.
michael@0: 
michael@0:     for (int i = 0; i < mPatchedFnsLen; i++) {
michael@0:       byteptr_t fn = mPatchedFns[i];
michael@0: 
michael@0:       // Ensure we can write to the code.
michael@0:       DWORD op;
michael@0:       if (!VirtualProtectEx(GetCurrentProcess(), fn, 2, PAGE_EXECUTE_READWRITE, &op)) {
michael@0:         // printf("VirtualProtectEx failed! %d\n", GetLastError());
michael@0:         continue;
michael@0:       }
michael@0: 
michael@0:       // mov edi, edi
michael@0:       *((uint16_t*)fn) = 0xff8b;
michael@0: 
michael@0:       // Restore the old protection.
michael@0:       VirtualProtectEx(GetCurrentProcess(), fn, 2, op, &op);
michael@0: 
michael@0:       // I don't think this is actually necessary, but it can't hurt.
michael@0:       FlushInstructionCache(GetCurrentProcess(),
michael@0:                             /* ignored */ nullptr,
michael@0:                             /* ignored */ 0);
michael@0:     }
michael@0:   }
michael@0: 
michael@0:   void Init(const char *modulename)
michael@0:   {
michael@0:     mModule = LoadLibraryExA(modulename, nullptr, 0);
michael@0:     if (!mModule) {
michael@0:       //printf("LoadLibraryEx for '%s' failed\n", modulename);
michael@0:       return;
michael@0:     }
michael@0:   }
michael@0: 
michael@0: #if defined(_M_IX86)
michael@0:   bool AddHook(const char *pname, intptr_t hookDest, void **origFunc)
michael@0:   {
michael@0:     if (!mModule)
michael@0:       return false;
michael@0: 
michael@0:     if (mPatchedFnsLen == maxPatchedFns) {
michael@0:       // printf ("No space for hook in mPatchedFns.\n");
michael@0:       return false;
michael@0:     }
michael@0: 
michael@0:     byteptr_t fn = reinterpret_cast<byteptr_t>(GetProcAddress(mModule, pname));
michael@0:     if (!fn) {
michael@0:       //printf ("GetProcAddress failed\n");
michael@0:       return false;
michael@0:     }
michael@0:   
michael@0:     // Ensure we can read and write starting at fn - 5 (for the long jmp we're
michael@0:     // going to write) and ending at fn + 2 (for the short jmp up to the long
michael@0:     // jmp).
michael@0:     DWORD op;
michael@0:     if (!VirtualProtectEx(GetCurrentProcess(), fn - 5, 7, PAGE_EXECUTE_READWRITE, &op)) {
michael@0:       //printf ("VirtualProtectEx failed! %d\n", GetLastError());
michael@0:       return false;
michael@0:     }
michael@0: 
michael@0:     bool rv = WriteHook(fn, hookDest, origFunc);
michael@0:     
michael@0:     // Re-protect, and we're done.
michael@0:     VirtualProtectEx(GetCurrentProcess(), fn - 5, 7, op, &op);
michael@0: 
michael@0:     if (rv) {
michael@0:       mPatchedFns[mPatchedFnsLen] = fn;
michael@0:       mPatchedFnsLen++;
michael@0:     }
michael@0: 
michael@0:     return rv;
michael@0:   }
michael@0: 
michael@0:   bool WriteHook(byteptr_t fn, intptr_t hookDest, void **origFunc)
michael@0:   {
michael@0:     // Check that the 5 bytes before fn are NOP's or INT 3's,
michael@0:     // and that the 2 bytes after fn are mov(edi, edi).
michael@0:     //
michael@0:     // It's safe to read fn[-5] because we set it to PAGE_EXECUTE_READWRITE
michael@0:     // before calling WriteHook.
michael@0: 
michael@0:     for (int i = -5; i <= -1; i++) {
michael@0:       if (fn[i] != 0x90 && fn[i] != 0xcc) // nop or int 3
michael@0:         return false;
michael@0:     }
michael@0: 
michael@0:     // mov edi, edi.  Yes, there are two ways to encode the same thing:
michael@0:     //
michael@0:     //   0x89ff == mov r/m, r
michael@0:     //   0x8bff == mov r, r/m
michael@0:     //
michael@0:     // where "r" is register and "r/m" is register or memory.  Windows seems to
michael@0:     // use 8bff; I include 89ff out of paranoia.
michael@0:     if ((fn[0] != 0x8b && fn[0] != 0x89) || fn[1] != 0xff) {
michael@0:       return false;
michael@0:     }
michael@0: 
michael@0:     // Write a long jump into the space above the function.
michael@0:     fn[-5] = 0xe9; // jmp
michael@0:     *((intptr_t*)(fn - 4)) = hookDest - (uintptr_t)(fn); // target displacement
michael@0: 
michael@0:     // Set origFunc here, because after this point, hookDest might be called,
michael@0:     // and hookDest might use the origFunc pointer.
michael@0:     *origFunc = fn + 2;
michael@0: 
michael@0:     // Short jump up into our long jump.
michael@0:     *((uint16_t*)(fn)) = 0xf9eb; // jmp $-5
michael@0: 
michael@0:     // I think this routine is safe without this, but it can't hurt.
michael@0:     FlushInstructionCache(GetCurrentProcess(),
michael@0:                           /* ignored */ nullptr,
michael@0:                           /* ignored */ 0);
michael@0: 
michael@0:     return true;
michael@0:   }
michael@0: #else
michael@0:   bool AddHook(const char *pname, intptr_t hookDest, void **origFunc)
michael@0:   {
michael@0:     // Not implemented except on x86-32.
michael@0:     return false;
michael@0:   }
michael@0: #endif
michael@0: };
michael@0: 
michael@0: class WindowsDllDetourPatcher
michael@0: {
michael@0:   typedef unsigned char *byteptr_t;
michael@0: public:
michael@0:   WindowsDllDetourPatcher() 
michael@0:     : mModule(0), mHookPage(0), mMaxHooks(0), mCurHooks(0)
michael@0:   {
michael@0:   }
michael@0: 
michael@0:   ~WindowsDllDetourPatcher()
michael@0:   {
michael@0:     int i;
michael@0:     byteptr_t p;
michael@0:     for (i = 0, p = mHookPage; i < mCurHooks; i++, p += kHookSize) {
michael@0: #if defined(_M_IX86)
michael@0:       size_t nBytes = 1 + sizeof(intptr_t);
michael@0: #elif defined(_M_X64)
michael@0:       size_t nBytes = 2 + sizeof(intptr_t);
michael@0: #else
michael@0: #error "Unknown processor type"
michael@0: #endif
michael@0:       byteptr_t origBytes = *((byteptr_t *)p);
michael@0:       // ensure we can modify the original code
michael@0:       DWORD op;
michael@0:       if (!VirtualProtectEx(GetCurrentProcess(), origBytes, nBytes, PAGE_EXECUTE_READWRITE, &op)) {
michael@0:         //printf ("VirtualProtectEx failed! %d\n", GetLastError());
michael@0:         continue;
michael@0:       }
michael@0:       // Remove the hook by making the original function jump directly
michael@0:       // in the trampoline.
michael@0:       intptr_t dest = (intptr_t)(p + sizeof(void *));
michael@0: #if defined(_M_IX86)
michael@0:       *((intptr_t*)(origBytes+1)) = dest - (intptr_t)(origBytes+5); // target displacement
michael@0: #elif defined(_M_X64)
michael@0:       *((intptr_t*)(origBytes+2)) = dest;
michael@0: #else
michael@0: #error "Unknown processor type"
michael@0: #endif
michael@0:       // restore protection; if this fails we can't really do anything about it
michael@0:       VirtualProtectEx(GetCurrentProcess(), origBytes, nBytes, op, &op);
michael@0:     }
michael@0:   }
michael@0: 
michael@0:   void Init(const char *modulename, int nhooks = 0)
michael@0:   {
michael@0:     if (mModule)
michael@0:       return;
michael@0: 
michael@0:     mModule = LoadLibraryExA(modulename, nullptr, 0);
michael@0:     if (!mModule) {
michael@0:       //printf("LoadLibraryEx for '%s' failed\n", modulename);
michael@0:       return;
michael@0:     }
michael@0: 
michael@0:     int hooksPerPage = 4096 / kHookSize;
michael@0:     if (nhooks == 0)
michael@0:       nhooks = hooksPerPage;
michael@0: 
michael@0:     mMaxHooks = nhooks + (hooksPerPage % nhooks);
michael@0: 
michael@0:     mHookPage = (byteptr_t) VirtualAllocEx(GetCurrentProcess(), nullptr,
michael@0:              mMaxHooks * kHookSize,
michael@0:              MEM_COMMIT | MEM_RESERVE,
michael@0:              PAGE_EXECUTE_READWRITE);
michael@0: 
michael@0:     if (!mHookPage) {
michael@0:       mModule = 0;
michael@0:       return;
michael@0:     }
michael@0:   }
michael@0: 
michael@0:   bool Initialized()
michael@0:   {
michael@0:     return !!mModule;
michael@0:   }
michael@0: 
michael@0:   void LockHooks()
michael@0:   {
michael@0:     if (!mModule)
michael@0:       return;
michael@0: 
michael@0:     DWORD op;
michael@0:     VirtualProtectEx(GetCurrentProcess(), mHookPage, mMaxHooks * kHookSize, PAGE_EXECUTE_READ, &op);
michael@0: 
michael@0:     mModule = 0;
michael@0:   }
michael@0: 
michael@0:   bool AddHook(const char *pname, intptr_t hookDest, void **origFunc)
michael@0:   {
michael@0:     if (!mModule)
michael@0:       return false;
michael@0: 
michael@0:     void *pAddr = (void *) GetProcAddress(mModule, pname);
michael@0:     if (!pAddr) {
michael@0:       //printf ("GetProcAddress failed\n");
michael@0:       return false;
michael@0:     }
michael@0: 
michael@0:     CreateTrampoline(pAddr, hookDest, origFunc);
michael@0:     if (!*origFunc) {
michael@0:       //printf ("CreateTrampoline failed\n");
michael@0:       return false;
michael@0:     }
michael@0: 
michael@0:     return true;
michael@0:   }
michael@0: 
michael@0: protected:
michael@0:   const static int kPageSize = 4096;
michael@0:   const static int kHookSize = 128;
michael@0: 
michael@0:   HMODULE mModule;
michael@0:   byteptr_t mHookPage;
michael@0:   int mMaxHooks;
michael@0:   int mCurHooks;
michael@0: 
michael@0:   void CreateTrampoline(void *origFunction,
michael@0:                         intptr_t dest,
michael@0:                         void **outTramp)
michael@0:   {
michael@0:     *outTramp = nullptr;
michael@0: 
michael@0:     byteptr_t tramp = FindTrampolineSpace();
michael@0:     if (!tramp)
michael@0:       return;
michael@0: 
michael@0:     byteptr_t origBytes = (byteptr_t) origFunction;
michael@0: 
michael@0:     int nBytes = 0;
michael@0:     int pJmp32 = -1;
michael@0: 
michael@0: #if defined(_M_IX86)
michael@0:     while (nBytes < 5) {
michael@0:       // Understand some simple instructions that might be found in a
michael@0:       // prologue; we might need to extend this as necessary.
michael@0:       //
michael@0:       // Note!  If we ever need to understand jump instructions, we'll
michael@0:       // need to rewrite the displacement argument.
michael@0:       if (origBytes[nBytes] >= 0x88 && origBytes[nBytes] <= 0x8B) {
michael@0:         // various MOVs
michael@0:         unsigned char b = origBytes[nBytes+1];
michael@0:         if (((b & 0xc0) == 0xc0) ||
michael@0:             (((b & 0xc0) == 0x00) &&
michael@0:              ((b & 0x07) != 0x04) && ((b & 0x07) != 0x05)))
michael@0:         {
michael@0:           // REG=r, R/M=r or REG=r, R/M=[r]
michael@0:           nBytes += 2;
michael@0:         } else if (((b & 0xc0) == 0x40) && ((b & 0x38) != 0x20)) {
michael@0:           // REG=r, R/M=[r + disp8]
michael@0:           nBytes += 3;
michael@0:         } else {
michael@0:           // complex MOV, bail
michael@0:           return;
michael@0:         }
michael@0:       } else if (origBytes[nBytes] == 0xB8) {
michael@0:         // MOV 0xB8: http://ref.x86asm.net/coder32.html#xB8
michael@0:         nBytes += 5;
michael@0:       } else if (origBytes[nBytes] == 0x83) {
michael@0:         // ADD|ODR|ADC|SBB|AND|SUB|XOR|CMP r/m, imm8
michael@0:         unsigned char b = origBytes[nBytes+1];
michael@0:         if ((b & 0xc0) == 0xc0) {
michael@0:           // ADD|ODR|ADC|SBB|AND|SUB|XOR|CMP r, imm8
michael@0:           nBytes += 3;
michael@0:         } else {
michael@0:           // bail
michael@0:           return;
michael@0:         }
michael@0:       } else if (origBytes[nBytes] == 0x68) {
michael@0:         // PUSH with 4-byte operand
michael@0:         nBytes += 5;
michael@0:       } else if ((origBytes[nBytes] & 0xf0) == 0x50) {
michael@0:         // 1-byte PUSH/POP
michael@0:         nBytes++;
michael@0:       } else if (origBytes[nBytes] == 0x6A) {
michael@0:         // PUSH imm8
michael@0:         nBytes += 2;
michael@0:       } else if (origBytes[nBytes] == 0xe9) {
michael@0:         pJmp32 = nBytes;
michael@0:         // jmp 32bit offset
michael@0:         nBytes += 5;
michael@0:       } else {
michael@0:         //printf ("Unknown x86 instruction byte 0x%02x, aborting trampoline\n", origBytes[nBytes]);
michael@0:         return;
michael@0:       }
michael@0:     }
michael@0: #elif defined(_M_X64)
michael@0:     byteptr_t directJmpAddr;
michael@0: 
michael@0:     while (nBytes < 13) {
michael@0: 
michael@0:       // if found JMP 32bit offset, next bytes must be NOP 
michael@0:       if (pJmp32 >= 0) {
michael@0:         if (origBytes[nBytes++] != 0x90)
michael@0:           return;
michael@0: 
michael@0:         continue;
michael@0:       } 
michael@0:       if (origBytes[nBytes] == 0x0f) {
michael@0:         nBytes++;
michael@0:         if (origBytes[nBytes] == 0x1f) {
michael@0:           // nop (multibyte)
michael@0:           nBytes++;
michael@0:           if ((origBytes[nBytes] & 0xc0) == 0x40 &&
michael@0:               (origBytes[nBytes] & 0x7) == 0x04) {
michael@0:             nBytes += 3;
michael@0:           } else {
michael@0:             return;
michael@0:           }
michael@0:         } else if (origBytes[nBytes] == 0x05) {
michael@0:           // syscall
michael@0:           nBytes++;
michael@0:         } else {
michael@0:           return;
michael@0:         }
michael@0:       } else if (origBytes[nBytes] == 0x41) {
michael@0:         // REX.B
michael@0:         nBytes++;
michael@0: 
michael@0:         if ((origBytes[nBytes] & 0xf0) == 0x50) {
michael@0:           // push/pop with Rx register
michael@0:           nBytes++;
michael@0:         } else if (origBytes[nBytes] >= 0xb8 && origBytes[nBytes] <= 0xbf) {
michael@0:           // mov r32, imm32
michael@0:           nBytes += 5;
michael@0:         } else {
michael@0:           return;
michael@0:         }
michael@0:       } else if (origBytes[nBytes] == 0x45) {
michael@0:         // REX.R & REX.B
michael@0:         nBytes++;
michael@0: 
michael@0:         if (origBytes[nBytes] == 0x33) {
michael@0:           // xor r32, r32
michael@0:           nBytes += 2;
michael@0:         } else {
michael@0:           return;
michael@0:         }
michael@0:       } else if ((origBytes[nBytes] & 0xfb) == 0x48) {
michael@0:         // REX.W | REX.WR
michael@0:         nBytes++;
michael@0: 
michael@0:         if (origBytes[nBytes] == 0x81 && (origBytes[nBytes+1] & 0xf8) == 0xe8) {
michael@0:           // sub r, dword
michael@0:           nBytes += 6;
michael@0:         } else if (origBytes[nBytes] == 0x83 &&
michael@0:                   (origBytes[nBytes+1] & 0xf8) == 0xe8) {
michael@0:           // sub r, byte
michael@0:           nBytes += 3;
michael@0:         } else if (origBytes[nBytes] == 0x83 &&
michael@0:                   (origBytes[nBytes+1] & 0xf8) == 0x60) {
michael@0:           // and [r+d], imm8
michael@0:           nBytes += 5;
michael@0:         } else if ((origBytes[nBytes] & 0xfd) == 0x89) {
michael@0:           // MOV r/m64, r64 | MOV r64, r/m64
michael@0:           if ((origBytes[nBytes+1] & 0xc0) == 0x40) {
michael@0:             if ((origBytes[nBytes+1] & 0x7) == 0x04) {
michael@0:               // R/M=[SIB+disp8], REG=r64
michael@0:               nBytes += 4;
michael@0:             } else {
michael@0:               // R/M=[r64+disp8], REG=r64
michael@0:               nBytes += 3;
michael@0:             }
michael@0:           } else if (((origBytes[nBytes+1] & 0xc0) == 0xc0) ||
michael@0:                      (((origBytes[nBytes+1] & 0xc0) == 0x00) &&
michael@0:                       ((origBytes[nBytes+1] & 0x07) != 0x04) && ((origBytes[nBytes+1] & 0x07) != 0x05))) {
michael@0:             // REG=r64, R/M=r64 or REG=r64, R/M=[r64]
michael@0:             nBytes += 2;
michael@0:           } else {
michael@0:             // complex MOV
michael@0:             return;
michael@0:           }
michael@0:         } else if (origBytes[nBytes] == 0xc7) {
michael@0:           // MOV r/m64, imm32
michael@0:           if (origBytes[nBytes + 1] == 0x44) {
michael@0:             // MOV [r64+disp8], imm32
michael@0:             // ModR/W + SIB + disp8 + imm32
michael@0:             nBytes += 8;
michael@0:           } else {
michael@0:             return;
michael@0:           }
michael@0:         } else if (origBytes[nBytes] == 0xff) {
michael@0:           pJmp32 = nBytes - 1;
michael@0:           // JMP /4
michael@0:           if ((origBytes[nBytes+1] & 0xc0) == 0x0 &&
michael@0:               (origBytes[nBytes+1] & 0x07) == 0x5) {
michael@0:             // [rip+disp32]
michael@0:             // convert JMP 32bit offset to JMP 64bit direct
michael@0:             directJmpAddr = (byteptr_t)*((uint64_t*)(origBytes + nBytes + 6 + (*((int32_t*)(origBytes + nBytes + 2)))));
michael@0:             nBytes += 6;
michael@0:           } else {
michael@0:             // not support yet!
michael@0:             return;
michael@0:           }
michael@0:         } else {
michael@0:           // not support yet!
michael@0:           return;
michael@0:         }
michael@0:       } else if ((origBytes[nBytes] & 0xf0) == 0x50) {
michael@0:         // 1-byte push/pop
michael@0:         nBytes++;
michael@0:       } else if (origBytes[nBytes] == 0x90) {
michael@0:         // nop
michael@0:         nBytes++;
michael@0:       } else if (origBytes[nBytes] == 0xb8) {
michael@0:         // MOV 0xB8: http://ref.x86asm.net/coder32.html#xB8
michael@0:         nBytes += 5;
michael@0:       } else if (origBytes[nBytes] == 0xc3) {
michael@0:         // ret
michael@0:         nBytes++;
michael@0:       } else if (origBytes[nBytes] == 0xe9) {
michael@0:         pJmp32 = nBytes;
michael@0:         // convert JMP 32bit offset to JMP 64bit direct
michael@0:         directJmpAddr = origBytes + pJmp32 + 5 + (*((int32_t*)(origBytes + pJmp32 + 1)));
michael@0:         // jmp 32bit offset
michael@0:         nBytes += 5;
michael@0:       } else if (origBytes[nBytes] == 0xff) {
michael@0:         nBytes++;
michael@0:         if ((origBytes[nBytes] & 0xf8) == 0xf0) {
michael@0:           // push r64
michael@0:           nBytes++;
michael@0:         } else {
michael@0:           return;
michael@0:         }
michael@0:       } else {
michael@0:         return;
michael@0:       }
michael@0:     }
michael@0: #else
michael@0: #error "Unknown processor type"
michael@0: #endif
michael@0: 
michael@0:     if (nBytes > 100) {
michael@0:       //printf ("Too big!");
michael@0:       return;
michael@0:     }
michael@0: 
michael@0:     // We keep the address of the original function in the first bytes of
michael@0:     // the trampoline buffer
michael@0:     *((void **)tramp) = origFunction;
michael@0:     tramp += sizeof(void *);
michael@0: 
michael@0:     memcpy(tramp, origFunction, nBytes);
michael@0: 
michael@0:     // OrigFunction+N, the target of the trampoline
michael@0:     byteptr_t trampDest = origBytes + nBytes;
michael@0: 
michael@0: #if defined(_M_IX86)
michael@0:     if (pJmp32 >= 0) {
michael@0:       // Jump directly to the original target of the jump instead of jumping to the
michael@0:       // original function.
michael@0:       // Adjust jump target displacement to jump location in the trampoline.
michael@0:       *((intptr_t*)(tramp+pJmp32+1)) += origBytes - tramp;
michael@0:     } else {
michael@0:       tramp[nBytes] = 0xE9; // jmp
michael@0:       *((intptr_t*)(tramp+nBytes+1)) = (intptr_t)trampDest - (intptr_t)(tramp+nBytes+5); // target displacement
michael@0:     }
michael@0: #elif defined(_M_X64)
michael@0:     // If JMP32 opcode found, we don't insert to trampoline jump 
michael@0:     if (pJmp32 >= 0) {
michael@0:       // mov r11, address
michael@0:       tramp[pJmp32]   = 0x49;
michael@0:       tramp[pJmp32+1] = 0xbb;
michael@0:       *((intptr_t*)(tramp+pJmp32+2)) = (intptr_t)directJmpAddr;
michael@0: 
michael@0:       // jmp r11
michael@0:       tramp[pJmp32+10] = 0x41;
michael@0:       tramp[pJmp32+11] = 0xff;
michael@0:       tramp[pJmp32+12] = 0xe3;
michael@0:     } else {
michael@0:       // mov r11, address
michael@0:       tramp[nBytes] = 0x49;
michael@0:       tramp[nBytes+1] = 0xbb;
michael@0:       *((intptr_t*)(tramp+nBytes+2)) = (intptr_t)trampDest;
michael@0: 
michael@0:       // jmp r11
michael@0:       tramp[nBytes+10] = 0x41;
michael@0:       tramp[nBytes+11] = 0xff;
michael@0:       tramp[nBytes+12] = 0xe3;
michael@0:     }
michael@0: #endif
michael@0: 
michael@0:     // The trampoline is now valid.
michael@0:     *outTramp = tramp;
michael@0: 
michael@0:     // ensure we can modify the original code
michael@0:     DWORD op;
michael@0:     if (!VirtualProtectEx(GetCurrentProcess(), origFunction, nBytes, PAGE_EXECUTE_READWRITE, &op)) {
michael@0:       //printf ("VirtualProtectEx failed! %d\n", GetLastError());
michael@0:       return;
michael@0:     }
michael@0: 
michael@0: #if defined(_M_IX86)
michael@0:     // now modify the original bytes
michael@0:     origBytes[0] = 0xE9; // jmp
michael@0:     *((intptr_t*)(origBytes+1)) = dest - (intptr_t)(origBytes+5); // target displacement
michael@0: #elif defined(_M_X64)
michael@0:     // mov r11, address
michael@0:     origBytes[0] = 0x49;
michael@0:     origBytes[1] = 0xbb;
michael@0: 
michael@0:     *((intptr_t*)(origBytes+2)) = dest;
michael@0: 
michael@0:     // jmp r11
michael@0:     origBytes[10] = 0x41;
michael@0:     origBytes[11] = 0xff;
michael@0:     origBytes[12] = 0xe3;
michael@0: #endif
michael@0: 
michael@0:     // restore protection; if this fails we can't really do anything about it
michael@0:     VirtualProtectEx(GetCurrentProcess(), origFunction, nBytes, op, &op);
michael@0:   }
michael@0: 
michael@0:   byteptr_t FindTrampolineSpace()
michael@0:   {
michael@0:     if (mCurHooks >= mMaxHooks)
michael@0:       return 0;
michael@0: 
michael@0:     byteptr_t p = mHookPage + mCurHooks*kHookSize;
michael@0: 
michael@0:     mCurHooks++;
michael@0: 
michael@0:     return p;
michael@0:   }
michael@0: };
michael@0: 
michael@0: } // namespace internal
michael@0: 
michael@0: class WindowsDllInterceptor
michael@0: {
michael@0:   internal::WindowsDllNopSpacePatcher mNopSpacePatcher;
michael@0:   internal::WindowsDllDetourPatcher mDetourPatcher;
michael@0: 
michael@0:   const char *mModuleName;
michael@0:   int mNHooks;
michael@0: 
michael@0: public:
michael@0:   WindowsDllInterceptor()
michael@0:     : mModuleName(nullptr)
michael@0:     , mNHooks(0)
michael@0:   {}
michael@0: 
michael@0:   void Init(const char *moduleName, int nhooks = 0)
michael@0:   {
michael@0:     if (mModuleName) {
michael@0:       return;
michael@0:     }
michael@0: 
michael@0:     mModuleName = moduleName;
michael@0:     mNHooks = nhooks;
michael@0:     mNopSpacePatcher.Init(moduleName);
michael@0: 
michael@0:     // Lazily initialize mDetourPatcher, since it allocates memory and we might
michael@0:     // not need it.
michael@0:   }
michael@0: 
michael@0:   void LockHooks()
michael@0:   {
michael@0:     if (mDetourPatcher.Initialized())
michael@0:       mDetourPatcher.LockHooks();
michael@0:   }
michael@0: 
michael@0:   bool AddHook(const char *pname, intptr_t hookDest, void **origFunc)
michael@0:   {
michael@0:     // Use a nop space patch if possible, otherwise fall back to a detour.
michael@0:     // This should be the preferred method for adding hooks.
michael@0: 
michael@0:     if (!mModuleName) {
michael@0:       return false;
michael@0:     }
michael@0: 
michael@0:     if (mNopSpacePatcher.AddHook(pname, hookDest, origFunc)) {
michael@0:       return true;
michael@0:     }
michael@0: 
michael@0:     return AddDetour(pname, hookDest, origFunc);
michael@0:   }
michael@0: 
michael@0:   bool AddDetour(const char *pname, intptr_t hookDest, void **origFunc)
michael@0:   {
michael@0:     // Generally, code should not call this method directly. Use AddHook unless
michael@0:     // there is a specific need to avoid nop space patches.
michael@0: 
michael@0:     if (!mModuleName) {
michael@0:       return false;
michael@0:     }
michael@0: 
michael@0:     if (!mDetourPatcher.Initialized()) {
michael@0:       mDetourPatcher.Init(mModuleName, mNHooks);
michael@0:     }
michael@0: 
michael@0:     return mDetourPatcher.AddHook(pname, hookDest, origFunc);
michael@0:   }
michael@0: };
michael@0: 
michael@0: } // namespace mozilla
michael@0: 
michael@0: #endif /* NS_WINDOWS_DLL_INTERCEPTOR_H_ */