Skip to content

Commit

Permalink
Enhancement: AVX2 SIMD pattern searching
Browse files Browse the repository at this point in the history
  • Loading branch information
@Ceiridge
Ceiridge committed Aug 2, 2021
1 parent c4c9ef9 commit f05162f
Show file tree
Hide file tree
Showing 8 changed files with 98 additions and 37 deletions.
4 changes: 2 additions & 2 deletions ChromePatcherDll/dllmain.cpp
View file
Expand Up @@ -31,8 +31,8 @@ BOOL APIENTRY ThreadMain(LPVOID lpModule) {
freopen_s(&ferr, winDirErr, "a", stderr);
}
#endif
std::cout << std::time(0) << std::endl; // Log time for debug purposes
std::cerr << std::time(0) << std::endl;
std::cout << std::time(nullptr) << std::endl; // Log time for debug purposes
std::cerr << std::time(nullptr) << std::endl;

std::wstring mutexStr = std::wstring(L"ChromeDllMutex") + std::to_wstring(GetCurrentProcessId());
HANDLE mutex = OpenMutex(MUTEX_ALL_ACCESS, FALSE, mutexStr.c_str()); // Never allow the dll to be injected twice
Expand Down
29 changes: 14 additions & 15 deletions ChromePatcherDll/patches.cpp
View file
Expand Up @@ -9,12 +9,12 @@ namespace ChromePatch {
std::ostream& operator<<(std::ostream& os, const Patch& patch) { // Write identifiable data to the output stream for debugging
const PatchPattern& firstPattern = patch.patterns[0];

os << "First Pattern: " << std::hex;
os << "(First Pattern: " << std::hex;
for (byte b : firstPattern.pattern) {
os << std::setw(2) << std::setfill('0') << (int)b << " ";
}

os << " with PatchByte " << static_cast<int>(patch.patchByte) << std::dec;
os << "with PatchByte " << static_cast<int>(patch.patchByte) << ")" << std::dec;
return os;
}

Expand Down Expand Up @@ -144,52 +144,51 @@ namespace ChromePatch {
// TODO: Externalize this function in different implementations (traditional and with SIMD support) and add multi-threading
int Patches::ApplyPatches() {
std::unique_ptr<PatternSearcher> patternSearcher;
if(SimdPatternSearcher::IsCpuSupported()) {
const bool simdCpuSupport = SimdPatternSearcher::IsCpuSupported();
std::cout << "SIMD support: " << simdCpuSupport << std::endl;

if(simdCpuSupport) {
patternSearcher = std::make_unique<SimdPatternSearcher>();
} else {
patternSearcher = std::make_unique<SimplePatternSearcher>();
}

int successfulPatches = 0;
std::cout << "Applying patches, please wait..." << std::endl;
HANDLE proc = GetCurrentProcess();
const HANDLE proc = GetCurrentProcess();
MODULEINFO chromeDllInfo;

GetModuleInformation(proc, chromeDll, &chromeDllInfo, sizeof(chromeDllInfo));
MEMORY_BASIC_INFORMATION mbi{};

for (uintptr_t i = (uintptr_t)chromeDll; i < (uintptr_t)chromeDll + (uintptr_t)chromeDllInfo.SizeOfImage; i++) {
if (VirtualQuery((LPCVOID)i, &mbi, sizeof(mbi))) {
if (mbi.Protect & (PAGE_GUARD | PAGE_NOCACHE | PAGE_NOACCESS) || !(mbi.State & MEM_COMMIT)) {
if (mbi.Protect & (PAGE_GUARD | PAGE_NOCACHE | PAGE_NOACCESS) || !(mbi.State & MEM_COMMIT) || !(mbi.Protect & (PAGE_EXECUTE_READ | PAGE_EXECUTE_READWRITE))) {
i += mbi.RegionSize; // Skip these regions
} else {
for (Patch& patch : patches) {
if (patch.finishedPatch) {
continue;
}

// TODO: Check for cpu features and use the respective searcher & multi-threading
byte* searchResult = patternSearcher->SearchBytePattern(patch, static_cast<byte*>(mbi.BaseAddress), mbi.RegionSize);

if(searchResult == nullptr) {
std::cerr << "Pattern not found for patch " << patch << std::endl;
patch.finishedPatch = true;
if(!searchResult) { // is null
continue;
}

int offsetAttempt = 0;
while(!patch.successfulPatch) {
byte* patchAddr = searchResult + patch.offsets[offsetAttempt];
std::cout << "Reading address " << std::hex << patchAddr << std::endl;
std::cout << "Reading address " << std::hex << (uintptr_t)patchAddr << std::endl;


if(patch.isSig) { // Add the offset found at the patchAddr (with a 4 byte rel. addr. offset) to the patchAddr
patchAddr += *reinterpret_cast<int*>(patchAddr) + 4 + patch.sigOffset;
std::cout << "New aftersig address: " << std::hex << patchAddr << std::endl;
std::cout << "New aftersig address: " << std::hex << (uintptr_t)patchAddr << std::endl;
}

if(patch.origByte == 0xFF || *patchAddr == patch.origByte) {
std::cout << "Patching byte " << std::hex << (int)*patchAddr << " to " << (int)patch.patchByte << " at " << patchAddr << std::endl;
std::cout << "Patching byte " << std::hex << (int)*patchAddr << " to " << (int)patch.patchByte << " at " << (uintptr_t)patchAddr << std::endl;
DWORD oldProtect;
VirtualProtect(mbi.BaseAddress, mbi.RegionSize, PAGE_EXECUTE_READWRITE, &oldProtect);

Expand All @@ -206,7 +205,7 @@ namespace ChromePatch {
patch.successfulPatch = true;
} else {
offsetAttempt++;
std::cerr << "Byte (" << std::hex << (int)*patchAddr << ") not original (" << (int)patch.origByte << ") at " << patchAddr << std::endl;
std::cerr << "Byte (" << std::hex << (int)*patchAddr << ") not original (" << (int)patch.origByte << ") at " << (uintptr_t)patchAddr << std::endl;

if(offsetAttempt == patch.offsets.size()) {
break; // Abort trying out offsets if none worked
Expand All @@ -221,7 +220,7 @@ namespace ChromePatch {
}
}
}

for (Patch& patch : patches) {
if (!patch.successfulPatch) {
std::cerr << "Couldn't patch " << patch << std::endl;
Expand Down
6 changes: 3 additions & 3 deletions ChromePatcherDll/patches.hpp
View file
Expand Up @@ -31,9 +31,9 @@ namespace ChromePatch {
ReadPatchResult ReadPatchFile();
int ApplyPatches();
private:
std::wstring MultibyteToWide(const std::string& str);
std::string ReadString(std::ifstream& file);
unsigned int ReadUInteger(std::ifstream& file);
static std::wstring MultibyteToWide(const std::string& str);
static std::string ReadString(std::ifstream& file);
static unsigned int ReadUInteger(std::ifstream& file);
};
inline Patches patches;

Expand Down
64 changes: 60 additions & 4 deletions ChromePatcherDll/simdpatternsearcher.cpp
View file
Expand Up @@ -3,13 +3,69 @@
#include "simdpatternsearcher.hpp"

namespace ChromePatch {
byte* SimdPatternSearcher::SearchBytePattern(Patch& patch, byte* startAddr, size_t length) {
// TODO
constexpr int SIMD_BYTE_COUNT = 32; // 256 bit

// Inspired by http://0x80.pl/articles/simd-strfind.html
byte* SimdPatternSearcher::SearchBytePattern(Patch& patch, byte* startAddr, const size_t length) {
for (PatchPattern& pattern : patch.patterns) {
const size_t patternSize = pattern.pattern.size();
const __m256i firstByte = _mm256_set1_epi8(pattern.pattern[0]); // Set first __m256i to the first byte
const __m256i lastByte = _mm256_set1_epi8(pattern.pattern[patternSize - 1]); // Set first __m256i to the last byte of the pattern

for (size_t i = 0; i < length; i += SIMD_BYTE_COUNT) {
const size_t lastBytesAdd = i + patternSize - 1;
if (lastBytesAdd + 32 > length) { // Prevent access violations
return nullptr;
}

const __m256i firstBytes = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(startAddr + i));
const __m256i lastBytes = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(startAddr + lastBytesAdd));

const __m256i equalFirst = _mm256_cmpeq_epi8(firstByte, firstBytes);
const __m256i equalLast = _mm256_cmpeq_epi8(lastByte, lastBytes);

int equalityMask = _mm256_movemask_epi8(_mm256_and_si256(equalFirst, equalLast)); // AND operation -> save FFs of the __mm256i into mask

while(equalityMask) { // Manually compare if the first and last byte were equal
DWORD bitPos;
_BitScanForward(&bitPos, equalityMask);

pattern.searchOffset = 1; // Skip first byte, because it has already been compared
for(size_t x = i + bitPos + 1; x < i + bitPos + patternSize - 1; x++) { // Also subtract one byte; same reason
const byte patternByte = pattern.pattern[pattern.searchOffset];
if(patternByte == 0xFF || startAddr[x] == patternByte) {
pattern.searchOffset++;
} else {
pattern.searchOffset = 0;
}

if (pattern.searchOffset == patternSize - 1) {
return startAddr + x - pattern.searchOffset + 1;
}
}

equalityMask = equalityMask & (equalityMask - 1); // Sets left-most bit to 0
}
}
}

return nullptr;
}

int SimdPatternSearcher::CountTrailingZeros(const int i) {
int bits = 0, ix = i;

if(ix) { // Check if i isn't 0
while((ix & 1) == 0) {
bits++;
ix >>= 1; // Bit shift to the right
}
}

return bits;
}

bool SimdPatternSearcher::IsCpuSupported() {
// TODO
return IsProcessorFeaturePresent(PF_AVX2_INSTRUCTIONS_AVAILABLE);
return IsProcessorFeaturePresent(PF_AVX_INSTRUCTIONS_AVAILABLE) && IsProcessorFeaturePresent(PF_AVX2_INSTRUCTIONS_AVAILABLE);
}
}
2 changes: 2 additions & 0 deletions ChromePatcherDll/simdpatternsearcher.hpp
View file
Expand Up @@ -9,5 +9,7 @@ namespace ChromePatch {
public:
byte* SearchBytePattern(Patch& patch, byte* startAddr, size_t length) override;
static bool IsCpuSupported();
private:
static int CountTrailingZeros(int i);
};
}
10 changes: 5 additions & 5 deletions ChromePatcherDll/threads.cpp
View file
Expand Up @@ -4,9 +4,9 @@
namespace ChromePatch {
// Partially taken from https://stackoverflow.com/questions/16684245/can-i-suspend-a-process-except-one-thread
void SuspendOtherThreads() {
DWORD pid = GetCurrentProcessId();
HANDLE snap = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
DWORD mine = GetCurrentThreadId();
const DWORD pid = GetCurrentProcessId();
const HANDLE snap = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
const DWORD mine = GetCurrentThreadId();

if (snap != INVALID_HANDLE_VALUE) {
THREADENTRY32 te;
Expand All @@ -17,7 +17,7 @@ namespace ChromePatch {
if (te.dwSize >= FIELD_OFFSET(THREADENTRY32, th32OwnerProcessID) + sizeof(te.th32OwnerProcessID)) {
if (te.th32ThreadID != mine && te.th32OwnerProcessID == pid)
{
HANDLE thread = OpenThread(THREAD_ALL_ACCESS, FALSE, te.th32ThreadID);
const HANDLE thread = OpenThread(THREAD_ALL_ACCESS, FALSE, te.th32ThreadID);
if (thread && thread != INVALID_HANDLE_VALUE) {
SuspendThread(thread);
CloseHandle(thread);
Expand All @@ -34,7 +34,7 @@ namespace ChromePatch {

void ResumeOtherThreads() {
for (DWORD tId : suspendedThreads) {
HANDLE thread = OpenThread(THREAD_ALL_ACCESS, FALSE, tId);
const HANDLE thread = OpenThread(THREAD_ALL_ACCESS, FALSE, tId);
if (thread && thread != INVALID_HANDLE_VALUE) {
ResumeThread(thread);
CloseHandle(thread);
Expand Down
16 changes: 10 additions & 6 deletions ChromePatcherDllUnitTests/patterntests.cpp
View file
Expand Up @@ -31,24 +31,28 @@ namespace ChromePatch {
public:
TEST_METHOD(SimplePatternSearcherTest) {
const TestBytes testBytes = CreateTestBytes();
Assert::AreEqual(testBytes.BytesLength, BYTE_ARRAY_SIZE);
Assert::AreEqual(BYTE_ARRAY_SIZE, testBytes.BytesLength);

auto simpleSearcher = std::make_unique<SimplePatternSearcher>();
std::vector<Patch> createdPatches = CreatePatches();

Assert::AreEqual(simpleSearcher->SearchBytePattern(createdPatches[0], testBytes.Bytes, testBytes.BytesLength), testBytes.Pattern1Ptr);
Assert::AreEqual(simpleSearcher->SearchBytePattern(createdPatches[1], testBytes.Bytes, testBytes.BytesLength), testBytes.Pattern2Ptr);
Assert::AreEqual(testBytes.Pattern1Ptr, simpleSearcher->SearchBytePattern(createdPatches[0], testBytes.Bytes, testBytes.BytesLength));
Assert::AreEqual(testBytes.Pattern2Ptr, simpleSearcher->SearchBytePattern(createdPatches[1], testBytes.Bytes, testBytes.BytesLength));
}

TEST_METHOD(SimdPatternSearcherTest) {
const TestBytes testBytes = CreateTestBytes();
Assert::AreEqual(testBytes.BytesLength, BYTE_ARRAY_SIZE);
Assert::AreEqual(BYTE_ARRAY_SIZE, testBytes.BytesLength);

auto simdSearcher = std::make_unique<SimdPatternSearcher>();
std::vector<Patch> createdPatches = CreatePatches();

Assert::AreEqual(simdSearcher->SearchBytePattern(createdPatches[0], testBytes.Bytes, testBytes.BytesLength), testBytes.Pattern1Ptr);
Assert::AreEqual(simdSearcher->SearchBytePattern(createdPatches[1], testBytes.Bytes, testBytes.BytesLength), testBytes.Pattern2Ptr);
Assert::AreEqual(testBytes.Pattern1Ptr, simdSearcher->SearchBytePattern(createdPatches[0], testBytes.Bytes, testBytes.BytesLength));
Assert::AreEqual(testBytes.Pattern2Ptr, simdSearcher->SearchBytePattern(createdPatches[1], testBytes.Bytes, testBytes.BytesLength));
}

TEST_METHOD(SimdCpuSupport) {
Assert::IsTrue(SimdPatternSearcher::IsCpuSupported());
}

private:
Expand Down
4 changes: 2 additions & 2 deletions patterns.xml
View file
Expand Up @@ -19,14 +19,14 @@
<!-- DevModeBubbleDelegate::ShouldIncludeExtension; ("ProxyOverriddenBubble.UserSelection" 2nd function in the vtable) -->
<BytePattern>56 48 83 EC 20 48 89 D6 48 89 D1 E8 ? ? ? ? 89 C1</BytePattern> <!-- Chr 79-84 -->

<BytePattern>56 48 83 EC ? 48 89 D6 48 89 D1 E8 ? ? ? ? 83 F8 ? 74 ?</BytePattern> <!-- Edg 81-84 -->
<BytePattern>56 48 83 EC ? 48 89 D6 48 89 D1 E8 ? ? ? ? 83 F8 ? 74</BytePattern> <!-- Edg 81-84 -->
</Pattern>

<Pattern name="Remove Debug Warning">
<!-- MaybeAddInfoBar; "MaybeAddInfoBar" in the function -->
<BytePattern>41 57 41 56 41 55 41 54 56 57 53 48 83 ec ? 48 8b 05 ? ? ? ? 48 31 e0 48 89 44 ? ? 80 79 ? ? 0f 85 ? ? ? ? 49 89 cd 48 83 c2 ? 48 8d 05 ? ? ? ? 48 89 d1 48 89 c2 e8 ? ? ? ? 48 8d 78 ? 48 85 c0 48 0f 44 f8 49 8b 55 ? 48 85 d2 74 ? 49 8d 45 ? 48 89 c1</BytePattern> <!-- Chr 90, only working for Chrome -->

<BytePattern>41 57 41 56 56 57 53 48 81 EC ? ? ? ? 48 8B 05 ? ? ? ? 48 31 E0 48 89 84 24 ? ? ? ? 80 79 ? ? 74 ? 48 8B 8C 24 ? ? ? ? 48 31 E1 E8 ? ? ? ? 90 48 81 C4 ? ? ? ? 5B 5F 5E 41 5E 41 5F C3 48 89 CE 48 83 C2 ?</BytePattern> <!-- Edg 84 -->
<BytePattern>41 57 41 56 56 57 53 48 81 EC ? ? ? ? 48 8B 05 ? ? ? ? 48 31 E0 48 89 84 24 ? ? ? ? 80 79 ? ? 74 ? 48 8B 8C 24 ? ? ? ? 48 31 E1 E8 ? ? ? ? 90 48 81 C4 ? ? ? ? 5B 5F 5E 41 5E 41 5F C3 48 89 CE 48 83 C2</BytePattern> <!-- Edg 84 -->
</Pattern>

<Pattern name="Remove Elision">
Expand Down

0 comments on commit f05162f

Please sign in to comment.