Skip to content
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
17 changes: 15 additions & 2 deletions rustmail/src/bot.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,21 +7,28 @@ use crate::prelude::panel_commands::*;
use crate::prelude::types::*;
use base64::Engine;
use rand::RngCore;
use serenity::all::{ClientBuilder, GatewayIntents};
use serenity::all::{ClientBuilder, GatewayIntents, ShardManager};
use serenity::cache::Settings as CacheSettings;
use serenity::prelude::TypeMapKey;
use std::collections::HashMap;
use std::process;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::Mutex;
use tokio::{select, spawn};

pub struct ShardManagerKey;

impl TypeMapKey for ShardManagerKey {
type Value = Arc<ShardManager>;
}

pub async fn init_bot_state() -> Arc<Mutex<BotState>> {
let pool = init_database().await.expect("An error occured!");
println!("Database connected!");

let mut bytes = [0u8; 32];
rand::thread_rng().fill_bytes(&mut bytes);
rand::rng().fill_bytes(&mut bytes);
Comment thread
Akinator31 marked this conversation as resolved.
let token = base64::engine::general_purpose::STANDARD.encode(&bytes);

let config = load_config("config.toml");
Expand Down Expand Up @@ -136,6 +143,7 @@ pub async fn run_bot(
registry.register_command(LogsCommand);
registry.register_command(TakeCommand);
registry.register_command(ReleaseCommand);
registry.register_command(PingCommand);

let registry = Arc::new(registry);

Expand Down Expand Up @@ -174,6 +182,11 @@ pub async fn run_bot(
state_lock.bot_http = Some(client.http.clone());
}

{
let mut data = client.data.write().await;
data.insert::<ShardManagerKey>(client.shard_manager.clone());
}

if let Err(e) = config.validate_servers(&client.http).await {
eprintln!("Configuration validation error: {}", e);
eprintln!(
Expand Down
2 changes: 2 additions & 0 deletions rustmail/src/commands/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ pub mod id;
pub mod logs;
pub mod move_thread;
pub mod new_thread;
pub mod ping;
pub mod recover;
pub mod release;
pub mod remove_reminder;
Expand All @@ -42,6 +43,7 @@ pub use id::*;
pub use logs::*;
pub use move_thread::*;
pub use new_thread::*;
pub use ping::*;
pub use recover::*;
pub use release::*;
pub use remove_reminder::*;
Expand Down
5 changes: 5 additions & 0 deletions rustmail/src/commands/ping/mod.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
mod slash_command;
mod text_command;

pub use slash_command::*;
pub use text_command::*;
3 changes: 3 additions & 0 deletions rustmail/src/commands/ping/slash_command/mod.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
pub mod ping;

pub use ping::*;
121 changes: 121 additions & 0 deletions rustmail/src/commands/ping/slash_command/ping.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,121 @@
use crate::bot::ShardManagerKey;
use crate::commands::{BoxFuture, CommunityRegistrable, RegistrableCommand};
use crate::config::Config;
use crate::errors::{DiscordError, ModmailError, ModmailResult};
use crate::handlers::InteractionHandler;
use crate::i18n::get_translated_message;
use crate::utils::{MessageBuilder, defer_response};
use chrono::Utc;
use serenity::FutureExt;
use serenity::all::{CommandInteraction, Context, CreateCommand, ResolvedOption};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use tokio::time::Instant;

pub struct PingCommand;

#[async_trait::async_trait]
impl RegistrableCommand for PingCommand {
fn as_community(&self) -> Option<&dyn CommunityRegistrable> {
None
}

fn name(&self) -> &'static str {
"ping"
}

fn doc<'a>(&self, config: &'a Config) -> BoxFuture<'a, String> {
async move { get_translated_message(config, "help.ping", None, None, None, None).await }
.boxed()
}

fn register(&self, config: &Config) -> BoxFuture<'_, Vec<CreateCommand>> {
let config = config.clone();

Box::pin(async move {
let cmd_desc = get_translated_message(
&config,
"slash_command.ping_command_desc",
None,
None,
None,
None,
)
.await;

vec![CreateCommand::new(self.name()).description(cmd_desc)]
})
}

fn run(
&self,
ctx: &Context,
command: &CommandInteraction,
_options: &[ResolvedOption<'_>],
config: &Config,
_handler: Arc<InteractionHandler>,
) -> BoxFuture<'_, ModmailResult<()>> {
let ctx = ctx.clone();
let command = command.clone();
let config = config.clone();

Box::pin(async move {
defer_response(&ctx, &command).await?;
let response = MessageBuilder::system_message(&ctx, &config)
.content("...")
.to_channel(command.channel_id)
.build_interaction_message_followup()
.await;

let shard_manager = ctx
.data
.read()
.await
.get::<ShardManagerKey>()
.cloned()
.ok_or(ModmailError::Discord(DiscordError::ShardManagerNotFound))?;

let time_before = Instant::now();
let mut res = command.create_followup(&ctx.http, response).await?;
let msg_send_ping = time_before.elapsed().as_millis();

let gateway_ping = {
let runners = shard_manager.runners.lock().await;
runners.get(&ctx.shard_id).and_then(|runner| runner.latency)
};

let start = Instant::now();
ctx.http.get_gateway().await?;
let api_ping = start.elapsed();

let mut params = HashMap::new();
params.insert(
"gateway_latency".to_string(),
format!(
"{:?}",
gateway_ping.unwrap_or(Duration::default()).as_millis()
),
);
params.insert(
"api_latency".to_string(),
format!("{:?}", api_ping.as_millis()),
);
params.insert(
"message_latency".to_string(),
format!("{:?}", msg_send_ping),
);

let response = MessageBuilder::system_message(&ctx, &config)
.translated_content("slash_command.ping_command", Some(&params), None, None)
.await
.to_channel(command.channel_id)
.build_edit_message()
.await;

res.edit(&ctx.http, response).await?;

Ok(())
})
}
}
3 changes: 3 additions & 0 deletions rustmail/src/commands/ping/text_command/mod.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
pub mod ping;

