From 74de70f2132f5b5bf00d67c61b8aaf4f00a8ffae Mon Sep 17 00:00:00 2001
From: Boris Zbarsky <bzbarsky@apple.com>
Date: Fri, 10 Jun 2022 11:12:05 -0400
Subject: [PATCH] Clarify the logic in ExchangeContext::OnSessionReleased.
 (#19383)

* 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

* Clarify comments
---
 src/messaging/ExchangeContext.cpp | 33 ++++++++++++++++++++++++-------
 src/messaging/ExchangeContext.h   |  5 +++--
 src/messaging/ExchangeDelegate.h  | 11 ++++++++++-
 3 files changed, 39 insertions(+), 10 deletions(-)

diff --git a/src/messaging/ExchangeContext.cpp b/src/messaging/ExchangeContext.cpp
index e5008e4bd91fa1..8573d9276e8d55 100644
--- a/src/messaging/ExchangeContext.cpp
+++ b/src/messaging/ExchangeContext.cpp
@@ -371,6 +371,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())
@@ -378,11 +380,25 @@ 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()
@@ -416,10 +432,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);
 
@@ -431,7 +447,10 @@ 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 db3d4e1489572c..b59c68cba17fb1 100644
--- a/src/messaging/ExchangeContext.h
+++ b/src/messaging/ExchangeContext.h
@@ -236,9 +236,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..fe870b6c3094a9 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 in this situation.
      *
      *  @param[in]    ec            A pointer to the ExchangeContext object.
      */