AuroraRuntime/Source/Process/ProcessSectionViewReserved.NT.cpp

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

    File: ProcessSectionViewReserved.NT.cpp
    Date: 2022-9-30
    Author: Reece
***/
#include <Source/RuntimeInternal.hpp>
#include "ProcessSectionViewReserved.NT.hpp"
#include "Process.hpp"
#include <Source/IO/FS/FileStream.NT.hpp>
#include "ProcessSectionFileMapView.NT.hpp"
#include "ProcessSectionView.NT.hpp"
#include <Source/IO/IPC/AuIPCHandle.hpp>
#include <Windows.h>

namespace Aurora::Process
{
    static PVOID (__stdcall *VirtualAlloc2_f)(
        HANDLE                 Process,
        PVOID                  BaseAddress,
        SIZE_T                 Size,
        ULONG                  AllocationType,
        ULONG                  PageProtection,
        MEM_EXTENDED_PARAMETER *ExtendedParameters,
        ULONG                  ParameterCount
    );

    static PVOID (__stdcall *MapViewOfFile3_f)(
        HANDLE                 FileMapping,
        HANDLE                 Process,
        PVOID                  BaseAddress,
        ULONG64                Offset,
        SIZE_T                 ViewSize,
        ULONG                  AllocationType,
        ULONG                  PageProtection,
        MEM_EXTENDED_PARAMETER *ExtendedParameters,
        ULONG                  ParameterCount
    );

    static PVOID(__stdcall *UnmapViewOfFile2_f)(
        HANDLE Process,
        PVOID  BaseAddress,
        ULONG  UnmapFlags
    );

    bool ProcessSectionViewReserved::Init(AuUInt uLength)
    {
        this->uMaxLength = uLength;

        this->pBaseAddress = AuReinterpretCast<decltype(this->pBaseAddress)>(VirtualAlloc2_f(
            nullptr,
            nullptr,
            this->uMaxLength,
            MEM_RESERVE | MEM_RESERVE_PLACEHOLDER,
            PAGE_NOACCESS,
            nullptr,
            0));

        return bool(this->pBaseAddress);
    }

    bool ProcessSectionViewReserved::AllocateAddress(AuUInt uOffset,
                                                     AuUInt uLength,
                                                     AuUInt &uAddressOut)
    {
        bool bCompleteBlock {};
        AuUInt uFoundOffset {};

        uAddressOut = 0;

        if (!this->GetAddress(uOffset, uLength, uFoundOffset, bCompleteBlock))
        {
            SysPushErrorMemory("Reserved address space has no available space");
            return false;
        }
        
        if (bCompleteBlock)
        {
            uAddressOut = uFoundOffset;
            return true;
        }

        if (!::VirtualFree(this->pBaseAddress + uFoundOffset,
                           uLength,
                           MEM_RELEASE | MEM_PRESERVE_PLACEHOLDER))
        {
            SysPushErrorMemory("System refused to split desired region");
            ReleaseAddress(uFoundOffset);
            uAddressOut = uFoundOffset;
            return false;
        }

        return true;
    }

