Fix AOB scan bounds and expose main-image ambiguity detection

.github/workflows/aob-tests.yml

@@ -0,0 +1,17 @@
+name: Native AOB regression tests
+on: [push, pull_request]
+permissions:
+  contents: read
+jobs:
+  test:
+    runs-on: windows-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Compile and run x86 boundary and ambiguity tests
+        shell: pwsh
+        run: |
+          $vswhere = "${env:ProgramFiles(x86)}/Microsoft Visual Studio/Installer/vswhere.exe"
+          $installation = & $vswhere -latest -products '*' -requires Microsoft.VisualStudio.Component.VC.Tools.x86.x64 -property installationPath
+          if (-not $installation) { throw 'x86 MSVC toolchain missing' }
+          cmd /c "`"$installation/VC/Auxiliary/Build/vcvars32.bat`" && tests\run-aob.cmd"
+          if ($LASTEXITCODE -ne 0) { exit $LASTEXITCODE }

AOB.cpp

@@ -4,6 +4,9 @@
 #include <regex>
 #include <string>
 #include <iostream>
+#include <algorithm>
+#include <cstdint>
+#include <stdexcept>
 
 #include "AOB.h"
 
@@ -19,7 +22,9 @@ namespace AOB {
 	}
 	DWORD FindPattern(DWORD dwAddress, DWORD dwLen, BYTE* bMask, char* szMask)
 	{
-		for (DWORD i = 0; i < (dwLen - strlen((char*)szMask)); i++)
+		const size_t length = strlen(szMask);
+		if (length == 0 || length > dwLen) return 0;
+		for (DWORD i = 0; i <= dwLen - length; i++)
 		{
 			if (bCompare((BYTE*)(dwAddress + i), bMask, szMask))
 			{
@@ -42,38 +47,34 @@ namespace AOB {
 		return bytes;
 	}
 
+	static bool ReadableImage(const MEMORY_BASIC_INFORMATION& region)
+	{
+		return region.State == MEM_COMMIT && region.Type == MEM_IMAGE &&
+			(region.Protect & (PAGE_NOACCESS | PAGE_GUARD)) == 0;
+	}
+
+	// Both bounds are inclusive. Merge adjacent readable regions so a pattern
+	// spanning a protection boundary is considered, without reading guarded pages.
 	DWORD Scan(char* content, char* mask, DWORD min, DWORD max)
 	{
-		SYSTEM_INFO si;
-		GetSystemInfo(&si);
-		_MEMORY_BASIC_INFORMATION32 mbi;
-		DWORD address = min;
-		int remainder = 0;
-
-		while (VirtualQuery((LPCVOID) address, ((MEMORY_BASIC_INFORMATION*)&mbi), sizeof(MEMORY_BASIC_INFORMATION)) != 0) {
-			if (mbi.State == MEM_COMMIT) {
-				if ((mbi.Type != MEM_MAPPED) && (mbi.Type != MEM_PRIVATE)) {
-					if ((mbi.Protect & PAGE_NOACCESS) == 0) {
-						// address = 0x401002
-						// mbi.BaseAddress = 0x401000
-						DWORD needle = FindPattern(address, mbi.RegionSize, (BYTE*)content, mask);
-						if (needle == 0) {
-							// address = 0x401000
-							address = mbi.BaseAddress;
-						}
-						else {
-							return needle;
-						}
-					}
+		const uint64_t stop = uint64_t(max) + 1;
+		uint64_t address = min, runStart = min;
+		while (address < stop) {
+			MEMORY_BASIC_INFORMATION mbi = {};
+			const bool queried = VirtualQuery(reinterpret_cast<LPCVOID>(uintptr_t(address)), &mbi, sizeof(mbi)) != 0;
+			const uint64_t end = queried ? uint64_t(reinterpret_cast<uintptr_t>(mbi.BaseAddress)) + mbi.RegionSize : address;
+			if (!queried || end <= address || !ReadableImage(mbi)) {
+				if (address > runStart) {
+					const DWORD found = FindPattern(DWORD(runStart), DWORD(address - runStart), reinterpret_cast<BYTE*>(content), mask);
+					if (found) return found;
 				}
+				if (!queried || end <= address) return 0;
+				runStart = (std::min)(end, stop);
 			}
-
-			// address = 0x59E000
-			address += mbi.RegionSize;
-			if (address > max) {
-				return 0;
-			}
+			address = (std::min)(end, stop);
 		}
+		if (address > runStart)
+			return FindPattern(DWORD(runStart), DWORD(address - runStart), reinterpret_cast<BYTE*>(content), mask);
 		return 0;
 	}
 
@@ -100,7 +101,46 @@ namespace AOB {
 			haystack = sm.suffix();
 		}
 
-		return Scan((char*)(&AOB::HexToBytes(content)[0]), (char*)mask.c_str(), min, max);
+		auto bytes = HexToBytes(content);
+		if (bytes.empty()) return 0;
+		return Scan(bytes.data(), const_cast<char*>(mask.c_str()), min, max);
+	}
+
+	DWORD FindInMainModule(std::string pattern, DWORD& second)
+	{
+		second = 0;
+		const HMODULE module = GetModuleHandleW(nullptr);
+		if (!module) throw std::runtime_error("Cannot locate the main executable");
+		uint64_t address = reinterpret_cast<uintptr_t>(module);
+		std::vector<std::pair<uint64_t, uint64_t>> ranges;
+		while (address <= MAXDWORD) {
+			MEMORY_BASIC_INFORMATION mbi = {};
+			if (!VirtualQuery(reinterpret_cast<LPCVOID>(uintptr_t(address)), &mbi, sizeof(mbi)))
+				throw std::runtime_error("Cannot query the main executable's memory");
+			if (mbi.AllocationBase != module) break;
+			const uint64_t end = uint64_t(reinterpret_cast<uintptr_t>(mbi.BaseAddress)) + mbi.RegionSize;
+			if (end <= address || end > uint64_t(MAXDWORD) + 1)
+				throw std::runtime_error("Invalid main executable memory range");
+			const DWORD executable = PAGE_EXECUTE | PAGE_EXECUTE_READ | PAGE_EXECUTE_READWRITE | PAGE_EXECUTE_WRITECOPY;
+			if (mbi.Protect & executable) {
+				if (!ReadableImage(mbi)) throw std::runtime_error("Main executable code is not accessible");
+				if (!ranges.empty() && ranges.back().second == address) ranges.back().second = end;
+				else ranges.emplace_back(address, end);
+			}
+			address = end;
+		}
+		DWORD first = 0;
+		for (const auto& range : ranges) {
+			const DWORD found = FindInRange(pattern, DWORD(range.first), DWORD(range.second - 1));
+			if (found) {
+				if (first) { second = found; return first; }
+				first = found;
+				if (uint64_t(found) + 1 < range.second)
+					second = FindInRange(pattern, found + 1, DWORD(range.second - 1));
+				if (second) return first;
+			}
+		}
+		return first;
 	}
 
 	// TODO: find all?

AOB.h

@@ -6,5 +6,7 @@ namespace AOB {
 	//Consider: https://github.com/CvX/hadesmem
 	DWORD FindPattern(DWORD dwAddress, DWORD dwLen, BYTE* bMask, char* szMask);
 	DWORD Find(std::string ucp_aob);
+	// First two matches in the main executable code, for ambiguity validation.
+	DWORD FindInMainModule(std::string pattern, DWORD& second);
 	DWORD FindInRange(std::string ucp_aob_spec, DWORD min, DWORD max);
 }

CodeFunctions.cpp

@@ -1115,3 +1115,17 @@ int luaScanForAOB(lua_State* L) {
 
 	return 1;
 }
+
+int luaScanForAOBInMainModule(lua_State* L) {
+	const char* pattern = luaL_checkstring(L, 1);
+	if (strlen(pattern) < 2 || !validateAOBQuery(pattern)) return luaL_error(L, "Invalid AOB format");
+	DWORD first = 0, second = 0;
+	// Leave C++ scopes before luaL_error can longjmp.
+	char failure[256] = {};
+	try { first = AOB::FindInMainModule(pattern, second); }
+	catch (const std::exception& error) { strncpy_s(failure, error.what(), _TRUNCATE); }
+	if (failure[0]) return luaL_error(L, "%s", failure);
+	if (first) lua_pushinteger(L, first); else lua_pushnil(L);
+	if (second) lua_pushinteger(L, second); else lua_pushnil(L);
+	return 2;
+}

CodeFunctions.h

@@ -21,6 +21,7 @@ int luaAllocateRWE(lua_State* L);
 int luaDeallocateRWE(lua_State* L);
 
 int luaScanForAOB(lua_State* L);
+int luaScanForAOBInMainModule(lua_State* L);
 
 int luaDetourCode(lua_State* L);
 

Directory.Build.props

@@ -1,7 +1,7 @@
 <Project>
   <PropertyGroup>
 	  <PackageId>RPS</PackageId>
-	  <Version>1.5.1</Version>
+	  <Version>1.5.3</Version>
 	  <Authors>gynt</Authors>
 	  <Company>gynt</Company>
     <Copyright>2023</Copyright>

README.md

@@ -102,11 +102,18 @@ detourCode(onDetour, 0xABCDEF, 7)
 scanForAOB(searchPattern[, min, max])
     searchPattern    scans the memory for this pattern
     min              address to start searching
-    max              address to stop searching
+    max              inclusive last byte; a match must fit fully inside the range
 
 searchPattern: hexadecimal array with question marks for wildcards.
     example: "FF A1 E? B? ?? 00"
 ```
