wifi: iwlwifi: mvm: track AP STA pointer and use it for MFP

Instead of relying on the AP STA ID and dereferencing our
firmware -> mac80211 STA array when we really only need
the STA, not the per-link information (and might not have
it, e.g. for a pairwise key installation where this is
relevant), keep track of the AP STA as a pointer to the
mac80211 STA, protected by the mutex. Then use it here in
the key code instead.

Signed-off-by: Johannes Berg <johannes.berg@intel.com>
Signed-off-by: Gregory Greenman <gregory.greenman@intel.com>
Link: https://lore.kernel.org/r/20230414130637.7db01b1508de.I88e6d252cdb0a6e4581d7c5fd7cbc037b4fd40ae@changeid
Signed-off-by: Johannes Berg <johannes.berg@intel.com>

drivers/net/wireless/intel/iwlwifi/mvm/mac80211.c

@@ -3669,6 +3669,7 @@ iwl_mvm_sta_state_notexist_to_none(struct iwl_mvm *mvm,
 				   struct ieee80211_sta *sta,
 				   struct iwl_mvm_sta_state_ops *callbacks)
 {
+	struct iwl_mvm_vif *mvmvif = iwl_mvm_vif_from_mac80211(vif);
 	unsigned int i;
 	int ret;
 
@@ -3704,6 +3705,9 @@ iwl_mvm_sta_state_notexist_to_none(struct iwl_mvm *mvm,
 	}
 	ieee80211_sta_recalc_aggregates(sta);
 
+	if (vif->type == NL80211_IFTYPE_STATION && !sta->tdls)
+		mvmvif->ap_sta = sta;
+
 	return 0;
 }
 
@@ -3931,8 +3935,10 @@ int iwl_mvm_mac_sta_state_common(struct ieee80211_hw *hw,
 		ret = 0;
 	} else if (old_state == IEEE80211_STA_NONE &&
 		   new_state == IEEE80211_STA_NOTEXIST) {
-		if (vif->type == NL80211_IFTYPE_STATION && !sta->tdls)
+		if (vif->type == NL80211_IFTYPE_STATION && !sta->tdls) {
 			iwl_mvm_stop_session_protection(mvm, vif);
+			mvmvif->ap_sta = NULL;
+		}
 		ret = callbacks->rm_sta(mvm, vif, sta);
 		if (sta->tdls) {
 			iwl_mvm_recalc_tdls_state(mvm, vif, false);

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

@@ -41,6 +41,8 @@ static u32 iwl_mvm_get_sec_flags(struct iwl_mvm *mvm,
 	struct iwl_mvm_vif *mvmvif = iwl_mvm_vif_from_mac80211(vif);
 	u32 flags = 0;
 
+	lockdep_assert_held(&mvm->mutex);
+
 	if (!(keyconf->flags & IEEE80211_KEY_FLAG_PAIRWISE))
 		flags |= IWL_SEC_KEY_FLAG_MCAST_KEY;
 
@@ -68,18 +70,11 @@ static u32 iwl_mvm_get_sec_flags(struct iwl_mvm *mvm,
 		break;
 	}
 
-	rcu_read_lock();
-	if (!sta && vif->type == NL80211_IFTYPE_STATION &&
-	    mvmvif->deflink.ap_sta_id != IWL_MVM_INVALID_STA) {
-		u8 sta_id = mvmvif->deflink.ap_sta_id;
-
-		sta = rcu_dereference_check(mvm->fw_id_to_mac_id[sta_id],
-					    lockdep_is_held(&mvm->mutex));
-	}
+	if (!sta && vif->type == NL80211_IFTYPE_STATION)
+		sta = mvmvif->ap_sta;
 
 	if (!IS_ERR_OR_NULL(sta) && sta->mfp)
 		flags |= IWL_SEC_KEY_FLAG_MFP;
-	rcu_read_unlock();
 
 	return flags;
 }

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

@@ -434,6 +434,8 @@ struct iwl_mvm_vif {
 	/* TCP Checksum Offload */
 	netdev_features_t features;
 
+	struct ieee80211_sta *ap_sta;
+
 	/* we can only have 2 GTK + 2 IGTK active at a time */
 	struct ieee80211_key_conf *ap_early_keys[4];