mac80211: init HT TX data before rate control

In case of connection, the station data is initialised from
the beacon/probe response first and then updated from the
association response. If the latter is different we update
the rate control algorithm and driver. Instead of doing it
this way, set the station data properly with data from the
association response before initializing rate control.

Also simplify the code by passing the station pointer.

Signed-off-by: Johannes Berg <johannes.berg@intel.com>
diff --git a/net/mac80211/mlme.c b/net/mac80211/mlme.c
index be11472..b97b615 100644
--- a/net/mac80211/mlme.c
+++ b/net/mac80211/mlme.c
@@ -175,6 +175,7 @@
 }
 
 static u32 ieee80211_config_ht_tx(struct ieee80211_sub_if_data *sdata,
+				  struct sta_info *sta,
 				  struct ieee80211_ht_operation *ht_oper,
 				  const u8 *bssid, bool reconfig)
 {
@@ -182,11 +183,13 @@
 	struct ieee80211_supported_band *sband;
 	struct ieee80211_chanctx_conf *chanctx_conf;
 	struct ieee80211_channel *chan;
-	struct sta_info *sta;
 	u32 changed = 0;
 	u16 ht_opmode;
 	bool disable_40 = false;
 
+	if (WARN_ON_ONCE(!sta))
+		return 0;
+
 	rcu_read_lock();
 	chanctx_conf = rcu_dereference(sdata->vif.chanctx_conf);
 	if (WARN_ON(!chanctx_conf)) {
@@ -216,28 +219,19 @@
 	if (!(ht_oper->ht_param & IEEE80211_HT_PARAM_CHAN_WIDTH_ANY))
 		disable_40 = true;
 
-	mutex_lock(&local->sta_mtx);
-	sta = sta_info_get(sdata, bssid);
-
-	if (WARN_ON_ONCE(!sta)) {
-		mutex_unlock(&local->sta_mtx);
-		return changed;
-	}
-
 	if (!(sta->sta.ht_cap.cap & IEEE80211_HT_CAP_SUP_WIDTH_20_40))
 		disable_40 = true;
 
-	if (!reconfig ||
-	    disable_40 != (sta->sta.bandwidth < IEEE80211_STA_RX_BW_40)) {
+	if (disable_40 != (sta->sta.bandwidth < IEEE80211_STA_RX_BW_40)) {
 		if (disable_40)
 			sta->sta.bandwidth = IEEE80211_STA_RX_BW_20;
 		else
 			sta->sta.bandwidth = ieee80211_sta_cur_vht_bw(sta);
 
-		rate_control_rate_update(local, sband, sta,
-					 IEEE80211_RC_BW_CHANGED);
+		if (reconfig)
+			rate_control_rate_update(local, sband, sta,
+						 IEEE80211_RC_BW_CHANGED);
 	}
-	mutex_unlock(&local->sta_mtx);
 
 	ht_opmode = le16_to_cpu(ht_oper->operation_mode);
 
@@ -2217,6 +2211,12 @@
 		ieee80211_vht_cap_ie_to_sta_vht_cap(sdata, sband,
 						    elems.vht_cap_elem, sta);
 
+	if (elems.ht_operation && elems.wmm_param &&
+	    !(ifmgd->flags & IEEE80211_STA_DISABLE_HT))
+		changed |= ieee80211_config_ht_tx(sdata, sta,
+						  elems.ht_operation,
+						  cbss->bssid, false);
+
 	rate_control_rate_init(sta);
 
 	if (ifmgd->flags & IEEE80211_STA_MFP_ENABLED)
@@ -2254,11 +2254,6 @@
 		ieee80211_set_wmm_default(sdata, false);
 	changed |= BSS_CHANGED_QOS;
 
-	if (elems.ht_operation && elems.wmm_param &&
-	    !(ifmgd->flags & IEEE80211_STA_DISABLE_HT))
-		changed |= ieee80211_config_ht_tx(sdata, elems.ht_operation,
-						  cbss->bssid, false);
-
 	/* set AID and assoc capability,
 	 * ieee80211_set_associated() will tell the driver */
 	bss_conf->aid = aid;
@@ -2734,10 +2729,14 @@
 			erp_valid, erp_value);
 
 
+	mutex_lock(&local->sta_mtx);
 	if (elems.ht_cap_elem && elems.ht_operation && elems.wmm_param &&
 	    !(ifmgd->flags & IEEE80211_STA_DISABLE_HT))
-		changed |= ieee80211_config_ht_tx(sdata, elems.ht_operation,
+		changed |= ieee80211_config_ht_tx(sdata,
+						  sta_info_get(sdata, bssid),
+						  elems.ht_operation,
 						  bssid, true);
+	mutex_unlock(&local->sta_mtx);
 
 	if (elems.country_elem && elems.pwr_constr_elem &&
 	    mgmt->u.probe_resp.capab_info &