image.cpp

微软提供的截取Win32 API函数的开发包和例子detours-src-1.2.rar

		SetLastError(ERROR_EXE_MARKED_INVALID);
		return FALSE;
	}
	CopyMemory(&m_SectionHeaders,
			   m_pMap + m_nSectionsOffset, 
			   sizeof(m_SectionHeaders[0]) * m_NtHeader.FileHeader.NumberOfSections);

	////////////////////////////////////////////////////////// Parse Sections.
	//
	DWORD rvaOriginalImageDirectory = 0;
	for (n = 0; n < m_NtHeader.FileHeader.NumberOfSections; n++) {
		if (strcmp((PCHAR)m_SectionHeaders[n].Name, ".detour") == 0) {
			DETOUR_SECTION_HEADER dh;
			CopyMemory(&dh,
					   m_pMap + m_SectionHeaders[n].PointerToRawData,
					   sizeof(dh));
			
			rvaOriginalImageDirectory = dh.nOriginalImportVirtualAddress;	
		}
	}

	//////////////////////////////////////////////////////// Get Import Table.
	//
	DWORD rvaImageDirectory = m_NtHeader.OptionalHeader
		.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress;
	PIMAGE_IMPORT_DESCRIPTOR iidp
		= (PIMAGE_IMPORT_DESCRIPTOR)RvaToVa(rvaImageDirectory);
	PIMAGE_IMPORT_DESCRIPTOR oidp
		= (PIMAGE_IMPORT_DESCRIPTOR)RvaToVa(rvaOriginalImageDirectory);
	if (oidp == NULL) {
		oidp = iidp;
	}
	if (iidp == NULL || oidp == NULL) {
		SetLastError(ERROR_EXE_MARKED_INVALID);
		return FALSE;
	}

	DWORD rvaIatBeg = m_NtHeader.OptionalHeader
		.DataDirectory[IMAGE_DIRECTORY_ENTRY_IAT].VirtualAddress;
	DWORD rvaIatEnd = rvaIatBeg + m_NtHeader.OptionalHeader
		.DataDirectory[IMAGE_DIRECTORY_ENTRY_IAT].Size;

	if (iidp) {
		for (DWORD nFiles = 0; iidp[nFiles].Characteristics != 0; nFiles++) {
		}
		
		CImageImportFile **ppLastFile = &m_pImportFiles;
		m_pImportFiles = NULL;

		for (DWORD n = 0; n < nFiles; n++, iidp++) {
			DWORD rvaName = iidp->Name;
			PCHAR pszName = (PCHAR)RvaToVa(rvaName);
			if (pszName == NULL) {
				SetLastError(ERROR_EXE_MARKED_INVALID);
				goto fail;
			}

			CImageImportFile *pImportFile = new CImageImportFile;
			if (pImportFile == NULL) {
				SetLastError(ERROR_OUTOFMEMORY);
				goto fail;
			}

			*ppLastFile = pImportFile;
			ppLastFile = &pImportFile->m_pNextFile;
			m_nImportFiles++;
			
			pImportFile->m_pszName = DuplicateString(pszName);
			if (pImportFile->m_pszName == NULL) {
				goto fail;
			}
			
			pImportFile->m_rvaOriginalFirstThunk
				= (PIMAGE_THUNK_DATA)iidp->OriginalFirstThunk;
			pImportFile->m_rvaFirstThunk = (PIMAGE_THUNK_DATA)iidp->FirstThunk;
			pImportFile->m_nForwarderChain = iidp->ForwarderChain;
			pImportFile->m_pImportNames = NULL;
			pImportFile->m_nImportNames = 0;
			pImportFile->m_bByway = FALSE;

			if ((ULONG)iidp->FirstThunk < rvaIatBeg ||
				(ULONG)iidp->FirstThunk >= rvaIatEnd) {
				
				pImportFile->m_pszOrig = NULL;
				pImportFile->m_bByway = TRUE;
				continue;
			}

			rvaName = oidp->Name;
			pszName = (PCHAR)RvaToVa(rvaName);
			if (pszName == NULL) {
				SetLastError(ERROR_EXE_MARKED_INVALID);
				goto fail;
			}
			pImportFile->m_pszOrig = DuplicateString(pszName);
			if (pImportFile->m_pszOrig == NULL) {
				goto fail;
			}
			
			DWORD rvaThunk = (DWORD)iidp->OriginalFirstThunk;
			PIMAGE_THUNK_DATA pThunk = (PIMAGE_THUNK_DATA)RvaToVa(rvaThunk);
			rvaThunk = (DWORD)oidp->OriginalFirstThunk;
			PIMAGE_THUNK_DATA pOrigThunk = (PIMAGE_THUNK_DATA)RvaToVa(rvaThunk);

			DWORD nNames = 0;
			if (pThunk) {
				for (; pThunk[nNames].u1.Ordinal; nNames++) {
				}
			}

			if (pThunk && nNames) {
				pImportFile->m_nImportNames = nNames;
				pImportFile->m_pImportNames = new CImageImportName [nNames];
				if (pImportFile->m_pImportNames == NULL) {
					SetLastError(ERROR_OUTOFMEMORY);
					goto fail;
				}

				CImageImportName *pImportName = &pImportFile->m_pImportNames[0];
				
				for (DWORD f = 0; f < nNames; f++, pImportName++) {
					pImportName->m_nOrdinal = 0;
					pImportName->m_nHint = 0;
					pImportName->m_pszName = NULL;
					pImportName->m_pszOrig = NULL;
					
					DWORD rvaName = pThunk[f].u1.Ordinal;
					if (rvaName & IMAGE_ORDINAL_FLAG) {
						pImportName->m_nOrdinal = IMAGE_ORDINAL(rvaName);
					}
					else {
						PIMAGE_IMPORT_BY_NAME pName
							= (PIMAGE_IMPORT_BY_NAME)RvaToVa(rvaName);
						if (pName) {
							pImportName->m_nHint = pName->Hint;
							pImportName->m_pszName = DuplicateString((PCHAR)pName->Name);
							if (pImportName->m_pszName == NULL) {
								goto fail;
							}
						}
						
						rvaName = pOrigThunk[f].u1.Ordinal;
						if (rvaName & IMAGE_ORDINAL_FLAG) {
							pImportName->m_nOrdinal = IMAGE_ORDINAL(rvaName);
						}
						else {
							pName = (PIMAGE_IMPORT_BY_NAME)RvaToVa(rvaName);
							if (pName) {
								pImportName->m_pszOrig
									= DuplicateString((PCHAR)pName->Name);
								if (pImportName->m_pszOrig == NULL) {
									goto fail;
								}
							}
						}
					}
				}
			}
			oidp++;
		}
	}
	
	////////////////////////////////////////////////////////// Parse Sections.
	//
	m_nExtraOffset = 0;
	for (n = 0; n < m_NtHeader.FileHeader.NumberOfSections; n++) {
		m_nExtraOffset = Max(m_SectionHeaders[n].PointerToRawData +
							 m_SectionHeaders[n].SizeOfRawData,
							 m_nExtraOffset);
		
		if (strcmp((PCHAR)m_SectionHeaders[n].Name, ".detour") == 0) {
			DETOUR_SECTION_HEADER dh;
			CopyMemory(&dh,
					   m_pMap + m_SectionHeaders[n].PointerToRawData,
					   sizeof(dh));

			if (dh.nDataOffset == 0) {
				dh.nDataOffset = dh.cbHeaderSize;
			}

			cbData = dh.cbDataSize - dh.nDataOffset;
			pbData = (m_pMap +
					  m_SectionHeaders[n].PointerToRawData +
					  dh.nDataOffset);
			
			m_nExtraOffset = Max(m_SectionHeaders[n].PointerToRawData +
								 m_SectionHeaders[n].SizeOfRawData,
								 m_nExtraOffset);
			
			m_NtHeader.FileHeader.NumberOfSections--;

			m_NtHeader.OptionalHeader
				.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress
				= dh.nOriginalImportVirtualAddress;
			m_NtHeader.OptionalHeader
				.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].Size
				= dh.nOriginalImportSize;
			
			m_NtHeader.OptionalHeader
				.DataDirectory[IMAGE_DIRECTORY_ENTRY_BOUND_IMPORT].VirtualAddress
				= dh.nOriginalBoundImportVirtualAddress;
			m_NtHeader.OptionalHeader
				.DataDirectory[IMAGE_DIRECTORY_ENTRY_BOUND_IMPORT].Size
				= dh.nOriginalBoundImportSize;
	
			m_NtHeader.OptionalHeader
				.DataDirectory[IMAGE_DIRECTORY_ENTRY_IAT].VirtualAddress
				= dh.nOriginalIatVirtualAddress;
			m_NtHeader.OptionalHeader
				.DataDirectory[IMAGE_DIRECTORY_ENTRY_IAT].Size
				= dh.nOriginalIatSize;

			m_NtHeader.OptionalHeader.CheckSum = 0;
			m_NtHeader.OptionalHeader.SizeOfImage
				= dh.nOriginalSizeOfImage;
		}
	}

	m_pImageData = new CImageData(pbData, cbData);
	if (m_pImageData == NULL) {
		SetLastError(ERROR_OUTOFMEMORY);
	}
	return TRUE;

