Add handle assignment validation to VM side.

2 files changed, 36 insertions, 9 deletions

diff --git a/src/vm/gchandleutilities.h b/src/vm/gchandleutilities.h
index f6d2d0c48d..41b6317084 100644
--- a/src/vm/gchandleutilities.h
+++ b/src/vm/gchandleutilities.h
@@ -30,7 +30,8 @@ private:
     GCHandleUtilities() = delete;
 };
 
-void ValidateHandleAndAppDomain(OBJECTHANDLE handle);
+void ValidateObjectAndAppDomain(OBJECTREF objRef, ADIndex appDomainIndex);
+void ValidateHandleAssignment(OBJECTHANDLE handle, OBJECTREF objRef);
 
 // Given a handle, returns an OBJECTREF for the object it refers to.
 inline OBJECTREF ObjectFromHandle(OBJECTHANDLE handle)
@@ -38,7 +39,10 @@ inline OBJECTREF ObjectFromHandle(OBJECTHANDLE handle)
     _ASSERTE(handle);
 
 #ifdef _DEBUG_IMPL
-    ValidateHandleAndAppDomain(handle);
+    DWORD context = (DWORD)GCHandleUtilities::GetGCHandleManager()->GetHandleContext(handle);
+    OBJECTREF objRef = ObjectToOBJECTREF(*(Object**)handle);
+
+    ValidateObjectAndAppDomain(objRef, ADIndex(context));
 #endif // _DEBUG_IMPL
 
     // Wrap the raw OBJECTREF and return it
@@ -173,16 +177,22 @@ inline OBJECTHANDLE CreateVariableHandle(IGCHandleStore* store, OBJECTREF object
 
 inline void StoreObjectInHandle(OBJECTHANDLE handle, OBJECTREF object)
 {
+    ValidateHandleAssignment(handle, object);
+
     GCHandleUtilities::GetGCHandleManager()->StoreObjectInHandle(handle, OBJECTREFToObject(object));
 }
 
 inline bool StoreFirstObjectInHandle(OBJECTHANDLE handle, OBJECTREF object)
 {
+    ValidateHandleAssignment(handle, object);
+
     return GCHandleUtilities::GetGCHandleManager()->StoreObjectInHandleIfNull(handle, OBJECTREFToObject(object));
 }
 
 inline void* InterlockedCompareExchangeObjectInHandle(OBJECTHANDLE handle, OBJECTREF object, OBJECTREF comparandObject)
 {
+    ValidateHandleAssignment(handle, object);
+
     return GCHandleUtilities::GetGCHandleManager()->CompareAndSwapObjectInHandle(handle, OBJECTREFToObject(object), OBJECTREFToObject(comparandObject));
 }
 
diff --git a/src/vm/gcheaputilities.cpp b/src/vm/gcheaputilities.cpp
index 4bce7f5592..cd7afede70 100644
--- a/src/vm/gcheaputilities.cpp
+++ b/src/vm/gcheaputilities.cpp
@@ -40,17 +40,12 @@ bool g_sw_ww_enabled_for_gc_heap = false;
 gc_alloc_context g_global_alloc_context = {};
 
 // Debug-only validation for handle.
-void ValidateHandleAndAppDomain(OBJECTHANDLE handle)
+
+void ValidateObjectAndAppDomain(OBJECTREF objRef, ADIndex appDomainIndex)
 {
 #ifdef _DEBUG_IMPL
-    OBJECTREF objRef = ObjectToOBJECTREF(*(Object**)handle);
     VALIDATEOBJECTREF(objRef);
 
-    IGCHandleManager *pHandleManager = GCHandleUtilities::GetGCHandleManager();
-
-    DWORD context = (DWORD)pHandleManager->GetHandleContext(handle);
-
-    ADIndex appDomainIndex = ADIndex(context);
     AppDomain *domain = SystemDomain::GetAppDomainAtIndex(appDomainIndex);
 
     // Access to a handle in an unloaded domain is not allowed
@@ -72,3 +67,25 @@ void ValidateHandleAndAppDomain(OBJECTHANDLE handle)
 #endif // CHECK_APP_DOMAIN_LEAKS
 #endif // _DEBUG_IMPL
 }
+
+void ValidateHandleAssignment(OBJECTHANDLE handle, OBJECTREF objRef)
+{
+#ifdef _DEBUG_IMPL
+    _ASSERTE(handle);
+
+#ifdef DEBUG_DestroyedHandleValue
+    // Verify that we are not trying to access a freed handle.
+    _ASSERTE("Attempt to access destroyed handle." && *(_UNCHECKED_OBJECTREF*)handle != DEBUG_DestroyedHandleValue);
+#endif
+
+    ADIndex appDomainIndex = HndGetHandleADIndex(handle);
+
+    AppDomain *unloadingDomain = SystemDomain::AppDomainBeingUnloaded();
+    if (unloadingDomain && unloadingDomain->GetIndex() == appDomainIndex && unloadingDomain->NoAccessToHandleTable())
+    {
+        _ASSERTE (!"Access to a handle in unloaded domain is not allowed");
+    }
+
+    ValidateObjectAndAppDomain(objRef, appDomainIndex);
+#endif // _DEBUG_IMPL
+}
\ No newline at end of file