From 521f95f608a99c6969e4207531883c901355ccee Mon Sep 17 00:00:00 2001
From: Abtin Keshavarzian <abtink@google.com>
Date: Thu, 6 Feb 2025 08:50:09 -0800
Subject: [PATCH] [border-agent] track `ForwardContext` and abort TMF txn on
 disconnect (#11216)

This commit updates the Border Agent so that the `CoapDtlsSession`
tracks the forwarded CoAP request and the allocated `ForwardContext`
to the leader using `Tmf::Agent` in a list.

If the CoAP session disconnects before the forwarded request finishes,
the pending transaction is explicitly aborted in the session's
`Cleanup()` method using `Tmf::Agent::AbortTransaction()`. This
ensures that the response handler, `HandleCoapResponse()`, is invoked
while the session remains valid.
---
 src/core/meshcop/border_agent.cpp | 11 ++++++++++-
 src/core/meshcop/border_agent.hpp |  7 ++++++-
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/src/core/meshcop/border_agent.cpp b/src/core/meshcop/border_agent.cpp
index ba9e63c7b..0257fd1fa 100644
--- a/src/core/meshcop/border_agent.cpp
+++ b/src/core/meshcop/border_agent.cpp
@@ -623,6 +623,13 @@ BorderAgent::CoapDtlsSession::CoapDtlsSession(Instance &aInstance, Dtls::Transpo
 
 void BorderAgent::CoapDtlsSession::Cleanup(void)
 {
+    while (!mForwardContexts.IsEmpty())
+    {
+        ForwardContext *forwardContext = mForwardContexts.Pop();
+
+        IgnoreError(Get<Tmf::Agent>().AbortTransaction(HandleCoapResponse, forwardContext));
+    }
+
     mTimer.Stop();
     IgnoreError(Get<Ip6::Udp>().RemoveReceiver(mUdpReceiver));
     Get<ThreadNetif>().RemoveUnicastAddress(mCommissionerAloc);
@@ -755,7 +762,7 @@ Error BorderAgent::CoapDtlsSession::ForwardToLeader(const Coap::Message    &aMes
     // will own it. We take back ownership from `HandleCoapResponse()`
     // callback.
 
-    forwardContext.Release();
+    mForwardContexts.Push(*forwardContext.Release());
 
     LogInfo("Forwarded request to leader on %s", PathForUri(aUri));
 
@@ -790,6 +797,8 @@ void BorderAgent::CoapDtlsSession::HandleCoapResponse(const ForwardContext &aFor
     Coap::Message *message = nullptr;
     Error          error;
 
+    IgnoreError(mForwardContexts.Remove(aForwardContext));
+
     SuccessOrExit(error = aResult);
     VerifyOrExit((message = NewPriorityMessage()) != nullptr, error = kErrorNoBufs);
 
diff --git a/src/core/meshcop/border_agent.hpp b/src/core/meshcop/border_agent.hpp
index feb8297cc..241658a9b 100644
--- a/src/core/meshcop/border_agent.hpp
+++ b/src/core/meshcop/border_agent.hpp
@@ -42,6 +42,7 @@
 
 #include "common/as_core_type.hpp"
 #include "common/heap_allocatable.hpp"
+#include "common/linked_list.hpp"
 #include "common/locator.hpp"
 #include "common/non_copyable.hpp"
 #include "common/notifier.hpp"
@@ -311,7 +312,9 @@ private:
         void  Cleanup(void);
 
     private:
-        class ForwardContext : public Heap::Allocatable<ForwardContext>, private ot::NonCopyable
+        class ForwardContext : public ot::LinkedListEntry<ForwardContext>,
+                               public Heap::Allocatable<ForwardContext>,
+                               private ot::NonCopyable
         {
             friend class Heap::Allocatable<ForwardContext>;
 
@@ -319,6 +322,7 @@ private:
             Error ToHeader(Coap::Message &aMessage, uint8_t aCode) const;
 
             CoapDtlsSession &mSession;
+            ForwardContext  *mNext;
             uint16_t         mMessageId;
             bool             mPetition : 1;
             bool             mSeparate : 1;
@@ -358,6 +362,7 @@ private:
         void        HandleTimer(void);
 
         bool                       mIsActiveCommissioner;
+        LinkedList<ForwardContext> mForwardContexts;
         TimerMilliContext          mTimer;
         Ip6::Udp::Receiver         mUdpReceiver;
         Ip6::Netif::UnicastAddress mCommissionerAloc;