Prevent SecretRule from picking invalid presets (#27456)

* Prevent SecretRule from picking invalid presets

* remove lonely semicolon

---------

Co-authored-by: Nemanja <98561806+EmoGarbage404@users.noreply.github.com>

Content.Server/GameTicking/GameTicker.RoundFlow.cs

@@ -173,6 +173,26 @@ namespace Content.Server.GameTicking
             return gridUids;
         }
 
+        public int ReadyPlayerCount()
+        {
+            var total = 0;
+            foreach (var (userId, status) in _playerGameStatuses)
+            {
+                if (LobbyEnabled && status == PlayerGameStatus.NotReadyToPlay)
+                    continue;
+
+                if (!_playerManager.TryGetSessionById(userId, out _))
+                    continue;
+
+                if (_banManager.GetRoleBans(userId) == null)
+                    continue;
+
+                total++;
+            }
+
+            return total;
+        }
+
         public void StartRound(bool force = false)
         {
 #if EXCEPTION_TOLERANCE
@@ -228,6 +248,8 @@ namespace Content.Server.GameTicking
                 readyPlayerProfiles.Add(userId, profile);
             }
 
+            DebugTools.AssertEqual(readyPlayers.Count, ReadyPlayerCount());
+
             // Just in case it hasn't been loaded previously we'll try loading it.
             LoadMaps();
 

Content.Server/GameTicking/Rules/SecretRuleSystem.cs

@@ -1,15 +1,17 @@
+using System.Diagnostics.CodeAnalysis;
+using System.Linq;
 using Content.Server.Administration.Logs;
 using Content.Server.GameTicking.Components;
 using Content.Server.Chat.Managers;
 using Content.Server.GameTicking.Presets;
 using Content.Server.GameTicking.Rules.Components;
 using Content.Shared.Random;
-using Content.Shared.Random.Helpers;
 using Content.Shared.CCVar;
 using Content.Shared.Database;
 using Robust.Shared.Prototypes;
 using Robust.Shared.Random;
 using Robust.Shared.Configuration;
+using Robust.Shared.Utility;
 
 namespace Content.Server.GameTicking.Rules;
 
@@ -20,11 +22,46 @@ public sealed class SecretRuleSystem : GameRuleSystem<SecretRuleComponent>
     [Dependency] private readonly IConfigurationManager _configurationManager = default!;
     [Dependency] private readonly IAdminLogManager _adminLogger = default!;
     [Dependency] private readonly IChatManager _chatManager = default!;
+    [Dependency] private readonly IComponentFactory _compFact = default!;
+
+    private string _ruleCompName = default!;
+
+    public override void Initialize()
+    {
+        base.Initialize();
+        _ruleCompName = _compFact.GetComponentName(typeof(GameRuleComponent));
+    }
 
     protected override void Added(EntityUid uid, SecretRuleComponent component, GameRuleComponent gameRule, GameRuleAddedEvent args)
     {
         base.Added(uid, component, gameRule, args);
-        PickRule(component);
+        var weights = _configurationManager.GetCVar(CCVars.SecretWeightPrototype);
+
+        if (!TryPickPreset(weights, out var preset))
+        {
+            Log.Error($"{ToPrettyString(uid)} failed to pick any preset. Removing rule.");
+            Del(uid);
+            return;
+        }
+
+        Log.Info($"Selected {preset.ID} as the secret preset.");
+        _adminLogger.Add(LogType.EventStarted, $"Selected {preset.ID} as the secret preset.");
+        _chatManager.SendAdminAnnouncement(Loc.GetString("rule-secret-selected-preset", ("preset", preset.ID)));
+
+        foreach (var rule in preset.Rules)
+        {
+            EntityUid ruleEnt;
+
+            // if we're pre-round (i.e. will only be added)
+            // then just add rules. if we're added in the middle of the round (or at any other point really)
+            // then we want to start them as well
+            if (GameTicker.RunLevel <= GameRunLevel.InRound)
+                ruleEnt = GameTicker.AddGameRule(rule);
+            else
+                GameTicker.StartGameRule(rule, out ruleEnt);
+
+            component.AdditionalGameRules.Add(ruleEnt);
+        }
     }
 
     protected override void Ended(EntityUid uid, SecretRuleComponent component, GameRuleComponent gameRule, GameRuleEndedEvent args)
@@ -37,32 +74,101 @@ public sealed class SecretRuleSystem : GameRuleSystem<SecretRuleComponent>
         }
     }
 
