[release/7.0] [NativeAOT] correctly initalize CONTEXT before failing fast (#81837)

* [NativeAOT] correctly initalize CONTEXT before failing fast

* Switch from using GetThreadContext to RtlCaptureContext.

* Add a better explination of why the function is unimplmented on Unix.

* Update src/coreclr/nativeaot/Runtime/PalRedhawk.h

Co-authored-by: Vladimir Sadov <[email protected]>

* Respond to feedback: unconditionally set CONTEXT_CONTROL

* Respond to feedback: consolidate setting of ContextFlags.

* On second thought, don't add a second layer of ifdef nesting.

---------

Co-authored-by: Austin Wise <[email protected]>
Co-authored-by: Vladimir Sadov <[email protected]>

Author: 3 people

src/coreclr/nativeaot/Runtime/EHHelpers.cpp

@@ -91,9 +91,19 @@ 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));
     CONTEXT* pContext = (CONTEXT *)pOSContext;
+
+#ifndef HOST_WASM
+
+    memset(pOSContext, 0, cbOSContext);
+    pContext->ContextFlags = CONTEXT_CONTROL | CONTEXT_INTEGER;
+
+    // Fill in CONTEXT_CONTROL registers that were not captured in PAL_LIMITED_CONTEXT.
+    PopulateControlSegmentRegisters(pContext);
+
+#endif // !HOST_WASM
+
 #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;
@@ -224,6 +229,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
 
@@ -267,6 +277,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
 
@@ -321,6 +337,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.
@@ -596,6 +617,11 @@ 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`.
+REDHAWK_PALIMPORT void REDHAWK_PALAPI PopulateControlSegmentRegisters(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

@@ -1100,6 +1100,14 @@ extern "C" int32_t _stricmp(const char *string1, const char *string2)
     return strcasecmp(string1, string2);
 }
 
+REDHAWK_PALIMPORT void REDHAWK_PALAPI PopulateControlSegmentRegisters(CONTEXT* pContext)
+{
+#if defined(TARGET_X86) || defined(TARGET_AMD64)
+    // Currently the CONTEXT is only used on Windows for RaiseFailFastException.
+    // So we punt on filling in SegCs and SegSs for now.
+#endif
+}
+
 uint32_t g_RhNumberOfProcessors;
 
 REDHAWK_PALEXPORT int32_t PalGetProcessCpuCount()

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

@@ -437,6 +437,18 @@ REDHAWK_PALEXPORT void REDHAWK_PALAPI PalRestoreContext(CONTEXT * pCtx)
     RtlRestoreContext(pCtx, NULL);
 }
 
+REDHAWK_PALIMPORT void REDHAWK_PALAPI PopulateControlSegmentRegisters(CONTEXT* pContext)
+{
+#if defined(TARGET_X86) || defined(TARGET_AMD64)
+    CONTEXT ctx;
+
+    RtlCaptureContext(&ctx);
+
+    pContext->SegCs = ctx.SegCs;
+    pContext->SegSs = ctx.SegSs;
+#endif //defined(TARGET_X86) || defined(TARGET_AMD64)
+}
+
 static PalHijackCallback g_pHijackCallback;
 
 REDHAWK_PALEXPORT UInt32_BOOL REDHAWK_PALAPI PalRegisterHijackCallback(_In_ PalHijackCallback callback)