Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Rewrite permission checks #325

Merged
merged 1 commit into from
Nov 13, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions src/builtins/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down
92 changes: 16 additions & 76 deletions src/dispatch/common.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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>(
Expand Down Expand Up @@ -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
Expand Down
1 change: 1 addition & 0 deletions src/dispatch/mod.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
//! Contains all code to dispatch incoming events onto framework commands

mod common;
mod permissions;
mod prefix;
mod slash;

Expand Down
82 changes: 82 additions & 0 deletions src/dispatch/permissions.rs
Original file line number Diff line number Diff line change
@@ -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))
}
23 changes: 23 additions & 0 deletions src/dispatch/permissions/application.rs
Original file line number Diff line number Diff line change
@@ -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),
}
}
62 changes: 62 additions & 0 deletions src/dispatch/permissions/prefix/cache.rs
Original file line number Diff line number Diff line change
@@ -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
}
}
40 changes: 40 additions & 0 deletions src/dispatch/permissions/prefix/http.rs
Original file line number Diff line number Diff line change
@@ -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,
})
}
12 changes: 12 additions & 0 deletions src/dispatch/permissions/prefix/mod.rs
Original file line number Diff line number Diff line change
@@ -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;
Loading
Loading