From 48d46d4cccb866a48993c0db1d9d2664791b5533 Mon Sep 17 00:00:00 2001
From: Sayaka <alice@lolita.technology>
Date: Thu, 14 Nov 2024 21:58:18 +0100
Subject: [PATCH] fix system dll enforcement

---
 LR2ArenaEx/src/overlay/dx9hook.cpp |  4 ++++
 LR2ArenaEx/src/utils/mem.cpp       | 14 +++++++++++++-
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/LR2ArenaEx/src/overlay/dx9hook.cpp b/LR2ArenaEx/src/overlay/dx9hook.cpp
index 91643b4..abb2976 100644
--- a/LR2ArenaEx/src/overlay/dx9hook.cpp
+++ b/LR2ArenaEx/src/overlay/dx9hook.cpp
@@ -159,6 +159,10 @@ int __cdecl hkShowCursor(int enabled) {
 
 void overlay::dx9hook::HookDX9() {
 	char* d3dPointer = mem::ScanModIn(d3dPattern, d3dMask, d3dName, true);
+	if (!d3dPointer) {
+		std::cout << "[!] d3d9.dll not found, aborting..." << std::endl;
+		return;
+	}
 	std::cout << "[i] D3D pointer: " << (int*)d3dPointer << std::endl;
 
 	uintptr_t d3dDeviceAddr = mem::FindDMAAddy((uintptr_t)d3dPointer + 0x4, { 0x0 });
diff --git a/LR2ArenaEx/src/utils/mem.cpp b/LR2ArenaEx/src/utils/mem.cpp
index 63859d1..5aac24b 100644
--- a/LR2ArenaEx/src/utils/mem.cpp
+++ b/LR2ArenaEx/src/utils/mem.cpp
@@ -1,6 +1,7 @@
 #include <vector>
 #include <iostream>
 #include <cstring>
+#include <filesystem>
 
 #include <framework.h>
 #include <ntdll.h>
@@ -226,7 +227,16 @@ LDR_DATA_TABLE_ENTRY* GetLDREntry(std::string name, bool forceSystem)
 
 			if (_stricmp(cName, name.c_str()) == 0)
 			{
-				if (!forceSystem || (forceSystem && std::wstring(mod->FullDllName.Buffer).rfind(system_dir, 0) == 0)) {
+				if (forceSystem) {
+					auto dllPath = std::filesystem::path(mod->FullDllName.Buffer);
+					auto systemPath = std::filesystem::path(system_dir);
+					auto rel = std::filesystem::relative(dllPath, systemPath);
+					if (!rel.empty() && rel.native()[0] != '.') { // If subpath
+						ldr = mod;
+						break;
+					}
+				}
+				else {
 					ldr = mod;
 					break;
 				}
@@ -241,6 +251,8 @@ LDR_DATA_TABLE_ENTRY* GetLDREntry(std::string name, bool forceSystem)
 char* mem::ScanModIn(char* pattern, char* mask, std::string modName, bool forceSystem)
 {
 	LDR_DATA_TABLE_ENTRY* ldr = GetLDREntry(modName, forceSystem);
+	if (!ldr)
+		return nullptr;
 	std::cout << "[i] DLL base for " << modName << ": " << ldr->DllBase << std::endl;
 	std::cout << "[i] Size of image: " << ldr->SizeOfImage << std::endl;
 	char* match = mem::ScanInternal(pattern, mask, (char*)ldr->DllBase, ldr->SizeOfImage);