From 24b97420c66a0538187ba7f509314db98edb314e Mon Sep 17 00:00:00 2001
From: Marc Lepage <mlepage@google.com>
Date: Mon, 27 Feb 2023 12:52:17 -0500
Subject: [PATCH 01/19] Break CASESession::SendSigma3 into fg/bg parts

This unblocks the main event loop while performance intensive
(e.g. crypto) parts of the process occur.
---
 src/protocols/secure_channel/CASESession.cpp | 241 ++++++++++++++-----
 src/protocols/secure_channel/CASESession.h   |  12 +-
 2 files changed, 194 insertions(+), 59 deletions(-)

diff --git a/src/protocols/secure_channel/CASESession.cpp b/src/protocols/secure_channel/CASESession.cpp
index 7c00272155624e..b26b00e7158c46 100644
--- a/src/protocols/secure_channel/CASESession.cpp
+++ b/src/protocols/secure_channel/CASESession.cpp
@@ -41,6 +41,7 @@
 #include <protocols/secure_channel/PairingSession.h>
 #include <protocols/secure_channel/SessionResumptionStorage.h>
 #include <protocols/secure_channel/StatusReport.h>
+#include <system/SystemClock.h>
 #include <system/TLVPacketBufferBackingStore.h>
 #include <trace/trace.h>
 #include <transport/SessionManager.h>
@@ -95,6 +96,31 @@ enum
 
 namespace chip {
 
+// Simple timing system
+const char* label = nullptr;
+System::Clock::Timestamp start;
+void StartTiming(const char* s) {
+    System::Clock::Timestamp now = System::SystemClock().GetMonotonicTimestamp();
+    if (label) {
+        uint32_t deltaMs = System::Clock::Milliseconds32(now - start).count();
+        ChipLogError(DeviceLayer, "Timing '%s' --> %" PRIu32 " ms", label, deltaMs);
+    }
+    label = s;
+    start = now;
+}
+void EndTiming() {
+    System::Clock::Timestamp now = System::SystemClock().GetMonotonicTimestamp();
+    if (label) {
+        uint32_t deltaMs = System::Clock::Milliseconds32(now - start).count();
+        ChipLogError(DeviceLayer, "Timing '%s' --> %" PRIu32 " ms", label, deltaMs);
+    }
+    label = nullptr;
+}
+
+}
+
+namespace chip {
+
 using namespace Crypto;
 using namespace Credentials;
 using namespace Messaging;
@@ -254,6 +280,10 @@ CHIP_ERROR CASESession::EstablishSession(SessionManager & sessionManager, Fabric
     ReturnErrorCodeIf(exchangeCtxt == nullptr, CHIP_ERROR_INVALID_ARGUMENT);
     ReturnErrorCodeIf(fabricTable == nullptr, CHIP_ERROR_INVALID_ARGUMENT);
 
+    // Sequence number used to coordinate foreground/background work for a
+    // particular session establishment.
+    mSequence++;
+
     // Use FabricTable directly to avoid situation of dangling index from stale FabricInfo
     // until we factor-out any FabricInfo direct usage.
     ReturnErrorCodeIf(peerScopedNodeId.GetFabricIndex() == kUndefinedFabricIndex, CHIP_ERROR_INVALID_ARGUMENT);
@@ -932,7 +962,7 @@ CHIP_ERROR CASESession::HandleSigma2_and_SendSigma3(System::PacketBufferHandle &
 {
     MATTER_TRACE_EVENT_SCOPE("HandleSigma2_and_SendSigma3", "CASESession");
     ReturnErrorOnFailure(HandleSigma2(std::move(msg)));
-    ReturnErrorOnFailure(SendSigma3());
+    ReturnErrorOnFailure(SendSigma3a());
 
     return CHIP_NO_ERROR;
 }
@@ -1112,26 +1142,26 @@ CHIP_ERROR CASESession::HandleSigma2(System::PacketBufferHandle && msg)
     return err;
 }
 
-CHIP_ERROR CASESession::SendSigma3()
+struct CASESession::SendSigma3Work
 {
-    MATTER_TRACE_EVENT_SCOPE("SendSigma3", "CASESession");
-    CHIP_ERROR err = CHIP_NO_ERROR;
-
-    MutableByteSpan messageDigestSpan(mMessageDigest);
-    System::PacketBufferHandle msg_R3;
-    size_t data_len;
+    // Status of background processing.
+    CHIP_ERROR status;
 
-    chip::Platform::ScopedMemoryBuffer<uint8_t> msg_R3_Encrypted;
-    size_t msg_r3_encrypted_len;
+    // Session to use after background processing.
+    CASESession * session;
 
-    uint8_t msg_salt[kIPKSize + kSHA256_Hash_Length];
+    // Sequence number used to coordinate foreground/background work for a
+    // particular session establishment.
+    int sequence;
 
-    uint8_t sr3k[CHIP_CRYPTO_SYMMETRIC_KEY_LENGTH_BYTES];
+    FabricTable * fabricTable;
+    FabricIndex fabricIndex;
 
     chip::Platform::ScopedMemoryBuffer<uint8_t> msg_R3_Signed;
     size_t msg_r3_signed_len;
 
-    P256ECDSASignature tbsData3Signature;
+    chip::Platform::ScopedMemoryBuffer<uint8_t> msg_R3_Encrypted;
+    size_t msg_r3_encrypted_len;
 
     chip::Platform::ScopedMemoryBuffer<uint8_t> icacBuf;
     MutableByteSpan icaCert;
@@ -1139,64 +1169,159 @@ CHIP_ERROR CASESession::SendSigma3()
     chip::Platform::ScopedMemoryBuffer<uint8_t> nocBuf;
     MutableByteSpan nocCert;
 
+    P256ECDSASignature tbsData3Signature;
+};
+
+CHIP_ERROR CASESession::SendSigma3a()
+{
+    MATTER_TRACE_EVENT_SCOPE("SendSigma3", "CASESession");
+    CHIP_ERROR err = CHIP_NO_ERROR;
+
     ChipLogDetail(SecureChannel, "Sending Sigma3");
 
-    VerifyOrExit(mEphemeralKey != nullptr, err = CHIP_ERROR_INTERNAL);
-    VerifyOrExit(icacBuf.Alloc(kMaxCHIPCertLength), err = CHIP_ERROR_NO_MEMORY);
-    icaCert = MutableByteSpan{ icacBuf.Get(), kMaxCHIPCertLength };
+    ChipLogError(DeviceLayer, "@@@@@@@@@@ SendSigma3a");
 
-    VerifyOrExit(nocBuf.Alloc(kMaxCHIPCertLength), err = CHIP_ERROR_NO_MEMORY);
-    nocCert = MutableByteSpan{ nocBuf.Get(), kMaxCHIPCertLength };
+    auto * workPtr = Platform::New<SendSigma3Work>();
+    VerifyOrExit(workPtr != nullptr, err = CHIP_ERROR_NO_MEMORY);
+    {
+        auto & work = *workPtr;
 
-    VerifyOrExit(mFabricsTable != nullptr, err = CHIP_ERROR_INCORRECT_STATE);
+        // Used to call back into the session after background event processing.
+        // It happens that there's only one pairing session (in CASEServer)
+        // so it will still be available for use. Use a sequence number to
+        // coordinate.
+        work.session  = this;
+        work.sequence = mSequence;
+
+        VerifyOrExit(mFabricsTable != nullptr, err = CHIP_ERROR_INCORRECT_STATE);
+        work.fabricTable = mFabricsTable;
+        work.fabricIndex = mFabricIndex;
+
+        VerifyOrExit(mEphemeralKey != nullptr, err = CHIP_ERROR_INTERNAL);
+
+        VerifyOrExit(work.icacBuf.Alloc(kMaxCHIPCertLength), err = CHIP_ERROR_NO_MEMORY);
+        work.icaCert = MutableByteSpan{ work.icacBuf.Get(), kMaxCHIPCertLength };
+
+        VerifyOrExit(work.nocBuf.Alloc(kMaxCHIPCertLength), err = CHIP_ERROR_NO_MEMORY);
+        work.nocCert = MutableByteSpan{ work.nocBuf.Get(), kMaxCHIPCertLength };
+
+        SuccessOrExit(err = mFabricsTable->FetchICACert(mFabricIndex, work.icaCert));
+        SuccessOrExit(err = mFabricsTable->FetchNOCCert(mFabricIndex, work.nocCert));
+
+        // Prepare Sigma3 TBS Data Blob
+        work.msg_r3_signed_len = TLV::EstimateStructOverhead(work.icaCert.size(), work.nocCert.size(), kP256_PublicKey_Length, kP256_PublicKey_Length);
+
+        VerifyOrExit(work.msg_R3_Signed.Alloc(work.msg_r3_signed_len), err = CHIP_ERROR_NO_MEMORY);
+
+        SuccessOrExit(err = ConstructTBSData(work.nocCert, work.icaCert, ByteSpan(mEphemeralKey->Pubkey(), mEphemeralKey->Pubkey().Length()),
+                                            ByteSpan(mRemotePubKey, mRemotePubKey.Length()), work.msg_R3_Signed.Get(), work.msg_r3_signed_len));
+
+        SuccessOrExit(
+            err = DeviceLayer::PlatformMgr().ScheduleBackgroundWork(
+                [](intptr_t arg) { SendSigma3b(*reinterpret_cast<SendSigma3Work *>(arg)); }, reinterpret_cast<intptr_t>(&work)));
+        workPtr = nullptr; // scheduling succeeded, so don't delete
+        mExchangeCtxt->WillSendMessage();
+        mState = State::kBackgroundPending;
 
-    SuccessOrExit(err = mFabricsTable->FetchICACert(mFabricIndex, icaCert));
-    SuccessOrExit(err = mFabricsTable->FetchNOCCert(mFabricIndex, nocCert));
+        // TODO decide if State::kBackgroundPending is sufficient
+        // or should have another state
+    }
 
-    // Prepare Sigma3 TBS Data Blob
-    msg_r3_signed_len = TLV::EstimateStructOverhead(icaCert.size(), nocCert.size(), kP256_PublicKey_Length, kP256_PublicKey_Length);
+exit:
+    Platform::Delete(workPtr);
 
-    VerifyOrExit(msg_R3_Signed.Alloc(msg_r3_signed_len), err = CHIP_ERROR_NO_MEMORY);
+    if (err != CHIP_NO_ERROR)
+    {
+        SendStatusReport(mExchangeCtxt, kProtocolCodeInvalidParam);
+        mState = State::kInitialized;
+    }
 
-    SuccessOrExit(err = ConstructTBSData(nocCert, icaCert, ByteSpan(mEphemeralKey->Pubkey(), mEphemeralKey->Pubkey().Length()),
-                                         ByteSpan(mRemotePubKey, mRemotePubKey.Length()), msg_R3_Signed.Get(), msg_r3_signed_len));
+    return err;
+}
+
+void CASESession::SendSigma3b(SendSigma3Work & work)
+{
+    CHIP_ERROR err = CHIP_NO_ERROR;
+
+    ChipLogError(DeviceLayer, "@@@@@@@@@@ SendSigma3b");
+
+    // TODO SignWithOpKeypair has moved here but check for race condition etc.
 
     // Generate a signature
-    err = mFabricsTable->SignWithOpKeypair(mFabricIndex, ByteSpan{ msg_R3_Signed.Get(), msg_r3_signed_len }, tbsData3Signature);
+    err = work.fabricTable->SignWithOpKeypair(work.fabricIndex, ByteSpan{ work.msg_R3_Signed.Get(), work.msg_r3_signed_len }, work.tbsData3Signature);
     SuccessOrExit(err);
 
     // Prepare Sigma3 TBE Data Blob
-    msg_r3_encrypted_len = TLV::EstimateStructOverhead(nocCert.size(), icaCert.size(), tbsData3Signature.Length());
+    work.msg_r3_encrypted_len = TLV::EstimateStructOverhead(work.nocCert.size(), work.icaCert.size(), work.tbsData3Signature.Length());
 
-    VerifyOrExit(msg_R3_Encrypted.Alloc(msg_r3_encrypted_len + CHIP_CRYPTO_AEAD_MIC_LENGTH_BYTES), err = CHIP_ERROR_NO_MEMORY);
+    VerifyOrExit(work.msg_R3_Encrypted.Alloc(work.msg_r3_encrypted_len + CHIP_CRYPTO_AEAD_MIC_LENGTH_BYTES), err = CHIP_ERROR_NO_MEMORY);
 
     {
         TLV::TLVWriter tlvWriter;
         TLV::TLVType outerContainerType = TLV::kTLVType_NotSpecified;
 
-        tlvWriter.Init(msg_R3_Encrypted.Get(), msg_r3_encrypted_len);
+        tlvWriter.Init(work.msg_R3_Encrypted.Get(), work.msg_r3_encrypted_len);
         SuccessOrExit(err = tlvWriter.StartContainer(TLV::AnonymousTag(), TLV::kTLVType_Structure, outerContainerType));
-        SuccessOrExit(err = tlvWriter.Put(TLV::ContextTag(kTag_TBEData_SenderNOC), nocCert));
-        if (!icaCert.empty())
+        SuccessOrExit(err = tlvWriter.Put(TLV::ContextTag(kTag_TBEData_SenderNOC), work.nocCert));
+        if (!work.icaCert.empty())
         {
-            SuccessOrExit(err = tlvWriter.Put(TLV::ContextTag(kTag_TBEData_SenderICAC), icaCert));
+            SuccessOrExit(err = tlvWriter.Put(TLV::ContextTag(kTag_TBEData_SenderICAC), work.icaCert));
         }
 
         // We are now done with ICAC and NOC certs so we can release the memory.
         {
-            icacBuf.Free();
-            icaCert = MutableByteSpan{};
+            work.icacBuf.Free();
+            work.icaCert = MutableByteSpan{};
 
-            nocBuf.Free();
-            nocCert = MutableByteSpan{};
+            work.nocBuf.Free();
+            work.nocCert = MutableByteSpan{};
         }
 
-        SuccessOrExit(err = tlvWriter.PutBytes(TLV::ContextTag(kTag_TBEData_Signature), tbsData3Signature.ConstBytes(),
-                                               static_cast<uint32_t>(tbsData3Signature.Length())));
+        SuccessOrExit(err = tlvWriter.PutBytes(TLV::ContextTag(kTag_TBEData_Signature), work.tbsData3Signature.ConstBytes(),
+                                            static_cast<uint32_t>(work.tbsData3Signature.Length())));
         SuccessOrExit(err = tlvWriter.EndContainer(outerContainerType));
         SuccessOrExit(err = tlvWriter.Finalize());
-        msg_r3_encrypted_len = static_cast<size_t>(tlvWriter.GetLengthWritten());
+        work.msg_r3_encrypted_len = static_cast<size_t>(tlvWriter.GetLengthWritten());
+    }
+
+exit:
+    work.status = err;
+
+    auto err2 = DeviceLayer::PlatformMgr().ScheduleWork(
+        [](intptr_t arg) {
+            auto & work2 = *reinterpret_cast<SendSigma3Work *>(arg);
+            work2.session->SendSigma3c(work2);
+        },
+        reinterpret_cast<intptr_t>(&work));
+
+    if (err2 != CHIP_NO_ERROR)
+    {
+        Platform::Delete(&work); // scheduling failed, so delete
     }
+}
+
+CHIP_ERROR CASESession::SendSigma3c(SendSigma3Work & work)
+{
+    CHIP_ERROR err = CHIP_NO_ERROR;
+
+    System::PacketBufferHandle msg_R3;
+    size_t data_len;
+
+    uint8_t msg_salt[kIPKSize + kSHA256_Hash_Length];
+
+    uint8_t sr3k[CHIP_CRYPTO_SYMMETRIC_KEY_LENGTH_BYTES];
+
+    ChipLogError(DeviceLayer, "@@@@@@@@@@ SendSigma3c");
+
+    // TODO add checks here (ignoreFailure?)
+
+    // Special case: if for whatever reason not in expected state or sequence,
+    // don't do anything, including sending a status report or aborting the
+    // pending establish.
+    VerifyOrExit(mState == State::kBackgroundPending, err = CHIP_ERROR_INCORRECT_STATE);
+    VerifyOrExit(mSequence == work.sequence, err = CHIP_ERROR_INCORRECT_STATE);
+
+    SuccessOrExit(err = work.status);
 
     // Generate S3K key
     {
@@ -1211,13 +1336,13 @@ CHIP_ERROR CASESession::SendSigma3()
     }
 
     // Generated Encrypted data blob
-    err = AES_CCM_encrypt(msg_R3_Encrypted.Get(), msg_r3_encrypted_len, nullptr, 0, sr3k, CHIP_CRYPTO_SYMMETRIC_KEY_LENGTH_BYTES,
-                          kTBEData3_Nonce, kTBEDataNonceLength, msg_R3_Encrypted.Get(),
-                          msg_R3_Encrypted.Get() + msg_r3_encrypted_len, CHIP_CRYPTO_AEAD_MIC_LENGTH_BYTES);
+    err = AES_CCM_encrypt(work.msg_R3_Encrypted.Get(), work.msg_r3_encrypted_len, nullptr, 0, sr3k, CHIP_CRYPTO_SYMMETRIC_KEY_LENGTH_BYTES,
+                        kTBEData3_Nonce, kTBEDataNonceLength, work.msg_R3_Encrypted.Get(),
+                        work.msg_R3_Encrypted.Get() + work.msg_r3_encrypted_len, CHIP_CRYPTO_AEAD_MIC_LENGTH_BYTES);
     SuccessOrExit(err);
 
     // Generate Sigma3 Msg
-    data_len = TLV::EstimateStructOverhead(CHIP_CRYPTO_AEAD_MIC_LENGTH_BYTES, msg_r3_encrypted_len);
+    data_len = TLV::EstimateStructOverhead(CHIP_CRYPTO_AEAD_MIC_LENGTH_BYTES, work.msg_r3_encrypted_len);
 
     msg_R3 = System::PacketBufferHandle::New(data_len);
     VerifyOrExit(!msg_R3.IsNull(), err = CHIP_ERROR_NO_MEMORY);
@@ -1229,8 +1354,8 @@ CHIP_ERROR CASESession::SendSigma3()
         tlvWriter.Init(std::move(msg_R3));
         err = tlvWriter.StartContainer(TLV::AnonymousTag(), TLV::kTLVType_Structure, outerContainerType);
         SuccessOrExit(err);
-        err = tlvWriter.PutBytes(TLV::ContextTag(1), msg_R3_Encrypted.Get(),
-                                 static_cast<uint32_t>(msg_r3_encrypted_len + CHIP_CRYPTO_AEAD_MIC_LENGTH_BYTES));
+        err = tlvWriter.PutBytes(TLV::ContextTag(1), work.msg_R3_Encrypted.Get(),
+                                static_cast<uint32_t>(work.msg_r3_encrypted_len + CHIP_CRYPTO_AEAD_MIC_LENGTH_BYTES));
         SuccessOrExit(err);
         err = tlvWriter.EndContainer(outerContainerType);
         SuccessOrExit(err);
@@ -1243,27 +1368,33 @@ CHIP_ERROR CASESession::SendSigma3()
 
     // Call delegate to send the Msg3 to peer
     err = mExchangeCtxt->SendMessage(Protocols::SecureChannel::MsgType::CASE_Sigma3, std::move(msg_R3),
-                                     SendFlags(SendMessageFlags::kExpectResponse));
+                                    SendFlags(SendMessageFlags::kExpectResponse));
     SuccessOrExit(err);
 
     ChipLogProgress(SecureChannel, "Sent Sigma3 msg");
 
-    err = mCommissioningHash.Finish(messageDigestSpan);
-    SuccessOrExit(err);
+    {
+        MutableByteSpan messageDigestSpan(mMessageDigest);
+        SuccessOrExit(err = mCommissioningHash.Finish(messageDigestSpan));
+    }
 
     mState = State::kSentSigma3;
 
 exit:
+    Platform::Delete(&work);
+
+    // TODO might need ignoreFailure to handle some paths
 
     if (err != CHIP_NO_ERROR)
     {
         SendStatusReport(mExchangeCtxt, kProtocolCodeInvalidParam);
         mState = State::kInitialized;
     }
+
     return err;
 }
 
-struct CASESession::Sigma3Work
+struct CASESession::HandleSigma3Work
 {
     // Status of background processing.
     CHIP_ERROR status;
@@ -1316,7 +1447,7 @@ CHIP_ERROR CASESession::HandleSigma3a(System::PacketBufferHandle && msg)
 
     ChipLogProgress(SecureChannel, "Received Sigma3 msg");
 
-    auto * workPtr = Platform::New<Sigma3Work>();
+    auto * workPtr = Platform::New<HandleSigma3Work>();
     VerifyOrExit(workPtr != nullptr, err = CHIP_ERROR_NO_MEMORY);
     {
         auto & work = *workPtr;
@@ -1442,7 +1573,7 @@ CHIP_ERROR CASESession::HandleSigma3a(System::PacketBufferHandle && msg)
 
         SuccessOrExit(
             err = DeviceLayer::PlatformMgr().ScheduleBackgroundWork(
-                [](intptr_t arg) { HandleSigma3b(*reinterpret_cast<Sigma3Work *>(arg)); }, reinterpret_cast<intptr_t>(&work)));
+                [](intptr_t arg) { HandleSigma3b(*reinterpret_cast<HandleSigma3Work *>(arg)); }, reinterpret_cast<intptr_t>(&work)));
         workPtr = nullptr; // scheduling succeeded, so don't delete
         mExchangeCtxt->WillSendMessage();
         mState = State::kBackgroundPending;
@@ -1459,7 +1590,7 @@ CHIP_ERROR CASESession::HandleSigma3a(System::PacketBufferHandle && msg)
     return err;
 }
 
