From e9a07592fc40f4aed2c3f1e113896ed361c0cf2c Mon Sep 17 00:00:00 2001
From: Nachiket Kukade <nachiket.kukade@espressif.com>
Date: Wed, 8 Jul 2020 16:40:40 +0530
Subject: [PATCH] wpa_supplicant: Fix WPA3 and WPA2 transition related failures

1. If Device is connected to AP in WPA3-PSK mode, AP switching
security to WPA2-PSK causes connection failures even after reset.
Fix is to not store WPA3's PMK in NVS for caching.

2. AP switching back to WPA3 causes even more connection failures.
This is due to device not clearing Supplicant level PMK Cache when
it is no longer valid. Fix is to clear the Cache when 4-way handshake
fails and to check Key Mgmt of Cache before using.

3. When AP switches from WPA3 to WPA2, device's PMF config in
Supplicant remains enabled. This may cause failures during
4-way handshake. So clear PMF config in when PMF is no longer used.
---
 .../wpa_supplicant/src/esp_supplicant/esp_wpa2.c     |  3 ++-
 .../wpa_supplicant/src/esp_supplicant/esp_wpa3.c     |  3 +--
 .../wpa_supplicant/src/esp_supplicant/esp_wpa_main.c |  1 +
 components/wpa_supplicant/src/rsn_supp/wpa.c         | 12 ++++++++++++
 components/wpa_supplicant/src/rsn_supp/wpa.h         |  1 +
 5 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/components/wpa_supplicant/src/esp_supplicant/esp_wpa2.c b/components/wpa_supplicant/src/esp_supplicant/esp_wpa2.c
index 707139cd8..fb3105835 100644
--- a/components/wpa_supplicant/src/esp_supplicant/esp_wpa2.c
+++ b/components/wpa_supplicant/src/esp_supplicant/esp_wpa2.c
@@ -695,7 +695,8 @@ static int wpa2_start_eapol_internal(void)
     if (!sm) {
         return ESP_FAIL;
     }
-    if (wpa_sta_is_cur_pmksa_set()) {
+
+    if (wpa_sta_cur_pmksa_matches_akm()) {
         wpa_printf(MSG_DEBUG,
                 "RSN: PMKSA caching - do not send EAPOL-Start");
         return ESP_FAIL;
diff --git a/components/wpa_supplicant/src/esp_supplicant/esp_wpa3.c b/components/wpa_supplicant/src/esp_supplicant/esp_wpa3.c
index 8abeb795a..3a106b5ca 100644
--- a/components/wpa_supplicant/src/esp_supplicant/esp_wpa3.c
+++ b/components/wpa_supplicant/src/esp_supplicant/esp_wpa3.c
@@ -32,7 +32,7 @@ static esp_err_t wpa3_build_sae_commit(u8 *bssid)
     u8 own_addr[ETH_ALEN];
     const u8 *pw;
 
-    if (wpa_sta_is_cur_pmksa_set()) {
+    if (wpa_sta_cur_pmksa_matches_akm()) {
         wpa_printf(MSG_INFO, "wpa3: Skip SAE and use cached PMK instead");
         return ESP_FAIL;
     }
@@ -201,7 +201,6 @@ static int wpa3_parse_sae_confirm(u8 *buf, u32 len)
     g_sae_data.state = SAE_ACCEPTED;
 
     wpa_set_pmk(g_sae_data.pmk, g_sae_data.pmkid, true);
-    memcpy(esp_wifi_sta_get_ap_info_prof_pmk_internal(), g_sae_data.pmk, PMK_LEN);
 
     return ESP_OK;
 }
diff --git a/components/wpa_supplicant/src/esp_supplicant/esp_wpa_main.c b/components/wpa_supplicant/src/esp_supplicant/esp_wpa_main.c
index f3288972a..75012a72e 100644
--- a/components/wpa_supplicant/src/esp_supplicant/esp_wpa_main.c
+++ b/components/wpa_supplicant/src/esp_supplicant/esp_wpa_main.c
@@ -189,6 +189,7 @@ static void wpa_sta_disconnected_cb(uint8_t reason_code)
         case WIFI_REASON_AUTH_FAIL:
         case WIFI_REASON_ASSOC_FAIL:
         case WIFI_REASON_CONNECTION_FAIL:
+        case WIFI_REASON_HANDSHAKE_TIMEOUT:
             esp_wpa3_free_sae_data();
             wpa_sta_clear_curr_pmksa();
             break;
diff --git a/components/wpa_supplicant/src/rsn_supp/wpa.c b/components/wpa_supplicant/src/rsn_supp/wpa.c
index 0402315f0..8b04126de 100644
--- a/components/wpa_supplicant/src/rsn_supp/wpa.c
+++ b/components/wpa_supplicant/src/rsn_supp/wpa.c
@@ -2130,6 +2130,9 @@ int wpa_set_bss(char *macddr, char * bssid, u8 pairwise_cipher, u8 group_cipher,
         esp_wifi_get_config(ESP_IF_WIFI_STA, &wifi_cfg);
         sm->pmf_cfg = wifi_cfg.sta.pmf_cfg;
         sm->mgmt_group_cipher = cipher_type_map_public_to_supp(esp_wifi_sta_get_mgmt_group_cipher());
+    } else {
+        memset(&sm->pmf_cfg, 0, sizeof(sm->pmf_cfg));
+        sm->mgmt_group_cipher = WPA_CIPHER_NONE;
     }
 #endif
  
@@ -2363,6 +2366,15 @@ bool wpa_sta_is_cur_pmksa_set(void) {
     return (pmksa_cache_get_current(sm) != NULL);
 }
 
+bool wpa_sta_cur_pmksa_matches_akm(void) {
+    struct wpa_sm *sm = &gWpaSm;
+    struct rsn_pmksa_cache_entry *pmksa;
+
+    pmksa = pmksa_cache_get_current(sm);
+    return (pmksa != NULL  &&
+            sm->key_mgmt == pmksa->akmp);
+}
+
 void wpa_sta_clear_curr_pmksa(void) {
     struct wpa_sm *sm = &gWpaSm;
 
diff --git a/components/wpa_supplicant/src/rsn_supp/wpa.h b/components/wpa_supplicant/src/rsn_supp/wpa.h
index 03ee4707b..87b13d50e 100644
--- a/components/wpa_supplicant/src/rsn_supp/wpa.h
+++ b/components/wpa_supplicant/src/rsn_supp/wpa.h
@@ -31,6 +31,7 @@ struct wpa_sm;
 int wpa_sm_rx_eapol(u8 *src_addr, u8 *buf, u32 len);
 bool wpa_sta_is_cur_pmksa_set(void);
 bool wpa_sta_in_4way_handshake(void);
+bool wpa_sta_cur_pmksa_matches_akm(void);
 
 #define WPA_ASSERT  assert