Add support for editing messages

src/database/models.rs

@@ -18,7 +18,7 @@ use twilight_model::{
     user::User as TwilightUser,
 };
 
-use crate::{constants::DEFAULT_PREFIX, groups::calls::DisplayProfile, utils::get_random_u64};
+use crate::{constants::DEFAULT_PREFIX, groups::calls::DisplayProfile, utils::{get_random_u64, i64_to_u64, u64_to_i64}};
 
 #[derive(Debug, Error)]
 pub enum DatabaseOrDiscordError {
@@ -44,34 +44,6 @@ pub enum DatabaseError {
     DeserializationFailed,
 }
 
-/// Copies the bytes from [`i64`] to [`u64`].
-///
-/// WARNING: This does NOT keep the same "represented" value. This only keeps the bytes intact,
-/// e.g. 0b0101 will still be 0b0101, but the value represented will change since the type changes.
-///
-/// Arguments:
-/// * `i` - The [`i64`] to convert.
-///
-/// Returns:
-/// [`u64`] - The resulting [`u64`].
-fn i64_to_u64(i: i64) -> u64 {
-    u64::from_ne_bytes(i.to_ne_bytes())
-}
-
-/// Copies the bytes from [`u64`] to [`i64`].
-///
-/// WARNING: This does NOT keep the same "represented" value. This only keeps the bytes intact,
-/// e.g. 0b0101 will still be 0b0101, but the value represented will change since the type changes.
-///
-/// Arguments:
-/// * `u` - The [`u64`] to convert.
-///
-/// Returns:
-/// [`i64`] - The resulting [`i64`].
-fn u64_to_i64(u: u64) -> i64 {
-    i64::from_ne_bytes(u.to_ne_bytes())
-}
-
 trait Object<'a> {
     type Row: DeserializeRow<'a, 'a>;
 
@@ -967,6 +939,11 @@ pub struct MessageBuilder {
     message: Option<TwilightMessage>,
     profile: Option<DisplayProfile>,
     call_id: Option<u64>,
+    is_deleted: bool,
+    is_censored: bool,
+    propagation_message_id: Option<u64>,
+    propagation_channel_id: Option<u64>,
+    replying_to_id: Option<u64>,
 }
 
 impl MessageBuilder {
@@ -975,6 +952,11 @@ impl MessageBuilder {
             message: None,
             profile: None,
             call_id: None,
+            is_deleted: false,
+            is_censored: false,
+            propagation_message_id: None,
+            propagation_channel_id: None,
+            replying_to_id: None,
         }
     }
 
@@ -993,6 +975,31 @@ impl MessageBuilder {
         self
     }
 
+    pub fn is_deleted(mut self, is_deleted: bool) -> Self {
+        self.is_deleted = is_deleted;
+        self
+    }
+
+    pub fn is_censored(mut self, is_censored: bool) -> Self {
+        self.is_censored = is_censored;
+        self
+    }
+
+    pub fn propagation_message_id(mut self, propagation_message_id: u64) -> Self {
+        self.propagation_message_id = Some(propagation_message_id);
+        self
+    }
+
+    pub fn propagation_channel_id(mut self, propagation_channel_id: u64) -> Self {
+        self.propagation_channel_id = Some(propagation_channel_id);
+        self
+    }
+
+    pub fn replying_to_id(mut self, replying_to_id: u64) -> Self {
+        self.replying_to_id = Some(replying_to_id);
+        self
+    }
+
     pub fn build(self) -> Result<Message, MessageBuilderError> {
         let message = self.message.ok_or(MessageBuilderError)?;
         let profile = self.profile.ok_or(MessageBuilderError)?;
@@ -1011,12 +1018,12 @@ impl MessageBuilder {
                 .sticker_items
                 .first()
                 .map(|s| u64_to_i64(s.id.get())),
-            propagation_message_id: None,
+            propagation_message_id: self.propagation_message_id.map(u64_to_i64),
-            propagation_channel_id: None,
+            propagation_channel_id: self.propagation_channel_id.map(u64_to_i64),
-            replying_to_id: None,
+            replying_to_id: self.replying_to_id.map(u64_to_i64),
-            is_deleted: false,
+            is_deleted: self.is_deleted,
             is_edited: message.edited_timestamp.is_some(),
-            is_censored: false,
+            is_censored: self.is_censored,
         };
 
         Ok(Message {
@@ -1063,6 +1070,46 @@ impl Message {
         MessageBuilder::new()
     }
 
+    pub async fn get_by_call_and_message_id(
+        call_id: u64,
+        message_id: u64,
+        session: &Session,
+    ) -> Result<Option<Self>, DatabaseError> {
+        let call_id = u64_to_i64(call_id);
+        let message_id = u64_to_i64(message_id);
+
+        let result = session
+            .query_unpaged(
+                r#"
+                    SELECT
+                        call_id,
+                        message_id,
+                        author_id,
+                        data
+                    FROM
+                        payphone.messages
+                    WHERE
+                        call_id = ? AND message_id = ?
+                    LIMIT 1;
+                "#,
+                (call_id, message_id),
+            )
+            .await
+            .map_err(DatabaseError::QueryError)?
+            .into_rows_result()
+            .map_err(|_| DatabaseError::ResultNotRows)?;
+
+        match result.single_row::<<Message as Object>::Row>() {
+            Ok(row) => Ok(Some(Self::from_row(row))),
+            Err(error) => match error {
+                SingleRowError::DeserializationFailed(_) | SingleRowError::TypeCheckFailed(_) => {
+                    Err(DatabaseError::DeserializationFailed)
+                }
+                SingleRowError::UnexpectedRowCount(_) => Ok(None),
+            },
+        }
+    }
+
     pub async fn insert(&self, session: &Session) -> Result<(), DatabaseError> {
         session
             .query_unpaged(

src/groups/calls/events.rs

@@ -2,13 +2,16 @@ use std::{error::Error, sync::Arc};
 
 use async_trait::async_trait;
 use itertools::Itertools;
-use twilight_model::{gateway::payload::incoming::MessageCreate, id::Id};
+use twilight_model::{
+    gateway::payload::incoming::{MessageCreate, MessageUpdate},
+    id::Id,
+};
 
 use crate::{
     database::models::{Call, Channel, Message, User, UserIncident},
     discord::sanitizing::SanitizeDiscordMarkdown,
     framework::{
-        events::traits::{EventResult, OnMessageHandler},
+        events::traits::{EventResult, OnMessageEditHandler, OnMessageHandler},
         interactions::components::ActionRow,
         runtime::Runtime,
     },
@@ -18,8 +21,62 @@ use crate::{
         moderation::{Action, Rule, escape_message},
     },
     state::State,
+    utils::i64_to_u64,
 };
 
+async fn penalize_user(
+    user: User,
+    message: Message,
+    broken_rules: Vec<Arc<dyn Rule>>,
+    rt: &Runtime<State>,
+) -> Result<(), Box<dyn Error>> {
+    let reason = format!(
+        "You broke the following rules during a call:\n{}",
+        broken_rules
+            .iter()
+            .map(|rule| format!("- {}", rule.name()))
+            .join("\n")
+    );
+
+    let new_incident = user
+        .incident(
+            reason.clone(),
+            message.data.clone(),
+            Some(user.id),
+            &rt.state.db,
+        )
+        .await?;
+
+    let dm_channel = rt
+        .http
+        .create_private_channel(Id::new(user.id))
+        .await?
+        .model()
+        .await?;
+
+    let embed = if new_incident.is_first_incident {
+        embeds::first_warning_embed(&rt.app, &message, broken_rules)
+    } else if new_incident.counting_incidents == 0 {
+        embeds::repeat_warning_embed(&rt.app, &message, broken_rules)
+    } else if new_incident.counting_incidents == 1 {
+        embeds::first_ban_embed(&rt.app, &message, broken_rules)
+    } else {
+        embeds::final_ban_embed(&rt.app, &message, broken_rules)
+    };
+
+    rt.http
+        .create_message(dm_channel.id)
+        .embeds(&[embed])
+        .components(&[ActionRow::new(&[
+            (None, Arc::new(CodeOfConductButton)),
+            (None, Arc::new(NeedHelpButton)),
+        ])
+        .build()])
+        .await?;
+
+    Ok(())
+}
+
 #[derive(Debug)]
 pub struct CallMessageHandler;
 
@@ -114,18 +171,73 @@ impl OnMessageHandler<State> for CallMessageHandler {
 
         let profile = get_display_profile(event);
 
-        let message = Message::build()
-            .call_id(call.id)
-            .message(event.0.clone())
-            .profile(profile.clone())
-            .build()
-            .expect(
-                "The message builder was missing either the call ID, the message, or the profile",
-            );
-        message.insert(&rt.state.db).await?;
-
         let (content, rules) = escape_message(event);
 
+        let message: Message;
+
+        if rules
+            .iter()
+            .any(|rule| rule.action() == &Action::BlockAndPenalty)
+        {
+            rt.http
+                .create_message(event.channel_id)
+                .content(
+                    "❌ Your message has been blocked because it broke Payphone's code of conduct. Check your DMs for more information.",
+                )
+                .await?;
+
+            rt.http
+                .create_message(Id::new(opposite_channel.id))
+                .content(&format!(
+                    "❌ **{}** Sent a message that was blocked.\n-# If you see an inappropiate message, please report it using `Message > Apps > Report Message`.",
+                    profile.name.sanitize(false)
+                ))
+                .await?;
+
+            message = Message::build()
+                .call_id(call.id)
+                .message(event.0.clone())
+                .profile(profile.clone())
+                .is_censored(!rules.is_empty())
+                .build()
+                .expect(
+                    "The message builder was missing either the call ID, the message, or the profile",
+                );
+        } else {
+            let message_sent = rt
+                .http
+                .execute_webhook(Id::new(webhook_id), &webhook_token)
+                .allowed_mentions(None)
+                .username(&profile.name)
+                .avatar_url(&profile.avatar_url)
+                .content(&content)
+                .wait()
+                .await
+                .inspect_err(|e| {
+                    log::error!(
+                        "Failed to send message to webhook for call {} in channel {}: {e:?}",
+                        call.id,
+                        opposite_channel.id
+                    );
+                })?
+                .model()
+                .await?;
+
+            message = Message::build()
+                .call_id(call.id)
+                .message(event.0.clone())
+                .profile(profile.clone())
+                .is_censored(!rules.is_empty())
+                .propagation_message_id(message_sent.id.get())
+                .propagation_channel_id(opposite_channel.id)
+                .build()
+                .expect(
+                    "The message builder was missing either the call ID, the message, or the profile",
+                );
+        }
+
+        message.insert(&rt.state.db).await?;
+
         if rules
             .iter()
             .any(|rule| rule.action() == &Action::BlockAndPenalty)
@@ -142,92 +254,94 @@ impl OnMessageHandler<State> for CallMessageHandler {
                 user
             };
 
-            rt.http
-                .create_message(event.channel_id)
-                .content(
-                    "❌ Your message has been blocked because it broke Payphone's code of conduct. Check your DMs for more information.",
-                )
-                .await?;
-
-            rt.http
-                .create_message(Id::new(opposite_channel.id))
-                .content(&format!(
-                    "❌ **{}** Sent a message that was blocked.\n-# If you see an inappropiate message, please report it using `Message > Apps > Report Message`.",
-                    profile.name.sanitize(false)
-                ))
-                .await?;
-
             penalize_user(user, message, rules, &rt).await?;
-        } else {
-            rt.http
-                .execute_webhook(Id::new(webhook_id), &webhook_token)
-                .allowed_mentions(None)
-                .username(&profile.name)
-                .avatar_url(&profile.avatar_url)
-                .content(&content)
-                .await
-                .inspect_err(|e| {
-                    log::error!(
-                        "Failed to send message to webhook for call {} in channel {}: {e:?}",
-                        call.id,
-                        opposite_channel.id
-                    );
-                })?;
         }
 
         Ok(())
     }
 }
 
-async fn penalize_user(
+#[derive(Debug)]
-    user: User,
+pub struct CallMessageEditHandler;
-    message: Message,
-    broken_rules: Vec<Arc<dyn Rule>>,
-    rt: &Runtime<State>,
-) -> Result<(), Box<dyn Error>> {
-    let reason = format!(
-        "You broke the following rules during a call:\n{}",
-        broken_rules
-            .iter()
-            .map(|rule| format!("- {}", rule.name()))
-            .join("\n")
-    );
 
-    let new_incident = user
+#[async_trait]
-        .incident(
+impl OnMessageEditHandler<State> for CallMessageEditHandler {
-            reason.clone(),
+    async fn on_message_edit(&self, event: &MessageUpdate, rt: Arc<Runtime<State>>) -> EventResult {
-            message.data.clone(),
+        let Some(channel) = Channel::get_by_id(event.channel_id.get(), &rt.state.db).await? else {
-            Some(user.id),
+            return Ok(());
-            &rt.state.db,
+        };
-        )
-        .await?;
 
-    let dm_channel = rt
+        let call = if let Some(call_id) = channel.call_id {
-        .http
+            if let Some(call) = Call::get_by_id(call_id, &rt.state.db).await? {
-        .create_private_channel(Id::new(user.id))
+                if !call.is_active() {
-        .await?
+                    return Ok(());
-        .model()
+                }
-        .await?;
 
-    let embed = if new_incident.is_first_incident {
+                call
-        embeds::first_warning_embed(&rt.app, &message, broken_rules)
+            } else {
-    } else if new_incident.counting_incidents == 0 {
+                return Ok(());
-        embeds::repeat_warning_embed(&rt.app, &message, broken_rules)
+            }
-    } else if new_incident.counting_incidents == 1 {
+        } else {
-        embeds::first_ban_embed(&rt.app, &message, broken_rules)
+            return Ok(());
-    } else {
+        };
-        embeds::final_ban_embed(&rt.app, &message, broken_rules)
-    };
 
-    rt.http
+        let Some(mut message) =
-        .create_message(dm_channel.id)
+            Message::get_by_call_and_message_id(call.id, event.id.get(), &rt.state.db).await?
-        .embeds(&[embed])
+        else {
-        .components(&[ActionRow::new(&[
+            return Ok(());
-            (None, Arc::new(CodeOfConductButton)),
+        };
-            (None, Arc::new(NeedHelpButton)),
-        ])
-        .build()])
-        .await?;
 
-    Ok(())
+        let opposite_channel_id = call.other_party_channel_id(event.channel_id.get());
+        let Some(opposite_channel) = Channel::get_by_id(opposite_channel_id, &rt.state.db).await?
+        else {
+            return Ok(());
+        };
+
+        if let Some(propagation_message_id) = message.data.propagation_message_id {
+            let (content, rules) = escape_message(event);
+
+            if rules
+                .iter()
+                .any(|rule| rule.action() == &Action::BlockAndPenalty)
+            {
+                rt.http
+                    .create_message(event.channel_id)
+                    .content(
+                        "❌ Your message's edit has been blocked because it broke Payphone's code of conduct. Check your DMs for more information.",
+                    )
+                    .reply(event.id)
+                    .await?;
+
+                let user = User::get_by_id(event.author.id.get(), &rt.state.db).await?;
+
+                if let Some(user) = user {
+                    penalize_user(user, message, rules, &rt).await?;
+                }
+
+                return Ok(());
+            }
+
+            if let Some((webhook_id, webhook_token)) = opposite_channel
+                .webhook_id
+                .zip(opposite_channel.webhook_token)
+            {
+                rt.http
+                    .update_webhook_message(
+                        Id::new(webhook_id),
+                        &webhook_token,
+                        Id::new(i64_to_u64(propagation_message_id)),
+                    )
+                    .content(Some(&content))
+                    .await?;
+
+                message.data.content = content;
+                message.data.is_edited = true;
+                message.data.is_censored = !rules.is_empty();
+
+                message.insert(&rt.state.db).await?;
+            }
+        }
+
+        Ok(())
+    }
 }

src/groups/calls/moderation.rs

@@ -88,7 +88,7 @@ use std::{collections::HashMap, sync::Arc};
 use linkify::LinkFinder;
 use once_cell::sync::Lazy;
 use regex::Regex;
-use twilight_model::gateway::payload::incoming::MessageCreate;
+use twilight_model::channel::Message;
 use unicode_segmentation::UnicodeSegmentation;
 use unidecode::unidecode;
 
@@ -103,7 +103,7 @@ use crate::discord::sanitizing::SanitizeDiscordMarkdown;
 /// Returns:
 /// [`(String, Vec<Arc<dyn Rule>>)`] - The escaped and censored message, and a vector of the rules
 /// broken.
-pub fn escape_message(event: &MessageCreate) -> (String, Vec<Arc<dyn Rule>>) {
+pub fn escape_message(event: &Message) -> (String, Vec<Arc<dyn Rule>>) {
     let mut message = event.content.to_string();
 
     for mention in &event.mentions {
@@ -274,10 +274,7 @@ type RuleAnalysis = ((Vec<String>, Vec<String>), Arc<dyn Rule>);
 /// Returns:
 /// [`Vec<RuleAnalysis>`] - The results of the analysis, including all detections (on the original
 /// message first, on the normalized message later) and their associated rules.
-pub fn analyze_message(
+pub fn analyze_message(original: &str, normalized: &str) -> Vec<RuleAnalysis> {
-    original: &str,
-    normalized: &str,
-) -> Vec<RuleAnalysis> {
     let mut results = Vec::new();
 
     for rule in RULES.iter() {
@@ -568,7 +565,10 @@ impl Rule for LinksRule {
     fn detect(&self, original: &str, _normalized: &str) -> (Vec<String>, Vec<String>) {
         let mut detections = Vec::new();
 
-        for link in LinkFinder::new().url_must_have_scheme(false).links(original) {
+        for link in LinkFinder::new()
+            .url_must_have_scheme(false)
+            .links(original)
+        {
             detections.push(link.as_str().to_string());
         }
 

src/main.rs

@@ -127,7 +127,7 @@ use crate::{
         about::{commands::HelpCommand, components::HelpGoBackButton},
         calls::{
             commands::{CallCommand, FriendRequestCommand, HangupCommand, SkipCommand},
-            events::CallMessageHandler,
+            events::{CallMessageEditHandler, CallMessageHandler},
         },
         gateway::events::{ReadyHandler, ReconnectHandler},
         servers::{
@@ -233,6 +233,7 @@ async fn start(token: String, database_addr: String, shard: ShardId) {
         .on_ready(ReadyHandler)
         .on_reconnect(ReconnectHandler)
         .on_message(CallMessageHandler)
+        .on_message_edit(CallMessageEditHandler)
         .on_server_create(ServerCreateHandler)
         .on_server_update(ServerUpdateHandler)
         .on_server_delete(ServerDeleteHandler);

src/utils.rs

@@ -72,3 +72,31 @@ pub fn ellipse(string: String, max_length: usize) -> String {
         format!("{}...", &string[..len])
     }
 }
+
+/// Copies the bytes from [`i64`] to [`u64`].
+///
+/// WARNING: This does NOT keep the same "represented" value. This only keeps the bytes intact,
+/// e.g. 0b0101 will still be 0b0101, but the value represented will change since the type changes.
+///
+/// Arguments:
+/// * `i` - The [`i64`] to convert.
+///
+/// Returns:
+/// [`u64`] - The resulting [`u64`].
+pub fn i64_to_u64(i: i64) -> u64 {
+    u64::from_ne_bytes(i.to_ne_bytes())
+}
+
+/// Copies the bytes from [`u64`] to [`i64`].
+///
+/// WARNING: This does NOT keep the same "represented" value. This only keeps the bytes intact,
+/// e.g. 0b0101 will still be 0b0101, but the value represented will change since the type changes.
+///
+/// Arguments:
+/// * `u` - The [`u64`] to convert.
+///
+/// Returns:
+/// [`i64`] - The resulting [`i64`].
+pub fn u64_to_i64(u: u64) -> i64 {
+    i64::from_ne_bytes(u.to_ne_bytes())
+}