From f83f378528d10a729d5ca1673295e819f64e929f Mon Sep 17 00:00:00 2001 From: Toby Harradine Date: Sun, 14 Jul 2019 10:47:40 +1000 Subject: [PATCH] [Core] Make Requires.verify() wait until rules are loaded (#2857) * Make Requires.verify() wait until rules are loaded Also ensures `Requires` objects are reset when unloaded, particularly in case a `Command` object manages to stay in memory between cog unload and load, and its permissions rules change between those events. Also, this PR re-ordered some of the event loop policy stuff, because it was required that the event loop policy be set before creating any `Requires` objects. This may or may not have an effect on other `get_event_loop()` calls elsewhere (either in our code, a dependency's, or asyncio's). Either way, these effects would be a *correction*, and any bugs that arise from it are likely to have been occurring silently beforehand. Signed-off-by: Toby Harradine * Remove calls to `remove_listener()` in permissions Signed-off-by: Toby Harradine * Fix adding rules for permissions cog/commands itself Also addresses feedback Signed-off-by: Toby Harradine * Clean up indentation when setting uvloop policy Signed-off-by: Toby Harradine * Use `set(walk_commands())` to traverse `Group` subcommands Signed-off-by: Toby Harradine --- redbot/__main__.py | 27 +++++++++-------- redbot/cogs/permissions/__init__.py | 11 ++++--- redbot/cogs/permissions/permissions.py | 36 ++++++++++++++--------- redbot/core/bot.py | 40 +++++++++++++++++--------- redbot/core/commands/requires.py | 19 ++++++++++++ 5 files changed, 88 insertions(+), 45 deletions(-) diff --git a/redbot/__main__.py b/redbot/__main__.py index 602a77489..6f7a6935d 100644 --- a/redbot/__main__.py +++ b/redbot/__main__.py @@ -10,6 +10,21 @@ import sys import discord +# Set the event loop policies here so any subsequent `get_event_loop()` +# calls, in particular those as a result of the following imports, +# return the correct loop object. +if sys.platform == "win32": + asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy()) +elif sys.implementation.name == "cpython": + # Let's not force this dependency, uvloop is much faster on cpython + try: + import uvloop + except ImportError: + uvloop = None + pass + else: + asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) + import redbot.logging from redbot.core.bot import Red, ExitCodes from redbot.core.cog_manager import CogManagerUI @@ -21,18 +36,6 @@ from redbot.core.dev_commands import Dev from redbot.core import __version__, modlog, bank, data_manager from signal import SIGTERM -# Let's not force this dependency, uvloop is much faster on cpython -if sys.implementation.name == "cpython": - try: - import uvloop - except ImportError: - uvloop = None - pass - else: - asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) - -if sys.platform == "win32": - asyncio.set_event_loop(asyncio.ProactorEventLoop()) log = logging.getLogger("red.main") diff --git a/redbot/cogs/permissions/__init__.py b/redbot/cogs/permissions/__init__.py index 2159535c4..fa080cb1f 100644 --- a/redbot/cogs/permissions/__init__.py +++ b/redbot/cogs/permissions/__init__.py @@ -4,10 +4,9 @@ from .permissions import Permissions async def setup(bot): cog = Permissions(bot) await cog.initialize() - # It's important that these listeners are added prior to load, so - # the permissions commands themselves have rules added. - # Automatic listeners being added in add_cog happen in arbitrary - # order, so we want to circumvent that. - bot.add_listener(cog.red_cog_added, "on_cog_add") - bot.add_listener(cog.red_command_added, "on_command_add") + # We should add the rules for the Permissions cog and its own commands *before* adding the cog. + # The actual listeners ought to skip the ones we're passing here. + await cog._on_cog_add(cog) + for command in cog.__cog_commands__: + await cog._on_command_add(command) bot.add_cog(cog) diff --git a/redbot/cogs/permissions/permissions.py b/redbot/cogs/permissions/permissions.py index 6e14cb650..bb27669b1 100644 --- a/redbot/cogs/permissions/permissions.py +++ b/redbot/cogs/permissions/permissions.py @@ -436,31 +436,41 @@ class Permissions(commands.Cog): await self._clear_rules(guild_id=ctx.guild.id) await ctx.tick() - async def red_cog_added(self, cog: commands.Cog) -> None: + @commands.Cog.listener() + async def on_cog_add(self, cog: commands.Cog) -> None: """Event listener for `cog_add`. This loads rules whenever a new cog is added. - - Do not convert to using Cog.listener decorator !! - This *must* be added manually prior to cog load, and removed at unload """ + if cog is self: + # This cog has its rules loaded manually in setup() + return + await self._on_cog_add(cog) + + @commands.Cog.listener() + async def on_command_add(self, command: commands.Command) -> None: + """Event listener for `command_add`. + + This loads rules whenever a new command is added. + """ + if command.cog is self: + # This cog's commands have their rules loaded manually in setup() + return + await self._on_command_add(command) + + async def _on_cog_add(self, cog: commands.Cog) -> None: self._load_rules_for( cog_or_command=cog, rule_dict=await self.config.custom(COG, cog.__class__.__name__).all(), ) + cog.requires.ready_event.set() - async def red_command_added(self, command: commands.Command) -> None: - """Event listener for `command_add`. - - This loads rules whenever a new command is added. - - Do not convert to using Cog.listener decorator !! - This *must* be added manually prior to cog load, and removed at unload - """ + async def _on_command_add(self, command: commands.Command) -> None: self._load_rules_for( cog_or_command=command, rule_dict=await self.config.custom(COMMAND, command.qualified_name).all(), ) + command.requires.ready_event.set() async def _add_rule( self, rule: bool, cog_or_cmd: CogOrCommand, model_id: int, guild_id: int @@ -708,8 +718,6 @@ class Permissions(commands.Cog): cog_or_command.deny_to(model_id, guild_id=guild_id) def cog_unload(self) -> None: - self.bot.remove_listener(self.red_cog_added, "on_cog_add") - self.bot.remove_listener(self.red_command_added, "on_command_add") self.bot.loop.create_task(self._unload_all_rules()) async def _unload_all_rules(self) -> None: diff --git a/redbot/core/bot.py b/redbot/core/bot.py index a6d48a0c9..b54e11f77 100644 --- a/redbot/core/bot.py +++ b/redbot/core/bot.py @@ -322,6 +322,8 @@ class RedBase(commands.GroupMixin, commands.bot.BotBase, RPCMixin): # pylint: d super().remove_cog(cogname) + cog.requires.reset() + for meth in self.rpc_handlers.pop(cogname.upper(), ()): self.unregister_rpc_handler(meth) @@ -424,21 +426,10 @@ class RedBase(commands.GroupMixin, commands.bot.BotBase, RPCMixin): # pylint: d self.add_permissions_hook(hook) added_hooks.append(hook) - for command in cog.__cog_commands__: - - if not isinstance(command, commands.Command): - raise RuntimeError( - f"The {cog.__class__.__name__} cog in the {cog.__module__} package," - " is not using Red's command module, and cannot be added. " - "If this is your cog, please use `from redbot.core import commands`" - "in place of `from discord.ext import commands`. For more details on " - "this requirement, see this page: " - "http://red-discordbot.readthedocs.io/en/v3-develop/framework_commands.html" - ) super().add_cog(cog) self.dispatch("cog_add", cog) - for command in cog.__cog_commands__: - self.dispatch("command_add", command) + if "permissions" not in self.extensions: + cog.requires.ready_event.set() except Exception: for hook in added_hooks: try: @@ -452,6 +443,29 @@ class RedBase(commands.GroupMixin, commands.bot.BotBase, RPCMixin): # pylint: d del cog raise + def add_command(self, command: commands.Command) -> None: + if not isinstance(command, commands.Command): + raise RuntimeError("Commands must be instances of `redbot.core.commands.Command`") + + super().add_command(command) + + permissions_not_loaded = "permissions" not in self.extensions + self.dispatch("command_add", command) + if permissions_not_loaded: + command.requires.ready_event.set() + if isinstance(command, commands.Group): + for subcommand in set(command.walk_commands()): + self.dispatch("command_add", subcommand) + if permissions_not_loaded: + command.requires.ready_event.set() + + def remove_command(self, name: str) -> None: + command = super().remove_command(name) + command.requires.reset() + if isinstance(command, commands.Group): + for subcommand in set(command.walk_commands()): + subcommand.requires.reset() + def clear_permission_rules(self, guild_id: Optional[int]) -> None: """Clear all permission overrides in a scope. diff --git a/redbot/core/commands/requires.py b/redbot/core/commands/requires.py index db06bc1d5..75a0e351c 100644 --- a/redbot/core/commands/requires.py +++ b/redbot/core/commands/requires.py @@ -272,6 +272,12 @@ class Requires: `user_perms` will be used exclusively, otherwise, for levels other than bot owner, the user can still run the command if they have the required `user_perms`. + ready_event : asyncio.Event + Event for when this Requires object has had its rules loaded. + If permissions is loaded, this should be set when permissions + has finished loading rules into this object. If permissions + is not loaded, it should be set as soon as the command or cog + is added. user_perms : Optional[discord.Permissions] The required permissions for users to execute the command. Can be ``None``, in which case the `privilege_level` will be used @@ -300,6 +306,7 @@ class Requires: ): self.checks: List[CheckPredicate] = checks self.privilege_level: Optional[PrivilegeLevel] = privilege_level + self.ready_event = asyncio.Event() if isinstance(user_perms, dict): self.user_perms: Optional[discord.Permissions] = discord.Permissions.none() @@ -413,6 +420,16 @@ class Requires: if default is not None: rules[self.DEFAULT] = default + def reset(self) -> None: + """Reset this Requires object to its original state. + + This will clear all rules, including defaults. It also resets + the `Requires.ready_event`. + """ + self._guild_rules.clear() # pylint: disable=no-member + self._global_rules.clear() # pylint: disable=no-member + self.ready_event.clear() + async def verify(self, ctx: "Context") -> bool: """Check if the given context passes the requirements. @@ -438,6 +455,8 @@ class Requires: Propogated from any permissions checks. """ + if not self.ready_event.is_set(): + await self.ready_event.wait() await self._verify_bot(ctx) # Owner should never be locked out of commands for user permissions.