From a9c128ba7144d1448722f824b69d448207dbd7f3 Mon Sep 17 00:00:00 2001
From: Rainer <r.sagitario@gmx.de>
Date: Sat, 12 Feb 2011 17:18:48 +0100
Subject: [PATCH 1/3] dll improvements

---
 src/core/dll_helper.d    |   52 +++++++++++++++++++++++++++++++++++++--------
 src/core/thread_helper.d |   11 ++-------
 2 files changed, 45 insertions(+), 18 deletions(-)

diff --git a/src/core/dll_helper.d b/src/core/dll_helper.d
index a5d7f40..010e867 100644
--- a/src/core/dll_helper.d
+++ b/src/core/dll_helper.d
@@ -52,8 +52,8 @@ private:
         alias extern(Windows)
         void* fnRtlAllocateHeap(void* HeapHandle, uint Flags, uint Size);
 
-        // find a code sequence and return the (relative) address that follows
-        static void* findCodeReference( void* adr, int len, ref ubyte[] pattern, bool relative )
+        // find a code sequence and return the address after the sequence
+        static void* findCodeSequence( void* adr, int len, ref ubyte[] pattern )
         {
             if( !adr )
                 return null;
@@ -64,14 +64,28 @@ private:
                 if( code[ p .. p + pattern.length ] == pattern[ 0 .. $ ] )
                 {
                     ubyte* padr = code + p + pattern.length;
-                    if( relative )
-                        return padr + 4 + *cast(int*) padr;
-                    return *cast(void**) padr;
+                    return padr;
                 }
             }
             return null;
         }
 
+        // find a code sequence and return the (relative) address that follows
+        static void* findCodeReference( void* adr, int len, ref ubyte[] pattern, bool relative )
+        {
+            if( !adr )
+                return null;
+
+            ubyte* padr = cast(ubyte*) findCodeSequence( adr, len, pattern );
+            if( padr )
+            {
+                if( relative )
+                    return padr + 4 + *cast(int*) padr;
+                return *cast(void**) padr;
+            }
+            return null;
+        }
+
         // crawl through ntdll to find function _LdrpAllocateTls@0 and references
         //  to _LdrpNumberOfTlsEntries, _NtdllBaseTag and _LdrpTlsList
         // LdrInitializeThunk
@@ -87,11 +101,14 @@ private:
 
         static __gshared ubyte[] jmp_LdrpInitialize = [ 0x33, 0xED, 0xE9 ]; // xor ebp,ebp; jmp _LdrpInitialize
         static __gshared ubyte[] jmp__LdrpInitialize = [ 0x5D, 0xE9 ]; // pop ebp; jmp __LdrpInitialize
+        static __gshared ubyte[] jmp__LdrpInitialize_xp64 = [ 0x5D, 0x90, 0x90, 0x90, 0x90, 0x90 ]; // pop ebp; nop; nop; nop; nop; nop;
         static __gshared ubyte[] call_LdrpInitializeThread = [ 0xFF, 0x75, 0x08, 0xE8 ]; // push [ebp+8]; call _LdrpInitializeThread
         static __gshared ubyte[] call_LdrpAllocateTls = [ 0x00, 0x00, 0xE8 ]; // jne 0xc3; call _LdrpAllocateTls
+        static __gshared ubyte[] call_LdrpAllocateTls_svr03 = [ 0x65, 0xfc, 0x00, 0xE8 ]; // and [ebp+fc], 0; call _LdrpAllocateTls
         static __gshared ubyte[] jne_LdrpAllocateTls = [ 0x0f, 0x85 ]; // jne body_LdrpAllocateTls
         static __gshared ubyte[] mov_LdrpNumberOfTlsEntries = [ 0x8B, 0x0D ]; // mov ecx, _LdrpNumberOfTlsEntries
         static __gshared ubyte[] mov_NtdllBaseTag = [ 0x51, 0x8B, 0x0D ]; // push ecx; mov ecx, _NtdllBaseTag
+        static __gshared ubyte[] mov_NtdllBaseTag_srv03 = [ 0x50, 0xA1 ]; // push eax; mov eax, _NtdllBaseTag
         static __gshared ubyte[] mov_LdrpTlsList = [ 0x8B, 0x3D ]; // mov edi, _LdrpTlsList
 
         static LdrpTlsListEntry* addTlsListEntry( void** peb, void* tlsstart, void* tlsend, void* tls_callbacks_a, int* tlsindex )
