Add panic handler (#140)

* Restructure dispatch::slash functions like prefix

Forgot some stuff from restructguring

* Catch panics during command execution

- Adds FrameworkError::CommandPanic
    - The default on_error impl spawns a nice error embed
- Wraps all instances of run_invocation/run_interaction/run_autocomplete in catch_unwind
- Adds a `div` command to the testing crate to test the panic catching (by triggering a div by zero)

* Replace let-else with match to please MSRV

* Gate catch_unwind call behind feature flag

---------

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

Cargo.toml

@@ -47,12 +47,17 @@ intl-memoizer = "0.5.1"
 fluent-syntax = "0.11"
 
 [features]
-default = ["serenity/rustls_backend", "cache", "chrono"]
+default = ["serenity/rustls_backend", "cache", "chrono", "handle_panics"]
 chrono = ["serenity/chrono"]
 cache = ["serenity/cache"]
 time = ["serenity/time"]
 # No-op feature because serenity/collector is now enabled by default
 collector = []
+# Enables support for handling panics inside commands via FrameworkError::CommandPanic.
+# This feature has no overhead and can always be enabled.
+# This feature exists because some users want to disable the mere possibility of catching panics at
+# build time for peace of mind.
+handle_panics = []
 
 [package.metadata.docs.rs]
 all-features = true

examples/testing/main.rs

@@ -12,7 +12,7 @@ pub struct Data {}
 async fn main() {
     let framework = poise::Framework::builder()
         .options(poise::FrameworkOptions {
-            commands: vec![inherit_checks::parent(), misc::paginate()],
+            commands: vec![inherit_checks::parent(), misc::paginate(), misc::div()],
             prefix_options: poise::PrefixFrameworkOptions {
                 prefix: Some("~".into()),
                 ..Default::default()

examples/testing/misc.rs

@@ -14,3 +14,9 @@ pub async fn paginate(ctx: Context<'_>) -> Result<(), Error> {
 
     Ok(())
 }
+
+#[poise::command(slash_command, prefix_command)]
+pub async fn div(ctx: Context<'_>, a: i32, b: i32) -> Result<(), Error> {
+    ctx.say((a / b).to_string()).await?;
+    Ok(())
+}

src/builtins/mod.rs

@@ -48,6 +48,18 @@ pub async fn on_error<U, E: std::fmt::Display + std::fmt::Debug>(
             eprintln!("An error occured in a command: {}", error);
             ctx.say(error).await?;
         }
+        crate::FrameworkError::CommandPanic { ctx, payload: _ } => {
+            // Not showing the payload to the user because it may contain sensitive info
+            ctx.send(|b| {
+                b.embed(|b| {
+                    b.title("Internal error")
+                        .color((255, 0, 0))
+                        .description("An unexpected internal error has occurred.")
+                })
+                .ephemeral(true)
+            })
+            .await?;
+        }
         crate::FrameworkError::ArgumentParse { ctx, input, error } => {
             // If we caught an argument parse error, give a helpful error message with the
             // command explanation if available

src/dispatch/prefix.rs

@@ -203,7 +203,12 @@ pub async fn dispatch_message<'a, U: Send + Sync, E>(
     )
     .await?
     {
-        run_invocation(ctx).await?;
+        crate::catch_unwind_maybe(run_invocation(ctx))
+            .await
+            .map_err(|payload| crate::FrameworkError::CommandPanic {
+                payload,
+                ctx: ctx.into(),
+            })??;
     }
     Ok(())
 }

src/dispatch/slash.rs

@@ -31,9 +31,11 @@ fn find_matching_command<'a, 'b, U, E>(
     })
 }
 
-/// Given an interaction, finds the matching framework command and checks if the user is allowed
-/// access
-pub async fn extract_command_and_run_checks<'a, U, E>(
+/// Parses an `Interaction` into a [`crate::ApplicationContext`] using some context data.
+///
+/// After this, the [`crate::ApplicationContext`] should be passed into [`run_command`] or
+/// [`run_autocomplete`].
+fn extract_command<'a, U, E>(
     framework: crate::FrameworkContext<'a, U, E>,
     ctx: &'a serenity::Context,
     interaction: crate::ApplicationCommandOrAutocompleteInteraction<'a>,
@@ -54,8 +56,8 @@ pub async fn extract_command_and_run_checks<'a, U, E>(
             interaction,
         })?;
 
