From 2f911fb1555e657784355d07ec1e12979b84221a Mon Sep 17 00:00:00 2001
From: Veerendranath Jakkam <quic_vjakkam@quicinc.com>
Date: Wed, 27 Sep 2023 11:27:13 +0530
Subject: [PATCH] SAE: Remove current PMKSA from driver after reauth threshold
 is passed

wpa_supplicant postpones expired PMKSA deletion untillassociation is
lost for SAE to avoid forced disconnection. But during this time the
driver may use the expired PMKSA for reassociation with the current
connected AP.

Remove the current PMKSA for SAE from the driver after reauth threshold
is passed when the driver takes care of BSS selection.

Signed-off-by: Veerendranath Jakkam <quic_vjakkam@quicinc.com>
---
 src/rsn_supp/pmksa_cache.c      | 22 +++++++++++++++++++++-
 src/rsn_supp/pmksa_cache.h      |  7 +++++++
 src/rsn_supp/wpa.c              |  8 ++++++++
 src/rsn_supp/wpa.h              |  2 ++
 src/rsn_supp/wpa_i.h            |  1 +
 wpa_supplicant/wpa_supplicant.c |  3 +++
 6 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/src/rsn_supp/pmksa_cache.c b/src/rsn_supp/pmksa_cache.c
index 8dd770f38..6c8775150 100644
--- a/src/rsn_supp/pmksa_cache.c
+++ b/src/rsn_supp/pmksa_cache.c
@@ -132,6 +132,23 @@ static void pmksa_cache_reauth(void *eloop_ctx, void *timeout_ctx)
 	if (!pmksa->sm)
 		return;
 
+	if (pmksa->sm->driver_bss_selection) {
+		struct rsn_pmksa_cache_entry *entry;
+
+		entry = pmksa->sm->cur_pmksa ?
+			pmksa->sm->cur_pmksa :
+			pmksa_cache_get(pmksa, pmksa->sm->bssid, NULL, NULL,
+					NULL, 0);
+		if (entry && wpa_key_mgmt_sae(entry->akmp)) {
+			wpa_printf(MSG_DEBUG,
+				   "RSN: remove reauth threshold passed PMKSA from the driver for SAE");
+			entry->sae_reauth_scheduled = true;
+			wpa_sm_remove_pmkid(pmksa->sm, entry->network_ctx,
+					    entry->aa, entry->pmkid, NULL);
+			return;
+		}
+	}
+
 	pmksa->sm->cur_pmksa = NULL;
 	eapol_sm_request_reauth(pmksa->sm->eapol);
 }
@@ -178,7 +195,10 @@ static void pmksa_cache_set_expiration(struct rsn_pmksa_cache *pmksa)
 
 	entry = pmksa->sm->cur_pmksa ? pmksa->sm->cur_pmksa :
 		pmksa_cache_get(pmksa, pmksa->sm->bssid, NULL, NULL, NULL, 0);
-	if (entry && !wpa_key_mgmt_sae(entry->akmp)) {
+	if (entry &&
+	    (!wpa_key_mgmt_sae(entry->akmp) ||
+	     (pmksa->sm->driver_bss_selection &&
+	      !entry->sae_reauth_scheduled))) {
 		sec = pmksa->pmksa->reauth_time - now.sec;
 		if (sec < 0)
 			sec = 0;
diff --git a/src/rsn_supp/pmksa_cache.h b/src/rsn_supp/pmksa_cache.h
index 37c116282..08af2e63b 100644
--- a/src/rsn_supp/pmksa_cache.h
+++ b/src/rsn_supp/pmksa_cache.h
@@ -45,6 +45,13 @@ struct rsn_pmksa_cache_entry {
 	void *network_ctx;
 	int opportunistic;
 	bool external;
+
+	/**
+	 * This field is used to avoid duplicate pmksa_cache_reauth() calls for
+	 * every 10 minutes during the periodic expiration check of the current
+	 * PMKSA for SAE.
+	 */
+	bool sae_reauth_scheduled;
 };
 
 struct rsn_pmksa_cache;
diff --git a/src/rsn_supp/wpa.c b/src/rsn_supp/wpa.c
index 2867d60af..856fe09e6 100644
--- a/src/rsn_supp/wpa.c
+++ b/src/rsn_supp/wpa.c
@@ -6511,3 +6511,11 @@ void wpa_sm_set_cur_pmksa(struct wpa_sm *sm,
 	if (sm)
 		sm->cur_pmksa = entry;
 }
+
+
+void wpa_sm_set_driver_bss_selection(struct wpa_sm *sm,
+				     bool driver_bss_selection)
+{
+	if (sm)
+		sm->driver_bss_selection = driver_bss_selection;
+}
diff --git a/src/rsn_supp/wpa.h b/src/rsn_supp/wpa.h
index bf7badb84..47a86b04b 100644
--- a/src/rsn_supp/wpa.h
+++ b/src/rsn_supp/wpa.h
@@ -611,5 +611,7 @@ struct rsn_pmksa_cache * wpa_sm_get_pmksa_cache(struct wpa_sm *sm);
 void wpa_sm_set_cur_pmksa(struct wpa_sm *sm,
 			  struct rsn_pmksa_cache_entry *entry);
 const u8 * wpa_sm_get_auth_addr(struct wpa_sm *sm);
+void wpa_sm_set_driver_bss_selection(struct wpa_sm *sm,
+				     bool driver_bss_selection);
 
 #endif /* WPA_H */
diff --git a/src/rsn_supp/wpa_i.h b/src/rsn_supp/wpa_i.h
index a0c135ec6..5fe6182ff 100644
--- a/src/rsn_supp/wpa_i.h
+++ b/src/rsn_supp/wpa_i.h
@@ -222,6 +222,7 @@ struct wpa_sm {
 	struct wpa_sm_mlo mlo;
 
 	bool wmm_enabled;
+	bool driver_bss_selection;
 };
 
 
diff --git a/wpa_supplicant/wpa_supplicant.c b/wpa_supplicant/wpa_supplicant.c
index 69f228919..ba68e8198 100644
--- a/wpa_supplicant/wpa_supplicant.c
+++ b/wpa_supplicant/wpa_supplicant.c
@@ -7161,6 +7161,9 @@ static int wpa_supplicant_init_iface(struct wpa_supplicant *wpa_s,
 #ifdef CONFIG_PASN
 	wpa_pasn_sm_set_caps(wpa_s->wpa, wpa_s->drv_flags2);
 #endif /* CONFIG_PASN */
+	wpa_sm_set_driver_bss_selection(wpa_s->wpa,
+					!!(wpa_s->drv_flags &
+					   WPA_DRIVER_FLAGS_BSS_SELECTION));
 	if (wpa_s->max_remain_on_chan == 0)
 		wpa_s->max_remain_on_chan = 1000;