AuroraRuntime/Source/Process/ProcessMap.cpp

/***
    Copyright (C) 2022 J Reece Wilson (a/k/a "Reece"). All rights reserved.

    File: ProcessMap.cpp
    Date: 2022-1-23
    Author: Reece
***/
#include <Source/RuntimeInternal.hpp>
#include "ProcessMap.hpp"

#if defined(AURORA_PLATFORM_WIN32)
    #include "ProcessMap.Win32.hpp"
#endif

#if defined(AURORA_IS_MODERNNT_DERIVED)
    #include "ProcessMap.NT.hpp"
#endif

namespace Aurora::Process
{
    struct ModuleBasePairHash
    {
        AuUInt operator()(const ModuleBasePair &in) const
        {
            return in.modBase;
        }
    };

    struct ModuleBasePairEq
    {
        constexpr bool operator()(const ModuleBasePair &lhs, const AuString &rhs) const
        {
            return lhs.module == rhs;
        }

        constexpr bool operator()(const ModuleBasePair &lhs, const AuUInt &rhs) const
        {
            return lhs.modBase == rhs;
        }

        constexpr bool operator()(const ModuleBasePair &lhs, const ModuleBasePair &rhs) const
        {
            return rhs.modBase ? lhs.modBase == rhs.modBase : (rhs.module.size() ? rhs.module == rhs.module : false);
        }
    };
    
    struct ModuleLookup : Section
    {
        ModuleLookup()
        {}

        ModuleLookup(const Section &s) : Section(s)
        {}
    };

    static AuBST<AuUInt, ModuleLookup> gModulePtrMap;

    static const auto kMinPageAlignment = 4096;
    static const auto kPageBufferPad = 20;

    static AuThreadPrimitives::MutexUnique_t gMutexUnique;
    static AuHashMap<ModuleBasePair, AuSPtr<PublicModule>, ModuleBasePairHash, ModuleBasePairEq> gModuleMap;

    static AuUInt ToLowestPageAlignment(AuUInt in)
    {
        return in & ~(kMinPageAlignment - 1);
    }

    static AuSPtr<PublicModule> GetModuleFromSectionCache(AuUInt pointer)
    {
        AU_LOCK_GUARD(gMutexUnique);
        auto itr = gModulePtrMap.find(pointer);
        if (itr == gModulePtrMap.end()) return {};
        return itr->second.moduleMeta.lock();
    }

    static AuOptional<Section>  GetSectionCache(AuUInt pointer)
    {
        AU_LOCK_GUARD(gMutexUnique);
        auto itr = gModulePtrMap.find(pointer);
        if (itr == gModulePtrMap.end()) return {};
        return itr->second;
    }

    bool IsInModuleCache(const ModuleBasePair &pair)
    {
        AU_LOCK_GUARD(gMutexUnique);
        return gModuleMap.find(pair) != gModuleMap.end();
    }

    void InsertModuleCache(const ModuleBasePair &pair, const AuSPtr<PublicModule> &mod)
    {
        for (auto &section : mod->sections)
        {
            AU_LOCK_GUARD(gMutexUnique);
            section.moduleMeta = mod;

            for (AuUInt i = section.baseVa; i < section.baseVa + section.size; i += (kMinPageAlignment * kPageBufferPad))
            {
                ModuleLookup a(section);
                gModulePtrMap[i] = a;
            }
        }

        {
            AU_LOCK_GUARD(gMutexUnique);
            gModuleMap[pair] = mod;
        }
    }

    void RemoveModuleCache(const ModuleBasePair &eitherOr)
    {
        AU_LOCK_GUARD(gMutexUnique);
        auto itr = gModuleMap.find(eitherOr);
        if (itr == gModuleMap.end()) return;
        auto mod = itr->second;

        for (const auto &section : mod->sections)
        {
            for (AuUInt i = section.baseVa; i < section.baseVa + section.size; i += (kMinPageAlignment * kPageBufferPad))
            {
                auto itr = gModulePtrMap.find(i);
                if (itr != gModulePtrMap.end())
                {
                    gModulePtrMap.erase(itr);
                }
            }
        }

        gModuleMap.erase(itr);
    }

    PublicModule GetFromModuleCache(AuUInt handle)
    {
        AU_LOCK_GUARD(gMutexUnique);
        auto itr = gModuleMap.find({"", handle});
        if (itr == gModuleMap.end()) return {};
        return *itr->second;
    }

    static AuOptional<Section>  FindInCache(AuUInt pointer)
    {
        auto curPtr = ToLowestPageAlignment(pointer);
        auto temp = GetSectionCache(curPtr);
        if (temp.has_value()) return temp;

        for (int i = 0; i < kPageBufferPad + 1; i++)
        {
            curPtr -= kMinPageAlignment;
            temp = GetSectionCache(curPtr); // TODO: i dont want to start from the top of the tree, thats stupid
            if (temp.has_value()) return temp;
        }

        return AuOptional<Section>{};
    }

    void InitProcessMap()
    {
        gMutexUnique = AuThreadPrimitives::MutexUnique();

    #if defined(AURORA_IS_MODERNNT_DERIVED)
        InitProcessMapNt();
    #endif

        TryRescanSlow();
    }

    AUKN_SYM void TryRescanSlow()
    {
    #if defined(AURORA_PLATFORM_WIN32)
        MakeToolHelp32Snapshot();
    #endif
    }

    void DeinitProcessMap()
    {
        gMutexUnique.reset();
        DeinitProcessMapNt();
    }

    AUKN_SYM AuOptional<Section>  GetSection(AuUInt pointer)
    {
        try 
        {
            auto ceg = FindInCache(pointer);
            if (ceg.has_value())
            {
                return ceg;
            }

        #if defined(AURORA_IS_MODERNNT_DERIVED)
            if (MakeAwarePtr(pointer))
            {
                return FindInCache(pointer);
            }
            else
            {
                return LookupArbitrarySection(pointer);
            }
        #endif

            return {};
        }
        catch (...)
        {
            return {};
        }
    }

    AUKN_SYM PublicModule DumpExecutableRoot()
    {
        try
        {
            return GetExecutableRoot();
        }
        catch (...)
        {
            return {};
        }
    }

    AUKN_SYM Sections DumpExecutableAll()
    {
        try
        {
            Sections ret;
            for (const auto &[meta, ptr] : gModuleMap)
            {
                ret.insert(ret.end(), ptr->sections.begin(), ptr->sections.end());
            }
            return ret;
        }
        catch (...)
        {
            return {};
        }
    }
}