commands: various tweaks and fixes

 - Fixed being able to bypass cooldowns when receiving a lot of commands at the same time
 - Added guild scoped cooldowns (only applied to mentionrole for the moment)
 - Guild scoped cooldowns are respected in custom commands execCC and execAdmin
 - Fix possibility of clashing cooldowns with commands in different containers but same names

Author: jogramming

commands/plugin_bot.go

@@ -84,9 +84,29 @@ func YAGCommandMiddleware(inner dcmd.RunFunc) dcmd.RunFunc {
 			data = data.WithContext(context.WithValue(data.Context(), CtxKeyMS, ms))
 		}
 
+		// Lock the command for execution
+		if !BlockingAddRunningCommand(data.Msg.GuildID, data.Msg.ChannelID, data.Msg.Author.ID, yc, time.Second*60) {
+			if atomic.LoadInt32(shuttingDown) == 1 {
+				return yc.Name + ": Bot is restarting, please try again in a couple seconds...", nil
+			}
+
+			return yc.Name + ": Gave up trying to run command after 60 seconds waiting for your previous instance of this command to finish", nil
+		}
+
+		defer removeRunningCommand(data.Msg.GuildID, data.Msg.ChannelID, data.Msg.Author.ID, yc)
+
 		// Check if the user can execute the command
 		canExecute, resp, settings, err := yc.checkCanExecuteCommand(data, data.CS)
