diff --git a/src/features/presence.rs b/src/features/presence.rs index dfba3a6d5..94411800a 100644 --- a/src/features/presence.rs +++ b/src/features/presence.rs @@ -140,6 +140,27 @@ impl<'a> Presence<'a> { Ok(()) } + /// Re-subscribe presence if the JID has an active subscription. + /// Does not modify the tracking set. + pub(crate) async fn re_subscribe_when_active(&self, jid: &Jid) -> Result<(), anyhow::Error> { + if !self + .client + .presence_subscriptions + .lock() + .await + .contains(jid) + { + return Ok(()); + } + + let node = self.build_subscription_node(jid).await; + self.client + .send_node(node) + .await + .map_err(anyhow::Error::from)?; + Ok(()) + } + /// Unsubscribe from a contact's presence updates. /// /// Sends a `` stanza to the target JID. @@ -204,14 +225,7 @@ impl Client { return; } - // Check membership before re-subscribing — a concurrent unsubscribe() - // call may have removed this JID while we were iterating. - if !self.presence_subscriptions.lock().await.contains(&jid) { - debug!("Skipping re-subscribe for {jid}: unsubscribed during iteration"); - continue; - } - - if let Err(err) = self.presence().subscribe(&jid).await { + if let Err(err) = self.presence().re_subscribe_when_active(&jid).await { warn!("Failed to re-subscribe to presence for {jid}: {err:?}"); } } diff --git a/src/handlers/notification.rs b/src/handlers/notification.rs index 8aadf415f..7e77583fb 100644 --- a/src/handlers/notification.rs +++ b/src/handlers/notification.rs @@ -546,7 +546,9 @@ async fn handle_privacy_token_notification(client: &Arc, node: &Node) { use wacore::iq::tctoken::parse_privacy_token_notification; use wacore::store::traits::TcTokenEntry; - // Resolve the sender to a LID JID for storage. + let from_jid = node.attrs().optional_jid("from"); + + // Resolve the sender to a LID key for storage. // WA Web uses `sender_lid` attr if present, otherwise resolves from `from`. let sender_lid = node .attrs() @@ -557,7 +559,7 @@ async fn handle_privacy_token_notification(client: &Arc, node: &Node) { Some(lid) if !lid.is_empty() => lid, _ => { // Fall back to resolving from the `from` JID via LID-PN cache - let from_jid = match node.attrs().optional_jid("from") { + let from = match &from_jid { Some(jid) => jid, None => { warn!(target: "Client/TcToken", "privacy_token notification missing 'from' attribute"); @@ -565,19 +567,19 @@ async fn handle_privacy_token_notification(client: &Arc, node: &Node) { } }; - if from_jid.is_lid() { - from_jid.user.clone() + if from.is_lid() { + from.user.clone() } else { // Try to resolve phone number to LID - match client.lid_pn_cache.get_current_lid(&from_jid.user).await { + match client.lid_pn_cache.get_current_lid(&from.user).await { Some(lid) => lid, None => { debug!( target: "Client/TcToken", "Cannot resolve LID for privacy_token sender {}, storing under PN", - from_jid + from ); - from_jid.user.clone() + from.user.clone() } } } @@ -599,10 +601,26 @@ async fn handle_privacy_token_notification(client: &Arc, node: &Node) { } let backend = client.persistence_manager.backend(); + let mut token_stored = false; for received in &received_tokens { match backend.get_tc_token(&sender_lid).await { Ok(Some(existing)) => { + // Skip if token bytes are identical and timestamp hasn't advanced + if existing.token == received.token { + if received.timestamp > existing.token_timestamp { + // Same bytes but newer timestamp — refresh to prevent premature pruning + let refreshed = TcTokenEntry { + token_timestamp: received.timestamp, + ..existing + }; + if let Err(e) = backend.put_tc_token(&sender_lid, &refreshed).await { + warn!(target: "Client/TcToken", "Failed to refresh tc_token timestamp for {}: {e}", sender_lid); + } + } + continue; + } + // Timestamp monotonicity guard: only store if incoming >= existing if received.timestamp < existing.token_timestamp { debug!( @@ -624,6 +642,7 @@ async fn handle_privacy_token_notification(client: &Arc, node: &Node) { warn!(target: "Client/TcToken", "Failed to update tc_token for {}: {e}", sender_lid); } else { debug!(target: "Client/TcToken", "Updated tc_token for {} (t={})", sender_lid, received.timestamp); + token_stored = true; } } Ok(None) => { @@ -638,6 +657,7 @@ async fn handle_privacy_token_notification(client: &Arc, node: &Node) { warn!(target: "Client/TcToken", "Failed to store tc_token for {}: {e}", sender_lid); } else { debug!(target: "Client/TcToken", "Stored new tc_token for {} (t={})", sender_lid, received.timestamp); + token_stored = true; } } Err(e) => { @@ -645,6 +665,14 @@ async fn handle_privacy_token_notification(client: &Arc, node: &Node) { } } } + + // Re-subscribe presence with the updated token. + if token_stored + && let Some(from) = &from_jid + && let Err(e) = client.presence().re_subscribe_when_active(from).await + { + debug!(target: "Client/TcToken", "Failed to re-subscribe presence for {from}: {e}"); + } } /// Handle business notification (WhatsApp Web: `WAWebHandleBusinessNotification`). diff --git a/tests/e2e/src/lib.rs b/tests/e2e/src/lib.rs index 0b556980d..1ffda52bf 100644 --- a/tests/e2e/src/lib.rs +++ b/tests/e2e/src/lib.rs @@ -1,5 +1,6 @@ use std::sync::Arc; +use wacore::store::traits::TcTokenEntry; use wacore::types::events::{Event, EventHandler}; use whatsapp_rust::Jid; use whatsapp_rust::bot::Bot; @@ -172,6 +173,46 @@ impl TestClient { .to_non_ad() } + /// Get the storage key used for this client's tcToken entries. + /// + /// Notification handling stores tcTokens under the sender's LID when it is + /// available, otherwise it falls back to the phone-number user part. + pub async fn tc_token_key(&self) -> anyhow::Result { + if let Some(lid) = self.client.get_lid().await { + return Ok(lid.user); + } + + self.client + .get_pn() + .await + .map(|jid| jid.user) + .ok_or_else(|| anyhow::anyhow!("Client should have a JID after connect")) + } + + /// Wait until a tcToken entry exists for the given storage key. + pub async fn wait_for_tc_token( + &self, + jid_key: &str, + timeout_secs: u64, + ) -> anyhow::Result { + let deadline = tokio::time::Instant::now() + tokio::time::Duration::from_secs(timeout_secs); + + loop { + if let Some(entry) = self.client.tc_token().get(jid_key).await? { + return Ok(entry); + } + + if tokio::time::Instant::now() >= deadline { + return Err(anyhow::anyhow!( + "Timed out waiting for tc_token entry for {}", + jid_key + )); + } + + tokio::time::sleep(tokio::time::Duration::from_millis(50)).await; + } + } + // ── Event waiting ─────────────────────────────────────────────────────── /// Wait for an event matching the predicate, with a timeout in seconds. diff --git a/tests/e2e/tests/privacy_tokens.rs b/tests/e2e/tests/privacy_tokens.rs new file mode 100644 index 000000000..ceae6628b --- /dev/null +++ b/tests/e2e/tests/privacy_tokens.rs @@ -0,0 +1,312 @@ +use e2e_tests::{TestClient, send_and_expect_text, text_msg}; +use log::info; +use wacore::iq::tctoken::tc_token_expiration_cutoff; +use wacore::store::traits::TcTokenEntry; +use whatsapp_rust::{NodeFilter, SendOptions}; + +fn unique_push_name(prefix: &str) -> String { + format!("{}_{}", prefix, uuid::Uuid::new_v4()) +} + +#[tokio::test] +async fn test_tc_token_notification_stores_token_for_sender() -> anyhow::Result<()> { + let _ = env_logger::builder().is_test(true).try_init(); + + let client_a = TestClient::connect("e2e_tctok_store_a").await?; + let mut client_b = TestClient::connect("e2e_tctok_store_b").await?; + + let jid_b = client_b.jid().await; + send_and_expect_text(&client_a.client, &mut client_b, &jid_b, "seed tc token", 30).await?; + + let key_a = client_a.tc_token_key().await?; + let entry = client_b.wait_for_tc_token(&key_a, 10).await?; + assert!(!entry.token.is_empty(), "tc token should contain bytes"); + assert!( + entry.token_timestamp > 0, + "tc token timestamp should be populated" + ); + assert_eq!( + entry.sender_timestamp, None, + "recipient-side storage should not set sender_timestamp yet" + ); + info!("B stored tc token for key {}", key_a); + + client_a.disconnect().await; + client_b.disconnect().await; + Ok(()) +} + +#[tokio::test] +async fn test_issue_tokens_api_delivers_notification_and_updates_index() -> anyhow::Result<()> { + let _ = env_logger::builder().is_test(true).try_init(); + + let client_a = TestClient::connect("e2e_tctok_issue_a").await?; + let client_b = TestClient::connect("e2e_tctok_issue_b").await?; + + let jid_b_lid = client_b + .client + .get_lid() + .await + .expect("B should have LID after connect"); + let issued = client_a + .client + .tc_token() + .issue_tokens(std::slice::from_ref(&jid_b_lid)) + .await?; + info!("issue_tokens returned {} token(s)", issued.len()); + + let key_a = client_a.tc_token_key().await?; + let stored = client_b.wait_for_tc_token(&key_a, 10).await?; + assert!( + !stored.token.is_empty(), + "issued tc token should be stored on recipient" + ); + + let all_jids = client_b.client.tc_token().get_all_jids().await?; + assert!( + all_jids.contains(&key_a), + "tc token index should include the sender key after explicit issuance" + ); + + client_a.disconnect().await; + client_b.disconnect().await; + Ok(()) +} + +#[tokio::test] +async fn test_reply_to_restricted_contact_uses_received_tc_token() -> anyhow::Result<()> { + let _ = env_logger::builder().is_test(true).try_init(); + + let restricted_name = format!("restricted:{}", unique_push_name("e2e_tctok_reply_a")); + let mut client_a = TestClient::connect_as("e2e_tctok_reply_a", &restricted_name).await?; + let mut client_b = TestClient::connect("e2e_tctok_reply_b").await?; + + let jid_a = client_a.jid().await; + let jid_b = client_b.jid().await; + + send_and_expect_text( + &client_a.client, + &mut client_b, + &jid_b, + "seed restricted reply path", + 30, + ) + .await?; + + let key_a = client_a.tc_token_key().await?; + let initial_entry = client_b.wait_for_tc_token(&key_a, 10).await?; + assert_eq!(initial_entry.sender_timestamp, None); + + let reply = "reply to restricted A"; + client_b + .client + .send_message(jid_a.clone(), text_msg(reply)) + .await?; + client_a.wait_for_text(reply, 30).await?; + + let updated_entry = client_b.wait_for_tc_token(&key_a, 5).await?; + assert!( + updated_entry.sender_timestamp.is_some(), + "using a valid tc token should set sender_timestamp" + ); + info!("B replied successfully to restricted A using stored tc token"); + + client_a.disconnect().await; + client_b.disconnect().await; + Ok(()) +} + +#[tokio::test] +async fn test_first_message_to_restricted_contact_receives_463_nack() -> anyhow::Result<()> { + let _ = env_logger::builder().is_test(true).try_init(); + + let restricted_name = format!("restricted:{}", unique_push_name("e2e_tctok_463_a")); + let mut client_a = TestClient::connect_as("e2e_tctok_463_a", &restricted_name).await?; + let client_b = TestClient::connect("e2e_tctok_463_b").await?; + + let jid_a = client_a.jid().await; + let msg_id = format!("E2E463{}", uuid::Uuid::new_v4().simple()); + let waiter = client_b.client.wait_for_node( + NodeFilter::tag("ack") + .attr("id", msg_id.clone()) + .attr("class", "message") + .attr("from", jid_a.to_string()) + .attr("error", "463"), + ); + + let returned_id = client_b + .client + .send_message_with_options( + jid_a.clone(), + text_msg("first contact to restricted account"), + SendOptions { + message_id: Some(msg_id.clone()), + ..Default::default() + }, + ) + .await?; + assert_eq!( + returned_id, msg_id, + "send should preserve caller-provided message ID" + ); + + let ack = tokio::time::timeout(tokio::time::Duration::from_secs(15), waiter) + .await + .map_err(|_| anyhow::anyhow!("Timed out waiting for 463 nack"))? + .map_err(|_| anyhow::anyhow!("463 nack waiter was canceled"))?; + assert_eq!(ack.tag, "ack"); + assert_eq!( + ack.attrs.get("error").map(|v| v.to_string()), + Some("463".to_string()) + ); + + client_a + .assert_no_event( + 5, + |e| matches!(e, wacore::types::events::Event::Message(msg, _) if msg.conversation.as_deref() == Some("first contact to restricted account")), + "restricted recipient should not receive first-contact message without tcToken", + ) + .await?; + + client_a.disconnect().await; + client_b.disconnect().await; + Ok(()) +} + +#[tokio::test] +async fn test_tc_token_notification_reaches_all_connected_devices() -> anyhow::Result<()> { + let _ = env_logger::builder().is_test(true).try_init(); + + let restricted_name = format!("restricted:{}", unique_push_name("e2e_tctok_multi_a")); + let shared_b_name = unique_push_name("e2e_tctok_multi_b"); + + let client_a = TestClient::connect_as("e2e_tctok_multi_a", &restricted_name).await?; + let mut client_b1 = TestClient::connect_as("e2e_tctok_multi_b1", &shared_b_name).await?; + let client_b2 = TestClient::connect_as("e2e_tctok_multi_b2", &shared_b_name).await?; + + let phone_b1 = client_b1.client.get_pn().await.expect("B1 should have JID"); + let phone_b2 = client_b2.client.get_pn().await.expect("B2 should have JID"); + assert_eq!( + phone_b1.user, phone_b2.user, + "B devices should share a phone" + ); + assert_ne!( + phone_b1.device, phone_b2.device, + "B devices should have different device IDs" + ); + + let jid_b = client_b1.jid().await; + + send_and_expect_text( + &client_a.client, + &mut client_b1, + &jid_b, + "seed multi-device tc token", + 30, + ) + .await?; + + let key_a = client_a.tc_token_key().await?; + client_b1.wait_for_tc_token(&key_a, 10).await?; + client_b2.wait_for_tc_token(&key_a, 10).await?; + info!("Both connected B devices stored A's tc token"); + + client_a.disconnect().await; + client_b1.disconnect().await; + client_b2.disconnect().await; + Ok(()) +} + +#[tokio::test] +async fn test_tc_token_survives_reconnect() -> anyhow::Result<()> { + let _ = env_logger::builder().is_test(true).try_init(); + + let restricted_name = format!("restricted:{}", unique_push_name("e2e_tctok_recon_a")); + let mut client_a = TestClient::connect_as("e2e_tctok_recon_a", &restricted_name).await?; + let mut client_b = TestClient::connect("e2e_tctok_recon_b").await?; + + let jid_a = client_a.jid().await; + let jid_b = client_b.jid().await; + + send_and_expect_text( + &client_a.client, + &mut client_b, + &jid_b, + "seed reconnect tc token", + 30, + ) + .await?; + + let key_a = client_a.tc_token_key().await?; + let initial_entry = client_b.wait_for_tc_token(&key_a, 10).await?; + + client_b.reconnect_and_wait().await?; + + let after_reconnect = client_b.wait_for_tc_token(&key_a, 5).await?; + assert_eq!( + after_reconnect.token, initial_entry.token, + "tc token bytes should survive reconnect" + ); + assert_eq!( + after_reconnect.token_timestamp, initial_entry.token_timestamp, + "tc token timestamp should survive reconnect" + ); + + let reply = "reply after reconnect"; + client_b + .client + .send_message(jid_a.clone(), text_msg(reply)) + .await?; + client_a.wait_for_text(reply, 30).await?; + info!("Stored tc token survived reconnect and still works for replies"); + + client_a.disconnect().await; + client_b.disconnect().await; + Ok(()) +} + +#[tokio::test] +async fn test_prune_expired_tc_tokens_removes_only_stale_entries() -> anyhow::Result<()> { + let _ = env_logger::builder().is_test(true).try_init(); + + let client = TestClient::connect("e2e_tctok_prune").await?; + let backend = client.client.persistence_manager().backend(); + let cutoff = tc_token_expiration_cutoff(); + let expired_key = format!("expired_{}", uuid::Uuid::new_v4()); + let fresh_key = format!("fresh_{}", uuid::Uuid::new_v4()); + + backend + .put_tc_token( + &expired_key, + &TcTokenEntry { + token: vec![0x01], + token_timestamp: cutoff - 1, + sender_timestamp: None, + }, + ) + .await?; + backend + .put_tc_token( + &fresh_key, + &TcTokenEntry { + token: vec![0x02], + token_timestamp: cutoff, + sender_timestamp: Some(cutoff), + }, + ) + .await?; + + let deleted = client.client.tc_token().prune_expired().await?; + assert_eq!(deleted, 1, "exactly one expired tc token should be pruned"); + assert!( + client.client.tc_token().get(&expired_key).await?.is_none(), + "expired tc token should be removed" + ); + assert!( + client.client.tc_token().get(&fresh_key).await?.is_some(), + "fresh tc token should be preserved" + ); + + client.disconnect().await; + Ok(()) +}