-void CASESession::HandleSigma3b(Sigma3Work & work)
+void CASESession::HandleSigma3b(HandleSigma3Work & work)
 {
     CHIP_ERROR err = CHIP_NO_ERROR;
 
@@ -1496,7 +1627,7 @@ void CASESession::HandleSigma3b(Sigma3Work & work)
 
     auto err2 = DeviceLayer::PlatformMgr().ScheduleWork(
         [](intptr_t arg) {
-            auto & work2 = *reinterpret_cast<Sigma3Work *>(arg);
+            auto & work2 = *reinterpret_cast<HandleSigma3Work *>(arg);
             work2.session->HandleSigma3c(work2);
         },
         reinterpret_cast<intptr_t>(&work));
@@ -1507,7 +1638,7 @@ void CASESession::HandleSigma3b(Sigma3Work & work)
     }
 }
 
-CHIP_ERROR CASESession::HandleSigma3c(Sigma3Work & work)
+CHIP_ERROR CASESession::HandleSigma3c(HandleSigma3Work & work)
 {
     CHIP_ERROR err     = CHIP_NO_ERROR;
     bool ignoreFailure = true;
diff --git a/src/protocols/secure_channel/CASESession.h b/src/protocols/secure_channel/CASESession.h
index 844425d655f41b..500b7ee4bbd0ca 100644
--- a/src/protocols/secure_channel/CASESession.h
+++ b/src/protocols/secure_channel/CASESession.h
@@ -223,11 +223,15 @@ class DLL_EXPORT CASESession : public Messaging::UnsolicitedMessageHandler,
     CHIP_ERROR HandleSigma2(System::PacketBufferHandle && msg);
     CHIP_ERROR HandleSigma2Resume(System::PacketBufferHandle && msg);
 
-    CHIP_ERROR SendSigma3();
-    struct Sigma3Work;
+    struct SendSigma3Work;
+    CHIP_ERROR SendSigma3a();
+    static void SendSigma3b(SendSigma3Work & work);
+    CHIP_ERROR SendSigma3c(SendSigma3Work & work);
+
+    struct HandleSigma3Work;
     CHIP_ERROR HandleSigma3a(System::PacketBufferHandle && msg);
-    static void HandleSigma3b(Sigma3Work & work);
-    CHIP_ERROR HandleSigma3c(Sigma3Work & work);
+    static void HandleSigma3b(HandleSigma3Work & work);
+    CHIP_ERROR HandleSigma3c(HandleSigma3Work & work);
 
     CHIP_ERROR SendSigma2Resume();
 

From 1187bee75a8149956b1f6fb1069ea433b5de05e2 Mon Sep 17 00:00:00 2001
From: Marc Lepage <mlepage@google.com>
Date: Mon, 27 Feb 2023 16:00:39 -0500
Subject: [PATCH 02/19] Fix host tests

---
 src/protocols/secure_channel/CASESession.cpp  | 40 ++++---------------
 src/protocols/secure_channel/CASESession.h    | 19 ++++-----
 .../secure_channel/tests/TestCASESession.cpp  | 18 ++++-----
 3 files changed, 26 insertions(+), 51 deletions(-)

diff --git a/src/protocols/secure_channel/CASESession.cpp b/src/protocols/secure_channel/CASESession.cpp
index 2303888c5b51f3..3b2c2f05e685ba 100644
--- a/src/protocols/secure_channel/CASESession.cpp
+++ b/src/protocols/secure_channel/CASESession.cpp
@@ -96,31 +96,6 @@ enum
 
 namespace chip {
 
-// Simple timing system
-const char* label = nullptr;
-System::Clock::Timestamp start;
-void StartTiming(const char* s) {
-    System::Clock::Timestamp now = System::SystemClock().GetMonotonicTimestamp();
-    if (label) {
-        uint32_t deltaMs = System::Clock::Milliseconds32(now - start).count();
-        ChipLogError(DeviceLayer, "Timing '%s' --> %" PRIu32 " ms", label, deltaMs);
-    }
-    label = s;
-    start = now;
-}
-void EndTiming() {
-    System::Clock::Timestamp now = System::SystemClock().GetMonotonicTimestamp();
-    if (label) {
-        uint32_t deltaMs = System::Clock::Milliseconds32(now - start).count();
-        ChipLogError(DeviceLayer, "Timing '%s' --> %" PRIu32 " ms", label, deltaMs);
-    }
-    label = nullptr;
-}
-
-}
-
-namespace chip {
-
 using namespace Crypto;
 using namespace Credentials;
 using namespace Messaging;
@@ -1207,10 +1182,7 @@ CHIP_ERROR CASESession::SendSigma3a()
                 [](intptr_t arg) { SendSigma3b(*reinterpret_cast<SendSigma3Work *>(arg)); }, reinterpret_cast<intptr_t>(&work)));
         workPtr = nullptr; // scheduling succeeded, so don't delete
         mExchangeCtxt->WillSendMessage();
-        mState = State::kBackgroundPending;
-
-        // TODO decide if State::kBackgroundPending is sufficient
-        // or should have another state
+        mState = State::kSendSigma3BackgroundPending;
     }
 
 exit:
@@ -1297,14 +1269,14 @@ CHIP_ERROR CASESession::SendSigma3c(SendSigma3Work & work)
 
     AutoReleaseSessionKey sr3k(*mSessionManager->GetSessionKeystore());
 
-    ChipLogError(DeviceLayer, "@@@@@@@@@@ SendSigma3c");
+    ChipLogError(DeviceLayer, "@@@@@@@@@@ SendSigma3c %d", (int)work.status.AsInteger());
 
     // TODO add checks here (ignoreFailure?)
 
     // Special case: if for whatever reason not in expected state or sequence,
     // don't do anything, including sending a status report or aborting the
     // pending establish.
-    VerifyOrExit(mState == State::kBackgroundPending, err = CHIP_ERROR_INCORRECT_STATE);
+    VerifyOrExit(mState == State::kSendSigma3BackgroundPending, err = CHIP_ERROR_INCORRECT_STATE);
     VerifyOrExit(mSequence == work.sequence, err = CHIP_ERROR_INCORRECT_STATE);
 
     SuccessOrExit(err = work.status);
@@ -1365,6 +1337,8 @@ CHIP_ERROR CASESession::SendSigma3c(SendSigma3Work & work)
 
     // TODO might need ignoreFailure to handle some paths
 
+    ChipLogError(DeviceLayer, "@@@@@@@@@@ SendSigma3c DONE %d", (int)err.AsInteger());
+
     if (err != CHIP_NO_ERROR)
     {
         SendStatusReport(mExchangeCtxt, kProtocolCodeInvalidParam);
@@ -1550,7 +1524,7 @@ CHIP_ERROR CASESession::HandleSigma3a(System::PacketBufferHandle && msg)
                 [](intptr_t arg) { HandleSigma3b(*reinterpret_cast<HandleSigma3Work *>(arg)); }, reinterpret_cast<intptr_t>(&work)));
         workPtr = nullptr; // scheduling succeeded, so don't delete
         mExchangeCtxt->WillSendMessage();
-        mState = State::kBackgroundPending;
+        mState = State::kHandleSigma3BackgroundPending;
     }
 
 exit:
@@ -1620,7 +1594,7 @@ CHIP_ERROR CASESession::HandleSigma3c(HandleSigma3Work & work)
     // Special case: if for whatever reason not in expected state or sequence,
     // don't do anything, including sending a status report or aborting the
     // pending establish.
-    VerifyOrExit(mState == State::kBackgroundPending, err = CHIP_ERROR_INCORRECT_STATE);
+    VerifyOrExit(mState == State::kHandleSigma3BackgroundPending, err = CHIP_ERROR_INCORRECT_STATE);
     VerifyOrExit(mSequence == work.sequence, err = CHIP_ERROR_INCORRECT_STATE);
 
     ignoreFailure = false;
