From 85f833fa34ffdbd1f24a89a8720ec98cfd57bb19 Mon Sep 17 00:00:00 2001
From: Stefano Avallone <stavallo@unina.it>
Date: Tue, 7 Feb 2023 11:29:37 +0100
Subject: [PATCH] wifi: STAs store AP's Transition Timeout in EMLSR manager

---
 src/wifi/model/eht/emlsr-manager.cc | 13 +++++++++++++
 src/wifi/model/eht/emlsr-manager.h  | 16 +++++++++++++++-
 src/wifi/model/sta-wifi-mac.cc      |  6 ++++++
 3 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/src/wifi/model/eht/emlsr-manager.cc b/src/wifi/model/eht/emlsr-manager.cc
index b7aaefc0fb..463a57e93 100644
--- a/src/wifi/model/eht/emlsr-manager.cc
+++ b/src/wifi/model/eht/emlsr-manager.cc
@@ -97,4 +97,17 @@ EmlsrManager::GetEhtFem(uint8_t linkId) const
     return StaticCast<EhtFrameExchangeManager>(m_staMac->GetFrameExchangeManager(linkId));
 }
 
+void
+EmlsrManager::SetTransitionTimeout(Time timeout)
+{
+    NS_LOG_FUNCTION(this << timeout.As(Time::US));
+    m_emlsrTransitionTimeout = timeout;
+}
+
+std::optional<Time>
+EmlsrManager::GetTransitionTimeout() const
+{
+    return m_emlsrTransitionTimeout;
+}
+
 } // namespace ns3
diff --git a/src/wifi/model/eht/emlsr-manager.h b/src/wifi/model/eht/emlsr-manager.h
index 569c4b7e9..2ab5cc04a 100644
--- a/src/wifi/model/eht/emlsr-manager.h
+++ b/src/wifi/model/eht/emlsr-manager.h
@@ -52,6 +52,18 @@ class EmlsrManager : public Object
      */
     void SetWifiMac(Ptr<StaWifiMac> mac);
 
+    /**
+     * Set the Transition Timeout advertised by the associated AP with EMLSR activated.
+     *
+     * \param timeout the advertised Transition Timeout
+     */
+    void SetTransitionTimeout(Time timeout);
+
+    /**
+     * \return the Transition Timeout, if advertised by the associated AP
+     */
+    std::optional<Time> GetTransitionTimeout() const;
+
   protected:
     void DoDispose() override;
 
@@ -70,7 +82,9 @@ class EmlsrManager : public Object
     Time m_emlsrTransitionDelay; //!< EMLSR Transition delay
 
   private:
-    Ptr<StaWifiMac> m_staMac; //!< the MAC of the managed non-AP MLD
+    Ptr<StaWifiMac> m_staMac;                     //!< the MAC of the managed non-AP MLD
+    std::optional<Time> m_emlsrTransitionTimeout; /**< Transition timeout advertised by APs with
+                                                       EMLSR activated */
 };
 
 } // namespace ns3
diff --git a/src/wifi/model/sta-wifi-mac.cc b/src/wifi/model/sta-wifi-mac.cc
index e293eb859..6ea2dc87c 100644
--- a/src/wifi/model/sta-wifi-mac.cc
+++ b/src/wifi/model/sta-wifi-mac.cc
@@ -1430,6 +1430,12 @@ StaWifiMac::UpdateApInfo(const MgtFrameType& frame,
         // TODO: once we support non constant rate managers, we should add checks here whether EHT
         // is supported by the peer
         GetWifiRemoteStationManager(linkId)->AddStationEhtCapabilities(apAddr, *ehtCapabilities);
+
+        if (const auto& mle = frame.template Get<MultiLinkElement>();
+            mle && mle->HasEmlCapabilities() && m_emlsrManager)
+        {
+            m_emlsrManager->SetTransitionTimeout(mle->GetTransitionTimeout());
+        }
     };
 
     // process Information Elements included in the current frame variant