From 44e9e950fc95d382d7a1038743e54d897da94107 Mon Sep 17 00:00:00 2001
From: Paul Gofman <pgofman@codeweavers.com>
Date: Thu, 23 Apr 2020 14:32:23 +0300
Subject: [PATCH] kernelbase: Maintain FLS storage list in PEB.

---
 dlls/kernel32/tests/fiber.c  | 35 ++++++++++++++++++++
 dlls/kernel32/tests/thread.c |  2 --
 dlls/kernelbase/thread.c     | 63 +++++++++++++++++++++++++++---------
 dlls/ntdll/loader.c          | 46 ++++++++++++++++++++++++++
 4 files changed, 129 insertions(+), 17 deletions(-)

diff --git a/dlls/kernel32/tests/fiber.c b/dlls/kernel32/tests/fiber.c
index 7501165dc87..2bcc3ee60c4 100644
--- a/dlls/kernel32/tests/fiber.c
+++ b/dlls/kernel32/tests/fiber.c
@@ -19,6 +19,7 @@
  */
 
 #include "wine/test.h"
+#include <winternl.h>
 
 static LPVOID (WINAPI *pCreateFiber)(SIZE_T,LPFIBER_START_ROUTINE,LPVOID);
 static LPVOID (WINAPI *pConvertThreadToFiber)(LPVOID);
@@ -171,8 +172,25 @@ static void test_FiberHandling(void)
     if (pIsThreadAFiber) ok(!pIsThreadAFiber(), "IsThreadAFiber reported TRUE\n");
 }
 
+static unsigned int list_size(const LIST_ENTRY *le)
+{
+    unsigned int count = 0;
+    LIST_ENTRY *entry;
+
+    for (entry = le->Flink; entry != le; entry = entry->Flink)
+        ++count;
+
+    return count;
+}
+
+#define FLS_TEST_INDEX_COUNT 2048
+
 static void test_FiberLocalStorage(void)
 {
+    static DWORD fls_indices[FLS_TEST_INDEX_COUNT];
+    TEB *teb = NtCurrentTeb();
+    unsigned int i, count;
+    PEB *peb = teb->Peb;
     DWORD fls, fls_2;
     BOOL ret;
     void* val;
@@ -183,6 +201,23 @@ static void test_FiberLocalStorage(void)
         return;
     }
 
