Add TLS protection

src/loader/loader.cpp

@@ -20,29 +20,35 @@ FARPROC loader::load(const utils::nt::module& module) const
 
 	if (source.get_optional_header()->DataDirectory[IMAGE_DIRECTORY_ENTRY_TLS].Size)
 	{
-		const IMAGE_TLS_DIRECTORY* target_tls = reinterpret_cast<PIMAGE_TLS_DIRECTORY>(module.get_ptr() + module
+		const auto target_tls = reinterpret_cast<PIMAGE_TLS_DIRECTORY>(module.get_ptr() + module
 		                                                                                  .get_optional_header()
 		                                                                                  ->
 		                                                                                  DataDirectory
 			[IMAGE_DIRECTORY_ENTRY_TLS].VirtualAddress);
-		const IMAGE_TLS_DIRECTORY* source_tls = reinterpret_cast<PIMAGE_TLS_DIRECTORY>(module.get_ptr() + source
+		const auto source_tls = reinterpret_cast<PIMAGE_TLS_DIRECTORY>(module.get_ptr() + source
 		                                                                                  .get_optional_header()
 		                                                                                  ->
 		                                                                                  DataDirectory
 			[IMAGE_DIRECTORY_ENTRY_TLS].VirtualAddress);
 
-		*reinterpret_cast<DWORD*>(source_tls->AddressOfIndex) = 0;
+		const auto tls_size = source_tls->EndAddressOfRawData - source_tls->StartAddressOfRawData;
+		const auto tls_index = *reinterpret_cast<DWORD*>(target_tls->AddressOfIndex);
+		*reinterpret_cast<DWORD*>(source_tls->AddressOfIndex) = tls_index;
+
+		if (tls_size > TLS_PAYLOAD_SIZE)
+		{
+			throw std::runtime_error(utils::string::va(
+				"TLS data is of size 0x%X, but we have only reserved 0x%X bytes!", tls_size, TLS_PAYLOAD_SIZE));
+		}
 
 		DWORD old_protect;
 		VirtualProtect(PVOID(target_tls->StartAddressOfRawData),
 		               source_tls->EndAddressOfRawData - source_tls->StartAddressOfRawData, PAGE_READWRITE,
 		               &old_protect);
 
-		const auto tls_base = *reinterpret_cast<LPVOID*>(__readfsdword(0x2C));
+		const auto tls_base = *reinterpret_cast<LPVOID*>(__readfsdword(0x2C) + 4 * tls_index);
-		std::memmove(tls_base, PVOID(source_tls->StartAddressOfRawData),
+		std::memmove(tls_base, PVOID(source_tls->StartAddressOfRawData), tls_size);
-		             source_tls->EndAddressOfRawData - source_tls->StartAddressOfRawData);
+		std::memmove(PVOID(target_tls->StartAddressOfRawData), PVOID(source_tls->StartAddressOfRawData), tls_size);
-		std::memmove(PVOID(target_tls->StartAddressOfRawData), PVOID(source_tls->StartAddressOfRawData),
-		             source_tls->EndAddressOfRawData - source_tls->StartAddressOfRawData);
 	}
 
 	DWORD oldProtect;

src/loader/tls_loader.cpp

@@ -1,99 +0,0 @@
-#include <std_include.hpp>
-#include "tls_loader.hpp"
-#include "utils/nt.hpp"
-
-std::mutex tls_loader::mutex_;
-std::vector<std::unique_ptr<tls_loader::tls_entry>> tls_loader::tls_entries_;
-
-//static thread_local tls_loader::tls_executer $;
-
-void tls_loader::handle(IMAGE_TLS_DIRECTORY* tls_directory)
-{
-	std::lock_guard _(mutex_);
-	tls_entries_.push_back(std::make_unique<tls_entry>(tls_directory));
-}
-
-void tls_loader::execute(const bool is_attaching)
-{
-	std::lock_guard _(mutex_);
-
-	for (auto& entry : tls_entries_)
-	{
-		entry->execute(is_attaching);
-	}
-}
-
-tls_loader::tls_entry::tls_entry(IMAGE_TLS_DIRECTORY* tls_directory) : tls_directory_(tls_directory)
-{
-	this->tls_index_ = 18;//TlsAlloc();
-	*reinterpret_cast<DWORD*>(this->tls_directory_->AddressOfIndex) = this->tls_index_;
-
-	this->execute(true);
-}
-
-tls_loader::tls_entry::~tls_entry()
-{
-	//TlsFree(this->tls_index_);
-	this->allocator_.clear();
-}
-
-void tls_loader::tls_entry::execute(const bool is_attaching)
-{
-	if (is_attaching) this->attach();
-	else this->detach();
-}
-
-void tls_loader::tls_entry::attach()
-{
-	const size_t data_size = this->tls_directory_->EndAddressOfRawData - this->tls_directory_->StartAddressOfRawData;
-	const size_t size = data_size + this->tls_directory_->SizeOfZeroFill;
-	const auto data = this->allocator_.allocate(size);
-
-	std::memcpy(data, PVOID(this->tls_directory_->StartAddressOfRawData), data_size);
-
-	const auto tls_indices = reinterpret_cast<LPVOID*>(__readfsdword(0x2C));
-	tls_indices[this->tls_index_] = data;
-
-	auto callbacks = reinterpret_cast<PIMAGE_TLS_CALLBACK*>(this->tls_directory_->AddressOfCallBacks);
-	if (callbacks)
-	{
-		utils::nt::module self;
-
-		while (*callbacks)
-		{
-			(*callbacks)(self.get_ptr(), DLL_THREAD_ATTACH, nullptr);
-			++callbacks;
-		}
-	}
-}
-
-void tls_loader::tls_entry::detach()
-{
-	auto callbacks = reinterpret_cast<PIMAGE_TLS_CALLBACK*>(this->tls_directory_->AddressOfCallBacks);
-	if (callbacks)
-	{
-		utils::nt::module self;
-
-		while (*callbacks)
-		{
-			(*callbacks)(self.get_ptr(), DLL_THREAD_DETACH, nullptr);
-			++callbacks;
-		}
-	}
-
-	const auto tls_indices = reinterpret_cast<LPVOID*>(__readfsdword(0x2C));
-	const auto data = tls_indices[this->tls_index_];
-
-	//const auto data = TlsGetValue(this->tls_index_);
-	//this->allocator_.free(data);
-}
-
-tls_loader::tls_executer::tls_executer()
-{
-	execute(true);
-}
-
-tls_loader::tls_executer::~tls_executer()
-{
-	execute(false);
-}

src/loader/tls_loader.hpp

@@ -1,38 +0,0 @@
-#pragma once
-#include "utils/memory.hpp"
-
-class tls_loader final
-{
-public:
-	class tls_executer final
-	{
-	public:
-		tls_executer();
-		~tls_executer();
-	};
-
-	static void handle(IMAGE_TLS_DIRECTORY* tls_directory);
-
-private:
-	class tls_entry final
-	{
-	public:
-		tls_entry(IMAGE_TLS_DIRECTORY* tls_directory);
-		~tls_entry();
-
-		void execute(bool is_attaching);
-
-	private:
-		utils::memory::allocator allocator_;
-		IMAGE_TLS_DIRECTORY* tls_directory_;
-		DWORD tls_index_;
-
-		void attach();
-		void detach();
-	};
-
-	static std::mutex mutex_;
-	static std::vector<std::unique_ptr<tls_entry>> tls_entries_;
-
-	static void execute(bool is_attaching);
-};

src/main.cpp

@@ -4,15 +4,39 @@
 #include "loader/module_loader.hpp"
 #include "game/game.hpp"
 #include "loader/binary_loader.hpp"
+#include "utils/string.hpp"
 
 //#define GENERATE_DIFFS
 
+__declspec(thread) char tls_data[TLS_PAYLOAD_SIZE];
+
 void exit_hook(const int code)
 {
 	module_loader::pre_destroy();
 	exit(code);
 }
 
+void verify_tls()
+{
+	utils::nt::module self;
+	const auto self_tls = reinterpret_cast<PIMAGE_TLS_DIRECTORY>(self.get_ptr() + self
+	                                                                              .get_optional_header()
+	                                                                              ->
+	                                                                              DataDirectory
+		[IMAGE_DIRECTORY_ENTRY_TLS].VirtualAddress);
+
+	const auto ref = DWORD(&tls_data);
+	const auto tls_index = *reinterpret_cast<DWORD*>(self_tls->AddressOfIndex);
+	const auto tls_vector = *reinterpret_cast<DWORD*>(__readfsdword(0x2C) + 4 * tls_index);
+	const auto offset = ref - tls_vector;
+
+	if (offset != 0 && offset != 8) // Actually 8 is bad, but I think msvc places custom stuff before
+	{
+		throw std::runtime_error(utils::string::va("TLS payload is at offset 0x%X, but should be at 0!",
+		                                           offset));
+	}
+}
+
 int main()
 {
 	FARPROC entry_point = nullptr;
@@ -24,6 +48,8 @@ int main()
 		return 0;
 #endif
 
+		verify_tls();
+
 		module_loader::post_start();
 
 		launcher launcher;

src/std_include.cpp

@@ -7,7 +7,8 @@
 #pragma comment(linker, "/section:.main,re")
 #pragma comment(linker, "/base:0x400000")
 
-__declspec(thread) char tls_data[0x2000];
+// Moved to main.cpp to enforce early linking
+//__declspec(thread) char tls_data[TLS_PAYLOAD_SIZE];
 
 #pragma bss_seg(".payload")
 char payload_data[BINARY_PAYLOAD_SIZE];

src/std_include.hpp

@@ -1,6 +1,7 @@
 #pragma once
 
 #define BINARY_PAYLOAD_SIZE 0x0A000000
+#define TLS_PAYLOAD_SIZE 0x2000
 
 #pragma warning(push)
 #pragma warning(disable: 4458)
@@ -38,12 +39,10 @@
 #include <patch.h>
 #include <tomcrypt.h>
 
-#pragma warning(pop)
-
-using namespace std::literals;
-
 #pragma comment (lib, "gdiplus.lib")
-
+#pragma warning(pop)
 #pragma warning(disable: 4100)
 
 #include "resource.hpp"
+
+using namespace std::literals;