From c8f2ddc72985172992570fd5f1d60300948b33ed Mon Sep 17 00:00:00 2001
From: blueDev2 <89804215+blueDev2@users.noreply.github.com>
Date: Thu, 12 Sep 2024 06:29:11 -0400
Subject: [PATCH] Add ReagentWhitelist to Injector component and system
 (#28262)

* Add reagent whitelist to Injector component and system

* Apply suggested Changes

* Remove LINQ function for performance

* Update Content.Server/Chemistry/EntitySystems/InjectorSystem.cs

Co-authored-by: Ed <96445749+TheShuEd@users.noreply.github.com>

---------

Co-authored-by: Ed <96445749+TheShuEd@users.noreply.github.com>
---
 .../Chemistry/EntitySystems/InjectorSystem.cs | 20 ++++++++++++++++++-
 .../Chemistry/Components/InjectorComponent.cs | 10 ++++++++++
 .../Chemistry/Components/Solution.cs          |  2 +-
 3 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/Content.Server/Chemistry/EntitySystems/InjectorSystem.cs b/Content.Server/Chemistry/EntitySystems/InjectorSystem.cs
index d94faa8123..c5c45daa5b 100644
--- a/Content.Server/Chemistry/EntitySystems/InjectorSystem.cs
+++ b/Content.Server/Chemistry/EntitySystems/InjectorSystem.cs
@@ -327,8 +327,23 @@ public sealed class InjectorSystem : SharedInjectorSystem
             return false;
         }
 
+        var applicableTargetSolution = targetSolution.Comp.Solution;
+        // If a whitelist exists, remove all non-whitelisted reagents from the target solution temporarily
+        var temporarilyRemovedSolution = new Solution();
+        if (injector.Comp.ReagentWhitelist is { } reagentWhitelist)
+        {
+            string[] reagentPrototypeWhitelistArray = new string[reagentWhitelist.Count];
+            var i = 0;
+            foreach (var reagent in reagentWhitelist)
+            {
+                reagentPrototypeWhitelistArray[i] = reagent;
+                ++i;
+            }
+            temporarilyRemovedSolution = applicableTargetSolution.SplitSolutionWithout(applicableTargetSolution.Volume, reagentPrototypeWhitelistArray);
+        }
+
         // Get transfer amount. May be smaller than _transferAmount if not enough room, also make sure there's room in the injector
-        var realTransferAmount = FixedPoint2.Min(injector.Comp.TransferAmount, targetSolution.Comp.Solution.Volume,
+        var realTransferAmount = FixedPoint2.Min(injector.Comp.TransferAmount, applicableTargetSolution.Volume,
             solution.AvailableVolume);
 
         if (realTransferAmount <= 0)
@@ -350,6 +365,9 @@ public sealed class InjectorSystem : SharedInjectorSystem
         // Move units from attackSolution to targetSolution
         var removedSolution = SolutionContainers.Draw(target.Owner, targetSolution, realTransferAmount);
 
+        // Add back non-whitelisted reagents to the target solution
+        applicableTargetSolution.AddSolution(temporarilyRemovedSolution, null);
+
         if (!SolutionContainers.TryAddSolution(soln.Value, removedSolution))
         {
             return false;
diff --git a/Content.Shared/Chemistry/Components/InjectorComponent.cs b/Content.Shared/Chemistry/Components/InjectorComponent.cs
index 1f2716356c..17a65ef1c1 100644
--- a/Content.Shared/Chemistry/Components/InjectorComponent.cs
+++ b/Content.Shared/Chemistry/Components/InjectorComponent.cs
@@ -1,7 +1,9 @@
 using Content.Shared.Chemistry.EntitySystems;
+using Content.Shared.Chemistry.Reagent;
 using Content.Shared.DoAfter;
 using Content.Shared.FixedPoint;
 using Robust.Shared.GameStates;
+using Robust.Shared.Prototypes;
 using Robust.Shared.Serialization;
 
 namespace Content.Shared.Chemistry.Components;
@@ -88,6 +90,14 @@ public sealed partial class InjectorComponent : Component
     [DataField]
     public InjectorToggleMode ToggleState = InjectorToggleMode.Draw;
 
+    /// <summary>
+    /// Reagents that are allowed to be within this injector.
+    /// If a solution has both allowed and non-allowed reagents, only allowed reagents will be drawn into this injector.
+    /// A null ReagentWhitelist indicates all reagents are allowed.
+    /// </summary>
+    [DataField]
+    public List<ProtoId<ReagentPrototype>>? ReagentWhitelist = null;
+
     #region Arguments for injection doafter
 
     /// <inheritdoc cref=DoAfterArgs.NeedHand>
diff --git a/Content.Shared/Chemistry/Components/Solution.cs b/Content.Shared/Chemistry/Components/Solution.cs
index fc25781005..c65ba0e80e 100644
--- a/Content.Shared/Chemistry/Components/Solution.cs
+++ b/Content.Shared/Chemistry/Components/Solution.cs
@@ -612,7 +612,7 @@ namespace Content.Shared.Chemistry.Components
         }
 
         /// <summary>
-        /// Splits a solution without the specified reagent prototypes.
+        /// Splits a solution with only the specified reagent prototypes.
         /// </summary>
         public Solution SplitSolutionWithOnly(FixedPoint2 toTake, params string[] includedPrototypes)
         {
-- 
2.52.0