wifi: Use the appropriate TA for BARs depending on whether RA is an MLD or link address

src/wifi/model/ap-wifi-mac.cc

@@ -1523,6 +1523,14 @@ ApWifiMac::IsAssociated(const Mac48Address& address) const
     return std::nullopt;
 }
 
+Mac48Address
+ApWifiMac::DoGetLocalAddress(const Mac48Address& remoteAddr) const
+{
+    auto linkId = IsAssociated(remoteAddr);
+    NS_ASSERT_MSG(linkId, remoteAddr << " is not associated");
+    return GetFrameExchangeManager(*linkId)->GetAddress();
+}
+
 void
 ApWifiMac::Receive(Ptr<const WifiMpdu> mpdu, uint8_t linkId)
 {

src/wifi/model/ap-wifi-mac.h

@@ -195,8 +195,9 @@ class ApWifiMac : public WifiMac
 
   private:
     std::unique_ptr<LinkEntity> CreateLinkEntity() const override;
-
+    Mac48Address DoGetLocalAddress(const Mac48Address& remoteAddr) const override;
     void Receive(Ptr<const WifiMpdu> mpdu, uint8_t linkId) override;
+
     /**
      * Check whether the supported rate set included in the received (Re)Association
      * Request frame is compatible with our Basic Rate Set. If so, record all the station's

src/wifi/model/qos-txop.cc

@@ -281,15 +281,17 @@ QosTxop::PrepareBlockAckRequest(Mac48Address recipient, uint8_t tid) const
     NS_LOG_FUNCTION(this << recipient << +tid);
     NS_ASSERT(QosUtilsMapTidToAc(tid) == m_ac);
 
+    auto recipientMld = m_mac->GetMldAddress(recipient);
+
     CtrlBAckRequestHeader reqHdr =
-        m_baManager->GetBlockAckReqHeader(m_mac->GetMldAddress(recipient).value_or(recipient), tid);
+        m_baManager->GetBlockAckReqHeader(recipientMld.value_or(recipient), tid);
     Ptr<Packet> bar = Create<Packet>();
     bar->AddHeader(reqHdr);
 
     WifiMacHeader hdr;
     hdr.SetType(WIFI_MAC_CTL_BACKREQ);
     hdr.SetAddr1(recipient);
-    hdr.SetAddr2(m_mac->GetAddress());
+    hdr.SetAddr2(m_mac->GetLocalAddress(recipient));
     hdr.SetDsNotTo();
     hdr.SetDsNotFrom();
     hdr.SetNoRetry();

src/wifi/model/sta-wifi-mac.cc

@@ -733,6 +733,15 @@ StaWifiMac::GetSetupLinkIds() const
     return linkIds;
 }
 
+Mac48Address
+StaWifiMac::DoGetLocalAddress(const Mac48Address& remoteAddr) const
+{
+    auto linkIds = GetSetupLinkIds();
+    NS_ASSERT_MSG(!linkIds.empty(), "Not associated");
+    uint8_t linkId = *linkIds.begin();
+    return GetFrameExchangeManager(linkId)->GetAddress();
+}
+
 bool
 StaWifiMac::CanForwardPacketsTo(Mac48Address to) const
 {

src/wifi/model/sta-wifi-mac.h

@@ -301,6 +301,7 @@ class StaWifiMac : public WifiMac
 
     void Receive(Ptr<const WifiMpdu> mpdu, uint8_t linkId) override;
     std::unique_ptr<LinkEntity> CreateLinkEntity() const override;
+    Mac48Address DoGetLocalAddress(const Mac48Address& remoteAddr) const override;
 
     /**
      * Process the Beacon frame received on the given link.

src/wifi/model/wifi-mac.cc

@@ -1244,6 +1244,40 @@ WifiMac::GetMldAddress(const Mac48Address& remoteAddr) const
     return std::nullopt;
 }
 
+Mac48Address
+WifiMac::GetLocalAddress(const Mac48Address& remoteAddr) const
+{
+    for (const auto& link : m_links)
+    {
+        if (auto mldAddress = link->stationManager->GetMldAddress(remoteAddr))
+        {
+            // this is a link setup with remote MLD
+            if (mldAddress != remoteAddr)
+            {
+                // the remote address is the address of a STA affiliated with the remote MLD
+                return link->feManager->GetAddress();
+            }
+            // we have to return our MLD address
+            return m_address;
+        }
+    }
+    // we get here if no ML setup was established between this device and the remote device,
+    // i.e., they are not both multi-link devices
+    if (GetNLinks() == 1)
+    {
+        // this is a single link device
+        return m_address;
+    }
+    // this is an MLD (hence the remote device is single link)
+    return DoGetLocalAddress(remoteAddr);
+}
+
+Mac48Address
+WifiMac::DoGetLocalAddress(const Mac48Address& remoteAddr [[maybe_unused]]) const
+{
+    return m_address;
+}
+
 WifiMac::OriginatorAgreementOptConstRef
 WifiMac::GetBaAgreementEstablishedAsOriginator(Mac48Address recipient, uint8_t tid) const
 {

src/wifi/model/wifi-mac.h

@@ -156,6 +156,22 @@ class WifiMac : public Object
      */
     std::optional<Mac48Address> GetMldAddress(const Mac48Address& remoteAddr) const;
 
+    /**
+     * Get the local MAC address used to communicate with a remote STA. Specifically:
+     * - If the given remote address is the address of a STA affiliated with a remote MLD
+     * and operating on a setup link, the address of the local STA operating on such a link
+     * is returned.
+     * - If the given remote address is the MLD address of a remote MLD (with which some link
+     * has been setup), the MLD address of this device is returned.
+     * - If this is a single link device, the unique MAC address of this device is returned.
+     * - Otherwise, return the MAC address of the affiliated STA (which must exists) that
+     * can be used to communicate with the remote device.
+     *
+     * \param remoteAddr the MAC address of the remote device
+     * \return the local MAC address used to communicate with the remote device
+     */
+    Mac48Address GetLocalAddress(const Mac48Address& remoteAddr) const;
+
     /**
      * Accessor for the Txop object
      *
@@ -787,6 +803,17 @@ class WifiMac : public Object
      */
     virtual std::unique_ptr<LinkEntity> CreateLinkEntity() const;
 
+    /**
+     * This method is called if this device is an MLD to determine the MAC address of
+     * the affiliated STA used to communicate with the single link device having the
+     * given MAC address. This method is overridden because its implementation depends
+     * on the type of station.
+     *
+     * \param remoteAddr the MAC address of the remote single link device
+     * \return the MAC address of the affiliated STA used to communicate with the remote device
+     */
+    virtual Mac48Address DoGetLocalAddress(const Mac48Address& remoteAddr) const;
+
     /**
      * Enable or disable ERP support for the given link.
      *