#pragma warning( disable : 4311 4312 )

#include <Windows.h>

#include <winnt.h>

#include "ntinternals.h"

#define DEBUG_OUTPUT 0

#ifdef DEBUG_OUTPUT

#include <stdio.h>

#include "MemoryModule.h"

typedef BOOL (WINAPI *DllEntryProc)(HINSTANCE hinstDLL, DWORD fdwReason, LPVOID lpReserved);

#define GET_NT_HEADER(module) ((PIMAGE_NT_HEADERS)&((const unsigned char *)(module))[((PIMAGE_DOS_HEADER)(module))->e_lfanew])

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

#define CALCULATE_REAL_ADDRESS(base, offset) (((unsigned char *)(base)) + (offset))

static DWORD ModuleCount = 0;

#ifdef DEBUG_OUTPUT

static void

OutputLastError(const char *msg)

{

LPVOID tmp;

char *tmpmsg;

FormatMessage(FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS,

NULL, GetLastError(), MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), (LPTSTR)&tmp, 0, NULL);

tmpmsg = (char *)LocalAlloc(LPTR, strlen(msg) + strlen(tmp) + 3);

sprintf(tmpmsg, "%s: %s", msg, tmp);

OutputDebugString(tmpmsg);

LocalFree(tmpmsg);

LocalFree(tmp);

}

static void

CopySections(const unsigned char *data, PIMAGE_NT_HEADERS old_headers, HMODULE module)

{

DWORD i, size;

LPVOID dest;

PIMAGE_SECTION_HEADER section;

section = IMAGE_FIRST_SECTION(GET_NT_HEADER(module));

for (i=0; i<GET_NT_HEADER(module)->FileHeader.NumberOfSections; i++, section++)

{

if (section->SizeOfRawData == 0)

{

// section doesn't contain data in the dll itself, but may define

// uninitialized data

size = old_headers->OptionalHeader.SectionAlignment;

if (size > 0)

{

dest = VirtualAlloc(CALCULATE_REAL_ADDRESS(module, section->VirtualAddress),

size,

MEM_COMMIT,

PAGE_READWRITE);

section->Misc.PhysicalAddress = (DWORD)dest;

memset(dest, 0, size);

}

// section is empty

continue;

}

// commit memory block and copy data from dll

dest = VirtualAlloc(CALCULATE_REAL_ADDRESS(module, section->VirtualAddress),

section->SizeOfRawData,

MEM_COMMIT,

PAGE_READWRITE);

memcpy(dest, data + section->PointerToRawData, section->SizeOfRawData);

section->Misc.PhysicalAddress = (DWORD)dest;

}

}

static 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},

},

};

static void

FinalizeSections(HMODULE module)

{

DWORD i;

PIMAGE_SECTION_HEADER section;

// loop through all sections and change access flags

section = IMAGE_FIRST_SECTION(GET_NT_HEADER(module));

for (i=0; i<GET_NT_HEADER(module)->FileHeader.NumberOfSections; i++, section++)

{

DWORD protect, oldProtect, size;

int executable = (section->Characteristics & IMAGE_SCN_MEM_EXECUTE) != 0;

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

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

if (section->Characteristics & IMAGE_SCN_MEM_DISCARDABLE)

{

// section is not needed any more and can safely be freed

VirtualFree((LPVOID)section->Misc.PhysicalAddress, section->SizeOfRawData, MEM_DECOMMIT);

continue;

}

// determine protection flags based on characteristics

protect = ProtectionFlags[executable][readable][writeable];

if (section->Characteristics & IMAGE_SCN_MEM_NOT_CACHED)

protect |= PAGE_NOCACHE;

// determine size of region

size = section->SizeOfRawData;

if (size == 0)

{

if (section->Characteristics & IMAGE_SCN_CNT_INITIALIZED_DATA)

size = GET_NT_HEADER(module)->OptionalHeader.SizeOfInitializedData;

else if (section->Characteristics & IMAGE_SCN_CNT_UNINITIALIZED_DATA)

size = GET_NT_HEADER(module)->OptionalHeader.SizeOfUninitializedData;

}

if (size > 0)

{

// change memory access flags

if (VirtualProtect((LPVOID)section->Misc.PhysicalAddress, section->SizeOfRawData, protect, &oldProtect) == 0)

#ifdef DEBUG_OUTPUT

OutputLastError("Error protecting memory page")

;

}

}

}

static void

PerformBaseRelocation(HMODULE module, DWORD delta)

