From e75438a0060c4b94fb61fbe3cd7003c58e70de92 Mon Sep 17 00:00:00 2001
From: Joachim Bauch <mail@joachim-bauch.de>
Date: Sat, 7 Dec 2013 18:16:14 +0100
Subject: [PATCH] Execute thread-local-storage callbacks (GitHub #11).

Pull request by @glysbaysb, cleaned up and adjusted to match code style.
---
 MemoryModule.c | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/MemoryModule.c b/MemoryModule.c
index 1eae4e7..e2b1fe8 100644
--- a/MemoryModule.c
+++ b/MemoryModule.c
@@ -184,6 +184,24 @@ FinalizeSections(PMEMORYMODULE module)
 #endif
 }
 
+static void
+ExecuteTLS(PMEMORYMODULE module) 
+{
+    unsigned char *codeBase = module->codeBase;
+    
+    PIMAGE_DATA_DIRECTORY directory = GET_HEADER_DICTIONARY(module, IMAGE_DIRECTORY_ENTRY_TLS);
+    if (directory->VirtualAddress > 0) {
+        PIMAGE_TLS_DIRECTORY tls = (PIMAGE_TLS_DIRECTORY) (codeBase + directory->VirtualAddress);
+        PIMAGE_TLS_CALLBACK* callback = (PIMAGE_TLS_CALLBACK *) tls->AddressOfCallBacks;
+        if (callback) {
+            while (*callback) {
+                (*callback)((LPVOID) codeBase, DLL_PROCESS_ATTACH, NULL);
+                callback++;
+            }
+        }
+    }
+}
+
 static void
 PerformBaseRelocation(PMEMORYMODULE module, SIZE_T delta)
 {
@@ -419,6 +437,9 @@ HMEMORYMODULE MemoryLoadLibraryEx(const void *data,
     // sections that are marked as "discardable"
     FinalizeSections(result);
 
+    // TLS callbacks are executed BEFORE the main loading
+    ExecuteTLS(result);
+
     // get entry point of loaded library
     if (result->headers->OptionalHeader.AddressOfEntryPoint != 0) {
         DllEntry = (DllEntryProc) (code + result->headers->OptionalHeader.AddressOfEntryPoint);