Revert "[wpa_supplicant] cumilative patch from commit 4b755c967"

Revert submission 26533062-Supplicant_merge_June24

Reason for revert: https://b.corp.google.com/issues/349780869

Reverted changes: /q/submissionid:26533062-Supplicant_merge_June24

Change-Id: I6c9b7a4323fa7edde47617da6c1e0d8f6e6d5101
diff --git a/wpa_supplicant/wnm_sta.c b/wpa_supplicant/wnm_sta.c
index ea79ae6..775e75b 100644
--- a/wpa_supplicant/wnm_sta.c
+++ b/wpa_supplicant/wnm_sta.c
@@ -418,7 +418,7 @@
 }
 
 
-void wnm_btm_reset(struct wpa_supplicant *wpa_s)
+void wnm_deallocate_memory(struct wpa_supplicant *wpa_s)
 {
 	int i;
 
@@ -431,17 +431,8 @@
 	os_free(wpa_s->wnm_neighbor_report_elements);
 	wpa_s->wnm_neighbor_report_elements = NULL;
 
-	wpa_s->wnm_cand_valid_until.sec = 0;
-	wpa_s->wnm_cand_valid_until.usec = 0;
-
-	wpa_s->wnm_mode = 0;
-	wpa_s->wnm_dialog_token = 0;
-	wpa_s->wnm_reply = 0;
-
-#ifdef CONFIG_MBO
-	wpa_s->wnm_mbo_trans_reason_present = 0;
-	wpa_s->wnm_mbo_transition_reason = 0;
-#endif /* CONFIG_MBO */
+	wpabuf_free(wpa_s->coloc_intf_elems);
+	wpa_s->coloc_intf_elems = NULL;
 }
 
 
@@ -790,11 +781,22 @@
 			}
 		}
 
