Re-implementation of the citation system

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"

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

src/cache.rs

@@ -0,0 +1,49 @@
+use anyhow::Context as _;
+use once_cell::sync::Lazy;
+use serenity::all::{ChannelId, GuildChannel, GuildId};
+use serenity::client::Context;
+
+pub static MESSAGE_PREVIEW_CHANNEL_CACHE: Lazy<moka::future::Cache<ChannelId, GuildChannel>> = {
+    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<GuildChannel> {
+    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
+    }
+}

src/config.rs

@@ -0,0 +1,63 @@
+use crate::get_env_config;
+
+pub static CONFIG: once_cell::sync::OnceCell<PreviewConfig> = once_cell::sync::OnceCell::new();
+
+#[derive(serde::Deserialize, Debug)]
+pub struct PreviewConfig {
+    // Enable optional features.
+    pub feature_flag: Option<String>,
+    // 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))
+    }
+}

src/event.rs

@@ -1 +1,2 @@
+pub mod preview;
 pub mod ready;

src/event/preview.rs

@@ -0,0 +1,103 @@
+use crate::config::PreviewConfig;
+use crate::message::MessageLinkIDs;
+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;
+
+#[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<String>,
+    pub channel_name: String,
+    pub create_at: serenity::model::Timestamp,
+    pub attachment_url: Option<String>,
+}
+
+#[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;
+        };
+
+        tracing::info!(
+            "Start processing citation requests from {} ({:?})",
+            request.author.name,
+            ids
+        );
+        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);
+        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.permission_overwrites.iter().any(|p| {
+                matches!(p.kind, PermissionOverwriteType::Role(_))
+                    && p.deny.contains(Permissions::VIEW_CHANNEL)
+            })
+        {
+            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);
+        }
+        tracing::info!("-- Preview sent successfully.");
+    }
+}
+
+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)
+}

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(),

src/main.rs

@@ -1,38 +1,34 @@
 #![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<String>,
     pub discord_api_token: String,
+    pub config_file_path: Option<String>,
 }
 
 pub fn get_env_config() -> &'static EnvConfig {
     static ENV_CONFIG: std::sync::OnceLock<EnvConfig> = std::sync::OnceLock::new();
     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(
@@ -41,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()
@@ -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.");