BOOL CImage::Write()

in Detours/image.cpp [1627:2060]


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);
#if 0
        m_nNextVirtAddr = Max(m_SectionHeaders[n].VirtualAddress +
                              m_SectionHeaders[n].Misc.VirtualSize,
                              m_nNextVirtAddr);
#else
        m_nNextVirtAddr = Max(m_SectionHeaders[n].VirtualAddress +
                              (m_SectionHeaders[n].Misc.VirtualSize
                               ? m_SectionHeaders[n].Misc.VirtualSize
                               : SectionAlign(m_SectionHeaders[n].SizeOfRawData)),
                              m_nNextVirtAddr);
#endif

        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;
}