development/Win32/vld/src/utility.cpp

   1 ////////////////////////////////////////////////////////////////////////////////
   2 //  $Id: utility.cpp,v 1.24 2006/11/18 03:12:35 dmouldin Exp $
   3 //
   4 //  Visual Leak Detector - Various Utility Functions
   5 //  Copyright (c) 2005-2006 Dan Moulding
   6 //
   7 //  This library is free software; you can redistribute it and/or
   8 //  modify it under the terms of the GNU Lesser General Public
   9 //  License as published by the Free Software Foundation; either
  10 //  version 2.1 of the License, or (at your option) any later version.
  11 //
  12 //  This library is distributed in the hope that it will be useful,
  13 //  but WITHOUT ANY WARRANTY; without even the implied warranty of
  14 //  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  15 //  Lesser General Public License for more details.
  16 //
  17 //  You should have received a copy of the GNU Lesser General Public
  18 //  License along with this library; if not, write to the Free Software
  19 //  Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
  20 //
  21 //  See COPYING.txt for the full terms of the GNU Lesser General Public License.
  22 //
  23 ////////////////////////////////////////////////////////////////////////////////
  24
  25 #include <cassert>
  26 #include <cstdio>
  27 #include <windows.h>
  28 #define __out_xcount(x) // Workaround for the specstrings.h bug in the Platform SDK.
  29 #define DBGHELP_TRANSLATE_TCHAR
  30 #include <dbghelp.h>    // Provides portable executable (PE) image access functions.
  31 #define VLDBUILD        // Declares that we are building Visual Leak Detector.
  32 #include "utility.h"    // Provides various utility functions and macros.
  33 #include "vldheap.h"    // Provides internal new and delete operators.
  34
  35 // Imported Global Variables
  36 extern CRITICAL_SECTION imagelock;
  37
  38 // Global variables.
  39 static BOOL        reportdelay = FALSE;     // If TRUE, we sleep for a bit after calling OutputDebugString to give the debugger time to catch up.
  40 static FILE       *reportfile = NULL;       // Pointer to the file, if any, to send the memory leak report to.
  41 static BOOL        reporttodebugger = TRUE; // If TRUE, a copy of the memory leak report will be sent to the debugger for display.
  42 static encoding_e  reportencoding = ascii;  // Output encoding of the memory leak report.
  43
  44 // dumpmemorya - Dumps a nicely formatted rendition of a region of memory.
  45 //   Includes both the hex value of each byte and its ASCII equivalent (if
  46 //   printable).
  47 //
  48 //  - address (IN): Pointer to the beginning of the memory region to dump.
  49 //
  50 //  - size (IN): The size, in bytes, of the region to dump.
  51 //
  52 //  Return Value:
  53 //
  54 //    None.
  55 //
  56 VOID dumpmemorya (LPCVOID address, SIZE_T size)
  57 {
  58     WCHAR  ascdump [18] = {0};
  59     SIZE_T ascindex;
  60     BYTE   byte;
  61     SIZE_T byteindex;
  62     SIZE_T bytesdone;
  63     SIZE_T dumplen;
  64     WCHAR  formatbuf [BYTEFORMATBUFFERLENGTH];
  65     WCHAR  hexdump [HEXDUMPLINELENGTH] = {0};
  66     SIZE_T hexindex;
  67
  68     // Each line of output is 16 bytes.
  69     if ((size % 16) == 0) {
  70         // No padding needed.
  71         dumplen = size;
  72     }
  73     else {
  74         // We'll need to pad the last line out to 16 bytes.
  75         dumplen = size + (16 - (size % 16));
  76     }
  77
  78     // For each byte of data, get both the ASCII equivalent (if it is a
  79     // printable character) and the hex representation.
  80     bytesdone = 0;
  81     for (byteindex = 0; byteindex < dumplen; byteindex++) {
  82         hexindex = 3 * ((byteindex % 16) + ((byteindex % 16) / 4)); // 3 characters per byte, plus a 3-character space after every 4 bytes.
  83         ascindex = (byteindex % 16) + (byteindex % 16) / 8;         // 1 character per byte, plus a 1-character space after every 8 bytes.
  84         if (byteindex < size) {
  85             byte = ((PBYTE)address)[byteindex];
  86             _snwprintf_s(formatbuf, BYTEFORMATBUFFERLENGTH, _TRUNCATE, L"%.2X ", byte);
  87             formatbuf[3] = '\0';
  88             wcsncpy_s(hexdump + hexindex, HEXDUMPLINELENGTH - hexindex, formatbuf, 4);
  89             if (isgraph(byte)) {
  90                 ascdump[ascindex] = (WCHAR)byte;
  91             }
  92             else {
  93                 ascdump[ascindex] = L'.';
  94             }
  95         }
  96         else {
  97             // Add padding to fill out the last line to 16 bytes.
  98             wcsncpy_s(hexdump + hexindex, HEXDUMPLINELENGTH - hexindex, L"   ", 4);
  99             ascdump[ascindex] = L'.';
 100         }
 101         bytesdone++;
 102         if ((bytesdone % 16) == 0) {
 103             // Print one line of data for every 16 bytes. Include the
 104             // ASCII dump and the hex dump side-by-side.
 105             report(L"    %s    %s\n", hexdump, ascdump);
 106         }
 107         else {
 108             if ((bytesdone % 8) == 0) {
 109                 // Add a spacer in the ASCII dump after every 8 bytes.
 110                 ascdump[ascindex + 1] = L' ';
 111             }
 112             if ((bytesdone % 4) == 0) {
 113                 // Add a spacer in the hex dump after every 4 bytes.
 114                 wcsncpy_s(hexdump + hexindex + 3, HEXDUMPLINELENGTH - hexindex - 3, L"   ", 4);
 115             }
 116         }
 117     }
 118 }
 119
 120 // dumpmemoryw - Dumps a nicely formatted rendition of a region of memory.
 121 //   Includes both the hex value of each byte and its Unicode equivalent.
 122 //
 123 //  - address (IN): Pointer to the beginning of the memory region to dump.
 124 //
 125 //  - size (IN): The size, in bytes, of the region to dump.
 126 //
 127 //  Return Value:
 128 //
 129 //    None.
 130 //
 131 VOID dumpmemoryw (LPCVOID address, SIZE_T size)
 132 {
 133     BYTE   byte;
 134     SIZE_T byteindex;
 135     SIZE_T bytesdone;
 136     SIZE_T dumplen;
 137     WCHAR  formatbuf [BYTEFORMATBUFFERLENGTH];
 138     WCHAR  hexdump [HEXDUMPLINELENGTH] = {0};
 139     SIZE_T hexindex;
 140     WORD   word;
 141     WCHAR  unidump [18] = {0};
 142     SIZE_T uniindex;
 143
 144     // Each line of output is 16 bytes.
 145     if ((size % 16) == 0) {
 146         // No padding needed.
 147         dumplen = size;
 148     }
 149     else {
 150         // We'll need to pad the last line out to 16 bytes.
 151         dumplen = size + (16 - (size % 16));
 152     }
 153
 154     // For each word of data, get both the Unicode equivalent and the hex
 155     // representation.
 156     bytesdone = 0;
 157     for (byteindex = 0; byteindex < dumplen; byteindex++) {
 158         hexindex = 3 * ((byteindex % 16) + ((byteindex % 16) / 4));   // 3 characters per byte, plus a 3-character space after every 4 bytes.
 159         uniindex = ((byteindex / 2) % 8) + ((byteindex / 2) % 8) / 8; // 1 character every other byte, plus a 1-character space after every 8 bytes.
 160         if (byteindex < size) {
 161             byte = ((PBYTE)address)[byteindex];
 162             _snwprintf_s(formatbuf, BYTEFORMATBUFFERLENGTH, _TRUNCATE, L"%.2X ", byte);
 163             formatbuf[BYTEFORMATBUFFERLENGTH - 1] = '\0';
 164             wcsncpy_s(hexdump + hexindex, HEXDUMPLINELENGTH - hexindex, formatbuf, 4);
 165             if (((byteindex % 2) == 0) && ((byteindex + 1) < dumplen)) {
 166                 // On every even byte, print one character.
 167                 word = ((PWORD)address)[byteindex / 2];
 168                 if ((word == 0x0000) || (word == 0x0020)) {
 169                     unidump[uniindex] = L'.';
 170                 }
 171                 else {
 172                     unidump[uniindex] = word;
 173                 }
 174             }
 175         }
 176         else {
 177             // Add padding to fill out the last line to 16 bytes.
 178             wcsncpy_s(hexdump + hexindex, HEXDUMPLINELENGTH - hexindex, L"   ", 4);
 179             unidump[uniindex] = L'.';
 180         }
 181         bytesdone++;
 182         if ((bytesdone % 16) == 0) {
 183             // Print one line of data for every 16 bytes. Include the
 184             // ASCII dump and the hex dump side-by-side.
 185             report(L"    %s    %s\n", hexdump, unidump);
 186         }
 187         else {
 188             if ((bytesdone % 8) == 0) {
 189                 // Add a spacer in the ASCII dump after every 8 bytes.
 190                 unidump[uniindex + 1] = L' ';
 191             }
 192             if ((bytesdone % 4) == 0) {
 193                 // Add a spacer in the hex dump after every 4 bytes.
 194                 wcsncpy_s(hexdump + hexindex + 3, HEXDUMPLINELENGTH - hexindex - 3, L"   ", 4);
 195             }
 196         }
 197     }
 198 }
 199
 200 // findimport - Determines if the specified module imports the named import
 201 //   from the named exporting module.
 202 //
 203 //  - importmodule (IN): Handle (base address) of the module to be searched to
 204 //      see if it imports the specified import.
 205 //
 206 //  - exportmodule (IN): Handle (base address) of the module that exports the
 207 //      import to be searched for.
 208 //
 209 //  - exportmodulename (IN): ANSI string containing the name of the module that
 210 //      exports the import to be searched for.
 211 //
 212 //  - importname (IN): ANSI string containing the name of the import to search
 213 //      for. May be an integer cast to a string if the import is exported by
 214 //      ordinal.
 215 //
 216 //  Return Value:
 217 //
 218 //    Returns TRUE if the module imports to the specified import. Otherwise
 219 //    returns FALSE.
 220 //
 221 BOOL findimport (HMODULE importmodule, HMODULE exportmodule, LPCSTR exportmodulename, LPCSTR importname)
 222 {
 223     IMAGE_THUNK_DATA        *iate;
 224     IMAGE_IMPORT_DESCRIPTOR *idte;
 225     FARPROC                  import;
 226     IMAGE_SECTION_HEADER    *section;
 227     ULONG                    size;
 228
 229     // Locate the importing module's Import Directory Table (IDT) entry for the
 230     // exporting module. The importing module actually can have several IATs --
 231     // one for each export module that it imports something from. The IDT entry
 232     // gives us the offset of the IAT for the module we are interested in.
 233     EnterCriticalSection(&imagelock);
 234     idte = (IMAGE_IMPORT_DESCRIPTOR*)ImageDirectoryEntryToDataEx((PVOID)importmodule, TRUE,
 235                                                                   IMAGE_DIRECTORY_ENTRY_IMPORT, &size, &section);
 236     LeaveCriticalSection(&imagelock);
 237     if (idte == NULL) {
 238         // This module has no IDT (i.e. it imports nothing).
 239         return FALSE;
 240     }
 241     while (idte->OriginalFirstThunk != 0x0) {
 242         if (_stricmp((PCHAR)R2VA(importmodule, idte->Name), exportmodulename) == 0) {
 243             // Found the IDT entry for the exporting module.
 244             break;
 245         }
 246         idte++;
 247     }
 248     if (idte->OriginalFirstThunk == 0x0) {
 249         // The importing module does not import anything from the exporting
 250         // module.
 251         return FALSE;
 252     }
 253
 254     // Get the *real* address of the import. If we find this address in the IAT,
 255     // then we've found that the module does import the named import.
 256     import = GetProcAddress(exportmodule, importname);
 257     assert(import != NULL); // Perhaps the named export module does not actually export the named import?
 258
 259     // Locate the import's IAT entry.
 260     iate = (IMAGE_THUNK_DATA*)R2VA(importmodule, idte->FirstThunk);
 261     while (iate->u1.Function != 0x0) {
 262         if (iate->u1.Function == (DWORD_PTR)import) {
 263             // Found the IAT entry. The module imports the named import.
 264             return TRUE;
 265         }
 266         iate++;
 267     }
 268
 269     // The module does not import the named import.
 270     return FALSE;
 271 }
 272
 273 // findpatch - Determines if the specified module has been patched to use the
 274 //   specified replacement.
 275 //
 276 //  - importmodule (IN): Handle (base address) of the module to be searched to
 277 //      see if it imports the specified replacement export.
 278 //
 279 //  - exportmodulename (IN): ANSI string containing the name of the module that
 280 //      normally exports that import that would have been patched to use the
 281 //      replacement export.
 282 //
 283 //  - replacement (IN): Address of the replacement, or destination, function or
 284 //      variable to search for.
 285 //
 286 //  Return Value:
 287 //
 288 //    Returns TRUE if the module has been patched to use the specified
 289 //    replacement export.
 290 //
 291 BOOL findpatch (HMODULE importmodule, LPCSTR exportmodulename, LPCVOID replacement)
 292 {
 293     IMAGE_THUNK_DATA        *iate;
 294     IMAGE_IMPORT_DESCRIPTOR *idte;
 295     IMAGE_SECTION_HEADER    *section;
 296     ULONG                    size;
 297
 298     // Locate the importing module's Import Directory Table (IDT) entry for the
 299     // exporting module. The importing module actually can have several IATs --
 300     // one for each export module that it imports something from. The IDT entry
 301     // gives us the offset of the IAT for the module we are interested in.
 302     EnterCriticalSection(&imagelock);
 303     idte = (IMAGE_IMPORT_DESCRIPTOR*)ImageDirectoryEntryToDataEx((PVOID)importmodule, TRUE,
 304                                                                   IMAGE_DIRECTORY_ENTRY_IMPORT, &size, &section);
 305     LeaveCriticalSection(&imagelock);
 306     if (idte == NULL) {
 307         // This module has no IDT (i.e. it imports nothing).
 308         return FALSE;
 309     }
 310     while (idte->OriginalFirstThunk != 0x0) {
 311         if (_stricmp((PCHAR)R2VA(importmodule, idte->Name), exportmodulename) == 0) {
 312             // Found the IDT entry for the exporting module.
 313             break;
 314         }
 315         idte++;
 316     }
 317     if (idte->OriginalFirstThunk == 0x0) {
 318         // The importing module does not import anything from the exporting
 319         // module.
 320         return FALSE;
 321     }
 322
 323     // Locate the replacement's IAT entry.
 324     iate = (IMAGE_THUNK_DATA*)R2VA(importmodule, idte->FirstThunk);
 325     while (iate->u1.Function != 0x0) {
 326         if (iate->u1.Function == (DWORD_PTR)replacement) {
 327             // Found the IAT entry for the replacement. This patch has been
 328             // installed.
 329             return TRUE;
 330         }
 331         iate++;
 332     }
 333
 334     // The module does not import the replacement. The patch has not been
 335     // installed.
 336     return FALSE;
 337 }
 338
 339 // insertreportdelay - Sets the report function to sleep for a bit after each
 340 //   call to OutputDebugString, in order to allow the debugger to catch up.
 341 //
 342 //  Return Value:
 343 //
 344 //    None.
 345 //
 346 VOID insertreportdelay ()
 347 {
 348     reportdelay = TRUE;
 349 }
 350
 351 // moduleispatched - Checks to see if any of the imports listed in the specified
 352 //   patch table have been patched into the specified importmodule.
 353 //
 354 //  - importmodule (IN): Handle (base address) of the module to be queried to
 355 //      determine if it has been patched.
 356 //
 357 //  - patchtable (IN): An array of patchentry_t structures specifying all of the
 358 //      import patches to search for.
 359 //
 360 //  - tablesize (IN): Size, in entries, of the specified patch table.
 361 //
 362 //  Return Value:
 363 //
 364 //    Returns TRUE if at least one of the patches listed in the patch table is
 365 //    installed in the importmodule. Otherwise returns FALSE.
 366 //
 367 BOOL moduleispatched (HMODULE importmodule, patchentry_t patchtable [], UINT tablesize)
 368 {
 369     patchentry_t *entry;
 370     BOOL          found = FALSE;
 371     UINT          index;
 372
 373     // Loop through the import patch table, individually checking each patch
 374     // entry to see if it is installed in the import module. If any patch entry
 375     // is installed in the import module, then the module has been patched.
 376     for (index = 0; index < tablesize; index++) {
 377         entry = &patchtable[index];
 378         found = findpatch(importmodule, entry->exportmodulename, entry->replacement);
 379         if (found == TRUE) {
 380             // Found one of the listed patches installed in the import module.
 381             return TRUE;
 382         }
 383     }
 384
 385     // No patches listed in the patch table were found in the import module.
 386     return FALSE;
 387 }
 388
 389 // patchimport - Patches all future calls to an imported function, or references
 390 //   to an imported variable, through to a replacement function or variable.
 391 //   Patching is done by replacing the import's address in the specified target
 392 //   module's Import Address Table (IAT) with the address of the replacement
 393 //   function or variable.
 394 //
 395 //  - importmodule (IN): Handle (base address) of the target module for which
 396 //      calls or references to the import should be patched.
 397 //
 398 //  - exportmodule (IN): Handle (base address) of the module that exports the
 399 //      the function or variable to be patched.
 400 //
 401 //  - exportmodulename (IN): ANSI string containing the name of the module that
 402 //      exports the function or variable to be patched.
 403 //
 404 //  - importname (IN): ANSI string containing the name of the imported function
 405 //      or variable to be patched. May be an integer cast to a string if the
 406 //      import is exported by ordinal.
 407 //
 408 //  - replacement (IN): Address of the function or variable to which future
 409 //      calls or references should be patched through to. This function or
 410 //      variable can be thought of as effectively replacing the original import
 411 //      from the point of view of the module specified by "importmodule".
 412 //
 413 //  Return Value:
 414 //
 415 //    Returns TRUE if the patch was installed into the import module. If the
 416 //    import module does not import the specified export, so nothing changed,
 417 //    then FALSE will be returned.
 418 //
 419 BOOL patchimport (HMODULE importmodule, HMODULE exportmodule, LPCSTR exportmodulename, LPCSTR importname,
 420                   LPCVOID replacement)
 421 {
 422     IMAGE_THUNK_DATA        *iate;
 423     IMAGE_IMPORT_DESCRIPTOR *idte;
 424     FARPROC                  import;
 425     DWORD                    protect;
 426     IMAGE_SECTION_HEADER    *section;
 427     ULONG                    size;
 428
 429     // Locate the importing module's Import Directory Table (IDT) entry for the
 430     // exporting module. The importing module actually can have several IATs --
 431     // one for each export module that it imports something from. The IDT entry
 432     // gives us the offset of the IAT for the module we are interested in.
 433     EnterCriticalSection(&imagelock);
 434     idte = (IMAGE_IMPORT_DESCRIPTOR*)ImageDirectoryEntryToDataEx((PVOID)importmodule, TRUE,
 435                                                                   IMAGE_DIRECTORY_ENTRY_IMPORT, &size, &section);
 436     LeaveCriticalSection(&imagelock);
 437     if (idte == NULL) {
 438         // This module has no IDT (i.e. it imports nothing).
 439         return FALSE;
 440     }
 441     while (idte->OriginalFirstThunk != 0x0) {
 442         if (_stricmp((PCHAR)R2VA(importmodule, idte->Name), exportmodulename) == 0) {
 443             // Found the IDT entry for the exporting module.
 444             break;
 445         }
 446         idte++;
 447     }
 448     if (idte->OriginalFirstThunk == 0x0) {
 449         // The importing module does not import anything from the exporting
 450         // module.
 451         return FALSE;
 452     }
 453
 454     // Get the *real* address of the import. If we find this address in the IAT,
 455     // then we've found the entry that needs to be patched.
 456     import = GetProcAddress(exportmodule, importname);
 457     assert(import != NULL); // Perhaps the named export module does not actually export the named import?
 458
 459     // Locate the import's IAT entry.
 460     iate = (IMAGE_THUNK_DATA*)R2VA(importmodule, idte->FirstThunk);
 461     while (iate->u1.Function != 0x0) {
 462         if (iate->u1.Function == (DWORD_PTR)import) {
 463             // Found the IAT entry. Overwrite the address stored in the IAT
 464             // entry with the address of the replacement. Note that the IAT
 465             // entry may be write-protected, so we must first ensure that it is
 466             // writable.
 467             VirtualProtect(&iate->u1.Function, sizeof(iate->u1.Function), PAGE_READWRITE, &protect);
 468             iate->u1.Function = (DWORD_PTR)replacement;
 469             VirtualProtect(&iate->u1.Function, sizeof(iate->u1.Function), protect, &protect);
 470
 471             // The patch has been installed in the import module.
 472             return TRUE;
 473         }
 474         iate++;
 475     }
 476
 477     // The import's IAT entry was not found. The importing module does not
 478     // import the specified import.
 479     return FALSE;
 480 }
 481
 482 // patchmodule - Patches all imports listed in the supplied patch table, and
 483 //   which are imported by the specified module, through to their respective
 484 //   replacement functions.
 485 //
 486 //   Note: If the specified module does not import any of the functions listed
 487 //     in the patch table, then nothing is changed for the specified module.
 488 //
 489 //  - importmodule (IN): Handle (base address) of the target module which is to
 490 //      have its imports patched.
 491 //
 492 //  - patchtable (IN): An array of patchentry_t structures specifying all of the
 493 //      imports to patch for the specified module.
 494 //
 495 //  - tablesize (IN): Size, in entries, of the specified patch table.
 496 //
 497 //  Return Value:
 498 //
 499 //    Returns TRUE if at least one of the patches listed in the patch table was
 500 //    installed in the importmodule. Otherwise returns FALSE.
 501 //
 502 BOOL patchmodule (HMODULE importmodule, patchentry_t patchtable [], UINT tablesize)
 503 {
 504     patchentry_t *entry;
 505     UINT          index;
 506     BOOL          patched = FALSE;
 507
 508     // Loop through the import patch table, individually patching each import
 509     // listed in the table.
 510     for (index = 0; index < tablesize; index++) {
 511         entry = &patchtable[index];
 512         if (patchimport(importmodule, (HMODULE)entry->modulebase, entry->exportmodulename, entry->importname,
 513                         entry->replacement) == TRUE) {
 514             patched = TRUE;
 515         }
 516     }
 517
 518     return patched;
 519 }
 520
 521 // report - Sends a printf-style formatted message to the debugger for display
 522 //   and/or to a file.
 523 //
 524 //   Note: A message longer than MAXREPORTLENGTH characters will be truncated
 525 //     to MAXREPORTLENGTH.
 526 //
 527 //  - format (IN): Specifies a printf-compliant format string containing the
 528 //      message to be sent to the debugger.
 529 //
 530 //  - ... (IN): Arguments to be formatted using the specified format string.
 531 //
 532 //  Return Value:
 533 //
 534 //    None.
 535 //
 536 VOID report (LPCWSTR format, ...)
 537 {
 538     va_list args;
 539     size_t  count;
 540     CHAR    messagea [MAXREPORTLENGTH + 1];
 541     WCHAR   messagew [MAXREPORTLENGTH + 1];
 542
 543     va_start(args, format);
 544     _vsnwprintf_s(messagew, MAXREPORTLENGTH + 1, _TRUNCATE, format, args);
 545     va_end(args);
 546     messagew[MAXREPORTLENGTH] = L'\0';
 547
 548     if (reportencoding == unicode) {
 549         if (reportfile != NULL) {
 550             // Send the report to the previously specified file.
 551             fwrite(messagew, sizeof(WCHAR), wcslen(messagew), reportfile);
 552         }
 553         if (reporttodebugger) {
 554             OutputDebugStringW(messagew);
 555         }
 556     }
 557     else {
 558         if (wcstombs_s(&count, messagea, MAXREPORTLENGTH + 1, messagew, _TRUNCATE) == -1) {
 559             // Failed to convert the Unicode message to ASCII.
 560             assert(FALSE);
 561             return;
 562         }
 563         messagea[MAXREPORTLENGTH] = '\0';
 564         if (reportfile != NULL) {
 565             // Send the report to the previously specified file.
 566             fwrite(messagea, sizeof(CHAR), strlen(messagea), reportfile);
 567         }
 568         if (reporttodebugger) {
 569             OutputDebugStringA(messagea);
 570         }
 571     }
 572
 573     if (reporttodebugger && (reportdelay == TRUE)) {
 574         Sleep(10); // Workaround the Visual Studio 6 bug where debug strings are sometimes lost if they're sent too fast.
 575     }
 576 }
 577
 578 // restoreimport - Restores the IAT entry for an import previously patched via
 579 //   a call to "patchimport" to the original address of the import.
 580 //
 581 //  - importmodule (IN): Handle (base address) of the target module for which
 582 //      calls or references to the import should be restored.
 583 //
 584 //  - exportmodule (IN): Handle (base address) of the module that exports the
 585 //      function or variable to be patched.
 586 //
 587 //  - exportmodulename (IN): ANSI string containing the name of the module that
 588 //      exports the function or variable to be patched.
 589 //
 590 //  - importname (IN): ANSI string containing the name of the imported function
 591 //      or variable to be restored. May be an integer cast to a string if the
 592 //      import is exported by ordinal.
 593 //
 594 //  - replacement (IN): Address of the function or variable which the import was
 595 //      previously patched through to via a call to "patchimport".
 596 //
 597 //  Return Value:
 598 //
 599 //    None.
 600 //
 601 VOID restoreimport (HMODULE importmodule, HMODULE exportmodule, LPCSTR exportmodulename, LPCSTR importname,
 602                     LPCVOID replacement)
 603 {
 604     IMAGE_THUNK_DATA        *iate;
 605     IMAGE_IMPORT_DESCRIPTOR *idte;
 606     FARPROC                  import;
 607     DWORD                    protect;
 608     IMAGE_SECTION_HEADER    *section;
 609     ULONG                    size;
 610
 611     // Locate the importing module's Import Directory Table (IDT) entry for the
 612     // exporting module. The importing module actually can have several IATs --
 613     // one for each export module that it imports something from. The IDT entry
 614     // gives us the offset of the IAT for the module we are interested in.
 615     EnterCriticalSection(&imagelock);
 616     idte = (IMAGE_IMPORT_DESCRIPTOR*)ImageDirectoryEntryToDataEx((PVOID)importmodule, TRUE,
 617                                                                   IMAGE_DIRECTORY_ENTRY_IMPORT, &size, &section);
 618     LeaveCriticalSection(&imagelock);
 619     if (idte == NULL) {
 620         // This module has no IDT (i.e. it imports nothing).
 621         return;
 622     }
 623     while (idte->OriginalFirstThunk != 0x0) {
 624         if (_stricmp((PCHAR)R2VA(importmodule, idte->Name), exportmodulename) == 0) {
 625             // Found the IDT entry for the exporting module.
 626             break;
 627         }
 628         idte++;
 629     }
 630     if (idte->OriginalFirstThunk == 0x0) {
 631         // The importing module does not import anything from the exporting
 632         // module.
 633         return;
 634     }
 635
 636     // Get the *real* address of the import.
 637     import = GetProcAddress(exportmodule, importname);
 638     assert(import != NULL); // Perhaps the named export module does not actually export the named import?
 639
 640     // Locate the import's original IAT entry (it currently has the replacement
 641     // address in it).
 642     iate = (IMAGE_THUNK_DATA*)R2VA(importmodule, idte->FirstThunk);
 643     while (iate->u1.Function != 0x0) {
 644         if (iate->u1.Function == (DWORD_PTR)replacement) {
 645             // Found the IAT entry. Overwrite the address stored in the IAT
 646             // entry with the import's real address. Note that the IAT entry may
 647             // be write-protected, so we must first ensure that it is writable.
 648             VirtualProtect(&iate->u1.Function, sizeof(iate->u1.Function), PAGE_READWRITE, &protect);
 649             iate->u1.Function = (DWORD_PTR)import;
 650             VirtualProtect(&iate->u1.Function, sizeof(iate->u1.Function), protect, &protect);
 651             break;
 652         }
 653         iate++;
 654     }
 655 }
 656
 657 // restoremodule - Restores all imports listed in the supplied patch table, and
 658 //   which are imported by the specified module, to their original functions.
 659 //
 660 //   Note: If the specified module does not import any of the functions listed
 661 //     in the patch table, then nothing is changed for the specified module.
 662 //
 663 //  - importmodule (IN): Handle (base address) of the target module which is to
 664 //      have its imports restored.
 665 //
 666 //  - patchtable (IN): Array of patchentry_t structures specifying all of the
 667 //      imports to restore for the specified module.
 668 //
 669 //  - tablesize (IN): Size, in entries, of the specified patch table.
 670 //
 671 //  Return Value:
 672 //
 673 //    None.
 674 //
 675 VOID restoremodule (HMODULE importmodule, patchentry_t patchtable [], UINT tablesize)
 676 {
 677     patchentry_t *entry;
 678     UINT          index;
 679
 680     // Loop through the import patch table, individually restoring each import
 681     // listed in the table.
 682     for (index = 0; index < tablesize; index++) {
 683         entry = &patchtable[index];
 684         restoreimport(importmodule, (HMODULE)entry->modulebase, entry->exportmodulename, entry->importname,
 685                       entry->replacement);
 686     }
 687 }
 688
 689 // setreportencoding - Sets the output encoding of report messages to either
 690 //   ASCII (the default) or Unicode.
 691 //
 692 //  - encoding (IN): Specifies either "ascii" or "unicode".
 693 //
 694 //  Return Value:
 695 //
 696 //    None.
 697 //
 698 VOID setreportencoding (encoding_e encoding)
 699 {
 700     switch (encoding) {
 701     case ascii:
 702     case unicode:
 703         reportencoding = encoding;
 704         break;
 705
 706     default:
 707         assert(FALSE);
 708     }
 709 }
 710
 711 // setreportfile - Sets a destination file to which all report messages should
 712 //   be sent. If this function is not called to set a destination file, then
 713 //   report messages will be sent to the debugger instead of to a file.
 714 //
 715 //  - file (IN): Pointer to an open file, to which future report messages should
 716 //      be sent.
 717 //
 718 //  - copydebugger (IN): If true, in addition to sending report messages to
 719 //      the specified file, a copy of each message will also be sent to the
 720 //      debugger.
 721 //
 722 //  Return Value:
 723 //
 724 //    None.
 725 //
 726 VOID setreportfile (FILE *file, BOOL copydebugger)
 727 {
 728     reportfile = file;
 729     reporttodebugger = copydebugger;
 730 }
 731
 732 // strapp - Appends the specified source string to the specified destination
 733 //   string. Allocates additional space so that the destination string "grows"
 734 //   as new strings are appended to it. This function is fairly infrequently
 735 //   used so efficiency is not a major concern.
 736 //
 737 //  - dest (IN/OUT): Address of the destination string. Receives the resulting
 738 //      combined string after the append operation.
 739 //
 740 //  - source (IN): Source string to be appended to the destination string.
 741 //
 742 //  Return Value:
 743 //
 744 //    None.
 745 //
 746 VOID strapp (LPWSTR *dest, LPCWSTR source)
 747 {
 748     SIZE_T length;
 749     LPWSTR temp;
 750
 751     temp = *dest;
 752     length = wcslen(*dest) + wcslen(source);
 753     *dest = new WCHAR [length + 1];
 754     wcsncpy_s(*dest, length + 1, temp, _TRUNCATE);
 755     wcsncat_s(*dest, length + 1, source, _TRUNCATE);
 756     delete [] temp;
 757 }
 758
 759 // strtobool - Converts string values (e.g. "yes", "no", "on", "off") to boolean
 760 //   values.
 761 //
 762 //  - s (IN): String value to convert.
 763 //
 764 //  Return Value:
 765 //
 766 //    Returns TRUE if the string is recognized as a "true" string. Otherwise
 767 //    returns FALSE.
 768 //
 769 BOOL strtobool (LPCWSTR s) {
 770     WCHAR *end;
 771
 772     if ((_wcsicmp(s, L"true") == 0) ||
 773         (_wcsicmp(s, L"yes") == 0) ||
 774         (_wcsicmp(s, L"on") == 0) ||
 775         (wcstol(s, &end, 10) == 1)) {
 776         return TRUE;
 777     }
 778     else {
 779         return FALSE;
 780     }
 781 }