From a7dd40b5673553d950b9d72e237367bd305ca4e0 Mon Sep 17 00:00:00 2001
From: Stefano Avallone <stavallo@unina.it>
Date: Wed, 12 Apr 2023 09:53:32 +0200
Subject: [PATCH] wifi: Pass reference to containing frame to MultiLinkElement

---
 src/wifi/model/eht/multi-link-element.cc |  9 +++----
 src/wifi/model/eht/multi-link-element.h  | 14 +++++++++--
 src/wifi/model/mgt-headers.cc            | 30 ++++++++++++++++++++++++
 3 files changed, 47 insertions(+), 6 deletions(-)

diff --git a/src/wifi/model/eht/multi-link-element.cc b/src/wifi/model/eht/multi-link-element.cc
index b00be838a..d64aab47e 100644
--- a/src/wifi/model/eht/multi-link-element.cc
+++ b/src/wifi/model/eht/multi-link-element.cc
@@ -222,14 +222,15 @@ CommonInfoBasicMle::DecodeEmlsrTransitionDelay(uint8_t value)
 /**
  * MultiLinkElement
  */
-MultiLinkElement::MultiLinkElement(WifiMacType frameType)
-    : m_frameType(frameType),
+MultiLinkElement::MultiLinkElement(WifiMacType frameType, ContainingFrame frame)
+    : m_containingFrame(frame),
+      m_frameType(frameType),
       m_commonInfo(std::in_place_type<std::monostate>) // initialize as UNSET
 {
 }
 
-MultiLinkElement::MultiLinkElement(Variant variant, WifiMacType frameType)
-    : MultiLinkElement(frameType)
+MultiLinkElement::MultiLinkElement(Variant variant, WifiMacType frameType, ContainingFrame frame)
+    : MultiLinkElement(frameType, frame)
 {
     NS_ASSERT(variant != UNSET);
     SetVariant(variant);
diff --git a/src/wifi/model/eht/multi-link-element.h b/src/wifi/model/eht/multi-link-element.h
index f93e5ebe6..b917e33ea 100644
--- a/src/wifi/model/eht/multi-link-element.h
+++ b/src/wifi/model/eht/multi-link-element.h
@@ -184,19 +184,27 @@ class MultiLinkElement : public WifiInformationElement
         PER_STA_PROFILE_SUBELEMENT_ID = 0
     };
 
+    /// Typedef for structure holding a reference to the containing frame
+    using ContainingFrame = std::variant<std::monostate,
+                                         std::reference_wrapper<const MgtAssocRequestHeader>,
+                                         std::reference_wrapper<const MgtReassocRequestHeader>,
+                                         std::reference_wrapper<const MgtAssocResponseHeader>>;
+
     /**
      * Construct a Multi-Link Element with no variant set.
      *
      * \param frameType the type of the frame containing the Multi-Link Element
+     * \param frame the management frame containing this Multi-Link Element
      */
-    MultiLinkElement(WifiMacType frameType);
+    MultiLinkElement(WifiMacType frameType, ContainingFrame frame = {});
     /**
      * Constructor
      *
      * \param variant the Multi-Link element variant (cannot be UNSET)
      * \param frameType the type of the frame containing the Multi-Link Element
+     * \param frame the management frame containing this Multi-Link Element
      */
-    MultiLinkElement(Variant variant, WifiMacType frameType);
+    MultiLinkElement(Variant variant, WifiMacType frameType, ContainingFrame frame = {});
 
     WifiInformationElementId ElementId() const override;
     WifiInformationElementId ElementIdExt() const override;
@@ -400,6 +408,8 @@ class MultiLinkElement : public WifiInformationElement
      */
     Time GetTransitionTimeout() const;
 
+    mutable ContainingFrame m_containingFrame; //!< reference to the mgt frame containing this MLE
+
     /**
      * \ingroup wifi
      * Per-STA Profile Subelement of Multi-Link element.
diff --git a/src/wifi/model/mgt-headers.cc b/src/wifi/model/mgt-headers.cc
index 8ee803a87..2ee5a5556 100644
--- a/src/wifi/model/mgt-headers.cc
+++ b/src/wifi/model/mgt-headers.cc
@@ -210,6 +210,11 @@ MgtAssocRequestHeader::Capabilities()
 uint32_t
 MgtAssocRequestHeader::GetSerializedSizeImpl() const
 {
+    if (auto& mle = Get<MultiLinkElement>())
+    {
+        mle->m_containingFrame = *this;
+    }
+
     uint32_t size = 0;
     size += m_capability.GetSerializedSize();
     size += 2; // listen interval
@@ -220,6 +225,11 @@ MgtAssocRequestHeader::GetSerializedSizeImpl() const
 void
 MgtAssocRequestHeader::SerializeImpl(Buffer::Iterator start) const
 {
+    if (auto& mle = Get<MultiLinkElement>())
+    {
+        mle->m_containingFrame = *this;
+    }
+
     Buffer::Iterator i = start;
     i = m_capability.Serialize(i);
     i.WriteHtolsbU16(m_listenInterval);
@@ -297,6 +307,11 @@ MgtReassocRequestHeader::SetCurrentApAddress(Mac48Address currentApAddr)
 uint32_t
 MgtReassocRequestHeader::GetSerializedSizeImpl() const
 {
+    if (auto& mle = Get<MultiLinkElement>())
+    {
+        mle->m_containingFrame = *this;
+    }
+
     uint32_t size = 0;
     size += m_capability.GetSerializedSize();
     size += 2; // listen interval
@@ -315,6 +330,11 @@ MgtReassocRequestHeader::PrintImpl(std::ostream& os) const
 void
 MgtReassocRequestHeader::SerializeImpl(Buffer::Iterator start) const
 {
+    if (auto& mle = Get<MultiLinkElement>())
+    {
+        mle->m_containingFrame = *this;
+    }
+
     Buffer::Iterator i = start;
     i = m_capability.Serialize(i);
     i.WriteHtolsbU16(m_listenInterval);
@@ -400,6 +420,11 @@ MgtAssocResponseHeader::GetAssociationId() const
 uint32_t
 MgtAssocResponseHeader::GetSerializedSizeImpl() const
 {
+    if (auto& mle = Get<MultiLinkElement>())
+    {
+        mle->m_containingFrame = *this;
+    }
+
     uint32_t size = 0;
     size += m_capability.GetSerializedSize();
     size += m_code.GetSerializedSize();
@@ -419,6 +444,11 @@ MgtAssocResponseHeader::PrintImpl(std::ostream& os) const
 void
 MgtAssocResponseHeader::SerializeImpl(Buffer::Iterator start) const
 {
+    if (auto& mle = Get<MultiLinkElement>())
+    {
+        mle->m_containingFrame = *this;
+    }
+
     Buffer::Iterator i = start;
     i = m_capability.Serialize(i);
     i = m_code.Serialize(i);