xmtp_mls/messages/
enrichment.rs

1use crate::messages::decoded_message::{DecodedMessage, DeletedBy, MessageBody};
2use prost::Message;
3use std::collections::HashMap;
4use thiserror::Error;
5use xmtp_common::{ErrorCode, RetryableError};
6use xmtp_db::DbQuery;
7use xmtp_db::group_message::{
8    ContentType as DbContentType, Deletable, RelationCounts, RelationQuery, StoredGroupMessage,
9};
10use xmtp_db::message_deletion::StoredMessageDeletion;
11use xmtp_proto::xmtp::mls::message_contents::{ContentTypeId, EncodedContent};
12
13use xmtp_proto::types::GroupId;
14/// Content type ID for deleted message placeholders shown in enriched message lists
15pub fn deleted_message_content_type() -> ContentTypeId {
16    ContentTypeId {
17        authority_id: "xmtp.org".to_string(),
18        type_id: "deletedMessage".to_string(),
19        version_major: 1,
20        version_minor: 0,
21    }
22}
23
24#[derive(Debug, Error, ErrorCode)]
25pub enum EnrichMessageError {
26    #[error("Storage error: {0}")]
27    #[error_code(inherit)]
28    Storage(#[from] xmtp_db::StorageError),
29    #[error("DB error: {0}")]
30    #[error_code(inherit)]
31    DbConnection(#[from] xmtp_db::ConnectionError),
32}
33
34impl From<xmtp_db::diesel::result::Error> for EnrichMessageError {
35    fn from(error: xmtp_db::diesel::result::Error) -> Self {
36        Self::Storage(error.into())
37    }
38}
39
40impl RetryableError for EnrichMessageError {
41    fn is_retryable(&self) -> bool {
42        match self {
43            Self::DbConnection(e) => e.is_retryable(),
44            Self::Storage(e) => e.is_retryable(),
45        }
46    }
47}
48
49// Mapping of reactions, keyed by the ID of the message being reacted to.
50type ReactionMap = HashMap<Vec<u8>, Vec<DecodedMessage>>;
51// Mapping of referenced messages, keyed by ID (stores both stored and decoded)
52type ReferencedMessageMap = HashMap<Vec<u8>, (StoredGroupMessage, DecodedMessage)>;
53// Mapping of deletions, keyed by the ID of the deleted message
54type DeletionMap = HashMap<Vec<u8>, Vec<StoredMessageDeletion>>;
55
56pub struct EnrichedStoredMessage {
57    pub delivery_cursor: Option<xmtp_db::delivery::DeliveryCursor>,
58    pub stored: StoredGroupMessage,
59    pub decoded: DecodedMessage,
60    pub parent_stored: Option<StoredGroupMessage>,
61}
62
63/// Check the stored wire type before a deletion uses a cached content type.
64// implements: CTYPE-018
65pub(crate) fn is_deletable_stored_message(message: &StoredGroupMessage) -> bool {
66    if !message.kind.is_deletable() {
67        return false;
68    }
69    let Ok(content) = EncodedContent::decode(message.decrypted_message_bytes.as_slice()) else {
70        return false;
71    };
72    let Some(identifier) = content.r#type else {
73        return false;
74    };
75    DbContentType::from_identifier(
76        &identifier.authority_id,
77        &identifier.type_id,
78        identifier.version_major,
79    )
80    .is_deletable()
81}
82
83/// Validates if a deletion should be applied. Checks group membership and authorization.
84// implements: PROC-037
85pub(crate) fn is_deletion_valid(
86    deletion: &StoredMessageDeletion,
87    message: &StoredGroupMessage,
88    group_id: &GroupId,
89) -> bool {
90    if deletion.deleted_message_id != message.id {
91        return false;
92    }
93
94    if deletion.group_id != *group_id || message.group_id != *group_id {
95        return false;
96    }
97
98    if !is_deletable_stored_message(message) {
99        return false;
100    }
101
102    let is_sender = deletion.deleted_by_inbox_id == message.sender_inbox_id;
103    is_sender || deletion.is_super_admin_deletion
104}
105
106#[xmtp_common::mls_span]
107pub fn enrich_messages(
108    conn: impl DbQuery,
109    group_id: &GroupId,
110    messages: Vec<StoredGroupMessage>,
111) -> Result<Vec<DecodedMessage>, EnrichMessageError> {
112    Ok(enrich_messages_with_stored(conn, group_id, messages)?
113        .into_iter()
114        .map(|message| message.decoded)
115        .collect())
116}
117
118/// Replace a message with its deletion placeholder. Every trace of the
119/// original content goes: body, raw bytes, nested evidence, type, and fallback.
120// implements: PROC-037
121fn apply_deletion(
122    message: &mut DecodedMessage,
123    deletion: &StoredMessageDeletion,
124    stored: &StoredGroupMessage,
125) {
126    message.content = MessageBody::DeletedMessage {
127        deleted_by: DeletedBy::new(&deletion.deleted_by_inbox_id, &stored.sender_inbox_id),
128    };
129    message.metadata.content_type = Some(deleted_message_content_type());
130    message.fallback_text = None;
131    message.reactions = Vec::new();
132    message.num_replies = 0;
133}
134
135/// Enrich every stored row. Content that fails to decode stays in the list
136/// as an undecodable body; only a database failure is an error.
137// implements: CTYPE-008
138pub fn enrich_messages_with_stored(
139    conn: impl DbQuery,
140    group_id: &GroupId,
141    messages: Vec<StoredGroupMessage>,
142) -> Result<Vec<EnrichedStoredMessage>, EnrichMessageError> {
143    let initial_message_ids: Vec<&[u8]> = messages.iter().map(|m| m.id.as_ref()).collect();
144
145    let reference_ids: Vec<&[u8]> = messages
146        .iter()
147        .filter_map(|m| m.reference_id.as_deref())
148        .collect();
149
150    let mut relations = get_relations(conn, group_id, &initial_message_ids, &reference_ids)?;
151
152    let messages: Vec<EnrichedStoredMessage> = messages
153        .into_iter()
154        .map(|stored_message| {
155            let mut decoded = DecodedMessage::from(stored_message.clone());
156            let mut parent_stored = None;
157
158            let valid_deletion =
159                relations
160                    .deletions
161                    .get(&decoded.metadata.id)
162                    .and_then(|deletions| {
163                        deletions
164                            .iter()
165                            .find(|deletion| is_deletion_valid(deletion, &stored_message, group_id))
166                    });
167
168            if let Some(deletion) = valid_deletion {
169                apply_deletion(&mut decoded, deletion, &stored_message);
170            } else {
171                decoded.reactions = relations
172                    .reactions
173                    .remove(&decoded.metadata.id)
174                    .unwrap_or_default();
175
176                decoded.num_replies = relations
177                    .reply_counts
178                    .get(&decoded.metadata.id)
179                    .cloned()
180                    .unwrap_or(0);
181
182                // Handle Reply messages - populate in_reply_to field
183                if let MessageBody::Reply(mut reply_body) = decoded.content {
184                    // The decoder accepted only a hex reference.
185                    if let Ok(id) = hex::decode(&reply_body.reference_id) {
186                        parent_stored = relations
187                            .referenced_messages
188                            .get(&id)
189                            .map(|(stored, _)| stored.clone());
190                        let mut in_reply_to = relations
191                            .referenced_messages
192                            .get(&id)
193                            .map(|(_, decoded)| decoded.clone());
194
195                        if let Some(msg) = in_reply_to.as_mut()
196                            && let Some(deletions) = relations.deletions.get(&id)
197                            && let Some((stored_msg, _)) = relations.referenced_messages.get(&id)
198                            && let Some(deletion) = deletions
199                                .iter()
200                                .find(|deletion| is_deletion_valid(deletion, stored_msg, group_id))
201                        {
202                            apply_deletion(msg, deletion, stored_msg);
203                        }
204                        reply_body.in_reply_to = in_reply_to.map(Box::new);
205                    }
206                    decoded.content = MessageBody::Reply(reply_body);
207                }
208            }
209
210            EnrichedStoredMessage {
211                delivery_cursor: None,
212                stored: stored_message,
213                decoded,
214                parent_stored,
215            }
216        })
217        .collect();
218
219    Ok(messages)
220}
221
222fn get_relations(
223    conn: impl DbQuery,
224    group_id: &GroupId,
225    message_ids: &[&[u8]],
226    reference_ids: &[&[u8]],
227) -> Result<GetRelationsResults, EnrichMessageError> {
228    if message_ids.is_empty() {
229        return Ok(GetRelationsResults {
230            reactions: HashMap::new(),
231            referenced_messages: HashMap::new(),
232            reply_counts: HashMap::new(),
233            deletions: HashMap::new(),
234        });
235    }
236
237    let reactions_relations_query = RelationQuery::builder()
238        .content_types(Some(vec![DbContentType::Reaction]))
239        .build()
240        .unwrap_or_default();
241
242    let replies_count_query = RelationQuery::builder()
243        .content_types(Some(vec![DbContentType::Reply]))
244        .build()
245        .unwrap_or_default();
246
247    let reactions = conn.get_inbound_relations(group_id, message_ids, reactions_relations_query)?;
248    let referenced_messages = conn.get_outbound_relations(group_id, reference_ids)?;
249    let reply_counts =
250        conn.get_inbound_relation_counts(group_id, message_ids, replies_count_query)?;
251
252    // Get deletions for all messages AND referenced messages in a single batch query.
253    // This ensures that if a reply references a deleted message, we can properly show
254    // the deletion state in the reply chain.
255    let mut all_ids: Vec<Vec<u8>> = message_ids.iter().map(|id| id.to_vec()).collect();
256    all_ids.extend(reference_ids.iter().map(|id| id.to_vec()));
257    let deletions = conn.get_deletions_for_messages(all_ids)?;
258
259    Ok(GetRelationsResults {
260        reactions: get_reactions(reactions),
261        referenced_messages: get_referenced_messages(referenced_messages),
262        reply_counts,
263        deletions: get_deletions(deletions),
264    })
265}
266
267struct GetRelationsResults {
268    reactions: ReactionMap,
269    referenced_messages: ReferencedMessageMap,
270    reply_counts: RelationCounts,
271    deletions: DeletionMap,
272}
273
274fn get_referenced_messages(messages: HashMap<Vec<u8>, StoredGroupMessage>) -> ReferencedMessageMap {
275    messages
276        .into_iter()
277        .map(|(id, stored_message)| {
278            let decoded = DecodedMessage::from(stored_message.clone());
279            (id, (stored_message, decoded))
280        })
281        .collect()
282}
283
284/// A parent's reaction summary lists the reactions that apply to it. A row
285/// with a reaction type but no decodable reaction stays in history as its
286/// own retained message; it cannot be applied here.
287fn get_reactions(messages: HashMap<Vec<u8>, Vec<StoredGroupMessage>>) -> ReactionMap {
288    messages
289        .into_iter()
290        .map(|(id, reaction_messages)| {
291            let mapped_reactions: Vec<DecodedMessage> = reaction_messages
292                .into_iter()
293                .map(DecodedMessage::from)
294                .filter(|decoded| matches!(decoded.content, MessageBody::Reaction(_)))
295                .collect();
296            (id, mapped_reactions)
297        })
298        .collect()
299}
300
301fn get_deletions(deletions: Vec<StoredMessageDeletion>) -> DeletionMap {
302    let mut by_message = DeletionMap::new();
303    for deletion in deletions {
304        by_message
305            .entry(deletion.deleted_message_id.clone())
306            .or_default()
307            .push(deletion);
308    }
309    for records in by_message.values_mut() {
310        records.sort_by(|a, b| a.deleted_at_ns.cmp(&b.deleted_at_ns).then(a.id.cmp(&b.id)));
311    }
312    by_message
313}
xmtp_mls/messages/enrichment.rs

xmtp_mls/messages/
enrichment.rs
