#include "stdafx.h"

#if defined(_M_ARM64)

#define HOST_MACHINE IMAGE_FILE_MACHINE_ARM64

#elif defined(_M_ARM)

#define HOST_MACHINE IMAGE_FILE_MACHINE_ARM

#elif defined(_WIN64)

#define HOST_MACHINE IMAGE_FILE_MACHINE_AMD64

#else

#define HOST_MACHINE IMAGE_FILE_MACHINE_I386

#endif

#define GET_HEADER_DICTIONARY(headers, idx) &headers->OptionalHeader.DataDirectory[idx]

#define AlignValueUp(value, alignment) ((size_t(value) + size_t(alignment) - 1) & ~(size_t(alignment) - 1))

#define OffsetPointer(data, offset) LPVOID(LPBYTE(data) + ptrdiff_t(offset))

// Protection flags for memory pages (Executable, Readable, Writeable)

static const int ProtectionFlags[2][2][2] = {

{

// not executable

{PAGE_NOACCESS, PAGE_WRITECOPY},

{PAGE_READONLY, PAGE_READWRITE},

}, {

// executable

{PAGE_EXECUTE, PAGE_EXECUTE_WRITECOPY},

{PAGE_EXECUTE_READ, PAGE_EXECUTE_READWRITE},

},

};

int MmpSizeOfImageHeadersUnsafe(PVOID BaseAddress) {

PIMAGE_DOS_HEADER dh = (PIMAGE_DOS_HEADER)BaseAddress;

PIMAGE_NT_HEADERS nh = (PIMAGE_NT_HEADERS)((LPBYTE)BaseAddress + dh->e_lfanew);

//https://learn.microsoft.com/en-us/windows/win32/api/winnt/ns-winnt-image_optional_header32

int sizeOfHeaders = dh->e_lfanew + // e_lfanew member of IMAGE_DOS_HEADER

4 + // 4 byte signature

sizeof(IMAGE_FILE_HEADER) + // size of IMAGE_FILE_HEADER

sizeof(IMAGE_OPTIONAL_HEADER) + // size of optional header

sizeof(IMAGE_SECTION_HEADER) * nh->FileHeader.NumberOfSections; // size of all section headers

return sizeOfHeaders;

}

PMEMORYMODULE WINAPI MapMemoryModuleHandle(HMEMORYMODULE hModule) {

if (!hModule)return nullptr;

PIMAGE_NT_HEADERS nh = RtlImageNtHeader(hModule);

if (!nh)return nullptr;

int sizeOfHeaders = MmpSizeOfImageHeadersUnsafe(hModule);

PMEMORYMODULE pModule = (PMEMORYMODULE)((LPBYTE)hModule + sizeOfHeaders);

if (pModule->Signature != MEMORY_MODULE_SIGNATURE || pModule->codeBase != (LPBYTE)hModule)return nullptr;

return pModule;

}

BOOL WINAPI IsValidMemoryModuleHandle(HMEMORYMODULE hModule) {

return MapMemoryModuleHandle(hModule) != nullptr;

}

NTSTATUS MmpInitializeStructure(DWORD ImageFileSize, LPCVOID ImageFileBuffer, PIMAGE_NT_HEADERS ImageHeaders) {

if (!ImageHeaders)return STATUS_ACCESS_VIOLATION;

//

// Make sure there have enough free space to embed our structure.

//

int sizeOfHeaders = MmpSizeOfImageHeadersUnsafe((PVOID)ImageHeaders->OptionalHeader.ImageBase);

PIMAGE_SECTION_HEADER pSections = IMAGE_FIRST_SECTION(ImageHeaders);

for (int i = 0; i < ImageHeaders->FileHeader.NumberOfSections; ++i) {

if (pSections[i].VirtualAddress < sizeOfHeaders + sizeof(MEMORYMODULE)) {

return STATUS_NOT_SUPPORTED;

}

}

//

// Setup MemoryModule structure.

//

PMEMORYMODULE hMemoryModule = (PMEMORYMODULE)(ImageHeaders->OptionalHeader.ImageBase + sizeOfHeaders);

RtlZeroMemory(hMemoryModule, sizeof(MEMORYMODULE));

hMemoryModule->codeBase = (PBYTE)ImageHeaders->OptionalHeader.ImageBase;

hMemoryModule->dwImageFileSize = ImageFileSize;

hMemoryModule->Signature = MEMORY_MODULE_SIGNATURE;

hMemoryModule->SizeofHeaders = ImageHeaders->OptionalHeader.SizeOfHeaders;

hMemoryModule->lpReserved = (LPVOID)ImageFileBuffer;

hMemoryModule->dwReferenceCount = 1;

return STATUS_SUCCESS;

}

