iwlwifi: mvm: set the TX disable bit when doing a chanctx switch

During a channel switch we should tell the firmware to disable TX
temporarily and re-enable it after the switch is done.

Signed-off-by: Luciano Coelho <luciano.coelho@intel.com>
Reviewed-by: Johannes Berg <johannes.berg@intel.com>
Signed-off-by: Emmanuel Grumbach <emmanuel.grumbach@intel.com>
diff --git a/drivers/net/wireless/iwlwifi/mvm/mac80211.c b/drivers/net/wireless/iwlwifi/mvm/mac80211.c
index cbe9ff0..3bbeedf 100644
--- a/drivers/net/wireless/iwlwifi/mvm/mac80211.c
+++ b/drivers/net/wireless/iwlwifi/mvm/mac80211.c
@@ -2748,6 +2748,19 @@
 		iwl_mvm_mac_ctxt_changed(mvm, vif, false);
 	}
 
+	if (vif->csa_active && 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);
+	}
+
 	goto out;
 
 out_remove_binding:
@@ -2779,6 +2792,7 @@
 {
 	struct iwl_mvm_vif *mvmvif = iwl_mvm_vif_from_mac80211(vif);
 	struct ieee80211_vif *disabled_vif = NULL;
+	struct iwl_mvm_sta *mvmsta;
 
 	lockdep_assert_held(&mvm->mutex);
 
@@ -2810,6 +2824,12 @@
 
 		disabled_vif = vif;
 
+		mvmsta = iwl_mvm_sta_from_staid_protected(mvm,
+							  mvmvif->ap_sta_id);
+
+		if (!WARN_ON(!mvmsta))
+			iwl_mvm_sta_modify_disable_tx(mvm, mvmsta, true);
+
 		iwl_mvm_mac_ctxt_changed(mvm, vif, true);
 		break;
 	default: