From 1f4b4cbfd67c43a5e2d3c6a9b22a8d510f9e9d81 Mon Sep 17 00:00:00 2001
From: Stefano Avallone <stavallo@unina.it>
Date: Wed, 10 May 2023 17:59:12 +0200
Subject: [PATCH] wifi: Add StaWifiMac methods to block transmissions on a link

---
 src/wifi/model/sta-wifi-mac.cc | 44 ++++++++++++++++++++++++++++++++++
 src/wifi/model/sta-wifi-mac.h  | 16 +++++++++++++
 2 files changed, 60 insertions(+)

diff --git a/src/wifi/model/sta-wifi-mac.cc b/src/wifi/model/sta-wifi-mac.cc
index c735a19ca..0d454d62f 100644
--- a/src/wifi/model/sta-wifi-mac.cc
+++ b/src/wifi/model/sta-wifi-mac.cc
@@ -1050,6 +1050,50 @@ StaWifiMac::Enqueue(Ptr<Packet> packet, Mac48Address to)
     }
 }
 
+void
+StaWifiMac::BlockTxOnLink(uint8_t linkId, WifiQueueBlockedReason reason)
+{
+    NS_LOG_FUNCTION(this << linkId << reason);
+
+    auto bssid = GetBssid(linkId);
+    auto apAddress = GetWifiRemoteStationManager(linkId)->GetMldAddress(bssid).value_or(bssid);
+
+    BlockUnicastTxOnLinks(reason, apAddress, {linkId});
+    // the only type of broadcast frames that a non-AP STA can send are management frames
+    for (const auto [acIndex, ac] : wifiAcList)
+    {
+        GetMacQueueScheduler()->BlockQueues(reason,
+                                            acIndex,
+                                            {WIFI_MGT_QUEUE},
+                                            Mac48Address::GetBroadcast(),
+                                            GetFrameExchangeManager(linkId)->GetAddress(),
+                                            {},
+                                            {linkId});
+    }
+}
+
+void
+StaWifiMac::UnblockTxOnLink(uint8_t linkId, WifiQueueBlockedReason reason)
+{
+    NS_LOG_FUNCTION(this << linkId << reason);
+
+    auto bssid = GetBssid(linkId);
+    auto apAddress = GetWifiRemoteStationManager(linkId)->GetMldAddress(bssid).value_or(bssid);
+
+    UnblockUnicastTxOnLinks(reason, apAddress, {linkId});
+    // the only type of broadcast frames that a non-AP STA can send are management frames
+    for (const auto [acIndex, ac] : wifiAcList)
+    {
+        GetMacQueueScheduler()->UnblockQueues(reason,
+                                              acIndex,
+                                              {WIFI_MGT_QUEUE},
+                                              Mac48Address::GetBroadcast(),
+                                              GetFrameExchangeManager(linkId)->GetAddress(),
+                                              {},
+                                              {linkId});
+    }
+}
+
 void
 StaWifiMac::Receive(Ptr<const WifiMpdu> mpdu, uint8_t linkId)
 {
diff --git a/src/wifi/model/sta-wifi-mac.h b/src/wifi/model/sta-wifi-mac.h
index 4e94c8753..43c762c9b 100644
--- a/src/wifi/model/sta-wifi-mac.h
+++ b/src/wifi/model/sta-wifi-mac.h
@@ -316,6 +316,22 @@ class StaWifiMac : public WifiMac
      */
     void NotifySwitchingEmlsrLink(Ptr<WifiPhy> phy, uint8_t linkId);
 
+    /**
+     * Block transmissions on the given link for the given reason.
+     *
+     * \param linkId the ID of the given link
+     * \param reason the reason for blocking transmissions on the given link
+     */
+    void BlockTxOnLink(uint8_t linkId, WifiQueueBlockedReason reason);
+
+    /**
+     * Unblock transmissions on the given link for the given reason.
+     *
+     * \param linkId the ID of the given link
+     * \param reason the reason for unblocking transmissions on the given link
+     */
+    void UnblockTxOnLink(uint8_t linkId, WifiQueueBlockedReason reason);
+
     /**
      * Assign a fixed random variable stream number to the random variables
      * used by this model.  Return the number of streams (possibly zero) that