From a7d511596bf7d1a66e1ab8fb4e041983c37e943c Mon Sep 17 00:00:00 2001
From: Boris Zbarsky <bzbarsky@apple.com>
Date: Thu, 9 Jun 2022 10:19:14 -0400
Subject: [PATCH] Clarify the logic in ExchangeContext::OnSessionReleased.

It was not clear who was responsible for releasing various refcounts.
Improve the code and documentation to make it clearer what's going on,
and to avoid the two calls into DoClose in the "waiting for response"
case.

Addresses a problem that was mentioned in
https://github.com/project-chip/connectedhomeip/issues/19359
---
 src/messaging/ExchangeContext.cpp | 29 ++++++++++++++++++++++-------
 src/messaging/ExchangeContext.h   |  5 +++--
 src/messaging/ExchangeDelegate.h  | 11 ++++++++++-
 3 files changed, 35 insertions(+), 10 deletions(-)

diff --git a/src/messaging/ExchangeContext.cpp b/src/messaging/ExchangeContext.cpp
index b0aeac785b9beb..83e8ec6e622fe1 100644
--- a/src/messaging/ExchangeContext.cpp
+++ b/src/messaging/ExchangeContext.cpp
@@ -369,6 +369,8 @@ void ExchangeContext::OnSessionReleased()
         return;
     }
 
+    // Hold a ref to ourselves so we can make calls into our delegate that might
+    // decrease our refcount without worrying about use-after-free.
     ExchangeHandle ref(*this);
 
     if (IsResponseExpected())
@@ -376,11 +378,22 @@ void ExchangeContext::OnSessionReleased()
         // If we're waiting on a response, we now know it's never going to show up
         // and we should notify our delegate accordingly.
         CancelResponseTimer();
-        SetResponseExpected(false);
-        NotifyResponseTimeout();
+        // We want to Abort, not just Close, so that RMP bits are cleared, so
+        // don't let NotifyResponseTimeout close us.
+        NotifyResponseTimeout(/* aCloseIfNeeded = */ false);
+        Abort();
+    } else {
+        // Either we're expecting a send or we are in our "just allocated, first
+        // send has not happened yet" state.
+        //
+        // Just mark ourselves as closed.  The consumer is responsible for
+        // releasing us.  See documentation for
+        // ExchangeDelegate::OnExchangeClosing.
+        if (IsSendExpected()) {
+            mFlags.Clear(Flags::kFlagWillSendMessage);
+        }
+        DoClose(true /* clearRetransTable */);
     }
-
-    DoClose(true /* clearRetransTable */);
 }
 
 CHIP_ERROR ExchangeContext::StartResponseTimer()
@@ -414,10 +427,10 @@ void ExchangeContext::HandleResponseTimeout(System::Layer * aSystemLayer, void *
     if (ec == nullptr)
         return;
 
-    ec->NotifyResponseTimeout();
+    ec->NotifyResponseTimeout(/* aCloseIfNeeded = */ true);
 }
 
-void ExchangeContext::NotifyResponseTimeout()
+void ExchangeContext::NotifyResponseTimeout(bool aCloseIfNeeded)
 {
     SetResponseExpected(false);
 
@@ -429,7 +442,9 @@ void ExchangeContext::NotifyResponseTimeout()
         delegate->OnResponseTimeout(this);
     }
 
-    MessageHandled();
+    if (aCloseIfNeeded) {
+        MessageHandled();
+    }
 }
 
 CHIP_ERROR ExchangeContext::HandleMessage(uint32_t messageCounter, const PayloadHeader & payloadHeader, MessageFlags msgFlags,
diff --git a/src/messaging/ExchangeContext.h b/src/messaging/ExchangeContext.h
index d1a8c943644561..5086f02e312baf 100644
--- a/src/messaging/ExchangeContext.h
+++ b/src/messaging/ExchangeContext.h
@@ -238,9 +238,10 @@ class DLL_EXPORT ExchangeContext : public ReliableMessageContext,
 
     /**
      * Notify our delegate, if any, that we have timed out waiting for a
-     * response.
+     * response.  If aCloseIfNeeded is true, check whether the exchange needs to
+     * be closed.
      */
-    void NotifyResponseTimeout();
+    void NotifyResponseTimeout(bool aCloseIfNeeded);
 
     CHIP_ERROR StartResponseTimer();
 
diff --git a/src/messaging/ExchangeDelegate.h b/src/messaging/ExchangeDelegate.h
index 1b66fd1c02b777..161502f64d3ffa 100644
--- a/src/messaging/ExchangeDelegate.h
+++ b/src/messaging/ExchangeDelegate.h
@@ -91,7 +91,16 @@ class DLL_EXPORT ExchangeDelegate
     /**
      * @brief
      *   This function is the protocol callback to invoke when the associated
-     *   exchange context is being closed
+     *   exchange context is being closed.
+     *
+     *   If the exchange was in a state where it was expecting a message to be
+     *   sent due to an earlier WillSendMessage call or because the exchange has
+     *   just been created as an initiator, the consumer is holding a reference
+     *   to the exchange and it's the consumer's responsibility to call
+     *   Release() on the exchange at some point.  The usual way this happens is
+     *   that the consumer tries to send its message, that fails, and the
+     *   consumer calls Close() on the exchange.  Calling Close() after an
+     *   OnExchangeClosing() notification is allowed.
      *
      *  @param[in]    ec            A pointer to the ExchangeContext object.
      */