Dynamic Libraries with Delayed Function Loading

#include <windows.h>
#include <stdio.h>

#include "delayedfunc.h"

const int FUNC_CODE_SIZE = 0x100; // TODO: Elaborate on this and calculate the size of the function.

FARPROC Delayed_Function_Load(const char * lib_name_sz, const char * proc_name_sz)
{
	FARPROC rtn = NULL;

	FILE *fp = NULL;
	errno_t ec = fopen_s(&fp, lib_name_sz, "rb");
	if (!fp || ec != 0) goto cleanup; // Cannot open DLL file for reading.

	IMAGE_DOS_HEADER dos_hdr;
	size_t read_count = fread(&dos_hdr, sizeof dos_hdr, 1, fp);
	if (read_count != 1) goto cleanup; // Cannot read the MS DOS header.

	if (dos_hdr.e_magic != IMAGE_DOS_SIGNATURE) goto cleanup; // Invalid PE file.
	
	int seek_res = fseek(fp, dos_hdr.e_lfanew, SEEK_SET);
	if (seek_res != 0) goto cleanup; // Cannot find the NT header.

	IMAGE_NT_HEADERS nt_header;
	read_count = fread(&nt_header, sizeof nt_header, 1, fp);
	if (read_count != 1) goto cleanup; // Cannot read the NT header.

	if (nt_header.Signature != IMAGE_NT_SIGNATURE) goto cleanup; // Invalid PE file.

	// Locate the export directory.
	DWORD export_dir_addr = nt_header
						    .OptionalHeader
						    .DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT]
						    .VirtualAddress;

	DWORD func_addr = 0;

	for (WORD k = 0; k < nt_header.FileHeader.NumberOfSections; k++)
	{
		IMAGE_SECTION_HEADER section_hdr;
		read_count = fread(&section_hdr, sizeof section_hdr, 1, fp);
		if (read_count != 1) goto cleanup; // Cannot read the section header.

		if (   section_hdr.VirtualAddress <= export_dir_addr
			&& export_dir_addr < section_hdr.VirtualAddress + section_hdr.Misc.VirtualSize)
		{
			DWORD section_base   = section_hdr.PointerToRawData 
								 - section_hdr.VirtualAddress;

			DWORD export_dir_pos = section_base
								 + export_dir_addr ;

			seek_res = fseek(fp, export_dir_pos, SEEK_SET);
			if (seek_res != 0) goto cleanup; // Cannot find the export directory.

			IMAGE_EXPORT_DIRECTORY export_dir;
   			read_count = fread(&export_dir, sizeof export_dir, 1, fp);
			if (read_count != 1) goto cleanup; // Cannot read the export directory.

			DWORD export_names_pos = section_base
								   + export_dir.AddressOfNames;
			
			for (DWORD n = 0; n < export_dir.NumberOfNames ; n++)
			{
				seek_res = fseek(fp, export_names_pos + n*sizeof DWORD, SEEK_SET);
				if (seek_res != 0) goto cleanup; // Cannot find function names.

				DWORD func_name_pos;
				read_count = fread(&func_name_pos, sizeof DWORD, 1, fp);
				if (read_count != 1) goto cleanup; // Cannot read function name name.

				func_name_pos += section_base;

				seek_res = fseek(fp, func_name_pos, SEEK_SET);
				if (seek_res != 0) goto cleanup; // Cannot find function name.

				char func_name[0x400];
				read_count = fread(func_name, 1, sizeof func_name, fp);
				if (read_count < 1) goto cleanup; // Cannot read function name.

				if (!strcmp(func_name, proc_name_sz))
				{
					DWORD func_addr_pos = section_base
									    + export_dir.AddressOfFunctions
									    + n * sizeof DWORD;

					seek_res = fseek(fp, func_addr_pos, SEEK_SET);
					if (seek_res != 0) goto cleanup; // Cannot find function address.

					read_count = fread(&func_addr, sizeof DWORD, 1, fp);
					if (read_count != 1) goto cleanup; // Cannot read function address.

					break;
				}
			}

			break;
		}
	}

	if (func_addr == 0) goto cleanup;

	seek_res = fseek(fp, dos_hdr.e_lfanew + sizeof nt_header, SEEK_SET);
	if (seek_res != 0) goto cleanup; // Cannot find function address.

	for (WORD k = 0; k < nt_header.FileHeader.NumberOfSections; k++)
	{
		IMAGE_SECTION_HEADER section_hdr;
		read_count = fread(&section_hdr, sizeof section_hdr, 1, fp);
		if (read_count != 1) goto cleanup; // Cannot read the section header.

		if (   section_hdr.VirtualAddress <= func_addr
			&& func_addr < section_hdr.VirtualAddress + section_hdr.Misc.VirtualSize)
		{
			// Found the function section.

			DWORD section_base   = section_hdr.PointerToRawData 
								 - section_hdr.VirtualAddress;

			func_addr += section_base;

			seek_res = fseek(fp, func_addr, SEEK_SET);
			if (seek_res != 0) goto cleanup; // Cannot find function code.

#ifdef _DEBUG
			// In debug mode the compiler used a JMP instruction to forward execution
			// to the real code. We have to follow this instruciton to locate it.
			// In a more realistic scenario one would know how to copy the correct
			// function code and all dependencies.

			char jmp_instruction[5];
			read_count = fread(jmp_instruction, 1, 5, fp);
			if (read_count != 5) goto cleanup;

			DWORD * jmp_offset = reinterpret_cast<DWORD *>(jmp_instruction+1);
			
			func_addr += 5;           // IP advances 5 bytes after the JMP instruction
			func_addr += *jmp_offset; // Near jump

			seek_res = fseek(fp, func_addr, SEEK_SET);
			if (seek_res != 0) goto cleanup;
#endif

			rtn = (FARPROC)::VirtualAlloc(NULL, FUNC_CODE_SIZE, MEM_COMMIT, PAGE_EXECUTE_READWRITE);
			if (!rtn) goto cleanup; // could not allocate memory.

			read_count = fread(rtn, 1, FUNC_CODE_SIZE, fp);
			if (read_count != FUNC_CODE_SIZE) goto cleanup; // Cannot read function code.

			break;
		}

	}

	if (fp) fclose(fp);
	return rtn;

cleanup:
	if (fp)  fclose(fp);
	if (rtn) ::VirtualFree(rtn, FUNC_CODE_SIZE, MEM_RELEASE);
	return NULL;
}

void Delayed_Function_Free(FARPROC & payload_proc)
{
	if (payload_proc)
	{
		BOOL res_b = ::VirtualFree(payload_proc, FUNC_CODE_SIZE, MEM_RELEASE);
		payload_proc = NULL;
	}
}