Skip to content

feat: Implement bidirectional ProtocolNode types for prekey responses #250

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 5 commits into from
Jan 26, 2026
Merged

feat: Implement bidirectional ProtocolNode types for prekey responses #250

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
07b2c26
feat: Implement bidirectional ProtocolNode types for prekey responses
jlucaso1 Jan 25, 2026
2aac7d4
feat: Enhance error handling for stream errors and improve transport …
jlucaso1 Jan 25, 2026
2a52187
feat: Improve prekey handling and validation in ProtocolNode implemen…
jlucaso1 Jan 25, 2026
5b4850e
fix: Simplify server JID check in handle_iq function
jlucaso1 Jan 26, 2026
b944de8
feat: Refactor blocklist and prekey handling for improved attribute p…
jlucaso1 Jan 26, 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
117 changes: 76 additions & 41 deletions src/client.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -724,6 +724,9 @@ impl Client {

/// Build and send an <ack/> node corresponding to the given stanza.
async fn send_ack_for(&self, node: &Node) -> Result<(), ClientError> {
if !self.is_connected() || self.expected_disconnect.load(Ordering::Relaxed) {
return Ok(());
}
Comment on lines +727 to +729

@coderabbitai coderabbitai Bot Jan 25, 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 | 🟡 Minor

Avoid dropping ACKs on transient lock contention.

Line 727 uses is_connected(), which relies on try_lock; a briefly held lock will return false and skip the ACK entirely. Prefer an async lock check here so ACKs aren’t silently dropped under load.

🛠️ Suggested fix 
-        if !self.is_connected() || self.expected_disconnect.load(Ordering::Relaxed) {
-            return Ok(());
-        }
+        if self.expected_disconnect.load(Ordering::Relaxed) {
+            return Ok(());
+        }
+        if self.noise_socket.lock().await.is_none() {
+            return Ok(());
+        }
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
if !self.is_connected() || self.expected_disconnect.load(Ordering::Relaxed) {
return Ok(());
}
if self.expected_disconnect.load(Ordering::Relaxed) {
return Ok(());
}
if self.noise_socket.lock().await.is_none() {
return Ok(());
}
🤖 Prompt for AI Agents 
In `@src/client.rs` around lines 727 - 729, The current check using is_connected()
(which uses try_lock) can falsely return false under transient contention and
drop ACKs; change this to use an async-aware lock check instead of try_lock —
either make is_connected async and await the lock or directly acquire the async
lock (e.g., await the same connection/state mutex/read-write lock used
elsewhere) before checking connectivity and expected_disconnect, so the ACK path
isn't skipped due to brief contention; update the call site where is_connected()
is used and ensure expected_disconnect.load(Ordering::Relaxed) is still
evaluated safely after acquiring the async lock (references: is_connected(),
expected_disconnect).

let id = match node.attrs.get("id") {
Some(v) => v.clone(),
None => return Ok(()),
Expand Down Expand Up @@ -1471,49 +1474,81 @@ impl Client {
.map(|n| n.attrs().optional_string("type").unwrap_or("").to_string())
.unwrap_or_default();

match (code, conflict_type.as_str()) {
("515", _) => {
// 515 is expected during registration/pairing phase - server closes stream after pairing
info!(target: "Client", "Got 515 stream error, server is closing stream (expected after pairing). Will auto-reconnect.");
self.expect_disconnect().await;
// Proactively disconnect transport since server may not close the connection
// Clone the transport Arc before spawning to avoid holding the lock
let transport_opt = self.transport.lock().await.clone();
if let Some(transport) = transport_opt {
// Spawn disconnect in background so we don't block the message loop
tokio::spawn(async move {
info!(target: "Client", "Disconnecting transport after 515");
transport.disconnect().await;
});
}
}
("401", "device_removed") | (_, "replaced") => {
info!(target: "Client", "Got stream error indicating client was removed or replaced. Logging out.");
self.expect_disconnect().await;
self.enable_auto_reconnect.store(false, Ordering::Relaxed);
if !conflict_type.is_empty() {
info!(
target: "Client",
"Got stream error indicating client was removed or replaced (conflict={}). Logging out.",
conflict_type
);
self.expect_disconnect().await;
self.enable_auto_reconnect.store(false, Ordering::Relaxed);

let event = if conflict_type == "replaced" {
Event::StreamReplaced(crate::types::events::StreamReplaced)
} else {
Event::LoggedOut(crate::types::events::LoggedOut {
on_connect: false,
reason: ConnectFailureReason::LoggedOut,
})
};
self.core.event_bus.dispatch(&event);
}
("503", _) => {
info!(target: "Client", "Got 503 service unavailable, will auto-reconnect.");
let event = if conflict_type == "replaced" {
Event::StreamReplaced(crate::types::events::StreamReplaced)
} else {
Event::LoggedOut(crate::types::events::LoggedOut {
on_connect: false,
reason: ConnectFailureReason::LoggedOut,
})
};
self.core.event_bus.dispatch(&event);

let transport_opt = self.transport.lock().await.clone();
if let Some(transport) = transport_opt {
tokio::spawn(async move {
info!(target: "Client", "Disconnecting transport after conflict");
transport.disconnect().await;
});
}
_ => {
error!(target: "Client", "Unknown stream error: {}", DisplayableNode(node));
self.expect_disconnect().await;
self.core.event_bus.dispatch(&Event::StreamError(
crate::types::events::StreamError {
code: code.to_string(),
raw: Some(node.clone()),
},
));
} else {
match code {
"515" => {
// 515 is expected during registration/pairing phase - server closes stream after pairing
info!(target: "Client", "Got 515 stream error, server is closing stream (expected after pairing). Will auto-reconnect.");
self.expect_disconnect().await;
// Proactively disconnect transport since server may not close the connection
// Clone the transport Arc before spawning to avoid holding the lock
let transport_opt = self.transport.lock().await.clone();
if let Some(transport) = transport_opt {
// Spawn disconnect in background so we don't block the message loop
tokio::spawn(async move {
info!(target: "Client", "Disconnecting transport after 515");
transport.disconnect().await;
});
}
}
"516" => {
info!(target: "Client", "Got 516 stream error (device removed). Logging out.");
self.expect_disconnect().await;
self.enable_auto_reconnect.store(false, Ordering::Relaxed);
self.core.event_bus.dispatch(&Event::LoggedOut(
crate::types::events::LoggedOut {
on_connect: false,
reason: ConnectFailureReason::LoggedOut,
},
));

let transport_opt = self.transport.lock().await.clone();
if let Some(transport) = transport_opt {
tokio::spawn(async move {
info!(target: "Client", "Disconnecting transport after 516");
transport.disconnect().await;
});
}
}
"503" => {
info!(target: "Client", "Got 503 service unavailable, will auto-reconnect.");
}
_ => {
error!(target: "Client", "Unknown stream error: {}", DisplayableNode(node));
self.expect_disconnect().await;
self.core.event_bus.dispatch(&Event::StreamError(
crate::types::events::StreamError {
code: code.to_string(),
raw: Some(node.clone()),
},
));
}
}
}

Expand Down
8 changes: 2 additions & 6 deletions src/handlers/message.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,12 +30,8 @@ impl StanzaHandler for MessageHandler {
// Extract the chat ID to serialize processing for this chat.
// This prevents race conditions where a later message is processed before
// the PreKey message that establishes the session.
let chat_id = match node.attrs().optional_string("from") {
Some(id) if !id.is_empty() => id.to_string(),
Some(_) => {
warn!("Message stanza has empty 'from' attribute");
return false;
}
let chat_id = match node.attrs().optional_jid("from") {
Some(jid) => jid.to_string(),
None => {
warn!("Message stanza missing required 'from' attribute");
return false;
Expand Down
7 changes: 3 additions & 4 deletions src/pair.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,12 +25,11 @@ pub fn make_qr_data(store: &crate::store::Device, ref_str: String) -> String {

pub async fn handle_iq(client: &Arc<Client>, node: &Node) -> bool {
// Server JID is "s.whatsapp.net" (no @ prefix for server-only JIDs)
if node
if !node
.attrs
.get("from")
.and_then(|s| s.as_str())
.unwrap_or_default()
!= SERVER_JID
.map(|from| from == SERVER_JID)
.unwrap_or(false)
{
return false;
}
Expand Down
77 changes: 22 additions & 55 deletions src/retry.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,11 +6,13 @@ use prost::Message;
use rand::TryRngCore;
use scopeguard;
use std::sync::Arc;
use wacore::iq::prekeys::{OneTimePreKeyNode, SignedPreKeyNode};
use wacore::libsignal::protocol::{
KeyPair, PreKeyBundle, PublicKey, UsePQRatchet, process_prekey_bundle,
};
use wacore::libsignal::store::PreKeyStore;
use wacore::libsignal::store::SessionStore;
use wacore::protocol::ProtocolNode;
use wacore::types::jid::JidExt;
use wacore_binary::builder::NodeBuilder;
use wacore_binary::jid::JidExt as _;
Expand Down Expand Up @@ -419,51 +421,27 @@ impl Client {
let identity_key = PublicKey::from_djb_public_key_bytes(identity_bytes)?;

// Extract prekey (optional in some cases).
let prekey_data = keys_node.get_optional_child("key").and_then(|key_node| {
let id_bytes = key_node
.get_optional_child("id")
.and_then(get_bytes_content)?;
let value_bytes = key_node
.get_optional_child("value")
.and_then(get_bytes_content)?;

// PreKey ID is 3 bytes big-endian.
let prekey_id = if id_bytes.len() >= 3 {
u32::from_be_bytes([0, id_bytes[0], id_bytes[1], id_bytes[2]])
} else {
return None;
};

let prekey_public = PublicKey::from_djb_public_key_bytes(value_bytes).ok()?;
Some((prekey_id.into(), prekey_public))
});
let prekey_node = keys_node
.get_optional_child("key")
.map(OneTimePreKeyNode::try_from_node)
.transpose()?;
let prekey_data = if let Some(prekey_node) = prekey_node {
let prekey_public = PublicKey::from_djb_public_key_bytes(&prekey_node.public_bytes)?;
Some((prekey_node.id.into(), prekey_public))
} else {
None
};

// Extract signed prekey.
let skey_node = keys_node
.get_optional_child("skey")
.ok_or_else(|| anyhow::anyhow!("Missing signed prekey in retry receipt"))?;

let skey_id_bytes = skey_node
.get_optional_child("id")
.and_then(get_bytes_content)
.ok_or_else(|| anyhow::anyhow!("Missing signed prekey ID"))?;
let skey_id = if skey_id_bytes.len() >= 3 {
u32::from_be_bytes([0, skey_id_bytes[0], skey_id_bytes[1], skey_id_bytes[2]])
} else {
return Err(anyhow::anyhow!("Invalid signed prekey ID length"));
};

let skey_value_bytes = skey_node
.get_optional_child("value")
.and_then(get_bytes_content)
.ok_or_else(|| anyhow::anyhow!("Missing signed prekey value"))?;
let skey_public = PublicKey::from_djb_public_key_bytes(skey_value_bytes)?;

let skey_sig_bytes = skey_node
.get_optional_child("signature")
.and_then(get_bytes_content)
.ok_or_else(|| anyhow::anyhow!("Missing signed prekey signature"))?;
let skey_signature: [u8; 64] = skey_sig_bytes
let signed_prekey = SignedPreKeyNode::try_from_node(skey_node)?;
let skey_public = PublicKey::from_djb_public_key_bytes(&signed_prekey.public_bytes)?;
let skey_signature: [u8; 64] = signed_prekey
.signature
.as_slice()
.try_into()
.map_err(|_| anyhow::anyhow!("Invalid signature length"))?;

Expand All @@ -472,7 +450,7 @@ impl Client {
registration_id,
u32::from(requester_jid.device).into(),
prekey_data,
skey_id.into(),
signed_prekey.id.into(),
skey_public,
skey_signature.into(),
identity_key.into(),
Expand Down Expand Up @@ -588,10 +566,9 @@ impl Client {
.public_key_bytes()
.to_vec();

let prekey_id_bytes = new_prekey_id.to_be_bytes()[1..].to_vec();
let prekey_value_bytes = new_prekey_keypair.public_key.public_key_bytes().to_vec();

let skey_id_bytes = 1u32.to_be_bytes()[1..].to_vec();
let skey_id = device_snapshot.signed_pre_key_id;
let skey_value_bytes = device_snapshot
.signed_pre_key
.public_key
Expand All @@ -614,19 +591,9 @@ impl Client {
NodeBuilder::new("identity")
.bytes(identity_key_bytes)
.build(),
NodeBuilder::new("key")
.children([
NodeBuilder::new("id").bytes(prekey_id_bytes).build(),
NodeBuilder::new("value").bytes(prekey_value_bytes).build(),
])
.build(),
NodeBuilder::new("skey")
.children([
NodeBuilder::new("id").bytes(skey_id_bytes).build(),
NodeBuilder::new("value").bytes(skey_value_bytes).build(),
NodeBuilder::new("signature").bytes(skey_sig_bytes).build(),
])
.build(),
OneTimePreKeyNode::new(new_prekey_id, prekey_value_bytes).into_node(),
SignedPreKeyNode::new(skey_id, skey_value_bytes, skey_sig_bytes)
.into_node(),
NodeBuilder::new("device-identity")
.bytes(device_identity_bytes)
.build(),
Expand Down
Loading