{

DWORD i;

PIMAGE_DATA_DIRECTORY directory;

directory = GET_HEADER_DICTIONARY(module, IMAGE_DIRECTORY_ENTRY_BASERELOC);

if (directory->Size > 0)

{

PIMAGE_BASE_RELOCATION relocation = (PIMAGE_BASE_RELOCATION)CALCULATE_REAL_ADDRESS(module, directory->VirtualAddress);

for (; relocation->VirtualAddress > 0; )

{

unsigned char *dest = (unsigned char *)CALCULATE_REAL_ADDRESS(module, relocation->VirtualAddress);

unsigned short *relInfo = (unsigned short *)CALCULATE_REAL_ADDRESS(relocation, IMAGE_SIZEOF_BASE_RELOCATION);

for (i=0; i<((relocation->SizeOfBlock-IMAGE_SIZEOF_BASE_RELOCATION) / 2); i++, relInfo++)

{

DWORD *patchAddrHL;

int type, offset;

// the upper 4 bits define the type of relocation

type = *relInfo >> 12;

// the lower 12 bits define the offset

offset = *relInfo & 0xfff;

switch (type)

{

case IMAGE_REL_BASED_ABSOLUTE:

// skip relocation

break;

case IMAGE_REL_BASED_HIGHLOW:

// change complete 32 bit address

patchAddrHL = (DWORD *)(dest + offset);

*patchAddrHL += delta;

break;

default:

//printf("Unknown relocation: %d\n", type);

break;

}

}

// advance to next relocation block

(char *)relocation += relocation->SizeOfBlock;

}

}

}

static int

BuildImportTable(HMODULE module)

{

int result=1;

PIMAGE_DATA_DIRECTORY directory = GET_HEADER_DICTIONARY(module, IMAGE_DIRECTORY_ENTRY_IMPORT);

if (directory->Size > 0)

{

PIMAGE_IMPORT_DESCRIPTOR importDesc = (PIMAGE_IMPORT_DESCRIPTOR)CALCULATE_REAL_ADDRESS(module, directory->VirtualAddress);

for (; !IsBadReadPtr(importDesc, sizeof(IMAGE_IMPORT_DESCRIPTOR)) && importDesc->Name; importDesc++)

{

DWORD *thunkRef, *funcRef;

HMODULE handle = LoadLibrary((LPCSTR)CALCULATE_REAL_ADDRESS(module, importDesc->Name));

if (handle == INVALID_HANDLE_VALUE)

{

SetLastError(ERROR_MOD_NOT_FOUND);

#if DEBUG_OUTPUT

OutputLastError("Can't load library");

result = 0;

break;

}

if (importDesc->OriginalFirstThunk)

{

thunkRef = (DWORD *)CALCULATE_REAL_ADDRESS(module, importDesc->OriginalFirstThunk);

funcRef = (DWORD *)CALCULATE_REAL_ADDRESS(module, importDesc->FirstThunk);

} else {

// no hint table

thunkRef = (DWORD *)CALCULATE_REAL_ADDRESS(module, importDesc->FirstThunk);

funcRef = (DWORD *)CALCULATE_REAL_ADDRESS(module, importDesc->FirstThunk);

}

for (; *thunkRef; thunkRef++, funcRef++)

{

if IMAGE_SNAP_BY_ORDINAL(*thunkRef)

*funcRef = (DWORD)GetProcAddress(handle, (LPCSTR)IMAGE_ORDINAL(*thunkRef));

else {

PIMAGE_IMPORT_BY_NAME thunkData = (PIMAGE_IMPORT_BY_NAME)CALCULATE_REAL_ADDRESS(module, *thunkRef);

*funcRef = (DWORD)GetProcAddress(handle, (LPCSTR)&thunkData->Name);

}

if (*funcRef == 0)

{

SetLastError(ERROR_PROC_NOT_FOUND);

result = 0;

break;

}

}

if (!result)

break;

}

}

return result;

}

static PPEB

GetPEB(void)

{

// XXX: is there a better or even documented way to do this?

__asm {

// get PEB

mov eax, dword ptr fs:[30h]

}

}

static HMODULE

FindLibraryInPEB(const unsigned char *name, int incLoadCount)

{

PPEB_LDR_DATA loaderData;

PLDR_MODULE loaderModule;

PWSTR longName;

size_t i;

HMODULE result=NULL;

if (name == NULL)

return NULL;

// convert name to long character name

longName = (PWSTR)calloc((strlen(name)+1)*2, 1);

for (i=0; i<strlen(name);i++)

longName[i] = name[i];

// search module with given name in PEB

loaderData = GetPEB()->LoaderData;

loaderModule = (PLDR_MODULE)(loaderData->InLoadOrderModuleList.Flink);

while (1)

{

if (wcsicmp(longName, loaderModule->BaseDllName.Buffer) == 0)

{

result = loaderModule->BaseAddress;

if (incLoadCount && loaderModule->LoadCount != 0xffff)

// we use this module, so increate the load count

loaderModule->LoadCount++;

goto exit;

}

// advance to next module

loaderModule = (PLDR_MODULE)(loaderModule->InLoadOrderModuleList.Flink);

if (loaderModule->BaseAddress == NULL || loaderModule == (PLDR_MODULE)(loaderData->InLoadOrderModuleList.Flink))

// we traversed through the complete list

// and didn't find the library

goto exit;

}

exit:

free(longName);

return result;

}

