fix(message)!: decrypt secret encrypted edits on receive

Cargo.toml

@@ -150,6 +150,7 @@ whatsapp-rust-ureq-http-client = { path = "./http_clients/ureq-client", version
 [dev-dependencies]
 aes = { workspace = true }
 cbc = { version = "0.2", features = ["alloc", "block-padding"] }
+flate2 = { workspace = true }
 hkdf = { workspace = true }
 hmac = { workspace = true }
 sha2 = { workspace = true }

src/history_sync.rs

@@ -1,7 +1,8 @@
 use crate::types::events::{Event, LazyHistorySync};
 use std::sync::Arc;
-use wacore::history_sync::{TcTokenCandidate, process_history_sync};
-use wacore::store::traits::TcTokenEntry;
+use wacore::history_sync::{HistoryMsgSecretRecord, TcTokenCandidate, process_history_sync};
+use wacore::store::traits::{MsgSecretEntry, TcTokenEntry};
+use wacore_binary::{Jid, JidExt as _};
 use waproto::whatsapp::message::HistorySyncNotification;
 
 use crate::client::Client;
@@ -145,6 +146,7 @@ impl Client {
         };
 
         let has_listeners = self.core.event_bus.has_handlers();
+        let retain_history_blob = has_listeners;
 
         // Small blobs (PushName, Recent): decode inline to avoid spawn_blocking overhead.
         // Large blobs: use blocking thread to avoid stalling the async runtime.
@@ -153,7 +155,7 @@ impl Client {
             Some(process_history_sync(
                 compressed_data,
                 own_user.as_deref(),
-                has_listeners,
+                retain_history_blob,
                 compressed_size_hint,
             ))
         } else {
@@ -162,7 +164,7 @@ impl Client {
                 let result = process_history_sync(
                     compressed_data,
                     own_user.as_deref(),
-                    has_listeners,
+                    retain_history_blob,
                     compressed_size_hint,
                 );
                 let _ = result_tx.send(result);
@@ -215,7 +217,9 @@ impl Client {
                     self.store_tc_token_candidate(candidate).await;
                 }
 
-                // Dispatch a single event with the full decompressed blob
+                self.store_history_sync_msg_secrets(sync_result.msg_secret_records)
+                    .await;
+
                 if let Some(decompressed) = sync_result.decompressed_bytes {
                     let lazy_hs = LazyHistorySync::new(
                         decompressed,
@@ -240,6 +244,77 @@ impl Client {
         }
     }
 
+    async fn store_history_sync_msg_secrets(&self, records: Vec<HistoryMsgSecretRecord>) -> usize {
+        const SECRET_LEN: usize = wacore::reporting_token::MESSAGE_SECRET_SIZE;
+
+        let device_snapshot = self.persistence_manager.get_device_snapshot().await;
+        let own_pn = device_snapshot.pn.as_ref().map(|j| j.to_non_ad());
+        let own_lid = device_snapshot.lid.as_ref().map(|j| j.to_non_ad());
+
+        let mut entries = Vec::new();
+        for record in records {
+            if record.secret.len() != SECRET_LEN {
+                continue;
+            }
+            let Ok(chat) = record.chat_id.parse::<Jid>() else {
+                continue;
+            };
+            let mut senders =
+                history_msg_secret_senders(&chat, &record, own_pn.as_ref(), own_lid.as_ref());
+            if chat.is_bot()
+                && let Some(lid) = own_lid.as_ref()
+            {
+                push_unique_sender(&mut senders, lid.to_non_ad());
+            }
+            if senders.is_empty() {
+                continue;
+            }
+
+            let sender_count = senders.len();
+            let mut chat_id = chat.to_non_ad_string();
+            let mut msg_id = record.msg_id;
+            let mut secret = record.secret;
+            for (idx, sender) in senders.into_iter().enumerate() {
+                let last_sender = idx + 1 == sender_count;
+                entries.push(MsgSecretEntry {
+                    chat: if last_sender {
+                        std::mem::take(&mut chat_id)
+                    } else {
+                        chat_id.clone()
+                    },
+                    sender: sender.to_non_ad_string(),
+                    msg_id: if last_sender {
+                        std::mem::take(&mut msg_id)
+                    } else {
+                        msg_id.clone()
+                    },
+                    secret: if last_sender {
+                        std::mem::take(&mut secret)
+                    } else {
+                        secret.clone()
+                    },
+                });
+            }
+        }
+
+        if entries.is_empty() {
+            return 0;
+        }
+
+        match self
+            .persistence_manager
+            .backend()
+            .put_msg_secrets(entries)
+            .await
+        {
+            Ok(stored) => stored,
+            Err(e) => {
+                log::warn!("failed to persist history-sync messageSecrets: {e:?}");
+                0
+            }
+        }
+    }
+
     /// Ask the phone to re-upload a history-sync blob whose download failed,
     /// by sending a `<receipt type="server-error" category="peer">` with the
     /// blob's `media_key`.
@@ -323,3 +398,184 @@ impl Client {
         }
     }
 }
+
+fn history_msg_secret_senders(
+    chat: &Jid,
+    record: &HistoryMsgSecretRecord,
+    own_pn: Option<&Jid>,
+    own_lid: Option<&Jid>,
+) -> Vec<Jid> {
+    let mut senders = Vec::with_capacity(2);
+
+    if record.from_me {
+        if let Some(lid) = own_lid {
+            push_unique_sender(&mut senders, lid.to_non_ad());
+        }
+        if let Some(pn) = own_pn {
+            push_unique_sender(&mut senders, pn.to_non_ad());
+        }
+        return senders;
+    }
+
+    if chat.is_pn() || chat.is_lid() || chat.is_bot() {
+        senders.push(chat.to_non_ad());
+        return senders;
+    }
+
+    if let Some(raw_sender) = record
+        .key_participant
+        .as_deref()
+        .or(record.web_msg_participant.as_deref())
+        && let Ok(sender) = raw_sender.parse::<Jid>()
+    {
+        senders.push(sender.to_non_ad());
+    }
+
+    senders
+}
+
+fn push_unique_sender(senders: &mut Vec<Jid>, sender: Jid) {
+    if !senders.contains(&sender) {
+        senders.push(sender);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use flate2::{Compression, write::ZlibEncoder};
+    use prost::Message as ProtoMessage;
+    use std::io::Write;
+    use std::sync::atomic::Ordering;
+    use waproto::whatsapp as wa;
+
+    fn compress_history_sync(history_sync: &wa::HistorySync) -> Vec<u8> {
+        let raw = history_sync.encode_to_vec();
+        let mut encoder = ZlibEncoder::new(Vec::new(), Compression::default());
+        encoder.write_all(&raw).expect("zlib write");
+        encoder.finish().expect("zlib finish")
+    }
+
+    #[tokio::test]
+    async fn process_history_sync_task_stores_message_secrets_without_handlers() {
+        let client = crate::test_utils::create_test_client_with_name("history_msg_secret").await;
+        client
+            .persistence_manager
+            .process_command(wacore::store::commands::DeviceCommand::SetId(Some(
+                "5511000000001:0@s.whatsapp.net".parse().unwrap(),
+            )))
+            .await;
+        client.is_running.store(true, Ordering::Relaxed);
+
+        let chat = "5511777776666@s.whatsapp.net";
+        let parent_id = "HIST_PARENT";
+        let secret = vec![0x44u8; 32];
+        let history_sync = wa::HistorySync {
+            sync_type: wa::history_sync::HistorySyncType::InitialBootstrap as i32,
+            conversations: vec![wa::Conversation {
+                id: chat.to_string(),
+                messages: vec![wa::HistorySyncMsg {
+                    message: Some(wa::WebMessageInfo {
+                        key: wa::MessageKey {
+                            remote_jid: Some(chat.to_string()),
+                            from_me: Some(false),
+                            id: Some(parent_id.to_string()),
+                            participant: None,
+                        },
+                        message: Some(wa::Message {
+                            conversation: Some("historical".to_string()),
+                            ..Default::default()
+                        }),
+                        message_secret: Some(secret.clone()),
+                        ..Default::default()
+                    }),
+                    msg_order_id: Some(1),
+                }],
+                ..Default::default()
+            }],
+            ..Default::default()
+        };
+        let compressed = compress_history_sync(&history_sync);
+        let notification = HistorySyncNotification {
+            file_length: Some(compressed.len() as u64),
+            sync_type: Some(wa::message::HistorySyncType::InitialBootstrap as i32),
+            initial_hist_bootstrap_inline_payload: Some(compressed),
+            ..Default::default()
+        };
+
+        client
+            .process_history_sync_task("HIST_SYNC_SECRET".to_string(), notification)
+            .await;
+
+        let got = client
+            .persistence_manager
+            .backend()
+            .get_msg_secret(chat, chat, parent_id)
+            .await
+            .unwrap();
+        assert_eq!(got, Some(secret));
+    }
+
+    #[tokio::test]
+    async fn process_history_sync_task_stores_bot_dm_secret_alias() {
+        let client =
+            crate::test_utils::create_test_client_with_name("history_bot_msg_secret").await;
+        client
+            .persistence_manager
+            .process_command(wacore::store::commands::DeviceCommand::SetLid(Some(
+                "999888777666555:0@lid".parse().unwrap(),
+            )))
+            .await;
+        client.is_running.store(true, Ordering::Relaxed);
+
+        let chat = "867051314767696@bot";
+        let parent_id = "HIST_BOT_PARENT";
+        let secret = vec![0x61u8; 32];
+        let history_sync = wa::HistorySync {
+            sync_type: wa::history_sync::HistorySyncType::InitialBootstrap as i32,
+            conversations: vec![wa::Conversation {
+                id: chat.to_string(),
+                messages: vec![wa::HistorySyncMsg {
+                    message: Some(wa::WebMessageInfo {
+                        key: wa::MessageKey {
+                            remote_jid: Some(chat.to_string()),
+                            from_me: Some(false),
+                            id: Some(parent_id.to_string()),
+                            participant: None,
+                        },
+                        message: Some(wa::Message {
+                            conversation: Some("bot historical".to_string()),
+                            ..Default::default()
+                        }),
+                        message_secret: Some(secret.clone()),
+                        ..Default::default()
+                    }),
+                    msg_order_id: Some(1),
+                }],
+                ..Default::default()
+            }],
+            ..Default::default()
+        };
+        let compressed = compress_history_sync(&history_sync);
+        let notification = HistorySyncNotification {
+            file_length: Some(compressed.len() as u64),
+            sync_type: Some(wa::message::HistorySyncType::InitialBootstrap as i32),
+            initial_hist_bootstrap_inline_payload: Some(compressed),
+            ..Default::default()
+        };
+
+        client
+            .process_history_sync_task("HIST_SYNC_BOT_SECRET".to_string(), notification)
+            .await;
+
+        let backend = client.persistence_manager.backend();
+        let primary = backend.get_msg_secret(chat, chat, parent_id).await.unwrap();
+        let alias = backend
+            .get_msg_secret(chat, "999888777666555@lid", parent_id)
+            .await
+            .unwrap();
+
+        assert_eq!(primary, Some(secret.clone()));
+        assert_eq!(alias, Some(secret));
+    }
+}