timeline: use the EncryptionInfo provided by the replacement event when processing edits

crates/matrix-sdk-ui/src/timeline/event_handler.rs

@@ -543,7 +543,21 @@ impl<'a, 'o> TimelineEventHandler<'a, 'o> {
             };
 
             trace!("Applying edit");
-            Some(event_item.with_content(new_content, edit_json))
+
+            if let EventTimelineItemKind::Remote(remote_event) = &event_item.kind {
+                let new_encryption_info = match &this.ctx.flow {
+                    Flow::Local { .. } => None,
+                    Flow::Remote { encryption_info, .. } => encryption_info.clone(),
+                };
+
+                Some(event_item.with_content(new_content, edit_json).with_kind(
+                    EventTimelineItemKind::Remote(
+                        remote_event.with_encryption_info(new_encryption_info),
+                    ),
+                ))
+            } else {
+                Some(event_item.with_content(new_content, edit_json))
+            }
         });
 
         if !found {

crates/matrix-sdk-ui/src/timeline/event_item/remote.rs

@@ -73,6 +73,11 @@ pub(in crate::timeline) struct RemoteEventTimelineItem {
 }
 
 impl RemoteEventTimelineItem {
+    /// Clone the current event item, and update its `encryption_info`.
+    pub fn with_encryption_info(&self, encryption_info: Option<EncryptionInfo>) -> Self {
+        Self { encryption_info, ..self.clone() }
+    }
+
     /// Clone the current event item, and update its `reactions`.
     pub fn with_reactions(&self, reactions: ReactionsByKeyBySender) -> Self {
         Self { reactions, ..self.clone() }

crates/matrix-sdk-ui/src/timeline/tests/edit.rs

@@ -12,10 +12,16 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::BTreeMap;
+
 use assert_matches2::assert_let;
 use eyeball_im::VectorDiff;
+use matrix_sdk::deserialized_responses::{
+    AlgorithmInfo, EncryptionInfo, VerificationLevel, VerificationState,
+};
 use matrix_sdk_test::{async_test, sync_timeline_event, ALICE};
 use ruma::{
+    event_id,
     events::room::message::{MessageType, RedactedRoomMessageEventContent},
     server_name, EventId,
 };
@@ -153,3 +159,66 @@ async fn test_aggregated_sanitized() {
     let day_divider = assert_next_matches!(stream, VectorDiff::PushFront { value } => value);
     assert!(day_divider.is_day_divider());
 }
+
+#[async_test]
+async fn test_edit_updates_encryption_info() {
+    let timeline = TestTimeline::new();
+    let event_factory = &timeline.factory;
+
+    let original_event_id = event_id!("$original_event");
+
+    let mut original_event = event_factory
+        .text_msg("**original** message")
+        .sender(*ALICE)
+        .event_id(original_event_id)
+        .into_sync();
+
+    let mut encryption_info = EncryptionInfo {
+        sender: (*ALICE).into(),
+        sender_device: None,
+        algorithm_info: AlgorithmInfo::MegolmV1AesSha2 {
+            curve25519_key: "123".to_owned(),
+            sender_claimed_keys: BTreeMap::new(),
+        },
+        verification_state: VerificationState::Verified,
+    };
+
+    original_event.encryption_info = Some(encryption_info.clone());
+
+    timeline.handle_live_event(original_event).await;
+
+    let items = timeline.inner.items().await;
+    let first_event = items[1].as_event().unwrap();
+
+    assert_eq!(
+        first_event.encryption_info().unwrap().verification_state,
+        VerificationState::Verified
+    );
+
+    assert_let!(TimelineItemContent::Message(message) = first_event.content());
+    assert_let!(MessageType::Text(text) = message.msgtype());
+    assert_eq!(text.body, "**original** message");
+
+    let mut edit_event = event_factory
+        .text_msg(" * !!edited!! **better** message")
+        .sender(*ALICE)
+        .edit(original_event_id, MessageType::text_plain("!!edited!! **better** message").into())
+        .into_sync();
+    encryption_info.verification_state =
+        VerificationState::Unverified(VerificationLevel::UnverifiedIdentity);
+    edit_event.encryption_info = Some(encryption_info);
+
+    timeline.handle_live_event(edit_event).await;
+
+    let items = timeline.inner.items().await;
+    let first_event = items[1].as_event().unwrap();
+
+    assert_eq!(
+        first_event.encryption_info().unwrap().verification_state,
+        VerificationState::Unverified(VerificationLevel::UnverifiedIdentity)
+    );
+
+    assert_let!(TimelineItemContent::Message(message) = first_event.content());
+    assert_let!(MessageType::Text(text) = message.msgtype());
+    assert_eq!(text.body, "!!edited!! **better** message");
+}