From 46241359ccc22fb6fb0db22608573d82a1aae371 Mon Sep 17 00:00:00 2001
From: Boris Zbarsky <bzbarsky@apple.com>
Date: Thu, 16 Dec 2021 10:01:14 -0500
Subject: [PATCH] Fix double-free bugs on failure to send a write message.
 (#13051)

There were two separate bugs here:

1) For a group write, we were violating the contract for
WriteClient::SendWriteRequest, which is that the caller must call
Shutdown on failure but the WriteClient itself will do it on success
(and call OnDone in the process).  We were unconditionally calling
Shutdown and OnDone inside SetWriteRequest, even on failure.

2) WriteInteraction was violating the contract for
WriteClientHandle::SendWriteRequest, which is different: it always
guarantees it will call OnDone. But the consumer was assuming that
OnDone would only be called if SendWriteRequest returned success.
---
 src/app/WriteClient.cpp           | 18 ++++++++++++------
 src/app/WriteClient.h             |  3 +++
 src/controller/WriteInteraction.h |  6 +++++-
 3 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/src/app/WriteClient.cpp b/src/app/WriteClient.cpp
index 4f4762631c8036..b14f2c4bc23e6a 100644
--- a/src/app/WriteClient.cpp
+++ b/src/app/WriteClient.cpp
@@ -248,14 +248,20 @@ CHIP_ERROR WriteClient::SendWriteRequest(const SessionHandle & session, System::
         ChipLogError(DataManagement, "Write client failed to SendWriteRequest");
         ClearExistingExchangeContext();
     }
-
-    if (session.IsGroupSession())
+    else
     {
-        // Always shutdown on Group communication
-        ChipLogDetail(DataManagement, "Closing on group Communication ");
+        // TODO: Ideally this would happen async, but to make sure that we
+        // handle this object dying (e.g. due to IM enging shutdown) while the
+        // async bits are pending we'd need to malloc some state bit that we can
+        // twiddle if we die.  For now just do the OnDone callback sync.
+        if (session.IsGroupSession())
+        {
+            // Always shutdown on Group communication
+            ChipLogDetail(DataManagement, "Closing on group Communication ");
 
-        // onDone is called
-        ShutdownInternal();
+            // onDone is called
+            ShutdownInternal();
+        }
     }
 
     return err;
diff --git a/src/app/WriteClient.h b/src/app/WriteClient.h
index e447d7d61ed302..35a5996e70959e 100644
--- a/src/app/WriteClient.h
+++ b/src/app/WriteClient.h
@@ -248,6 +248,9 @@ class WriteClientHandle
     /**
      *  Finalize the message and send it to the desired node. The underlying write object will always be released, and the user
      * should not use this object after calling this function.
+     *
+     * Note: In failure cases this will _synchronously_ invoke OnDone on the
+     * WriteClient::Callback before returning.
      */
     CHIP_ERROR SendWriteRequest(const SessionHandle & session, System::Clock::Timeout timeout = kImMessageTimeout);
 
diff --git a/src/controller/WriteInteraction.h b/src/controller/WriteInteraction.h
index e110132b14563d..a79d820616f954 100644
--- a/src/controller/WriteInteraction.h
+++ b/src/controller/WriteInteraction.h
@@ -106,6 +106,11 @@ CHIP_ERROR WriteAttribute(const SessionHandle & sessionHandle, chip::EndpointId
     VerifyOrReturnError(callback != nullptr, CHIP_ERROR_NO_MEMORY);
 
     ReturnErrorOnFailure(app::InteractionModelEngine::GetInstance()->NewWriteClient(handle, callback.get(), aTimedWriteTimeoutMs));
+
+    // At this point the handle will ensure our callback's OnDone is always
+    // called.
+    callback.release();
+
     if (sessionHandle.IsGroupSession())
     {
         ReturnErrorOnFailure(
@@ -119,7 +124,6 @@ CHIP_ERROR WriteAttribute(const SessionHandle & sessionHandle, chip::EndpointId
 
     ReturnErrorOnFailure(handle.SendWriteRequest(sessionHandle));
 
-    callback.release();
     return CHIP_NO_ERROR;
 }