Full SEH with VEH

seh_via_veh.cpp

// deprecated for manual_exception_handling repo

#pragma once

#include <Windows.h>
#include <cstdint>
#include <memory>
#include <iostream>
#include <vector>
#include <cassert>
#include <ehdata.h>

#include "ntos.h"
#include "files.h"
#include "pe.h"

namespace {
	/* exception handling */
	namespace exceptions {
		uint64_t image_base;
		size_t image_size;

		IMAGE_RUNTIME_FUNCTION_ENTRY* exception_table;
		size_t exception_size;

#define GetUnwindCodeEntry(info, index) \
    ((info)->UnwindCode[index])

#define GetLanguageSpecificDataPtr(info) \
    ((PVOID)&GetUnwindCodeEntry((info),((info)->CountOfCodes + 1) & ~1))

#define GetExceptionHandler(base, info) \
    ((PEXCEPTION_ROUTINE)((base) + *(PULONG)GetLanguageSpecificDataPtr(info)))

#define GetChainedFunctionEntry(base, info) \
    ((PRUNTIME_FUNCTION)((base) + *(PULONG)GetLanguageSpecificDataPtr(info)))

#define GetExceptionDataPtr(info) \
    ((PULONG)((PULONG)GetLanguageSpecificDataPtr(info) + 1))

		bool set;
		LONG NTAPI exception_handler(_In_ PEXCEPTION_POINTERS exception) {
			// verify exception address to our image boundaries
			auto exception_address = exception->ContextRecord->Rip; // exception->ExceptionRecord->ExceptionInformation[2] is NULL, and worthless...
			if (exception_address >= image_base && exception_address <= (image_base + image_size)) {
				// enumerate exception table
				if (exception_table) {
					for (size_t i = 0; i < (exception_size / sizeof(IMAGE_RUNTIME_FUNCTION_ENTRY)); i++) {
						auto entry = exception_table[i];
						auto unwind_info = reinterpret_cast<UNWIND_INFO*>(exceptions::image_base + entry.UnwindData);
						auto frame_count = GetExceptionDataPtr(unwind_info)[0];
							
						// check if exception occured in this runtime entry, as well as if it has a handler, and if it has a frame
						auto start = image_base + entry.BeginAddress, end = image_base + entry.EndAddress;
						if ((exception_address >= start && exception_address <= end) && (unwind_info->Flags & UNW_FLAG_EHANDLER || unwind_info->Flags & UNW_FLAG_UHANDLER) && frame_count > 0) {
							for (size_t i = 0; i < frame_count; i++) { // enum frames
								auto frame_index = 1 + (i * 4);
								auto exception_frame = exceptions::image_base + GetExceptionDataPtr(unwind_info)[frame_index + 3];

								auto start = exceptions::image_base + entry.BeginAddress, end = exceptions::image_base + entry.EndAddress;
								auto frame_start = exceptions::image_base + GetExceptionDataPtr(unwind_info)[frame_index], frame_end = exceptions::image_base + GetExceptionDataPtr(unwind_info)[frame_index + 1];
								if (exception_address >= frame_start && exception_address <= frame_end) { // hop to __except statement/s
									// winapi stuff
									SetLastError(exception->ExceptionRecord->ExceptionCode);

									// hop
									exception->ContextRecord->Rip = exception_frame;
									return EXCEPTION_CONTINUE_EXECUTION;
								}
							}
						}
					}
				}
			}

			return EXCEPTION_CONTINUE_SEARCH;
		}
	};

	bool enable_exceptions(void* dll) { // limits: only supports single target module
		auto dos = reinterpret_cast<IMAGE_DOS_HEADER*>(dll);
		auto nt = reinterpret_cast<IMAGE_NT_HEADERS*>(reinterpret_cast<uint64_t>(dll) + dos->e_lfanew);
		if (dos->e_magic != IMAGE_DOS_SIGNATURE || nt->Signature != IMAGE_NT_SIGNATURE)
			return false; // :(

		exceptions::image_base = NULL;
		exceptions::image_size = NULL;
		exceptions::exception_table = NULL;
		exceptions::exception_size = NULL;

		auto exception_dir = nt->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXCEPTION];
		auto exception_table = reinterpret_cast<IMAGE_RUNTIME_FUNCTION_ENTRY*>(reinterpret_cast<uint64_t>(dll) + exception_dir.VirtualAddress); auto exception_size = exception_dir.Size;
		if (exception_table) {
			exceptions::image_base = reinterpret_cast<uint64_t>(dll);
			exceptions::image_size = nt->OptionalHeader.SizeOfImage;
			exceptions::exception_table = exception_table;
			exceptions::exception_size = exception_size;
				
			// add veh
			if (exceptions::set != true) {
				exceptions::set = true;
				AddVectoredExceptionHandler(0, &exceptions::exception_handler);
			}

			return exceptions::set;
		}

		return false;
	}
}