Fix handling of a response that comes after MRP resends end. (#29640)

* Fix handling of a response that comes after MRP resends end.

After #29173 we can get into the following situation:

1. A message is sent.
2. Before we get an ack or response, all MRP retries happen, MRP gives up, but
   the exchange response timer has not been hit yet.
3. We get an actual response.
4. Because our exchange is marked as having an un-acked message, but the
   incoming message is not treated as an ack (because the MRP state that would
   do that has been torn down), we do not clear our "have un-acked message"
   state and end up discarding the incoming message.

The fix is as follows:

* Rename things to make it clear that what we really have is "waiting for an
  ack" state, which in fact _does_ get cleared when we run out of MRP retries,
  not an "un-acked message" state.
* Have a separate state bit for tracking that we ran out of MRP retries on a
  message we sent.

* Address review comment.

Author: bzbarsky-apple

src/messaging/ExchangeContext.cpp

@@ -79,6 +79,7 @@ bool ExchangeContext::IsResponseExpected() const
 void ExchangeContext::SetResponseExpected(bool inResponseExpected)
 {
     mFlags.Set(Flags::kFlagResponseExpected, inResponseExpected);
+    SetWaitingForResponseOrAck(inResponseExpected);
 }
 
 void ExchangeContext::UseSuggestedResponseTimeout(Timeout applicationProcessingTimeout)
@@ -489,6 +490,9 @@ void ExchangeContext::NotifyResponseTimeout(bool aCloseIfNeeded)
     }
 #endif // CONFIG_DEVICE_LAYER && CHIP_CONFIG_ENABLE_ICD_SERVER
 
+    // Grab the value of WaitingForResponseOrAck() before we mess with our state.
+    bool gotMRPAck = !WaitingForResponseOrAck();
+
     SetResponseExpected(false);
 
     // Hold a ref to ourselves so we can make calls into our delegate that might
