iwlwifi: mvm: Use CS tx block bit for AP/GO

An AP/GO may perform the channel switch slightly before its stations.
This scenario may result in packet loss, since the transmission may start
before the client is actually on a new channel. In order to prevent
potential packet loss disable tx to all the stations when the channel
switch flow starts. Clear the disable_tx bit when a station is seen on a
target channel, or after IWL_MVM_CS_UNBLOCK_TX_TIMEOUT beacons on a new
channel. In addition call ieee80211_sta_block_awake in order to inform
mac80211 that the frames for this station should be buffered.

Signed-off-by: Andrei Otcheretianski <andrei.otcheretianski@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/sta.c b/drivers/net/wireless/iwlwifi/mvm/sta.c
index d3a6cf7..8128139 100644
--- a/drivers/net/wireless/iwlwifi/mvm/sta.c
+++ b/drivers/net/wireless/iwlwifi/mvm/sta.c
@@ -1468,3 +1468,57 @@
 	if (ret)
 		IWL_ERR(mvm, "Failed to send ADD_STA command (%d)\n", ret);
 }
+
+void iwl_mvm_sta_modify_disable_tx_ap(struct iwl_mvm *mvm,
+				      struct ieee80211_sta *sta,
+				      bool disable)
+{
+	struct iwl_mvm_sta *mvm_sta = iwl_mvm_sta_from_mac80211(sta);
+
+	spin_lock_bh(&mvm_sta->lock);
+
+	if (mvm_sta->disable_tx == disable) {
+		spin_unlock_bh(&mvm_sta->lock);
+		return;
+	}
+
+	mvm_sta->disable_tx = disable;
+
+	/*
+	 * Tell mac80211 to start/stop queueing tx for this station,
+	 * but don't stop queueing if there are still pending frames
+	 * for this station.
+	 */
+	if (disable || !atomic_read(&mvm->pending_frames[mvm_sta->sta_id]))
+		ieee80211_sta_block_awake(mvm->hw, sta, disable);
+
+	iwl_mvm_sta_modify_disable_tx(mvm, mvm_sta, disable);
+
+	spin_unlock_bh(&mvm_sta->lock);
+}
+
+void iwl_mvm_modify_all_sta_disable_tx(struct iwl_mvm *mvm,
+				       struct iwl_mvm_vif *mvmvif,
+				       bool disable)
+{
+	struct ieee80211_sta *sta;
+	struct iwl_mvm_sta *mvm_sta;
+	int i;
+
+	lockdep_assert_held(&mvm->mutex);
+
+	/* Block/unblock all the stations of the given mvmvif */
+	for (i = 0; i < IWL_MVM_STATION_COUNT; i++) {
+		sta = rcu_dereference_protected(mvm->fw_id_to_mac_id[i],
+						lockdep_is_held(&mvm->mutex));
+		if (IS_ERR_OR_NULL(sta))
+			continue;
+
+		mvm_sta = iwl_mvm_sta_from_mac80211(sta);
+		if (mvm_sta->mac_id_n_color !=
+		    FW_CMD_ID_AND_COLOR(mvmvif->id, mvmvif->color))
+			continue;
+
+		iwl_mvm_sta_modify_disable_tx_ap(mvm, sta, disable);
+	}
+}