From 816d86561d7295bf3d1618525fcb451b1e092528 Mon Sep 17 00:00:00 2001
From: NathanKell <nathankell@gmail.com>
Date: Thu, 5 Oct 2023 22:13:36 -0700
Subject: [PATCH] Patch ModuleRCS to not allow throttling to get around control
 locker (and prevent any actuation when non-axial or from MJ)

---
 Source/RP0/Avionics/ControlLocker.cs | 10 ++++
 Source/RP0/Harmony/ModuleRCS.cs      | 70 ++++++++++++++++++++++++++++
 Source/RP0/RP0.csproj                |  1 +
 3 files changed, 81 insertions(+)
 create mode 100644 Source/RP0/Harmony/ModuleRCS.cs

diff --git a/Source/RP0/Avionics/ControlLocker.cs b/Source/RP0/Avionics/ControlLocker.cs
index 645352d4476..44c42d0e0b1 100644
--- a/Source/RP0/Avionics/ControlLocker.cs
+++ b/Source/RP0/Avionics/ControlLocker.cs
@@ -36,6 +36,8 @@ class ControlLocker : MonoBehaviour
         private static PropertyInfo _mjDeactivateControl = null;
         private object _masterMechJeb = null;
 
+        public static ControlLocker Instance { get; private set; }
+
         private void Awake()
         {
             if (!_isFirstLoad) return;
@@ -52,6 +54,11 @@ private void Awake()
                 RP0Debug.Log($"{ModTag} MechJeb methods found", true);
             else
                 RP0Debug.Log($"{ModTag} MJ assembly or methods NOT found", true);
+
+            if (Instance != null)
+                Destroy(Instance);
+
+            Instance = this;
         }
 
         private void Start()
@@ -175,6 +182,9 @@ public void OnDestroy()
             GameEvents.onVesselSwitching.Remove(OnVesselSwitchingHandler);
             GameEvents.onVesselGoOnRails.Remove(OnRailsHandler);
             GameEvents.onVesselGoOffRails.Remove(OffRailsHandler);
+
+            if (Instance == this)
+                Instance = null;
         }
 
         private void DisableAutopilot()
diff --git a/Source/RP0/Harmony/ModuleRCS.cs b/Source/RP0/Harmony/ModuleRCS.cs
new file mode 100644
index 00000000000..aeb858c3d98
--- /dev/null
+++ b/Source/RP0/Harmony/ModuleRCS.cs
@@ -0,0 +1,70 @@
+using HarmonyLib;
+
+namespace RP0.Harmony
+{
+    [HarmonyPatch(typeof(ModuleRCS))]
+    internal class PatchModuleRCS
+    {
+        private static bool _resetState = false;
+        private static readonly FlightCtrlState _state = new FlightCtrlState();
+
+        [HarmonyPrefix]
+        [HarmonyPatch("Update")]
+        internal static void Prefix_Update(ModuleRCS __instance)
+        {
+            if (!HighLogic.LoadedSceneIsFlight || __instance.vessel == null || __instance.vessel != FlightGlobals.ActiveVessel)
+                return;
+
+            if (ControlLocker.Instance == null)
+                return;
+
+            var level = ControlLocker.Instance.ShouldLock();
+            if (level == ControlLockerUtils.LockLevel.Unlocked)
+                return;
+
+            _resetState = true;
+            _state.X = __instance.vessel.ctrlState.X;
+            _state.Y = __instance.vessel.ctrlState.Y;
+            _state.Z = __instance.vessel.ctrlState.Z;
+            _state.pitch = __instance.vessel.ctrlState.pitch;
+            _state.yaw = __instance.vessel.ctrlState.yaw;
+            _state.roll = __instance.vessel.ctrlState.roll;
+            _state.mainThrottle = __instance.vessel.ctrlState.mainThrottle;
+
+            __instance.vessel.ctrlState.X = __instance.vessel.ctrlState.Y = __instance.vessel.ctrlState.pitch = __instance.vessel.ctrlState.yaw = __instance.vessel.ctrlState.roll = 0f;
+
+            if (level == ControlLockerUtils.LockLevel.Axial)
+            {
+                if (__instance.vessel.ctrlState.Z < 0f)
+                    __instance.vessel.ctrlState.Z = -1f;
+                else if (__instance.vessel.ctrlState.Z > 0f)
+                    __instance.vessel.ctrlState.Z = 1f;
+
+                if (__instance.vessel.ctrlState.mainThrottle > 0f)
+                    __instance.vessel.ctrlState.mainThrottle = 1f;
+            }
+            else
+            {
+                __instance.vessel.ctrlState.Z = 0f;
+                __instance.vessel.ctrlState.mainThrottle = 0f;
+            }
+        }
+
+        [HarmonyPostfix]
+        [HarmonyPatch("Update")]
+        internal static void Postfix_Update(ModuleRCS __instance)
+        {
+            if (!_resetState)
+                return;
+
+            _resetState = false;
+            __instance.vessel.ctrlState.X = _state.X;
+            __instance.vessel.ctrlState.Y = _state.Y;
+            __instance.vessel.ctrlState.Z = _state.Z;
+            __instance.vessel.ctrlState.pitch = _state.pitch;
+            __instance.vessel.ctrlState.yaw = _state.yaw;
+            __instance.vessel.ctrlState.roll = _state.roll;
+            __instance.vessel.ctrlState.mainThrottle = _state.mainThrottle;
+        }
+    }
+}
diff --git a/Source/RP0/RP0.csproj b/Source/RP0/RP0.csproj
index 672db358a24..acf61e92b29 100644
--- a/Source/RP0/RP0.csproj
+++ b/Source/RP0/RP0.csproj
@@ -38,6 +38,7 @@
   <ItemGroup>
     <Compile Include="Avionics\ControlLockerUtils.cs" />
     <Compile Include="Avionics\EditorBinder.cs" />
+    <Compile Include="Harmony\ModuleRCS.cs" />
     <Compile Include="Harmony\KSPWheel.cs" />
     <Compile Include="UI\ProceduralAvionicsWindow.cs" />
     <Compile Include="CareerLog\CareerEvent.cs" />