From 177ca6b627eb0141895d26fb1ddc0c78e7963629 Mon Sep 17 00:00:00 2001
From: Pieter-Jan Briers <pieterjan.briers@gmail.com>
Date: Sat, 15 Jul 2023 18:09:57 +0200
Subject: [PATCH] Sandboxing performance improvements.

Don't leave file handles dangling.
Prefetch verifying assembly images to speed stuff up.
---
 RELEASE-NOTES.md                              |  1 +
 .../ContentPack/AssemblyTypeChecker.cs        | 28 ++++++++++++++++---
 Robust.Shared/ContentPack/ModLoader.cs        | 10 +++++--
 3 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md
index e95db920331..83a085bab24 100644
--- a/RELEASE-NOTES.md
+++ b/RELEASE-NOTES.md
@@ -50,6 +50,7 @@ END TEMPLATE-->
 
 * BQL `with` now includes paused entities.
 * The game loop now times more accurately and avoids sleeping more than necessary.
+* Sandboxing (and thus, client startup) should be much faster when ran from the launcher.
 
 ### Internal
 
diff --git a/Robust.Shared/ContentPack/AssemblyTypeChecker.cs b/Robust.Shared/ContentPack/AssemblyTypeChecker.cs
index a16bd3a4617..d9715814661 100644
--- a/Robust.Shared/ContentPack/AssemblyTypeChecker.cs
+++ b/Robust.Shared/ContentPack/AssemblyTypeChecker.cs
@@ -59,7 +59,7 @@ public AssemblyTypeChecker(IResourceManager res, ISawmill sawmill)
             _config = Task.Run(() => LoadConfig(sawmill));
         }
 
-        private Resolver CreateResolver()
+        internal Resolver CreateResolver()
         {
             var dotnetDir = Path.GetDirectoryName(typeof(int).Assembly.Location)!;
             var ourPath = typeof(AssemblyTypeChecker).Assembly.Location;
@@ -100,6 +100,19 @@ private Resolver CreateResolver()
         /// <param name="assembly">Assembly to load.</param>
         /// <returns></returns>
         public bool CheckAssembly(Stream assembly)
+        {
+            using var resolver = CreateResolver();
+
+            return CheckAssembly(assembly, resolver);
+        }
+
+        /// <summary>
+        ///     Check the assembly for any illegal types. Any types not on the white list
+        ///     will cause the assembly to be rejected.
+        /// </summary>
+        /// <param name="assembly">Assembly to load.</param>
+        /// <returns></returns>
+        public bool CheckAssembly(Stream assembly, Resolver resolver)
         {
             if (WouldNoOp)
             {
@@ -110,8 +123,7 @@ public bool CheckAssembly(Stream assembly)
             _sawmill.Debug("Checking assembly...");
             var fullStopwatch = Stopwatch.StartNew();
 
-            var resolver = CreateResolver();
-            using var peReader = ModLoader.MakePEReader(assembly, leaveOpen: true);
+            using var peReader = ModLoader.MakePEReader(assembly, leaveOpen: true, PEStreamOptions.PrefetchEntireImage);
             var reader = peReader.GetMetadataReader();
 
             var asmName = reader.GetString(reader.GetAssemblyDefinition().Name);
@@ -856,7 +868,7 @@ public bool CheckAssembly(string diskPath)
             return handle.IsNil ? null : reader.GetString(handle);
         }
 
-        private sealed class Resolver : IResolver
+        internal sealed class Resolver : IResolver, IDisposable
         {
             private readonly ConcurrentDictionary<string, PEReader?> _dictionary = new();
             private readonly AssemblyTypeChecker _parent;
@@ -910,6 +922,14 @@ public Resolver(AssemblyTypeChecker parent, string[] diskLoadPaths, ResPath[] re
             {
                 return _dictionary.GetOrAdd(simpleName, ResolveCore);
             }
+
+            public void Dispose()
+            {
+                foreach (var reader in _dictionary.Values)
+                {
+                    reader?.Dispose();
+                }
+            }
         }
 
         internal sealed class TypeProvider : ISignatureTypeProvider<MType, int>
diff --git a/Robust.Shared/ContentPack/ModLoader.cs b/Robust.Shared/ContentPack/ModLoader.cs
index 056c2eaf9d8..b7f92c886b0 100644
--- a/Robust.Shared/ContentPack/ModLoader.cs
+++ b/Robust.Shared/ContentPack/ModLoader.cs
@@ -118,13 +118,14 @@ public bool TryLoadModules(IEnumerable<ResPath> paths)
                 var checkerSw = Stopwatch.StartNew();
 
                 var typeChecker = MakeTypeChecker();
+                var resolver = typeChecker.CreateResolver();
 
                 Parallel.ForEach(files, pair =>
                 {
                     var (name, (path, _)) = pair;
 
                     using var stream = _res.ContentFileRead(path);
-                    if (!typeChecker.CheckAssembly(stream))
+                    if (!typeChecker.CheckAssembly(stream, resolver))
                     {
                         throw new TypeCheckFailedException($"Assembly {name} failed type checks.");
                     }
@@ -426,12 +427,15 @@ private AssemblyTypeChecker MakeTypeChecker()
             };
         }
 
-        internal static PEReader MakePEReader(Stream stream, bool leaveOpen=false)
+        internal static PEReader MakePEReader(Stream stream, bool leaveOpen=false, PEStreamOptions options=PEStreamOptions.Default)
         {
             if (!stream.CanSeek)
                 stream = leaveOpen ? stream.CopyToMemoryStream() : stream.ConsumeToMemoryStream();
 
-            return new PEReader(stream, leaveOpen ? PEStreamOptions.LeaveOpen : default);
+            if (leaveOpen)
+                options |= PEStreamOptions.LeaveOpen;
+
+            return new PEReader(stream, options);
         }
     }
 }