From 18052cc0fc29f7ff327e0202b25d0c58b8429535 Mon Sep 17 00:00:00 2001
From: Shyamal Khachane <shyamal.khachane@espressif.com>
Date: Thu, 5 Jan 2023 14:19:30 +0530
Subject: [PATCH] wpa_supplicant: Set PMK from PMKSA incase of caching

---
 components/wpa_supplicant/src/common/defs.h  |  7 ++++
 components/wpa_supplicant/src/rsn_supp/wpa.c | 36 ++++++++------------
 2 files changed, 21 insertions(+), 22 deletions(-)

diff --git a/components/wpa_supplicant/src/common/defs.h b/components/wpa_supplicant/src/common/defs.h
index eabd2c0681..03ce20de03 100644
--- a/components/wpa_supplicant/src/common/defs.h
+++ b/components/wpa_supplicant/src/common/defs.h
@@ -118,6 +118,13 @@ static inline int wpa_key_mgmt_cckm(int akm)
 	return akm == WPA_KEY_MGMT_CCKM;
 }
 
+#ifdef ESP_SUPPLICANT
+static inline int wpa_key_mgmt_supports_caching(int akm)
+{
+        return wpa_key_mgmt_wpa_ieee8021x(akm) ||
+		wpa_key_mgmt_sae(akm);
+}
+#endif
 
 #define WPA_PROTO_WPA BIT(0)
 #define WPA_PROTO_RSN BIT(1)
diff --git a/components/wpa_supplicant/src/rsn_supp/wpa.c b/components/wpa_supplicant/src/rsn_supp/wpa.c
index 27ca0fca05..9390128645 100644
--- a/components/wpa_supplicant/src/rsn_supp/wpa.c
+++ b/components/wpa_supplicant/src/rsn_supp/wpa.c
@@ -2326,7 +2326,7 @@ int wpa_set_bss(char *macddr, char * bssid, u8 pairwise_cipher, u8 group_cipher,
 {
     int res = 0;
     struct wpa_sm *sm = &gWpaSm;
-    bool use_pmk_cache = true;
+    bool use_pmk_cache = !esp_wifi_skip_supp_pmkcaching();
     u8 assoc_rsnxe[20];
     size_t assoc_rsnxe_len = sizeof(assoc_rsnxe);
 
@@ -2351,27 +2351,19 @@ int wpa_set_bss(char *macddr, char * bssid, u8 pairwise_cipher, u8 group_cipher,
     sm->use_ext_key_id = (sm->proto == WPA_PROTO_WPA);
     pmksa_cache_clear_current(sm);
 
-    if (sm->key_mgmt == WPA_KEY_MGMT_SAE ||
-        is_wpa2_enterprise_connection()) {
-        if (!esp_wifi_skip_supp_pmkcaching() && use_pmk_cache) {
-            if (pmksa_cache_set_current(sm, NULL, (const u8*) bssid, 0, 0) == 0) {
-                struct rsn_pmksa_cache_entry *pmksa = pmksa_cache_get_current(sm);
-                if (pmksa && (pmksa->akmp != sm->key_mgmt)) {
-                    pmksa_cache_clear_current(sm);
-                    pmksa_cache_flush(sm->pmksa, NULL, pmksa->pmk, pmksa->pmk_len);
-                }
-            } else {
-                wpa_sm_set_pmk_from_pmksa(sm);
-            }
-        } else {
-            struct rsn_pmksa_cache_entry *entry = NULL;
-
-            if (sm->pmksa) {
-                entry = pmksa_cache_get(sm->pmksa, (const u8 *)bssid, NULL, NULL);
-            }
-            if (entry) {
-                pmksa_cache_flush(sm->pmksa, NULL, entry->pmk, entry->pmk_len);
-            }
+    struct rsn_pmksa_cache_entry *pmksa = NULL;
+    if (use_pmk_cache) {
+        pmksa = pmksa_cache_get(sm->pmksa, (const u8 *)bssid, NULL, NULL);
+        if (pmksa && (pmksa->akmp != sm->key_mgmt)) {
+            use_pmk_cache = false;
+        }
+    }
+    if (wpa_key_mgmt_supports_caching(sm->key_mgmt) && use_pmk_cache) {
+        pmksa_cache_set_current(sm, NULL, (const u8*) bssid, 0, 0);
+        wpa_sm_set_pmk_from_pmksa(sm);
+    } else {
+        if (pmksa) {
+            pmksa_cache_flush(sm->pmksa, NULL, pmksa->pmk, pmksa->pmk_len);
         }
     }