fail:
	return FALSE;
}

static inline BOOL strneq(PCHAR pszOne, PCHAR pszTwo)
{
	if (pszOne == pszTwo) {
		return FALSE;
	}
	if (!pszOne || !pszTwo) {
		return TRUE;
	}
	return (strcmp(pszOne, pszTwo) != 0);
}

BOOL CImage::CheckImportsNeeded(DWORD *pnTables, DWORD *pnThunks, DWORD *pnChars)
{
	DWORD nTables = 0;
	DWORD nThunks = 0;
	DWORD nChars = 0;
	BOOL bNeedDetourSection = FALSE;

	for (CImageImportFile *pImportFile = m_pImportFiles;
		 pImportFile != NULL; pImportFile = pImportFile->m_pNextFile) {
		
		nChars += strlen(pImportFile->m_pszName) + 1;
		nChars += nChars & 1;

		if (pImportFile->m_bByway) {
			bNeedDetourSection = TRUE;
			nThunks++;
		}
		else {
			if (!bNeedDetourSection &&
				strneq(pImportFile->m_pszName, pImportFile->m_pszOrig)) {
				
				bNeedDetourSection = TRUE;
			}
			for (DWORD n = 0; n < pImportFile->m_nImportNames; n++) {
				CImageImportName *pImportName = &pImportFile->m_pImportNames[n];

				if (!bNeedDetourSection &&
					strneq(pImportName->m_pszName, pImportName->m_pszOrig)) {
					
					bNeedDetourSection = TRUE;
				}
				
				if (pImportName->m_pszName) {
					nChars += sizeof(WORD);				// Hint
					nChars += strlen(pImportName->m_pszName) + 1;
					nChars += nChars & 1;
				}
				nThunks++;
			}
		}
		nThunks++;
		nTables++;
	}
	nTables++;

	*pnTables = nTables;
	*pnThunks = nThunks;
	*pnChars = nChars;

	return bNeedDetourSection;
}