NTSTATUS MemorySetSectionProtection(

_In_ LPBYTE base,

_In_ PIMAGE_NT_HEADERS lpNtHeaders) {

NTSTATUS status = STATUS_SUCCESS;

PIMAGE_SECTION_HEADER section = IMAGE_FIRST_SECTION(lpNtHeaders);

for (DWORD i = 0; i < lpNtHeaders->FileHeader.NumberOfSections; ++i, ++section) {

LPVOID address = LPBYTE(base) + section->VirtualAddress;

SIZE_T size = AlignValueUp(section->Misc.VirtualSize, lpNtHeaders->OptionalHeader.SectionAlignment);

BOOL executable = (section->Characteristics & IMAGE_SCN_MEM_EXECUTE) != 0,

readable = (section->Characteristics & IMAGE_SCN_MEM_READ) != 0,

writeable = (section->Characteristics & IMAGE_SCN_MEM_WRITE) != 0;

DWORD protect = ProtectionFlags[executable][readable][writeable], oldProtect;

if (section->Characteristics & IMAGE_SCN_MEM_NOT_CACHED) protect |= PAGE_NOCACHE;

status = NtProtectVirtualMemory(NtCurrentProcess(), &address, &size, protect, &oldProtect);

if (!NT_SUCCESS(status))break;

}

return status;

}

