attempt to fix #82, add basic fls support.

3 files changed, 46 insertions, 5 deletions

diff --git a/peloader/pe_linker.c b/peloader/pe_linker.c
index aaca663..7357b17 100644
--- a/peloader/pe_linker.c
+++ b/peloader/pe_linker.c
@@ -92,6 +92,7 @@ static const char *image_directory_name[] = {
 extern struct wrap_export crt_exports[];
 
 uintptr_t LocalStorage[1024] = {0};
+PFLS_CALLBACK_FUNCTION FlsCallbacks[1024] = {0};
 
 static ULONG TlsBitmapData[32];
 static RTL_BITMAP TlsBitmap = {
diff --git a/peloader/winapi/TlsAlloc.c b/peloader/winapi/TlsAlloc.c
index e9ae9e4..17e506a 100644
--- a/peloader/winapi/TlsAlloc.c
+++ b/peloader/winapi/TlsAlloc.c
@@ -19,6 +19,7 @@
 // Index zero is reserved for .tls
 static int TlsIndex = 1;
 extern uintptr_t LocalStorage[1024];
+extern PFLS_CALLBACK_FUNCTION FlsCallbacks[1024];
 
 STATIC DWORD WINAPI TlsAlloc(void)
 {
@@ -62,13 +63,50 @@ STATIC BOOL WINAPI TlsFree(DWORD dwTlsIndex)
     return FALSE;
 }
 
+static DWORD WINAPI FlsAlloc(PVOID lpCallback)
+{
+    DWORD Result;
+    DebugLog("%p", lpCallback);
+
+    // The primary API difference is the availability of callbacks for fibers.
+    if ((Result = TlsAlloc()) != TLS_OUT_OF_INDEXES) {
+        FlsCallbacks[Result] = lpCallback;
+    }
+
+    return Result;
+}
+
+static DWORD WINAPI FlsSetValue(DWORD dwFlsIndex, PVOID lpFlsData)
+{
+    DebugLog("%#x, %p", dwFlsIndex, lpFlsData);
+
+    return TlsSetValue(dwFlsIndex, lpFlsData);
+}
+
+static DWORD WINAPI FlsGetValue(DWORD dwFlsIndex)
+{
+    DebugLog("%#x", dwFlsIndex);
+
+    return TlsGetValue(dwFlsIndex);
+}
+
+static BOOL WINAPI FlsFree(DWORD dwFlsIndex)
+{
+    DebugLog("%#x", dwFlsIndex);
+
+    if (FlsCallbacks[dwFlsIndex]) {
+        FlsCallbacks[dwFlsIndex]((PVOID)(TlsGetValue(dwFlsIndex)));
+    }
+
+    return TlsFree(dwFlsIndex);
+}
+
 DECLARE_CRT_EXPORT("TlsFree", TlsFree);
 DECLARE_CRT_EXPORT("TlsAlloc", TlsAlloc);
 DECLARE_CRT_EXPORT("TlsSetValue", TlsSetValue);
 DECLARE_CRT_EXPORT("TlsGetValue", TlsGetValue);
 
-// These deliberately don't resolve, mpengine redirects to Tls variants.
-//DECLARE_CRT_EXPORT("FlsFree", NULL);
-//DECLARE_CRT_EXPORT("FlsAlloc", NULL);
-//DECLARE_CRT_EXPORT("FlsSetValue", NULL);
-//DECLARE_CRT_EXPORT("FlsGetValue", NULL);
+DECLARE_CRT_EXPORT("FlsFree", FlsFree);
+DECLARE_CRT_EXPORT("FlsAlloc", FlsAlloc);
+DECLARE_CRT_EXPORT("FlsSetValue", FlsSetValue);
+DECLARE_CRT_EXPORT("FlsGetValue", FlsGetValue);
diff --git a/peloader/winnt_types.h b/peloader/winnt_types.h
index cc80c08..bf197fd 100644
--- a/peloader/winnt_types.h
+++ b/peloader/winnt_types.h
@@ -1325,6 +1325,8 @@ struct object_attributes {
         void *security_qos;
 };
 
+typedef void (*PFLS_CALLBACK_FUNCTION)(PVOID lpFlsData) wstdcall;
+
 typedef void (*PCALLBACK_FUNCTION)(void *context, void *arg1,
                                    void *arg2) wstdcall;