[04/28] wifi: mac80211: recalc station aggregate data during link switch

diff --git a/net/mac80211/link.c b/net/mac80211/link.c
index e309708abae8..d1f5a9f7c647 100644
--- a/net/mac80211/link.c
+++ b/net/mac80211/link.c
@@ -357,6 +357,11 @@  static int _ieee80211_set_active_links(struct ieee80211_sub_if_data *sdata,
 	list_for_each_entry(sta, &local->sta_list, list) {
 		if (sdata != sta->sdata)
 			continue;
+
+		/* this is very temporary, but do it anyway */
+		__ieee80211_sta_recalc_aggregates(sta,
+						  old_active | active_links);
+
 		ret = drv_change_sta_links(local, sdata, &sta->sta,
 					   old_active,
 					   old_active | active_links);
@@ -369,10 +374,22 @@  static int _ieee80211_set_active_links(struct ieee80211_sub_if_data *sdata,
 	list_for_each_entry(sta, &local->sta_list, list) {
 		if (sdata != sta->sdata)
 			continue;
+
+		__ieee80211_sta_recalc_aggregates(sta, active_links);
+
 		ret = drv_change_sta_links(local, sdata, &sta->sta,
 					   old_active | active_links,
 					   active_links);
 		WARN_ON_ONCE(ret);
+
+		/*
+		 * Do it again, just in case - the driver might very
+		 * well have called ieee80211_sta_recalc_aggregates()
+		 * from there when filling in the new links, which
+		 * would set it wrong since the vif's active links are
+		 * not switched yet...
+		 */
+		__ieee80211_sta_recalc_aggregates(sta, active_links);
 	}
 
 	for_each_set_bit(link_id, &add, IEEE80211_MLD_MAX_NUM_LINKS) {
diff --git a/net/mac80211/sta_info.c b/net/mac80211/sta_info.c
index e6beaea4075e..2bb6a71c72ef 100644
--- a/net/mac80211/sta_info.c
+++ b/net/mac80211/sta_info.c
@@ -2151,22 +2151,30 @@  void ieee80211_sta_register_airtime(struct ieee80211_sta *pubsta, u8 tid,
 }
 EXPORT_SYMBOL(ieee80211_sta_register_airtime);
 
-void ieee80211_sta_recalc_aggregates(struct ieee80211_sta *pubsta)
+void __ieee80211_sta_recalc_aggregates(struct sta_info *sta, u16 active_links)
 {
-	struct sta_info *sta = container_of(pubsta, struct sta_info, sta);
-	struct ieee80211_link_sta *link_sta;
-	int link_id, i;
 	bool first = true;
+	int link_id;
 
-	if (!pubsta->valid_links || !pubsta->mlo) {
-		pubsta->cur = &pubsta->deflink.agg;
+	if (!sta->sta.valid_links || !sta->sta.mlo) {
+		sta->sta.cur = &sta->sta.deflink.agg;
 		return;
 	}
 
 	rcu_read_lock();
-	for_each_sta_active_link(&sta->sdata->vif, pubsta, link_sta, link_id) {
+	for (link_id = 0; link_id < ARRAY_SIZE((sta)->link); link_id++) {
+		struct ieee80211_link_sta *link_sta;
+		int i;
+
+		if (!(active_links & BIT(link_id)))
+			continue;
+
+		link_sta = rcu_dereference(sta->sta.link[link_id]);
+		if (!link_sta)
+			continue;
+
 		if (first) {
-			sta->cur = pubsta->deflink.agg;
+			sta->cur = sta->sta.deflink.agg;
 			first = false;
 			continue;
 		}
@@ -2185,7 +2193,14 @@  void ieee80211_sta_recalc_aggregates(struct ieee80211_sta *pubsta)
 	}
 	rcu_read_unlock();
 
-	pubsta->cur = &sta->cur;
+	sta->sta.cur = &sta->cur;
+}
+
+void ieee80211_sta_recalc_aggregates(struct ieee80211_sta *pubsta)
+{
+	struct sta_info *sta = container_of(pubsta, struct sta_info, sta);
+
+	__ieee80211_sta_recalc_aggregates(sta, sta->sdata->vif.active_links);
 }
 EXPORT_SYMBOL(ieee80211_sta_recalc_aggregates);
 
diff --git a/net/mac80211/sta_info.h b/net/mac80211/sta_info.h
index 6e672bf9c79d..69820b551668 100644
--- a/net/mac80211/sta_info.h
+++ b/net/mac80211/sta_info.h
@@ -927,6 +927,8 @@  void ieee80211_sta_set_max_amsdu_subframes(struct sta_info *sta,
 					   const u8 *ext_capab,
 					   unsigned int ext_capab_len);
 
+void __ieee80211_sta_recalc_aggregates(struct sta_info *sta, u16 active_links);
+
 enum sta_stats_type {
 	STA_STATS_RATE_TYPE_INVALID = 0,
 	STA_STATS_RATE_TYPE_LEGACY,