From e345ad196bffaa69ff3d84233f672520a8d0a20f Mon Sep 17 00:00:00 2001
From: Davide Magrin <magrin.davide@gmail.com>
Date: Wed, 1 Mar 2023 23:12:20 +0100
Subject: [PATCH] wifi: Track Power Management mode in WifiRemoteStationManager

---
 src/wifi/model/wifi-remote-station-manager.cc | 13 +++++++++++++
 src/wifi/model/wifi-remote-station-manager.h  | 17 +++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/src/wifi/model/wifi-remote-station-manager.cc b/src/wifi/model/wifi-remote-station-manager.cc
index bc89e15d8..af34129ea 100644
--- a/src/wifi/model/wifi-remote-station-manager.cc
+++ b/src/wifi/model/wifi-remote-station-manager.cc
@@ -552,6 +552,18 @@ WifiRemoteStationManager::GetStaId(Mac48Address address, const WifiTxVector& txV
     return staId;
 }
 
+bool
+WifiRemoteStationManager::IsInPsMode(const Mac48Address& address) const
+{
+    return LookupState(address)->m_isInPsMode;
+}
+
+void
+WifiRemoteStationManager::SetPsMode(const Mac48Address& address, bool isInPsMode)
+{
+    LookupState(address)->m_isInPsMode = isInPsMode;
+}
+
 void
 WifiRemoteStationManager::SetMldAddress(const Mac48Address& address, const Mac48Address& mldAddress)
 {
@@ -1381,6 +1393,7 @@ WifiRemoteStationManager::LookupState(Mac48Address address) const
     state->m_ness = 0;
     state->m_aggregation = false;
     state->m_qosSupported = false;
+    state->m_isInPsMode = false;
     const_cast<WifiRemoteStationManager*>(this)->m_states.insert({address, state});
     NS_LOG_DEBUG("WifiRemoteStationManager::LookupState returning new state");
     return state;
diff --git a/src/wifi/model/wifi-remote-station-manager.h b/src/wifi/model/wifi-remote-station-manager.h
index bb8af3bed..1828d0866 100644
--- a/src/wifi/model/wifi-remote-station-manager.h
+++ b/src/wifi/model/wifi-remote-station-manager.h
@@ -122,6 +122,7 @@ struct WifiRemoteStationState
     bool m_shortPreamble;     //!< Flag if short PHY preamble is supported by the remote station
     bool m_shortSlotTime;     //!< Flag if short ERP slot time is supported by the remote station
     bool m_qosSupported;      //!< Flag if QoS is supported by the station
+    bool m_isInPsMode;        //!< Flag if the STA is currently in PS mode
 };
 
 /**
@@ -733,6 +734,22 @@ class WifiRemoteStationManager : public Object
      */
     void RecordAssocRefused(Mac48Address address);
 
+    /**
+     * Return whether the STA is currently in Power Save mode.
+     *
+     * \param address the address of the station
+     *
+     * \return true if the station is in Power Save mode, false otherwise
+     */
+    bool IsInPsMode(const Mac48Address& address) const;
+    /**
+     * Register whether the STA is in Power Save mode or not.
+     *
+     * \param address the address of the station
+     * \param isInPsMode whether the STA is in PS mode or not
+     */
+    void SetPsMode(const Mac48Address& address, bool isInPsMode);
+
     /**
      * Set the address of the MLD the given station is affiliated with.
      *