[NativeAOT] correctly initalize CONTEXT before failing fast

Author: AustinWise

src/coreclr/nativeaot/Runtime/EHHelpers.cpp

@@ -90,9 +90,18 @@ COOP_PINVOKE_HELPER(int32_t, RhGetModuleFileName, (HANDLE moduleHandle, _Out_ co
 
 COOP_PINVOKE_HELPER(void, RhpCopyContextFromExInfo, (void * pOSContext, int32_t cbOSContext, PAL_LIMITED_CONTEXT * pPalContext))
 {
-    UNREFERENCED_PARAMETER(cbOSContext);
     ASSERT((size_t)cbOSContext >= sizeof(CONTEXT));
+    memset(pOSContext, 0, cbOSContext);
     CONTEXT* pContext = (CONTEXT *)pOSContext;
+
+#ifndef HOST_WASM
+    if (TryPopulateControlSegmentRegisters(pContext))
+    {
+        pContext->ContextFlags |= CONTEXT_CONTROL;
+    }
+    pContext->ContextFlags |= CONTEXT_INTEGER;
+#endif
+
 #if defined(UNIX_AMD64_ABI)
     pContext->Rip = pPalContext->IP;
     pContext->Rsp = pPalContext->Rsp;

src/coreclr/nativeaot/Runtime/PalRedhawk.h

@@ -112,6 +112,11 @@ struct FILETIME
 
 #ifdef HOST_AMD64
 
+#define CONTEXT_AMD64   0x00100000L
+
+#define CONTEXT_CONTROL         (CONTEXT_AMD64 | 0x00000001L)
+#define CONTEXT_INTEGER         (CONTEXT_AMD64 | 0x00000002L)
+
 typedef struct DECLSPEC_ALIGN(16) _XSAVE_FORMAT {
     uint16_t  ControlWord;
     uint16_t  StatusWord;
@@ -232,6 +237,11 @@ typedef struct DECLSPEC_ALIGN(16) _CONTEXT {
 } CONTEXT, *PCONTEXT;
 #elif defined(HOST_ARM)
 
+#define CONTEXT_ARM   0x00200000L
+
+#define CONTEXT_CONTROL (CONTEXT_ARM | 0x1L)
+#define CONTEXT_INTEGER (CONTEXT_ARM | 0x2L)
+
 #define ARM_MAX_BREAKPOINTS     8
 #define ARM_MAX_WATCHPOINTS     1
 
@@ -275,6 +285,12 @@ typedef struct DECLSPEC_ALIGN(8) _CONTEXT {
 } CONTEXT, *PCONTEXT;
 
 #elif defined(HOST_X86)
+
+#define CONTEXT_i386    0x00010000L
+
+#define CONTEXT_CONTROL         (CONTEXT_i386 | 0x00000001L) // SS:SP, CS:IP, FLAGS, BP
+#define CONTEXT_INTEGER         (CONTEXT_i386 | 0x00000002L) // AX, BX, CX, DX, SI, DI
+
 #define SIZE_OF_80387_REGISTERS      80
 #define MAXIMUM_SUPPORTED_EXTENSION  512
 
@@ -329,6 +345,11 @@ typedef struct _CONTEXT {
 
 #elif defined(HOST_ARM64)
 
+#define CONTEXT_ARM64   0x00400000L
+
+#define CONTEXT_CONTROL (CONTEXT_ARM64 | 0x1L)
+#define CONTEXT_INTEGER (CONTEXT_ARM64 | 0x2L)
+
 // Specify the number of breakpoints and watchpoints that the OS
 // will track. Architecturally, ARM64 supports up to 16. In practice,
 // however, almost no one implements more than 4 of each.
@@ -484,8 +505,6 @@ typedef enum _EXCEPTION_DISPOSITION {
 //#endif // !DACCESS_COMPILE
 #endif // !_INC_WINDOWS
 
-
-
 #ifndef DACCESS_COMPILE
 #ifndef _INC_WINDOWS
 
@@ -614,6 +633,14 @@ REDHAWK_PALIMPORT bool REDHAWK_PALAPI PalGetCompleteThreadContext(HANDLE hThread
 REDHAWK_PALIMPORT bool REDHAWK_PALAPI PalSetThreadContext(HANDLE hThread, _Out_ CONTEXT * pCtx);
 REDHAWK_PALIMPORT void REDHAWK_PALAPI PalRestoreContext(CONTEXT * pCtx);
 
+// For platforms that have segment registers in the CONTEXT_CONTROL set that
+// are not saved in PAL_LIMITED_CONTEXT, this captures them from the current
+// thread and saves them in `pContext`.
+// This function returns true if the current plateform has no such registers or
+// if the registers are successfully saved in `pContext`.
+// Other false is returned.
+REDHAWK_PALIMPORT bool REDHAWK_PALAPI TryPopulateControlSegmentRegisters(CONTEXT* pContext);
+
 REDHAWK_PALIMPORT int32_t REDHAWK_PALAPI PalGetProcessCpuCount();
 
 // Retrieves the entire range of memory dedicated to the calling thread's stack.  This does

src/coreclr/nativeaot/Runtime/unix/PalRedhawkUnix.cpp

@@ -1068,6 +1068,16 @@ extern "C" int32_t _stricmp(const char *string1, const char *string2)
     return strcasecmp(string1, string2);
 }
 
+REDHAWK_PALIMPORT bool REDHAWK_PALAPI TryPopulateControlSegmentRegisters(CONTEXT* pContext)
+{
+#if defined(TARGET_X86) || defined(TARGET_AMD64)
+    // TODO: attempt to fill in SegCs and SegSs
+    return false;
+#else
+    return true;
+#endif
+}
+
 uint32_t g_RhNumberOfProcessors;
 
 REDHAWK_PALEXPORT int32_t PalGetProcessCpuCount()

src/coreclr/nativeaot/Runtime/windows/PalRedhawkMinWin.cpp

@@ -446,6 +446,27 @@ REDHAWK_PALEXPORT void REDHAWK_PALAPI PalRestoreContext(CONTEXT * pCtx)
     RtlRestoreContext(pCtx, NULL);
 }
 
+REDHAWK_PALIMPORT bool REDHAWK_PALAPI TryPopulateControlSegmentRegisters(CONTEXT* pContext)
+{
+#if defined(TARGET_X86) || defined(TARGET_AMD64)
+    HANDLE hThread = GetCurrentThread();
+
+    CONTEXT ctx;
+    memset(&ctx, 0, sizeof(ctx));
+    ctx.ContextFlags = CONTEXT_CONTROL;
+
+    if (!GetThreadContext(hThread, &ctx))
+    {
+        return false;
+    }
+
+    pContext->SegCs = ctx.SegCs;
+    pContext->SegSs = ctx.SegSs;
+#endif //defined(TARGET_X86) || defined(TARGET_AMD64)
+
+    return true;
+}
+
 static PalHijackCallback g_pHijackCallback;
 
 #ifdef FEATURE_SPECIAL_USER_MODE_APC