Implement message editing for Discord and Matrix

src/bridgers.rs

@@ -1,12 +1,12 @@
-use std::{cell::RefCell, str::FromStr, sync::Mutex};
+use std::{cell::RefCell, collections::HashSet, str::FromStr, sync::Mutex};
 
 use sled::Db;
 
 use crate::{
     channels::ChannelReference,
-    discord::{self, forward_to_discord},
+    discord::{self, edit_on_discord, forward_to_discord},
-    matrix::{self, forward_to_matrix},
+    matrix::{self, edit_on_matrix, forward_to_matrix},
-    messages::{MessageReference, SentMessage},
+    messages::{EditedMessage, MessageReference, SentMessage},
 };
 
 pub struct Bridgers {
@@ -102,15 +102,15 @@ impl Bridgers {
         None
     }
 
-    fn get_related_matrix_message(&self, source: &MessageReference) -> Option<(String, String)> {
+    fn get_related_matrix_message(&self, source: &MessageReference) -> Option<matrix::EventId> {
-        if let MessageReference::Matrix(room_id, event_id) = source {
+        if let MessageReference::Matrix(_, event_id) = source {
-            return Some((room_id.to_string(), event_id.to_string()));
+            return Some(matrix::EventId::from_str(event_id).unwrap());
         }
 
         if let Some(relations) = self.get_related_messages(source) {
             for relation in relations {
-                if let MessageReference::Matrix(room_id, event_id) = relation {
+                if let MessageReference::Matrix(_, event_id) = relation {
-                    return Some((room_id, event_id));
+                    return Some(matrix::EventId::from_str(&event_id).unwrap());
                 }
             }
         }
@@ -154,7 +154,28 @@ impl Bridgers {
         }
     }
 
-    pub async fn send_message(&self, message: SentMessage) -> Vec<MessageReference> {
+    fn store_related_messages(&self, related_messages: &[MessageReference]) {
+        let tree = self
+            .db
+            .open_tree("message_relations")
+            .expect("Failed to open relations tree");
+
+        for source in related_messages {
+            let relations = related_messages
+                .iter()
+                .filter(|r| r != &source)
+                .collect::<Vec<_>>();
+
+            let key = bincode::serialize(source).expect("Failed to serialize message reference");
+            let value =
+                bincode::serialize(&relations).expect("Failed to serialize message relations");
+
+            tree.insert(key, value)
+                .expect("Failed to store message relations");
+        }
+    }
+
+    pub async fn send_message(&self, message: SentMessage) {
         let mut related_messages = vec![message.source.clone()];
 
         if let Some(discord) = self.discord.lock().unwrap().borrow().as_ref() {
@@ -183,25 +204,46 @@ impl Bridgers {
             }
         }
 
-        let tree = self
+        self.store_related_messages(&related_messages);
-            .db
-            .open_tree("message_relations")
-            .expect("Failed to open relations tree");
-
-        for source in related_messages.iter() {
-            let relations = related_messages
-                .iter()
-                .filter(|r| r != &source)
-                .collect::<Vec<_>>();
-
-            let key = bincode::serialize(source).expect("Failed to serialize message reference");
-            let value =
-                bincode::serialize(&relations).expect("Failed to serialize message relations");
-
-            tree.insert(key, value)
-                .expect("Failed to store message relations");
     }
 
-        related_messages
+    pub async fn edit_message(&self, message: EditedMessage) {
+        if let Some(related_messages) = self.get_related_messages(&message.replacing) {
+            let mut new_related_messages = HashSet::new();
+            related_messages.iter().for_each(|r| {
+                new_related_messages.insert(r.clone());
+            });
+
+            for related_message in related_messages.iter() {
+                match related_message {
+                    MessageReference::Discord(channel_id, message_id) => {
+                        if let Some(discord) = self.discord.lock().unwrap().borrow().as_ref() {
+                            let channel_id = discord::ChannelId(*channel_id);
+                            let message_id = discord::MessageId(*message_id);
+
+                            if let Some(m) =
+                                edit_on_discord(discord, channel_id, message_id, &message).await
+                            {
+                                new_related_messages.insert(m);
+                            }
+                        }
+                    }
+                    MessageReference::Matrix(room_id, event_id) => {
+                        if let Some(matrix) = self.matrix.lock().unwrap().borrow().as_ref() {
+                            let room_id = matrix::RoomId::from_str(room_id).unwrap();
+                            let event_id = matrix::EventId::from_str(event_id).unwrap();
+
+                            if let Some(m) =
+                                edit_on_matrix(matrix, room_id, event_id, &message).await
+                            {
+                                new_related_messages.insert(m);
+                            }
+                        }
+                    }
+                }
+            }
+
+            self.store_related_messages(&new_related_messages.into_iter().collect::<Vec<_>>())
+        }
     }
 }

src/discord.rs

@@ -5,11 +5,11 @@ use tokio::sync::mpsc;
 use crate::{
     channels::ChannelReference,
     message_ast::{self, format_discord},
-    messages::{MessageAuthor, MessageEvent, MessageReference, SentMessage},
+    messages::{EditedMessage, MessageAuthor, MessageEvent, MessageReference, SentMessage},
 };
 
 pub use serenity::client::Context;
-pub use serenity::model::id::ChannelId;
+pub use serenity::model::id::{ChannelId, MessageId};
 
 impl From<&Message> for MessageReference {
     fn from(message: &Message) -> Self {
@@ -72,13 +72,45 @@ impl EventHandler for DiscordHandler {
             content,
             author: MessageAuthor {
                 display_name: message
-                    .author_nick(&ctx.http)
+                    .author_nick(&ctx)
                     .await
                     .unwrap_or(message.author.name),
             },
             replies_to,
         }));
     }
+
+    async fn message_update(
+        &self,
+        ctx: Context,
+        _old_if_available: Option<Message>,
+        new: Option<Message>,
+        event: MessageUpdateEvent,
+    ) {
+        if let Ok(new_message) = {
+            if let Some(m) = new {
+                Ok(m)
+            } else {
+                event.channel_id.message(&ctx, event.id).await
+            }
+        } {
+            let message_ref = MessageReference::from(&new_message);
+
+            let content = discord_message_format::parse(&new_message.content);
+            let content = message_ast::convert_discord(&content);
+
+            let _ = self.event_tx.send(MessageEvent::Edit(EditedMessage {
+                replacing: message_ref,
+                content,
+                author: MessageAuthor {
+                    display_name: new_message
+                        .author_nick(&ctx)
+                        .await
+                        .unwrap_or(new_message.author.name),
+                },
+            }));
+        }
+    }
 }
 
 pub async fn forward_to_discord(
@@ -88,7 +120,7 @@ pub async fn forward_to_discord(
     reply: Option<(u64, u64)>,
 ) -> Option<MessageReference> {
     channel
-        .send_message(&discord_ctx.http, |m| {
+        .send_message(&discord_ctx, |m| {
             let b = m.content(format_discord(&message.content));
 
             if let Some((channel_id, message_id)) = reply {
@@ -103,6 +135,22 @@ pub async fn forward_to_discord(
         .map(MessageReference::from)
 }
 
+pub async fn edit_on_discord(
+    discord_ctx: &Context,
+    channel_id: ChannelId,
+    message_id: MessageId,
+    message: &EditedMessage,
+) -> Option<MessageReference> {
+    channel_id
+        .edit_message(&discord_ctx, &message_id, |m| {
+            m.content(format_discord(&message.content))
+        })
+        .await
+        .as_ref()
+        .ok()
+        .map(MessageReference::from)
+}
+
 pub async fn create_discord_client(
     ctx_tx: mpsc::UnboundedSender<Context>,
     message_tx: mpsc::UnboundedSender<MessageEvent>,

src/main.rs

@@ -88,10 +88,10 @@ async fn main() {
     while let Some(event) = event_rx.recv().await {
         match event {
             MessageEvent::Send(sent_message) => {
-                let _ = bridgers.send_message(sent_message).await;
+                bridgers.send_message(sent_message).await;
             }
-            MessageEvent::Edit(_edited_message) => {
+            MessageEvent::Edit(edited_message) => {
-                todo!();
+                bridgers.edit_message(edited_message).await;
             }
             MessageEvent::AdminLinkChannels(channels) => {
                 bridgers.link_channels(&channels);

src/matrix.rs

@@ -1,23 +1,29 @@
-use std::{str::FromStr, sync::Arc};
+use std::sync::Arc;
 
 use matrix_sdk::{
     room::{Joined, Room},
     ruma::{
+        api,
         events::{
+            self,
             room::{
                 message::{
                     FormattedBody, MessageEventContent, MessageFormat, MessageType, Relation,
+                    Replacement,
                 },
                 redaction::RedactionEventContent,
             },
             AnyMessageEvent, AnyMessageEventContent, AnyRoomEvent, AnySyncRoomEvent,
             SyncMessageEvent,
         },
-        EventId, UserId,
+        UserId,
     },
     ClientConfig, SyncSettings,
 };
-pub use matrix_sdk::{ruma::RoomId, Client};
+pub use matrix_sdk::{
+    ruma::{EventId, RoomId},
+    Client,
+};
 
 use log::info;
 use tokio::sync::mpsc;
@@ -28,7 +34,7 @@ use crate::{
         convert_matrix, convert_plain, format_discord, format_matrix, MessageComponent,
         MessageContent,
     },
-    messages::{MessageAuthor, MessageEvent, MessageReference, SentMessage},
+    messages::{EditedMessage, MessageAuthor, MessageEvent, MessageReference, SentMessage},
 };
 
 impl From<(&RoomId, &EventId)> for MessageReference {
@@ -102,6 +108,28 @@ impl MatrixHandler {
     }
 }
 
+async fn get_room_message_event(
+    room: &Joined,
+    event_id: &EventId,
+) -> Option<events::MessageEvent<MessageEventContent>> {
+    let event = room
+        .event(api::client::r0::room::get_room_event::Request::new(
+            room.room_id(),
+            event_id,
+        ))
+        .await
+        .ok()?
+        .event
+        .deserialize()
+        .ok()?;
+
+    if let AnyRoomEvent::Message(AnyMessageEvent::RoomMessage(message_event)) = event {
+        Some(message_event)
+    } else {
+        None
+    }
+}
+
 async fn on_room_message_event(
     ctx: Arc<MatrixHandler>,
     event: SyncMessageEvent<MessageEventContent>,
@@ -112,20 +140,26 @@ async fn on_room_message_event(
     }
 
     if let Room::Joined(room) = room {
+        if let Some(Relation::Replacement(replacement)) = &event.content.relates_to {
+            on_message_edited(ctx, &event, room, replacement).await;
+        } else {
+            on_message_sent(ctx, &event, room).await;
+        }
+    }
+}
+
+async fn on_message_sent(
+    ctx: Arc<MatrixHandler>,
+    event: &SyncMessageEvent<MessageEventContent>,
+    room: Joined,
+) {
     let message_ref = MessageReference::from((room.room_id(), &event.event_id));
 
-        let message_type =
+    if let Some(author) = ctx.get_message_author(&room, &event.sender).await {
-            if let Some(Relation::Replacement(replacement)) = &event.content.relates_to {
+        let message_event = match &event.content.msgtype {
-                &replacement.new_content.msgtype
-            } else {
-                &event.content.msgtype
-            };
-
-        match message_type {
             MessageType::Text(text) => {
                 let content = ctx.get_content(&text.body, &text.formatted);
 
-                if let Some(author) = ctx.get_message_author(&room, &event.sender).await {
                 let replies_to =
                     if let Some(Relation::Reply { in_reply_to }) = &event.content.relates_to {
                         Some(MessageReference::from((
@@ -136,65 +170,68 @@ async fn on_room_message_event(
                         None
                     };
 
-                    let _ = ctx.message_tx.send(MessageEvent::Send(SentMessage {
+                Some(MessageEvent::Send(SentMessage {
                     source: message_ref,
                     content,
                     author,
                     replies_to,
-                    }));
+                }))
-                }
             }
 
             MessageType::Emote(emote) => {
                 let mut content = ctx.get_content(&emote.body, &emote.formatted);
                 content.insert(0, MessageComponent::Plain("* ".to_string()));
 
-                if let Some(author) = ctx.get_message_author(&room, &event.sender).await {
+                Some(MessageEvent::Send(SentMessage {
-                    let _ = ctx.message_tx.send(MessageEvent::Send(SentMessage {
                     source: message_ref,
                     content,
                     author,
                     replies_to: None,
-                    }));
+                }))
-                }
             }
 
             // TODO: Handle reactions, uploads (audio, video, image, file), and any other types of event
-            _ => {}
+            _ => None,
         };
 
+        if let Some(e) = message_event {
+            let _ = ctx.message_tx.send(e);
+        }
+    }
+
     let _ = room.read_receipt(&event.event_id).await;
 }
+
+async fn on_message_edited(
+    ctx: Arc<MatrixHandler>,
+    event: &SyncMessageEvent<MessageEventContent>,
+    room: Joined,
+    replacement: &Replacement,
+) {
+    let message_ref = MessageReference::from((room.room_id(), &replacement.event_id));
+
+    if let MessageType::Text(text) = &replacement.new_content.msgtype {
+        let content = ctx.get_content(&text.body, &text.formatted);
+
+        if let Some(author) = ctx.get_message_author(&room, &event.sender).await {
+            let _ = ctx.message_tx.send(MessageEvent::Edit(EditedMessage {
+                replacing: message_ref,
+                content,
+                author,
+            }));
+        }
+    }
 }
 
 pub async fn forward_to_matrix(
     client: &Client,
     room_id: RoomId,
     message: &SentMessage,
-    reply: Option<(String, String)>,
+    replying_to: Option<EventId>,
 ) -> Option<MessageReference> {
     if let Some(room) = client.get_joined_room(&room_id) {
-        let replied_message_event = if let Some((room_id, event_id)) = reply {
+        let replied_message_event = if let Some(event_id) = replying_to {
-            let event = room
+            get_room_message_event(&room, &event_id).await
-                .event(
-                    matrix_sdk::ruma::api::client::r0::room::get_room_event::Request::new(
-                        &RoomId::from_str(&room_id).unwrap(),
-                        &EventId::from_str(&event_id).unwrap(),
-                    ),
-                )
-                .await
-                .unwrap()
-                .event
-                .deserialize()
-                .unwrap();
-
-            if let AnyRoomEvent::Message(AnyMessageEvent::RoomMessage(original_message_event)) =
-                event
-            {
-                Some(original_message_event)
-            } else {
-                None
-            }
         } else {
             None
         };
@@ -223,6 +260,64 @@ pub async fn forward_to_matrix(
     None
 }
 
+pub async fn edit_on_matrix(
+    client: &Client,
+    room_id: RoomId,
+    event_id: EventId,
+    message: &EditedMessage,
+) -> Option<MessageReference> {
+    if let Some(room) = client.get_joined_room(&room_id) {
+        if let Some(original_message_event) = get_room_message_event(&room, &event_id).await {
+            if let Some(Relation::Replacement(_)) = &original_message_event.content.relates_to {
+                return None;
+            }
+
+            let replied_message_event = if let Some(Relation::Reply { in_reply_to }) =
+                &original_message_event.content.relates_to
+            {
+                get_room_message_event(&room, &in_reply_to.event_id).await
+            } else {
+                None
+            };
+
+            let plain_reply = format_discord(&message.content);
+            let html_reply = format_matrix(&message.content);
+
+            let mut edit_content = if let Some(replied_message_event) = &replied_message_event {
+                MessageEventContent::text_reply_html(
+                    format!("* {}", &plain_reply),
+                    format!("* {}", html_reply),
+                    replied_message_event,
+                )
+            } else {
+                MessageEventContent::text_html(
+                    format!("* {}", &plain_reply),
+                    format!("* {}", html_reply),
+                )
+            };
+
+            let basic_content = if let Some(replied_message_event) = &replied_message_event {
+                MessageEventContent::text_reply_html(plain_reply, html_reply, replied_message_event)
+            } else {
+                MessageEventContent::text_html(plain_reply, html_reply)
+            };
+
+            edit_content.relates_to = Some(Relation::Replacement(Replacement::new(
+                event_id,
+                Box::new(basic_content),
+            )));
+
+            let new_event = room
+                .send(AnyMessageEventContent::RoomMessage(edit_content), None)
+                .await
+                .ok()?;
+
+            return Some(MessageReference::from((&room_id, &new_event.event_id)));
+        }
+    }
+    None
+}
+
 pub async fn create_matrix_client(
     homeserver_url: String,
     username: String,

src/messages.rs

@@ -2,7 +2,7 @@ use serde::{Deserialize, Serialize};
 
 use crate::{channels::ChannelReference, message_ast::MessageContent};
 
-#[derive(Serialize, Deserialize, Clone, PartialEq, Eq)]
+#[derive(Serialize, Deserialize, Clone, PartialEq, Eq, Hash)]
 pub enum MessageReference {
     Discord(u64, u64),
     Matrix(String, String),