diff --git a/src/client.rs b/src/client.rs index 5de901c..572302f 100644 --- a/src/client.rs +++ b/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 { 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 { 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] diff --git a/src/commands/minecraft/enchantment.rs b/src/commands/minecraft/enchantment.rs index d25f8ff..1b48997 100644 --- a/src/commands/minecraft/enchantment.rs +++ b/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::().expect("Failed to get store"); diff --git a/src/commands/minecraft/item.rs b/src/commands/minecraft/item.rs index df4bd46..496e173 100644 --- a/src/commands/minecraft/item.rs +++ b/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::().expect("Failed to get store"); diff --git a/src/commands/misc/about.rs b/src/commands/misc/about.rs index b3ef1d4..7c301de 100644 --- a/src/commands/misc/about.rs +++ b/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| { diff --git a/src/commands/misc/pekofy.rs b/src/commands/misc/pekofy.rs index 18feeef..3bc72bd 100644 --- a/src/commands/misc/pekofy.rs +++ b/src/commands/misc/pekofy.rs @@ -21,6 +21,7 @@ static PEKOS: &[&str] = &[ #[usage("()")] #[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(); diff --git a/src/commands/misc/ping.rs b/src/commands/misc/ping.rs index 78a2fcd..f9567e9 100644 --- a/src/commands/misc/ping.rs +++ b/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?; diff --git a/src/commands/misc/qalc.rs b/src/commands/misc/qalc.rs index 06ab58d..5903796 100644 --- a/src/commands/misc/qalc.rs +++ b/src/commands/misc/qalc.rs @@ -13,6 +13,7 @@ static QALC_HELP: &[&str] = &["help", "--help", "-h", "h"]; #[min_args(1)] #[usage("")] #[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! { diff --git a/src/commands/misc/sauce.rs b/src/commands/misc/sauce.rs index b6efd66..3d114d0 100644 --- a/src/commands/misc/sauce.rs +++ b/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?; diff --git a/src/commands/misc/stats.rs b/src/commands/misc/stats.rs index 6bca0f7..1c5b777 100644 --- a/src/commands/misc/stats.rs +++ b/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(); diff --git a/src/commands/misc/time.rs b/src/commands/misc/time.rs index c9c9113..6e4b077 100644 --- a/src/commands/misc/time.rs +++ b/src/commands/misc/time.rs @@ -10,6 +10,7 @@ use serenity::model::channel::Message; #[min_args(1)] #[max_args(3)] #[usage("<%H:%M/now> () ()")] +#[bucket("general")] async fn time(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult { let when = args.single::().unwrap_or("now".to_string()); let first_timezone = args.single::().ok(); diff --git a/src/commands/misc/timezones.rs b/src/commands/misc/timezones.rs index 3c1fdb6..49d2eca 100644 --- a/src/commands/misc/timezones.rs +++ b/src/commands/misc/timezones.rs @@ -8,6 +8,7 @@ use serenity::model::channel::Message; #[min_args(1)] #[usage("")] #[example("Europe Berlin")] +#[bucket("general")] async fn timezones(ctx: &Context, msg: &Message, mut args: Args) -> CommandResult { let query = args .iter::() diff --git a/src/commands/music/clear_queue.rs b/src/commands/music/clear_queue.rs index 3e240a7..78bd902 100644 --- a/src/commands/music/clear_queue.rs +++ b/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(); diff --git a/src/commands/music/current.rs b/src/commands/music/current.rs index 5aad683..57ae17d 100644 --- a/src/commands/music/current.rs +++ b/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(); diff --git a/src/commands/music/join.rs b/src/commands/music/join.rs index 65543d7..9301c03 100644 --- a/src/commands/music/join.rs +++ b/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)?; diff --git a/src/commands/music/leave.rs b/src/commands/music/leave.rs index 08043d1..23a3c7e 100644 --- a/src/commands/music/leave.rs +++ b/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); diff --git a/src/commands/music/lyrics.rs b/src/commands/music/lyrics.rs index 82416ae..6244dcc 100644 --- a/src/commands/music/lyrics.rs +++ b/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); diff --git a/src/commands/music/pause.rs b/src/commands/music/pause.rs index 97624a8..c3fbfe6 100644 --- a/src/commands/music/pause.rs +++ b/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); diff --git a/src/commands/music/play.rs b/src/commands/music/play.rs index 2dc78b2..e53884a 100644 --- a/src/commands/music/play.rs +++ b/src/commands/music/play.rs @@ -16,6 +16,7 @@ use crate::providers::settings::{get_setting, Setting}; #[usage("()")] #[min_args(1)] #[aliases("p")] +#[bucket("music_api")] async fn play(ctx: &Context, msg: &Message, args: Args) -> CommandResult { let query = args.message(); diff --git a/src/commands/music/play_next.rs b/src/commands/music/play_next.rs index be6769a..23825ad 100644 --- a/src/commands/music/play_next.rs +++ b/src/commands/music/play_next.rs @@ -15,6 +15,7 @@ use crate::commands::music::{ #[usage("")] #[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(); diff --git a/src/commands/music/playlists.rs b/src/commands/music/playlists.rs index 92ca932..2ebe6ac 100644 --- a/src/commands/music/playlists.rs +++ b/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); diff --git a/src/commands/music/queue.rs b/src/commands/music/queue.rs index cd984ec..481e71b 100644 --- a/src/commands/music/queue.rs +++ b/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); diff --git a/src/commands/music/save_playlist.rs b/src/commands/music/save_playlist.rs index 55ee8da..46f2b33 100644 --- a/src/commands/music/save_playlist.rs +++ b/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(); diff --git a/src/commands/music/shuffle.rs b/src/commands/music/shuffle.rs index 27c5267..cc0321b 100644 --- a/src/commands/music/shuffle.rs +++ b/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(); diff --git a/src/commands/music/skip.rs b/src/commands/music/skip.rs index 4e2dcd6..d32ab2b 100644 --- a/src/commands/music/skip.rs +++ b/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? { diff --git a/src/commands/settings/get.rs b/src/commands/settings/get.rs index 2effdec..da4669c 100644 --- a/src/commands/settings/get.rs +++ b/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(); diff --git a/src/commands/settings/set.rs b/src/commands/settings/set.rs index f4b6f55..62b38a0 100644 --- a/src/commands/settings/set.rs +++ b/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::().unwrap().to_lowercase(); let all_settings: Vec = ALL_SETTINGS.iter().map(|s| s.to_string()).collect();