    bool ProcessSectionViewReserved::ReleaseAndCoaleceAddress(AuUInt uOffset, AuUInt uLength)
    {
        AU_LOCK_GUARD(this->spinlock);
        bool bSuccess {};

        for (AuUInt i = 0; i < this->allocations.size(); i++)
        {
            auto &a = this->allocations[i];
            if (a.first != uOffset)
            {
                continue;
            }

            if (uLength != a.second)
            {
                break;
            }

            AuUInt uPrevFree {};
            AuUInt uNextFree {};

            {
                if (i != 0)
                {
                    auto eh = this->allocations[i - 1];
                    uPrevFree = eh.first + eh.second;
                }
                else
                {
                    uPrevFree = uOffset;
                }
            }

            {
                AuUInt uStart {};
                if (i != this->allocations.size() - 1)
                {
                    uNextFree = this->allocations[i + 1].first;
                }
                else if (this->allocations.size() == 1)
                {
                    uNextFree = this->uMaxLength;
                }
                else
                {
                    uNextFree = uOffset + uLength;
                }
            }

            if (!UnmapViewOfFile2_f(INVALID_HANDLE_VALUE,
                                    this->pBaseAddress + uOffset,
                                    MEM_PRESERVE_PLACEHOLDER))
            {
                SysPushErrorMemory("System refused to restore placeholder: {}", GetLastError());
                break;
            }

        #if 0
            if (!::VirtualFree(this->pBaseAddress + uOffset, 
                               0,
                               MEM_RELEASE | MEM_PRESERVE_PLACEHOLDER))
            {
                SysPushErrorMemory("System refused to restore placeholder: {}", GetLastError());
                break;
            }
        #endif

            auto uFreeRange = uNextFree - uPrevFree;
            if (uFreeRange != uLength)
            {
                if (!::VirtualFree(this->pBaseAddress + uPrevFree,
                                   uFreeRange,
                                   MEM_RELEASE | MEM_COALESCE_PLACEHOLDERS))
                {
                    SysPushErrorMemory("System refused to coalece memory segments");
                    break;
                }
            }
            
            bSuccess = true;
            break;
        }

        if (bSuccess)
        {
            AuTryRemoveByTupleN<0>(this->allocations, uOffset);
        }

        return bSuccess;
    }

    AuSPtr<IProcessSectionMapView> ProcessSectionViewReserved::Allocate(AuUInt uLength)
    {
        PageTable table {};
        table.NX = true;
        table.readable = true;
        table.writable = true;
        return this->AllocateEx2(uLength, -1, table);
    }

    AuSPtr<IProcessSectionMapView> ProcessSectionViewReserved::MapFileByPath(const AuString &str,
                                                                             AuUInt64 uOffset,
                                                                             AuUInt uLength,
                                                                             AuFS::EFileOpenMode mode,
                                                                             AuFS::EFileAdvisoryLockLevel sectionLock)
    {
        return this->MapFileByPathEx(-1, str, uOffset, uLength, mode, sectionLock);
    }

    AuSPtr<IProcessSectionMapView> ProcessSectionViewReserved::MapFileByObject(const AuSPtr<AuFS::IFileStream> &pStream,
                                                                               AuUInt64 uOffset,
                                                                               AuUInt uLength,
                                                                               AuFS::EFileOpenMode mode,
                                                                               AuFS::EFileAdvisoryLockLevel processLockLevel)
    {
        return this->MapFileByObjectEx(-1, pStream, uOffset, uLength, mode, processLockLevel);
    }

    AuSPtr<IProcessSectionMapView> ProcessSectionViewReserved::MapIPCMemory(const AuString &handleString,
                                                                            AuUInt64 uOffset,
                                                                            AuUInt uLength,
                                                                            AuFS::EFileOpenMode mode)
    {
        return this->MapIPCMemoryEx(-1, handleString, uOffset, uLength, mode);
    }

    AuSPtr<IProcessSectionMapView> ProcessSectionViewReserved::AllocateEx(AuUInt uLength,
                                                                          AuUInt uOffset)
    {
        PageTable table {};
        table.NX = true;
        table.readable = true;
        table.writable = true;
        return this->AllocateEx2(uLength, uOffset, table);
    }

