whatsapp_rust/
request.rs

1use crate::client::Client;
2use crate::client::ClientError;
3use crate::client::ResponseWaiter;
4use crate::socket::error::{EncryptSendError, SocketError};
5use futures::FutureExt;
6use std::num::NonZeroU64;
7use std::sync::Arc;
8use std::sync::atomic::Ordering;
9use std::time::Duration;
10use thiserror::Error;
11use wacore::runtime::timeout as rt_timeout;
12use wacore_binary::Node;
13
14pub use wacore::request::{InfoQuery, InfoQueryType, RequestUtils};
15
16/// How long an IQ waits for its answer when the caller does not pass a timeout
17/// of its own. Callers may pass a longer one, so this bounds the default path
18/// rather than every request; app-state derives its reservation wait from it on
19/// that basis, since the sends it waits behind take the default.
20pub(crate) const DEFAULT_IQ_TIMEOUT: Duration = Duration::from_secs(75);
21const IQ_ID_ATTR: &str = "id";
22const IQ_TAG: &str = "iq";
23
24/// Type-erased send future handed to [`Client::send_and_wait_iq`]. Boxing it
25/// keeps that function non-generic so it isn't re-monomorphized per `IqSpec`.
26/// `Send` on native (IQ awaits happen inside spawned handler tasks); dropped
27/// on wasm where the runtime is single-threaded.
28#[cfg(not(target_arch = "wasm32"))]
29type IqSendFuture<'a> =
30    std::pin::Pin<Box<dyn Future<Output = Result<(), ClientError>> + Send + 'a>>;
31#[cfg(target_arch = "wasm32")]
32type IqSendFuture<'a> =
33    std::pin::Pin<Box<dyn std::future::Future<Output = Result<(), ClientError>> + 'a>>;
34
35/// Runs once the request is on the wire and before the response is awaited.
36///
37/// A caller that must hold a lock across the send — the waiter has to be in the
38/// map before the stanza leaves, or a fast response is dropped as unmatched —
39/// releases it here instead of holding it through the whole round trip. Boxed
40/// rather than generic for the same reason [`IqSendFuture`] is, and `None` on
41/// every other path costs nothing.
42#[cfg(not(target_arch = "wasm32"))]
43type IqOnSent<'a> = Box<dyn FnOnce() + Send + 'a>;
44#[cfg(target_arch = "wasm32")]
45type IqOnSent<'a> = Box<dyn FnOnce() + 'a>;
46
47/// Removes a pending `response_waiters` entry when dropped.
48///
49/// `send_and_wait_iq` can be cancelled mid-await — e.g. the losing side of a
50/// `futures::try_join!` is dropped the instant its sibling errors. Without this
51/// guard the registered waiter would linger in the map: the explicit cleanups
52/// only fired on the send-fail / timeout / shutdown paths, never on
53/// cancellation-via-drop, and a lingering waiter suppresses keepalives for the
54/// life of the connection. Dropping the guard removes the entry on every exit
55/// path; on success `resolve_waiters` already removed it, so it's a no-op.
56pub(crate) struct ResponseWaiterGuard {
57    waiters: Arc<std::sync::Mutex<crate::client::ResponseWaiterMap>>,
58    req_id: String,
59    cleanup_generation: NonZeroU64,
60}
61
62impl ResponseWaiterGuard {
63    pub(crate) fn new(
64        waiters: Arc<std::sync::Mutex<crate::client::ResponseWaiterMap>>,
65        req_id: String,
66        cleanup_generation: NonZeroU64,
67    ) -> Self {
68        Self {
69            waiters,
70            req_id,
71            cleanup_generation,
72        }
73    }
74}
75
76impl Drop for ResponseWaiterGuard {
77    fn drop(&mut self) {
78        self.waiters
79            .lock()
80            .unwrap_or_else(|p| p.into_inner())
81            .remove_guarded(&self.req_id, self.cleanup_generation);
82    }
83}
84
85/// Outcome of the per-spec encode/build step in [`Client::execute`]. Owned (no
86/// spec type parameter) so the send/wait tail behind it stays non-generic.
87enum PreparedIq {
88    /// Fully binary-encoded stanza from `encode_iq_direct` (fast path).
89    Encoded(Vec<u8>),
90    /// Fallback: an `InfoQuery` from `build_iq`, still to be marshalled.
91    /// Boxed to keep the enum small (`InfoQuery` is ~200 bytes vs the
92    /// fast-path `Vec`'s 24) — one alloc per fallback IQ, control-plane only.
93    Query(Box<InfoQuery<'static>>),
94}
95
96#[derive(Debug, Error)]
97#[non_exhaustive]
98pub enum IqError {
99    #[error("IQ request timed out")]
100    Timeout,
101    #[error("client is not connected")]
102    NotConnected,
103    #[error("socket error")]
104    Socket(#[from] SocketError),
105    #[error("encrypted send pipeline failed")]
106    EncryptSend(#[from] EncryptSendError),
107    // Boxed to break the `ClientError::Iq(IqError)` <-> `IqError::ClientState`
108    // type cycle (both would otherwise be infinitely sized).
109    #[error("client state prevented send")]
110    ClientState(#[source] Box<ClientError>),
111    #[error("received disconnect node during IQ wait: {0:?}")]
112    Disconnected(Box<Node>),
113    #[error("received a server error response: code={code}, text='{text}'")]
114    ServerError {
115        code: u16,
116        text: String,
117        /// XMPP error class from the `type` attr; `None` if absent.
118        error_type: Option<String>,
119        /// Server-directed retry delay in seconds from the `backoff` attr; `None` if absent.
120        backoff: Option<u32>,
121    },
122    #[error("received unexpected IQ response type: {got:?}")]
123    UnexpectedResponseType { got: Option<String> },
124    #[error("internal channel closed unexpectedly")]
125    InternalChannelClosed,
126    #[error("IQ request ID is already in flight: {0}")]
127    DuplicateRequestId(String),
128    #[error("failed to encode IQ request")]
129    EncodeError(#[source] anyhow::Error),
130    #[error("failed to parse IQ response")]
131    ParseError(#[from] anyhow::Error),
132}
133
134impl IqError {
135    pub(crate) fn is_transport_unavailable(&self) -> bool {
136        match self {
137            IqError::NotConnected | IqError::Disconnected(_) | IqError::InternalChannelClosed => {
138                true
139            }
140            IqError::EncryptSend(error) => error.is_transport_unavailable(),
141            IqError::ClientState(client) => client.is_transport_unavailable(),
142            _ => false,
143        }
144    }
145
146    /// The request went out and no answer came back in time.
147    ///
148    /// Matched exhaustively so a new variant has to be classified here rather
149    /// than defaulting to "not a timeout" unnoticed.
150    pub(crate) fn is_timeout(&self) -> bool {
151        match self {
152            IqError::Timeout => true,
153            IqError::NotConnected
154            | IqError::Socket(_)
155            | IqError::EncryptSend(_)
156            | IqError::ClientState(_)
157            | IqError::Disconnected(_)
158            | IqError::ServerError { .. }
159            | IqError::UnexpectedResponseType { .. }
160            | IqError::InternalChannelClosed
161            | IqError::DuplicateRequestId(_)
162            | IqError::EncodeError(_)
163            | IqError::ParseError(_) => false,
164        }
165    }
166}
167
168impl From<wacore::request::IqError> for IqError {
169    fn from(err: wacore::request::IqError) -> Self {
170        match err {
171            wacore::request::IqError::Timeout => Self::Timeout,
172            wacore::request::IqError::NotConnected => Self::NotConnected,
173            wacore::request::IqError::Disconnected(node) => Self::Disconnected(node),
174            wacore::request::IqError::ServerError {
175                code,
176                text,
177                error_type,
178                backoff,
179            } => Self::ServerError {
180                code,
181                text,
182                error_type,
183                backoff,
184            },
185            wacore::request::IqError::UnexpectedResponseType { got } => {
186                Self::UnexpectedResponseType { got }
187            }
188            wacore::request::IqError::InternalChannelClosed => Self::InternalChannelClosed,
189            // wacore::IqError is #[non_exhaustive]; a new upstream variant should
190            // get its own arm above. Until then treat it as an unexpected internal error.
191            _ => Self::InternalChannelClosed,
192        }
193    }
194}
195
196impl Client {
197    pub(crate) fn generate_request_id(&self) -> String {
198        self.get_request_utils().generate_request_id()
199    }
200
201    /// Generates a unique message ID that conforms to the WhatsApp protocol format.
202    ///
203    /// This is an advanced function that allows library users to generate message IDs
204    /// that are compatible with the WhatsApp protocol. The generated ID includes
205    /// timestamp, user JID, and random components to ensure uniqueness.
206    ///
207    /// # Advanced Use Case
208    ///
209    /// This function is intended for advanced users who need to build custom protocol
210    /// interactions or manage message IDs manually. Most users should use higher-level
211    /// methods like `send_message` which handle ID generation automatically.
212    ///
213    /// # Returns
214    ///
215    /// A string containing the generated message ID in the format expected by WhatsApp.
216    pub fn generate_message_id(&self) -> String {
217        self.generate_message_id_at(wacore::time::now_secs_u64())
218    }
219
220    /// Same as [`Self::generate_message_id`], but against a caller-supplied
221    /// second (see [`RequestUtils::generate_message_id_at`]).
222    pub(crate) fn generate_message_id_at(&self, unix_secs: u64) -> String {
223        let device_snapshot = self.persistence_manager.get_device_snapshot();
224        // Associated function on purpose: building a RequestUtils here cloned
225        // the unique id per message, and the derivation never reads it.
226        RequestUtils::message_id_at(device_snapshot.pn.as_ref(), unix_secs)
227    }
228
229    fn get_request_utils(&self) -> RequestUtils {
230        RequestUtils::with_counter(self.unique_id.clone(), self.id_counter.clone())
231    }
232
233    /// Sends a custom IQ (Info/Query) stanza to the WhatsApp server.
234    ///
235    /// This is an advanced function that allows library users to send custom IQ stanzas
236    /// for protocol interactions that are not covered by higher-level methods. Common
237    /// use cases include live location updates, custom presence management, or other
238    /// advanced WhatsApp features.
239    ///
240    /// # Advanced Use Case
241    ///
242    /// This function bypasses some of the higher-level abstractions and safety checks
243    /// provided by other client methods. Users should be familiar with the WhatsApp
244    /// protocol and IQ stanza format before using this function.
245    ///
246    /// # Arguments
247    ///
248    /// * `query` - The IQ query to send, containing the stanza type, namespace, content, and optional timeout
249    ///
250    /// # Returns
251    ///
252    /// * `Ok(Arc<OwnedNodeRef>)` - The response node from the server (zero-copy, borrowed from decode buffer)
253    /// * `Err(IqError)` - Various error conditions including timeout, connection issues, or server errors
254    ///
255    /// # Example
256    ///
257    /// ```rust,no_run
258    /// use wacore::request::{InfoQuery, InfoQueryType};
259    /// use wacore_binary::builder::NodeBuilder;
260    /// use wacore_binary::NodeContent;
261    /// use wacore_binary::{Jid, Server};
262    ///
263    /// // This is a simplified example - real usage requires proper setup
264    /// # async fn example(client: &whatsapp_rust::Client) -> Result<(), Box<dyn std::error::Error>> {
265    /// let query_node = NodeBuilder::new("presence")
266    ///     .attr("type", "available")
267    ///     .build();
268    ///
269    /// let server_jid = Jid::new("", Server::Pn);
270    ///
271    /// let query = InfoQuery {
272    ///     query_type: InfoQueryType::Set,
273    ///     namespace: "presence",
274    ///     to: server_jid,
275    ///     target: None,
276    ///     content: Some(NodeContent::Nodes(vec![query_node])),
277    ///     id: None,
278    ///     timeout: None,
279    /// };
280    ///
281    /// let response = client.send_iq(query).await?;
282    /// // Access the node via response.get()
283    /// # Ok(())
284    /// # }
285    /// ```
286    #[cfg_attr(
287        feature = "tracing",
288        tracing::instrument(
289            name = "wa.iq",
290            level = "debug",
291            skip_all,
292            fields(
293                ns = %query.namespace,
294                kind = ?query.query_type,
295                lid = tracing::field::Empty,
296                pn = tracing::field::Empty
297            ),
298            err(Debug)
299        )
300    )]
301    pub async fn send_iq(
302        &self,
303        query: InfoQuery<'_>,
304    ) -> Result<Arc<wacore_binary::OwnedNodeRef>, IqError> {
305        #[cfg(feature = "tracing")]
306        self.record_identity_on_span(&tracing::Span::current());
307
308        let iq_timeout = query.timeout.unwrap_or(DEFAULT_IQ_TIMEOUT);
309        let req_id = query
310            .id
311            .clone()
312            .unwrap_or_else(|| self.generate_request_id());
313
314        let request_utils = self.get_request_utils();
315        let node = request_utils.build_iq_node(query, Some(req_id.clone()));
316
317        self.send_and_wait_iq(
318            req_id,
319            iq_timeout,
320            Box::pin(async { self.send_node(node).await }),
321            None,
322        )
323        .await
324    }
325
326    /// Sends a fully constructed IQ stanza and waits for its matching response.
327    ///
328    /// The stanza ID is preserved when supplied and generated otherwise. The
329    /// same waiter, cancellation, timeout and response validation path used by
330    /// typed IQ specifications handles the request.
331    pub async fn send_iq_node(
332        &self,
333        node: Node,
334        timeout: Option<Duration>,
335    ) -> Result<Arc<wacore_binary::OwnedNodeRef>, IqError> {
336        self.send_iq_node_then(node, timeout, None).await
337    }
338
339    /// [`Self::send_iq_node`] with a hook that runs between the send and the
340    /// wait. See [`IqOnSent`] for why a caller would want one.
341    #[cfg_attr(
342        feature = "tracing",
343        tracing::instrument(name = "wa.iq.node", level = "debug", skip_all, err(Debug))
344    )]
345    pub(crate) async fn send_iq_node_then(
346        &self,
347        mut node: Node,
348        timeout: Option<Duration>,
349        on_sent: Option<IqOnSent<'_>>,
350    ) -> Result<Arc<wacore_binary::OwnedNodeRef>, IqError> {
351        #[cfg(feature = "tracing")]
352        self.record_identity_on_span(&tracing::Span::current());
353
354        if node.tag.as_ref() != IQ_TAG {
355            return Err(IqError::ParseError(anyhow::anyhow!(
356                "expected an <iq> stanza, got <{}>",
357                node.tag
358            )));
359        }
360
361        let req_id = node
362            .attrs
363            .get(IQ_ID_ATTR)
364            .map(|value| value.as_str().into_owned())
365            .filter(|value| !value.is_empty())
366            .unwrap_or_else(|| self.generate_request_id());
367        node.attrs.insert(IQ_ID_ATTR, req_id.clone());
368
369        self.send_and_wait_iq(
370            req_id,
371            timeout.unwrap_or(DEFAULT_IQ_TIMEOUT),
372            Box::pin(async { self.send_node(node).await }),
373            on_sent,
374        )
375        .await
376    }
377
378    /// Executes an IQ specification and returns the typed response.
379    ///
380    /// This is a convenience method that combines building the IQ request,
381    /// sending it, and parsing the response into a single operation.
382    ///
383    /// # Example
384    ///
385    /// ```ignore
386    /// use wacore::iq::groups::GroupQueryIq;
387    ///
388    /// let group_info = client.execute(GroupQueryIq::new(&group_jid)).await?;
389    /// println!("Group subject: {}", group_info.subject);
390    /// ```
391    pub async fn execute<S>(&self, spec: S) -> Result<S::Response, IqError>
392    where
393        S: wacore::iq::spec::IqSpec,
394    {
395        // Only the three spec calls live in this generic body; the send/wait
396        // machinery sits behind the non-generic `execute_prepared` so it isn't
397        // re-stamped for every IqSpec instantiation (~55 of them).
398        let req_id = self.generate_request_id();
399        let mut buf = Vec::new();
400        let prepared = match spec.encode_iq_direct(&req_id, &mut buf) {
401            Ok(true) => PreparedIq::Encoded(buf),
402            Ok(false) => PreparedIq::Query(Box::new(spec.build_iq())),
403            Err(e) => return Err(IqError::EncodeError(e)),
404        };
405
406        let response = self.execute_prepared(req_id, prepared).await?;
407        spec.parse_response(response.get())
408            .map_err(IqError::ParseError)
409    }
410
411    /// Non-generic tail of [`Client::execute`]: sends the already-prepared IQ
412    /// and waits for the response node.
413    async fn execute_prepared(
414        &self,
415        req_id: String,
416        prepared: PreparedIq,
417    ) -> Result<Arc<wacore_binary::OwnedNodeRef>, IqError> {
418        match prepared {
419            // Direct-encode fast path: skip the Node tree for hot IQ specs
420            // (e.g. PreKeyUploadSpec). Fixed 75s timeout — specs needing a
421            // custom timeout don't opt into this path.
422            PreparedIq::Encoded(buf) => {
423                self.send_and_wait_iq(
424                    req_id,
425                    DEFAULT_IQ_TIMEOUT,
426                    Box::pin(async { self.send_raw_bytes(buf).await }),
427                    None,
428                )
429                .await
430            }
431            PreparedIq::Query(iq) => {
432                let mut iq = *iq;
433                // Reuse the id already generated for the fast-path attempt so
434                // send_iq doesn't mint a second one.
435                if iq.id.is_none() {
436                    iq.id = Some(req_id);
437                }
438                self.send_iq(iq).await
439            }
440        }
441    }
442
443    /// Centralizes waiter registration and shutdown/timeout handling.
444    ///
445    /// `send_fn` is type-erased (boxed) rather than a generic `F`: it's only
446    /// awaited once, inline, and `execute<S>` would otherwise stamp out a
447    /// fresh copy of this whole waiter/timeout body per IqSpec (the send
448    /// closure's type is distinct per `S`). One box allocation per IQ — all
449    /// control-plane, never the message hot path — collapses ~15 monomorphized
450    /// copies into one.
451    async fn send_and_wait_iq(
452        &self,
453        req_id: String,
454        timeout: Duration,
455        send_fn: IqSendFuture<'_>,
456        on_sent: Option<IqOnSent<'_>>,
457    ) -> Result<Arc<wacore_binary::OwnedNodeRef>, IqError> {
458        let _t = wacore::telemetry::timer(wacore::telemetry::IQ_DURATION);
459        if !self.is_running.load(Ordering::Relaxed) {
460            wacore::telemetry::iq("error");
461            return Err(IqError::NotConnected);
462        }
463
464        let (tx, rx) = futures::channel::oneshot::channel();
465        let cleanup_generation = {
466            let mut waiters = self.response_waiters_guard();
467            // Explicit IDs are accepted by both InfoQuery and send_iq_node. Never
468            // overwrite an older waiter. The per-registration generation also
469            // prevents an older guard from removing a later reuse of this ID.
470            let Some(cleanup_generation) =
471                waiters.try_insert_guarded(req_id.clone(), ResponseWaiter::Iq(tx))
472            else {
473                wacore::telemetry::iq("error");
474                return Err(IqError::DuplicateRequestId(req_id));
475            };
476            cleanup_generation
477        };
478        // RAII cleanup covers every exit below — including this future being
479        // dropped mid-await (cancellation), which the explicit paths can't
480        // catch. So the send-fail / timeout / shutdown arms no longer remove
481        // the waiter by hand; the guard does it on drop.
482        let _waiter_guard =
483            ResponseWaiterGuard::new(self.response_waiters.clone(), req_id, cleanup_generation);
484
485        // Per-connection: pending IQ requests are bound to the current socket;
486        // a reconnect aborts them (sender retries on the new connection).
487        let shutdown = wacore::runtime::wait_for_shutdown(&self.connection_shutdown_signal());
488
489        if !self.is_running.load(Ordering::Acquire) {
490            wacore::telemetry::iq("error");
491            return Err(IqError::NotConnected);
492        }
493
494        if let Err(e) = send_fn.await {
495            wacore::telemetry::iq("error");
496            return match e {
497                ClientError::Socket(s_err) => Err(IqError::Socket(s_err)),
498                ClientError::EncryptSend(es_err) => Err(IqError::EncryptSend(es_err)),
499                ClientError::NotConnected => Err(IqError::NotConnected),
500                // The send future only ever yields the transport/state errors
501                // above; any other (incl. future #[non_exhaustive]) variant is
502                // surfaced as a client-state failure.
503                other => Err(IqError::ClientState(Box::new(other))),
504            };
505        }
506
507        if let Some(on_sent) = on_sent {
508            on_sent();
509        }
510
511        let request_utils = self.get_request_utils();
512        let result = futures::select! {
513            result = rt_timeout(&*self.runtime, timeout, rx).fuse() => {
514                match result {
515                    Ok(Ok(response_node)) => match request_utils.parse_iq_response(response_node.get()) {
516                        Ok(()) => Ok(response_node),
517                        Err(e) => Err(e.into()),
518                    },
519                    Ok(Err(_)) => Err(IqError::InternalChannelClosed),
520                    Err(_) => Err(IqError::Timeout),
521                }
522            }
523            _ = shutdown.fuse() => Err(IqError::NotConnected),
524        };
525        wacore::telemetry::iq(match &result {
526            Ok(_) => "ok",
527            Err(IqError::Timeout) => "timeout",
528            Err(_) => "error",
529        });
530        result
531    }
532}
533
534#[cfg(test)]
535mod tests {
536    use super::{IQ_ID_ATTR, IQ_TAG, IqError, ResponseWaiterGuard};
537    use crate::client::{ResponseWaiter, ResponseWaiterMap};
538    use std::sync::atomic::Ordering;
539    use std::sync::{Arc, Mutex};
540    use wacore_binary::builder::NodeBuilder;
541
542    #[tokio::test]
543    async fn send_iq_node_rejects_non_iq_stanzas() {
544        let client = crate::test_utils::create_test_client_with_name("invalid_iq_node").await;
545        let error = client
546            .send_iq_node(NodeBuilder::new("message").build(), None)
547            .await
548            .expect_err("a non-IQ stanza must be rejected before transport");
549        assert!(matches!(error, IqError::ParseError(_)));
550    }
551
552    #[tokio::test]
553    async fn send_iq_node_rejects_duplicate_in_flight_id() {
554        let client = crate::test_utils::create_test_client_with_name("duplicate_iq_id").await;
555        client.is_running.store(true, Ordering::Release);
556        let request_id = "duplicate-request";
557        let (tx, _rx) = futures::channel::oneshot::channel();
558        client
559            .response_waiters_guard()
560            .insert(request_id.to_owned(), ResponseWaiter::Iq(tx));
561
562        let error = client
563            .send_iq_node(
564                NodeBuilder::new(IQ_TAG)
565                    .attr(IQ_ID_ATTR, request_id)
566                    .build(),
567                None,
568            )
569            .await
570            .expect_err("a duplicate ID must not replace an existing waiter");
571        assert!(matches!(error, IqError::DuplicateRequestId(id) if id == request_id));
572        assert!(client.response_waiters_guard().contains_key(request_id));
573
574        client.response_waiters_guard().remove(request_id);
575        client.is_running.store(false, Ordering::Release);
576    }
577
578    #[test]
579    fn converts_unexpected_response_type() {
580        let err = IqError::from(wacore::request::IqError::UnexpectedResponseType {
581            got: Some("get".to_string()),
582        });
583
584        match err {
585            IqError::UnexpectedResponseType { got } => assert_eq!(got.as_deref(), Some("get")),
586            other => panic!("expected UnexpectedResponseType, got {other:?}"),
587        }
588    }
589
590    // Cancellation cleanup: dropping a `send_and_wait_iq` future mid-await (e.g.
591    // the loser of a `try_join!`) must remove its still-pending waiter, or a
592    // leaked entry suppresses keepalives for the life of the connection.
593    #[test]
594    fn waiter_guard_removes_pending_entry_on_drop() {
595        let waiters: Arc<Mutex<ResponseWaiterMap>> =
596            Arc::new(Mutex::new(ResponseWaiterMap::default()));
597        let (tx, _rx) = futures::channel::oneshot::channel();
598        let cleanup_generation = waiters
599            .lock()
600            .unwrap()
601            .try_insert_guarded("req-1".to_string(), ResponseWaiter::Iq(tx))
602            .expect("unique request ID");
603        assert!(waiters.lock().unwrap().contains_key("req-1"));
604
605        {
606            let _guard = ResponseWaiterGuard {
607                waiters: waiters.clone(),
608                req_id: "req-1".to_string(),
609                cleanup_generation,
610            };
611        }
612        assert!(
613            !waiters.lock().unwrap().contains_key("req-1"),
614            "dropping the guard must remove the pending waiter"
615        );
616    }
617
618    // On the success path the resolver already removed the entry before the
619    // guard drops, so the guard's removal must be a harmless no-op.
620    #[test]
621    fn waiter_guard_drop_is_noop_when_already_resolved() {
622        let waiters: Arc<Mutex<ResponseWaiterMap>> =
623            Arc::new(Mutex::new(ResponseWaiterMap::default()));
624        let (tx, _rx) = futures::channel::oneshot::channel();
625        let cleanup_generation = waiters
626            .lock()
627            .unwrap()
628            .try_insert_guarded("req-1".to_string(), ResponseWaiter::Iq(tx))
629            .expect("unique request ID");
630        // Map empty = resolver already delivered + removed this request's waiter.
631        waiters.lock().unwrap().remove("req-1");
632        {
633            let _guard = ResponseWaiterGuard {
634                waiters: waiters.clone(),
635                req_id: "req-1".to_string(),
636                cleanup_generation,
637            };
638        }
639        assert!(waiters.lock().unwrap().is_empty());
640    }
641
642    #[test]
643    fn stale_waiter_guard_preserves_a_reused_request_id() {
644        let waiters = Arc::new(Mutex::new(ResponseWaiterMap::default()));
645        let (old_tx, _old_rx) = futures::channel::oneshot::channel();
646        let old_generation = waiters
647            .lock()
648            .unwrap()
649            .try_insert_guarded("reused-id".to_string(), ResponseWaiter::Iq(old_tx))
650            .expect("initial request ID");
651        let old_guard = ResponseWaiterGuard {
652            waiters: waiters.clone(),
653            req_id: "reused-id".to_string(),
654            cleanup_generation: old_generation,
655        };
656
657        // Simulate response delivery removing the old sender, followed by a
658        // new explicit-ID request registering before the old future is dropped.
659        waiters.lock().unwrap().remove("reused-id");
660        let (new_tx, _new_rx) = futures::channel::oneshot::channel();
661        waiters
662            .lock()
663            .unwrap()
664            .try_insert_guarded("reused-id".to_string(), ResponseWaiter::Iq(new_tx))
665            .expect("reused request ID");
666
667        drop(old_guard);
668        assert!(
669            waiters.lock().unwrap().contains_key("reused-id"),
670            "an old guard must not remove the newer registration"
671        );
672    }
673
674    #[test]
675    fn disconnected_waiter_guard_preserves_a_reused_request_id() {
676        let waiters = Arc::new(Mutex::new(ResponseWaiterMap::default()));
677        let (old_tx, _old_rx) = futures::channel::oneshot::channel();
678        let old_generation = waiters
679            .lock()
680            .unwrap()
681            .try_insert_guarded("reused-id".to_string(), ResponseWaiter::Iq(old_tx))
682            .expect("initial request ID");
683        let old_guard = ResponseWaiterGuard {
684            waiters: waiters.clone(),
685            req_id: "reused-id".to_string(),
686            cleanup_generation: old_generation,
687        };
688
689        // Disconnect drains the old sender but its request future (and guard)
690        // may not be polled and dropped until after a reconnect reuses the ID.
691        waiters.lock().unwrap().clear();
692        let (new_tx, _new_rx) = futures::channel::oneshot::channel();
693        waiters
694            .lock()
695            .unwrap()
696            .try_insert_guarded("reused-id".to_string(), ResponseWaiter::Iq(new_tx))
697            .expect("reused request ID");
698
699        drop(old_guard);
700        assert!(
701            waiters.lock().unwrap().contains_key("reused-id"),
702            "a pre-disconnect guard must not remove the post-reconnect waiter"
703        );
704    }
705}
whatsapp_rust/request.rs

whatsapp_rust/
request.rs