-		/*
-		 * TODO: Could consider allowing transition to another ESS if
-		 * PMF was enabled for the association.
-		 */
-		if (!wpa_scan_res_match(wpa_s, 0, target, wpa_s->current_ssid,
+		if (bss->ssid_len != target->ssid_len ||
+		    os_memcmp(bss->ssid, target->ssid, bss->ssid_len) != 0) {
+			/*
+			 * TODO: Could consider allowing transition to another
+			 * ESS if PMF was enabled for the association.
+			 */
+			wpa_printf(MSG_DEBUG, "Candidate BSS " MACSTR
+				   " (pref %d) in different ESS",
+				   MAC2STR(nei->bssid),
+				   nei->preference_present ? nei->preference :
+				   -1);
+			continue;
+		}
+
+		if (wpa_s->current_ssid &&
+		    !wpa_scan_res_match(wpa_s, 0, target, wpa_s->current_ssid,
 					1, 0)) {
 			wpa_printf(MSG_DEBUG, "Candidate BSS " MACSTR
 				   " (pref %d) does not match the current network profile",
@@ -804,6 +806,14 @@
 			continue;
 		}
 
+		if (wpa_is_bss_tmp_disallowed(wpa_s, target)) {
+			wpa_printf(MSG_DEBUG,
+				   "MBO: Candidate BSS " MACSTR
+				   " retry delay is not over yet",
+				   MAC2STR(nei->bssid));
+			continue;
+		}
+
 		if (target->level < bss->level && target->level < -80) {
 			wpa_printf(MSG_DEBUG, "Candidate BSS " MACSTR
 				   " (pref %d) does not have sufficient signal level (%d)",
@@ -1035,8 +1045,8 @@
 
 #define BTM_RESP_MIN_SIZE	5 + ETH_ALEN
 
-static int wnm_send_bss_transition_mgmt_resp(
-	struct wpa_supplicant *wpa_s,
+static void wnm_send_bss_transition_mgmt_resp(
+	struct wpa_supplicant *wpa_s, u8 dialog_token,
 	enum bss_trans_mgmt_status_code status,
 	enum mbo_transition_reject_reason reason,
 	u8 delay, const u8 *target_bssid)
@@ -1044,24 +1054,21 @@
 	struct wpabuf *buf;
 	int res;
 
-	wpa_s->wnm_reply = 0;
-
 	wpa_printf(MSG_DEBUG,
 		   "WNM: Send BSS Transition Management Response to " MACSTR
 		   " dialog_token=%u status=%u reason=%u delay=%d",
-		   MAC2STR(wpa_s->bssid), wpa_s->wnm_dialog_token, status,
-		   reason, delay);
+		   MAC2STR(wpa_s->bssid), dialog_token, status, reason, delay);
 	if (!wpa_s->current_bss) {
 		wpa_printf(MSG_DEBUG,
 			   "WNM: Current BSS not known - drop response");
-		return -1;
+		return;
 	}
 
 	buf = wpabuf_alloc(BTM_RESP_MIN_SIZE);
 	if (!buf) {
 		wpa_printf(MSG_DEBUG,
 			   "WNM: Failed to allocate memory for BTM response");
-		return -1;
+		return;
 	}
 
 	wpa_s->bss_tm_status = status;
@@ -1069,7 +1076,7 @@
 
 	wpabuf_put_u8(buf, WLAN_ACTION_WNM);
 	wpabuf_put_u8(buf, WNM_BSS_TRANS_MGMT_RESP);
-	wpabuf_put_u8(buf, wpa_s->wnm_dialog_token);
+	wpabuf_put_u8(buf, dialog_token);
 	wpabuf_put_u8(buf, status);
 	wpabuf_put_u8(buf, delay);
 	if (target_bssid) {
@@ -1083,7 +1090,7 @@
 		wpabuf_put_data(buf, "\0\0\0\0\0\0", ETH_ALEN);
 	}
 
-	if (status == WNM_BSS_TM_ACCEPT && target_bssid)
+	if (status == WNM_BSS_TM_ACCEPT && !wpa_s->wnm_link_removal)
 		wnm_add_cand_list(wpa_s, &buf);
 
 #ifdef CONFIG_MBO
@@ -1099,7 +1106,7 @@
 				wpabuf_free(buf);
 				wpa_printf(MSG_DEBUG,
 					   "WNM: Failed to allocate memory for MBO IE");
-				return -1;
+				return;
 			}
 
 			wpabuf_put_data(buf, mbo, ret);
@@ -1116,8 +1123,6 @@
 	}
 
 	wpabuf_free(buf);
-
-	return res;
 }
 
 
@@ -1135,24 +1140,19 @@
 
 	/* Send the BSS Management Response - Accept */
 	if (wpa_s->wnm_reply) {
-		wpa_s->wnm_target_bss = bss;
+		wpa_s->wnm_reply = 0;
 		wpa_printf(MSG_DEBUG,
 			   "WNM: Sending successful BSS Transition Management Response");
-
-		/* This function will be called again from the TX handler to
-		 * start the actual reassociation after this response has been
-		 * delivered to the current AP. */
-		if (wnm_send_bss_transition_mgmt_resp(
-			    wpa_s, WNM_BSS_TM_ACCEPT,
-			    MBO_TRANSITION_REJECT_REASON_UNSPECIFIED, 0,
-			    bss->bssid) >= 0)
-			return;
+		wnm_send_bss_transition_mgmt_resp(
+			wpa_s, wpa_s->wnm_dialog_token, WNM_BSS_TM_ACCEPT,
+			MBO_TRANSITION_REJECT_REASON_UNSPECIFIED, 0,
+			bss->bssid);
 	}
 
 	if (bss == wpa_s->current_bss) {
 		wpa_printf(MSG_DEBUG,
 			   "WNM: Already associated with the preferred candidate");
-		wnm_btm_reset(wpa_s);
+		wnm_deallocate_memory(wpa_s);
 		return;
 	}
 
@@ -1168,10 +1168,11 @@
 	 */
 	if (!already_connecting && radio_work_pending(wpa_s, "sme-connect"))
 		wpa_s->bss_trans_mgmt_in_progress = true;
+	wnm_deallocate_memory(wpa_s);
 }
 
 
-int wnm_scan_process(struct wpa_supplicant *wpa_s, bool pre_scan_check)
+int wnm_scan_process(struct wpa_supplicant *wpa_s, int reply_on_fail)
 {
 	struct wpa_bss *bss;
 	struct wpa_ssid *ssid = wpa_s->current_ssid;
@@ -1179,51 +1180,27 @@
 	enum mbo_transition_reject_reason reason =
 		MBO_TRANSITION_REJECT_REASON_UNSPECIFIED;
 
-	if (!wpa_s->wnm_dialog_token)
+	if (!wpa_s->wnm_neighbor_report_elements)
 		return 0;
 
 	wpa_dbg(wpa_s, MSG_DEBUG,
 		"WNM: Process scan results for BSS Transition Management");
-	if (!pre_scan_check &&
-	    os_reltime_initialized(&wpa_s->wnm_cand_valid_until) &&
-	    os_reltime_before(&wpa_s->wnm_cand_valid_until,
+	if (os_reltime_before(&wpa_s->wnm_cand_valid_until,
 			      &wpa_s->scan_trigger_time)) {
 		wpa_printf(MSG_DEBUG, "WNM: Previously stored BSS transition candidate list is not valid anymore - drop it");
-		goto send_bss_resp_fail;
+		wnm_deallocate_memory(wpa_s);
+		return 0;
+	}
+
+	if (!wpa_s->current_bss ||
+	    !ether_addr_equal(wpa_s->wnm_cand_from_bss,
+			      wpa_s->current_bss->bssid)) {
+		wpa_printf(MSG_DEBUG, "WNM: Stored BSS transition candidate list not from the current BSS - ignore it");
+		return 0;
 	}
 
 	/* Compare the Neighbor Report and scan results */
 	bss = compare_scan_neighbor_results(wpa_s, 0, &reason);
-
-	/*
-	 * If this is a pre-scan check, returning 0 will trigger a scan and
-	 * another call. In that case, reject "bad" candidates in the hope of
-	 * finding a better candidate after scanning.
-	 *
-	 * Use a simple heuristic to check whether the selection is reasonable
-	 * or a scan is a good idea. For that, we need to have found a
-	 * candidate BSS (which might be the current one), it is up-to-date,
-	 * and we don't want to immediately roam back again.
-	 */
-	if (pre_scan_check) {
-		struct os_reltime age;
-
-		if (!bss)
-			return 0;
-
-		os_reltime_age(&bss->last_update, &age);
-		if (age.sec >= 10)
-			return 0;
-
-#ifndef CONFIG_NO_ROAMING
-		if (wpa_s->current_bss && bss != wpa_s->current_bss &&
-		    wpa_supplicant_need_to_roam_within_ess(wpa_s,
-							   wpa_s->current_bss,
-							   bss))
-			return 0;
-#endif /* CONFIG_NO_ROAMING */
-	}
-
 	if (!bss) {
 		wpa_printf(MSG_DEBUG, "WNM: No BSS transition candidate match found");
 		status = WNM_BSS_TM_REJECT_NO_SUITABLE_CANDIDATES;
@@ -1235,13 +1212,18 @@
 	return 1;
 
 send_bss_resp_fail:
+	if (!reply_on_fail)
+		return 0;
+
 	/* Send reject response for all the failures */
 
-	if (wpa_s->wnm_reply)
-		wnm_send_bss_transition_mgmt_resp(wpa_s, status, reason,
-						  0, NULL);
-
-	wnm_btm_reset(wpa_s);
+	if (wpa_s->wnm_reply) {
+		wpa_s->wnm_reply = 0;
+		wnm_send_bss_transition_mgmt_resp(wpa_s,
+						  wpa_s->wnm_dialog_token,
+						  status, reason, 0, NULL);
+	}
+	wnm_deallocate_memory(wpa_s);
 
 	return 0;
 }
@@ -1348,10 +1330,6 @@
 		struct neighbor_report *nei;
 
 		nei = &wpa_s->wnm_neighbor_report_elements[i];
-
-		if (nei->preference_present && nei->preference == 0)
-			continue;
-
 		if (nei->freq <= 0) {
 			wpa_printf(MSG_DEBUG,
 				   "WNM: Unknown neighbor operating frequency for "
@@ -1376,6 +1354,79 @@
 }
 
 
+static int wnm_fetch_scan_results(struct wpa_supplicant *wpa_s)
+{
+	struct wpa_scan_results *scan_res;
+	struct wpa_bss *bss;
+	struct wpa_ssid *ssid = wpa_s->current_ssid;
+	u8 i, found = 0;
+	size_t j;
+
+	wpa_dbg(wpa_s, MSG_DEBUG,
+		"WNM: Fetch current scan results from the driver for checking transition candidates");
+	scan_res = wpa_drv_get_scan_results2(wpa_s);
+	if (!scan_res) {
+		wpa_dbg(wpa_s, MSG_DEBUG, "WNM: Failed to get scan results");
+		return 0;
+	}
+
+	if (scan_res->fetch_time.sec == 0)
+		os_get_reltime(&scan_res->fetch_time);
+
+	filter_scan_res(wpa_s, scan_res);
+
+	for (i = 0; i < wpa_s->wnm_num_neighbor_report; i++) {
+		struct neighbor_report *nei;
+
+		nei = &wpa_s->wnm_neighbor_report_elements[i];
+		if (nei->preference_present && nei->preference == 0)
+			continue;
+
+		for (j = 0; j < scan_res->num; j++) {
+			struct wpa_scan_res *res;
+			const u8 *ssid_ie;
+
+			res = scan_res->res[j];
+			if (!ether_addr_equal(nei->bssid, res->bssid) ||
+			    res->age > WNM_SCAN_RESULT_AGE * 1000)
+				continue;
+			bss = wpa_s->current_bss;
+			ssid_ie = wpa_scan_get_ie(res, WLAN_EID_SSID);
+			if (bss && ssid_ie && ssid_ie[1] &&
+			    (bss->ssid_len != ssid_ie[1] ||
+			     os_memcmp(bss->ssid, ssid_ie + 2,
+				       bss->ssid_len) != 0))
+				continue; /* Skip entries for other ESSs */
+
+			/* Potential candidate found */
+			found = 1;
+			scan_snr(res);
+			scan_est_throughput(wpa_s, res);
+			wpa_bss_update_scan_res(wpa_s, res,
+						&scan_res->fetch_time);
+		}
+	}
+
+	wpa_scan_results_free(scan_res);
+	if (!found) {
+		wpa_dbg(wpa_s, MSG_DEBUG,
+			"WNM: No transition candidate matches existing scan results");
+		return 0;
+	}
+
+	bss = compare_scan_neighbor_results(wpa_s, WNM_SCAN_RESULT_AGE, NULL);
+	if (!bss) {
+		wpa_dbg(wpa_s, MSG_DEBUG,
+			"WNM: Comparison of scan results against transition candidates did not find matches");
+		return 0;
+	}
+
+	/* Associate to the network */
+	wnm_bss_tm_connect(wpa_s, bss, ssid, 0);
+	return 1;
+}
+
+
 static void ieee802_11_rx_bss_trans_mgmt_req(struct wpa_supplicant *wpa_s,
 					     const u8 *pos, const u8 *end,
 					     int reply)
@@ -1394,6 +1445,8 @@
 		return;
 
 #ifdef CONFIG_MBO
+	wpa_s->wnm_mbo_trans_reason_present = 0;
+	wpa_s->wnm_mbo_transition_reason = 0;
 	wpa_s->wnm_mbo_cell_pref_present = 0;
 	wpa_s->wnm_mbo_cell_preference = 0;
 	wpa_s->wnm_mbo_assoc_retry_delay_present = 0;
@@ -1405,8 +1458,6 @@
 	else
 		beacon_int = 100; /* best guess */
 
-	wnm_btm_reset(wpa_s);
-
 	wpa_s->wnm_dialog_token = pos[0];
 	wpa_s->wnm_mode = pos[1];
 	wpa_s->wnm_dissoc_timer = WPA_GET_LE16(pos + 2);
@@ -1426,7 +1477,8 @@
 			   "WNM: Testing - reject BSS Transition Management Request: reject_btm_req_reason=%d",
 			   wpa_s->reject_btm_req_reason);
 		wnm_send_bss_transition_mgmt_resp(
-			wpa_s, wpa_s->reject_btm_req_reason,
+			wpa_s, wpa_s->wnm_dialog_token,
+			wpa_s->reject_btm_req_reason,
 			MBO_TRANSITION_REJECT_REASON_UNSPECIFIED, 0, NULL);
 		return;
 	}
@@ -1489,34 +1541,15 @@
 	 * set to 1, and the BSS Termination Included field is set to 1, only
 	 * one of the links is removed and the other links remain associated.
 	 * Ignore the Disassociation Imminent field in such a case.
-	 *
-	 * TODO: We should check if the AP has more than one link.
-	 * TODO: We should pass the RX link and use that
 	 */
-	if (disassoc_imminent && wpa_s->valid_links &&
+	if (disassoc_imminent &&
+	    (wpa_s->valid_links & (wpa_s->valid_links - 1)) != 0 &&
 	    (wpa_s->wnm_mode & WNM_BSS_TM_REQ_LINK_REMOVAL_IMMINENT) &&
 	    (wpa_s->wnm_mode & WNM_BSS_TM_REQ_BSS_TERMINATION_INCLUDED)) {
-		/* If we still have a link, then just accept the request */
-		if (wpa_s->valid_links & (wpa_s->valid_links - 1)) {
-			wpa_printf(MSG_INFO,
-				   "WNM: BTM request for a single MLO link - ignore disassociation imminent since other links remain associated");
-			disassoc_imminent = false;
-
-			wnm_send_bss_transition_mgmt_resp(
-				wpa_s, WNM_BSS_TM_ACCEPT, 0, 0, NULL);
-
-			return;
-		}
-
-		/* The last link is being removed (which must be the assoc link)
-		 */
+		wpa_printf(MSG_INFO,
+			   "WNM: BTM request for a single MLO link - ignore disassociation imminent since other links remain associated");
+		disassoc_imminent = false;
 		wpa_s->wnm_link_removal = true;
-		os_memcpy(wpa_s->wnm_dissoc_addr,
-			  wpa_s->links[wpa_s->mlo_assoc_link_id].bssid,
-			  ETH_ALEN);
-	} else {
-		os_memcpy(wpa_s->wnm_dissoc_addr, wpa_s->valid_links ?
-			  wpa_s->ap_mld_addr : wpa_s->bssid, ETH_ALEN);
 	}
 
 	if (disassoc_imminent) {
@@ -1533,6 +1566,7 @@
 		unsigned int valid_ms;
 
 		wpa_msg(wpa_s, MSG_INFO, "WNM: Preferred List Available");
+		wnm_deallocate_memory(wpa_s);
 		wpa_s->wnm_neighbor_report_elements = os_calloc(
 			WNM_MAX_NEIGHBOR_REPORT,
 			sizeof(struct neighbor_report));
@@ -1580,7 +1614,8 @@
 			wpa_printf(MSG_DEBUG,
 				   "WNM: Candidate list included bit is set, but no candidates found");
 			wnm_send_bss_transition_mgmt_resp(
-				wpa_s, WNM_BSS_TM_REJECT_NO_SUITABLE_CANDIDATES,
+				wpa_s, wpa_s->wnm_dialog_token,
+				WNM_BSS_TM_REJECT_NO_SUITABLE_CANDIDATES,
 				MBO_TRANSITION_REJECT_REASON_UNSPECIFIED, 0,
 				NULL);
 			return;
@@ -1590,7 +1625,8 @@
 			wpa_printf(MSG_DEBUG,
 				   "WNM: Configuration prevents roaming (BSSID set)");
 			wnm_send_bss_transition_mgmt_resp(
-				wpa_s, WNM_BSS_TM_REJECT_NO_SUITABLE_CANDIDATES,
+				wpa_s, wpa_s->wnm_dialog_token,
+				WNM_BSS_TM_REJECT_NO_SUITABLE_CANDIDATES,
 				MBO_TRANSITION_REJECT_REASON_UNSPECIFIED, 0,
 				NULL);
 			return;
@@ -1607,21 +1643,35 @@
 		wpa_s->wnm_cand_valid_until.sec +=
 			wpa_s->wnm_cand_valid_until.usec / 1000000;
 		wpa_s->wnm_cand_valid_until.usec %= 1000000;
+		os_memcpy(wpa_s->wnm_cand_from_bss, wpa_s->bssid, ETH_ALEN);
 
 		/*
-		* Try fetching the latest scan results from the kernel.
-		* This can help in finding more up-to-date information should
-		* the driver have done some internal scanning operations after
-		* the last scan result update in wpa_supplicant.
-		*
-		* It is not a new scan, this does not update the last_scan
-		* timestamp nor will it expire old BSSs.
-		*/
-		wpa_supplicant_update_scan_results(wpa_s, NULL);
-		if (wnm_scan_process(wpa_s, true) > 0)
+		 * Fetch the latest scan results from the kernel and check for
+		 * candidates based on those results first. This can help in
+		 * finding more up-to-date information should the driver has
+		 * done some internal scanning operations after the last scan
+		 * result update in wpa_supplicant.
+		 */
+		if (wnm_fetch_scan_results(wpa_s) > 0)
 			return;
-		wpa_printf(MSG_DEBUG,
-			   "WNM: No valid match in previous scan results - try a new scan");
+
+		/*
+		 * Try to use previously received scan results, if they are
+		 * recent enough to use for a connection.
+		 */
+		if (wpa_s->last_scan_res_used > 0) {
+			struct os_reltime now;
+
+			os_get_reltime(&now);
+			if (!os_reltime_expired(&now, &wpa_s->last_scan, 10)) {
+				wpa_printf(MSG_DEBUG,
+					   "WNM: Try to use recent scan results");
+				if (wnm_scan_process(wpa_s, 0) > 0)
+					return;
+				wpa_printf(MSG_DEBUG,
+					   "WNM: No match in previous scan results - try a new scan");
+			}
+		}
 
 		wnm_set_scan_freqs(wpa_s);
 		if (wpa_s->wnm_num_neighbor_report == 1) {
@@ -1644,45 +1694,12 @@
 			status = WNM_BSS_TM_REJECT_UNSPECIFIED;
 		}
 		wnm_send_bss_transition_mgmt_resp(
-			wpa_s, status,
+			wpa_s, wpa_s->wnm_dialog_token, status,
 			MBO_TRANSITION_REJECT_REASON_UNSPECIFIED, 0, NULL);
 	}
 }
 
 
-int wnm_btm_resp_tx_status(struct wpa_supplicant *wpa_s, const u8 *data,
-			   size_t data_len)
-{
-	const struct ieee80211_mgmt *frame =
-		(const struct ieee80211_mgmt *) data;
-
-	if (data_len <
-	    IEEE80211_HDRLEN + sizeof(frame->u.action.u.bss_tm_resp) ||
-	    frame->u.action.category != WLAN_ACTION_WNM ||
-	    frame->u.action.u.bss_tm_resp.action != WNM_BSS_TRANS_MGMT_RESP ||
-	    frame->u.action.u.bss_tm_resp.status_code != WNM_BSS_TM_ACCEPT)
-		return -1;
-
-	/*
-	 * If disassoc imminent bit was set in the request, the response may
-	 * indicate accept even if no candidate was found, so bail out here.
-	 */
-	if (!wpa_s->wnm_target_bss) {
-		wpa_printf(MSG_DEBUG, "WNM: Target BSS is not set");
-		return 0;
-	}
-
-	if (!wpa_s->current_ssid)
-		return 0;
-
-	wnm_bss_tm_connect(wpa_s, wpa_s->wnm_target_bss, wpa_s->current_ssid,
-			   0);
-
-	wpa_s->wnm_target_bss = NULL;
-	return 0;
-}
-
-
 #define BTM_QUERY_MIN_SIZE	4
 
 int wnm_send_bss_transition_mgmt_query(struct wpa_supplicant *wpa_s,
@@ -2035,14 +2052,14 @@
 void wnm_set_coloc_intf_elems(struct wpa_supplicant *wpa_s,
 			      struct wpabuf *elems)
 {
+	wpabuf_free(wpa_s->coloc_intf_elems);
 	if (elems && wpabuf_len(elems) == 0) {
 		wpabuf_free(elems);
 		elems = NULL;
 	}
+	wpa_s->coloc_intf_elems = elems;
 
-	/* NOTE: The elements are not stored as they are only send out once */
-
-	if (wpa_s->conf->coloc_intf_reporting && elems &&
+	if (wpa_s->conf->coloc_intf_reporting && wpa_s->coloc_intf_elems &&
 	    wpa_s->coloc_intf_dialog_token &&
 	    (wpa_s->coloc_intf_auto_report == 1 ||
 	     wpa_s->coloc_intf_auto_report == 3)) {
@@ -2051,10 +2068,8 @@
 		 */
 		wnm_send_coloc_intf_report(wpa_s,
 					   wpa_s->coloc_intf_dialog_token,
-					   elems);
+					   wpa_s->coloc_intf_elems);
 	}
-
-	wpabuf_free(elems);
 }
 
 
@@ -2067,6 +2082,8 @@
 
 bool wnm_is_bss_excluded(struct wpa_supplicant *wpa_s, struct wpa_bss *bss)
 {
+	unsigned int i;
+
 	if (!(wpa_s->wnm_mode & WNM_BSS_TM_REQ_DISASSOC_IMMINENT))
 		return false;
 
@@ -2074,14 +2091,26 @@
 	 * In case disassociation imminent is set, do no try to use a BSS to
 	 * which we are connected.
 	 */
-	if (wpa_s->wnm_link_removal ||
-	    !(wpa_s->drv_flags2 & WPA_DRIVER_FLAGS2_MLO) ||
-	    is_zero_ether_addr(bss->mld_addr)) {
-		if (ether_addr_equal(bss->bssid, wpa_s->wnm_dissoc_addr))
+
+	if (wpa_s->current_bss &&
+	    ether_addr_equal(wpa_s->current_bss->bssid, bss->bssid)) {
+		wpa_dbg(wpa_s, MSG_DEBUG,
+			"WNM: Disassociation imminent: current BSS");
+		return true;
+	}
+
+	if (!wpa_s->valid_links)
+		return false;
+
+	for (i = 0; i < MAX_NUM_MLD_LINKS; i++) {
+		if (!(wpa_s->valid_links & BIT(i)))
+			continue;
+
+		if (ether_addr_equal(wpa_s->links[i].bssid, bss->bssid)) {
+			wpa_dbg(wpa_s, MSG_DEBUG,
+				"WNM: MLD: Disassociation imminent: current link");
 			return true;
-	} else {
-		if (ether_addr_equal(bss->mld_addr, wpa_s->wnm_dissoc_addr))
-			return true;
+		}
 	}
 
 	return false;