Rewrite permission checks (#325)

Co-authored-by: jamesbt365 <[email protected]>

src/builtins/mod.rs

@@ -173,6 +173,10 @@ pub async fn on_error<U, E: std::fmt::Display + std::fmt::Debug>(
             ctx.send(CreateReply::default().content(response).ephemeral(true))
                 .await?;
         }
+        crate::FrameworkError::PermissionFetchFailed { ctx } => {
+            ctx.say("An error occurred when fetching permissions.")
+                .await?;
+        }
         crate::FrameworkError::NotAnOwner { ctx } => {
             let response = "Only bot owners can call this command";
             ctx.send(CreateReply::default().content(response).ephemeral(true))

src/dispatch/common.rs

@@ -2,60 +2,6 @@
 
 use crate::serenity_prelude as serenity;
 
-/// Retrieves user permissions in the given channel. If unknown, returns None. If in DMs, returns
-/// `Permissions::all()`.
-async fn user_permissions(
-    ctx: &serenity::Context,
-    guild_id: Option<serenity::GuildId>,
-    channel_id: serenity::ChannelId,
-    user_id: serenity::UserId,
-) -> Option<serenity::Permissions> {
-    let guild_id = match guild_id {
-        Some(x) => x,
-        None => return Some(serenity::Permissions::all()), // no permission checks in DMs
-    };
-
-    let guild = guild_id.to_partial_guild(ctx).await.ok()?;
-
-    // Use to_channel so that it can fallback on HTTP for threads (which aren't in cache usually)
-    let channel = match channel_id.to_channel(ctx).await {
-        Ok(serenity::Channel::Guild(channel)) => channel,
-        Ok(_other_channel) => {
-            tracing::warn!(
-                "guild message was supposedly sent in a non-guild channel. Denying invocation"
-            );
-            return None;
-        }
-        Err(_) => return None,
-    };
-
-    let member = guild.member(ctx, user_id).await.ok()?;
-
-    Some(guild.user_permissions_in(&channel, &member))
-}
-
-/// Retrieves the set of permissions that are lacking, relative to the given required permission set
-///
-/// Returns None if permissions couldn't be retrieved
-async fn missing_permissions<U, E>(
-    ctx: crate::Context<'_, U, E>,
-    user: serenity::UserId,
-    required_permissions: serenity::Permissions,
-) -> Option<serenity::Permissions> {
-    if required_permissions.is_empty() {
-        return Some(serenity::Permissions::empty());
-    }
-
-    let permissions = user_permissions(
-        ctx.serenity_context(),
-        ctx.guild_id(),
-        ctx.channel_id(),
-        user,
-    )
-    .await;
-    Some(required_permissions - permissions?)
-}
-
 /// See [`check_permissions_and_cooldown`]. Runs the check only for a single command. The caller
 /// should call this multiple time for each parent command to achieve the check inheritance logic.
 async fn check_permissions_and_cooldown_single<'a, U, E>(
@@ -111,35 +57,29 @@ async fn check_permissions_and_cooldown_single<'a, U, E>(
     }
 
     // Make sure that user has required permissions
-    match missing_permissions(ctx, ctx.author().id, cmd.required_permissions).await {
-        Some(missing_permissions) if missing_permissions.is_empty() => {}
-        Some(missing_permissions) => {
-            return Err(crate::FrameworkError::MissingUserPermissions {
-                ctx,
-                missing_permissions: Some(missing_permissions),
-            })
-        }
-        // Better safe than sorry: when perms are unknown, restrict access
-        None => {
+    if let Some((user_missing_permissions, bot_missing_permissions)) =
+        super::permissions::calculate_missing(
+            ctx,
+            cmd.required_permissions,
+            cmd.required_bot_permissions,
+        )
+        .await
+    {
+        if !user_missing_permissions.is_empty() {
             return Err(crate::FrameworkError::MissingUserPermissions {
                 ctx,
-                missing_permissions: None,
-            })
+                missing_permissions: Some(user_missing_permissions),
+            });
         }
-    }
 
-    // Before running any pre-command checks, make sure the bot has the permissions it needs
-    match missing_permissions(ctx, ctx.framework().bot_id(), cmd.required_bot_permissions).await {
-        Some(missing_permissions) if missing_permissions.is_empty() => {}
-        Some(missing_permissions) => {
+        if !bot_missing_permissions.is_empty() {
             return Err(crate::FrameworkError::MissingBotPermissions {
                 ctx,
-                missing_permissions,
-            })
+                missing_permissions: bot_missing_permissions,
+            });
         }
-        // When in doubt, just let it run. Not getting fancy missing permissions errors is better
-        // than the command not executing at all
-        None => {}
+    } else {
+        return Err(crate::FrameworkError::PermissionFetchFailed { ctx });
     }
 
     // Only continue if command checks returns true