@@ -105,13 +122,19 @@ private:
             {
                 void* pLdrpInitialize = findCodeReference( fn, 20, jmp_LdrpInitialize, true );
                 void* p_LdrpInitialize = findCodeReference( pLdrpInitialize, 40, jmp__LdrpInitialize, true );
+                if( !p_LdrpInitialize )
+                    p_LdrpInitialize = findCodeSequence( pLdrpInitialize, 40, jmp__LdrpInitialize_xp64 );
                 void* pLdrpInitializeThread = findCodeReference( p_LdrpInitialize, 200, call_LdrpInitializeThread, true );
                 void* pLdrpAllocateTls = findCodeReference( pLdrpInitializeThread, 40, call_LdrpAllocateTls, true );
+                if(!pLdrpAllocateTls)
+                    pLdrpAllocateTls = findCodeReference( pLdrpInitializeThread, 100, call_LdrpAllocateTls_svr03, true );
                 void* pBodyAllocateTls = findCodeReference( pLdrpAllocateTls, 40, jne_LdrpAllocateTls, true );
 
-                int* pLdrpNumberOfTlsEntries = cast(int*) findCodeReference( pBodyAllocateTls, 20, mov_LdrpNumberOfTlsEntries, false );
+                int* pLdrpNumberOfTlsEntries = cast(int*) findCodeReference( pBodyAllocateTls, 60, mov_LdrpNumberOfTlsEntries, false );
                 pNtdllBaseTag = cast(int*) findCodeReference( pBodyAllocateTls, 30, mov_NtdllBaseTag, false );
-                LdrpTlsListEntry* pLdrpTlsList = cast(LdrpTlsListEntry*)findCodeReference( pBodyAllocateTls, 60, mov_LdrpTlsList, false );
+                if(!pNtdllBaseTag)
+                    pNtdllBaseTag = cast(int*) findCodeReference( pBodyAllocateTls, 30, mov_NtdllBaseTag_srv03, false );
+                LdrpTlsListEntry* pLdrpTlsList = cast(LdrpTlsListEntry*)findCodeReference( pBodyAllocateTls, 80, mov_LdrpTlsList, false );
 
                 if( !pLdrpNumberOfTlsEntries || !pNtdllBaseTag || !pLdrpTlsList )
                     return null;
@@ -384,20 +407,29 @@ public:
     static this() { tlsCtorRun = true; }
 
     // to be called from DllMain with reason DLL_THREAD_ATTACH
-    void dll_thread_attach( bool attach_thread = true, bool initTls = true )
+    bool dll_thread_attach( bool attach_thread = true, bool initTls = true )
     {
+        // if the OS has not prepared TLS for us, don't attach to the thread 
+	//  (happened when running under x64 OS)
+        if( !GetTlsDataAddress( GetCurrentThreadId() ) )
+            return false;
         if( attach_thread && !thread_findByAddr( GetCurrentThreadId() ) )
             thread_attachThis();
         if( initTls && !tlsCtorRun ) // avoid duplicate calls
             _moduleTlsCtor();
+        return true;
     }
 
     // to be called from DllMain with reason DLL_THREAD_DETACH
-    void dll_thread_detach( bool detach_thread = true, bool exitTls = true )
+    bool dll_thread_detach( bool detach_thread = true, bool exitTls = true )
     {
+        // if the OS has not prepared TLS for us, we did not attach to the thread
+        if( !GetTlsDataAddress( GetCurrentThreadId() ) )
+             return false;
         if( exitTls )
             _moduleTlsDtor();
         if( detach_thread )
-            thread_detachThis();
+            thread_detachThis(); 
+        return true;
     }
 }
diff --git a/src/core/thread_helper.d b/src/core/thread_helper.d
index c7fe6bb..19c8835 100644
--- a/src/core/thread_helper.d
+++ b/src/core/thread_helper.d
@@ -201,7 +201,7 @@ private:
                 return;
             }
 
-            // temporarily set current TLS data pointer to the data pointer of the referenced thread
+            // temporarily set current TLS array pointer to the array pointer of the referenced thread
             void** curteb = getTEB();
             void** teb    = getTEB( id );
             assert( teb && curteb );
@@ -211,14 +211,9 @@ private:
             if( !curtlsarray || !tlsarray )
                 return;
 
-            void** curtlsdata = cast(void**) curtlsarray[_tls_index];
-            void** tlsdata    = cast(void**) tlsarray[_tls_index];
-            if( !curtlsdata || !tlsdata )
-                return;
-
-            curtlsarray[_tls_index] = tlsdata;
+            curteb[11] = tlsarray;
             fn();
-            curtlsarray[_tls_index] = curtlsdata;
+            curteb[11] = curtlsarray;
         }
     }
 
-- 
1.7.3.1.msysgit.0

