CopyBot/Red-DiscordBot
1
0
Fork 0
mirror of https://github.com/Cog-Creators/Red-DiscordBot.git synced 2025-11-06 03:08:55 -05:00
Code Issues Packages Projects Releases Wiki Activity

[Core] Make Requires.verify() wait until rules are loaded (#2857)

Browse Source 
* 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 <tobyharradine@gmail.com>

* Remove calls to `remove_listener()` in permissions

Signed-off-by: Toby Harradine <tobyharradine@gmail.com>

* Fix adding rules for permissions cog/commands itself

Also addresses feedback

Signed-off-by: Toby Harradine <tobyharradine@gmail.com>

* Clean up indentation when setting uvloop policy

Signed-off-by: Toby Harradine <tobyharradine@gmail.com>

* Use `set(walk_commands())` to traverse `Group` subcommands

Signed-off-by: Toby Harradine <tobyharradine@gmail.com>
This commit is contained in:
Toby Harradine 2019-07-14 10:47:40 +10:00 committed by Michael H
parent 21a6384ebf
commit f83f378528
5 changed files with 88 additions and 45 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
redbot/__main__.py
View File

@ -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")

11
redbot/cogs/permissions/__init__.py
View File

@ -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)

36
redbot/cogs/permissions/permissions.py
View File

@ -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:

40
redbot/core/bot.py
View File

@ -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.

19
redbot/core/commands/requires.py
View File

@ -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.