From da9965197e970e5ee111db5c0704908b8557d026 Mon Sep 17 00:00:00 2001
From: Stefano Avallone <stavallo@unina.it>
Date: Tue, 30 Aug 2022 16:03:33 +0200
Subject: [PATCH] wifi: Recipient BA agreements are in effect if ADDBA_RESPONSE
 is acknowledged

---
 src/wifi/model/ht/ht-frame-exchange-manager.cc | 16 +++++++++++++++-
 src/wifi/model/ht/ht-frame-exchange-manager.h  | 11 ++++++++---
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/src/wifi/model/ht/ht-frame-exchange-manager.cc b/src/wifi/model/ht/ht-frame-exchange-manager.cc
index 08bd7cdc0..b3e7b8882 100644
--- a/src/wifi/model/ht/ht-frame-exchange-manager.cc
+++ b/src/wifi/model/ht/ht-frame-exchange-manager.cc
@@ -67,6 +67,7 @@ HtFrameExchangeManager::DoDispose (void)
 {
   NS_LOG_FUNCTION (this);
   m_agreements.clear ();
+  m_pendingAgreements.clear ();
   m_msduAggregator = 0;
   m_mpduAggregator = 0;
   m_psdu = 0;
@@ -317,7 +318,7 @@ HtFrameExchangeManager::CreateBlockAckAgreement (const MgtAddBaResponseHeader *r
                                                          this, originator, tid, false);
     }
 
-  m_agreements.insert ({{originator, tid}, agreement});
+  m_pendingAgreements.insert_or_assign ({originator, tid}, agreement);
 }
 
 void
@@ -574,6 +575,19 @@ HtFrameExchangeManager::NotifyReceivedNormalAck (Ptr<WifiMacQueueItem> mpdu)
                                    &QosTxop::AddBaResponseTimeout, edca,
                                    mpdu->GetHeader ().GetAddr1 (), addBa.GetTid ());
             }
+          else if (actionHdr.GetAction ().blockAck == WifiActionHeader::BLOCK_ACK_ADDBA_RESPONSE)
+            {
+              // The recipient Block Ack agreement can be moved from the pending queue
+              // to the queue of established Block Ack agreements
+              MgtAddBaResponseHeader addBa;
+              p->PeekHeader (addBa);
+              AgreementKey key (mpdu->GetHeader ().GetAddr1 (), addBa.GetTid ());
+              auto nh = m_pendingAgreements.extract (key);
+              NS_ASSERT_MSG (!nh.empty (), "Pending agreement {" << key.first
+                                           << ", " << +key.second << "} not found");
+              m_agreements.erase (nh.key ());
+              m_agreements.insert (std::move (nh));
+            }
         }
     }
   QosFrameExchangeManager::NotifyReceivedNormalAck (mpdu);
diff --git a/src/wifi/model/ht/ht-frame-exchange-manager.h b/src/wifi/model/ht/ht-frame-exchange-manager.h
index fcce481a1..b71aeaccd 100644
--- a/src/wifi/model/ht/ht-frame-exchange-manager.h
+++ b/src/wifi/model/ht/ht-frame-exchange-manager.h
@@ -376,9 +376,14 @@ protected:
   /// agreement key typedef (MAC address and TID)
   typedef std::pair<Mac48Address, uint8_t> AgreementKey;
 
-  std::map<AgreementKey, RecipientBlockAckAgreement> m_agreements; //!< agreements
-  Ptr<MsduAggregator> m_msduAggregator;                            //!< A-MSDU aggregator
-  Ptr<MpduAggregator> m_mpduAggregator;                            //!< A-MPDU aggregator
+  /// typedef for map of recipient Block Ack agreements
+  using RecipientBlockAckAgreementMap = std::map<AgreementKey, RecipientBlockAckAgreement>;
+
+  RecipientBlockAckAgreementMap m_agreements;        //!< Block Ack agreements
+  RecipientBlockAckAgreementMap m_pendingAgreements; //!< pending Block Ack agreements (waiting
+                                                     //!< for Ack in response to ADDBA_RESPONSE)
+  Ptr<MsduAggregator> m_msduAggregator;              //!< A-MSDU aggregator
+  Ptr<MpduAggregator> m_mpduAggregator;              //!< A-MPDU aggregator
 
 private:
   /**