From 1a8b37184f74a08e84dec35235ac53ae7aae01f4 Mon Sep 17 00:00:00 2001
From: Stefano Avallone <stavallo@unina.it>
Date: Mon, 20 Mar 2023 10:19:06 +0100
Subject: [PATCH] wifi: Ensure mgt and ctrl frames are protected if needed

---
 .../model/ht/ht-frame-exchange-manager.cc     | 47 +++++++------------
 src/wifi/model/ht/ht-frame-exchange-manager.h |  8 +++-
 2 files changed, 24 insertions(+), 31 deletions(-)

diff --git a/src/wifi/model/ht/ht-frame-exchange-manager.cc b/src/wifi/model/ht/ht-frame-exchange-manager.cc
index 77a359460..4ea4465f7 100644
--- a/src/wifi/model/ht/ht-frame-exchange-manager.cc
+++ b/src/wifi/model/ht/ht-frame-exchange-manager.cc
@@ -133,14 +133,16 @@ HtFrameExchangeManager::NeedSetupBlockAck(Mac48Address recipient, uint8_t tid)
     return establish;
 }
 
-void
+bool
 HtFrameExchangeManager::SendAddBaRequest(Mac48Address dest,
                                          uint8_t tid,
                                          uint16_t startingSeq,
                                          uint16_t timeout,
-                                         bool immediateBAck)
+                                         bool immediateBAck,
+                                         Time availableTime)
 {
-    NS_LOG_FUNCTION(this << dest << +tid << startingSeq << timeout << immediateBAck);
+    NS_LOG_FUNCTION(this << dest << +tid << startingSeq << timeout << immediateBAck
+                         << availableTime);
     NS_LOG_DEBUG("Send ADDBA request to " << dest);
 
     WifiMacHeader hdr;
@@ -193,12 +195,16 @@ HtFrameExchangeManager::SendAddBaRequest(Mac48Address dest,
     WifiTxParameters txParams;
     txParams.m_txVector =
         GetWifiRemoteStationManager()->GetDataTxVector(mpdu->GetHeader(), m_allowedWidth);
-    txParams.m_protection = std::unique_ptr<WifiProtection>(new WifiNoProtection);
-    txParams.m_acknowledgment = GetAckManager()->TryAddMpdu(mpdu, txParams);
+    if (!TryAddMpdu(mpdu, txParams, availableTime))
+    {
+        NS_LOG_DEBUG("Not enough time to send the ADDBA Request frame");
+        return false;
+    }
 
     // Wifi MAC queue scheduler is expected to prioritize management frames
     m_mac->GetQosTxop(tid)->GetWifiMacQueue()->Enqueue(mpdu);
     SendMpduWithProtection(mpdu, txParams);
+    return true;
 }
 
 void
@@ -381,12 +387,12 @@ HtFrameExchangeManager::StartFrameExchange(Ptr<QosTxop> edca, Time availableTime
             (hdr.IsRetry()
                  ? hdr.GetSequenceNumber()
                  : m_txMiddle->GetNextSeqNumberByTidAndAddress(hdr.GetQosTid(), hdr.GetAddr1()));
-        SendAddBaRequest(hdr.GetAddr1(),
-                         hdr.GetQosTid(),
-                         startingSeq,
-                         edca->GetBlockAckInactivityTimeout(),
-                         true);
-        return true;
+        return SendAddBaRequest(hdr.GetAddr1(),
+                                hdr.GetQosTid(),
+                                startingSeq,
+                                edca->GetBlockAckInactivityTimeout(),
+                                true,
+                                availableTime);
     }
 
     // Use SendDataFrame if we can try aggregation
@@ -570,25 +576,8 @@ HtFrameExchangeManager::SendMpduFromBaManager(Ptr<WifiMpdu> mpdu,
     WifiTxParameters txParams;
     txParams.m_txVector =
         GetWifiRemoteStationManager()->GetDataTxVector(mpdu->GetHeader(), m_allowedWidth);
-    txParams.m_protection = std::unique_ptr<WifiProtection>(new WifiNoProtection);
-    txParams.m_acknowledgment = GetAckManager()->TryAddMpdu(mpdu, txParams);
 
-    NS_ABORT_IF(txParams.m_acknowledgment->method != WifiAcknowledgment::BLOCK_ACK);
-
-    WifiBlockAck* blockAcknowledgment = static_cast<WifiBlockAck*>(txParams.m_acknowledgment.get());
-    CalculateAcknowledgmentTime(blockAcknowledgment);
-    // the BlockAckReq frame is sent using the same TXVECTOR as the BlockAck frame
-    txParams.m_txVector = blockAcknowledgment->blockAckTxVector;
-
-    Time barTxDuration = m_phy->CalculateTxDuration(mpdu->GetSize(),
-                                                    blockAcknowledgment->blockAckTxVector,
-                                                    m_phy->GetPhyBand());
-
-    // if the available time is limited and we are not transmitting the initial
-    // frame of the TXOP, we have to check that this frame and its response fit
-    // within the given time limits
-    if (availableTime != Time::Min() && !initialFrame &&
-        barTxDuration + m_phy->GetSifs() + blockAcknowledgment->acknowledgmentTime > availableTime)
+    if (!TryAddMpdu(mpdu, txParams, availableTime))
     {
         NS_LOG_DEBUG("Not enough time to send the BAR frame returned by the Block Ack Manager");
         return false;
diff --git a/src/wifi/model/ht/ht-frame-exchange-manager.h b/src/wifi/model/ht/ht-frame-exchange-manager.h
index 21898681d..0264eb926 100644
--- a/src/wifi/model/ht/ht-frame-exchange-manager.h
+++ b/src/wifi/model/ht/ht-frame-exchange-manager.h
@@ -317,12 +317,16 @@ class HtFrameExchangeManager : public QosFrameExchangeManager
      * \param startingSeq the BA agreement starting sequence number
      * \param timeout timeout value.
      * \param immediateBAck flag to indicate whether immediate BlockAck is used.
+     * \param availableTime the amount of time allowed for the frame exchange. Equals
+     *                      Time::Min() in case the TXOP limit is null
+     * \return true if ADDBA Request frame is transmitted, false otherwise
      */
-    void SendAddBaRequest(Mac48Address recipient,
+    bool SendAddBaRequest(Mac48Address recipient,
                           uint8_t tid,
                           uint16_t startingSeq,
                           uint16_t timeout,
-                          bool immediateBAck);
+                          bool immediateBAck,
+                          Time availableTime);
 
     /**
      * Create a BlockAck frame with header equal to <i>blockAck</i> and start its transmission.