From 1cd6a6009dcc3df8cebcb57be810593928a66661 Mon Sep 17 00:00:00 2001 From: Supositware Date: Tue, 12 Dec 2023 21:20:48 +0100 Subject: [PATCH] Adding a limit to how many time a command can be executed at the same time --- events/client/interactionCreate.js | 25 +++++++++++++++++- events/client/messageCreate.js | 24 ++++++++++++++++- utils/ratelimiter.js | 42 ++++++++++++++++++++++++++++++ 3 files changed, 89 insertions(+), 2 deletions(-) diff --git a/events/client/interactionCreate.js b/events/client/interactionCreate.js index a0572122..4bb21a84 100644 --- a/events/client/interactionCreate.js +++ b/events/client/interactionCreate.js @@ -11,10 +11,17 @@ export default { if (interaction.type !== InteractionType.ApplicationCommand) return; const globalBlacklist = await db.Blacklists.findOne({ where: { type:'global', uid:interaction.user.id } }); + // const serverBlacklist = await db.Blacklists.findOne({ where: { type:'guild', uid:interaction.guild.id } }); const commandBlacklist = await db.Blacklists.findOne({ where: { type:interaction.commandName, uid:interaction.user.id } }); + if (globalBlacklist) { return interaction.reply({ content: `You are globally blacklisted for the following reason: \`${globalBlacklist.reason}\``, ephemeral: true }); } + /* Server blacklist is untested + else if (serverBlacklist) { + return interaction.reply({ content: `This guild has been blacklisted for the following reason: \`${serverBlacklist.reason}\``, ephemeral: true }); + } + */ else if (commandBlacklist) { return interaction.reply({ content: `You are blacklisted for the following reason: \`${commandBlacklist.reason}\``, ephemeral: true }); } @@ -64,6 +71,18 @@ export default { } */ + // Check if the limit of parallel execution has been reached + if (command.parallelLimit) { + console.log('Command has a parallel limit'); + const doParallelLimit = ratelimiter.checkParallel(interaction.user, commandName, command); + console.log(doParallelLimit); + if (doParallelLimit) { + return await interaction.reply({ content: doParallelLimit, ephemeral: true }); + } + + ratelimiter.addParallel(commandName); + } + // Check the ratelimit const doRateLimit = ratelimiter.check(interaction.user, commandName, command); if (doRateLimit) { @@ -91,7 +110,11 @@ export default { console.log(`\x1b[33m⤷\x1b[0m with args ${JSON.stringify(args)}`); } - await command.execute(interaction, args, client); + await command.execute(interaction, args, client) + .then(() => { + const hasPrallelLimit = ratelimiter.checkParallel(interaction.user, commandName, command); + if (hasPrallelLimit) ratelimiter.removeParallel(commandName); + }); } catch (error) { console.error(error); diff --git a/events/client/messageCreate.js b/events/client/messageCreate.js index 3d048899..62862016 100644 --- a/events/client/messageCreate.js +++ b/events/client/messageCreate.js @@ -281,11 +281,17 @@ export default { if (!command) return; const globalBlacklist = await db.Blacklists.findOne({ where: { type:'global', uid:message.author.id } }); + // const serverBlacklist = await db.Blacklists.findOne({ where: { type:'guild', uid:message.guild.id } }); const commandBlacklist = await db.Blacklists.findOne({ where: { type:commandName, uid:message.author.id } }); if (globalBlacklist) { return message.reply({ content: `You are globally blacklisted for the following reason: \`${globalBlacklist.reason}\``, ephemeral: true }); } + /* Server blacklist is untested + else if (serverBlacklist) { + return message.reply({ content: `This guild has been blacklisted for the following reason: \`${serverBlacklist.reason}\``, ephemeral: true }); + } + */ else if (commandBlacklist) { return message.reply({ content: `You are blacklisted for the following reason: \`${commandBlacklist.reason}\``, ephemeral: true }); } @@ -327,6 +333,18 @@ export default { } } + // Check if the limit of parallel execution has been reached + if (command.parallelLimit) { + console.log('Command has a parallel limit'); + const doParallelLimit = ratelimiter.checkParallel(message.author, commandName, command); + console.log(doParallelLimit); + if (doParallelLimit) { + return await message.reply({ content: doParallelLimit, ephemeral: true }); + } + + ratelimiter.addParallel(commandName); + } + // Check the ratelimit const doRateLimit = ratelimiter.check(message.author, commandName, command); if (doRateLimit) { @@ -414,7 +432,11 @@ export default { console.log(`\x1b[33m⤷\x1b[0m with args ${JSON.stringify(args)}`); } - await command.execute(message, args, client); + await command.execute(message, args, client) + .then(() => { + const hasPrallelLimit = ratelimiter.checkParallel(message.author, commandName, command); + if (hasPrallelLimit) ratelimiter.removeParallel(commandName); + }); } catch (error) { console.error(error); diff --git a/utils/ratelimiter.js b/utils/ratelimiter.js index 9ae6b8b8..3fd1b9f4 100644 --- a/utils/ratelimiter.js +++ b/utils/ratelimiter.js @@ -1,10 +1,14 @@ const ratelimit = {}; +const parallelLimit = {}; const { ownerId } = process.env; import db from '../models/index.js'; export default { check, + addParallel, + removeParallel, + checkParallel, }; function check(user, commandName, commands) { const userID = user.id; @@ -51,5 +55,43 @@ function check(user, commandName, commands) { ratelimit[userID][commandName] = { limit: ratelimit[userID][commandName] ? ratelimit[userID][commandName].limit + 1 : 1, cooldown: date }; } + return false; +} + +function addParallel(commandName) { + console.log(`[ADD] Adding parallel to ${commandName}`); + if (!parallelLimit[commandName]) parallelLimit[commandName] = 0; + + const prevNumber = parallelLimit[commandName]; + + console.log(`[ADD] Previous parallel executions: ${prevNumber}`); + console.log(`[ADD] Current parallel executions: ${JSON.stringify(parallelLimit)}`); + parallelLimit[commandName] = prevNumber + 1; +} + +function removeParallel(commandName) { + console.log(`[REMOVE] Removing parallel to ${commandName}`); + + // This shouldn't be possible + if (!parallelLimit[commandName]) parallelLimit[commandName] = 0; + + const prevNumber = parallelLimit[commandName]; + + console.log(`[REMOVE] previous number: ${prevNumber}`); + console.log(`[REMOVE] previous parallel limit: ${JSON.stringify(parallelLimit)}`); + parallelLimit[commandName] = prevNumber - 1; + console.log(`[REMOVE] current parallel limit: ${JSON.stringify(parallelLimit)}`); +} + +function checkParallel(user, commandName, command) { + // Don't apply the rate limit to bot owner + // if (user.id === ownerId) return false; + + console.log(`[CHECK] command limit: ${command.parallelLimit}`); + console.log(`[CHECK] current parallel executions: ${parallelLimit[commandName]}`); + if (parallelLimit[commandName] >= command.parallelLimit) { + return 'There are currently too many parallel execution of this command, please wait before retrying.'; + } + return false; } \ No newline at end of file