Index: trunk/src/kWorker/kWorker.c
===================================================================
--- trunk/src/kWorker/kWorker.c	(revision 2901)
+++ trunk/src/kWorker/kWorker.c	(revision 2902)
@@ -443,4 +443,16 @@
 
 
+/** Pointer to a heap (HeapCreate) tracker entry. */
+typedef struct KWHEAP *PKWHEAP;
+/**
+ * Tracking an heap (HeapCreate)
+ */
+typedef struct KWHEAP
+{
+    PKWHEAP             pNext;
+    HANDLE              hHeap;
+} KWHEAP;
+
+
 /** Pointer to a FlsAlloc/TlsAlloc tracker entry. */
 typedef struct KWLOCALSTORAGE *PKWLOCALSTORAGE;
@@ -453,4 +465,18 @@
     KU32                idx;
 } KWLOCALSTORAGE;
+
+
+/** Pointer to an at exit callback record */
+typedef struct KWEXITCALLACK *PKWEXITCALLACK;
+/**
+ * At exit callback record.
+ */
+typedef struct KWEXITCALLACK
+{
+    PKWEXITCALLACK      pNext;
+    _onexit_t           pfnCallback;
+    /** At exit doesn't have an exit code. */
+    KBOOL               fAtExit;
+} KWEXITCALLACK;
 
 
@@ -469,4 +495,5 @@
     KWTOOLHINT_NONE,
     KWTOOLHINT_VISUAL_CPP_CL,
+    KWTOOLHINT_VISUAL_CPP_LINK,
     KWTOOLHINT_END
 } KWTOOLHINT;
@@ -577,8 +604,15 @@
     /** Head of the virtual alloc allocations. */
     PKWVIRTALLOC    pVirtualAllocHead;
+    /** Head of the heap list (HeapCreate).
+     * This is only done from images we forcibly restore.  */
+    PKWHEAP         pHeapHead;
     /** Head of the FlsAlloc indexes. */
     PKWLOCALSTORAGE pFlsAllocHead;
     /** Head of the TlsAlloc indexes. */
     PKWLOCALSTORAGE pTlsAllocHead;
+
+    /** The at exit callback head.
+     * This is only done from images we forcibly restore.  */
+    PKWEXITCALLACK  pExitCallbackHead;
 
     UNICODE_STRING  SavedCommandLine;
@@ -627,4 +661,7 @@
     /** The replacement function or data address. */
     KUPTR       pfnReplacement;
+    /** Only replace in the executable.
+     * @todo fix the reinitialization of non-native DLLs!  */
+    KBOOL       fOnlyExe;
 } KWREPLACEMENTFUNCTION;
 typedef KWREPLACEMENTFUNCTION const *PCKWREPLACEMENTFUNCTION;
@@ -1759,6 +1796,10 @@
                     || kHlpStrICompAscii(g_aSandboxReplacements[i].pszModule, &pImpMod->pszPath[pImpMod->offFilename]) == 0)
                 {
-                    KW_LOG(("replacing %s!%s\n", &pImpMod->pszPath[pImpMod->offFilename], g_aSandboxReplacements[i].pszFunction));
-                    *puValue = g_aSandboxReplacements[i].pfnReplacement;
+                    if (   pCurMod->fExe
+                        || !g_aSandboxReplacements[i].fOnlyExe)
+                    {
+                        KW_LOG(("replacing %s!%s\n",&pImpMod->pszPath[pImpMod->offFilename], g_aSandboxReplacements[i].pszFunction));
+                        *puValue = g_aSandboxReplacements[i].pfnReplacement;
+                    }
                     break;
                 }
@@ -2292,4 +2333,6 @@
                 if (kHlpStrICompAscii(pToolFsObj->pszName, "cl.exe") == 0)
                     pTool->u.Sandboxed.enmHint = KWTOOLHINT_VISUAL_CPP_CL;
+                else if (kHlpStrICompAscii(pToolFsObj->pszName, "link.exe") == 0)
+                    pTool->u.Sandboxed.enmHint = KWTOOLHINT_VISUAL_CPP_LINK;
                 else
                     pTool->u.Sandboxed.enmHint = KWTOOLHINT_NONE;