diff --git a/src/protocols/secure_channel/CASESession.h b/src/protocols/secure_channel/CASESession.h
index 863b6a4d4f026c..483b1a480fdf20 100644
--- a/src/protocols/secure_channel/CASESession.h
+++ b/src/protocols/secure_channel/CASESession.h
@@ -184,15 +184,16 @@ class DLL_EXPORT CASESession : public Messaging::UnsolicitedMessageHandler,
     friend class TestCASESession;
     enum class State : uint8_t
     {
-        kInitialized       = 0,
-        kSentSigma1        = 1,
-        kSentSigma2        = 2,
-        kSentSigma3        = 3,
-        kSentSigma1Resume  = 4,
-        kSentSigma2Resume  = 5,
-        kFinished          = 6,
-        kFinishedViaResume = 7,
-        kBackgroundPending = 8,
+        kInitialized                   = 0,
+        kSentSigma1                    = 1,
+        kSentSigma2                    = 2,
+        kSentSigma3                    = 3,
+        kSentSigma1Resume              = 4,
+        kSentSigma2Resume              = 5,
+        kFinished                      = 6,
+        kFinishedViaResume             = 7,
+        kSendSigma3BackgroundPending   = 8,
+        kHandleSigma3BackgroundPending = 9,
     };
 
     /*
diff --git a/src/protocols/secure_channel/tests/TestCASESession.cpp b/src/protocols/secure_channel/tests/TestCASESession.cpp
index a863b3ca57dae8..2e753ac996ca2d 100644
--- a/src/protocols/secure_channel/tests/TestCASESession.cpp
+++ b/src/protocols/secure_channel/tests/TestCASESession.cpp
@@ -60,16 +60,16 @@ namespace {
 
 void ServiceEvents(TestContext & ctx)
 {
-    // Service any messages
-    ctx.DrainAndServiceIO();
-
-    // Messages may have scheduled work, so service them
-    chip::DeviceLayer::PlatformMgr().ScheduleWork([](intptr_t) -> void { chip::DeviceLayer::PlatformMgr().StopEventLoopTask(); },
-                                                  (intptr_t) nullptr);
-    chip::DeviceLayer::PlatformMgr().RunEventLoop();
+    // Takes a few rounds of this because handling IO messages may schedule work,
+    // and scheduled work may queue messages for sending...
+    for (int i = 0; i < 3; ++i)
+    {
+        ctx.DrainAndServiceIO();
 
-    // Work may have sent messages, so service them
-    ctx.DrainAndServiceIO();
+        chip::DeviceLayer::PlatformMgr().ScheduleWork([](intptr_t) -> void { chip::DeviceLayer::PlatformMgr().StopEventLoopTask(); },
+                                                    (intptr_t) nullptr);
+        chip::DeviceLayer::PlatformMgr().RunEventLoop();
+    }
 }
 
 class TemporarySessionManager

From 52f4a6f04c0228bbc272f23124700b988389ccd6 Mon Sep 17 00:00:00 2001
From: Marc Lepage <mlepage@google.com>
Date: Wed, 15 Mar 2023 15:07:15 -0400
Subject: [PATCH 03/19] Remove temp log statements

---
 src/protocols/secure_channel/CASESession.cpp | 21 +++++---------------
 1 file changed, 5 insertions(+), 16 deletions(-)

diff --git a/src/protocols/secure_channel/CASESession.cpp b/src/protocols/secure_channel/CASESession.cpp
index 69f5e515821932..712edce6677b25 100644
--- a/src/protocols/secure_channel/CASESession.cpp
+++ b/src/protocols/secure_channel/CASESession.cpp
@@ -1141,8 +1141,6 @@ CHIP_ERROR CASESession::SendSigma3a()
 
     ChipLogDetail(SecureChannel, "Sending Sigma3");
 
-    ChipLogError(DeviceLayer, "@@@@@@@@@@ SendSigma3a");
-
     auto * workPtr = Platform::New<SendSigma3Work>();
     VerifyOrExit(workPtr != nullptr, err = CHIP_ERROR_NO_MEMORY);
     {
@@ -1202,10 +1200,6 @@ void CASESession::SendSigma3b(SendSigma3Work & work)
 {
     CHIP_ERROR err = CHIP_NO_ERROR;
 
-    ChipLogError(DeviceLayer, "@@@@@@@@@@ SendSigma3b");
-
-    // TODO SignWithOpKeypair has moved here but check for race condition etc.
-
     // Generate a signature
     err = work.fabricTable->SignWithOpKeypair(work.fabricIndex, ByteSpan{ work.msg_R3_Signed.Get(), work.msg_r3_signed_len }, work.tbsData3Signature);
     SuccessOrExit(err);
@@ -1261,7 +1255,8 @@ void CASESession::SendSigma3b(SendSigma3Work & work)
 
 CHIP_ERROR CASESession::SendSigma3c(SendSigma3Work & work)
 {
-    CHIP_ERROR err = CHIP_NO_ERROR;
+    CHIP_ERROR err     = CHIP_NO_ERROR;
+    bool ignoreFailure = true;
 
     System::PacketBufferHandle msg_R3;
     size_t data_len;
@@ -1270,16 +1265,14 @@ CHIP_ERROR CASESession::SendSigma3c(SendSigma3Work & work)
 
     AutoReleaseSessionKey sr3k(*mSessionManager->GetSessionKeystore());
 
-    ChipLogError(DeviceLayer, "@@@@@@@@@@ SendSigma3c %d", (int)work.status.AsInteger());
-
-    // TODO add checks here (ignoreFailure?)
-
     // Special case: if for whatever reason not in expected state or sequence,
     // don't do anything, including sending a status report or aborting the
     // pending establish.
     VerifyOrExit(mState == State::kSendSigma3BackgroundPending, err = CHIP_ERROR_INCORRECT_STATE);
     VerifyOrExit(mSequence == work.sequence, err = CHIP_ERROR_INCORRECT_STATE);
 
+    ignoreFailure = false;
+
     SuccessOrExit(err = work.status);
 
     // Generate S3K key
@@ -1336,11 +1329,7 @@ CHIP_ERROR CASESession::SendSigma3c(SendSigma3Work & work)
 exit:
     Platform::Delete(&work);
 
-    // TODO might need ignoreFailure to handle some paths
-
-    ChipLogError(DeviceLayer, "@@@@@@@@@@ SendSigma3c DONE %d", (int)err.AsInteger());
-
-    if (err != CHIP_NO_ERROR)
+    if (err != CHIP_NO_ERROR && !ignoreFailure)
     {
         SendStatusReport(mExchangeCtxt, kProtocolCodeInvalidParam);
         mState = State::kInitialized;

From 7bbd6f8bcf8eefabca9308e19c404c4b164a16e4 Mon Sep 17 00:00:00 2001
From: Marc Lepage <mlepage@google.com>
Date: Wed, 15 Mar 2023 15:54:18 -0400
Subject: [PATCH 04/19] Restyle

---
 src/protocols/secure_channel/CASESession.cpp  | 36 +++++++++++--------
 .../secure_channel/tests/TestCASESession.cpp  |  4 +--
 2 files changed, 23 insertions(+), 17 deletions(-)

diff --git a/src/protocols/secure_channel/CASESession.cpp b/src/protocols/secure_channel/CASESession.cpp
index 712edce6677b25..d23397035695f5 100644
--- a/src/protocols/secure_channel/CASESession.cpp
+++ b/src/protocols/secure_channel/CASESession.cpp
@@ -1169,12 +1169,14 @@ CHIP_ERROR CASESession::SendSigma3a()
         SuccessOrExit(err = mFabricsTable->FetchNOCCert(mFabricIndex, work.nocCert));
 
         // Prepare Sigma3 TBS Data Blob
-        work.msg_r3_signed_len = TLV::EstimateStructOverhead(work.icaCert.size(), work.nocCert.size(), kP256_PublicKey_Length, kP256_PublicKey_Length);
+        work.msg_r3_signed_len =
+            TLV::EstimateStructOverhead(work.icaCert.size(), work.nocCert.size(), kP256_PublicKey_Length, kP256_PublicKey_Length);
 
         VerifyOrExit(work.msg_R3_Signed.Alloc(work.msg_r3_signed_len), err = CHIP_ERROR_NO_MEMORY);
 
-        SuccessOrExit(err = ConstructTBSData(work.nocCert, work.icaCert, ByteSpan(mEphemeralKey->Pubkey(), mEphemeralKey->Pubkey().Length()),
-                                            ByteSpan(mRemotePubKey, mRemotePubKey.Length()), work.msg_R3_Signed.Get(), work.msg_r3_signed_len));
+        SuccessOrExit(err = ConstructTBSData(
+                          work.nocCert, work.icaCert, ByteSpan(mEphemeralKey->Pubkey(), mEphemeralKey->Pubkey().Length()),
+                          ByteSpan(mRemotePubKey, mRemotePubKey.Length()), work.msg_R3_Signed.Get(), work.msg_r3_signed_len));
 
         SuccessOrExit(
             err = DeviceLayer::PlatformMgr().ScheduleBackgroundWork(
@@ -1201,13 +1203,16 @@ void CASESession::SendSigma3b(SendSigma3Work & work)
     CHIP_ERROR err = CHIP_NO_ERROR;
 
     // Generate a signature
-    err = work.fabricTable->SignWithOpKeypair(work.fabricIndex, ByteSpan{ work.msg_R3_Signed.Get(), work.msg_r3_signed_len }, work.tbsData3Signature);
+    err = work.fabricTable->SignWithOpKeypair(work.fabricIndex, ByteSpan{ work.msg_R3_Signed.Get(), work.msg_r3_signed_len },
+                                              work.tbsData3Signature);
     SuccessOrExit(err);
 
     // Prepare Sigma3 TBE Data Blob
-    work.msg_r3_encrypted_len = TLV::EstimateStructOverhead(work.nocCert.size(), work.icaCert.size(), work.tbsData3Signature.Length());
+    work.msg_r3_encrypted_len =
+        TLV::EstimateStructOverhead(work.nocCert.size(), work.icaCert.size(), work.tbsData3Signature.Length());
 
-    VerifyOrExit(work.msg_R3_Encrypted.Alloc(work.msg_r3_encrypted_len + CHIP_CRYPTO_AEAD_MIC_LENGTH_BYTES), err = CHIP_ERROR_NO_MEMORY);
+    VerifyOrExit(work.msg_R3_Encrypted.Alloc(work.msg_r3_encrypted_len + CHIP_CRYPTO_AEAD_MIC_LENGTH_BYTES),
+                 err = CHIP_ERROR_NO_MEMORY);
 
     {
         TLV::TLVWriter tlvWriter;
@@ -1231,7 +1236,7 @@ void CASESession::SendSigma3b(SendSigma3Work & work)
         }
 
         SuccessOrExit(err = tlvWriter.PutBytes(TLV::ContextTag(kTag_TBEData_Signature), work.tbsData3Signature.ConstBytes(),
-                                            static_cast<uint32_t>(work.tbsData3Signature.Length())));
+                                               static_cast<uint32_t>(work.tbsData3Signature.Length())));
         SuccessOrExit(err = tlvWriter.EndContainer(outerContainerType));
         SuccessOrExit(err = tlvWriter.Finalize());
         work.msg_r3_encrypted_len = static_cast<size_t>(tlvWriter.GetLengthWritten());
@@ -1283,9 +1288,10 @@ CHIP_ERROR CASESession::SendSigma3c(SendSigma3Work & work)
     }
 
     // Generated Encrypted data blob
-    SuccessOrExit(err = AES_CCM_encrypt(work.msg_R3_Encrypted.Get(), work.msg_r3_encrypted_len, nullptr, 0, sr3k.KeyHandle(), kTBEData3_Nonce,
-                                        kTBEDataNonceLength, work.msg_R3_Encrypted.Get(), work.msg_R3_Encrypted.Get() + work.msg_r3_encrypted_len,
-                                        CHIP_CRYPTO_AEAD_MIC_LENGTH_BYTES));
+    SuccessOrExit(err =
+                      AES_CCM_encrypt(work.msg_R3_Encrypted.Get(), work.msg_r3_encrypted_len, nullptr, 0, sr3k.KeyHandle(),
+                                      kTBEData3_Nonce, kTBEDataNonceLength, work.msg_R3_Encrypted.Get(),
+                                      work.msg_R3_Encrypted.Get() + work.msg_r3_encrypted_len, CHIP_CRYPTO_AEAD_MIC_LENGTH_BYTES));
 
     // Generate Sigma3 Msg
     data_len = TLV::EstimateStructOverhead(CHIP_CRYPTO_AEAD_MIC_LENGTH_BYTES, work.msg_r3_encrypted_len);
@@ -1301,7 +1307,7 @@ CHIP_ERROR CASESession::SendSigma3c(SendSigma3Work & work)
         err = tlvWriter.StartContainer(TLV::AnonymousTag(), TLV::kTLVType_Structure, outerContainerType);
         SuccessOrExit(err);
         err = tlvWriter.PutBytes(TLV::ContextTag(1), work.msg_R3_Encrypted.Get(),
-                                static_cast<uint32_t>(work.msg_r3_encrypted_len + CHIP_CRYPTO_AEAD_MIC_LENGTH_BYTES));
+                                 static_cast<uint32_t>(work.msg_r3_encrypted_len + CHIP_CRYPTO_AEAD_MIC_LENGTH_BYTES));
         SuccessOrExit(err);
         err = tlvWriter.EndContainer(outerContainerType);
         SuccessOrExit(err);
@@ -1314,7 +1320,7 @@ CHIP_ERROR CASESession::SendSigma3c(SendSigma3Work & work)
 
     // Call delegate to send the Msg3 to peer
     err = mExchangeCtxt->SendMessage(Protocols::SecureChannel::MsgType::CASE_Sigma3, std::move(msg_R3),
-                                    SendFlags(SendMessageFlags::kExpectResponse));
+                                     SendFlags(SendMessageFlags::kExpectResponse));
     SuccessOrExit(err);
 
     ChipLogProgress(SecureChannel, "Sent Sigma3 msg");
@@ -1509,9 +1515,9 @@ CHIP_ERROR CASESession::HandleSigma3a(System::PacketBufferHandle && msg)
             }
         }
 
-        SuccessOrExit(
-            err = DeviceLayer::PlatformMgr().ScheduleBackgroundWork(
-                [](intptr_t arg) { HandleSigma3b(*reinterpret_cast<HandleSigma3Work *>(arg)); }, reinterpret_cast<intptr_t>(&work)));
+        SuccessOrExit(err = DeviceLayer::PlatformMgr().ScheduleBackgroundWork(
+                          [](intptr_t arg) { HandleSigma3b(*reinterpret_cast<HandleSigma3Work *>(arg)); },
+                          reinterpret_cast<intptr_t>(&work)));
         workPtr = nullptr; // scheduling succeeded, so don't delete
         mExchangeCtxt->WillSendMessage();
         mState = State::kHandleSigma3BackgroundPending;
diff --git a/src/protocols/secure_channel/tests/TestCASESession.cpp b/src/protocols/secure_channel/tests/TestCASESession.cpp
index 2e753ac996ca2d..11d5000f71df03 100644
--- a/src/protocols/secure_channel/tests/TestCASESession.cpp
+++ b/src/protocols/secure_channel/tests/TestCASESession.cpp
@@ -66,8 +66,8 @@ void ServiceEvents(TestContext & ctx)
     {
         ctx.DrainAndServiceIO();
 
-        chip::DeviceLayer::PlatformMgr().ScheduleWork([](intptr_t) -> void { chip::DeviceLayer::PlatformMgr().StopEventLoopTask(); },
-                                                    (intptr_t) nullptr);
+        chip::DeviceLayer::PlatformMgr().ScheduleWork(
+            [](intptr_t) -> void { chip::DeviceLayer::PlatformMgr().StopEventLoopTask(); }, (intptr_t) nullptr);
         chip::DeviceLayer::PlatformMgr().RunEventLoop();
     }
 }

From ebc6de0449422e584da6bd62e02bf778fe6c8be1 Mon Sep 17 00:00:00 2001
From: Marc Lepage <mlepage@google.com>
Date: Mon, 27 Mar 2023 15:06:34 -0400
Subject: [PATCH 05/19] Refactor CASESession::SendSigma3

Add more structure to manage multiple steps of work.
---
 src/protocols/secure_channel/CASESession.cpp | 384 ++++++++++++-------
 src/protocols/secure_channel/CASESession.h   |  35 +-
 2 files changed, 273 insertions(+), 146 deletions(-)

diff --git a/src/protocols/secure_channel/CASESession.cpp b/src/protocols/secure_channel/CASESession.cpp
index 3b0ff3de8147cf..b83a836f6332c0 100644
--- a/src/protocols/secure_channel/CASESession.cpp
+++ b/src/protocols/secure_channel/CASESession.cpp
@@ -25,7 +25,9 @@
  */
 #include <protocols/secure_channel/CASESession.h>
 
+#include <atomic>
 #include <inttypes.h>
+#include <memory>
 #include <string.h>
 
 #include <lib/core/CHIPEncoding.h>
@@ -130,14 +132,198 @@ static constexpr ExchangeContext::Timeout kExpectedLowProcessingTime    = System
 static constexpr ExchangeContext::Timeout kExpectedSigma1ProcessingTime = kExpectedLowProcessingTime;
 static constexpr ExchangeContext::Timeout kExpectedHighProcessingTime   = System::Clock::Seconds16(30);
 
