fix(transaction): add the 2xx ACK route before sending the 2xx

src/transaction/tests/test_server_invite_ack.rs

@@ -255,3 +255,95 @@ async fn test_server_invite_ignores_ack_with_other_cseq_over_tcp() {
     token.cancel();
     serve_handle.abort();
 }
+
+/// The peer can ACK a final response as soon as it leaves the socket,
+/// before `respond()` has moved the transaction on (issue #192). The send is
+/// held in a full channel while the ACK arrives: the ACK must still reach
+/// the transaction, for a 2xx (new branch, routed by dialog and CSeq) as for
+/// a non-2xx (INVITE branch). A failed or cancelled send leaves no ACK
+/// route behind.
+#[tokio::test]
+async fn test_server_invite_ack_during_final_response_send() {
+    use crate::transaction::key::{TransactionKey, TransactionRole};
+    use crate::transaction::transaction::Transaction;
+    use crate::transport::{channel::ChannelConnection, SipAddr};
+    use std::future::Future;
+    use std::task::{Context, Waker};
+
+    let endpoint = EndpointBuilder::new().build();
+    let addr = SipAddr::from("127.0.0.1:5060".parse::<std::net::SocketAddr>().unwrap());
+    for (cseq, status, outcome, after_ack) in [
+        (3, StatusCode::OK, "send fails", TransactionState::Trying),
+        (5, StatusCode::OK, "cancelled", TransactionState::Trying),
+        (2, StatusCode::OK, "ACK", TransactionState::Terminated),
+        (4, StatusCode::BusyHere, "ACK", TransactionState::Confirmed),
+    ] {
+        let invite = make_reinvite("UDP", cseq);
+        let (_in_tx, in_rx) = tokio::sync::mpsc::unbounded_channel();
+        let (out_tx, mut out_rx) = tokio::sync::mpsc::channel(1);
+        let conn: SipConnection =
+            ChannelConnection::create_connection_bounded(in_rx, out_tx.clone(), addr.clone(), None)
+                .await
+                .unwrap()
+                .into();
+        let key = TransactionKey::from_request(&invite, TransactionRole::Server).unwrap();
+        let mut tx = Transaction::new_server(
+            key,
+            invite.clone(),
+            endpoint.inner.clone(),
+            Some(conn.clone()),
+        );
+        if outcome == "send fails" {
+            drop(out_rx);
+            assert!(tx.reply(status).await.is_err());
+            assert_eq!(tx.state, after_ack);
+            assert_eq!(endpoint.inner.waiting_ack_cseq.len(), 0);
+            assert_eq!(endpoint.inner.waiting_ack.len(), 0);
+            continue;
+        }
+        let ack = if status == StatusCode::OK {
+            make_2xx_ack(&invite, cseq)
+        } else {
+            let mut ack = invite.clone();
+            ack.method = Method::Ack;
+            for h in ack.headers.iter_mut() {
+                if let crate::sip::Header::CSeq(c) = h {
+                    *c = CSeq::new(format!("{} ACK", cseq));
+                }
+            }
+            ack
+        };
+        // Fill the channel so the response send waits for room.
+        let filler = crate::transport::TransportEvent::Incoming(
+            invite.clone().into(),
+            conn.clone(),
+            addr.clone(),
+        );
+        out_tx.try_send(filler).unwrap();
+        let mut reply = Box::pin(tx.reply(status.clone()));
+        let mut cx = Context::from_waker(Waker::noop());
+        assert!(reply.as_mut().poll(&mut cx).is_pending());
+        if outcome == "cancelled" {
+            drop(reply);
+            assert_eq!(tx.state, after_ack);
+            assert_eq!(endpoint.inner.waiting_ack_cseq.len(), 0);
+            continue;
+        }
+        // The ACK arrives while the response is being sent.
+        endpoint
+            .inner
+            .on_received_message(ack.into(), conn.clone(), &addr)
+            .await
+            .unwrap();
+        out_rx.recv().await.unwrap();
+        reply.await.expect("reply");
+        assert!(out_rx.recv().await.is_some(), "the response was sent");
+        let msg = timeout(Duration::from_millis(200), tx.receive())
+            .await
+            .unwrap_or_else(|_| panic!("{status}: the ACK must reach the transaction"))
+            .expect("transaction ended before the ACK");
+        assert!(matches!(msg, SipMessage::Request(ref r) if r.method == Method::Ack));
+        assert_eq!(tx.state, after_ack);
+        assert_eq!(endpoint.inner.waiting_ack_cseq.len(), 0);
+    }
+}