#define AppendToChain(module, list, chain) { \

(module)->##chain##.Flink = (list)->##chain##.Flink; \

(module)->##chain##.Blink = (list)->##chain##.Blink; \

((PLDR_MODULE)((list)->##chain##.Blink))->##chain##.Flink = &(module)->##chain##; \

(list)->##chain##.Blink = &(module)->##chain##; \

};

static PLDR_MODULE

InsertModuleInPEB(HMODULE module, unsigned char *name, unsigned char *baseName, DWORD locationDelta)

{

PLDR_MODULE loaderModule;

PPEB_LDR_DATA loaderData = GetPEB()->LoaderData;

DWORD entry = GET_NT_HEADER(module)->OptionalHeader.AddressOfEntryPoint;

size_t i;

loaderModule = HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, sizeof(LDR_MODULE));

if (loaderModule == NULL)

return NULL;

loaderModule->BaseAddress = module;

loaderModule->EntryPoint = (PVOID)(entry ? CALCULATE_REAL_ADDRESS(module, entry) : 0);

loaderModule->SizeOfImage = GET_NT_HEADER(module)->OptionalHeader.SizeOfImage;

loaderModule->LoadCount = 1;

loaderModule->BaseDllName.Buffer = HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, (strlen(baseName)+1)*2);

if (loaderModule->BaseDllName.Buffer == NULL)

{

HeapFree(GetProcessHeap(), 0, loaderModule);

return NULL;

}

loaderModule->BaseDllName.Length = (USHORT)strlen(baseName)*2;

loaderModule->BaseDllName.MaximumLength = (USHORT)HeapSize(GetProcessHeap(), 0, loaderModule->BaseDllName.Buffer);

for (i=0; i<strlen(baseName); i++)

loaderModule->BaseDllName.Buffer[i] = baseName[i];

loaderModule->FullDllName.Buffer = HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, (strlen(name)+1)*2);

if (loaderModule->BaseDllName.Buffer == NULL)

{

HeapFree(GetProcessHeap(), 0, loaderModule->BaseDllName.Buffer);

HeapFree(GetProcessHeap(), 0, loaderModule);

return NULL;

}

loaderModule->FullDllName.Length = (USHORT)strlen(name)*2;

loaderModule->FullDllName.MaximumLength = (USHORT)HeapSize(GetProcessHeap(), 0, loaderModule->FullDllName.Buffer);

for (i=0; i<strlen(name); i++)

loaderModule->FullDllName.Buffer[i] = name[i];

// XXX: are these the correct flags?

loaderModule->Flags = IMAGE_DLL | ENTRY_PROCESSED;

if (locationDelta != 0)

loaderModule->Flags |= IMAGE_NOT_AT_BASE;

loaderModule->TimeDateStamp = GET_NT_HEADER(module)->FileHeader.TimeDateStamp;

// XXX: do we need more set the hash table?

//loaderModule->HashTableEntry.Flink = &loaderModule->HashTableEntry;

//loaderModule->HashTableEntry.Blink = &loaderModule->HashTableEntry;

AppendToChain(loaderModule, loaderData, InLoadOrderModuleList);

AppendToChain(loaderModule, loaderData, InInitializationOrderModuleList);

// XXX: insert at the correct position in the chain

AppendToChain(loaderModule, loaderData, InMemoryOrderModuleList);

return loaderModule;

}

HMODULE MemoryLoadLibrary(const void *data, unsigned char *name)