+// Helper for managing a session's outstanding work.
+// Holds work data which is provided to a scheduled work callback (standalone),
+// then (if not canceled) to a scheduled after work callback (on the session).
+template <class DATA>
+class CASESession::WorkHelper
+{
+public:
+    // Work callback, processed in the background via `PlatformManager::ScheduleBackgroundWork`.
+    // This is a non-member function which does not use the associated session.
+    // The return value is passed to the after work callback (called afterward).
+    // Set `cancel` to true if calling the after work callback is not necessary.
+    typedef CHIP_ERROR (*WorkCallback)(DATA & data, bool & cancel);
+
+    // After work callback, processed in the main CHIP task via `PlatformManager::ScheduleWork`.
+    // This is a member function to be called on the associated session after the work callback.
+    // The `status` value is the result of the work callback (called beforehand).
+    typedef CHIP_ERROR (CASESession::*AfterWorkCallback)(DATA & data, CHIP_ERROR status);
+
+    // Create a work helper using the specified session, work callback, after work callback, and data (template arg).
+    // Lifetime is not managed, see `Create` for that option.
+    WorkHelper(CASESession & session, WorkCallback workCallback, AfterWorkCallback afterWorkCallback) : mSession(&session), mWorkCallback(workCallback), mAfterWorkCallback(afterWorkCallback)
+    {
+    }
+
+    ~WorkHelper()
+    {
+        ChipLogError(SecureChannel, "@@@@@ WorkHelper ~dtor %p", this);
+    }
+
+    // Create a work helper using the specified session, work callback, after work callback, and data (template arg).
+    // Lifetime is managed by sharing between the caller (typically the session) and the helper itself (while work is scheduled).
+    static std::shared_ptr<WorkHelper> Create(CASESession & session, WorkCallback workCallback, AfterWorkCallback afterWorkCallback)
+    {
+        std::shared_ptr<WorkHelper> ptr(Platform::New<WorkHelper>(session, workCallback, afterWorkCallback), Platform::Delete<WorkHelper>);
+        if (ptr)
+        {
+            ptr->mWeakPtr = ptr; // used by `ScheduleWork`
+        }
+        return ptr;
+    }
+
+    // Schedule the work after configuring the data.
+    // If lifetime is managed, the helper shares management while work is outstanding.
+    CHIP_ERROR ScheduleWork()
+    {
+        if (!mSession || !mWorkCallback || !mAfterWorkCallback)
+        {
+            return CHIP_ERROR_INCORRECT_STATE;
+        }
+        mStrongPtr = mWeakPtr.lock(); // set in `Create`
+        ChipLogError(SecureChannel, "@@@@@ scheduling WorkHandler %p", this);
+        auto status = DeviceLayer::PlatformMgr().ScheduleBackgroundWork(WorkHandler, reinterpret_cast<intptr_t>(this));
+        if (status != CHIP_NO_ERROR)
+        {
+            ChipLogError(SecureChannel, "@@@@@ ScheduleWork releasing helper %p", this);
+            mStrongPtr.reset();
+        }
+        return status;
+    }
+ 
+    // Cancel the work, by clearing the associated session.
+    void CancelWork()
+    {
+        mSession.store(nullptr);
+    }
+
+private:
+    // Handler for the work callback.
+    static void WorkHandler(intptr_t arg)
+    {
+        WorkHelper* helper = reinterpret_cast<WorkHelper *>(arg);
+        ChipLogError(SecureChannel, "@@@@@ WorkHandler %p", helper);
+        bool cancel = false;
+        VerifyOrExit(helper->mSession.load(), ;); // cancelled by `CancelWork`?
+        ChipLogError(SecureChannel, "@@@@@ calling mWorkCallback %p", helper);
+        helper->mStatus = helper->mWorkCallback(helper->mData, cancel);
+        ChipLogError(SecureChannel, "@@@@@ done mWorkCallback %p", helper);
+        VerifyOrExit(!cancel, ;); // canceled by `mWorkCallback`?
+        VerifyOrExit(helper->mSession.load(), ;); // cancelled by `CancelWork`?
+        ChipLogError(SecureChannel, "@@@@@ scheduling AfterWorkHandler %p", helper);
+        SuccessOrExit(DeviceLayer::PlatformMgr().ScheduleWork(AfterWorkHandler, reinterpret_cast<intptr_t>(helper)));
+        return;
+    exit:
+        ChipLogError(SecureChannel, "@@@@@ WorkHandler releasing helper %p", helper);
+        helper->mStrongPtr.reset();
+    }
+
+    // Handler for the after work callback.
+    static void AfterWorkHandler(intptr_t arg)
+    {
+        // TODO some notes here about how we should be on the main event thread
+        // so shouldn't be deleting the session out from under while this executes
+        // (double check that's true, that timers don't run on some other thread, etc.)
+        WorkHelper* helper = reinterpret_cast<WorkHelper *>(arg);
+        ChipLogError(SecureChannel, "@@@@@ AfterWorkHandler %p", helper);
+        if (auto * session = helper->mSession.load())
+        {
+            ChipLogError(SecureChannel, "@@@@@ calling mAfterWorkCallback %p", helper);
+            (session->*(helper->mAfterWorkCallback))(helper->mData, helper->mStatus);
+            ChipLogError(SecureChannel, "@@@@@ done mAfterWorkCallback %p", helper);
+        }
+        else
+        {
+            ChipLogError(SecureChannel, "@@@@@ session missing %p", helper);
+        }
+        ChipLogError(SecureChannel, "@@@@@ AfterWorkHandler releasing helper %p", helper);
+        helper->mStrongPtr.reset();
+    }
+
+private:
+    // Lifetime management: `ScheduleWork` sets `mStrongPtr` from `mWeakPtr`.
+    std::weak_ptr<WorkHelper> mWeakPtr;
+
+    // Lifetime management: `ScheduleWork` sets `mStrongPtr` from `mWeakPtr`.
+    std::shared_ptr<WorkHelper> mStrongPtr;
+
+    // Associated session, cleared by `CancelWork`.
+    std::atomic<CASESession *> mSession;
+
+    // Work callback, called by `WorkHandler`.
+    WorkCallback mWorkCallback;
+
+    // After work callback, called by `AfterWorkHandler`.
+    AfterWorkCallback mAfterWorkCallback;
+
+    // Return value of `mWorkCallback`, passed to `mAfterWorkCallback`.
+    CHIP_ERROR mStatus;
+
+public:
+    // Data passed to `mWorkCallback` and `mAfterWorkCallback`.
+    DATA mData;
+};
+
+struct CASESession::SendSigma3Data
+{
+    FabricTable * fabricTable;
+    FabricIndex fabricIndex;
+
+    chip::Platform::ScopedMemoryBuffer<uint8_t> msg_R3_Signed;
+    size_t msg_r3_signed_len;
+
+    chip::Platform::ScopedMemoryBuffer<uint8_t> msg_R3_Encrypted;
+    size_t msg_r3_encrypted_len;
+
+    chip::Platform::ScopedMemoryBuffer<uint8_t> icacBuf;
+    MutableByteSpan icaCert;
+
+    chip::Platform::ScopedMemoryBuffer<uint8_t> nocBuf;
+    MutableByteSpan nocCert;
+
+    P256ECDSASignature tbsData3Signature;
+};
+
+struct CASESession::HandleSigma3Work
+{
+    // Status of background processing.
+    CHIP_ERROR status;
+
+    // Session to use after background processing.
+    CASESession * session;
+
+    // Sequence number used to coordinate foreground/background work for a
+    // particular session establishment.
+    int sequence;
+
+    chip::Platform::ScopedMemoryBuffer<uint8_t> msg_R3_Signed;
+    size_t msg_r3_signed_len;
+
+    ByteSpan initiatorNOC;
+    ByteSpan initiatorICAC;
+
+    uint8_t rootCertBuf[kMaxCHIPCertLength];
+    ByteSpan fabricRCAC;
+
+    P256ECDSASignature tbsData3Signature;
+
+    FabricId fabricId;
+    NodeId initiatorNodeId;
+
+    ValidationContext validContext;
+};
+
 CASESession::~CASESession()
 {
+    ChipLogError(SecureChannel, "@@@@@ CASESession ~dtor %p", this);
     // Let's clear out any security state stored in the object, before destroying it.
     Clear();
 }
 
 void CASESession::OnSessionReleased()
 {
+    ChipLogError(SecureChannel, "@@@@@ OnSessionReleased %p", this);
     // Call into our super-class before we clear our state.
     PairingSession::OnSessionReleased();
     Clear();
@@ -145,6 +331,14 @@ void CASESession::OnSessionReleased()
 
 void CASESession::Clear()
 {
+    ChipLogError(SecureChannel, "@@@@@ Clear %p", this);
+    // Cancel any outstanding work.
+    if (mSendSigma3Helper)
+    {
+        mSendSigma3Helper->CancelWork();
+        mSendSigma3Helper.reset();
+    }
+
     // This function zeroes out and resets the memory used by the object.
     // It's done so that no security related information will be leaked.
     mCommissioningHash.Clear();
@@ -246,6 +440,8 @@ CHIP_ERROR CASESession::EstablishSession(SessionManager & sessionManager, Fabric
     MATTER_TRACE_EVENT_SCOPE("EstablishSession", "CASESession");
     CHIP_ERROR err = CHIP_NO_ERROR;
 
+    ChipLogError(SecureChannel, "@@@@@ EstablishSession %p", this);
+
     // Return early on error here, as we have not initialized any state yet
     ReturnErrorCodeIf(exchangeCtxt == nullptr, CHIP_ERROR_INVALID_ARGUMENT);
     ReturnErrorCodeIf(fabricTable == nullptr, CHIP_ERROR_INVALID_ARGUMENT);
@@ -294,6 +490,7 @@ CHIP_ERROR CASESession::EstablishSession(SessionManager & sessionManager, Fabric
     {
         Clear();
     }
+    ChipLogError(SecureChannel, "@@@@@ done EstablishSession %p", this);
     return err;
 }
 
@@ -311,6 +508,7 @@ void CASESession::OnResponseTimeout(ExchangeContext * ec)
 
 void CASESession::AbortPendingEstablish(CHIP_ERROR err)
 {
+    ChipLogError(SecureChannel, "@@@@@ AbortPendingEstablish %p", this);
     Clear();
     // Do this last in case the delegate frees us.
     NotifySessionEstablishmentError(err);
@@ -1104,36 +1302,6 @@ CHIP_ERROR CASESession::HandleSigma2(System::PacketBufferHandle && msg)
     return err;
 }
 
-struct CASESession::SendSigma3Work
-{
-    // Status of background processing.
-    CHIP_ERROR status;
-
-    // Session to use after background processing.
-    CASESession * session;
-
-    // Sequence number used to coordinate foreground/background work for a
-    // particular session establishment.
-    int sequence;
-
-    FabricTable * fabricTable;
-    FabricIndex fabricIndex;
-
-    chip::Platform::ScopedMemoryBuffer<uint8_t> msg_R3_Signed;
-    size_t msg_r3_signed_len;
-
-    chip::Platform::ScopedMemoryBuffer<uint8_t> msg_R3_Encrypted;
-    size_t msg_r3_encrypted_len;
-
-    chip::Platform::ScopedMemoryBuffer<uint8_t> icacBuf;
-    MutableByteSpan icaCert;
-
-    chip::Platform::ScopedMemoryBuffer<uint8_t> nocBuf;
-    MutableByteSpan nocCert;
-
-    P256ECDSASignature tbsData3Signature;
-};
-
 CHIP_ERROR CASESession::SendSigma3a()
 {
     MATTER_TRACE_EVENT_SCOPE("SendSigma3", "CASESession");
@@ -1141,54 +1309,45 @@ CHIP_ERROR CASESession::SendSigma3a()
 
     ChipLogDetail(SecureChannel, "Sending Sigma3");
 
-    auto * workPtr = Platform::New<SendSigma3Work>();
-    VerifyOrExit(workPtr != nullptr, err = CHIP_ERROR_NO_MEMORY);
+    auto helper = WorkHelper<SendSigma3Data>::Create(*this, &SendSigma3b, &CASESession::SendSigma3c);
+    VerifyOrExit(helper, err = CHIP_ERROR_NO_MEMORY);
     {
-        auto & work = *workPtr;
+        auto & data = helper->mData;
 
-        // Used to call back into the session after background event processing.
-        // It happens that there's only one pairing session (in CASEServer)
-        // so it will still be available for use. Use a sequence number to
-        // coordinate.
-        work.session  = this;
-        work.sequence = mSequence;
+        ChipLogError(SecureChannel, "@@@@@ session %p helper %p", this, helper.get());
 
         VerifyOrExit(mFabricsTable != nullptr, err = CHIP_ERROR_INCORRECT_STATE);
-        work.fabricTable = mFabricsTable;
-        work.fabricIndex = mFabricIndex;
+        data.fabricTable = mFabricsTable;
+        data.fabricIndex = mFabricIndex;
 
         VerifyOrExit(mEphemeralKey != nullptr, err = CHIP_ERROR_INTERNAL);
 
-        VerifyOrExit(work.icacBuf.Alloc(kMaxCHIPCertLength), err = CHIP_ERROR_NO_MEMORY);
-        work.icaCert = MutableByteSpan{ work.icacBuf.Get(), kMaxCHIPCertLength };
+        VerifyOrExit(data.icacBuf.Alloc(kMaxCHIPCertLength), err = CHIP_ERROR_NO_MEMORY);
+        data.icaCert = MutableByteSpan{ data.icacBuf.Get(), kMaxCHIPCertLength };
 
-        VerifyOrExit(work.nocBuf.Alloc(kMaxCHIPCertLength), err = CHIP_ERROR_NO_MEMORY);
-        work.nocCert = MutableByteSpan{ work.nocBuf.Get(), kMaxCHIPCertLength };
+        VerifyOrExit(data.nocBuf.Alloc(kMaxCHIPCertLength), err = CHIP_ERROR_NO_MEMORY);
+        data.nocCert = MutableByteSpan{ data.nocBuf.Get(), kMaxCHIPCertLength };
 
-        SuccessOrExit(err = mFabricsTable->FetchICACert(mFabricIndex, work.icaCert));
-        SuccessOrExit(err = mFabricsTable->FetchNOCCert(mFabricIndex, work.nocCert));
+        SuccessOrExit(err = mFabricsTable->FetchICACert(mFabricIndex, data.icaCert));
+        SuccessOrExit(err = mFabricsTable->FetchNOCCert(mFabricIndex, data.nocCert));
 
         // Prepare Sigma3 TBS Data Blob
-        work.msg_r3_signed_len =
-            TLV::EstimateStructOverhead(work.icaCert.size(), work.nocCert.size(), kP256_PublicKey_Length, kP256_PublicKey_Length);
+        data.msg_r3_signed_len =
+            TLV::EstimateStructOverhead(data.icaCert.size(), data.nocCert.size(), kP256_PublicKey_Length, kP256_PublicKey_Length);
 
-        VerifyOrExit(work.msg_R3_Signed.Alloc(work.msg_r3_signed_len), err = CHIP_ERROR_NO_MEMORY);
+        VerifyOrExit(data.msg_R3_Signed.Alloc(data.msg_r3_signed_len), err = CHIP_ERROR_NO_MEMORY);
 
         SuccessOrExit(err = ConstructTBSData(
-                          work.nocCert, work.icaCert, ByteSpan(mEphemeralKey->Pubkey(), mEphemeralKey->Pubkey().Length()),
-                          ByteSpan(mRemotePubKey, mRemotePubKey.Length()), work.msg_R3_Signed.Get(), work.msg_r3_signed_len));
+                          data.nocCert, data.icaCert, ByteSpan(mEphemeralKey->Pubkey(), mEphemeralKey->Pubkey().Length()),
+                          ByteSpan(mRemotePubKey, mRemotePubKey.Length()), data.msg_R3_Signed.Get(), data.msg_r3_signed_len));
 
-        SuccessOrExit(
-            err = DeviceLayer::PlatformMgr().ScheduleBackgroundWork(
-                [](intptr_t arg) { SendSigma3b(*reinterpret_cast<SendSigma3Work *>(arg)); }, reinterpret_cast<intptr_t>(&work)));
-        workPtr = nullptr; // scheduling succeeded, so don't delete
+        SuccessOrExit(err = helper->ScheduleWork());
+        mSendSigma3Helper = helper;
         mExchangeCtxt->WillSendMessage();
-        mState = State::kSendSigma3BackgroundPending;
+        mState = State::kSendSigma3Pending;
     }
 
 exit:
-    Platform::Delete(workPtr);
-
     if (err != CHIP_NO_ERROR)
     {
         SendStatusReport(mExchangeCtxt, kProtocolCodeInvalidParam);
@@ -1198,67 +1357,55 @@ CHIP_ERROR CASESession::SendSigma3a()
     return err;
 }
 
-void CASESession::SendSigma3b(SendSigma3Work & work)
+CHIP_ERROR CASESession::SendSigma3b(SendSigma3Data & data, bool & cancel)
 {
     CHIP_ERROR err = CHIP_NO_ERROR;
 
     // Generate a signature
-    err = work.fabricTable->SignWithOpKeypair(work.fabricIndex, ByteSpan{ work.msg_R3_Signed.Get(), work.msg_r3_signed_len },
-                                              work.tbsData3Signature);
+    err = data.fabricTable->SignWithOpKeypair(data.fabricIndex, ByteSpan{ data.msg_R3_Signed.Get(), data.msg_r3_signed_len },
+                                              data.tbsData3Signature);
     SuccessOrExit(err);
 
     // Prepare Sigma3 TBE Data Blob
-    work.msg_r3_encrypted_len =
-        TLV::EstimateStructOverhead(work.nocCert.size(), work.icaCert.size(), work.tbsData3Signature.Length());
+    data.msg_r3_encrypted_len =
+        TLV::EstimateStructOverhead(data.nocCert.size(), data.icaCert.size(), data.tbsData3Signature.Length());
 
-    VerifyOrExit(work.msg_R3_Encrypted.Alloc(work.msg_r3_encrypted_len + CHIP_CRYPTO_AEAD_MIC_LENGTH_BYTES),
+    VerifyOrExit(data.msg_R3_Encrypted.Alloc(data.msg_r3_encrypted_len + CHIP_CRYPTO_AEAD_MIC_LENGTH_BYTES),
                  err = CHIP_ERROR_NO_MEMORY);
 
     {
         TLV::TLVWriter tlvWriter;
         TLV::TLVType outerContainerType = TLV::kTLVType_NotSpecified;
 
-        tlvWriter.Init(work.msg_R3_Encrypted.Get(), work.msg_r3_encrypted_len);
+        tlvWriter.Init(data.msg_R3_Encrypted.Get(), data.msg_r3_encrypted_len);
         SuccessOrExit(err = tlvWriter.StartContainer(TLV::AnonymousTag(), TLV::kTLVType_Structure, outerContainerType));
-        SuccessOrExit(err = tlvWriter.Put(TLV::ContextTag(kTag_TBEData_SenderNOC), work.nocCert));
-        if (!work.icaCert.empty())
+        SuccessOrExit(err = tlvWriter.Put(TLV::ContextTag(kTag_TBEData_SenderNOC), data.nocCert));
+        if (!data.icaCert.empty())
         {
-            SuccessOrExit(err = tlvWriter.Put(TLV::ContextTag(kTag_TBEData_SenderICAC), work.icaCert));
+            SuccessOrExit(err = tlvWriter.Put(TLV::ContextTag(kTag_TBEData_SenderICAC), data.icaCert));
         }
 
         // We are now done with ICAC and NOC certs so we can release the memory.
         {
-            work.icacBuf.Free();
-            work.icaCert = MutableByteSpan{};
+            data.icacBuf.Free();
+            data.icaCert = MutableByteSpan{};
 
-            work.nocBuf.Free();
-            work.nocCert = MutableByteSpan{};
+            data.nocBuf.Free();
+            data.nocCert = MutableByteSpan{};
         }
 
-        SuccessOrExit(err = tlvWriter.PutBytes(TLV::ContextTag(kTag_TBEData_Signature), work.tbsData3Signature.ConstBytes(),
-                                               static_cast<uint32_t>(work.tbsData3Signature.Length())));
+        SuccessOrExit(err = tlvWriter.PutBytes(TLV::ContextTag(kTag_TBEData_Signature), data.tbsData3Signature.ConstBytes(),
+                                               static_cast<uint32_t>(data.tbsData3Signature.Length())));
         SuccessOrExit(err = tlvWriter.EndContainer(outerContainerType));
         SuccessOrExit(err = tlvWriter.Finalize());
-        work.msg_r3_encrypted_len = static_cast<size_t>(tlvWriter.GetLengthWritten());
+        data.msg_r3_encrypted_len = static_cast<size_t>(tlvWriter.GetLengthWritten());
     }
 
 exit:
-    work.status = err;
-
-    auto err2 = DeviceLayer::PlatformMgr().ScheduleWork(
-        [](intptr_t arg) {
-            auto & work2 = *reinterpret_cast<SendSigma3Work *>(arg);
-            work2.session->SendSigma3c(work2);
-        },
-        reinterpret_cast<intptr_t>(&work));
-
-    if (err2 != CHIP_NO_ERROR)
-    {
-        Platform::Delete(&work); // scheduling failed, so delete
-    }
+    return err;
 }
 
-CHIP_ERROR CASESession::SendSigma3c(SendSigma3Work & work)
+CHIP_ERROR CASESession::SendSigma3c(SendSigma3Data & data, CHIP_ERROR status)
 {
     CHIP_ERROR err     = CHIP_NO_ERROR;
     bool ignoreFailure = true;
@@ -1270,15 +1417,17 @@ CHIP_ERROR CASESession::SendSigma3c(SendSigma3Work & work)
 
     AutoReleaseSessionKey sr3k(*mSessionManager->GetSessionKeystore());
 
+    // TODO hopefully can get rid of ignoreFailure and this state check
+    // by ensuring that work is cancelled in cases where state would be unexpected
+
     // Special case: if for whatever reason not in expected state or sequence,
     // don't do anything, including sending a status report or aborting the
     // pending establish.
-    VerifyOrExit(mState == State::kSendSigma3BackgroundPending, err = CHIP_ERROR_INCORRECT_STATE);
-    VerifyOrExit(mSequence == work.sequence, err = CHIP_ERROR_INCORRECT_STATE);
+    VerifyOrExit(mState == State::kSendSigma3Pending, err = CHIP_ERROR_INCORRECT_STATE);
 
     ignoreFailure = false;
 
-    SuccessOrExit(err = work.status);
+    SuccessOrExit(err = status);
 
     // Generate S3K key
     {
@@ -1289,12 +1438,12 @@ CHIP_ERROR CASESession::SendSigma3c(SendSigma3Work & work)
 
     // Generated Encrypted data blob
     SuccessOrExit(err =
-                      AES_CCM_encrypt(work.msg_R3_Encrypted.Get(), work.msg_r3_encrypted_len, nullptr, 0, sr3k.KeyHandle(),
-                                      kTBEData3_Nonce, kTBEDataNonceLength, work.msg_R3_Encrypted.Get(),
-                                      work.msg_R3_Encrypted.Get() + work.msg_r3_encrypted_len, CHIP_CRYPTO_AEAD_MIC_LENGTH_BYTES));
+                      AES_CCM_encrypt(data.msg_R3_Encrypted.Get(), data.msg_r3_encrypted_len, nullptr, 0, sr3k.KeyHandle(),
+                                      kTBEData3_Nonce, kTBEDataNonceLength, data.msg_R3_Encrypted.Get(),
+                                      data.msg_R3_Encrypted.Get() + data.msg_r3_encrypted_len, CHIP_CRYPTO_AEAD_MIC_LENGTH_BYTES));
 
     // Generate Sigma3 Msg
-    data_len = TLV::EstimateStructOverhead(CHIP_CRYPTO_AEAD_MIC_LENGTH_BYTES, work.msg_r3_encrypted_len);
+    data_len = TLV::EstimateStructOverhead(CHIP_CRYPTO_AEAD_MIC_LENGTH_BYTES, data.msg_r3_encrypted_len);
 
     msg_R3 = System::PacketBufferHandle::New(data_len);
     VerifyOrExit(!msg_R3.IsNull(), err = CHIP_ERROR_NO_MEMORY);
@@ -1306,8 +1455,8 @@ CHIP_ERROR CASESession::SendSigma3c(SendSigma3Work & work)
         tlvWriter.Init(std::move(msg_R3));
         err = tlvWriter.StartContainer(TLV::AnonymousTag(), TLV::kTLVType_Structure, outerContainerType);
         SuccessOrExit(err);
-        err = tlvWriter.PutBytes(TLV::ContextTag(1), work.msg_R3_Encrypted.Get(),
-                                 static_cast<uint32_t>(work.msg_r3_encrypted_len + CHIP_CRYPTO_AEAD_MIC_LENGTH_BYTES));
+        err = tlvWriter.PutBytes(TLV::ContextTag(1), data.msg_R3_Encrypted.Get(),
+                                 static_cast<uint32_t>(data.msg_r3_encrypted_len + CHIP_CRYPTO_AEAD_MIC_LENGTH_BYTES));
         SuccessOrExit(err);
         err = tlvWriter.EndContainer(outerContainerType);
         SuccessOrExit(err);
@@ -1333,7 +1482,8 @@ CHIP_ERROR CASESession::SendSigma3c(SendSigma3Work & work)
     mState = State::kSentSigma3;
 
 exit:
-    Platform::Delete(&work);
+    ChipLogError(SecureChannel, "@@@@@ session %p releasing helper %p", this, mSendSigma3Helper.get());
+    mSendSigma3Helper.reset();
 
     if (err != CHIP_NO_ERROR && !ignoreFailure)
     {
@@ -1341,38 +1491,10 @@ CHIP_ERROR CASESession::SendSigma3c(SendSigma3Work & work)
         mState = State::kInitialized;
     }
 
+    ChipLogError(SecureChannel, "@@@@@ done SendSigma3c %p", this);
     return err;
 }
 
-struct CASESession::HandleSigma3Work
-{
-    // Status of background processing.
-    CHIP_ERROR status;
-
-    // Session to use after background processing.
-    CASESession * session;
-
-    // Sequence number used to coordinate foreground/background work for a
-    // particular session establishment.
-    int sequence;
-
-    chip::Platform::ScopedMemoryBuffer<uint8_t> msg_R3_Signed;
-    size_t msg_r3_signed_len;
-
-    ByteSpan initiatorNOC;
-    ByteSpan initiatorICAC;
-
-    uint8_t rootCertBuf[kMaxCHIPCertLength];
-    ByteSpan fabricRCAC;
-
-    P256ECDSASignature tbsData3Signature;
-
-    FabricId fabricId;
-    NodeId initiatorNodeId;
-
-    ValidationContext validContext;
-};
-
 CHIP_ERROR CASESession::HandleSigma3a(System::PacketBufferHandle && msg)
 {
     MATTER_TRACE_EVENT_SCOPE("HandleSigma3", "CASESession");
@@ -1520,7 +1642,7 @@ CHIP_ERROR CASESession::HandleSigma3a(System::PacketBufferHandle && msg)
                           reinterpret_cast<intptr_t>(&work)));
         workPtr = nullptr; // scheduling succeeded, so don't delete
         mExchangeCtxt->WillSendMessage();
-        mState = State::kHandleSigma3BackgroundPending;
+        mState = State::kHandleSigma3Pending;
     }
 
 exit:
@@ -1590,7 +1712,7 @@ CHIP_ERROR CASESession::HandleSigma3c(HandleSigma3Work & work)
     // Special case: if for whatever reason not in expected state or sequence,
     // don't do anything, including sending a status report or aborting the
     // pending establish.
-    VerifyOrExit(mState == State::kHandleSigma3BackgroundPending, err = CHIP_ERROR_INCORRECT_STATE);
+    VerifyOrExit(mState == State::kHandleSigma3Pending, err = CHIP_ERROR_INCORRECT_STATE);
     VerifyOrExit(mSequence == work.sequence, err = CHIP_ERROR_INCORRECT_STATE);
 
     ignoreFailure = false;
diff --git a/src/protocols/secure_channel/CASESession.h b/src/protocols/secure_channel/CASESession.h
index 2c362eb7c4d1a8..ec80c1dbed1a16 100644
--- a/src/protocols/secure_channel/CASESession.h
+++ b/src/protocols/secure_channel/CASESession.h
@@ -25,6 +25,8 @@
 
 #pragma once
 
+#include <memory>
+
 #include <credentials/CHIPCert.h>
 #include <credentials/CertificateValidityPolicy.h>
 #include <credentials/FabricTable.h>
@@ -66,7 +68,7 @@ class DLL_EXPORT CASESession : public Messaging::UnsolicitedMessageHandler,
 
     /**
      * @brief
-     *   Initialize using configured fabrics and wait for session establishment requests.
+     *   Initialize using configured fabrics and wait for session establishment requests (as a responder).
      *
      * @param sessionManager                session manager from which to allocate a secure session object
      * @param fabricTable                   Table of fabrics that are currently configured on the device
@@ -88,7 +90,7 @@ class DLL_EXPORT CASESession : public Messaging::UnsolicitedMessageHandler,
 
     /**
      * @brief
-     *   Create and send session establishment request using device's operational credentials.
+     *   Create and send session establishment request (as an initiator) using device's operational credentials.
      *
      * @param sessionManager                session manager from which to allocate a secure session object
      * @param fabricTable                   The fabric table that contains a fabric in common with the peer
@@ -194,16 +196,16 @@ class DLL_EXPORT CASESession : public Messaging::UnsolicitedMessageHandler,
     friend class TestCASESession;
     enum class State : uint8_t
     {
-        kInitialized                   = 0,
-        kSentSigma1                    = 1,
-        kSentSigma2                    = 2,
-        kSentSigma3                    = 3,
-        kSentSigma1Resume              = 4,
-        kSentSigma2Resume              = 5,
-        kFinished                      = 6,
-        kFinishedViaResume             = 7,
-        kSendSigma3BackgroundPending   = 8,
-        kHandleSigma3BackgroundPending = 9,
+        kInitialized         = 0,
+        kSentSigma1          = 1,
+        kSentSigma2          = 2,
+        kSentSigma3          = 3,
+        kSentSigma1Resume    = 4,
+        kSentSigma2Resume    = 5,
+        kFinished            = 6,
+        kFinishedViaResume   = 7,
+        kSendSigma3Pending   = 8,
+        kHandleSigma3Pending = 9,
     };
 
     /*
@@ -234,10 +236,10 @@ class DLL_EXPORT CASESession : public Messaging::UnsolicitedMessageHandler,
     CHIP_ERROR HandleSigma2(System::PacketBufferHandle && msg);
     CHIP_ERROR HandleSigma2Resume(System::PacketBufferHandle && msg);
 
-    struct SendSigma3Work;
+    struct SendSigma3Data;
     CHIP_ERROR SendSigma3a();
-    static void SendSigma3b(SendSigma3Work & work);
-    CHIP_ERROR SendSigma3c(SendSigma3Work & work);
+    static CHIP_ERROR SendSigma3b(SendSigma3Data & data, bool & cancel);
+    CHIP_ERROR SendSigma3c(SendSigma3Data & data, CHIP_ERROR status);
 
     struct HandleSigma3Work;
     CHIP_ERROR HandleSigma3a(System::PacketBufferHandle && msg);
@@ -306,6 +308,9 @@ class DLL_EXPORT CASESession : public Messaging::UnsolicitedMessageHandler,
     // particular session establishment.
     int mSequence = 0;
 
+    template <class DATA> class WorkHelper;
+    std::shared_ptr<WorkHelper<SendSigma3Data> > mSendSigma3Helper;
+
     State mState;
 
 #if CONFIG_BUILD_FOR_HOST_UNIT_TEST

From 56ca1d5cde27cce8dc30c556e11c79abbac73d2a Mon Sep 17 00:00:00 2001
From: Marc Lepage <mlepage@google.com>
Date: Tue, 28 Mar 2023 15:39:55 -0400
Subject: [PATCH 06/19] Remove temporary logging

---
 src/protocols/secure_channel/CASESession.cpp | 31 --------------------
 1 file changed, 31 deletions(-)

diff --git a/src/protocols/secure_channel/CASESession.cpp b/src/protocols/secure_channel/CASESession.cpp
index b83a836f6332c0..ba5d6caa5240e8 100644
--- a/src/protocols/secure_channel/CASESession.cpp
+++ b/src/protocols/secure_channel/CASESession.cpp
@@ -156,11 +156,6 @@ class CASESession::WorkHelper
     {
     }
 
-    ~WorkHelper()
-    {
-        ChipLogError(SecureChannel, "@@@@@ WorkHelper ~dtor %p", this);
-    }
-
     // Create a work helper using the specified session, work callback, after work callback, and data (template arg).
     // Lifetime is managed by sharing between the caller (typically the session) and the helper itself (while work is scheduled).
     static std::shared_ptr<WorkHelper> Create(CASESession & session, WorkCallback workCallback, AfterWorkCallback afterWorkCallback)
@@ -182,11 +177,9 @@ class CASESession::WorkHelper
             return CHIP_ERROR_INCORRECT_STATE;
         }
         mStrongPtr = mWeakPtr.lock(); // set in `Create`
-        ChipLogError(SecureChannel, "@@@@@ scheduling WorkHandler %p", this);
         auto status = DeviceLayer::PlatformMgr().ScheduleBackgroundWork(WorkHandler, reinterpret_cast<intptr_t>(this));
         if (status != CHIP_NO_ERROR)
         {
-            ChipLogError(SecureChannel, "@@@@@ ScheduleWork releasing helper %p", this);
             mStrongPtr.reset();
         }
         return status;
@@ -203,19 +196,14 @@ class CASESession::WorkHelper
     static void WorkHandler(intptr_t arg)
     {
         WorkHelper* helper = reinterpret_cast<WorkHelper *>(arg);
-        ChipLogError(SecureChannel, "@@@@@ WorkHandler %p", helper);
         bool cancel = false;
         VerifyOrExit(helper->mSession.load(), ;); // cancelled by `CancelWork`?
-        ChipLogError(SecureChannel, "@@@@@ calling mWorkCallback %p", helper);
         helper->mStatus = helper->mWorkCallback(helper->mData, cancel);
-        ChipLogError(SecureChannel, "@@@@@ done mWorkCallback %p", helper);
         VerifyOrExit(!cancel, ;); // canceled by `mWorkCallback`?
         VerifyOrExit(helper->mSession.load(), ;); // cancelled by `CancelWork`?
-        ChipLogError(SecureChannel, "@@@@@ scheduling AfterWorkHandler %p", helper);
         SuccessOrExit(DeviceLayer::PlatformMgr().ScheduleWork(AfterWorkHandler, reinterpret_cast<intptr_t>(helper)));
         return;
     exit:
-        ChipLogError(SecureChannel, "@@@@@ WorkHandler releasing helper %p", helper);
         helper->mStrongPtr.reset();
     }
 
@@ -226,18 +214,10 @@ class CASESession::WorkHelper
         // so shouldn't be deleting the session out from under while this executes
         // (double check that's true, that timers don't run on some other thread, etc.)
         WorkHelper* helper = reinterpret_cast<WorkHelper *>(arg);
-        ChipLogError(SecureChannel, "@@@@@ AfterWorkHandler %p", helper);
         if (auto * session = helper->mSession.load())
         {
-            ChipLogError(SecureChannel, "@@@@@ calling mAfterWorkCallback %p", helper);
             (session->*(helper->mAfterWorkCallback))(helper->mData, helper->mStatus);
-            ChipLogError(SecureChannel, "@@@@@ done mAfterWorkCallback %p", helper);
-        }
-        else
-        {
-            ChipLogError(SecureChannel, "@@@@@ session missing %p", helper);
         }
-        ChipLogError(SecureChannel, "@@@@@ AfterWorkHandler releasing helper %p", helper);
         helper->mStrongPtr.reset();
     }
 
@@ -316,14 +296,12 @@ struct CASESession::HandleSigma3Work
 
 CASESession::~CASESession()
 {
-    ChipLogError(SecureChannel, "@@@@@ CASESession ~dtor %p", this);
     // Let's clear out any security state stored in the object, before destroying it.
     Clear();
 }
 
 void CASESession::OnSessionReleased()
 {
-    ChipLogError(SecureChannel, "@@@@@ OnSessionReleased %p", this);
     // Call into our super-class before we clear our state.
     PairingSession::OnSessionReleased();
     Clear();
@@ -331,7 +309,6 @@ void CASESession::OnSessionReleased()
 
 void CASESession::Clear()
 {
-    ChipLogError(SecureChannel, "@@@@@ Clear %p", this);
     // Cancel any outstanding work.
     if (mSendSigma3Helper)
     {
@@ -440,8 +417,6 @@ CHIP_ERROR CASESession::EstablishSession(SessionManager & sessionManager, Fabric
     MATTER_TRACE_EVENT_SCOPE("EstablishSession", "CASESession");
     CHIP_ERROR err = CHIP_NO_ERROR;
 
-    ChipLogError(SecureChannel, "@@@@@ EstablishSession %p", this);
-
     // Return early on error here, as we have not initialized any state yet
     ReturnErrorCodeIf(exchangeCtxt == nullptr, CHIP_ERROR_INVALID_ARGUMENT);
     ReturnErrorCodeIf(fabricTable == nullptr, CHIP_ERROR_INVALID_ARGUMENT);
@@ -490,7 +465,6 @@ CHIP_ERROR CASESession::EstablishSession(SessionManager & sessionManager, Fabric
     {
         Clear();
     }
-    ChipLogError(SecureChannel, "@@@@@ done EstablishSession %p", this);
     return err;
 }
 
@@ -508,7 +482,6 @@ void CASESession::OnResponseTimeout(ExchangeContext * ec)
 
 void CASESession::AbortPendingEstablish(CHIP_ERROR err)
 {
-    ChipLogError(SecureChannel, "@@@@@ AbortPendingEstablish %p", this);
     Clear();
     // Do this last in case the delegate frees us.
     NotifySessionEstablishmentError(err);
@@ -1314,8 +1287,6 @@ CHIP_ERROR CASESession::SendSigma3a()
     {
         auto & data = helper->mData;
 
-        ChipLogError(SecureChannel, "@@@@@ session %p helper %p", this, helper.get());
-
         VerifyOrExit(mFabricsTable != nullptr, err = CHIP_ERROR_INCORRECT_STATE);
         data.fabricTable = mFabricsTable;
         data.fabricIndex = mFabricIndex;
@@ -1482,7 +1453,6 @@ CHIP_ERROR CASESession::SendSigma3c(SendSigma3Data & data, CHIP_ERROR status)
     mState = State::kSentSigma3;
 
 exit:
-    ChipLogError(SecureChannel, "@@@@@ session %p releasing helper %p", this, mSendSigma3Helper.get());
     mSendSigma3Helper.reset();
 
     if (err != CHIP_NO_ERROR && !ignoreFailure)
@@ -1491,7 +1461,6 @@ CHIP_ERROR CASESession::SendSigma3c(SendSigma3Data & data, CHIP_ERROR status)
         mState = State::kInitialized;
     }
 
-    ChipLogError(SecureChannel, "@@@@@ done SendSigma3c %p", this);
     return err;
 }
 

From 85bba1e5ae45681cb3540717620ca6178c4bd09d Mon Sep 17 00:00:00 2001
From: Marc Lepage <mlepage@google.com>
Date: Tue, 28 Mar 2023 15:43:32 -0400
Subject: [PATCH 07/19] Restyle

---
 src/protocols/secure_channel/CASESession.cpp | 26 +++++++++-----------
 src/protocols/secure_channel/CASESession.h   |  5 ++--
 2 files changed, 15 insertions(+), 16 deletions(-)

diff --git a/src/protocols/secure_channel/CASESession.cpp b/src/protocols/secure_channel/CASESession.cpp
index ba5d6caa5240e8..a8bc7ab119229b 100644
--- a/src/protocols/secure_channel/CASESession.cpp
+++ b/src/protocols/secure_channel/CASESession.cpp
@@ -152,15 +152,16 @@ class CASESession::WorkHelper
 
     // Create a work helper using the specified session, work callback, after work callback, and data (template arg).
     // Lifetime is not managed, see `Create` for that option.
-    WorkHelper(CASESession & session, WorkCallback workCallback, AfterWorkCallback afterWorkCallback) : mSession(&session), mWorkCallback(workCallback), mAfterWorkCallback(afterWorkCallback)
-    {
-    }
+    WorkHelper(CASESession & session, WorkCallback workCallback, AfterWorkCallback afterWorkCallback) :
+        mSession(&session), mWorkCallback(workCallback), mAfterWorkCallback(afterWorkCallback)
+    {}
 
     // Create a work helper using the specified session, work callback, after work callback, and data (template arg).
     // Lifetime is managed by sharing between the caller (typically the session) and the helper itself (while work is scheduled).
     static std::shared_ptr<WorkHelper> Create(CASESession & session, WorkCallback workCallback, AfterWorkCallback afterWorkCallback)
     {
-        std::shared_ptr<WorkHelper> ptr(Platform::New<WorkHelper>(session, workCallback, afterWorkCallback), Platform::Delete<WorkHelper>);
+        std::shared_ptr<WorkHelper> ptr(Platform::New<WorkHelper>(session, workCallback, afterWorkCallback),
+                                        Platform::Delete<WorkHelper>);
         if (ptr)
         {
             ptr->mWeakPtr = ptr; // used by `ScheduleWork`
@@ -176,7 +177,7 @@ class CASESession::WorkHelper
         {
             return CHIP_ERROR_INCORRECT_STATE;
         }
-        mStrongPtr = mWeakPtr.lock(); // set in `Create`
+        mStrongPtr  = mWeakPtr.lock(); // set in `Create`
         auto status = DeviceLayer::PlatformMgr().ScheduleBackgroundWork(WorkHandler, reinterpret_cast<intptr_t>(this));
         if (status != CHIP_NO_ERROR)
         {
@@ -184,22 +185,19 @@ class CASESession::WorkHelper
         }
         return status;
     }
- 
+
     // Cancel the work, by clearing the associated session.
-    void CancelWork()
-    {
-        mSession.store(nullptr);
-    }
+    void CancelWork() { mSession.store(nullptr); }
 
 private:
     // Handler for the work callback.
     static void WorkHandler(intptr_t arg)
     {
-        WorkHelper* helper = reinterpret_cast<WorkHelper *>(arg);
-        bool cancel = false;
+        WorkHelper * helper = reinterpret_cast<WorkHelper *>(arg);
+        bool cancel         = false;
         VerifyOrExit(helper->mSession.load(), ;); // cancelled by `CancelWork`?
         helper->mStatus = helper->mWorkCallback(helper->mData, cancel);
-        VerifyOrExit(!cancel, ;); // canceled by `mWorkCallback`?
+        VerifyOrExit(!cancel, ;);                 // canceled by `mWorkCallback`?
         VerifyOrExit(helper->mSession.load(), ;); // cancelled by `CancelWork`?
         SuccessOrExit(DeviceLayer::PlatformMgr().ScheduleWork(AfterWorkHandler, reinterpret_cast<intptr_t>(helper)));
         return;
@@ -213,7 +211,7 @@ class CASESession::WorkHelper
         // TODO some notes here about how we should be on the main event thread
         // so shouldn't be deleting the session out from under while this executes
         // (double check that's true, that timers don't run on some other thread, etc.)
-        WorkHelper* helper = reinterpret_cast<WorkHelper *>(arg);
+        WorkHelper * helper = reinterpret_cast<WorkHelper *>(arg);
         if (auto * session = helper->mSession.load())
         {
             (session->*(helper->mAfterWorkCallback))(helper->mData, helper->mStatus);
diff --git a/src/protocols/secure_channel/CASESession.h b/src/protocols/secure_channel/CASESession.h
index ec80c1dbed1a16..95d98375e83072 100644
--- a/src/protocols/secure_channel/CASESession.h
+++ b/src/protocols/secure_channel/CASESession.h
@@ -308,8 +308,9 @@ class DLL_EXPORT CASESession : public Messaging::UnsolicitedMessageHandler,
     // particular session establishment.
     int mSequence = 0;
 
-    template <class DATA> class WorkHelper;
-    std::shared_ptr<WorkHelper<SendSigma3Data> > mSendSigma3Helper;
+    template <class DATA>
+    class WorkHelper;
+    std::shared_ptr<WorkHelper<SendSigma3Data>> mSendSigma3Helper;
 
     State mState;
 

From 6410b30488678ce243cc23b291950a1bf067d0da Mon Sep 17 00:00:00 2001
From: Marc Lepage <mlepage@google.com>
Date: Tue, 28 Mar 2023 16:05:47 -0400
Subject: [PATCH 08/19] Minor cleanup

---
 src/protocols/secure_channel/CASESession.cpp | 17 +++--------------
 1 file changed, 3 insertions(+), 14 deletions(-)

diff --git a/src/protocols/secure_channel/CASESession.cpp b/src/protocols/secure_channel/CASESession.cpp
index a8bc7ab119229b..74e83d2caa437e 100644
--- a/src/protocols/secure_channel/CASESession.cpp
+++ b/src/protocols/secure_channel/CASESession.cpp
@@ -208,9 +208,7 @@ class CASESession::WorkHelper
     // Handler for the after work callback.
     static void AfterWorkHandler(intptr_t arg)
     {
-        // TODO some notes here about how we should be on the main event thread
-        // so shouldn't be deleting the session out from under while this executes
-        // (double check that's true, that timers don't run on some other thread, etc.)
+        // Since this runs in the main CHIP thread, the session shouldn't be otherwise used (messages, timers, etc.)
         WorkHelper * helper = reinterpret_cast<WorkHelper *>(arg);
         if (auto * session = helper->mSession.load())
         {
@@ -1377,7 +1375,6 @@ CHIP_ERROR CASESession::SendSigma3b(SendSigma3Data & data, bool & cancel)
 CHIP_ERROR CASESession::SendSigma3c(SendSigma3Data & data, CHIP_ERROR status)
 {
     CHIP_ERROR err     = CHIP_NO_ERROR;
-    bool ignoreFailure = true;
 
     System::PacketBufferHandle msg_R3;
     size_t data_len;
@@ -1386,15 +1383,7 @@ CHIP_ERROR CASESession::SendSigma3c(SendSigma3Data & data, CHIP_ERROR status)
 
     AutoReleaseSessionKey sr3k(*mSessionManager->GetSessionKeystore());
 
-    // TODO hopefully can get rid of ignoreFailure and this state check
-    // by ensuring that work is cancelled in cases where state would be unexpected
-
-    // Special case: if for whatever reason not in expected state or sequence,
-    // don't do anything, including sending a status report or aborting the
-    // pending establish.
-    VerifyOrExit(mState == State::kSendSigma3Pending, err = CHIP_ERROR_INCORRECT_STATE);
-
-    ignoreFailure = false;
+    VerifyOrDieWithMsg(mState == State::kSendSigma3Pending, SecureChannel, "Bad internal state.");
 
     SuccessOrExit(err = status);
 
@@ -1453,7 +1442,7 @@ CHIP_ERROR CASESession::SendSigma3c(SendSigma3Data & data, CHIP_ERROR status)
 exit:
     mSendSigma3Helper.reset();
 
-    if (err != CHIP_NO_ERROR && !ignoreFailure)
+    if (err != CHIP_NO_ERROR)
     {
         SendStatusReport(mExchangeCtxt, kProtocolCodeInvalidParam);
         mState = State::kInitialized;

From 239fce10553ba0012b2187b929bf02e9aa7a2895 Mon Sep 17 00:00:00 2001
From: Marc Lepage <mlepage@google.com>
Date: Tue, 28 Mar 2023 16:08:24 -0400
Subject: [PATCH 09/19] Minor cleanup

---
 src/protocols/secure_channel/CASESession.cpp | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/src/protocols/secure_channel/CASESession.cpp b/src/protocols/secure_channel/CASESession.cpp
index 74e83d2caa437e..a8f5f555fe30a3 100644
--- a/src/protocols/secure_channel/CASESession.cpp
+++ b/src/protocols/secure_channel/CASESession.cpp
@@ -417,10 +417,6 @@ CHIP_ERROR CASESession::EstablishSession(SessionManager & sessionManager, Fabric
     ReturnErrorCodeIf(exchangeCtxt == nullptr, CHIP_ERROR_INVALID_ARGUMENT);
     ReturnErrorCodeIf(fabricTable == nullptr, CHIP_ERROR_INVALID_ARGUMENT);
 
-    // Sequence number used to coordinate foreground/background work for a
-    // particular session establishment.
-    mSequence++;
-
     // Use FabricTable directly to avoid situation of dangling index from stale FabricInfo
     // until we factor-out any FabricInfo direct usage.
     ReturnErrorCodeIf(peerScopedNodeId.GetFabricIndex() == kUndefinedFabricIndex, CHIP_ERROR_INVALID_ARGUMENT);

From 5d0f7165aee0b0849b8144a2da9efab234dae447 Mon Sep 17 00:00:00 2001
From: Marc Lepage <mlepage@google.com>
Date: Wed, 29 Mar 2023 11:27:00 -0400
Subject: [PATCH 10/19] Restyle

---
 src/protocols/secure_channel/CASESession.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/protocols/secure_channel/CASESession.cpp b/src/protocols/secure_channel/CASESession.cpp
index a8f5f555fe30a3..b3303ee76f73be 100644
--- a/src/protocols/secure_channel/CASESession.cpp
+++ b/src/protocols/secure_channel/CASESession.cpp
@@ -1370,7 +1370,7 @@ CHIP_ERROR CASESession::SendSigma3b(SendSigma3Data & data, bool & cancel)
 
 CHIP_ERROR CASESession::SendSigma3c(SendSigma3Data & data, CHIP_ERROR status)
 {
-    CHIP_ERROR err     = CHIP_NO_ERROR;
+    CHIP_ERROR err = CHIP_NO_ERROR;
 
     System::PacketBufferHandle msg_R3;
     size_t data_len;

From c9d12f0ab965af4878e357b368452b0746b17c8d Mon Sep 17 00:00:00 2001
From: Marc Lepage <mlepage@google.com>
Date: Tue, 4 Apr 2023 11:08:00 -0400
Subject: [PATCH 11/19] Use Platform::SharedPtr

Also add alias template for Platform::WeakPtr.
---
 src/lib/support/CHIPMem.h                    |  3 +++
 src/protocols/secure_channel/CASESession.cpp | 10 +++++-----
 src/protocols/secure_channel/CASESession.h   |  5 ++---
 3 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/src/lib/support/CHIPMem.h b/src/lib/support/CHIPMem.h
index a1e27c20b73ce3..b4b78aca647e3a 100644
--- a/src/lib/support/CHIPMem.h
+++ b/src/lib/support/CHIPMem.h
@@ -193,6 +193,9 @@ inline SharedPtr<T> MakeShared(Args &&... args)
     return SharedPtr<T>(New<T>(std::forward<Args>(args)...), Deleter<T>());
 }
 
+template <typename T>
+using WeakPtr = std::weak_ptr<T>;
+
 // See MemoryDebugCheckPointer().
 extern bool MemoryInternalCheckPointer(const void * p, size_t min_size);
 
diff --git a/src/protocols/secure_channel/CASESession.cpp b/src/protocols/secure_channel/CASESession.cpp
index b3303ee76f73be..008b7c372d264b 100644
--- a/src/protocols/secure_channel/CASESession.cpp
+++ b/src/protocols/secure_channel/CASESession.cpp
@@ -158,10 +158,10 @@ class CASESession::WorkHelper
 
     // Create a work helper using the specified session, work callback, after work callback, and data (template arg).
     // Lifetime is managed by sharing between the caller (typically the session) and the helper itself (while work is scheduled).
-    static std::shared_ptr<WorkHelper> Create(CASESession & session, WorkCallback workCallback, AfterWorkCallback afterWorkCallback)
+    static Platform::SharedPtr<WorkHelper> Create(CASESession & session, WorkCallback workCallback,
+                                                  AfterWorkCallback afterWorkCallback)
     {
-        std::shared_ptr<WorkHelper> ptr(Platform::New<WorkHelper>(session, workCallback, afterWorkCallback),
-                                        Platform::Delete<WorkHelper>);
+        auto ptr = Platform::MakeShared<WorkHelper>(session, workCallback, afterWorkCallback);
         if (ptr)
         {
             ptr->mWeakPtr = ptr; // used by `ScheduleWork`
@@ -219,10 +219,10 @@ class CASESession::WorkHelper
 
 private:
     // Lifetime management: `ScheduleWork` sets `mStrongPtr` from `mWeakPtr`.
-    std::weak_ptr<WorkHelper> mWeakPtr;
+    Platform::WeakPtr<WorkHelper> mWeakPtr;
 
     // Lifetime management: `ScheduleWork` sets `mStrongPtr` from `mWeakPtr`.
-    std::shared_ptr<WorkHelper> mStrongPtr;
+    Platform::SharedPtr<WorkHelper> mStrongPtr;
 
     // Associated session, cleared by `CancelWork`.
     std::atomic<CASESession *> mSession;
diff --git a/src/protocols/secure_channel/CASESession.h b/src/protocols/secure_channel/CASESession.h
index 95d98375e83072..21578e592ba907 100644
--- a/src/protocols/secure_channel/CASESession.h
+++ b/src/protocols/secure_channel/CASESession.h
@@ -25,8 +25,6 @@
 
 #pragma once
 
-#include <memory>
-
 #include <credentials/CHIPCert.h>
 #include <credentials/CertificateValidityPolicy.h>
 #include <credentials/FabricTable.h>
@@ -35,6 +33,7 @@
 #include <lib/core/ScopedNodeId.h>
 #include <lib/core/TLV.h>
 #include <lib/support/Base64.h>
+#include <lib/support/CHIPMem.h>
 #include <messaging/ExchangeContext.h>
 #include <messaging/ExchangeDelegate.h>
 #include <messaging/ReliableMessageProtocolConfig.h>
@@ -310,7 +309,7 @@ class DLL_EXPORT CASESession : public Messaging::UnsolicitedMessageHandler,
 
     template <class DATA>
     class WorkHelper;
-    std::shared_ptr<WorkHelper<SendSigma3Data>> mSendSigma3Helper;
+    Platform::SharedPtr<WorkHelper<SendSigma3Data>> mSendSigma3Helper;
 
     State mState;
 

From 54586ca72f59a1c5420c2045bc6a513027da6261 Mon Sep 17 00:00:00 2001
From: Marc Lepage <mlepage@google.com>
Date: Thu, 6 Apr 2023 14:47:27 -0400
Subject: [PATCH 12/19] Add mutex to FabricTable

This supports locking during SignWithOpKeypair, and other operations
that may alter state in the fabric table, while CASESession is
performing work in the background during session establishment.

CASESession registers as a fabric table listener, and when a fabric
is removed (e.g. timeout) it attempts to cancel any outstanding work,
and also clears out the fabric index in the work helper data.

Therefore, if outstanding work has made it into SignWithOpKeypair,
it should be OK until complete.

It still relies on other tasks not altering FabricInfo, or the
configured OperationalKeystore, but that would have had to have been
true before anyways.

The mutex was not made recursive. It's omitted from a few functions,
which should be OK for now, and there should be cleanup on a subsequent
commit (and probably fix up const-ness of member functions, and
factoring of API vs. impl functions). This commit is to flush out
build/test errors on all CI platforms, and to discuss/review/comment on
the general approach.
---
 src/credentials/FabricTable.cpp              | 38 ++++++++++++++++++++
 src/credentials/FabricTable.h                |  3 ++
 src/protocols/secure_channel/CASESession.cpp |  3 +-
 3 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/src/credentials/FabricTable.cpp b/src/credentials/FabricTable.cpp
index 63f7afebd7d0ba..5e2f117df7517b 100644
--- a/src/credentials/FabricTable.cpp
+++ b/src/credentials/FabricTable.cpp
@@ -910,6 +910,8 @@ FabricTable::AddOrUpdateInner(FabricIndex fabricIndex, bool isAddition, Crypto::
 
 CHIP_ERROR FabricTable::Delete(FabricIndex fabricIndex)
 {
+    std::unique_lock<std::mutex> lock(mMutex);
+
     VerifyOrReturnError(mStorage != nullptr, CHIP_ERROR_INVALID_ARGUMENT);
     VerifyOrReturnError(IsValidFabricIndex(fabricIndex), CHIP_ERROR_INVALID_ARGUMENT);
 
@@ -1019,6 +1021,8 @@ CHIP_ERROR FabricTable::Delete(FabricIndex fabricIndex)
 
 void FabricTable::DeleteAllFabrics()
 {
+    // Rely on mutex in impl methods
+
     static_assert(kMaxValidFabricIndex <= UINT8_MAX, "Cannot create more fabrics than UINT8_MAX");
 
     RevertPendingFabricData();
@@ -1031,6 +1035,8 @@ void FabricTable::DeleteAllFabrics()
 
 CHIP_ERROR FabricTable::Init(const FabricTable::InitParams & initParams)
 {
+    // Mutex not necessary while initializing
+
     VerifyOrReturnError(initParams.storage != nullptr, CHIP_ERROR_INVALID_ARGUMENT);
     VerifyOrReturnError(initParams.opCertStore != nullptr, CHIP_ERROR_INVALID_ARGUMENT);
 
@@ -1105,6 +1111,8 @@ CHIP_ERROR FabricTable::Init(const FabricTable::InitParams & initParams)
 
 void FabricTable::Forget(FabricIndex fabricIndex)
 {
+    std::unique_lock<std::mutex> lock(mMutex);
+
     ChipLogProgress(FabricProvisioning, "Forgetting fabric 0x%x", static_cast<unsigned>(fabricIndex));
 
     auto * fabricInfo = GetMutableFabricByIndex(fabricIndex);
@@ -1116,6 +1124,8 @@ void FabricTable::Forget(FabricIndex fabricIndex)
 
 void FabricTable::Shutdown()
 {
+    std::unique_lock<std::mutex> lock(mMutex);
+
     VerifyOrReturn(mStorage != nullptr);
     ChipLogProgress(FabricProvisioning, "Shutting down FabricTable");
 
@@ -1141,6 +1151,8 @@ void FabricTable::Shutdown()
 
 FabricIndex FabricTable::GetDeletedFabricFromCommitMarker()
 {
+    std::unique_lock<std::mutex> lock(mMutex);
+
     FabricIndex retVal = mDeletedFabricIndexFromInit;
 
     // Reset for next read
@@ -1151,6 +1163,8 @@ FabricIndex FabricTable::GetDeletedFabricFromCommitMarker()
 
 CHIP_ERROR FabricTable::AddFabricDelegate(FabricTable::Delegate * delegate)
 {
+    // Mutex not necessary while adding/removing delegates
+
     VerifyOrReturnError(delegate != nullptr, CHIP_ERROR_INVALID_ARGUMENT);
     for (FabricTable::Delegate * iter = mDelegateListRoot; iter != nullptr; iter = iter->next)
     {
@@ -1166,6 +1180,8 @@ CHIP_ERROR FabricTable::AddFabricDelegate(FabricTable::Delegate * delegate)
 
 void FabricTable::RemoveFabricDelegate(FabricTable::Delegate * delegateToRemove)
 {
+    // Mutex not necessary while adding/removing delegates
+
     VerifyOrReturn(delegateToRemove != nullptr);
 
     if (delegateToRemove == mDelegateListRoot)
@@ -1197,6 +1213,8 @@ void FabricTable::RemoveFabricDelegate(FabricTable::Delegate * delegateToRemove)
 
 CHIP_ERROR FabricTable::SetLastKnownGoodChipEpochTime(System::Clock::Seconds32 lastKnownGoodChipEpochTime)
 {
+    std::unique_lock<std::mutex> lock(mMutex);
+
     CHIP_ERROR err = CHIP_NO_ERROR;
     // Find our latest NotBefore time for any installed certificate.
     System::Clock::Seconds32 latestNotBefore = System::Clock::Seconds32(0);
@@ -1390,6 +1408,8 @@ CHIP_ERROR FabricTable::ReadFabricInfo(TLV::ContiguousBufferTLVReader & reader)
 
 Crypto::P256Keypair * FabricTable::AllocateEphemeralKeypairForCASE()
 {
+    std::unique_lock<std::mutex> lock(mMutex);
+
     if (mOperationalKeystore != nullptr)
     {
         return mOperationalKeystore->AllocateEphemeralKeypairForCASE();
@@ -1400,6 +1420,8 @@ Crypto::P256Keypair * FabricTable::AllocateEphemeralKeypairForCASE()
 
 void FabricTable::ReleaseEphemeralKeypair(Crypto::P256Keypair * keypair)
 {
+    std::unique_lock<std::mutex> lock(mMutex);
+
     if (mOperationalKeystore != nullptr)
     {
         mOperationalKeystore->ReleaseEphemeralKeypair(keypair);
@@ -1481,6 +1503,8 @@ bool FabricTable::HasOperationalKeyForFabric(FabricIndex fabricIndex) const
 
 CHIP_ERROR FabricTable::SignWithOpKeypair(FabricIndex fabricIndex, ByteSpan message, P256ECDSASignature & outSignature) const
 {
+    std::unique_lock<std::mutex> lock(mMutex);
+
     const FabricInfo * fabricInfo = FindFabricWithIndex(fabricIndex);
     VerifyOrReturnError(fabricInfo != nullptr, CHIP_ERROR_KEY_NOT_FOUND);
 
@@ -1525,6 +1549,8 @@ bool FabricTable::SetPendingDataFabricIndex(FabricIndex fabricIndex)
 
 CHIP_ERROR FabricTable::AllocatePendingOperationalKey(Optional<FabricIndex> fabricIndex, MutableByteSpan & outputCsr)
 {
+    std::unique_lock<std::mutex> lock(mMutex);
+
     // We can only manage commissionable pending fail-safe state if we have a keystore
     VerifyOrReturnError(mOperationalKeystore != nullptr, CHIP_ERROR_INCORRECT_STATE);
 
@@ -1567,6 +1593,8 @@ CHIP_ERROR FabricTable::AllocatePendingOperationalKey(Optional<FabricIndex> fabr
 
 CHIP_ERROR FabricTable::AddNewPendingTrustedRootCert(const ByteSpan & rcac)
 {
+    std::unique_lock<std::mutex> lock(mMutex);
+
     VerifyOrReturnError(mOpCertStore != nullptr, CHIP_ERROR_INCORRECT_STATE);
 
     // We should not already have pending NOC chain elements when we get here
@@ -1644,6 +1672,8 @@ CHIP_ERROR FabricTable::AddNewPendingFabricCommon(const ByteSpan & noc, const By
                                                   Crypto::P256Keypair * existingOpKey, bool isExistingOpKeyExternallyOwned,
                                                   FabricIndex * outNewFabricIndex)
 {
+    std::unique_lock<std::mutex> lock(mMutex);
+
     VerifyOrReturnError(mOpCertStore != nullptr, CHIP_ERROR_INCORRECT_STATE);
     VerifyOrReturnError(outNewFabricIndex != nullptr, CHIP_ERROR_INVALID_ARGUMENT);
     static_assert(kMaxValidFabricIndex <= UINT8_MAX, "Cannot create more fabrics than UINT8_MAX");
@@ -1714,6 +1744,8 @@ CHIP_ERROR FabricTable::AddNewPendingFabricCommon(const ByteSpan & noc, const By
 CHIP_ERROR FabricTable::UpdatePendingFabricCommon(FabricIndex fabricIndex, const ByteSpan & noc, const ByteSpan & icac,
                                                   Crypto::P256Keypair * existingOpKey, bool isExistingOpKeyExternallyOwned)
 {
+    std::unique_lock<std::mutex> lock(mMutex);
+
     VerifyOrReturnError(mOpCertStore != nullptr, CHIP_ERROR_INCORRECT_STATE);
     VerifyOrReturnError(IsValidFabricIndex(fabricIndex), CHIP_ERROR_INVALID_ARGUMENT);
 
@@ -1772,6 +1804,8 @@ CHIP_ERROR FabricTable::UpdatePendingFabricCommon(FabricIndex fabricIndex, const
 
 CHIP_ERROR FabricTable::CommitPendingFabricData()
 {
+    std::unique_lock<std::mutex> lock(mMutex);
+
     VerifyOrReturnError((mStorage != nullptr) && (mOpCertStore != nullptr), CHIP_ERROR_INCORRECT_STATE);
 
     bool haveNewTrustedRoot      = mStateFlags.Has(StateFlags::kIsTrustedRootPending);
@@ -2050,6 +2084,8 @@ void FabricTable::RevertPendingOpCertsExceptRoot()
 
 CHIP_ERROR FabricTable::SetFabricLabel(FabricIndex fabricIndex, const CharSpan & fabricLabel)
 {
+    std::unique_lock<std::mutex> lock(mMutex);
+
     VerifyOrReturnError(mStorage != nullptr, CHIP_ERROR_INCORRECT_STATE);
     VerifyOrReturnError(IsValidFabricIndex(fabricIndex), CHIP_ERROR_INVALID_FABRIC_INDEX);
 
@@ -2073,6 +2109,8 @@ CHIP_ERROR FabricTable::SetFabricLabel(FabricIndex fabricIndex, const CharSpan &
 
 CHIP_ERROR FabricTable::GetFabricLabel(FabricIndex fabricIndex, CharSpan & outFabricLabel)
 {
+    std::unique_lock<std::mutex> lock(mMutex);
+
     const FabricInfo * fabricInfo = FindFabricWithIndex(fabricIndex);
     VerifyOrReturnError(fabricInfo != nullptr, CHIP_ERROR_INVALID_FABRIC_INDEX);
 
diff --git a/src/credentials/FabricTable.h b/src/credentials/FabricTable.h
index 42f88024b877c6..2a19f5bd97d60f 100644
--- a/src/credentials/FabricTable.h
+++ b/src/credentials/FabricTable.h
@@ -22,6 +22,7 @@
 #pragma once
 
 #include <algorithm>
+#include <mutex>
 
 #include <app/util/basic-types.h>
 #include <credentials/CHIPCert.h>
@@ -1146,6 +1147,8 @@ class DLL_EXPORT FabricTable
     uint8_t mFabricCount = 0;
 
     BitFlags<StateFlags> mStateFlags;
+
+    mutable std::mutex mMutex;
 };
 
 } // namespace chip
diff --git a/src/protocols/secure_channel/CASESession.cpp b/src/protocols/secure_channel/CASESession.cpp
index 008b7c372d264b..36a372069405bb 100644
--- a/src/protocols/secure_channel/CASESession.cpp
+++ b/src/protocols/secure_channel/CASESession.cpp
@@ -244,7 +244,7 @@ class CASESession::WorkHelper
 struct CASESession::SendSigma3Data
 {
     FabricTable * fabricTable;
-    FabricIndex fabricIndex;
+    std::atomic<FabricIndex> fabricIndex;
 
     chip::Platform::ScopedMemoryBuffer<uint8_t> msg_R3_Signed;
     size_t msg_r3_signed_len;
@@ -308,6 +308,7 @@ void CASESession::Clear()
     // Cancel any outstanding work.
     if (mSendSigma3Helper)
     {
+        mSendSigma3Helper->mData.fabricIndex = kUndefinedFabricIndex;
         mSendSigma3Helper->CancelWork();
         mSendSigma3Helper.reset();
     }

From 43e4e63c5427a12cd74d24661a159ecd7186167b Mon Sep 17 00:00:00 2001
From: Marc Lepage <mlepage@google.com>
Date: Wed, 12 Apr 2023 15:18:57 -0400
Subject: [PATCH 13/19] Remove mutex, only async sometimes

It's tricky to async background the signing operation, because of the
two ways operational signing occurs.

Legacy way: opkeypair manually added for fabric info
Recommended way: opkeystore handles everything

Removed std::mutex because it wasn't supported by all platforms.

Instead, made background signing occur only if using the operational
keystore (recommended way), since implementors can perform any needed
mutual exclusion in the operational keystore. If using manually added
operational keypairs (legacy way), keep signing in the foreground,
since it's not feasible to mutex the entire fabric table and typically
the operations is simpler anyways.
---
 src/credentials/FabricTable.cpp              | 38 ------------
 src/credentials/FabricTable.h                | 14 ++++-
 src/protocols/secure_channel/CASESession.cpp | 64 +++++++++++++++++---
 3 files changed, 68 insertions(+), 48 deletions(-)

diff --git a/src/credentials/FabricTable.cpp b/src/credentials/FabricTable.cpp
index 5e2f117df7517b..63f7afebd7d0ba 100644
--- a/src/credentials/FabricTable.cpp
+++ b/src/credentials/FabricTable.cpp
@@ -910,8 +910,6 @@ FabricTable::AddOrUpdateInner(FabricIndex fabricIndex, bool isAddition, Crypto::
 
 CHIP_ERROR FabricTable::Delete(FabricIndex fabricIndex)
 {
-    std::unique_lock<std::mutex> lock(mMutex);
-
     VerifyOrReturnError(mStorage != nullptr, CHIP_ERROR_INVALID_ARGUMENT);
     VerifyOrReturnError(IsValidFabricIndex(fabricIndex), CHIP_ERROR_INVALID_ARGUMENT);
 
@@ -1021,8 +1019,6 @@ CHIP_ERROR FabricTable::Delete(FabricIndex fabricIndex)
 
 void FabricTable::DeleteAllFabrics()
 {
-    // Rely on mutex in impl methods
-
     static_assert(kMaxValidFabricIndex <= UINT8_MAX, "Cannot create more fabrics than UINT8_MAX");
 
     RevertPendingFabricData();
@@ -1035,8 +1031,6 @@ void FabricTable::DeleteAllFabrics()
 
 CHIP_ERROR FabricTable::Init(const FabricTable::InitParams & initParams)
 {
-    // Mutex not necessary while initializing
-
     VerifyOrReturnError(initParams.storage != nullptr, CHIP_ERROR_INVALID_ARGUMENT);
     VerifyOrReturnError(initParams.opCertStore != nullptr, CHIP_ERROR_INVALID_ARGUMENT);
 
@@ -1111,8 +1105,6 @@ CHIP_ERROR FabricTable::Init(const FabricTable::InitParams & initParams)
 
 void FabricTable::Forget(FabricIndex fabricIndex)
 {
-    std::unique_lock<std::mutex> lock(mMutex);
-
     ChipLogProgress(FabricProvisioning, "Forgetting fabric 0x%x", static_cast<unsigned>(fabricIndex));
 
     auto * fabricInfo = GetMutableFabricByIndex(fabricIndex);
@@ -1124,8 +1116,6 @@ void FabricTable::Forget(FabricIndex fabricIndex)
 
 void FabricTable::Shutdown()
 {
-    std::unique_lock<std::mutex> lock(mMutex);
-
     VerifyOrReturn(mStorage != nullptr);
     ChipLogProgress(FabricProvisioning, "Shutting down FabricTable");
 
@@ -1151,8 +1141,6 @@ void FabricTable::Shutdown()
 
 FabricIndex FabricTable::GetDeletedFabricFromCommitMarker()
 {
-    std::unique_lock<std::mutex> lock(mMutex);
-
     FabricIndex retVal = mDeletedFabricIndexFromInit;
 
     // Reset for next read
@@ -1163,8 +1151,6 @@ FabricIndex FabricTable::GetDeletedFabricFromCommitMarker()
 
 CHIP_ERROR FabricTable::AddFabricDelegate(FabricTable::Delegate * delegate)
 {
-    // Mutex not necessary while adding/removing delegates
-
     VerifyOrReturnError(delegate != nullptr, CHIP_ERROR_INVALID_ARGUMENT);
     for (FabricTable::Delegate * iter = mDelegateListRoot; iter != nullptr; iter = iter->next)
     {
@@ -1180,8 +1166,6 @@ CHIP_ERROR FabricTable::AddFabricDelegate(FabricTable::Delegate * delegate)
 
 void FabricTable::RemoveFabricDelegate(FabricTable::Delegate * delegateToRemove)
 {
-    // Mutex not necessary while adding/removing delegates
-
     VerifyOrReturn(delegateToRemove != nullptr);
 
     if (delegateToRemove == mDelegateListRoot)
@@ -1213,8 +1197,6 @@ void FabricTable::RemoveFabricDelegate(FabricTable::Delegate * delegateToRemove)
 
 CHIP_ERROR FabricTable::SetLastKnownGoodChipEpochTime(System::Clock::Seconds32 lastKnownGoodChipEpochTime)
 {
-    std::unique_lock<std::mutex> lock(mMutex);
-
     CHIP_ERROR err = CHIP_NO_ERROR;
     // Find our latest NotBefore time for any installed certificate.
     System::Clock::Seconds32 latestNotBefore = System::Clock::Seconds32(0);
@@ -1408,8 +1390,6 @@ CHIP_ERROR FabricTable::ReadFabricInfo(TLV::ContiguousBufferTLVReader & reader)
 
 Crypto::P256Keypair * FabricTable::AllocateEphemeralKeypairForCASE()
 {
-    std::unique_lock<std::mutex> lock(mMutex);
-
     if (mOperationalKeystore != nullptr)
     {
         return mOperationalKeystore->AllocateEphemeralKeypairForCASE();
@@ -1420,8 +1400,6 @@ Crypto::P256Keypair * FabricTable::AllocateEphemeralKeypairForCASE()
 
 void FabricTable::ReleaseEphemeralKeypair(Crypto::P256Keypair * keypair)
 {
-    std::unique_lock<std::mutex> lock(mMutex);
-
     if (mOperationalKeystore != nullptr)
     {
         mOperationalKeystore->ReleaseEphemeralKeypair(keypair);
@@ -1503,8 +1481,6 @@ bool FabricTable::HasOperationalKeyForFabric(FabricIndex fabricIndex) const
 
 CHIP_ERROR FabricTable::SignWithOpKeypair(FabricIndex fabricIndex, ByteSpan message, P256ECDSASignature & outSignature) const
 {
-    std::unique_lock<std::mutex> lock(mMutex);
-
     const FabricInfo * fabricInfo = FindFabricWithIndex(fabricIndex);
     VerifyOrReturnError(fabricInfo != nullptr, CHIP_ERROR_KEY_NOT_FOUND);
 
@@ -1549,8 +1525,6 @@ bool FabricTable::SetPendingDataFabricIndex(FabricIndex fabricIndex)
 
 CHIP_ERROR FabricTable::AllocatePendingOperationalKey(Optional<FabricIndex> fabricIndex, MutableByteSpan & outputCsr)
 {
-    std::unique_lock<std::mutex> lock(mMutex);
-
     // We can only manage commissionable pending fail-safe state if we have a keystore
     VerifyOrReturnError(mOperationalKeystore != nullptr, CHIP_ERROR_INCORRECT_STATE);
 
@@ -1593,8 +1567,6 @@ CHIP_ERROR FabricTable::AllocatePendingOperationalKey(Optional<FabricIndex> fabr
 
 CHIP_ERROR FabricTable::AddNewPendingTrustedRootCert(const ByteSpan & rcac)
 {
-    std::unique_lock<std::mutex> lock(mMutex);
-
     VerifyOrReturnError(mOpCertStore != nullptr, CHIP_ERROR_INCORRECT_STATE);
 
     // We should not already have pending NOC chain elements when we get here
@@ -1672,8 +1644,6 @@ CHIP_ERROR FabricTable::AddNewPendingFabricCommon(const ByteSpan & noc, const By
                                                   Crypto::P256Keypair * existingOpKey, bool isExistingOpKeyExternallyOwned,
                                                   FabricIndex * outNewFabricIndex)
 {
-    std::unique_lock<std::mutex> lock(mMutex);
-
     VerifyOrReturnError(mOpCertStore != nullptr, CHIP_ERROR_INCORRECT_STATE);
     VerifyOrReturnError(outNewFabricIndex != nullptr, CHIP_ERROR_INVALID_ARGUMENT);
     static_assert(kMaxValidFabricIndex <= UINT8_MAX, "Cannot create more fabrics than UINT8_MAX");
@@ -1744,8 +1714,6 @@ CHIP_ERROR FabricTable::AddNewPendingFabricCommon(const ByteSpan & noc, const By
 CHIP_ERROR FabricTable::UpdatePendingFabricCommon(FabricIndex fabricIndex, const ByteSpan & noc, const ByteSpan & icac,
                                                   Crypto::P256Keypair * existingOpKey, bool isExistingOpKeyExternallyOwned)
 {
-    std::unique_lock<std::mutex> lock(mMutex);
-
     VerifyOrReturnError(mOpCertStore != nullptr, CHIP_ERROR_INCORRECT_STATE);
     VerifyOrReturnError(IsValidFabricIndex(fabricIndex), CHIP_ERROR_INVALID_ARGUMENT);
 
@@ -1804,8 +1772,6 @@ CHIP_ERROR FabricTable::UpdatePendingFabricCommon(FabricIndex fabricIndex, const
 
 CHIP_ERROR FabricTable::CommitPendingFabricData()
 {
-    std::unique_lock<std::mutex> lock(mMutex);
-
     VerifyOrReturnError((mStorage != nullptr) && (mOpCertStore != nullptr), CHIP_ERROR_INCORRECT_STATE);
 
     bool haveNewTrustedRoot      = mStateFlags.Has(StateFlags::kIsTrustedRootPending);
@@ -2084,8 +2050,6 @@ void FabricTable::RevertPendingOpCertsExceptRoot()
 
 CHIP_ERROR FabricTable::SetFabricLabel(FabricIndex fabricIndex, const CharSpan & fabricLabel)
 {
-    std::unique_lock<std::mutex> lock(mMutex);
-
     VerifyOrReturnError(mStorage != nullptr, CHIP_ERROR_INCORRECT_STATE);
     VerifyOrReturnError(IsValidFabricIndex(fabricIndex), CHIP_ERROR_INVALID_FABRIC_INDEX);
 
@@ -2109,8 +2073,6 @@ CHIP_ERROR FabricTable::SetFabricLabel(FabricIndex fabricIndex, const CharSpan &
 
 CHIP_ERROR FabricTable::GetFabricLabel(FabricIndex fabricIndex, CharSpan & outFabricLabel)
 {
-    std::unique_lock<std::mutex> lock(mMutex);
-
     const FabricInfo * fabricInfo = FindFabricWithIndex(fabricIndex);
     VerifyOrReturnError(fabricInfo != nullptr, CHIP_ERROR_INVALID_FABRIC_INDEX);
 
diff --git a/src/credentials/FabricTable.h b/src/credentials/FabricTable.h
index 2a19f5bd97d60f..5407b627378755 100644
--- a/src/credentials/FabricTable.h
+++ b/src/credentials/FabricTable.h
@@ -22,7 +22,6 @@
 #pragma once
 
 #include <algorithm>
-#include <mutex>
 
 #include <app/util/basic-types.h>
 #include <credentials/CHIPCert.h>
@@ -718,6 +717,17 @@ class DLL_EXPORT FabricTable
      */
     bool HasOperationalKeyForFabric(FabricIndex fabricIndex) const;
 
+    /**
+     * @brief Returns the operational keystore. This is used for
+     *        CASE and the only way the keystore should be used.
+     *
+     * @return The operational keystore, nullptr otherwise.
+     */
+    Crypto::OperationalKeystore * GetOperationalKeystore()
+    {
+        return mOperationalKeystore;
+    }
+
     /**
      * @brief Add a pending trusted root certificate for the next fabric created with `AddNewPendingFabric*` methods.
      *
@@ -1147,8 +1157,6 @@ class DLL_EXPORT FabricTable
     uint8_t mFabricCount = 0;
 
     BitFlags<StateFlags> mStateFlags;
-
-    mutable std::mutex mMutex;
 };
 
 } // namespace chip
diff --git a/src/protocols/secure_channel/CASESession.cpp b/src/protocols/secure_channel/CASESession.cpp
index 36a372069405bb..fbdf593a87ac44 100644
--- a/src/protocols/secure_channel/CASESession.cpp
+++ b/src/protocols/secure_channel/CASESession.cpp
@@ -169,6 +169,24 @@ class CASESession::WorkHelper
         return ptr;
     }
 
+    // Do the work immediately.
+    // No scheduling, no outstanding work, no shared lifetime management.
+    CHIP_ERROR DoWork()
+    {
+        if (!mSession || !mWorkCallback || !mAfterWorkCallback)
+        {
+            return CHIP_ERROR_INCORRECT_STATE;
+        }
+        WorkHelper * helper = this;
+        bool cancel = false;
+        helper->mStatus = helper->mWorkCallback(helper->mData, cancel);
+        if (!cancel)
+        {
+            (helper->mSession->*(helper->mAfterWorkCallback))(helper->mData, helper->mStatus);
+        }
+        return CHIP_NO_ERROR;
+    }
+
     // Schedule the work after configuring the data.
     // If lifetime is managed, the helper shares management while work is outstanding.
     CHIP_ERROR ScheduleWork()
@@ -246,6 +264,8 @@ struct CASESession::SendSigma3Data
     FabricTable * fabricTable;
     std::atomic<FabricIndex> fabricIndex;
 
+    Crypto::OperationalKeystore * keystore;
+
     chip::Platform::ScopedMemoryBuffer<uint8_t> msg_R3_Signed;
     size_t msg_r3_signed_len;
 
@@ -1284,6 +1304,19 @@ CHIP_ERROR CASESession::SendSigma3a()
         data.fabricTable = mFabricsTable;
         data.fabricIndex = mFabricIndex;
 
+        // If an operational keystore is used, work will be performed in the background.
+        // Otherwise, legacy signing will be performed in the foreground.
+        data.keystore = nullptr;
+        {
+            const FabricInfo * fabricInfo = mFabricsTable->FindFabricWithIndex(mFabricIndex);
+            VerifyOrExit(fabricInfo != nullptr, err = CHIP_ERROR_KEY_NOT_FOUND);
+            if (!fabricInfo->HasOperationalKey())
+            {
+                data.keystore = mFabricsTable->GetOperationalKeystore();
+                VerifyOrExit(data.keystore != nullptr, err = CHIP_ERROR_KEY_NOT_FOUND);
+            }
+        }
+
         VerifyOrExit(mEphemeralKey != nullptr, err = CHIP_ERROR_INTERNAL);
 
         VerifyOrExit(data.icacBuf.Alloc(kMaxCHIPCertLength), err = CHIP_ERROR_NO_MEMORY);
@@ -1305,10 +1338,17 @@ CHIP_ERROR CASESession::SendSigma3a()
                           data.nocCert, data.icaCert, ByteSpan(mEphemeralKey->Pubkey(), mEphemeralKey->Pubkey().Length()),
                           ByteSpan(mRemotePubKey, mRemotePubKey.Length()), data.msg_R3_Signed.Get(), data.msg_r3_signed_len));
 
-        SuccessOrExit(err = helper->ScheduleWork());
-        mSendSigma3Helper = helper;
-        mExchangeCtxt->WillSendMessage();
-        mState = State::kSendSigma3Pending;
+        if (data.keystore != nullptr)
+        {
+            SuccessOrExit(err = helper->ScheduleWork());
+            mSendSigma3Helper = helper;
+            mExchangeCtxt->WillSendMessage();
+            mState = State::kSendSigma3Pending;
+        }
+        else
+        {
+            SuccessOrExit(err = helper->DoWork());
+        }
     }
 
 exit:
@@ -1326,8 +1366,18 @@ CHIP_ERROR CASESession::SendSigma3b(SendSigma3Data & data, bool & cancel)
     CHIP_ERROR err = CHIP_NO_ERROR;
 
     // Generate a signature
-    err = data.fabricTable->SignWithOpKeypair(data.fabricIndex, ByteSpan{ data.msg_R3_Signed.Get(), data.msg_r3_signed_len },
-                                              data.tbsData3Signature);
+    if (data.keystore != nullptr)
+    {
+        // Recommended case: delegate to operational keystore
+        err = data.keystore->SignWithOpKeypair(data.fabricIndex, ByteSpan{ data.msg_R3_Signed.Get(), data.msg_r3_signed_len },
+                                               data.tbsData3Signature);
+    }
+    else
+    {
+        // Legacy case: delegate to fabric table fabric info
+        err = data.fabricTable->SignWithOpKeypair(data.fabricIndex, ByteSpan{ data.msg_R3_Signed.Get(), data.msg_r3_signed_len },
+                                                data.tbsData3Signature);
+    }
     SuccessOrExit(err);
 
     // Prepare Sigma3 TBE Data Blob
@@ -1380,7 +1430,7 @@ CHIP_ERROR CASESession::SendSigma3c(SendSigma3Data & data, CHIP_ERROR status)
 
     AutoReleaseSessionKey sr3k(*mSessionManager->GetSessionKeystore());
 
-    VerifyOrDieWithMsg(mState == State::kSendSigma3Pending, SecureChannel, "Bad internal state.");
+    VerifyOrDieWithMsg(data.keystore == nullptr || mState == State::kSendSigma3Pending, SecureChannel, "Bad internal state.");
 
     SuccessOrExit(err = status);
 

From 94c1f182ee908a3a820da46c8a8e0bd7b80ff29c Mon Sep 17 00:00:00 2001
From: Marc Lepage <mlepage@google.com>
Date: Wed, 12 Apr 2023 17:06:27 -0400
Subject: [PATCH 14/19] Clean up error handling

---
 src/protocols/secure_channel/CASESession.cpp | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/src/protocols/secure_channel/CASESession.cpp b/src/protocols/secure_channel/CASESession.cpp
index fbdf593a87ac44..b133f30dbe61bb 100644
--- a/src/protocols/secure_channel/CASESession.cpp
+++ b/src/protocols/secure_channel/CASESession.cpp
@@ -182,9 +182,9 @@ class CASESession::WorkHelper
         helper->mStatus = helper->mWorkCallback(helper->mData, cancel);
         if (!cancel)
         {
-            (helper->mSession->*(helper->mAfterWorkCallback))(helper->mData, helper->mStatus);
+            helper->mStatus = (helper->mSession->*(helper->mAfterWorkCallback))(helper->mData, helper->mStatus);
         }
-        return CHIP_NO_ERROR;
+        return helper->mStatus;
     }
 
     // Schedule the work after configuring the data.
@@ -1304,7 +1304,7 @@ CHIP_ERROR CASESession::SendSigma3a()
         data.fabricTable = mFabricsTable;
         data.fabricIndex = mFabricIndex;
 
-        // If an operational keystore is used, work will be performed in the background.
+        // If an operational keystore is used, signing will be performed in the background.
         // Otherwise, legacy signing will be performed in the foreground.
         data.keystore = nullptr;
         {
@@ -1489,10 +1489,14 @@ CHIP_ERROR CASESession::SendSigma3c(SendSigma3Data & data, CHIP_ERROR status)
 exit:
     mSendSigma3Helper.reset();
 
-    if (err != CHIP_NO_ERROR)
+    // If data.keystore is set, processing occurred in the background, so if an error occurred,
+    // need to send status report (normally occurs in SendSigma3a), and discard exchange and
+    // abort pending establish (normally occurs in OnMessageReceived).
+    if (data.keystore != nullptr && err != CHIP_NO_ERROR)
     {
         SendStatusReport(mExchangeCtxt, kProtocolCodeInvalidParam);
-        mState = State::kInitialized;
+        DiscardExchange();
+        AbortPendingEstablish(err);
     }
 
     return err;

From e719521a57e93eb686cc3f98032e81c87d1fcf0b Mon Sep 17 00:00:00 2001
From: Marc Lepage <mlepage@google.com>
Date: Wed, 12 Apr 2023 17:08:36 -0400
Subject: [PATCH 15/19] Restyle

---
 src/credentials/FabricTable.h                | 5 +----
 src/protocols/secure_channel/CASESession.cpp | 6 +++---
 2 files changed, 4 insertions(+), 7 deletions(-)

diff --git a/src/credentials/FabricTable.h b/src/credentials/FabricTable.h
index 5407b627378755..c27f2027ae6b48 100644
--- a/src/credentials/FabricTable.h
+++ b/src/credentials/FabricTable.h
@@ -723,10 +723,7 @@ class DLL_EXPORT FabricTable
      *
      * @return The operational keystore, nullptr otherwise.
      */
-    Crypto::OperationalKeystore * GetOperationalKeystore()
-    {
-        return mOperationalKeystore;
-    }
+    Crypto::OperationalKeystore * GetOperationalKeystore() { return mOperationalKeystore; }
 
     /**
      * @brief Add a pending trusted root certificate for the next fabric created with `AddNewPendingFabric*` methods.
diff --git a/src/protocols/secure_channel/CASESession.cpp b/src/protocols/secure_channel/CASESession.cpp
index b133f30dbe61bb..286bd92e3bace8 100644
--- a/src/protocols/secure_channel/CASESession.cpp
+++ b/src/protocols/secure_channel/CASESession.cpp
@@ -178,8 +178,8 @@ class CASESession::WorkHelper
             return CHIP_ERROR_INCORRECT_STATE;
         }
         WorkHelper * helper = this;
-        bool cancel = false;
-        helper->mStatus = helper->mWorkCallback(helper->mData, cancel);
+        bool cancel         = false;
+        helper->mStatus     = helper->mWorkCallback(helper->mData, cancel);
         if (!cancel)
         {
             helper->mStatus = (helper->mSession->*(helper->mAfterWorkCallback))(helper->mData, helper->mStatus);
@@ -1376,7 +1376,7 @@ CHIP_ERROR CASESession::SendSigma3b(SendSigma3Data & data, bool & cancel)
     {
         // Legacy case: delegate to fabric table fabric info
         err = data.fabricTable->SignWithOpKeypair(data.fabricIndex, ByteSpan{ data.msg_R3_Signed.Get(), data.msg_r3_signed_len },
-                                                data.tbsData3Signature);
+                                                  data.tbsData3Signature);
     }
     SuccessOrExit(err);
 

From 647999793d98861553510b4c4d9eca4f6d933edc Mon Sep 17 00:00:00 2001
From: Marc Lepage <mlepage@google.com>
Date: Mon, 17 Apr 2023 15:05:54 -0400
Subject: [PATCH 16/19] Only store data.fabricTable if fg case

Store only one of data.fabricTable or data.keystore.
---
 src/protocols/secure_channel/CASESession.cpp | 17 +++++++++++------
 1 file changed, 11 insertions(+), 6 deletions(-)

diff --git a/src/protocols/secure_channel/CASESession.cpp b/src/protocols/secure_channel/CASESession.cpp
index 286bd92e3bace8..70146409192202 100644
--- a/src/protocols/secure_channel/CASESession.cpp
+++ b/src/protocols/secure_channel/CASESession.cpp
@@ -261,9 +261,10 @@ class CASESession::WorkHelper
 
 struct CASESession::SendSigma3Data
 {
-    FabricTable * fabricTable;
     std::atomic<FabricIndex> fabricIndex;
 
+    // Use one or the other
+    FabricTable * fabricTable;
     Crypto::OperationalKeystore * keystore;
 
     chip::Platform::ScopedMemoryBuffer<uint8_t> msg_R3_Signed;
@@ -1301,17 +1302,21 @@ CHIP_ERROR CASESession::SendSigma3a()
         auto & data = helper->mData;
 
         VerifyOrExit(mFabricsTable != nullptr, err = CHIP_ERROR_INCORRECT_STATE);
-        data.fabricTable = mFabricsTable;
         data.fabricIndex = mFabricIndex;
+        data.fabricTable = nullptr;
+        data.keystore    = nullptr;
 
-        // If an operational keystore is used, signing will be performed in the background.
-        // Otherwise, legacy signing will be performed in the foreground.
-        data.keystore = nullptr;
         {
             const FabricInfo * fabricInfo = mFabricsTable->FindFabricWithIndex(mFabricIndex);
             VerifyOrExit(fabricInfo != nullptr, err = CHIP_ERROR_KEY_NOT_FOUND);
-            if (!fabricInfo->HasOperationalKey())
+            if (fabricInfo->HasOperationalKey())
+            {
+                // NOTE: used to sign in foreground.
+                data.fabricTable = mFabricsTable;
+            }
+            else
             {
+                // NOTE: used to sign in background.
                 data.keystore = mFabricsTable->GetOperationalKeystore();
                 VerifyOrExit(data.keystore != nullptr, err = CHIP_ERROR_KEY_NOT_FOUND);
             }

From db76a70792bbec1f91c45270d7387f23cf39b2ce Mon Sep 17 00:00:00 2001
From: Marc Lepage <mlepage@google.com>
Date: Wed, 19 Apr 2023 13:14:38 -0400
Subject: [PATCH 17/19] Declare wither signing in background is supported

OperationalKeystore declares whether it supports this capability.
If so, then CASE session establishment may take advantage of it.
If not, then CASE session establishment must use foreground.
---
 src/crypto/OperationalKeystore.h             | 14 +++++++++++++-
 src/protocols/secure_channel/CASESession.cpp | 12 ++++++------
 2 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/src/crypto/OperationalKeystore.h b/src/crypto/OperationalKeystore.h
index bfa846b9d1a8be..6af92629174810 100644
--- a/src/crypto/OperationalKeystore.h
+++ b/src/crypto/OperationalKeystore.h
@@ -148,6 +148,19 @@ class OperationalKeystore
     virtual void RevertPendingKeypair() = 0;
 
     // ==== Primary operation required: signature
+    /**
+     * @brief Whether `SignWithOpKeypair` may be performed in the background.
+     *
+     * If true, `CASESession` may attempt to perform `SignWithOpKeypair` in the
+     * background. In this case, `OperationalKeystore` should protect itself,
+     * e.g. with a mutex, as the signing could occur at any time during session
+     * establishment.
+     *
+     * @retval true if `SignWithOpKeypair` may be performed in the background
+     * @retval false if `SignWithOpKeypair` may NOT be performed in the background
+     */
+    virtual bool SupportsSignWithOpKeypairInBackground() const { return false; }
+
     /**
      * @brief Sign a message with a fabric's currently-active operational keypair.
      *
@@ -164,7 +177,6 @@ class OperationalKeystore
      * @retval CHIP_ERROR_INVALID_FABRIC_INDEX if no active key is found for the given `fabricIndex` or if
      *                                         `fabricIndex` is invalid.
      * @retval other CHIP_ERROR value on internal crypto engine errors
-     *
      */
     virtual CHIP_ERROR SignWithOpKeypair(FabricIndex fabricIndex, const ByteSpan & message,
                                          Crypto::P256ECDSASignature & outSignature) const = 0;
diff --git a/src/protocols/secure_channel/CASESession.cpp b/src/protocols/secure_channel/CASESession.cpp
index 70146409192202..512ed319623b92 100644
--- a/src/protocols/secure_channel/CASESession.cpp
+++ b/src/protocols/secure_channel/CASESession.cpp
@@ -1309,16 +1309,16 @@ CHIP_ERROR CASESession::SendSigma3a()
         {
             const FabricInfo * fabricInfo = mFabricsTable->FindFabricWithIndex(mFabricIndex);
             VerifyOrExit(fabricInfo != nullptr, err = CHIP_ERROR_KEY_NOT_FOUND);
-            if (fabricInfo->HasOperationalKey())
+            auto * keystore = mFabricsTable->GetOperationalKeystore();
+            if (!fabricInfo->HasOperationalKey() && keystore != nullptr && keystore->SupportsSignWithOpKeypairInBackground())
             {
-                // NOTE: used to sign in foreground.
-                data.fabricTable = mFabricsTable;
+                // NOTE: used to sign in background.
+                data.keystore = keystore;
             }
             else
             {
-                // NOTE: used to sign in background.
-                data.keystore = mFabricsTable->GetOperationalKeystore();
-                VerifyOrExit(data.keystore != nullptr, err = CHIP_ERROR_KEY_NOT_FOUND);
+                // NOTE: used to sign in foreground.
+                data.fabricTable = mFabricsTable;
             }
         }
 

From 633b57f7baf5b7a6ba6f191f8418fce652f011d9 Mon Sep 17 00:00:00 2001
From: Marc Lepage <mlepage@google.com>
Date: Thu, 27 Apr 2023 10:47:41 -0400
Subject: [PATCH 18/19] Make some variables const

---
 src/credentials/FabricTable.h                | 2 +-
 src/protocols/secure_channel/CASESession.cpp | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/credentials/FabricTable.h b/src/credentials/FabricTable.h
index c27f2027ae6b48..1b01e1f4fd4a4f 100644
--- a/src/credentials/FabricTable.h
+++ b/src/credentials/FabricTable.h
@@ -723,7 +723,7 @@ class DLL_EXPORT FabricTable
      *
      * @return The operational keystore, nullptr otherwise.
      */
-    Crypto::OperationalKeystore * GetOperationalKeystore() { return mOperationalKeystore; }
+    const Crypto::OperationalKeystore * GetOperationalKeystore() { return mOperationalKeystore; }
 
     /**
      * @brief Add a pending trusted root certificate for the next fabric created with `AddNewPendingFabric*` methods.
diff --git a/src/protocols/secure_channel/CASESession.cpp b/src/protocols/secure_channel/CASESession.cpp
index 512ed319623b92..a710c02963d494 100644
--- a/src/protocols/secure_channel/CASESession.cpp
+++ b/src/protocols/secure_channel/CASESession.cpp
@@ -264,8 +264,8 @@ struct CASESession::SendSigma3Data
     std::atomic<FabricIndex> fabricIndex;
 
     // Use one or the other
-    FabricTable * fabricTable;
-    Crypto::OperationalKeystore * keystore;
+    const FabricTable * fabricTable;
+    const Crypto::OperationalKeystore * keystore;
 
     chip::Platform::ScopedMemoryBuffer<uint8_t> msg_R3_Signed;
     size_t msg_r3_signed_len;

From b75e07960418d92a8587c59867cd4f0ed6d2562e Mon Sep 17 00:00:00 2001
From: Marc Lepage <mlepage@google.com>
Date: Thu, 27 Apr 2023 10:50:37 -0400
Subject: [PATCH 19/19] Clean up a few comments

---
 src/protocols/secure_channel/CASESession.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/protocols/secure_channel/CASESession.cpp b/src/protocols/secure_channel/CASESession.cpp
index a710c02963d494..20b3567cb5d74b 100644
--- a/src/protocols/secure_channel/CASESession.cpp
+++ b/src/protocols/secure_channel/CASESession.cpp
@@ -145,7 +145,7 @@ class CASESession::WorkHelper
     // Set `cancel` to true if calling the after work callback is not necessary.
     typedef CHIP_ERROR (*WorkCallback)(DATA & data, bool & cancel);
 
-    // After work callback, processed in the main CHIP task via `PlatformManager::ScheduleWork`.
+    // After work callback, processed in the main Matter task via `PlatformManager::ScheduleWork`.
     // This is a member function to be called on the associated session after the work callback.
     // The `status` value is the result of the work callback (called beforehand).
     typedef CHIP_ERROR (CASESession::*AfterWorkCallback)(DATA & data, CHIP_ERROR status);
@@ -226,7 +226,7 @@ class CASESession::WorkHelper
     // Handler for the after work callback.
     static void AfterWorkHandler(intptr_t arg)
     {
-        // Since this runs in the main CHIP thread, the session shouldn't be otherwise used (messages, timers, etc.)
+        // Since this runs in the main Matter thread, the session shouldn't be otherwise used (messages, timers, etc.)
         WorkHelper * helper = reinterpret_cast<WorkHelper *>(arg);
         if (auto * session = helper->mSession.load())
         {