home *** CD-ROM | disk | FTP | other *** search
/ Tricks of the Windows Gam…ming Gurus (2nd Edition) / Disc2.iso / vc98 / include / delayhlp.cpp < prev    next >
C/C++ Source or Header  |  1998-05-05  |  8KB  |  284 lines
  1. #define STRICT
  2. #include <windows.h>
  3. #pragma hdrstop
  4. #include "delayImp.h"
  5.  
  6. extern "C"
  7. PUnloadInfo __puiHead = 0;
  8.  
  9. struct ULI : public UnloadInfo {
  10.     ULI(PCImgDelayDescr pidd_) {
  11.         pidd = pidd_;
  12.         Link();
  13.         }
  14.  
  15.     ~ULI() {
  16.         Unlink();
  17.         }
  18.  
  19.     void *
  20.     operator new(unsigned int cb) {
  21.         return ::LocalAlloc(LPTR, cb);
  22.         }
  23.  
  24.     void
  25.     operator delete(void * pv) {
  26.         ::LocalFree(pv);
  27.         }
  28.  
  29.     void
  30.     Unlink() {
  31.         PUnloadInfo *   ppui = &__puiHead;
  32.  
  33.         while (*ppui && *ppui != this) {
  34.             ppui = &((*ppui)->puiNext);
  35.             }
  36.         if (*ppui == this) {
  37.             *ppui = puiNext;
  38.             }
  39.         }
  40.  
  41.     void
  42.     Link() {
  43.         puiNext = __puiHead;
  44.         __puiHead = this;
  45.         }
  46.     };
  47.  
  48. static inline
  49. PIMAGE_NT_HEADERS WINAPI
  50. PinhFromImageBase(HMODULE);
  51.  
  52. static inline
  53. DWORD WINAPI
  54. TimeStampOfImage(PIMAGE_NT_HEADERS);
  55.  
  56. static inline
  57. void WINAPI
  58. OverlayIAT(PImgThunkData pitdDst, PCImgThunkData pitdSrc);
  59.  
  60. static inline
  61. bool WINAPI
  62. FLoadedAtPreferredAddress(PIMAGE_NT_HEADERS, HMODULE);
  63.  
  64. extern "C"
  65. FARPROC WINAPI
  66. __delayLoadHelper(
  67.     PCImgDelayDescr pidd,
  68.     FARPROC *       ppfnIATEntry
  69.     ) {
  70.  
  71.     // Set up some data we use for the hook procs but also useful for
  72.     // our own use
  73.     //
  74.     DelayLoadInfo   dli = {
  75.         sizeof DelayLoadInfo,
  76.         pidd,
  77.         ppfnIATEntry,
  78.         pidd->szName,
  79.             { 0 },
  80.         0,
  81.         0,
  82.         0
  83.         };
  84.  
  85.     HMODULE hmod = *(pidd->phmod);
  86.  
  87.     // Calculate the index for the name in the import name table.
  88.     // N.B. it is ordered the same as the IAT entries so the calculation
  89.     // comes from the IAT side.
  90.     //
  91.     unsigned        iINT;
  92.     iINT = IndexFromPImgThunkData(PCImgThunkData(ppfnIATEntry), pidd->pIAT);
  93.  
  94.     PCImgThunkData  pitd = &((pidd->pINT)[iINT]);
  95.  
  96.     if (dli.dlp.fImportByName = ((pitd->u1.Ordinal & IMAGE_ORDINAL_FLAG) == 0)) {
  97.         dli.dlp.szProcName = LPCSTR(pitd->u1.AddressOfData->Name);
  98.         }
  99.     else {
  100.         dli.dlp.dwOrdinal = IMAGE_ORDINAL(pitd->u1.Ordinal);
  101.         }
  102.  
  103.     // Call the initial hook.  If it exists and returns a function pointer,
  104.     // abort the rest of the processing and just return it for the call.
  105.     //
  106.     FARPROC pfnRet = NULL;
  107.  
  108.     if (__pfnDliNotifyHook) {
  109.         if (pfnRet = ((*__pfnDliNotifyHook)(dliStartProcessing, &dli))) {
  110.             goto HookBypass;
  111.             }
  112.         }
  113.  
  114.     if (hmod == 0) {
  115.         if (__pfnDliNotifyHook) {
  116.             hmod = HMODULE(((*__pfnDliNotifyHook)(dliNotePreLoadLibrary, &dli)));
  117.             }
  118.         if (hmod == 0) {
  119.             hmod = ::LoadLibrary(dli.szDll);
  120.             }
  121.         if (hmod == 0) {
  122.             dli.dwLastError = ::GetLastError();
  123.             if (__pfnDliFailureHook) {
  124.                 // when the hook is called on LoadLibrary failure, it will
  125.                 // return 0 for failure and an hmod for the lib if it fixed
  126.                 // the problem.
  127.                 //
  128.                 hmod = HMODULE((*__pfnDliFailureHook)(dliFailLoadLib, &dli));
  129.                 }
  130.  
  131.             if (hmod == 0) {
  132.                 PDelayLoadInfo  pdli = &dli;
  133.  
  134.                 RaiseException(
  135.                     VcppException(ERROR_SEVERITY_ERROR, ERROR_MOD_NOT_FOUND),
  136.                     0,
  137.                     1,
  138.                     PDWORD(&pdli)
  139.                     );
  140.                 
  141.                 // If we get to here, we blindly assume that the handler of the exception
  142.                 // has magically fixed everything up and left the function pointer in 
  143.                 // dli.pfnCur.
  144.                 //
  145.                 return dli.pfnCur;
  146.                 }
  147.             }
  148.  
  149.         // Store the library handle.  If it is already there, we infer
  150.         // that another thread got there first, and we need to do a
  151.         // FreeLibrary() to reduce the refcount
  152.         //
  153.         HMODULE hmodT = HMODULE(::InterlockedExchange(LPLONG(pidd->phmod), LONG(hmod)));
  154.         if (hmodT != hmod) {
  155.             // add lib to unload list if we have unload data
  156.             if (pidd->pUnloadIAT) {
  157.                 ULI *   puli = new ULI(pidd);
  158.                 (void *)puli;
  159.                 }
  160.             }
  161.         else {
  162.             ::FreeLibrary(hmod);
  163.             }
  164.         
  165.         }
  166.  
  167.     // Go for the procedure now.
  168.     dli.hmodCur = hmod;
  169.     if (__pfnDliNotifyHook) {
  170.         pfnRet = (*__pfnDliNotifyHook)(dliNotePreGetProcAddress, &dli);
  171.         }
  172.     if (pfnRet == 0) {
  173.         if (pidd->pBoundIAT && pidd->dwTimeStamp) {
  174.             // bound imports exist...check the timestamp from the target image
  175.             PIMAGE_NT_HEADERS   pinh(PinhFromImageBase(hmod));
  176.  
  177.             if (pinh->Signature == IMAGE_NT_SIGNATURE &&
  178.                 TimeStampOfImage(pinh) == pidd->dwTimeStamp &&
  179.                 FLoadedAtPreferredAddress(pinh, hmod)) {
  180.  
  181.                 OverlayIAT(pidd->pIAT, pidd->pBoundIAT);
  182.                 pfnRet = FARPROC(pidd->pIAT[iINT].u1.Function);
  183.                 goto HookBypass;
  184.                 }
  185.             }
  186.  
  187.         pfnRet = ::GetProcAddress(hmod, dli.dlp.szProcName);
  188.         }
  189.  
  190.     if (pfnRet == 0) {
  191.         dli.dwLastError = ::GetLastError();
  192.         if (__pfnDliFailureHook) {
  193.             // when the hook is called on GetProcAddress failure, it will
  194.             // return 0 on failure and a valid proc address on success
  195.             //
  196.             pfnRet = (*__pfnDliFailureHook)(dliFailGetProc, &dli);
  197.             }
  198.         if (pfnRet == 0) {
  199.             PDelayLoadInfo  pdli = &dli;
  200.  
  201.             RaiseException(
  202.                 VcppException(ERROR_SEVERITY_ERROR, ERROR_PROC_NOT_FOUND),
  203.                 0,
  204.                 1,
  205.                 PDWORD(&pdli)
  206.                 );
  207.  
  208.             // If we get to here, we blindly assume that the handler of the exception
  209.             // has magically fixed everything up and left the function pointer in 
  210.             // dli.pfnCur.
  211.             //
  212.             pfnRet = dli.pfnCur;
  213.             }
  214.         }
  215.  
  216.  
  217.     *ppfnIATEntry = pfnRet;
  218.  
  219. HookBypass:
  220.     if (__pfnDliNotifyHook) {
  221.         dli.dwLastError = 0;
  222.         dli.hmodCur = hmod;
  223.         dli.pfnCur = pfnRet;
  224.         (*__pfnDliNotifyHook)(dliNoteEndProcessing, &dli);
  225.         }
  226.     return pfnRet;
  227.     }
  228.  
  229.  
  230. #pragma intrinsic(strlen,memcmp,memcpy)
  231.  
  232. extern "C"
  233. BOOL WINAPI
  234. __FUnloadDelayLoadedDLL(LPCSTR szDll) {
  235.     
  236.     BOOL        fRet = FALSE;
  237.     PUnloadInfo pui = __puiHead;
  238.     
  239.     for (pui = __puiHead; pui; pui = pui->puiNext) {
  240.         if (memcmp(szDll, pui->pidd->szName, strlen(pui->pidd->szName)) == 0) {
  241.             break;
  242.             }
  243.         }
  244.  
  245.     if (pui && pui->pidd->pUnloadIAT) {
  246.         PCImgDelayDescr pidd = pui->pidd;
  247.         HMODULE         hmod = *pidd->phmod;
  248.  
  249.         OverlayIAT(pidd->pIAT, pidd->pUnloadIAT);
  250.         ::FreeLibrary(hmod);
  251.         *pidd->phmod = NULL;
  252.         
  253.         delete reinterpret_cast<ULI*> (pui);
  254.  
  255.         fRet = TRUE;
  256.         }
  257.  
  258.     return fRet;
  259.     }
  260.  
  261. static inline
  262. PIMAGE_NT_HEADERS WINAPI
  263. PinhFromImageBase(HMODULE hmod) {
  264.     return PIMAGE_NT_HEADERS(PCHAR(hmod) + PIMAGE_DOS_HEADER(hmod)->e_lfanew);
  265.     }
  266.  
  267. static inline
  268. void WINAPI
  269. OverlayIAT(PImgThunkData pitdDst, PCImgThunkData pitdSrc) {
  270.     memcpy(pitdDst, pitdSrc, CountOfImports(pitdDst) * sizeof IMAGE_THUNK_DATA);
  271.     }
  272.  
  273. static inline
  274. DWORD WINAPI
  275. TimeStampOfImage(PIMAGE_NT_HEADERS pinh) {
  276.     return pinh->FileHeader.TimeDateStamp;
  277.     }
  278.  
  279. static inline
  280. bool WINAPI
  281. FLoadedAtPreferredAddress(PIMAGE_NT_HEADERS pinh, HMODULE hmod) {
  282.     return DWORD(hmod) == pinh->OptionalHeader.ImageBase;
  283.     }
  284.