+`scanForAOBInMainModule(searchPattern)` returns the first two matches (or `nil`
+for either absent result) in committed executable pages of the main process
+image. Use the second result to reject ambiguous instruction signatures. It
+excludes DLLs, heaps and non-executable data, handles overlapping matches and
+protection boundaries, and reports inaccessible executable pages as errors.
+Call during initialization before patching the code being identified.
+
 #### Data functions
 ```
 allocate(size)

RuntimePatchingSystem.cpp

@@ -69,6 +69,7 @@ const struct luaL_Reg RPS_LIB[] = {
 	{"registerString", registerString},
 
 	{"scanForAOB", luaScanForAOB},
+	{"scanForAOBInMainModule", luaScanForAOBInMainModule},
 	{NULL, NULL} /* end of array */
 };
 

RuntimePatchingSystem.nuspec

@@ -2,7 +2,7 @@
 <package xmlns="http://schemas.microsoft.com/packaging/2011/08/nuspec.xsd">
   <metadata>
     <id>RPS</id>
-    <version>1.5.2</version>
+    <version>1.5.3</version>
     <title>Runtime Patching System</title>
     <authors>gynt</authors>
     <readme>build\README.md</readme>

tests/.gitignore

@@ -0,0 +1,2 @@
+/build/
+__pycache__/

tests/aob_test.cpp

@@ -0,0 +1,84 @@
+#include "../framework.h"
+#include <string>
+#include <cstring>
+#include <stdexcept>
+#include <iostream>
+#include "../AOB.h"
+
+#pragma section(".aobdata", read, write)
+#pragma section(".aobexec", read, execute)
+__declspec(allocate(".aobdata")) __declspec(align(4096)) unsigned char data[16384];
+__declspec(allocate(".aobexec")) __declspec(align(4096)) unsigned char code[16384];
+
+static DWORD address(const void* p) { return DWORD(reinterpret_cast<uintptr_t>(p)); }
+static void check(bool ok, const char* message) {
+    if (!ok) throw std::runtime_error(message);
+}
+static void protect(void* p, DWORD flags) {
+    DWORD old;
+    check(VirtualProtect(p, 4096, flags, &old) != 0, "VirtualProtect");
+}
+
+int main() {
+    try {
+        const char* text = "UniqueAobFixture!";
+        const std::string pattern = "55 6E 69 71 75 65 41 6F 62 46 69 78 74 75 72 65 21";
+        const size_t size = strlen(text);
+        memset(data, 0, sizeof(data));
+        memcpy(data + 100, text, size);
+        const DWORD begin = address(data);
+        check(AOB::FindInRange(pattern, begin + 100, begin + 100 + DWORD(size) - 1) == begin + 100, "exact-size range");
+        check(AOB::FindInRange(pattern, begin, begin + 100 + DWORD(size) - 1) == begin + 100, "last legal start");
+        check(!AOB::FindInRange(pattern, begin, begin + 100 + DWORD(size) - 2), "upper bound clips pattern");
+        check(!AOB::FindInRange(pattern, begin + 101, begin + 200), "lower bound clips pattern");
+        check(!AOB::FindInRange(pattern, begin + 200, begin + 199), "inverted range");
+        check(!AOB::FindInRange(pattern, begin + 100, begin + 100), "pattern longer than range");
+        check(AOB::FindInRange("55 ? 69 ? 75 65", begin + 100, begin + 105) == begin + 100, "wildcards");
+
+        // A last-byte match next to an inaccessible page must not over-read it.
+        data[4095] = 0xDA;
+        protect(data + 4096, PAGE_NOACCESS);
+        check(AOB::FindInRange("DA", begin + 4095, begin + 8191) == begin + 4095, "last byte before noaccess");
+        check(!AOB::FindInRange(pattern, begin + 4095, begin + 8191), "interior region start before noaccess");
+        protect(data + 4096, PAGE_READWRITE | PAGE_GUARD);
+        check(!AOB::FindInRange(pattern, begin + 4095, begin + 8191), "guard page skipped");
+        MEMORY_BASIC_INFORMATION info = {};
+        VirtualQuery(data + 4096, &info, sizeof(info));
+        check((info.Protect & PAGE_GUARD) != 0, "guard not consumed");
+        protect(data + 4096, PAGE_READWRITE);
+
+        // Protection changes do not make an otherwise readable match disappear.
+        memcpy(data + 4096 - 8, text, size);
+        protect(data, PAGE_READONLY);
+        check(AOB::FindInRange(pattern, begin + 4000, begin + 4200) == begin + 4096 - 8, "cross-region match");
+        protect(data, PAGE_READWRITE);
+
+        #ifndef AOB_TEST_LEGACY
+        DWORD second = 0;
+        check(!AOB::FindInMainModule(pattern, second) && !second, "data and literals excluded");
+        memcpy(code + 4096 - 8, text, size);
+        protect(code, PAGE_EXECUTE_READ);
+        check(AOB::FindInMainModule(pattern, second) == address(code + 4096 - 8) && !second, "unique executable cross-region match");
+        memcpy(code + 8200, text, size);
+        check(AOB::FindInMainModule(pattern, second) == address(code + 4096 - 8) && second == address(code + 8200), "two executable matches");
+        protect(code + 4096, PAGE_READWRITE);
+        check(AOB::FindInMainModule(pattern, second) == address(code + 8200) && !second, "non-executable gap not searched");
+        protect(code + 4096, PAGE_EXECUTE_READWRITE);
+        protect(code, PAGE_EXECUTE_READWRITE);
+        memset(code, 0, sizeof(code));
+        memset(code + 100, 0xDA, 18);
+        const std::string overlapping = "DA DA DA DA DA DA DA DA DA DA DA DA DA DA DA DA DA";
+        check(AOB::FindInMainModule(overlapping, second) == address(code + 100) && second == address(code + 101), "overlapping duplicates");
+        protect(code, PAGE_EXECUTE_READWRITE | PAGE_GUARD);
+        bool rejected = false;
+        try { AOB::FindInMainModule(pattern, second); } catch (const std::runtime_error&) { rejected = true; }
+        protect(code, PAGE_EXECUTE_READWRITE);
+        check(rejected, "guarded executable code reports failure");
+        #endif
+        std::cout << "AOB native boundary and ambiguity checks passed\n";
+        return 0;
+    } catch (const std::exception& error) {
+        std::cerr << error.what() << '\n';
+        return 1;
+    }
+}

tests/run-aob.cmd

@@ -0,0 +1,9 @@
+@echo off
+rem Run from an x86 Visual Studio developer command prompt. No Lua/game needed.
+setlocal
+cd /d "%~dp0.."
+if not exist tests\build mkdir tests\build
+cl /nologo /EHsc /W4 /std:c++14 /Fe:tests\build\aob.exe /Fo:tests\build\ AOB.cpp tests\aob_test.cpp /link /SECTION:.aobexec,ERW
+if errorlevel 1 exit /b 1
+tests\build\aob.exe
+exit /b %errorlevel%

tests/test-aob.lua

@@ -0,0 +1,27 @@
+local lunatest = require("tests.lunatest.lunatest")
+local rps = require("RPS")
+local test_aob = {}
+
+function test_aob.test_overlapping_main_image_results()
+  local first, second = rps.scanForAOBInMainModule('? ? ?')
+  lunatest.assert_true(type(first) == 'number')
+  lunatest.assert_equal(first + 1, second)
+end
+
+function test_aob.test_inclusive_exact_and_short_ranges()
+  local first = rps.scanForAOBInMainModule('? ? ?')
+  lunatest.assert_equal(first, rps.scanForAOB('? ? ?', first, first + 2))
+  lunatest.assert_nil(rps.scanForAOB('? ? ?', first, first + 1))
+end
+
+function test_aob.test_missing_result_shape()
+  local first, second = rps.scanForAOBInMainModule(string.rep('37 9B E1 46 02 AF DC 85 ', 32))
+  lunatest.assert_nil(first)
+  lunatest.assert_nil(second)
+end
+
+function test_aob.test_invalid_pattern()
+  lunatest.assert_false(pcall(rps.scanForAOBInMainModule, 'ZZ'))
+end
+
+return test_aob