//////////////////////////////////////////////////////////////////////////////
//
CImageImportFile * CImage::NewByway(PCHAR pszName)
{
	CImageImportFile *pImportFile = new CImageImportFile;
	if (pImportFile == NULL) {
		SetLastError(ERROR_OUTOFMEMORY);
		goto fail;
	}

	if (pImportFile) {
		pImportFile->m_pNextFile = NULL;
		pImportFile->m_bByway = TRUE;
		
		pImportFile->m_pszName = DuplicateString(pszName);
		if (pImportFile->m_pszName == NULL) {
			goto fail;
		}
		
		pImportFile->m_rvaOriginalFirstThunk = 0;
		pImportFile->m_rvaFirstThunk = 0;
		pImportFile->m_nForwarderChain = -1;
		pImportFile->m_pImportNames = NULL;
		pImportFile->m_nImportNames = 0;

		m_nImportFiles++;
	}
	return pImportFile;

fail:
	if (pImportFile) {
		delete pImportFile;
		pImportFile = NULL;
	}
	return NULL;
}

BOOL CImage::EditImports(PVOID pContext,
						 PF_DETOUR_BINARY_BYWAY_CALLBACK pfBywayCallback,
						 PF_DETOUR_BINARY_FILE_CALLBACK pfFileCallback,
						 PF_DETOUR_BINARY_SYMBOL_CALLBACK pfSymbolCallback,
						 PF_DETOUR_BINARY_FINAL_CALLBACK pfFinalCallback)
{
	CImageImportFile *pImportFile = NULL;
	CImageImportFile **ppLastFile = &m_pImportFiles;

	SetLastError(ERROR_CALL_NOT_IMPLEMENTED);
	
	while ((pImportFile = *ppLastFile) != NULL) {
		
		if (pfBywayCallback) {
			PCHAR pszFile = NULL;
			if (!(*pfBywayCallback)(pContext, pszFile, &pszFile)) {
				goto fail;
			}
			
			if (pszFile) {
				// Insert a new Byway.
				CImageImportFile *pByway = NewByway(pszFile);
				if (pByway == NULL)
					return FALSE;
				
				pByway->m_pNextFile = pImportFile;
				*ppLastFile = pByway;
				ppLastFile = &pByway->m_pNextFile;
				continue;								// Retry after Byway.
			}
		}

		if (pImportFile->m_bByway) {
			if (pfBywayCallback) {
				PCHAR pszFile = pImportFile->m_pszName;

				if (!(*pfBywayCallback)(pContext, pszFile, &pszFile)) {
					goto fail;
				}
				
				if (pszFile) {					// Replace? Byway
					if (ReplaceString(&pImportFile->m_pszName, pszFile) == NULL) {
						goto fail;
					}
				}
				else {								// Delete Byway
					*ppLastFile = pImportFile->m_pNextFile;
					pImportFile->m_pNextFile = NULL;
					delete pImportFile;
					pImportFile = *ppLastFile;
					m_nImportFiles--;
					continue;						// Retry after delete.
				}
			}
		}
		else {
			if (pfFileCallback) {
				PCHAR pszFile = pImportFile->m_pszName;
				
				if (!(*pfFileCallback)(pContext, pImportFile->m_pszOrig,
									   pszFile, &pszFile)) {
					goto fail;
				}
				
				if (pszFile != NULL) {
					if (ReplaceString(&pImportFile->m_pszName, pszFile) == NULL) {
						goto fail;
					}
				}
			}

			if (pfSymbolCallback) {
				for (DWORD n = 0; n < pImportFile->m_nImportNames; n++) {
					CImageImportName *pImportName = &pImportFile->m_pImportNames[n];
					
					PCHAR pszName = pImportName->m_pszName;
					if (!(*pfSymbolCallback)(pContext,
											pImportName->m_nOrdinal,
											pImportName->m_pszOrig,
											pszName,
											&pszName)) {
						goto fail;
					}
					
					if (pszName != NULL) {
						if (ReplaceString(&pImportName->m_pszName, pszName) == NULL) {
							goto fail;
						}
					}
				}
			}
		}
		
		ppLastFile = &pImportFile->m_pNextFile;
		pImportFile = pImportFile->m_pNextFile;
	}

	for (;;) {
		if (pfBywayCallback) {
			PCHAR pszFile = NULL;
			if (!(*pfBywayCallback)(pContext, NULL, &pszFile)) {
				goto fail;
			}
			if (pszFile) {
				// Insert a new Byway.
				CImageImportFile *pByway = NewByway(pszFile);
				if (pByway == NULL)
					return FALSE;
				
				pByway->m_pNextFile = pImportFile;
				*ppLastFile = pByway;
				ppLastFile = &pByway->m_pNextFile;
				continue;								// Retry after Byway.
			}
		}
		break;
	}

	if (pfFinalCallback) {
		if (!(*pfFinalCallback)(pContext)) {
			goto fail;
		}
	}
	
	SetLastError(NO_ERROR);
	return TRUE;

fail:
	return FALSE;
}

BOOL CImage::Write(HANDLE hFile)
{
	if (hFile == INVALID_HANDLE_VALUE) {
		SetLastError(ERROR_INVALID_HANDLE);
		return FALSE;
	}

	//////////////////////////////////////////////////////////////////////////
	//
	m_nNextFileAddr = 0;
	m_nNextVirtAddr = 0;
	
	//////////////////////////////////////////////////////////// Copy Headers.