333 lines
		
	
	
		
			9.5 KiB
		
	
	
	
		
			C++
		
	
	
	
		
		
			
		
	
	
			333 lines
		
	
	
		
			9.5 KiB
		
	
	
	
		
			C++
		
	
	
	
|  | //////////////////////////////////////////////////////////////////////////////
 | ||
|  | //
 | ||
|  | //  Detours Test Program (setdll.cpp of setdll.exe)
 | ||
|  | //
 | ||
|  | //  Microsoft Research Detours Package
 | ||
|  | //
 | ||
|  | //  Copyright (c) Microsoft Corporation.  All rights reserved.
 | ||
|  | //
 | ||
|  | #include <stdio.h>
 | ||
|  | #include <stdlib.h>
 | ||
|  | #include <windows.h>
 | ||
|  | #include <shellapi.h>
 | ||
|  | #include <detours.h>
 | ||
|  | #pragma warning(push)
 | ||
|  | #if _MSC_VER > 1400
 | ||
|  | #pragma warning(disable:6102 6103) // /analyze warnings
 | ||
|  | #endif
 | ||
|  | #include <strsafe.h>
 | ||
|  | #pragma warning(pop)
 | ||
|  | 
 | ||
|  | ////////////////////////////////////////////////////////////// Error Messages.
 | ||
|  | //
 | ||
|  | VOID AssertMessage(PCSTR szMsg, PCSTR szFile, DWORD nLine) | ||
|  | { | ||
|  |     printf("ASSERT(%s) failed in %s, line %ld.", szMsg, szFile, nLine); | ||
|  | } | ||
|  | 
 | ||
|  | #define ASSERT(x)   \
 | ||
