using Microsoft.EntityFrameworkCore; using NadekoBot.Common.ModuleBehaviors; using NadekoBot.Db; namespace NadekoBot.Modules.Permissions.Services; public sealed class CmdCdService : IExecPreCommand, IReadyExecutor, INService { private readonly DbService _db; private readonly ConcurrentDictionary> _settings = new(); private readonly ConcurrentDictionary<(ulong, string), ConcurrentDictionary> _activeCooldowns = new(); public int Priority => 0; public CmdCdService(Bot bot, DbService db) { _db = db; _settings = bot .AllGuildConfigs .ToDictionary(x => x.GuildId, x => x.CommandCooldowns .DistinctBy(x => x.CommandName.ToLowerInvariant()) .ToDictionary(c => c.CommandName, c => c.Seconds) .ToConcurrent()) .ToConcurrent(); } public Task ExecPreCommandAsync(ICommandContext context, string moduleName, CommandInfo command) => TryBlock(context.Guild, context.User, command.Name.ToLowerInvariant()); public Task TryBlock(IGuild? guild, IUser user, string commandName) { if (guild is null) return Task.FromResult(false); if (!_settings.TryGetValue(guild.Id, out var cooldownSettings)) return Task.FromResult(false); if (!cooldownSettings.TryGetValue(commandName, out var cdSeconds)) return Task.FromResult(false); var cooldowns = _activeCooldowns.GetOrAdd( (guild.Id, commandName), static _ => new()); // if user is not already on cooldown, add if (cooldowns.TryAdd(user.Id, DateTime.UtcNow)) { return Task.FromResult(false); } // if there is an entry, maybe it expired. Try to check if it expired and don't fail if it did // - just update if (cooldowns.TryGetValue(user.Id, out var oldValue)) { var diff = DateTime.UtcNow - oldValue; if (diff.TotalSeconds > cdSeconds) { if (cooldowns.TryUpdate(user.Id, DateTime.UtcNow, oldValue)) return Task.FromResult(false); } } return Task.FromResult(true); } public async Task OnReadyAsync() { using var timer = new PeriodicTimer(TimeSpan.FromHours(1)); while (await timer.WaitForNextTickAsync()) { // once per hour delete expired entries foreach (var ((guildId, commandName), dict) in _activeCooldowns) { // if this pair no longer has associated config, that means it has been removed. // remove all cooldowns if (!_settings.TryGetValue(guildId, out var inner) || !inner.TryGetValue(commandName, out var cdSeconds)) { _activeCooldowns.Remove((guildId, commandName), out _); continue; } Cleanup(dict, cdSeconds); } } } private void Cleanup(ConcurrentDictionary dict, int cdSeconds) { var now = DateTime.UtcNow; foreach (var (key, _) in dict.Where(x => (now - x.Value).TotalSeconds > cdSeconds).ToArray()) { dict.TryRemove(key, out _); } } public void ClearCooldowns(ulong guildId, string cmdName) { if (_settings.TryGetValue(guildId, out var dict)) dict.TryRemove(cmdName, out _); _activeCooldowns.TryRemove((guildId, cmdName), out _); using var ctx = _db.GetDbContext(); var gc = ctx.GuildConfigsForId(guildId, x => x.Include(x => x.CommandCooldowns)); gc.CommandCooldowns.RemoveWhere(x => x.CommandName == cmdName); ctx.SaveChanges(); } public void AddCooldown(ulong guildId, string name, int secs) { if (secs <= 0) throw new ArgumentOutOfRangeException(nameof(secs)); var sett = _settings.GetOrAdd(guildId, static _ => new()); sett[name] = secs; // force cleanup if (_activeCooldowns.TryGetValue((guildId, name), out var dict)) Cleanup(dict, secs); using var ctx = _db.GetDbContext(); var gc = ctx.GuildConfigsForId(guildId, x => x.Include(x => x.CommandCooldowns)); gc.CommandCooldowns.RemoveWhere(x => x.CommandName == name); gc.CommandCooldowns.Add(new() { Seconds = secs, CommandName = name }); ctx.SaveChanges(); } public IReadOnlyCollection<(string CommandName, int Seconds)> GetCommandCooldowns(ulong guildId) { if (!_settings.TryGetValue(guildId, out var dict)) return Array.Empty<(string, int)>(); return dict.Select(x => (x.Key, x.Value)).ToArray(); } }