in src/UndockedRegFreeWinRT/detours/image.cpp [1606:2035]
BOOL CImage::Write(HANDLE hFile)
{
DWORD cbDone;
if (hFile == INVALID_HANDLE_VALUE) {
SetLastError(ERROR_INVALID_HANDLE);
return FALSE;
}
m_nNextFileAddr = 0;
m_nNextVirtAddr = 0;
DWORD nTables = 0;
DWORD nThunks = 0;
DWORD nChars = 0;
BOOL fNeedDetourSection = CheckImportsNeeded(&nTables, &nThunks, &nChars);
//////////////////////////////////////////////////////////// Copy Headers.
//
if (SetFilePointer(hFile, 0, NULL, FILE_BEGIN) == ~0u) {
return FALSE;
}
if (!CopyFileData(hFile, 0, m_NtHeader.OptionalHeader.SizeOfHeaders)) {
return FALSE;
}
if (fNeedDetourSection || !m_pImageData->IsEmpty()) {
// Replace the file's DOS header with our own.
m_nPeOffset = sizeof(m_DosHeader) + sizeof(s_rbDosCode);
m_nSectionsOffset = m_nPeOffset
+ sizeof(m_NtHeader.Signature)
+ sizeof(m_NtHeader.FileHeader)
+ m_NtHeader.FileHeader.SizeOfOptionalHeader;
m_DosHeader.e_lfanew = m_nPeOffset;
if (SetFilePointer(hFile, 0, NULL, FILE_BEGIN) == ~0u) {
return FALSE;
}
if (!WriteFile(hFile, &m_DosHeader, sizeof(m_DosHeader), &cbDone)) {
return FALSE;
}
if (!WriteFile(hFile, &s_rbDosCode, sizeof(s_rbDosCode), &cbDone)) {
return FALSE;
}
}
else {
// Restore the file's original DOS header.
if (m_nPrePE != 0) {
m_nPeOffset = m_cbPrePE;
m_nSectionsOffset = m_nPeOffset
+ sizeof(m_NtHeader.Signature)
+ sizeof(m_NtHeader.FileHeader)
+ m_NtHeader.FileHeader.SizeOfOptionalHeader;
m_DosHeader.e_lfanew = m_nPeOffset;
if (SetFilePointer(hFile, 0, NULL, FILE_BEGIN) == ~0u) {
return FALSE;
}
if (!CopyFileData(hFile, m_nPrePE, m_cbPrePE)) {
return FALSE;
}
}
}
m_nNextFileAddr = m_NtHeader.OptionalHeader.SizeOfHeaders;
m_nNextVirtAddr = 0;
if (!AlignFileData(hFile)) {
return FALSE;
}
/////////////////////////////////////////////////////////// Copy Sections.
//
DWORD n = 0;
for (; n < m_NtHeader.FileHeader.NumberOfSections; n++) {
if (m_SectionHeaders[n].SizeOfRawData) {
if (SetFilePointer(hFile,
m_SectionHeaders[n].PointerToRawData,
NULL, FILE_BEGIN) == ~0u) {
return FALSE;
}
if (!CopyFileData(hFile,
m_SectionHeaders[n].PointerToRawData,
m_SectionHeaders[n].SizeOfRawData)) {
return FALSE;
}
}
m_nNextFileAddr = Max(m_SectionHeaders[n].PointerToRawData +
m_SectionHeaders[n].SizeOfRawData,
m_nNextFileAddr);
// Old images have VirtualSize == 0 as a matter of course, e.g. NT 3.1.
// In which case, use SizeOfRawData instead.
m_nNextVirtAddr = Max(m_SectionHeaders[n].VirtualAddress +
(m_SectionHeaders[n].Misc.VirtualSize
? m_SectionHeaders[n].Misc.VirtualSize
: SectionAlign(m_SectionHeaders[n].SizeOfRawData)),
m_nNextVirtAddr);
m_nExtraOffset = Max(m_nNextFileAddr, m_nExtraOffset);
if (!AlignFileData(hFile)) {
return FALSE;
}
}
if (fNeedDetourSection || !m_pImageData->IsEmpty()) {
if (m_NtHeader.FileHeader.NumberOfSections >= ARRAYSIZE(m_SectionHeaders)) {
SetLastError(ERROR_EXE_MARKED_INVALID);
return FALSE;
}
////////////////////////////////////////////// Insert .detour Section.
//
DWORD nSection = m_NtHeader.FileHeader.NumberOfSections++;
DETOUR_SECTION_HEADER dh;
ZeroMemory(&dh, sizeof(dh));
ZeroMemory(&m_SectionHeaders[nSection], sizeof(m_SectionHeaders[nSection]));
dh.cbHeaderSize = sizeof(DETOUR_SECTION_HEADER);
dh.nSignature = DETOUR_SECTION_HEADER_SIGNATURE;
dh.nOriginalImportVirtualAddress = m_NtHeader.OptionalHeader
.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress;
dh.nOriginalImportSize = m_NtHeader.OptionalHeader
.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].Size;
dh.nOriginalBoundImportVirtualAddress
= m_NtHeader.OptionalHeader
.DataDirectory[IMAGE_DIRECTORY_ENTRY_BOUND_IMPORT].VirtualAddress;
dh.nOriginalBoundImportSize = m_NtHeader.OptionalHeader
.DataDirectory[IMAGE_DIRECTORY_ENTRY_BOUND_IMPORT].Size;
dh.nOriginalIatVirtualAddress = m_NtHeader.OptionalHeader
.DataDirectory[IMAGE_DIRECTORY_ENTRY_IAT].VirtualAddress;
dh.nOriginalIatSize = m_NtHeader.OptionalHeader
.DataDirectory[IMAGE_DIRECTORY_ENTRY_IAT].Size;
dh.nOriginalSizeOfImage = m_NtHeader.OptionalHeader.SizeOfImage;
DWORD clrAddr = m_NtHeader.OptionalHeader
.DataDirectory[IMAGE_DIRECTORY_ENTRY_COM_DESCRIPTOR].VirtualAddress;
DWORD clrSize = m_NtHeader.OptionalHeader
.DataDirectory[IMAGE_DIRECTORY_ENTRY_COM_DESCRIPTOR].Size;
if (clrAddr && clrSize) {
PDETOUR_CLR_HEADER pHdr = (PDETOUR_CLR_HEADER)RvaToVa(clrAddr);
if (pHdr != NULL) {
DETOUR_CLR_HEADER hdr;
hdr = *pHdr;
dh.nOriginalClrFlags = hdr.Flags;
}
}
HRESULT hrRet = StringCchCopyA((PCHAR)m_SectionHeaders[nSection].Name, IMAGE_SIZEOF_SHORT_NAME , ".detour");
if (FAILED(hrRet))
return FALSE;
m_SectionHeaders[nSection].Characteristics
= IMAGE_SCN_CNT_INITIALIZED_DATA | IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_WRITE;
m_nOutputVirtAddr = m_nNextVirtAddr;
m_nOutputVirtSize = 0;
m_nOutputFileAddr = m_nNextFileAddr;
dh.nDataOffset = 0; // pbData
dh.cbDataSize = m_pImageData->m_cbData;
dh.cbPrePE = m_cbPrePE;
//////////////////////////////////////////////////////////////////////////
//
DWORD rvaImportTable = 0;
DWORD rvaLookupTable = 0;
DWORD rvaBoundTable = 0;
DWORD rvaNameTable = 0;
DWORD nImportTableSize = nTables * sizeof(IMAGE_IMPORT_DESCRIPTOR);
if (!SizeOutputBuffer(QuadAlign(sizeof(dh))
+ m_cbPrePE
+ QuadAlign(m_pImageData->m_cbData)
+ QuadAlign(sizeof(IMAGE_THUNK_DATA) * nThunks)
+ QuadAlign(sizeof(IMAGE_THUNK_DATA) * nThunks)
+ QuadAlign(nChars)
+ QuadAlign(nImportTableSize))) {
return FALSE;
}
DWORD vaHead = 0;
PBYTE pbHead = NULL;
DWORD vaPrePE = 0;
PBYTE pbPrePE = NULL;
DWORD vaData = 0;
PBYTE pbData = NULL;
if ((pbHead = AllocateOutput(sizeof(dh), &vaHead)) == NULL) {
return FALSE;
}
if ((pbPrePE = AllocateOutput(m_cbPrePE, &vaPrePE)) == NULL) {
return FALSE;
}
CImageThunks lookupTable(this, nThunks, &rvaLookupTable);
CImageThunks boundTable(this, nThunks, &rvaBoundTable);
CImageChars nameTable(this, nChars, &rvaNameTable);
if ((pbData = AllocateOutput(m_pImageData->m_cbData, &vaData)) == NULL) {
return FALSE;
}
dh.nDataOffset = vaData - vaHead;
dh.cbDataSize = dh.nDataOffset + m_pImageData->m_cbData;
CopyMemory(pbHead, &dh, sizeof(dh));
CopyMemory(pbPrePE, m_pMap + m_nPrePE, m_cbPrePE);
CopyMemory(pbData, m_pImageData->m_pbData, m_pImageData->m_cbData);
PIMAGE_IMPORT_DESCRIPTOR piidDst = (PIMAGE_IMPORT_DESCRIPTOR)
AllocateOutput(nImportTableSize, &rvaImportTable);
if (piidDst == NULL) {
return FALSE;
}
//////////////////////////////////////////////// Step Through Imports.
//
for (CImageImportFile *pImportFile = m_pImportFiles;
pImportFile != NULL; pImportFile = pImportFile->m_pNextFile) {
ZeroMemory(piidDst, sizeof(piidDst));
nameTable.Allocate(pImportFile->m_pszName, (DWORD *)&piidDst->Name);
piidDst->TimeDateStamp = 0;
piidDst->ForwarderChain = pImportFile->m_nForwarderChain;
if (pImportFile->m_fByway) {
ULONG rvaIgnored;
lookupTable.Allocate(IMAGE_ORDINAL_FLAG+1,
(DWORD *)&piidDst->OriginalFirstThunk);
boundTable.Allocate(IMAGE_ORDINAL_FLAG+1,
(DWORD *)&piidDst->FirstThunk);
lookupTable.Allocate(0, &rvaIgnored);
boundTable.Allocate(0, &rvaIgnored);
}
else {
ULONG rvaIgnored;
piidDst->FirstThunk = (ULONG)pImportFile->m_rvaFirstThunk;
lookupTable.Current((DWORD *)&piidDst->OriginalFirstThunk);
for (n = 0; n < pImportFile->m_nImportNames; n++) {
CImageImportName *pImportName = &pImportFile->m_pImportNames[n];
if (pImportName->m_pszName) {
ULONG nDstName = 0;
nameTable.Allocate(pImportName->m_pszName,
pImportName->m_nHint,
&nDstName);
lookupTable.Allocate(nDstName, &rvaIgnored);
}
else {
lookupTable.Allocate(IMAGE_ORDINAL_FLAG + pImportName->m_nOrdinal,
&rvaIgnored);
}
}
lookupTable.Allocate(0, &rvaIgnored);
}
piidDst++;
}
ZeroMemory(piidDst, sizeof(piidDst));
//////////////////////////////////////////////////////////////////////////
//
m_nNextVirtAddr += m_nOutputVirtSize;
m_nNextFileAddr += FileAlign(m_nOutputVirtSize);
if (!AlignFileData(hFile)) {
return FALSE;
}
//////////////////////////////////////////////////////////////////////////
//
m_SectionHeaders[nSection].VirtualAddress = m_nOutputVirtAddr;
m_SectionHeaders[nSection].Misc.VirtualSize = m_nOutputVirtSize;
m_SectionHeaders[nSection].PointerToRawData = m_nOutputFileAddr;
m_SectionHeaders[nSection].SizeOfRawData = FileAlign(m_nOutputVirtSize);
m_NtHeader.OptionalHeader
.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress
= rvaImportTable;
m_NtHeader.OptionalHeader
.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].Size
= nImportTableSize;
m_NtHeader.OptionalHeader
.DataDirectory[IMAGE_DIRECTORY_ENTRY_BOUND_IMPORT].VirtualAddress = 0;
m_NtHeader.OptionalHeader
.DataDirectory[IMAGE_DIRECTORY_ENTRY_BOUND_IMPORT].Size = 0;
//////////////////////////////////////////////////////////////////////////
//
if (SetFilePointer(hFile, m_SectionHeaders[nSection].PointerToRawData,
NULL, FILE_BEGIN) == ~0u) {
return FALSE;
}
if (!WriteFile(hFile, m_pbOutputBuffer, m_SectionHeaders[nSection].SizeOfRawData,
&cbDone)) {
return FALSE;
}
}
///////////////////////////////////////////////////// Adjust Extra Data.
//
LONG nExtraAdjust = m_nNextFileAddr - m_nExtraOffset;
for (n = 0; n < m_NtHeader.FileHeader.NumberOfSections; n++) {
if (m_SectionHeaders[n].PointerToRawData > m_nExtraOffset) {
m_SectionHeaders[n].PointerToRawData += nExtraAdjust;
}
if (m_SectionHeaders[n].PointerToRelocations > m_nExtraOffset) {
m_SectionHeaders[n].PointerToRelocations += nExtraAdjust;
}
if (m_SectionHeaders[n].PointerToLinenumbers > m_nExtraOffset) {
m_SectionHeaders[n].PointerToLinenumbers += nExtraAdjust;
}
}
if (m_NtHeader.FileHeader.PointerToSymbolTable > m_nExtraOffset) {
m_NtHeader.FileHeader.PointerToSymbolTable += nExtraAdjust;
}
m_NtHeader.OptionalHeader.CheckSum = 0;
m_NtHeader.OptionalHeader.SizeOfImage = m_nNextVirtAddr;
////////////////////////////////////////////////// Adjust Debug Directory.
//
DWORD debugAddr = m_NtHeader.OptionalHeader
.DataDirectory[IMAGE_DIRECTORY_ENTRY_DEBUG].VirtualAddress;
DWORD debugSize = m_NtHeader.OptionalHeader
.DataDirectory[IMAGE_DIRECTORY_ENTRY_DEBUG].Size;
if (debugAddr && debugSize) {
DWORD nFileOffset = RvaToFileOffset(debugAddr);
if (SetFilePointer(hFile, nFileOffset, NULL, FILE_BEGIN) == ~0u) {
return FALSE;
}
PIMAGE_DEBUG_DIRECTORY pDir = (PIMAGE_DEBUG_DIRECTORY)RvaToVa(debugAddr);
if (pDir == NULL) {
return FALSE;
}
DWORD nEntries = debugSize / sizeof(*pDir);
for (n = 0; n < nEntries; n++) {
IMAGE_DEBUG_DIRECTORY dir = pDir[n];
if (dir.PointerToRawData > m_nExtraOffset) {
dir.PointerToRawData += nExtraAdjust;
}
if (!WriteFile(hFile, &dir, sizeof(dir), &cbDone)) {
return FALSE;
}
}
}
/////////////////////////////////////////////////////// Adjust CLR Header.
//
DWORD clrAddr = m_NtHeader.OptionalHeader
.DataDirectory[IMAGE_DIRECTORY_ENTRY_COM_DESCRIPTOR].VirtualAddress;
DWORD clrSize = m_NtHeader.OptionalHeader
.DataDirectory[IMAGE_DIRECTORY_ENTRY_COM_DESCRIPTOR].Size;
if (clrAddr && clrSize && fNeedDetourSection) {
DWORD nFileOffset = RvaToFileOffset(clrAddr);
if (SetFilePointer(hFile, nFileOffset, NULL, FILE_BEGIN) == ~0u) {
return FALSE;
}
PDETOUR_CLR_HEADER pHdr = (PDETOUR_CLR_HEADER)RvaToVa(clrAddr);
if (pHdr == NULL) {
return FALSE;
}
DETOUR_CLR_HEADER hdr;
hdr = *pHdr;
hdr.Flags &= 0xfffffffe; // Clear the IL_ONLY flag.
if (!WriteFile(hFile, &hdr, sizeof(hdr), &cbDone)) {
return FALSE;
}
}
///////////////////////////////////////////////// Copy Left-over Data.
//
if (m_nFileSize > m_nExtraOffset) {
if (SetFilePointer(hFile, m_nNextFileAddr, NULL, FILE_BEGIN) == ~0u) {
return FALSE;
}
if (!CopyFileData(hFile, m_nExtraOffset, m_nFileSize - m_nExtraOffset)) {
return FALSE;
}
}
//////////////////////////////////////////////////// Finalize Headers.
//
if (SetFilePointer(hFile, m_nPeOffset, NULL, FILE_BEGIN) == ~0u) {
return FALSE;
}
if (!WriteFile(hFile, &m_NtHeader, sizeof(m_NtHeader), &cbDone)) {
return FALSE;
}
if (SetFilePointer(hFile, m_nSectionsOffset, NULL, FILE_BEGIN) == ~0u) {
return FALSE;
}
if (!WriteFile(hFile, &m_SectionHeaders,
sizeof(m_SectionHeaders[0])
* m_NtHeader.FileHeader.NumberOfSections,
&cbDone)) {
return FALSE;
}
m_cbPostPE = SetFilePointer(hFile, 0, NULL, FILE_CURRENT);
if (m_cbPostPE == ~0u) {
return FALSE;
}
m_cbPostPE = m_NtHeader.OptionalHeader.SizeOfHeaders - m_cbPostPE;
return TRUE;
}