pub use ping::*;
67 changes: 67 additions & 0 deletions rustmail/src/commands/ping/text_command/ping.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,67 @@
use crate::bot::ShardManagerKey;
use crate::prelude::config::*;
use crate::prelude::errors::*;
use crate::prelude::handlers::*;
use crate::utils::MessageBuilder;
use chrono::Utc;
use serenity::all::{Context, Message};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use tokio::time::Instant;

pub async fn ping(
ctx: Context,
msg: Message,
config: &Config,
_handler: Arc<GuildMessagesHandler>,
) -> ModmailResult<()> {
let shard_manager = ctx
.data
.read()
.await
.get::<ShardManagerKey>()
.cloned()
.ok_or(ModmailError::Discord(DiscordError::ShardManagerNotFound))?;

let time_before = Instant::now();
let mut res = msg.reply(&ctx.http, "...").await?;
let msg_send_ping = time_before.elapsed().as_millis();

let gateway_ping = {
let runners = shard_manager.runners.lock().await;
runners.get(&ctx.shard_id).and_then(|runner| runner.latency)
};

let start = Instant::now();
ctx.http.get_gateway().await?;
let api_ping = start.elapsed();

let mut params = HashMap::new();
params.insert(
"gateway_latency".to_string(),
format!(
"{:?}",
gateway_ping.unwrap_or(Duration::default()).as_millis()
),
);
params.insert(
"api_latency".to_string(),
format!("{:?}", api_ping.as_millis()),
);
params.insert(
"message_latency".to_string(),
format!("{:?}", msg_send_ping),
);

let edited_msg = MessageBuilder::system_message(&ctx, &config)
.translated_content("slash_command.ping_command", Some(&params), None, None)
.await
.to_channel(msg.channel_id)
.build_edit_message()
.await;

res.edit(&ctx.http, edited_msg).await?;

Ok(())
}
20 changes: 15 additions & 5 deletions rustmail/src/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,8 @@ pub struct Config {

pub db_pool: Option<SqlitePool>,
pub error_handler: Option<Arc<ErrorHandler>>,
pub thread_locks: Arc<std::sync::Mutex<std::collections::HashMap<u64, Arc<tokio::sync::Mutex<()>>>>>,
pub thread_locks:
Arc<std::sync::Mutex<std::collections::HashMap<u64, Arc<tokio::sync::Mutex<()>>>>>,
}

