From e71a0167d471e12f9abb4875323f326f58460799 Mon Sep 17 00:00:00 2001
From: Stefano Avallone <stavallo@unina.it>
Date: Thu, 7 Jan 2021 18:17:58 +0100
Subject: [PATCH] wifi: BlockAckManager also stores MU-BAR Trigger Frames

---
 src/wifi/model/block-ack-manager.cc | 50 ++++++++++++++++++++++-------
 src/wifi/model/block-ack-manager.h  | 18 +++++++----
 2 files changed, 50 insertions(+), 18 deletions(-)

diff --git a/src/wifi/model/block-ack-manager.cc b/src/wifi/model/block-ack-manager.cc
index 0553a4c45..b4a6357b5 100644
--- a/src/wifi/model/block-ack-manager.cc
+++ b/src/wifi/model/block-ack-manager.cc
@@ -274,7 +274,7 @@ BlockAckManager::StorePacket (Ptr<WifiMacQueueItem> mpdu)
 }
 
 Ptr<const WifiMacQueueItem>
-BlockAckManager::GetBar (bool remove)
+BlockAckManager::GetBar (bool remove, uint8_t tid, Mac48Address address)
 {
   Ptr<const WifiMacQueueItem> bar;
   // remove all expired MPDUs in the retransmission queue, so that Block Ack Requests
@@ -285,9 +285,18 @@ BlockAckManager::GetBar (bool remove)
 
   while (nextBar != m_bars.end ())
     {
+      Mac48Address recipient = nextBar->bar->GetHeader ().GetAddr1 ();
+
+      if (address != Mac48Address::GetBroadcast () && tid != 8
+            && (!nextBar->bar->GetHeader ().IsBlockAckReq ()
+                  || address != recipient || tid != nextBar->tid))
+        {
+          // we can only return a BAR addressed to the given station and for the given TID
+          nextBar++;
+          continue;
+        }
       if (nextBar->bar->GetHeader ().IsBlockAckReq ())
         {
-          Mac48Address recipient = nextBar->bar->GetHeader ().GetAddr1 ();
           AgreementsI it = m_agreements.find (std::make_pair (recipient, nextBar->tid));
           if (it == m_agreements.end ())
             {
@@ -653,21 +662,38 @@ void
 BlockAckManager::ScheduleBar (Ptr<const WifiMacQueueItem> bar, bool skipIfNoDataQueued)
 {
   NS_LOG_FUNCTION (this << *bar);
-  NS_ASSERT (bar->GetHeader ().IsBlockAckReq ());
+  NS_ASSERT (bar->GetHeader ().IsBlockAckReq () || bar->GetHeader ().IsTrigger ());
 
-  CtrlBAckRequestHeader reqHdr;
-  bar->GetPacket ()->PeekHeader (reqHdr);
-  uint8_t tid = reqHdr.GetTidInfo ();
+  uint8_t tid = 0;
+  if (bar->GetHeader ().IsBlockAckReq ())
+    {
+      CtrlBAckRequestHeader reqHdr;
+      bar->GetPacket ()->PeekHeader (reqHdr);
+      tid = reqHdr.GetTidInfo ();
+    }
+#ifdef NS3_BUILD_PROFILE_DEBUG
+  else
+    {
+      CtrlTriggerHeader triggerHdr;
+      bar->GetPacket ()->PeekHeader (triggerHdr);
+      NS_ASSERT (triggerHdr.IsMuBar ());
+    }
+#endif
   Bar request (bar, tid, skipIfNoDataQueued);
 
   // if a BAR for the given agreement is present, replace it with the new one
-  for (std::list<Bar>::const_iterator i = m_bars.begin (); i != m_bars.end (); i++)
+  std::list<Bar>::const_iterator i = m_bars.end ();
+
+  if (bar->GetHeader ().IsBlockAckReq ())
     {
-      if (i->bar->GetHeader ().GetAddr1 () == bar->GetHeader ().GetAddr1 () && i->tid == tid)
+      for (i = m_bars.begin (); i != m_bars.end (); i++)
         {
-          i = m_bars.erase (i);
-          m_bars.insert (i, request);
-          return;
+          if (i->bar->GetHeader ().IsBlockAckReq ()
+              && i->bar->GetHeader ().GetAddr1 () == bar->GetHeader ().GetAddr1 () && i->tid == tid)
+            {
+              i = m_bars.erase (i);
+              break;
+            }
         }
     }
 
@@ -677,7 +703,7 @@ BlockAckManager::ScheduleBar (Ptr<const WifiMacQueueItem> bar, bool skipIfNoData
     }
   else
     {
-      m_bars.push_back (request);
+      m_bars.insert (i, request);
     }
 }
 
diff --git a/src/wifi/model/block-ack-manager.h b/src/wifi/model/block-ack-manager.h
index 02b9ee9d3..429439185 100644
--- a/src/wifi/model/block-ack-manager.h
+++ b/src/wifi/model/block-ack-manager.h
@@ -51,16 +51,16 @@ struct Bar
 {
   Bar ();
   /**
-   * Store a BlockAckRequest along with the corresponding TID.
+   * Store a BlockAckRequest along with the corresponding TID or a MU-BAR Trigger Frame.
    *
    * \param bar the BAR
    * \param tid the Traffic ID
    * \param skipIfNoDataQueued true to hold this BAR if there is no data queued
    */
   Bar (Ptr<const WifiMacQueueItem> bar, uint8_t tid, bool skipIfNoDataQueued = false);
-  Ptr<const WifiMacQueueItem> bar;  ///< BlockAckRequest
-  uint8_t tid;                      ///< TID
-  bool skipIfNoDataQueued;          ///< do not send if there is no data queued
+  Ptr<const WifiMacQueueItem> bar;  ///< BlockAckRequest or MU-BAR Trigger Frame
+  uint8_t tid;                      ///< TID (unused if MU-BAR)
+  bool skipIfNoDataQueued;          ///< do not send if there is no data queued (unused if MU-BAR)
 };
 
 
@@ -152,13 +152,19 @@ public:
    */
   void StorePacket (Ptr<WifiMacQueueItem> mpdu);
   /**
-   * Returns the next BlockAckRequest to send, if any.
+   * Returns the next BlockAckRequest or MU-BAR Trigger Frame to send, if any.
+   * If the given recipient is not the broadcast address and the given TID is less
+   * than 8, then only return a BlockAckRequest, if any, addressed to that recipient
+   * and for the given TID.
    *
    * \param remove true if the BAR has to be removed from the queue
+   * \param tid the TID
+   * \param recipient the recipient of the BAR
    *
    * \return the next BAR to be sent, if any
    */
-  Ptr<const WifiMacQueueItem> GetBar (bool remove = true);
+  Ptr<const WifiMacQueueItem> GetBar (bool remove = true, uint8_t tid = 8,
+                                      Mac48Address recipient = Mac48Address::GetBroadcast ());
   /**
    * Returns true if there are packets that need of retransmission or at least a
    * BAR is scheduled. Returns false otherwise.