From a3ceae1ea71eef5a5dcd872097cab360949c8c77 Mon Sep 17 00:00:00 2001
From: Stefano Avallone <stavallo@unina.it>
Date: Wed, 7 Dec 2022 14:51:50 +0100
Subject: [PATCH] wifi: Add an AP MAC method to get the MLD or link address
 from the AID

---
 src/wifi/model/ap-wifi-mac.cc | 16 ++++++++++++++++
 src/wifi/model/ap-wifi-mac.h  | 11 +++++++++++
 2 files changed, 27 insertions(+)

diff --git a/src/wifi/model/ap-wifi-mac.cc b/src/wifi/model/ap-wifi-mac.cc
index 0649aa89f..01fe89019 100644
--- a/src/wifi/model/ap-wifi-mac.cc
+++ b/src/wifi/model/ap-wifi-mac.cc
@@ -1177,6 +1177,11 @@ ApWifiMac::SetAid(MgtAssocResponseHeader& assoc, const LinkIdStaAddrMap& linkIdS
         aid = GetNextAssociationId(linkIds);
     }
 
+    // store the MLD or link address in the AID-to-address map
+    const auto& [linkId, staAddr] = *linkIdStaAddrMap.cbegin();
+    m_aidToMldOrLinkAddress[aid] =
+        GetWifiRemoteStationManager(linkId)->GetMldAddress(staAddr).value_or(staAddr);
+
     for (const auto& [id, staAddr] : linkIdStaAddrMap)
     {
         auto remoteStationManager = GetWifiRemoteStationManager(id);
@@ -1531,6 +1536,17 @@ ApWifiMac::DoGetLocalAddress(const Mac48Address& remoteAddr) const
     return GetFrameExchangeManager(*linkId)->GetAddress();
 }
 
+std::optional<Mac48Address>
+ApWifiMac::GetMldOrLinkAddressByAid(uint16_t aid) const
+{
+    if (const auto staIt = m_aidToMldOrLinkAddress.find(aid);
+        staIt != m_aidToMldOrLinkAddress.cend())
+    {
+        return staIt->second;
+    }
+    return std::nullopt;
+}
+
 void
 ApWifiMac::Receive(Ptr<const WifiMpdu> mpdu, uint8_t linkId)
 {
diff --git a/src/wifi/model/ap-wifi-mac.h b/src/wifi/model/ap-wifi-mac.h
index 5f02b1f9e..c86923754 100644
--- a/src/wifi/model/ap-wifi-mac.h
+++ b/src/wifi/model/ap-wifi-mac.h
@@ -126,6 +126,13 @@ class ApWifiMac : public WifiMac
      */
     std::optional<uint8_t> IsAssociated(const Mac48Address& address) const;
 
+    /**
+     * \param aid the given AID
+     * \return the MLD address (in case of MLD) or link address (in case of single link device)
+     *         of the STA having the given AID, if any
+     */
+    std::optional<Mac48Address> GetMldOrLinkAddressByAid(uint16_t aid) const;
+
     /**
      * Return the value of the Queue Size subfield of the last QoS Data or QoS Null
      * frame received from the station with the given MAC address and belonging to
@@ -193,6 +200,10 @@ class ApWifiMac : public WifiMac
      */
     ApLinkEntity& GetLink(uint8_t linkId) const;
 
+    std::map<uint16_t, Mac48Address>
+        m_aidToMldOrLinkAddress; //!< Maps AIDs to MLD addresses (for MLDs) or link addresses (in
+                                 //!< case of single link devices)
+
   private:
     std::unique_ptr<LinkEntity> CreateLinkEntity() const override;
     Mac48Address DoGetLocalAddress(const Mac48Address& remoteAddr) const override;