toolkit/xre/nsWindowsDllInterceptor.h

   1 /* -*- Mode: C++; tab-width: 40; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
   2 /* ***** BEGIN LICENSE BLOCK *****
   3  * Version: MPL 1.1/GPL 2.0/LGPL 2.1
   4  *
   5  * The contents of this file are subject to the Mozilla Public License Version
   6  * 1.1 (the "License"); you may not use this file except in compliance with
   7  * the License. You may obtain a copy of the License at
   8  * http://www.mozilla.org/MPL/
   9  *
  10  * Software distributed under the License is distributed on an "AS IS" basis,
  11  * WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License
  12  * for the specific language governing rights and limitations under the
  13  * License.
  14  *
  15  * The Original Code is mozilla.org code.
  16  *
  17  * The Initial Developer of the Original Code is
  18  *   Mozilla Corp
  19  * Portions created by the Initial Developer are Copyright (C) 2009
  20  * the Initial Developer. All Rights Reserved.
  21  *
  22  * Contributor(s):
  23  *  Vladimir Vukicevic <vladimir@pobox.com>
  24  *
  25  * Alternatively, the contents of this file may be used under the terms of
  26  * either the GNU General Public License Version 2 or later (the "GPL"), or
  27  * the GNU Lesser General Public License Version 2.1 or later (the "LGPL"),
  28  * in which case the provisions of the GPL or the LGPL are applicable instead
  29  * of those above. If you wish to allow use of your version of this file only
  30  * under the terms of either the GPL or the LGPL, and not to allow others to
  31  * use your version of this file under the terms of the MPL, indicate your
  32  * decision by deleting the provisions above and replace them with the notice
  33  * and other provisions required by the GPL or the LGPL. If you do not delete
  34  * the provisions above, a recipient may use your version of this file under
  35  * the terms of any one of the MPL, the GPL or the LGPL.
  36  *
  37  * ***** END LICENSE BLOCK ***** */
  38
  39 #ifndef NS_WINDOWS_DLL_INTERCEPTOR_H_
  40 #define NS_WINDOWS_DLL_INTERCEPTOR_H_
  41 #include <windows.h>
  42 #include <winternl.h>
  43
  44 /*
  45  * Simple trampoline interception
  46  *
  47  * 1. Save first N bytes of OrigFunction to trampoline, where N is a
  48  *    number of bytes >= 5 that are instruction aligned.
  49  *
  50  * 2. Replace first 5 bytes of OrigFunction with a jump to the Hook
  51  *    function.
  52  *
  53  * 3. After N bytes of the trampoline, add a jump to OrigFunction+N to
  54  *    continue original program flow.
  55  *
  56  * 4. Hook function needs to call the trampoline during its execution,
  57  *    to invoke the original function (so address of trampoline is
  58  *    returned).
  59  *
  60  */
  61
  62 class WindowsDllInterceptor
  63 {
  64   typedef unsigned char *byteptr_t;
  65 public:
  66   WindowsDllInterceptor()
  67     : mModule(0)
  68   {
  69   }
  70
  71   WindowsDllInterceptor(const char *modulename, int nhooks = 0) {
  72     Init(modulename, nhooks);
  73   }
  74
  75   void Init(const char *modulename, int nhooks = 0) {
  76     if (mModule)
  77       return;
  78
  79     mModule = LoadLibraryExA(modulename, NULL, 0);
  80     if (!mModule) {
  81       //printf("LoadLibraryEx for '%s' failed\n", modulename);
  82       return;
  83     }
  84
  85     int hooksPerPage = 4096 / kHookSize;
  86     if (nhooks == 0)
  87       nhooks = hooksPerPage;
  88
  89     mMaxHooks = nhooks + (hooksPerPage % nhooks);
  90     mCurHooks = 0;
  91
  92     mHookPage = (byteptr_t) VirtualAllocEx(GetCurrentProcess(), NULL, mMaxHooks * kHookSize,
  93              MEM_COMMIT | MEM_RESERVE,
  94              PAGE_EXECUTE_READWRITE);
  95
  96     if (!mHookPage) {
  97       mModule = 0;
  98       return;
  99     }
 100   }
 101
 102   void LockHooks() {
 103     if (!mModule)
 104       return;
 105
 106     DWORD op;
 107     VirtualProtectEx(GetCurrentProcess(), mHookPage, mMaxHooks * kHookSize, PAGE_EXECUTE_READ, &op);
 108
 109     mModule = 0;
 110   }
 111
 112   bool AddHook(const char *pname,
 113          void *hookDest,
 114          void **origFunc)
 115   {
 116     if (!mModule)
 117       return false;
 118
 119     void *pAddr = (void *) GetProcAddress(mModule, pname);
 120     if (!pAddr) {
 121       //printf ("GetProcAddress failed\n");
 122       return false;
 123     }
 124
 125     void *tramp = CreateTrampoline(pAddr, hookDest);
 126     if (!tramp) {
 127       //printf ("CreateTrampoline failed\n");
 128       return false;
 129     }
 130
 131     *origFunc = tramp;
 132
 133     return true;
 134   }
 135
 136 protected:
 137   const static int kPageSize = 4096;
 138   const static int kHookSize = 128;
 139
 140   HMODULE mModule;
 141   byteptr_t mHookPage;
 142   int mMaxHooks;
 143   int mCurHooks;
 144
 145   byteptr_t CreateTrampoline(void *origFunction,
 146            void *dest)
 147   {
 148     byteptr_t tramp = FindTrampolineSpace();
 149     if (!tramp)
 150       return 0;
 151
 152     byteptr_t origBytes = (byteptr_t) origFunction;
 153
 154     int nBytes = 0;
 155     while (nBytes < 5) {
 156       // Understand some simple instructions that might be found in a
 157       // prologue; we might need to extend this as necessary.
 158       //
 159       // Note!  If we ever need to understand jump instructions, we'll
 160       // need to rewrite the displacement argument.
 161       if (origBytes[nBytes] >= 0x88 && origBytes[nBytes] <= 0x8B) {
 162         // various MOVs; but only handle the case where it truly is a 2-byte instruction
 163         unsigned char b = origBytes[nBytes+1];
 164         if (((b & 0xc0) == 0xc0) ||
 165             (((b & 0xc0) == 0x00) &&
 166              ((b & 0x38) != 0x20) && ((b & 0x38) != 0x28)))
 167         {
 168           nBytes += 2;
 169         } else {
 170           // complex MOV, bail
 171           return 0;
 172         }
 173       } else if (origBytes[nBytes] == 0x68) {
 174         // PUSH with 4-byte operand
 175         nBytes += 5;
 176       } else if ((origBytes[nBytes] & 0xf0) == 0x50) {
 177         // 1-byte PUSH/POP
 178         nBytes++;
 179       } else {
 180         //printf ("Unknown x86 instruction byte 0x%02x, aborting trampoline\n", origBytes[nBytes]);
 181         return 0;
 182       }
 183     }
 184
 185     if (nBytes > 100) {
 186       //printf ("Too big!");
 187       return 0;
 188     }
 189
 190     memcpy(tramp, origFunction, nBytes);
 191
 192     // OrigFunction+N, the target of the trampoline
 193     byteptr_t trampDest = origBytes + nBytes;
 194
 195     tramp[nBytes] = 0xE9; // jmp
 196     *((intptr_t*)(tramp+nBytes+1)) = (intptr_t)trampDest - (intptr_t)(tramp+nBytes+5); // target displacement
 197
 198     // ensure we can modify the original code
 199     DWORD op;
 200     if (!VirtualProtectEx(GetCurrentProcess(), origFunction, nBytes, PAGE_EXECUTE_READWRITE, &op)) {
 201       //printf ("VirtualProtectEx failed! %d\n", GetLastError());
 202       return 0;
 203     }
 204
 205     // now modify the original bytes
 206     origBytes[0] = 0xE9; // jmp
 207     *((intptr_t*)(origBytes+1)) = (intptr_t)dest - (intptr_t)(origBytes+5); // target displacement
 208
 209     // restore protection; if this fails we can't really do anything about it
 210     VirtualProtectEx(GetCurrentProcess(), origFunction, nBytes, op, &op);
 211
 212     return tramp;
 213   }
 214
 215   byteptr_t FindTrampolineSpace() {
 216     if (mCurHooks >= mMaxHooks)
 217       return 0;
 218
 219     byteptr_t p = mHookPage + mCurHooks*kHookSize;
 220
 221     mCurHooks++;
 222
 223     return p;
 224   }
 225 };
 226
 227
 228 #endif /* NS_WINDOWS_DLL_INTERCEPTOR_H_ */