Add WhitelistManager that tracks the whitelist status of connected players (#44732)

* Initial commit

* Fix the command and update better description

* Add comment

* Oops

Content.Server/IoC/ServerContentIoC.cs

@@ -19,6 +19,7 @@ using Content.Server.NodeContainer.NodeGroups;
 using Content.Server.Players.JobWhitelist;
 using Content.Server.Players.PlayTimeTracking;
 using Content.Server.Players.RateLimiting;
+using Content.Server.Players.Whitelist;
 using Content.Server.Preferences.Managers;
 using Content.Server.ServerInfo;
 using Content.Server.ServerUpdates;
@@ -81,5 +82,6 @@ internal static class ServerContentIoC
         deps.Register<DiscordChatLink>();
         deps.Register<ServerFeedbackManager>();
         deps.Register<ISharedFeedbackManager, ServerFeedbackManager>();
+        deps.Register<WhitelistManager>();
     }
 }

Content.Server/Players/Whitelist/WhitelistManager.cs

@@ -0,0 +1,85 @@
+using System.Threading;
+using System.Threading.Tasks;
+using Content.Server.Database;
+using Robust.Shared.Network;
+using Robust.Shared.Player;
+
+namespace Content.Server.Players.Whitelist;
+
+/// <summary>
+/// Managed for caching whitelist statuses for connected players on the server.
+/// If trying to get the data of disconnected players, use <see cref="IServerDbManager.GetWhitelistStatusAsync"/> directly.
+/// </summary>
+// TODO: This manager can likely be converted into a more general purpose manager once we expand the IServerDbManager's functions.
+public sealed partial class WhitelistManager : IPostInjectInit
+{
+    [Dependency] private IServerDbManager _db = default!;
+    [Dependency] private UserDbDataManager _userDb = default!;
+    [Dependency] private ILogManager _logManager = default!;
+
+    private readonly Dictionary<NetUserId, bool> _whitelistStatus = new();
+    private ISawmill _sawmill = default!;
+
+    private async Task LoadData(ICommonSession session, CancellationToken cancel)
+    {
+        var whitelistStatus = await _db.GetWhitelistStatusAsync(session.UserId);
+        _whitelistStatus[session.UserId] = whitelistStatus;
+    }
+
+    private void ClientDisconnected(ICommonSession session)
+    {
+        _whitelistStatus.Remove(session.UserId);
+    }
+
+    /// <inheritdoc cref="IsConnectedWhitelisted(Robust.Shared.Network.NetUserId)"/>
+    public bool IsConnectedWhitelisted(ICommonSession session)
+    {
+        return IsConnectedWhitelisted(session.UserId);
+    }
+
+    /// <summary>
+    /// Returns true if the player has whitelist status on the server, false if they don't.
+    /// Returns false and raises an error if the player is not in the cached database, meaning they are either not connected or database failed.
+    /// </summary>
+    public bool IsConnectedWhitelisted(NetUserId userId)
+    {
+        if (!_whitelistStatus.TryGetValue(userId, out var whitelistStatus))
+        {
+            _sawmill.Error("{Player} is either not connected, or the database failed to load their whitelist status. Stack trace:\\n{StackTrace}",
+                userId,
+                Environment.StackTrace);
+            return false;
+        }
+
+        return whitelistStatus;
+    }
+
+    /// <summary>
+    /// Adds a player to the whitelist and updates the tracker.
+    /// </summary>
+    public async void AddWhitelist(NetUserId player)
+    {
+        if (_whitelistStatus.ContainsKey(player))
+            _whitelistStatus[player] = true;
+
+        await _db.AddToWhitelistAsync(player);
+    }
+
+    /// <summary>
+    /// Removes a player from the whitelist and updates the tracker.
+    /// </summary>
+    public async void RemoveWhitelist(NetUserId player)
+    {
+        if (_whitelistStatus.ContainsKey(player))
+            _whitelistStatus[player] = false;
+
+        await _db.RemoveFromWhitelistAsync(player);
+    }
+
+    void IPostInjectInit.PostInject()
+    {
+        _userDb.AddOnLoadPlayer(LoadData);
+        _userDb.AddOnPlayerDisconnect(ClientDisconnected);
+        _sawmill = _logManager.GetSawmill("server_whitelist");
+    }
+}

Content.Server/Voting/Managers/VoteManager.DefaultVotes.cs

@@ -344,7 +344,7 @@ namespace Content.Server.Voting.Managers
                 _votingSystem = _entityManager.SystemOrNull<VotingSystem>();
 
             // Check that the initiator is actually allowed to do a votekick.