    AuSPtr<IProcessSectionMapView> ProcessSectionViewReserved::AllocateEx2(AuUInt uLength,
                                                                           AuUInt uOffset,
                                                                           PageTable permissions)
    {
        HANDLE hFileMap;
        DWORD uPageFlags {};
        AuUInt uAddressOut;

        if (!uLength)
        {
            SysPushErrorArg("invalid uLength");
            return {};
        }

        if (permissions.writable && permissions.NX)
        {
            uPageFlags = PAGE_READWRITE;
        }
        else if (permissions.readable && permissions.NX)
        {
            uPageFlags = PAGE_READONLY;
        }
        else if (permissions.writable)
        {
            uPageFlags = PAGE_EXECUTE_READWRITE;
        }
        else if (permissions.readable)
        {
            uPageFlags = PAGE_EXECUTE_READ;
        }

        if (!this->AllocateAddress(uOffset, uLength, uAddressOut))
        {
            return {};
        }

        hFileMap = ::CreateFileMappingA(INVALID_HANDLE_VALUE,
                                        nullptr,
                                        uPageFlags,
                                     #if defined(AURORA_IS_64BIT)
                                        AuBitsToHigher(uLength),
                                        AuBitsToLower(uLength),
                                     #else
                                        0,
                                        uLength,
                                     #endif
                                        nullptr);
        if ((hFileMap == INVALID_HANDLE_VALUE) ||
            (!hFileMap))
        {
            SysPushErrorIO("Couldn't create file map");
            return {};
        }

        auto map = MapViewOfFile3_f(hFileMap,
                                    INVALID_HANDLE_VALUE,
                                    this->pBaseAddress + uAddressOut,
                                    0,
                                    uLength,
                                    MEM_REPLACE_PLACEHOLDER,
                                    uPageFlags,
                                    nullptr,
                                    0);
        if (!map)
        {
            SysPushErrorIO("Couldn't create allocation of section");
            AuWin32CloseHandle(hFileMap);
            return {};
        }

        auto pNewObject = AuMakeShared<ProcessSectionFileMapView>(AuUInt(map), hFileMap);
        if (!pNewObject)
        {
            SysPushErrorMem();
            AuWin32CloseHandle(hFileMap);
            return {};
        }

        pNewObject->pSharedSectionHint = AuSharedFromThis();
        pNewObject->uOffset = uAddressOut;
        pNewObject->uLength = uLength;

        return pNewObject;
    }

    AuSPtr<IProcessSectionMapView> ProcessSectionViewReserved::MapFileByPathEx(AuUInt viewOffset,
                                                                               const AuString &str,
                                                                               AuUInt64 uOffset,
                                                                               AuUInt uLength,
                                                                               Aurora::IO::FS::EFileOpenMode mode,
                                                                               Aurora::IO::FS::EFileAdvisoryLockLevel sectionLock)
    {
        auto file = AuFS::OpenShared(str, mode, AuFS::EFileAdvisoryLockLevel::eNoSafety);
        return file ? this->MapFileByObjectEx(viewOffset, file, uOffset, uLength, mode, sectionLock) : AuSPtr<IProcessSectionMapView> {};
    }