@@ -502,9 +506,8 @@ void ExchangeContext::NotifyResponseTimeout(bool aCloseIfNeeded)
     {
         // If we timed out _after_ getting an ack for the message, that means
         // the session is probably fine (since our message and the ack got
-        // through), so don't mark the session defunct unless we have an
-        // un-acked message here.
-        if (IsMessageNotAcked())
+        // through), so don't mark the session defunct if we got an MRP ack.
+        if (!gotMRPAck)
         {
             if (mSession->IsSecureSession() && mSession->AsSecureSession()->IsCASESession())
             {
@@ -596,7 +599,7 @@ CHIP_ERROR ExchangeContext::HandleMessage(uint32_t messageCounter, const Payload
         return CHIP_NO_ERROR;
     }
 
-    if (IsMessageNotAcked())
+    if (IsWaitingForAck())
     {
         // The only way we can get here is a spec violation on the other side:
         // we sent a message that needs an ack, and the other side responded

src/messaging/ReliableMessageContext.cpp

@@ -55,14 +55,14 @@ ReliableMessageMgr * ReliableMessageContext::GetReliableMessageMgr()
     return static_cast<ExchangeContext *>(this)->GetExchangeMgr()->GetReliableMessageMgr();
 }
 
-void ReliableMessageContext::SetMessageNotAcked(bool messageNotAcked)
+void ReliableMessageContext::SetWaitingForAck(bool waitingForAck)
 {
-    mFlags.Set(Flags::kFlagMessageNotAcked, messageNotAcked);
+    mFlags.Set(Flags::kFlagWaitingForAck, waitingForAck);
 
 #if CONFIG_DEVICE_LAYER && CHIP_CONFIG_ENABLE_ICD_SERVER
     DeviceLayer::ChipDeviceEvent event;
     event.Type                = DeviceLayer::DeviceEventType::kICDMsgAckSyncEvent;
-    event.AckSync.awaitingAck = messageNotAcked;
+    event.AckSync.awaitingAck = waitingForAck;
     CHIP_ERROR status         = DeviceLayer::PlatformMgr().PostEvent(&event);
     if (status != CHIP_NO_ERROR)
     {
@@ -103,7 +103,11 @@ CHIP_ERROR ReliableMessageContext::FlushAcks()
 void ReliableMessageContext::HandleRcvdAck(uint32_t ackMessageCounter)
 {
     // Msg is an Ack; Check Retrans Table and remove message context
-    if (!GetReliableMessageMgr()->CheckAndRemRetransTable(this, ackMessageCounter))
+    if (GetReliableMessageMgr()->CheckAndRemRetransTable(this, ackMessageCounter))
+    {
+        SetWaitingForResponseOrAck(false);
+    }
+    else
     {
         // This can happen quite easily due to a packet with a piggyback ack
         // being lost and retransmitted.

src/messaging/ReliableMessageContext.h

@@ -114,11 +114,11 @@ class ReliableMessageContext
      */
     bool IsAckPending() const;
 
-    /// Determine whether there is message hasn't been acknowledged.
-    bool IsMessageNotAcked() const;
+    /// Determine whether the reliable message context is waiting for an ack.
+    bool IsWaitingForAck() const;
 
-    /// Set whether there is a message hasn't been acknowledged.
-    void SetMessageNotAcked(bool messageNotAcked);
+    /// Set whether the reliable message context is waiting for an ack.
+    void SetWaitingForAck(bool waitingForAck);
 
     /// Set if this exchange is requesting Sleepy End Device active mode
     void SetRequestingActiveMode(bool activeMode);
@@ -136,6 +136,9 @@ class ReliableMessageContext
     ReliableMessageMgr * GetReliableMessageMgr();
 
 protected:
+    bool WaitingForResponseOrAck() const;
+    void SetWaitingForResponseOrAck(bool waitingForResponseOrAck);
+
     enum class Flags : uint16_t
     {
         /// When set, signifies that this context is the initiator of the exchange.
@@ -147,8 +150,10 @@ class ReliableMessageContext
         /// When set, automatically request an acknowledgment whenever a message is sent via UDP.
         kFlagAutoRequestAck = (1u << 2),
 
-        /// When set, signifies there is a message which hasn't been acknowledged.
-        kFlagMessageNotAcked = (1u << 3),
+        /// When set, signifies the reliable message context is waiting for an
+        /// ack: a message that needs an ack has been sent, no ack has been
+        /// received, and we have not yet run out of MRP retries.
+        kFlagWaitingForAck = (1u << 3),
 
         /// When set, signifies that there is an acknowledgment pending to be sent back.
         kFlagAckPending = (1u << 4),
@@ -172,6 +177,13 @@ class ReliableMessageContext
 
         /// When set, ignore session being released, because we are releasing it ourselves.
         kFlagIgnoreSessionRelease = (1u << 10),
+
+        /// When set:
+        ///
+        /// (1) We sent a message that expected a response (hence
+        ///     IsResponseExpected() is true).
+        /// (2) We have received neither a response nor an ack for that message.
+        kFlagWaitingForResponseOrAck = (1u << 11),
     };
 
     BitFlags<Flags> mFlags; // Internal state flags
@@ -216,9 +228,9 @@ inline bool ReliableMessageContext::IsAckPending() const
     return mFlags.Has(Flags::kFlagAckPending);
 }
 
-inline bool ReliableMessageContext::IsMessageNotAcked() const
+inline bool ReliableMessageContext::IsWaitingForAck() const
 {
-    return mFlags.Has(Flags::kFlagMessageNotAcked);
+    return mFlags.Has(Flags::kFlagWaitingForAck);
 }
 
 inline bool ReliableMessageContext::HasPiggybackAckPending() const
@@ -251,5 +263,15 @@ inline bool ReliableMessageContext::IsEphemeralExchange() const
     return mFlags.Has(Flags::kFlagEphemeralExchange);
 }
 
+inline bool ReliableMessageContext::WaitingForResponseOrAck() const
+{
+    return mFlags.Has(Flags::kFlagWaitingForResponseOrAck);
+}
+
+inline void ReliableMessageContext::SetWaitingForResponseOrAck(bool waitingForResponseOrAck)
+{
+    mFlags.Set(Flags::kFlagWaitingForResponseOrAck, waitingForResponseOrAck);
+}
+
 } // namespace Messaging
 } // namespace chip

src/messaging/ReliableMessageMgr.cpp

@@ -49,12 +49,12 @@ namespace Messaging {
 ReliableMessageMgr::RetransTableEntry::RetransTableEntry(ReliableMessageContext * rc) :
     ec(*rc->GetExchangeContext()), nextRetransTime(0), sendCount(0)
 {
-    ec->SetMessageNotAcked(true);
+    ec->SetWaitingForAck(true);
 }
 
 ReliableMessageMgr::RetransTableEntry::~RetransTableEntry()
 {
-    ec->SetMessageNotAcked(false);
+    ec->SetWaitingForAck(false);
 }
 
 ReliableMessageMgr::ReliableMessageMgr(ObjectPool<ExchangeContext, CHIP_CONFIG_MAX_EXCHANGE_CONTEXTS> & contextPool) :
@@ -158,12 +158,6 @@ void ReliableMessageMgr::ExecuteActions()
             // Do not StartTimer, we will schedule the timer at the end of the timer handler.
             mRetransTable.ReleaseObject(entry);
 
-            // Dropping our entry marked the exchange as not having an un-acked
-            // message... but of course it _does_ have an un-acked message and
-            // we have just given up on waiting for the ack.
-
-            ec->GetReliableMessageContext()->SetMessageNotAcked(true);
-
             return Loop::Continue;
         }
 
@@ -204,7 +198,7 @@ void ReliableMessageMgr::Timeout(System::Layer * aSystemLayer, void * aAppState)
 
 CHIP_ERROR ReliableMessageMgr::AddToRetransTable(ReliableMessageContext * rc, RetransTableEntry ** rEntry)
 {
-    VerifyOrDie(!rc->IsMessageNotAcked());
+    VerifyOrDie(!rc->IsWaitingForAck());
 
     *rEntry = mRetransTable.CreateObject(rc);
     if (*rEntry == nullptr)

src/messaging/tests/MessagingContext.cpp

@@ -163,12 +163,26 @@ CHIP_ERROR MessagingContext::CreateSessionBobToAlice()
                                                         mBobFabricIndex, mAliceAddress, CryptoContext::SessionRole::kInitiator);
 }
 
+CHIP_ERROR MessagingContext::CreateCASESessionBobToAlice()
+{
+    return mSessionManager.InjectCaseSessionWithTestKey(mSessionBobToAlice, kBobKeyId, kAliceKeyId, GetBobFabric()->GetNodeId(),
+                                                        GetAliceFabric()->GetNodeId(), mBobFabricIndex, mAliceAddress,
+                                                        CryptoContext::SessionRole::kInitiator);
+}
+
 CHIP_ERROR MessagingContext::CreateSessionAliceToBob()
 {
     return mSessionManager.InjectPaseSessionWithTestKey(mSessionAliceToBob, kAliceKeyId, GetBobFabric()->GetNodeId(), kBobKeyId,
                                                         mAliceFabricIndex, mBobAddress, CryptoContext::SessionRole::kResponder);
 }
 
+CHIP_ERROR MessagingContext::CreateCASESessionAliceToBob()
+{
+    return mSessionManager.InjectCaseSessionWithTestKey(mSessionAliceToBob, kAliceKeyId, kBobKeyId, GetAliceFabric()->GetNodeId(),
+                                                        GetBobFabric()->GetNodeId(), mAliceFabricIndex, mBobAddress,
+                                                        CryptoContext::SessionRole::kResponder);
+}
+
 CHIP_ERROR MessagingContext::CreatePASESessionCharlieToDavid()
 {
     return mSessionManager.InjectPaseSessionWithTestKey(mSessionCharlieToDavid, kCharlieKeyId, 0xdeadbeef, kDavidKeyId,

src/messaging/tests/MessagingContext.h

@@ -139,9 +139,11 @@ class MessagingContext : public PlatformMemoryUser
     const FabricInfo * GetAliceFabric() { return mFabricTable.FindFabricWithIndex(mAliceFabricIndex); }
     const FabricInfo * GetBobFabric() { return mFabricTable.FindFabricWithIndex(mBobFabricIndex); }
 
-    CHIP_ERROR CreateSessionBobToAlice();
-    CHIP_ERROR CreateSessionAliceToBob();
-    CHIP_ERROR CreateSessionBobToFriends();
+    CHIP_ERROR CreateSessionBobToAlice(); // Creates PASE session
+    CHIP_ERROR CreateCASESessionBobToAlice();
+    CHIP_ERROR CreateSessionAliceToBob(); // Creates PASE session
+    CHIP_ERROR CreateCASESessionAliceToBob();
+    CHIP_ERROR CreateSessionBobToFriends(); // Creates PASE session
     CHIP_ERROR CreatePASESessionCharlieToDavid();
     CHIP_ERROR CreatePASESessionDavidToCharlie();