relocate.h

Go to the documentation of this file.

 
#pragma once
#include <windows.h>
#include <unordered_set>
 
#include "peconv/buffer_util.h"
 
namespace peconv {
 
    typedef struct _BASE_RELOCATION_ENTRY {
        WORD Offset : 12;
        WORD Type : 4;
    } BASE_RELOCATION_ENTRY;
 
    class RelocBlockCallback
    {
    public:
        RelocBlockCallback(bool _is64bit)
            : is64bit(_is64bit)
        {
        }
 
        virtual ~RelocBlockCallback() {}
 
        virtual bool processRelocField(ULONG_PTR relocField) = 0;
 
    protected:
        bool is64bit;
    };
 
    //--
 
    class CollectRelocs : public peconv::RelocBlockCallback
    {
    public:
        CollectRelocs(const BYTE* pe_buffer, size_t buffer_size, IN bool _is64bit, OUT std::unordered_set<ULONGLONG>& _relocs)
            : RelocBlockCallback(_is64bit), relocs(_relocs),
            peBuffer(pe_buffer), bufferSize(buffer_size)
        {
        }
 
        virtual bool processRelocField(ULONG_PTR relocField)
        {
            ULONGLONG rva = 0;
            if (is64bit) {
                
                ULONGLONG* relocateAddr = (ULONGLONG*)((ULONG_PTR)relocField);
                if (!validate_ptr(peBuffer, bufferSize, relocateAddr, sizeof(ULONGLONG))) {
                    return false;
                }
                rva = (*relocateAddr);
            }
            else {
                DWORD* relocateAddr = (DWORD*)((ULONG_PTR)relocField);
                if (!validate_ptr(peBuffer, bufferSize, relocateAddr, sizeof(DWORD))) {
                    return false;
                }
                rva = ULONGLONG(*relocateAddr);
            }
            relocs.insert(rva);
            return true;
        }
 
    protected:
        std::unordered_set<ULONGLONG>& relocs;
 
        const BYTE* peBuffer;
        size_t bufferSize;
    };
 
    // Processs the relocation table and make your own callback on each relocation field
    bool process_relocation_table(IN PVOID modulePtr, IN SIZE_T moduleSize, IN RelocBlockCallback *callback);
 
    bool relocate_module(IN PBYTE modulePtr, IN SIZE_T moduleSize, IN ULONGLONG newBase, IN ULONGLONG oldBase = 0);
 
    bool has_valid_relocation_table(IN const PBYTE modulePtr, IN const size_t moduleSize);
 
    bool virtual_addr_to_rva(IN const PBYTE imgBase, IN const size_t imgSize, IN ULONGLONG virtualAddr, OUT DWORD& outRVA, IN std::unordered_set<ULONGLONG>* relocs=nullptr);
 
};//namespace peconv