    AuSPtr<IProcessSectionMapView> ProcessSectionViewReserved::MapFileByObjectEx(AuUInt viewOffset,
                                                                                 const AuSPtr<Aurora::IO::FS::IFileStream> &pStream,
                                                                                 AuUInt64 uOffset,
                                                                                 AuUInt uLength,
                                                                                 Aurora::IO::FS::EFileOpenMode mode,
                                                                                 Aurora::IO::FS::EFileAdvisoryLockLevel processLockLevel)
    {
        HANDLE hFileMap;
        ULONG desiredAccess {}, pageAttributes {};
        AuUInt uAddressOut;

        if (!pStream)
        {
            SysPushErrorArg();
            return {};
        }

        if (!uLength)
        {
            SysPushErrorArg("invalid uLength");
            return {};
        }

        if (!this->AllocateAddress(viewOffset, uLength, uAddressOut))
        {
            return {};
        }

        auto pFileStream = AuStaticCast<AuFS::WinFileStream>(pStream);
        if (processLockLevel != AuFS::EFileAdvisoryLockLevel::eNoSafety)
        {
            DWORD dwFlags {};
            OVERLAPPED overlapped {};

            if (processLockLevel == AuFS::EFileAdvisoryLockLevel::eBlockReadWrite)
            {
                dwFlags |= LOCKFILE_EXCLUSIVE_LOCK;
            }

            dwFlags |= LOCKFILE_FAIL_IMMEDIATELY;

            overlapped.Offset = AuBitsToLower(uOffset);
            overlapped.OffsetHigh = AuBitsToHigher(uOffset);

            if (!::LockFileEx(pFileStream->GetHandle(),
                              dwFlags,
                              0,
                              AuBitsToLower(uLength),
                              AuBitsToHigher(uLength),
                              &overlapped))
            {
                SysPushErrorIO("No Lock");
                return {};
            }
        }
   
        switch (mode)
        {
        case AuFS::EFileOpenMode::eRead:
        {
            desiredAccess  = SECTION_MAP_READ;
            pageAttributes = PAGE_READONLY;
            break;
        }
        case AuFS::EFileOpenMode::eWrite:
        case AuFS::EFileOpenMode::eReadWrite:
        {
            desiredAccess  = SECTION_MAP_READ | SECTION_MAP_WRITE;
            pageAttributes = PAGE_READWRITE;
            break;
        }
        default:
            SysPushErrorGeneric();
            return {};
        };
        
        hFileMap = ::CreateFileMappingA(pFileStream->GetHandle(),
                                        nullptr,
                                        pageAttributes,
                                     #if defined(AURORA_IS_64BIT)
                                        AuBitsToHigher(uLength),
                                        AuBitsToLower(uLength),
                                     #else
                                        0,
                                        uLength,
                                     #endif
                                        nullptr);
        if ((hFileMap == INVALID_HANDLE_VALUE) ||
            (!hFileMap))
        {
            SysPushErrorIO("Couldn't create file map. Is the requeted access mode too high for the given object?");
            return {};
        }

        auto map = MapViewOfFile3_f(hFileMap,
                                    INVALID_HANDLE_VALUE,
                                    this->pBaseAddress + uAddressOut,
                                    uOffset,
                                    uLength,
                                    MEM_REPLACE_PLACEHOLDER,
                                    pageAttributes,
                                    nullptr,
                                    0);
        if (!map)
        {
            SysPushErrorIO("Couldn't create map of section: {}", GetLastError());
            AuWin32CloseHandle(hFileMap);
            return {};
        }

        auto pNewObject = AuMakeShared<ProcessSectionFileMapView>(AuUInt(map), hFileMap);
        if (!pNewObject)
        {
            SysPushErrorMem();
            AuWin32CloseHandle(hFileMap);
            ::UnmapViewOfFile(map);
            return {};
        }

        pNewObject->pSharedSectionHint = AuSharedFromThis();
        pNewObject->uOffset = uAddressOut;
        pNewObject->uLength = uLength;

        return pNewObject;
    }

    AuSPtr<IProcessSectionMapView> ProcessSectionViewReserved::MapIPCMemoryEx(AuUInt viewOffset,
                                                                              const AuString &handleString,
                                                                              AuUInt64 uOffset,
                                                                              AuUInt uLength,
                                                                              Aurora::IO::FS::EFileOpenMode mode)
    {
        AuIPC::IPCHandle handle;
        HANDLE hFileMap;
        ULONG desiredAccess {}, pageAttributes {};
        AuUInt uAddressOut;

        if (!uLength)
        {
            SysPushErrorArg("invalid uLength");
            return {};
        }

        if (!handle.FromString(handleString))
        {
            SysPushErrorParseError("{}", handleString);
            return {};
        }

        auto token = handle.GetToken(AuIPC::EIPCHandleType::eIPCMemory, 0);
        if (!token)
        {
            SysPushErrorParseError();
            return {};
        }

        auto actualLength = token->token.word;
        auto path         = token->token.ToNTPath();

        if (actualLength < uOffset + uLength)
        {
            SysPushErrorIO("Out of range");
            return {};
        }

        switch (mode)
        {
        case AuFS::EFileOpenMode::eRead:
        {
            desiredAccess  = SECTION_MAP_READ;
            pageAttributes = PAGE_READONLY;
            break;
        }
        case AuFS::EFileOpenMode::eWrite:
        case AuFS::EFileOpenMode::eReadWrite:
        {
            pageAttributes = PAGE_READWRITE;
            break;
        }
        default:
            SysPushErrorGeneric();
            return {};
        };

        if (!this->AllocateAddress(viewOffset, uLength, uAddressOut))
        {
            return {};
        }

        hFileMap = ::OpenFileMappingA(desiredAccess,
                                       FALSE,
                                       path.c_str());
        if ((hFileMap == INVALID_HANDLE_VALUE) ||
            (!hFileMap))
        {
            SysPushErrorIO("Couldn't create IPC map (handle: {})", handleString);
            return {};
        }
        
        auto map = MapViewOfFile3_f(hFileMap,
                                    INVALID_HANDLE_VALUE,
                                    this->pBaseAddress + uAddressOut,
                                    uOffset,
                                    uLength,
                                    MEM_REPLACE_PLACEHOLDER,
                                    pageAttributes,
                                    nullptr,
                                    0);
        if (!map)
        {
            SysPushErrorIO("Couldn't create map of IPC section (handle: {})", handleString);
            AuWin32CloseHandle(hFileMap);
            return {};
        }

        auto pNewObject = AuMakeShared<ProcessSectionFileMapView>(AuUInt(map), hFileMap);
        if (!pNewObject)
        {
            SysPushErrorMem();
            AuWin32CloseHandle(hFileMap);
            ::UnmapViewOfFile(map);
            return {};
        }

        pNewObject->pSharedSectionHint = AuSharedFromThis();
        pNewObject->uOffset = uAddressOut;
        pNewObject->uLength = uLength;

        return pNewObject;
    }