-            if (_votingSystem != null && !await _votingSystem.CheckVotekickInitEligibility(initiator))
+            if (_votingSystem != null && !_votingSystem.CheckVotekickInitEligibility(initiator))
             {
                 _logManager.GetSawmill("admin.votekick").Warning($"User {initiator} attempted a votekick, despite not being eligible to!");
                 _adminLogger.Add(LogType.Vote, LogImpact.Extreme, $"Votekick attempted by {initiator}, but they are not eligible to votekick!");

Content.Server/Voting/VotingSystem.cs

@@ -13,6 +13,7 @@ using Robust.Shared.Network;
 using Robust.Shared.Player;
 using Robust.Shared.Timing;
 using System.Threading.Tasks;
+using Content.Server.Players.Whitelist;
 using Content.Shared.Players.PlayTimeTracking;
 
 namespace Content.Server.Voting;
@@ -22,12 +23,12 @@ public sealed partial class VotingSystem : EntitySystem
 
     [Dependency] private IPlayerManager _playerManager = default!;
     [Dependency] private IAdminManager _adminManager = default!;
-    [Dependency] private IServerDbManager _dbManager = default!;
     [Dependency] private IGameTiming _gameTiming = default!;
     [Dependency] private IConfigurationManager _cfg = default!;
     [Dependency] private JobSystem _jobs = default!;
     [Dependency] private GameTicker _gameTicker = default!;
     [Dependency] private ISharedPlaytimeManager _playtimeManager = default!;
+    [Dependency] private WhitelistManager _whitelistManager = default!;
 
     public override void Initialize()
     {
@@ -36,9 +37,9 @@ public sealed partial class VotingSystem : EntitySystem
         SubscribeNetworkEvent<VotePlayerListRequestEvent>(OnVotePlayerListRequestEvent);
     }
 
-    private async void OnVotePlayerListRequestEvent(VotePlayerListRequestEvent msg, EntitySessionEventArgs args)
+    private void OnVotePlayerListRequestEvent(VotePlayerListRequestEvent msg, EntitySessionEventArgs args)
     {
-        if (!await CheckVotekickInitEligibility(args.SenderSession))
+        if (!CheckVotekickInitEligibility(args.SenderSession))
         {
             var deniedResponse = new VotePlayerListResponseEvent(new (NetUserId, NetEntity, string)[0], true);
             RaiseNetworkEvent(deniedResponse, args.SenderSession.Channel);
@@ -86,7 +87,7 @@ public sealed partial class VotingSystem : EntitySystem
     /// Used to check whether the player initiating a votekick is allowed to do so serverside.
     /// </summary>
     /// <param name="initiator">The session initiating the votekick.</param>
-    public async Task<bool> CheckVotekickInitEligibility(ICommonSession? initiator)
+    public bool CheckVotekickInitEligibility(ICommonSession? initiator)
     {
         if (initiator == null)
             return false;
@@ -111,7 +112,7 @@ public sealed partial class VotingSystem : EntitySystem
         }
 
         // Must be whitelisted
-        if (!await _dbManager.GetWhitelistStatusAsync(initiator.UserId) && _cfg.GetCVar(CCVars.VotekickInitiatorWhitelistedRequirement))
+        if (!_whitelistManager.IsConnectedWhitelisted(initiator.UserId) && _cfg.GetCVar(CCVars.VotekickInitiatorWhitelistedRequirement))
             return false;
 
         // Must be eligible to vote

Content.Server/Whitelist/WhitelistCommands.cs

@@ -1,5 +1,6 @@
 using Content.Server.Administration;
 using Content.Server.Database;
+using Content.Server.Players.Whitelist;
 using Content.Shared.Administration;
 using Content.Shared.CCVar;
 using Robust.Server.Player;
@@ -14,6 +15,7 @@ public sealed partial class AddWhitelistCommand : LocalizedCommands
 {
     [Dependency] private IPlayerLocator _locator = default!;
     [Dependency] private IServerDbManager _dbManager = default!;
+    [Dependency] private WhitelistManager _whitelistManager = default!;
     public override string Command => "whitelistadd";
 
     public override async void Execute(IConsoleShell shell, string argStr, string[] args)
@@ -31,14 +33,13 @@ public sealed partial class AddWhitelistCommand : LocalizedCommands
         if (data != null)
         {
             var guid = data.UserId;
-            var isWhitelisted = await _dbManager.GetWhitelistStatusAsync(guid);
-            if (isWhitelisted)
+            if (await _dbManager.GetWhitelistStatusAsync(guid))
             {
                 shell.WriteLine(Loc.GetString("cmd-whitelistadd-existing", ("username", data.Username)));
                 return;
             }
 
-            await _dbManager.AddToWhitelistAsync(guid);
+            _whitelistManager.AddWhitelist(guid);
             shell.WriteLine(Loc.GetString("cmd-whitelistadd-added", ("username", data.Username)));
             return;
         }
@@ -62,6 +63,7 @@ public sealed partial class RemoveWhitelistCommand : LocalizedCommands
 {
     [Dependency] private IPlayerLocator _locator = default!;
     [Dependency] private IServerDbManager _dbManager = default!;
+    [Dependency] private WhitelistManager _whitelistManager = default!;
 
     public override string Command => "whitelistremove";
 
@@ -80,14 +82,13 @@ public sealed partial class RemoveWhitelistCommand : LocalizedCommands
         if (data != null)
         {
             var guid = data.UserId;
-            var isWhitelisted = await _dbManager.GetWhitelistStatusAsync(guid);
-            if (!isWhitelisted)
+            if (!await _dbManager.GetWhitelistStatusAsync(guid))
             {
                 shell.WriteLine(Loc.GetString("cmd-whitelistremove-existing", ("username", data.Username)));
                 return;
             }
 
-            await _dbManager.RemoveFromWhitelistAsync(guid);
+            _whitelistManager.RemoveWhitelist(guid);
             shell.WriteLine(Loc.GetString("cmd-whitelistremove-removed", ("username", data.Username)));
             return;
         }
@@ -133,6 +134,8 @@ public sealed partial class KickNonWhitelistedCommand : LocalizedCommands
             if (await _dbManager.GetAdminDataForAsync(session.UserId) is not null)
                 continue;
 
+            // We let this one query the whitelist to be 100% certain it's kicking out non-whitelisted players +
+            // it's mostly adding/removing whitelists that needs to go through the cache.
             if (!await _dbManager.GetWhitelistStatusAsync(session.UserId))
                 _netManager.DisconnectChannel(session.Channel, Loc.GetString("whitelist-not-whitelisted"));
         }