Fix async continuation hijacking for x86 NativeAOT

src/coreclr/inc/eetwain.h

@@ -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
 /*

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/gcdecoder.cpp

@@ -207,7 +207,8 @@ PTR_CBYTE FASTCALL decodeHeader(PTR_CBYTE table, UINT32 version, InfoHdr* header
                 // encoding here always corresponds to codes in InfoHdrAdjust2 set
                 if (encoding <= SET_RET_KIND_MAX)
                 {
-                    header->returnKind = (ReturnKind)encoding;
+                    header->returnKind = (ReturnKind)encoding & 3;
+                    header->isAsync = (encoding & 4) != 0;
                 }
                 else if (encoding < FFFF_NOGCREGION_CNT)
                 {

src/coreclr/inc/gcinfotypes.h

@@ -305,7 +305,7 @@ 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 = 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 +358,8 @@ 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 = SET_RETURNKIND + SET_RET_KIND_MAX + 1,        // 0x08
+    FFFF_NOGCREGION_CNT = SET_NOGCREGIONS_CNT + SET_NOGCREGIONS_MAX + 1 // 0x0D There is a count (>SET_NOGCREGIONS_MAX) after the header encoding
 };
 
 #define HAS_UNTRACKED               ((unsigned int) -1)
@@ -410,6 +410,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/jit/gcencode.cpp

@@ -877,11 +877,12 @@ 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;
+        encoding          = header.returnKind | (header.isAsync ? 4 : 0);
         _ASSERTE(encoding <= SET_RET_KIND_MAX);
         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;
@@ -1572,6 +1573,7 @@ size_t GCInfo::gcInfoBlockHdrSave(
     _ASSERTE(!IsStructReturnKind(returnKind) && "Struct Return Kinds Unexpected for JIT32");
     _ASSERTE(((int)returnKind <= (int)SET_RET_KIND_MAX) && "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/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/vm/eetwain.cpp

@@ -1638,7 +1638,7 @@ size_t EECodeManager::GetFunctionSize(GCInfoToken gcInfoToken)
 *  returns true.
 *  If hijacking is not possible for some reason, it return false.
 */
-bool EECodeManager::GetReturnAddressHijackInfo(GCInfoToken gcInfoToken X86_ARG(ReturnKind * returnKind))
+bool EECodeManager::GetReturnAddressHijackInfo(GCInfoToken gcInfoToken X86_ARG(ReturnKind * returnKind) X86_ARG(bool* hasAsyncRet))
 {
     CONTRACTL{
         NOTHROW;
@@ -1658,6 +1658,7 @@ bool EECodeManager::GetReturnAddressHijackInfo(GCInfoToken gcInfoToken X86_ARG(R
     }
 
     *returnKind = info.returnKind;
+    *hasAsyncRet = info.isAsync;
     return true;
 #else // !USE_GC_INFO_DECODER
 

src/coreclr/vm/gc_unwind_x86.inl

@@ -219,6 +219,7 @@ size_t DecodeGCHdrInfo(GCInfoToken gcInfoToken,
     infoPtr->ebpFrame        = header.ebpFrame;
     infoPtr->interruptible   = header.interruptible;
     infoPtr->returnKind      = (ReturnKind) header.returnKind;
+    infoPtr->isAsync         = header.isAsync;
 
     infoPtr->prologSize      = header.prologSize;
     infoPtr->epilogSize      = header.epilogSize;

src/coreclr/vm/threadsuspend.cpp

@@ -4838,13 +4838,7 @@ static bool GetReturnAddressHijackInfo(EECodeInfo *pCodeInfo X86_ARG(ReturnKind
 {
     X86_ONLY(*hasAsyncRet = false);
     GCInfoToken gcInfoToken = pCodeInfo->GetGCInfoToken();
-    if (!pCodeInfo->GetCodeManager()->GetReturnAddressHijackInfo(gcInfoToken X86_ARG(returnKind)))
-        return false;
-
-    MethodDesc* pMD = pCodeInfo->GetMethodDesc();
-    X86_ONLY(*hasAsyncRet = pMD->IsAsyncMethod());
-
-    return true;
+    return pCodeInfo->GetCodeManager()->GetReturnAddressHijackInfo(gcInfoToken X86_ARG(returnKind) X86_ARG(hasAsyncRet));
 }
 
 #ifndef TARGET_UNIX