-    private void PickRule(SecretRuleComponent component)
+    private bool TryPickPreset(ProtoId<WeightedRandomPrototype> weights, [NotNullWhen(true)] out GamePresetPrototype? preset)
     {
-        // TODO: This doesn't consider what can't start due to minimum player count,
+        var options = _prototypeManager.Index(weights).Weights.ShallowClone();
-        // but currently there's no way to know anyway as they use cvars.
+        var players = GameTicker.ReadyPlayerCount();
-        var presetString = _configurationManager.GetCVar(CCVars.SecretWeightPrototype);
-        var preset = _prototypeManager.Index<WeightedRandomPrototype>(presetString).Pick(_random);
-        Log.Info($"Selected {preset} for secret.");
-        _adminLogger.Add(LogType.EventStarted, $"Selected {preset} for secret.");
-        _chatManager.SendAdminAnnouncement(Loc.GetString("rule-secret-selected-preset", ("preset", preset)));
 
-        var rules = _prototypeManager.Index<GamePresetPrototype>(preset).Rules;
+        GamePresetPrototype? selectedPreset = null;
-        foreach (var rule in rules)
+        var sum = options.Values.Sum();
+        while (options.Count > 0)
         {
-            EntityUid ruleEnt;
+            var accumulated = 0f;
+            var rand = _random.NextFloat(sum);
+            foreach (var (key, weight) in options)
+            {
+                accumulated += weight;
+                if (accumulated < rand)
+                    continue;
 
-            // if we're pre-round (i.e. will only be added)
+                if (!_prototypeManager.TryIndex(key, out selectedPreset))
-            // then just add rules. if we're added in the middle of the round (or at any other point really)
+                    Log.Error($"Invalid preset {selectedPreset} in secret rule weights: {weights}");
-            // then we want to start them as well
+
-            if (GameTicker.RunLevel <= GameRunLevel.InRound)
+                options.Remove(key);
-                ruleEnt = GameTicker.AddGameRule(rule);
+                sum -= weight;
-            else
+                break;
-            {
-                GameTicker.StartGameRule(rule, out ruleEnt);
             }
 
-            component.AdditionalGameRules.Add(ruleEnt);
+            if (CanPick(selectedPreset, players))
+            {
+                preset = selectedPreset;
+                return true;
             }
+
+            if (selectedPreset != null)
+                Log.Info($"Excluding {selectedPreset.ID} from secret preset selection.");
+        }
+
+        preset = null;
+        return false;
+    }
+
+    public bool CanPickAny()
+    {
+        var secretPresetId = _configurationManager.GetCVar(CCVars.SecretWeightPrototype);
+        return CanPickAny(secretPresetId);
+    }
+
+    /// <summary>
+    /// Can any of the given presets be picked, taking into account the currently available player count?
+    /// </summary>
+    public bool CanPickAny(ProtoId<WeightedRandomPrototype> weightedPresets)
+    {
+        var ids = _prototypeManager.Index(weightedPresets).Weights.Keys
+            .Select(x => new ProtoId<GamePresetPrototype>(x));
+
+        return CanPickAny(ids);
+    }
+
+    /// <summary>
+    /// Can any of the given presets be picked, taking into account the currently available player count?
+    /// </summary>
+    public bool CanPickAny(IEnumerable<ProtoId<GamePresetPrototype>> protos)
+    {
+        var players = GameTicker.ReadyPlayerCount();
+        foreach (var id in protos)
+        {
+            if (!_prototypeManager.TryIndex(id, out var selectedPreset))
+                Log.Error($"Invalid preset {selectedPreset} in secret rule weights: {id}");
+
+            if (CanPick(selectedPreset, players))
+                return true;
+        }
+
+        return false;
+    }
+
+    /// <summary>
+    /// Can the given preset be picked, taking into account the currently available player count?
+    /// </summary>
+    private bool CanPick([NotNullWhen(true)] GamePresetPrototype? selected, int players)
+    {
+        if (selected == null)
+            return false;
+
+        foreach (var ruleId in selected.Rules)
+        {
+            if (!_prototypeManager.TryIndex(ruleId, out EntityPrototype? rule)
+                || !rule.TryGetComponent(_ruleCompName, out GameRuleComponent? ruleComp))
+            {
+                Log.Error($"Encountered invalid rule {ruleId} in preset {selected.ID}");
+                return false;
+            }
+
+            if (ruleComp.MinPlayers > players)
+                return false;
+        }
+
+        return true;
     }
 }