Browse Source

Fix potential conflicts in snowflake keys

This can happen on really old channels with the same ID as the guild ID
and having a command with both a role and a channel.
pull/7492/head
Rapptz 3 years ago
parent
commit
cdb7b3728e
  1. 4
      discord/app_commands/commands.py
  2. 63
      discord/app_commands/namespace.py
  3. 10
      discord/app_commands/tree.py

4
discord/app_commands/commands.py

@ -50,7 +50,7 @@ import re
from .enums import AppCommandOptionType, AppCommandType from .enums import AppCommandOptionType, AppCommandType
from ..interactions import Interaction from ..interactions import Interaction
from ..enums import ChannelType, try_enum from ..enums import ChannelType, try_enum
from .models import Choice from .models import AppCommandChannel, AppCommandThread, Choice
from .errors import CommandSignatureMismatch, CommandAlreadyRegistered from .errors import CommandSignatureMismatch, CommandAlreadyRegistered
from ..utils import resolve_annotation, MISSING, is_inside_class from ..utils import resolve_annotation, MISSING, is_inside_class
from ..user import User from ..user import User
@ -195,6 +195,8 @@ annotation_to_option_type: Dict[Any, AppCommandOptionType] = {
User: AppCommandOptionType.user, User: AppCommandOptionType.user,
Member: AppCommandOptionType.user, Member: AppCommandOptionType.user,
Role: AppCommandOptionType.role, Role: AppCommandOptionType.role,
AppCommandChannel: AppCommandOptionType.channel,
AppCommandThread: AppCommandOptionType.channel,
# StageChannel: AppCommandOptionType.channel, # StageChannel: AppCommandOptionType.channel,
# StoreChannel: AppCommandOptionType.channel, # StoreChannel: AppCommandOptionType.channel,
# VoiceChannel: AppCommandOptionType.channel, # VoiceChannel: AppCommandOptionType.channel,

63
discord/app_commands/namespace.py

@ -24,7 +24,7 @@ DEALINGS IN THE SOFTWARE.
from __future__ import annotations from __future__ import annotations
from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Tuple from typing import TYPE_CHECKING, Any, Dict, Iterable, List, NamedTuple, Tuple
from ..interactions import Interaction from ..interactions import Interaction
from ..member import Member from ..member import Member
from ..object import Object from ..object import Object
@ -32,11 +32,35 @@ from ..role import Role
from ..message import Message, Attachment from ..message import Message, Attachment
from ..channel import PartialMessageable from ..channel import PartialMessageable
from .models import AppCommandChannel, AppCommandThread from .models import AppCommandChannel, AppCommandThread
from .enums import AppCommandOptionType
if TYPE_CHECKING: if TYPE_CHECKING:
from ..types.interactions import ResolvedData, ApplicationCommandInteractionDataOption from ..types.interactions import ResolvedData, ApplicationCommandInteractionDataOption
class ResolveKey(NamedTuple):
id: str
# CommandOptionType does not use 0 or negative numbers so those can be safe for library
# internal use, if necessary. Likewise, only 6, 7, 8, and 11 are actually in use.
type: int
@classmethod
def any_with(cls, id: str) -> ResolveKey:
return ResolveKey(id=id, type=-1)
def __eq__(self, o: object) -> bool:
if not isinstance(o, ResolveKey):
return NotImplemented
if self.type == -1 or o.type == -1:
return self.id == o.id
return (self.id, self.type) == (o.id, o.type)
def __hash__(self) -> int:
# Most of the time an ID lookup is all that is necessary
# In case of collision then we look up both the ID and the type.
return hash(self.id)
class Namespace: class Namespace:
"""An object that holds the parameters being passed to a command in a mostly raw state. """An object that holds the parameters being passed to a command in a mostly raw state.
@ -103,47 +127,62 @@ class Namespace:
elif opt_type in (6, 7, 8, 9, 11): elif opt_type in (6, 7, 8, 9, 11):
# Remaining ones should be snowflake based ones with resolved data # Remaining ones should be snowflake based ones with resolved data
snowflake: str = option['value'] # type: ignore -- Key is there snowflake: str = option['value'] # type: ignore -- Key is there
value = completed.get(snowflake) if opt_type == 9: # Mentionable
# Mentionable is User | Role, these do not cause any conflict
key = ResolveKey.any_with(snowflake)
else:
# The remaining keys can conflict, for example, a role and a channel
# could end up with the same ID in very old guilds since they used to default
# to sharing the guild ID. Old general channels no longer exist, but some old
# servers will still have them so this needs to be handled.
key = ResolveKey(id=snowflake, type=opt_type)
value = completed.get(key)
self.__dict__[name] = value self.__dict__[name] = value
@classmethod @classmethod
def _get_resolved_items(cls, interaction: Interaction, resolved: ResolvedData) -> Dict[str, Any]: def _get_resolved_items(cls, interaction: Interaction, resolved: ResolvedData) -> Dict[ResolveKey, Any]:
completed: Dict[str, Any] = {} completed: Dict[ResolveKey, Any] = {}
state = interaction._state state = interaction._state
members = resolved.get('members', {}) members = resolved.get('members', {})
guild_id = interaction.guild_id guild_id = interaction.guild_id
guild = (state._get_guild(guild_id) or Object(id=guild_id)) if guild_id is not None else None guild = (state._get_guild(guild_id) or Object(id=guild_id)) if guild_id is not None else None
type = AppCommandOptionType.user.value
for (user_id, user_data) in resolved.get('users', {}).items(): for (user_id, user_data) in resolved.get('users', {}).items():
try: try:
member_data = members[user_id] member_data = members[user_id]
except KeyError: except KeyError:
completed[user_id] = state.create_user(user_data) completed[ResolveKey(id=user_id, type=type)] = state.create_user(user_data)
else: else:
member_data['user'] = user_data member_data['user'] = user_data
# Guild ID can't be None in this case. # Guild ID can't be None in this case.
# There's a type mismatch here that I don't actually care about # There's a type mismatch here that I don't actually care about
member = Member(state=state, guild=guild, data=member_data) # type: ignore member = Member(state=state, guild=guild, data=member_data) # type: ignore
completed[user_id] = member completed[ResolveKey(id=user_id, type=type)] = member
type = AppCommandOptionType.role.value
completed.update( completed.update(
{ {
# The guild ID can't be None in this case. # The guild ID can't be None in this case.
role_id: Role(guild=guild, state=state, data=role_data) # type: ignore ResolveKey(id=role_id, type=type): Role(guild=guild, state=state, data=role_data) # type: ignore
for role_id, role_data in resolved.get('roles', {}).items() for role_id, role_data in resolved.get('roles', {}).items()
} }
) )
type = AppCommandOptionType.channel.value
for (channel_id, channel_data) in resolved.get('channels', {}).items(): for (channel_id, channel_data) in resolved.get('channels', {}).items():
key = ResolveKey(id=channel_id, type=type)
if channel_data['type'] in (10, 11, 12): if channel_data['type'] in (10, 11, 12):
# The guild ID can't be none in this case # The guild ID can't be none in this case
completed[channel_id] = AppCommandThread(state=state, data=channel_data, guild_id=guild_id) # type: ignore completed[key] = AppCommandThread(state=state, data=channel_data, guild_id=guild_id) # type: ignore
else: else:
# The guild ID can't be none in this case # The guild ID can't be none in this case
completed[channel_id] = AppCommandChannel(state=state, data=channel_data, guild_id=guild_id) # type: ignore completed[key] = AppCommandChannel(state=state, data=channel_data, guild_id=guild_id) # type: ignore
type = AppCommandOptionType.attachment.value
completed.update( completed.update(
{ {
attachment_id: Attachment(data=attachment_data, state=state) ResolveKey(id=attachment_id, type=type): Attachment(data=attachment_data, state=state)
for attachment_id, attachment_data in resolved.get('attachments', {}).items() for attachment_id, attachment_data in resolved.get('attachments', {}).items()
} }
) )
@ -157,7 +196,9 @@ class Namespace:
channel = guild.get_channel_or_thread(channel_id) or PartialMessageable(state=state, id=channel_id) channel = guild.get_channel_or_thread(channel_id) or PartialMessageable(state=state, id=channel_id)
# Type checker doesn't understand this due to failure to narrow # Type checker doesn't understand this due to failure to narrow
completed[message_id] = Message(state=state, channel=channel, data=message_data) # type: ignore message = Message(state=state, channel=channel, data=message_data) # type: ignore
key = ResolveKey(id=message_id, type=-1)
completed[key] = message
return completed return completed

10
discord/app_commands/tree.py

@ -27,7 +27,7 @@ import inspect
from typing import Callable, Dict, List, Literal, Optional, TYPE_CHECKING, Tuple, Type, Union, overload from typing import Callable, Dict, List, Literal, Optional, TYPE_CHECKING, Tuple, Type, Union, overload
from .namespace import Namespace from .namespace import Namespace, ResolveKey
from .models import AppCommand from .models import AppCommand
from .commands import Command, ContextMenu, Group, _shorten from .commands import Command, ContextMenu, Group, _shorten
from .enums import AppCommandType from .enums import AppCommandType
@ -532,8 +532,14 @@ class CommandTree:
raise CommandNotFound(name, [], AppCommandType(type)) raise CommandNotFound(name, [], AppCommandType(type))
resolved = Namespace._get_resolved_items(interaction, data.get('resolved', {})) resolved = Namespace._get_resolved_items(interaction, data.get('resolved', {}))
target_id = data.get('target_id')
# Right now, the only types are message and user
# Therefore, there's no conflict with snowflakes
# This will always work at runtime # This will always work at runtime
value = resolved.get(data.get('target_id')) # type: ignore key = ResolveKey.any_with(target_id) # type: ignore
value = resolved.get(key)
if ctx_menu.type.value != type: if ctx_menu.type.value != type:
raise CommandSignatureMismatch(ctx_menu) raise CommandSignatureMismatch(ctx_menu)

Loading…
Cancel
Save