From fb53a454b92f4e59ed5fa57bc98258b432dc444f Mon Sep 17 00:00:00 2001
From: "kapil.gupta" <kapil.gupta@espressif.com>
Date: Tue, 19 May 2020 16:20:04 +0530
Subject: [PATCH] ESP-WIFI: Optimize 4way handshake failure time

In case of wrong passpharse, AP will keep on sending 1/4 multiple
times which may take around 10 secs to disconnect and detect
wrong password event.

Add changes to reject EAPOL1 after 3 consecutive reception
---
 components/esp_wifi/lib                                |  2 +-
 .../src/esp_supplicant/esp_wifi_driver.h               |  1 +
 components/wpa_supplicant/src/rsn_supp/wpa.c           | 10 ++++++++++
 components/wpa_supplicant/src/rsn_supp/wpa_i.h         |  1 +
 4 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/components/esp_wifi/lib b/components/esp_wifi/lib
index fc602534d3..6d5c01d883 160000
--- a/components/esp_wifi/lib
+++ b/components/esp_wifi/lib
@@ -1 +1 @@
-Subproject commit fc602534d328ef33085d45ce97c81a8942a06aad
+Subproject commit 6d5c01d883a2081d869d0baa8eedd4ef40228bdc
diff --git a/components/wpa_supplicant/src/esp_supplicant/esp_wifi_driver.h b/components/wpa_supplicant/src/esp_supplicant/esp_wifi_driver.h
index 9999251e2a..2b2a7ca008 100644
--- a/components/wpa_supplicant/src/esp_supplicant/esp_wifi_driver.h
+++ b/components/wpa_supplicant/src/esp_supplicant/esp_wifi_driver.h
@@ -233,5 +233,6 @@ esp_err_t esp_wifi_set_wps_start_flag_internal(bool start);
 uint16_t esp_wifi_sta_pmf_enabled(void);
 wifi_cipher_type_t esp_wifi_sta_get_mgmt_group_cipher(void);
 int esp_wifi_set_igtk_internal(uint8_t if_index, const wifi_wpa_igtk_t *igtk);
+esp_err_t esp_wifi_internal_issue_disconnect(uint8_t reason_code);
 
 #endif /* _ESP_WIFI_DRIVER_H_ */
diff --git a/components/wpa_supplicant/src/rsn_supp/wpa.c b/components/wpa_supplicant/src/rsn_supp/wpa.c
index 7e66dc952f..5a3f438a53 100644
--- a/components/wpa_supplicant/src/rsn_supp/wpa.c
+++ b/components/wpa_supplicant/src/rsn_supp/wpa.c
@@ -48,6 +48,7 @@
 #define WPA_TX_MSG_BUFF_MAXLEN 200
 
 #define ASSOC_IE_LEN 24 + 2 + PMKID_LEN + RSN_SELECTOR_LEN
+#define MAX_EAPOL_RETRIES 3
 u8 assoc_ie_buf[ASSOC_IE_LEN+2]; 
 
 void set_assoc_ie(u8 * assoc_buf);
@@ -1938,6 +1939,14 @@ int   wpa_sm_rx_eapol(u8 *src_addr, u8 *buf, u32 len)
             wpa_supplicant_process_3_of_4(sm, key, ver);
         } else {
             /* 1/4 4-Way Handshake */
+            sm->eapol1_count++;
+            if (sm->eapol1_count > MAX_EAPOL_RETRIES) {
+#ifdef DEBUG_PRINT
+                wpa_printf(MSG_INFO, "EAPOL1 received for %d times, sending deauth", sm->eapol1_count);
+#endif
+                esp_wifi_internal_issue_disconnect(WLAN_REASON_4WAY_HANDSHAKE_TIMEOUT);
+                goto out;
+            }
             wpa_supplicant_process_1_of_4(sm, src_addr, key,
                               ver);
         }
@@ -2123,6 +2132,7 @@ int wpa_set_bss(char *macddr, char * bssid, u8 pairwise_cipher, u8 group_cipher,
         wpa_sm_set_pmk_from_pmksa(sm);
     }
 
+    sm->eapol1_count = 0;
 #ifdef CONFIG_IEEE80211W
     if (esp_wifi_sta_pmf_enabled()) {
         wifi_config_t wifi_cfg;
diff --git a/components/wpa_supplicant/src/rsn_supp/wpa_i.h b/components/wpa_supplicant/src/rsn_supp/wpa_i.h
index 09b0072dff..831810a5d5 100644
--- a/components/wpa_supplicant/src/rsn_supp/wpa_i.h
+++ b/components/wpa_supplicant/src/rsn_supp/wpa_i.h
@@ -90,6 +90,7 @@ struct wpa_sm {
     u16 txcb_flags;
     bool   ap_notify_completed_rsne;
     wifi_pmf_config_t pmf_cfg;
+    u8 eapol1_count;
 };
 
 /**