    AuList<AuPair<AuUInt, AuUInt>> ProcessSectionViewReserved::GetAllocations()
    {
        AU_LOCK_GUARD(this->spinlock);
        return this->allocations;
    }

    AUKN_SYM AuSPtr<IProcessSectionView> ReserveAddressSpace(AuUInt uLength)
    {
        auto &platform = AuSwInfo::GetPlatformInfo();

        if (!AuSwInfo::IsWindows10OrGreater() || // < win10 or
            ((platform.uKernelMajor == 10) && (platform.uKernelPatch < 1803)) ||
             !VirtualAlloc2_f ||
             !MapViewOfFile3_f) // < RS4
        {

        #define WIN_7_WARN "WARNING: ADDRESS SPACE CANNOT BE RESERVED ON OLDER NT KERNELS. \r\n" \
                           "AuProcess::ReserveAddressSpace(AuUInt uOffset) is about to lie about reserving the address space, yield the entire address space, and leave a note to terminate the application if an explicit fixed-offset request is made."

            SysPushErrorUnimplemented("Win7_ReserveAddressSpace_RS4_REQ");
            AuLogWarn(WIN_7_WARN);
            static ProcessSectionView gSingleton;
            gSingleton.bPanicOnEx = true;
            return AuUnsafeRaiiToShared(&gSingleton);
        }

        auto pSectionView = AuMakeShared<ProcessSectionViewReserved>();
        if (!pSectionView)
        {
            SysPushErrorMemory();
            return {};
        }

        if (!pSectionView->Init(uLength))
        {
            SysPushErrorMemory();
            return {};
        }

        return pSectionView;
    }

    void LoadProcessSectionViewSymbol()
    {
        auto hKernelHandle = LoadLibraryA("Kernel32.dll");
        auto hKernelHandle2 = LoadLibraryA("KernelBase.dll");
        VirtualAlloc2_f = AuReinterpretCast<decltype(VirtualAlloc2_f)>(::GetProcAddress(hKernelHandle, "VirtualAlloc2"));
        if (!VirtualAlloc2_f)
        {
            VirtualAlloc2_f = AuReinterpretCast<decltype(VirtualAlloc2_f)>(::GetProcAddress(hKernelHandle2, "VirtualAlloc2"));
        }

        MapViewOfFile3_f = AuReinterpretCast<decltype(MapViewOfFile3_f)>(::GetProcAddress(hKernelHandle, "MapViewOfFile3"));
        if (!MapViewOfFile3_f)
        {
            MapViewOfFile3_f = AuReinterpretCast<decltype(MapViewOfFile3_f)>(::GetProcAddress(hKernelHandle2, "MapViewOfFile3"));
        }

        UnmapViewOfFile2_f = AuReinterpretCast<decltype(UnmapViewOfFile2_f)>(::GetProcAddress(hKernelHandle, "UnmapViewOfFile2"));
        if (!UnmapViewOfFile2_f)
        {
            UnmapViewOfFile2_f = AuReinterpretCast<decltype(UnmapViewOfFile2_f)>(::GetProcAddress(hKernelHandle2, "UnmapViewOfFile2"));
        }
    }
}