Overhaul Windows TLS emulation, multithreading fixes

examples/src/windows/tls_multithread/CMakeLists.txt

@@ -0,0 +1,25 @@
+cmake_minimum_required(VERSION 3.10)
+project(x8664_tls_multithread C)
+
+set(CMAKE_C_STANDARD 99)
+set(CMAKE_C_STANDARD_REQUIRED ON)
+
+add_executable(x8664_tls_multithread main.c)
+
+if(MINGW)
+    target_compile_options(x8664_tls_multithread PRIVATE
+        -Os
+        -fno-asynchronous-unwind-tables
+        -fno-ident
+        -fno-stack-protector
+    )
+
+    target_link_options(x8664_tls_multithread PRIVATE
+        -nostdlib
+        -eEntryPoint
+        -s
+        -Wl,--gc-sections
+    )
+
+    target_link_libraries(x8664_tls_multithread PRIVATE kernel32 user32)
+endif()

examples/src/windows/tls_multithread/main.c

@@ -0,0 +1,111 @@
+#include <windows.h>
+
+static DWORD g_tlsIndex = TLS_OUT_OF_INDEXES;
+
+static volatile LONG g_hasError = 0;
+
+static DWORD StringLengthA(const char* str) {
+    DWORD len = 0;
+    while(str && str[len] != '\0') {
+        len++;
+    }
+    return len;
+}
+
+static void PrintMessage(const char* message) {
+    HANDLE hStdout = GetStdHandle(STD_OUTPUT_HANDLE);
+    if(hStdout != INVALID_HANDLE_VALUE && hStdout != NULL && message != NULL) {
+        DWORD bytesWritten;
+        WriteFile(hStdout, message, (DWORD)StringLengthA(message), &bytesWritten, NULL);
+    }
+}
+
+DWORD WINAPI ThreadWorker(LPVOID lpParam) {
+    int threadId = (int)(UINT_PTR)lpParam;
+    void* expectedData = (void*)(UINT_PTR)(threadId + 100);
+
+    void* initialData = TlsGetValue(g_tlsIndex);
+    if(GetLastError() != ERROR_SUCCESS || initialData != NULL) {
+        char buf[128];
+        wsprintfA(buf, "[Thread %d] FAIL Initial TlsGetValue failed or not NULL.\n", threadId);
+        PrintMessage(buf);
+        InterlockedExchange(&g_hasError, 1);
+        return 1;
+    }
+
+    if(!TlsSetValue(g_tlsIndex, expectedData)) {
+        char buf[128];
+        wsprintfA(buf, "[Thread %d] FAIL TlsSetValue failed with error: %lu\n", threadId, GetLastError());
+        PrintMessage(buf);
+        InterlockedExchange(&g_hasError, 1);
+        return 1;
+    }
+
+    void* retrievedData = TlsGetValue(g_tlsIndex);
+    if(GetLastError() != ERROR_SUCCESS || retrievedData != expectedData) {
+        char buf[128];
+        wsprintfA(buf, "[Thread %d] FAIL TlsGetValue returned unexpected data.\n", threadId);
+        PrintMessage(buf);
+        InterlockedExchange(&g_hasError, 1);
+        return 1;
+    }
+
+    return 0;
+}
+
+void EntryPoint(void) {
+    int exitCode = 0;
+
+    g_tlsIndex = TlsAlloc();
+    if(g_tlsIndex == TLS_OUT_OF_INDEXES) {
+        char buf[128];
+        wsprintfA(buf, "FAIL TlsAlloc failed with error: %lu\n", GetLastError());
+        PrintMessage(buf);
+        ExitProcess((UINT)-1);
+    }
+
+    #define NUM_THREADS 5
+    HANDLE threads[NUM_THREADS];
+
+    for(DWORD i = 0; i < NUM_THREADS; ++i) {
+        threads[i] = CreateThread(
+            NULL,
+            0,
+            ThreadWorker,
+            (LPVOID)(UINT_PTR)i,
+            0,
+            NULL
+        );
+
+        if(threads[i] == NULL) {
+            char buf[128];
+            wsprintfA(buf, "FAIL CreateThread failed with error: %lu\n", GetLastError());
+            PrintMessage(buf);
+            TlsFree(g_tlsIndex);
+            ExitProcess((UINT)-1);
+        }
+    }
+
+    WaitForMultipleObjects(NUM_THREADS, threads, TRUE, INFINITE);
+
+    for(DWORD i = 0; i < NUM_THREADS; ++i) {
+        CloseHandle(threads[i]);
+    }
+
+    if(!TlsFree(g_tlsIndex)) {
+        char buf[128];
+        wsprintfA(buf, "FAIL TlsFree failed with error: %lu\n", GetLastError());
+        PrintMessage(buf);
+        ExitProcess((UINT)-1);
+    }
+
+    if(InterlockedCompareExchange(&g_hasError, 0, 0) != 0) {
+        PrintMessage("FAIL TLS test completed with errors.\n");
+        exitCode = -1;
+    } else {
+        PrintMessage("SUCCESS All TLS operations succeeded across all threads.\n");
+        exitCode = 0;
+    }
+
+    ExitProcess((UINT)exitCode);
+}

