From bb7d590e8bb6214fe4164611fd52c63da50d8525 Mon Sep 17 00:00:00 2001
From: Stefano Avallone <stavallo@unina.it>
Date: Fri, 9 Jul 2021 12:04:37 +0200
Subject: [PATCH] wifi: Store the MLD MAC Address in WifiRemoteStationState
 objects

---
 src/wifi/model/wifi-remote-station-manager.cc | 33 +++++++++++++++++++
 src/wifi/model/wifi-remote-station-manager.h  | 29 ++++++++++++++++
 2 files changed, 62 insertions(+)

diff --git a/src/wifi/model/wifi-remote-station-manager.cc b/src/wifi/model/wifi-remote-station-manager.cc
index c75e71217..6303fe253 100644
--- a/src/wifi/model/wifi-remote-station-manager.cc
+++ b/src/wifi/model/wifi-remote-station-manager.cc
@@ -517,6 +517,39 @@ WifiRemoteStationManager::GetStaId (Mac48Address address, const WifiTxVector& tx
   return staId;
 }
 
+void
+WifiRemoteStationManager::SetMldAddress (const Mac48Address& address, const Mac48Address& mldAddress)
+{
+  NS_LOG_FUNCTION (this << address << mldAddress);
+
+  WifiRemoteStationState *state = LookupState (address);
+  state->m_mldAddress = mldAddress;
+  // insert another entry in m_states indexed by the MLD address and pointing to the same state
+  const_cast<WifiRemoteStationManager *> (this)->m_states.insert ({mldAddress, state});
+}
+
+std::optional<Mac48Address>
+WifiRemoteStationManager::GetMldAddress (const Mac48Address& address) const
+{
+  return LookupState (address)->m_mldAddress;
+}
+
+std::optional<Mac48Address>
+WifiRemoteStationManager::GetAffiliatedStaAddress (const Mac48Address& mldAddress) const
+{
+  auto stateIt = m_states.find (mldAddress);
+
+  if (stateIt == m_states.end ())
+    {
+      // MLD address not found
+      return std::optional<Mac48Address> ();
+    }
+
+  NS_ASSERT (stateIt->second->m_mldAddress.has_value ()
+             && *stateIt->second->m_mldAddress == mldAddress);
+  return stateIt->second->m_address;
+}
+
 WifiTxVector
 WifiRemoteStationManager::GetDataTxVector (const WifiMacHeader &header, uint16_t allowedWidth)
 {
diff --git a/src/wifi/model/wifi-remote-station-manager.h b/src/wifi/model/wifi-remote-station-manager.h
index 3e948828b..5586531f3 100644
--- a/src/wifi/model/wifi-remote-station-manager.h
+++ b/src/wifi/model/wifi-remote-station-manager.h
@@ -22,6 +22,7 @@
 #define WIFI_REMOTE_STATION_MANAGER_H
 
 #include <array>
+#include <optional>
 #include <unordered_map>
 #include "ns3/traced-callback.h"
 #include "ns3/object.h"
@@ -96,6 +97,8 @@ struct WifiRemoteStationState
   Mac48Address m_address;            //!< Mac48Address of the remote station
   uint16_t m_aid;                    /**< AID of the remote station (unused if this object
                                           is installed on a non-AP station) */
+  std::optional<Mac48Address> m_mldAddress; /**< MLD MAC address if the remote station is
+                                                 affiliated with an MLD */
   WifiRemoteStationInfo m_info;      //!< remote station info
   bool m_dsssSupported;              //!< Flag if DSSS is supported by the remote station
   bool m_erpOfdmSupported;           //!< Flag if ERP OFDM is supported by the remote station
@@ -708,6 +711,32 @@ public:
    */
   void RecordDisassociated (Mac48Address address);
 
+  /**
+   * Set the address of the MLD the given station is affiliated with.
+   *
+   * \param address the MAC address of the remote station
+   * \param mldAddress the MAC address of the MLD the remote station is
+   *                   affiliated with
+   */
+  void SetMldAddress (const Mac48Address& address, const Mac48Address& mldAddress);
+  /**
+   * Get the address of the MLD the given station is affiliated with, if any,
+   * or the broadcast address, otherwise.
+   *
+   * \param address the MAC address of the remote station
+   * \return the address of the MLD the given station is affiliated with, if any
+   */
+  std::optional<Mac48Address> GetMldAddress (const Mac48Address& address) const;
+  /**
+   * Get the address of the remote station operating on this link and affiliated
+   * with the MLD having the given MAC address, if any.
+   *
+   * \param mldAddress the MLD MAC address
+   * \return the address of the remote station operating on this link and
+   *         affiliated with the MLD, if any
+   */
+  std::optional<Mac48Address> GetAffiliatedStaAddress (const Mac48Address& mldAddress) const;
+
   /**
    * \param header MAC header
    * \param allowedWidth the allowed width in MHz to send this packet