Merge pull request #57 from rust-lang/protection

add protected commands and protected help menu items

Author: technetos

src/ban.rs

@@ -79,43 +79,41 @@ pub(crate) fn start_unban_thread(cx: Context) {
 ///
 /// Requires the ban members permission
 pub(crate) fn temp_ban(args: Args) -> Result<()> {
-    if api::is_mod(&args)? {
-        let user_id = parse_username(
-            &args
-                .params
-                .get("user")
-                .ok_or("unable to retrieve user param")?,
-        )
-        .ok_or("unable to retrieve user id")?;
+    let user_id = parse_username(
+        &args
+            .params
+            .get("user")
+            .ok_or("unable to retrieve user param")?,
+    )
+    .ok_or("unable to retrieve user id")?;
 
-        use std::str::FromStr;
+    use std::str::FromStr;
 
-        let hours = u64::from_str(
-            args.params
-                .get("hours")
-                .ok_or("unable to retrieve hours param")?,
-        )?;
+    let hours = u64::from_str(
+        args.params
+            .get("hours")
+            .ok_or("unable to retrieve hours param")?,
+    )?;
 
-        let reason = args
-            .params
-            .get("reason")
-            .ok_or("unable to retrieve reason param")?;
+    let reason = args
+        .params
+        .get("reason")
+        .ok_or("unable to retrieve reason param")?;
 
-        if let Some(guild) = args.msg.guild(&args.cx) {
-            info!("Banning user from guild");
-            let user = UserId::from(user_id);
+    if let Some(guild) = args.msg.guild(&args.cx) {
+        info!("Banning user from guild");
+        let user = UserId::from(user_id);
 
-            user.create_dm_channel(args.cx)?
-                .say(args.cx, ban_message(reason, hours))?;
+        user.create_dm_channel(args.cx)?
+            .say(args.cx, ban_message(reason, hours))?;
 
-            guild.read().ban(args.cx, &user, &"all")?;
+        guild.read().ban(args.cx, &user, &"all")?;
 
-            save_ban(
-                format!("{}", user_id),
-                format!("{}", guild.read().id),
-                hours,
-            )?;
-        }
+        save_ban(
+            format!("{}", user_id),
+            format!("{}", guild.read().id),
+            hours,
+        )?;
     }
     Ok(())
 }

src/commands.rs

@@ -1,11 +1,29 @@
-use crate::state_machine::{CharacterSet, StateMachine};
+use crate::{
+    api,
+    state_machine::{CharacterSet, StateMachine},
+};
 use reqwest::blocking::Client as HttpClient;
 use serenity::{model::channel::Message, prelude::Context};
 use std::{collections::HashMap, sync::Arc};
 
 const PREFIX: &'static str = "?";
 pub(crate) type Result<T> = std::result::Result<T, Box<dyn std::error::Error>>;
-pub(crate) type CmdPtr = Arc<dyn for<'m> Fn(Args<'m>) -> Result<()> + Send + Sync>;
+pub(crate) type GuardFn = fn(&Args) -> Result<bool>;
+
+struct Command {
+    guard: GuardFn,
+    ptr: Box<dyn for<'m> Fn(Args<'m>) -> Result<()> + Send + Sync>,
+}
+
+impl Command {
+    fn authorize(&self, args: &Args) -> Result<bool> {
+        (self.guard)(&args)
+    }
+
+    fn call(&self, args: Args) -> Result<()> {
+        (self.ptr)(args)
+    }
+}
 
 pub struct Args<'m> {
     pub http: &'m HttpClient,
@@ -15,24 +33,33 @@ pub struct Args<'m> {
 }
 
 pub(crate) struct Commands {
-    state_machine: StateMachine,
+    state_machine: StateMachine<Arc<Command>>,
     client: HttpClient,
-    menu: HashMap<&'static str, &'static str>,
+    menu: Option<HashMap<&'static str, (&'static str, GuardFn)>>,
 }
 
 impl Commands {
     pub(crate) fn new() -> Self {
         Self {
             state_machine: StateMachine::new(),
             client: HttpClient::new(),
-            menu: HashMap::new(),
+            menu: Some(HashMap::new()),
         }
     }
 
     pub(crate) fn add(
         &mut self,
         command: &'static str,
         handler: impl Fn(Args) -> Result<()> + Send + Sync + 'static,
+    ) {
+        self.add_protected(command, handler, |_| Ok(true));
+    }
+
+    pub(crate) fn add_protected(
+        &mut self,
+        command: &'static str,
+        handler: impl Fn(Args) -> Result<()> + Send + Sync + 'static,
+        guard: GuardFn,
     ) {
         info!("Adding command {}", &command);
         let mut state = 0;
@@ -89,7 +116,10 @@ impl Commands {
                 }
             });
 