fn get_local_ip() -> Option<String> {
Expand Down Expand Up @@ -50,21 +51,30 @@ pub fn load_config(path: &str) -> Option<Config> {
}

if u64::from_str_radix(&config_response.thread.user_message_color, 16).is_err() {
eprintln!("Incorrect user message color in the config.toml! Please put a color in hex format!");
eprintln!(
"Incorrect user message color in the config.toml! Please put a color in hex format!"
);
return None;
}

if u64::from_str_radix(&config_response.thread.staff_message_color, 16).is_err() {
eprintln!("Incorrect staff message color in the config.toml! Please put a color in hex format!");
eprintln!(
"Incorrect staff message color in the config.toml! Please put a color in hex format!"
);
return None;
}

if u64::from_str_radix(&config_response.reminders.embed_color, 16).is_err() {
eprintln!("Incorrect reminder embed color in the config.toml! Please put a color in hex format!");
eprintln!(
"Incorrect reminder embed color in the config.toml! Please put a color in hex format!"
);
return None;
}

if !config_response.language.is_language_supported(config_response.language.get_default_language()) {
if !config_response
.language
.is_language_supported(config_response.language.get_default_language())
{
eprintln!(
"Warning: Default language '{}' is not in supported languages list",
config_response.language.default_language
Expand Down
3 changes: 3 additions & 0 deletions rustmail/src/errors/dictionary.rs
Original file line number Diff line number Diff line change
Expand Up @@ -188,6 +188,9 @@ impl DictionaryManager {
DiscordError::UserIsABot => ("discord.user_is_a_bot".to_string(), None),
DiscordError::PermissionDenied => ("discord.permission_denied".to_string(), None),
DiscordError::DmCreationFailed => ("discord.dm_creation_failed".to_string(), None),
DiscordError::ShardManagerNotFound => {
("discord.shard_manager_not_found".to_string(), None)
}
_ => ("discord.api_error".to_string(), None),
},
ModmailError::Command(cmd_err) => match cmd_err {
Expand Down
2 changes: 2 additions & 0 deletions rustmail/src/errors/types.rs
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@ pub enum DiscordError {
DmCreationFailed,
FailedToFetchCategories,
FailedToMoveChannel,
ShardManagerNotFound,
}

#[derive(Debug, Clone)]
Expand Down Expand Up @@ -198,6 +199,7 @@ impl fmt::Display for DiscordError {
DiscordError::DmCreationFailed => write!(f, "Failed to create DM channel"),
DiscordError::FailedToFetchCategories => write!(f, "Failed to fetch categories"),
DiscordError::FailedToMoveChannel => write!(f, "Failed to move_thread channel"),
DiscordError::ShardManagerNotFound => write!(f, "Shard manager not found"),
}
}
}
Expand Down
1 change: 1 addition & 0 deletions rustmail/src/handlers/guild_messages_handler.rs
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,7 @@ impl GuildMessagesHandler {
wrap_command!(lock, "logs", logs);
wrap_command!(lock, "take", take);
wrap_command!(lock, "release", release);
wrap_command!(lock, "ping", ping);

drop(lock);
h
Expand Down
16 changes: 16 additions & 0 deletions rustmail/src/i18n/language/en.rs
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,10 @@ pub fn load_english_messages(dict: &mut ErrorDictionary) {
"discord.user_is_a_bot".to_string(),
DictionaryMessage::new("The specified user is a rustmail."),
);
dict.messages.insert(
"discord.shard_manager_not_found".to_string(),
DictionaryMessage::new("Shard manager not found."),
);
dict.messages.insert(
"command.invalid_format".to_string(),
DictionaryMessage::new("Invalid command format")
Expand Down Expand Up @@ -883,6 +887,10 @@ pub fn load_english_messages(dict: &mut ErrorDictionary) {
"help.release".to_string(),
DictionaryMessage::new("Releases ownership of a ticket previously taken with the `!take` command. To release a ticket, use `!release` in the ticket."),
);
dict.messages.insert(
"help.ping".to_string(),
DictionaryMessage::new("Shows the actual latency of the bot."),
);
dict.messages.insert(
"add_reminder.helper".to_string(),
DictionaryMessage::new(
Expand Down Expand Up @@ -928,4 +936,12 @@ pub fn load_english_messages(dict: &mut ErrorDictionary) {
"slash_command.help_command_argument_desc".to_string(),
DictionaryMessage::new("The command to get help with"),
);
dict.messages.insert(
"slash_command.ping_command_desc".to_string(),
DictionaryMessage::new("Check the Discord bot latency."),
);
dict.messages.insert(
"slash_command.ping_command".to_string(),
DictionaryMessage::new("## Latency\n\nGateway latency: **{gateway_latency}** ms\nMinimal REST latency (GET /gateway): **{api_latency}** ms\nREST latency (message send): **{message_latency}** ms"),
);
}
Loading