Fix async continuation hijacking for x86 NativeAOT

src/coreclr/inc/eetwain.h

@@ -262,7 +262,7 @@ virtual size_t GetFunctionSize(GCInfoToken gcInfoToken) = 0;
 *  returns true.
 *  If hijacking is not possible for some reason, it return false.
 */
-virtual bool GetReturnAddressHijackInfo(GCInfoToken gcInfoToken X86_ARG(ReturnKind * returnKind)) = 0;
+virtual bool GetReturnAddressHijackInfo(GCInfoToken gcInfoToken X86_ARG(ReturnKind * returnKind) X86_ARG(bool* hasAsyncRet)) = 0;
 
 #ifndef USE_GC_INFO_DECODER
 /*
@@ -478,7 +478,7 @@ size_t GetFunctionSize(GCInfoToken gcInfoToken);
 *  returns true.
 *  If hijacking is not possible for some reason, it return false.
 */
-virtual bool GetReturnAddressHijackInfo(GCInfoToken gcInfoToken X86_ARG(ReturnKind * returnKind));
+virtual bool GetReturnAddressHijackInfo(GCInfoToken gcInfoToken X86_ARG(ReturnKind * returnKind) X86_ARG(bool* hasAsyncRet));
 
 #ifndef USE_GC_INFO_DECODER
 /*
@@ -627,7 +627,7 @@ bool IsInPrologOrEpilog(
 virtual
 size_t GetFunctionSize(GCInfoToken gcInfoToken);
 
-virtual bool GetReturnAddressHijackInfo(GCInfoToken gcInfoToken X86_ARG(ReturnKind * returnKind))
+virtual bool GetReturnAddressHijackInfo(GCInfoToken gcInfoToken X86_ARG(ReturnKind * returnKind) X86_ARG(bool* hasAsyncRet))
 {
     // Interpreter-TODO: Implement this if needed
     _ASSERTE(FALSE);

src/coreclr/inc/gc_unwind_x86.h

@@ -331,6 +331,7 @@ struct hdrInfo
     unsigned char       epilogCnt;
     bool                epilogEnd;      // is the epilog at the end of the method
 
+    bool                isAsync;        // is this an async function with async continuation arg/return
     bool                ebpFrame;       // locals and arguments addressed relative to EBP
     bool                doubleAlign;    // is the stack double-aligned? locals addressed relative to ESP, and arguments relative to EBP
     bool                interruptible;  // intr. at all times (excluding prolog/epilog), not just call sites

src/coreclr/inc/gcinfo.h

@@ -37,7 +37,7 @@ const unsigned   this_OFFSET_FLAG  = 0x2;  // the offset is "this"
 // The current GCInfo Version
 //-----------------------------------------------------------------------------
 
-#define GCINFO_VERSION 4
+#define GCINFO_VERSION 5
 
 #ifdef SOS_INCLUDE
 extern bool IsRuntimeVersionAtLeast(DWORD major);
@@ -46,7 +46,11 @@ inline int GCInfoVersion()
     // In SOS we only care about ability to parse/dump the GC Info.
     // Since v2 and v3 had the same file format and v1 is no longer supported,
     // we can assume that everything before net10.0 uses GCInfo v3.
-    return IsRuntimeVersionAtLeast(10) ? 4 : 3;
+    if (IsRuntimeVersionAtLeast(11))
+        return 5;
+    if (IsRuntimeVersionAtLeast(10))
+        return 4;
+    return 3;
 }
 #endif
 
@@ -80,6 +84,9 @@ struct GCInfoToken
     // Keep this in sync with GetR2RGCInfoVersion in cDac (ExecutionManagerCore.ReadyToRunJitManager.cs)
     static uint32_t ReadyToRunVersionToGcInfoVersion(uint32_t readyToRunMajorVersion, uint32_t readyToRunMinorVersion)
     {
+        if (readyToRunMajorVersion >= 21)
+            return 5;
+
         if (readyToRunMajorVersion >= 11)
             return 4;
 

src/coreclr/inc/gcinfotypes.h

@@ -305,7 +305,8 @@ enum infoHdrAdjustConstants {
     SET_EPILOGSIZE_MAX = 10,  // Change to 6
     SET_EPILOGCNT_MAX = 4,
     SET_UNTRACKED_MAX = 3,
-    SET_RET_KIND_MAX = 3,   // 2 bits for ReturnKind
+    SET_RET_KIND_MAX_V4 = 3,   // 2 bits for ReturnKind
+    SET_RET_KIND_MAX_V5 = 7,   // 3 bits for ReturnKind + isAsync
     SET_NOGCREGIONS_MAX = 4,
     ADJ_ENCODING_MAX = 0x7f, // Maximum valid encoding in a byte
                              // Also used to mask off next bit from each encoding byte.
@@ -358,8 +359,10 @@ enum infoHdrAdjust {
 // Second set of opcodes, when first code is 0x4F
 enum infoHdrAdjust2 {
     SET_RETURNKIND = 0,  // 0x00-SET_RET_KIND_MAX Set ReturnKind to value
-    SET_NOGCREGIONS_CNT = SET_RETURNKIND + SET_RET_KIND_MAX + 1,        // 0x04
-    FFFF_NOGCREGION_CNT = SET_NOGCREGIONS_CNT + SET_NOGCREGIONS_MAX + 1 // 0x09 There is a count (>SET_NOGCREGIONS_MAX) after the header encoding
+    SET_NOGCREGIONS_CNT_V4 = SET_RETURNKIND + SET_RET_KIND_MAX_V4 + 1,        // 0x04
+    FFFF_NOGCREGION_CNT_V4 = SET_NOGCREGIONS_CNT_V4 + SET_NOGCREGIONS_MAX + 1, // 0x09 There is a count (>SET_NOGCREGIONS_MAX) after the header encoding
+    SET_NOGCREGIONS_CNT_V5 = SET_RETURNKIND + SET_RET_KIND_MAX_V5 + 1,        // 0x08
+    FFFF_NOGCREGION_CNT_V5 = SET_NOGCREGIONS_CNT_V5 + SET_NOGCREGIONS_MAX + 1, // 0x0D There is a count (>SET_NOGCREGIONS_MAX) after the header encoding
 };
 
 #define HAS_UNTRACKED               ((unsigned int) -1)
@@ -410,6 +413,7 @@ struct InfoHdrSmall {
     unsigned char  genericsContext : 1;//4 [1]      function reports a generics context parameter is present
     unsigned char  genericsContextIsMethodDesc : 1;//4[2]
     unsigned char  returnKind : 2; // 4 [4]  Available GcInfo v2 onwards, previously undefined
+    unsigned char  isAsync : 1;    // 4 [5]
     unsigned short argCount;          // 5,6        in bytes
     unsigned int   frameSize;         // 7,8,9,10   in bytes
     unsigned int   untrackedCnt;      // 11,12,13,14

src/coreclr/inc/readytorun.h

@@ -19,10 +19,10 @@
 //  src/coreclr/nativeaot/Runtime/inc/ModuleHeaders.h
 // If you update this, ensure you run `git grep MINIMUM_READYTORUN_MAJOR_VERSION`
 // and handle pending work.
-#define READYTORUN_MAJOR_VERSION 20
+#define READYTORUN_MAJOR_VERSION 21
 #define READYTORUN_MINOR_VERSION 0x0000
 
-#define MINIMUM_READYTORUN_MAJOR_VERSION 20
+#define MINIMUM_READYTORUN_MAJOR_VERSION 21
 
 // R2R Version 2.1 adds the InliningInfo section
 // R2R Version 2.2 adds the ProfileDataInfo section
@@ -60,6 +60,7 @@
 // R2R Version 18.7 adds READYTORUN_HELPER_R2RToInterpreter
 // R2R Version 19 removes the READYTORUN_HELPER_ByRefWriteBarrier helper
 // R2R Version 20 changes NativeVarInfo encoding to include ASYNC_CONTINUATION_ILNUM
+// R2R Version 21 updates GC info version to 5 which adds isAsync to x86 GC info
 
 struct READYTORUN_CORE_HEADER
 {

src/coreclr/jit/gcencode.cpp

@@ -877,12 +877,13 @@ BYTE FASTCALL encodeHeaderNext(const InfoHdr& header, InfoHdr* state, BYTE& code
         goto DO_RETURN;
     }
 
-    if (state->returnKind != header.returnKind)
+    if ((state->returnKind != header.returnKind) || (state->isAsync != header.isAsync))
     {
         state->returnKind = header.returnKind;
+        state->isAsync    = header.isAsync;
         codeSet           = 2; // Two byte encoding
-        encoding          = header.returnKind;
-        _ASSERTE(encoding <= SET_RET_KIND_MAX);
+        encoding          = header.returnKind | (header.isAsync ? 4 : 0);
+        _ASSERTE(encoding <= SET_RET_KIND_MAX_V5);
         goto DO_RETURN;
     }
 
@@ -955,14 +956,14 @@ BYTE FASTCALL encodeHeaderNext(const InfoHdr& header, InfoHdr* state, BYTE& code
         {
             state->noGCRegionCnt = header.noGCRegionCnt;
             codeSet              = 2;
-            encoding             = (BYTE)(SET_NOGCREGIONS_CNT + header.noGCRegionCnt);
+            encoding             = (BYTE)(SET_NOGCREGIONS_CNT_V5 + header.noGCRegionCnt);
             goto DO_RETURN;
         }
         else if (state->noGCRegionCnt != HAS_NOGCREGIONS)
         {
             state->noGCRegionCnt = HAS_NOGCREGIONS;
             codeSet              = 2;
-            encoding             = FFFF_NOGCREGION_CNT;
+            encoding             = FFFF_NOGCREGION_CNT_V5;
             goto DO_RETURN;
         }
     }
@@ -1187,9 +1188,9 @@ static int measureDistance(const InfoHdr& header, const InfoHdrSmall* p, int clo
             return distance;
     }
 
-    if (p->returnKind != header.returnKind)
+    if ((p->returnKind != header.returnKind) || (p->isAsync != header.isAsync))
     {
-        // Setting the ReturnKind requires two bytes of encoding.
+        // Setting the ReturnKind/isAsync requires two bytes of encoding.
         distance += 2;
         if (distance >= closeness)
             return distance;
@@ -1570,8 +1571,9 @@ size_t GCInfo::gcInfoBlockHdrSave(
     ReturnKind returnKind = getReturnKind();
     _ASSERTE(IsValidReturnKind(returnKind) && "Return Kind must be valid");
     _ASSERTE(!IsStructReturnKind(returnKind) && "Struct Return Kinds Unexpected for JIT32");
-    _ASSERTE(((int)returnKind <= (int)SET_RET_KIND_MAX) && "ReturnKind has no legal encoding");
+    _ASSERTE(((int)returnKind <= (int)SET_RET_KIND_MAX_V5) && "ReturnKind has no legal encoding");
     header->returnKind = returnKind;
+    header->isAsync    = m_compiler->compIsAsync();
 
     header->gsCookieOffset = INVALID_GS_COOKIE_OFFSET;
     if (m_compiler->getNeedsGSSecurityCookie())

src/coreclr/nativeaot/Runtime/ICodeManager.h

@@ -37,17 +37,18 @@ enum GCRefKind : unsigned char
 static_assert(PTFF_RAX_IS_GCREF == ((uint64_t)GCRK_Object << 16));
 static_assert(PTFF_RAX_IS_BYREF == ((uint64_t)GCRK_Byref << 16));
 
-inline uintptr_t ReturnKindToTransitionFrameFlags(GCRefKind returnKind)
+inline uintptr_t ReturnKindToTransitionFrameFlags(GCRefKind returnKind, bool isAsync)
 {
     // just need to report gc ref bits here.
     // appropriate PTFF_SAVE_ bits will be added by the frame building routine.
-    return ((uintptr_t)returnKind << 16);
+    return ((uintptr_t)returnKind << 16) | (isAsync ? PTFF_RCX_IS_GCREF : 0);
 }
 
-inline GCRefKind TransitionFrameFlagsToReturnKind(uintptr_t transFrameFlags)
+inline GCRefKind TransitionFrameFlagsToReturnKind(uintptr_t transFrameFlags, bool* isAsync)
 {
     GCRefKind returnKind = (GCRefKind)((transFrameFlags & (PTFF_RAX_IS_GCREF | PTFF_RAX_IS_BYREF)) >> 16);
     ASSERT((returnKind == GCRK_Scalar) || (transFrameFlags & PTFF_SAVE_RAX));
+    *isAsync = (transFrameFlags & PTFF_RCX_IS_GCREF) != 0;
     return returnKind;
 }
 
@@ -168,7 +169,8 @@ class ICodeManager
 
 #ifdef TARGET_X86
     virtual GCRefKind GetReturnValueKind(MethodInfo *   pMethodInfo,
-                                  REGDISPLAY *   pRegisterSet) PURE_VIRTUAL
+                                  REGDISPLAY *   pRegisterSet,
+                                  bool* isAsync) PURE_VIRTUAL
 #endif
 
     virtual PTR_VOID RemapHardwareFaultToGCSafePoint(MethodInfo * pMethodInfo, PTR_VOID controlPC) PURE_VIRTUAL

src/coreclr/nativeaot/Runtime/StackFrameIterator.cpp

@@ -117,6 +117,7 @@ void StackFrameIterator::EnterInitialInvalidState(Thread * pThreadToWalk)
 #ifdef TARGET_X86
     m_pHijackedReturnValue = NULL;
     m_HijackedReturnValueKind = GCRK_Unknown;
+    m_pHijackedAsyncContinuation = NULL;
 #endif
     m_pConservativeStackRangeLowerBound = NULL;
     m_pConservativeStackRangeUpperBound = NULL;
@@ -336,12 +337,18 @@ void StackFrameIterator::InternalInit(Thread * pThreadToWalk, PInvokeTransitionF
 #endif // TARGET_AMD64
 
 #ifdef TARGET_X86
-    GCRefKind retValueKind = TransitionFrameFlagsToReturnKind(pFrame->m_Flags);
+    bool isAsync;
+    GCRefKind retValueKind = TransitionFrameFlagsToReturnKind(pFrame->m_Flags, &isAsync);
     if (retValueKind != GCRK_Scalar)
     {
         m_pHijackedReturnValue = (PTR_OBJECTREF)m_RegDisplay.pRax;
         m_HijackedReturnValueKind = retValueKind;
     }
+
+    if (isAsync)
+    {
+        m_pHijackedAsyncContinuation = (PTR_OBJECTREF)m_RegDisplay.pRcx;
+    }
 #endif
 
 #endif // TARGET_ARM
@@ -1793,6 +1800,7 @@ void StackFrameIterator::NextInternal()
 #ifdef TARGET_X86
     m_pHijackedReturnValue = NULL;
     m_HijackedReturnValueKind = GCRK_Unknown;
+    m_pHijackedAsyncContinuation = NULL;
 #endif
 
 #ifdef _DEBUG
@@ -2223,6 +2231,15 @@ bool StackFrameIterator::GetHijackedReturnValueLocation(PTR_OBJECTREF * pLocatio
     *pKind = m_HijackedReturnValueKind;
     return true;
 }
+
+bool StackFrameIterator::GetHijackedAsyncContinuation(PTR_OBJECTREF * pLocation)
+{
+    if (m_pHijackedAsyncContinuation == NULL)
+        return false;
+
+    *pLocation = m_pHijackedAsyncContinuation;
+    return true;
+}
 #endif
 
 void StackFrameIterator::SetControlPC(PTR_VOID controlPC)

src/coreclr/nativeaot/Runtime/StackFrameIterator.h

@@ -53,6 +53,7 @@ class StackFrameIterator
     bool             IsActiveStackFrame();
 #ifdef TARGET_X86
     bool             GetHijackedReturnValueLocation(PTR_OBJECTREF * pLocation, GCRefKind * pKind);
+    bool             GetHijackedAsyncContinuation(PTR_OBJECTREF * pLocation);
 #endif
     void             SetControlPC(PTR_VOID controlPC);
     PTR_VOID         GetControlPC() { return m_ControlPC; }
@@ -240,6 +241,7 @@ class StackFrameIterator
 #ifdef TARGET_X86
     PTR_OBJECTREF       m_pHijackedReturnValue;
     GCRefKind           m_HijackedReturnValueKind;
+    PTR_OBJECTREF       m_pHijackedAsyncContinuation;
 #endif
     PTR_uintptr_t       m_pConservativeStackRangeLowerBound;
     PTR_uintptr_t       m_pConservativeStackRangeUpperBound;

src/coreclr/nativeaot/Runtime/i386/AsmOffsetsCpu.h

@@ -7,28 +7,28 @@
 //
 // NOTE: the offsets MUST be in hex notation WITHOUT the 0x prefix
 
-PLAT_ASM_SIZEOF(c4, ExInfo)
+PLAT_ASM_SIZEOF(c8, ExInfo)
 PLAT_ASM_OFFSET(0, ExInfo, m_pPrevExInfo)
 PLAT_ASM_OFFSET(4, ExInfo, m_pExContext)
 PLAT_ASM_OFFSET(8, ExInfo, m_exception)
 PLAT_ASM_OFFSET(0c, ExInfo, m_kind)
 PLAT_ASM_OFFSET(0d, ExInfo, m_passNumber)
 PLAT_ASM_OFFSET(10, ExInfo, m_idxCurClause)
 PLAT_ASM_OFFSET(14, ExInfo, m_frameIter)
-PLAT_ASM_OFFSET(c0, ExInfo, m_notifyDebuggerSP)
+PLAT_ASM_OFFSET(c4, ExInfo, m_notifyDebuggerSP)
 
 PLAT_ASM_OFFSET(0, PInvokeTransitionFrame, m_RIP)
 PLAT_ASM_OFFSET(4, PInvokeTransitionFrame, m_FramePointer)
 PLAT_ASM_OFFSET(8, PInvokeTransitionFrame, m_pThread)
 PLAT_ASM_OFFSET(0c, PInvokeTransitionFrame, m_Flags)
 PLAT_ASM_OFFSET(10, PInvokeTransitionFrame, m_PreservedRegs)
 
-PLAT_ASM_SIZEOF(ac, StackFrameIterator)
+PLAT_ASM_SIZEOF(b0, StackFrameIterator)
 PLAT_ASM_OFFSET(08, StackFrameIterator, m_FramePointer)
 PLAT_ASM_OFFSET(0c, StackFrameIterator, m_ControlPC)
 PLAT_ASM_OFFSET(10, StackFrameIterator, m_RegDisplay)
-PLAT_ASM_OFFSET(a4, StackFrameIterator, m_OriginalControlPC)
-PLAT_ASM_OFFSET(a8, StackFrameIterator, m_pPreviousTransitionFrame)
+PLAT_ASM_OFFSET(a8, StackFrameIterator, m_OriginalControlPC)
+PLAT_ASM_OFFSET(ac, StackFrameIterator, m_pPreviousTransitionFrame)
 
 PLAT_ASM_SIZEOF(1c, PAL_LIMITED_CONTEXT)
 PLAT_ASM_OFFSET(0, PAL_LIMITED_CONTEXT, IP)

src/coreclr/nativeaot/Runtime/inc/ModuleHeaders.h

@@ -11,7 +11,7 @@ struct ReadyToRunHeaderConstants
 {
     static const uint32_t Signature = 0x00525452; // 'RTR'
 
-    static const uint32_t CurrentMajorVersion = 20;
+    static const uint32_t CurrentMajorVersion = 21;
     static const uint32_t CurrentMinorVersion = 0;
 };
 

src/coreclr/nativeaot/Runtime/inc/rhbinder.h

@@ -240,6 +240,7 @@ enum PInvokeTransitionFrameFlags
 #if defined(TARGET_X86)
     PTFF_RAX_IS_GCREF   = 0x00010000,   // used by hijack handler to report return value of hijacked method
     PTFF_RAX_IS_BYREF   = 0x00020000,
+    PTFF_RCX_IS_GCREF   = 0x00040000,
 #endif
 
     PTFF_THREAD_HIJACK  = 0x00100000,   // indicates that this is a frame for a hijacked call

src/coreclr/nativeaot/Runtime/thread.cpp

@@ -481,6 +481,12 @@ void Thread::GcScanRootsWorker(ScanFunc * pfnEnumCallback, ScanContext * pvCallb
             EnumGcRef(pHijackedReturnValue, returnKind, pfnEnumCallback, pvCallbackData);
         }
     }
+
+    PTR_OBJECTREF pHijackedAsyncContinuation = NULL;
+    if (frameIterator.GetHijackedAsyncContinuation(&pHijackedAsyncContinuation))
+    {
+        EnumGcRef(pHijackedAsyncContinuation, GCRK_Object, pfnEnumCallback, pvCallbackData);
+    }
 #endif
 
 #ifndef DACCESS_COMPILE
@@ -817,9 +823,9 @@ void Thread::HijackReturnAddressWorker(StackFrameIterator* frameIterator, Hijack
         m_ppvHijackedReturnAddressLocation = ppvRetAddrLocation;
         m_pvHijackedReturnAddress = pvRetAddr;
 #if defined(TARGET_X86)
-        m_uHijackedReturnValueFlags = ReturnKindToTransitionFrameFlags(
-            frameIterator->GetCodeManager()->GetReturnValueKind(frameIterator->GetMethodInfo(),
-                                                                frameIterator->GetRegisterSet()));
+        bool isAsync;
+        GCRefKind retKind = frameIterator->GetCodeManager()->GetReturnValueKind(frameIterator->GetMethodInfo(), frameIterator->GetRegisterSet(), &isAsync);
+        m_uHijackedReturnValueFlags = ReturnKindToTransitionFrameFlags(retKind, isAsync);
 #endif
 
         *ppvRetAddrLocation = (void*)pfnHijackFunction;

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

@@ -981,14 +981,15 @@ bool CoffNativeCodeManager::GetReturnAddressHijackInfo(MethodInfo *    pMethodIn
 }
 
 #ifdef TARGET_X86
-GCRefKind CoffNativeCodeManager::GetReturnValueKind(MethodInfo *   pMethodInfo, REGDISPLAY *   pRegisterSet)
+GCRefKind CoffNativeCodeManager::GetReturnValueKind(MethodInfo *   pMethodInfo, REGDISPLAY *   pRegisterSet, bool* isAsync)
 {
     PTR_uint8_t gcInfo;
     uint32_t codeOffset = GetCodeOffset(pMethodInfo, (PTR_VOID)pRegisterSet->IP, &gcInfo);
     hdrInfo infoBuf;
     size_t infoSize = DecodeGCHdrInfo(GCInfoToken(gcInfo), codeOffset, &infoBuf);
 
     ASSERT(infoBuf.returnKind != RT_Float); // See TODO above
+    *isAsync = infoBuf.isAsync;
     return (GCRefKind)infoBuf.returnKind;
 }
 #endif

src/coreclr/nativeaot/Runtime/windows/CoffNativeCodeManager.h

@@ -96,7 +96,8 @@ class CoffNativeCodeManager : public ICodeManager
 
 #ifdef TARGET_X86
     GCRefKind GetReturnValueKind(MethodInfo *   pMethodInfo,
-                                    REGDISPLAY *   pRegisterSet);
+                                 REGDISPLAY *   pRegisterSet,
+                                 bool* isAsync);
 #endif
 
     PTR_VOID RemapHardwareFaultToGCSafePoint(MethodInfo * pMethodInfo, PTR_VOID controlPC);

src/coreclr/tools/Common/Internal/Runtime/ModuleHeaders.cs

@@ -15,7 +15,7 @@ internal struct ReadyToRunHeaderConstants
     {
         public const uint Signature = 0x00525452; // 'RTR'
 
-        public const ushort CurrentMajorVersion = 20;
+        public const ushort CurrentMajorVersion = 21;
         public const ushort CurrentMinorVersion = 0;
     }
 #if READYTORUN