development/Win32/vld/src/utility.cpp

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