+		if err != nil {
+			yc.Logger(data).WithError(err).Error("An error occured while checking if we could run command")
+		}
+
 		if resp != "" {
+			if resp == ReasonCooldown {
+				cdLeft, _ := yc.LongestCooldownLeft(data.ContainerChain, data.Msg.Author.ID, data.Msg.GuildID)
+				return fmt.Sprintf("This command is on cooldown for another %d seconds", cdLeft), nil
+			}
+
 			// yc.PostCommandExecuted(settings, data, "", errors.WithMessage(err, "checkCanExecuteCommand"))
 			// m, err := common.BotSession.ChannelMessageSend(cState.ID(), resp)
 			// go yc.deleteResponse([]*discordgo.Message{m})
@@ -103,17 +123,6 @@ func YAGCommandMiddleware(inner dcmd.RunFunc) dcmd.RunFunc {
 
 		data = data.WithContext(context.WithValue(data.Context(), CtxKeyCmdSettings, settings))
 
-		// Lock the command for execution
-		if !BlockingAddRunningCommand(data.Msg.GuildID, data.Msg.ChannelID, data.Msg.Author.ID, yc, time.Second*60) {
-			if atomic.LoadInt32(shuttingDown) == 1 {
-				return yc.Name + ": Bot is shutting down or restarting, please try again in a couple seconds...", nil
-			}
-
-			return yc.Name + ": Gave up trying to run command after 60 seconds waiting for your previous instance of this command to finish", nil
-		}
-
-		defer removeRunningCommand(data.Msg.GuildID, data.Msg.ChannelID, data.Msg.Author.ID, yc)
-
 		err = dcmd.ParseCmdArgs(data)
 		if err != nil {
 			if dcmd.IsUserError(err) {

commands/tmplexec.go

@@ -78,7 +78,7 @@ type cmdExecFunc func(cmd string, args ...interface{}) (interface{}, error)
 func TmplExecCmdFuncs(ctx *templates.Context, maxExec int, dryRun bool) (userCtxCommandExec cmdExecFunc, botCtxCommandExec cmdExecFunc) {
 	execUser := func(cmd string, args ...interface{}) (interface{}, error) {
 		messageCopy := *ctx.Msg
-		if ctx.CS != nil { //Check if CS is not a nil pointer 
+		if ctx.CS != nil { //Check if CS is not a nil pointer
 			messageCopy.ChannelID = ctx.CS.ID
 		}
 		mc := &discordgo.MessageCreate{&messageCopy}
@@ -96,7 +96,7 @@ func TmplExecCmdFuncs(ctx *templates.Context, maxExec int, dryRun bool) (userCtx
 
 		messageCopy := *ctx.Msg
 		messageCopy.Author = &botUserCopy
-		if ctx.CS != nil { //Check if CS is not a nil pointer 
+		if ctx.CS != nil { //Check if CS is not a nil pointer
 			messageCopy.ChannelID = ctx.CS.ID
 		}
 
@@ -219,13 +219,24 @@ func execCmd(tmplCtx *templates.Context, dryRun bool, m *discordgo.MessageCreate
 		}
 		runFunc = data.ContainerChain[len(data.ContainerChain)-1-i].BuildMiddlewareChain(runFunc, foundCmd)
 	}
-	// foundCmd.Trigger.
+
+	// Check guild scope cooldown
+	cd, err := cast.GuildScopeCooldownLeft(data.ContainerChain, tmplCtx.GS.ID)
+	if err != nil {
+		return "", errors.WithStackIf(err)
+	}
+
+	if cd > 0 {
+		return "", errors.NewPlain("this command is on guild scope cooldown")
+	}
 
 	resp, err := runFunc(data)
 	if err != nil {
 		return "", errors.WithMessage(err, "exec/execadmin, run")
 	}
 
+	cast.SetCooldownGuild(data.ContainerChain, tmplCtx.GS.ID)
+
 	switch v := resp.(type) {
 	case error:
 		return "Error: " + v.Error(), nil

commands/yagcommmand.go

@@ -60,8 +60,9 @@ var (
 )
 
 var (
-	RKeyCommandCooldown = func(uID int64, cmd string) string { return "cmd_cd:" + discordgo.StrID(uID) + ":" + cmd }
-	RKeyCommandLock     = func(uID int64, cmd string) string { return "cmd_lock:" + discordgo.StrID(uID) + ":" + cmd }
+	RKeyCommandCooldown      = func(uID int64, cmd string) string { return "cmd_cd:" + discordgo.StrID(uID) + ":" + cmd }
+	RKeyCommandCooldownGuild = func(gID int64, cmd string) string { return "cmd_guild_cd:" + discordgo.StrID(gID) + ":" + cmd }
+	RKeyCommandLock          = func(uID int64, cmd string) string { return "cmd_lock:" + discordgo.StrID(uID) + ":" + cmd }
 
 	CommandExecTimeout = time.Minute
 
@@ -98,8 +99,9 @@ type YAGCommand struct {
 	CustomEnabled        bool   // Set to true to handle the enable check itself
 	Default              bool   // The default enabled state of this command
 
-	Cooldown    int // Cooldown in seconds before user can use it again
-	CmdCategory *dcmd.Category
+	Cooldown           int // Cooldown in seconds before user can use it again
+	CmdCategory        *dcmd.Category
+	GuildScopeCooldown int
 
 	RunInDM      bool // Set to enable this commmand in DM's
 	HideFromHelp bool // Set to hide from help
@@ -190,7 +192,7 @@ func (yc *YAGCommand) Run(data *dcmd.Data) (interface{}, error) {
 
 	// Log errors
 	if cmdErr == nil {
-		err := yc.SetCooldown(data.Msg.Author.ID)
+		err := yc.SetCooldowns(data.ContainerChain, data.Msg.Author.ID, data.Msg.GuildID)
 		if err != nil {
 			logger.WithError(err).Error("Failed setting cooldown")
 		}
@@ -305,6 +307,15 @@ func (yc *YAGCommand) PostCommandExecuted(settings *CommandSettings, cmdData *dc
 	return
 }
 
+const (
+	ReasonError                    = "An error occured"
+	ReasonCommandDisabaledSettings = "Command is disabled in the settings"
+	ReasonMissingRole              = "Missing a required role for this command"
+	ReasonIgnoredRole              = "Has a ignored role for this command"
+	ReasonUserMissingPerms         = "User is missing one or more permissions to run this command"
+	ReasonCooldown                 = "This command is on cooldown"
+)
+
 // checks if the specified user can execute the command, and if so returns the settings for said command
 func (yc *YAGCommand) checkCanExecuteCommand(data *dcmd.Data, cState *dstate.ChannelState) (canExecute bool, resp string, settings *CommandSettings, err error) {
 	// Check guild specific settings if not triggered from a DM
@@ -316,7 +327,8 @@ func (yc *YAGCommand) checkCanExecuteCommand(data *dcmd.Data, cState *dstate.Cha
 		guild = cState.Guild
 
 		if guild == nil {
-			resp = "You're not on a server?"
+			err = errors.NewPlain("Not on a guild")
+			resp = ReasonError
 			return
 		}
 
@@ -329,12 +341,12 @@ func (yc *YAGCommand) checkCanExecuteCommand(data *dcmd.Data, cState *dstate.Cha
 		settings, err = yc.GetSettings(data.ContainerChain, cState.ID, cop.ParentID, guild.ID)
 		if err != nil {
 			err = errors.WithMessage(err, "cs.GetSettings")
-			resp = "Bot is having isssues, contact the bot owner."
+			resp = ReasonError
 			return
 		}
 
 		if !settings.Enabled {
-			resp = fmt.Sprintf("The %q command is currently disabled on this server or channel. *(Control panel to enable/disable <https://%s>)*", yc.Name, common.ConfHost.GetString())
+			resp = ReasonCommandDisabaledSettings
 			return
 		}
 
@@ -350,14 +362,14 @@ func (yc *YAGCommand) checkCanExecuteCommand(data *dcmd.Data, cState *dstate.Cha
 			}
 
 			if !found {
-				resp = "Missing a required role set up by the server admins for this command."
+				resp = ReasonMissingRole
 				return
 			}
 		}
 
 		for _, ignored := range settings.IgnoreRoles {
 			if common.ContainsInt64Slice(member.Roles, ignored) {
-				resp = "One of your roles is set up to be ignored by the server admins."
+				resp = ReasonIgnoredRole
 				return
 			}
 		}
@@ -367,7 +379,7 @@ func (yc *YAGCommand) checkCanExecuteCommand(data *dcmd.Data, cState *dstate.Cha
 			var perms int
 			perms, err = cState.Guild.MemberPermissionsMS(true, cState.ID, member)
 			if err != nil {
-				resp = "Unable to check permissions"
+				resp = ReasonError
 				return
 			}
 
@@ -380,7 +392,7 @@ func (yc *YAGCommand) checkCanExecuteCommand(data *dcmd.Data, cState *dstate.Cha
 			}
 
 			if !foundMatch {
-				resp = "Missing required permissions to use this command (" + yc.humanizedRequiredPerms() + ")"
+				resp = ReasonUserMissingPerms
 				return
 			}
 		}
@@ -391,14 +403,14 @@ func (yc *YAGCommand) checkCanExecuteCommand(data *dcmd.Data, cState *dstate.Cha
 	}
 
 	// Check the command cooldown
-	cdLeft, err := yc.CooldownLeft(data.Msg.Author.ID)
+	cdLeft, err := yc.LongestCooldownLeft(data.ContainerChain, data.Msg.Author.ID, data.Msg.GuildID)
 	if err != nil {
 		// Just pretend the cooldown is off...
-		logger.WithError(err).WithField("author", data.Msg.Author.ID).Error("Failed checking command cooldown")
+		yc.Logger(data).Error("Failed checking command cooldown")
 	}
 
 	if cdLeft > 0 {
-		resp = fmt.Sprintf("**%q:** You need to wait %d seconds before you can use the %q command again", common.EscapeSpecialMentions(data.Msg.Author.Username), cdLeft, yc.Name)
+		resp = ReasonCooldown
 		return
 	}
 
@@ -595,34 +607,94 @@ OUTER:
 	}
 }
 
-// CooldownLeft returns the number of seconds before a command can be used again
-func (cs *YAGCommand) CooldownLeft(userID int64) (int, error) {
-	if cs.Cooldown < 1 || common.Testing {
+// LongestCooldownLeft returns the longest cooldown for this command, either user scoped or guild scoped
+func (cs *YAGCommand) LongestCooldownLeft(cc []*dcmd.Container, userID int64, guildID int64) (int, error) {
+	cdUser, err := cs.UserScopeCooldownLeft(cc, userID)
+	if err != nil {
+		return 0, err
+	}
+
+	cdGuild, err := cs.GuildScopeCooldownLeft(cc, guildID)
+	if err != nil {
+		return 0, err
+	}
+
+	if cdUser > cdGuild {
+		return cdUser, nil
+	}
+
+	return cdGuild, nil
+}
+
+// UserScopeCooldownLeft returns the number of seconds before a command can be used again by this user
+func (cs *YAGCommand) UserScopeCooldownLeft(cc []*dcmd.Container, userID int64) (int, error) {
+	if cs.Cooldown < 1 {
 		return 0, nil
 	}
 
 	var ttl int
-	err := common.RedisPool.Do(retryableredis.Cmd(&ttl, "TTL", RKeyCommandCooldown(userID, cs.Name)))
-	if ttl < 1 {
+	err := common.RedisPool.Do(retryableredis.Cmd(&ttl, "TTL", RKeyCommandCooldown(userID, cs.FindNameFromContainerChain(cc))))
+	if err != nil {
+		return 0, errors.WithStackIf(err)
+	}
+
+	return ttl, nil
+}
+
+// GuildScopeCooldownLeft returns the number of seconds before a command can be used again on this server
+func (cs *YAGCommand) GuildScopeCooldownLeft(cc []*dcmd.Container, guildID int64) (int, error) {
+	if cs.GuildScopeCooldown < 1 {
 		return 0, nil
 	}
 
-	return ttl, err
+	var ttl int
+	err := common.RedisPool.Do(retryableredis.Cmd(&ttl, "TTL", RKeyCommandCooldownGuild(guildID, cs.FindNameFromContainerChain(cc))))
+	if err != nil {
+		return 0, errors.WithStackIf(err)
+	}
+
+	return ttl, nil
+}
+
+// SetCooldowns is a helper that serts both User and Guild cooldown
+func (cs *YAGCommand) SetCooldowns(cc []*dcmd.Container, userID int64, guildID int64) error {
+	err := cs.SetCooldownUser(cc, userID)
+	if err != nil {
+		return errors.WithStackIf(err)
+	}
+
+	err = cs.SetCooldownGuild(cc, guildID)
+	if err != nil {
+		return errors.WithStackIf(err)
+	}
+
+	return nil
 }
 
-// SetCooldown sets the cooldown of the command as it's defined in the struct
-func (cs *YAGCommand) SetCooldown(userID int64) error {
+// SetCooldownUser sets the user scoped cooldown of the command as it's defined in the struct
+func (cs *YAGCommand) SetCooldownUser(cc []*dcmd.Container, userID int64) error {
 	if cs.Cooldown < 1 {
 		return nil
 	}
 	now := time.Now().Unix()
 
-	err := common.RedisPool.Do(retryableredis.FlatCmd(nil, "SET", RKeyCommandCooldown(userID, cs.Name), now, "EX", cs.Cooldown))
-	return err
+	err := common.RedisPool.Do(retryableredis.FlatCmd(nil, "SET", RKeyCommandCooldown(userID, cs.FindNameFromContainerChain(cc)), now, "EX", cs.Cooldown))
+	return errors.WithStackIf(err)
+}
+
+// SetCooldownGuild sets the guild scoped cooldown of the command as it's defined in the struct
+func (cs *YAGCommand) SetCooldownGuild(cc []*dcmd.Container, guildID int64) error {
+	if cs.GuildScopeCooldown < 1 {
+		return nil
+	}
+
+	now := time.Now().Unix()
+	err := common.RedisPool.Do(retryableredis.FlatCmd(nil, "SET", RKeyCommandCooldownGuild(guildID, cs.FindNameFromContainerChain(cc)), now, "EX", cs.GuildScopeCooldown))
+	return errors.WithStackIf(err)
 }
 
 func (yc *YAGCommand) Logger(data *dcmd.Data) *logrus.Entry {
-	l := logger.WithField("cmd", yc.Name)
+	l := logger.WithField("cmd", yc.FindNameFromContainerChain(data.ContainerChain))
 	if data != nil {
 		if data.Msg != nil {
 			l = l.WithField("user_n", data.Msg.Author.Username)
@@ -632,6 +704,10 @@ func (yc *YAGCommand) Logger(data *dcmd.Data) *logrus.Entry {
 		if data.CS != nil {
 			l = l.WithField("channel", data.CS.ID)
 		}
+
+		if data.GS != nil {
+			l = l.WithField("guild", data.GS.ID)
+		}
 	}
 
 	return l
@@ -720,3 +796,24 @@ func removeRunningCommand(guildID, channelID, authorID int64, cmd *YAGCommand) {
 
 	return
 }
+
+func (yc *YAGCommand) FindNameFromContainerChain(cc []*dcmd.Container) string {
+	name := ""
+	for _, v := range cc {
+		if len(v.Names) < 1 {
+			continue
+		}
+
+		if name != "" {
+			name += " "
+		}
+
+		name += v.Names[0]
+	}
+
+	if name != "" {
+		name += " "
+	}
+
+	return name + yc.Name
+}

stdcommands/mentionrole/mentionrole.go

@@ -70,7 +70,8 @@ var Command = &commands.YAGCommand{
 	ArgSwitches: []*dcmd.ArgDef{
 		&dcmd.ArgDef{Switch: "channel", Help: "Optional channel to send in", Type: dcmd.Channel},
 	},
-	RunFunc: cmdFuncMentionRole,
+	RunFunc:            cmdFuncMentionRole,
+	GuildScopeCooldown: 10,
 }
 
 func cmdFuncMentionRole(data *dcmd.Data) (interface{}, error) {