From 87344be00018b4e7669b72d4024b7f1440673e8a Mon Sep 17 00:00:00 2001 From: Sho Sakuma Date: Fri, 24 Jan 2025 02:30:03 +0900 Subject: [PATCH 1/3] feat: Rebuild configuration system --- config/config.toml | 34 ++++++++++++++----------- src/config.rs | 63 ++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 83 insertions(+), 14 deletions(-) create mode 100644 src/config.rs diff --git a/config/config.toml b/config/config.toml index 5e39e6a..8f73db4 100644 --- a/config/config.toml +++ b/config/config.toml @@ -1,18 +1,24 @@ -### babyrite configuration ### -# This file is used to configure babyrite. -# You can change the settings here to customize babyrite's behavior. -# For more information, see the babyrite documentation at https://babyrite.m1sk9.dev/configuration.html +### babyrite configuration file ### +# This file is used for babyrite configuration. +# By changing this configuration, you can modify the behavior of babyrite. Please restart babyrite after changing the settings. +# The settings are written in TOML format. To use the configuration file, specify the recursive path to the configuration file in the environment variable "CONFIG_FILE_PATH" injected into the container. +# Note that babyrite will run with default settings if not configured. -## == General settings == ## +# Startup settings (Feature Flags) +## Flags to change the behavior of babyrite. Specify them separated by commas. The default is empty. +## The valid flags are as follows: +## - "json_logging": Outputs logs in JSON format (useful when integrating with log collection tools like Grafana Loki) +feature_flag = "" -# Sets the format of babyrite log output. -# * `compact`: output logs in normal format. -# * `json`: Output logs in JSON format. This is the recommended setting if you are using Grafana Loki or similar. -logger_format='compact' +# Mention settings for preview generation (default is enabled) +## Mentions the request sender when generating a preview. +is_mention = true -## == Preview settings == ## +# Preview deletion feature (default is enabled) +## Sets whether to enable the deletion of previews. If enabled, a trash can reaction is added to each preview, and pressing the reaction will delete it. +## If disabled, the reaction is not added, and pressing the reaction will not delete the preview. +is_deletable = true -[preview] - -# Whether to enable mentions in quoted messages. -is_mention=true +# Allow NSFW (Not Safe For Work) content (default is disabled) +## Sets whether to allow the generation of messages from channels marked as NSFW. It is strongly recommended to disable this setting on community servers. +is_allow_nsfw = false diff --git a/src/config.rs b/src/config.rs new file mode 100644 index 0000000..0763886 --- /dev/null +++ b/src/config.rs @@ -0,0 +1,63 @@ +use crate::get_env_config; + +pub static CONFIG: once_cell::sync::OnceCell = once_cell::sync::OnceCell::new(); + +#[derive(serde::Deserialize, Debug)] +pub struct PreviewConfig { + // Enable optional features. + pub feature_flag: Option, + // If enabled, previews are generated with mentions. + pub is_mention: bool, + // If enabled, preview can be deleted. + pub is_deletable: bool, + // If enabled, allow preview generation of NSFW content. + pub is_allow_nsfw: bool, +} + +impl Default for PreviewConfig { + fn default() -> Self { + Self { + feature_flag: None, + is_mention: true, + is_deletable: true, + is_allow_nsfw: false, + } + } +} + +#[derive(thiserror::Error, Debug)] +pub enum PreviewConfigError { + #[error("Failed to read configuration file.")] + FailedToReadConfig, + #[error("Failed to parse configuration file.")] + FailedToParseConfig, +} + +impl PreviewConfig { + pub fn init() -> anyhow::Result<(), PreviewConfigError> { + let envs = get_env_config(); + match &envs.config_file_path { + Some(p) => { + let buffer = &std::fs::read_to_string(p) + .map_err(|_| PreviewConfigError::FailedToReadConfig)?; + let config: PreviewConfig = + toml::from_str(buffer).map_err(|_| PreviewConfigError::FailedToParseConfig)?; + Ok(CONFIG + .set(config) + .map_err(|_| PreviewConfigError::FailedToParseConfig)?) + } + None => Ok(CONFIG + .set(PreviewConfig::default()) + .map_err(|_| PreviewConfigError::FailedToParseConfig)?), + } + } + + pub fn get_config() -> &'static PreviewConfig { + CONFIG.get().expect("Failed to get configuration.") + } + + pub fn get_feature_flag(flag: &str) -> bool { + let c = Self::get_config(); + c.feature_flag.as_ref().is_some_and(|f| f.contains(flag)) + } +} From 678d55e9be0ca0c033f4e879fc05b73f486a2391 Mon Sep 17 00:00:00 2001 From: Sho Sakuma Date: Fri, 24 Jan 2025 02:30:36 +0900 Subject: [PATCH 2/3] feat: Rebuild the citation system --- src/cache.rs | 49 +++++++++++++++++++++++ src/event.rs | 1 + src/event/preview.rs | 94 ++++++++++++++++++++++++++++++++++++++++++++ src/event/ready.rs | 4 +- src/main.rs | 28 ++++++------- src/message.rs | 88 +++++++++++++++++++++++++++++++++++++++++ 6 files changed, 246 insertions(+), 18 deletions(-) create mode 100644 src/cache.rs create mode 100644 src/event/preview.rs create mode 100644 src/message.rs diff --git a/src/cache.rs b/src/cache.rs new file mode 100644 index 0000000..ae5919f --- /dev/null +++ b/src/cache.rs @@ -0,0 +1,49 @@ +use anyhow::Context as _; +use once_cell::sync::Lazy; +use serenity::all::{Channel, ChannelId, Guild, GuildChannel, GuildId}; +use serenity::client::Context; + +pub static MESSAGE_PREVIEW_CHANNEL_CACHE: Lazy> = { + Lazy::new(|| { + moka::future::CacheBuilder::new(1000) + .name("message_preview_channel_cache") + .time_to_idle(std::time::Duration::from_secs(3600)) + .build() + }) +}; + +pub async fn get_channel_from_cache( + id: ChannelId, + guild: GuildId, + ctx: &Context, +) -> anyhow::Result { + let channel = if let Some(c) = MESSAGE_PREVIEW_CHANNEL_CACHE.get(&id).await { + Ok(c) + } else { + let channels = guild + .channels(&ctx.http) + .await + .context("Failed to get channels.")?; + if let Some(channel) = channels.get(&id) { + Ok(channel.clone()) + } else { + let guild_threads = guild + .get_active_threads(&ctx.http) + .await + .context("Failed to get active threads.")?; + guild_threads + .threads + .iter() + .find(|c| c.id == id) + .cloned() + .context("Failed to find channel.") + } + }; + + if let Ok(c) = channel { + MESSAGE_PREVIEW_CHANNEL_CACHE.insert(id, c.clone()).await; + Ok(c) + } else { + channel + } +} diff --git a/src/event.rs b/src/event.rs index 849debc..59a02b6 100644 --- a/src/event.rs +++ b/src/event.rs @@ -1 +1,2 @@ +pub mod preview; pub mod ready; diff --git a/src/event/preview.rs b/src/event/preview.rs new file mode 100644 index 0000000..90f8512 --- /dev/null +++ b/src/event/preview.rs @@ -0,0 +1,94 @@ +use crate::config::PreviewConfig; +use crate::message::MessageLinkIDs; +use serenity::all::{Context, CreateEmbed, Message, ReactionType}; +use serenity::builder::{ + CreateAllowedMentions, CreateEmbedAuthor, CreateEmbedFooter, CreateMessage, +}; +use serenity::prelude::EventHandler; + +pub struct PreviewHandler; + +#[derive(thiserror::Error, Debug)] +pub enum PreviewHandlerError { + #[error("Failed to get preview: {0}")] + FailedToGetPreview(#[from] anyhow::Error), +} + +#[derive(Debug, typed_builder::TypedBuilder)] +pub struct PreviewEmbedArgs { + pub content: String, + pub author_name: String, + pub author_avatar: Option, + pub channel_name: String, + pub create_at: serenity::model::Timestamp, + pub attachment_url: Option, +} + +#[serenity::async_trait] +impl EventHandler for PreviewHandler { + async fn message(&self, ctx: Context, request: Message) { + // check if the message is command or bot + if request.content.starts_with("b!") || request.author.bot { + return; + }; + + let config = PreviewConfig::get_config(); + + let Some(ids) = MessageLinkIDs::parse_url(&request.content) else { + return; + }; + + let preview = match ids.get_message(&ctx).await { + Ok(p) => p, + Err(e) => { + tracing::error!("Failed to get preview: {:?}", e); + return; + } + }; + let (message, channel) = (preview.preview_message, preview.preview_channel); + + if channel.nsfw && !config.is_allow_nsfw + || channel + .permissions + .is_some_and(|p| !p.read_message_history()) + { + return; + } + + let args = PreviewEmbedArgs::builder() + .content(message.content.clone()) + .author_name(message.author.name.clone()) + .author_avatar(message.author.avatar_url().clone()) + .channel_name(channel.name) + .create_at(message.timestamp) + .attachment_url(message.attachments.first().map(|a| a.url.clone())) + .build(); + let preview = CreateMessage::default() + .embed(generate_preview(args)) + .reference_message(&request) + .reactions(match config.is_deletable { + true => vec![ReactionType::Unicode("🗑️".to_string())], + false => vec![], + }) + .allowed_mentions(CreateAllowedMentions::new().replied_user(config.is_mention)); + if let Err(e) = request.channel_id.send_message(&ctx.http, preview).await { + tracing::error!("Failed to send preview: {:?}", e); + } + } +} + +fn generate_preview(args: PreviewEmbedArgs) -> CreateEmbed { + CreateEmbed::default() + .description(args.content) + .author( + CreateEmbedAuthor::new(&args.author_name).icon_url( + args.author_avatar + .as_deref() + .unwrap_or("https://cdn.discordapp.com/embed/avatars/0.png"), + ), + ) + .footer(CreateEmbedFooter::new(&args.channel_name)) + .timestamp(args.create_at) + .image(args.attachment_url.unwrap_or_default()) + .color(0x7A4AFF) +} diff --git a/src/event/ready.rs b/src/event/ready.rs index 9993822..e9fadf5 100644 --- a/src/event/ready.rs +++ b/src/event/ready.rs @@ -1,8 +1,10 @@ use serenity::all::{ActivityData, Context, Ready}; use serenity::prelude::EventHandler; +pub struct ReadyHandler; + #[serenity::async_trait] -impl EventHandler for crate::Handler { +impl EventHandler for ReadyHandler { async fn ready(&self, ctx: Context, bot: Ready) { ctx.set_activity( ActivityData::custom(format!("Running v{}", env!("CARGO_PKG_VERSION"))).into(), diff --git a/src/main.rs b/src/main.rs index 19050bb..7c9421d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,16 +1,20 @@ #![deny(clippy::all)] +mod cache; +mod config; mod event; +mod message; +use crate::config::PreviewConfig; +use crate::event::preview::PreviewHandler; +use crate::event::ready::ReadyHandler; use serenity::all::GatewayIntents; use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; -pub struct Handler; - #[derive(serde::Deserialize)] pub struct EnvConfig { - pub feature_flag: Option, pub discord_api_token: String, + pub config_file_path: Option, } pub fn get_env_config() -> &'static EnvConfig { @@ -18,21 +22,13 @@ pub fn get_env_config() -> &'static EnvConfig { ENV_CONFIG.get_or_init(|| envy::from_env().expect("Failed to load environment configuration.")) } -impl EnvConfig { - pub fn get_feature(&self, feature: &str) -> bool { - self.feature_flag - .as_ref() - .is_some_and(|f| f.contains(feature)) - } -} - #[tokio::main] async fn main() -> anyhow::Result<()> { dotenvy::dotenv().ok(); + PreviewConfig::init()?; let envs = get_env_config(); - // FIXME: I want to refactor this to be more idiomatic. - match envs.get_feature("json_logging") { + match PreviewConfig::get_feature_flag("json_logging") { true => { tracing_subscriber::registry() .with( @@ -54,14 +50,12 @@ async fn main() -> anyhow::Result<()> { } } - // TODO: Add configuration logic here. - // ... - let mut client = serenity::Client::builder( &envs.discord_api_token, GatewayIntents::MESSAGE_CONTENT | GatewayIntents::GUILD_MESSAGES, ) - .event_handler(Handler) + .event_handler(ReadyHandler) + .event_handler(PreviewHandler) .await .expect("Failed to initialize client."); diff --git a/src/message.rs b/src/message.rs new file mode 100644 index 0000000..6f00e13 --- /dev/null +++ b/src/message.rs @@ -0,0 +1,88 @@ +use crate::cache::get_channel_from_cache; +use once_cell::sync::Lazy; +use regex::Regex; +use serenity::all::{ChannelId, GuildChannel, GuildId, Message, MessageId}; +use url::Url; + +pub static MESSAGE_LINK_REGEX: Lazy = Lazy::new(|| { + Regex::new(r"https://(?:ptb\.|canary\.)?discord\.com/channels/(\d+)/(\d+)/(\d+)").unwrap() +}); + +#[derive(serde::Deserialize, Debug)] +pub struct MessageLinkIDs { + pub guild_id: GuildId, + pub channel_id: ChannelId, + pub message_id: MessageId, +} + +#[derive(Debug)] +pub struct MessagePreview { + pub preview_message: Message, + pub preview_channel: GuildChannel, +} + +impl MessageLinkIDs { + pub fn parse_url(text: &str) -> Option { + if !MESSAGE_LINK_REGEX.is_match(text) { + return None; + } + + if let Some(captures) = MESSAGE_LINK_REGEX.captures(text) { + let url = Url::parse(captures.get(0)?.as_str()).ok()?; + + if !matches!( + url.domain(), + Some("discord.com") | Some("canary.discord.com") | Some("ptb.discord.com") + ) { + return None; + } + + let guild_id = GuildId::new(captures.get(1)?.as_str().parse().ok()?); + let channel_id = ChannelId::new(captures.get(2)?.as_str().parse().ok()?); + let message_id = MessageId::new(captures.get(3)?.as_str().parse().ok()?); + + Some(MessageLinkIDs { + guild_id, + channel_id, + message_id, + }) + } else { + None + } + } + + pub async fn get_message( + &self, + ctx: &serenity::prelude::Context, + ) -> anyhow::Result { + let guild = self.guild_id; + let channel = get_channel_from_cache(self.channel_id, guild, ctx).await?; + let message = channel.message(&ctx.http, self.message_id).await?; + + Ok(MessagePreview { + preview_message: message, + preview_channel: channel, + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_parse_url() { + let url = "https://canary.discord.com/channels/1331992336129069118/1331992336560947271/1332012065077854368"; + let ids = MessageLinkIDs::parse_url(url).unwrap(); + assert_eq!(ids.guild_id, GuildId::new(1331992336129069118)); + assert_eq!(ids.channel_id, ChannelId::new(1331992336560947271)); + assert_eq!(ids.message_id, MessageId::new(1332012065077854368)); + } + + #[test] + fn test_invalid_domain() { + let url = "https://discord.gg/invite"; + let ids = MessageLinkIDs::parse_url(url); + assert!(ids.is_none()); + } +} From af0eae7a427b6e1b33d508b783ca84f2deaaad5d Mon Sep 17 00:00:00 2001 From: Sho Sakuma Date: Fri, 24 Jan 2025 02:31:20 +0900 Subject: [PATCH 3/3] deps: Add `typed-builder` --- Cargo.lock | 21 +++++++++++++++++++++ Cargo.toml | 1 + src/cache.rs | 2 +- src/event/preview.rs | 19 ++++++++++++++----- src/main.rs | 2 +- 5 files changed, 38 insertions(+), 7 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 0c85260..f969b4d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -86,6 +86,7 @@ dependencies = [ "toml", "tracing", "tracing-subscriber", + "typed-builder", "url", ] @@ -1830,6 +1831,26 @@ dependencies = [ "utf-8", ] +[[package]] +name = "typed-builder" +version = "0.20.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e14ed59dc8b7b26cacb2a92bad2e8b1f098806063898ab42a3bd121d7d45e75" +dependencies = [ + "typed-builder-macro", +] + +[[package]] +name = "typed-builder-macro" +version = "0.20.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "560b82d656506509d43abe30e0ba64c56b1953ab3d4fe7ba5902747a7a3cedd5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "typemap_rev" version = "0.3.0" diff --git a/Cargo.toml b/Cargo.toml index 81f3d45..835cc5b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,6 +15,7 @@ tracing = { version = "0.1.40" } thiserror = { version = "2.0.0" } url = { version = "2.5.2" } toml = { version = "0.8.19" } +typed-builder = { version = "0.20.0" } [dependencies.serenity] version = "0.12.1" diff --git a/src/cache.rs b/src/cache.rs index ae5919f..b4f9e49 100644 --- a/src/cache.rs +++ b/src/cache.rs @@ -1,6 +1,6 @@ use anyhow::Context as _; use once_cell::sync::Lazy; -use serenity::all::{Channel, ChannelId, Guild, GuildChannel, GuildId}; +use serenity::all::{ChannelId, GuildChannel, GuildId}; use serenity::client::Context; pub static MESSAGE_PREVIEW_CHANNEL_CACHE: Lazy> = { diff --git a/src/event/preview.rs b/src/event/preview.rs index 90f8512..3414f9d 100644 --- a/src/event/preview.rs +++ b/src/event/preview.rs @@ -1,9 +1,10 @@ use crate::config::PreviewConfig; use crate::message::MessageLinkIDs; -use serenity::all::{Context, CreateEmbed, Message, ReactionType}; +use serenity::all::{Context, CreateEmbed, Message, PermissionOverwriteType, ReactionType}; use serenity::builder::{ CreateAllowedMentions, CreateEmbedAuthor, CreateEmbedFooter, CreateMessage, }; +use serenity::model::Permissions; use serenity::prelude::EventHandler; pub struct PreviewHandler; @@ -33,11 +34,15 @@ impl EventHandler for PreviewHandler { }; let config = PreviewConfig::get_config(); - let Some(ids) = MessageLinkIDs::parse_url(&request.content) else { return; }; + tracing::info!( + "Start processing citation requests from {} ({:?})", + request.author.name, + ids + ); let preview = match ids.get_message(&ctx).await { Ok(p) => p, Err(e) => { @@ -46,11 +51,14 @@ impl EventHandler for PreviewHandler { } }; let (message, channel) = (preview.preview_message, preview.preview_channel); + tracing::debug!("Message: {:?}, Channel: {:?}", message, channel); + // Verify that: @everyone on the previewer channel does not have read permission (i.e. limit channel) if channel.nsfw && !config.is_allow_nsfw - || channel - .permissions - .is_some_and(|p| !p.read_message_history()) + || channel.permission_overwrites.iter().any(|p| { + matches!(p.kind, PermissionOverwriteType::Role(_)) + && p.deny.contains(Permissions::VIEW_CHANNEL) + }) { return; } @@ -74,6 +82,7 @@ impl EventHandler for PreviewHandler { if let Err(e) = request.channel_id.send_message(&ctx.http, preview).await { tracing::error!("Failed to send preview: {:?}", e); } + tracing::info!("-- Preview sent successfully."); } } diff --git a/src/main.rs b/src/main.rs index 7c9421d..09f4865 100644 --- a/src/main.rs +++ b/src/main.rs @@ -37,7 +37,7 @@ async fn main() -> anyhow::Result<()> { ) .with(tracing_subscriber::fmt::layer().json()) .init(); - tracing::info!("JSON logging enabled."); + tracing::info!("Feature Flag : Log output in JSON format is now enabled."); } false => { tracing_subscriber::registry()