wifi: iwlwifi: mld: support iwl_omi_send_status_notif version 2

The firmware provides the station id, use it since it makes our lives
easier. No need to assume we have a single BSS vif, and look up the
station id to whom the OMI was sent.

Signed-off-by: Emmanuel Grumbach <emmanuel.grumbach@intel.com>
Signed-off-by: Miri Korenblit <miriam.rachel.korenblit@intel.com>
Link: https://patch.msgid.link/20250711183056.7d2cd878855f.I8625ebb2c4e1fb484aafd16a07549f2eeb506e08@changeid

3 files changed, 43 insertions, 3 deletions

diff --git a/drivers/net/wireless/intel/iwlwifi/fw/api/datapath.h b/drivers/net/wireless/intel/iwlwifi/fw/api/datapath.h
index 6c8e6874a5e7..ee822a87c42c 100644
--- a/drivers/net/wireless/intel/iwlwifi/fw/api/datapath.h
+++ b/drivers/net/wireless/intel/iwlwifi/fw/api/datapath.h
@@ -700,12 +700,23 @@ struct iwl_sec_key_cmd {
 } __packed; /* SEC_KEY_CMD_API_S_VER_1 */
 
 /**
+ * struct iwl_omi_send_status_notif_v1 - OMI status notification
+ * @success: indicates that the OMI was sent successfully
+ *	(currently always set)
+ */
+struct iwl_omi_send_status_notif_v1 {
+	__le32 success;
+} __packed; /* OMI_SEND_STATUS_NTFY_API_S_VER_1 */
+
+/**
  * struct iwl_omi_send_status_notif - OMI status notification
  * @success: indicates that the OMI was sent successfully
  *	(currently always set)
+ * @sta_id: sta_id to which the OMI was sent
  */
 struct iwl_omi_send_status_notif {
 	__le32 success;
-} __packed; /* OMI_SEND_STATUS_NTFY_API_S_VER_1 */
+	__le32 sta_id;
+} __packed; /* OMI_SEND_STATUS_NTFY_API_S_VER_2 */
 
 #endif /* __iwl_fw_api_datapath_h__ */
diff --git a/drivers/net/wireless/intel/iwlwifi/mld/link.c b/drivers/net/wireless/intel/iwlwifi/mld/link.c
index ff237f78a468..c48cc3909637 100644
--- a/drivers/net/wireless/intel/iwlwifi/mld/link.c
+++ b/drivers/net/wireless/intel/iwlwifi/mld/link.c
@@ -649,11 +649,39 @@ void iwl_mld_omi_ap_changed_bw(struct iwl_mld *mld,
 void iwl_mld_handle_omi_status_notif(struct iwl_mld *mld,
 				     struct iwl_rx_packet *pkt)
 {
+	int ver = iwl_fw_lookup_notif_ver(mld->fw, DATA_PATH_GROUP,
+					  OMI_SEND_STATUS_NOTIF, 1);
 	struct ieee80211_link_sta *link_sta;
 	struct iwl_mld_link *mld_link;
 	struct ieee80211_vif *vif;
 
-	vif = iwl_mld_get_omi_bw_reduction_pointers(mld, &link_sta, &mld_link);
+	if (ver == 2) {
+		const struct iwl_omi_send_status_notif *notif =
+			(const void *)pkt->data;
+		u32 sta_id = le32_to_cpu(notif->sta_id);
+		struct iwl_mld_vif *mld_vif;
+
+		if (IWL_FW_CHECK(mld, sta_id >= mld->fw->ucode_capa.num_stations,
+				 "Invalid station %d\n", sta_id))
+			return;
+
+		link_sta = wiphy_dereference(mld->wiphy,
+					     mld->fw_id_to_link_sta[sta_id]);
+		if (IWL_FW_CHECK(mld, !link_sta, "Station does not exist\n"))
+			return;
+
+		vif = iwl_mld_sta_from_mac80211(link_sta->sta)->vif;
+		mld_vif = iwl_mld_vif_from_mac80211(vif);
+
+		mld_link = iwl_mld_link_dereference_check(mld_vif,
+							  link_sta->link_id);
+		if (WARN(!mld_link, "Link %d does not exist\n",
+			 link_sta->link_id))
+			return;
+	} else {
+		vif = iwl_mld_get_omi_bw_reduction_pointers(mld, &link_sta,
+							    &mld_link);
+	}
 	if (IWL_FW_CHECK(mld, !vif, "unexpected OMI notification\n"))
 		return;
 
diff --git a/drivers/net/wireless/intel/iwlwifi/mld/notif.c b/drivers/net/wireless/intel/iwlwifi/mld/notif.c
index 4bfed90c30c5..262d8e25e62a 100644
--- a/drivers/net/wireless/intel/iwlwifi/mld/notif.c
+++ b/drivers/net/wireless/intel/iwlwifi/mld/notif.c
@@ -349,7 +349,8 @@ CMD_VERSIONS(time_msmt_notif,
 CMD_VERSIONS(time_sync_confirm_notif,
 	     CMD_VER_ENTRY(1, iwl_time_msmt_cfm_notify))
 CMD_VERSIONS(omi_status_notif,
-	     CMD_VER_ENTRY(1, iwl_omi_send_status_notif))
+	     CMD_VER_ENTRY(1, iwl_omi_send_status_notif_v1)
+	     CMD_VER_ENTRY(2, iwl_omi_send_status_notif))
 CMD_VERSIONS(ftm_resp_notif, CMD_VER_ENTRY(10, iwl_tof_range_rsp_ntfy))
 CMD_VERSIONS(beacon_filter_notif, CMD_VER_ENTRY(2, iwl_beacon_filter_notif))