+From 94dd7ce1885e530a7b10bbe50d5d68ba1bb99e6e Mon Sep 17 00:00:00 2001
+From: Martin Kaistra <martin.kaistra@linutronix.de>
+Date: Mon, 5 Feb 2024 10:30:40 +0100
+Subject: [PATCH] wifi: rtl8xxxu: update rate mask per sta
+
+Until now, rtl8xxxu_watchdog_callback() only fetches RSSI and updates
+the rate mask in station mode. This means, in AP mode only the default
+rate mask is used.
+
+In order to have the rate mask reflect the actual connection quality,
+extend rtl8xxxu_watchdog_callback() to iterate over every sta. Like in
+the rtw88 driver, add a function to collect all currently present stas
+and then iterate over a list of copies to ensure no RCU lock problems
+for register access via USB. Remove the existing RCU lock in
+rtl8xxxu_refresh_rate_mask().
+
+Since the currently used ieee80211_ave_rssi() is only for 'vif', add
+driver-level tracking of RSSI per sta.
+
+Signed-off-by: Martin Kaistra <martin.kaistra@linutronix.de>
+Reviewed-by: Ping-Ke Shih <pkshih@realtek.com>
+Signed-off-by: Kalle Valo <kvalo@kernel.org>
+Link: https://msgid.link/20240205093040.1941140-1-martin.kaistra@linutronix.de
+---
+ .../net/wireless/realtek/rtl8xxxu/rtl8xxxu.h | 8 +-
+ .../wireless/realtek/rtl8xxxu/rtl8xxxu_core.c | 188 ++++++++++++++----
+ 2 files changed, 158 insertions(+), 38 deletions(-)
+
+--- a/drivers/net/wireless/realtek/rtl8xxxu/rtl8xxxu.h
++++ b/drivers/net/wireless/realtek/rtl8xxxu/rtl8xxxu.h
+@@ -6,6 +6,7 @@
+ */
+
+ #include <asm/byteorder.h>
++#include <linux/average.h>
+
+ #define RTL8XXXU_DEBUG_REG_WRITE 0x01
+ #define RTL8XXXU_DEBUG_REG_READ 0x02
+@@ -1858,6 +1859,8 @@ struct rtl8xxxu_priv {
+ int next_mbox;
+ int nr_out_eps;
+
++ /* Ensure no added or deleted stas while iterating */
++ struct mutex sta_mutex;
+ struct mutex h2c_mutex;
+ /* Protect the indirect register accesses of RTL8710BU. */
+ struct mutex syson_indirect_access_mutex;
+@@ -1892,7 +1895,6 @@ struct rtl8xxxu_priv {
+ u8 pi_enabled:1;
+ u8 no_pape:1;
+ u8 int_buf[USB_INTR_CONTENT_LENGTH];
+- u8 rssi_level;
+ DECLARE_BITMAP(tx_aggr_started, IEEE80211_NUM_TIDS);
+ DECLARE_BITMAP(tid_tx_operational, IEEE80211_NUM_TIDS);
+
+@@ -1913,11 +1915,15 @@ struct rtl8xxxu_priv {
+ DECLARE_BITMAP(cam_map, RTL8XXXU_MAX_SEC_CAM_NUM);
+ };
+
++DECLARE_EWMA(rssi, 10, 16);
++
+ struct rtl8xxxu_sta_info {
+ struct ieee80211_sta *sta;
+ struct ieee80211_vif *vif;
+
+ u8 macid;
++ struct ewma_rssi avg_rssi;
++ u8 rssi_level;
+ };
+
+ struct rtl8xxxu_vif {
+--- a/drivers/net/wireless/realtek/rtl8xxxu/rtl8xxxu_core.c
++++ b/drivers/net/wireless/realtek/rtl8xxxu/rtl8xxxu_core.c
+@@ -4991,10 +4991,11 @@ rtl8xxxu_bss_info_changed(struct ieee802
+ struct rtl8xxxu_vif *rtlvif = (struct rtl8xxxu_vif *)vif->drv_priv;
+ struct rtl8xxxu_priv *priv = hw->priv;
+ struct device *dev = &priv->udev->dev;
++ struct rtl8xxxu_sta_info *sta_info;
+ struct ieee80211_sta *sta;
+ struct rtl8xxxu_ra_report *rarpt;
++ u8 val8, macid;
+ u32 val32;
+- u8 val8;
+
+ rarpt = &priv->ra_report;
+
+@@ -5017,6 +5018,7 @@ rtl8xxxu_bss_info_changed(struct ieee802
+ rcu_read_unlock();
+ goto error;
+ }
++ macid = rtl8xxxu_get_macid(priv, sta);
+
+ if (sta->deflink.ht_cap.ht_supported)
+ dev_info(dev, "%s: HT supported\n", __func__);
+@@ -5037,14 +5039,15 @@ rtl8xxxu_bss_info_changed(struct ieee802
+ bw = RATE_INFO_BW_40;
+ else
+ bw = RATE_INFO_BW_20;
++
++ sta_info = (struct rtl8xxxu_sta_info *)sta->drv_priv;
++ sta_info->rssi_level = RTL8XXXU_RATR_STA_INIT;
+ rcu_read_unlock();
+
+ rtl8xxxu_update_ra_report(rarpt, highest_rate, sgi, bw);
+
+- priv->rssi_level = RTL8XXXU_RATR_STA_INIT;
+-
+ priv->fops->update_rate_mask(priv, ramask, 0, sgi,
+- bw == RATE_INFO_BW_40, 0);
++ bw == RATE_INFO_BW_40, macid);
+
+ rtl8xxxu_write8(priv, REG_BCN_MAX_ERR, 0xff);
+
+@@ -6317,6 +6320,76 @@ static void rtl8188e_c2hcmd_callback(str
+ }
+ }
+
++#define rtl8xxxu_iterate_vifs_atomic(priv, iterator, data) \
++ ieee80211_iterate_active_interfaces_atomic((priv)->hw, \
++ IEEE80211_IFACE_ITER_NORMAL, iterator, data)
++
++struct rtl8xxxu_rx_update_rssi_data {
++ struct rtl8xxxu_priv *priv;
++ struct ieee80211_hdr *hdr;
++ struct ieee80211_rx_status *rx_status;
++ u8 *bssid;
++};
++
++static void rtl8xxxu_rx_update_rssi_iter(void *data, u8 *mac,
++ struct ieee80211_vif *vif)
++{
++ struct rtl8xxxu_rx_update_rssi_data *iter_data = data;
++ struct ieee80211_sta *sta;
++ struct ieee80211_hdr *hdr = iter_data->hdr;
++ struct rtl8xxxu_priv *priv = iter_data->priv;
++ struct rtl8xxxu_sta_info *sta_info;
++ struct ieee80211_rx_status *rx_status = iter_data->rx_status;
++ u8 *bssid = iter_data->bssid;
++
++ if (!ether_addr_equal(vif->bss_conf.bssid, bssid))
++ return;
++
++ if (!(ether_addr_equal(vif->addr, hdr->addr1) ||
++ ieee80211_is_beacon(hdr->frame_control)))
++ return;
++
++ sta = ieee80211_find_sta_by_ifaddr(priv->hw, hdr->addr2,
++ vif->addr);
++ if (!sta)
++ return;
++
++ sta_info = (struct rtl8xxxu_sta_info *)sta->drv_priv;
++ ewma_rssi_add(&sta_info->avg_rssi, -rx_status->signal);
++}
++
++static inline u8 *get_hdr_bssid(struct ieee80211_hdr *hdr)
++{
++ __le16 fc = hdr->frame_control;
++ u8 *bssid;
++
++ if (ieee80211_has_tods(fc))
++ bssid = hdr->addr1;
++ else if (ieee80211_has_fromds(fc))
++ bssid = hdr->addr2;
++ else
++ bssid = hdr->addr3;
++
++ return bssid;
++}
++
++static void rtl8xxxu_rx_update_rssi(struct rtl8xxxu_priv *priv,
++ struct ieee80211_rx_status *rx_status,
++ struct ieee80211_hdr *hdr)
++{
++ struct rtl8xxxu_rx_update_rssi_data data = {};
++
++ if (ieee80211_is_ctl(hdr->frame_control))
++ return;
++
++ data.priv = priv;
++ data.hdr = hdr;
++ data.rx_status = rx_status;
++ data.bssid = get_hdr_bssid(hdr);
++
++ rtl8xxxu_iterate_vifs_atomic(priv, rtl8xxxu_rx_update_rssi_iter, &data);
++}
++
+ int rtl8xxxu_parse_rxdesc16(struct rtl8xxxu_priv *priv, struct sk_buff *skb)
+ {
+ struct ieee80211_hw *hw = priv->hw;
+@@ -6376,18 +6449,26 @@ int rtl8xxxu_parse_rxdesc16(struct rtl8x
+ skb_queue_tail(&priv->c2hcmd_queue, skb);
+ schedule_work(&priv->c2hcmd_work);
+ } else {
++ struct ieee80211_hdr *hdr;
++
+ phy_stats = (struct rtl8723au_phy_stats *)skb->data;
+
+ skb_pull(skb, drvinfo_sz + desc_shift);
+
+ skb_trim(skb, pkt_len);
+
+- if (rx_desc->phy_stats)
++ hdr = (struct ieee80211_hdr *)skb->data;
++ if (rx_desc->phy_stats) {
+ priv->fops->parse_phystats(
+ priv, rx_status, phy_stats,
+ rx_desc->rxmcs,
+- (struct ieee80211_hdr *)skb->data,
++ hdr,
+ rx_desc->crc32 || rx_desc->icverr);
++ if (!rx_desc->crc32 && !rx_desc->icverr)
++ rtl8xxxu_rx_update_rssi(priv,
++ rx_status,
++ hdr);
++ }
+
+ rx_status->mactime = rx_desc->tsfl;
+ rx_status->flag |= RX_FLAG_MACTIME_START;
+@@ -6484,10 +6565,15 @@ int rtl8xxxu_parse_rxdesc24(struct rtl8x
+ } else {
+ struct ieee80211_hdr *hdr = (struct ieee80211_hdr *)skb->data;
+
+- if (rx_desc->phy_stats)
++ if (rx_desc->phy_stats) {
+ priv->fops->parse_phystats(priv, rx_status, phy_stats,
+ rx_desc->rxmcs, hdr,
+ rx_desc->crc32 || rx_desc->icverr);
++ if (!rx_desc->crc32 && !rx_desc->icverr)
++ rtl8xxxu_rx_update_rssi(priv,
++ rx_status,
++ hdr);
++ }
+
+ rx_status->mactime = rx_desc->tsfl;
+ rx_status->flag |= RX_FLAG_MACTIME_START;
+@@ -7111,6 +7197,7 @@ static void rtl8xxxu_refresh_rate_mask(s
+ int signal, struct ieee80211_sta *sta,
+ bool force)
+ {
++ struct rtl8xxxu_sta_info *sta_info = (struct rtl8xxxu_sta_info *)sta->drv_priv;
+ struct ieee80211_hw *hw = priv->hw;
+ u16 wireless_mode;
+ u8 rssi_level, ratr_idx;
+@@ -7119,7 +7206,7 @@ static void rtl8xxxu_refresh_rate_mask(s
+ u8 go_up_gap = 5;
+ u8 macid = rtl8xxxu_get_macid(priv, sta);
+
+- rssi_level = priv->rssi_level;
++ rssi_level = sta_info->rssi_level;
+ snr = rtl8xxxu_signal_to_snr(signal);
+ snr_thresh_high = RTL8XXXU_SNR_THRESH_HIGH;
+ snr_thresh_low = RTL8XXXU_SNR_THRESH_LOW;
+@@ -7144,18 +7231,16 @@ static void rtl8xxxu_refresh_rate_mask(s
+ else
+ rssi_level = RTL8XXXU_RATR_STA_LOW;
+
+- if (rssi_level != priv->rssi_level || force) {
++ if (rssi_level != sta_info->rssi_level || force) {
+ int sgi = 0;
+ u32 rate_bitmap = 0;
+
+- rcu_read_lock();
+ rate_bitmap = (sta->deflink.supp_rates[0] & 0xfff) |
+ (sta->deflink.ht_cap.mcs.rx_mask[0] << 12) |
+ (sta->deflink.ht_cap.mcs.rx_mask[1] << 20);
+ if (sta->deflink.ht_cap.cap &
+ (IEEE80211_HT_CAP_SGI_40 | IEEE80211_HT_CAP_SGI_20))
+ sgi = 1;
+- rcu_read_unlock();
+
+ wireless_mode = rtl8xxxu_wireless_mode(hw, sta);
+ switch (wireless_mode) {
+@@ -7236,7 +7321,7 @@ static void rtl8xxxu_refresh_rate_mask(s
+ break;
+ }
+
+- priv->rssi_level = rssi_level;
++ sta_info->rssi_level = rssi_level;
+ priv->fops->update_rate_mask(priv, rate_bitmap, ratr_idx, sgi, txbw_40mhz, macid);
+ }
+ }
+@@ -7329,40 +7414,60 @@ static void rtl8xxxu_track_cfo(struct rt
+ rtl8xxxu_set_atc_status(priv, abs(cfo_average) >= CFO_TH_ATC);
+ }
+
+-static void rtl8xxxu_watchdog_callback(struct work_struct *work)
++static void rtl8xxxu_ra_iter(void *data, struct ieee80211_sta *sta)
+ {
+- struct ieee80211_vif *vif;
+- struct rtl8xxxu_priv *priv;
+- int i;
++ struct rtl8xxxu_sta_info *sta_info = (struct rtl8xxxu_sta_info *)sta->drv_priv;
++ struct rtl8xxxu_priv *priv = data;
++ int signal = -ewma_rssi_read(&sta_info->avg_rssi);
+
+- priv = container_of(work, struct rtl8xxxu_priv, ra_watchdog.work);
+- for (i = 0; i < ARRAY_SIZE(priv->vifs); i++) {
+- vif = priv->vifs[i];
++ priv->fops->report_rssi(priv, rtl8xxxu_get_macid(priv, sta),
++ rtl8xxxu_signal_to_snr(signal));
++ rtl8xxxu_refresh_rate_mask(priv, signal, sta, false);
++}
++
++struct rtl8xxxu_stas_entry {
++ struct list_head list;
++ struct ieee80211_sta *sta;
++};
+
+- if (!vif || vif->type != NL80211_IFTYPE_STATION)
+- continue;
++struct rtl8xxxu_iter_stas_data {
++ struct rtl8xxxu_priv *priv;
++ struct list_head list;
++};
+
+- int signal;
+- struct ieee80211_sta *sta;
++static void rtl8xxxu_collect_sta_iter(void *data, struct ieee80211_sta *sta)
++{
++ struct rtl8xxxu_iter_stas_data *iter_stas = data;
++ struct rtl8xxxu_stas_entry *stas_entry;
+
+- rcu_read_lock();
+- sta = ieee80211_find_sta(vif, vif->bss_conf.bssid);
+- if (!sta) {
+- struct device *dev = &priv->udev->dev;
++ stas_entry = kmalloc(sizeof(*stas_entry), GFP_ATOMIC);
++ if (!stas_entry)
++ return;
+
+- dev_dbg(dev, "%s: no sta found\n", __func__);
+- rcu_read_unlock();
+- continue;
+- }
+- rcu_read_unlock();
++ stas_entry->sta = sta;
++ list_add_tail(&stas_entry->list, &iter_stas->list);
++}
+
+- signal = ieee80211_ave_rssi(vif);
++static void rtl8xxxu_watchdog_callback(struct work_struct *work)
++{
+
+- priv->fops->report_rssi(priv, rtl8xxxu_get_macid(priv, sta),
+- rtl8xxxu_signal_to_snr(signal));
++ struct rtl8xxxu_iter_stas_data iter_data;
++ struct rtl8xxxu_stas_entry *sta_entry, *tmp;
++ struct rtl8xxxu_priv *priv;
+
+- rtl8xxxu_refresh_rate_mask(priv, signal, sta, false);
++ priv = container_of(work, struct rtl8xxxu_priv, ra_watchdog.work);
++ iter_data.priv = priv;
++ INIT_LIST_HEAD(&iter_data.list);
++
++ mutex_lock(&priv->sta_mutex);
++ ieee80211_iterate_stations_atomic(priv->hw, rtl8xxxu_collect_sta_iter,
++ &iter_data);
++ list_for_each_entry_safe(sta_entry, tmp, &iter_data.list, list) {
++ list_del_init(&sta_entry->list);
++ rtl8xxxu_ra_iter(priv, sta_entry->sta);
++ kfree(sta_entry);
+ }
++ mutex_unlock(&priv->sta_mutex);
+
+ if (priv->fops->set_crystal_cap)
+ rtl8xxxu_track_cfo(priv);
+@@ -7504,10 +7609,15 @@ static int rtl8xxxu_sta_add(struct ieee8
+ struct rtl8xxxu_vif *rtlvif = (struct rtl8xxxu_vif *)vif->drv_priv;
+ struct rtl8xxxu_priv *priv = hw->priv;
+
++ mutex_lock(&priv->sta_mutex);
++ ewma_rssi_init(&sta_info->avg_rssi);
+ if (vif->type == NL80211_IFTYPE_AP) {
++ sta_info->rssi_level = RTL8XXXU_RATR_STA_INIT;
+ sta_info->macid = rtl8xxxu_acquire_macid(priv);
+- if (sta_info->macid >= RTL8XXXU_MAX_MAC_ID_NUM)
++ if (sta_info->macid >= RTL8XXXU_MAX_MAC_ID_NUM) {
++ mutex_unlock(&priv->sta_mutex);
+ return -ENOSPC;
++ }
+
+ rtl8xxxu_refresh_rate_mask(priv, 0, sta, true);
+ priv->fops->report_connect(priv, sta_info->macid, H2C_MACID_ROLE_STA, true);
+@@ -7523,6 +7633,7 @@ static int rtl8xxxu_sta_add(struct ieee8
+ break;
+ }
+ }
++ mutex_unlock(&priv->sta_mutex);
+
+ return 0;
+ }
+@@ -7534,8 +7645,10 @@ static int rtl8xxxu_sta_remove(struct ie
+ struct rtl8xxxu_sta_info *sta_info = (struct rtl8xxxu_sta_info *)sta->drv_priv;
+ struct rtl8xxxu_priv *priv = hw->priv;
+
++ mutex_lock(&priv->sta_mutex);
+ if (vif->type == NL80211_IFTYPE_AP)
+ rtl8xxxu_release_macid(priv, sta_info->macid);
++ mutex_unlock(&priv->sta_mutex);
+
+ return 0;
+ }
+@@ -7766,6 +7879,7 @@ static int rtl8xxxu_probe(struct usb_int
+ mutex_init(&priv->usb_buf_mutex);
+ mutex_init(&priv->syson_indirect_access_mutex);
+ mutex_init(&priv->h2c_mutex);
++ mutex_init(&priv->sta_mutex);
+ INIT_LIST_HEAD(&priv->tx_urb_free_list);
+ spin_lock_init(&priv->tx_urb_lock);
+ INIT_LIST_HEAD(&priv->rx_urb_pending_list);