ModulesPe.h

#ifndef ROSE_Partitioner2_ModulesPe_H
#define ROSE_Partitioner2_ModulesPe_H

#include <Partitioner2/Function.h>
#include <Partitioner2/Modules.h>
#include <Sawyer/Map.h>

namespace Rose {
namespace BinaryAnalysis {
namespace Partitioner2 {

namespace ModulesPe {

std::string systemFunctionName(const std::string&);

typedef Sawyer::Container::Map<rose_addr_t, SgAsmPEImportItem*> ImportIndex;

std::vector<Function::Ptr> findExportFunctions(const Partitioner&, SgAsmPEFileHeader*);
std::vector<Function::Ptr> findExportFunctions(const Partitioner&, SgAsmInterpretation*);
size_t findExportFunctions(const Partitioner&, SgAsmInterpretation*, std::vector<Function::Ptr>&);
std::vector<Function::Ptr> findImportFunctions(const Partitioner&, SgAsmPEFileHeader*);
std::vector<Function::Ptr> findImportFunctions(const Partitioner&, SgAsmInterpretation*);
size_t findImportFunctions(const Partitioner&, SgAsmPEFileHeader*, const ImportIndex&, std::vector<Function::Ptr>&);
ImportIndex getImportIndex(const Partitioner&, SgAsmPEFileHeader*);
ImportIndex getImportIndex(const Partitioner&, SgAsmInterpretation*);
size_t getImportIndex(const Partitioner&, SgAsmPEFileHeader*, ImportIndex&);
void rebaseImportAddressTables(Partitioner &partitioner, const ImportIndex &index);

void nameImportThunks(const Partitioner&, SgAsmInterpretation*);

void buildMayReturnLists(Partitioner&);

class PeDescrambler: public BasicBlockCallback {
public:
    typedef Sawyer::SharedPointer<PeDescrambler> Ptr;

    struct DispatchEntry {                              // one entry of PEScrambler's dispatch table
        uint32_t returnVa;                              // return address for dispatcher call
        uint32_t calleeVa;                              // address if function that should have been called
        DispatchEntry(): returnVa(0), calleeVa(0) {}
        DispatchEntry(uint32_t returnVa, uint32_t calleeVa): returnVa(returnVa), calleeVa(calleeVa) {}
    };

    typedef std::vector<DispatchEntry> DispatchTable;

private:
    rose_addr_t dispatcherVa_;                          // address of PEScrambler's indirection decoder and dispatcher
    rose_addr_t dispatchTableVa_;                       // address of PEScrambler's dispatch table
    bool reachedEndOfTable_;                            // true when we cannot read any more table entries
    bool checkedPreconditions_;                         // true after we did some first-call precondition checking
    static const size_t sizeofDispatcherFunction = 65;  // default size of PEscrambler dispatcher function in bytes
    static const size_t bitsPerWord = 32;               // this callback only works for 32-bit specimens
    DispatchTable dispatchTable_;                       // the entire dispatch table (and possibly more)

protected:
    // Non-subclass users: please use the instance() method instead; these objects are reference counted.
    PeDescrambler(rose_addr_t dispatcherVa, rose_addr_t dispatchTableVa)
        : dispatcherVa_(dispatcherVa), dispatchTableVa_(dispatchTableVa), reachedEndOfTable_(false),
          checkedPreconditions_(false) {}

public:
    static Ptr instance(rose_addr_t dispatcherVa, rose_addr_t dispatchTableVa) {
        return Ptr(new PeDescrambler(dispatcherVa, dispatchTableVa));
    }

    static Ptr instance(rose_addr_t dispatcherVa) {
        return Ptr(new PeDescrambler(dispatcherVa, dispatcherVa + sizeofDispatcherFunction));
    }

    void nameKeyAddresses(Partitioner&);

    rose_addr_t dispatcherVa() const { return dispatcherVa_; }

    rose_addr_t dispatchTableVa() const { return dispatchTableVa_; }

    const DispatchTable& dispatchTable() const { return dispatchTable_; }
    DispatchTable& dispatchTable() { return dispatchTable_; }
    // Callback invoked by the partitioner each time an instruction is appended to a basic block.
    virtual bool operator()(bool chain, const Args&) ROSE_OVERRIDE;

private:
    // True if the only CFG successor for the specified block is the PEScrambler dispatcher and the block ends with a CALL
    // instruction.
    bool basicBlockCallsDispatcher(const Partitioner&, const BasicBlock::Ptr&);

    // Look up the return address in the PEScrambler dispatch table as if we were the dispatcher and return it if found.
    Sawyer::Optional<rose_addr_t> findCalleeAddress(const Partitioner&, rose_addr_t returnVa);
};

} // namespace
} // namespace
} // namespace
} // namespace

#endif