From 07b2c26b464048d92d71446e2b7e2296c5804626 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= Date: Sun, 25 Jan 2026 18:54:47 -0300 Subject: [PATCH 1/5] feat: Implement bidirectional ProtocolNode types for prekey responses - Added SignedPreKeyNode and OneTimePreKeyNode structs for handling signed and one-time prekeys. - Introduced PreKeyBundleUserNode to represent complete prekey bundle responses. - Implemented ProtocolNode trait for these types to facilitate XML serialization and deserialization. - Added helper functions for byte manipulation related to prekey IDs. - Enhanced tests for new prekey structures and their round-trip serialization. feat: Extend A/B property handling in props module - Introduced SamplingProp struct for A/B sampling properties. - Updated AbPropConfig enum to include both AbProp and SamplingProp. - Modified PropsResponse to handle a vector of AbPropConfig instead of just AbProp. - Enhanced parsing logic in PropsSpec to utilize the new ProtocolNode implementations. - Added tests for new A/B property structures and their serialization. feat: Add message stanza types with ProtocolNode derive macro - Created MessageStanza struct for typed message stanzas with JID-aware attributes. - Implemented ProtocolNode derive macro for MessageStanza to facilitate XML serialization. - Updated stanza module to include message handling. fix: Update lib.rs to use self as wacore --- src/handlers/message.rs | 8 +- src/retry.rs | 75 ++---- wacore/derive/src/lib.rs | 163 +++++++++-- wacore/src/iq/mediaconn.rs | 311 +++++++++++++++++++++ wacore/src/iq/prekeys.rs | 506 +++++++++++++++++++++++++++++++++++ wacore/src/iq/props.rs | 403 ++++++++++++++++++++++++---- wacore/src/lib.rs | 1 + wacore/src/stanza/message.rs | 52 ++++ wacore/src/stanza/mod.rs | 2 + 9 files changed, 1387 insertions(+), 134 deletions(-) create mode 100644 wacore/src/stanza/message.rs diff --git a/src/handlers/message.rs b/src/handlers/message.rs index 6cabed3d3..232b6c722 100644 --- a/src/handlers/message.rs +++ b/src/handlers/message.rs @@ -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; diff --git a/src/retry.rs b/src/retry.rs index 405abb9f6..7b6c47d0a 100644 --- a/src/retry.rs +++ b/src/retry.rs @@ -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 _; @@ -419,51 +421,25 @@ 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_data = keys_node + .get_optional_child("key") + .and_then(|key_node| OneTimePreKeyNode::try_from_node(key_node).ok()) + .and_then(|prekey_node| { + let prekey_public = + PublicKey::from_djb_public_key_bytes(&prekey_node.public_bytes).ok()?; + Some((prekey_node.id.into(), prekey_public)) + }); // 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"))?; @@ -472,7 +448,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(), @@ -588,10 +564,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 = 1u32; let skey_value_bytes = device_snapshot .signed_pre_key .public_key @@ -614,19 +589,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(), diff --git a/wacore/derive/src/lib.rs b/wacore/derive/src/lib.rs index a592b8d4e..7599f08b8 100644 --- a/wacore/derive/src/lib.rs +++ b/wacore/derive/src/lib.rs @@ -36,17 +36,28 @@ use syn::{Data, DeriveInput, Fields, parse_macro_input}; /// # Attributes /// /// - `#[protocol(tag = "tagname")]` - Required. Specifies the XML tag name. -/// - `#[attr(name = "attrname")]` - Marks a field as an XML attribute. +/// - `#[attr(name = "attrname")]` - Marks a String field as an XML attribute. /// - `#[attr(name = "attrname", default = "value")]` - Attribute with default value. +/// - `#[attr(name = "attrname", jid)]` - Marks a Jid field as a JID attribute (required). +/// - `#[attr(name = "attrname", jid, optional)]` - Marks an Option field as optional. /// /// # Example /// /// ```ignore /// #[derive(ProtocolNode)] -/// #[protocol(tag = "query")] -/// pub struct QueryRequest { -/// #[attr(name = "request", default = "interactive")] -/// pub request_type: String, +/// #[protocol(tag = "message")] +/// pub struct MessageStanza { +/// #[attr(name = "from", jid)] +/// pub from: Jid, +/// +/// #[attr(name = "to", jid)] +/// pub to: Jid, +/// +/// #[attr(name = "id")] +/// pub id: String, +/// +/// #[attr(name = "sender_lid", jid, optional)] +/// pub sender_lid: Option, /// } /// ``` #[proc_macro_derive(ProtocolNode, attributes(protocol, attr))] @@ -105,8 +116,36 @@ pub fn derive_protocol_node(input: TokenStream) -> TokenStream { .map(|info| { let field_ident = &info.field_ident; let attr_name = &info.attr_name; - quote! { - .attr(#attr_name, self.#field_ident.to_string()) + + match (&info.attr_type, info.optional) { + (AttrType::Jid, true) => { + // Option - only insert if Some + quote! { + if let Some(jid) = self.#field_ident { + builder = builder.jid_attr(#attr_name, jid); + } + } + } + (AttrType::Jid, false) => { + // Required Jid - always insert + quote! { + builder = builder.jid_attr(#attr_name, self.#field_ident); + } + } + (AttrType::String, true) => { + // Option - only insert if Some + quote! { + if let Some(s) = self.#field_ident { + builder = builder.attr(#attr_name, s); + } + } + } + (AttrType::String, false) => { + // Required String - always insert + quote! { + builder = builder.attr(#attr_name, self.#field_ident); + } + } } }) .collect(); @@ -116,32 +155,71 @@ pub fn derive_protocol_node(input: TokenStream) -> TokenStream { .map(|info| { let field_ident = &info.field_ident; let attr_name = &info.attr_name; - if let Some(default) = &info.default { - quote! { - #field_ident: node.attrs().optional_string(#attr_name) - .map(|s| s.to_string()) - .unwrap_or_else(|| #default.to_string()) + + match (&info.attr_type, info.optional, &info.default) { + (AttrType::Jid, false, _) => { + // Required Jid + quote! { + #field_ident: node.attrs().optional_jid(#attr_name) + .ok_or_else(|| ::anyhow::anyhow!("missing required attribute '{}'", #attr_name))? + } + } + (AttrType::Jid, true, _) => { + // Optional Jid + quote! { + #field_ident: node.attrs().optional_jid(#attr_name) + } + } + (AttrType::String, false, Some(default)) => { + // String with default + quote! { + #field_ident: node.attrs().optional_string(#attr_name) + .map(|s| s.to_string()) + .unwrap_or_else(|| #default.to_string()) + } } - } else { - quote! { - #field_ident: node.attrs().optional_string(#attr_name) - .ok_or_else(|| ::anyhow::anyhow!("missing required attribute '{}'", #attr_name))? - .to_string() + (AttrType::String, false, None) => { + // Required String + quote! { + #field_ident: node.attrs().required_string(#attr_name)?.to_string() + } + } + (AttrType::String, true, Some(default)) => { + // Optional String with default + quote! { + #field_ident: node.attrs().optional_string(#attr_name) + .map(|s| Some(s.to_string())) + .unwrap_or_else(|| Some(#default.to_string())) + } + } + (AttrType::String, true, None) => { + // Optional String + quote! { + #field_ident: node.attrs().optional_string(#attr_name).map(|s| s.to_string()) + } } } }) .collect(); - // Only generate Default impl if all fields have defaults - let all_have_defaults = attr_fields.iter().all(|info| info.default.is_some()); + // Only generate Default impl if all fields have defaults or are optional + let all_have_defaults = attr_fields + .iter() + .all(|info| info.default.is_some() || info.optional); let default_impl = if all_have_defaults { let default_fields: Vec<_> = attr_fields .iter() .map(|info| { let field_ident = &info.field_ident; - let default = info.default.as_ref().unwrap(); - quote! { #field_ident: #default.to_string() } + match (&info.attr_type, info.optional, &info.default) { + (_, true, Some(default)) => quote! { #field_ident: Some(#default.to_string()) }, + (_, true, None) => quote! { #field_ident: None }, + (AttrType::String, false, Some(default)) => { + quote! { #field_ident: #default.to_string() } + } + _ => quote! {}, // This shouldn't happen if all_have_defaults is true + } }) .collect(); @@ -165,9 +243,9 @@ pub fn derive_protocol_node(input: TokenStream) -> TokenStream { } fn into_node(self) -> ::wacore_binary::node::Node { - ::wacore_binary::builder::NodeBuilder::new(#tag) - #(#attr_setters)* - .build() + let mut builder = ::wacore_binary::builder::NodeBuilder::new(#tag); + #(#attr_setters)* + builder.build() } fn try_from_node(node: &::wacore_binary::node::Node) -> ::anyhow::Result { @@ -248,9 +326,16 @@ fn generate_empty_impl(name: &syn::Ident, tag: &str) -> proc_macro2::TokenStream } } +enum AttrType { + String, + Jid, +} + struct AttrFieldInfo { field_ident: syn::Ident, attr_name: String, + attr_type: AttrType, + optional: bool, default: Option, } @@ -279,10 +364,15 @@ fn extract_attr_info(field: &syn::Field) -> Result, syn::E None => return Ok(None), }; + // Check if field type is Option + let is_optional = is_option_type(&field.ty); + for attr in &field.attrs { if attr.path().is_ident("attr") { let mut attr_name = None; let mut default = None; + let mut is_jid = false; + let mut explicit_optional = false; attr.parse_nested_meta(|meta| { if meta.path.is_ident("name") { @@ -291,15 +381,30 @@ fn extract_attr_info(field: &syn::Field) -> Result, syn::E } else if meta.path.is_ident("default") { let value: syn::LitStr = meta.value()?.parse()?; default = Some(value.value()); + } else if meta.path.is_ident("jid") { + is_jid = true; + } else if meta.path.is_ident("optional") { + explicit_optional = true; } Ok(()) })?; match attr_name { Some(name) => { + let attr_type = if is_jid { + AttrType::Jid + } else { + AttrType::String + }; + + // Determine if optional: either explicit marker or Option type + let optional = explicit_optional || is_optional; + return Ok(Some(AttrFieldInfo { field_ident, attr_name: name, + attr_type, + optional, default, })); } @@ -315,6 +420,16 @@ fn extract_attr_info(field: &syn::Field) -> Result, syn::E Ok(None) } +/// Check if a type is Option +fn is_option_type(ty: &syn::Type) -> bool { + if let syn::Type::Path(type_path) = ty + && let Some(segment) = type_path.path.segments.last() + { + return segment.ident == "Option"; + } + false +} + /// Derive macro for enums with string representations. /// /// Automatically implements: diff --git a/wacore/src/iq/mediaconn.rs b/wacore/src/iq/mediaconn.rs index 3eddcd5ad..c63380776 100644 --- a/wacore/src/iq/mediaconn.rs +++ b/wacore/src/iq/mediaconn.rs @@ -17,6 +17,7 @@ //! ``` use crate::iq::spec::IqSpec; +use crate::protocol::ProtocolNode; use crate::request::InfoQuery; use anyhow::anyhow; use wacore_binary::builder::NodeBuilder; @@ -29,6 +30,174 @@ pub struct MediaConnHost { pub hostname: String, } +/// Extended media connection host with all attributes (for server-side responses). +#[derive(Debug, Clone)] +pub struct MediaConnHostExtended { + pub hostname: String, + pub host_type: String, // "primary" or "fallback" + pub fallback_hostname: Option, + pub ip4: Option, + pub ip6: Option, + pub fallback_ip4: Option, + pub fallback_ip6: Option, + pub upload: bool, + pub download: bool, + pub download_categories: Vec, + pub download_buckets: Vec, +} + +impl MediaConnHostExtended { + /// Create a simple host (for fallback hosts). + pub fn simple(hostname: String, host_type: String) -> Self { + Self { + hostname, + host_type, + fallback_hostname: None, + ip4: None, + ip6: None, + fallback_ip4: None, + fallback_ip6: None, + upload: false, + download: false, + download_categories: Vec::new(), + download_buckets: Vec::new(), + } + } + + /// Create a primary host with upload/download capabilities. + pub fn primary( + hostname: String, + fallback_hostname: String, + ip4: String, + ip6: String, + download_categories: Vec, + download_buckets: Vec, + ) -> Self { + Self { + hostname: hostname.clone(), + host_type: "primary".to_string(), + fallback_hostname: Some(fallback_hostname), + ip4: Some(ip4.clone()), + ip6: Some(ip6.clone()), + fallback_ip4: Some(ip4), + fallback_ip6: Some(ip6), + upload: true, + download: true, + download_categories, + download_buckets, + } + } +} + +impl ProtocolNode for MediaConnHostExtended { + fn tag(&self) -> &'static str { + "host" + } + + fn into_node(self) -> Node { + let mut builder = NodeBuilder::new("host") + .attr("hostname", &self.hostname) + .attr("type", &self.host_type); + + if let Some(ref fallback_hostname) = self.fallback_hostname { + builder = builder.attr("fallback_hostname", fallback_hostname); + } + if let Some(ref ip4) = self.ip4 { + builder = builder.attr("ip4", ip4); + } + if let Some(ref ip6) = self.ip6 { + builder = builder.attr("ip6", ip6); + } + if let Some(ref fallback_ip4) = self.fallback_ip4 { + builder = builder.attr("fallback_ip4", fallback_ip4); + } + if let Some(ref fallback_ip6) = self.fallback_ip6 { + builder = builder.attr("fallback_ip6", fallback_ip6); + } + + // Build children nodes if upload/download are enabled + let mut children = Vec::new(); + + if self.upload { + children.push(NodeBuilder::new("upload").build()); + } + + if self.download { + // Build download categories + let download_cat_nodes: Vec = self + .download_categories + .iter() + .map(|cat| NodeBuilder::new(cat).build()) + .collect(); + + // Build download buckets + let download_bucket_nodes: Vec = self + .download_buckets + .iter() + .map(|bucket| NodeBuilder::new(bucket).build()) + .collect(); + + children.push( + NodeBuilder::new("download") + .children(download_cat_nodes) + .build(), + ); + + if !download_bucket_nodes.is_empty() { + children.push( + NodeBuilder::new("download_buckets") + .children(download_bucket_nodes) + .build(), + ); + } + } + + if !children.is_empty() { + builder = builder.children(children); + } + + builder.build() + } + + fn try_from_node(node: &Node) -> Result { + if node.tag != "host" { + return Err(anyhow!("expected , got <{}>", node.tag)); + } + + let mut attrs = node.attrs(); + let hostname = attrs + .optional_string("hostname") + .ok_or_else(|| anyhow!("missing hostname attribute"))? + .to_string(); + let host_type = attrs + .optional_string("type") + .unwrap_or("primary") + .to_string(); + + Ok(Self { + hostname, + host_type, + fallback_hostname: attrs.optional_string("fallback_hostname").map(String::from), + ip4: attrs.optional_string("ip4").map(String::from), + ip6: attrs.optional_string("ip6").map(String::from), + fallback_ip4: attrs.optional_string("fallback_ip4").map(String::from), + fallback_ip6: attrs.optional_string("fallback_ip6").map(String::from), + upload: node.get_optional_child("upload").is_some(), + download: node.get_optional_child("download").is_some(), + download_categories: node + .get_optional_child("download") + .and_then(|d| d.children()) + .map(|children| children.iter().map(|c| c.tag.clone()).collect()) + .unwrap_or_default(), + download_buckets: node + .get_optional_child("download_buckets") + .and_then(|d| d.children()) + .map(|children| children.iter().map(|c| c.tag.clone()).collect()) + .unwrap_or_default(), + }) + } +} + /// Media connection response containing auth token and hosts. #[derive(Debug, Clone)] pub struct MediaConnResponse { @@ -38,6 +207,95 @@ pub struct MediaConnResponse { pub hosts: Vec, } +/// Extended media connection response with all server-side attributes. +#[derive(Debug, Clone)] +pub struct MediaConnResponseExtended { + pub auth: String, + pub ttl: u64, + pub auth_ttl: Option, + pub max_buckets: Option, + pub ip_token: Option, + pub set_ip_token: Option, + pub hosts: Vec, +} + +impl MediaConnResponseExtended { + /// Create a simple media conn response for mock servers. + pub fn mock(auth: String, ttl: u64, hosts: Vec) -> Self { + Self { + auth, + ttl, + auth_ttl: Some(21600), // 6 hours + max_buckets: Some(12), + ip_token: Some("MOCK_IP_TOKEN".to_string()), + set_ip_token: Some(1), + hosts, + } + } +} + +impl ProtocolNode for MediaConnResponseExtended { + fn tag(&self) -> &'static str { + "media_conn" + } + + fn into_node(self) -> Node { + let mut builder = NodeBuilder::new("media_conn") + .attr("auth", &self.auth) + .attr("ttl", self.ttl.to_string()); + + if let Some(auth_ttl) = self.auth_ttl { + builder = builder.attr("auth_ttl", auth_ttl.to_string()); + } + if let Some(max_buckets) = self.max_buckets { + builder = builder.attr("max_buckets", max_buckets.to_string()); + } + if let Some(ref ip_token) = self.ip_token { + builder = builder.attr("ip_token", ip_token); + } + if let Some(set_ip_token) = self.set_ip_token { + builder = builder.attr("set_ip_token", set_ip_token.to_string()); + } + + let host_nodes: Vec = self.hosts.into_iter().map(|h| h.into_node()).collect(); + builder = builder.children(host_nodes); + + builder.build() + } + + fn try_from_node(node: &Node) -> Result { + if node.tag != "media_conn" { + return Err(anyhow!("expected , got <{}>", node.tag)); + } + + let mut attrs = node.attrs(); + let auth = attrs + .optional_string("auth") + .ok_or_else(|| anyhow!("missing auth attribute"))? + .to_string(); + let ttl = attrs.optional_u64("ttl").unwrap_or(0); + let auth_ttl = attrs.optional_u64("auth_ttl"); + let max_buckets = attrs.optional_u64("max_buckets"); + let ip_token = attrs.optional_string("ip_token").map(String::from); + let set_ip_token = attrs.optional_u64("set_ip_token"); + + let mut hosts = Vec::new(); + for host_node in node.get_children_by_tag("host") { + hosts.push(MediaConnHostExtended::try_from_node(host_node)?); + } + + Ok(Self { + auth, + ttl, + auth_ttl, + max_buckets, + ip_token, + set_ip_token, + hosts, + }) + } +} + /// Requests media server connection details (auth token and hosts). #[derive(Debug, Clone, Default)] pub struct MediaConnSpec; @@ -155,4 +413,57 @@ mod tests { let result = spec.parse_response(&response); assert!(result.is_err()); } + + #[test] + fn test_media_conn_host_extended_round_trip() { + let host = MediaConnHostExtended::primary( + "127.0.0.1:3000".to_string(), + "127.0.0.1:3000".to_string(), + "127.0.0.1".to_string(), + "::1".to_string(), + vec!["image".to_string(), "video".to_string()], + vec!["0".to_string()], + ); + + let node = host.clone().into_node(); + assert_eq!(node.tag, "host"); + + let parsed = MediaConnHostExtended::try_from_node(&node).unwrap(); + assert_eq!(parsed.hostname, host.hostname); + assert_eq!(parsed.host_type, "primary"); + assert!(parsed.upload); + assert!(parsed.download); + assert_eq!(parsed.download_categories.len(), 2); + assert_eq!(parsed.download_buckets.len(), 1); + } + + #[test] + fn test_media_conn_response_extended_round_trip() { + let hosts = vec![ + MediaConnHostExtended::primary( + "localhost:3000".to_string(), + "localhost:3000".to_string(), + "127.0.0.1".to_string(), + "::1".to_string(), + vec!["image".to_string()], + vec!["0".to_string()], + ), + MediaConnHostExtended::simple("localhost:3000".to_string(), "fallback".to_string()), + ]; + + let response = MediaConnResponseExtended::mock("test-auth".to_string(), 300, hosts); + + let node = response.clone().into_node(); + assert_eq!(node.tag, "media_conn"); + + let parsed = MediaConnResponseExtended::try_from_node(&node).unwrap(); + assert_eq!(parsed.auth, "test-auth"); + assert_eq!(parsed.ttl, 300); + assert_eq!(parsed.auth_ttl, Some(21600)); + assert_eq!(parsed.max_buckets, Some(12)); + assert_eq!(parsed.ip_token, Some("MOCK_IP_TOKEN".to_string())); + assert_eq!(parsed.hosts.len(), 2); + assert_eq!(parsed.hosts[0].host_type, "primary"); + assert_eq!(parsed.hosts[1].host_type, "fallback"); + } } diff --git a/wacore/src/iq/prekeys.rs b/wacore/src/iq/prekeys.rs index 87f547e14..ab6d7e62b 100644 --- a/wacore/src/iq/prekeys.rs +++ b/wacore/src/iq/prekeys.rs @@ -37,8 +37,10 @@ //! //! ``` +use crate::iq::node::{required_attr, required_child}; use crate::iq::spec::IqSpec; use crate::prekeys::PreKeyUtils; +use crate::protocol::ProtocolNode; use crate::request::InfoQuery; use anyhow::anyhow; use wacore_binary::builder::NodeBuilder; @@ -275,6 +277,362 @@ impl IqSpec for PreKeyUploadSpec { } } +// ============================================================================ +// Bidirectional ProtocolNode Types for PreKey Responses +// ============================================================================ +// +// These types implement ProtocolNode for both building (server-side) and +// parsing (client-side) prekey bundle responses. The mock-server uses +// `into_node()` to build responses, while the client can use `try_from_node()` +// to parse them. + +/// Helper function to truncate u32 to 3-byte big-endian representation. +fn truncate_to_3bytes(id: u32) -> Vec { + id.to_be_bytes()[1..].to_vec() +} + +/// Helper function to expand 3-byte big-endian to u32. +fn expand_from_3bytes(bytes: &[u8]) -> Result { + if bytes.len() != 3 { + return Err(anyhow!("Expected 3 bytes for ID, got {}", bytes.len())); + } + Ok(u32::from_be_bytes([0, bytes[0], bytes[1], bytes[2]])) +} + +/// Signed prekey node: `` +/// +/// Wire format: +/// ```xml +/// +/// [3-byte BE u32] +/// [32-byte public key] +/// [64-byte signature] +/// +/// ``` +#[derive(Debug, Clone)] +pub struct SignedPreKeyNode { + pub id: u32, + pub public_bytes: Vec, + pub signature: Vec, +} + +impl SignedPreKeyNode { + pub fn new(id: u32, public_bytes: Vec, signature: Vec) -> Self { + Self { + id, + public_bytes, + signature, + } + } +} + +impl ProtocolNode for SignedPreKeyNode { + fn tag(&self) -> &'static str { + "skey" + } + + fn into_node(self) -> Node { + NodeBuilder::new("skey") + .children([ + NodeBuilder::new("id") + .bytes(truncate_to_3bytes(self.id)) + .build(), + NodeBuilder::new("value").bytes(self.public_bytes).build(), + NodeBuilder::new("signature").bytes(self.signature).build(), + ]) + .build() + } + + fn try_from_node(node: &Node) -> Result { + if node.tag != "skey" { + return Err(anyhow!("expected , got <{}>", node.tag)); + } + + let id_node = required_child(node, "id")?; + let id_bytes = id_node + .content + .as_ref() + .and_then(|c| match c { + NodeContent::Bytes(b) => Some(b), + _ => None, + }) + .ok_or_else(|| anyhow!("missing bytes in "))?; + let id = expand_from_3bytes(id_bytes)?; + + let value_node = required_child(node, "value")?; + let public_bytes = value_node + .content + .as_ref() + .and_then(|c| match c { + NodeContent::Bytes(b) => Some(b.clone()), + _ => None, + }) + .ok_or_else(|| anyhow!("missing bytes in "))?; + + let sig_node = required_child(node, "signature")?; + let signature = sig_node + .content + .as_ref() + .and_then(|c| match c { + NodeContent::Bytes(b) => Some(b.clone()), + _ => None, + }) + .ok_or_else(|| anyhow!("missing bytes in "))?; + + Ok(Self { + id, + public_bytes, + signature, + }) + } +} + +/// One-time prekey node: `` +/// +/// Wire format: +/// ```xml +/// +/// [3-byte BE u32] +/// [32-byte public key] +/// +/// ``` +#[derive(Debug, Clone)] +pub struct OneTimePreKeyNode { + pub id: u32, + pub public_bytes: Vec, +} + +impl OneTimePreKeyNode { + pub fn new(id: u32, public_bytes: Vec) -> Self { + Self { id, public_bytes } + } +} + +impl ProtocolNode for OneTimePreKeyNode { + fn tag(&self) -> &'static str { + "key" + } + + fn into_node(self) -> Node { + NodeBuilder::new("key") + .children([ + NodeBuilder::new("id") + .bytes(truncate_to_3bytes(self.id)) + .build(), + NodeBuilder::new("value").bytes(self.public_bytes).build(), + ]) + .build() + } + + fn try_from_node(node: &Node) -> Result { + if node.tag != "key" { + return Err(anyhow!("expected , got <{}>", node.tag)); + } + + let id_node = required_child(node, "id")?; + let id_bytes = id_node + .content + .as_ref() + .and_then(|c| match c { + NodeContent::Bytes(b) => Some(b), + _ => None, + }) + .ok_or_else(|| anyhow!("missing bytes in "))?; + let id = expand_from_3bytes(id_bytes)?; + + let value_node = required_child(node, "value")?; + let public_bytes = value_node + .content + .as_ref() + .and_then(|c| match c { + NodeContent::Bytes(b) => Some(b.clone()), + _ => None, + }) + .ok_or_else(|| anyhow!("missing bytes in "))?; + + Ok(Self { id, public_bytes }) + } +} + +/// Complete prekey bundle user node: `...` +/// +/// Wire format: +/// ```xml +/// +/// [4-byte BE u32] +/// [0x05 = Curve25519] +/// [32-byte raw public key] +/// ... +/// ... +/// ... +/// +/// ``` +/// +/// This node represents a complete prekey bundle response for a single user/device. +/// Both client and server can use this type: +/// - **Server**: Use `from_bundle()` + `into_node()` to build responses +/// - **Client**: Use `try_from_node()` to parse responses +#[derive(Debug, Clone)] +pub struct PreKeyBundleUserNode { + pub jid: Jid, + pub registration_id: u32, + pub identity_key: Vec, // Must be 32 bytes (raw key, not 33-byte serialized) + pub signed_pre_key: SignedPreKeyNode, + pub one_time_pre_key: Option, + pub device_identity: Option>, // ADVSignedDeviceIdentity protobuf (companion only) +} + +impl PreKeyBundleUserNode { + /// Create a PreKeyBundleUserNode from a PreKeyBundle. + /// + /// The `device_identity` parameter should be provided for companion devices + /// (device ID != 0) and contains the ADVSignedDeviceIdentity protobuf bytes. + pub fn from_bundle( + jid: Jid, + bundle: &PreKeyBundle, + device_identity: Option>, + ) -> Result { + let registration_id = bundle.registration_id()?; + + // Identity key must be 32 bytes (raw key), not 33 bytes (serialized with 0x05 prefix) + let identity_key = bundle + .identity_key()? + .public_key() + .public_key_bytes() + .to_vec(); + + let signed_pre_key_id: u32 = bundle.signed_pre_key_id()?.into(); + let signed_pre_key_public = bundle.signed_pre_key_public()?.public_key_bytes().to_vec(); + let signed_pre_key_signature = bundle.signed_pre_key_signature()?.to_vec(); + + let signed_pre_key = SignedPreKeyNode::new( + signed_pre_key_id, + signed_pre_key_public, + signed_pre_key_signature, + ); + + // Optional one-time prekey + let one_time_pre_key = match (bundle.pre_key_id()?, bundle.pre_key_public()?) { + (Some(id), Some(pk)) => { + let pre_key_id: u32 = id.into(); + let public_bytes = pk.public_key_bytes().to_vec(); + Some(OneTimePreKeyNode::new(pre_key_id, public_bytes)) + } + _ => None, + }; + + Ok(Self { + jid, + registration_id, + identity_key, + signed_pre_key, + one_time_pre_key, + device_identity, + }) + } +} + +impl ProtocolNode for PreKeyBundleUserNode { + fn tag(&self) -> &'static str { + "user" + } + + fn into_node(self) -> Node { + let mut children = vec![ + // Registration ID (4 bytes big-endian) + NodeBuilder::new("registration") + .bytes(self.registration_id.to_be_bytes().to_vec()) + .build(), + // Key type: 0x05 = Curve25519 + NodeBuilder::new("type").bytes(vec![5]).build(), + // Identity key (32 bytes raw public key) + NodeBuilder::new("identity") + .bytes(self.identity_key) + .build(), + // Signed prekey + self.signed_pre_key.into_node(), + ]; + + // Add optional one-time prekey + if let Some(otpk) = self.one_time_pre_key { + children.push(otpk.into_node()); + } + + // Add optional device identity (required for companion devices) + if let Some(dev_id) = self.device_identity { + children.push(NodeBuilder::new("device-identity").bytes(dev_id).build()); + } + + NodeBuilder::new("user") + .attr("jid", self.jid.to_string()) + .attr("type", "result") + .children(children) + .build() + } + + fn try_from_node(node: &Node) -> Result { + if node.tag != "user" { + return Err(anyhow!("expected , got <{}>", node.tag)); + } + + let jid_str = required_attr(node, "jid")?; + let jid = jid_str.parse()?; + + // Parse registration ID (4 bytes big-endian) + let reg_node = required_child(node, "registration")?; + let reg_bytes = reg_node + .content + .as_ref() + .and_then(|c| match c { + NodeContent::Bytes(b) => Some(b), + _ => None, + }) + .ok_or_else(|| anyhow!("missing bytes in "))?; + if reg_bytes.len() != 4 { + return Err(anyhow!("registration ID must be 4 bytes")); + } + let registration_id = + u32::from_be_bytes([reg_bytes[0], reg_bytes[1], reg_bytes[2], reg_bytes[3]]); + + // Parse identity key (32 bytes) + let identity_node = required_child(node, "identity")?; + let identity_key = identity_node + .content + .as_ref() + .and_then(|c| match c { + NodeContent::Bytes(b) => Some(b.clone()), + _ => None, + }) + .ok_or_else(|| anyhow!("missing bytes in "))?; + + // Parse signed prekey + let skey_node = required_child(node, "skey")?; + let signed_pre_key = SignedPreKeyNode::try_from_node(skey_node)?; + + // Parse optional one-time prekey + let one_time_pre_key = node + .get_optional_child("key") + .and_then(|n| OneTimePreKeyNode::try_from_node(n).ok()); + + // Parse optional device identity + let device_identity = + node.get_optional_child("device-identity") + .and_then(|n| match &n.content { + Some(NodeContent::Bytes(b)) => Some(b.clone()), + _ => None, + }); + + Ok(Self { + jid, + registration_id, + identity_key, + signed_pre_key, + one_time_pre_key, + device_identity, + }) + } +} + #[cfg(test)] mod tests { use super::*; @@ -447,4 +805,152 @@ mod tests { let result = spec.parse_response(&response); assert!(result.is_ok()); } + + // ======================================================================== + // Tests for bidirectional ProtocolNode types + // ======================================================================== + + #[test] + fn test_truncate_to_3bytes() { + assert_eq!(truncate_to_3bytes(0x12345678), vec![0x34, 0x56, 0x78]); + assert_eq!(truncate_to_3bytes(0x00000001), vec![0x00, 0x00, 0x01]); + assert_eq!(truncate_to_3bytes(0xFFABCDEF), vec![0xAB, 0xCD, 0xEF]); + } + + #[test] + fn test_expand_from_3bytes() { + assert_eq!(expand_from_3bytes(&[0x34, 0x56, 0x78]).unwrap(), 0x00345678); + assert_eq!(expand_from_3bytes(&[0x00, 0x00, 0x01]).unwrap(), 0x00000001); + assert_eq!(expand_from_3bytes(&[0xAB, 0xCD, 0xEF]).unwrap(), 0x00ABCDEF); + } + + #[test] + fn test_signed_prekey_node_round_trip() { + // Use an ID that fits in 3 bytes (0x00XXXXXX) + let original = SignedPreKeyNode::new(0x00345678, vec![1; 32], vec![2; 64]); + let node = original.clone().into_node(); + + assert_eq!(node.tag, "skey"); + + let parsed = SignedPreKeyNode::try_from_node(&node).unwrap(); + assert_eq!(parsed.id, original.id); + assert_eq!(parsed.public_bytes, original.public_bytes); + assert_eq!(parsed.signature, original.signature); + } + + #[test] + fn test_onetime_prekey_node_round_trip() { + let original = OneTimePreKeyNode::new(0xABCDEF, vec![3; 32]); + let node = original.clone().into_node(); + + assert_eq!(node.tag, "key"); + + let parsed = OneTimePreKeyNode::try_from_node(&node).unwrap(); + assert_eq!(parsed.id, original.id); + assert_eq!(parsed.public_bytes, original.public_bytes); + } + + #[test] + fn test_prekey_bundle_user_node_structure() { + let jid: Jid = "1234567890:33@s.whatsapp.net".parse().unwrap(); + let user_node = PreKeyBundleUserNode { + jid: jid.clone(), + registration_id: 12345, + identity_key: vec![1; 32], + signed_pre_key: SignedPreKeyNode::new(100, vec![2; 32], vec![3; 64]), + one_time_pre_key: Some(OneTimePreKeyNode::new(200, vec![4; 32])), + device_identity: Some(vec![5; 128]), + }; + + let node = user_node.into_node(); + + assert_eq!(node.tag, "user"); + assert_eq!( + node.attrs().optional_string("jid"), + Some(jid.to_string().as_str()) + ); + assert_eq!(node.attrs().optional_string("type"), Some("result")); + + // Verify children count (registration, type, identity, skey, key, device-identity) + if let Some(children) = node.children() { + assert_eq!(children.len(), 6); + assert_eq!(children[0].tag, "registration"); + assert_eq!(children[1].tag, "type"); + assert_eq!(children[2].tag, "identity"); + assert_eq!(children[3].tag, "skey"); + assert_eq!(children[4].tag, "key"); + assert_eq!(children[5].tag, "device-identity"); + } else { + panic!("Expected user node to have children"); + } + } + + #[test] + fn test_prekey_bundle_user_node_round_trip() { + let jid: Jid = "1234567890:33@s.whatsapp.net".parse().unwrap(); + let original = PreKeyBundleUserNode { + jid: jid.clone(), + registration_id: 12345, + identity_key: vec![1; 32], + signed_pre_key: SignedPreKeyNode::new(100, vec![2; 32], vec![3; 64]), + one_time_pre_key: Some(OneTimePreKeyNode::new(200, vec![4; 32])), + device_identity: Some(vec![5; 128]), + }; + + let node = original.clone().into_node(); + let parsed = PreKeyBundleUserNode::try_from_node(&node).unwrap(); + + assert_eq!(parsed.jid, original.jid); + assert_eq!(parsed.registration_id, original.registration_id); + assert_eq!(parsed.identity_key, original.identity_key); + assert_eq!(parsed.signed_pre_key.id, original.signed_pre_key.id); + assert_eq!( + parsed.signed_pre_key.public_bytes, + original.signed_pre_key.public_bytes + ); + assert_eq!( + parsed.signed_pre_key.signature, + original.signed_pre_key.signature + ); + assert!(parsed.one_time_pre_key.is_some()); + assert_eq!( + parsed.one_time_pre_key.as_ref().unwrap().id, + original.one_time_pre_key.as_ref().unwrap().id + ); + assert_eq!(parsed.device_identity, original.device_identity); + } + + #[test] + fn test_prekey_bundle_user_node_without_optional_fields() { + let jid: Jid = "1234567890:0@s.whatsapp.net".parse().unwrap(); + let original = PreKeyBundleUserNode { + jid: jid.clone(), + registration_id: 54321, + identity_key: vec![9; 32], + signed_pre_key: SignedPreKeyNode::new(500, vec![8; 32], vec![7; 64]), + one_time_pre_key: None, + device_identity: None, + }; + + let node = original.clone().into_node(); + + // Verify structure + if let Some(children) = node.children() { + // Should have 4 children (registration, type, identity, skey) - no key, no device-identity + assert_eq!(children.len(), 4); + assert_eq!(children[0].tag, "registration"); + assert_eq!(children[1].tag, "type"); + assert_eq!(children[2].tag, "identity"); + assert_eq!(children[3].tag, "skey"); + } else { + panic!("Expected user node to have children"); + } + + // Round-trip test + let parsed = PreKeyBundleUserNode::try_from_node(&node).unwrap(); + assert_eq!(parsed.jid, original.jid); + assert_eq!(parsed.registration_id, original.registration_id); + assert!(parsed.one_time_pre_key.is_none()); + assert!(parsed.device_identity.is_none()); + } } diff --git a/wacore/src/iq/props.rs b/wacore/src/iq/props.rs index ef44fd3f0..f4af828e3 100644 --- a/wacore/src/iq/props.rs +++ b/wacore/src/iq/props.rs @@ -13,15 +13,18 @@ //! //! //! +//! //! //! ... //! //! //! ``` //! -//! Verified against WhatsApp Web JS (WASmaxOutAbPropsGetExperimentConfigRequest). +//! Verified against WhatsApp Web JS (WASmaxOutAbPropsGetExperimentConfigRequest, +//! WASmaxInAbPropsConfigs). use crate::iq::spec::IqSpec; +use crate::protocol::ProtocolNode; use crate::request::InfoQuery; use wacore_binary::builder::NodeBuilder; use wacore_binary::jid::{Jid, SERVER_JID}; @@ -33,7 +36,7 @@ pub const PROPS_NAMESPACE: &str = "abt"; /// Protocol version for props requests. pub const PROPS_PROTOCOL_VERSION: &str = "1"; -/// A/B property returned from the server. +/// A/B experiment property returned from the server. #[derive(Debug, Clone)] pub struct AbProp { /// The config code (property identifier). @@ -44,6 +47,146 @@ pub struct AbProp { pub config_expo_key: Option, } +impl crate::protocol::ProtocolNode for AbProp { + fn tag(&self) -> &'static str { + "prop" + } + + fn into_node(self) -> Node { + let mut builder = NodeBuilder::new("prop") + .attr("config_code", self.config_code.to_string()) + .attr("config_value", &self.config_value); + + if let Some(expo_key) = self.config_expo_key { + builder = builder.attr("config_expo_key", expo_key.to_string()); + } + + builder.build() + } + + fn try_from_node(node: &Node) -> Result { + use crate::iq::node::optional_attr; + + if node.tag != "prop" { + return Err(anyhow::anyhow!("expected , got <{}>", node.tag)); + } + + let config_code: u32 = optional_attr(node, "config_code") + .ok_or_else(|| anyhow::anyhow!("missing config_code in prop"))? + .parse()?; + let config_value = optional_attr(node, "config_value") + .ok_or_else(|| anyhow::anyhow!("missing config_value in prop"))? + .to_string(); + let config_expo_key = optional_attr(node, "config_expo_key").and_then(|s| s.parse().ok()); + + Ok(Self { + config_code, + config_value, + config_expo_key, + }) + } +} + +/// A/B sampling property returned from the server. +#[derive(Debug, Clone)] +pub struct SamplingProp { + /// The event code (sampling identifier). + pub event_code: u32, + /// The sampling weight (typically -10000..=10000). + pub sampling_weight: i32, +} + +impl crate::protocol::ProtocolNode for SamplingProp { + fn tag(&self) -> &'static str { + "prop" + } + + fn into_node(self) -> Node { + NodeBuilder::new("prop") + .attr("event_code", self.event_code.to_string()) + .attr("sampling_weight", self.sampling_weight.to_string()) + .build() + } + + fn try_from_node(node: &Node) -> Result { + use crate::iq::node::optional_attr; + + if node.tag != "prop" { + return Err(anyhow::anyhow!("expected , got <{}>", node.tag)); + } + + let event_code: u32 = optional_attr(node, "event_code") + .ok_or_else(|| anyhow::anyhow!("missing event_code in prop"))? + .parse()?; + if event_code == 0 { + return Err(anyhow::anyhow!("event_code must be >= 1")); + } + + let sampling_weight: i32 = optional_attr(node, "sampling_weight") + .ok_or_else(|| anyhow::anyhow!("missing sampling_weight in prop"))? + .parse()?; + if !(-10000..=10000).contains(&sampling_weight) { + return Err(anyhow::anyhow!( + "sampling_weight out of range (-10000..=10000): {}", + sampling_weight + )); + } + + Ok(Self { + event_code, + sampling_weight, + }) + } +} + +/// A/B config entry, which can be an experiment or sampling config. +#[derive(Debug, Clone)] +pub enum AbPropConfig { + Experiment(AbProp), + Sampling(SamplingProp), +} + +impl crate::protocol::ProtocolNode for AbPropConfig { + fn tag(&self) -> &'static str { + "prop" + } + + fn into_node(self) -> Node { + match self { + Self::Experiment(prop) => prop.into_node(), + Self::Sampling(prop) => prop.into_node(), + } + } + + fn try_from_node(node: &Node) -> Result { + if node.tag != "prop" { + return Err(anyhow::anyhow!("expected , got <{}>", node.tag)); + } + + let experiment = AbProp::try_from_node(node); + if let Ok(prop) = experiment { + return Ok(Self::Experiment(prop)); + } + + let sampling = SamplingProp::try_from_node(node); + if let Ok(prop) = sampling { + return Ok(Self::Sampling(prop)); + } + + let experiment_err = experiment + .err() + .unwrap_or_else(|| anyhow::anyhow!("unknown error")); + let sampling_err = sampling + .err() + .unwrap_or_else(|| anyhow::anyhow!("unknown error")); + Err(anyhow::anyhow!( + "prop did not match experiment or sampling config: experiment_err={}; sampling_err={}", + experiment_err, + sampling_err + )) + } +} + /// Response from props query. #[derive(Debug, Clone, Default)] pub struct PropsResponse { @@ -57,8 +200,67 @@ pub struct PropsResponse { pub refresh_id: Option, /// Whether this is a delta update. pub delta_update: bool, - /// The properties. - pub props: Vec, + /// The properties (experiment or sampling configs). + pub props: Vec, +} + +impl crate::protocol::ProtocolNode for PropsResponse { + fn tag(&self) -> &'static str { + "props" + } + + fn into_node(self) -> Node { + let mut builder = NodeBuilder::new("props").attr("protocol", PROPS_PROTOCOL_VERSION); + + if let Some(ref ab_key) = self.ab_key { + builder = builder.attr("ab_key", ab_key); + } + if let Some(ref hash) = self.hash { + builder = builder.attr("hash", hash); + } + if let Some(refresh) = self.refresh { + builder = builder.attr("refresh", refresh.to_string()); + } + if let Some(refresh_id) = self.refresh_id { + builder = builder.attr("refresh_id", refresh_id.to_string()); + } + builder = builder.attr("delta_update", self.delta_update.to_string()); + + let prop_nodes: Vec = self.props.into_iter().map(|p| p.into_node()).collect(); + builder = builder.children(prop_nodes); + + builder.build() + } + + fn try_from_node(node: &Node) -> Result { + use crate::iq::node::optional_attr; + + if node.tag != "props" { + return Err(anyhow::anyhow!("expected , got <{}>", node.tag)); + } + + let ab_key = optional_attr(node, "ab_key").map(str::to_string); + let hash = optional_attr(node, "hash").map(str::to_string); + let refresh = optional_attr(node, "refresh").and_then(|s| s.parse().ok()); + let refresh_id = optional_attr(node, "refresh_id").and_then(|s| s.parse().ok()); + let delta_update = optional_attr(node, "delta_update") + .map(|s| s == "true") + .unwrap_or(false); + + let mut props = Vec::new(); + for child in node.get_children_by_tag("prop") { + props.push(AbPropConfig::try_from_node(child)?); + } + + Ok(Self { + ab_key, + hash, + refresh, + refresh_id, + delta_update, + props, + }) + } } /// Fetches A/B testing properties from the server. @@ -115,46 +317,11 @@ impl IqSpec for PropsSpec { } fn parse_response(&self, response: &Node) -> Result { - use crate::iq::node::{optional_attr, required_child}; + use crate::iq::node::required_child; - // Find the props child node + // Find the props child node and parse it using ProtocolNode let props_node = required_child(response, "props")?; - - let ab_key = optional_attr(props_node, "ab_key").map(str::to_string); - let hash = optional_attr(props_node, "hash").map(str::to_string); - let refresh = optional_attr(props_node, "refresh").and_then(|s| s.parse().ok()); - let refresh_id = optional_attr(props_node, "refresh_id").and_then(|s| s.parse().ok()); - let delta_update = optional_attr(props_node, "delta_update") - .map(|s| s == "true") - .unwrap_or(false); - - // Parse individual prop children - let mut props = Vec::new(); - for child in props_node.get_children_by_tag("prop") { - let config_code: u32 = optional_attr(child, "config_code") - .ok_or_else(|| anyhow::anyhow!("missing config_code in prop"))? - .parse()?; - let config_value = optional_attr(child, "config_value") - .ok_or_else(|| anyhow::anyhow!("missing config_value in prop"))? - .to_string(); - let config_expo_key = - optional_attr(child, "config_expo_key").and_then(|s| s.parse().ok()); - - props.push(AbProp { - config_code, - config_value, - config_expo_key, - }); - } - - Ok(PropsResponse { - ab_key, - hash, - refresh, - refresh_id, - delta_update, - props, - }) + PropsResponse::try_from_node(props_node) } } @@ -232,6 +399,10 @@ mod tests { .attr("config_code", "100") .attr("config_value", "enabled") .build(), + NodeBuilder::new("prop") + .attr("event_code", "5138") + .attr("sampling_weight", "-1") + .build(), NodeBuilder::new("prop") .attr("config_code", "200") .attr("config_value", "disabled") @@ -247,13 +418,30 @@ mod tests { assert_eq!(result.refresh, Some(3600)); assert_eq!(result.refresh_id, Some(123)); assert!(!result.delta_update); - assert_eq!(result.props.len(), 2); - assert_eq!(result.props[0].config_code, 100); - assert_eq!(result.props[0].config_value, "enabled"); - assert!(result.props[0].config_expo_key.is_none()); - assert_eq!(result.props[1].config_code, 200); - assert_eq!(result.props[1].config_value, "disabled"); - assert_eq!(result.props[1].config_expo_key, Some(5)); + assert_eq!(result.props.len(), 3); + match &result.props[0] { + AbPropConfig::Experiment(prop) => { + assert_eq!(prop.config_code, 100); + assert_eq!(prop.config_value, "enabled"); + assert!(prop.config_expo_key.is_none()); + } + _ => panic!("Expected Experiment prop"), + } + match &result.props[1] { + AbPropConfig::Sampling(prop) => { + assert_eq!(prop.event_code, 5138); + assert_eq!(prop.sampling_weight, -1); + } + _ => panic!("Expected Sampling prop"), + } + match &result.props[2] { + AbPropConfig::Experiment(prop) => { + assert_eq!(prop.config_code, 200); + assert_eq!(prop.config_value, "disabled"); + assert_eq!(prop.config_expo_key, Some(5)); + } + _ => panic!("Expected Experiment prop"), + } } #[test] @@ -270,4 +458,121 @@ mod tests { let result = spec.parse_response(&response).unwrap(); assert!(result.delta_update); } + + #[test] + fn test_ab_prop_protocol_node_round_trip() { + let prop = AbProp { + config_code: 123, + config_value: "test_value".to_string(), + config_expo_key: Some(456), + }; + + let node = prop.clone().into_node(); + let parsed = AbProp::try_from_node(&node).unwrap(); + + assert_eq!(parsed.config_code, prop.config_code); + assert_eq!(parsed.config_value, prop.config_value); + assert_eq!(parsed.config_expo_key, prop.config_expo_key); + } + + #[test] + fn test_ab_prop_protocol_node_no_expo_key() { + let prop = AbProp { + config_code: 789, + config_value: "another_value".to_string(), + config_expo_key: None, + }; + + let node = prop.clone().into_node(); + let parsed = AbProp::try_from_node(&node).unwrap(); + + assert_eq!(parsed.config_code, prop.config_code); + assert_eq!(parsed.config_value, prop.config_value); + assert_eq!(parsed.config_expo_key, None); + } + + #[test] + fn test_props_response_protocol_node_round_trip() { + let response = PropsResponse { + ab_key: Some("test_ab_key".to_string()), + hash: Some("hash123".to_string()), + refresh: Some(7200), + refresh_id: Some(42), + delta_update: true, + props: vec![ + AbPropConfig::Experiment(AbProp { + config_code: 100, + config_value: "value1".to_string(), + config_expo_key: None, + }), + AbPropConfig::Sampling(SamplingProp { + event_code: 9001, + sampling_weight: -5, + }), + AbPropConfig::Experiment(AbProp { + config_code: 200, + config_value: "value2".to_string(), + config_expo_key: Some(99), + }), + ], + }; + + let node = response.clone().into_node(); + let parsed = PropsResponse::try_from_node(&node).unwrap(); + + assert_eq!(parsed.ab_key, response.ab_key); + assert_eq!(parsed.hash, response.hash); + assert_eq!(parsed.refresh, response.refresh); + assert_eq!(parsed.refresh_id, response.refresh_id); + assert_eq!(parsed.delta_update, response.delta_update); + assert_eq!(parsed.props.len(), response.props.len()); + match &parsed.props[0] { + AbPropConfig::Experiment(prop) => assert_eq!(prop.config_code, 100), + _ => panic!("Expected Experiment prop"), + } + match &parsed.props[1] { + AbPropConfig::Sampling(prop) => assert_eq!(prop.event_code, 9001), + _ => panic!("Expected Sampling prop"), + } + match &parsed.props[2] { + AbPropConfig::Experiment(prop) => assert_eq!(prop.config_expo_key, Some(99)), + _ => panic!("Expected Experiment prop"), + } + } + + #[test] + fn test_props_response_protocol_node_minimal() { + let response = PropsResponse { + ab_key: None, + hash: None, + refresh: None, + refresh_id: None, + delta_update: false, + props: vec![], + }; + + let node = response.clone().into_node(); + let parsed = PropsResponse::try_from_node(&node).unwrap(); + + assert_eq!(parsed.ab_key, None); + assert_eq!(parsed.hash, None); + assert_eq!(parsed.refresh, None); + assert_eq!(parsed.refresh_id, None); + assert!(!parsed.delta_update); + assert_eq!(parsed.props.len(), 0); + } + + #[test] + fn test_sampling_prop_protocol_node_round_trip() { + let prop = SamplingProp { + event_code: 5138, + sampling_weight: -1, + }; + + let node = prop.clone().into_node(); + let parsed = SamplingProp::try_from_node(&node).unwrap(); + + assert_eq!(parsed.event_code, prop.event_code); + assert_eq!(parsed.sampling_weight, prop.sampling_weight); + } } diff --git a/wacore/src/lib.rs b/wacore/src/lib.rs index e03b97280..5db7d2f29 100644 --- a/wacore/src/lib.rs +++ b/wacore/src/lib.rs @@ -1,4 +1,5 @@ #![feature(portable_simd)] +extern crate self as wacore; pub use aes_gcm; pub use wacore_appstate as appstate; diff --git a/wacore/src/stanza/message.rs b/wacore/src/stanza/message.rs new file mode 100644 index 000000000..32043ec22 --- /dev/null +++ b/wacore/src/stanza/message.rs @@ -0,0 +1,52 @@ +//! Message stanza types with ProtocolNode derive macro. +//! +//! Provides type-safe message stanzas with JID-aware attributes. + +use crate::ProtocolNode; +use wacore_binary::jid::Jid; + +/// Typed 1-to-1 or generic message stanza with JID-safe attributes. +/// +/// Wire format: +/// ```xml +/// +/// ``` +#[derive(Debug, Clone, ProtocolNode)] +#[protocol(tag = "message")] +pub struct MessageStanza { + /// Sender JID (required) + #[attr(name = "from", jid)] + pub from: Jid, + + /// Recipient JID (optional; not always present) + #[attr(name = "to", jid, optional)] + pub to: Option, + + /// Message ID (required) + #[attr(name = "id")] + pub id: String, + + /// Timestamp (required, seconds since epoch) + #[attr(name = "t")] + pub timestamp: String, + + /// Message type (default: "text") + #[attr(name = "type", default = "text")] + pub msg_type: String, + + /// Optional sender LID (JID) + #[attr(name = "sender_lid", jid, optional)] + pub sender_lid: Option, + + /// Optional participant JID (group sender) + #[attr(name = "participant", jid, optional)] + pub participant: Option, + + /// Optional participant phone JID + #[attr(name = "participant_pn", jid, optional)] + pub participant_pn: Option, + + /// Optional addressing mode (e.g., "lid") + #[attr(name = "addressing_mode", optional)] + pub addressing_mode: Option, +} diff --git a/wacore/src/stanza/mod.rs b/wacore/src/stanza/mod.rs index 434948c99..c6a96e30d 100644 --- a/wacore/src/stanza/mod.rs +++ b/wacore/src/stanza/mod.rs @@ -3,5 +3,7 @@ //! This module contains type-safe parsers for incoming notification stanzas. pub mod devices; +pub mod message; pub use devices::*; +pub use message::*; From 2aac7d4b61dbfd09f305a987034fbd1d61a910da Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= Date: Sun, 25 Jan 2026 19:51:30 -0300 Subject: [PATCH 2/5] feat: Enhance error handling for stream errors and improve transport disconnection logic --- src/client.rs | 117 ++++++++++++++++++++++++++++++++------------------ 1 file changed, 76 insertions(+), 41 deletions(-) diff --git a/src/client.rs b/src/client.rs index 851e1fdd7..04358902c 100644 --- a/src/client.rs +++ b/src/client.rs @@ -724,6 +724,9 @@ impl Client { /// Build and send an 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(()); + } let id = match node.attrs.get("id") { Some(v) => v.clone(), None => return Ok(()), @@ -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()), + }, + )); + } } } From 2a52187c69c5a648ddb2df515ccc4c437d299d6c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= Date: Sun, 25 Jan 2026 20:22:10 -0300 Subject: [PATCH 3/5] feat: Improve prekey handling and validation in ProtocolNode implementations --- src/retry.rs | 18 ++++++++++-------- wacore/derive/src/lib.rs | 9 +++++---- wacore/src/iq/mediaconn.rs | 2 +- wacore/src/iq/prekeys.rs | 37 ++++++++++++++++++++++++++----------- wacore/src/iq/props.rs | 3 +++ 5 files changed, 45 insertions(+), 24 deletions(-) diff --git a/src/retry.rs b/src/retry.rs index 7b6c47d0a..759738df8 100644 --- a/src/retry.rs +++ b/src/retry.rs @@ -421,14 +421,16 @@ impl Client { let identity_key = PublicKey::from_djb_public_key_bytes(identity_bytes)?; // Extract prekey (optional in some cases). - let prekey_data = keys_node + let prekey_node = keys_node .get_optional_child("key") - .and_then(|key_node| OneTimePreKeyNode::try_from_node(key_node).ok()) - .and_then(|prekey_node| { - let prekey_public = - PublicKey::from_djb_public_key_bytes(&prekey_node.public_bytes).ok()?; - Some((prekey_node.id.into(), prekey_public)) - }); + .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 @@ -566,7 +568,7 @@ impl Client { let prekey_value_bytes = new_prekey_keypair.public_key.public_key_bytes().to_vec(); - let skey_id = 1u32; + let skey_id = device_snapshot.signed_pre_key_id; let skey_value_bytes = device_snapshot .signed_pre_key .public_key diff --git a/wacore/derive/src/lib.rs b/wacore/derive/src/lib.rs index 7599f08b8..d5452775e 100644 --- a/wacore/derive/src/lib.rs +++ b/wacore/derive/src/lib.rs @@ -38,6 +38,7 @@ use syn::{Data, DeriveInput, Fields, parse_macro_input}; /// - `#[protocol(tag = "tagname")]` - Required. Specifies the XML tag name. /// - `#[attr(name = "attrname")]` - Marks a String field as an XML attribute. /// - `#[attr(name = "attrname", default = "value")]` - Attribute with default value. +/// For `Option` fields, a default always yields `Some(default)`. /// - `#[attr(name = "attrname", jid)]` - Marks a Jid field as a JID attribute (required). /// - `#[attr(name = "attrname", jid, optional)]` - Marks an Option field as optional. /// @@ -185,11 +186,11 @@ pub fn derive_protocol_node(input: TokenStream) -> TokenStream { } } (AttrType::String, true, Some(default)) => { - // Optional String with default + // Optional String with default (always Some) quote! { #field_ident: node.attrs().optional_string(#attr_name) - .map(|s| Some(s.to_string())) - .unwrap_or_else(|| Some(#default.to_string())) + .map(|s| s.to_string()) + .or_else(|| Some(#default.to_string())) } } (AttrType::String, true, None) => { @@ -218,7 +219,7 @@ pub fn derive_protocol_node(input: TokenStream) -> TokenStream { (AttrType::String, false, Some(default)) => { quote! { #field_ident: #default.to_string() } } - _ => quote! {}, // This shouldn't happen if all_have_defaults is true + _ => unreachable!("all_have_defaults check should prevent this branch"), } }) .collect(); diff --git a/wacore/src/iq/mediaconn.rs b/wacore/src/iq/mediaconn.rs index c63380776..d20174ffb 100644 --- a/wacore/src/iq/mediaconn.rs +++ b/wacore/src/iq/mediaconn.rs @@ -74,7 +74,7 @@ impl MediaConnHostExtended { download_buckets: Vec, ) -> Self { Self { - hostname: hostname.clone(), + hostname, host_type: "primary".to_string(), fallback_hostname: Some(fallback_hostname), ip4: Some(ip4.clone()), diff --git a/wacore/src/iq/prekeys.rs b/wacore/src/iq/prekeys.rs index ab6d7e62b..1feff1b76 100644 --- a/wacore/src/iq/prekeys.rs +++ b/wacore/src/iq/prekeys.rs @@ -288,6 +288,7 @@ impl IqSpec for PreKeyUploadSpec { /// Helper function to truncate u32 to 3-byte big-endian representation. fn truncate_to_3bytes(id: u32) -> Vec { + debug_assert!(id <= 0x00FF_FFFF, "prekey id exceeds 3-byte range: {id}"); id.to_be_bytes()[1..].to_vec() } @@ -368,6 +369,9 @@ impl ProtocolNode for SignedPreKeyNode { _ => None, }) .ok_or_else(|| anyhow!("missing bytes in "))?; + if public_bytes.len() != 32 { + return Err(anyhow!("signed prekey public key must be 32 bytes")); + } let sig_node = required_child(node, "signature")?; let signature = sig_node @@ -378,6 +382,9 @@ impl ProtocolNode for SignedPreKeyNode { _ => None, }) .ok_or_else(|| anyhow!("missing bytes in "))?; + if signature.len() != 64 { + return Err(anyhow!("signed prekey signature must be 64 bytes")); + } Ok(Self { id, @@ -449,6 +456,9 @@ impl ProtocolNode for OneTimePreKeyNode { _ => None, }) .ok_or_else(|| anyhow!("missing bytes in "))?; + if public_bytes.len() != 32 { + return Err(anyhow!("one-time prekey public key must be 32 bytes")); + } Ok(Self { id, public_bytes }) } @@ -604,23 +614,28 @@ impl ProtocolNode for PreKeyBundleUserNode { _ => None, }) .ok_or_else(|| anyhow!("missing bytes in "))?; + if identity_key.len() != 32 { + return Err(anyhow!("identity key must be 32 bytes")); + } // Parse signed prekey let skey_node = required_child(node, "skey")?; let signed_pre_key = SignedPreKeyNode::try_from_node(skey_node)?; // Parse optional one-time prekey - let one_time_pre_key = node - .get_optional_child("key") - .and_then(|n| OneTimePreKeyNode::try_from_node(n).ok()); + let one_time_pre_key = match node.get_optional_child("key") { + Some(n) => Some(OneTimePreKeyNode::try_from_node(n)?), + None => None, + }; // Parse optional device identity - let device_identity = - node.get_optional_child("device-identity") - .and_then(|n| match &n.content { - Some(NodeContent::Bytes(b)) => Some(b.clone()), - _ => None, - }); + let device_identity = match node.get_optional_child("device-identity") { + Some(n) => match &n.content { + Some(NodeContent::Bytes(b)) => Some(b.clone()), + _ => return Err(anyhow!("device-identity must be bytes")), + }, + None => None, + }; Ok(Self { jid, @@ -812,9 +827,9 @@ mod tests { #[test] fn test_truncate_to_3bytes() { - assert_eq!(truncate_to_3bytes(0x12345678), vec![0x34, 0x56, 0x78]); + assert_eq!(truncate_to_3bytes(0x00345678), vec![0x34, 0x56, 0x78]); assert_eq!(truncate_to_3bytes(0x00000001), vec![0x00, 0x00, 0x01]); - assert_eq!(truncate_to_3bytes(0xFFABCDEF), vec![0xAB, 0xCD, 0xEF]); + assert_eq!(truncate_to_3bytes(0x00ABCDEF), vec![0xAB, 0xCD, 0xEF]); } #[test] diff --git a/wacore/src/iq/props.rs b/wacore/src/iq/props.rs index f4af828e3..ef8ec8b22 100644 --- a/wacore/src/iq/props.rs +++ b/wacore/src/iq/props.rs @@ -74,6 +74,9 @@ impl crate::protocol::ProtocolNode for AbProp { let config_code: u32 = optional_attr(node, "config_code") .ok_or_else(|| anyhow::anyhow!("missing config_code in prop"))? .parse()?; + if config_code == 0 { + return Err(anyhow::anyhow!("config_code must be >= 1")); + } let config_value = optional_attr(node, "config_value") .ok_or_else(|| anyhow::anyhow!("missing config_value in prop"))? .to_string(); From 5b4850e16a59eabf91e9b4cf5087080f04b89073 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= Date: Sun, 25 Jan 2026 21:02:11 -0300 Subject: [PATCH 4/5] fix: Simplify server JID check in handle_iq function --- src/pair.rs | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/pair.rs b/src/pair.rs index 170daf663..fe5d305cb 100644 --- a/src/pair.rs +++ b/src/pair.rs @@ -25,12 +25,11 @@ pub fn make_qr_data(store: &crate::store::Device, ref_str: String) -> String { pub async fn handle_iq(client: &Arc, 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; } From b944de88d7a56afc61dc80fb795c13e19302a637 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= Date: Sun, 25 Jan 2026 21:39:08 -0300 Subject: [PATCH 5/5] feat: Refactor blocklist and prekey handling for improved attribute parsing and normalization --- wacore/src/iq/blocklist.rs | 23 ++++++++---- wacore/src/prekeys.rs | 77 +++++++++++++++++++++++++++++++++++++- 2 files changed, 92 insertions(+), 8 deletions(-) diff --git a/wacore/src/iq/blocklist.rs b/wacore/src/iq/blocklist.rs index 7ac1ea375..5067f03ec 100644 --- a/wacore/src/iq/blocklist.rs +++ b/wacore/src/iq/blocklist.rs @@ -4,7 +4,7 @@ //! the `ProtocolNode` pattern defined in `wacore/src/protocol.rs`. use crate::StringEnum; -use crate::iq::node::{optional_attr, optional_child, optional_u64}; +use crate::iq::node::{optional_child, optional_u64}; use crate::iq::spec::IqSpec; use crate::protocol::ProtocolNode; use crate::request::InfoQuery; @@ -66,11 +66,16 @@ impl ProtocolNode for BlocklistItemRequest { return Err(anyhow!("expected , got <{}>", node.tag)); } - let action_str = - optional_attr(node, "action").ok_or_else(|| anyhow!("missing action attribute"))?; + let mut attrs = node.attrs(); + let action_str = attrs + .optional_string("action") + .ok_or_else(|| anyhow!("missing action attribute"))?; let action = BlocklistAction::try_from(action_str)?; - let jid_str = optional_attr(node, "jid").ok_or_else(|| anyhow!("missing jid attribute"))?; - let jid = jid_str.parse()?; + let jid = attrs.optional_jid("jid"); + if let Err(e) = attrs.finish() { + return Err(anyhow!("{e}")); + } + let jid = jid.ok_or_else(|| anyhow!("missing jid attribute"))?; Ok(Self { jid, action }) } @@ -102,8 +107,12 @@ impl ProtocolNode for BlocklistEntry { return Err(anyhow!("expected , got <{}>", node.tag)); } - let jid_str = optional_attr(node, "jid").ok_or_else(|| anyhow!("missing jid attribute"))?; - let jid = jid_str.parse()?; + let mut attrs = node.attrs(); + let jid = attrs.optional_jid("jid"); + if let Err(e) = attrs.finish() { + return Err(anyhow!("{e}")); + } + let jid = jid.ok_or_else(|| anyhow!("missing jid attribute"))?; let timestamp = optional_u64(node, "t"); Ok(Self { jid, timestamp }) diff --git a/wacore/src/prekeys.rs b/wacore/src/prekeys.rs index 4e4be0ec8..4b4996796 100644 --- a/wacore/src/prekeys.rs +++ b/wacore/src/prekeys.rs @@ -86,7 +86,21 @@ impl PreKeyUtils { continue; } let mut attrs = user_node.attrs(); - let jid = attrs.jid("jid"); + let mut jid = attrs.jid("jid"); + if jid.server == wacore_binary::jid::DEFAULT_USER_SERVER + || jid.server == wacore_binary::jid::HIDDEN_USER_SERVER + { + jid.agent = 0; + } + if jid.device == 0 + && (jid.server == wacore_binary::jid::DEFAULT_USER_SERVER + || jid.server == wacore_binary::jid::HIDDEN_USER_SERVER) + && let Some((user_base, device_str)) = jid.user.split_once(':') + && let Ok(device) = device_str.parse::() + { + jid.user = user_base.to_string(); + jid.device = device; + } let bundle = match Self::node_to_pre_key_bundle(&jid, user_node) { Ok(b) => b, Err(_e) => { @@ -242,3 +256,64 @@ impl PreKeyUtils { Ok((id, public_key_bytes, sig_arr)) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::iq::prekeys::PreKeyBundleUserNode; + use crate::libsignal::protocol::{IdentityKeyPair, KeyPair}; + use crate::protocol::ProtocolNode; + use rand::TryRngCore; + use wacore_binary::node::NodeValue; + + fn create_mock_bundle(device_id: u32) -> PreKeyBundle { + let mut rng = rand::rngs::OsRng.unwrap_err(); + let identity_pair = IdentityKeyPair::generate(&mut rng); + let signed_prekey_pair = KeyPair::generate(&mut rng); + let prekey_pair = KeyPair::generate(&mut rng); + + PreKeyBundle::new( + 1, + device_id.into(), + Some((1u32.into(), prekey_pair.public_key)), + 2u32.into(), + signed_prekey_pair.public_key, + vec![0u8; 64], + *identity_pair.identity_key(), + ) + .expect("Failed to create PreKeyBundle") + } + + #[test] + fn test_parse_prekeys_response_normalizes_lid_device_jid() { + let base_jid = Jid::lid_device("100000012345678", 33); + let bundle = create_mock_bundle(33); + let mut user_node = PreKeyBundleUserNode::from_bundle(base_jid.clone(), &bundle, None) + .expect("build bundle node") + .into_node(); + + let raw_jid = Jid { + user: "100000012345678:33".to_string(), + server: "lid".to_string(), + agent: 1, + device: 0, + integrator: 0, + }; + user_node + .attrs + .insert("jid".to_string(), NodeValue::Jid(raw_jid.clone())); + + let response = NodeBuilder::new("iq") + .children([NodeBuilder::new("list").children([user_node]).build()]) + .build(); + + let bundles = PreKeyUtils::parse_prekeys_response(&response).expect("parse bundles"); + assert!(bundles.contains_key(&base_jid)); + assert!(!bundles.contains_key(&raw_jid)); + + let parsed_jid = bundles.keys().next().expect("parsed jid"); + assert_eq!(parsed_jid.user, base_jid.user); + assert_eq!(parsed_jid.device, base_jid.device); + assert_eq!(parsed_jid.agent, 0); + } +}