wifi: iwlwifi: mvm: factor out iwl_mvm_sta_fw_id_mask()

We are going to need this in more places than just the
key code, so factor out the functionality of getting
the FW station ID mask (filtered to a specific link if
needed) to a separate function that can now be called
both under RCU and mvm->mutex protection.
Signed-off-by: Johannes Berg <johannes.berg@intel.com>
Signed-off-by: Gregory Greenman <gregory.greenman@intel.com>
Link: https://lore.kernel.org/r/20230416154301.eff92b93025d.I2c50290a0537d5db3d3460f4d57c78a4712ffb75@changeidSigned-off-by: Johannes Berg <johannes.berg@intel.com>

drivers/net/wireless/intel/iwlwifi/mvm/mld-key.c

@@ -15,10 +15,6 @@ static u32 iwl_mvm_get_sec_sta_mask(struct iwl_mvm *mvm,
 {
 	struct iwl_mvm_vif *mvmvif = iwl_mvm_vif_from_mac80211(vif);
 	struct iwl_mvm_vif_link_info *link_info = &mvmvif->deflink;
-	struct iwl_mvm_link_sta *link_sta;
-	struct iwl_mvm_sta *mvmsta;
-	u32 result = 0;
-	int link_id;
 
 	lockdep_assert_held(&mvm->mutex);
 
@@ -49,33 +45,10 @@ static u32 iwl_mvm_get_sec_sta_mask(struct iwl_mvm *mvm,
 	if (!sta && (keyconf->link_id >= 0 || !vif->valid_links))
 		return BIT(link_info->ap_sta_id);
 
-	/* this shouldn't happen now */
-	if (!sta)
-		return 0;
-
-	mvmsta = iwl_mvm_sta_from_mac80211(sta);
-
-	/* it's easy when the STA is not an MLD */
-	if (!sta->valid_links)
-		return BIT(mvmsta->deflink.sta_id);
-
-	/* but if it is an MLD, get the mask of all the FW STAs it has ... */
-	for (link_id = 0; link_id < ARRAY_SIZE(mvmsta->link); link_id++) {
-		/* unless we have a specific link in mind (GTK on client) */
-		if (keyconf->link_id >= 0 &&
-		    keyconf->link_id != link_id)
-			continue;
-
-		link_sta =
-			rcu_dereference_protected(mvmsta->link[link_id],
-						  lockdep_is_held(&mvm->mutex));
-		if (!link_sta)
-			continue;
-
-		result |= BIT(link_sta->sta_id);
-	}
+	/* STA should be non-NULL now, but iwl_mvm_sta_fw_id_mask() checks */
 
-	return result;
+	/* pass link_id to filter by it if not -1 (GTK on client) */
+	return iwl_mvm_sta_fw_id_mask(mvm, sta, keyconf->link_id);
 }
 
 static u32 iwl_mvm_get_sec_flags(struct iwl_mvm *mvm,

drivers/net/wireless/intel/iwlwifi/mvm/mld-sta.c

@@ -4,6 +4,43 @@
  */
 #include "mvm.h"
 #include "time-sync.h"
+#include "sta.h"
+
+u32 iwl_mvm_sta_fw_id_mask(struct iwl_mvm *mvm, struct ieee80211_sta *sta,
+			   int filter_link_id)
+{
+	struct iwl_mvm_sta *mvmsta;
+	unsigned int link_id;
+	u32 result = 0;
+
+	if (!sta)
+		return 0;
+
+	mvmsta = iwl_mvm_sta_from_mac80211(sta);
+
+	/* it's easy when the STA is not an MLD */
+	if (!sta->valid_links)
+		return BIT(mvmsta->deflink.sta_id);
+
+	/* but if it is an MLD, get the mask of all the FW STAs it has ... */
+	for (link_id = 0; link_id < ARRAY_SIZE(mvmsta->link); link_id++) {
+		struct iwl_mvm_link_sta *link_sta;
+
+		/* unless we have a specific link in mind */
+		if (filter_link_id >= 0 && link_id != filter_link_id)
+			continue;
+
+		link_sta =
+			rcu_dereference_check(mvmsta->link[link_id],
+					      lockdep_is_held(&mvm->mutex));
+		if (!link_sta)
+			continue;
+
+		result |= BIT(link_sta->sta_id);
+	}
+
+	return result;
+}
 
 static int iwl_mvm_mld_send_sta_cmd(struct iwl_mvm *mvm,
 				    struct iwl_mvm_sta_cfg_cmd *cmd)

drivers/net/wireless/intel/iwlwifi/mvm/sta.h

@@ -642,6 +642,8 @@ int iwl_mvm_mld_update_sta_links(struct iwl_mvm *mvm,
 				 struct ieee80211_vif *vif,
 				 struct ieee80211_sta *sta,
 				 u16 old_links, u16 new_links);
+u32 iwl_mvm_sta_fw_id_mask(struct iwl_mvm *mvm, struct ieee80211_sta *sta,
+			   int filter_link_id);
 
 /* Queues */
 void iwl_mvm_mld_modify_all_sta_disable_tx(struct iwl_mvm *mvm,