qiling/loader/pe.py

@@ -80,6 +80,7 @@ class Process:
     pe_image_address: int
     stack_address: int
     stack_size: int
+    tls_templates: List[bytes]
 
     dlls: MutableMapping[str, int]
     import_address_table: MutableMapping[str, Mapping]
@@ -399,6 +400,13 @@ def load_dll(self, name: str, is_driver: bool = False) -> int:
             self.ql.log.debug(f'Init imports for {dll_name}')
             self.init_imports(dll, is_driver)
 
+            if not is_driver:
+                self.ql.log.debug(f'Init TLS for {dll_name}')
+                self.init_module_tls(dll)
+
+                self.ql.os.thread_manager.init_tls(self.tls_templates)
+                self.ql.os.thread_manager.update_tls_address()
+
             # calling DllMain is essential for dlls to initialize properly. however
             # DllMain of system libraries may fail due to incomplete or inaccurate
             # mock implementation. due to unicorn limitations, recovering from such
@@ -493,6 +501,35 @@ def set_cmdline(self, name: bytes, address: int, memory: bytearray):
 
         return {"name": name, "address": address}
 
+    def init_module_tls(self, module: pefile.PE):
+        if hasattr(module, "DIRECTORY_ENTRY_TLS"):
+            # Not every module has a TLS directory entry,
+            # so we only bother with allocating a TLS block
+            # for the module if the module actually uses TLS.
+
+            tls_struct = module.DIRECTORY_ENTRY_TLS.struct
+
+            tls_data_block_size = tls_struct.EndAddressOfRawData - tls_struct.StartAddressOfRawData
+
+            # The Windows loader places a zero-initialized block immediately
+            # after the data block, with size specified by SizeOfZeroFill.
+            # So the total size of the TLS block is the size of the data block
+            # plus the size of the zero-initialized block.
+            tls_zero_block_size = tls_struct.SizeOfZeroFill
+
+            tls_data_block = bytes(self.ql.mem.read(tls_struct.StartAddressOfRawData, tls_data_block_size))
+            tls_zero_block = bytes([0] for _ in range(tls_zero_block_size))
+            tls_block = tls_data_block + tls_zero_block
+
+            module_tls_index = len(self.tls_templates)
+            self.tls_templates.append(tls_block)
+
+            tls_index_bytes = module_tls_index.to_bytes(4, byteorder='little')
+
+            # The module expects to receive its index in the TLS vector
+            # at the location specified by AddressOfIndex
+            self.ql.mem.write(tls_struct.AddressOfIndex, tls_index_bytes)
+
     def init_teb(self):
         teb_struct = make_teb(self.ql.arch.bits)
 
@@ -825,6 +862,7 @@ def __init__(self, ql: Qiling, libcache: bool):
         self.ql       = ql
         self.path     = self.ql.path
         self.libcache = QlPeCache() if libcache else None
