From aee8728418ca86bebec452491ea3446fe1307fe8 Mon Sep 17 00:00:00 2001
From: Doug <douglase@element.io>
Date: Thu, 1 Aug 2024 12:01:26 +0100
Subject: [PATCH] sdk-ui: Fix a bug where local echoes were showing unencrypted
 shields.

---
 .../src/timeline/event_item/mod.rs            |  2 +-
 .../src/timeline/tests/shields.rs             | 80 ++++++++++++++++++-
 2 files changed, 78 insertions(+), 4 deletions(-)

diff --git a/crates/matrix-sdk-ui/src/timeline/event_item/mod.rs b/crates/matrix-sdk-ui/src/timeline/event_item/mod.rs
index d23148928..1adb0b46f 100644
--- a/crates/matrix-sdk-ui/src/timeline/event_item/mod.rs
+++ b/crates/matrix-sdk-ui/src/timeline/event_item/mod.rs
@@ -370,7 +370,7 @@ impl EventTimelineItem {
     /// Gets the [`ShieldState`] which can be used to decorate messages in the
     /// recommended way.
     pub fn get_shield(&self, strict: bool) -> Option<ShieldState> {
-        if self.is_room_encrypted != Some(true) {
+        if self.is_room_encrypted != Some(true) || self.is_local_echo() {
             return None;
         }
 
diff --git a/crates/matrix-sdk-ui/src/timeline/tests/shields.rs b/crates/matrix-sdk-ui/src/timeline/tests/shields.rs
index e865729e0..3c0c8004d 100644
--- a/crates/matrix-sdk-ui/src/timeline/tests/shields.rs
+++ b/crates/matrix-sdk-ui/src/timeline/tests/shields.rs
@@ -1,10 +1,14 @@
+use assert_matches::assert_matches;
 use eyeball_im::VectorDiff;
 use matrix_sdk_base::deserialized_responses::ShieldState;
-use matrix_sdk_test::{async_test, ALICE};
-use ruma::events::room::message::RoomMessageEventContent;
+use matrix_sdk_test::{async_test, sync_timeline_event, ALICE};
+use ruma::{
+    event_id,
+    events::{room::message::RoomMessageEventContent, AnyMessageLikeEventContent},
+};
 use stream_assert::assert_next_matches;
 
-use crate::timeline::tests::TestTimeline;
+use crate::timeline::{tests::TestTimeline, EventSendState};
 
 #[async_test]
 async fn test_no_shield_in_unencrypted_room() {
@@ -39,3 +43,73 @@ async fn test_sent_in_clear_shield() {
     let shield = item.as_event().unwrap().get_shield(false);
     assert_eq!(shield, Some(ShieldState::Grey { message: "Sent in clear." }));
 }
+
+#[async_test]
+/// Note: This is a regression test for a bug where the SentInClear shield was
+/// incorrectly shown on a local echo whilst sending and was then removed once
+/// the local echo came back. There doesn't appear to be an obvious way to mock
+/// trust in the timeline, so this test uses a message sent in the clear, making
+/// sure the shield only appears once the remote echo is received.
+async fn test_local_sent_in_clear_shield() {
+    // Given an encrypted timeline.
+    let timeline = TestTimeline::with_is_room_encrypted(true);
+    let mut stream = timeline.subscribe().await;
+
+    // When sending an unencrypted event.
+    let event_id = event_id!("$W6mZSLWMmfuQQ9jhZWeTxFIM");
+    let txn_id = timeline
+        .handle_local_event(AnyMessageLikeEventContent::RoomMessage(
+            RoomMessageEventContent::text_plain("Local message"),
+        ))
+        .await;
+    let item = assert_next_matches!(stream, VectorDiff::PushBack { value } => value);
+    let event_item = item.as_event().unwrap();
+
+    // Then it's local echo should not have a shield (as encryption info isn't
+    // available).
+    assert!(event_item.is_local_echo());
+    let shield = event_item.get_shield(false);
+    assert_eq!(shield, None);
+
+    {
+        // The day divider comes in late.
+        let day_divider = assert_next_matches!(stream, VectorDiff::PushFront { value } => value);
+        assert!(day_divider.is_day_divider());
+    }
+
+    // When the event is sent (but without a remote echo).
+    timeline
+        .inner
+        .update_event_send_state(&txn_id, EventSendState::Sent { event_id: event_id.to_owned() })
+        .await;
+    let item = assert_next_matches!(stream, VectorDiff::Set { value, index: 1 } => value);
+    let event_item = item.as_event().unwrap();
+    let timestamp = event_item.timestamp();
+    assert_matches!(event_item.send_state(), Some(EventSendState::Sent { .. }));
+
+    // Then the local echo still should not have a shield.
+    assert!(event_item.is_local_echo());
+    let shield = event_item.get_shield(false);
+    assert_eq!(shield, None);
+
+    // When the remote echo comes in.
+    timeline
+        .handle_live_event(sync_timeline_event!({
+            "content": {
+                "body": "Local message",
+                "msgtype": "m.text",
+            },
+            "sender": &*ALICE,
+            "event_id": event_id,
+            "origin_server_ts": timestamp,
+            "type": "m.room.message",
+        }))
+        .await;
+    let item = assert_next_matches!(stream, VectorDiff::Set { index: 1, value } => value);
+    let event_item = item.as_event().unwrap();
+
+    // Then the remote echo should now be showing the shield.
+    assert!(!event_item.is_local_echo());
+    let shield = event_item.get_shield(false);
+    assert_eq!(shield, Some(ShieldState::Grey { message: "Sent in clear." }));
+}