From c0e535c7dfb5854fc0c7cfda797f378e95b1130a Mon Sep 17 00:00:00 2001
From: Stefano Avallone <stavallo@unina.it>
Date: Wed, 5 Oct 2022 13:34:30 +0200
Subject: [PATCH] wifi: A non-AP MLD associated with a single link AP
 advertises the link MAC address

---
 src/wifi/model/wifi-net-device.cc | 29 +++++++++++++++++++++++++++--
 1 file changed, 27 insertions(+), 2 deletions(-)

diff --git a/src/wifi/model/wifi-net-device.cc b/src/wifi/model/wifi-net-device.cc
index ba6a7b46a..60a24d9e4 100644
--- a/src/wifi/model/wifi-net-device.cc
+++ b/src/wifi/model/wifi-net-device.cc
@@ -19,7 +19,8 @@
 
 #include "wifi-net-device.h"
 
-#include "wifi-mac.h"
+#include "frame-exchange-manager.h"
+#include "sta-wifi-mac.h"
 #include "wifi-phy.h"
 
 #include "ns3/channel.h"
@@ -374,6 +375,30 @@ WifiNetDevice::SetAddress(Address address)
 Address
 WifiNetDevice::GetAddress() const
 {
+    Ptr<StaWifiMac> staMac;
+    std::set<uint8_t> linkIds;
+
+    /**
+     * Normally, the MAC address that the network device has to advertise to upper layers is
+     * the MLD address, if this device is an MLD, or the unique MAC address, otherwise.
+     * Advertising the MAC address returned by WifiMac::GetAddress() is therefore the right
+     * thing to do in both cases. However, there is an exception: if this device is a non-AP MLD
+     * associated with a single link AP (hence, no ML setup was done), we need to advertise the
+     * MAC address of the link used to communicate with the AP. In fact, if we advertised the
+     * MLD address, the AP could not forward a frame to us because it would not recognize our
+     * MLD address as the MAC address of an associated station.
+     */
+
+    // Handle the exception first
+    if (m_mac->GetTypeOfStation() == STA &&
+        (staMac = StaticCast<StaWifiMac>(m_mac))->IsAssociated() && m_mac->GetNLinks() > 1 &&
+        (linkIds = staMac->GetSetupLinkIds()).size() == 1 &&
+        !GetRemoteStationManager(*linkIds.begin())
+             ->GetMldAddress(m_mac->GetBssid(*linkIds.begin())))
+    {
+        return m_mac->GetFrameExchangeManager(*linkIds.begin())->GetAddress();
+    }
+
     return m_mac->GetAddress();
 }
 
@@ -504,7 +529,7 @@ WifiNetDevice::ForwardUp(Ptr<const Packet> packet, Mac48Address from, Mac48Addre
     {
         type = NetDevice::PACKET_MULTICAST;
     }
-    else if (to == m_mac->GetAddress())
+    else if (to == GetAddress())
     {
         type = NetDevice::PACKET_HOST;
     }