+        self.tls_templates = []
 
     def run(self):
         self.init_dlls = (
@@ -927,10 +965,13 @@ def load(self, pe: Optional[pefile.PE]):
                 self.init_peb()
                 self.init_ldr_data()
                 self.init_exports(pe)
+                self.init_module_tls(pe)
 
                 # add image to ldr table
                 self.add_ldr_data_table_entry(image_name)
 
+                self.ql.os.thread_manager.init_tls(self.tls_templates)
+
             self.init_ki_user_shared_data()
 
             pe.parse_data_directories()
@@ -1032,6 +1073,10 @@ def load(self, pe: Optional[pefile.PE]):
             for each in self.init_dlls:
                 super().load_dll(each, self.is_driver)
 
+        # set the TLS address in the TEB to the main thread's
+        # TLS vector
+        self.ql.os.thread_manager.update_tls_address()
+
         # move entry_point to ql.os
         self.ql.os.entry_point = self.entry_point
         self.init_sp = self.ql.arch.regs.arch_sp

qiling/os/windows/dlls/kernel32/consoleapi.py

@@ -5,6 +5,7 @@
 
 from qiling import Qiling
 from qiling.os.windows.api import *
+from qiling.os.windows.const import *
 from qiling.os.windows.fncc import *
 
 # BOOL WINAPI GetConsoleMode(
@@ -17,3 +18,62 @@
 })
 def hook_GetConsoleMode(ql: Qiling, address: int, params):
     return 1
+
+def _WriteConsole(ql: Qiling, address: int, params, unicode: bool) -> int:
+    hConsoleOutput = params["hConsoleOutput"]
+    lpBuffer = params["lpBuffer"]
+    nNumberOfCharsToWrite = params["nNumberOfCharsToWrite"]
+    lpNumberOfCharsWritten = params["lpNumberOfCharsWritten"]
+
+    handle = ql.os.handle_manager.get(hConsoleOutput)
+
+    if handle is None:
+        ql.os.last_error = ERROR_INVALID_HANDLE
+        return 0
+
+    character_size = 2 if unicode else 1
+    encoding = 'utf-16le' if unicode else 'utf-8'
+
+    text = ql.mem.read(lpBuffer, nNumberOfCharsToWrite * character_size).decode(encoding)
+
+    if hConsoleOutput in (STD_OUTPUT_HANDLE, STD_ERROR_HANDLE):
+        ql.os.stats.log_string(text)
+
+    written = handle.obj.write(bytes(text, encoding='utf-8'))
+    ql.mem.write_ptr(lpNumberOfCharsWritten, written, 4)
+
+    return 1
+
+# BOOL WriteConsoleW(
+#   _In_       HANDLE  hConsoleOutput,
+#   _In_       const VOID    *lpBuffer,
+#   _In_       DWORD   nNumberOfCharsToWrite,
+#   _Out_opt_  LPDWORD lpNumberOfCharsWritten,
+#   _Reserved_ LPVOID  lpReserved
+# );
+@winsdkapi(cc=STDCALL, params={
+    'hConsoleOutput'          : HANDLE,
+    'lpBuffer'                : LPCVOID,
+    'nNumberOfCharsToWrite'   : DWORD,
+    'lpNumberOfCharsWritten' : LPDWORD,
+    'lpReserved'              : LPVOID
+})
+def hook_WriteConsoleW(ql: Qiling, address: int, params):
+    return _WriteConsole(ql, address, params, True)
+
+# BOOL WriteConsoleA(
+#   _In_       HANDLE  hConsoleOutput,
+#   _In_       const VOID    *lpBuffer,
+#   _In_       DWORD   nNumberOfCharsToWrite,
+#   _Out_opt_  LPDWORD lpNumberOfCharsWritten,
+#   _Reserved_ LPVOID  lpReserved
+# );
+@winsdkapi(cc=STDCALL, params={
+    'hConsoleOutput'          : HANDLE,
+    'lpBuffer'                : LPCVOID,
+    'nNumberOfCharsToWrite'   : DWORD,
+    'lpNumberOfCharsWritten' : LPDWORD,
+    'lpReserved'              : LPVOID
+})
+def hook_WriteConsoleA(ql: Qiling, address: int, params):
+    return _WriteConsole(ql, address, params, False)

qiling/os/windows/dlls/kernel32/errhandlingapi.py

@@ -42,6 +42,15 @@ def hook_GetLastError(ql: Qiling, address: int, params):
     'dwErrCode' : DWORD
 })
 def hook_SetLastError(ql: Qiling, address: int, params):
+    # TODO: Make this compatible with the new TLS emulation.
+    # Under the hood, the Windows implementation uses TLS to store last error
+    # codes.
+    # 
+    # Alternatively, these hooks could be removed as the native implementations
+    # appear to work with the new TLS emulation, however any Qiling hooks that
+    # rely on setting ql.os.last_error would have to be rewritten to access the
+    # last error stored in TLS instead.
+
     ql.os.last_error = params['dwErrCode']
 
 # UINT SetErrorMode(