Add rate limits to commands

Signed-off-by: trivernis <trivernis@protonmail.com>

src/client.rs

@@ -15,6 +15,7 @@ use crate::handler::Handler;
 use crate::utils::context_data::{DatabaseContainer, Store, StoreData};
 use crate::utils::error::{BotError, BotResult};
 use bot_serenityutils::menu::EventDrivenMessageContainer;
+use serenity::framework::standard::buckets::LimitedFor;
 use std::sync::Arc;
 use tokio::sync::Mutex;
 
@@ -24,7 +25,7 @@ pub async fn get_client() -> BotResult<Client> {
 
     let client = Client::builder(token)
         .event_handler(Handler)
-        .framework(get_framework())
+        .framework(get_framework().await)
         .register_songbird()
         .await?;
     {
@@ -37,7 +38,7 @@ pub async fn get_client() -> BotResult<Client> {
     Ok(client)
 }
 
-pub fn get_framework() -> StandardFramework {
+pub async fn get_framework() -> StandardFramework {
     let mut owners = HashSet::new();
     if let Some(owner) = dotenv::var("BOT_OWNER").ok().and_then(|o| o.parse().ok()) {
         owners.insert(UserId(owner));
@@ -61,6 +62,22 @@ pub fn get_framework() -> StandardFramework {
         .before(before_hook)
         .on_dispatch_error(dispatch_error)
         .help(&HELP)
+        .bucket("music_api", |b| {
+            b.delay(1)
+                .time_span(60)
+                .limit(30)
+                .limit_for(LimitedFor::User)
+        })
+        .await
+        .bucket("sauce_api", |b| {
+            b.delay(1)
+                .time_span(60)
+                .limit(10)
+                .limit_for(LimitedFor::User)
+        })
+        .await
+        .bucket("general", |b| b.time_span(10).limit(5))
+        .await
 }
 
 #[hook]

src/commands/minecraft/enchantment.rs

@@ -11,6 +11,7 @@ use crate::utils::context_data::Store;
 #[example("unbreaking")]
 #[min_args(1)]
 #[aliases("ench")]
+#[bucket("general")]
 pub(crate) async fn enchantment(ctx: &Context, msg: &Message, args: Args) -> CommandResult {
     let data = ctx.data.read().await;
     let store = data.get::<Store>().expect("Failed to get store");

src/commands/minecraft/item.rs

@@ -11,6 +11,7 @@ use crate::utils::context_data::Store;
 #[example("bread")]
 #[min_args(1)]
 #[aliases("i")]
+#[bucket("general")]
 pub(crate) async fn item(ctx: &Context, msg: &Message, args: Args) -> CommandResult {
     let data = ctx.data.read().await;
     let store = data.get::<Store>().expect("Failed to get store");

src/commands/misc/about.rs

@@ -5,6 +5,7 @@ use serenity::model::channel::Message;
 
 #[command]
 #[description("Displays information about the bot")]
+#[bucket("general")]
 async fn about(ctx: &Context, msg: &Message) -> CommandResult {
     msg.channel_id
         .send_message(ctx, |m| {

src/commands/misc/pekofy.rs

@@ -21,6 +21,7 @@ static PEKOS: &[&str] = &[
 #[usage("(<content>)")]
 #[example("Hello")]
 #[aliases("peko")]
+#[bucket("general")]
 async fn pekofy(ctx: &Context, msg: &Message, args: Args) -> CommandResult {
     let mut reference_message = msg.id;
     let mut content = args.message().to_string();

src/commands/misc/ping.rs

@@ -6,6 +6,7 @@ use serenity::model::channel::Message;
 #[command]
 #[description("Simple ping test command")]
 #[usage("")]
+#[bucket("general")]
 async fn ping(ctx: &Context, msg: &Message) -> CommandResult {
     msg.reply(ctx, "Pong!").await?;
 

src/commands/misc/qalc.rs

@@ -13,6 +13,7 @@ static QALC_HELP: &[&str] = &["help", "--help", "-h", "h"];
 #[min_args(1)]
 #[usage("<expression>")]
 #[example("1 * 1 + 1 / sqrt(2)")]
+#[bucket("general")]
 async fn qalc(ctx: &Context, msg: &Message, args: Args) -> CommandResult {
     let expression = args.message();
     lazy_static::lazy_static! {

src/commands/misc/sauce.rs

@@ -14,6 +14,7 @@ use crate::utils::get_previous_message_or_reply;
 #[description("Searches for the source of a previously posted image or an image replied to.")]
 #[usage("")]
 #[aliases("source")]
+#[bucket("sauce_api")]
 async fn sauce(ctx: &Context, msg: &Message) -> CommandResult {
     log::debug!("Got sauce command");
     let source_msg = get_previous_message_or_reply(ctx, msg).await?;

src/commands/misc/stats.rs

@@ -15,6 +15,7 @@ const VERSION: &'static str = env!("CARGO_PKG_VERSION");
 #[command]
 #[description("Shows some statistics about the bot")]
 #[usage("")]
+#[bucket("general")]
 async fn stats(ctx: &Context, msg: &Message) -> CommandResult {
     log::debug!("Reading system stats");
     let mut system = sysinfo::System::new_all();

src/commands/misc/time.rs

@@ -10,6 +10,7 @@ use serenity::model::channel::Message;
 #[min_args(1)]
 #[max_args(3)]
 #[usage("<%H:%M/now> (<from-timezone>) (<to-timezone>)")]
+#[bucket("general")]
 async fn time(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult {
     let when = args.single::<String>().unwrap_or("now".to_string());
     let first_timezone = args.single::<String>().ok();

src/commands/misc/timezones.rs

@@ -8,6 +8,7 @@ use serenity::model::channel::Message;
 #[min_args(1)]
 #[usage("<query...>")]
 #[example("Europe Berlin")]
+#[bucket("general")]
 async fn timezones(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult {
     let query = args
         .iter::<String>()

src/commands/music/clear_queue.rs

@@ -11,6 +11,7 @@ use crate::commands::music::{get_queue_for_guild, is_dj};
 #[description("Clears the queue")]
 #[usage("")]
 #[aliases("cl")]
+#[bucket("general")]
 async fn clear_queue(ctx: &Context, msg: &Message) -> CommandResult {
     let guild = msg.guild(&ctx.cache).await.unwrap();
 

src/commands/music/current.rs

@@ -14,6 +14,7 @@ use crate::messages::music::NowPlayingMessage;
 #[description("Displays the currently playing song")]
 #[usage("")]
 #[aliases("nowplaying", "np")]
+#[bucket("general")]
 async fn current(ctx: &Context, msg: &Message) -> CommandResult {
     let guild = msg.guild(&ctx.cache).await.unwrap();
 

src/commands/music/join.rs

@@ -10,6 +10,7 @@ use crate::commands::music::{get_channel_for_author, join_channel};
 #[only_in(guilds)]
 #[description("Joins a voice channel")]
 #[usage("")]
+#[bucket("general")]
 async fn join(ctx: &Context, msg: &Message) -> CommandResult {
     let guild = msg.guild(&ctx.cache).await.unwrap();
     let channel_id = get_channel_for_author(&msg.author.id, &guild)?;

src/commands/music/leave.rs

@@ -11,6 +11,7 @@ use crate::commands::music::{get_queue_for_guild, get_voice_manager, is_dj};
 #[description("Leaves a voice channel")]
 #[usage("")]
 #[aliases("stop")]
+#[bucket("general")]
 async fn leave(ctx: &Context, msg: &Message) -> CommandResult {
     let guild = msg.guild(&ctx.cache).await.unwrap();
     log::debug!("Leave request received for guild {}", guild.id);

src/commands/music/lyrics.rs

@@ -11,6 +11,7 @@ use crate::providers::music::lyrics::get_lyrics;
 #[only_in(guilds)]
 #[description("Shows the lyrics for the currently playing song")]
 #[usage("")]
+#[bucket("general")]
 async fn lyrics(ctx: &Context, msg: &Message) -> CommandResult {
     let guild = msg.guild(&ctx.cache).await.unwrap();
     log::debug!("Fetching lyrics for song playing in {}", guild.id);

src/commands/music/pause.rs

@@ -10,6 +10,7 @@ use crate::commands::music::{get_queue_for_guild, is_dj};
 #[only_in(guilds)]
 #[description("Pauses playback")]
 #[usage("")]
+#[bucket("general")]
 async fn pause(ctx: &Context, msg: &Message) -> CommandResult {
     let guild = msg.guild(&ctx.cache).await.unwrap();
     log::debug!("Pausing playback for guild {}", guild.id);

src/commands/music/play.rs

@@ -16,6 +16,7 @@ use crate::providers::settings::{get_setting, Setting};
 #[usage("(<spotify_url,youtube_url,query>)")]
 #[min_args(1)]
 #[aliases("p")]
+#[bucket("music_api")]
 async fn play(ctx: &Context, msg: &Message, args: Args) -> CommandResult {
     let query = args.message();
 

src/commands/music/play_next.rs

@@ -15,6 +15,7 @@ use crate::commands::music::{
 #[usage("<song-url>")]
 #[min_args(1)]
 #[aliases("pn", "play-next")]
+#[bucket("music_api")]
 async fn play_next(ctx: &Context, msg: &Message, args: Args) -> CommandResult {
     let query = args.message();
 

src/commands/music/playlists.rs

@@ -10,6 +10,7 @@ use crate::utils::context_data::get_database_from_context;
 #[only_in(guilds)]
 #[description("Displays a list of all saved playlists")]
 #[usage("")]
+#[bucket("general")]
 async fn playlists(ctx: &Context, msg: &Message) -> CommandResult {
     let guild = msg.guild(&ctx.cache).await.unwrap();
     log::debug!("Displaying playlists for guild {}", guild.id);

src/commands/music/queue.rs

@@ -13,6 +13,7 @@ use crate::commands::music::get_queue_for_guild;
 #[description("Shows the song queue")]
 #[usage("")]
 #[aliases("q")]
+#[bucket("general")]
 async fn queue(ctx: &Context, msg: &Message) -> CommandResult {
     let guild = msg.guild(&ctx.cache).await.unwrap();
     log::trace!("Displaying queue for guild {}", guild.id);

src/commands/music/save_playlist.rs

@@ -13,6 +13,7 @@ use crate::utils::context_data::get_database_from_context;
 #[example("anime https://www.youtube.com/playlist?list=PLqaM77H_o5hykROCe3uluvZEaPo6bZj-C")]
 #[min_args(2)]
 #[aliases("add-playlist", "save-playlist")]
+#[bucket("general")]
 async fn save_playlist(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult {
     let guild = msg.guild(&ctx.cache).await.unwrap();
 

src/commands/music/shuffle.rs

@@ -11,6 +11,7 @@ use crate::commands::music::{get_queue_for_guild, is_dj};
 #[description("Shuffles the queue")]
 #[usage("")]
 #[aliases("sh")]
+#[bucket("general")]
 async fn shuffle(ctx: &Context, msg: &Message) -> CommandResult {
     let guild = msg.guild(&ctx.cache).await.unwrap();
 

src/commands/music/skip.rs

@@ -11,6 +11,7 @@ use crate::commands::music::{get_queue_for_guild, is_dj};
 #[description("Skips to the next song")]
 #[usage("")]
 #[aliases("next")]
+#[bucket("general")]
 async fn skip(ctx: &Context, msg: &Message) -> CommandResult {
     let guild = msg.guild(&ctx.cache).await.unwrap();
     if !is_dj(ctx, guild.id, &msg.author).await? {

src/commands/settings/get.rs

@@ -14,6 +14,7 @@ use crate::utils::context_data::get_database_from_context;
 #[min_args(0)]
 #[max_args(1)]
 #[required_permissions("MANAGE_GUILD")]
+#[bucket("general")]
 async fn get(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult {
     let database = get_database_from_context(ctx).await;
     let guild = msg.guild(&ctx.cache).await.unwrap();

src/commands/settings/set.rs

@@ -16,6 +16,7 @@ use crate::utils::context_data::get_database_from_context;
 #[min_args(1)]
 #[max_args(2)]
 #[required_permissions("MANAGE_GUILD")]
+#[bucket("general")]
 async fn set(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult {
     let key = args.single::<String>().unwrap().to_lowercase();
     let all_settings: Vec<String> = ALL_SETTINGS.iter().map(|s| s.to_string()).collect();