|  | do { if (!(x)) { AssertMessage(#x, __FILE__, __LINE__); DebugBreak(); }} while (0) | ||
|  |     ; | ||
|  | 
 | ||
|  | 
 | ||
|  | //////////////////////////////////////////////////////////////////////////////
 | ||
|  | //
 | ||
|  | static BOOLEAN  s_fRemove = FALSE; | ||
|  | static CHAR     s_szDllPath[MAX_PATH] = ""; | ||
|  | 
 | ||
|  | //////////////////////////////////////////////////////////////////////////////
 | ||
|  | //
 | ||
|  | //  This code verifies that the named DLL has been configured correctly
 | ||
|  | //  to be imported into the target process.  DLLs must export a function with
 | ||
|  | //  ordinal #1 so that the import table touch-up magic works.
 | ||
|  | //
 | ||
|  | static BOOL CALLBACK ExportCallback(_In_opt_ PVOID pContext, | ||
|  |                                     _In_ ULONG nOrdinal, | ||
|  |                                     _In_opt_ LPCSTR pszName, | ||
|  |                                     _In_opt_ PVOID pCode) | ||
|  | { | ||
|  |     (void)pContext; | ||
|  |     (void)pCode; | ||
|  |     (void)pszName; | ||
|  | 
 | ||
|  |     if (nOrdinal == 1) { | ||
|  |         *((BOOL *)pContext) = TRUE; | ||
|  |     } | ||
|  |     return TRUE; | ||
|  | } | ||
|  | 
 | ||
|  | BOOL DoesDllExportOrdinal1(PCHAR pszDllPath) | ||
|  | { | ||
|  |     HMODULE hDll = LoadLibraryExA(pszDllPath, NULL, DONT_RESOLVE_DLL_REFERENCES); | ||
|  |     if (hDll == NULL) { | ||
|  |         printf("setdll.exe: LoadLibraryEx(%s) failed with error %ld.\n", | ||
|  |                pszDllPath, | ||
|  |                GetLastError()); | ||
|  |         return FALSE; | ||
|  |     } | ||
|  | 
 | ||
|  |     BOOL validFlag = FALSE; | ||
|  |     DetourEnumerateExports(hDll, &validFlag, ExportCallback); | ||
|  |     FreeLibrary(hDll); | ||
|  |     return validFlag; | ||
|  | } | ||
|  | 
 | ||
|  | //////////////////////////////////////////////////////////////////////////////
 | ||
|  | //
 | ||
|  | static BOOL CALLBACK ListBywayCallback(_In_opt_ PVOID pContext, | ||
|  |                                        _In_opt_ LPCSTR pszFile, | ||
|  |                                        _Outptr_result_maybenull_ LPCSTR *ppszOutFile) | ||
|  | { | ||
|  |     (void)pContext; | ||
|  | 
 | ||
|  |     *ppszOutFile = pszFile; | ||
|  |     if (pszFile) { | ||
|  |         printf("    %s\n", pszFile); | ||
|  |     } | ||
|  |     return TRUE; | ||
|  | } | ||
|  | 
 | ||
|  | static BOOL CALLBACK ListFileCallback(_In_opt_ PVOID pContext, | ||
|  |                                       _In_ LPCSTR pszOrigFile, | ||
|  |                                       _In_ LPCSTR pszFile, | ||
|  |                                       _Outptr_result_maybenull_ LPCSTR *ppszOutFile) | ||
|  | { | ||
|  |     (void)pContext; | ||
|  | 
 | ||
|  |     *ppszOutFile = pszFile; | ||
|  |     printf("    %s -> %s\n", pszOrigFile, pszFile); | ||
|  |     return TRUE; | ||
|  | } | ||
|  | 
 | ||
|  | static BOOL CALLBACK AddBywayCallback(_In_opt_ PVOID pContext, | ||
|  |                                       _In_opt_ LPCSTR pszFile, | ||
|  |                                       _Outptr_result_maybenull_ LPCSTR *ppszOutFile) | ||
|  | { | ||
|  |     PBOOL pbAddedDll = (PBOOL)pContext; | ||
|  |     if (!pszFile && !*pbAddedDll) {                     // Add new byway.
 | ||
|  |         *pbAddedDll = TRUE; | ||
|  |         *ppszOutFile = s_szDllPath; | ||
|  |     } | ||
|  |     return TRUE; | ||
|  | } | ||
|  | 
 | ||
|  | BOOL SetFile(PCHAR pszPath) | ||
|  | { | ||
|  |     BOOL bGood = TRUE; | ||
|  |     HANDLE hOld = INVALID_HANDLE_VALUE; | ||
|  |     HANDLE hNew = INVALID_HANDLE_VALUE; | ||
|  |     PDETOUR_BINARY pBinary = NULL; | ||
|  | 
 | ||
|  |     CHAR szOrg[MAX_PATH]; | ||
|  |     CHAR szNew[MAX_PATH]; | ||
|  |     CHAR szOld[MAX_PATH]; | ||
|  | 
 | ||
|  |     szOld[0] = '\0'; | ||
|  |     szNew[0] = '\0'; | ||
|  | 
 | ||
|  |     StringCchCopyA(szOrg, sizeof(szOrg), pszPath); | ||
|  |     StringCchCopyA(szNew, sizeof(szNew), szOrg); | ||
|  |     StringCchCatA(szNew, sizeof(szNew), "#"); | ||
|  |     StringCchCopyA(szOld, sizeof(szOld), szOrg); | ||
|  |     StringCchCatA(szOld, sizeof(szOld), "~"); | ||
|  |     printf("  %s:\n", pszPath); | ||
|  | 
 | ||
|  |     hOld = CreateFileA(szOrg, | ||
|  |                        GENERIC_READ, | ||
|  |                        FILE_SHARE_READ, | ||
|  |                        NULL, | ||
|  |                        OPEN_EXISTING, | ||
|  |                        FILE_ATTRIBUTE_NORMAL, | ||
|  |                        NULL); | ||
|  | 
 | ||
|  |     if (hOld == INVALID_HANDLE_VALUE) { | ||
|  |         printf("Couldn't open input file: %s, error: %ld\n", | ||
|  |                szOrg, GetLastError()); | ||
|  |         bGood = FALSE; | ||
|  |         goto end; | ||
|  |     } | ||
|  | 
 | ||
|  |     hNew = CreateFileA(szNew, | ||
|  |                        GENERIC_WRITE | GENERIC_READ, 0, NULL, CREATE_ALWAYS, | ||
|  |                        FILE_ATTRIBUTE_NORMAL | FILE_FLAG_SEQUENTIAL_SCAN, NULL); | ||
|  |     if (hNew == INVALID_HANDLE_VALUE) { | ||
|  |         printf("Couldn't open output file: %s, error: %ld\n", | ||
|  |                szNew, GetLastError()); | ||
|  |         bGood = FALSE; | ||
|  |         goto end; | ||
|  |     } | ||
|  | 
 | ||
|  |     if ((pBinary = DetourBinaryOpen(hOld)) == NULL) { | ||
|  |         printf("DetourBinaryOpen failed: %ld\n", GetLastError()); | ||
|  |         goto end; | ||
|  |     } | ||
|  | 
 | ||
|  |     if (hOld != INVALID_HANDLE_VALUE) { | ||
|  |         CloseHandle(hOld); | ||
|  |         hOld = INVALID_HANDLE_VALUE; | ||
|  |     } | ||
|  | 
 | ||
|  |     { | ||
|  |         BOOL bAddedDll = FALSE; | ||
|  | 
 | ||
|  |         DetourBinaryResetImports(pBinary); | ||
|  | 
 | ||
|  |         if (!s_fRemove) { | ||
|  |             if (!DetourBinaryEditImports(pBinary, | ||
|  |                                          &bAddedDll, | ||
|  |                                          AddBywayCallback, NULL, NULL, NULL)) { | ||
|  |                 printf("DetourBinaryEditImports failed: %ld\n", GetLastError()); | ||
|  |             } | ||
|  |         } | ||
|  | 
 | ||
|  |         if (!DetourBinaryEditImports(pBinary, NULL, | ||
|  |                                      ListBywayCallback, ListFileCallback, | ||
|  |                                      NULL, NULL)) { | ||
|  | 
 | ||
|  |             printf("DetourBinaryEditImports failed: %ld\n", GetLastError()); | ||
|  |         } | ||
|  | 
 | ||
|  |         if (!DetourBinaryWrite(pBinary, hNew)) { | ||
|  |             printf("DetourBinaryWrite failed: %ld\n", GetLastError()); | ||
|  |             bGood = FALSE; | ||
|  |         } | ||
|  | 
 | ||
|  |         DetourBinaryClose(pBinary); | ||
|  |         pBinary = NULL; | ||
|  | 
 | ||
|  |         if (hNew != INVALID_HANDLE_VALUE) { | ||
|  |             CloseHandle(hNew); | ||
|  |             hNew = INVALID_HANDLE_VALUE; | ||
|  |         } | ||
|  | 
 | ||
|  |         if (bGood) { | ||
|  |             if (!DeleteFileA(szOld)) { | ||
|  |                 DWORD dwError = GetLastError(); | ||
|  |                 if (dwError != ERROR_FILE_NOT_FOUND) { | ||
|  |                     printf("Warning: Couldn't delete %s: %ld\n", szOld, dwError); | ||
|  |                     bGood = FALSE; | ||
|  |                 } | ||
|  |             } | ||
|  |             if (!MoveFileA(szOrg, szOld)) { | ||
|  |                 printf("Error: Couldn't back up %s to %s: %ld\n", | ||
|  |                        szOrg, szOld, GetLastError()); | ||
|  |                 bGood = FALSE; | ||
|  |             } | ||
|  |             if (!MoveFileA(szNew, szOrg)) { | ||
|  |                 printf("Error: Couldn't install %s as %s: %ld\n", | ||
|  |                        szNew, szOrg, GetLastError()); | ||
|  |                 bGood = FALSE; | ||
|  |             } | ||
|  |         } | ||
|  | 
 | ||
|  |         DeleteFileA(szNew); | ||
|  |     } | ||
|  | 
 | ||
|  | 
 | ||
|  |   end: | ||
|  |     if (pBinary) { | ||
|  |         DetourBinaryClose(pBinary); | ||
|  |         pBinary = NULL; | ||
|  |     } | ||
|  |     if (hNew != INVALID_HANDLE_VALUE) { | ||
|  |         CloseHandle(hNew); | ||
|  |         hNew = INVALID_HANDLE_VALUE; | ||
|  |     } | ||
|  |     if (hOld != INVALID_HANDLE_VALUE) { | ||
|  |         CloseHandle(hOld); | ||
|  |         hOld = INVALID_HANDLE_VALUE; | ||
|  |     } | ||
|  |     return bGood; | ||
|  | } | ||
|  | 
 | ||
|  | //////////////////////////////////////////////////////////////////////////////
 | ||
|  | //
 | ||
|  | void PrintUsage(void) | ||
|  | { | ||
|  |     printf("Usage:\n" | ||
|  |            "    setdll [options] binary_files\n" | ||
|  |            "Options:\n" | ||
|  |            "    /d:file.dll  : Add file.dll binary files\n" | ||
|  |            "    /r           : Remove extra DLLs from binary files\n" | ||
|  |            "    /?           : This help screen.\n"); | ||
|  | } | ||
|  | 
 | ||
|  | //////////////////////////////////////////////////////////////////////// main.
 | ||
|  | //
 | ||
|  | int CDECL main(int argc, char **argv) | ||
|  | { | ||
|  |     BOOL fNeedHelp = FALSE; | ||
|  |     PCHAR pszFilePart = NULL; | ||
|  | 
 | ||
|  |     int arg = 1; | ||
|  |     for (; arg < argc; arg++) { | ||
|  |         if (argv[arg][0] == '-' || argv[arg][0] == '/') { | ||
|  |             CHAR *argn = argv[arg] + 1; | ||
|  |             CHAR *argp = argn; | ||
|  |             while (*argp && *argp != ':' && *argp != '=') | ||
|  |                 argp++; | ||
|  |             if (*argp == ':' || *argp == '=') | ||
|  |                 *argp++ = '\0'; | ||
|  | 
 | ||
|  |             switch (argn[0]) { | ||
|  | 
 | ||
|  |               case 'd':                                 // Set DLL
 | ||
|  |               case 'D': | ||
|  |                 if ((strchr(argp, ':') != NULL || strchr(argp, '\\') != NULL) && | ||
|  |                     GetFullPathNameA(argp, sizeof(s_szDllPath), s_szDllPath, &pszFilePart)) { | ||
|  |                 } | ||
|  |                 else { | ||
|  |                     StringCchPrintfA(s_szDllPath, sizeof(s_szDllPath), "%s", argp); | ||
|  |                 } | ||
|  |                 break; | ||
|  | 
 | ||
|  |               case 'r':                                 // Remove extra set DLLs.
 | ||
|  |               case 'R': | ||
|  |                 s_fRemove = TRUE; | ||
|  |                 break; | ||
|  | 
 | ||
|  |               case '?':                                 // Help
 | ||
|  |                 fNeedHelp = TRUE; | ||
|  |                 break; | ||
|  | 
 | ||
|  |               default: | ||
|  |                 fNeedHelp = TRUE; | ||
|  |                 printf("Bad argument: %s:%s\n", argn, argp); | ||
|  |                 break; | ||
|  |             } | ||
|  |         } | ||
|  |     } | ||
|  |     if (argc == 1) { | ||
|  |         fNeedHelp = TRUE; | ||
|  |     } | ||
|  |     if (!s_fRemove && s_szDllPath[0] == 0) { | ||
|  |         fNeedHelp = TRUE; | ||
|  |     } | ||
|  |     if (fNeedHelp) { | ||
|  |         PrintUsage(); | ||
|  |         return 1; | ||
|  |     } | ||
|  | 
 | ||
|  | 
 | ||
|  |     if (s_fRemove) { | ||
|  |         printf("Removing extra DLLs from binary files.\n"); | ||
|  |     } | ||
|  |     else { | ||
|  |         if (!DoesDllExportOrdinal1(s_szDllPath)) { | ||
|  |             printf("Error: %hs does not export function with ordinal #1.\n", | ||
|  |                    s_szDllPath); | ||
|  |             return 2; | ||
|  |         } | ||
|  |         printf("Adding %hs to binary files.\n", s_szDllPath); | ||
|  |     } | ||
|  | 
 | ||
|  |     for (arg = 1; arg < argc; arg++) { | ||
|  |         if (argv[arg][0] != '-' && argv[arg][0] != '/') { | ||
|  |             SetFile(argv[arg]); | ||
|  |         } | ||
|  |     } | ||
|  |     return 0; | ||
|  | } | ||
|  | 
 | ||
|  | // End of File
 |