+    ok(!!teb->FlsSlots, "Got NULL teb->FlsSlots.\n");
+    ok(!!peb->FlsCallback, "Got NULL peb->FlsCallback.\n");
+
+    count = list_size(&peb->FlsListHead);
+    ok(count == 1, "Got unexpected FLS list size %u.\n", count);
+
+    for (i = 0; i < FLS_TEST_INDEX_COUNT; ++i)
+    {
+        if ((fls_indices[i] = pFlsAlloc(NULL)) == FLS_OUT_OF_INDEXES)
+            break;
+    }
+    count = i;
+    ok(count <= 127, "Could allocate %u FLS indices.\n", count);
+
+    for (i = 0; i < count; ++i)
+        pFlsFree(fls_indices[i]);
+
     /* Test an unallocated index
      * FlsFree should fail
      * FlsGetValue and FlsSetValue should succeed
diff --git a/dlls/kernel32/tests/thread.c b/dlls/kernel32/tests/thread.c
index b91d279e327..924f19703bf 100644
--- a/dlls/kernel32/tests/thread.c
+++ b/dlls/kernel32/tests/thread.c
@@ -2445,7 +2445,6 @@ START_TEST(thread)
        }
        return;
    }
-
    test_thread_info();
    test_reserved_tls();
    test_CreateRemoteThread();
@@ -2473,6 +2472,5 @@ START_TEST(thread)
    test_thread_fpu_cw();
    test_thread_actctx();
    test_thread_description();
-
    test_threadpool();
 }
diff --git a/dlls/kernelbase/thread.c b/dlls/kernelbase/thread.c
index f20d460d07d..3c77b6a42f4 100644
--- a/dlls/kernelbase/thread.c
+++ b/dlls/kernelbase/thread.c
@@ -38,6 +38,10 @@
 
 WINE_DEFAULT_DEBUG_CHANNEL(thread);
 
+static inline void **fls_addr_from_index(void *fls_slots, DWORD index)
+{
+    return (void **)((BYTE *)fls_slots + sizeof(LIST_ENTRY)) + index;
+}
 
 /***********************************************************************
  * Threads
@@ -915,6 +919,19 @@ static void init_fiber_context( struct fiber_data *fiber )
 #endif
 }
 
+typedef void WINAPI (*FLS_LOCK_FUNCTION)(void *);
+
+static void lock_fls_section(void)
+{
+    if (NtCurrentTeb()->Peb->FlsCallback)
+        ((FLS_LOCK_FUNCTION)NtCurrentTeb()->Peb->FlsCallback[0])(NULL);
+}
+
+static void unlock_fls_section(void)
+{
+    if (NtCurrentTeb()->Peb->FlsCallback)
+        ((FLS_LOCK_FUNCTION)NtCurrentTeb()->Peb->FlsCallback[1])(NULL);
+}
 
 /***********************************************************************
  *           CreateFiber   (kernelbase.@)
@@ -1022,7 +1039,16 @@ void WINAPI DECLSPEC_HOTPATCH DeleteFiber( LPVOID fiber_ptr )
         RtlExitUserThread( 1 );
     }
     RtlFreeUserStack( fiber->stack_allocation );
-    HeapFree( GetProcessHeap(), 0, fiber->fls_slots );
+
+    if (fiber->fls_slots)
+    {
+        lock_fls_section();
+        RemoveEntryList( (LIST_ENTRY *)fiber->fls_slots );
+        unlock_fls_section();
+
+        HeapFree( GetProcessHeap(), 0, fiber->fls_slots );
+    }
+
     HeapFree( GetProcessHeap(), 0, fiber );
 }
 
@@ -1044,6 +1070,19 @@ void WINAPI DECLSPEC_HOTPATCH SwitchToFiber( LPVOID fiber )
     struct fiber_data *new_fiber = fiber;
     struct fiber_data *current_fiber = NtCurrentTeb()->Tib.u.FiberData;
 
+    if (!new_fiber->fls_slots)
+    {
+        TEB *teb = NtCurrentTeb();
+        PEB *peb = teb->Peb;
+
+        new_fiber->fls_slots = RtlAllocateHeap( GetProcessHeap(), HEAP_ZERO_MEMORY,
+                sizeof(LIST_ENTRY) + 8 * sizeof(peb->FlsBitmapBits) * sizeof(void*) );
+
+        lock_fls_section();
+        InsertTailList(&peb->FlsListHead, (LIST_ENTRY *)new_fiber->fls_slots);
+        unlock_fls_section();
+    }
+
     current_fiber->except      = NtCurrentTeb()->Tib.ExceptionList;
     current_fiber->stack_limit = NtCurrentTeb()->Tib.StackLimit;
     current_fiber->fls_slots   = NtCurrentTeb()->FlsSlots;
@@ -1069,9 +1108,7 @@ DWORD WINAPI DECLSPEC_HOTPATCH FlsAlloc( PFLS_CALLBACK_FUNCTION callback )
     PEB * const peb = NtCurrentTeb()->Peb;
 
     RtlAcquirePebLock();
-    if (!peb->FlsCallback &&
-        !(peb->FlsCallback = HeapAlloc( GetProcessHeap(), HEAP_ZERO_MEMORY,
-                                        8 * sizeof(peb->FlsBitmapBits) * sizeof(void*) )))
+    if (!peb->FlsCallback)
     {
         SetLastError( ERROR_NOT_ENOUGH_MEMORY );
         index = FLS_OUT_OF_INDEXES;
@@ -1081,9 +1118,7 @@ DWORD WINAPI DECLSPEC_HOTPATCH FlsAlloc( PFLS_CALLBACK_FUNCTION callback )
         index = RtlFindClearBitsAndSet( peb->FlsBitmap, 1, 1 );
         if (index != ~0U)
         {
-            if (!NtCurrentTeb()->FlsSlots &&
-                !(NtCurrentTeb()->FlsSlots = HeapAlloc( GetProcessHeap(), HEAP_ZERO_MEMORY,
-                                                        8 * sizeof(peb->FlsBitmapBits) * sizeof(void*) )))
+            if (!NtCurrentTeb()->FlsSlots)
             {
                 RtlClearBits( peb->FlsBitmap, index, 1 );
                 index = FLS_OUT_OF_INDEXES;
@@ -1091,8 +1126,8 @@ DWORD WINAPI DECLSPEC_HOTPATCH FlsAlloc( PFLS_CALLBACK_FUNCTION callback )
             }
             else
             {
-                NtCurrentTeb()->FlsSlots[index] = 0; /* clear the value */
-                peb->FlsCallback[index] = callback;
+                *fls_addr_from_index(NtCurrentTeb()->FlsSlots, index) = NULL; /* clear the value */
+                peb->FlsCallback[index + 2] = callback;
             }
         }
         else SetLastError( ERROR_NO_MORE_ITEMS );
