diff --git a/src/client.rs b/src/client.rs index 48074c0bc..a9e1ab72a 100644 --- a/src/client.rs +++ b/src/client.rs @@ -89,6 +89,8 @@ pub struct Client { pub(crate) unique_id: String, pub(crate) id_counter: Arc, + pub(crate) unified_session: crate::unified_session::UnifiedSessionManager, + /// Per-device session locks for Signal protocol operations. /// Prevents race conditions when multiple messages from the same sender /// are processed concurrently across different chats. @@ -223,6 +225,7 @@ impl Client { response_waiters: Arc::new(Mutex::new(HashMap::new())), unique_id: format!("{}.{}", unique_id_bytes[0], unique_id_bytes[1]), id_counter: Arc::new(AtomicU64::new(0)), + unified_session: crate::unified_session::UnifiedSessionManager::new(), session_locks: Cache::builder() .time_to_live(Duration::from_secs(300)) // 5 minute TTL @@ -826,6 +829,8 @@ impl Client { *self.last_successful_connect.lock().await = Some(chrono::Utc::now()); self.auto_reconnect_errors.store(0, Ordering::Relaxed); + self.update_server_time_offset(node); + if let Some(lid_str) = node.attrs.get("lid") { if let Ok(lid) = lid_str.parse::() { let device_snapshot = self.persistence_manager.get_device_snapshot().await; @@ -879,6 +884,9 @@ impl Client { return; } + check_generation!(); + client_clone.send_unified_session().await; + // === Establish session with primary phone for PDO === // This must happen BEFORE we exit passive mode (before offline messages arrive). // PDO needs a session with device 0 to request decrypted content from our phone. @@ -1570,6 +1578,26 @@ impl Client { self.is_logged_in.load(Ordering::Relaxed) } + pub(crate) fn update_server_time_offset(&self, node: &wacore_binary::node::Node) { + self.unified_session.update_server_time_offset(node); + } + + pub(crate) async fn send_unified_session(&self) { + if !self.is_connected() { + debug!(target: "Client/UnifiedSession", "Skipping: not connected"); + return; + } + + let Some((node, _sequence)) = self.unified_session.prepare_send().await else { + return; + }; + + if let Err(e) = self.send_node(node).await { + debug!(target: "Client/UnifiedSession", "Send failed: {e}"); + self.unified_session.clear_last_sent().await; + } + } + /// Waits for the noise socket to be established. /// /// Returns `Ok(())` when the socket is ready, or `Err` on timeout. @@ -2693,4 +2721,222 @@ mod tests { info!("✅ test_immediate_session_does_not_wait_for_offline_sync passed"); } + + #[test] + fn test_unified_session_id_calculation() { + // Test the mathematical calculation of the unified session ID. + // Formula: (now_ms + server_offset_ms + 3_days_ms) % 7_days_ms + + const DAY_MS: i64 = 24 * 60 * 60 * 1000; + const WEEK_MS: i64 = 7 * DAY_MS; + const OFFSET_MS: i64 = 3 * DAY_MS; + + // Helper function matching the implementation + fn calculate_session_id(now_ms: i64, server_offset_ms: i64) -> i64 { + let adjusted_now = now_ms + server_offset_ms; + (adjusted_now + OFFSET_MS) % WEEK_MS + } + + // Test 1: Zero offset + let now_ms = 1706000000000_i64; // Some arbitrary timestamp + let id = calculate_session_id(now_ms, 0); + assert!( + (0..WEEK_MS).contains(&id), + "Session ID should be in [0, WEEK_MS)" + ); + + // Test 2: Positive server offset (server is ahead) + let id_with_positive_offset = calculate_session_id(now_ms, 5000); + assert!( + (0..WEEK_MS).contains(&id_with_positive_offset), + "Session ID should be in [0, WEEK_MS)" + ); + // The ID should be different from zero offset (unless wrap-around) + // Not testing exact value as it depends on the offset + + // Test 3: Negative server offset (server is behind) + let id_with_negative_offset = calculate_session_id(now_ms, -5000); + assert!( + (0..WEEK_MS).contains(&id_with_negative_offset), + "Session ID should be in [0, WEEK_MS)" + ); + + // Test 4: Verify modulo wrap-around + // If adjusted_now + OFFSET_MS >= WEEK_MS, it should wrap + let wrap_test_now = WEEK_MS - OFFSET_MS + 1000; // Should produce small result + let wrapped_id = calculate_session_id(wrap_test_now, 0); + assert_eq!(wrapped_id, 1000, "Should wrap around correctly"); + + // Test 5: Edge case - at exact boundary + let boundary_now = WEEK_MS - OFFSET_MS; + let boundary_id = calculate_session_id(boundary_now, 0); + assert_eq!(boundary_id, 0, "At exact boundary should be 0"); + } + + #[tokio::test] + async fn test_server_time_offset_extraction() { + use wacore_binary::builder::NodeBuilder; + + let backend = Arc::new( + crate::store::SqliteStore::new(":memory:") + .await + .expect("Failed to create in-memory backend for test"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("persistence manager should initialize"), + ); + let (client, _rx) = Client::new( + pm, + Arc::new(crate::transport::mock::MockTransportFactory::new()), + Arc::new(MockHttpClient), + None, + ) + .await; + + // Initially, offset should be 0 + assert_eq!( + client.unified_session.server_time_offset_ms(), + 0, + "Initial offset should be 0" + ); + + // Create a node with a 't' attribute + let server_time = chrono::Utc::now().timestamp() + 10; // Server is 10 seconds ahead + let node = NodeBuilder::new("success") + .attr("t", server_time.to_string()) + .build(); + + // Update the offset + client.update_server_time_offset(&node); + + // The offset should be approximately 10 * 1000 = 10000 ms + // Allow some tolerance for timing differences during the test + let offset = client.unified_session.server_time_offset_ms(); + assert!( + (offset - 10000).abs() < 1000, // Allow 1 second tolerance + "Offset should be approximately 10000ms, got {}", + offset + ); + + // Test with no 't' attribute - should not change offset + let node_no_t = NodeBuilder::new("success").build(); + client.update_server_time_offset(&node_no_t); + let offset_after = client.unified_session.server_time_offset_ms(); + assert!( + (offset_after - offset).abs() < 100, // Should be same (or very close) + "Offset should not change when 't' is missing" + ); + + // Test with invalid 't' attribute - should not change offset + let node_invalid = NodeBuilder::new("success") + .attr("t", "not_a_number") + .build(); + client.update_server_time_offset(&node_invalid); + let offset_after_invalid = client.unified_session.server_time_offset_ms(); + assert!( + (offset_after_invalid - offset).abs() < 100, + "Offset should not change when 't' is invalid" + ); + + // Test with negative/zero 't' - should not change offset + let node_zero = NodeBuilder::new("success").attr("t", "0").build(); + client.update_server_time_offset(&node_zero); + let offset_after_zero = client.unified_session.server_time_offset_ms(); + assert!( + (offset_after_zero - offset).abs() < 100, + "Offset should not change when 't' is 0" + ); + + info!("✅ test_server_time_offset_extraction passed"); + } + + #[tokio::test] + async fn test_unified_session_manager_integration() { + // Test the unified session manager through the client + + let backend = Arc::new( + crate::store::SqliteStore::new(":memory:") + .await + .expect("Failed to create in-memory backend for test"), + ); + let pm = Arc::new( + PersistenceManager::new(backend) + .await + .expect("persistence manager should initialize"), + ); + let (client, _rx) = Client::new( + pm, + Arc::new(crate::transport::mock::MockTransportFactory::new()), + Arc::new(MockHttpClient), + None, + ) + .await; + + // Initially, sequence should be 0 + assert_eq!( + client.unified_session.sequence(), + 0, + "Initial sequence should be 0" + ); + + // First prepare_send should succeed and return sequence 1 (pre-increment like WhatsApp Web) + let result = client.unified_session.prepare_send().await; + assert!(result.is_some(), "First send should succeed"); + let (node, seq) = result.unwrap(); + assert_eq!(node.tag, "ib", "Should be an IB stanza"); + assert_eq!(seq, 1, "First sequence should be 1 (pre-increment)"); + + // Sequence counter should now be 1 + assert_eq!(client.unified_session.sequence(), 1); + + // Second prepare_send should be blocked (duplicate prevention) + let result2 = client.unified_session.prepare_send().await; + assert!(result2.is_none(), "Duplicate should be prevented"); + + // Sequence should still be 1 (not incremented for duplicates) + assert_eq!(client.unified_session.sequence(), 1); + + // Clear last sent and try again - sequence resets on "new" session ID + client.unified_session.clear_last_sent().await; + let result3 = client.unified_session.prepare_send().await; + assert!(result3.is_some(), "Should succeed after clearing"); + let (_, seq3) = result3.unwrap(); + assert_eq!(seq3, 1, "Sequence resets when session ID changes"); + assert_eq!(client.unified_session.sequence(), 1); + + info!("✅ test_unified_session_manager_integration passed"); + } + + #[test] + fn test_unified_session_protocol_node() { + // Test the type-safe protocol node implementation + use wacore::ib::{IbStanza, UnifiedSession}; + use wacore::protocol::ProtocolNode; + + // Create a unified session + let session = UnifiedSession::new("123456789"); + assert_eq!(session.id, "123456789"); + assert_eq!(session.tag(), "unified_session"); + + // Convert to node + let node = session.into_node(); + assert_eq!(node.tag, "unified_session"); + assert_eq!(node.attrs.get("id"), Some(&"123456789".to_string())); + + // Create an IB stanza + let stanza = IbStanza::unified_session(UnifiedSession::new("987654321")); + assert_eq!(stanza.tag(), "ib"); + + // Convert to node and verify structure + let ib_node = stanza.into_node(); + assert_eq!(ib_node.tag, "ib"); + let children = ib_node.children().expect("IB stanza should have children"); + assert_eq!(children.len(), 1); + assert_eq!(children[0].tag, "unified_session"); + assert_eq!(children[0].attrs.get("id"), Some(&"987654321".to_string())); + + info!("✅ test_unified_session_protocol_node passed"); + } } diff --git a/src/features/presence.rs b/src/features/presence.rs index 930166ef7..d28879be1 100644 --- a/src/features/presence.rs +++ b/src/features/presence.rs @@ -51,6 +51,10 @@ impl<'a> Presence<'a> { )); } + if status == PresenceStatus::Available { + self.client.send_unified_session().await; + } + let presence_type = status.as_str(); let node = NodeBuilder::new("presence") diff --git a/src/lib.rs b/src/lib.rs index c2e7e894d..2e55e69b7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -29,6 +29,7 @@ pub mod pdo; pub mod prekeys; pub mod receipt; pub mod retry; +pub mod unified_session; pub mod appstate_sync; pub mod history_sync; diff --git a/src/pair.rs b/src/pair.rs index 6158edf29..00a2558eb 100644 --- a/src/pair.rs +++ b/src/pair.rs @@ -128,6 +128,8 @@ async fn handle_pair_success(client: &Arc, request_node: &Node, success_ // Clear pair code state if active *client.pair_code_state.lock().await = wacore::pair_code::PairCodeState::Completed; + client.update_server_time_offset(request_node); + let req_id = match request_node.attrs.get("id") { Some(id) => id.to_string(), None => { @@ -272,6 +274,11 @@ async fn handle_pair_success(client: &Arc, request_node: &Node, success_ return; } + let client_for_unified = client.clone(); + tokio::spawn(async move { + client_for_unified.send_unified_session().await; + }); + // --- START: FIX --- // Set the flag to trigger a full sync on the next successful connection. client diff --git a/src/unified_session.rs b/src/unified_session.rs new file mode 100644 index 000000000..eff99044a --- /dev/null +++ b/src/unified_session.rs @@ -0,0 +1,222 @@ +//! Unified session telemetry manager. +//! +//! Sends `` stanzas to match WhatsApp Web behavior. +//! Features: server time sync, duplicate prevention, sequence counter. + +use log::debug; +use std::sync::Arc; +use std::sync::atomic::{AtomicI64, AtomicU64, Ordering}; +use tokio::sync::Mutex; +use wacore::ib::{IbStanza, UnifiedSession}; +use wacore::protocol::ProtocolNode; +use wacore_binary::node::Node; + +/// Manager for unified session telemetry. +pub struct UnifiedSessionManager { + server_time_offset_ms: Arc, + last_sent_id: Arc>>, + sequence: Arc, +} + +impl Default for UnifiedSessionManager { + fn default() -> Self { + Self::new() + } +} + +impl UnifiedSessionManager { + pub fn new() -> Self { + Self { + server_time_offset_ms: Arc::new(AtomicI64::new(0)), + last_sent_id: Arc::new(Mutex::new(None)), + sequence: Arc::new(AtomicU64::new(0)), + } + } + + pub fn server_time_offset_ms(&self) -> i64 { + self.server_time_offset_ms.load(Ordering::Relaxed) + } + + pub fn sequence(&self) -> u64 { + self.sequence.load(Ordering::Relaxed) + } + + /// Update server time offset from node's `t` attribute (Unix timestamp in seconds). + pub fn update_server_time_offset(&self, node: &Node) { + if let Some(t_str) = node.attrs.get("t") + && let Ok(server_time) = t_str.parse::() + && server_time > 0 + { + let local_time = chrono::Utc::now().timestamp(); + let offset_ms = (server_time - local_time) * 1000; + self.server_time_offset_ms + .store(offset_ms, Ordering::Relaxed); + debug!(target: "UnifiedSession", "Server time offset: {}ms", offset_ms); + } + } + + pub fn calculate_session_id(&self) -> String { + let offset = self.server_time_offset_ms.load(Ordering::Relaxed); + UnifiedSession::calculate_id(offset) + } + + /// Prepare to send unified session. Returns None if duplicate (already sent this ID). + pub async fn prepare_send(&self) -> Option<(Node, u64)> { + let session_id = self.calculate_session_id(); + + { + let mut last_id = self.last_sent_id.lock().await; + if let Some(ref prev_id) = *last_id + && prev_id == &session_id + { + debug!(target: "UnifiedSession", "Skipping duplicate id={}", session_id); + return None; + } + + // Reset sequence when session ID changes (matches WhatsApp Web behavior) + if last_id.as_ref() != Some(&session_id) { + self.sequence.store(0, Ordering::Relaxed); + } + *last_id = Some(session_id.clone()); + } + + // Pre-increment to return 1 on first call (matches WhatsApp Web's ++$2) + let sequence = self.sequence.fetch_add(1, Ordering::Relaxed) + 1; + let stanza = IbStanza::unified_session(UnifiedSession::new(&session_id)); + let node = stanza.into_node(); + + debug!(target: "UnifiedSession", "Sending id={}, seq={}", session_id, sequence); + + Some((node, sequence)) + } + + /// Clear last sent ID to allow retry on failure. + pub async fn clear_last_sent(&self) { + *self.last_sent_id.lock().await = None; + } + + /// Reset state on disconnect (keeps sequence counter). + pub async fn reset(&self) { + self.server_time_offset_ms.store(0, Ordering::Relaxed); + *self.last_sent_id.lock().await = None; + } +} + +#[cfg(test)] +mod tests { + use super::*; + use wacore_binary::builder::NodeBuilder; + + #[test] + fn test_manager_default() { + let manager = UnifiedSessionManager::new(); + assert_eq!(manager.server_time_offset_ms(), 0); + assert_eq!(manager.sequence(), 0); + } + + #[test] + fn test_update_server_time_offset() { + let manager = UnifiedSessionManager::new(); + + let server_time = chrono::Utc::now().timestamp() + 10; + let node = NodeBuilder::new("success") + .attr("t", server_time.to_string()) + .build(); + + manager.update_server_time_offset(&node); + + let offset = manager.server_time_offset_ms(); + assert!( + (offset - 10000).abs() < 1000, + "Offset should be ~10000ms, got {}", + offset + ); + } + + #[test] + fn test_update_server_time_offset_invalid() { + let manager = UnifiedSessionManager::new(); + + let node = NodeBuilder::new("success").build(); + manager.update_server_time_offset(&node); + assert_eq!(manager.server_time_offset_ms(), 0); + + let node = NodeBuilder::new("success") + .attr("t", "not_a_number") + .build(); + manager.update_server_time_offset(&node); + assert_eq!(manager.server_time_offset_ms(), 0); + + let node = NodeBuilder::new("success").attr("t", "0").build(); + manager.update_server_time_offset(&node); + assert_eq!(manager.server_time_offset_ms(), 0); + } + + #[test] + fn test_calculate_session_id() { + let manager = UnifiedSessionManager::new(); + let id = manager.calculate_session_id(); + + let id_num: i64 = id.parse().expect("Should be a valid number"); + const WEEK_MS: i64 = 7 * 24 * 60 * 60 * 1000; + assert!((0..WEEK_MS).contains(&id_num)); + } + + #[tokio::test] + async fn test_prepare_send() { + let manager = UnifiedSessionManager::new(); + + let result = manager.prepare_send().await; + assert!(result.is_some()); + let (node, seq) = result.unwrap(); + assert_eq!(node.tag, "ib"); + assert_eq!( + seq, 1, + "First sequence should be 1 (pre-increment like WhatsApp Web)" + ); + + let result2 = manager.prepare_send().await; + assert!(result2.is_none(), "Duplicate should be prevented"); + assert_eq!(manager.sequence(), 1); + } + + #[tokio::test] + async fn test_clear_last_sent() { + let manager = UnifiedSessionManager::new(); + + let (_, seq1) = manager.prepare_send().await.unwrap(); + assert_eq!(seq1, 1); + assert_eq!(manager.sequence(), 1); + + manager.clear_last_sent().await; + + // After clear, it's treated as a new session -> sequence resets + let result = manager.prepare_send().await; + assert!(result.is_some()); + let (_, seq2) = result.unwrap(); + assert_eq!(seq2, 1, "Sequence resets when session ID changes"); + assert_eq!(manager.sequence(), 1); + } + + #[tokio::test] + async fn test_reset() { + let manager = UnifiedSessionManager::new(); + + let node = NodeBuilder::new("success") + .attr("t", (chrono::Utc::now().timestamp() + 10).to_string()) + .build(); + manager.update_server_time_offset(&node); + let (_, seq1) = manager.prepare_send().await.unwrap(); + assert_eq!(seq1, 1); + + manager.reset().await; + + assert_eq!(manager.server_time_offset_ms(), 0); + // Sequence persists until next prepare_send detects new session ID + assert_eq!(manager.sequence(), 1); + + // After reset, next send will reset sequence since session ID changed + let (_, seq2) = manager.prepare_send().await.unwrap(); + assert_eq!(seq2, 1, "Sequence resets on new session"); + } +} diff --git a/wacore/src/ib.rs b/wacore/src/ib.rs new file mode 100644 index 000000000..9eb3b7791 --- /dev/null +++ b/wacore/src/ib.rs @@ -0,0 +1,179 @@ +//! IB stanza types for unified session telemetry. +//! +//! Wire format: `` + +use crate::protocol::ProtocolNode; +use anyhow::Result; +use wacore_binary::builder::NodeBuilder; +use wacore_binary::node::Node; + +/// Unified session telemetry node. +/// +/// Session ID formula: `(now_ms + server_offset_ms + 3_DAYS_MS) % 7_DAYS_MS` +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct UnifiedSession { + pub id: String, +} + +impl UnifiedSession { + pub fn new(id: impl Into) -> Self { + Self { id: id.into() } + } + + /// Calculate session ID from server time offset. + pub fn calculate_id(server_time_offset_ms: i64) -> String { + const DAY_MS: i64 = 24 * 60 * 60 * 1000; + const WEEK_MS: i64 = 7 * DAY_MS; + const OFFSET_MS: i64 = 3 * DAY_MS; + + let now = chrono::Utc::now().timestamp_millis(); + let adjusted_now = now + server_time_offset_ms; + let id = (adjusted_now + OFFSET_MS) % WEEK_MS; + id.to_string() + } + + pub fn from_offset(server_time_offset_ms: i64) -> Self { + Self::new(Self::calculate_id(server_time_offset_ms)) + } +} + +impl ProtocolNode for UnifiedSession { + fn tag(&self) -> &'static str { + "unified_session" + } + + fn into_node(self) -> Node { + NodeBuilder::new("unified_session") + .attr("id", self.id) + .build() + } + + fn try_from_node(node: &Node) -> Result { + if node.tag != "unified_session" { + return Err(anyhow::anyhow!( + "expected , got <{}>", + node.tag + )); + } + let id = node.attrs.get("id").cloned().unwrap_or_default(); + Ok(Self { id }) + } +} + +/// IB stanza content types. +#[derive(Debug, Clone)] +pub enum IbContent { + UnifiedSession(UnifiedSession), +} + +impl IbContent { + pub fn into_node(self) -> Node { + match self { + IbContent::UnifiedSession(us) => us.into_node(), + } + } +} + +/// IB (Information Broadcast) stanza container. +#[derive(Debug, Clone)] +pub struct IbStanza { + pub content: IbContent, +} + +impl IbStanza { + pub fn unified_session(session: UnifiedSession) -> Self { + Self { + content: IbContent::UnifiedSession(session), + } + } +} + +impl ProtocolNode for IbStanza { + fn tag(&self) -> &'static str { + "ib" + } + + fn into_node(self) -> Node { + NodeBuilder::new("ib") + .children([self.content.into_node()]) + .build() + } + + fn try_from_node(node: &Node) -> Result { + if node.tag != "ib" { + return Err(anyhow::anyhow!("expected , got <{}>", node.tag)); + } + + if let Some(children) = node.children() { + for child in children { + if child.tag == "unified_session" { + return Ok(Self::unified_session(UnifiedSession::try_from_node(child)?)); + } + } + } + + Err(anyhow::anyhow!("unknown or missing content")) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_unified_session_into_node() { + let session = UnifiedSession::new("123456789"); + let node = session.into_node(); + + assert_eq!(node.tag, "unified_session"); + assert_eq!(node.attrs.get("id"), Some(&"123456789".to_string())); + } + + #[test] + fn test_unified_session_try_from_node() { + let node = NodeBuilder::new("unified_session") + .attr("id", "123456789") + .build(); + + let session = UnifiedSession::try_from_node(&node).unwrap(); + assert_eq!(session.id, "123456789"); + } + + #[test] + fn test_ib_stanza_into_node() { + let stanza = IbStanza::unified_session(UnifiedSession::new("123456789")); + let node = stanza.into_node(); + + assert_eq!(node.tag, "ib"); + let children = node.children().unwrap(); + assert_eq!(children.len(), 1); + assert_eq!(children[0].tag, "unified_session"); + assert_eq!(children[0].attrs.get("id"), Some(&"123456789".to_string())); + } + + #[test] + fn test_unified_session_calculate_id() { + const WEEK_MS: i64 = 7 * 24 * 60 * 60 * 1000; + + let id = UnifiedSession::calculate_id(0); + let id_num: i64 = id.parse().unwrap(); + assert!(id_num >= 0); + assert!(id_num < WEEK_MS); + + let id_positive = UnifiedSession::calculate_id(5000); + let id_positive_num: i64 = id_positive.parse().unwrap(); + assert!(id_positive_num >= 0); + assert!(id_positive_num < WEEK_MS); + + let id_negative = UnifiedSession::calculate_id(-5000); + let id_negative_num: i64 = id_negative.parse().unwrap(); + assert!(id_negative_num >= 0); + assert!(id_negative_num < WEEK_MS); + } + + #[test] + fn test_unified_session_from_offset() { + let session = UnifiedSession::from_offset(1000); + assert!(!session.id.is_empty()); + } +} diff --git a/wacore/src/lib.rs b/wacore/src/lib.rs index b74e28f89..90ea17d26 100644 --- a/wacore/src/lib.rs +++ b/wacore/src/lib.rs @@ -14,6 +14,7 @@ pub mod protocol; pub use wacore_noise::framing; pub mod handshake; pub mod history_sync; +pub mod ib; pub use wacore_libsignal as libsignal; pub mod messages; pub mod net;