@@ -2592,4 +2635,46 @@
 
 
+/** CRT - _onexit   */
+static _onexit_t __cdecl kwSandbox_msvcrt__onexit(_onexit_t pfnFunc)
+{
+    PKWEXITCALLACK pCallback;
+    KW_LOG(("_onexit(%p)\n", pfnFunc));
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
+
+    pCallback = kHlpAlloc(sizeof(*pCallback));
+    if (pCallback)
+    {
+        pCallback->pfnCallback = pfnFunc;
+        pCallback->fAtExit     = K_FALSE;
+        pCallback->pNext       = g_Sandbox.pExitCallbackHead;
+        g_Sandbox.pExitCallbackHead = pCallback;
+        return pfnFunc;
+    }
+
+    return NULL;
+}
+
+
+/** CRT - atexit   */
+static int __cdecl kwSandbox_msvcrt_atexit(int (__cdecl *pfnFunc)(void))
+{
+    PKWEXITCALLACK pCallback;
+    KW_LOG(("_onexit(%p)\n", pfnFunc));
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
+
+    pCallback = kHlpAlloc(sizeof(*pCallback));
+    if (pCallback)
+    {
+        pCallback->pfnCallback = (_onexit_t)pfnFunc;
+        pCallback->fAtExit     = K_TRUE;
+        pCallback->pNext       = g_Sandbox.pExitCallbackHead;
+        g_Sandbox.pExitCallbackHead = pCallback;
+        return 0;
+    }
+
+    return -1;
+}
+
+
 /** Kernel32 - SetConsoleCtrlHandler(). */
 static BOOL WINAPI kwSandbox_Kernel32_SetConsoleCtrlHandler(PHANDLER_ROUTINE pfnHandler, BOOL fAdd)