src/transaction/transaction.rs

@@ -428,8 +428,10 @@ impl Transaction {
     /// # Cancel safety
     ///
     /// **NOT cancel-safe.** The mutation order is: store `last_response`
-    /// → `await` the transport send → transition the state machine. If the
-    /// future is dropped during the `await`, `last_response` will already
+    /// (and, for a server INVITE 2xx, add its ACK route) → `await` the
+    /// transport send → transition the state machine. If the
+    /// future is dropped during the `await`, the ACK route is removed
+    /// again but `last_response` will already
     /// be updated while the response may have been only partially
     /// transmitted (or not at all) and the state did not advance. If it
     /// is dropped between the send completing and the transition, the
@@ -496,10 +498,35 @@ impl Transaction {
             SipMessage::Response(resp) => self.last_response.replace(resp),
             _ => None,
         };
+        // The peer can ACK the 2xx before the transition below runs: add
+        // the route the ACK takes (new branch, matched by dialog and CSeq)
+        // before sending, so that ACK waits in this transaction's queue.
+        // The guard removes it again if the send fails or is cancelled.
+        let ack_route = self
+            .last_response
+            .as_ref()
+            .filter(|_| {
+                new_state == TransactionState::Accepted && self.state != TransactionState::Accepted
+            })
+            .and_then(|resp| DialogId::try_from((resp, TransactionRole::Server)).ok())
+            .and_then(|dialog_id| self.add_ack_cseq_route(&dialog_id));
+        let mut guard = AckRouteGuard(
+            ack_route.map(|route| (self.endpoint_inner.clone(), route, self.key.clone())),
+        );
         connection.send(response, self.destination.as_ref()).await?;
+        guard.0 = None;
         self.transition(new_state).map(|_| ())
     }
 
+    fn add_ack_cseq_route(&self, dialog_id: &DialogId) -> Option<(DialogId, u32)> {
+        let seq = self.original.cseq_header().and_then(|c| c.seq()).ok()?;
+        let route = (dialog_id.clone(), seq);
+        self.endpoint_inner
+            .waiting_ack_cseq
+            .insert(route.clone(), self.key.clone());
+        Some(route)
+    }
+
     fn can_transition(&self, target: &TransactionState) -> Result<()> {
         match (&self.state, target) {
             (&TransactionState::Nothing, &TransactionState::Calling)
@@ -1361,13 +1388,10 @@ impl Transaction {
                             // one was answered, and must reach its own
                             // transaction (RFC 3261 §13.2.2.4). A non-2xx ACK
                             // matches by branch (§17.2.3) — unreachable here,
-                            // Accepted only carries 2xx.
-                            let seq = self.original.cseq_header().and_then(|c| c.seq()).ok();
-                            if let Some(seq) = seq {
-                                self.endpoint_inner
-                                    .waiting_ack_cseq
-                                    .insert((dialog_id.clone(), seq), self.key.clone());
-                            }
+                            // Accepted only carries 2xx. `respond` adds the
+                            // CSeq route before sending the 2xx (an ACK can
+                            // arrive before this transition).
+                            self.add_ack_cseq_route(&dialog_id);
                             self.endpoint_inner
                                 .waiting_ack
                                 .insert(dialog_id, self.key.clone());
@@ -1634,6 +1658,22 @@ impl Transaction {
     }
 }
 
+/// A 2xx ACK route added by `respond` before the send; removed on drop
+/// unless disarmed once the send completed.
+struct AckRouteGuard(Option<(EndpointInnerRef, (DialogId, u32), TransactionKey)>);
+
+impl Drop for AckRouteGuard {
+    fn drop(&mut self) {
+        if let Some((endpoint_inner, route, key)) = self.0.take() {
+            endpoint_inner.waiting_ack_cseq.with_mut(|m| {
+                if m.get(&route) == Some(&key) {
+                    m.remove(&route);
+                }
+            });
+        }
+    }
+}
+
 impl Drop for Transaction {
     fn drop(&mut self) {
         self.cleanup();