feat: Implement bidirectional ProtocolNode types for prekey responses

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;

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(),

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<Jid> 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<Jid>,
 /// }
 /// ```
 #[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<Jid> - 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<String> - 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<Self> {
@@ -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<String>,
 }
 
@@ -279,10 +364,15 @@ fn extract_attr_info(field: &syn::Field) -> Result<Option<AttrFieldInfo>, syn::E
         None => return Ok(None),
     };
 
+    // Check if field type is Option<T>
+    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<Option<AttrFieldInfo>, 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<T> 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<Option<AttrFieldInfo>, syn::E
     Ok(None)
 }
 
+/// Check if a type is Option<T>
+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: