From 5549017cb4a3651148907085f3d96e803d9e1550 Mon Sep 17 00:00:00 2001
From: Boris Zbarsky <bzbarsky@apple.com>
Date: Tue, 28 Feb 2023 11:20:58 -0500
Subject: [PATCH] Make PASE setup a bit more robust if multiple clients race.
 (#25352)

Without this change, we can end up in the following situation:

1) A PBKDFParamRequest comes in.  We start a PASE handshake.
2) While we are in the middle of that, a PBKDFParamRequest from some other
   entity comes in on a different exchange.
3) Since we are not expecting PBKDFParamRequest, we not only respond with
   failure to the new message, but also cancel the exising handshake.

So if two clients race to establish PASE, they can keep canceling each other and
neither will complete.

The fix is to stop listening for PBKDFParamRequest while we are in the middle of
a handshake.
---
 src/app/server/CommissioningWindowManager.cpp | 30 ++++++++++++++++++-
 src/app/server/CommissioningWindowManager.h   | 12 ++++++--
 2 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/src/app/server/CommissioningWindowManager.cpp b/src/app/server/CommissioningWindowManager.cpp
index a899179c85e545..3ddde11a931acb 100644
--- a/src/app/server/CommissioningWindowManager.cpp
+++ b/src/app/server/CommissioningWindowManager.cpp
@@ -239,7 +239,7 @@ CHIP_ERROR CommissioningWindowManager::AdvertiseAndListenForPASE()
 #endif
 
     ReturnErrorOnFailure(mServer->GetExchangeManager().RegisterUnsolicitedMessageHandlerForType(
-        Protocols::SecureChannel::MsgType::PBKDFParamRequest, &mPairingSession));
+        Protocols::SecureChannel::MsgType::PBKDFParamRequest, this));
     mListeningForPASE = true;
 
     if (mUseECM)
@@ -584,4 +584,32 @@ void CommissioningWindowManager::UpdateOpenerFabricIndex(Nullable<FabricIndex> a
     mOpenerFabricIndex = aNewOpenerFabricIndex;
 }
 
+CHIP_ERROR CommissioningWindowManager::OnUnsolicitedMessageReceived(const PayloadHeader & payloadHeader,
+                                                                    Messaging::ExchangeDelegate *& newDelegate)
+{
+    using Protocols::SecureChannel::MsgType;
+
+    // Must be a PBKDFParamRequest message.  Stop listening to new
+    // PBKDFParamRequest messages and hand it off to mPairingSession.  If
+    // mPairingSession's OnMessageReceived fails, it will call our
+    // OnSessionEstablishmentError, and that will either start listening for a
+    // new PBKDFParamRequest or not, depending on how many failures we had seen.
+    //
+    // It's very important that we stop listening here, so that new incoming
+    // PASE establishment attempts don't interrupt our existing establishment.
+    mServer->GetExchangeManager().UnregisterUnsolicitedMessageHandlerForType(MsgType::PBKDFParamRequest);
+    newDelegate = &mPairingSession;
+    return CHIP_NO_ERROR;
+}
+
+void CommissioningWindowManager::OnExchangeCreationFailed(Messaging::ExchangeDelegate * delegate)
+{
+    using Protocols::SecureChannel::MsgType;
+
+    // We couldn't create an exchange, so didn't manage to call
+    // OnMessageReceived on mPairingSession.  Just go back to listening for
+    // PBKDFParamRequest messages.
+    mServer->GetExchangeManager().RegisterUnsolicitedMessageHandlerForType(MsgType::PBKDFParamRequest, this);
+}
+
 } // namespace chip
diff --git a/src/app/server/CommissioningWindowManager.h b/src/app/server/CommissioningWindowManager.h
index cfd13eecb06c11..cc2eda2ea29199 100644
--- a/src/app/server/CommissioningWindowManager.h
+++ b/src/app/server/CommissioningWindowManager.h
@@ -24,6 +24,7 @@
 #include <lib/core/CHIPVendorIdentifiers.hpp>
 #include <lib/core/DataModelTypes.h>
 #include <lib/dnssd/Advertiser.h>
+#include <messaging/ExchangeDelegate.h>
 #include <platform/CHIPDeviceConfig.h>
 #include <protocols/secure_channel/RendezvousParameters.h>
 #include <system/SystemClock.h>
@@ -38,7 +39,8 @@ enum class CommissioningWindowAdvertisement
 
 class Server;
 
-class CommissioningWindowManager : public SessionEstablishmentDelegate,
+class CommissioningWindowManager : public Messaging::UnsolicitedMessageHandler,
+                                   public SessionEstablishmentDelegate,
                                    public app::CommissioningModeProvider,
                                    public SessionDelegate
 {
@@ -104,6 +106,11 @@ class CommissioningWindowManager : public SessionEstablishmentDelegate,
     // CommissioningModeProvider implementation.
     Dnssd::CommissioningMode GetCommissioningMode() const override;
 
+    //// UnsolicitedMessageHandler Implementation ////
+    CHIP_ERROR OnUnsolicitedMessageReceived(const PayloadHeader & payloadHeader,
+                                            Messaging::ExchangeDelegate *& newDelegate) override;
+    void OnExchangeCreationFailed(Messaging::ExchangeDelegate * delegate) override;
+
     //////////// SessionEstablishmentDelegate Implementation ///////////////
     void OnSessionEstablishmentError(CHIP_ERROR error) override;
     void OnSessionEstablishmentStarted() override;
@@ -195,7 +202,8 @@ class CommissioningWindowManager : public SessionEstablishmentDelegate,
     Spake2pVerifier mECMPASEVerifier;
     uint16_t mECMDiscriminator = 0;
     // mListeningForPASE is true only when we are listening for
-    // PBKDFParamRequest messages.
+    // PBKDFParamRequest messages or when we're in the middle of a PASE
+    // handshake.
     bool mListeningForPASE = false;
     // Boolean that tracks whether we have a live commissioning timeout timer.
     bool mCommissioningTimeoutTimerArmed = false;