From 97a28d443162c8fe0fd08fa27a1e61dfc2c26a5c Mon Sep 17 00:00:00 2001
From: Wojtek Figat <wojtek@figat.pl>
Date: Mon, 13 Nov 2023 15:48:12 +0100
Subject: [PATCH] Add security lockers for managed typeinfo access

---
 Source/Engine/Scripting/BinaryModule.cpp   |  1 +
 Source/Engine/Scripting/Runtime/DotNet.cpp | 26 +++++++++++++++++++---
 Source/Engine/Scripting/Scripting.cpp      |  1 +
 3 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/Source/Engine/Scripting/BinaryModule.cpp b/Source/Engine/Scripting/BinaryModule.cpp
index 679ead4e3..5bbaf8d0b 100644
--- a/Source/Engine/Scripting/BinaryModule.cpp
+++ b/Source/Engine/Scripting/BinaryModule.cpp
@@ -906,6 +906,7 @@ void ManagedBinaryModule::OnLoaded(MAssembly* assembly)
 #if !COMPILE_WITHOUT_CSHARP
     PROFILE_CPU();
     ASSERT(ClassToTypeIndex.IsEmpty());
+    ScopeLock lock(Locker);
 
     const auto& classes = assembly->GetClasses();
 
diff --git a/Source/Engine/Scripting/Runtime/DotNet.cpp b/Source/Engine/Scripting/Runtime/DotNet.cpp
index 200efa650..c4371736d 100644
--- a/Source/Engine/Scripting/Runtime/DotNet.cpp
+++ b/Source/Engine/Scripting/Runtime/DotNet.cpp
@@ -719,6 +719,7 @@ void GetAssemblyName(void* assemblyHandle, StringAnsi& name, StringAnsi& fullnam
 
 DEFINE_INTERNAL_CALL(void) NativeInterop_CreateClass(NativeClassDefinitions* managedClass, void* assemblyHandle)
 {
+    ScopeLock lock(BinaryModule::Locker);
     MAssembly* assembly = GetAssembly(assemblyHandle);
     if (assembly == nullptr)
     {
@@ -732,7 +733,18 @@ DEFINE_INTERNAL_CALL(void) NativeInterop_CreateClass(NativeClassDefinitions* man
     MClass* klass = New<MClass>(assembly, managedClass->typeHandle, managedClass->name, managedClass->fullname, managedClass->namespace_, managedClass->typeAttributes);
     if (assembly != nullptr)
     {
-        const_cast<MAssembly::ClassesDictionary&>(assembly->GetClasses()).Add(klass->GetFullName(), klass);
+        auto& classes = const_cast<MAssembly::ClassesDictionary&>(assembly->GetClasses());
+        MClass* oldKlass;
+        if (classes.TryGet(klass->GetFullName(), oldKlass))
+        {
+            LOG(Warning, "Class '{0}' was already added to assembly '{1}'", String(klass->GetFullName()), String(assembly->GetName()));
+            Delete(klass);
+            klass = oldKlass;
+        }
+        else
+        {
+            classes.Add(klass->GetFullName(), klass);
+        }
     }
     managedClass->nativePointer = klass;
 }
@@ -873,7 +885,7 @@ MClass::MClass(const MAssembly* parentAssembly, void* handle, const char* name,
     static void* TypeIsEnumPtr = GetStaticMethodPointer(TEXT("TypeIsEnum"));
     _isEnum = CallStaticMethod<bool, void*>(TypeIsEnumPtr, handle);
 
-    CachedClassHandles.Add(handle, this);
+    CachedClassHandles[handle] = this;
 }
 
 bool MAssembly::ResolveMissingFile(String& assemblyPath) const
@@ -1551,6 +1563,7 @@ const Array<MObject*>& MProperty::GetAttributes() const
 
 MAssembly* GetAssembly(void* assemblyHandle)
 {
+    ScopeLock lock(BinaryModule::Locker);
     MAssembly* assembly;
     if (CachedAssemblyHandles.TryGet(assemblyHandle, assembly))
         return assembly;
@@ -1559,6 +1572,7 @@ MAssembly* GetAssembly(void* assemblyHandle)
 
 MClass* GetClass(MType* typeHandle)
 {
+    ScopeLock lock(BinaryModule::Locker);
     MClass* klass = nullptr;
     CachedClassHandles.TryGet(typeHandle, klass);
     return nullptr;
@@ -1568,6 +1582,7 @@ MClass* GetOrCreateClass(MType* typeHandle)
 {
     if (!typeHandle)
         return nullptr;
+    ScopeLock lock(BinaryModule::Locker);
     MClass* klass;
     if (!CachedClassHandles.TryGet(typeHandle, klass))
     {
@@ -1579,7 +1594,12 @@ MClass* GetOrCreateClass(MType* typeHandle)
         klass = New<MClass>(assembly, classInfo.typeHandle, classInfo.name, classInfo.fullname, classInfo.namespace_, classInfo.typeAttributes);
         if (assembly != nullptr)
         {
-            const_cast<MAssembly::ClassesDictionary&>(assembly->GetClasses()).Add(klass->GetFullName(), klass);
+            auto& classes = const_cast<MAssembly::ClassesDictionary&>(assembly->GetClasses());
+            if (classes.ContainsKey(klass->GetFullName()))
+            {
+                LOG(Warning, "Class '{0}' was already added to assembly '{1}'", String(klass->GetFullName()), String(assembly->GetName()));
+            }
+            classes[klass->GetFullName()] = klass;
         }
 
         if (typeHandle != classInfo.typeHandle)
diff --git a/Source/Engine/Scripting/Scripting.cpp b/Source/Engine/Scripting/Scripting.cpp
index bdf9f60f5..e9dc0421c 100644
--- a/Source/Engine/Scripting/Scripting.cpp
+++ b/Source/Engine/Scripting/Scripting.cpp
@@ -436,6 +436,7 @@ bool Scripting::Load()
     PROFILE_CPU();
     // Note: this action can be called from main thread (due to Mono problems with assemblies actions from other threads)
     ASSERT(IsInMainThread());
+    ScopeLock lock(BinaryModule::Locker);
 
 #if USE_CSHARP
     // Load C# core assembly