From e5091bb7af9e6532af1f23abca6b7baeda1f2e6d Mon Sep 17 00:00:00 2001
From: Kapil Gupta <kapil.gupta@espressif.com>
Date: Thu, 28 Oct 2021 14:23:34 +0530
Subject: [PATCH] wpa_supplicant: Optimize BTM request scan

---
 components/wpa_supplicant/src/common/scan.c   |  4 +++
 .../wpa_supplicant/src/common/wnm_sta.c       | 29 ++++++++++++++-----
 2 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/components/wpa_supplicant/src/common/scan.c b/components/wpa_supplicant/src/common/scan.c
index 7625de274f..46dab15be1 100644
--- a/components/wpa_supplicant/src/common/scan.c
+++ b/components/wpa_supplicant/src/common/scan.c
@@ -38,6 +38,10 @@ void wpa_supplicant_req_scan(struct wpa_supplicant *wpa_s, int sec, int usec)
 	}
 	params = os_zalloc(sizeof(*params));
 
+	if (!params) {
+		wpa_printf(MSG_ERROR, "Memory allocation failed");
+		return;
+	}
 	if (wpa_s->wnm_mode) {
 		/* Use the same memory */
 		params->ssids[0].ssid = wpa_s->current_bss->ssid;
diff --git a/components/wpa_supplicant/src/common/wnm_sta.c b/components/wpa_supplicant/src/common/wnm_sta.c
index 33a58b7351..9a05d57b30 100644
--- a/components/wpa_supplicant/src/common/wnm_sta.c
+++ b/components/wpa_supplicant/src/common/wnm_sta.c
@@ -580,15 +580,15 @@ static void wnm_send_bss_transition_mgmt_resp(
 	struct wpabuf *buf;
 	int res;
 
-	wpa_printf(MSG_DEBUG,
-		   "WNM: Send BSS Transition Management Response to " MACSTR
-		   " dialog_token=%u status=%u reason=%u delay=%d",
-		   MAC2STR(wpa_s->current_bss->bssid), dialog_token, status, reason, delay);
 	if (!wpa_s->current_bss) {
 		wpa_printf(MSG_DEBUG,
 			   "WNM: Current BSS not known - drop response");
 		return;
 	}
+	wpa_printf(MSG_DEBUG,
+		   "WNM: Send BSS Transition Management Response to " MACSTR
+		   " dialog_token=%u status=%u reason=%u delay=%d",
+		   MAC2STR(wpa_s->current_bss->bssid), dialog_token, status, reason, delay);
 
 	buf = wpabuf_alloc(BTM_RESP_MIN_SIZE);
 	if (!buf) {
@@ -784,7 +784,7 @@ static void wnm_dump_cand_list(struct wpa_supplicant *wpa_s)
 static void wnm_set_scan_freqs(struct wpa_supplicant *wpa_s)
 {
 	unsigned int i;
-	int num_chan;
+	int num_chan = 0;
 	u8 chan = 0;
 
 	wpa_s->next_scan_chan = 0;
@@ -802,8 +802,10 @@ static void wnm_set_scan_freqs(struct wpa_supplicant *wpa_s)
 				   MAC2STR(nei->bssid));
 			return;
 		}
-		if (nei->channel_number != chan)
+		if (nei->channel_number != chan) {
+			chan = nei->channel_number;
 			num_chan++;
+		}
 	}
 
 	if (num_chan == 1) {
@@ -826,6 +828,9 @@ static void ieee802_11_rx_bss_trans_mgmt_req(struct wpa_supplicant *wpa_s,
 #ifdef CONFIG_MBO
 	const u8 *vendor;
 #endif /* CONFIG_MBO */
+#ifdef ESP_SUPPLICANT
+	bool scan_required = false;
+#endif
 
 	if (wpa_s->disable_mbo_oce || wpa_s->disable_btm)
 		return;
@@ -904,9 +909,11 @@ static void ieee802_11_rx_bss_trans_mgmt_req(struct wpa_supplicant *wpa_s,
 #ifdef ESP_SUPPLICANT
 			os_memset(wpa_s->next_scan_bssid, 0, ETH_ALEN);
 			wpa_s->next_scan_chan = 0;
-#endif
+			scan_required = true;
+#else
 			wpa_printf(MSG_DEBUG, "Trying to find another BSS");
 			wpa_supplicant_req_scan(wpa_s, 0, 0);
+#endif
 		}
 	}
 
@@ -1018,6 +1025,12 @@ static void ieee802_11_rx_bss_trans_mgmt_req(struct wpa_supplicant *wpa_s,
 				   "WNM: Scan only for a specific BSSID since there is only a single candidate "
 				   MACSTR, MAC2STR(wpa_s->next_scan_bssid));
 		}
+#ifdef ESP_SUPPLICANT
+		scan_required = true;
+	}
+	if (scan_required) {
+		wpa_printf(MSG_DEBUG, "Trying to find another BSS");
+#endif
 		wpa_supplicant_req_scan(wpa_s, 0, 0);
 	} else if (reply) {
 		enum bss_trans_mgmt_status_code status;
@@ -1108,7 +1121,7 @@ void ieee802_11_rx_wnm_action(struct wpa_supplicant *wpa_s,
 	switch (act) {
 	case WNM_BSS_TRANS_MGMT_REQ:
 		ieee802_11_rx_bss_trans_mgmt_req(wpa_s, pos, end,
-						 !(sender[0] & 0x01));
+						 0x01);
 		break;
 	default:
 		wpa_printf(MSG_ERROR, "WNM: Unknown request");