-        let handler = Arc::new(handler);
+        let handler = Arc::new(Command {
+            guard,
+            ptr: Box::new(handler),
+        });
 
         if opt_lambda_state.is_some() {
             opt_final_states.iter().for_each(|state| {
@@ -107,35 +137,69 @@ impl Commands {
         cmd: &'static str,
         desc: &'static str,
         handler: impl Fn(Args) -> Result<()> + Send + Sync + 'static,
+    ) {
+        self.help_protected(cmd, desc, handler, |_| Ok(true));
+    }
+
+    pub(crate) fn help_protected(
+        &mut self,
+        cmd: &'static str,
+        desc: &'static str,
+        handler: impl Fn(Args) -> Result<()> + Send + Sync + 'static,
+        guard: GuardFn,
     ) {
         let base_cmd = &cmd[1..];
         info!("Adding command ?help {}", &base_cmd);
         let mut state = 0;
 
-        self.menu.insert(cmd, desc);
-        state = add_help_menu(&mut self.state_machine, base_cmd, state);
+        self.menu.as_mut().map(|menu| {
+            menu.insert(cmd, (desc, guard));
+            menu
+        });
 
+        state = add_help_menu(&mut self.state_machine, base_cmd, state);
         self.state_machine.set_final_state(state);
-        self.state_machine.set_handler(state, Arc::new(handler));
+        self.state_machine.set_handler(
+            state,
+            Arc::new(Command {
+                guard,
+                ptr: Box::new(handler),
+            }),
+        );
     }
 
-    pub(crate) fn menu(&mut self) -> &HashMap<&'static str, &'static str> {
-        &self.menu
+    pub(crate) fn menu(&mut self) -> Option<HashMap<&'static str, (&'static str, GuardFn)>> {
+        self.menu.take()
     }
 
     pub(crate) fn execute<'m>(&'m self, cx: Context, msg: Message) {
         let message = &msg.content;
         if !msg.is_own(&cx) && message.starts_with(PREFIX) {
             self.state_machine.process(message).map(|matched| {
-                info!("Executing command {}", message);
+                info!("Processing command: {}", message);
                 let args = Args {
                     http: &self.client,
                     cx: &cx,
                     msg: &msg,
                     params: matched.params,
                 };
-                if let Err(e) = (matched.handler)(args) {
-                    println!("{}", e);
+                info!("Checking permissions");
+                match matched.handler.authorize(&args) {
+                    Ok(true) => {
+                        info!("Executing command");
+                        if let Err(e) = matched.handler.call(args) {
+                            error!("{}", e);
+                        }
+                    }
+                    Ok(false) => {
+                        info!("Not executing command, unauthorized");
+                        if let Err(e) =
+                            api::send_reply(&args, "You do not have permission to run this command")
+                        {
+                            error!("{}", e);
+                        }
+                    }
+                    Err(e) => error!("{}", e),
                 }
             });
         }
@@ -156,7 +220,7 @@ fn key_value_pair(s: &'static str) -> Option<&'static str> {
         .flatten()
 }
 
-fn add_space(state_machine: &mut StateMachine, mut state: usize, i: usize) -> usize {
+fn add_space<T>(state_machine: &mut StateMachine<T>, mut state: usize, i: usize) -> usize {
     if i > 0 {
         let mut char_set = CharacterSet::from_char(' ');
         char_set.insert('\n');
@@ -167,8 +231,8 @@ fn add_space(state_machine: &mut StateMachine, mut state: usize, i: usize) -> us
     state
 }
 
-fn add_help_menu(
-    mut state_machine: &mut StateMachine,
+fn add_help_menu<T>(
+    mut state_machine: &mut StateMachine<T>,
     cmd: &'static str,
     mut state: usize,
 ) -> usize {
@@ -183,8 +247,8 @@ fn add_help_menu(
     state
 }
 
-fn add_dynamic_segment(
-    state_machine: &mut StateMachine,
+fn add_dynamic_segment<T>(
+    state_machine: &mut StateMachine<T>,
     name: &'static str,
     mut state: usize,
 ) -> usize {
@@ -198,8 +262,8 @@ fn add_dynamic_segment(
     state
 }
 
-fn add_remaining_segment(
-    state_machine: &mut StateMachine,
+fn add_remaining_segment<T>(
+    state_machine: &mut StateMachine<T>,
     name: &'static str,
     mut state: usize,
 ) -> usize {
@@ -212,8 +276,8 @@ fn add_remaining_segment(
     state
 }
 
-fn add_code_segment_multi_line(
-    state_machine: &mut StateMachine,
+fn add_code_segment_multi_line<T>(
+    state_machine: &mut StateMachine<T>,
     name: &'static str,
     mut state: usize,
 ) -> usize {
@@ -246,8 +310,8 @@ fn add_code_segment_multi_line(
     state
 }
 
-fn add_code_segment_single_line(
-    state_machine: &mut StateMachine,
+fn add_code_segment_single_line<T>(
+    state_machine: &mut StateMachine<T>,
     name: &'static str,
     mut state: usize,
     n_backticks: usize,
@@ -266,7 +330,11 @@ fn add_code_segment_single_line(
     state
 }
 
-fn add_key_value(state_machine: &mut StateMachine, name: &'static str, mut state: usize) -> usize {
+fn add_key_value<T>(
+    state_machine: &mut StateMachine<T>,
+    name: &'static str,
+    mut state: usize,
+) -> usize {
     name.chars().for_each(|c| {
         state = state_machine.add(state, CharacterSet::from_char(c));
     });