-    let ctx = crate::ApplicationContext {
-        data: framework.user_data().await,
+    Ok(crate::ApplicationContext {
+        data: framework.user_data,
         serenity_context: ctx,
         framework,
         interaction,
@@ -65,36 +67,39 @@ pub async fn extract_command_and_run_checks<'a, U, E>(
         has_sent_initial_response,
         invocation_data,
         __non_exhaustive: (),
-    };
-
-    super::common::check_permissions_and_cooldown(ctx.into()).await?;
-
-    Ok(ctx)
+    })
 }
 
-/// Dispatches this interaction onto framework commands, i.e. runs the associated command
-pub async fn dispatch_interaction<'a, U, E>(
+/// Given an interaction, finds the matching framework command and checks if the user is allowed
+/// access
+pub async fn extract_command_and_run_checks<'a, U, E>(
     framework: crate::FrameworkContext<'a, U, E>,
     ctx: &'a serenity::Context,
-    interaction: &'a serenity::ApplicationCommandInteraction,
-    // Need to pass this in from outside because of lifetime issues
+    interaction: crate::ApplicationCommandOrAutocompleteInteraction<'a>,
     has_sent_initial_response: &'a std::sync::atomic::AtomicBool,
-    // Need to pass this in from outside because of lifetime issues
     invocation_data: &'a tokio::sync::Mutex<Box<dyn std::any::Any + Send + Sync>>,
-    // Need to pass this in from outside because of lifetime issues
     parent_commands: &'a mut Vec<&'a crate::Command<U, E>>,
-) -> Result<(), crate::FrameworkError<'a, U, E>> {
-    let ctx = extract_command_and_run_checks(
+) -> Result<crate::ApplicationContext<'a, U, E>, crate::FrameworkError<'a, U, E>> {
+    let ctx = extract_command(
         framework,
         ctx,
-        crate::ApplicationCommandOrAutocompleteInteraction::ApplicationCommand(interaction),
+        interaction,
         has_sent_initial_response,
         invocation_data,
         parent_commands,
-    )
-    .await?;
+    )?;
+    super::common::check_permissions_and_cooldown(ctx.into()).await?;
+    Ok(ctx)
+}
+
+/// Given the extracted application command data from [`extract_command`], runs the command,
+/// including all the before and after code like checks.
+async fn run_command<U, E>(
+    ctx: crate::ApplicationContext<'_, U, E>,
+) -> Result<(), crate::FrameworkError<'_, U, E>> {
+    super::common::check_permissions_and_cooldown(ctx.into()).await?;
 
-    (framework.options.pre_command)(crate::Context::Application(ctx)).await;
+    (ctx.framework.options.pre_command)(crate::Context::Application(ctx)).await;
 
     // Check which interaction type we received and grab the command action and, if context menu,
     // the resolved click target, and execute the action
@@ -103,7 +108,7 @@ pub async fn dispatch_interaction<'a, U, E>(
         description: "received interaction type but command contained no \
                 matching action or interaction contained no matching context menu object",
     };