src/dispatch/mod.rs

@@ -1,6 +1,7 @@
 //! Contains all code to dispatch incoming events onto framework commands
 
 mod common;
+mod permissions;
 mod prefix;
 mod slash;
 

src/dispatch/permissions.rs

@@ -0,0 +1,82 @@
+//! Module for calculating permission checks for commands
+
+use crate::serenity_prelude as serenity;
+
+mod application;
+mod prefix;
+
+/// Simple POD type to hold the results of permission lookups.
+struct PermissionsInfo {
+    /// The Permissions of the author, if requested.
+    author_permissions: Option<serenity::Permissions>,
+    /// The Permissions of the bot, if requested.
+    bot_permissions: Option<serenity::Permissions>,
+}
+
+impl PermissionsInfo {
+    /// Returns the fake permissions info from a DM.
+    fn dm_permissions() -> Self {
+        Self {
+            author_permissions: Some(serenity::Permissions::dm_permissions()),
+            bot_permissions: Some(serenity::Permissions::dm_permissions()),
+        }
+    }
+}
+
+/// Retrieves the permissions for the context author and the bot.
+async fn get_author_and_bot_permissions<U, E>(
+    ctx: crate::Context<'_, U, E>,
+    skip_author: bool,
+    skip_bot: bool,
+) -> Option<PermissionsInfo> {
+    // No permission checks in DMs.
+    let Some(guild_id) = ctx.guild_id() else {
+        return Some(PermissionsInfo::dm_permissions());
+    };
+
+    match ctx {
+        crate::Context::Application(ctx) => {
+            Some(application::get_author_and_bot_permissions(ctx.interaction))
+        }
+        crate::Context::Prefix(ctx) => {
+            prefix::get_author_and_bot_permissions(ctx, guild_id, skip_author, skip_bot).await
+        }
+    }
+}
+
+/// Retrieves the set of permissions that are lacking, relative to the given required permission set
+///
+/// Returns None if permissions couldn't be retrieved.
+pub(super) async fn calculate_missing<U, E>(
+    ctx: crate::Context<'_, U, E>,
+    author_required_permissions: serenity::Permissions,
+    bot_required_permissions: serenity::Permissions,
+) -> Option<(serenity::Permissions, serenity::Permissions)> {
+    // If both user and bot are None, return empty permissions
+    if author_required_permissions.is_empty() && bot_required_permissions.is_empty() {
+        return Some((
+            serenity::Permissions::empty(),
+            serenity::Permissions::empty(),
+        ));
+    }
+
+    // Fetch permissions, returning None if an error occurred
+    let permissions = get_author_and_bot_permissions(
+        ctx,
+        author_required_permissions.is_empty(),
+        bot_required_permissions.is_empty(),
+    )
+    .await?;
+
+    let author_missing_perms = permissions
+        .author_permissions
+        .map(|permissions| author_required_permissions - permissions)
+        .unwrap_or_default();
+
+    let bot_missing_perms = permissions
+        .bot_permissions
+        .map(|permissions| bot_required_permissions - permissions)
+        .unwrap_or_default();
+
+    Some((author_missing_perms, bot_missing_perms))
+}

src/dispatch/permissions/application.rs

@@ -0,0 +1,23 @@
+//! Application command permissions calculation
+use crate::serenity_prelude as serenity;
+
+use super::PermissionsInfo;
+
+/// Gets the permissions of the ctx author and the bot.
+pub(super) fn get_author_and_bot_permissions(
+    interaction: &serenity::CommandInteraction,
+) -> PermissionsInfo {
+    let err = "member is Some if interaction is in guild";
+    let author_member = interaction.member.as_ref().expect(err);
+
+    let err = "should always be some as inside interaction";
+    let author_permissions = author_member.permissions.expect(err);
+
+    let err = "should always be some according to discord docs";
+    let bot_permissions = interaction.app_permissions.expect(err);
+
+    PermissionsInfo {
+        author_permissions: Some(author_permissions),
+        bot_permissions: Some(bot_permissions),
+    }
+}

src/dispatch/permissions/prefix/cache.rs