@@ -2604,4 +2689,5 @@
                                                   int dowildcard, int const *piNewMode)
 {
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     *pargc = g_Sandbox.cArgs;
     *pargv = g_Sandbox.papszArgs;
@@ -2617,4 +2703,5 @@
                                                    int dowildcard, int const *piNewMode)
 {
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     *pargc = g_Sandbox.cArgs;
     *pargv = g_Sandbox.papwszArgs;
@@ -2630,4 +2717,5 @@
 static LPCSTR /*LPSTR*/ WINAPI kwSandbox_Kernel32_GetCommandLineA(VOID)
 {
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     return g_Sandbox.pszCmdLine;
 }
@@ -2637,4 +2725,5 @@
 static LPCWSTR /*LPWSTR*/ WINAPI kwSandbox_Kernel32_GetCommandLineW(VOID)
 {
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     return g_Sandbox.pwszCmdLine;
 }
@@ -2668,4 +2757,5 @@
 static int * __cdecl kwSandbox_msvcrt___p___argc(void)
 {
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     return &g_Sandbox.cArgs;
 }
@@ -2675,4 +2765,5 @@
 static char *** __cdecl kwSandbox_msvcrt___p___argv(void)
 {
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     return &g_Sandbox.papszArgs;
 }
@@ -2682,4 +2773,5 @@
 static wchar_t *** __cdecl kwSandbox_msvcrt___p___wargv(void)
 {
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     return &g_Sandbox.papwszArgs;
 }
@@ -2689,4 +2781,5 @@
 static char ** __cdecl kwSandbox_msvcrt___p__acmdln(void)
 {
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     return (char **)&g_Sandbox.pszCmdLine;
 }
@@ -2696,4 +2789,5 @@
 static wchar_t ** __cdecl kwSandbox_msvcrt___p__wcmdln(void)
 {
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     return &g_Sandbox.pwszCmdLine;
 }
@@ -2703,4 +2797,5 @@
 static char ** __cdecl kwSandbox_msvcrt___p__pgmptr(void)
 {
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     return &g_Sandbox.pgmptr;
 }
@@ -2710,4 +2805,5 @@
 static wchar_t ** __cdecl kwSandbox_msvcrt___p__wpgmptr(void)
 {
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     return &g_Sandbox.wpgmptr;
 }
@@ -2717,4 +2813,5 @@
 static errno_t __cdecl kwSandbox_msvcrt__get_pgmptr(char **ppszValue)
 {
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     *ppszValue = g_Sandbox.pgmptr;
     return 0;
@@ -2725,4 +2822,5 @@
 static errno_t __cdecl kwSandbox_msvcrt__get_wpgmptr(wchar_t **ppwszValue)
 {
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     *ppwszValue = g_Sandbox.wpgmptr;
     return 0;
@@ -2732,4 +2830,5 @@
 static void kwSandbox_msvcrt__wincmdln(void)
 {
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     KWFS_TODO();
 }
@@ -2739,4 +2838,5 @@
 static void kwSandbox_msvcrt__wwincmdln(void)
 {
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     KWFS_TODO();
 }
@@ -2747,6 +2847,17 @@
                                                      DWORD fFlags, PDWORD pidThread)
 {
-    KWFS_TODO();
-    return NULL;
+    HANDLE hThread = NULL;
+    KW_LOG(("CreateThread: pSecAttr=%p (inh=%d) cbStack=%#x pfnThreadProc=%p pvUser=%p fFlags=%#x pidThread=%p\n",
+            pSecAttr, pSecAttr ? pSecAttr->bInheritHandle : 0, cbStack, pfnThreadProc, pvUser, fFlags, pidThread));
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
+    if (g_Sandbox.pTool->u.Sandboxed.enmHint == KWTOOLHINT_VISUAL_CPP_LINK)
+    {
+        /* Allow link::DbgThread. */
+        hThread = CreateThread(pSecAttr, cbStack, pfnThreadProc, pvUser, fFlags, pidThread);
+        KW_LOG(("CreateThread -> %p, *pidThread=%#x\n", hThread, pidThread ? *pidThread : 0));
+    }
+    else
+        KWFS_TODO();
+    return hThread;
 }
 
@@ -2755,4 +2866,5 @@
 static uintptr_t __cdecl kwSandbox_msvcrt__beginthread(void (__cdecl *pfnThreadProc)(void *), unsigned cbStack, void *pvUser)
 {
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     KWFS_TODO();
     return 0;
@@ -2765,4 +2877,5 @@
                                                          unsigned fCreate, unsigned *pidThread)
 {
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     KWFS_TODO();
     return 0;
@@ -2782,9 +2895,11 @@
 {
     char *pszzEnv;
-
-    /* Figure how space much we need first.  */
     char *pszCur;
     KSIZE cbNeeded = 1;
     KSIZE iVar = 0;
+
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
+
+    /* Figure how space much we need first.  */
     while ((pszCur = g_Sandbox.papszEnvVars[iVar++]) != NULL)
         cbNeeded += kHlpStrLen(pszCur) + 1;
@@ -2837,9 +2952,11 @@
 {
     wchar_t *pwszzEnv;
-
-    /* Figure how space much we need first.  */
     wchar_t *pwszCur;
     KSIZE    cwcNeeded = 1;
     KSIZE    iVar = 0;
+
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
+
+    /* Figure how space much we need first.  */
     while ((pwszCur = g_Sandbox.papwszEnvVars[iVar++]) != NULL)
         cwcNeeded += kwUtf16Len(pwszCur) + 1;
@@ -2870,4 +2987,5 @@
 {
     KW_LOG(("FreeEnvironmentStringsA: %p -> TRUE\n", pszzEnv));
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     kHlpFree(pszzEnv);
     return TRUE;
@@ -2879,4 +2997,5 @@
 {
     KW_LOG(("FreeEnvironmentStringsW: %p -> TRUE\n", pwszzEnv));
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     kHlpFree(pwszzEnv);
     return TRUE;
@@ -3206,5 +3325,8 @@
 static DWORD WINAPI kwSandbox_Kernel32_GetEnvironmentVariableA(LPCSTR pszVar, LPSTR pszValue, DWORD cbValue)
 {
-    char *pszFoundValue = kwSandboxDoGetEnvA(&g_Sandbox, pszVar, kHlpStrLen(pszVar));
+    char *pszFoundValue;
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
+
+    pszFoundValue = kwSandboxDoGetEnvA(&g_Sandbox, pszVar, kHlpStrLen(pszVar));
     if (pszFoundValue)
     {
@@ -3222,5 +3344,8 @@
 static DWORD WINAPI kwSandbox_Kernel32_GetEnvironmentVariableW(LPCWSTR pwszVar, LPWSTR pwszValue, DWORD cwcValue)
 {
-    wchar_t *pwszFoundValue = kwSandboxDoGetEnvW(&g_Sandbox, pwszVar, kwUtf16Len(pwszVar));
+    wchar_t *pwszFoundValue;
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
+
+    pwszFoundValue = kwSandboxDoGetEnvW(&g_Sandbox, pwszVar, kwUtf16Len(pwszVar));
     if (pwszFoundValue)
     {
@@ -3239,4 +3364,6 @@
 {
     int rc;
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
+
     if (pszValue)
         rc = kwSandboxDoSetEnvA(&g_Sandbox, pszVar, kHlpStrLen(pszVar), pszValue);
@@ -3261,4 +3388,6 @@
 {
     int rc;
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
+
     if (pwszValue)
         rc = kwSandboxDoSetEnvW(&g_Sandbox, pwszVar, kwUtf16Len(pwszVar), pwszValue);
@@ -3282,4 +3411,5 @@
 static DWORD WINAPI kwSandbox_Kernel32_ExpandEnvironmentStringsA(LPCSTR pszSrc, LPSTR pwszDst, DWORD cbDst)
 {
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     KWFS_TODO();
     return 0;
@@ -3290,4 +3420,5 @@
 static DWORD WINAPI kwSandbox_Kernel32_ExpandEnvironmentStringsW(LPCWSTR pwszSrc, LPWSTR pwszDst, DWORD cbDst)
 {
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     KWFS_TODO();
     return 0;
@@ -3299,5 +3430,8 @@
 {
     int rc;
-    char const *pszEqual = kHlpStrChr(pszVarEqualValue, '=');
+    char const *pszEqual;
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
+
+    pszEqual = kHlpStrChr(pszVarEqualValue, '=');
     if (pszEqual)
     {
@@ -3322,5 +3456,8 @@
 {
     int rc;
-    wchar_t const *pwszEqual = wcschr(pwszVarEqualValue, '=');
+    wchar_t const *pwszEqual;
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
+
+    pwszEqual = wcschr(pwszVarEqualValue, '=');
     if (pwszEqual)
     {
@@ -3344,5 +3481,8 @@
 static errno_t __cdecl kwSandbox_msvcrt__putenv_s(const char *pszVar, const char *pszValue)
 {
-    char const *pszEqual = kHlpStrChr(pszVar, '=');
+    char const *pszEqual;
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
+
+    pszEqual = kHlpStrChr(pszVar, '=');
     if (pszEqual == NULL)
     {
@@ -3373,5 +3513,8 @@
 static errno_t __cdecl kwSandbox_msvcrt__wputenv_s(const wchar_t *pwszVar, const wchar_t *pwszValue)
 {
-    wchar_t const *pwszEqual = wcschr(pwszVar, '=');
+    wchar_t const *pwszEqual;
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
+
+    pwszEqual = wcschr(pwszVar, '=');
     if (pwszEqual == NULL)
     {
@@ -3403,4 +3546,5 @@
 {
     KW_LOG(("__p___initenv\n"));
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     KWFS_TODO();
     return &g_Sandbox.initenv;
@@ -3412,4 +3556,5 @@
 {
     KW_LOG(("__p___winitenv\n"));
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     KWFS_TODO();
     return &g_Sandbox.winitenv;
@@ -3421,4 +3566,5 @@
 {
     KW_LOG(("__p__environ\n"));
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     return &g_Sandbox.environ;
 }
@@ -3429,4 +3575,5 @@
 {
     KW_LOG(("__p__wenviron\n"));
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     return &g_Sandbox.wenviron;
 }
@@ -3438,4 +3585,5 @@
 {
     KWFS_TODO(); /** @todo check the callers expectations! */
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     *ppapszEnviron = g_Sandbox.environ;
     return 0;
@@ -3448,4 +3596,5 @@
 {
     KWFS_TODO(); /** @todo check the callers expectations! */
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     *ppapwszEnviron = g_Sandbox.wenviron;
     return 0;
@@ -3585,4 +3734,5 @@
     PKWMODULE   pMod;
     int         rc;
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
 
     /*
@@ -3732,4 +3882,5 @@
 {
     /* Ignored, we like to keep everything loaded. */
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     return TRUE;
 }
@@ -3741,4 +3892,5 @@
     KSIZE i;
     KSIZE cchModule;
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
 
     /*
@@ -3771,4 +3923,5 @@
     KSIZE i;
     KSIZE cwcModule;
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
 
     /*
@@ -3799,5 +3952,9 @@
 static UINT WINAPI kwSandbox_BreakIntoDebugger(void *pv1, void *pv2, void *pv3, void *pv4)
 {
+#ifdef _MSC_VER
+    __debugbreak();
+#else
     KWFS_TODO();
+#endif
     return -1;
 }
@@ -3808,9 +3965,11 @@
 {
     KSIZE i;
+    PKWMODULE pMod;
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
 
     /*
      * Try locate the module.
      */
-    PKWMODULE pMod = kwToolLocateModuleByHandle(g_Sandbox.pTool, hmod);
+    pMod = kwToolLocateModuleByHandle(g_Sandbox.pTool, hmod);
     if (pMod)
     {
@@ -3858,5 +4017,8 @@
 static DWORD WINAPI kwSandbox_Kernel32_GetModuleFileNameA(HMODULE hmod, LPSTR pszFilename, DWORD cbFilename)
 {
-    PKWMODULE pMod = kwToolLocateModuleByHandle(g_Sandbox.pTool, hmod);
+    PKWMODULE pMod;
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
+
+    pMod = kwToolLocateModuleByHandle(g_Sandbox.pTool, hmod);
     if (pMod != NULL)
     {
@@ -3873,5 +4035,8 @@
 static DWORD WINAPI kwSandbox_Kernel32_GetModuleFileNameW(HMODULE hmod, LPWSTR pwszFilename, DWORD cbFilename)
 {
-    PKWMODULE pMod = kwToolLocateModuleByHandle(g_Sandbox.pTool, hmod);
+    PKWMODULE pMod;
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
+
+    pMod = kwToolLocateModuleByHandle(g_Sandbox.pTool, hmod);
     if (pMod)
     {
@@ -4366,4 +4531,5 @@
     MY_UNICODE_STRING       UniStr;
     MY_NTSTATUS             rcNt;
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
 
     /*
@@ -4538,5 +4704,8 @@
                     {
                         KFSLOOKUPERROR enmError;
-                        PKFSOBJ pFsObj = kFsCacheLookupNoMissingA(g_pFsCache, pszFilename, &enmError);
+                        PKFSOBJ pFsObj;
+                        kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
+
+                        pFsObj = kFsCacheLookupNoMissingA(g_pFsCache, pszFilename, &enmError);
                         if (pFsObj)
                         {
@@ -4617,5 +4786,5 @@
                     if (kwFsIsCachablePathExtensionW(pwszFilename, K_FALSE /*fAttrQuery*/))
                     {
-                        /** @todo rewrite to pure UTF-16. */
+                        /** @todo rewrite in pure UTF-16. */
                         char szTmp[2048];
                         KSIZE cch = kwUtf16ToStr(pwszFilename, szTmp, sizeof(szTmp));
@@ -4648,4 +4817,5 @@
 {
     KUPTR const idxHandle = KW_HANDLE_TO_INDEX(hFile);
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     if (idxHandle < g_Sandbox.cHandles)
     {
@@ -4728,4 +4898,5 @@
 {
     KUPTR const idxHandle = KW_HANDLE_TO_INDEX(hFile);
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     if (idxHandle < g_Sandbox.cHandles)
     {
@@ -4807,4 +4978,5 @@
 {
     KUPTR const idxHandle = KW_HANDLE_TO_INDEX(hFile);
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     if (idxHandle < g_Sandbox.cHandles)
     {
@@ -4922,4 +5094,5 @@
 {
     KUPTR const idxHandle = KW_HANDLE_TO_INDEX(hFile);
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     if (idxHandle < g_Sandbox.cHandles)
     {
@@ -4998,4 +5171,5 @@
 {
     KUPTR const idxHandle = KW_HANDLE_TO_INDEX(hFile);
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     if (idxHandle < g_Sandbox.cHandles)
     {
@@ -5103,4 +5277,5 @@
 {
     KUPTR const idxHandle = KW_HANDLE_TO_INDEX(hFile);
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     if (idxHandle < g_Sandbox.cHandles)
     {
@@ -5121,4 +5296,5 @@
 {
     KUPTR const idxHandle = KW_HANDLE_TO_INDEX(hFile);
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     if (idxHandle < g_Sandbox.cHandles)
     {
@@ -5167,4 +5343,5 @@
 {
     KUPTR const idxHandle = KW_HANDLE_TO_INDEX(hFile);
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     if (idxHandle < g_Sandbox.cHandles)
     {
@@ -5194,4 +5371,5 @@
 {
     KUPTR const idxHandle = KW_HANDLE_TO_INDEX(hFile);
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     if (idxHandle < g_Sandbox.cHandles)
     {
@@ -5229,4 +5407,5 @@
 {
     KUPTR const idxHandle = KW_HANDLE_TO_INDEX(hFile);
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     if (idxHandle < g_Sandbox.cHandles)
     {
@@ -5265,4 +5444,5 @@
 {
     KUPTR const idxHandle = KW_HANDLE_TO_INDEX(hFile);
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     if (idxHandle < g_Sandbox.cHandles)
     {
@@ -5304,4 +5484,5 @@
 {
     KUPTR const idxHandle = KW_HANDLE_TO_INDEX(hSection);
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     if (idxHandle < g_Sandbox.cHandles)
     {
@@ -5389,4 +5570,5 @@
     /* Is this one of our temporary mappings? */
     PKWFSTEMPFILE pCur = g_Sandbox.pTempFileHead;
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     while (pCur)
     {
@@ -5415,4 +5597,5 @@
     BOOL        fRet;
     KUPTR const idxHandle = KW_HANDLE_TO_INDEX(hObject);
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     if (   idxHandle < g_Sandbox.cHandles
         && g_Sandbox.papHandles[idxHandle] != NULL)
@@ -5454,5 +5637,8 @@
     {
         KFSLOOKUPERROR enmError;
-        PKFSOBJ pFsObj = kFsCacheLookupNoMissingA(g_pFsCache, pszFilename, &enmError);
+        PKFSOBJ pFsObj;
+        kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
+
+        pFsObj = kFsCacheLookupNoMissingA(g_pFsCache, pszFilename, &enmError);
         if (pFsObj)
         {
@@ -5484,5 +5670,8 @@
     {
         KFSLOOKUPERROR enmError;
-        PKFSOBJ pFsObj = kFsCacheLookupNoMissingW(g_pFsCache, pwszFilename, &enmError);
+        PKFSOBJ pFsObj;
+        kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
+
+        pFsObj = kFsCacheLookupNoMissingW(g_pFsCache, pwszFilename, &enmError);
         if (pFsObj)
         {
@@ -5516,5 +5705,8 @@
     {
         KFSLOOKUPERROR enmError;
-        PKFSOBJ pObj = kFsCacheLookupW(g_pFsCache, pwszLongPath, &enmError);
+        PKFSOBJ pObj;
+        kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
+
+        pObj = kFsCacheLookupW(g_pFsCache, pwszLongPath, &enmError);
         if (pObj)
         {
@@ -5568,4 +5760,5 @@
         && kwFsIsClTempFileW(pwszFilename))
     {
+        kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
         KWFS_LOG(("DeleteFileW(%s) -> TRUE [temp]\n", pwszFilename));
         fRc = TRUE;
@@ -5978,4 +6171,5 @@
     BOOL                    fRc;
     PKWCONSOLEOUTPUTLINE    pLineBuf;
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
 
     if (hConOutput == g_Sandbox.StdErr.hOutput)
@@ -6009,4 +6203,5 @@
     BOOL                    fRc;
     PKWCONSOLEOUTPUTLINE    pLineBuf;
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
 
     if (hConOutput == g_Sandbox.StdErr.hOutput)
@@ -6054,5 +6249,8 @@
         && pvMem)
     {
-        PKWVIRTALLOC pTracker = g_Sandbox.pVirtualAllocHead;
+        PKWVIRTALLOC pTracker;
+        kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
+
+        pTracker = g_Sandbox.pVirtualAllocHead;
         while (   pTracker
                && (KUPTR)pvMem - (KUPTR)pTracker->pvAlloc >= pTracker->cbAlloc)
@@ -6083,4 +6281,5 @@
     if (g_Sandbox.pTool->u.Sandboxed.enmHint == KWTOOLHINT_VISUAL_CPP_CL)
     {
+        kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
         if (dwFreeType & MEM_RELEASE)
         {
@@ -6112,4 +6311,64 @@
 
 
+/** Kernel32 - HeapCreate / NtDll - RTlCreateHeap  */
+HANDLE WINAPI kwSandbox_Kernel32_HeapCreate(DWORD fOptions, SIZE_T cbInitial, SIZE_T cbMax)
+{
+    HANDLE hHeap;
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
+
+    hHeap = HeapCreate(fOptions, cbInitial, cbMax);
+    if (hHeap != NULL)
+    {
+        DWORD dwErr = GetLastError();
+        PKWHEAP pTracker = (PKWHEAP)kHlpAlloc(sizeof(*pTracker));
+        if (pTracker)
+        {
+            pTracker->hHeap = hHeap;
+            pTracker->pNext = g_Sandbox.pHeapHead;
+            g_Sandbox.pHeapHead = pTracker;
+        }
+
+        SetLastError(dwErr);
+    }
+    return hHeap;
+
+}
+
+
+/** Kernel32 - HeapDestroy / NtDll - RTlDestroyHeap */
+BOOL WINAPI kwSandbox_Kernel32_HeapDestroy(HANDLE hHeap)
+{
+    BOOL fRc = HeapDestroy(hHeap);
+    KW_LOG(("HeapDestroy: hHeap=%p -> %d\n", hHeap, fRc));
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
+    if (fRc)
+    {
+        PKWHEAP pTracker = g_Sandbox.pHeapHead;
+        if (pTracker)
+        {
+            if (pTracker->hHeap == hHeap)
+                g_Sandbox.pHeapHead = pTracker->pNext;
+            else
+            {
+                PKWHEAP pPrev;
+                do
+                {
+                    pPrev = pTracker;
+                    pTracker = pTracker->pNext;
+                } while (pTracker && pTracker->hHeap == hHeap);
+                if (pTracker)
+                    pPrev->pNext = pTracker->pNext;
+            }
+            if (pTracker)
+                kHlpFree(pTracker);
+            else
+                KW_LOG(("HeapDestroy: pvAddr=%p not found!\n", hHeap));
+        }
+    }
+
+    return fRc;
+}
+
+
 
 /*
@@ -6136,4 +6395,5 @@
         if (pTracker)
         {
+            kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
             pTracker->idx = idxFls;
             pTracker->pNext = g_Sandbox.pFlsAllocHead;
@@ -6152,5 +6412,8 @@
     if (fRc)
     {
-        PKWLOCALSTORAGE pTracker = g_Sandbox.pFlsAllocHead;
+        PKWLOCALSTORAGE pTracker;
+        kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
+
+        pTracker = g_Sandbox.pFlsAllocHead;
         if (pTracker)
         {
@@ -6215,4 +6478,5 @@
                     if (pHash)
                     {
+                        kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
                         pHash->uMagic        = KWHASHMD5_MAGIC;
                         pHash->cbHashed      = 0;
@@ -6258,4 +6522,5 @@
     BOOL        fRc;
     PKWHASHMD5  pHash = g_Sandbox.pHashHead;
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     while (pHash && (KUPTR)pHash != hHash)
         pHash = pHash->pNext;
@@ -6373,4 +6638,5 @@
     BOOL        fRc;
     PKWHASHMD5  pHash = g_Sandbox.pHashHead;
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     while (pHash && (KUPTR)pHash != hHash)
         pHash = pHash->pNext;
@@ -6530,4 +6796,5 @@
     PKWHASHMD5  pPrev = NULL;
     PKWHASHMD5  pHash = g_Sandbox.pHashHead;
+    kHlpAssert(GetCurrentThreadId() == g_Sandbox.idMainThread);
     while (pHash && (KUPTR)pHash != hHash)
     {
@@ -6668,4 +6935,7 @@
     { TUPLE("VirtualFree"),                 NULL,       (KUPTR)kwSandbox_Kernel32_VirtualFree },
 
+    { TUPLE("HeapCreate"),                  NULL,       (KUPTR)kwSandbox_Kernel32_HeapCreate,       K_TRUE /*fOnlyExe*/ },
+    { TUPLE("HeapDestroy"),                 NULL,       (KUPTR)kwSandbox_Kernel32_HeapDestroy,      K_TRUE /*fOnlyExe*/ },
+
     { TUPLE("FlsAlloc"),                    NULL,       (KUPTR)kwSandbox_Kernel32_FlsAlloc },
     { TUPLE("FlsFree"),                     NULL,       (KUPTR)kwSandbox_Kernel32_FlsFree },
@@ -6689,4 +6959,8 @@
     { TUPLE("_amsg_exit"),                  NULL,       (KUPTR)kwSandbox_msvcrt__amsg_exit },
     { TUPLE("terminate"),                   NULL,       (KUPTR)kwSandbox_msvcrt_terminate },
+
+    { TUPLE("onexit"),                      NULL,       (KUPTR)kwSandbox_msvcrt__onexit,            K_TRUE /*fOnlyExe*/ },
+    { TUPLE("_onexit"),                     NULL,       (KUPTR)kwSandbox_msvcrt__onexit,            K_TRUE /*fOnlyExe*/ },
+    { TUPLE("atexit"),                      NULL,       (KUPTR)kwSandbox_msvcrt_atexit,             K_TRUE /*fOnlyExe*/ },
 
     { TUPLE("_beginthread"),                NULL,       (KUPTR)kwSandbox_msvcrt__beginthread },
@@ -7200,4 +7474,5 @@
     PROCESS_MEMORY_COUNTERS     MemInfo;
     PKWVIRTALLOC                pTracker;
+    PKWHEAP                     pHeap;
     PKWLOCALSTORAGE             pLocalStorage;
 #ifdef WITH_HASH_MD5_CACHE
@@ -7206,5 +7481,31 @@
 #ifdef WITH_TEMP_MEMORY_FILES
     PKWFSTEMPFILE               pTempFile;
-
+#endif
+    PKWEXITCALLACK              pExitCallback;
+
+    /* Do exit callback first. */
+    pExitCallback = g_Sandbox.pExitCallbackHead;
+    g_Sandbox.pExitCallbackHead = NULL;
+    while (pExitCallback)
+    {
+        PKWEXITCALLACK  pNext = pExitCallback->pNext;
+        KW_LOG(("kwSandboxCleanupLate: calling %p %sexit handler\n",
+                pExitCallback->pfnCallback, pExitCallback->fAtExit ? "at" : "_on"));
+        __try
+        {
+            pExitCallback->pfnCallback();
+        }
+        __except (EXCEPTION_EXECUTE_HANDLER)
+        {
+            KW_LOG(("kwSandboxCleanupLate: %sexit handler %p threw an exception!\n",
+                    pExitCallback->fAtExit ? "at" : "_on", pExitCallback->pfnCallback));
+            kHlpAssertFailed();
+        }
+        kHlpFree(pExitCallback);
+        pExitCallback = pNext;
+    }
+
+
+#ifdef WITH_TEMP_MEMORY_FILES
     /* The temporary files aren't externally visible, they're all in memory. */
     pTempFile = pSandbox->pTempFileHead;
@@ -7234,4 +7535,15 @@
         kHlpFree(pTracker);
         pTracker = pNext;
+    }
+
+    /* Free left behind HeapCreate leaks. */
+    pHeap = g_Sandbox.pHeapHead;
+    g_Sandbox.pHeapHead = NULL;
+    while (pHeap != NULL)
+    {
+        PKWHEAP pNext = pHeap->pNext;
+        KW_LOG(("Freeing HeapCreate leak %p\n", pHeap->hHeap));
+        HeapDestroy(pHeap->hHeap);
+        pHeap = pNext;
     }
 
