harden the contexts initialization path

Assertion on server side for client calling with wrong args or at wrong state to help diagnose the issue

lib/Backend/NativeCodeGenerator.cpp

@@ -2007,12 +2007,6 @@ NativeCodeGenerator::UpdateJITState()
 {
     if (JITManager::GetJITManager()->IsOOPJITEnabled())
     {
-        // ensure jit contexts have been set up
-        if (!scriptContext->GetRemoteScriptAddr())
-        {
-            scriptContext->InitializeRemoteScriptContext();
-        }
-
         // update all property records on server that have been changed since last jit
         ThreadContext::PropertyMap * pendingProps = scriptContext->GetThreadContext()->GetPendingJITProperties();
         PropertyRecordIDL ** newPropArray = nullptr;

lib/Backend/ServerThreadContext.cpp

@@ -22,6 +22,8 @@ ServerThreadContext::ServerThreadContext(ThreadContextDataIDL * data) :
 #endif
     m_jitCRTBaseAddress((intptr_t)GetModuleHandle(UCrtC99MathApis::LibraryName))
 {
+    m_pid = GetProcessId((HANDLE)data->processHandle);
+
 #if !_M_X64_OR_ARM64 && _CONTROL_FLOW_GUARD
     m_codeGenAlloc.canCreatePreReservedSegment = data->allowPrereserveAlloc != FALSE;
 #endif

lib/Backend/ServerThreadContext.h

@@ -65,6 +65,8 @@ private:
     CodeGenAllocators m_codeGenAlloc;
 
     ThreadContextDataIDL m_threadContextData;
+
+    DWORD m_pid; //save client process id for easier diagnose
     
     intptr_t m_jitChakraBaseAddress;
     intptr_t m_jitCRTBaseAddress;

lib/JITServer/JITServer.cpp

@@ -159,19 +159,16 @@ ServerCleanupThreadContext(
     ServerThreadContext * threadContextInfo = (ServerThreadContext*)DecodePointer((void*)threadContextRoot);
     if (threadContextInfo == nullptr)
     {
+        Assert(false);
         return RPC_S_INVALID_ARG;
     }
 
-    if (!ServerContextManager::IsThreadContextAlive(threadContextInfo))
+    if (ServerContextManager::IsThreadContextAlive(threadContextInfo))
     {
-        return E_ACCESSDENIED;
+        AutoReleaseContext<ServerThreadContext> autoThreadContext(threadContextInfo);
+        threadContextInfo->Close();
+        ServerContextManager::UnRegisterThreadContext(threadContextInfo);
     }
-
-    AutoReleaseContext<ServerThreadContext> autoThreadContext(threadContextInfo);
-
-    threadContextInfo->Close();
-    ServerContextManager::UnRegisterThreadContext(threadContextInfo);
-
     return S_OK;
 }
 
@@ -187,11 +184,13 @@ ServerUpdatePropertyRecordMap(
 
     if (threadContextInfo == nullptr)
     {
+        Assert(false);
         return RPC_S_INVALID_ARG;
     }
 
     if (!ServerContextManager::IsThreadContextAlive(threadContextInfo))
     {
+        Assert(false);
         return E_ACCESSDENIED;
     }
 
@@ -225,12 +224,14 @@ ServerAddDOMFastPathHelper(
 
     if (scriptContextInfo == nullptr)
     {
+        Assert(false);
         return RPC_S_INVALID_ARG;
     }
 
     if (!ServerContextManager::IsScriptContextAlive(scriptContextInfo))
     {
-        return RPC_S_INVALID_ARG;
+        Assert(false);
+        return E_ACCESSDENIED;
     }
 
     AutoReleaseContext<ServerScriptContext> autoScriptContext(scriptContextInfo);
@@ -253,12 +254,14 @@ ServerAddModuleRecordInfo(
     ServerScriptContext * serverScriptContext = (ServerScriptContext*)DecodePointer((void*)scriptContextInfoAddress);
     if (serverScriptContext == nullptr)
     {
+        Assert(false);
         return RPC_S_INVALID_ARG;
     }
 
     if (!ServerContextManager::IsScriptContextAlive(serverScriptContext))
     {
-        return RPC_S_INVALID_ARG;
+        Assert(false);
+        return E_ACCESSDENIED;
     }
 
     AutoReleaseContext<ServerScriptContext> autoScriptContext(serverScriptContext);
@@ -280,11 +283,13 @@ ServerSetWellKnownHostTypeId(
 
     if (threadContextInfo == nullptr)
     {
+        Assert(false);
         return RPC_S_INVALID_ARG;
     }
 
     if (!ServerContextManager::IsThreadContextAlive(threadContextInfo))
     {
+        Assert(false);
         return E_ACCESSDENIED;
     }
 
@@ -307,11 +312,13 @@ ServerInitializeScriptContext(
 
     if (threadContextInfo == nullptr)
     {
-        return E_ACCESSDENIED;
+        Assert(false);
+        return RPC_S_INVALID_ARG;
     }
 
     if (!ServerContextManager::IsThreadContextAlive(threadContextInfo))
     {
+        Assert(false);
         return E_ACCESSDENIED;
     }
 
@@ -340,26 +347,26 @@ ServerCloseScriptContext(
 
     if (scriptContextInfo == nullptr)
     {
+        Assert(false);
         return RPC_S_INVALID_ARG;
     }
 
-    if (!ServerContextManager::IsScriptContextAlive(scriptContextInfo))
+    if (ServerContextManager::IsScriptContextAlive(scriptContextInfo))
     {
-        return E_ACCESSDENIED;
-    }
-
-    AutoReleaseContext<ServerScriptContext> autoScriptContext(scriptContextInfo);
+        AutoReleaseContext<ServerScriptContext> autoScriptContext(scriptContextInfo);
 
 #ifdef PROFILE_EXEC
-    auto profiler = scriptContextInfo->GetCodeGenProfiler();
-    if (profiler && profiler->IsInitialized())
-    {
-        profiler->ProfilePrint(Js::Configuration::Global.flags.Profile.GetFirstPhase());
-    }
+        auto profiler = scriptContextInfo->GetCodeGenProfiler();
+        if (profiler && profiler->IsInitialized())
+        {
+            profiler->ProfilePrint(Js::Configuration::Global.flags.Profile.GetFirstPhase());
+        }
 #endif
 
-    scriptContextInfo->Close();
-    ServerContextManager::UnRegisterScriptContext(scriptContextInfo);
+        scriptContextInfo->Close();
+        ServerContextManager::UnRegisterScriptContext(scriptContextInfo);
+    }
+
     return S_OK;
 }
 
@@ -372,6 +379,7 @@ ServerCleanupScriptContext(
 
     if (scriptContextInfo == nullptr)
     {
+        Assert(false);
         return RPC_S_INVALID_ARG;
     }
 
@@ -396,20 +404,18 @@ ServerFreeAllocation(
 
     if (context == nullptr)
     {
+        Assert(false);
         return RPC_S_INVALID_ARG;
     }
 
-    if (!ServerContextManager::IsThreadContextAlive(context))
-    {
-        return E_ACCESSDENIED;
-    }
-    AutoReleaseContext<ServerThreadContext> autoThreadContext(context);
-    return ServerCallWrapper(context, [&]()->HRESULT
+    if (ServerContextManager::IsThreadContextAlive(context))
     {
+        AutoReleaseContext<ServerThreadContext> autoThreadContext(context);
         context->SetValidCallTargetForCFG((PVOID)address, false);
-        bool succeeded = context->GetCodeGenAllocators()->emitBufferManager.FreeAllocation((void*)address);
-        return succeeded ? S_OK : E_FAIL;
-    });
+        context->GetCodeGenAllocators()->emitBufferManager.FreeAllocation((void*)address);
+    }
+ 
+    return S_OK;
 }
 
 HRESULT
@@ -424,12 +430,14 @@ ServerIsNativeAddr(
     if (context == nullptr)
     {
         *result = false;
+        Assert(false);
         return RPC_S_INVALID_ARG;
     }
 
     if (!ServerContextManager::IsThreadContextAlive(context))
     {
         *result = false;
+        Assert(false);
         return E_ACCESSDENIED;
     }
 
@@ -465,21 +473,19 @@ ServerSetIsPRNGSeeded(
 
     if (scriptContextInfo == nullptr)
     {
+        Assert(false);
         return RPC_S_INVALID_ARG;
     }
 
     if (!ServerContextManager::IsScriptContextAlive(scriptContextInfo))
     {
-        return RPC_S_INVALID_ARG;
+        Assert(false);
+        return E_ACCESSDENIED;
     }
 
-    AutoReleaseContext<ServerScriptContext> autoScriptContext(scriptContextInfo);    
-
-    return ServerCallWrapper(scriptContextInfo, [&]()->HRESULT
-    {
-        scriptContextInfo->SetIsPRNGSeeded(value != FALSE);
-        return S_OK;
-    });
+    AutoReleaseContext<ServerScriptContext> autoScriptContext(scriptContextInfo);
+    scriptContextInfo->SetIsPRNGSeeded(value != FALSE);
+    return S_OK;
 }
 
 HRESULT
@@ -501,11 +507,13 @@ ServerRemoteCodeGen(
 
     if (scriptContextInfo == nullptr)
     {
+        Assert(false);
         return RPC_S_INVALID_ARG;
     }
 
     if (!ServerContextManager::IsScriptContextAlive(scriptContextInfo))
     {
+        Assert(false);
         return E_ACCESSDENIED;
     }
 
@@ -686,7 +694,7 @@ HRESULT ServerCallWrapper(ServerThreadContext* threadContextInfo, Fn fn)
         // like VM is in a state of removing but not completed yet
         hr = HRESULT_FROM_WIN32(ex.LastError);
 
-        DWORD exitCode;
+        DWORD exitCode = STILL_ACTIVE;
         if (!GetExitCodeProcess(threadContextInfo->GetProcessHandle(), &exitCode))
         {
             Assert(false); // fail to check target process
@@ -696,6 +704,12 @@ HRESULT ServerCallWrapper(ServerThreadContext* threadContextInfo, Fn fn)
         {
             threadContextInfo->Close();
             ServerContextManager::UnRegisterThreadContext(threadContextInfo);
+            return hr;
+        }
+        else
+        {
+            // The content process is still alive, the falure most likely an OOM
+            return E_OUTOFMEMORY;
         }
     }
 
@@ -708,6 +722,7 @@ HRESULT ServerCallWrapper(ServerScriptContext* scriptContextInfo, Fn fn)
     ServerThreadContext* threadContextInfo = scriptContextInfo->GetThreadContext();
     if (!ServerContextManager::IsThreadContextAlive(threadContextInfo))
     {
+        Assert(false);
         return E_ACCESSDENIED;
     }
     AutoReleaseContext<ServerThreadContext> autoThreadContext(threadContextInfo);

lib/Runtime/Base/ScriptContext.cpp

@@ -4508,7 +4508,8 @@ void ScriptContext::RegisterPrototypeChainEnsuredToHaveOnlyWritableDataPropertie
 #endif
         this->GetThreadContext()->EnsureJITThreadContext(allowPrereserveAlloc);
 
-        JITManager::GetJITManager()->InitializeScriptContext(&contextData, this->GetThreadContext()->GetRemoteThreadContextAddr(), &m_remoteScriptContextAddr);
+        HRESULT hr = JITManager::GetJITManager()->InitializeScriptContext(&contextData, this->GetThreadContext()->GetRemoteThreadContextAddr(), &m_remoteScriptContextAddr);
+        JITManager::HandleServerCallResult(hr);
     }
 #endif
 

lib/Runtime/Base/ScriptContext.h

@@ -913,7 +913,14 @@ private:
         void SetDirectHostTypeId(TypeId typeId) {directHostTypeId = typeId; }
         TypeId GetDirectHostTypeId() const { return directHostTypeId; }
 
-        intptr_t GetRemoteScriptAddr() { return m_remoteScriptContextAddr; }
+        intptr_t GetRemoteScriptAddr() 
+        {
+            if (!m_remoteScriptContextAddr)
+            {
+                InitializeRemoteScriptContext();
+            }
+            return m_remoteScriptContextAddr; 
+        }
 
         char16 const * GetUrl() const { return url; }
         void SetUrl(BSTR bstr);

lib/Runtime/Base/ThreadContext.cpp

@@ -2008,7 +2008,8 @@ ThreadContext::EnsureJITThreadContext(bool allowPrereserveAlloc)
     m_reclaimedJITProperties = HeapNew(PropertyList, &HeapAllocator::Instance);
     m_pendingJITProperties = propertyMap->Clone();
 
-    JITManager::GetJITManager()->InitializeThreadContext(&contextData, &m_remoteThreadContextInfo, &m_prereservedRegionAddr);
+    HRESULT hr = JITManager::GetJITManager()->InitializeThreadContext(&contextData, &m_remoteThreadContextInfo, &m_prereservedRegionAddr);
+    JITManager::HandleServerCallResult(hr);
 }
 #endif
 

lib/Runtime/Base/ThreadContext.h

@@ -554,8 +554,9 @@ public:
     static void SetJITConnectionInfo(HANDLE processHandle, void* serverSecurityDescriptor, UUID connectionId);
     void EnsureJITThreadContext(bool allowPrereserveAlloc);
 
-    intptr_t GetRemoteThreadContextAddr() const
+    intptr_t GetRemoteThreadContextAddr()
     {
+        Assert(m_remoteThreadContextInfo);
         return m_remoteThreadContextInfo;
     }
 #endif

lib/Runtime/Language/SourceTextModuleRecord.cpp

@@ -841,10 +841,6 @@ namespace Js
 #if ENABLE_NATIVE_CODEGEN
             if (JITManager::GetJITManager()->IsOOPJITEnabled())
             {
-                if (!scriptContext->GetRemoteScriptAddr())
-                {
-                    scriptContext->InitializeRemoteScriptContext();
-                }
                 HRESULT hr = JITManager::GetJITManager()->AddModuleRecordInfo(
                     scriptContext->GetRemoteScriptAddr(),
                     this->GetModuleId(),

lib/Runtime/Library/JavascriptLibrary.cpp

@@ -6588,11 +6588,8 @@ namespace Js
 #if ENABLE_NATIVE_CODEGEN
         if (JITManager::GetJITManager()->IsOOPJITEnabled())
         {
-            if (!GetScriptContext()->GetRemoteScriptAddr())
-            {
-                GetScriptContext()->InitializeRemoteScriptContext();
-            }
-            JITManager::GetJITManager()->SetIsPRNGSeeded(GetScriptContext()->GetRemoteScriptAddr(), val);
+            HRESULT hr = JITManager::GetJITManager()->SetIsPRNGSeeded(GetScriptContext()->GetRemoteScriptAddr(), val);
+            JITManager::HandleServerCallResult(hr);
         }
 #endif
     }