From bdfa41118385b025f0ae4eadab0c6d761844a1f2 Mon Sep 17 00:00:00 2001
From: Stefano Avallone <stavallo@unina.it>
Date: Tue, 17 Jan 2023 13:20:44 +0100
Subject: [PATCH] wifi: Adjust HePhy and HePpdu to support EHT MU PPDUs

EHT SU transmissions are carried in EHT MU PPDUs, which have the EHT-SIG field.
Makes no difference for HE PPDUs.
---
 src/wifi/model/he/he-phy.cc  | 18 ++++++++++++++----
 src/wifi/model/he/he-ppdu.cc |  4 ++--
 2 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/src/wifi/model/he/he-phy.cc b/src/wifi/model/he/he-phy.cc
index f996a273c..360d69c27 100644
--- a/src/wifi/model/he/he-phy.cc
+++ b/src/wifi/model/he/he-phy.cc
@@ -217,7 +217,7 @@ HePhy::GetSigADuration(WifiPreamble preamble) const
 Time
 HePhy::GetSigBDuration(const WifiTxVector& txVector) const
 {
-    if (txVector.IsDlMu()) // See section 27.3.10.8 of IEEE 802.11ax draft 4.0.
+    if (ns3::IsDlMu(txVector.GetPreambleType())) // See section 27.3.11.8 of IEEE 802.11ax-2021
     {
         NS_ASSERT(txVector.GetModulationClass() >= WIFI_MOD_CLASS_HE);
 
@@ -707,8 +707,18 @@ HePhy::ProcessSigB(Ptr<Event> event, PhyFieldRxStatus status)
             return PhyFieldRxStatus(false, FILTERED, DROP);
         }
     }
-    m_currentMuPpduUid =
-        event->GetPpdu()->GetUid(); // to be able to correctly schedule start of OFDMA payload
+    if (event->GetTxVector().IsDlMu())
+    {
+        // When including a Trigger Frame, a DL MU PPDU solicits a TB PPDU.
+        // NOTE that the 'if' condition above is not needed for HE because SIG-B is only
+        // included in HE MU PPDUs, but it is necessary for EHT to avoid that a non-AP
+        // STA receiving a Trigger Frame sent as an EHT SU transmission (which carries
+        // the EHT-SIG field) stores the PPDU UID and uses it later to schedule the
+        // reception of the OFDMA payload of the TB PPDU (see HePhy::StartReceivePreamble())
+        // despite it lacks the TRIGVECTOR.
+        m_currentMuPpduUid =
+            event->GetPpdu()->GetUid(); // to be able to correctly schedule start of OFDMA payload
+    }
     return status;
 }
 
@@ -1713,7 +1723,7 @@ uint32_t
 HePhy::GetSigBFieldSize(const WifiTxVector& txVector)
 {
     NS_ASSERT(txVector.GetModulationClass() >= WIFI_MOD_CLASS_HE);
-    NS_ASSERT(txVector.IsDlMu());
+    NS_ASSERT(ns3::IsDlMu(txVector.GetPreambleType()));
 
     // Compute the number of bits used by common field.
     // Assume that compression bit in HE-SIG-A is not set (i.e. not
diff --git a/src/wifi/model/he/he-ppdu.cc b/src/wifi/model/he/he-ppdu.cc
index 6ddc30f06..63b58c846 100644
--- a/src/wifi/model/he/he-ppdu.cc
+++ b/src/wifi/model/he/he-ppdu.cc
@@ -123,12 +123,12 @@ HePpdu::SetPhyHeaders(const WifiTxVector& txVector, Time ppduDuration)
                         3) -
                        3 - m);
     m_lSig.SetLength(length);
-    if (txVector.IsDlMu())
+    if (ns3::IsDlMu(m_preamble))
     {
         m_heSig.SetMuFlag(true);
         m_heSig.SetMcs(txVector.GetSigBMode().GetMcsValue());
     }
-    else if (!txVector.IsUlMu())
+    else if (!ns3::IsUlMu(m_preamble))
     {
         m_heSig.SetMcs(txVector.GetMode().GetMcsValue());
         m_heSig.SetNStreams(txVector.GetNss());