Cleanup CASESession and PASESession lifetime

src/app/OperationalDeviceProxy.cpp

@@ -221,7 +221,7 @@ void OperationalDeviceProxy::HandleCASEConnectionFailure(void * context, CASECli
 
     device->DequeueConnectionSuccessCallbacks(/* executeCallback */ false);
     device->DequeueConnectionFailureCallbacks(error, /* executeCallback */ true);
-    device->DeferCloseCASESession();
+    device->CloseCASESession();
 }
 
 void OperationalDeviceProxy::HandleCASEConnected(void * context, CASEClient * client)
@@ -242,7 +242,7 @@ void OperationalDeviceProxy::HandleCASEConnected(void * context, CASEClient * cl
 
         device->DequeueConnectionFailureCallbacks(CHIP_NO_ERROR, /* executeCallback */ false);
         device->DequeueConnectionSuccessCallbacks(/* executeCallback */ true);
-        device->DeferCloseCASESession();
+        device->CloseCASESession();
     }
 }
 
@@ -280,22 +280,15 @@ void OperationalDeviceProxy::Clear()
     mInitParams = DeviceProxyInitParams();
 }
 
-void OperationalDeviceProxy::CloseCASESessionTask(System::Layer * layer, void * context)
+void OperationalDeviceProxy::CloseCASESession()
 {
-    OperationalDeviceProxy * device = static_cast<OperationalDeviceProxy *>(context);
-    if (device->mCASEClient)
+    if (this->mCASEClient)
     {
-        device->mInitParams.clientPool->Release(device->mCASEClient);
-        device->mCASEClient = nullptr;
+        this->mInitParams.clientPool->Release(this->mCASEClient);
+        this->mCASEClient = nullptr;
     }
 }
 
-void OperationalDeviceProxy::DeferCloseCASESession()
-{
-    // Defer the release for the pending Ack to be sent
-    mSystemLayer->ScheduleWork(CloseCASESessionTask, this);
-}
-
 void OperationalDeviceProxy::OnSessionReleased(const SessionHandle & session)
 {
     VerifyOrReturn(mSecureSession.Contains(session),

src/app/OperationalDeviceProxy.h

@@ -235,7 +235,7 @@ class DLL_EXPORT OperationalDeviceProxy : public DeviceProxy, SessionReleaseDele
 
     static void CloseCASESessionTask(System::Layer * layer, void * context);
 
-    void DeferCloseCASESession();
+    void CloseCASESession();
 
     void EnqueueConnectionCallbacks(Callback::Callback<OnDeviceConnected> * onConnection,
                                     Callback::Callback<OnDeviceConnectionFailure> * onFailure);

src/protocols/secure_channel/CASESession.cpp

@@ -252,11 +252,13 @@ void CASESession::OnResponseTimeout(ExchangeContext * ec)
     VerifyOrReturn(mExchangeCtxt == ec, ChipLogError(SecureChannel, "CASESession::OnResponseTimeout exchange doesn't match"));
     ChipLogError(SecureChannel, "CASESession timed out while waiting for a response from the peer. Current state was %" PRIu8,
                  mState);
-    mDelegate->OnSessionEstablishmentError(CHIP_ERROR_TIMEOUT);
     // Null out mExchangeCtxt so that Clear() doesn't try closing it.  The
     // exchange will handle that.
+    mExchangeCtxt->SetDelegate(nullptr);
     mExchangeCtxt = nullptr;
     Clear();
+    // Do this last in case the delegate frees us.
+    mDelegate->OnSessionEstablishmentError(CHIP_ERROR_TIMEOUT);
 }
 
 CHIP_ERROR CASESession::DeriveSecureSession(CryptoContext & session, CryptoContext::SessionRole role)
@@ -684,9 +686,11 @@ CHIP_ERROR CASESession::HandleSigma2Resume(System::PacketBufferHandle && msg)
     mCASESessionEstablished = true;
 
     // Forget our exchange, as no additional messages are expected from the peer
+    mExchangeCtxt->SetDelegate(nullptr);
     mExchangeCtxt = nullptr;
 
     // Call delegate to indicate session establishment is successful
+    // Do this last in case the delegate frees us.
     mDelegate->OnSessionEstablished();
 
 exit:
@@ -1118,9 +1122,11 @@ CHIP_ERROR CASESession::HandleSigma3(System::PacketBufferHandle && msg)
     mCASESessionEstablished = true;
 
     // Forget our exchange, as no additional messages are expected from the peer
+    mExchangeCtxt->SetDelegate(nullptr);
     mExchangeCtxt = nullptr;
 
     // Call delegate to indicate session establishment is successful
+    // Do this last in case the delegate frees us.
     mDelegate->OnSessionEstablished();
 
 exit:
@@ -1302,15 +1308,17 @@ void CASESession::OnSuccessStatusReport()
     mCASESessionEstablished = true;
 
     // Forget our exchange, as no additional messages are expected from the peer
+    mExchangeCtxt->SetDelegate(nullptr);
     mExchangeCtxt = nullptr;
 
-    // Call delegate to indicate pairing completion
-    mDelegate->OnSessionEstablished();
-
     mState = kInitialized;
 
     // TODO: Set timestamp on the new session, to allow selecting a least-recently-used session for eviction
     // on running out of session contexts.
+
+    // Call delegate to indicate pairing completion.
+    // Do this last in case the delegate frees us.
+    mDelegate->OnSessionEstablished();
 }
 
 CHIP_ERROR CASESession::OnFailureStatusReport(Protocols::SecureChannel::GeneralStatusCode generalCode, uint16_t protocolCode)
