From 018d4aa2c10dac12a6913ad01230bd3320a05f96 Mon Sep 17 00:00:00 2001
From: otavepto <153766569+otavepto@users.noreply.github.com>
Date: Tue, 13 Feb 2024 02:08:14 +0200
Subject: [PATCH] sloppy code is problematic! (at least verify PE signature)

---
 CHANGELOG.md                                     |  6 ++++++
 helpers/pe_helpers.cpp                           | 16 +++++++++++-----
 .../steamclient_loader/win/ColdClientLoader.cpp  |  8 +++-----
 3 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 27728173..3f292fc0 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,3 +1,9 @@
+## 2024/2/13
+
+* cold client loader: validate the PE signature before attempting to detect arch
+
+---
+
 ## 2024/2/10
 
 * a hacky fix for the overlay on directx12, currently very slow when loading images
diff --git a/helpers/pe_helpers.cpp b/helpers/pe_helpers.cpp
index e573abbe..e57e83d6 100644
--- a/helpers/pe_helpers.cpp
+++ b/helpers/pe_helpers.cpp
@@ -24,19 +24,23 @@ PIMAGE_NT_HEADERS pe_helpers::get_nt_header(HMODULE hModule)
 {
     // https://dev.to/wireless90/validating-the-pe-signature-my-av-flagged-me-windows-pe-internals-2m5o/
     PIMAGE_DOS_HEADER dosHeader = (PIMAGE_DOS_HEADER)(char*)hModule;
+    if (dosHeader->e_magic != 0x5A4D) { // "MZ" 
+        return nullptr;
+    }
     LONG newExeHeaderOffset = dosHeader->e_lfanew;
-
     return (PIMAGE_NT_HEADERS)((char*)hModule + newExeHeaderOffset);
 }
 
 PIMAGE_FILE_HEADER pe_helpers::get_file_header(HMODULE hModule)
 {
-    return &get_nt_header(hModule)->FileHeader;
+    auto nt_header = get_nt_header(hModule);
+    return nt_header ? &nt_header->FileHeader : nullptr;
 }
 
 PIMAGE_OPTIONAL_HEADER pe_helpers::get_optional_header(HMODULE hModule)
 {
-    return &get_nt_header(hModule)->OptionalHeader;
+    auto nt_header = get_nt_header(hModule);
+    return nt_header ? &nt_header->OptionalHeader : nullptr;
 }
 
 uint8_t* pe_helpers::search_memory(uint8_t *mem, size_t size, const std::string &search_patt)
@@ -234,12 +238,14 @@ std::string pe_helpers::get_err_string(DWORD code)
 
 bool pe_helpers::is_module_64(HMODULE hModule)
 {
-    return (get_file_header(hModule)->Machine == IMAGE_FILE_MACHINE_AMD64);
+    auto file_header = get_file_header(hModule);
+    return file_header ? (file_header->Machine == IMAGE_FILE_MACHINE_AMD64) : false;
 }
 
 bool pe_helpers::is_module_32(HMODULE hModule)
 {
-    return (get_file_header(hModule)->Machine == IMAGE_FILE_MACHINE_I386);
+    auto file_header = get_file_header(hModule);
+    return file_header ? (file_header->Machine == IMAGE_FILE_MACHINE_I386) : false;
 }
 
 pe_helpers::SectionHeadersResult pe_helpers::get_section_headers(HMODULE hModule)
diff --git a/tools/steamclient_loader/win/ColdClientLoader.cpp b/tools/steamclient_loader/win/ColdClientLoader.cpp
index cffc104b..1849e625 100644
--- a/tools/steamclient_loader/win/ColdClientLoader.cpp
+++ b/tools/steamclient_loader/win/ColdClientLoader.cpp
@@ -58,8 +58,7 @@ std::wstring get_ini_value(LPCWSTR section, LPCWSTR key, LPCWSTR default_val = N
 
 static std::vector<uint8_t> get_pe_header(const std::wstring &filepath)
 {
-    try
-    {
+    try {
         std::ifstream file(filepath, std::ios::binary);
         if (!file.is_open()) throw;
 
@@ -71,9 +70,8 @@ static std::vector<uint8_t> get_pe_header(const std::wstring &filepath)
         file.close();
 
         return data;
-    }
-    catch(const std::exception& e)
-    {
+    } catch(const std::exception& e) {
+        dbg_log::write(std::string("Error reading PE header: ") + e.what());
         return std::vector<uint8_t>();
     }
 }