From 2f12eab7c7f7a9a174388db5c8b34c69f74b7dea Mon Sep 17 00:00:00 2001
From: Yufeng Wang <yufengwang@google.com>
Date: Sun, 17 Sep 2023 21:26:43 -0700
Subject: [PATCH 1/6] Move private enum State to public as CaseSessionState

---
 src/app/OperationalSessionSetup.cpp | 76 ++++++++++++++---------------
 src/app/OperationalSessionSetup.h   | 38 ++++++++-------
 2 files changed, 58 insertions(+), 56 deletions(-)

diff --git a/src/app/OperationalSessionSetup.cpp b/src/app/OperationalSessionSetup.cpp
index 7b5d70919f2a8a..d6c8cd60c4207c 100644
--- a/src/app/OperationalSessionSetup.cpp
+++ b/src/app/OperationalSessionSetup.cpp
@@ -46,16 +46,16 @@ using chip::AddressResolve::ResolveResult;
 
 namespace chip {
 
-void OperationalSessionSetup::MoveToState(State aTargetState)
+void OperationalSessionSetup::MoveToState(CaseSessionState aTargetState)
 {
     if (mState != aTargetState)
     {
-        ChipLogDetail(Discovery, "OperationalSessionSetup[%u:" ChipLogFormatX64 "]: State change %d --> %d",
+        ChipLogDetail(Discovery, "OperationalSessionSetup[%u:" ChipLogFormatX64 "]: CaseSessionState change %d --> %d",
                       mPeerId.GetFabricIndex(), ChipLogValueX64(mPeerId.GetNodeId()), to_underlying(mState),
                       to_underlying(aTargetState));
 
 #if CHIP_DEVICE_CONFIG_ENABLE_AUTOMATIC_CASE_RETRIES
-        if (mState == State::WaitingForRetry)
+        if (mState == CaseSessionState::WaitingForRetry)
         {
             CancelSessionSetupReattempt();
         }
@@ -63,7 +63,7 @@ void OperationalSessionSetup::MoveToState(State aTargetState)
 
         mState = aTargetState;
 
-        if (aTargetState != State::Connecting)
+        if (aTargetState != CaseSessionState::Connecting)
         {
             CleanupCASEClient();
         }
@@ -72,8 +72,8 @@ void OperationalSessionSetup::MoveToState(State aTargetState)
 
 bool OperationalSessionSetup::AttachToExistingSecureSession()
 {
-    VerifyOrReturnError(mState == State::NeedsAddress || mState == State::ResolvingAddress || mState == State::HasAddress ||
-                            mState == State::WaitingForRetry,
+    VerifyOrReturnError(mState == CaseSessionState::NeedsAddress || mState == CaseSessionState::ResolvingAddress ||
+                            mState == CaseSessionState::HasAddress || mState == CaseSessionState::WaitingForRetry,
                         false);
 
     auto sessionHandle =
@@ -106,40 +106,40 @@ void OperationalSessionSetup::Connect(Callback::Callback<OnDeviceConnected> * on
 
     switch (mState)
     {
-    case State::Uninitialized:
+    case CaseSessionState::Uninitialized:
         err = CHIP_ERROR_INCORRECT_STATE;
         break;
 
-    case State::NeedsAddress:
+    case CaseSessionState::NeedsAddress:
         isConnected = AttachToExistingSecureSession();
         if (!isConnected)
         {
             // LookupPeerAddress could perhaps call back with a result
             // synchronously, so do our state update first.
-            MoveToState(State::ResolvingAddress);
+            MoveToState(CaseSessionState::ResolvingAddress);
             err = LookupPeerAddress();
             if (err != CHIP_NO_ERROR)
             {
                 // Roll back the state change, since we are presumably not in
                 // the middle of a lookup.
-                MoveToState(State::NeedsAddress);
+                MoveToState(CaseSessionState::NeedsAddress);
             }
         }
 
         break;
 
-    case State::ResolvingAddress:
-    case State::WaitingForRetry:
+    case CaseSessionState::ResolvingAddress:
+    case CaseSessionState::WaitingForRetry:
         isConnected = AttachToExistingSecureSession();
         break;
 
-    case State::HasAddress:
+    case CaseSessionState::HasAddress:
         isConnected = AttachToExistingSecureSession();
         if (!isConnected)
         {
-            // We should not actually every be in be in State::HasAddress. This
-            // is because in the same call that we moved to State::HasAddress
-            // we either move to State::Connecting or call
+            // We should not actually every be in be in CaseSessionState::HasAddress. This
+            // is because in the same call that we moved to CaseSessionState::HasAddress
+            // we either move to CaseSessionState::Connecting or call
             // DequeueConnectionCallbacks with an error thus releasing
             // ourselves before any call would reach this section of code.
             err = CHIP_ERROR_INCORRECT_STATE;
@@ -147,10 +147,10 @@ void OperationalSessionSetup::Connect(Callback::Callback<OnDeviceConnected> * on
 
         break;
 
-    case State::Connecting:
+    case CaseSessionState::Connecting:
         break;
 
-    case State::SecureConnected:
+    case CaseSessionState::SecureConnected:
         isConnected = true;
         break;
 
@@ -160,7 +160,7 @@ void OperationalSessionSetup::Connect(Callback::Callback<OnDeviceConnected> * on
 
     if (isConnected)
     {
-        MoveToState(State::SecureConnected);
+        MoveToState(CaseSessionState::SecureConnected);
     }
 
     //
@@ -180,7 +180,7 @@ void OperationalSessionSetup::Connect(Callback::Callback<OnDeviceConnected> * on
 
 void OperationalSessionSetup::UpdateDeviceData(const Transport::PeerAddress & addr, const ReliableMessageProtocolConfig & config)
 {
-    if (mState == State::Uninitialized)
+    if (mState == CaseSessionState::Uninitialized)
     {
         return;
     }
@@ -202,7 +202,7 @@ void OperationalSessionSetup::UpdateDeviceData(const Transport::PeerAddress & ad
         mCASEClient->SetRemoteMRPIntervals(config);
     }
 
-    if (mState != State::ResolvingAddress)
+    if (mState != CaseSessionState::ResolvingAddress)
     {
         ChipLogError(Discovery, "Received UpdateDeviceData in incorrect state");
         DequeueConnectionCallbacks(CHIP_ERROR_INCORRECT_STATE);
@@ -211,7 +211,7 @@ void OperationalSessionSetup::UpdateDeviceData(const Transport::PeerAddress & ad
         return;
     }
 
-    MoveToState(State::HasAddress);
+    MoveToState(CaseSessionState::HasAddress);
     mInitParams.sessionManager->UpdateAllSessionsPeerAddress(mPeerId, addr);
 
     if (mPerformingAddressUpdate)
@@ -233,7 +233,7 @@ void OperationalSessionSetup::UpdateDeviceData(const Transport::PeerAddress & ad
 
     // Move to the ResolvingAddress state, in case we have more results,
     // since we expect to receive results in that state.
-    MoveToState(State::ResolvingAddress);
+    MoveToState(CaseSessionState::ResolvingAddress);
     if (CHIP_NO_ERROR == Resolver::Instance().TryNextResult(mAddressLookupHandle))
     {
         // No need to NotifyRetryHandlers, since we never actually
@@ -257,7 +257,7 @@ CHIP_ERROR OperationalSessionSetup::EstablishConnection(const ReliableMessagePro
         return err;
     }
 
-    MoveToState(State::Connecting);
+    MoveToState(CaseSessionState::Connecting);
 
     return CHIP_NO_ERROR;
 }
@@ -357,7 +357,7 @@ void OperationalSessionSetup::NotifyConnectionCallbacks(Cancelable & failureRead
 
 void OperationalSessionSetup::OnSessionEstablishmentError(CHIP_ERROR error)
 {
-    VerifyOrReturn(mState == State::Connecting,
+    VerifyOrReturn(mState == CaseSessionState::Connecting,
                    ChipLogError(Discovery, "OnSessionEstablishmentError was called while we were not connecting"));
 
     if (CHIP_ERROR_TIMEOUT == error)
@@ -370,7 +370,7 @@ void OperationalSessionSetup::OnSessionEstablishmentError(CHIP_ERROR error)
 
         // Move to the ResolvingAddress state, in case we have more results,
         // since we expect to receive results in that state.
-        MoveToState(State::ResolvingAddress);
+        MoveToState(CaseSessionState::ResolvingAddress);
         if (CHIP_NO_ERROR == Resolver::Instance().TryNextResult(mAddressLookupHandle))
         {
 #if CHIP_DEVICE_CONFIG_ENABLE_AUTOMATIC_CASE_RETRIES
@@ -383,7 +383,7 @@ void OperationalSessionSetup::OnSessionEstablishmentError(CHIP_ERROR error)
         // Moving back to the Connecting state would be a bit of a lie, since we
         // don't have an mCASEClient.  Just go back to NeedsAddress, since
         // that's really where we are now.
-        MoveToState(State::NeedsAddress);
+        MoveToState(CaseSessionState::NeedsAddress);
 
 #if CHIP_DEVICE_CONFIG_ENABLE_AUTOMATIC_CASE_RETRIES
         if (mRemainingAttempts > 0)
@@ -392,7 +392,7 @@ void OperationalSessionSetup::OnSessionEstablishmentError(CHIP_ERROR error)
             CHIP_ERROR err = ScheduleSessionSetupReattempt(reattemptDelay);
             if (err == CHIP_NO_ERROR)
             {
-                MoveToState(State::WaitingForRetry);
+                MoveToState(CaseSessionState::WaitingForRetry);
                 NotifyRetryHandlers(error, remoteMprConfig, reattemptDelay);
                 return;
             }
@@ -406,7 +406,7 @@ void OperationalSessionSetup::OnSessionEstablishmentError(CHIP_ERROR error)
 
 void OperationalSessionSetup::OnSessionEstablished(const SessionHandle & session)
 {
-    VerifyOrReturn(mState == State::Connecting,
+    VerifyOrReturn(mState == CaseSessionState::Connecting,
                    ChipLogError(Discovery, "OnSessionEstablished was called while we were not connecting"));
 
     if (!mSecureSession.Grab(session))
@@ -419,7 +419,7 @@ void OperationalSessionSetup::OnSessionEstablished(const SessionHandle & session
         return;
     }
 
-    MoveToState(State::SecureConnected);
+    MoveToState(CaseSessionState::SecureConnected);
 
     DequeueConnectionCallbacks(CHIP_NO_ERROR);
 }
@@ -482,7 +482,7 @@ CHIP_ERROR OperationalSessionSetup::LookupPeerAddress()
 #endif // CHIP_DEVICE_CONFIG_ENABLE_AUTOMATIC_CASE_RETRIES
 
     // NOTE: This is public API that can be used to update our stored peer
-    // address even when we are in State::Connected, so we do not make any
+    // address even when we are in CaseSessionState::Connected, so we do not make any
     // MoveToState calls in this method.
     if (mAddressLookupHandle.IsActive())
     {
@@ -514,11 +514,11 @@ void OperationalSessionSetup::PerformAddressUpdate()
     }
 
     // We must be newly-allocated to handle this address lookup, so must be in the NeedsAddress state.
-    VerifyOrDie(mState == State::NeedsAddress);
+    VerifyOrDie(mState == CaseSessionState::NeedsAddress);
 
     // We are doing an address lookup whether we have an active session for this peer or not.
     mPerformingAddressUpdate = true;
-    MoveToState(State::ResolvingAddress);
+    MoveToState(CaseSessionState::ResolvingAddress);
     CHIP_ERROR err = LookupPeerAddress();
     if (err != CHIP_NO_ERROR)
     {
@@ -545,10 +545,10 @@ void OperationalSessionSetup::OnNodeAddressResolutionFailed(const PeerId & peerI
     // where multicast DNS does not actually work and hence only the initial
     // unicast DNS-SD queries get a response.
     //
-    // We check for State::ResolvingAddress just in case in the meantime
+    // We check for CaseSessionState::ResolvingAddress just in case in the meantime
     // something weird happened and we are no longer trying to resolve an
     // address.
-    if (mState == State::ResolvingAddress && mResolveAttemptsAllowed > 0)
+    if (mState == CaseSessionState::ResolvingAddress && mResolveAttemptsAllowed > 0)
     {
         ChipLogProgress(Discovery, "Retrying operational DNS-SD discovery. Attempts remaining: %u", mResolveAttemptsAllowed);
 
@@ -592,7 +592,7 @@ void OperationalSessionSetup::UpdateAttemptCount(uint8_t attemptCount)
         return;
     }
 
-    if (mState != State::NeedsAddress)
+    if (mState != CaseSessionState::NeedsAddress)
     {
         // We're in the middle of an attempt already, so decrement attemptCount
         // by 1 to account for that.
@@ -620,7 +620,7 @@ CHIP_ERROR OperationalSessionSetup::ScheduleSessionSetupReattempt(System::Clock:
         return CHIP_ERROR_INCORRECT_STATE;
     }
 
-    MoveToState(State::NeedsAddress);
+    MoveToState(CaseSessionState::NeedsAddress);
     // Stop exponential backoff before our delays get too large.
     //
     // Note that mAttemptsDone is always > 0 here, because we have
@@ -671,7 +671,7 @@ void OperationalSessionSetup::TrySetupAgain(System::Layer * systemLayer, void *
 {
     auto * self = static_cast<OperationalSessionSetup *>(state);
 
-    self->MoveToState(State::ResolvingAddress);
+    self->MoveToState(CaseSessionState::ResolvingAddress);
     CHIP_ERROR err = self->LookupPeerAddress();
     if (err == CHIP_NO_ERROR)
     {
diff --git a/src/app/OperationalSessionSetup.h b/src/app/OperationalSessionSetup.h
index bfa76fbc41d730..acb81a2f69f442 100644
--- a/src/app/OperationalSessionSetup.h
+++ b/src/app/OperationalSessionSetup.h
@@ -109,6 +109,18 @@ class OperationalDeviceProxy : public DeviceProxy
     ScopedNodeId mPeerScopedNodeId;
 };
 
+enum class CaseSessionState : uint8_t
+{
+    Uninitialized,    // Error state: OperationalSessionSetup is useless
+    NeedsAddress,     // No address known, lookup not started yet.
+    ResolvingAddress, // Address lookup in progress.
+    HasAddress,       // Have an address, CASE handshake not started yet.
+    Connecting,       // CASE handshake in progress.
+    SecureConnected,  // CASE session established.
+    WaitingForRetry,  // No address known, but a retry is pending.  Added at
+                      // end to make logs easier to understand.
+};
+
 /**
  * @brief Callback prototype when secure session is established.
  *
@@ -163,14 +175,14 @@ class DLL_EXPORT OperationalSessionSetup : public SessionEstablishmentDelegate,
         mInitParams = params;
         if (params.Validate() != CHIP_NO_ERROR || clientPool == nullptr || releaseDelegate == nullptr)
         {
-            mState = State::Uninitialized;
+            mState = CaseSessionState::Uninitialized;
             return;
         }
 
         mClientPool      = clientPool;
         mPeerId          = peerId;
         mReleaseDelegate = releaseDelegate;
-        mState           = State::NeedsAddress;
+        mState           = CaseSessionState::NeedsAddress;
         mAddressLookupHandle.SetListener(this);
     }
 
@@ -200,6 +212,8 @@ class DLL_EXPORT OperationalSessionSetup : public SessionEstablishmentDelegate,
 
     ScopedNodeId GetPeerId() const { return mPeerId; }
 
+    CaseSessionState GetCaseSessionState() const { return mState; }
+
     static Transport::PeerAddress ToPeerAddress(const Dnssd::ResolvedNodeData & nodeData)
     {
         Inet::InterfaceId interfaceId = Inet::InterfaceId::Null();
@@ -237,23 +251,11 @@ class DLL_EXPORT OperationalSessionSetup : public SessionEstablishmentDelegate,
 #endif // CHIP_DEVICE_CONFIG_ENABLE_AUTOMATIC_CASE_RETRIES
 
 private:
-    enum class State : uint8_t
-    {
-        Uninitialized,    // Error state: OperationalSessionSetup is useless
-        NeedsAddress,     // No address known, lookup not started yet.
-        ResolvingAddress, // Address lookup in progress.
-        HasAddress,       // Have an address, CASE handshake not started yet.
-        Connecting,       // CASE handshake in progress.
-        SecureConnected,  // CASE session established.
-        WaitingForRetry,  // No address known, but a retry is pending.  Added at
-                          // end to make logs easier to understand.
-    };
-
     CASEClientInitParams mInitParams;
     CASEClientPoolDelegate * mClientPool = nullptr;
 
-    // mCASEClient is only non-null if we are in State::Connecting or just
-    // allocated it as part of an attempt to enter State::Connecting.
+    // mCASEClient is only non-null if we are in CaseSessionState::Connecting or just
+    // allocated it as part of an attempt to enter CaseSessionState::Connecting.
     CASEClient * mCASEClient = nullptr;
 
     ScopedNodeId mPeerId;
@@ -270,7 +272,7 @@ class DLL_EXPORT OperationalSessionSetup : public SessionEstablishmentDelegate,
     /// This is used when a node address is required.
     chip::AddressResolve::NodeLookupHandle mAddressLookupHandle;
 
-    State mState = State::Uninitialized;
+    CaseSessionState mState = CaseSessionState::Uninitialized;
 
     bool mPerformingAddressUpdate = false;
 
@@ -283,7 +285,7 @@ class DLL_EXPORT OperationalSessionSetup : public SessionEstablishmentDelegate,
     Callback::CallbackDeque mConnectionRetry;
 #endif // CHIP_DEVICE_CONFIG_ENABLE_AUTOMATIC_CASE_RETRIES
 
-    void MoveToState(State aTargetState);
+    void MoveToState(CaseSessionState aTargetState);
 
     CHIP_ERROR EstablishConnection(const ReliableMessageProtocolConfig & config);
 

From e0dc8e970ca859992ef9ca3a4d89487723e20580 Mon Sep 17 00:00:00 2001
From: Yufeng Wang <yufengwang@google.com>
Date: Tue, 19 Sep 2023 12:57:38 -0700
Subject: [PATCH 2/6] Address review comments

---
 src/app/OperationalSessionSetup.cpp | 86 ++++++++++++++---------------
 src/app/OperationalSessionSetup.h   | 62 +++++++++++++--------
 2 files changed, 82 insertions(+), 66 deletions(-)

diff --git a/src/app/OperationalSessionSetup.cpp b/src/app/OperationalSessionSetup.cpp
index d6c8cd60c4207c..3b4a5588c32a28 100644
--- a/src/app/OperationalSessionSetup.cpp
+++ b/src/app/OperationalSessionSetup.cpp
@@ -46,16 +46,16 @@ using chip::AddressResolve::ResolveResult;
 
 namespace chip {
 
-void OperationalSessionSetup::MoveToState(CaseSessionState aTargetState)
+void OperationalSessionSetup::MoveToState(State aTargetState)
 {
     if (mState != aTargetState)
     {
-        ChipLogDetail(Discovery, "OperationalSessionSetup[%u:" ChipLogFormatX64 "]: CaseSessionState change %d --> %d",
+        ChipLogDetail(Discovery, "OperationalSessionSetup[%u:" ChipLogFormatX64 "]: State change %d --> %d",
                       mPeerId.GetFabricIndex(), ChipLogValueX64(mPeerId.GetNodeId()), to_underlying(mState),
                       to_underlying(aTargetState));
 
 #if CHIP_DEVICE_CONFIG_ENABLE_AUTOMATIC_CASE_RETRIES
-        if (mState == CaseSessionState::WaitingForRetry)
+        if (mState == State::WaitingForRetry)
         {
             CancelSessionSetupReattempt();
         }
@@ -63,7 +63,7 @@ void OperationalSessionSetup::MoveToState(CaseSessionState aTargetState)
 
         mState = aTargetState;
 
-        if (aTargetState != CaseSessionState::Connecting)
+        if (aTargetState != State::Connecting)
         {
             CleanupCASEClient();
         }
@@ -72,8 +72,8 @@ void OperationalSessionSetup::MoveToState(CaseSessionState aTargetState)
 
 bool OperationalSessionSetup::AttachToExistingSecureSession()
 {
-    VerifyOrReturnError(mState == CaseSessionState::NeedsAddress || mState == CaseSessionState::ResolvingAddress ||
-                            mState == CaseSessionState::HasAddress || mState == CaseSessionState::WaitingForRetry,
+    VerifyOrReturnError(mState == State::NeedsAddress || mState == State::ResolvingAddress || mState == State::HasAddress ||
+                            mState == State::WaitingForRetry,
                         false);
 
     auto sessionHandle =
@@ -91,8 +91,8 @@ bool OperationalSessionSetup::AttachToExistingSecureSession()
     return true;
 }
 
-void OperationalSessionSetup::Connect(Callback::Callback<OnDeviceConnected> * onConnection,
-                                      Callback::Callback<OnDeviceConnectionFailure> * onFailure)
+void OperationalSessionSetup::Connect(Callback::Callback<OnSetupSuccessful> * onConnection,
+                                      Callback::Callback<OnSetupFailure> * onFailure)
 {
     CHIP_ERROR err   = CHIP_NO_ERROR;
     bool isConnected = false;
@@ -106,40 +106,40 @@ void OperationalSessionSetup::Connect(Callback::Callback<OnDeviceConnected> * on
 
     switch (mState)
     {
-    case CaseSessionState::Uninitialized:
+    case State::Uninitialized:
         err = CHIP_ERROR_INCORRECT_STATE;
         break;
 
-    case CaseSessionState::NeedsAddress:
+    case State::NeedsAddress:
         isConnected = AttachToExistingSecureSession();
         if (!isConnected)
         {
             // LookupPeerAddress could perhaps call back with a result
             // synchronously, so do our state update first.
-            MoveToState(CaseSessionState::ResolvingAddress);
+            MoveToState(State::ResolvingAddress);
             err = LookupPeerAddress();
             if (err != CHIP_NO_ERROR)
             {
                 // Roll back the state change, since we are presumably not in
                 // the middle of a lookup.
-                MoveToState(CaseSessionState::NeedsAddress);
+                MoveToState(State::NeedsAddress);
             }
         }
 
         break;
 
-    case CaseSessionState::ResolvingAddress:
-    case CaseSessionState::WaitingForRetry:
+    case State::ResolvingAddress:
+    case State::WaitingForRetry:
         isConnected = AttachToExistingSecureSession();
         break;
 
-    case CaseSessionState::HasAddress:
+    case State::HasAddress:
         isConnected = AttachToExistingSecureSession();
         if (!isConnected)
         {
-            // We should not actually every be in be in CaseSessionState::HasAddress. This
-            // is because in the same call that we moved to CaseSessionState::HasAddress
-            // we either move to CaseSessionState::Connecting or call
+            // We should not actually every be in be in State::HasAddress. This
+            // is because in the same call that we moved to State::HasAddress
+            // we either move to State::Connecting or call
             // DequeueConnectionCallbacks with an error thus releasing
             // ourselves before any call would reach this section of code.
             err = CHIP_ERROR_INCORRECT_STATE;
@@ -147,10 +147,10 @@ void OperationalSessionSetup::Connect(Callback::Callback<OnDeviceConnected> * on
 
         break;
 
-    case CaseSessionState::Connecting:
+    case State::Connecting:
         break;
 
-    case CaseSessionState::SecureConnected:
+    case State::SecureConnected:
         isConnected = true;
         break;
 
@@ -160,7 +160,7 @@ void OperationalSessionSetup::Connect(Callback::Callback<OnDeviceConnected> * on
 
     if (isConnected)
     {
-        MoveToState(CaseSessionState::SecureConnected);
+        MoveToState(State::SecureConnected);
     }
 
     //
@@ -180,7 +180,7 @@ void OperationalSessionSetup::Connect(Callback::Callback<OnDeviceConnected> * on
 
 void OperationalSessionSetup::UpdateDeviceData(const Transport::PeerAddress & addr, const ReliableMessageProtocolConfig & config)
 {
-    if (mState == CaseSessionState::Uninitialized)
+    if (mState == State::Uninitialized)
     {
         return;
     }
@@ -202,7 +202,7 @@ void OperationalSessionSetup::UpdateDeviceData(const Transport::PeerAddress & ad
         mCASEClient->SetRemoteMRPIntervals(config);
     }
 
-    if (mState != CaseSessionState::ResolvingAddress)
+    if (mState != State::ResolvingAddress)
     {
         ChipLogError(Discovery, "Received UpdateDeviceData in incorrect state");
         DequeueConnectionCallbacks(CHIP_ERROR_INCORRECT_STATE);
@@ -211,7 +211,7 @@ void OperationalSessionSetup::UpdateDeviceData(const Transport::PeerAddress & ad
         return;
     }
 
-    MoveToState(CaseSessionState::HasAddress);
+    MoveToState(State::HasAddress);
     mInitParams.sessionManager->UpdateAllSessionsPeerAddress(mPeerId, addr);
 
     if (mPerformingAddressUpdate)
@@ -233,7 +233,7 @@ void OperationalSessionSetup::UpdateDeviceData(const Transport::PeerAddress & ad
 
     // Move to the ResolvingAddress state, in case we have more results,
     // since we expect to receive results in that state.
-    MoveToState(CaseSessionState::ResolvingAddress);
+    MoveToState(State::ResolvingAddress);
     if (CHIP_NO_ERROR == Resolver::Instance().TryNextResult(mAddressLookupHandle))
     {
         // No need to NotifyRetryHandlers, since we never actually
@@ -257,13 +257,13 @@ CHIP_ERROR OperationalSessionSetup::EstablishConnection(const ReliableMessagePro
         return err;
     }
 
-    MoveToState(CaseSessionState::Connecting);
+    MoveToState(State::Connecting);
 
     return CHIP_NO_ERROR;
 }
 
-void OperationalSessionSetup::EnqueueConnectionCallbacks(Callback::Callback<OnDeviceConnected> * onConnection,
-                                                         Callback::Callback<OnDeviceConnectionFailure> * onFailure)
+void OperationalSessionSetup::EnqueueConnectionCallbacks(Callback::Callback<OnSetupSuccessful> * onConnection,
+                                                         Callback::Callback<OnSetupFailure> * onFailure)
 {
     if (onConnection != nullptr)
     {
@@ -357,7 +357,7 @@ void OperationalSessionSetup::NotifyConnectionCallbacks(Cancelable & failureRead
 
 void OperationalSessionSetup::OnSessionEstablishmentError(CHIP_ERROR error)
 {
-    VerifyOrReturn(mState == CaseSessionState::Connecting,
+    VerifyOrReturn(mState == State::Connecting,
                    ChipLogError(Discovery, "OnSessionEstablishmentError was called while we were not connecting"));
 
     if (CHIP_ERROR_TIMEOUT == error)
@@ -370,7 +370,7 @@ void OperationalSessionSetup::OnSessionEstablishmentError(CHIP_ERROR error)
 
         // Move to the ResolvingAddress state, in case we have more results,
         // since we expect to receive results in that state.
-        MoveToState(CaseSessionState::ResolvingAddress);
+        MoveToState(State::ResolvingAddress);
         if (CHIP_NO_ERROR == Resolver::Instance().TryNextResult(mAddressLookupHandle))
         {
 #if CHIP_DEVICE_CONFIG_ENABLE_AUTOMATIC_CASE_RETRIES
@@ -383,7 +383,7 @@ void OperationalSessionSetup::OnSessionEstablishmentError(CHIP_ERROR error)
         // Moving back to the Connecting state would be a bit of a lie, since we
         // don't have an mCASEClient.  Just go back to NeedsAddress, since
         // that's really where we are now.
-        MoveToState(CaseSessionState::NeedsAddress);
+        MoveToState(State::NeedsAddress);
 
 #if CHIP_DEVICE_CONFIG_ENABLE_AUTOMATIC_CASE_RETRIES
         if (mRemainingAttempts > 0)
@@ -392,7 +392,7 @@ void OperationalSessionSetup::OnSessionEstablishmentError(CHIP_ERROR error)
             CHIP_ERROR err = ScheduleSessionSetupReattempt(reattemptDelay);
             if (err == CHIP_NO_ERROR)
             {
-                MoveToState(CaseSessionState::WaitingForRetry);
+                MoveToState(State::WaitingForRetry);
                 NotifyRetryHandlers(error, remoteMprConfig, reattemptDelay);
                 return;
             }
@@ -406,7 +406,7 @@ void OperationalSessionSetup::OnSessionEstablishmentError(CHIP_ERROR error)
 
 void OperationalSessionSetup::OnSessionEstablished(const SessionHandle & session)
 {
-    VerifyOrReturn(mState == CaseSessionState::Connecting,
+    VerifyOrReturn(mState == State::Connecting,
                    ChipLogError(Discovery, "OnSessionEstablished was called while we were not connecting"));
 
     if (!mSecureSession.Grab(session))
@@ -419,7 +419,7 @@ void OperationalSessionSetup::OnSessionEstablished(const SessionHandle & session
         return;
     }
 
-    MoveToState(CaseSessionState::SecureConnected);
+    MoveToState(State::SecureConnected);
 
     DequeueConnectionCallbacks(CHIP_NO_ERROR);
 }
@@ -482,7 +482,7 @@ CHIP_ERROR OperationalSessionSetup::LookupPeerAddress()
 #endif // CHIP_DEVICE_CONFIG_ENABLE_AUTOMATIC_CASE_RETRIES
 
     // NOTE: This is public API that can be used to update our stored peer
-    // address even when we are in CaseSessionState::Connected, so we do not make any
+    // address even when we are in State::Connected, so we do not make any
     // MoveToState calls in this method.
     if (mAddressLookupHandle.IsActive())
     {
@@ -514,11 +514,11 @@ void OperationalSessionSetup::PerformAddressUpdate()
     }
 
     // We must be newly-allocated to handle this address lookup, so must be in the NeedsAddress state.
-    VerifyOrDie(mState == CaseSessionState::NeedsAddress);
+    VerifyOrDie(mState == State::NeedsAddress);
 
     // We are doing an address lookup whether we have an active session for this peer or not.
     mPerformingAddressUpdate = true;
-    MoveToState(CaseSessionState::ResolvingAddress);
+    MoveToState(State::ResolvingAddress);
     CHIP_ERROR err = LookupPeerAddress();
     if (err != CHIP_NO_ERROR)
     {
@@ -545,10 +545,10 @@ void OperationalSessionSetup::OnNodeAddressResolutionFailed(const PeerId & peerI
     // where multicast DNS does not actually work and hence only the initial
     // unicast DNS-SD queries get a response.
     //
-    // We check for CaseSessionState::ResolvingAddress just in case in the meantime
+    // We check for State::ResolvingAddress just in case in the meantime
     // something weird happened and we are no longer trying to resolve an
     // address.
-    if (mState == CaseSessionState::ResolvingAddress && mResolveAttemptsAllowed > 0)
+    if (mState == State::ResolvingAddress && mResolveAttemptsAllowed > 0)
     {
         ChipLogProgress(Discovery, "Retrying operational DNS-SD discovery. Attempts remaining: %u", mResolveAttemptsAllowed);
 
@@ -592,7 +592,7 @@ void OperationalSessionSetup::UpdateAttemptCount(uint8_t attemptCount)
         return;
     }
 
-    if (mState != CaseSessionState::NeedsAddress)
+    if (mState != State::NeedsAddress)
     {
         // We're in the middle of an attempt already, so decrement attemptCount
         // by 1 to account for that.
@@ -620,7 +620,7 @@ CHIP_ERROR OperationalSessionSetup::ScheduleSessionSetupReattempt(System::Clock:
         return CHIP_ERROR_INCORRECT_STATE;
     }
 
-    MoveToState(CaseSessionState::NeedsAddress);
+    MoveToState(State::NeedsAddress);
     // Stop exponential backoff before our delays get too large.
     //
     // Note that mAttemptsDone is always > 0 here, because we have
@@ -671,7 +671,7 @@ void OperationalSessionSetup::TrySetupAgain(System::Layer * systemLayer, void *
 {
     auto * self = static_cast<OperationalSessionSetup *>(state);
 
-    self->MoveToState(CaseSessionState::ResolvingAddress);
+    self->MoveToState(State::ResolvingAddress);
     CHIP_ERROR err = self->LookupPeerAddress();
     if (err == CHIP_NO_ERROR)
     {
@@ -683,7 +683,7 @@ void OperationalSessionSetup::TrySetupAgain(System::Layer * systemLayer, void *
     // Do not touch `self` instance anymore; it has been destroyed in DequeueConnectionCallbacks.
 }
 
-void OperationalSessionSetup::AddRetryHandler(Callback::Callback<OnDeviceConnectionRetry> * onRetry)
+void OperationalSessionSetup::AddRetryHandler(Callback::Callback<OnSetupRetry> * onRetry)
 {
     mConnectionRetry.Enqueue(onRetry->Cancel());
 }
diff --git a/src/app/OperationalSessionSetup.h b/src/app/OperationalSessionSetup.h
index acb81a2f69f442..dfbd6a9a9325c8 100644
--- a/src/app/OperationalSessionSetup.h
+++ b/src/app/OperationalSessionSetup.h
@@ -109,18 +109,6 @@ class OperationalDeviceProxy : public DeviceProxy
     ScopedNodeId mPeerScopedNodeId;
 };
 
-enum class CaseSessionState : uint8_t
-{
-    Uninitialized,    // Error state: OperationalSessionSetup is useless
-    NeedsAddress,     // No address known, lookup not started yet.
-    ResolvingAddress, // Address lookup in progress.
-    HasAddress,       // Have an address, CASE handshake not started yet.
-    Connecting,       // CASE handshake in progress.
-    SecureConnected,  // CASE session established.
-    WaitingForRetry,  // No address known, but a retry is pending.  Added at
-                      // end to make logs easier to understand.
-};
-
 /**
  * @brief Callback prototype when secure session is established.
  *
@@ -167,6 +155,34 @@ typedef void (*OnDeviceConnectionRetry)(void * context, const ScopedNodeId & pee
 class DLL_EXPORT OperationalSessionSetup : public SessionEstablishmentDelegate, public AddressResolve::NodeListener
 {
 public:
+    enum class State : uint8_t
+    {
+        Uninitialized,    // Error state: OperationalSessionSetup is useless
+        NeedsAddress,     // No address known, lookup not started yet.
+        ResolvingAddress, // Address lookup in progress.
+        HasAddress,       // Have an address, CASE handshake not started yet.
+        Connecting,       // CASE handshake in progress.
+        SecureConnected,  // CASE session established.
+        WaitingForRetry,  // No address known, but a retry is pending.  Added at
+                          // end to make logs easier to understand.
+    };
+
+    struct ConnnectionFailureInfo
+    {
+        const ScopedNodeId & peerId;
+        CHIP_ERROR error;
+        OperationalSessionSetup::State sessionState;
+
+        ConnnectionFailureInfo(const ScopedNodeId & peer, CHIP_ERROR err, OperationalSessionSetup::State state) :
+            peerId(peer), error(err), sessionState(state)
+        {}
+    };
+
+    using OnSetupSuccessful      = OnDeviceConnected;
+    using OnSetupFailure         = OnDeviceConnectionFailure;
+    using OnSetupRetry           = OnDeviceConnectionRetry;
+    using OnExtendedSetupFailure = void (*)(void * context, const ConnnectionFailureInfo & failureInfo);
+
     ~OperationalSessionSetup() override;
 
     OperationalSessionSetup(const CASEClientInitParams & params, CASEClientPoolDelegate * clientPool, ScopedNodeId peerId,
@@ -175,14 +191,14 @@ class DLL_EXPORT OperationalSessionSetup : public SessionEstablishmentDelegate,
         mInitParams = params;
         if (params.Validate() != CHIP_NO_ERROR || clientPool == nullptr || releaseDelegate == nullptr)
         {
-            mState = CaseSessionState::Uninitialized;
+            mState = State::Uninitialized;
             return;
         }
 
         mClientPool      = clientPool;
         mPeerId          = peerId;
         mReleaseDelegate = releaseDelegate;
-        mState           = CaseSessionState::NeedsAddress;
+        mState           = State::NeedsAddress;
         mAddressLookupHandle.SetListener(this);
     }
 
@@ -202,7 +218,7 @@ class DLL_EXPORT OperationalSessionSetup : public SessionEstablishmentDelegate,
      * cases that are detected synchronously (e.g. inability to start an address
      * lookup).
      */
-    void Connect(Callback::Callback<OnDeviceConnected> * onConnection, Callback::Callback<OnDeviceConnectionFailure> * onFailure);
+    void Connect(Callback::Callback<OnSetupSuccessful> * onConnection, Callback::Callback<OnSetupFailure> * onFailure);
 
     bool IsForAddressUpdate() const { return mPerformingAddressUpdate; }
 
@@ -212,7 +228,7 @@ class DLL_EXPORT OperationalSessionSetup : public SessionEstablishmentDelegate,
 
     ScopedNodeId GetPeerId() const { return mPeerId; }
 
-    CaseSessionState GetCaseSessionState() const { return mState; }
+    State GetOperationalSessionSetupState() const { return mState; }
 
     static Transport::PeerAddress ToPeerAddress(const Dnssd::ResolvedNodeData & nodeData)
     {
@@ -247,15 +263,15 @@ class DLL_EXPORT OperationalSessionSetup : public SessionEstablishmentDelegate,
     void UpdateAttemptCount(uint8_t attemptCount);
 
     // Add a retry handler for this session setup.
-    void AddRetryHandler(Callback::Callback<OnDeviceConnectionRetry> * onRetry);
+    void AddRetryHandler(Callback::Callback<OnSetupRetry> * onRetry);
 #endif // CHIP_DEVICE_CONFIG_ENABLE_AUTOMATIC_CASE_RETRIES
 
 private:
     CASEClientInitParams mInitParams;
     CASEClientPoolDelegate * mClientPool = nullptr;
 
-    // mCASEClient is only non-null if we are in CaseSessionState::Connecting or just
-    // allocated it as part of an attempt to enter CaseSessionState::Connecting.
+    // mCASEClient is only non-null if we are in State::Connecting or just
+    // allocated it as part of an attempt to enter State::Connecting.
     CASEClient * mCASEClient = nullptr;
 
     ScopedNodeId mPeerId;
@@ -272,7 +288,7 @@ class DLL_EXPORT OperationalSessionSetup : public SessionEstablishmentDelegate,
     /// This is used when a node address is required.
     chip::AddressResolve::NodeLookupHandle mAddressLookupHandle;
 
-    CaseSessionState mState = CaseSessionState::Uninitialized;
+    State mState = State::Uninitialized;
 
     bool mPerformingAddressUpdate = false;
 
@@ -285,7 +301,7 @@ class DLL_EXPORT OperationalSessionSetup : public SessionEstablishmentDelegate,
     Callback::CallbackDeque mConnectionRetry;
 #endif // CHIP_DEVICE_CONFIG_ENABLE_AUTOMATIC_CASE_RETRIES
 
-    void MoveToState(CaseSessionState aTargetState);
+    void MoveToState(State aTargetState);
 
     CHIP_ERROR EstablishConnection(const ReliableMessageProtocolConfig & config);
 
@@ -300,8 +316,8 @@ class DLL_EXPORT OperationalSessionSetup : public SessionEstablishmentDelegate,
 
     void CleanupCASEClient();
 
-    void EnqueueConnectionCallbacks(Callback::Callback<OnDeviceConnected> * onConnection,
-                                    Callback::Callback<OnDeviceConnectionFailure> * onFailure);
+    void EnqueueConnectionCallbacks(Callback::Callback<OnSetupSuccessful> * onConnection,
+                                    Callback::Callback<OnSetupFailure> * onFailure);
 
     enum class ReleaseBehavior
     {

From 87d717516bfd9fcf1534386a5d55c3bb3b226792 Mon Sep 17 00:00:00 2001
From: Yufeng Wang <yufengwang@google.com>
Date: Wed, 20 Sep 2023 10:27:26 -0700
Subject: [PATCH 3/6] Rename OnExtendedSetupFailure to OnSetupFailure

---
 src/app/OperationalSessionSetup.cpp | 75 ++++++++++++++++++++---------
 src/app/OperationalSessionSetup.h   | 51 ++++++++++++++------
 2 files changed, 89 insertions(+), 37 deletions(-)

diff --git a/src/app/OperationalSessionSetup.cpp b/src/app/OperationalSessionSetup.cpp
index 3b4a5588c32a28..b7a85dfe547fca 100644
--- a/src/app/OperationalSessionSetup.cpp
+++ b/src/app/OperationalSessionSetup.cpp
@@ -91,8 +91,15 @@ bool OperationalSessionSetup::AttachToExistingSecureSession()
     return true;
 }
 
-void OperationalSessionSetup::Connect(Callback::Callback<OnSetupSuccessful> * onConnection,
-                                      Callback::Callback<OnSetupFailure> * onFailure)
+void OperationalSessionSetup::Connect(Callback::Callback<OnDeviceConnected> * onConnection,
+                                      Callback::Callback<OnDeviceConnectionFailure> * onFailure)
+{
+    Connect(onConnection, onFailure, nullptr);
+}
+
+void OperationalSessionSetup::Connect(Callback::Callback<OnDeviceConnected> * onConnection,
+                                      Callback::Callback<OnDeviceConnectionFailure> * onFailure,
+                                      Callback::Callback<OnSetupFailure> * onSetupFailure)
 {
     CHIP_ERROR err   = CHIP_NO_ERROR;
     bool isConnected = false;
@@ -102,7 +109,7 @@ void OperationalSessionSetup::Connect(Callback::Callback<OnSetupSuccessful> * on
     // If anything goes wrong below, we'll trigger failures (including any queued from
     // a previous iteration which in theory shouldn't happen, but this is written to be more defensive)
     //
-    EnqueueConnectionCallbacks(onConnection, onFailure);
+    EnqueueConnectionCallbacks(onConnection, onFailure, onSetupFailure);
 
     switch (mState)
     {
@@ -170,7 +177,7 @@ void OperationalSessionSetup::Connect(Callback::Callback<OnSetupSuccessful> * on
     //
     if (err != CHIP_NO_ERROR || isConnected)
     {
-        DequeueConnectionCallbacks(err);
+        DequeueConnectionCallbacks(err, mState);
         // Do not touch `this` instance anymore; it has been destroyed in DequeueConnectionCallbacks.
         // While it is odd to have an explicit return here at the end of the function, we do so
         // as a precaution in case someone later on adds something to the end of this function.
@@ -205,7 +212,7 @@ void OperationalSessionSetup::UpdateDeviceData(const Transport::PeerAddress & ad
     if (mState != State::ResolvingAddress)
     {
         ChipLogError(Discovery, "Received UpdateDeviceData in incorrect state");
-        DequeueConnectionCallbacks(CHIP_ERROR_INCORRECT_STATE);
+        DequeueConnectionCallbacks(CHIP_ERROR_INCORRECT_STATE, mState);
         // Do not touch `this` instance anymore; it has been destroyed in
         // DequeueConnectionCallbacks.
         return;
@@ -217,7 +224,7 @@ void OperationalSessionSetup::UpdateDeviceData(const Transport::PeerAddress & ad
     if (mPerformingAddressUpdate)
     {
         // Nothing else to do here.
-        DequeueConnectionCallbacks(CHIP_NO_ERROR);
+        DequeueConnectionCallbacks(CHIP_NO_ERROR, State::HasAddress);
         // Do not touch `this` instance anymore; it has been destroyed in DequeueConnectionCallbacks.
         return;
     }
@@ -241,7 +248,7 @@ void OperationalSessionSetup::UpdateDeviceData(const Transport::PeerAddress & ad
         return;
     }
 
-    DequeueConnectionCallbacks(err);
+    DequeueConnectionCallbacks(err, State::ResolvingAddress);
     // Do not touch `this` instance anymore; it has been destroyed in DequeueConnectionCallbacks.
 }
 
@@ -262,8 +269,9 @@ CHIP_ERROR OperationalSessionSetup::EstablishConnection(const ReliableMessagePro
     return CHIP_NO_ERROR;
 }
 
-void OperationalSessionSetup::EnqueueConnectionCallbacks(Callback::Callback<OnSetupSuccessful> * onConnection,
-                                                         Callback::Callback<OnSetupFailure> * onFailure)
+void OperationalSessionSetup::EnqueueConnectionCallbacks(Callback::Callback<OnDeviceConnected> * onConnection,
+                                                         Callback::Callback<OnDeviceConnectionFailure> * onFailure,
+                                                         Callback::Callback<OnSetupFailure> * onSetupFailure)
 {
     if (onConnection != nullptr)
     {
@@ -274,11 +282,16 @@ void OperationalSessionSetup::EnqueueConnectionCallbacks(Callback::Callback<OnSe
     {
         mConnectionFailure.Enqueue(onFailure->Cancel());
     }
+
+    if (onSetupFailure != nullptr)
+    {
+        mSetupFailure.Enqueue(onSetupFailure->Cancel());
+    }
 }
 
-void OperationalSessionSetup::DequeueConnectionCallbacks(CHIP_ERROR error, ReleaseBehavior releaseBehavior)
+void OperationalSessionSetup::DequeueConnectionCallbacks(CHIP_ERROR error, State state, ReleaseBehavior releaseBehavior)
 {
-    Cancelable failureReady, successReady;
+    Cancelable failureReady, setupFailureReady, successReady;
 
     //
     // Dequeue both failure and success callback lists into temporary stack args before invoking either of them.
@@ -286,6 +299,7 @@ void OperationalSessionSetup::DequeueConnectionCallbacks(CHIP_ERROR error, Relea
     // since the callee may destroy this object as part of that callback.
     //
     mConnectionFailure.DequeueAll(failureReady);
+    mSetupFailure.DequeueAll(setupFailureReady);
     mConnectionSuccess.DequeueAll(successReady);
 
 #if CHIP_DEVICE_CONFIG_ENABLE_AUTOMATIC_CASE_RETRIES
@@ -311,11 +325,12 @@ void OperationalSessionSetup::DequeueConnectionCallbacks(CHIP_ERROR error, Relea
 
     // DO NOT touch any members of this object after this point.  It's dead.
 
-    NotifyConnectionCallbacks(failureReady, successReady, error, peerId, performingAddressUpdate, exchangeMgr,
-                              optionalSessionHandle);
+    NotifyConnectionCallbacks(failureReady, setupFailureReady, successReady, error, state, peerId, performingAddressUpdate,
+                              exchangeMgr, optionalSessionHandle);
 }
 
-void OperationalSessionSetup::NotifyConnectionCallbacks(Cancelable & failureReady, Cancelable & successReady, CHIP_ERROR error,
+void OperationalSessionSetup::NotifyConnectionCallbacks(Cancelable & failureReady, Cancelable & setupFailureReady,
+                                                        Cancelable & successReady, CHIP_ERROR error, State state,
                                                         const ScopedNodeId & peerId, bool performingAddressUpdate,
                                                         Messaging::ExchangeManager * exchangeMgr,
                                                         const Optional<SessionHandle> & optionalSessionHandle)
@@ -339,6 +354,22 @@ void OperationalSessionSetup::NotifyConnectionCallbacks(Cancelable & failureRead
         }
     }
 
+    while (setupFailureReady.mNext != &setupFailureReady)
+    {
+        // We expect that we only have callbacks if we are not performing just address update.
+        VerifyOrDie(!performingAddressUpdate);
+        Callback::Callback<OnSetupFailure> * cb = Callback::Callback<OnSetupFailure>::FromCancelable(setupFailureReady.mNext);
+
+        cb->Cancel();
+
+        if (error != CHIP_NO_ERROR)
+        {
+            // Initialize the ConnnectionFailureInfo object
+            ConnnectionFailureInfo failureInfo(peerId, error, state);
+            cb->mCall(cb->mContext, failureInfo);
+        }
+    }
+
     while (successReady.mNext != &successReady)
     {
         // We expect that we only have callbacks if we are not performing just address update.
@@ -400,7 +431,7 @@ void OperationalSessionSetup::OnSessionEstablishmentError(CHIP_ERROR error)
 #endif // CHIP_DEVICE_CONFIG_ENABLE_AUTOMATIC_CASE_RETRIES
     }
 
-    DequeueConnectionCallbacks(error);
+    DequeueConnectionCallbacks(error, mState);
     // Do not touch `this` instance anymore; it has been destroyed in DequeueConnectionCallbacks.
 }
 
@@ -413,7 +444,7 @@ void OperationalSessionSetup::OnSessionEstablished(const SessionHandle & session
     {
         // Got an invalid session, just dispatch an error.  We have to do this
         // so we don't leak.
-        DequeueConnectionCallbacks(CHIP_ERROR_INCORRECT_STATE);
+        DequeueConnectionCallbacks(CHIP_ERROR_INCORRECT_STATE, State::Connecting);
 
         // Do not touch `this` instance anymore; it has been destroyed in DequeueConnectionCallbacks.
         return;
@@ -421,7 +452,7 @@ void OperationalSessionSetup::OnSessionEstablished(const SessionHandle & session
 
     MoveToState(State::SecureConnected);
 
-    DequeueConnectionCallbacks(CHIP_NO_ERROR);
+    DequeueConnectionCallbacks(CHIP_NO_ERROR, State::SecureConnected);
 }
 
 void OperationalSessionSetup::CleanupCASEClient()
@@ -461,7 +492,7 @@ OperationalSessionSetup::~OperationalSessionSetup()
     CancelSessionSetupReattempt();
 #endif // CHIP_DEVICE_CONFIG_ENABLE_AUTOMATIC_CASE_RETRIES
 
-    DequeueConnectionCallbacks(CHIP_ERROR_CANCELLED, ReleaseBehavior::DoNotRelease);
+    DequeueConnectionCallbacks(CHIP_ERROR_CANCELLED, mState, ReleaseBehavior::DoNotRelease);
 }
 
 CHIP_ERROR OperationalSessionSetup::LookupPeerAddress()
@@ -523,7 +554,7 @@ void OperationalSessionSetup::PerformAddressUpdate()
     if (err != CHIP_NO_ERROR)
     {
         ChipLogError(Discovery, "Failed to look up peer address: %" CHIP_ERROR_FORMAT, err.Format());
-        DequeueConnectionCallbacks(err);
+        DequeueConnectionCallbacks(err, State::ResolvingAddress);
         // Do not touch `this` instance anymore; it has been destroyed in DequeueConnectionCallbacks.
         return;
     }
@@ -579,7 +610,7 @@ void OperationalSessionSetup::OnNodeAddressResolutionFailed(const PeerId & peerI
 #endif
 
     // No need to modify any variables in `this` since call below releases `this`.
-    DequeueConnectionCallbacks(reason);
+    DequeueConnectionCallbacks(reason, mState);
     // Do not touch `this` instance anymore; it has been destroyed in DequeueConnectionCallbacks.
 }
 
@@ -679,11 +710,11 @@ void OperationalSessionSetup::TrySetupAgain(System::Layer * systemLayer, void *
     }
 
     // Give up; we could not start a lookup.
-    self->DequeueConnectionCallbacks(err);
+    self->DequeueConnectionCallbacks(err, State::ResolvingAddress);
     // Do not touch `self` instance anymore; it has been destroyed in DequeueConnectionCallbacks.
 }
 
-void OperationalSessionSetup::AddRetryHandler(Callback::Callback<OnSetupRetry> * onRetry)
+void OperationalSessionSetup::AddRetryHandler(Callback::Callback<OnDeviceConnectionRetry> * onRetry)
 {
     mConnectionRetry.Enqueue(onRetry->Cancel());
 }
diff --git a/src/app/OperationalSessionSetup.h b/src/app/OperationalSessionSetup.h
index dfbd6a9a9325c8..847b3cea0d303f 100644
--- a/src/app/OperationalSessionSetup.h
+++ b/src/app/OperationalSessionSetup.h
@@ -178,10 +178,7 @@ class DLL_EXPORT OperationalSessionSetup : public SessionEstablishmentDelegate,
         {}
     };
 
-    using OnSetupSuccessful      = OnDeviceConnected;
-    using OnSetupFailure         = OnDeviceConnectionFailure;
-    using OnSetupRetry           = OnDeviceConnectionRetry;
-    using OnExtendedSetupFailure = void (*)(void * context, const ConnnectionFailureInfo & failureInfo);
+    using OnSetupFailure = void (*)(void * context, const ConnnectionFailureInfo & failureInfo);
 
     ~OperationalSessionSetup() override;
 
@@ -218,7 +215,26 @@ class DLL_EXPORT OperationalSessionSetup : public SessionEstablishmentDelegate,
      * cases that are detected synchronously (e.g. inability to start an address
      * lookup).
      */
-    void Connect(Callback::Callback<OnSetupSuccessful> * onConnection, Callback::Callback<OnSetupFailure> * onFailure);
+    void Connect(Callback::Callback<OnDeviceConnected> * onConnection, Callback::Callback<OnDeviceConnectionFailure> * onFailure);
+
+    /*
+     * This function can be called to establish a secure session with the device.
+     *
+     * The device is expected to have been commissioned, A CASE session
+     * setup will be triggered.
+     *
+     * On establishing the session, the callback function `onConnection` will be called. If the
+     * session setup fails, `onFailure` and `onSetupFailure` will be called.
+     *
+     * If the session already exists, `onConnection` will be called immediately,
+     * before the Connect call returns.
+     *
+     * `onFailure` and `onSetupFailure` may be called before the Connect call returns, for error
+     * cases that are detected synchronously (e.g. inability to start an address
+     * lookup).
+     */
+    void Connect(Callback::Callback<OnDeviceConnected> * onConnection, Callback::Callback<OnDeviceConnectionFailure> * onFailure,
+                 Callback::Callback<OnSetupFailure> * onSetupFailure);
 
     bool IsForAddressUpdate() const { return mPerformingAddressUpdate; }
 
@@ -263,7 +279,7 @@ class DLL_EXPORT OperationalSessionSetup : public SessionEstablishmentDelegate,
     void UpdateAttemptCount(uint8_t attemptCount);
 
     // Add a retry handler for this session setup.
-    void AddRetryHandler(Callback::Callback<OnSetupRetry> * onRetry);
+    void AddRetryHandler(Callback::Callback<OnDeviceConnectionRetry> * onRetry);
 #endif // CHIP_DEVICE_CONFIG_ENABLE_AUTOMATIC_CASE_RETRIES
 
 private:
@@ -282,6 +298,7 @@ class DLL_EXPORT OperationalSessionSetup : public SessionEstablishmentDelegate,
 
     Callback::CallbackDeque mConnectionSuccess;
     Callback::CallbackDeque mConnectionFailure;
+    Callback::CallbackDeque mSetupFailure;
 
     OperationalSessionReleaseDelegate * mReleaseDelegate;
 
@@ -316,8 +333,9 @@ class DLL_EXPORT OperationalSessionSetup : public SessionEstablishmentDelegate,
 
     void CleanupCASEClient();
 
-    void EnqueueConnectionCallbacks(Callback::Callback<OnSetupSuccessful> * onConnection,
-                                    Callback::Callback<OnSetupFailure> * onFailure);
+    void EnqueueConnectionCallbacks(Callback::Callback<OnDeviceConnected> * onConnection,
+                                    Callback::Callback<OnDeviceConnectionFailure> * onFailure,
+                                    Callback::Callback<OnSetupFailure> * onSetupFailure);
 
     enum class ReleaseBehavior
     {
@@ -326,11 +344,13 @@ class DLL_EXPORT OperationalSessionSetup : public SessionEstablishmentDelegate,
     };
 
     /*
-     * This dequeues all failure and success callbacks and appropriately
-     * invokes either set depending on the value of error.
+     * This dequeues all failure and success callbacks and appropriately invokes either set depending
+     * on the value of error.
+     *
+     * If error == CHIP_NO_ERROR, only success callbacks are invoked. Otherwise, only failure callbacks are invoked.
      *
-     * If error == CHIP_NO_ERROR, only success callbacks are invoked.
-     * Otherwise, only failure callbacks are invoked.
+     * The state offers additional context regarding the failure, indicating the specific state in which
+     * the error occurs. It is only relayed through failure callbacks when the error is not equal to CHIP_NO_ERROR.
      *
      * If releaseBehavior is Release, this uses mReleaseDelegate to release
      * ourselves (aka `this`). As a result any caller should return right away
@@ -338,15 +358,16 @@ class DLL_EXPORT OperationalSessionSetup : public SessionEstablishmentDelegate,
      *
      * Setting releaseBehavior to DoNotRelease is meant for use from the destructor
      */
-    void DequeueConnectionCallbacks(CHIP_ERROR error, ReleaseBehavior releaseBehavior = ReleaseBehavior::Release);
+    void DequeueConnectionCallbacks(CHIP_ERROR error, State state, ReleaseBehavior releaseBehavior = ReleaseBehavior::Release);
 
     /**
      * Helper for DequeueConnectionCallbacks that handles the actual callback
      * notifications. This happens after the object has been released, if it's
      * being released.
      */
-    static void NotifyConnectionCallbacks(Callback::Cancelable & failureReady, Callback::Cancelable & successReady,
-                                          CHIP_ERROR error, const ScopedNodeId & peerId, bool performingAddressUpdate,
+    static void NotifyConnectionCallbacks(Callback::Cancelable & failureReady, Callback::Cancelable & setupFailureReady,
+                                          Callback::Cancelable & successReady, CHIP_ERROR error, State state,
+                                          const ScopedNodeId & peerId, bool performingAddressUpdate,
                                           Messaging::ExchangeManager * exchangeMgr,
                                           const Optional<SessionHandle> & optionalSessionHandle);
 

From cbf64236ce0174e5f25451d54d5b9ffd96a771fc Mon Sep 17 00:00:00 2001
From: Yufeng Wang <yufengwang@google.com>
Date: Tue, 3 Oct 2023 11:31:20 -0700
Subject: [PATCH 4/6] Use Session State to augument errorCode

---
 src/app/OperationalSessionSetup.cpp           | 28 ++++++-------
 src/app/OperationalSessionSetup.h             | 42 ++++++++++---------
 src/protocols/secure_channel/CASESession.cpp  | 26 +++++++++++-
 src/protocols/secure_channel/CASESession.h    |  2 +
 .../secure_channel/PairingSession.cpp         |  4 +-
 src/protocols/secure_channel/PairingSession.h |  9 ++--
 .../SessionEstablishmentDelegate.h            | 21 ++++++++++
 7 files changed, 93 insertions(+), 39 deletions(-)

diff --git a/src/app/OperationalSessionSetup.cpp b/src/app/OperationalSessionSetup.cpp
index b7a85dfe547fca..7907f3ef7d1b87 100644
--- a/src/app/OperationalSessionSetup.cpp
+++ b/src/app/OperationalSessionSetup.cpp
@@ -177,7 +177,7 @@ void OperationalSessionSetup::Connect(Callback::Callback<OnDeviceConnected> * on
     //
     if (err != CHIP_NO_ERROR || isConnected)
     {
-        DequeueConnectionCallbacks(err, mState);
+        DequeueConnectionCallbacks(err);
         // Do not touch `this` instance anymore; it has been destroyed in DequeueConnectionCallbacks.
         // While it is odd to have an explicit return here at the end of the function, we do so
         // as a precaution in case someone later on adds something to the end of this function.
@@ -212,7 +212,7 @@ void OperationalSessionSetup::UpdateDeviceData(const Transport::PeerAddress & ad
     if (mState != State::ResolvingAddress)
     {
         ChipLogError(Discovery, "Received UpdateDeviceData in incorrect state");
-        DequeueConnectionCallbacks(CHIP_ERROR_INCORRECT_STATE, mState);
+        DequeueConnectionCallbacks(CHIP_ERROR_INCORRECT_STATE);
         // Do not touch `this` instance anymore; it has been destroyed in
         // DequeueConnectionCallbacks.
         return;
@@ -224,7 +224,7 @@ void OperationalSessionSetup::UpdateDeviceData(const Transport::PeerAddress & ad
     if (mPerformingAddressUpdate)
     {
         // Nothing else to do here.
-        DequeueConnectionCallbacks(CHIP_NO_ERROR, State::HasAddress);
+        DequeueConnectionCallbacks(CHIP_NO_ERROR);
         // Do not touch `this` instance anymore; it has been destroyed in DequeueConnectionCallbacks.
         return;
     }
@@ -248,7 +248,7 @@ void OperationalSessionSetup::UpdateDeviceData(const Transport::PeerAddress & ad
         return;
     }
 
-    DequeueConnectionCallbacks(err, State::ResolvingAddress);
+    DequeueConnectionCallbacks(err);
     // Do not touch `this` instance anymore; it has been destroyed in DequeueConnectionCallbacks.
 }
 
@@ -289,7 +289,7 @@ void OperationalSessionSetup::EnqueueConnectionCallbacks(Callback::Callback<OnDe
     }
 }
 
-void OperationalSessionSetup::DequeueConnectionCallbacks(CHIP_ERROR error, State state, ReleaseBehavior releaseBehavior)
+void OperationalSessionSetup::DequeueConnectionCallbacks(CHIP_ERROR error, SessionState state, ReleaseBehavior releaseBehavior)
 {
     Cancelable failureReady, setupFailureReady, successReady;
 
@@ -330,7 +330,7 @@ void OperationalSessionSetup::DequeueConnectionCallbacks(CHIP_ERROR error, State
 }
 
 void OperationalSessionSetup::NotifyConnectionCallbacks(Cancelable & failureReady, Cancelable & setupFailureReady,
-                                                        Cancelable & successReady, CHIP_ERROR error, State state,
+                                                        Cancelable & successReady, CHIP_ERROR error, SessionState state,
                                                         const ScopedNodeId & peerId, bool performingAddressUpdate,
                                                         Messaging::ExchangeManager * exchangeMgr,
                                                         const Optional<SessionHandle> & optionalSessionHandle)
@@ -386,7 +386,7 @@ void OperationalSessionSetup::NotifyConnectionCallbacks(Cancelable & failureRead
     }
 }
 
-void OperationalSessionSetup::OnSessionEstablishmentError(CHIP_ERROR error)
+void OperationalSessionSetup::OnSessionEstablishmentError(CHIP_ERROR error, SessionState state)
 {
     VerifyOrReturn(mState == State::Connecting,
                    ChipLogError(Discovery, "OnSessionEstablishmentError was called while we were not connecting"));
@@ -431,7 +431,7 @@ void OperationalSessionSetup::OnSessionEstablishmentError(CHIP_ERROR error)
 #endif // CHIP_DEVICE_CONFIG_ENABLE_AUTOMATIC_CASE_RETRIES
     }
 
-    DequeueConnectionCallbacks(error, mState);
+    DequeueConnectionCallbacks(error, state);
     // Do not touch `this` instance anymore; it has been destroyed in DequeueConnectionCallbacks.
 }
 
@@ -444,7 +444,7 @@ void OperationalSessionSetup::OnSessionEstablished(const SessionHandle & session
     {
         // Got an invalid session, just dispatch an error.  We have to do this
         // so we don't leak.
-        DequeueConnectionCallbacks(CHIP_ERROR_INCORRECT_STATE, State::Connecting);
+        DequeueConnectionCallbacks(CHIP_ERROR_INCORRECT_STATE);
 
         // Do not touch `this` instance anymore; it has been destroyed in DequeueConnectionCallbacks.
         return;
@@ -452,7 +452,7 @@ void OperationalSessionSetup::OnSessionEstablished(const SessionHandle & session
 
     MoveToState(State::SecureConnected);
 
-    DequeueConnectionCallbacks(CHIP_NO_ERROR, State::SecureConnected);
+    DequeueConnectionCallbacks(CHIP_NO_ERROR);
 }
 
 void OperationalSessionSetup::CleanupCASEClient()
@@ -492,7 +492,7 @@ OperationalSessionSetup::~OperationalSessionSetup()
     CancelSessionSetupReattempt();
 #endif // CHIP_DEVICE_CONFIG_ENABLE_AUTOMATIC_CASE_RETRIES
 
-    DequeueConnectionCallbacks(CHIP_ERROR_CANCELLED, mState, ReleaseBehavior::DoNotRelease);
+    DequeueConnectionCallbacks(CHIP_ERROR_CANCELLED, ReleaseBehavior::DoNotRelease);
 }
 
 CHIP_ERROR OperationalSessionSetup::LookupPeerAddress()
@@ -554,7 +554,7 @@ void OperationalSessionSetup::PerformAddressUpdate()
     if (err != CHIP_NO_ERROR)
     {
         ChipLogError(Discovery, "Failed to look up peer address: %" CHIP_ERROR_FORMAT, err.Format());
-        DequeueConnectionCallbacks(err, State::ResolvingAddress);
+        DequeueConnectionCallbacks(err);
         // Do not touch `this` instance anymore; it has been destroyed in DequeueConnectionCallbacks.
         return;
     }
@@ -610,7 +610,7 @@ void OperationalSessionSetup::OnNodeAddressResolutionFailed(const PeerId & peerI
 #endif
 
     // No need to modify any variables in `this` since call below releases `this`.
-    DequeueConnectionCallbacks(reason, mState);
+    DequeueConnectionCallbacks(reason);
     // Do not touch `this` instance anymore; it has been destroyed in DequeueConnectionCallbacks.
 }
 
@@ -710,7 +710,7 @@ void OperationalSessionSetup::TrySetupAgain(System::Layer * systemLayer, void *
     }
 
     // Give up; we could not start a lookup.
-    self->DequeueConnectionCallbacks(err, State::ResolvingAddress);
+    self->DequeueConnectionCallbacks(err);
     // Do not touch `self` instance anymore; it has been destroyed in DequeueConnectionCallbacks.
 }
 
diff --git a/src/app/OperationalSessionSetup.h b/src/app/OperationalSessionSetup.h
index 847b3cea0d303f..7998548cd99ed0 100644
--- a/src/app/OperationalSessionSetup.h
+++ b/src/app/OperationalSessionSetup.h
@@ -155,25 +155,13 @@ typedef void (*OnDeviceConnectionRetry)(void * context, const ScopedNodeId & pee
 class DLL_EXPORT OperationalSessionSetup : public SessionEstablishmentDelegate, public AddressResolve::NodeListener
 {
 public:
-    enum class State : uint8_t
-    {
-        Uninitialized,    // Error state: OperationalSessionSetup is useless
-        NeedsAddress,     // No address known, lookup not started yet.
-        ResolvingAddress, // Address lookup in progress.
-        HasAddress,       // Have an address, CASE handshake not started yet.
-        Connecting,       // CASE handshake in progress.
-        SecureConnected,  // CASE session established.
-        WaitingForRetry,  // No address known, but a retry is pending.  Added at
-                          // end to make logs easier to understand.
-    };
-
     struct ConnnectionFailureInfo
     {
         const ScopedNodeId & peerId;
         CHIP_ERROR error;
-        OperationalSessionSetup::State sessionState;
+        SessionState sessionState;
 
-        ConnnectionFailureInfo(const ScopedNodeId & peer, CHIP_ERROR err, OperationalSessionSetup::State state) :
+        ConnnectionFailureInfo(const ScopedNodeId & peer, CHIP_ERROR err, SessionState state) :
             peerId(peer), error(err), sessionState(state)
         {}
     };
@@ -240,12 +228,10 @@ class DLL_EXPORT OperationalSessionSetup : public SessionEstablishmentDelegate,
 
     //////////// SessionEstablishmentDelegate Implementation ///////////////
     void OnSessionEstablished(const SessionHandle & session) override;
-    void OnSessionEstablishmentError(CHIP_ERROR error) override;
+    void OnSessionEstablishmentError(CHIP_ERROR error, SessionState state) override;
 
     ScopedNodeId GetPeerId() const { return mPeerId; }
 
-    State GetOperationalSessionSetupState() const { return mState; }
-
     static Transport::PeerAddress ToPeerAddress(const Dnssd::ResolvedNodeData & nodeData)
     {
         Inet::InterfaceId interfaceId = Inet::InterfaceId::Null();
@@ -283,6 +269,18 @@ class DLL_EXPORT OperationalSessionSetup : public SessionEstablishmentDelegate,
 #endif // CHIP_DEVICE_CONFIG_ENABLE_AUTOMATIC_CASE_RETRIES
 
 private:
+    enum class State : uint8_t
+    {
+        Uninitialized,    // Error state: OperationalSessionSetup is useless
+        NeedsAddress,     // No address known, lookup not started yet.
+        ResolvingAddress, // Address lookup in progress.
+        HasAddress,       // Have an address, CASE handshake not started yet.
+        Connecting,       // CASE handshake in progress.
+        SecureConnected,  // CASE session established.
+        WaitingForRetry,  // No address known, but a retry is pending.  Added at
+                          // end to make logs easier to understand.
+    };
+
     CASEClientInitParams mInitParams;
     CASEClientPoolDelegate * mClientPool = nullptr;
 
@@ -358,7 +356,13 @@ class DLL_EXPORT OperationalSessionSetup : public SessionEstablishmentDelegate,
      *
      * Setting releaseBehavior to DoNotRelease is meant for use from the destructor
      */
-    void DequeueConnectionCallbacks(CHIP_ERROR error, State state, ReleaseBehavior releaseBehavior = ReleaseBehavior::Release);
+    void DequeueConnectionCallbacks(CHIP_ERROR error, SessionState state,
+                                    ReleaseBehavior releaseBehavior = ReleaseBehavior::Release);
+
+    void DequeueConnectionCallbacks(CHIP_ERROR error, ReleaseBehavior releaseBehavior = ReleaseBehavior::Release)
+    {
+        this->DequeueConnectionCallbacks(error, SessionState::kUndefined, releaseBehavior);
+    }
 
     /**
      * Helper for DequeueConnectionCallbacks that handles the actual callback
@@ -366,7 +370,7 @@ class DLL_EXPORT OperationalSessionSetup : public SessionEstablishmentDelegate,
      * being released.
      */
     static void NotifyConnectionCallbacks(Callback::Cancelable & failureReady, Callback::Cancelable & setupFailureReady,
-                                          Callback::Cancelable & successReady, CHIP_ERROR error, State state,
+                                          Callback::Cancelable & successReady, CHIP_ERROR error, SessionState state,
                                           const ScopedNodeId & peerId, bool performingAddressUpdate,
                                           Messaging::ExchangeManager * exchangeMgr,
                                           const Optional<SessionHandle> & optionalSessionHandle);
diff --git a/src/protocols/secure_channel/CASESession.cpp b/src/protocols/secure_channel/CASESession.cpp
index 2afc40feacaf02..dc8e49f407e6c4 100644
--- a/src/protocols/secure_channel/CASESession.cpp
+++ b/src/protocols/secure_channel/CASESession.cpp
@@ -555,9 +555,10 @@ void CASESession::OnResponseTimeout(ExchangeContext * ec)
 
 void CASESession::AbortPendingEstablish(CHIP_ERROR err)
 {
+    SessionState state = MapCASEStateToSessionState(mState);
     Clear();
     // Do this last in case the delegate frees us.
-    NotifySessionEstablishmentError(err);
+    NotifySessionEstablishmentError(err, state);
 }
 
 CHIP_ERROR CASESession::DeriveSecureSession(CryptoContext & session) const
@@ -2267,4 +2268,27 @@ bool CASESession::InvokeBackgroundWorkWatchdog()
     return watchdogFired;
 }
 
+// Helper function to map CASESession::State to SessionState
+SessionState CASESession::MapCASEStateToSessionState(State caseState)
+{
+    switch (caseState)
+    {
+    case State::kInitialized:
+        return SessionState::kUndefined;
+    case State::kSentSigma1:
+    case State::kSentSigma1Resume:
+        return SessionState::kSentSigma1;
+    case State::kSentSigma2:
+    case State::kSentSigma2Resume:
+    case State::kSendSigma3Pending:
+    case State::kHandleSigma3Pending:
+        return SessionState::kSentSigma2;
+    case State::kSentSigma3:
+        return SessionState::kSentSigma3;
+    // Add more mappings here for other states
+    default:
+        return SessionState::kUndefined; // Default mapping
+    }
+}
+
 } // namespace chip
diff --git a/src/protocols/secure_channel/CASESession.h b/src/protocols/secure_channel/CASESession.h
index 7453b6b5002dc4..8305eb1ca9f81c 100644
--- a/src/protocols/secure_channel/CASESession.h
+++ b/src/protocols/secure_channel/CASESession.h
@@ -320,6 +320,8 @@ class DLL_EXPORT CASESession : public Messaging::UnsolicitedMessageHandler,
 #if CONFIG_BUILD_FOR_HOST_UNIT_TEST
     Optional<State> mStopHandshakeAtState = Optional<State>::Missing();
 #endif // CONFIG_BUILD_FOR_HOST_UNIT_TEST
+
+    SessionState MapCASEStateToSessionState(State caseState);
 };
 
 } // namespace chip
diff --git a/src/protocols/secure_channel/PairingSession.cpp b/src/protocols/secure_channel/PairingSession.cpp
index af555bfeee5046..a30b0eb5d52c64 100644
--- a/src/protocols/secure_channel/PairingSession.cpp
+++ b/src/protocols/secure_channel/PairingSession.cpp
@@ -187,7 +187,7 @@ void PairingSession::Clear()
     mSessionManager = nullptr;
 }
 
-void PairingSession::NotifySessionEstablishmentError(CHIP_ERROR error)
+void PairingSession::NotifySessionEstablishmentError(CHIP_ERROR error, SessionState state)
 {
     if (mDelegate == nullptr)
     {
@@ -197,7 +197,7 @@ void PairingSession::NotifySessionEstablishmentError(CHIP_ERROR error)
 
     auto * delegate = mDelegate;
     mDelegate       = nullptr;
-    delegate->OnSessionEstablishmentError(error);
+    delegate->OnSessionEstablishmentError(error, state);
 }
 
 void PairingSession::OnSessionReleased()
diff --git a/src/protocols/secure_channel/PairingSession.h b/src/protocols/secure_channel/PairingSession.h
index c604cd0662dfd3..f03c901b26ead8 100644
--- a/src/protocols/secure_channel/PairingSession.h
+++ b/src/protocols/secure_channel/PairingSession.h
@@ -216,10 +216,13 @@ class DLL_EXPORT PairingSession : public SessionDelegate
     void Clear();
 
     /**
-     * Notify our delegate about a session establishment error, if we have not
-     * notified it of an error or success before.
+     * Notify our delegate about a session establishment error if we have not
+     * already notified it of an error or success for the given session state.
+     *
+     * @param error The error code to report.
+     * @param state The state of the session, defaults to State::kUndefined.
      */
-    void NotifySessionEstablishmentError(CHIP_ERROR error);
+    void NotifySessionEstablishmentError(CHIP_ERROR error, SessionState state = SessionState::kUndefined);
 
 protected:
     CryptoContext::SessionRole mRole;
diff --git a/src/protocols/secure_channel/SessionEstablishmentDelegate.h b/src/protocols/secure_channel/SessionEstablishmentDelegate.h
index dc73a0ffe6997d..956af2885a66d9 100644
--- a/src/protocols/secure_channel/SessionEstablishmentDelegate.h
+++ b/src/protocols/secure_channel/SessionEstablishmentDelegate.h
@@ -25,6 +25,7 @@
 
 #pragma once
 
+#include <protocols/secure_channel/PairingSession.h>
 #include <system/SystemPacketBuffer.h>
 #include <transport/Session.h>
 #include <transport/raw/MessageHeader.h>
@@ -32,6 +33,19 @@
 
 namespace chip {
 
+enum class SessionState : uint8_t
+{
+    kUndefined              = 0,
+    kSentPBKDFParamRequest  = 1,
+    kSentPBKDFParamResponse = 2,
+    kSentPake1              = 3,
+    kSentPake2              = 4,
+    kSentPake3              = 5,
+    kSentSigma1             = 6,
+    kSentSigma2             = 7,
+    kSentSigma3             = 8,
+};
+
 class DLL_EXPORT SessionEstablishmentDelegate
 {
 public:
@@ -42,6 +56,13 @@ class DLL_EXPORT SessionEstablishmentDelegate
      */
     virtual void OnSessionEstablishmentError(CHIP_ERROR error) {}
 
+    /**
+     *   Called when session establishment fails with an error and state at the
+     *   failure. This will be called at most once per session establishment and
+     *   will not be called if OnSessionEstablished is called.
+     */
+    virtual void OnSessionEstablishmentError(CHIP_ERROR error, SessionState state) { OnSessionEstablishmentError(error); }
+
     /**
      *   Called on start of session establishment process
      */

From 55e44046ffc0b6786233acb5aefbe941a32e9ac2 Mon Sep 17 00:00:00 2001
From: Yufeng Wang <yufengwang@google.com>
Date: Tue, 10 Oct 2023 13:38:43 -0700
Subject: [PATCH 5/6] Address review comments

---
 src/app/OperationalSessionSetup.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/app/OperationalSessionSetup.h b/src/app/OperationalSessionSetup.h
index 7998548cd99ed0..30553a1cf009b5 100644
--- a/src/app/OperationalSessionSetup.h
+++ b/src/app/OperationalSessionSetup.h
@@ -157,7 +157,7 @@ class DLL_EXPORT OperationalSessionSetup : public SessionEstablishmentDelegate,
 public:
     struct ConnnectionFailureInfo
     {
-        const ScopedNodeId & peerId;
+        const ScopedNodeId peerId;
         CHIP_ERROR error;
         SessionState sessionState;
 

From 04556d33eb8281882176e497212bd0426839cc1a Mon Sep 17 00:00:00 2001
From: Yufeng Wang <yufengwang@google.com>
Date: Wed, 11 Oct 2023 07:20:30 -0700
Subject: [PATCH 6/6] If we have onSetupFailure, remove onFailure

---
 src/app/OperationalSessionSetup.cpp           | 189 +++++++++---------
 src/app/OperationalSessionSetup.h             |  33 +--
 src/protocols/secure_channel/CASESession.cpp  |  20 +-
 src/protocols/secure_channel/CASESession.h    |   2 +-
 .../secure_channel/PairingSession.cpp         |   4 +-
 src/protocols/secure_channel/PairingSession.h |   9 +-
 .../SessionEstablishmentDelegate.h            |  25 +--
 7 files changed, 147 insertions(+), 135 deletions(-)

diff --git a/src/app/OperationalSessionSetup.cpp b/src/app/OperationalSessionSetup.cpp
index 7907f3ef7d1b87..30ec0af8590fe7 100644
--- a/src/app/OperationalSessionSetup.cpp
+++ b/src/app/OperationalSessionSetup.cpp
@@ -98,91 +98,9 @@ void OperationalSessionSetup::Connect(Callback::Callback<OnDeviceConnected> * on
 }
 
 void OperationalSessionSetup::Connect(Callback::Callback<OnDeviceConnected> * onConnection,
-                                      Callback::Callback<OnDeviceConnectionFailure> * onFailure,
                                       Callback::Callback<OnSetupFailure> * onSetupFailure)
 {
-    CHIP_ERROR err   = CHIP_NO_ERROR;
-    bool isConnected = false;
-
-    //
-    // Always enqueue our user provided callbacks into our callback list.
-    // If anything goes wrong below, we'll trigger failures (including any queued from
-    // a previous iteration which in theory shouldn't happen, but this is written to be more defensive)
-    //
-    EnqueueConnectionCallbacks(onConnection, onFailure, onSetupFailure);
-
-    switch (mState)
-    {
-    case State::Uninitialized:
-        err = CHIP_ERROR_INCORRECT_STATE;
-        break;
-
-    case State::NeedsAddress:
-        isConnected = AttachToExistingSecureSession();
-        if (!isConnected)
-        {
-            // LookupPeerAddress could perhaps call back with a result
-            // synchronously, so do our state update first.
-            MoveToState(State::ResolvingAddress);
-            err = LookupPeerAddress();
-            if (err != CHIP_NO_ERROR)
-            {
-                // Roll back the state change, since we are presumably not in
-                // the middle of a lookup.
-                MoveToState(State::NeedsAddress);
-            }
-        }
-
-        break;
-
-    case State::ResolvingAddress:
-    case State::WaitingForRetry:
-        isConnected = AttachToExistingSecureSession();
-        break;
-
-    case State::HasAddress:
-        isConnected = AttachToExistingSecureSession();
-        if (!isConnected)
-        {
-            // We should not actually every be in be in State::HasAddress. This
-            // is because in the same call that we moved to State::HasAddress
-            // we either move to State::Connecting or call
-            // DequeueConnectionCallbacks with an error thus releasing
-            // ourselves before any call would reach this section of code.
-            err = CHIP_ERROR_INCORRECT_STATE;
-        }
-
-        break;
-
-    case State::Connecting:
-        break;
-
-    case State::SecureConnected:
-        isConnected = true;
-        break;
-
-    default:
-        err = CHIP_ERROR_INCORRECT_STATE;
-    }
-
-    if (isConnected)
-    {
-        MoveToState(State::SecureConnected);
-    }
-
-    //
-    // Dequeue all our callbacks on either encountering an error
-    // or if we successfully connected. Both should not be set
-    // simultaneously.
-    //
-    if (err != CHIP_NO_ERROR || isConnected)
-    {
-        DequeueConnectionCallbacks(err);
-        // Do not touch `this` instance anymore; it has been destroyed in DequeueConnectionCallbacks.
-        // While it is odd to have an explicit return here at the end of the function, we do so
-        // as a precaution in case someone later on adds something to the end of this function.
-        return;
-    }
+    Connect(onConnection, nullptr, onSetupFailure);
 }
 
 void OperationalSessionSetup::UpdateDeviceData(const Transport::PeerAddress & addr, const ReliableMessageProtocolConfig & config)
@@ -269,6 +187,94 @@ CHIP_ERROR OperationalSessionSetup::EstablishConnection(const ReliableMessagePro
     return CHIP_NO_ERROR;
 }
 
+void OperationalSessionSetup::Connect(Callback::Callback<OnDeviceConnected> * onConnection,
+                                      Callback::Callback<OnDeviceConnectionFailure> * onFailure,
+                                      Callback::Callback<OnSetupFailure> * onSetupFailure)
+{
+    CHIP_ERROR err   = CHIP_NO_ERROR;
+    bool isConnected = false;
+
+    //
+    // Always enqueue our user provided callbacks into our callback list.
+    // If anything goes wrong below, we'll trigger failures (including any queued from
+    // a previous iteration which in theory shouldn't happen, but this is written to be more defensive)
+    //
+    EnqueueConnectionCallbacks(onConnection, onFailure, onSetupFailure);
+
+    switch (mState)
+    {
+    case State::Uninitialized:
+        err = CHIP_ERROR_INCORRECT_STATE;
+        break;
+
+    case State::NeedsAddress:
+        isConnected = AttachToExistingSecureSession();
+        if (!isConnected)
+        {
+            // LookupPeerAddress could perhaps call back with a result
+            // synchronously, so do our state update first.
+            MoveToState(State::ResolvingAddress);
+            err = LookupPeerAddress();
+            if (err != CHIP_NO_ERROR)
+            {
+                // Roll back the state change, since we are presumably not in
+                // the middle of a lookup.
+                MoveToState(State::NeedsAddress);
+            }
+        }
+
+        break;
+
+    case State::ResolvingAddress:
+    case State::WaitingForRetry:
+        isConnected = AttachToExistingSecureSession();
+        break;
+
+    case State::HasAddress:
+        isConnected = AttachToExistingSecureSession();
+        if (!isConnected)
+        {
+            // We should not actually every be in be in State::HasAddress. This
+            // is because in the same call that we moved to State::HasAddress
+            // we either move to State::Connecting or call
+            // DequeueConnectionCallbacks with an error thus releasing
+            // ourselves before any call would reach this section of code.
+            err = CHIP_ERROR_INCORRECT_STATE;
+        }
+
+        break;
+
+    case State::Connecting:
+        break;
+
+    case State::SecureConnected:
+        isConnected = true;
+        break;
+
+    default:
+        err = CHIP_ERROR_INCORRECT_STATE;
+    }
+
+    if (isConnected)
+    {
+        MoveToState(State::SecureConnected);
+    }
+
+    //
+    // Dequeue all our callbacks on either encountering an error
+    // or if we successfully connected. Both should not be set
+    // simultaneously.
+    //
+    if (err != CHIP_NO_ERROR || isConnected)
+    {
+        DequeueConnectionCallbacks(err);
+        // Do not touch `this` instance anymore; it has been destroyed in DequeueConnectionCallbacks.
+        // While it is odd to have an explicit return here at the end of the function, we do so
+        // as a precaution in case someone later on adds something to the end of this function.
+        return;
+    }
+}
+
 void OperationalSessionSetup::EnqueueConnectionCallbacks(Callback::Callback<OnDeviceConnected> * onConnection,
                                                          Callback::Callback<OnDeviceConnectionFailure> * onFailure,
                                                          Callback::Callback<OnSetupFailure> * onSetupFailure)
@@ -289,7 +295,8 @@ void OperationalSessionSetup::EnqueueConnectionCallbacks(Callback::Callback<OnDe
     }
 }
 
-void OperationalSessionSetup::DequeueConnectionCallbacks(CHIP_ERROR error, SessionState state, ReleaseBehavior releaseBehavior)
+void OperationalSessionSetup::DequeueConnectionCallbacks(CHIP_ERROR error, SessionEstablishmentStage stage,
+                                                         ReleaseBehavior releaseBehavior)
 {
     Cancelable failureReady, setupFailureReady, successReady;
 
@@ -325,14 +332,14 @@ void OperationalSessionSetup::DequeueConnectionCallbacks(CHIP_ERROR error, Sessi
 
     // DO NOT touch any members of this object after this point.  It's dead.
 
-    NotifyConnectionCallbacks(failureReady, setupFailureReady, successReady, error, state, peerId, performingAddressUpdate,
+    NotifyConnectionCallbacks(failureReady, setupFailureReady, successReady, error, stage, peerId, performingAddressUpdate,
                               exchangeMgr, optionalSessionHandle);
 }
 
 void OperationalSessionSetup::NotifyConnectionCallbacks(Cancelable & failureReady, Cancelable & setupFailureReady,
-                                                        Cancelable & successReady, CHIP_ERROR error, SessionState state,
-                                                        const ScopedNodeId & peerId, bool performingAddressUpdate,
-                                                        Messaging::ExchangeManager * exchangeMgr,
+                                                        Cancelable & successReady, CHIP_ERROR error,
+                                                        SessionEstablishmentStage stage, const ScopedNodeId & peerId,
+                                                        bool performingAddressUpdate, Messaging::ExchangeManager * exchangeMgr,
                                                         const Optional<SessionHandle> & optionalSessionHandle)
 {
     //
@@ -365,7 +372,7 @@ void OperationalSessionSetup::NotifyConnectionCallbacks(Cancelable & failureRead
         if (error != CHIP_NO_ERROR)
         {
             // Initialize the ConnnectionFailureInfo object
-            ConnnectionFailureInfo failureInfo(peerId, error, state);
+            ConnnectionFailureInfo failureInfo(peerId, error, stage);
             cb->mCall(cb->mContext, failureInfo);
         }
     }
@@ -386,7 +393,7 @@ void OperationalSessionSetup::NotifyConnectionCallbacks(Cancelable & failureRead
     }
 }
 
-void OperationalSessionSetup::OnSessionEstablishmentError(CHIP_ERROR error, SessionState state)
+void OperationalSessionSetup::OnSessionEstablishmentError(CHIP_ERROR error, SessionEstablishmentStage stage)
 {
     VerifyOrReturn(mState == State::Connecting,
                    ChipLogError(Discovery, "OnSessionEstablishmentError was called while we were not connecting"));
@@ -431,7 +438,7 @@ void OperationalSessionSetup::OnSessionEstablishmentError(CHIP_ERROR error, Sess
 #endif // CHIP_DEVICE_CONFIG_ENABLE_AUTOMATIC_CASE_RETRIES
     }
 
-    DequeueConnectionCallbacks(error, state);
+    DequeueConnectionCallbacks(error, stage);
     // Do not touch `this` instance anymore; it has been destroyed in DequeueConnectionCallbacks.
 }
 
diff --git a/src/app/OperationalSessionSetup.h b/src/app/OperationalSessionSetup.h
index 30553a1cf009b5..acf3f1fb143c0b 100644
--- a/src/app/OperationalSessionSetup.h
+++ b/src/app/OperationalSessionSetup.h
@@ -159,10 +159,10 @@ class DLL_EXPORT OperationalSessionSetup : public SessionEstablishmentDelegate,
     {
         const ScopedNodeId peerId;
         CHIP_ERROR error;
-        SessionState sessionState;
+        SessionEstablishmentStage sessionStage;
 
-        ConnnectionFailureInfo(const ScopedNodeId & peer, CHIP_ERROR err, SessionState state) :
-            peerId(peer), error(err), sessionState(state)
+        ConnnectionFailureInfo(const ScopedNodeId & peer, CHIP_ERROR err, SessionEstablishmentStage stage) :
+            peerId(peer), error(err), sessionStage(stage)
         {}
     };
 
@@ -193,8 +193,8 @@ class DLL_EXPORT OperationalSessionSetup : public SessionEstablishmentDelegate,
      * The device is expected to have been commissioned, A CASE session
      * setup will be triggered.
      *
-     * On establishing the session, the callback function `onConnection` will be called. If the
-     * session setup fails, `onFailure` will be called.
+     * On establishing the session, if the session setup succeeds, the callback function `onConnection` will be called.
+     * If the session setup fails, `onFailure` will be called.
      *
      * If the session already exists, `onConnection` will be called immediately,
      * before the Connect call returns.
@@ -211,24 +211,22 @@ class DLL_EXPORT OperationalSessionSetup : public SessionEstablishmentDelegate,
      * The device is expected to have been commissioned, A CASE session
      * setup will be triggered.
      *
-     * On establishing the session, the callback function `onConnection` will be called. If the
-     * session setup fails, `onFailure` and `onSetupFailure` will be called.
+     * On establishing the session, if the session setup succeeds, the callback function `onConnection` will be called.
+     * If the session setup fails, `onSetupFailure` will be called.
      *
      * If the session already exists, `onConnection` will be called immediately,
      * before the Connect call returns.
      *
-     * `onFailure` and `onSetupFailure` may be called before the Connect call returns, for error
-     * cases that are detected synchronously (e.g. inability to start an address
-     * lookup).
+     * `onSetupFailure` may be called before the Connect call returns, for error cases that are detected synchronously
+     * (e.g. inability to start an address lookup).
      */
-    void Connect(Callback::Callback<OnDeviceConnected> * onConnection, Callback::Callback<OnDeviceConnectionFailure> * onFailure,
-                 Callback::Callback<OnSetupFailure> * onSetupFailure);
+    void Connect(Callback::Callback<OnDeviceConnected> * onConnection, Callback::Callback<OnSetupFailure> * onSetupFailure);
 
     bool IsForAddressUpdate() const { return mPerformingAddressUpdate; }
 
     //////////// SessionEstablishmentDelegate Implementation ///////////////
     void OnSessionEstablished(const SessionHandle & session) override;
-    void OnSessionEstablishmentError(CHIP_ERROR error, SessionState state) override;
+    void OnSessionEstablishmentError(CHIP_ERROR error, SessionEstablishmentStage stage) override;
 
     ScopedNodeId GetPeerId() const { return mPeerId; }
 
@@ -331,6 +329,9 @@ class DLL_EXPORT OperationalSessionSetup : public SessionEstablishmentDelegate,
 
     void CleanupCASEClient();
 
+    void Connect(Callback::Callback<OnDeviceConnected> * onConnection, Callback::Callback<OnDeviceConnectionFailure> * onFailure,
+                 Callback::Callback<OnSetupFailure> * onSetupFailure);
+
     void EnqueueConnectionCallbacks(Callback::Callback<OnDeviceConnected> * onConnection,
                                     Callback::Callback<OnDeviceConnectionFailure> * onFailure,
                                     Callback::Callback<OnSetupFailure> * onSetupFailure);
@@ -356,12 +357,12 @@ class DLL_EXPORT OperationalSessionSetup : public SessionEstablishmentDelegate,
      *
      * Setting releaseBehavior to DoNotRelease is meant for use from the destructor
      */
-    void DequeueConnectionCallbacks(CHIP_ERROR error, SessionState state,
+    void DequeueConnectionCallbacks(CHIP_ERROR error, SessionEstablishmentStage stage,
                                     ReleaseBehavior releaseBehavior = ReleaseBehavior::Release);
 
     void DequeueConnectionCallbacks(CHIP_ERROR error, ReleaseBehavior releaseBehavior = ReleaseBehavior::Release)
     {
-        this->DequeueConnectionCallbacks(error, SessionState::kUndefined, releaseBehavior);
+        this->DequeueConnectionCallbacks(error, SessionEstablishmentStage::kNotInKeyExchange, releaseBehavior);
     }
 
     /**
@@ -370,7 +371,7 @@ class DLL_EXPORT OperationalSessionSetup : public SessionEstablishmentDelegate,
      * being released.
      */
     static void NotifyConnectionCallbacks(Callback::Cancelable & failureReady, Callback::Cancelable & setupFailureReady,
-                                          Callback::Cancelable & successReady, CHIP_ERROR error, SessionState state,
+                                          Callback::Cancelable & successReady, CHIP_ERROR error, SessionEstablishmentStage stage,
                                           const ScopedNodeId & peerId, bool performingAddressUpdate,
                                           Messaging::ExchangeManager * exchangeMgr,
                                           const Optional<SessionHandle> & optionalSessionHandle);
diff --git a/src/protocols/secure_channel/CASESession.cpp b/src/protocols/secure_channel/CASESession.cpp
index dc8e49f407e6c4..c9519a710beeaa 100644
--- a/src/protocols/secure_channel/CASESession.cpp
+++ b/src/protocols/secure_channel/CASESession.cpp
@@ -555,7 +555,7 @@ void CASESession::OnResponseTimeout(ExchangeContext * ec)
 
 void CASESession::AbortPendingEstablish(CHIP_ERROR err)
 {
-    SessionState state = MapCASEStateToSessionState(mState);
+    SessionEstablishmentStage state = MapCASEStateToSessionEstablishmentStage(mState);
     Clear();
     // Do this last in case the delegate frees us.
     NotifySessionEstablishmentError(err, state);
@@ -2268,26 +2268,28 @@ bool CASESession::InvokeBackgroundWorkWatchdog()
     return watchdogFired;
 }
 
-// Helper function to map CASESession::State to SessionState
-SessionState CASESession::MapCASEStateToSessionState(State caseState)
+// Helper function to map CASESession::State to SessionEstablishmentStage
+SessionEstablishmentStage CASESession::MapCASEStateToSessionEstablishmentStage(State caseState)
 {
     switch (caseState)
     {
     case State::kInitialized:
-        return SessionState::kUndefined;
+        return SessionEstablishmentStage::kNotInKeyExchange;
     case State::kSentSigma1:
     case State::kSentSigma1Resume:
-        return SessionState::kSentSigma1;
+        return SessionEstablishmentStage::kSentSigma1;
     case State::kSentSigma2:
     case State::kSentSigma2Resume:
+        return SessionEstablishmentStage::kSentSigma2;
     case State::kSendSigma3Pending:
-    case State::kHandleSigma3Pending:
-        return SessionState::kSentSigma2;
+        return SessionEstablishmentStage::kReceivedSigma2;
     case State::kSentSigma3:
-        return SessionState::kSentSigma3;
+        return SessionEstablishmentStage::kSentSigma3;
+    case State::kHandleSigma3Pending:
+        return SessionEstablishmentStage::kReceivedSigma3;
     // Add more mappings here for other states
     default:
-        return SessionState::kUndefined; // Default mapping
+        return SessionEstablishmentStage::kUnknown; // Default mapping
     }
 }
 
diff --git a/src/protocols/secure_channel/CASESession.h b/src/protocols/secure_channel/CASESession.h
index 8305eb1ca9f81c..6fc58dfb90dc83 100644
--- a/src/protocols/secure_channel/CASESession.h
+++ b/src/protocols/secure_channel/CASESession.h
@@ -321,7 +321,7 @@ class DLL_EXPORT CASESession : public Messaging::UnsolicitedMessageHandler,
     Optional<State> mStopHandshakeAtState = Optional<State>::Missing();
 #endif // CONFIG_BUILD_FOR_HOST_UNIT_TEST
 
-    SessionState MapCASEStateToSessionState(State caseState);
+    SessionEstablishmentStage MapCASEStateToSessionEstablishmentStage(State caseState);
 };
 
 } // namespace chip
diff --git a/src/protocols/secure_channel/PairingSession.cpp b/src/protocols/secure_channel/PairingSession.cpp
index a30b0eb5d52c64..c61acbbc23d27e 100644
--- a/src/protocols/secure_channel/PairingSession.cpp
+++ b/src/protocols/secure_channel/PairingSession.cpp
@@ -187,7 +187,7 @@ void PairingSession::Clear()
     mSessionManager = nullptr;
 }
 
-void PairingSession::NotifySessionEstablishmentError(CHIP_ERROR error, SessionState state)
+void PairingSession::NotifySessionEstablishmentError(CHIP_ERROR error, SessionEstablishmentStage stage)
 {
     if (mDelegate == nullptr)
     {
@@ -197,7 +197,7 @@ void PairingSession::NotifySessionEstablishmentError(CHIP_ERROR error, SessionSt
 
     auto * delegate = mDelegate;
     mDelegate       = nullptr;
-    delegate->OnSessionEstablishmentError(error, state);
+    delegate->OnSessionEstablishmentError(error, stage);
 }
 
 void PairingSession::OnSessionReleased()
diff --git a/src/protocols/secure_channel/PairingSession.h b/src/protocols/secure_channel/PairingSession.h
index f03c901b26ead8..a9d3ffa64099fd 100644
--- a/src/protocols/secure_channel/PairingSession.h
+++ b/src/protocols/secure_channel/PairingSession.h
@@ -216,13 +216,14 @@ class DLL_EXPORT PairingSession : public SessionDelegate
     void Clear();
 
     /**
-     * Notify our delegate about a session establishment error if we have not
-     * already notified it of an error or success for the given session state.
+     * Notify our delegate about a session establishment error and the stage when the error occurs
+     * if we have not already notified it of an error or success before.
      *
      * @param error The error code to report.
-     * @param state The state of the session, defaults to State::kUndefined.
+     * @param stage The stage of the session when the error occurs, defaults to kNotInKeyExchange.
      */
-    void NotifySessionEstablishmentError(CHIP_ERROR error, SessionState state = SessionState::kUndefined);
+    void NotifySessionEstablishmentError(CHIP_ERROR error,
+                                         SessionEstablishmentStage stage = SessionEstablishmentStage::kNotInKeyExchange);
 
 protected:
     CryptoContext::SessionRole mRole;
diff --git a/src/protocols/secure_channel/SessionEstablishmentDelegate.h b/src/protocols/secure_channel/SessionEstablishmentDelegate.h
index 956af2885a66d9..a074e5ee074c12 100644
--- a/src/protocols/secure_channel/SessionEstablishmentDelegate.h
+++ b/src/protocols/secure_channel/SessionEstablishmentDelegate.h
@@ -25,7 +25,6 @@
 
 #pragma once
 
-#include <protocols/secure_channel/PairingSession.h>
 #include <system/SystemPacketBuffer.h>
 #include <transport/Session.h>
 #include <transport/raw/MessageHeader.h>
@@ -33,17 +32,16 @@
 
 namespace chip {
 
-enum class SessionState : uint8_t
+enum class SessionEstablishmentStage : uint8_t
 {
-    kUndefined              = 0,
-    kSentPBKDFParamRequest  = 1,
-    kSentPBKDFParamResponse = 2,
-    kSentPake1              = 3,
-    kSentPake2              = 4,
-    kSentPake3              = 5,
-    kSentSigma1             = 6,
-    kSentSigma2             = 7,
-    kSentSigma3             = 8,
+    kUnknown          = 0,
+    kNotInKeyExchange = 1,
+    kSentSigma1       = 2,
+    kReceivedSigma1   = 3,
+    kSentSigma2       = 4,
+    kReceivedSigma2   = 5,
+    kSentSigma3       = 6,
+    kReceivedSigma3   = 7,
 };
 
 class DLL_EXPORT SessionEstablishmentDelegate
@@ -61,7 +59,10 @@ class DLL_EXPORT SessionEstablishmentDelegate
      *   failure. This will be called at most once per session establishment and
      *   will not be called if OnSessionEstablished is called.
      */
-    virtual void OnSessionEstablishmentError(CHIP_ERROR error, SessionState state) { OnSessionEstablishmentError(error); }
+    virtual void OnSessionEstablishmentError(CHIP_ERROR error, SessionEstablishmentStage stage)
+    {
+        OnSessionEstablishmentError(error);
+    }
 
     /**
      *   Called on start of session establishment process