@@ -0,0 +1,62 @@
+//! The cache variant of prefix permissions calculation
+
+use crate::{serenity_prelude as serenity, PrefixContext};
+
+use crate::dispatch::permissions::PermissionsInfo;
+
+/// Gets the permissions of the ctx author and the bot.
+pub(in crate::dispatch::permissions) async fn get_author_and_bot_permissions<U, E>(
+    ctx: PrefixContext<'_, U, E>,
+    guild_id: serenity::GuildId,
+    skip_author: bool,
+    skip_bot: bool,
+) -> Option<PermissionsInfo> {
+    // Should only fail if the guild is not cached, which is fair to bail on.
+    let guild = ctx.cache().guild(guild_id)?;
+
+    let author_permissions = if skip_author {
+        None
+    } else {
+        let err_msg = "should only fail if the guild is not cached";
+        Some(ctx.msg.author_permissions(ctx.cache()).expect(err_msg))
+    };
+
+    let bot_permissions = if skip_bot {
+        None
+    } else {
+        let channel_id = ctx.channel_id();
+        let bot_user_id = ctx.framework.bot_id();
+        Some(get_bot_permissions(&guild, channel_id, bot_user_id)?)
+    };
+
+    Some(PermissionsInfo {
+        author_permissions,
+        bot_permissions,
+    })
+}
+
+/// Gets the permissions for the bot.
+fn get_bot_permissions(
+    guild: &serenity::Guild,
+    channel_id: serenity::ChannelId,
+    bot_id: serenity::UserId,
+) -> Option<serenity::Permissions> {
+    // Should never fail, as the bot member is always cached
+    let bot_member = guild.members.get(&bot_id)?;
+
+    if let Some(channel) = guild.channels.get(&channel_id) {
+        Some(guild.user_permissions_in(channel, bot_member))
+    } else if let Some(thread) = guild.threads.iter().find(|th| th.id == channel_id) {
+        let err = "parent id should always be Some for thread";
+        let parent_channel_id = thread.parent_id.expect(err);
+
+        let parent_channel = guild.channels.get(&parent_channel_id)?;
+        Some(guild.user_permissions_in(parent_channel, bot_member))
+    } else {
+        // The message was either:
+        // - Sent in a guild with broken caching
+        // - Not set in a channel or thread?
+        tracing::warn!("Could not find channel/thread ({channel_id}) for permissions check in cache for guild: {}", guild.id);
+        None
+    }
+}

src/dispatch/permissions/prefix/http.rs

@@ -0,0 +1,40 @@
+//! The cache variant of prefix permissions calculation
+
+use crate::{serenity_prelude as serenity, PrefixContext};
+
+use crate::dispatch::permissions::PermissionsInfo;
+
+/// Gets the permissions of the ctx author and the bot.
+pub(in crate::dispatch::permissions) async fn get_author_and_bot_permissions<U, E>(
+    ctx: PrefixContext<'_, U, E>,
+    guild_id: serenity::GuildId,
+    skip_author: bool,
+    skip_bot: bool,
+) -> Option<PermissionsInfo> {
+    let http = ctx.http();
+    let guild = guild_id.to_partial_guild(http).await.ok()?;
+    let guild_channel = {
+        let channel = ctx.http().get_channel(ctx.channel_id()).await.ok()?;
+        channel.guild().expect("channel should be a guild channel")
+    };
+
+    let bot_permissions = if skip_bot {
+        None
+    } else {
+        let bot_member = guild.id.member(http, ctx.framework.bot_id).await.ok()?;
+        Some(guild.user_permissions_in(&guild_channel, &bot_member))
+    };
+
+    let author_permissions = if skip_author {
+        None
+    } else {
+        let err = "should always be Some in MessageCreateEvent";
+        let author_member = ctx.msg.member.as_ref().expect(err);
+        Some(guild.partial_member_permissions_in(&guild_channel, ctx.author().id, author_member))
+    };
+
+    Some(PermissionsInfo {
+        author_permissions,
+        bot_permissions,
+    })
+}

src/dispatch/permissions/prefix/mod.rs

@@ -0,0 +1,12 @@
+//! Prefix command permissions calculation
+
+#[cfg(feature = "cache")]
+mod cache;
+#[cfg(not(feature = "cache"))]
+mod http;
+
+#[cfg(feature = "cache")]
+pub(super) use cache::get_author_and_bot_permissions;
+
+#[cfg(not(feature = "cache"))]
+pub(super) use http::get_author_and_bot_permissions;