@@ -1524,8 +1532,10 @@ CHIP_ERROR CASESession::OnMessageReceived(ExchangeContext * ec, const PayloadHea
     {
         // Null out mExchangeCtxt so that Clear() doesn't try closing it.  The
         // exchange will handle that.
+        mExchangeCtxt->SetDelegate(nullptr);
         mExchangeCtxt = nullptr;
         Clear();
+        // Do this last in case the delegate frees us.
         mDelegate->OnSessionEstablishmentError(err);
     }
     return err;

src/protocols/secure_channel/PASESession.cpp

@@ -343,11 +343,13 @@ void PASESession::OnResponseTimeout(ExchangeContext * ec)
     ChipLogError(SecureChannel,
                  "PASESession timed out while waiting for a response from the peer. Expected message type was %" PRIu8,
                  to_underlying(mNextExpectedMsg));
-    mDelegate->OnSessionEstablishmentError(CHIP_ERROR_TIMEOUT);
     // Null out mExchangeCtxt so that Clear() doesn't try closing it.  The
     // exchange will handle that.
+    mExchangeCtxt->SetDelegate(nullptr);
     mExchangeCtxt = nullptr;
     Clear();
+    // Do this last in case the delegate frees us.
+    mDelegate->OnSessionEstablishmentError(CHIP_ERROR_TIMEOUT);
 }
 
 CHIP_ERROR PASESession::DeriveSecureSession(CryptoContext & session, CryptoContext::SessionRole role)
@@ -824,9 +826,11 @@ CHIP_ERROR PASESession::HandleMsg3(System::PacketBufferHandle && msg)
     mPairingComplete = true;
 
     // Forget our exchange, as no additional messages are expected from the peer
+    mExchangeCtxt->SetDelegate(nullptr);
     mExchangeCtxt = nullptr;
 
     // Call delegate to indicate pairing completion
+    // Do this last in case the delegate frees us.
     mDelegate->OnSessionEstablished();
 
 exit:
@@ -843,9 +847,11 @@ void PASESession::OnSuccessStatusReport()
     mPairingComplete = true;
 
     // Forget our exchange, as no additional messages are expected from the peer
+    mExchangeCtxt->SetDelegate(nullptr);
     mExchangeCtxt = nullptr;
 
     // Call delegate to indicate pairing completion
+    // Do this last in case the delegate frees us.
     mDelegate->OnSessionEstablished();
 }
 
@@ -938,9 +944,11 @@ CHIP_ERROR PASESession::OnMessageReceived(ExchangeContext * exchange, const Payl
     {
         // Null out mExchangeCtxt so that Clear() doesn't try closing it.  The
         // exchange will handle that.
+        mExchangeCtxt->SetDelegate(nullptr);
         mExchangeCtxt = nullptr;
         Clear();
         ChipLogError(SecureChannel, "Failed during PASE session setup. %s", ErrorStr(err));
+        // Do this last in case the delegate frees us.
         mDelegate->OnSessionEstablishmentError(err);
     }
     return err;