Added execution critical section handler to injector file modifications

IPA.Injector/Injector.cs

@@ -83,6 +83,8 @@ namespace IPA.Injector
                     return;
                 }
 
+                CriticalSection.Configure();
+
                 loader.Debug("Prepping bootstrapper");
                 
                 // updates backup
@@ -141,6 +143,9 @@ namespace IPA.Injector
                 var unityPath = Path.Combine(managedPath,
                     "UnityEngine.CoreModule.dll");
 
+                // this is a critical section because if you exit in here, CoreModule can die
+                CriticalSection.EnterExecuteSection();
+
                 var unityAsmDef = AssemblyDefinition.ReadAssembly(unityPath, new ReaderParameters
                 {
                     ReadWrite = false,
@@ -220,8 +225,8 @@ namespace IPA.Injector
                     bkp?.Add(unityPath);
                     unityAsmDef.Write(unityPath);
                 }
-                /*else
-                    return; // shortcut*/
+
+                CriticalSection.ExitExecuteSection();
             }
 
             #endregion Insert patch into UnityEngine.CoreModule.dll
@@ -235,15 +240,20 @@ namespace IPA.Injector
                 #region Virtualize Assembly-CSharp.dll
 
                 {
+                    CriticalSection.EnterExecuteSection();
+
                     var ascModule = VirtualizedModule.Load(ascPath);
                     ascModule.Virtualize(cAsmName, () => bkp?.Add(ascPath));
+
+                    CriticalSection.ExitExecuteSection();
                 }
 
                 #endregion Virtualize Assembly-CSharp.dll
 
                 #region Anti-Yeet
 
-                //if (SelfConfig.SelfConfigRef.Value.ApplyAntiYeet)
+                CriticalSection.EnterExecuteSection();
+
                 try
                 {
                     loader.Debug("Applying anti-yeet patch");
@@ -266,6 +276,8 @@ namespace IPA.Injector
                     // ignore
                 }
 
+                CriticalSection.ExitExecuteSection();
+
                 #endregion
             }
         }

IPA.Loader/IPA.Loader.csproj

@@ -13,7 +13,6 @@
     <Deterministic>true</Deterministic>
     <PathMap>$(SolutionDir)=C:\</PathMap>
     <DebugType>portable</DebugType>
-    
     <ShowTargetFrameworks>false</ShowTargetFrameworks>
     <ShowTargetFrameworks Condition=" '$(BuildingInsideVisualStudio)' != 'true' Or '$(CI)' == 'True' ">true</ShowTargetFrameworks>
     <TargetFrameworks Condition=" '$(ShowTargetFrameworks)' == 'true' ">net461;net35</TargetFrameworks>
@@ -79,7 +78,7 @@
     </Reference>
   </ItemGroup>
   <ItemGroup>
-    <ProjectReference Include="..\Net3-Proxy\Net3-Proxy.csproj" Condition=" '$(Platform)' == 'Net3' " >
+    <ProjectReference Include="..\Net3-Proxy\Net3-Proxy.csproj" Condition=" '$(Platform)' == 'Net3' ">
       <Project>{642F52DA-90F9-40E3-8784-6964F36752FB}</Project>
       <Name>Net3-Proxy</Name>
     </ProjectReference>
@@ -122,6 +121,7 @@
     <Compile Include="JsonConverters\SemverVersionConverter.cs" />
     <Compile Include="Utilities\BeatSaber.cs" />
     <Compile Include="Utilities\AlmostVersion.cs" />
+    <Compile Include="Utilities\CriticalSection.cs" />
     <Compile Include="Utilities\Ref.cs" />
     <Compile Include="Utilities\ReflectionUtil.cs" />
     <Compile Include="Loader\Composite\CompositeIPAPlugin.cs" />

IPA.Loader/Utilities/BeatSaber.cs

@@ -1,4 +1,5 @@
-using System;
+using IPA.Config;
+using System;
 using System.IO;
 using System.Reflection;
 using UnityEngine;
@@ -18,20 +19,30 @@ namespace IPA.Utilities
         /// Provides the current game version.
         /// </summary>
         /// <value>the SemVer version of the game</value>
-        public static AlmostVersion GameVersion => _gameVersion ?? (_gameVersion = new AlmostVersion(Application.version));
+        public static AlmostVersion GameVersion => _gameVersion ?? (_gameVersion = new AlmostVersion(ApplicationVersionProxy));
 
         internal static void SetEarlyGameVersion(AlmostVersion ver)
         {
             _gameVersion = ver;
             Logging.Logger.log.Debug($"GameVersion set early to {ver}");
         }
+        private static string ApplicationVersionProxy => Application.version;
         internal static void EnsureRuntimeGameVersion()
         {
-            var rtVer = new AlmostVersion(Application.version);
-            if (!rtVer.Equals(_gameVersion)) // this actually uses stricter equality than == for AlmostVersion
+            try
             {
-                Logging.Logger.log.Warn($"Early version {_gameVersion} parsed from game files doesn't match runtime version {rtVer}!");
-                _gameVersion = rtVer;
+                var rtVer = new AlmostVersion(ApplicationVersionProxy);
+                if (!rtVer.Equals(_gameVersion)) // this actually uses stricter equality than == for AlmostVersion
+                {
+                    Logging.Logger.log.Warn($"Early version {_gameVersion} parsed from game files doesn't match runtime version {rtVer}!");
+                    _gameVersion = rtVer;
+                }
+            }
+            catch (MissingMethodException e)
+            {
+                Logging.Logger.log.Error("Application.version was not found! Cannot check early parsed version");
+                if (SelfConfig.Debug_.ShowHandledErrorStackTraces_)
+                    Logging.Logger.log.Error(e);
             }
         }
 

IPA.Loader/Utilities/CriticalSection.cs

@@ -0,0 +1,74 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Runtime.InteropServices;
+using System.Text;
+using System.Threading.Tasks;
+using IPA.Logging;
+
+namespace IPA.Utilities
+{
+    public static class CriticalSection
+    {
+
+        internal static void Configure()
+        {
+            Logger.log.Debug("Configuring exit handlers");
+
+            ResetExitHandlers();
+        }
+
+        #region Execute section
+
+        private static readonly EventHandler registeredHandler = HandleExit;
+        internal static void ResetExitHandlers()
+        {
+            SetConsoleCtrlHandler(registeredHandler, false);
+            SetConsoleCtrlHandler(registeredHandler, true);
+        }
+
+        [DllImport("Kernel32")]
+        private static extern bool SetConsoleCtrlHandler(EventHandler handler, bool add);
+
+        private delegate bool EventHandler(CtrlType sig);
+        private static EventHandler _handler = null;
+
+        private static bool HandleExit(CtrlType type)
+        {
+            if (_handler != null)
+                return _handler(type);
+
+            return false;
+        } 
+
+        enum CtrlType
+        {
+            CTRL_C_EVENT = 0,
+            CTRL_BREAK_EVENT = 1,
+            CTRL_CLOSE_EVENT = 2,
+            CTRL_LOGOFF_EVENT = 5,
+            CTRL_SHUTDOWN_EVENT = 6
+        }
+
+        private static volatile bool exitRecieved = false;
+
+        public static void EnterExecuteSection()
+        {
+            ResetExitHandlers();
+
+            exitRecieved = false;
+            _handler = sig => exitRecieved = true;
+        }
+
+        public static void ExitExecuteSection()
+        {
+            _handler = null;
+
+            if (exitRecieved)
+                Environment.Exit(1);
+        }
+
+        #endregion
+
+    }
+}