iwlwifi: mvm: finalize on post_switch instead of unassign

Now that mac80211 waits for the first beacon on the new channel before
calling post_channel_switch, the reconfiguration of the firmware
should be done in the post_channel_switch operation instead of when
assigning the vif to the new context.

Signed-off-by: Luciano Coelho <luciano.coelho@intel.com>
diff --git a/drivers/net/wireless/iwlwifi/mvm/mac80211.c b/drivers/net/wireless/iwlwifi/mvm/mac80211.c
index 3ce5831..c85c0f1 100644
--- a/drivers/net/wireless/iwlwifi/mvm/mac80211.c
+++ b/drivers/net/wireless/iwlwifi/mvm/mac80211.c
@@ -2841,23 +2841,13 @@
 	}
 
 	/* Handle binding during CSA */
-	if ((vif->type == NL80211_IFTYPE_AP) ||
-	    (switching_chanctx && (vif->type == NL80211_IFTYPE_STATION))) {
+	if (vif->type == NL80211_IFTYPE_AP) {
 		iwl_mvm_update_quotas(mvm, NULL);
 		iwl_mvm_mac_ctxt_changed(mvm, vif, false, NULL);
 	}
 
 	if (switching_chanctx && vif->type == NL80211_IFTYPE_STATION) {
-		struct iwl_mvm_sta *mvmsta;
-
-		mvmsta = iwl_mvm_sta_from_staid_protected(mvm,
-							  mvmvif->ap_sta_id);
-
-		if (WARN_ON(!mvmsta))
-			goto out;
-
-		/* TODO: only re-enable after the first beacon */
-		iwl_mvm_sta_modify_disable_tx(mvm, mvmsta, false);
+		iwl_mvm_update_quotas(mvm, NULL);
 	}
 
 	goto out;
@@ -3228,10 +3218,27 @@
 
 	mutex_lock(&mvm->mutex);
 
+	if (vif->type == NL80211_IFTYPE_STATION) {
+		struct iwl_mvm_sta *mvmsta;
+
+		mvmsta = iwl_mvm_sta_from_staid_protected(mvm,
+							  mvmvif->ap_sta_id);
+
+		if (WARN_ON(!mvmsta)) {
+			ret = -EIO;
+			goto out_unlock;
+		}
+
+		iwl_mvm_sta_modify_disable_tx(mvm, mvmsta, false);
+
+		iwl_mvm_mac_ctxt_changed(mvm, vif, false, NULL);
+	}
+
 	mvmvif->ps_disabled = false;
 
 	ret = iwl_mvm_power_update_ps(mvm);
 
+out_unlock:
 	mutex_unlock(&mvm->mutex);
 
 	return ret;