Capture the hidden argument in TheUMEntryPrestubWorker for interpreted reverse p/invoke stubs and ensure that the worker runs on every reverse p/invoke, so that the interpreter can recover the hidden argument and make it available via GetStubContext.

Add INTOP_GETSTUBCONTEXT and always generate it for the GetStubContext intrinsic, even if other intrinsics are disabled

Author: kg

src/coreclr/interpreter/compiler.cpp

@@ -2306,6 +2306,14 @@ bool InterpCompiler::EmitNamedIntrinsicCall(NamedIntrinsic ni, CORINFO_CLASS_HAN
             AddIns(INTOP_THROW_PNSE);
             return true;
 
+        case NI_System_StubHelpers_GetStubContext:
+        {
+            AddIns(INTOP_GETSTUBCONTEXT);
+            PushStackType(StackTypeI, NULL);
+            m_pLastNewIns->SetDVar(m_pStackPointer[-1].var);
+            return true;
+        }
+
         case NI_System_Runtime_CompilerServices_StaticsHelpers_VolatileReadAsByref:
         {
             CHECK_STACK(1);
@@ -2817,7 +2825,7 @@ void InterpCompiler::EmitCall(CORINFO_RESOLVED_TOKEN* pConstrainedToken, bool re
             //  intrinsics for the recursive call. Otherwise we will just recurse infinitely and overflow stack.
             //  This expansion can produce value that is inconsistent with the value seen by JIT/R2R code that can
             //  cause user code to misbehave. This is by design. One-off method Interpretation is for internal use only.
-            bool isMustExpand = (callInfo.hMethod == m_methodHnd);
+            bool isMustExpand = (callInfo.hMethod == m_methodHnd) || (ni == NI_System_StubHelpers_GetStubContext);
             if ((InterpConfig.InterpMode() == 3) || isMustExpand)
             {
                 if (EmitNamedIntrinsicCall(ni, resolvedCallToken.hClass, callInfo.hMethod, callInfo.sig))

src/coreclr/interpreter/intops.def

@@ -415,6 +415,7 @@ OPDEF(INTOP_LOAD_FRAMEVAR, "load.framevar", 2, 1, 0, InterpOpNoArgs)
 // Intrinsics
 OPDEF(INTOP_COMPARE_EXCHANGE_I4, "compare.exchange.i4", 5, 1, 3, InterpOpNoArgs)
 OPDEF(INTOP_COMPARE_EXCHANGE_I8, "compare.exchange.i8", 5, 1, 3, InterpOpNoArgs)
+OPDEF(INTOP_GETSTUBCONTEXT, "getstubcontext", 2, 1, 0, InterpOpNoArgs)
 
 // All instructions after this point are IROPS, instructions that are not emitted/executed
 OPDEF(INTOP_NOP, "nop", 1, 0, 0, InterpOpNoArgs)

src/coreclr/interpreter/intrinsics.cpp

@@ -44,6 +44,8 @@ NamedIntrinsic GetNamedIntrinsic(COMP_HANDLE compHnd, CORINFO_METHOD_HANDLE comp
         {
             if (!strcmp(methodName, "NextCallReturnAddress"))
                 return NI_System_StubHelpers_NextCallReturnAddress;
+            else if (!strcmp(methodName, "GetStubContext"))
+                return NI_System_StubHelpers_GetStubContext;
         }
     }
     else if (!strcmp(namespaceName, "System.Numerics"))

src/coreclr/vm/dllimportcallback.cpp

@@ -210,6 +210,15 @@ VOID CallbackOnCollectedDelegate(UMEntryThunkData* pEntryThunkData)
     EEPOLICY_HANDLE_FATAL_ERROR_WITH_MESSAGE(COR_E_FAILFAST, message.GetUnicode());
 }
 
+PLATFORM_THREAD_LOCAL UMEntryThunkData * t_MostRecentUMEntryThunkData;
+
+UMEntryThunkData * GetMostRecentUMEntryThunkData()
+{
+    UMEntryThunkData * result = t_MostRecentUMEntryThunkData;
+    t_MostRecentUMEntryThunkData = nullptr;
+    return result;
+}
+
 PCODE TheUMEntryPrestubWorker(UMEntryThunkData * pUMEntryThunkData)
 {
     STATIC_CONTRACT_THROWS;
@@ -229,18 +238,23 @@ PCODE TheUMEntryPrestubWorker(UMEntryThunkData * pUMEntryThunkData)
     if (pUMEntryThunkData->IsCollectedDelegate())
         CallbackOnCollectedDelegate(pUMEntryThunkData);
 
+    PCODE result = 0;
+
     INSTALL_MANAGED_EXCEPTION_DISPATCHER;
     // this method is called by stubs which are called by managed code,
     // so we need an unwind and continue handler so that our internal
     // exceptions don't leak out into managed code.
     INSTALL_UNWIND_AND_CONTINUE_HANDLER;
 
-    pUMEntryThunkData->RunTimeInit();
+    // HACK: Stash the entry thunk data address so that the interpreter can recover it.
+    // The InterpreterStub overwrites the register that normally would contain this address.
+    t_MostRecentUMEntryThunkData = pUMEntryThunkData;
+    result = pUMEntryThunkData->RunTimeInit();
 
     UNINSTALL_UNWIND_AND_CONTINUE_HANDLER;
     UNINSTALL_MANAGED_EXCEPTION_DISPATCHER;
 
-    return (PCODE)pUMEntryThunkData->GetCode();
+    return result;
 }
 
 UMEntryThunkData* UMEntryThunkData::CreateUMEntryThunk()
@@ -263,7 +277,7 @@ UMEntryThunkData* UMEntryThunkData::CreateUMEntryThunk()
         LoaderAllocator *pLoaderAllocator = SystemDomain::GetGlobalLoaderAllocator();
         AllocMemTracker amTracker;
         AllocMemTracker *pamTracker = &amTracker;
-       
+
         pData = (UMEntryThunkData *)pamTracker->Track(pLoaderAllocator->GetLowFrequencyHeap()->AllocMem(S_SIZE_T(sizeof(UMEntryThunkData))));
         UMEntryThunk* pThunk = (UMEntryThunk*)pamTracker->Track(pLoaderAllocator->GetNewStubPrecodeHeap()->AllocStub());
 #ifdef FEATURE_PERFMAP

src/coreclr/vm/dllimportcallback.h

@@ -134,7 +134,7 @@ class UMEntryThunk : private StubPrecode
     PTR_UMEntryThunkData GetData() const
     {
         LIMITED_METHOD_CONTRACT;
-        
+
         return dac_cast<PTR_UMEntryThunkData>(GetSecretParam());
     }
 };
@@ -214,7 +214,7 @@ class UMEntryThunkData
 
     void Terminate();
 
-    VOID RunTimeInit()
+    PCODE RunTimeInit()
     {
         STANDARD_VM_CONTRACT;
 
@@ -227,11 +227,29 @@ class UMEntryThunkData
         if (m_pObjectHandle == NULL && m_pManagedTarget == (TADDR)0)
             m_pManagedTarget = m_pMD->GetMultiCallableAddrOfCode();
 
-        m_pUMEntryThunk->SetTargetUnconditional(m_pUMThunkMarshInfo->GetExecStubEntryPoint());
+#ifdef FEATURE_INTERPRETER
+        // HACK: For interpreted stubs we need to ensure that TheUMEntryPrestubWorker runs for every
+        //  unmanaged-to-managed invocation in order to populate the TLS variable every time.
+        PCODE pILStub = m_pUMThunkMarshInfo->GetExecStubEntryPoint();
+        Precode* pStubPrecode = (Precode*)(void*)pILStub;
+        if (pStubPrecode->GetType() == PRECODE_INTERPRETER)
+        {
+        }
+        else
+#endif
+        {
+            m_pUMEntryThunk->SetTargetUnconditional(pILStub);
+        }
 
 #ifdef _DEBUG
         m_state = kRunTimeInited;
 #endif // _DEBUG
+
+#ifdef FEATURE_INTERPRETER
+        if (pStubPrecode->GetType() == PRECODE_INTERPRETER)
+            return pILStub;
+#endif
+        return GetCode();
     }
 
     PCODE GetManagedTarget() const

src/coreclr/vm/interpexec.cpp

@@ -576,6 +576,10 @@ void InterpExecMethod(InterpreterFrame *pInterpreterFrame, InterpMethodContextFr
                     MemoryBarrier();
                     ip++;
                     break;
+                case INTOP_GETSTUBCONTEXT:
+                    LOCAL_VAR(ip[1], void*) = pFrame->hiddenArgument;
+                    ip += 2;
+                    break;
                 case INTOP_LDC_I4:
                     LOCAL_VAR(ip[1], int32_t) = ip[2];
                     ip += 3;

src/coreclr/vm/interpexec.h

@@ -35,6 +35,7 @@ struct InterpMethodContextFrame
     int8_t *pRetVal;
     const int32_t *ip; // This ip is updated only when execution can leave the frame
     PTR_InterpMethodContextFrame pNext;
+    void *hiddenArgument;
 
 #ifndef DACCESS_COMPILE
     void ReInit(InterpMethodContextFrame *pParent, InterpByteCodeStart* startIp, int8_t *pRetVal, int8_t *pStack)

src/coreclr/vm/prestub.cpp

@@ -2016,6 +2016,8 @@ static InterpThreadContext* GetInterpThreadContext()
 
 EXTERN_C void STDCALL ReversePInvokeBadTransition();
 
+UMEntryThunkData * GetMostRecentUMEntryThunkData();
+
 extern "C" void* STDCALL ExecuteInterpretedMethod(TransitionBlock* pTransitionBlock, TADDR byteCodeAddr, void* retBuff)
 {
     // Argument registers are in the TransitionBlock
@@ -2055,6 +2057,7 @@ extern "C" void* STDCALL ExecuteInterpretedMethod(TransitionBlock* pTransitionBl
     frames.interpMethodContextFrame.startIp = dac_cast<PTR_InterpByteCodeStart>(byteCodeAddr);
     frames.interpMethodContextFrame.pStack = sp;
     frames.interpMethodContextFrame.pRetVal = (retBuff != NULL) ? (int8_t*)retBuff : sp;
+    frames.interpMethodContextFrame.hiddenArgument = GetMostRecentUMEntryThunkData();
 
     InterpExecMethod(&frames.interpreterFrame, &frames.interpMethodContextFrame, threadContext);