wifi: Add function in WifiMac to store capabilities read from information elements

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

@@ -149,10 +149,6 @@ class StaWifiMac : public WifiMac
     friend class ::MultiLinkOperationsTestBase;
     friend class ::ProbeExchTest;
 
-    /// type of the management frames used to get info about APs
-    using MgtFrameType =
-        std::variant<MgtBeaconHeader, MgtProbeResponseHeader, MgtAssocResponseHeader>;
-
     /**
      * Struct to hold information regarding observed AP through
      * active/passive scanning

src/wifi/model/wifi-mac.cc

@@ -2590,4 +2590,91 @@ WifiMac::GetRobustAVStreamingSupported() const
     return m_robustAVStreamingSupported;
 }
 
+void
+WifiMac::RecordCapabilities(const MgtFrameType& frame, const Mac48Address& from, uint8_t linkId)
+{
+    NS_LOG_FUNCTION(this << frame.index() << from << linkId);
+    auto remoteStationManager = GetWifiRemoteStationManager(linkId);
+    auto phy = GetWifiPhy(linkId);
+
+    // lambda processing Information Elements included in all frame types
+    auto recordFromIes = [&](auto&& frame) {
+        NS_ASSERT(frame.template Get<SupportedRates>());
+        const auto rates = AllSupportedRates{*frame.template Get<SupportedRates>(),
+                                             frame.template Get<ExtendedSupportedRatesIE>()};
+        for (const auto& mode : phy->GetModeList())
+        {
+            if (rates.IsSupportedRate(mode.GetDataRate(phy->GetChannelWidth())))
+            {
+                remoteStationManager->AddSupportedMode(from, mode);
+                if (rates.IsBasicRate(mode.GetDataRate(phy->GetChannelWidth())))
+                {
+                    remoteStationManager->AddBasicMode(mode);
+                }
+            }
+        }
+
+        // we do not return if HT is not supported because HE STAs operating in
+        // the 6 GHz band do not support VHT
+        if (GetHtSupported(linkId))
+        {
+            /* HT station */
+            if (const auto& htCapabilities = frame.template Get<HtCapabilities>())
+            {
+                remoteStationManager->AddStationHtCapabilities(from, *htCapabilities);
+            }
+            if (const auto& extendedCapabilities = frame.template Get<ExtendedCapabilities>())
+            {
+                remoteStationManager->AddStationExtendedCapabilities(from, *extendedCapabilities);
+            }
+        }
+
+        // we do not return if VHT is not supported because HE STAs operating in
+        // the 2.4 GHz or 6 GHz band do not support VHT
+        if (GetVhtSupported(linkId))
+        {
+            const auto& vhtCapabilities = frame.template Get<VhtCapabilities>();
+            // we will always fill in RxHighestSupportedLgiDataRate field at TX, so this can be used
+            // to check whether it supports VHT
+            if (vhtCapabilities.has_value() &&
+                vhtCapabilities->GetRxHighestSupportedLgiDataRate() > 0)
+            {
+                remoteStationManager->AddStationVhtCapabilities(from, *vhtCapabilities);
+            }
+        }
+
+        if (!GetHeSupported())
+        {
+            return;
+        }
+
+        const auto& heCapabilities = frame.template Get<HeCapabilities>();
+        if (heCapabilities.has_value() && heCapabilities->GetSupportedMcsAndNss() != 0)
+        {
+            remoteStationManager->AddStationHeCapabilities(from, *heCapabilities);
+        }
+
+        if (Is6GhzBand(linkId))
+        {
+            if (const auto& he6GhzCapabilities = frame.template Get<He6GhzBandCapabilities>())
+            {
+                remoteStationManager->AddStationHe6GhzCapabilities(from, *he6GhzCapabilities);
+            }
+        }
+
+        if (!GetEhtSupported())
+        {
+            return;
+        }
+
+        if (const auto& ehtCapabilities = frame.template Get<EhtCapabilities>())
+        {
+            remoteStationManager->AddStationEhtCapabilities(from, *ehtCapabilities);
+        }
+    };
+
+    // process Information Elements included in the current frame variant
+    std::visit(recordFromIes, frame);
+}
+
 } // namespace ns3

src/wifi/model/wifi-mac.h

@@ -102,6 +102,10 @@ class WifiMac : public Object
     WifiMac(const WifiMac&) = delete;
     WifiMac& operator=(const WifiMac&) = delete;
 
+    /// type of the management frames
+    using MgtFrameType =
+        std::variant<MgtBeaconHeader, MgtProbeResponseHeader, MgtAssocResponseHeader>;
+
     /**
      * Assign a fixed random variable stream number to the random variables
      * used by this model.  Return the number of streams (possibly zero) that
@@ -983,6 +987,15 @@ class WifiMac : public Object
                                 WifiDirection dir,
                                 const WifiTidLinkMapping& mapping);
 
+    /**
+     * Update capabilities information from the given management frame.
+     *
+     * @param frame the body of the given management frame
+     * @param addr MAC address of the sender
+     * @param linkId ID of the link the management frame was received over
+     */
+    void RecordCapabilities(const MgtFrameType& frame, const Mac48Address& addr, uint8_t linkId);
+
     UniformRandomBitGenerator m_shuffleLinkIdsGen; //!< random number generator to shuffle link IDs
 
     Ptr<MacRxMiddle> m_rxMiddle; //!< RX middle (defragmentation etc.)