@@ -1116,7 +1151,7 @@ BOOL WINAPI DECLSPEC_HOTPATCH FlsFree( DWORD index )
     {
         /* FIXME: call Fls callback */
         /* FIXME: add equivalent of ThreadZeroTlsCell here */
-        if (NtCurrentTeb()->FlsSlots) NtCurrentTeb()->FlsSlots[index] = 0;
+        if (NtCurrentTeb()->FlsSlots) *fls_addr_from_index(NtCurrentTeb()->FlsSlots, index) = NULL;
     }
     else SetLastError( ERROR_INVALID_PARAMETER );
     RtlReleasePebLock();
@@ -1135,7 +1170,7 @@ PVOID WINAPI DECLSPEC_HOTPATCH FlsGetValue( DWORD index )
         return NULL;
     }
     SetLastError( ERROR_SUCCESS );
-    return NtCurrentTeb()->FlsSlots[index];
+    return *fls_addr_from_index(NtCurrentTeb()->FlsSlots, index);
 }
 
 
@@ -1149,14 +1184,12 @@ BOOL WINAPI DECLSPEC_HOTPATCH FlsSetValue( DWORD index, PVOID data )
         SetLastError( ERROR_INVALID_PARAMETER );
         return FALSE;
     }
-    if (!NtCurrentTeb()->FlsSlots &&
-        !(NtCurrentTeb()->FlsSlots = HeapAlloc( GetProcessHeap(), HEAP_ZERO_MEMORY,
-                                        8 * sizeof(NtCurrentTeb()->Peb->FlsBitmapBits) * sizeof(void*) )))
+    if (!NtCurrentTeb()->FlsSlots)
     {
         SetLastError( ERROR_NOT_ENOUGH_MEMORY );
         return FALSE;
     }
-    NtCurrentTeb()->FlsSlots[index] = data;
+    *fls_addr_from_index(NtCurrentTeb()->FlsSlots, index) = data;
     return TRUE;
 }
 
diff --git a/dlls/ntdll/loader.c b/dlls/ntdll/loader.c
index efacc0c081e..e61c6e5abc7 100644
--- a/dlls/ntdll/loader.c
+++ b/dlls/ntdll/loader.c
@@ -167,6 +167,15 @@ static RTL_BITMAP tls_bitmap;
 static RTL_BITMAP tls_expansion_bitmap;
 static RTL_BITMAP fls_bitmap;
 
