tokio-xmpp: Ensure id is added only to stanza

Maxime “pep” Buquet created 2 years ago

The previous commit didn't fix a bug where @id would be added to
elements that didn't need it / where it was invalid (e.g., stream
management).

Signed-off-by: Maxime “pep” Buquet <[e-mail]>

Change summary

tokio-xmpp/src/client/async_client.rs  |  9 +++------
tokio-xmpp/src/client/simple_client.rs | 12 ++++++------
tokio-xmpp/src/component/mod.rs        |  8 ++------
tokio-xmpp/src/xmpp_stream.rs          | 15 ++++++++++++++-
4 files changed, 25 insertions(+), 19 deletions(-)

Detailed changes

tokio-xmpp/src/client/async_client.rs

@@ -17,7 +17,7 @@ use crate::event::Event;
 use crate::happy_eyeballs::{connect_to_host, connect_with_srv};
 use crate::starttls::starttls;
 use crate::xmpp_codec::Packet;
-use crate::xmpp_stream::{self, make_id};
+use crate::xmpp_stream::{self, add_stanza_id};
 use crate::{Error, ProtocolError};
 
 /// XMPP client connection and state
@@ -161,11 +161,8 @@ impl Client {
 
     /// Send stanza
     pub async fn send_stanza(&mut self, stanza: Element) -> Result<(), Error> {
-        let mut el: Element = stanza;
-        if el.attr("id").is_none() {
-            el.set_attr("id", make_id());
-        }
-        self.send(Packet::Stanza(el)).await
+        self.send(Packet::Stanza(add_stanza_id(stanza, ns::JABBER_CLIENT)))
+            .await
     }
 
     /// End connection by sending `</stream:stream>`

tokio-xmpp/src/client/simple_client.rs

@@ -17,7 +17,7 @@ use super::bind::bind;
 use crate::happy_eyeballs::connect_with_srv;
 use crate::starttls::starttls;
 use crate::xmpp_codec::Packet;
-use crate::xmpp_stream::{self, make_id};
+use crate::xmpp_stream::{self, add_stanza_id};
 use crate::{Error, ProtocolError};
 
 /// A simple XMPP client connection
@@ -98,11 +98,11 @@ impl Client {
     where
         E: Into<Element>,
     {
-        let mut el: Element = stanza.into();
-        if el.attr("id").is_none() {
-            el.set_attr("id", make_id());
-        }
-        self.send(Packet::Stanza(el.into())).await
+        self.send(Packet::Stanza(add_stanza_id(
+            stanza.into(),
+            ns::JABBER_CLIENT,
+        )))
+        .await
     }
 
     /// End connection by sending `</stream:stream>`

tokio-xmpp/src/component/mod.rs

@@ -12,7 +12,7 @@ use super::happy_eyeballs::connect_to_host;
 use super::xmpp_codec::Packet;
 use super::xmpp_stream;
 use super::Error;
-use crate::xmpp_stream::make_id;
+use crate::xmpp_stream::add_stanza_id;
 
 mod auth;
 
@@ -54,11 +54,7 @@ impl Component {
 
     /// Send stanza
     pub async fn send_stanza(&mut self, stanza: Element) -> Result<(), Error> {
-        let mut el: Element = stanza;
-        if el.attr("id").is_none() {
-            el.set_attr("id", make_id());
-        }
-        self.send(el).await
+        self.send(add_stanza_id(stanza, ns::COMPONENT_ACCEPT)).await
     }
 
     /// End connection

tokio-xmpp/src/xmpp_stream.rs

@@ -14,11 +14,24 @@ use crate::stream_start;
 use crate::xmpp_codec::{Packet, XMPPCodec};
 use crate::Error;
 
-pub(crate) fn make_id() -> String {
+fn make_id() -> String {
     let id: u64 = thread_rng().gen();
     format!("{}", id)
 }
 
+pub(crate) fn add_stanza_id(mut stanza: Element, default_ns: &str) -> Element {
+    if stanza.is("iq", default_ns)
+        || stanza.is("message", default_ns)
+        || stanza.is("presence", default_ns)
+    {
+        if stanza.attr("id").is_none() {
+            stanza.set_attr("id", make_id());
+        }
+    }
+
+    stanza
+}
+
 /// Wraps a binary stream (tokio's `AsyncRead + AsyncWrite`) to decode
 /// and encode XMPP packets.
 ///