-    let action_result = match interaction.data.kind {
+    let action_result = match ctx.interaction.data().kind {
         serenity::CommandType::ChatInput => {
             let action = ctx
                 .command
@@ -112,7 +117,10 @@ pub async fn dispatch_interaction<'a, U, E>(
             action(ctx).await
         }
         serenity::CommandType::User => {
-            match (ctx.command.context_menu_action, &interaction.data.target()) {
+            match (
+                ctx.command.context_menu_action,
+                &ctx.interaction.data().target(),
+            ) {
                 (
                     Some(crate::ContextMenuCommandAction::User(action)),
                     Some(serenity::ResolvedTarget::User(user, _)),
@@ -121,7 +129,10 @@ pub async fn dispatch_interaction<'a, U, E>(
             }
         }
         serenity::CommandType::Message => {
-            match (ctx.command.context_menu_action, &interaction.data.target()) {
+            match (
+                ctx.command.context_menu_action,
+                &ctx.interaction.data().target(),
+            ) {
                 (
                     Some(crate::ContextMenuCommandAction::Message(action)),
                     Some(serenity::ResolvedTarget::Message(message)),
@@ -136,33 +147,48 @@ pub async fn dispatch_interaction<'a, U, E>(
     };
     action_result?;
 
-    (framework.options.post_command)(crate::Context::Application(ctx)).await;
+    (ctx.framework.options.post_command)(crate::Context::Application(ctx)).await;
 
     Ok(())
 }
 
-/// Dispatches this interaction onto framework commands, i.e. runs the associated autocomplete
-/// callback
-pub async fn dispatch_autocomplete<'a, U, E>(
+/// Dispatches this interaction onto framework commands, i.e. runs the associated command
+pub async fn dispatch_interaction<'a, U, E>(
     framework: crate::FrameworkContext<'a, U, E>,
     ctx: &'a serenity::Context,
-    interaction: &'a serenity::AutocompleteInteraction,
+    interaction: &'a serenity::ApplicationCommandInteraction,
     // Need to pass this in from outside because of lifetime issues
     has_sent_initial_response: &'a std::sync::atomic::AtomicBool,
     // Need to pass this in from outside because of lifetime issues
     invocation_data: &'a tokio::sync::Mutex<Box<dyn std::any::Any + Send + Sync>>,
     // Need to pass this in from outside because of lifetime issues
     parent_commands: &'a mut Vec<&'a crate::Command<U, E>>,
 ) -> Result<(), crate::FrameworkError<'a, U, E>> {
-    let ctx = extract_command_and_run_checks(
+    let ctx = extract_command(
         framework,
         ctx,
-        crate::ApplicationCommandOrAutocompleteInteraction::Autocomplete(interaction),
+        crate::ApplicationCommandOrAutocompleteInteraction::ApplicationCommand(interaction),
         has_sent_initial_response,
         invocation_data,
         parent_commands,
-    )
-    .await?;
+    )?;
+
+    crate::catch_unwind_maybe(run_command(ctx))
+        .await
+        .map_err(|payload| crate::FrameworkError::CommandPanic {
+            payload,
+            ctx: ctx.into(),
+        })??;
+
+    Ok(())
+}
+
+/// Given the extracted application command data from [`extract_command`], runs the autocomplete
+/// callbacks, including all the before and after code like checks.
+async fn run_autocomplete<U, E>(
+    ctx: crate::ApplicationContext<'_, U, E>,
+) -> Result<(), crate::FrameworkError<'_, U, E>> {
+    super::common::check_permissions_and_cooldown(ctx.into()).await?;
 
     // Find which parameter is focused by the user
     let focused_option = match ctx.args.iter().find(|o| o.focused) {
@@ -211,6 +237,14 @@ pub async fn dispatch_autocomplete<'a, U, E>(
         }
     };
 
+    let interaction = match ctx.interaction {
+        crate::ApplicationCommandOrAutocompleteInteraction::Autocomplete(x) => x,
+        _ => {
+            log::warn!("a non-autocomplete interaction was given to run_autocomplete()");
+            return Ok(());
+        }
+    };
+
     // Send the generates autocomplete response
     if let Err(e) = interaction
         .create_autocomplete_response(&ctx.serenity_context.http, |b| {
@@ -224,3 +258,35 @@ pub async fn dispatch_autocomplete<'a, U, E>(
 
     Ok(())
 }
+
+/// Dispatches this interaction onto framework commands, i.e. runs the associated autocomplete
+/// callback
+pub async fn dispatch_autocomplete<'a, U, E>(
+    framework: crate::FrameworkContext<'a, U, E>,
+    ctx: &'a serenity::Context,
+    interaction: &'a serenity::AutocompleteInteraction,
+    // Need to pass this in from outside because of lifetime issues
+    has_sent_initial_response: &'a std::sync::atomic::AtomicBool,
+    // Need to pass this in from outside because of lifetime issues
+    invocation_data: &'a tokio::sync::Mutex<Box<dyn std::any::Any + Send + Sync>>,
+    // Need to pass this in from outside because of lifetime issues
+    parent_commands: &'a mut Vec<&'a crate::Command<U, E>>,
+) -> Result<(), crate::FrameworkError<'a, U, E>> {
+    let ctx = extract_command(
+        framework,
+        ctx,
+        crate::ApplicationCommandOrAutocompleteInteraction::Autocomplete(interaction),
+        has_sent_initial_response,
+        invocation_data,
+        parent_commands,
+    )?;
+
+    crate::catch_unwind_maybe(run_autocomplete(ctx))
+        .await
+        .map_err(|payload| crate::FrameworkError::CommandPanic {
+            payload,
+            ctx: ctx.into(),
+        })??;
+
+    Ok(())
+}

src/lib.rs

@@ -466,3 +466,14 @@ use serenity_prelude as serenity; // private alias for crate root docs intradoc-
 ///
 /// An owned future has the `'static` lifetime.
 pub type BoxFuture<'a, T> = std::pin::Pin<Box<dyn std::future::Future<Output = T> + Send + 'a>>;
+
+/// Internal wrapper function for catch_unwind that respects the `handle_panics` feature flag
+async fn catch_unwind_maybe<T>(
+    fut: impl std::future::Future<Output = T>,
+) -> Result<T, Box<dyn std::any::Any + Send + 'static>> {
+    #[cfg(feature = "handle_panics")]
+    let res = futures_util::FutureExt::catch_unwind(std::panic::AssertUnwindSafe(fut)).await;
+    #[cfg(not(feature = "handle_panics"))]
+    let res = Ok(fut.await);
+    res
+}

src/structs/framework_error.rs

@@ -42,6 +42,17 @@ pub enum FrameworkError<'a, U, E> {
         /// General context
         ctx: crate::Context<'a, U, E>,
     },
+    /// Panic occured at any phase of command execution after constructing the `crate::Context`.
+    ///
+    /// This feature is intended as a last-resort safeguard to gracefully print an error message to
+    /// the user on a panic. Panics should only be thrown for bugs in the code, don't use this for
+    /// normal errors!
+    CommandPanic {
+        /// Panic payload which was thrown in the command code
+        payload: Box<dyn std::any::Any + Send + 'static>,
+        /// Command context
+        ctx: crate::Context<'a, U, E>,
+    },
     /// A command argument failed to parse from the Discord message or interaction content
     ArgumentParse {
         /// Error which was thrown by the parameter type's parsing routine
@@ -170,6 +181,7 @@ impl<'a, U, E> FrameworkError<'a, U, E> {
             Self::Setup { ctx, .. } => ctx,
             Self::EventHandler { ctx, .. } => ctx,
             Self::Command { ctx, .. } => ctx.serenity_context(),
+            Self::CommandPanic { ctx, .. } => ctx.serenity_context(),
             Self::ArgumentParse { ctx, .. } => ctx.serenity_context(),
             Self::CommandStructureMismatch { ctx, .. } => ctx.serenity_context,
             Self::CooldownHit { ctx, .. } => ctx.serenity_context(),
@@ -191,6 +203,7 @@ impl<'a, U, E> FrameworkError<'a, U, E> {
     pub fn ctx(&self) -> Option<crate::Context<'a, U, E>> {
         Some(match *self {
             Self::Command { ctx, .. } => ctx,
+            Self::CommandPanic { ctx, .. } => ctx,
             Self::ArgumentParse { ctx, .. } => ctx,
             Self::CommandStructureMismatch { ctx, .. } => crate::Context::Application(ctx),
             Self::CooldownHit { ctx, .. } => ctx,
@@ -245,6 +258,9 @@ impl<U, E: std::fmt::Display> std::fmt::Display for FrameworkError<'_, U, E> {
             Self::Command { error: _, ctx } => {
                 write!(f, "error in command `{}`", full_command_name!(ctx))
             }
+            Self::CommandPanic { ctx, payload: _ } => {
+                write!(f, "panic in command `{}`", full_command_name!(ctx))
+            }
             Self::ArgumentParse {
                 error: _,
                 input,
@@ -343,6 +359,7 @@ impl<'a, U: std::fmt::Debug, E: std::error::Error + 'static> std::error::Error
             Self::Setup { error, .. } => Some(error),
             Self::EventHandler { error, .. } => Some(error),
             Self::Command { error, .. } => Some(error),
+            Self::CommandPanic { .. } => None,
             Self::ArgumentParse { error, .. } => Some(&**error),
             Self::CommandStructureMismatch { .. } => None,
             Self::CooldownHit { .. } => None,