+static CRITICAL_SECTION fls_section;
+static CRITICAL_SECTION_DEBUG fls_critsect_debug =
+{
+    0, 0, &fls_section,
+    { &fls_critsect_debug.ProcessLocksList, &fls_critsect_debug.ProcessLocksList },
+      0, 0, { (DWORD_PTR)(__FILE__ ": fls_section") }
+};
+static CRITICAL_SECTION fls_section = { &fls_critsect_debug, -1, 0, 0, 0, 0 };;
+
 static WINE_MODREF *cached_modref;
 static WINE_MODREF *current_modref;
 static WINE_MODREF *last_failed_modref;
@@ -207,6 +216,16 @@ static RTL_UNLOAD_EVENT_TRACE unload_traces[RTL_UNLOAD_EVENT_TRACE_NUMBER];
 static RTL_UNLOAD_EVENT_TRACE *unload_trace_ptr;
 static unsigned int unload_trace_seq;
 
+static void WINAPI lock_fls_section( PVOID dummy )
+{
+    RtlEnterCriticalSection( &fls_section );
+}
+
+static void WINAPI unlock_fls_section( PVOID dummy )
+{
+    RtlLeaveCriticalSection( &fls_section );
+}
+
 static void module_push_unload_trace( const LDR_DATA_TABLE_ENTRY *ldr )
 {
     RTL_UNLOAD_EVENT_TRACE *ptr = &unload_traces[unload_trace_seq];
@@ -3201,6 +3220,13 @@ void WINAPI LdrShutdownThread(void)
     /* don't do any detach calls if process is exiting */
     if (process_detaching) return;
 
+    if (NtCurrentTeb()->FlsSlots)
+    {
+        lock_fls_section( NULL );
+        RemoveEntryList( (LIST_ENTRY *)NtCurrentTeb()->FlsSlots );
+        unlock_fls_section( NULL );
+    }
+
     RtlEnterCriticalSection( &loader_section );
     wm = get_modref( NtCurrentTeb()->Peb->ImageBaseAddress );
 
@@ -3414,6 +3440,7 @@ PIMAGE_NT_HEADERS WINAPI RtlImageNtHeader(HMODULE hModule)
  */
 void WINAPI LdrInitializeThunk( CONTEXT *context, void **entry, ULONG_PTR unknown3, ULONG_PTR unknown4 )
 {
+    static const unsigned int fls_slot_count = 8 * sizeof(NtCurrentTeb()->Peb->FlsBitmapBits);
     static int attach_done;
     int i;
     NTSTATUS status;
@@ -3449,6 +3476,25 @@ void WINAPI LdrInitializeThunk( CONTEXT *context, void **entry, ULONG_PTR unknow
     InsertHeadList( &tls_links, &NtCurrentTeb()->TlsLinks );
     RtlReleasePebLock();
 
+    if (!NtCurrentTeb()->Peb->FlsCallback)
+    {
+        NtCurrentTeb()->Peb->FlsCallback = RtlAllocateHeap( GetProcessHeap(), HEAP_ZERO_MEMORY,
+                 (fls_slot_count + 2) * sizeof(void*) );
+        if (!NtCurrentTeb()->Peb->FlsCallback)
+            ERR( "No memory for FLS callbacks.\n" );
+
+        if (NtCurrentTeb()->Peb->FlsCallback)
+        {
+            NtCurrentTeb()->Peb->FlsCallback[0] = lock_fls_section;
+            NtCurrentTeb()->Peb->FlsCallback[1] = unlock_fls_section;
+        }
+    }
+    NtCurrentTeb()->FlsSlots = RtlAllocateHeap( GetProcessHeap(), HEAP_ZERO_MEMORY,
+             sizeof(LIST_ENTRY) + fls_slot_count * sizeof(void*) );
+    lock_fls_section( NULL );
+    InsertTailList(&NtCurrentTeb()->Peb->FlsListHead, (LIST_ENTRY *)NtCurrentTeb()->FlsSlots);
+    unlock_fls_section( NULL );
+
     if (!attach_done)  /* first time around */
     {
         attach_done = 1;
-- 
2.27.0

