From bfcc5b71fa78e28ba25ce8470998c82234f86cd0 Mon Sep 17 00:00:00 2001
From: Stefano Avallone <stavallo@unina.it>
Date: Wed, 7 Dec 2022 11:35:42 +0100
Subject: [PATCH] wifi: MLDs are indexed by MLD address in the Buffer Status
 map

---
 src/wifi/model/he/rr-multi-user-scheduler.cc | 13 +++++++------
 src/wifi/model/qos-frame-exchange-manager.cc |  7 ++++---
 2 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/src/wifi/model/he/rr-multi-user-scheduler.cc b/src/wifi/model/he/rr-multi-user-scheduler.cc
index 9fa5444a7..5ad41fe2e 100644
--- a/src/wifi/model/he/rr-multi-user-scheduler.cc
+++ b/src/wifi/model/he/rr-multi-user-scheduler.cc
@@ -394,22 +394,23 @@ RrMultiUserScheduler::TrySendingBasicTf()
 
     for (const auto& candidate : txVector.GetHeMuUserInfoMap())
     {
-        auto staIt = m_apMac->GetStaList().find(candidate.first);
-        NS_ASSERT(staIt != m_apMac->GetStaList().end());
-        uint8_t queueSize = m_apMac->GetMaxBufferStatus(staIt->second);
+        auto address = m_apMac->GetMldOrLinkAddressByAid(candidate.first);
+        NS_ASSERT_MSG(address, "AID " << candidate.first << " not found");
+
+        uint8_t queueSize = m_apMac->GetMaxBufferStatus(*address);
         if (queueSize == 255)
         {
-            NS_LOG_DEBUG("Buffer status of station " << staIt->second << " is unknown");
+            NS_LOG_DEBUG("Buffer status of station " << *address << " is unknown");
             maxBufferSize = std::max(maxBufferSize, m_ulPsduSize);
         }
         else if (queueSize == 254)
         {
-            NS_LOG_DEBUG("Buffer status of station " << staIt->second << " is not limited");
+            NS_LOG_DEBUG("Buffer status of station " << *address << " is not limited");
             maxBufferSize = 0xffffffff;
         }
         else
         {
-            NS_LOG_DEBUG("Buffer status of station " << staIt->second << " is " << +queueSize);
+            NS_LOG_DEBUG("Buffer status of station " << *address << " is " << +queueSize);
             maxBufferSize = std::max(maxBufferSize, static_cast<uint32_t>(queueSize * 256));
         }
     }
diff --git a/src/wifi/model/qos-frame-exchange-manager.cc b/src/wifi/model/qos-frame-exchange-manager.cc
index 34766713e..2f91a44a2 100644
--- a/src/wifi/model/qos-frame-exchange-manager.cc
+++ b/src/wifi/model/qos-frame-exchange-manager.cc
@@ -658,9 +658,10 @@ QosFrameExchangeManager::PreProcessFrame(Ptr<const WifiPsdu> psdu, const WifiTxV
                 NS_LOG_DEBUG("Station " << hdr.GetAddr2() << " reported a buffer status of "
                                         << +hdr.GetQosQueueSize()
                                         << " for tid=" << +hdr.GetQosTid());
-                StaticCast<ApWifiMac>(m_mac)->SetBufferStatus(hdr.GetQosTid(),
-                                                              hdr.GetAddr2(),
-                                                              hdr.GetQosQueueSize());
+                StaticCast<ApWifiMac>(m_mac)->SetBufferStatus(
+                    hdr.GetQosTid(),
+                    mpdu->GetOriginal()->GetHeader().GetAddr2(),
+                    hdr.GetQosQueueSize());
             }
         }
     }