{

HMODULE result;

PIMAGE_DOS_HEADER dos_header;

PIMAGE_NT_HEADERS old_header;

LPVOID headers;

unsigned char *baseName;

unsigned char fullname[MAX_DLL_NAME_LENGTH], tempname[MAX_DLL_NAME_LENGTH];

DWORD locationDelta;

DllEntryProc DllEntry;

BOOL successfull;

DWORD hasFullName;

PLDR_MODULE loaderModule=NULL;

// make sure we have a module name

if (name == NULL || strlen(name) == 0)

{

sprintf(tempname, "memorymodule%d", ModuleCount);

name = (unsigned char *)&tempname;

}

// maybe a module with the given name has been loaded already

hasFullName = GetFullPathName(name, sizeof(fullname), (LPSTR)&fullname, &baseName);

// search for module in PEB

result = FindLibraryInPEB(hasFullName ? baseName : name, 1);

if (result != NULL)

// already loaded this module

goto exit;

if (hasFullName)

// use complete filename as module name

name = (unsigned char *)&fullname;

else

baseName = name;

dos_header = (PIMAGE_DOS_HEADER)data;

if (dos_header->e_magic != IMAGE_DOS_SIGNATURE)

{

#if DEBUG_OUTPUT

OutputDebugString("Not a valid executable file.\n");

goto error;

}

old_header = GET_NT_HEADER(data);

if (old_header->Signature != IMAGE_NT_SIGNATURE)

{

#if DEBUG_OUTPUT

OutputDebugString("No PE header found.\n");

goto error;

}

// reserve memory for image of library

result = (HMODULE)VirtualAlloc((LPVOID)(old_header->OptionalHeader.ImageBase),

old_header->OptionalHeader.SizeOfImage,

MEM_RESERVE,

PAGE_READWRITE);

if (result == NULL)

// try to allocate memory at arbitrary position

result = (HMODULE)VirtualAlloc(NULL,

old_header->OptionalHeader.SizeOfImage,

MEM_RESERVE,

PAGE_READWRITE);

if (result == NULL)

{

SetLastError(ERROR_OUTOFMEMORY);

goto error;

}

// XXX: is it correct to commit the complete memory region at once?

// calling DllEntry raises an exception if we don't...

VirtualAlloc(result,

old_header->OptionalHeader.SizeOfImage,

MEM_COMMIT,

PAGE_READWRITE);

// commit memory for headers

headers = VirtualAlloc(result,

dos_header->e_lfanew + old_header->OptionalHeader.SizeOfHeaders,

MEM_COMMIT,

PAGE_READWRITE);

// copy PE header to code

memcpy(headers, dos_header, dos_header->e_lfanew + old_header->OptionalHeader.SizeOfHeaders);

// update position

GET_NT_HEADER(result)->OptionalHeader.ImageBase = (DWORD)result;

// copy sections from DLL file block to new memory location

CopySections(data, old_header, result);

// adjust base address of imported data

locationDelta = (DWORD)((DWORD)result - old_header->OptionalHeader.ImageBase);

if (locationDelta != 0)

PerformBaseRelocation(result, locationDelta);

// load required dlls and adjust function table of imports

if (!BuildImportTable(result))

goto error;

// mark memory pages depending on section headers and release

// sections that are marked as "discardable"

FinalizeSections(result);

// Add loaded module to PEB

if (!(loaderModule = InsertModuleInPEB(result, name, baseName, locationDelta)))

goto error;

// get entry point of loaded library

if (GET_NT_HEADER(result)->OptionalHeader.AddressOfEntryPoint != 0)

{

// notify library about attaching to process

DllEntry = (DllEntryProc)CALCULATE_REAL_ADDRESS(result, GET_NT_HEADER(result)->OptionalHeader.AddressOfEntryPoint);

successfull = (*DllEntry)(result, DLL_PROCESS_ATTACH, 0);

if (!successfull)

{

SetLastError(ERROR_DLL_INIT_FAILED);

goto error;

}

loaderModule->Flags |= PROCESS_ATTACH_CALLED;

}

ModuleCount++;

goto exit;

error:

// perform some cleanup...

if (loaderModule != NULL)

{

if ((loaderModule->Flags & PROCESS_ATTACH_CALLED) != 0)

(*DllEntry)(result, DLL_PROCESS_DETACH, 0);

// remove from module chains

loaderModule->InInitializationOrderModuleList.Flink->Blink = loaderModule->InInitializationOrderModuleList.Blink;

loaderModule->InInitializationOrderModuleList.Blink->Flink = loaderModule->InInitializationOrderModuleList.Flink;

loaderModule->InLoadOrderModuleList.Flink->Blink = loaderModule->InLoadOrderModuleList.Blink;

loaderModule->InLoadOrderModuleList.Blink->Flink = loaderModule->InLoadOrderModuleList.Flink;

loaderModule->InMemoryOrderModuleList.Flink->Blink = loaderModule->InMemoryOrderModuleList.Blink;

loaderModule->InMemoryOrderModuleList.Blink->Flink = loaderModule->InMemoryOrderModuleList.Flink;

// free memory for PEB structures

HeapFree(GetProcessHeap(), 0, loaderModule->BaseDllName.Buffer);

HeapFree(GetProcessHeap(), 0, loaderModule->FullDllName.Buffer);

HeapFree(GetProcessHeap(), 0, loaderModule);

}

if (result != NULL)

VirtualFree(result, 0, MEM_RELEASE);

result = NULL;

exit:

return result;

}