From 19e840aa53507ead0e295036dba40dbdc77a5453 Mon Sep 17 00:00:00 2001
From: Hrudaynath Dhabe <hrudaynath.dhabe@espressif.com>
Date: Mon, 11 Nov 2019 23:22:14 +0800
Subject: [PATCH] wpa_supplicant: Set assoc_ie_len based on generated RSN/WPA
 IE

---
 .../src/esp_supplicant/esp_wpa_main.c         |  9 +++++---
 components/wpa_supplicant/src/rsn_supp/wpa.c  | 21 ++++++++++++-------
 .../wpa_supplicant/src/rsn_supp/wpa_i.h       |  2 +-
 3 files changed, 20 insertions(+), 12 deletions(-)

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 c576982213..af6ef8ae3b 100644
--- a/components/wpa_supplicant/src/esp_supplicant/esp_wpa_main.c
+++ b/components/wpa_supplicant/src/esp_supplicant/esp_wpa_main.c
@@ -83,13 +83,14 @@ void  wpa_config_profile()
 
 int wpa_config_bss(uint8_t *bssid)
 {
+    int ret = 0;
     struct wifi_ssid *ssid = esp_wifi_sta_get_prof_ssid_internal();
     u8 mac[6];
 
     esp_wifi_get_macaddr_internal(0, mac);
-    wpa_set_bss((char *)mac, (char *)bssid, esp_wifi_sta_get_pairwise_cipher_internal(), esp_wifi_sta_get_group_cipher_internal(),
+    ret = wpa_set_bss((char *)mac, (char *)bssid, esp_wifi_sta_get_pairwise_cipher_internal(), esp_wifi_sta_get_group_cipher_internal(),
                 (char *)esp_wifi_sta_get_prof_password_internal(), ssid->ssid, ssid->len);
-    return ESP_OK;
+    return ret;
 }
 
 void  wpa_config_assoc_ie(u8 proto, u8 *assoc_buf, u32 assoc_wpa_ie_len)
@@ -151,8 +152,10 @@ bool  wpa_deattach(void)
 
 void  wpa_sta_connect(uint8_t *bssid)
 {
+    int ret = 0;
     wpa_config_profile();
-    wpa_config_bss(bssid);
+    ret = wpa_config_bss(bssid);
+    WPA_ASSERT(ret == 0);
 }
 
 int cipher_type_map(int wpa_cipher)
diff --git a/components/wpa_supplicant/src/rsn_supp/wpa.c b/components/wpa_supplicant/src/rsn_supp/wpa.c
index bc68cf2b2a..cb9220fb74 100644
--- a/components/wpa_supplicant/src/rsn_supp/wpa.c
+++ b/components/wpa_supplicant/src/rsn_supp/wpa.c
@@ -528,7 +528,10 @@ void   wpa_supplicant_process_1_of_4(struct wpa_sm *sm,
     if (res)
         goto failed;
 
-    pmksa_cache_set_current(sm, NULL, sm->bssid, 0, 0);
+    if (esp_wifi_sta_prof_is_wpa2_internal() 
+            && esp_wifi_sta_get_prof_authmode_internal() == WPA2_AUTH_ENT) {
+        pmksa_cache_set_current(sm, NULL, sm->bssid, 0, 0);
+    }
 
     if (sm->renew_snonce) {
         if (os_get_random(sm->snonce, WPA_NONCE_LEN)) {
@@ -1988,9 +1991,9 @@ void wpa_set_pmk(uint8_t *pmk)
     sm->pmk_len = PMK_LEN;
 }
 
-void  
-wpa_set_bss(char *macddr, char * bssid, u8 pairwise_cipher, u8 group_cipher, char *passphrase, u8 *ssid, size_t ssid_len)
+int wpa_set_bss(char *macddr, char * bssid, u8 pairwise_cipher, u8 group_cipher, char *passphrase, u8 *ssid, size_t ssid_len)
 {
+    int res = 0;
     struct wpa_sm *sm = &gWpaSm;
     
     sm->pairwise_cipher = BIT(pairwise_cipher);
@@ -2008,11 +2011,13 @@ wpa_set_bss(char *macddr, char * bssid, u8 pairwise_cipher, u8 group_cipher, cha
         pmksa_cache_set_current(sm, NULL, (const u8*) bssid, 0, 0);
         wpa_sm_set_pmk_from_pmksa(sm);
     }
-
     set_assoc_ie(assoc_ie_buf); /* use static buffer */
-
-    wpa_gen_wpa_ie(sm, sm->assoc_wpa_ie, sm->assoc_wpa_ie_len); //TODO: NEED TO DEBUG!!
+    res = wpa_gen_wpa_ie(sm, sm->assoc_wpa_ie, sm->assoc_wpa_ie_len);
+    if (res < 0) 
+        return -1;
+    sm->assoc_wpa_ie_len = res;
     wpa_set_passphrase(passphrase, ssid, ssid_len);
+    return 0;
 }
 
 /*
@@ -2063,8 +2068,8 @@ set_assoc_ie(u8 * assoc_buf)
     if ( sm->proto == WPA_PROTO_WPA) 
          sm->assoc_wpa_ie_len = ASSOC_IE_LEN;
     else 
-         sm->assoc_wpa_ie_len = ASSOC_IE_LEN - 2;    
-    
+         sm->assoc_wpa_ie_len = ASSOC_IE_LEN - 2;
+
     sm->config_assoc_ie(sm->proto, assoc_buf, sm->assoc_wpa_ie_len);    
 }
 
diff --git a/components/wpa_supplicant/src/rsn_supp/wpa_i.h b/components/wpa_supplicant/src/rsn_supp/wpa_i.h
index ab99145af3..51ce03aa2f 100644
--- a/components/wpa_supplicant/src/rsn_supp/wpa_i.h
+++ b/components/wpa_supplicant/src/rsn_supp/wpa_i.h
@@ -163,7 +163,7 @@ void eapol_txcb(void *eb);
 
 void wpa_set_profile(u32 wpa_proto, u8 auth_mode);
 
-void wpa_set_bss(char *macddr, char * bssid, u8 pairwise_cipher, u8 group_cipher, char *passphrase, u8 *ssid, size_t ssid_len);
+int wpa_set_bss(char *macddr, char * bssid, u8 pairwise_cipher, u8 group_cipher, char *passphrase, u8 *ssid, size_t ssid_len);
 
 int wpa_sm_rx_eapol(u8 *src_addr, u8 *buf, u32 len);
 #endif /* WPA_I_H */