diff --git a/src/features/message_edit.rs b/src/features/message_edit.rs new file mode 100644 index 000000000..b33cd2903 --- /dev/null +++ b/src/features/message_edit.rs @@ -0,0 +1,456 @@ +//! Decryption of E2E message-edit envelopes (`secret_encrypted_message` +//! with `secret_enc_type = MESSAGE_EDIT`). +//! +//! See [`wacore::message_edit`] for the cryptographic primitives. This +//! module is the high-level surface: it takes typed [`Jid`]s, normalises +//! them the same way WA Web does (strip device suffix, optional LID↔PN +//! fallback) and returns the decrypted inner [`wa::Message`]. +//! +//! ### Integration +//! +//! The library does not auto-decrypt edits on the dispatch path because +//! doing so requires a callback into the consumer's message store to +//! fetch the parent's `messageContextInfo.messageSecret`. Consumers: +//! +//! 1. Observe `Event::Message` for messages whose +//! `message.secret_encrypted_message.secret_enc_type == MessageEdit`. +//! 2. Detect the envelope with [`extract_envelope`]. +//! 3. Look up the targeted message via `target_message_key`. +//! 4. Call [`decrypt`] with the parent's `messageSecret`. +//! 5. Optionally call [`rewrap_as_legacy_edit`] so downstream code that +//! already handles `protocol_message.edited_message` sees one shape. +//! +//! Mirrors the existing flow for poll vote decryption (`Polls::decrypt_vote`). + +use anyhow::{Result, anyhow}; +use log::warn; +use wacore::message_edit::{self, MessageEditContext}; +use wacore_binary::Jid; +use waproto::whatsapp as wa; + +/// Decrypt a `secret_encrypted_message` MESSAGE_EDIT envelope. +/// +/// JIDs may carry their device suffix — they are normalised before being +/// fed into the HKDF info buffer (matching WA Web's `widToUserJid`). +/// +/// Returns the inner [`wa::Message`]; the new content is at +/// `result.protocol_message.edited_message`. +/// +/// Implementation notes: +/// - HKDF: `salt = zeros[32]`, `ikm = message_secret`, +/// `info = original_msg_id || original_sender_jid || editor_jid || "Message Edit"`, +/// `L = 32`. +/// - AAD: empty. WA Web's `WAWebAddonEncryption` (function `g`) only binds +/// `stanzaId\0sender` into AAD for PollVote/EventResponse; everything +/// else, including MessageEdit, uses an empty AAD. +/// - IV must be exactly 12 bytes (matches WA Web's +/// `WAWebParseMessageEditEncryptedMessageProto`). +pub fn decrypt( + enc_payload: &[u8], + enc_iv: &[u8], + message_secret: &[u8], + original_msg_id: &str, + original_sender_jid: &Jid, + editor_jid: &Jid, +) -> Result { + let primary_orig = original_sender_jid.to_non_ad().to_string(); + let primary_editor = editor_jid.to_non_ad().to_string(); + let primary = MessageEditContext { + original_msg_id, + original_sender_jid: &primary_orig, + editor_jid: &primary_editor, + }; + message_edit::decrypt_message_edit(enc_payload, enc_iv, message_secret, &primary) +} + +/// Same as [`decrypt`] but tries a fallback addressing combination if +/// the first attempt fails its GCM tag check. +/// +/// `fallback_original_sender` / `fallback_editor` are typically the LID +/// form when the primary attempt used PN form (or vice versa). Mirrors +/// `WAWebAddonEncryption.decryptAddOn`, which falls back across LID/PN +/// to handle cross-addressing edits between newer and legacy clients. +#[allow(clippy::too_many_arguments)] +pub fn decrypt_with_fallback( + enc_payload: &[u8], + enc_iv: &[u8], + message_secret: &[u8], + original_msg_id: &str, + original_sender_jid: &Jid, + editor_jid: &Jid, + fallback_original_sender: Option<&Jid>, + fallback_editor: Option<&Jid>, +) -> Result { + let primary_orig = original_sender_jid.to_non_ad().to_string(); + let primary_editor = editor_jid.to_non_ad().to_string(); + let primary = MessageEditContext { + original_msg_id, + original_sender_jid: &primary_orig, + editor_jid: &primary_editor, + }; + + let fb_orig = fallback_original_sender.map(|j| j.to_non_ad().to_string()); + let fb_editor = fallback_editor.map(|j| j.to_non_ad().to_string()); + let fb_orig_resolved = fb_orig.as_deref().unwrap_or(primary.original_sender_jid); + let fb_editor_resolved = fb_editor.as_deref().unwrap_or(primary.editor_jid); + // Skip the retry when the fallback would key the HKDF identically to + // primary — covers both "no fallback supplied" and "fallback normalises + // to the same JIDs". Avoids a guaranteed-failing duplicate decrypt. + let fallback_ctx = if fb_orig_resolved == primary.original_sender_jid + && fb_editor_resolved == primary.editor_jid + { + None + } else { + Some(MessageEditContext { + original_msg_id, + original_sender_jid: fb_orig_resolved, + editor_jid: fb_editor_resolved, + }) + }; + + message_edit::decrypt_message_edit_with_fallback( + enc_payload, + enc_iv, + message_secret, + &primary, + fallback_ctx.as_ref(), + ) +} + +/// Pull `enc_payload` / `enc_iv` / `target_message_key` out of a received +/// [`wa::Message`] if it carries a MESSAGE_EDIT envelope. Returns `None` +/// if the message is not an encrypted edit, or if the envelope is +/// malformed (missing fields, IV not 12 bytes). +/// +/// Malformed-but-tagged envelopes emit a `log::warn!` so the gap is +/// visible without exposing the encrypted payload. +pub fn extract_envelope(msg: &wa::Message) -> Option> { + let sec = msg.secret_encrypted_message.as_ref()?; + let enc_type = sec.secret_enc_type(); + if enc_type != wa::message::secret_encrypted_message::SecretEncType::MessageEdit { + return None; + } + let target_key = sec.target_message_key.as_ref(); + let enc_payload = sec.enc_payload.as_deref(); + let enc_iv = sec.enc_iv.as_deref(); + + match (target_key, enc_payload, enc_iv) { + (Some(tk), Some(payload), Some(iv)) if iv.len() == 12 => Some(EncryptedEdit { + enc_payload: payload, + enc_iv: iv, + target_message_key: tk, + }), + (tk, payload, iv) => { + warn!( + "secret_encrypted_message MESSAGE_EDIT malformed: target_id={:?} has_payload={} iv_len={:?} (expected 12)", + tk.and_then(|t| t.id.as_deref()), + payload.is_some(), + iv.map(|b| b.len()), + ); + None + } + } +} + +/// Rewrap a decrypted edit `inner` into the same shape produced by the +/// legacy `protocol_message.edited_message` path so downstream consumers +/// can use one code path: +/// +/// ```text +/// Message { protocol_message: { edited_message: } } +/// ``` +/// +/// `inner` is the value returned by [`decrypt`]. Returns `None` if the +/// decrypted message did not contain `protocol_message.edited_message` +/// (caller should log + skip). +pub fn rewrap_as_legacy_edit(inner: wa::Message) -> Option { + let pm = inner.protocol_message?; + let edited = pm.edited_message?; + Some(wa::Message { + protocol_message: Some(Box::new(wa::message::ProtocolMessage { + key: pm.key, + r#type: Some(wa::message::protocol_message::Type::MessageEdit as i32), + edited_message: Some(edited), + timestamp_ms: pm.timestamp_ms, + ..Default::default() + })), + ..Default::default() + }) +} + +/// Extracted edit-envelope fields ready to feed into [`decrypt`]. +#[derive(Debug, Clone, Copy)] +pub struct EncryptedEdit<'a> { + pub enc_payload: &'a [u8], + pub enc_iv: &'a [u8], + pub target_message_key: &'a wa::MessageKey, +} + +impl<'a> EncryptedEdit<'a> { + /// Convenience: returns the targeted message id. + pub fn target_id(&self) -> Option<&str> { + self.target_message_key.id.as_deref() + } + + /// Resolve the original sender JID from the target message key. + /// + /// `my_jid` is the receiver's own JID in the addressing mode of the + /// chat (PN or LID). It is needed because for self-sent edits — e.g. + /// edits to our own messages that arrive via device sync — + /// `target_message_key` has `from_me = true` and its `remote_jid` + /// points to the *other* party, not us. WA Web's + /// `MsgGetters.getOriginalSender` reads `originalSelfAuthor || sender` + /// from its materialised msg-row store; we have no row here, so we + /// reconstruct the same fact from `from_me` + own jid. + /// + /// Resolution order: + /// 1. `participant` if present (always set in groups). + /// 2. `my_jid` if `from_me == Some(true)` (self-sent edit sync). + /// 3. `remote_jid` (1:1 incoming edit; the chat is the other party). + pub fn original_sender_jid(&self, my_jid: &Jid) -> Result { + if let Some(p) = self.target_message_key.participant.as_deref() { + return p + .parse::() + .map_err(|e| anyhow!("invalid participant jid in target key: {e}")); + } + if self.target_message_key.from_me == Some(true) { + return Ok(my_jid.to_non_ad()); + } + let raw = self + .target_message_key + .remote_jid + .as_deref() + .ok_or_else(|| anyhow!("target message key missing participant and remote_jid"))?; + raw.parse::() + .map_err(|e| anyhow!("invalid remote_jid in target key: {e}")) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use wacore::message_edit::encrypt_message_edit; + + fn inner(text: &str) -> wa::Message { + wa::Message { + protocol_message: Some(Box::new(wa::message::ProtocolMessage { + key: Some(wa::MessageKey { + remote_jid: Some("123@s.whatsapp.net".to_string()), + from_me: Some(false), + id: Some("AC1".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_700_000_000_000), + ..Default::default() + })), + ..Default::default() + } + } + + #[test] + fn decrypt_normalises_device_suffix() { + let secret = [0x55u8; 32]; + // Encrypt with the non-AD form, the only form WA actually feeds to HKDF. + let ctx = MessageEditContext { + original_msg_id: "AC1", + original_sender_jid: "5511999@s.whatsapp.net", + editor_jid: "5511999@s.whatsapp.net", + }; + let (enc, iv) = encrypt_message_edit(&inner("hi"), &secret, &ctx).unwrap(); + + // Caller passes JIDs with device numbers — they should be stripped. + let with_device = "5511999:13@s.whatsapp.net".parse::().unwrap(); + let m = decrypt(&enc, &iv, &secret, "AC1", &with_device, &with_device).unwrap(); + assert_eq!( + m.protocol_message + .as_ref() + .and_then(|pm| pm.edited_message.as_ref()) + .and_then(|e| e.conversation.as_deref()), + Some("hi") + ); + } + + #[test] + fn extract_envelope_recognises_message_edit() { + let msg = wa::Message { + secret_encrypted_message: Some(wa::message::SecretEncryptedMessage { + target_message_key: Some(wa::MessageKey { + remote_jid: Some("g@g.us".to_string()), + from_me: Some(false), + id: Some("AC1".to_string()), + participant: Some("5511999@s.whatsapp.net".to_string()), + }), + enc_payload: Some(vec![0u8; 32]), + enc_iv: Some(vec![0u8; 12]), + secret_enc_type: Some( + wa::message::secret_encrypted_message::SecretEncType::MessageEdit as i32, + ), + remote_key_id: None, + }), + ..Default::default() + }; + let env = extract_envelope(&msg).expect("recognised"); + assert_eq!(env.target_id(), Some("AC1")); + // Group: participant takes priority over my_jid and remote_jid. + let my_jid = "999@s.whatsapp.net".parse::().unwrap(); + assert_eq!( + env.original_sender_jid(&my_jid).unwrap().to_string(), + "5511999@s.whatsapp.net" + ); + } + + #[test] + fn original_sender_jid_uses_my_jid_for_self_sent_edits() { + let msg = wa::Message { + secret_encrypted_message: Some(wa::message::SecretEncryptedMessage { + target_message_key: Some(wa::MessageKey { + remote_jid: Some("5510000@s.whatsapp.net".to_string()), + from_me: Some(true), + id: Some("AC1".to_string()), + participant: None, + }), + enc_payload: Some(vec![0u8; 32]), + enc_iv: Some(vec![0u8; 12]), + secret_enc_type: Some( + wa::message::secret_encrypted_message::SecretEncType::MessageEdit as i32, + ), + remote_key_id: None, + }), + ..Default::default() + }; + let env = extract_envelope(&msg).expect("recognised"); + let my_jid = "5511999:13@s.whatsapp.net".parse::().unwrap(); + // Must return my_jid (stripped of device), NOT remote_jid (the other party). + assert_eq!( + env.original_sender_jid(&my_jid).unwrap().to_string(), + "5511999@s.whatsapp.net" + ); + } + + #[test] + fn original_sender_jid_falls_back_to_remote_jid_for_incoming_one_to_one_edit() { + let msg = wa::Message { + secret_encrypted_message: Some(wa::message::SecretEncryptedMessage { + target_message_key: Some(wa::MessageKey { + remote_jid: Some("5510000@s.whatsapp.net".to_string()), + from_me: Some(false), + id: Some("AC1".to_string()), + participant: None, + }), + enc_payload: Some(vec![0u8; 32]), + enc_iv: Some(vec![0u8; 12]), + secret_enc_type: Some( + wa::message::secret_encrypted_message::SecretEncType::MessageEdit as i32, + ), + remote_key_id: None, + }), + ..Default::default() + }; + let env = extract_envelope(&msg).expect("recognised"); + let my_jid = "5511999@s.whatsapp.net".parse::().unwrap(); + assert_eq!( + env.original_sender_jid(&my_jid).unwrap().to_string(), + "5510000@s.whatsapp.net" + ); + } + + #[test] + fn extract_envelope_rejects_non_edit_secret_enc_type() { + let msg = wa::Message { + secret_encrypted_message: Some(wa::message::SecretEncryptedMessage { + target_message_key: Some(wa::MessageKey::default()), + enc_payload: Some(vec![0u8; 32]), + enc_iv: Some(vec![0u8; 12]), + secret_enc_type: Some( + wa::message::secret_encrypted_message::SecretEncType::EventEdit as i32, + ), + remote_key_id: None, + }), + ..Default::default() + }; + assert!(extract_envelope(&msg).is_none()); + } + + #[test] + fn extract_envelope_rejects_invalid_iv_size() { + let msg = wa::Message { + secret_encrypted_message: Some(wa::message::SecretEncryptedMessage { + target_message_key: Some(wa::MessageKey::default()), + enc_payload: Some(vec![0u8; 32]), + enc_iv: Some(vec![0u8; 11]), + secret_enc_type: Some( + wa::message::secret_encrypted_message::SecretEncType::MessageEdit as i32, + ), + remote_key_id: None, + }), + ..Default::default() + }; + assert!(extract_envelope(&msg).is_none()); + } + + #[test] + fn fallback_normalising_to_primary_jids_is_skipped() { + // wacore::message_edit::decrypt_message_edit_with_fallback returns the + // bare primary error when no fallback is run, or a combined + // "edit decrypt failed: primary=...; fallback=..." when both attempts + // run. We use that to assert the dedup path. + let secret = [0xAAu8; 32]; + let real_ctx = MessageEditContext { + original_msg_id: "ID", + original_sender_jid: "5511777@s.whatsapp.net", + editor_jid: "5511777@s.whatsapp.net", + }; + let (enc, iv) = encrypt_message_edit(&inner("hi"), &secret, &real_ctx).unwrap(); + + // Wrong primary JID so decrypt fails; fallback is a device-suffixed + // form of the *same* wrong jid → normalises identical → must be skipped. + let wrong = "5511000@s.whatsapp.net".parse::().unwrap(); + let wrong_with_device = "5511000:5@s.whatsapp.net".parse::().unwrap(); + + let err = decrypt_with_fallback( + &enc, + &iv, + &secret, + "ID", + &wrong, + &wrong, + Some(&wrong_with_device), + Some(&wrong_with_device), + ) + .expect_err("decryption should fail"); + assert!( + !err.to_string().contains("fallback="), + "no-op fallback must be skipped, got: {err}" + ); + } + + #[test] + fn rewrap_yields_legacy_shape() { + let dec = inner("edited"); + let rewrap = rewrap_as_legacy_edit(dec).expect("present"); + let edited = rewrap + .protocol_message + .as_ref() + .and_then(|pm| pm.edited_message.as_ref()) + .and_then(|m| m.conversation.as_deref()); + assert_eq!(edited, Some("edited")); + assert_eq!( + rewrap.protocol_message.as_ref().and_then(|pm| pm.r#type), + Some(wa::message::protocol_message::Type::MessageEdit as i32) + ); + } + + #[test] + fn rewrap_returns_none_when_inner_missing_edit() { + let m = wa::Message { + protocol_message: Some(Box::new(wa::message::ProtocolMessage::default())), + ..Default::default() + }; + assert!(rewrap_as_legacy_edit(m).is_none()); + } +} diff --git a/src/features/mod.rs b/src/features/mod.rs index 7dcf97af7..61dabc1ea 100644 --- a/src/features/mod.rs +++ b/src/features/mod.rs @@ -5,6 +5,7 @@ mod community; mod contacts; mod groups; mod media_reupload; +pub mod message_edit; mod mex; pub(crate) mod newsletter; mod polls; @@ -37,6 +38,8 @@ pub use groups::{ pub use media_reupload::{MediaRetryResult, MediaReupload, MediaReuploadRequest}; +pub use message_edit::EncryptedEdit; + pub use mex::{Mex, MexError, MexErrorExtensions, MexGraphQLError, MexRequest, MexResponse}; pub use newsletter::{ diff --git a/src/features/polls.rs b/src/features/polls.rs index 16dbb51f4..de3488dbe 100644 --- a/src/features/polls.rs +++ b/src/features/polls.rs @@ -127,16 +127,14 @@ impl<'a> Polls<'a> { .map(|name| poll::compute_option_hash(name).to_vec()) .collect(); - let key = poll::derive_vote_encryption_key( + let (enc_payload, iv) = poll::encrypt_poll_vote_with_secret( + &selected_hashes, message_secret, poll_msg_id, &creator_jid_str, &voter_jid_str, )?; - let (enc_payload, iv) = - poll::encrypt_poll_vote(&selected_hashes, &key, poll_msg_id, &voter_jid_str)?; - let from_me = my_base.is_same_user_as(poll_creator_jid); let poll_update = wa::message::PollUpdateMessage { @@ -178,8 +176,14 @@ impl<'a> Polls<'a> { ) -> Result>> { let creator = poll_creator_jid.to_non_ad().to_string(); let voter = voter_jid.to_non_ad().to_string(); - let key = poll::derive_vote_encryption_key(message_secret, poll_msg_id, &creator, &voter)?; - poll::decrypt_poll_vote(enc_payload, enc_iv, &key, poll_msg_id, &voter) + poll::decrypt_poll_vote_with_secret( + enc_payload, + enc_iv, + message_secret, + poll_msg_id, + &creator, + &voter, + ) } /// Decrypts each vote and tallies per-option results. @@ -205,19 +209,14 @@ impl<'a> Polls<'a> { let mut latest_votes: HashMap>> = HashMap::with_capacity(votes.len()); for (voter_jid, enc_payload, enc_iv) in votes { let voter_str = voter_jid.to_non_ad().to_string(); - let key = match poll::derive_vote_encryption_key( + match poll::decrypt_poll_vote_with_secret( + enc_payload, + enc_iv, message_secret, poll_msg_id, &creator_str, &voter_str, ) { - Ok(k) => k, - Err(e) => { - log::warn!("Failed to derive vote key for {voter_jid}: {e}"); - continue; - } - }; - match poll::decrypt_poll_vote(enc_payload, enc_iv, &key, poll_msg_id, &voter_str) { Ok(selected_hashes) => { if selected_hashes.is_empty() { // Empty selection = voter cleared their vote diff --git a/src/lib.rs b/src/lib.rs index 39b1fb9af..f54810dbc 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -66,17 +66,18 @@ pub mod features; pub use features::{ BatchGroupResult, Blocking, BlocklistEntry, ChatActions, ChatStateType, Chatstate, Community, CommunitySubgroup, Contacts, CreateCommunityOptions, CreateCommunityResult, CreateGroupResult, - GroupCreateOptions, GroupDescription, GroupJoinError, GroupMetadata, GroupParticipant, - GroupParticipantOptions, GroupProfilePicture, GroupSubject, GroupType, Groups, GrowthLockInfo, - InviteInfoError, IsOnWhatsAppResult, JoinGroupResult, LinkSubgroupsResult, MediaRetryResult, - MediaReupload, MediaReuploadRequest, MemberAddMode, MemberLinkMode, MemberShareHistoryMode, - MembershipApprovalMode, MembershipRequest, Mex, MexError, MexErrorExtensions, MexRequest, - MexResponse, Newsletter, NewsletterMessage, NewsletterMessageType, NewsletterMetadata, - NewsletterReactionCount, NewsletterRole, NewsletterState, NewsletterVerification, - ParticipantChangeResponse, ParticipantType, PictureType, Presence, PresenceError, - PresenceStatus, Profile, ProfilePicture, SetProfilePictureResponse, Signal, Status, - StatusPrivacySetting, StatusSendOptions, SyncActionMessageRange, TcToken, - UnlinkSubgroupsResult, UserInfo, group_type, message_key, message_range, + EncryptedEdit, GroupCreateOptions, GroupDescription, GroupJoinError, GroupMetadata, + GroupParticipant, GroupParticipantOptions, GroupProfilePicture, GroupSubject, GroupType, + Groups, GrowthLockInfo, InviteInfoError, IsOnWhatsAppResult, JoinGroupResult, + LinkSubgroupsResult, MediaRetryResult, MediaReupload, MediaReuploadRequest, MemberAddMode, + MemberLinkMode, MemberShareHistoryMode, MembershipApprovalMode, MembershipRequest, Mex, + MexError, MexErrorExtensions, MexRequest, MexResponse, Newsletter, NewsletterMessage, + NewsletterMessageType, NewsletterMetadata, NewsletterReactionCount, NewsletterRole, + NewsletterState, NewsletterVerification, ParticipantChangeResponse, ParticipantType, + PictureType, Presence, PresenceError, PresenceStatus, Profile, ProfilePicture, + SetProfilePictureResponse, Signal, Status, StatusPrivacySetting, StatusSendOptions, + SyncActionMessageRange, TcToken, UnlinkSubgroupsResult, UserInfo, group_type, message_key, + message_range, }; pub mod bot; diff --git a/wacore/src/lib.rs b/wacore/src/lib.rs index 8f1fd0a4c..e689af9d8 100644 --- a/wacore/src/lib.rs +++ b/wacore/src/lib.rs @@ -20,6 +20,7 @@ pub mod history_sync; pub mod ib; pub use wacore_libsignal as libsignal; pub mod media_retry; +pub mod message_edit; pub mod message_processing; pub mod messages; pub mod net; @@ -31,6 +32,7 @@ pub mod proto_helpers; pub mod reporting_token; pub mod request; pub mod runtime; +pub mod secret_enc_addon; pub mod send; pub mod session; pub mod stanza; diff --git a/wacore/src/message_edit.rs b/wacore/src/message_edit.rs new file mode 100644 index 000000000..da89a09d6 --- /dev/null +++ b/wacore/src/message_edit.rs @@ -0,0 +1,289 @@ +//! E2E-encrypted message edit envelope (`secret_encrypted_message` with +//! `secret_enc_type = MESSAGE_EDIT`). +//! +//! Introduced by WhatsApp in 2026: newer clients wrap message edits in an +//! E2EE envelope keyed by the original message's +//! `messageContextInfo.messageSecret`, replacing the older in-clear +//! `protocolMessage.editedMessage` path. +//! +//! Verified against `docs/captured-js/`: +//! +//! - `WA/Use/CaseSecret.js` — use-case literal `"Message Edit"` and the +//! HKDF info ordering (`stanzaId || parentOrigSender || editor || usecase`). +//! - `WAWeb/Parse/MessageEditEncryptedMessageProto.js` — envelope detection +//! and IV-length validation (`u.length !== 12`). +//! - `WAWeb/Process/EncryptedMessageEditMsgs.js` — invocation contract. +//! - `WAWeb/Addon/Encryption.js` function `g` — AAD is empty for the +//! MessageEdit branch (only `PollVote` and `EventResponse` bind stanza/sender +//! into AAD). +//! +//! The plaintext is a `Message` proto whose `protocolMessage.editedMessage` +//! carries the new content — same shape as the legacy edit path, so callers +//! that already handle `protocolMessage.editedMessage` can reuse their code. + +use anyhow::{Result, anyhow}; +use prost::Message; + +use crate::secret_enc_addon::{AddonContext, ModificationType, decrypt_addon, encrypt_addon}; + +const IV_SIZE: usize = 12; + +/// Inputs for decrypting / encrypting a `MESSAGE_EDIT` envelope. +/// +/// JIDs must be passed in the addressing they were received under. The +/// caller is responsible for any LID↔PN normalisation; see +/// [`decrypt_message_edit_with_fallback`]. +#[derive(Debug, Clone, Copy)] +pub struct MessageEditContext<'a> { + /// ID of the *original* (target) message being edited — this is the + /// stanza id WA Web feeds into the HKDF info buffer. + pub original_msg_id: &'a str, + /// JID of the sender of the original message (`participant` for groups, + /// `remote_jid` for 1:1). + pub original_sender_jid: &'a str, + /// JID of the user performing the edit (sender of the envelope). + pub editor_jid: &'a str, +} + +impl<'a> MessageEditContext<'a> { + fn as_addon_ctx(&self) -> AddonContext<'a> { + AddonContext { + stanza_id: self.original_msg_id, + parent_msg_original_sender: self.original_sender_jid, + modification_sender: self.editor_jid, + modification_type: ModificationType::MessageEdit, + } + } +} + +/// Encrypt an edit payload for testing or for client-side outgoing edits. +/// +/// `inner_message` is the full `Message` proto whose `protocolMessage.editedMessage` +/// carries the new content; it gets serialised and encrypted in one shot. +pub fn encrypt_message_edit( + inner_message: &waproto::whatsapp::Message, + message_secret: &[u8], + ctx: &MessageEditContext<'_>, +) -> Result<(Vec, [u8; IV_SIZE])> { + let mut plaintext = Vec::new(); + inner_message.encode(&mut plaintext)?; + encrypt_addon(&plaintext, message_secret, &ctx.as_addon_ctx()) +} + +/// Decrypt a `secret_encrypted_message` MESSAGE_EDIT envelope. +/// +/// Returns the inner `Message` whose `protocolMessage.editedMessage` carries +/// the new content. The caller can re-wrap it in `protocolMessage.editedMessage` +/// for consumer parity with the legacy edit path. +/// +/// `iv` must be exactly 12 bytes (matches WA Web's `u.length !== 12` check). +pub fn decrypt_message_edit( + enc_payload: &[u8], + iv: &[u8], + message_secret: &[u8], + ctx: &MessageEditContext<'_>, +) -> Result { + if iv.len() != IV_SIZE { + return Err(anyhow!( + "Invalid edit IV length: expected {IV_SIZE}, got {}", + iv.len() + )); + } + let plaintext = decrypt_addon(enc_payload, iv, message_secret, &ctx.as_addon_ctx())?; + let msg = waproto::whatsapp::Message::decode(&plaintext[..]) + .map_err(|e| anyhow!("Failed to decode inner edit Message: {e}"))?; + Ok(msg) +} + +/// Try decryption with up to two JID combinations, mirroring WA Web's +/// `decryptAddOn` resilience for cross-addressing edits. +/// +/// The pair `(originals, alternates)` is the (as-received, normalised) +/// combination — caller decides which is which. Returns the first success. +/// +/// WA Web tries 3 attempts (LID→LID, PN→PN, originals) but two are enough +/// in practice when the caller already knows the canonical pair. +pub fn decrypt_message_edit_with_fallback( + enc_payload: &[u8], + iv: &[u8], + message_secret: &[u8], + primary: &MessageEditContext<'_>, + fallback: Option<&MessageEditContext<'_>>, +) -> Result { + match decrypt_message_edit(enc_payload, iv, message_secret, primary) { + Ok(m) => Ok(m), + Err(primary_err) => match fallback { + Some(fb) => { + decrypt_message_edit(enc_payload, iv, message_secret, fb).map_err(|fb_err| { + anyhow!("edit decrypt failed: primary={primary_err}; fallback={fb_err}") + }) + } + None => Err(primary_err), + }, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use waproto::whatsapp as wa; + + fn make_inner_edit(new_text: &str) -> wa::Message { + wa::Message { + protocol_message: Some(Box::new(wa::message::ProtocolMessage { + key: Some(wa::MessageKey { + remote_jid: Some("g@g.us".to_string()), + from_me: Some(true), + id: Some("AC1234567890ABCDEF".to_string()), + participant: None, + }), + r#type: Some(wa::message::protocol_message::Type::MessageEdit as i32), + edited_message: Some(Box::new(wa::Message { + conversation: Some(new_text.to_string()), + ..Default::default() + })), + timestamp_ms: Some(1_700_000_000_000), + ..Default::default() + })), + ..Default::default() + } + } + + #[test] + fn encrypt_decrypt_roundtrip_text() { + let secret = [0x07u8; 32]; + let ctx = MessageEditContext { + original_msg_id: "AC1234567890ABCDEF", + original_sender_jid: "5511999999999@s.whatsapp.net", + editor_jid: "5511999999999@s.whatsapp.net", + }; + let inner = make_inner_edit("edited text"); + let (enc, iv) = encrypt_message_edit(&inner, &secret, &ctx).unwrap(); + + let decoded = decrypt_message_edit(&enc, &iv, &secret, &ctx).unwrap(); + let edited = decoded + .protocol_message + .as_ref() + .and_then(|pm| pm.edited_message.as_ref()) + .expect("inner edited message present"); + assert_eq!(edited.conversation.as_deref(), Some("edited text")); + } + + #[test] + fn encrypt_decrypt_roundtrip_with_lid_jids() { + let secret = [0x42u8; 32]; + let ctx = MessageEditContext { + original_msg_id: "AC1191FE0A25A0E319BEA72064819280", + original_sender_jid: "260661598801930@lid", + editor_jid: "260661598801930@lid", + }; + let inner = make_inner_edit("B"); + let (enc, iv) = encrypt_message_edit(&inner, &secret, &ctx).unwrap(); + let decoded = decrypt_message_edit(&enc, &iv, &secret, &ctx).unwrap(); + assert_eq!( + decoded + .protocol_message + .as_ref() + .and_then(|pm| pm.edited_message.as_ref()) + .and_then(|m| m.conversation.as_deref()), + Some("B") + ); + } + + #[test] + fn wrong_editor_jid_fails() { + // GCM tag should fail when the editor JID feeding HKDF differs. + let secret = [0x07u8; 32]; + let ctx = MessageEditContext { + original_msg_id: "AC1", + original_sender_jid: "a@s.whatsapp.net", + editor_jid: "a@s.whatsapp.net", + }; + let (enc, iv) = encrypt_message_edit(&make_inner_edit("x"), &secret, &ctx).unwrap(); + + let bad = MessageEditContext { + editor_jid: "b@s.whatsapp.net", + ..ctx + }; + assert!(decrypt_message_edit(&enc, &iv, &secret, &bad).is_err()); + } + + #[test] + fn wrong_message_secret_fails() { + let ctx = MessageEditContext { + original_msg_id: "AC1", + original_sender_jid: "a@s.whatsapp.net", + editor_jid: "a@s.whatsapp.net", + }; + let (enc, iv) = encrypt_message_edit(&make_inner_edit("x"), &[0x07u8; 32], &ctx).unwrap(); + assert!(decrypt_message_edit(&enc, &iv, &[0x08u8; 32], &ctx).is_err()); + } + + #[test] + fn invalid_iv_length_rejected() { + let ctx = MessageEditContext { + original_msg_id: "AC1", + original_sender_jid: "a@s.whatsapp.net", + editor_jid: "a@s.whatsapp.net", + }; + let (enc, _iv) = encrypt_message_edit(&make_inner_edit("x"), &[0x07u8; 32], &ctx).unwrap(); + // WA Web enforces 12-byte IV; we surface a typed error here. + assert!(decrypt_message_edit(&enc, &[0u8; 11], &[0x07u8; 32], &ctx).is_err()); + assert!(decrypt_message_edit(&enc, &[0u8; 16], &[0x07u8; 32], &ctx).is_err()); + } + + #[test] + fn fallback_recovers_on_alternate_jid_form() { + let secret = [0x09u8; 32]; + // Encrypted under PN form... + let pn_ctx = MessageEditContext { + original_msg_id: "ID", + original_sender_jid: "5511999@s.whatsapp.net", + editor_jid: "5511999@s.whatsapp.net", + }; + // ...but consumer first guesses LID. + let lid_ctx = MessageEditContext { + original_msg_id: "ID", + original_sender_jid: "12345@lid", + editor_jid: "12345@lid", + }; + + let (enc, iv) = encrypt_message_edit(&make_inner_edit("hello"), &secret, &pn_ctx).unwrap(); + + // Primary (LID) fails, fallback (PN) succeeds. + let m = decrypt_message_edit_with_fallback(&enc, &iv, &secret, &lid_ctx, Some(&pn_ctx)) + .expect("fallback should rescue"); + assert_eq!( + m.protocol_message + .as_ref() + .and_then(|pm| pm.edited_message.as_ref()) + .and_then(|m| m.conversation.as_deref()), + Some("hello") + ); + } + + #[test] + fn fallback_returns_combined_error_when_both_fail() { + let secret = [0x09u8; 32]; + let pn_ctx = MessageEditContext { + original_msg_id: "ID", + original_sender_jid: "5511999@s.whatsapp.net", + editor_jid: "5511999@s.whatsapp.net", + }; + let (enc, iv) = encrypt_message_edit(&make_inner_edit("x"), &secret, &pn_ctx).unwrap(); + + let wrong1 = MessageEditContext { + editor_jid: "evil1@s.whatsapp.net", + ..pn_ctx + }; + let wrong2 = MessageEditContext { + editor_jid: "evil2@s.whatsapp.net", + ..pn_ctx + }; + let err = decrypt_message_edit_with_fallback(&enc, &iv, &secret, &wrong1, Some(&wrong2)) + .expect_err("both fail"); + let s = err.to_string(); + assert!(s.contains("primary=")); + assert!(s.contains("fallback=")); + } +} diff --git a/wacore/src/poll.rs b/wacore/src/poll.rs index 64289294e..49de390bf 100644 --- a/wacore/src/poll.rs +++ b/wacore/src/poll.rs @@ -1,16 +1,31 @@ -//! Poll vote encryption/decryption (AES-256-GCM + HKDF-SHA256). +//! Poll vote encryption/decryption. //! -//! Matches WAWebPollVoteEncryptMsgData / WAUseCaseSecret. +//! Thin wrapper over [`secret_enc_addon`] specialised for the +//! `PollVoteMessage` proto and the `"Poll Vote"` use-case. use anyhow::{Result, anyhow}; -use hkdf::Hkdf; use sha2::{Digest, Sha256}; -use crate::libsignal::crypto::{aes_256_gcm_decrypt, aes_256_gcm_encrypt}; +use crate::secret_enc_addon::{ + AddonContext, ModificationType, build_aad, decrypt_addon, encrypt_addon, +}; const GCM_IV_SIZE: usize = 12; const GCM_TAG_SIZE: usize = 16; +fn poll_vote_addon_ctx<'a>( + stanza_id: &'a str, + poll_creator_jid: &'a str, + voter_jid: &'a str, +) -> AddonContext<'a> { + AddonContext { + stanza_id, + parent_msg_original_sender: poll_creator_jid, + modification_sender: voter_jid, + modification_type: ModificationType::PollVote, + } +} + /// Votes reference options by SHA-256 hash, not by name. pub fn compute_option_hash(option_name: &str) -> [u8; 32] { let mut hasher = Sha256::new(); @@ -19,64 +34,48 @@ pub fn compute_option_hash(option_name: &str) -> [u8; 32] { } /// HKDF-SHA256: info = stanzaId || pollCreator || voter || "Poll Vote", no salt. -/// Matches WA Web's `createUseCaseSecret()` with UseCase = "Poll Vote". +/// +/// Kept as a public API for backwards compatibility; new code should call +/// [`secret_enc_addon::derive_use_case_secret`] with `ModificationType::PollVote`. pub fn derive_vote_encryption_key( message_secret: &[u8], stanza_id: &str, poll_creator_jid: &str, voter_jid: &str, ) -> Result<[u8; 32]> { - if message_secret.len() != 32 { - return Err(anyhow!( - "Invalid messageSecret size: expected 32, got {}", - message_secret.len() - )); - } - - let mut info = Vec::new(); - info.extend_from_slice(stanza_id.as_bytes()); - info.extend_from_slice(poll_creator_jid.as_bytes()); - info.extend_from_slice(voter_jid.as_bytes()); - info.extend_from_slice(b"Poll Vote"); - - let hk = Hkdf::::new(None, message_secret); - let mut key = [0u8; 32]; - hk.expand(&info, &mut key) - .map_err(|e| anyhow!("HKDF expand failed: {e}"))?; - - Ok(key) -} - -/// AAD = stanzaId + "\0" + voterJid (WAWebAddonEncryption.js:10) -fn build_vote_aad(stanza_id: &str, voter_jid: &str) -> Vec { - let mut aad = Vec::with_capacity(stanza_id.len() + 1 + voter_jid.len()); - aad.extend_from_slice(stanza_id.as_bytes()); - aad.push(0); - aad.extend_from_slice(voter_jid.as_bytes()); - aad + crate::secret_enc_addon::derive_use_case_secret( + message_secret, + &poll_vote_addon_ctx(stanza_id, poll_creator_jid, voter_jid), + ) } -/// Returns `(encrypted_payload_with_tag, iv)`. +/// Encrypt a poll vote with a pre-derived 32-byte key, symmetric to +/// [`decrypt_poll_vote`]. Returns `(payload_with_tag, iv)`. +/// +/// Kept for callers that built their own key via [`derive_vote_encryption_key`]. +/// New code should prefer [`encrypt_poll_vote_with_secret`], which derives +/// the key in a single step from the parent poll's `messageSecret`. pub fn encrypt_poll_vote( selected_option_hashes: &[Vec], encryption_key: &[u8; 32], stanza_id: &str, voter_jid: &str, ) -> Result<(Vec, [u8; GCM_IV_SIZE])> { + use crate::libsignal::crypto::aes_256_gcm_encrypt; use prost::Message; use rand::Rng; let vote_msg = waproto::whatsapp::message::PollVoteMessage { selected_options: selected_option_hashes.to_vec(), }; - let mut plaintext = Vec::new(); vote_msg.encode(&mut plaintext)?; let mut iv = [0u8; GCM_IV_SIZE]; rand::make_rng::().fill_bytes(&mut iv); - let aad = build_vote_aad(stanza_id, voter_jid); + // poll_creator_jid is not part of the AAD; supply an empty placeholder. + let aad = build_aad(&poll_vote_addon_ctx(stanza_id, "", voter_jid)); let mut payload = Vec::with_capacity(plaintext.len() + GCM_TAG_SIZE); aes_256_gcm_encrypt(encryption_key, &iv, &aad, &plaintext, &mut payload) @@ -85,7 +84,34 @@ pub fn encrypt_poll_vote( Ok((payload, iv)) } +/// Encrypt a poll vote given the parent poll's `messageSecret`. Returns +/// `(payload_with_tag, iv)`. +pub fn encrypt_poll_vote_with_secret( + selected_option_hashes: &[Vec], + message_secret: &[u8], + stanza_id: &str, + poll_creator_jid: &str, + voter_jid: &str, +) -> Result<(Vec, [u8; GCM_IV_SIZE])> { + use prost::Message; + + let vote_msg = waproto::whatsapp::message::PollVoteMessage { + selected_options: selected_option_hashes.to_vec(), + }; + let mut plaintext = Vec::new(); + vote_msg.encode(&mut plaintext)?; + + encrypt_addon( + &plaintext, + message_secret, + &poll_vote_addon_ctx(stanza_id, poll_creator_jid, voter_jid), + ) +} + /// Returns the selected option hashes (each 32 bytes). +/// +/// Kept for backwards compatibility with callers that pre-derived the key. +/// New code should call [`decrypt_poll_vote_with_secret`]. pub fn decrypt_poll_vote( enc_payload: &[u8], iv: &[u8], @@ -93,12 +119,12 @@ pub fn decrypt_poll_vote( stanza_id: &str, voter_jid: &str, ) -> Result>> { + use crate::libsignal::crypto::aes_256_gcm_decrypt; use prost::Message as _; let nonce: &[u8; GCM_IV_SIZE] = iv .try_into() .map_err(|_| anyhow!("Invalid IV size: expected {GCM_IV_SIZE}, got {}", iv.len()))?; - if enc_payload.len() < GCM_TAG_SIZE { return Err(anyhow!( "Encrypted payload too short: need at least {GCM_TAG_SIZE} bytes for tag, got {}", @@ -106,7 +132,8 @@ pub fn decrypt_poll_vote( )); } - let aad = build_vote_aad(stanza_id, voter_jid); + // poll_creator_jid is not part of the AAD; supply an empty placeholder. + let aad = build_aad(&poll_vote_addon_ctx(stanza_id, "", voter_jid)); let mut plaintext = Vec::with_capacity(enc_payload.len().saturating_sub(GCM_TAG_SIZE)); aes_256_gcm_decrypt(encryption_key, nonce, &aad, enc_payload, &mut plaintext) @@ -116,6 +143,28 @@ pub fn decrypt_poll_vote( Ok(vote_msg.selected_options) } +/// Decrypt a poll vote given the poll's `messageSecret` directly. Preferred +/// over the legacy two-step path that splits derive+decrypt. +pub fn decrypt_poll_vote_with_secret( + enc_payload: &[u8], + iv: &[u8], + message_secret: &[u8], + stanza_id: &str, + poll_creator_jid: &str, + voter_jid: &str, +) -> Result>> { + use prost::Message as _; + + let plaintext = decrypt_addon( + enc_payload, + iv, + message_secret, + &poll_vote_addon_ctx(stanza_id, poll_creator_jid, voter_jid), + )?; + let vote_msg = waproto::whatsapp::message::PollVoteMessage::decode(&plaintext[..])?; + Ok(vote_msg.selected_options) +} + #[cfg(test)] mod tests { use super::*; @@ -130,35 +179,6 @@ mod tests { assert_eq!(h1.len(), 32); } - #[test] - fn vote_key_derivation_deterministic() { - let secret = [0xABu8; 32]; - let k1 = derive_vote_encryption_key( - &secret, - "msg123", - "creator@s.whatsapp.net", - "voter@s.whatsapp.net", - ) - .unwrap(); - let k2 = derive_vote_encryption_key( - &secret, - "msg123", - "creator@s.whatsapp.net", - "voter@s.whatsapp.net", - ) - .unwrap(); - assert_eq!(k1, k2); - - let k3 = derive_vote_encryption_key( - &secret, - "msg123", - "creator@s.whatsapp.net", - "other@s.whatsapp.net", - ) - .unwrap(); - assert_ne!(k1, k3); - } - #[test] fn vote_encrypt_decrypt_roundtrip() { let secret = [0xCDu8; 32]; @@ -166,80 +186,79 @@ mod tests { let creator = "creator@s.whatsapp.net"; let voter = "voter@s.whatsapp.net"; - let key = derive_vote_encryption_key(&secret, stanza_id, creator, voter).unwrap(); - - let option_hashes = vec![ + let hashes = vec![ compute_option_hash("Yes").to_vec(), compute_option_hash("No").to_vec(), ]; - let (enc_payload, iv) = encrypt_poll_vote(&option_hashes, &key, stanza_id, voter).unwrap(); - let decrypted = decrypt_poll_vote(&enc_payload, &iv, &key, stanza_id, voter).unwrap(); - - assert_eq!(decrypted, option_hashes); + let (enc, iv) = + encrypt_poll_vote_with_secret(&hashes, &secret, stanza_id, creator, voter).unwrap(); + let out = + decrypt_poll_vote_with_secret(&enc, &iv, &secret, stanza_id, creator, voter).unwrap(); + assert_eq!(out, hashes); } #[test] - fn vote_decrypt_wrong_key_fails() { + fn legacy_decrypt_path_still_works() { let secret = [0xCDu8; 32]; let stanza_id = "3EB0ABCD1234"; let creator = "creator@s.whatsapp.net"; let voter = "voter@s.whatsapp.net"; + let hashes = vec![compute_option_hash("Yes").to_vec()]; + let (enc, iv) = + encrypt_poll_vote_with_secret(&hashes, &secret, stanza_id, creator, voter).unwrap(); + let key = derive_vote_encryption_key(&secret, stanza_id, creator, voter).unwrap(); - let option_hashes = vec![compute_option_hash("Yes").to_vec()]; - let (enc_payload, iv) = encrypt_poll_vote(&option_hashes, &key, stanza_id, voter).unwrap(); + let out = decrypt_poll_vote(&enc, &iv, &key, stanza_id, voter).unwrap(); + assert_eq!(out, hashes); + } + + #[test] + fn wrong_voter_fails() { + let secret = [0xEFu8; 32]; + let (enc, iv) = encrypt_poll_vote_with_secret( + &[compute_option_hash("Yes").to_vec()], + &secret, + "id", + "c@s.whatsapp.net", + "v@s.whatsapp.net", + ) + .unwrap(); - let wrong_key = - derive_vote_encryption_key(&secret, stanza_id, creator, "wrong@s.whatsapp.net") - .unwrap(); assert!( - decrypt_poll_vote( - &enc_payload, + decrypt_poll_vote_with_secret( + &enc, &iv, - &wrong_key, - stanza_id, + &secret, + "id", + "c@s.whatsapp.net", "wrong@s.whatsapp.net" ) .is_err() ); } - #[test] - fn vote_decrypt_wrong_aad_fails() { - let secret = [0xCDu8; 32]; - let stanza_id = "3EB0ABCD1234"; - let creator = "creator@s.whatsapp.net"; - let voter = "voter@s.whatsapp.net"; - - let key = derive_vote_encryption_key(&secret, stanza_id, creator, voter).unwrap(); - let option_hashes = vec![compute_option_hash("Yes").to_vec()]; - let (enc_payload, iv) = encrypt_poll_vote(&option_hashes, &key, stanza_id, voter).unwrap(); - - assert!(decrypt_poll_vote(&enc_payload, &iv, &key, "wrong_stanza", voter).is_err()); - } - #[test] fn empty_vote_roundtrip() { let secret = [0xEFu8; 32]; - let key = derive_vote_encryption_key(&secret, "id", "c@s.whatsapp.net", "v@s.whatsapp.net") - .unwrap(); - - let (enc, iv) = encrypt_poll_vote(&[], &key, "id", "v@s.whatsapp.net").unwrap(); - let dec = decrypt_poll_vote(&enc, &iv, &key, "id", "v@s.whatsapp.net").unwrap(); - assert!(dec.is_empty()); - } - - #[test] - fn vote_decrypt_invalid_iv_length_fails() { - let secret = [0xCDu8; 32]; - let key = derive_vote_encryption_key(&secret, "id", "c@s.whatsapp.net", "v@s.whatsapp.net") - .unwrap(); - let option_hashes = vec![compute_option_hash("Yes").to_vec()]; - let (enc_payload, _iv) = - encrypt_poll_vote(&option_hashes, &key, "id", "v@s.whatsapp.net").unwrap(); - - let bad_iv = [0u8; 8]; // wrong length - assert!(decrypt_poll_vote(&enc_payload, &bad_iv, &key, "id", "v@s.whatsapp.net").is_err()); + let (enc, iv) = encrypt_poll_vote_with_secret( + &[], + &secret, + "id", + "c@s.whatsapp.net", + "v@s.whatsapp.net", + ) + .unwrap(); + let out = decrypt_poll_vote_with_secret( + &enc, + &iv, + &secret, + "id", + "c@s.whatsapp.net", + "v@s.whatsapp.net", + ) + .unwrap(); + assert!(out.is_empty()); } } diff --git a/wacore/src/secret_enc_addon.rs b/wacore/src/secret_enc_addon.rs new file mode 100644 index 000000000..8c12d4cf4 --- /dev/null +++ b/wacore/src/secret_enc_addon.rs @@ -0,0 +1,383 @@ +//! Secret-encrypted addon envelope (HKDF-SHA256 + AES-256-GCM). +//! +//! Mirrors `WAWebAddonEncryption.decryptAddOn` / `WAUseCaseSecret.createUseCaseSecret` +//! from the captured WA Web bundle. Used for the addon family that piggybacks on a +//! parent message's `messageContextInfo.messageSecret`: +//! +//! - Poll Vote, Poll Edit, Poll Add Option +//! - Event Response, Event Edit +//! - Enc Reaction, Enc Comment +//! - **Message Edit** (added by WA in 2026) +//! +//! Key derivation (per `WAUseCaseSecret.createUseCaseSecret`): +//! +//! ```text +//! info = stanzaId || parentMsgOriginalSender || modificationSender || +//! key = HKDF-SHA256(salt = zeros[32], ikm = messageSecret, info, L = 32) +//! ``` +//! +//! AAD (per `WAWebAddonEncryption.js` function `g`): +//! +//! - `PollVote` / `EventResponse` → `stanzaId || 0x00 || modificationSenderJid` +//! - everything else (edits, reactions, comments, poll add option) → empty + +use anyhow::{Result, anyhow}; +use hkdf::Hkdf; +use sha2::Sha256; + +use crate::libsignal::crypto::{aes_256_gcm_decrypt, aes_256_gcm_encrypt}; + +const GCM_IV_SIZE: usize = 12; +const GCM_TAG_SIZE: usize = 16; +const KEY_SIZE: usize = 32; + +/// Use-case literal that goes into the HKDF `info` buffer. +/// +/// Source of truth: `docs/captured-js/WA/Use/CaseSecret.js` `UseCaseSecretModificationType`. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ModificationType { + PollVote, + EncReaction, + EncComment, + ReportToken, + EventResponse, + EventEdit, + PollEdit, + PollAddOption, + MessageEdit, +} + +impl ModificationType { + pub const fn as_str(self) -> &'static str { + match self { + Self::PollVote => "Poll Vote", + Self::EncReaction => "Enc Reaction", + Self::EncComment => "Enc Comment", + Self::ReportToken => "Report Token", + Self::EventResponse => "Event Response", + Self::EventEdit => "Event Edit", + Self::PollEdit => "Poll Edit", + Self::PollAddOption => "Poll Add Option", + Self::MessageEdit => "Message Edit", + } + } + + /// Only PollVote and EventResponse bind the stanza+sender into AAD. + /// Everything else (edits, reactions, comments, add-option) uses empty AAD. + pub const fn aad_mode(self) -> AadMode { + match self { + Self::PollVote | Self::EventResponse => AadMode::StanzaAndSender, + _ => AadMode::Empty, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AadMode { + /// AAD = `stanzaId || 0x00 || modificationSenderJid` + StanzaAndSender, + /// AAD = empty buffer + Empty, +} + +/// Inputs threaded through every addon (en|de)crypt. +/// +/// `parent_msg_original_sender` is the sender of the *targeted* message (the poll +/// creator, the event creator, the edited message's original author, etc.). +/// `modification_sender` is the user performing the addon action (voter, editor, +/// reactor, ...). +#[derive(Debug, Clone, Copy)] +pub struct AddonContext<'a> { + pub stanza_id: &'a str, + pub parent_msg_original_sender: &'a str, + pub modification_sender: &'a str, + pub modification_type: ModificationType, +} + +/// HKDF derivation, matching `WAUseCaseSecret.createUseCaseSecret`. +pub fn derive_use_case_secret( + message_secret: &[u8], + ctx: &AddonContext<'_>, +) -> Result<[u8; KEY_SIZE]> { + if message_secret.len() != KEY_SIZE { + return Err(anyhow!( + "Invalid messageSecret size: expected {KEY_SIZE}, got {}", + message_secret.len() + )); + } + + let mut info = Vec::with_capacity( + ctx.stanza_id.len() + + ctx.parent_msg_original_sender.len() + + ctx.modification_sender.len() + + ctx.modification_type.as_str().len(), + ); + info.extend_from_slice(ctx.stanza_id.as_bytes()); + info.extend_from_slice(ctx.parent_msg_original_sender.as_bytes()); + info.extend_from_slice(ctx.modification_sender.as_bytes()); + info.extend_from_slice(ctx.modification_type.as_str().as_bytes()); + + let hk = Hkdf::::new(None, message_secret); + let mut key = [0u8; KEY_SIZE]; + hk.expand(&info, &mut key) + .map_err(|e| anyhow!("HKDF expand failed: {e}"))?; + Ok(key) +} + +pub(crate) fn build_aad(ctx: &AddonContext<'_>) -> Vec { + match ctx.modification_type.aad_mode() { + AadMode::Empty => Vec::new(), + AadMode::StanzaAndSender => { + let mut aad = + Vec::with_capacity(ctx.stanza_id.len() + 1 + ctx.modification_sender.len()); + aad.extend_from_slice(ctx.stanza_id.as_bytes()); + aad.push(0); + aad.extend_from_slice(ctx.modification_sender.as_bytes()); + aad + } + } +} + +/// AES-256-GCM encrypt of `plaintext` with addon use-case key and AAD. +/// Returns `(payload_with_tag, iv)`. +pub fn encrypt_addon( + plaintext: &[u8], + message_secret: &[u8], + ctx: &AddonContext<'_>, +) -> Result<(Vec, [u8; GCM_IV_SIZE])> { + use rand::Rng; + + let key = derive_use_case_secret(message_secret, ctx)?; + let aad = build_aad(ctx); + + let mut iv = [0u8; GCM_IV_SIZE]; + rand::make_rng::().fill_bytes(&mut iv); + + let mut payload = Vec::with_capacity(plaintext.len() + GCM_TAG_SIZE); + aes_256_gcm_encrypt(&key, &iv, &aad, plaintext, &mut payload) + .map_err(|e| anyhow!("AES-GCM encrypt failed: {e}"))?; + + Ok((payload, iv)) +} + +/// AES-256-GCM decrypt with key derived from `message_secret` + `ctx`. +/// +/// Returns the raw plaintext; the caller is responsible for decoding the +/// underlying protobuf (the inner shape depends on `modification_type`). +pub fn decrypt_addon( + enc_payload: &[u8], + iv: &[u8], + message_secret: &[u8], + ctx: &AddonContext<'_>, +) -> Result> { + let nonce: &[u8; GCM_IV_SIZE] = iv + .try_into() + .map_err(|_| anyhow!("Invalid IV size: expected {GCM_IV_SIZE}, got {}", iv.len()))?; + if enc_payload.len() < GCM_TAG_SIZE { + return Err(anyhow!( + "Encrypted payload too short: need at least {GCM_TAG_SIZE} bytes for tag, got {}", + enc_payload.len() + )); + } + + let key = derive_use_case_secret(message_secret, ctx)?; + let aad = build_aad(ctx); + + let mut plaintext = Vec::with_capacity(enc_payload.len().saturating_sub(GCM_TAG_SIZE)); + aes_256_gcm_decrypt(&key, nonce, &aad, enc_payload, &mut plaintext) + .map_err(|_| anyhow!("Addon GCM tag verification failed"))?; + Ok(plaintext) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn ctx<'a>( + modification_type: ModificationType, + stanza_id: &'a str, + parent: &'a str, + sender: &'a str, + ) -> AddonContext<'a> { + AddonContext { + stanza_id, + parent_msg_original_sender: parent, + modification_sender: sender, + modification_type, + } + } + + #[test] + fn use_case_literals_match_wa_web() { + // Mirror WAUseCaseSecret enum exactly. + assert_eq!(ModificationType::PollVote.as_str(), "Poll Vote"); + assert_eq!(ModificationType::EncReaction.as_str(), "Enc Reaction"); + assert_eq!(ModificationType::EncComment.as_str(), "Enc Comment"); + assert_eq!(ModificationType::ReportToken.as_str(), "Report Token"); + assert_eq!(ModificationType::EventResponse.as_str(), "Event Response"); + assert_eq!(ModificationType::EventEdit.as_str(), "Event Edit"); + assert_eq!(ModificationType::PollEdit.as_str(), "Poll Edit"); + assert_eq!(ModificationType::PollAddOption.as_str(), "Poll Add Option"); + assert_eq!(ModificationType::MessageEdit.as_str(), "Message Edit"); + } + + #[test] + fn aad_mode_matches_wa_web_function_g() { + // Per WAWebAddonEncryption.js function `g`: AAD only for PollVote and EventResponse. + assert_eq!( + ModificationType::PollVote.aad_mode(), + AadMode::StanzaAndSender + ); + assert_eq!( + ModificationType::EventResponse.aad_mode(), + AadMode::StanzaAndSender + ); + for mt in [ + ModificationType::EncReaction, + ModificationType::EncComment, + ModificationType::ReportToken, + ModificationType::EventEdit, + ModificationType::PollEdit, + ModificationType::PollAddOption, + ModificationType::MessageEdit, + ] { + assert_eq!(mt.aad_mode(), AadMode::Empty, "{mt:?} must use empty AAD"); + } + } + + #[test] + fn derive_key_invalid_secret_size() { + let bad = [0u8; 16]; + let c = ctx(ModificationType::MessageEdit, "id", "a", "b"); + assert!(derive_use_case_secret(&bad, &c).is_err()); + } + + #[test] + fn derive_key_changes_with_each_input() { + let secret = [0xAAu8; 32]; + let base = ctx(ModificationType::MessageEdit, "id1", "alice@s", "bob@s"); + let k0 = derive_use_case_secret(&secret, &base).unwrap(); + + // Different stanza id + let k1 = derive_use_case_secret( + &secret, + &ctx(ModificationType::MessageEdit, "id2", "alice@s", "bob@s"), + ) + .unwrap(); + // Different parent + let k2 = derive_use_case_secret( + &secret, + &ctx(ModificationType::MessageEdit, "id1", "carol@s", "bob@s"), + ) + .unwrap(); + // Different sender + let k3 = derive_use_case_secret( + &secret, + &ctx(ModificationType::MessageEdit, "id1", "alice@s", "carol@s"), + ) + .unwrap(); + // Different use-case literal + let k4 = derive_use_case_secret( + &secret, + &ctx(ModificationType::PollEdit, "id1", "alice@s", "bob@s"), + ) + .unwrap(); + + for other in [k1, k2, k3, k4] { + assert_ne!(k0, other); + } + } + + #[test] + fn encrypt_decrypt_message_edit_roundtrip() { + let secret = [0x11u8; 32]; + let c = ctx( + ModificationType::MessageEdit, + "AC1234567890ABCDEF", + "5511999999999@s.whatsapp.net", + "5511999999999@s.whatsapp.net", + ); + let pt = b"hello world plaintext"; + let (ct, iv) = encrypt_addon(pt, &secret, &c).unwrap(); + let out = decrypt_addon(&ct, &iv, &secret, &c).unwrap(); + assert_eq!(out, pt); + } + + #[test] + fn encrypt_decrypt_poll_vote_roundtrip_uses_aad() { + let secret = [0x22u8; 32]; + let c = ctx( + ModificationType::PollVote, + "stanza", + "creator@s.whatsapp.net", + "voter@s.whatsapp.net", + ); + let pt = b"vote payload"; + let (ct, iv) = encrypt_addon(pt, &secret, &c).unwrap(); + + // Tamper modification_sender — AAD differs → GCM tag fails. + let bad = ctx( + ModificationType::PollVote, + "stanza", + "creator@s.whatsapp.net", + "attacker@s.whatsapp.net", + ); + assert!(decrypt_addon(&ct, &iv, &secret, &bad).is_err()); + + let out = decrypt_addon(&ct, &iv, &secret, &c).unwrap(); + assert_eq!(out, pt); + } + + #[test] + fn aad_mismatch_under_same_key_fails_decrypt() { + // Prove the AAD branch is actually load-bearing in the GCM check, not + // just an unused field. Bypass derive (which would also change with + // modification_type) and hand-roll an encrypt with one AAD shape, then + // try to decrypt with the other under the same key. + use crate::libsignal::crypto::{aes_256_gcm_decrypt, aes_256_gcm_encrypt}; + + let secret = [0x33u8; 32]; + let pv_ctx = ctx(ModificationType::PollVote, "stanza", "p@s", "s@s"); + let me_ctx = ctx(ModificationType::MessageEdit, "stanza", "p@s", "s@s"); + let aad_pv = build_aad(&pv_ctx); + let aad_me = build_aad(&me_ctx); + assert!(!aad_pv.is_empty(), "PollVote AAD must bind stanza+sender"); + assert!(aad_me.is_empty(), "MessageEdit AAD must be empty"); + + // Pick one key (PollVote's) and use it for both encrypt and decrypt + // attempts so the only thing that varies is the AAD. + let key = derive_use_case_secret(&secret, &pv_ctx).unwrap(); + let iv = [0u8; GCM_IV_SIZE]; + let plaintext = b"vote payload"; + + let mut ct = Vec::with_capacity(plaintext.len() + GCM_TAG_SIZE); + aes_256_gcm_encrypt(&key, &iv, &aad_pv, plaintext, &mut ct).unwrap(); + + // Same key, PollVote AAD → ok. + let mut out = Vec::new(); + aes_256_gcm_decrypt(&key, &iv, &aad_pv, &ct, &mut out).unwrap(); + assert_eq!(out, plaintext); + + // Same key, MessageEdit AAD (empty) → must fail despite identical key. + let mut out2 = Vec::new(); + assert!(aes_256_gcm_decrypt(&key, &iv, &aad_me, &ct, &mut out2).is_err()); + } + + #[test] + fn decrypt_invalid_iv_size() { + let secret = [0x44u8; 32]; + let c = ctx(ModificationType::MessageEdit, "id", "p@s", "s@s"); + let (ct, _iv) = encrypt_addon(b"x", &secret, &c).unwrap(); + let bad_iv = [0u8; 8]; + assert!(decrypt_addon(&ct, &bad_iv, &secret, &c).is_err()); + } + + #[test] + fn decrypt_payload_too_short() { + let secret = [0x55u8; 32]; + let c = ctx(ModificationType::MessageEdit, "id", "p@s", "s@s"); + let iv = [0u8; GCM_IV_SIZE]; + let too_short = vec![0u8; GCM_TAG_SIZE - 1]; + assert!(decrypt_addon(&too_short, &iv, &secret, &c).is_err()); + } +}