NTSTATUS MemoryLoadLibrary(

_Out_ HMEMORYMODULE* MemoryModuleHandle,

_In_ LPCVOID data,

_In_ DWORD size) {

PIMAGE_DOS_HEADER dos_header = nullptr;

PIMAGE_NT_HEADERS old_header = nullptr;

BOOLEAN CorImage = FALSE;

NTSTATUS status = STATUS_SUCCESS;

//

// Check parameters

//

__try {

*MemoryModuleHandle = nullptr;

//

// Check dos magic

//

dos_header = (PIMAGE_DOS_HEADER)data;

if (dos_header->e_magic != IMAGE_DOS_SIGNATURE) {

status = STATUS_INVALID_IMAGE_FORMAT;

__leave;

}

//

// Check nt headers

//

old_header = (PIMAGE_NT_HEADERS)((size_t)data + dos_header->e_lfanew);

if (old_header->Signature != IMAGE_NT_SIGNATURE ||

old_header->OptionalHeader.SectionAlignment & 1) {

status = STATUS_INVALID_IMAGE_FORMAT;

__leave;

}

//

// Match machine type

//

if (old_header->FileHeader.Machine != HOST_MACHINE) {

status = STATUS_IMAGE_MACHINE_TYPE_MISMATCH;

__leave;

}

//

// Only dll image support

//

if (!(old_header->FileHeader.Characteristics & IMAGE_FILE_DLL)) {

status = STATUS_NOT_SUPPORTED;

__leave;

}

}

__except (EXCEPTION_EXECUTE_HANDLER) {

status = GetExceptionCode();

}

if (!NT_SUCCESS(status) || status == STATUS_IMAGE_MACHINE_TYPE_MISMATCH)return status;

//

// Reserve the address range of image

//

LPBYTE base = nullptr;

if ((old_header->OptionalHeader.DllCharacteristics & IMAGE_DLLCHARACTERISTICS_DYNAMIC_BASE) == 0) {

base = (LPBYTE)VirtualAlloc(

LPVOID(old_header->OptionalHeader.ImageBase),

old_header->OptionalHeader.SizeOfImage,

MEM_RESERVE,

PAGE_EXECUTE_READWRITE

);

}

if (!base) {

base = (LPBYTE)VirtualAlloc(

nullptr,

old_header->OptionalHeader.SizeOfImage,

MEM_RESERVE,

PAGE_EXECUTE_READWRITE

);

if (!base) status = STATUS_NO_MEMORY;

}

if (!NT_SUCCESS(status)) {

return status;

}

//

// Allocate memory for image headers

//

size_t alignedHeadersSize = (DWORD)AlignValueUp(old_header->OptionalHeader.SizeOfHeaders + sizeof(MEMORYMODULE), MmpGlobalDataPtr->SystemInfo.dwPageSize);

if (!VirtualAlloc(base, alignedHeadersSize, MEM_COMMIT, PAGE_READWRITE)) {

VirtualFree(base, 0, MEM_RELEASE);

status = STATUS_NO_MEMORY;

return status;

}

//

// Copy headers

//

PIMAGE_DOS_HEADER new_dos_header = (PIMAGE_DOS_HEADER)base;

PIMAGE_NT_HEADERS new_header = (PIMAGE_NT_HEADERS)(base + dos_header->e_lfanew);

RtlCopyMemory(

new_dos_header,

dos_header,

old_header->OptionalHeader.SizeOfHeaders

);

new_header->OptionalHeader.ImageBase = (size_t)base;

do {

//

// Setup MEMORYMODULE structure.

//

status = MmpInitializeStructure(size, data, new_header);

if (!NT_SUCCESS(status)) break;

//

// Allocate and copy sections

//

PIMAGE_SECTION_HEADER section = IMAGE_FIRST_SECTION(new_header);

for (DWORD i = 0; i < new_header->FileHeader.NumberOfSections; ++i, ++section) {

DWORD size = AlignValueUp(

section->Misc.VirtualSize,

new_header->OptionalHeader.SectionAlignment

);

if (size < section->SizeOfRawData) {

status = STATUS_INVALID_IMAGE_FORMAT;

break;

}

LPVOID dest = VirtualAlloc(

(LPSTR)new_header->OptionalHeader.ImageBase + section->VirtualAddress,

size,

MEM_COMMIT,

PAGE_READWRITE

);

if (!dest) {

status = STATUS_NO_MEMORY;

break;

}

if (section->SizeOfRawData) {

RtlCopyMemory(

dest,

LPBYTE(data) + section->PointerToRawData,

section->SizeOfRawData

);

}

}

if (!NT_SUCCESS(status))break;

//

// Rebase image

//

auto locationDelta = new_header->OptionalHeader.ImageBase - old_header->OptionalHeader.ImageBase;

if (locationDelta) {

typedef struct _REBASE_INFO {

USHORT Offset : 12;

USHORT Type : 4;

}REBASE_INFO, * PREBASE_INFO;

typedef struct _IMAGE_BASE_RELOCATION_HEADER {

DWORD VirtualAddress;

DWORD SizeOfBlock;

REBASE_INFO TypeOffset[ANYSIZE_ARRAY];

DWORD TypeOffsetCount()const {

return (this->SizeOfBlock - 8) / sizeof(_REBASE_INFO);

}

}IMAGE_BASE_RELOCATION_HEADER, * PIMAGE_BASE_RELOCATION_HEADER;

PIMAGE_DATA_DIRECTORY dir = GET_HEADER_DICTIONARY(new_header, IMAGE_DIRECTORY_ENTRY_BASERELOC);

PIMAGE_BASE_RELOCATION_HEADER relocation = (PIMAGE_BASE_RELOCATION_HEADER)(LPBYTE(base) + dir->VirtualAddress);

if (dir->Size && dir->VirtualAddress) {

while ((LPBYTE(relocation) < LPBYTE(base) + dir->VirtualAddress + dir->Size) && relocation->VirtualAddress > 0) {

auto relInfo = (_REBASE_INFO*)&relocation->TypeOffset;

for (DWORD i = 0; i < relocation->TypeOffsetCount(); ++i, ++relInfo) {

switch (relInfo->Type) {

case IMAGE_REL_BASED_HIGHLOW: *(DWORD*)(base + relocation->VirtualAddress + relInfo->Offset) += (DWORD)locationDelta; break;

#ifdef _WIN64

case IMAGE_REL_BASED_DIR64: *(ULONGLONG*)(base + relocation->VirtualAddress + relInfo->Offset) += (ULONGLONG)locationDelta; break;

#endif

case IMAGE_REL_BASED_ABSOLUTE:

default: break;

}

}

// advance to next relocation block

//relocation->VirtualAddress += module->headers_align;

relocation = decltype(relocation)(OffsetPointer(relocation, relocation->SizeOfBlock));

}

}

}

if (!NT_SUCCESS(status))break;

__try {

*MemoryModuleHandle = (HMEMORYMODULE)base;

}

__except (EXCEPTION_EXECUTE_HANDLER) {

status = GetExceptionCode();

break;

}

return status;

} while (false);

MemoryFreeLibrary((HMEMORYMODULE)base);

return status;

}

BOOL MemoryFreeLibrary(HMEMORYMODULE mod) {

PMEMORYMODULE module = MapMemoryModuleHandle(mod);

PIMAGE_NT_HEADERS headers = RtlImageNtHeader(mod);

if (!module) return FALSE;

if (module->loadFromLdrLoadDllMemory && !module->underUnload)return FALSE;

if (module->hModulesList)MemoryFreeImportTable(module);

if (module->codeBase) VirtualFree(mod, 0, MEM_RELEASE);

return TRUE;

}