Skip to content

fix: complete identity change handler to match WA Web #490

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jlucaso1 merged 3 commits into from
Apr 6, 2026
Merged

fix: complete identity change handler to match WA Web #490

Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
fbff40e
fix: complete identity change handler to match WA Web
jlucaso1 Apr 6, 2026
1b8966c
fix: flush sender key deletion to DB to survive restarts
jlucaso1 Apr 6, 2026
4aed404
test: assert identity key deletion in identity change test
jlucaso1 Apr 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
168 changes: 159 additions & 9 deletions src/handlers/notification.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -319,15 +319,16 @@ fn handle_digest_key(client: &Arc<Client>) {
/// <identity/>
/// </notification>
/// ```
/// Clears device record (sessions + sender keys) for the user so that
/// fresh sessions are established and SKDM is redistributed on next send.
///
/// WA Web defers this when offline. We process immediately because all cleanup
/// is local-only, and `ensure_e2e_sessions` self-defers via `wait_for_offline_delivery_end`.
async fn handle_identity_change(client: &Arc<Client>, node: &Node) {
let Some(from_jid) = node.attrs().optional_jid("from") else {
warn!("Identity change notification missing 'from' attribute");
return;
};

// WA Web ignores companion devices (device != 0) — only primary identity matters
// Only primary device identity changes matter
if from_jid.device != 0 {
debug!(
"Ignoring identity change from companion device {}",
Expand All @@ -336,8 +337,7 @@ async fn handle_identity_change(client: &Arc<Client>, node: &Node) {
return;
}

// WA Web: isMePrimary(h) → handleSelfPrimaryIdentityChange (different flow).
// We must not clear our own device record.
// Self-identity changes use a different flow; clearing our own record would break sessions
let device_snapshot = client.persistence_manager.get_device_snapshot().await;
let is_me = device_snapshot
.pn
Expand All @@ -357,23 +357,64 @@ async fn handle_identity_change(client: &Arc<Client>, node: &Node) {
from_jid.user
);

// Load existing record to pass to clear_device_record
// Deletes non-primary sessions + all sender key device tracking
if let Some(record) = client.load_device_record(&from_jid.user).await {
client
.clear_device_record(&from_jid.user, &from_jid.server, &record)
.await;
}

// Invalidate device cache so next send triggers fresh usync
// Delete primary device session + identity key so a fresh session can be established
{
use wacore::types::jid::JidExt;
let resolved = client.resolve_encryption_jid(&from_jid).await;
let addr = resolved.to_protocol_address();
client.signal_cache.delete_session(&addr).await;
client.signal_cache.delete_identity(&addr).await;
Comment on lines +376 to +377

@chatgpt-codex-connector chatgpt-codex-connector Bot Apr 6, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Acquire session lock before deleting primary session

This new primary-session reset path deletes session and identity directly without taking the per-address session mutex, so it is not serialized with in-flight encrypt/decrypt work that uses session_lock_for(...) (see the existing locked deletion pattern in SignalFeature::delete_sessions). If a send is already using the old session when this handler runs, it can repopulate that stale session after these deletes are flushed, and then ensure_e2e_sessions will treat the session as present and skip prekey re-establishment, effectively undoing the identity-change reset for that user.

Useful? React with 👍 / 👎.

if let Err(e) = client.flush_signal_cache().await {
warn!("Identity change: failed to flush after primary session deletion: {e}");
}
}

// Force status sender key rotation for forward secrecy (clear_device_record
// only cleared device tracking, not the key itself)
{
use wacore::libsignal::store::sender_key_name::SenderKeyName;
use wacore::types::jid::JidExt;
let status_group = "status@broadcast";
for own_jid in device_snapshot.pn.iter().chain(device_snapshot.lid.iter()) {
let sk_name = SenderKeyName::new(
status_group.to_string(),
own_jid.to_protocol_address().to_string(),
);
client
.signal_cache
.delete_sender_key(sk_name.cache_key())
.await;
Comment on lines +385 to +388

@chatgpt-codex-connector chatgpt-codex-connector Bot Apr 6, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Flush status sender-key deletion before returning

delete_sender_key() only marks an in-memory cache entry dirty, but this handler never calls flush_signal_cache() after deleting status@broadcast keys. If ensure_e2e_sessions cannot establish a session (e.g., missing prekeys/offline) or the process restarts before another flush path runs, the old sender key remains in the DB and can be reused after restart, which breaks the intended forward-secrecy rotation on identity change.

Useful? React with 👍 / 👎.

}
Comment on lines +367 to +389

@coderabbitai coderabbitai Bot Apr 6, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Refresh the local device snapshot before rotating status@broadcast sender keys.

device_snapshot is frozen at Line 341 and then reused here after several awaits to build the SenderKeyNames. If our PN/LID changes in that window, this loop deletes the old cache keys and leaves the active sender key intact, so the forward-secrecy rotation silently misses the current identity.

♻️ Proposed fix 
-        let status_group = "status@broadcast";
-        for own_jid in device_snapshot.pn.iter().chain(device_snapshot.lid.iter()) {
+        let own_device_snapshot = client.persistence_manager.get_device_snapshot().await;
+        let status_group = "status@broadcast";
+        for own_jid in own_device_snapshot
+            .pn
+            .iter()
+            .chain(own_device_snapshot.lid.iter())
+        {
             let sk_name = SenderKeyName::new(
                 status_group.to_string(),
                 own_jid.to_protocol_address().to_string(),
             );
🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@src/handlers/notification.rs` around lines 367 - 389, The code uses a stale
device_snapshot when rotating status@broadcast sender keys which can miss
current PN/LID after awaiting; before building SenderKeyName and calling
client.signal_cache.delete_sender_key, re-fetch or refresh the local device
snapshot (the source used to populate device_snapshot.pn and
device_snapshot.lid) so that you iterate over the up-to-date PN/LID values;
specifically, obtain an updated snapshot after the await boundaries (before the
for own_jid in device_snapshot.pn.iter().chain(device_snapshot.lid.iter()) loop)
and use that fresh list to construct
SenderKeyName::new(status_group.to_string(),
own_jid.to_protocol_address().to_string()) and call
client.signal_cache.delete_sender_key(...) to ensure you delete the active
sender keys.

}
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Outdated

// Force fresh usync on next send
client.invalidate_device_cache(&from_jid.user).await;

// Dispatch event so application layer can show security code change
let session_jid = from_jid.clone();
client.core.event_bus.dispatch(&Event::IdentityChange(
crate::types::events::IdentityChange {
user: from_jid,
lid_user: node.attrs().optional_jid("lid"),
},
));

// Re-establish session in background (self-defers when offline)
let client_clone = client.clone();
client
.runtime
.spawn(Box::pin(async move {
if let Err(e) = client_clone.ensure_e2e_sessions(&[session_jid]).await {
warn!("Identity change: failed to re-establish session: {e}");
}
}))
.detach();
Comment thread
jlucaso1 marked this conversation as resolved.
}

/// Handle device list change notifications.
Expand Down Expand Up @@ -2021,7 +2062,6 @@ mod tests {
let collector = Arc::new(TestEventCollector::default());
client.register_handler(collector.clone());

// Companion device (device=5) — should be ignored per WA Web
let node = NodeBuilder::new("notification")
.attr("type", "encrypt")
.attr("from", "5511999999999:5@s.whatsapp.net")
Expand All @@ -2035,4 +2075,114 @@ mod tests {
"companion device identity change should be ignored"
);
}

#[tokio::test]
async fn test_identity_change_deletes_primary_session() {
use wacore::libsignal::protocol::SessionRecord;
use wacore::types::jid::JidExt;

let client = create_test_client().await;

let target_jid: Jid = "5511888888888@s.whatsapp.net".parse().unwrap();
let addr = target_jid.to_protocol_address();

// Pre-populate a session for the primary device
client
.signal_cache
.put_session(&addr, SessionRecord::new_fresh())
.await;
client.signal_cache.put_identity(&addr, &[0u8; 32]).await;

let node = NodeBuilder::new("notification")
.attr("type", "encrypt")
.attr("from", "5511888888888@s.whatsapp.net")
.attr("id", "identity-change-3")
.children([NodeBuilder::new("identity").build()])
.build();
handle_notification_impl(&client, &node).await;

let backend = client.persistence_manager.backend();
let has = client
.signal_cache
.has_session(&addr, &*backend)
.await
.unwrap();
assert!(
!has,
"primary session should be deleted after identity change"
);
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.

#[tokio::test]
async fn test_identity_change_rotates_status_sender_key() {
use wacore::libsignal::store::sender_key_name::SenderKeyName;
use wacore::types::jid::JidExt;

let client = create_test_client().await;

// Set our own JID so sender key deletion knows which namespaces to check
let own_jid: Jid = "5511777777777@s.whatsapp.net".parse().unwrap();
client
.persistence_manager
.modify_device(|d| {
d.pn = Some(own_jid.clone());
})
.await;
Comment on lines +2128 to +2133

@coderabbitai coderabbitai Bot Apr 6, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion | 🟠 Major

Avoid direct Device mutation in the new test setup.

This seeds pn through modify_device, which bypasses the same mutation path production code uses to preserve device-state invariants. Please set this up via DeviceCommand/process_command() so the test exercises the supported surface.

As per coding guidelines "Never modify Device state directly; use DeviceCommand + PersistenceManager::process_command() and read via get_device_snapshot()".

🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@src/handlers/notification.rs` around lines 2128 - 2133, The test is directly
mutating Device.pn via client.persistence_manager.modify_device, which bypasses
invariants; instead construct the appropriate DeviceCommand to set the pn (e.g.,
a command that updates the pn field) and send it through
PersistenceManager::process_command (using the same process_command() path),
then verify the change by calling
client.persistence_manager.get_device_snapshot() rather than reading the device
directly; replace the modify_device(...) closure call with creating the
DeviceCommand, processing it, and reading the snapshot to exercise the supported
surface.


// Pre-populate a sender key for status@broadcast
let sk_name = SenderKeyName::new(
"status@broadcast".to_string(),
own_jid.to_protocol_address().to_string(),
);
let sk_record = wacore::libsignal::protocol::SenderKeyRecord::new_empty();
client
.signal_cache
.put_sender_key(&sk_name, sk_record)
.await;

// Fire identity change for a different user
let node = NodeBuilder::new("notification")
.attr("type", "encrypt")
.attr("from", "5511888888888@s.whatsapp.net")
.attr("id", "identity-change-4")
.children([NodeBuilder::new("identity").build()])
.build();
handle_notification_impl(&client, &node).await;

let backend = client.persistence_manager.backend();
let sk = client
.signal_cache
.get_sender_key(&sk_name, &*backend)
.await
.unwrap();
assert!(
sk.is_none(),
"status@broadcast sender key should be deleted for forward secrecy"
);
}

#[tokio::test]
async fn test_identity_change_with_offline_attribute() {
let client = create_test_client().await;
let collector = Arc::new(TestEventCollector::default());
client.register_handler(collector.clone());

// Notification with offline attribute should still be processed
let node = NodeBuilder::new("notification")
.attr("type", "encrypt")
.attr("from", "5511888888888@s.whatsapp.net")
.attr("id", "identity-change-5")
.attr("offline", "1")
.children([NodeBuilder::new("identity").build()])
.build();
handle_notification_impl(&client, &node).await;

assert!(
collector
.events()
.iter()
.any(|e| matches!(e, Event::IdentityChange(_))),
"offline identity change should still dispatch event"
);
}
}
Loading