diff --git a/src/wifi/model/sta-wifi-mac.cc b/src/wifi/model/sta-wifi-mac.cc
index 3c6452340..34cb875cf 100644
--- a/src/wifi/model/sta-wifi-mac.cc
+++ b/src/wifi/model/sta-wifi-mac.cc
@@ -714,6 +714,25 @@ StaWifiMac::IsWaitAssocResp() const
     return m_state == WAIT_ASSOC_RESP;
 }
 
+std::set<uint8_t>
+StaWifiMac::GetSetupLinkIds() const
+{
+    if (!IsAssociated())
+    {
+        return {};
+    }
+
+    std::set<uint8_t> linkIds;
+    for (uint8_t linkId = 0; linkId < GetNLinks(); linkId++)
+    {
+        if (GetLink(linkId).bssid)
+        {
+            linkIds.insert(linkId);
+        }
+    }
+    return linkIds;
+}
+
 bool
 StaWifiMac::CanForwardPacketsTo(Mac48Address to) const
 {
@@ -826,7 +845,12 @@ StaWifiMac::Receive(Ptr<const WifiMpdu> mpdu, uint8_t linkId)
             NotifyRxDrop(packet);
             return;
         }
-        if (hdr->GetAddr2() != GetBssid(0)) // TODO use appropriate linkId
+        std::set<Mac48Address> apAddresses; // link addresses of AP
+        for (auto id : GetSetupLinkIds())
+        {
+            apAddresses.insert(GetBssid(id));
+        }
+        if (apAddresses.count(mpdu->GetHeader().GetAddr2()) == 0)
         {
             NS_LOG_LOGIC("Received data frame not from the BSS we are associated with: ignore");
             NotifyRxDrop(packet);
@@ -836,7 +860,7 @@ StaWifiMac::Receive(Ptr<const WifiMpdu> mpdu, uint8_t linkId)
         {
             if (hdr->IsQosAmsdu())
             {
-                NS_ASSERT(hdr->GetAddr3() == GetBssid(0)); // TODO use appropriate linkId
+                NS_ASSERT(apAddresses.count(mpdu->GetHeader().GetAddr3()) != 0);
                 DeaggregateAmsduAndForward(mpdu);
                 packet = nullptr;
             }
diff --git a/src/wifi/model/sta-wifi-mac.h b/src/wifi/model/sta-wifi-mac.h
index b2b1d5703..66a15c8e7 100644
--- a/src/wifi/model/sta-wifi-mac.h
+++ b/src/wifi/model/sta-wifi-mac.h
@@ -25,6 +25,7 @@
 #include "mgt-headers.h"
 #include "wifi-mac.h"
 
+#include <set>
 #include <variant>
 
 class TwoLevelAggregationTest;
@@ -202,6 +203,13 @@ class StaWifiMac : public WifiMac
      */
     bool IsAssociated() const;
 
+    /**
+     * Get the IDs of the setup links (if any).
+     *
+     * \return the IDs of the setup links
+     */
+    std::set<uint8_t> GetSetupLinkIds() const;
+
     /**
      * Return the association ID.
      *