diff --git a/Cargo.lock b/Cargo.lock index 78e7811ef..1efe8f9f2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2396,6 +2396,7 @@ dependencies = [ "chrono", "env_logger", "event-listener", + "flate2", "futures", "hex", "hkdf", diff --git a/Cargo.toml b/Cargo.toml index d5d931d77..330614926 100644 --- a/Cargo.toml +++ b/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 } diff --git a/src/history_sync.rs b/src/history_sync.rs index 693c6b494..767ad1059 100644 --- a/src/history_sync.rs +++ b/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) -> 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::() 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 `` 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 { + 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::() + { + senders.push(sender.to_non_ad()); + } + + senders +} + +fn push_unique_sender(senders: &mut Vec, 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 { + 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)); + } +} diff --git a/src/message.rs b/src/message.rs index b5691cf6d..df669d1cb 100644 --- a/src/message.rs +++ b/src/message.rs @@ -99,16 +99,18 @@ impl Client { msg.get_base_message().get_ephemeral_expiration(); } - // Awaited (not spawned) so the messageSecret is durably stored before - // this chat's worker dequeues the next stanza. A bot reply queued right - // behind its own fanout (offline replay) would otherwise race the write - // and hit MissingMessageSecret. + // Keep this ordered with dispatch; add-on messages can immediately + // reference the secret from the stanza just processed. self.maybe_capture_inbound_msg_secret(&msg, &info).await; + let dispatch_msg = self + .maybe_decrypt_secret_encrypted_message(&msg, &info) + .await + .unwrap_or(msg); self.ack_received_message(&info); self.core .event_bus - .dispatch(Event::Message(Arc::new(msg), info)); + .dispatch(Event::Message(Arc::new(dispatch_msg), info)); } /// Acknowledge a received message so the server drops it from the offline @@ -145,66 +147,278 @@ impl Client { }); } - /// Capture an embedded `MessageContextInfo.message_secret` from any - /// bot-targeted message (fanout from us OR reply from the bot) so a - /// future `` referencing this id can decrypt. - /// Mirrors WA Web `processRenderableMessages`: - /// `$ && (P || N || w || A) && !isForwarded → addMsmsgMsgSecretToCache`. + /// Capture embedded `MessageContextInfo.message_secret` for add-on + /// decrypts. Bot DMs keep the legacy LID key as a second entry. pub(crate) async fn maybe_capture_inbound_msg_secret( self: &Arc, msg: &wa::Message, info: &Arc, ) { use wacore::proto_helpers::MessageExt; - const SECRET_LEN: usize = wacore::reporting_token::MESSAGE_SECRET_SIZE; let mci = msg.message_context_info.as_ref(); - let chat_is_bot = info.source.chat.server == wacore_binary::Server::Bot; - let mentions_bot = msg.mentions_any_bot(); - // `MessageContextInfo.bot_metadata` is the bot-invocation envelope WA - // Web reads; it's present on bot prompts (incl. our own group prompt) - // even when no JID is mentioned, covering WA Web's `w`/`A` gates. - let has_bot_metadata = mci.is_some_and(|m| m.bot_metadata.is_some()); - let Some(secret_bytes) = mci.and_then(|m| m.message_secret.as_deref()) else { return; }; - let Ok(secret_arr) = <&[u8; SECRET_LEN]>::try_from(secret_bytes) else { + if msg.is_forwarded() { return; + } + + self.persist_msg_secret_bytes( + &info.source.chat, + &info.source.sender, + &info.id, + secret_bytes, + ) + .await; + + let chat_is_bot = info.source.chat.server == wacore_binary::Server::Bot; + if chat_is_bot + && let Some(sender) = self.dm_sender_identity_for(&info.source.chat).await + && sender.to_non_ad() != info.source.sender.to_non_ad() + { + self.persist_msg_secret_bytes(&info.source.chat, &sender, &info.id, secret_bytes) + .await; + } + } + + pub(crate) async fn persist_msg_secret_bytes( + &self, + chat: &Jid, + sender: &Jid, + msg_id: &str, + secret_bytes: &[u8], + ) -> bool { + const SECRET_LEN: usize = wacore::reporting_token::MESSAGE_SECRET_SIZE; + + let Ok(secret) = <&[u8; SECRET_LEN]>::try_from(secret_bytes) else { + return false; }; - // WA Web `processRenderableMessages`: `$ && (P || N || w || A) && !fwd`. - // P=chat is bot, N=mentions a bot, w/A=bot group participant (we proxy - // those via bot_metadata presence, the bot-invocation envelope). - if !chat_is_bot && !mentions_bot && !has_bot_metadata { - return; + let chat_str = chat.to_non_ad_string(); + let sender_str = sender.to_non_ad_string(); + match self + .persistence_manager + .backend() + .put_msg_secret(&chat_str, &sender_str, msg_id, secret) + .await + { + Ok(()) => true, + Err(e) => { + log::warn!("[msg:{msg_id}] failed to persist messageSecret: {e:?}"); + false + } } - if msg.is_forwarded() { - return; + } + + async fn own_jid_for_secret_encrypted(&self, info: &MessageInfo) -> Option { + use wacore::types::message::AddressingMode; + + if info.source.is_from_me { + return Some(info.source.sender.to_non_ad()); } - // Key the secret under the identity the bot reply will echo in - // `` at GET time: - // * bot DM: `info.source.sender` is our PN device JID, but the reply - // echoes our LID — resolve via dm_sender_identity_for. - // * group / regular: `info.source.sender` is already the author's - // addressing identity (our LID in a LID group, the other - // participant's JID for their prompt) — exactly what the reply - // echoes. Use it directly. - // alternate_msg_secret_lookup still bridges any residual LID↔PN skew. - let sender = if chat_is_bot { - match self.dm_sender_identity_for(&info.source.chat).await { - Some(j) => j, - None => return, + + match info.source.addressing_mode { + Some(AddressingMode::Lid) => match self.get_lid().await { + Some(jid) => Some(jid), + None => self.get_pn().await, + }, + Some(AddressingMode::Pn) => match self.get_pn().await { + Some(jid) => Some(jid), + None => self.get_lid().await, + }, + None if info.source.sender.is_lid() || info.source.chat.is_lid() => { + match self.get_lid().await { + Some(jid) => Some(jid), + None => self.get_pn().await, + } } - } else { - info.source.sender.clone() + None => match self.get_pn().await { + Some(jid) => Some(jid), + None => self.get_lid().await, + }, + } + } + + async fn maybe_decrypt_secret_encrypted_message( + self: &Arc, + msg: &wa::Message, + info: &Arc, + ) -> Option { + use crate::features::message_edit::{self, SecretEncKind}; + + let env = message_edit::extract_secret_encrypted(msg)?; + let target_id = env.target_id()?; + + let my_jid = self.own_jid_for_secret_encrypted(info).await?; + let original_sender = match env.original_sender_jid(&my_jid) { + Ok(jid) => jid, + Err(_) => return None, }; - log::debug!( - "[msg:{}] cached bot messageSecret under sender={}", - info.id, - sender.to_non_ad_string() - ); - self.persist_outbound_msg_secret(&info.source.chat, &sender, &info.id, secret_arr) + + let backend = self.persistence_manager.backend(); + let chat_for_lookup = info.source.chat.to_non_ad_string(); + let original_sender_str = original_sender.to_non_ad_string(); + let fallback_original_sender = self + .alternate_msg_secret_jid(&backend, &original_sender) + .await + .unwrap_or_default(); + + let primary = backend + .get_msg_secret(&chat_for_lookup, &original_sender_str, target_id) + .await; + let secret = match primary { + Ok(Some(secret)) => secret, + Ok(None) => match self + .alternate_msg_secret_lookup( + &backend, + &chat_for_lookup, + &original_sender, + target_id, + ) + .await + { + Ok(Some(secret)) => secret, + Ok(None) => return None, + Err(e) => { + log::warn!( + "[msg:{}] secret_encrypted_message alternate secret lookup failed: {e:?}", + info.id + ); + return None; + } + }, + Err(e) => { + log::warn!( + "[msg:{}] backend error reading secret_encrypted_message secret: {e:?}", + info.id + ); + return None; + } + }; + + let fallback_editor = match info.source.sender_alt.clone() { + Some(jid) => Some(jid), + None => self + .alternate_msg_secret_jid(&backend, &info.source.sender) + .await + .unwrap_or_default(), + }; + + let inner = match message_edit::decrypt_secret_encrypted( + env.enc_payload, + env.enc_iv, + &secret, + env.kind, + target_id, + &original_sender, + &info.source.sender, + ) { + Ok(inner) => inner, + Err(primary_err) => { + let mut last_err = primary_err; + let mut decrypted = None; + + if let Some(fallback_original) = fallback_original_sender.as_ref() { + match message_edit::decrypt_secret_encrypted( + env.enc_payload, + env.enc_iv, + &secret, + env.kind, + target_id, + fallback_original, + &info.source.sender, + ) { + Ok(inner) => decrypted = Some(inner), + Err(e) => last_err = e, + } + } + + if decrypted.is_none() + && let Some(fallback_editor) = fallback_editor.as_ref() + { + match message_edit::decrypt_secret_encrypted( + env.enc_payload, + env.enc_iv, + &secret, + env.kind, + target_id, + &original_sender, + fallback_editor, + ) { + Ok(inner) => decrypted = Some(inner), + Err(e) => last_err = e, + } + } + + if decrypted.is_none() + && let (Some(fallback_original), Some(fallback_editor)) = + (fallback_original_sender.as_ref(), fallback_editor.as_ref()) + { + match message_edit::decrypt_secret_encrypted( + env.enc_payload, + env.enc_iv, + &secret, + env.kind, + target_id, + fallback_original, + fallback_editor, + ) { + Ok(inner) => decrypted = Some(inner), + Err(e) => last_err = e, + } + } + + match decrypted { + Some(inner) => inner, + None => { + log::warn!( + "[msg:{}] secret_encrypted_message {:?} decrypt failed: {last_err:?}", + info.id, + env.kind + ); + return None; + } + } + } + }; + + if let Some(secret_bytes) = inner + .message_context_info + .as_ref() + .and_then(|m| m.message_secret.as_deref()) + { + self.persist_msg_secret_bytes( + &info.source.chat, + &original_sender, + target_id, + secret_bytes, + ) .await; + if let Some(alternate_sender) = fallback_original_sender.as_ref() { + self.persist_msg_secret_bytes( + &info.source.chat, + alternate_sender, + target_id, + secret_bytes, + ) + .await; + } + } + + if env.kind != SecretEncKind::MessageEdit { + return Some(inner); + } + + match message_edit::rewrap_as_legacy_edit(inner) { + Some(rewrapped) => Some(rewrapped), + None => { + log::warn!( + "[msg:{}] decrypted MESSAGE_EDIT missing protocol_message.edited_message", + info.id + ); + None + } + } } /// Decrypt and dispatch a `` bot reply. Looks up the @@ -444,31 +658,39 @@ impl Client { /// `lid_pn_mapping` store and retry. Returns `Ok(None)` when no mapping /// is known or the alternate row is absent — the caller treats that as /// a terminal miss. - async fn alternate_msg_secret_lookup( + async fn alternate_msg_secret_jid( &self, backend: &Arc, - chat_for_lookup: &str, primary_sender: &Jid, - target_id: &str, - ) -> Result>, crate::store::error::StoreError> { - let alternate_user = match primary_sender.server { + ) -> Result, crate::store::error::StoreError> { + let alternate = match primary_sender.server { wacore_binary::Server::Lid => backend .get_lid_mapping(&primary_sender.user) .await? - .map(|m| (m.phone_number, wacore_binary::Server::Pn)), + .map(|m| Jid::new(m.phone_number, wacore_binary::Server::Pn)), wacore_binary::Server::Pn => backend .get_pn_mapping(&primary_sender.user) .await? - .map(|m| (m.lid, wacore_binary::Server::Lid)), + .map(|m| Jid::new(m.lid, wacore_binary::Server::Lid)), _ => None, }; - let Some((user, server)) = alternate_user else { + Ok(alternate) + } + + async fn alternate_msg_secret_lookup( + &self, + backend: &Arc, + chat_for_lookup: &str, + primary_sender: &Jid, + target_id: &str, + ) -> Result>, crate::store::error::StoreError> { + let Some(alternate) = self + .alternate_msg_secret_jid(backend, primary_sender) + .await? + else { return Ok(None); }; - let mut alternate_str = String::with_capacity(user.len() + 1 + server.as_str().len()); - alternate_str.push_str(&user); - alternate_str.push('@'); - alternate_str.push_str(server.as_str()); + let alternate_str = alternate.to_non_ad_string(); backend .get_msg_secret(chat_for_lookup, &alternate_str, target_id) .await @@ -8259,6 +8481,412 @@ mod tests { None } + fn legacy_edit_text(msg: &wa::Message) -> Option<&str> { + msg.protocol_message + .as_ref() + .and_then(|pm| pm.edited_message.as_ref()) + .and_then(|edited| edited.conversation.as_deref()) + } + + fn inner_message_edit(text: &str, next_secret: Option>) -> wa::Message { + wa::Message { + protocol_message: Some(Box::new(wa::message::ProtocolMessage { + key: Some(wa::MessageKey { + remote_jid: Some("5511777776666@s.whatsapp.net".to_string()), + from_me: Some(false), + id: Some("PARENT_EDIT".to_string()), + participant: None, + }), + r#type: Some(wa::message::protocol_message::Type::MessageEdit as i32), + edited_message: Some(Box::new(wa::Message { + conversation: Some(text.to_string()), + ..Default::default() + })), + timestamp_ms: Some(1_770_000_000_000), + ..Default::default() + })), + message_context_info: next_secret.map(|secret| wa::MessageContextInfo { + message_secret: Some(secret), + ..Default::default() + }), + ..Default::default() + } + } + + fn encrypted_message_edit( + target_key: wa::MessageKey, + original_sender: &str, + editor: &str, + parent_id: &str, + secret: &[u8], + text: &str, + next_secret: Option>, + ) -> wa::Message { + let ctx = wacore::message_edit::MessageEditContext { + original_msg_id: parent_id, + original_sender_jid: original_sender, + editor_jid: editor, + }; + let (enc_payload, enc_iv) = wacore::message_edit::encrypt_message_edit( + &inner_message_edit(text, next_secret), + secret, + &ctx, + ) + .expect("test edit encryption"); + + wa::Message { + secret_encrypted_message: Some(wa::message::SecretEncryptedMessage { + target_message_key: Some(target_key), + enc_payload: Some(enc_payload), + enc_iv: Some(enc_iv.to_vec()), + secret_enc_type: Some( + wa::message::secret_encrypted_message::SecretEncType::MessageEdit as i32, + ), + remote_key_id: None, + }), + ..Default::default() + } + } + + #[tokio::test] + async fn secret_encrypted_message_edit_dispatches_legacy_edit() { + let (client, _transport) = capturing_client("secret_edit_dispatch").await; + let collector = Arc::new(crate::test_utils::TestEventCollector::default()); + client.register_handler(collector.clone()); + + let chat = "5511777776666@s.whatsapp.net"; + let parent_id = "PARENT_EDIT"; + let edit_id = "EDIT_1"; + let secret = [0x42u8; 32]; + client + .persistence_manager + .backend() + .put_msg_secret(chat, chat, parent_id, &secret) + .await + .unwrap(); + + let info = Arc::new(MessageInfo { + id: edit_id.into(), + source: crate::types::message::MessageSource { + chat: chat.parse().unwrap(), + sender: chat.parse().unwrap(), + ..Default::default() + }, + ..Default::default() + }); + let target_key = wa::MessageKey { + remote_jid: Some(chat.to_string()), + from_me: Some(false), + id: Some(parent_id.to_string()), + participant: None, + }; + let msg = + encrypted_message_edit(target_key, chat, chat, parent_id, &secret, "edited", None); + + client.dispatch_parsed_message(msg, &info).await; + + let got = collect_event( + &client, + collector, + |e| { + matches!(e, wacore::types::events::Event::Message(msg, info) + if info.id == edit_id + && legacy_edit_text(msg.as_ref()) == Some("edited") + && msg.secret_encrypted_message.is_none()) + }, + 500, + ) + .await; + assert!(got.is_some(), "encrypted edit must dispatch as legacy edit"); + } + + #[tokio::test] + async fn decrypted_message_edit_recaptures_secret_for_next_edit() { + let (client, _transport) = capturing_client("secret_edit_chain").await; + let collector = Arc::new(crate::test_utils::TestEventCollector::default()); + client.register_handler(collector.clone()); + + let chat = "5511777776666@s.whatsapp.net"; + let parent_id = "PARENT_EDIT"; + let first_secret = [0x11u8; 32]; + let second_secret = [0x22u8; 32]; + client + .persistence_manager + .backend() + .put_msg_secret(chat, chat, parent_id, &first_secret) + .await + .unwrap(); + + let target_key = wa::MessageKey { + remote_jid: Some(chat.to_string()), + from_me: Some(false), + id: Some(parent_id.to_string()), + participant: None, + }; + let first_info = Arc::new(MessageInfo { + id: "EDIT_CHAIN_1".into(), + source: crate::types::message::MessageSource { + chat: chat.parse().unwrap(), + sender: chat.parse().unwrap(), + ..Default::default() + }, + ..Default::default() + }); + let first_msg = encrypted_message_edit( + target_key.clone(), + chat, + chat, + parent_id, + &first_secret, + "first", + Some(second_secret.to_vec()), + ); + client.dispatch_parsed_message(first_msg, &first_info).await; + + let stored = client + .persistence_manager + .backend() + .get_msg_secret(chat, chat, parent_id) + .await + .unwrap(); + assert_eq!(stored.as_deref(), Some(&second_secret[..])); + + let second_info = Arc::new(MessageInfo { + id: "EDIT_CHAIN_2".into(), + source: crate::types::message::MessageSource { + chat: chat.parse().unwrap(), + sender: chat.parse().unwrap(), + ..Default::default() + }, + ..Default::default() + }); + let second_msg = encrypted_message_edit( + target_key, + chat, + chat, + parent_id, + &second_secret, + "second", + None, + ); + client + .dispatch_parsed_message(second_msg, &second_info) + .await; + + let got = collect_event( + &client, + collector, + |e| { + matches!(e, wacore::types::events::Event::Message(msg, info) + if info.id == "EDIT_CHAIN_2" + && legacy_edit_text(msg.as_ref()) == Some("second")) + }, + 500, + ) + .await; + assert!(got.is_some(), "second edit must use the re-captured secret"); + } + + #[tokio::test] + async fn secret_encrypted_message_edit_uses_lid_pn_fallback_in_group() { + use wacore::store::traits::LidPnMappingEntry; + + let (client, _transport) = capturing_client("secret_edit_alt_group").await; + let collector = Arc::new(crate::test_utils::TestEventCollector::default()); + client.register_handler(collector.clone()); + + let chat = "120363021033254949@g.us"; + let parent_id = "GROUP_PARENT_EDIT"; + let sender_lid = "236395184570386@lid"; + let sender_pn = "5511777776666@s.whatsapp.net"; + let secret = [0x77u8; 32]; + + client + .persistence_manager + .backend() + .put_lid_mapping(&LidPnMappingEntry { + lid: "236395184570386".into(), + phone_number: "5511777776666".into(), + created_at: 0, + updated_at: 0, + learning_source: "test".into(), + }) + .await + .unwrap(); + client + .persistence_manager + .backend() + .put_msg_secret(chat, sender_pn, parent_id, &secret) + .await + .unwrap(); + + let info = Arc::new(MessageInfo { + id: "GROUP_EDIT_1".into(), + source: crate::types::message::MessageSource { + chat: chat.parse().unwrap(), + sender: sender_lid.parse().unwrap(), + is_group: true, + addressing_mode: Some(wacore::types::message::AddressingMode::Lid), + ..Default::default() + }, + ..Default::default() + }); + let target_key = wa::MessageKey { + remote_jid: Some(chat.to_string()), + from_me: Some(false), + id: Some(parent_id.to_string()), + participant: Some(sender_lid.to_string()), + }; + let msg = encrypted_message_edit( + target_key, + sender_pn, + sender_lid, + parent_id, + &secret, + "group edited", + None, + ); + + client.dispatch_parsed_message(msg, &info).await; + + let got = collect_event( + &client, + collector, + |e| { + matches!(e, wacore::types::events::Event::Message(msg, info) + if info.id == "GROUP_EDIT_1" + && legacy_edit_text(msg.as_ref()) == Some("group edited")) + }, + 500, + ) + .await; + assert!( + got.is_some(), + "group edit must decrypt when the stored secret is under PN" + ); + } + + #[tokio::test] + async fn decrypted_message_edit_refreshes_alternate_secret_alias() { + use wacore::store::traits::LidPnMappingEntry; + + let (client, _transport) = capturing_client("secret_edit_alt_refresh").await; + let collector = Arc::new(crate::test_utils::TestEventCollector::default()); + client.register_handler(collector.clone()); + + let chat = "120363021033254949@g.us"; + let parent_id = "GROUP_PARENT_EDIT"; + let sender_lid = "236395184570386@lid"; + let sender_pn = "5511777776666@s.whatsapp.net"; + let first_secret = [0x31u8; 32]; + let second_secret = [0x32u8; 32]; + + client + .persistence_manager + .backend() + .put_lid_mapping(&LidPnMappingEntry { + lid: "236395184570386".into(), + phone_number: "5511777776666".into(), + created_at: 0, + updated_at: 0, + learning_source: "test".into(), + }) + .await + .unwrap(); + for sender in [sender_lid, sender_pn] { + client + .persistence_manager + .backend() + .put_msg_secret(chat, sender, parent_id, &first_secret) + .await + .unwrap(); + } + + let first_info = Arc::new(MessageInfo { + id: "GROUP_EDIT_REFRESH_1".into(), + source: crate::types::message::MessageSource { + chat: chat.parse().unwrap(), + sender: sender_lid.parse().unwrap(), + is_group: true, + addressing_mode: Some(wacore::types::message::AddressingMode::Lid), + ..Default::default() + }, + ..Default::default() + }); + let first_target_key = wa::MessageKey { + remote_jid: Some(chat.to_string()), + from_me: Some(false), + id: Some(parent_id.to_string()), + participant: Some(sender_lid.to_string()), + }; + let first_msg = encrypted_message_edit( + first_target_key, + sender_lid, + sender_lid, + parent_id, + &first_secret, + "first", + Some(second_secret.to_vec()), + ); + client.dispatch_parsed_message(first_msg, &first_info).await; + + for sender in [sender_lid, sender_pn] { + let stored = client + .persistence_manager + .backend() + .get_msg_secret(chat, sender, parent_id) + .await + .unwrap(); + assert_eq!(stored.as_deref(), Some(&second_secret[..])); + } + + let second_info = Arc::new(MessageInfo { + id: "GROUP_EDIT_REFRESH_2".into(), + source: crate::types::message::MessageSource { + chat: chat.parse().unwrap(), + sender: sender_pn.parse().unwrap(), + is_group: true, + addressing_mode: Some(wacore::types::message::AddressingMode::Pn), + ..Default::default() + }, + ..Default::default() + }); + let second_target_key = wa::MessageKey { + remote_jid: Some(chat.to_string()), + from_me: Some(false), + id: Some(parent_id.to_string()), + participant: Some(sender_pn.to_string()), + }; + let second_msg = encrypted_message_edit( + second_target_key, + sender_pn, + sender_pn, + parent_id, + &second_secret, + "second", + None, + ); + client + .dispatch_parsed_message(second_msg, &second_info) + .await; + + let got = collect_event( + &client, + collector, + |e| { + matches!(e, wacore::types::events::Event::Message(msg, info) + if info.id == "GROUP_EDIT_REFRESH_2" + && legacy_edit_text(msg.as_ref()) == Some("second")) + }, + 500, + ) + .await; + assert!( + got.is_some(), + "chained edit must use the refreshed alternate alias" + ); + } + /// Round-trip: store an outbound messageSecret, build a fake bot reply /// whose payload we encrypt with the symmetric helper, route it through /// classify, and assert the decrypted `wa::Message` lands on the bus. @@ -8956,8 +9584,8 @@ mod tests { } #[tokio::test] - async fn maybe_capture_inbound_msg_secret_skips_non_bot_chats() { - let (client, _transport) = capturing_client("capture_skip_dm").await; + async fn maybe_capture_inbound_msg_secret_persists_for_non_bot_chats() { + let (client, _transport) = capturing_client("capture_regular_dm").await; let info = Arc::new(MessageInfo { id: "DM_1".into(), source: crate::types::message::MessageSource { @@ -8978,9 +9606,6 @@ mod tests { }; client.maybe_capture_inbound_msg_secret(&msg, &info).await; - for _ in 0..16 { - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - } let got = client .persistence_manager .backend() @@ -8991,10 +9616,7 @@ mod tests { ) .await .unwrap(); - assert!( - got.is_none(), - "non-bot chats must not persist embedded msg_secrets" - ); + assert_eq!(got.as_deref(), Some(&[0xCDu8; 32][..])); } /// Group invocation: user mentions @MetaAI in a group → chat is the diff --git a/storages/sqlite-storage/src/sqlite_store.rs b/storages/sqlite-storage/src/sqlite_store.rs index 698d15c63..b4cd7b038 100644 --- a/storages/sqlite-storage/src/sqlite_store.rs +++ b/storages/sqlite-storage/src/sqlite_store.rs @@ -2552,6 +2552,58 @@ impl MsgSecretStore for SqliteStore { .await } + async fn put_msg_secrets(&self, entries: Vec) -> Result { + if entries.is_empty() { + return Ok(0); + } + + let device_id = self.device_id; + let entries: Arc<[MsgSecretEntry]> = Arc::from(entries); + let now = wacore::time::now_secs(); + self.with_retry("put_msg_secrets", || { + let entries = Arc::clone(&entries); + Box::new(move |conn: &mut SqliteConnection| { + let records: Vec<_> = entries + .iter() + .map(|entry| { + ( + msg_secrets::chat.eq(entry.chat.as_str()), + msg_secrets::sender.eq(entry.sender.as_str()), + msg_secrets::msg_id.eq(entry.msg_id.as_str()), + msg_secrets::secret.eq(entry.secret.as_slice()), + msg_secrets::device_id.eq(device_id), + msg_secrets::created_at.eq(now), + ) + }) + .collect(); + + const CHUNK_SIZE: usize = 100; + + conn.immediate_transaction(|conn| { + let mut stored = 0usize; + for chunk in records.chunks(CHUNK_SIZE) { + stored += diesel::insert_into(msg_secrets::table) + .values(chunk) + .on_conflict(( + msg_secrets::chat, + msg_secrets::sender, + msg_secrets::msg_id, + msg_secrets::device_id, + )) + .do_update() + .set(( + msg_secrets::secret.eq(excluded(msg_secrets::secret)), + msg_secrets::created_at.eq(excluded(msg_secrets::created_at)), + )) + .execute(conn)?; + } + Ok(stored) + }) + }) + }) + .await + } + async fn get_msg_secret( &self, chat: &str, @@ -3269,6 +3321,44 @@ mod tests { } } + #[tokio::test] + async fn msg_secret_batch_upserts_in_one_call() { + let store = create_test_store().await; + let stored = store + .put_msg_secrets(vec![ + MsgSecretEntry { + chat: "c".into(), + sender: "s".into(), + msg_id: "M1".into(), + secret: vec![1u8; 32], + }, + MsgSecretEntry { + chat: "c".into(), + sender: "s".into(), + msg_id: "M2".into(), + secret: vec![2u8; 32], + }, + MsgSecretEntry { + chat: "c".into(), + sender: "s".into(), + msg_id: "M1".into(), + secret: vec![9u8; 32], + }, + ]) + .await + .unwrap(); + + assert_eq!(stored, 3); + assert_eq!( + store.get_msg_secret("c", "s", "M1").await.unwrap().unwrap(), + vec![9u8; 32] + ); + assert_eq!( + store.get_msg_secret("c", "s", "M2").await.unwrap().unwrap(), + vec![2u8; 32] + ); + } + #[tokio::test] async fn delete_expired_msg_secrets_deletes_only_below_cutoff() { let store = create_test_store().await; diff --git a/wacore/src/history_sync.rs b/wacore/src/history_sync.rs index 277244cda..be7436b94 100644 --- a/wacore/src/history_sync.rs +++ b/wacore/src/history_sync.rs @@ -22,6 +22,7 @@ pub struct HistorySyncResult { pub conversations_processed: usize, /// Tctoken candidates extracted from 1:1 conversations during streaming. pub tc_token_candidates: Vec, + pub msg_secret_records: Vec, /// The full decompressed protobuf blob, only retained when event /// listeners exist. Wrapped in `LazyHistorySync` for on-demand decoding. pub decompressed_bytes: Option, @@ -36,11 +37,9 @@ mod wire_type { /// Decompress and process a history sync blob. /// -/// **Memory strategy**: Decompresses the entire blob into a single `Bytes` buffer, -/// then extracts conversation fields as zero-copy `Bytes::slice()` sub-views. -/// This trades a slightly higher peak (full decompressed blob in memory) for -/// **zero per-conversation heap allocations** — each conversation is just an -/// Arc refcount increment on the shared buffer. +/// **Memory strategy**: Decompresses the entire blob into a single `Bytes` +/// buffer, then scans top-level fields and partially decodes only the +/// conversation fields needed for internal caches. /// /// After decompression, the compressed input is dropped immediately, so peak /// memory = max(compressed, decompressed) + small overhead, not both. @@ -65,6 +64,7 @@ pub fn process_history_sync( nct_salt: None, conversations_processed: 0, tc_token_candidates: Vec::new(), + msg_secret_records: Vec::new(), decompressed_bytes: if retain_blob { Some(buf.clone()) } else { None }, }; @@ -83,9 +83,13 @@ pub fn process_history_sync( let end = checked_end(pos, len, buf.len(), "conversation")?; result.conversations_processed += 1; - if let Some(candidate) = extract_tc_token_fields(&buf[pos..end]) { + let extracted = extract_conversation_fields(&buf[pos..end]); + if let Some(candidate) = extracted.tc_token_candidate { result.tc_token_candidates.push(candidate); } + result + .msg_secret_records + .extend(extracted.msg_secret_records); pos = end; } @@ -250,11 +254,14 @@ fn extract_own_pushname(data: &[u8], own_user: &str) -> Option { if id_match { pushname } else { None } } -/// Prost partial decode — only tctoken fields, skips heavy `messages`. +/// Prost partial decode for internal HistorySync fields we persist while +/// streaming conversations. #[derive(Clone, PartialEq, prost::Message)] -pub(crate) struct ConversationTcTokenFields { +pub(crate) struct ConversationInternalFields { #[prost(string, required, tag = "1")] pub id: String, + #[prost(message, repeated, tag = "2")] + pub messages: Vec, #[prost(bytes = "vec", optional, tag = "21")] pub tc_token: Option>, #[prost(uint64, optional, tag = "22")] @@ -263,14 +270,283 @@ pub(crate) struct ConversationTcTokenFields { pub tc_token_sender_timestamp: Option, } -/// Extract tctoken candidate from a raw Conversation proto. -/// Uses prost partial decode (only fields 1/21/22/28, skips messages). -/// Returns `None` for groups, newsletters, bots, or conversations without tctokens. -pub(crate) fn extract_tc_token_fields(data: &[u8]) -> Option { +#[derive(Clone, PartialEq, prost::Message)] +pub(crate) struct HistorySyncMsgInternalFields { + #[prost(message, optional, tag = "1")] + pub message: Option, +} + +#[derive(Clone, PartialEq, prost::Message)] +pub(crate) struct WebMessageInfoInternalFields { + #[prost(message, optional, tag = "1")] + pub key: Option, + #[prost(message, optional, tag = "2")] + pub message: Option, + #[prost(string, optional, tag = "5")] + pub participant: Option, + #[prost(bytes = "vec", optional, tag = "49")] + pub message_secret: Option>, +} + +#[derive(Clone, PartialEq, prost::Message)] +pub(crate) struct MessageKeyInternalFields { + #[prost(bool, optional, tag = "2")] + pub from_me: Option, + #[prost(string, optional, tag = "3")] + pub id: Option, + #[prost(string, optional, tag = "4")] + pub participant: Option, +} + +#[derive(Clone, PartialEq, prost::Message)] +pub(crate) struct MessageInternalFields { + #[prost(message, optional, tag = "3")] + pub image_message: Option, + #[prost(message, optional, tag = "4")] + pub contact_message: Option, + #[prost(message, optional, tag = "5")] + pub location_message: Option, + #[prost(message, optional, tag = "6")] + pub extended_text_message: Option, + #[prost(message, optional, tag = "7")] + pub document_message: Option, + #[prost(message, optional, tag = "8")] + pub audio_message: Option, + #[prost(message, optional, tag = "9")] + pub video_message: Option, + #[prost(message, optional, tag = "13")] + pub contacts_array_message: Option, + #[prost(message, optional, tag = "18")] + pub live_location_message: Option, + #[prost(message, optional, tag = "25")] + pub template_message: Option, + #[prost(message, optional, tag = "26")] + pub sticker_message: Option, + #[prost(message, optional, tag = "28")] + pub group_invite_message: Option, + #[prost(message, optional, tag = "29")] + pub template_button_reply_message: Option, + #[prost(message, optional, tag = "30")] + pub product_message: Option, + #[prost(message, optional, tag = "31")] + pub device_sent_message: Option, + #[prost(message, optional, tag = "35")] + pub message_context_info: Option, + #[prost(message, optional, tag = "36")] + pub list_message: Option, + #[prost(message, optional, tag = "37")] + pub view_once_message: Option, + #[prost(message, optional, tag = "38")] + pub order_message: Option, + #[prost(message, optional, tag = "39")] + pub list_response_message: Option, + #[prost(message, optional, tag = "40")] + pub ephemeral_message: Option, + #[prost(message, optional, tag = "42")] + pub buttons_message: Option, + #[prost(message, optional, tag = "43")] + pub buttons_response_message: Option, + #[prost(message, optional, tag = "45")] + pub interactive_message: Option, + #[prost(message, optional, tag = "48")] + pub interactive_response_message: Option, + #[prost(message, optional, tag = "49")] + pub poll_creation_message: Option, + #[prost(message, optional, tag = "53")] + pub document_with_caption_message: Option, + #[prost(message, optional, tag = "55")] + pub view_once_message_v2: Option, + #[prost(message, optional, tag = "58")] + pub edited_message: Option, + #[prost(message, optional, tag = "60")] + pub poll_creation_message_v2: Option, + #[prost(message, optional, tag = "64")] + pub poll_creation_message_v3: Option, + #[prost(message, optional, tag = "75")] + pub event_message: Option, + #[prost(message, optional, tag = "78")] + pub newsletter_admin_invite_message: Option, + #[prost(message, optional, tag = "86")] + pub sticker_pack_message: Option, +} + +#[derive(Clone, PartialEq, prost::Message)] +pub(crate) struct MessageContextInfoInternalFields { + #[prost(bytes = "vec", optional, tag = "3")] + pub message_secret: Option>, +} + +macro_rules! define_context_info_carrier { + ($name:ident, $tag:literal) => { + #[derive(Clone, PartialEq, prost::Message)] + pub(crate) struct $name { + #[prost(message, optional, tag = $tag)] + pub context_info: Option, + } + + impl $name { + fn is_forwarded(&self) -> bool { + self.context_info + .as_ref() + .and_then(|ctx| ctx.is_forwarded) + .unwrap_or(false) + } + } + }; +} + +define_context_info_carrier!(ContextInfoTag1InternalFields, "1"); +define_context_info_carrier!(ContextInfoTag3InternalFields, "3"); +define_context_info_carrier!(ContextInfoTag4InternalFields, "4"); +define_context_info_carrier!(ContextInfoTag5InternalFields, "5"); +define_context_info_carrier!(ContextInfoTag6InternalFields, "6"); +define_context_info_carrier!(ContextInfoTag7InternalFields, "7"); +define_context_info_carrier!(ContextInfoTag8InternalFields, "8"); +define_context_info_carrier!(ContextInfoTag11InternalFields, "11"); +define_context_info_carrier!(ContextInfoTag15InternalFields, "15"); +define_context_info_carrier!(ContextInfoTag17InternalFields, "17"); + +#[derive(Clone, PartialEq, prost::Message)] +pub(crate) struct ContextInfoInternalFields { + #[prost(bool, optional, tag = "22")] + pub is_forwarded: Option, +} + +#[derive(Clone, PartialEq, prost::Message)] +pub(crate) struct DeviceSentMessageInternalFields { + #[prost(message, optional, boxed, tag = "2")] + pub message: Option>, +} + +#[derive(Clone, PartialEq, prost::Message)] +pub(crate) struct FutureProofMessageInternalFields { + #[prost(message, optional, boxed, tag = "1")] + pub message: Option>, +} + +impl MessageInternalFields { + fn base_message(&self) -> &Self { + let mut current = self; + loop { + let next = current + .device_sent_message + .as_ref() + .and_then(|m| m.message.as_deref()) + .or_else(|| { + current + .ephemeral_message + .as_ref() + .and_then(|m| m.message.as_deref()) + }) + .or_else(|| { + current + .view_once_message + .as_ref() + .and_then(|m| m.message.as_deref()) + }) + .or_else(|| { + current + .view_once_message_v2 + .as_ref() + .and_then(|m| m.message.as_deref()) + }) + .or_else(|| { + current + .document_with_caption_message + .as_ref() + .and_then(|m| m.message.as_deref()) + }) + .or_else(|| { + current + .edited_message + .as_ref() + .and_then(|m| m.message.as_deref()) + }); + + match next { + Some(msg) => current = msg, + None => return current, + } + } + } + + fn is_forwarded(&self) -> bool { + let base = self.base_message(); + macro_rules! any_forwarded { + ($($field:ident),+ $(,)?) => { + false $(|| base.$field.as_ref().map(|m| m.is_forwarded()).unwrap_or(false))+ + }; + } + + any_forwarded!( + extended_text_message, + image_message, + video_message, + audio_message, + document_message, + sticker_message, + location_message, + live_location_message, + contact_message, + contacts_array_message, + buttons_message, + buttons_response_message, + list_message, + list_response_message, + template_message, + template_button_reply_message, + interactive_message, + interactive_response_message, + poll_creation_message, + poll_creation_message_v2, + poll_creation_message_v3, + product_message, + order_message, + group_invite_message, + event_message, + sticker_pack_message, + newsletter_admin_invite_message, + ) + } +} + +struct ConversationExtraction { + tc_token_candidate: Option, + msg_secret_records: Vec, +} + +/// Message-secret data extracted from a conversation during streaming. +#[derive(Debug)] +pub struct HistoryMsgSecretRecord { + pub chat_id: String, + pub from_me: bool, + pub key_participant: Option, + pub web_msg_participant: Option, + pub msg_id: String, + pub secret: Vec, +} + +fn extract_conversation_fields(data: &[u8]) -> ConversationExtraction { use prost::Message; - let conv = ConversationTcTokenFields::decode(data).ok()?; + let Ok(conv) = ConversationInternalFields::decode(data) else { + return ConversationExtraction { + tc_token_candidate: None, + msg_secret_records: Vec::new(), + }; + }; + + let tc_token_candidate = extract_tc_token_fields(&conv); + let msg_secret_records = extract_msg_secret_records(&conv); + ConversationExtraction { + tc_token_candidate, + msg_secret_records, + } +} + +/// Returns `None` for groups, newsletters, bots, or conversations without tctokens. +fn extract_tc_token_fields(conv: &ConversationInternalFields) -> Option { // Early-out for non-1:1 conversations if let Some(parts) = wacore_binary::jid::parse_jid_fast(&conv.id) && (parts.server == "g.us" || parts.server == "newsletter" || parts.server == "bot") @@ -278,17 +554,58 @@ pub(crate) fn extract_tc_token_fields(data: &[u8]) -> Option { return None; } - let tc_token = conv.tc_token.filter(|t| !t.is_empty())?; + let tc_token = conv.tc_token.as_ref().filter(|t| !t.is_empty())?.clone(); let tc_token_timestamp = conv.tc_token_timestamp?; Some(TcTokenCandidate { - id: conv.id, + id: conv.id.clone(), tc_token, tc_token_timestamp, tc_token_sender_timestamp: conv.tc_token_sender_timestamp, }) } +fn extract_msg_secret_records(conv: &ConversationInternalFields) -> Vec { + let mut records = Vec::new(); + + for history_msg in &conv.messages { + let Some(web_msg) = history_msg.message.as_ref() else { + continue; + }; + let Some(key) = web_msg.key.as_ref() else { + continue; + }; + let Some(msg_id) = key.id.as_ref() else { + continue; + }; + if let Some(message) = web_msg.message.as_ref() + && message.is_forwarded() + { + continue; + } + let Some(secret) = web_msg.message_secret.as_ref().or_else(|| { + web_msg + .message + .as_ref() + .and_then(|m| m.message_context_info.as_ref()) + .and_then(|mci| mci.message_secret.as_ref()) + }) else { + continue; + }; + + records.push(HistoryMsgSecretRecord { + chat_id: conv.id.clone(), + from_me: key.from_me == Some(true), + key_participant: key.participant.clone(), + web_msg_participant: web_msg.participant.clone(), + msg_id: msg_id.clone(), + secret: secret.clone(), + }); + } + + records +} + /// Tctoken data extracted from a conversation during streaming. #[derive(Debug)] pub struct TcTokenCandidate { @@ -362,4 +679,177 @@ mod tests { assert_eq!(result.nct_salt, Some(salt)); assert_eq!(result.own_pushname.as_deref(), Some("TestUser")); } + + #[test] + fn test_message_secrets_extracted_from_history_sync() { + let chat = "5511777776666@s.whatsapp.net"; + let participant = "5511888889999@s.whatsapp.net"; + let top_level_secret = vec![0x44u8; 32]; + let context_secret = vec![0x55u8; 32]; + let hs = 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("HIST_TOP_LEVEL".to_string()), + participant: Some(participant.to_string()), + }, + message_secret: Some(top_level_secret.clone()), + ..Default::default() + }), + ..Default::default() + }, + wa::HistorySyncMsg { + message: Some(wa::WebMessageInfo { + key: wa::MessageKey { + remote_jid: Some(chat.to_string()), + from_me: Some(true), + id: Some("HIST_CONTEXT".to_string()), + participant: None, + }, + message: Some(wa::Message { + message_context_info: Some(wa::MessageContextInfo { + message_secret: Some(context_secret.clone()), + ..Default::default() + }), + ..Default::default() + }), + ..Default::default() + }), + ..Default::default() + }, + ], + ..Default::default() + }], + ..Default::default() + }; + + let compressed = encode_and_compress(&hs); + let result = process_history_sync(compressed, None, false, None).unwrap(); + + assert_eq!(result.msg_secret_records.len(), 2); + assert_eq!(result.msg_secret_records[0].chat_id, chat); + assert_eq!(result.msg_secret_records[0].msg_id, "HIST_TOP_LEVEL"); + assert_eq!( + result.msg_secret_records[0].key_participant.as_deref(), + Some(participant) + ); + assert_eq!(result.msg_secret_records[0].secret, top_level_secret); + assert_eq!(result.msg_secret_records[1].msg_id, "HIST_CONTEXT"); + assert!(result.msg_secret_records[1].from_me); + assert_eq!(result.msg_secret_records[1].secret, context_secret); + } + + #[test] + fn test_forwarded_message_secrets_skipped_from_history_sync() { + let chat = "5511000000001@s.whatsapp.net"; + let hs = 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("HIST_FORWARDED".to_string()), + ..Default::default() + }, + message: Some(wa::Message { + extended_text_message: Some(Box::new( + wa::message::ExtendedTextMessage { + text: Some("forwarded".into()), + context_info: Some(Box::new(wa::ContextInfo { + is_forwarded: Some(true), + ..Default::default() + })), + ..Default::default() + }, + )), + message_context_info: Some(wa::MessageContextInfo { + message_secret: Some(vec![0x66u8; 32]), + ..Default::default() + }), + ..Default::default() + }), + ..Default::default() + }), + ..Default::default() + }], + ..Default::default() + }], + ..Default::default() + }; + + let compressed = encode_and_compress(&hs); + let result = process_history_sync(compressed, None, false, None).unwrap(); + + assert!(result.msg_secret_records.is_empty()); + } + + #[test] + fn test_nested_forwarded_message_secrets_skipped_from_history_sync() { + let chat = "5511000000002@s.whatsapp.net"; + let hs = 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("HIST_NESTED_FORWARDED".to_string()), + ..Default::default() + }, + message: Some(wa::Message { + view_once_message: Some(Box::new(wa::message::FutureProofMessage { + message: Some(Box::new(wa::Message { + ephemeral_message: Some(Box::new( + wa::message::FutureProofMessage { + message: Some(Box::new(wa::Message { + extended_text_message: Some(Box::new( + wa::message::ExtendedTextMessage { + text: Some("nested".into()), + context_info: Some(Box::new( + wa::ContextInfo { + is_forwarded: Some(true), + ..Default::default() + }, + )), + ..Default::default() + }, + )), + ..Default::default() + })), + }, + )), + ..Default::default() + })), + })), + message_context_info: Some(wa::MessageContextInfo { + message_secret: Some(vec![0x77u8; 32]), + ..Default::default() + }), + ..Default::default() + }), + ..Default::default() + }), + ..Default::default() + }], + ..Default::default() + }], + ..Default::default() + }; + + let compressed = encode_and_compress(&hs); + let result = process_history_sync(compressed, None, false, None).unwrap(); + + assert!(result.msg_secret_records.is_empty()); + } } diff --git a/wacore/src/store/in_memory.rs b/wacore/src/store/in_memory.rs index 382c71413..0bf208f1f 100644 --- a/wacore/src/store/in_memory.rs +++ b/wacore/src/store/in_memory.rs @@ -608,6 +608,19 @@ impl MsgSecretStore for InMemoryBackend { Ok(()) } + async fn put_msg_secrets(&self, entries: Vec) -> Result { + let now = crate::time::now_secs(); + let stored = entries.len(); + let mut state = self.state.lock().await; + for entry in entries { + state.msg_secrets.insert( + (entry.chat, entry.sender, entry.msg_id), + (entry.secret, now), + ); + } + Ok(stored) + } + async fn get_msg_secret( &self, chat: &str, @@ -760,6 +773,52 @@ mod tests { ); } + #[tokio::test] + async fn msg_secret_batch_round_trip_and_overwrite() { + let backend = InMemoryBackend::new(); + let stored = backend + .put_msg_secrets(vec![ + MsgSecretEntry { + chat: "chat".into(), + sender: "sender".into(), + msg_id: "M1".into(), + secret: vec![1u8; 32], + }, + MsgSecretEntry { + chat: "chat".into(), + sender: "sender".into(), + msg_id: "M2".into(), + secret: vec![2u8; 32], + }, + MsgSecretEntry { + chat: "chat".into(), + sender: "sender".into(), + msg_id: "M1".into(), + secret: vec![9u8; 32], + }, + ]) + .await + .unwrap(); + + assert_eq!(stored, 3); + assert_eq!( + backend + .get_msg_secret("chat", "sender", "M1") + .await + .unwrap() + .unwrap(), + vec![9u8; 32] + ); + assert_eq!( + backend + .get_msg_secret("chat", "sender", "M2") + .await + .unwrap() + .unwrap(), + vec![2u8; 32] + ); + } + #[tokio::test] async fn delete_expired_msg_secrets_removes_only_old_rows() { let backend = InMemoryBackend::new(); diff --git a/wacore/src/store/traits.rs b/wacore/src/store/traits.rs index f494a00aa..c317d4e72 100644 --- a/wacore/src/store/traits.rs +++ b/wacore/src/store/traits.rs @@ -50,6 +50,15 @@ pub struct TcTokenEntry { pub sender_timestamp: Option, } +/// Message-secret write entry keyed by chat, sender, and message ID. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct MsgSecretEntry { + pub chat: String, + pub sender: String, + pub msg_id: String, + pub secret: Vec, +} + /// Device information for registry tracking. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct DeviceInfo { @@ -381,6 +390,17 @@ pub trait MsgSecretStore: Send + Sync { secret: &[u8], ) -> Result<()>; + /// Batched variant of `put_msg_secret`. + async fn put_msg_secrets(&self, entries: Vec) -> Result { + let mut stored = 0usize; + for entry in entries { + self.put_msg_secret(&entry.chat, &entry.sender, &entry.msg_id, &entry.secret) + .await?; + stored += 1; + } + Ok(stored) + } + /// Fetch the persisted secret; returns `None` if absent. async fn get_msg_secret( &self,