From e746fc0debe0567bd148aa5d4b4abfa1fca29110 Mon Sep 17 00:00:00 2001
From: Shreyas Sheth <shreyas.sheth@espressif.com>
Date: Mon, 9 Oct 2023 09:26:32 +0530
Subject: [PATCH 1/8] feat(wifi): Refactor and update wpa_supplicant with
 upstream

---
 .../esp_supplicant/src/esp_eap_client.c       |  2 +-
 .../esp_supplicant/src/esp_wpas_glue.c        |  2 +-
 .../esp_supplicant/src/esp_wpas_glue.h        |  2 +-
 components/wpa_supplicant/src/common/sae.c    | 60 +++++++------
 components/wpa_supplicant/src/common/sae.h    |  8 +-
 .../wpa_supplicant/src/common/wpa_common.c    | 89 ++++++++++++++-----
 .../wpa_supplicant/src/common/wpa_common.h    |  4 +-
 components/wpa_supplicant/src/rsn_supp/wpa.c  | 58 ++++++------
 .../wpa_supplicant/src/rsn_supp/wpa_ie.c      | 29 +++---
 9 files changed, 148 insertions(+), 106 deletions(-)

diff --git a/components/wpa_supplicant/esp_supplicant/src/esp_eap_client.c b/components/wpa_supplicant/esp_supplicant/src/esp_eap_client.c
index c2255e1119..3d46f7fe12 100644
--- a/components/wpa_supplicant/esp_supplicant/src/esp_eap_client.c
+++ b/components/wpa_supplicant/esp_supplicant/src/esp_eap_client.c
@@ -590,7 +590,7 @@ static int eap_sm_rx_eapol_internal(u8 *src_addr, u8 *buf, u32 len, uint8_t *bss
         break;
     case EAP_CODE_SUCCESS:
         if (sm->eapKeyData) {
-            wpa_set_pmk(sm->eapKeyData, NULL, false);
+            wpa_set_pmk(sm->eapKeyData, 0, NULL, false);
             os_free(sm->eapKeyData);
             sm->eapKeyData = NULL;
             wpa_printf(MSG_INFO, ">>>>>EAP FINISH");
diff --git a/components/wpa_supplicant/esp_supplicant/src/esp_wpas_glue.c b/components/wpa_supplicant/esp_supplicant/src/esp_wpas_glue.c
index 4f1da5d792..9a74c4e6b7 100644
--- a/components/wpa_supplicant/esp_supplicant/src/esp_wpas_glue.c
+++ b/components/wpa_supplicant/esp_supplicant/src/esp_wpas_glue.c
@@ -93,7 +93,7 @@ int hostapd_send_eapol(const u8 *source, const u8 *sta_addr,
 
 }
 
-void wpa_supplicant_transition_disable(void *sm, u8 bitmap)
+void wpa_supplicant_transition_disable(struct wpa_sm *sm, u8 bitmap)
 {
     wpa_printf(MSG_DEBUG, "TRANSITION_DISABLE %02x", bitmap);
 
diff --git a/components/wpa_supplicant/esp_supplicant/src/esp_wpas_glue.h b/components/wpa_supplicant/esp_supplicant/src/esp_wpas_glue.h
index 11f5b3ddff..e18ad28cdd 100644
--- a/components/wpa_supplicant/esp_supplicant/src/esp_wpas_glue.h
+++ b/components/wpa_supplicant/esp_supplicant/src/esp_wpas_glue.h
@@ -31,7 +31,7 @@ void wpa_free_eapol(u8 *buffer);
 int wpa_ether_send(void *ctx, const u8 *dest, u16 proto,
                    const u8 *data, size_t data_len);
 
-void wpa_supplicant_transition_disable(void *sm, u8 bitmap);
+void wpa_supplicant_transition_disable(struct wpa_sm *sm, u8 bitmap);
 
 int hostapd_send_eapol(const u8 *source, const u8 *sta_addr,
 		       const u8 *data, size_t data_len);
diff --git a/components/wpa_supplicant/src/common/sae.c b/components/wpa_supplicant/src/common/sae.c
index e97f99d449..bc394c80db 100644
--- a/components/wpa_supplicant/src/common/sae.c
+++ b/components/wpa_supplicant/src/common/sae.c
@@ -170,7 +170,7 @@ static int sae_test_pwd_seed_ecc(struct sae_data *sae, const u8 *pwd_seed,
 		return ESP_FAIL;
 
 	res = dragonfly_is_quadratic_residue_blind(sae->tmp->ec, qr, qnr,
-                           y_sqr);
+						   y_sqr);
 	crypto_bignum_deinit(y_sqr, 1);
 
 	if (res < 0) {
@@ -219,7 +219,6 @@ static int sae_test_pwd_seed_ffc(struct sae_data *sae, const u8 *pwd_seed,
 
 	res = -1;
 	a = crypto_bignum_init_set(pwd_value, sae->tmp->prime_len);
-
 	if (!a)
 		goto fail;
 
@@ -240,10 +239,9 @@ static int sae_test_pwd_seed_ffc(struct sae_data *sae, const u8 *pwd_seed,
 		if (b == NULL ||
 		    crypto_bignum_sub(sae->tmp->prime, b, b) < 0 ||
 		    crypto_bignum_div(b, sae->tmp->order, b) < 0)
-		    goto fail;
+			goto fail;
 	}
 
-
 	if (!b)
 		goto fail;
 
@@ -314,7 +312,7 @@ static int sae_derive_pwe_ecc(struct sae_data *sae, const u8 *addr1,
 	if (dragonfly_get_random_qr_qnr(sae->tmp->prime, &qr, &qnr) < 0 ||
 	    crypto_bignum_to_bin(qr, qr_bin, sizeof(qr_bin), prime_len) < 0 ||
 	    crypto_bignum_to_bin(qnr, qnr_bin, sizeof(qnr_bin), prime_len) < 0)
-		    goto fail;
+		goto fail;
 
 	wpa_hexdump_ascii_key(MSG_DEBUG, "SAE: password",
 			      password, password_len);
@@ -432,6 +430,7 @@ fail:
 	crypto_bignum_deinit(x, 1);
 	os_memset(x_bin, 0, sizeof(x_bin));
 	os_memset(x_cand_bin, 0, sizeof(x_cand_bin));
+
 	return res;
 }
 
@@ -450,6 +449,7 @@ static int sae_derive_pwe_ffc(struct sae_data *sae, const u8 *addr1,
 	struct crypto_bignum *pwe;
 	size_t prime_len = sae->tmp->prime_len * 8;
 	u8 *pwe_buf;
+
 	crypto_bignum_deinit(sae->tmp->pwe_ffc, 1);
 	sae->tmp->pwe_ffc = NULL;
 
@@ -478,7 +478,7 @@ static int sae_derive_pwe_ffc(struct sae_data *sae, const u8 *addr1,
 	k = dragonfly_min_pwe_loop_iter(sae->group);
 
 	for (counter = 1; counter <= k || !found; counter++) {
-        u8 pwd_seed[SHA256_MAC_LEN];
+		u8 pwd_seed[SHA256_MAC_LEN];
 		int res;
 
 		if (counter > 200) {
@@ -570,6 +570,7 @@ static void debug_print_bignum(const char *title, const struct crypto_bignum *a,
 	bin_clear_free(bin, prime_len);
 }
 
+
 static struct crypto_ec_point * sswu(struct crypto_ec *ec, int group,
 				     const struct crypto_bignum *u)
 {
@@ -1480,9 +1481,9 @@ static int sae_derive_keys(struct sae_data *sae, const u8 *k)
 	 * KCK || PMK = KDF-Hash-Length(keyseed, "SAE KCK and PMK",
 	 *                      (commit-scalar + peer-commit-scalar) modulo r)
 	 * PMKID = L((commit-scalar + peer-commit-scalar) modulo r, 0, 128)
-	*
-	* When SAE_PK is used,
-	* KCK || PMK || KEK = KDF-Hash-Length(keyseed, "SAE-PK keys", context)
+	 *
+	 * When SAE-PK is used,
+	 * KCK || PMK || KEK = KDF-Hash-Length(keyseed, "SAE-PK keys", context)
 	 */
 	if (!sae->h2e)
 		hash_len = SHA256_MAC_LEN;
@@ -1542,7 +1543,7 @@ static int sae_derive_keys(struct sae_data *sae, const u8 *k)
 	 * zero padding it from left to the length of the order (in full
 	 * octets). */
 	if (crypto_bignum_to_bin(tmp, val, sizeof(val),
-		sae->tmp->order_len) < 0) {
+				 sae->tmp->order_len) < 0) {
 		goto fail;
 	}
 	wpa_hexdump(MSG_DEBUG, "SAE: PMKID", val, SAE_PMKID_LEN);
@@ -1550,13 +1551,13 @@ static int sae_derive_keys(struct sae_data *sae, const u8 *k)
 #ifdef CONFIG_SAE_PK
 	if (sae->pk) {
 		if (sae_kdf_hash(hash_len, keyseed, "SAE-PK keys",
-			   val, sae->tmp->order_len,
-			   keys, 2 * hash_len + SAE_PMK_LEN) < 0)
+				 val, sae->tmp->order_len,
+				 keys, 2 * hash_len + SAE_PMK_LEN) < 0)
 			goto fail;
 	} else {
 		if (sae_kdf_hash(hash_len, keyseed, "SAE KCK and PMK",
-			   val, sae->tmp->order_len,
-			   keys, hash_len + SAE_PMK_LEN) < 0)
+				 val, sae->tmp->order_len,
+				 keys, hash_len + SAE_PMK_LEN) < 0)
 			goto fail;
 	}
 #else /* CONFIG_SAE_PK */
@@ -1568,20 +1569,22 @@ static int sae_derive_keys(struct sae_data *sae, const u8 *k)
 
 	forced_memzero(keyseed, sizeof(keyseed));
 	os_memcpy(sae->tmp->kck, keys, hash_len);
+	sae->tmp->kck_len = hash_len;
 	os_memcpy(sae->pmk, keys + hash_len, SAE_PMK_LEN);
 	os_memcpy(sae->pmkid, val, SAE_PMKID_LEN);
 
 #ifdef CONFIG_SAE_PK
 	if (sae->pk) {
-		os_memcpy(sae->tmp->kek, keys + hash_len + SAE_PMK_LEN, hash_len);
+		os_memcpy(sae->tmp->kek, keys + hash_len + SAE_PMK_LEN,
+			  hash_len);
 		sae->tmp->kek_len = hash_len;
 		wpa_hexdump_key(MSG_DEBUG, "SAE: KEK for SAE-PK",
-				   sae->tmp->kek, sae->tmp->kek_len);
+				sae->tmp->kek, sae->tmp->kek_len);
 	}
 #endif /* CONFIG_SAE_PK */
 	forced_memzero(keys, sizeof(keys));
 	wpa_hexdump_key(MSG_DEBUG, "SAE: KCK",
-			sae->tmp->kck, SAE_KCK_LEN);
+			sae->tmp->kck, sae->tmp->kck_len);
 	wpa_hexdump_key(MSG_DEBUG, "SAE: PMK", sae->pmk, SAE_PMK_LEN);
 
 	ret = 0;
@@ -1658,6 +1661,7 @@ int sae_write_commit(struct sae_data *sae, struct wpabuf *buf,
 		wpa_printf(MSG_DEBUG, "SAE: own Password Identifier: %s",
 			   identifier);
 	}
+
 	if (sae->h2e && sae->tmp->own_rejected_groups) {
 		wpa_hexdump_buf(MSG_DEBUG, "SAE: own Rejected Groups",
 				sae->tmp->own_rejected_groups);
@@ -1725,7 +1729,7 @@ u16 sae_group_allowed(struct sae_data *sae, int *allowed_groups, u16 group)
 
 static int sae_is_password_id_elem(const u8 *pos, const u8 *end)
 {
-	return  end - pos >= 3 &&
+	return end - pos >= 3 &&
 		pos[0] == WLAN_EID_EXTENSION &&
 		pos[1] >= 1 &&
 		end - pos - 2 >= pos[1] &&
@@ -1842,6 +1846,7 @@ static u16 sae_parse_commit_scalar(struct sae_data *sae, const u8 **pos,
 		return WLAN_STATUS_UNSPECIFIED_FAILURE;
 	}
 
+
 	crypto_bignum_deinit(sae->peer_commit_scalar, 0);
 	sae->peer_commit_scalar = peer_scalar;
 	wpa_hexdump(MSG_DEBUG, "SAE: Peer commit-scalar",
@@ -2144,7 +2149,7 @@ static int sae_cn_confirm(struct sae_data *sae, const u8 *sc,
 	len[3] = sae->tmp->prime_len;
 	addr[4] = element2;
 	len[4] = element2_len;
-	return hkdf_extract(SAE_KCK_LEN, sae->tmp->kck, SAE_KCK_LEN,
+	return hkdf_extract(sae->tmp->kck_len, sae->tmp->kck, sae->tmp->kck_len,
 			    5, addr, len, confirm);
 }
 
@@ -2171,7 +2176,7 @@ static int sae_cn_confirm_ecc(struct sae_data *sae, const u8 *sc,
 	}
 
 	sae_cn_confirm(sae, sc, scalar1, element_b1, 2 * sae->tmp->prime_len,
-		       scalar2, element_b2, 2 * sae->tmp->prime_len, confirm);
+		      scalar2, element_b2, 2 * sae->tmp->prime_len, confirm);
 	return ESP_OK;
 }
 
@@ -2187,7 +2192,7 @@ static int sae_cn_confirm_ffc(struct sae_data *sae, const u8 *sc,
 	u8 element_b2[SAE_MAX_PRIME_LEN];
 
 	if (crypto_bignum_to_bin(element1, element_b1, sizeof(element_b1),
-			         sae->tmp->prime_len) < 0) {
+				 sae->tmp->prime_len) < 0) {
 		wpa_printf(MSG_ERROR, "SAE: failed bignum op while generating SAE confirm - e1");
 		return ESP_FAIL;
 	}
@@ -2207,10 +2212,13 @@ static int sae_cn_confirm_ffc(struct sae_data *sae, const u8 *sc,
 int sae_write_confirm(struct sae_data *sae, struct wpabuf *buf)
 {
 	const u8 *sc;
+	size_t hash_len;
 
 	if (sae->tmp == NULL)
 		return ESP_FAIL;
 
+	hash_len = sae->tmp->kck_len;
+
 	/* Send-Confirm */
 	sc = wpabuf_put(buf, 0);
 	wpabuf_put_le16(buf, sae->send_confirm);
@@ -2222,7 +2230,7 @@ int sae_write_confirm(struct sae_data *sae, struct wpabuf *buf)
 					 sae->tmp->own_commit_element_ecc,
 					 sae->peer_commit_scalar,
 					 sae->tmp->peer_commit_element_ecc,
-					 wpabuf_put(buf, SHA256_MAC_LEN))) {
+					 wpabuf_put(buf, hash_len))) {
 			wpa_printf(MSG_ERROR, "SAE: failed generate SAE confirm (ecc)");
 			return ESP_FAIL;
 		}
@@ -2231,7 +2239,7 @@ int sae_write_confirm(struct sae_data *sae, struct wpabuf *buf)
 					 sae->tmp->own_commit_element_ffc,
 					 sae->peer_commit_scalar,
 					 sae->tmp->peer_commit_element_ffc,
-					 wpabuf_put(buf, SHA256_MAC_LEN))) {
+					 wpabuf_put(buf, hash_len))) {
 			wpa_printf(MSG_ERROR, "SAE: failed generate SAE confirm (ffc)");
 			return ESP_FAIL;
 		}
@@ -2243,11 +2251,12 @@ int sae_write_confirm(struct sae_data *sae, struct wpabuf *buf)
 int sae_check_confirm(struct sae_data *sae, const u8 *data, size_t len)
 {
 	u8 verifier[SAE_MAX_HASH_LEN];
-	size_t hash_len= SAE_KCK_LEN;
+	size_t hash_len;
 
 	if (!sae->tmp)
 		return ESP_FAIL;
 
+	hash_len = sae->tmp->kck_len;
 	if (len < 2 + hash_len) {
 		wpa_printf(MSG_DEBUG, "SAE: Too short confirm message");
 		return ESP_FAIL;
@@ -2298,13 +2307,14 @@ int sae_check_confirm(struct sae_data *sae, const u8 *data, size_t len)
 
 #ifdef CONFIG_SAE_PK
 	if (sae_check_confirm_pk(sae, data + 2 + hash_len,
-		   len - 2 - hash_len) != ESP_OK)
+				 len - 2 - hash_len) != ESP_OK)
 		return ESP_FAIL;
 #endif /* CONFIG_SAE_PK */
 
 	return ESP_OK;
 }
 
+
 const char * sae_state_txt(enum sae_state state)
 {
 	switch (state) {
diff --git a/components/wpa_supplicant/src/common/sae.h b/components/wpa_supplicant/src/common/sae.h
index 77a834cdcd..0f266bc1c2 100644
--- a/components/wpa_supplicant/src/common/sae.h
+++ b/components/wpa_supplicant/src/common/sae.h
@@ -42,7 +42,8 @@ struct sae_pk {
 };
 
 struct sae_temporary_data {
-	u8 kck[SAE_KCK_LEN];
+	u8 kck[SAE_MAX_HASH_LEN];
+	size_t kck_len;
 	struct crypto_bignum *own_commit_scalar;
 	struct crypto_bignum *own_commit_element_ffc;
 	struct crypto_ec_point *own_commit_element_ecc;
@@ -109,15 +110,14 @@ struct sae_data {
 	u16 send_confirm;
 	u8 pmk[SAE_PMK_LEN];
 	u8 pmkid[SAE_PMKID_LEN];
-
 	struct crypto_bignum *peer_commit_scalar;
+	struct crypto_bignum *peer_commit_scalar_accepted;
 	int group;
 	unsigned int sync; /* protocol instance variable: Sync */
 	u16 rc; /* protocol instance variable: Rc (received send-confirm) */
-	struct sae_temporary_data *tmp;
-	struct crypto_bignum *peer_commit_scalar_accepted;
 	unsigned int h2e:1;
 	unsigned int pk:1;
+	struct sae_temporary_data *tmp;
 };
 
 int sae_set_group(struct sae_data *sae, int group);
diff --git a/components/wpa_supplicant/src/common/wpa_common.c b/components/wpa_supplicant/src/common/wpa_common.c
index 05e00463aa..cbf1fa630f 100644
--- a/components/wpa_supplicant/src/common/wpa_common.c
+++ b/components/wpa_supplicant/src/common/wpa_common.c
@@ -235,27 +235,39 @@ int wpa_ft_parse_ies(const u8 *ies, size_t ies_len,
 #endif /* CONFIG_IEEE80211R */
 
 
-static unsigned int wpa_kck_len(int akmp)
+static unsigned int wpa_kck_len(int akmp, size_t pmk_len)
 {
-	if (akmp == WPA_KEY_MGMT_IEEE8021X_SUITE_B_192)
+	switch (akmp) {
+	case WPA_KEY_MGMT_IEEE8021X_SUITE_B_192:
 		return 24;
-	return 16;
+	case WPA_KEY_MGMT_OWE:
+		return pmk_len / 2;
+	default:
+		return 16;
+	}
 }
 
-static unsigned int wpa_kek_len(int akmp)
+static unsigned int wpa_kek_len(int akmp, size_t pmk_len)
 {
-	if (akmp == WPA_KEY_MGMT_IEEE8021X_SUITE_B_192)
+	switch (akmp) {
+	case WPA_KEY_MGMT_IEEE8021X_SUITE_B_192:
 		return 32;
-	return 16;
+	case WPA_KEY_MGMT_OWE:
+		return pmk_len <= 32 ? 16 : 32;
+	default:
+		return 16;
+	}
 }
 
 
-unsigned int wpa_mic_len(int akmp)
+unsigned int wpa_mic_len(int akmp, size_t pmk_len)
 {
-	if (akmp == WPA_KEY_MGMT_IEEE8021X_SUITE_B_192)
+	switch (akmp) {
+	case WPA_KEY_MGMT_IEEE8021X_SUITE_B_192:
 		return 24;
-
-	return 16;
+	default:
+		return 16;
+	}
 }
 
 static int rsn_selector_to_bitfield(const u8 *s)
@@ -783,8 +795,8 @@ int wpa_pmk_r1_to_ptk(const u8 *pmk_r1, const u8 *snonce, const u8 *anonce,
 	os_memcpy(pos, sta_addr, ETH_ALEN);
 	pos += ETH_ALEN;
 
-	ptk->kck_len = wpa_kck_len(akmp);
-	ptk->kek_len = wpa_kek_len(akmp);
+	ptk->kck_len = wpa_kck_len(akmp, PMK_LEN);
+	ptk->kek_len = wpa_kek_len(akmp, PMK_LEN);
 	ptk->tk_len = wpa_cipher_key_len(cipher);
 	ptk_len = ptk->kck_len + ptk->kek_len + ptk->tk_len;
 
@@ -835,10 +847,9 @@ int wpa_pmk_r1_to_ptk(const u8 *pmk_r1, const u8 *snonce, const u8 *anonce,
 */
 
 int wpa_use_akm_defined(int akmp){
-	int ret = 0;
-	if(wpa_key_mgmt_sae(akmp))
-		ret = 1;
-	return ret;
+	return akmp == WPA_KEY_MGMT_OWE ||
+		wpa_key_mgmt_sae(akmp) ||
+		wpa_key_mgmt_suite_b(akmp);
 }
 
 int wpa_use_aes_key_wrap(int akmp)
@@ -936,6 +947,35 @@ int wpa_eapol_key_mic(const u8 *key, size_t key_len, int akmp, int ver,
 	return 0;
 }
 
+u32 wpa_akm_to_suite(int akm)
+{
+#ifdef CONFIG_IEEE80211R
+	if (akm & WPA_KEY_MGMT_FT_IEEE8021X)
+		return RSN_AUTH_KEY_MGMT_FT_802_1X;
+	if (akm & WPA_KEY_MGMT_FT_PSK)
+		return RSN_AUTH_KEY_MGMT_FT_PSK;
+#endif /* CONFIG_IEEE80211R */
+	if (akm & WPA_KEY_MGMT_IEEE8021X_SHA256)
+		return RSN_AUTH_KEY_MGMT_802_1X_SHA256;
+	if (akm & WPA_KEY_MGMT_IEEE8021X)
+		return RSN_AUTH_KEY_MGMT_UNSPEC_802_1X;
+	if (akm & WPA_KEY_MGMT_PSK_SHA256)
+		return RSN_AUTH_KEY_MGMT_PSK_SHA256;
+	if (akm & WPA_KEY_MGMT_PSK)
+		return RSN_AUTH_KEY_MGMT_PSK_OVER_802_1X;
+	if (akm & WPA_KEY_MGMT_IEEE8021X_SUITE_B)
+		return RSN_AUTH_KEY_MGMT_802_1X_SUITE_B;
+	if (akm & WPA_KEY_MGMT_IEEE8021X_SUITE_B_192)
+		return RSN_AUTH_KEY_MGMT_802_1X_SUITE_B_192;
+	if (akm & WPA_KEY_MGMT_SAE)
+		return RSN_AUTH_KEY_MGMT_SAE;
+	if (akm & WPA_KEY_MGMT_FT_SAE)
+		return RSN_AUTH_KEY_MGMT_FT_SAE;
+	if (akm & WPA_KEY_MGMT_OWE)
+		return RSN_AUTH_KEY_MGMT_OWE;
+	return 0;
+}
+
 int wpa_compare_rsn_ie(int ft_initial_assoc,
 		       const u8 *ie1, size_t ie1len,
 		       const u8 *ie2, size_t ie2len)
@@ -1099,6 +1139,7 @@ int wpa_pmk_to_ptk(const u8 *pmk, size_t pmk_len, const char *label,
 		return -1;
 	}
 	u8 data[2 * ETH_ALEN + 2 * WPA_NONCE_LEN];
+	size_t data_len = 2 * ETH_ALEN + 2 * WPA_NONCE_LEN;
 	u8 tmp[WPA_KCK_MAX_LEN + WPA_KEK_MAX_LEN + WPA_TK_MAX_LEN];
 	size_t ptk_len;
 
@@ -1120,25 +1161,25 @@ int wpa_pmk_to_ptk(const u8 *pmk, size_t pmk_len, const char *label,
 			  WPA_NONCE_LEN);
 	}
 
-	ptk->kck_len = wpa_kck_len(akmp);
-	ptk->kek_len = wpa_kek_len(akmp);
+	ptk->kck_len = wpa_kck_len(akmp, pmk_len);
+	ptk->kek_len = wpa_kek_len(akmp, pmk_len);
 	ptk->tk_len = wpa_cipher_key_len(cipher);
 	ptk_len = ptk->kck_len + ptk->kek_len + ptk->tk_len;
 
 #if defined(CONFIG_SUITEB192)
 	if (wpa_key_mgmt_sha384(akmp)) {
 		wpa_printf(MSG_DEBUG, "WPA: PTK derivation using PRF(SHA384)");
-		if (sha384_prf(pmk, pmk_len, label, data, sizeof(data),
+		if (sha384_prf(pmk, pmk_len, label, data, data_len,
 					tmp, ptk_len) < 0)
 			return -1;
 	} else
 #endif
-	if (wpa_key_mgmt_sha256(akmp))
-		sha256_prf(pmk, pmk_len, label, data, sizeof(data),
+	if (wpa_key_mgmt_sha256(akmp)) {
+		sha256_prf(pmk, pmk_len, label, data, data_len,
 			   tmp, ptk_len);
-	else
-		sha1_prf(pmk, pmk_len, label, data, sizeof(data), tmp, ptk_len);
-
+	} else {
+		sha1_prf(pmk, pmk_len, label, data, data_len, tmp, ptk_len);
+	}
 	wpa_printf(MSG_DEBUG, "WPA: PTK derivation - A1=" MACSTR " A2=" MACSTR"\n",
 		   MAC2STR(addr1), MAC2STR(addr2));
 
diff --git a/components/wpa_supplicant/src/common/wpa_common.h b/components/wpa_supplicant/src/common/wpa_common.h
index 4462c81f00..1f9adb76a2 100644
--- a/components/wpa_supplicant/src/common/wpa_common.h
+++ b/components/wpa_supplicant/src/common/wpa_common.h
@@ -395,7 +395,7 @@ int wpa_parse_wpa_ie(const u8 *wpa_ie, size_t wpa_ie_len,
 		     struct wpa_ie_data *data);
 int wpa_parse_wpa_ie_rsnxe(const u8 *rsnxe_ie, size_t rsnxe_ie_len,
 			    struct wpa_ie_data *data);
-
+u32 wpa_akm_to_suite(int akm);
 int wpa_compare_rsn_ie(int ft_initial_assoc,
 		       const u8 *ie1, size_t ie1len,
 		       const u8 *ie2, size_t ie2len);
@@ -447,7 +447,7 @@ int wpa_parse_wpa_ie_wpa(const u8 *wpa_ie, size_t wpa_ie_len,
            struct wpa_ie_data *data);
 
 int rsn_cipher_put_suites(u8 *pos, int ciphers);
-unsigned int wpa_mic_len(int akmp);
+unsigned int wpa_mic_len(int akmp, size_t pmk_len);
 int wpa_use_akm_defined(int akmp);
 int wpa_use_aes_key_wrap(int akmp);
 
diff --git a/components/wpa_supplicant/src/rsn_supp/wpa.c b/components/wpa_supplicant/src/rsn_supp/wpa.c
index e581538990..d23da66d29 100644
--- a/components/wpa_supplicant/src/rsn_supp/wpa.c
+++ b/components/wpa_supplicant/src/rsn_supp/wpa.c
@@ -260,7 +260,7 @@ void wpa_eapol_key_send(struct wpa_sm *sm, const u8 *kck, size_t kck_len,
         goto out;
     }
     wpa_hexdump_key(MSG_DEBUG, "WPA: KCK", kck, kck_len);
-    wpa_hexdump(MSG_DEBUG, "WPA: Derived Key MIC", key_mic, wpa_mic_len(sm->key_mgmt));
+    wpa_hexdump(MSG_DEBUG, "WPA: Derived Key MIC", key_mic, wpa_mic_len(sm->key_mgmt, sm->pmk_len));
     wpa_hexdump(MSG_MSGDUMP, "WPA: TX EAPOL-Key", msg, msg_len);
     wpa_sm_ether_send(sm, dest, proto, msg, msg_len);
 out:
@@ -302,7 +302,7 @@ static void wpa_sm_key_request(struct wpa_sm *sm, int error, int pairwise)
         return;
     }
 
-    mic_len = wpa_mic_len(sm->key_mgmt);
+    mic_len = wpa_mic_len(sm->key_mgmt, sm->pmk_len);
     hdrlen = mic_len == 24 ? sizeof(*reply192) : sizeof(*reply);
     rbuf = wpa_sm_alloc_eapol(sm, IEEE802_1X_TYPE_EAPOL_KEY, NULL,
                   hdrlen, &rlen, (void *) &reply);
@@ -424,11 +424,23 @@ static int wpa_supplicant_get_pmk(struct wpa_sm *sm,
         //eapol_sm_notify_cached(sm->eapol);
 #ifdef CONFIG_IEEE80211R
         sm->xxkey_len = 0;
+#ifdef CONFIG_WPA3_SAE
+        if ((sm->key_mgmt == WPA_KEY_MGMT_FT_SAE) &&
+             sm->pmk_len == PMK_LEN) {
+            /* Need to allow FT key derivation to proceed with
+             * PMK from SAE being used as the XXKey in cases where
+             * the PMKID in msg 1/4 matches the PMKSA entry that was
+             * just added based on SAE authentication for the
+             * initial mobility domain association. */
+            os_memcpy(sm->xxkey, sm->pmk, sm->pmk_len);
+            sm->xxkey_len = sm->pmk_len;
+        }
+#endif /* CONFIG_WPA3_SAE */
 #endif /* CONFIG_IEEE80211R */
     } else if (wpa_key_mgmt_wpa_ieee8021x(sm->key_mgmt)) {
         int res = 0, pmk_len;
         /* For ESP_SUPPLICANT this is already set using wpa_set_pmk*/
-        //res = eapol_sm_get_key(sm->eapol, sm->pmk, PMK_LEN);
+        //res = eapol_sm_get_key(sm->eapol, 0, sm->pmk, PMK_LEN);
         if (wpa_key_mgmt_sha384(sm->key_mgmt))
             pmk_len = PMK_LEN_SUITE_B_192;
         else
@@ -586,7 +598,7 @@ int   wpa_supplicant_send_2_of_4(struct wpa_sm *sm, const unsigned char *dst,
 #endif /* CONFIG_IEEE80211R */
     wpa_hexdump(MSG_MSGDUMP, "WPA: WPA IE for msg 2/4\n", wpa_ie, wpa_ie_len);
 
-    mic_len = wpa_mic_len(sm->key_mgmt);
+    mic_len = wpa_mic_len(sm->key_mgmt, sm->pmk_len);
     hdrlen = mic_len == 24 ? sizeof(*reply192) : sizeof(*reply);
     rbuf = wpa_sm_alloc_eapol(sm, IEEE802_1X_TYPE_EAPOL_KEY,
                   NULL, hdrlen + wpa_ie_len,
@@ -1205,7 +1217,7 @@ static int wpa_supplicant_send_4_of_4(struct wpa_sm *sm, const unsigned char *ds
     struct wpa_eapol_key_192 *reply192;
     u8 *rbuf, *key_mic;
 
-    mic_len = wpa_mic_len(sm->key_mgmt);
+    mic_len = wpa_mic_len(sm->key_mgmt, sm->pmk_len);
     hdrlen = mic_len == 24 ? sizeof(*reply192) : sizeof(*reply);
 
     rbuf = wpa_sm_alloc_eapol(sm, IEEE802_1X_TYPE_EAPOL_KEY, NULL,
@@ -1574,7 +1586,7 @@ static int wpa_supplicant_send_2_of_2(struct wpa_sm *sm,
     struct wpa_eapol_key_192 *reply192;
     u8 *rbuf, *key_mic;
 
-    mic_len = wpa_mic_len(sm->key_mgmt);
+    mic_len = wpa_mic_len(sm->key_mgmt, sm->pmk_len);
     hdrlen = mic_len == 24 ? sizeof(*reply192) : sizeof(*reply);
 
     rbuf = wpa_sm_alloc_eapol(sm, IEEE802_1X_TYPE_EAPOL_KEY, NULL,
@@ -1683,7 +1695,7 @@ static int wpa_supplicant_verify_eapol_key_mic(struct wpa_sm *sm,
 {
     u8 mic[WPA_EAPOL_KEY_MIC_MAX_LEN];
     int ok = 0;
-    size_t mic_len = wpa_mic_len(sm->key_mgmt);
+    size_t mic_len = wpa_mic_len(sm->key_mgmt, sm->pmk_len);
 
     os_memcpy(mic, key->key_mic, mic_len);
     if (sm->tptk_set) {
@@ -1754,10 +1766,7 @@ static int wpa_supplicant_decrypt_key_data(struct wpa_sm *sm,
         }
     } else if (ver == WPA_KEY_INFO_TYPE_HMAC_SHA1_AES ||
                ver == WPA_KEY_INFO_TYPE_AES_128_CMAC ||
-               sm->key_mgmt == WPA_KEY_MGMT_OSEN ||
-               wpa_key_mgmt_suite_b(sm->key_mgmt) ||
-               sm->key_mgmt == WPA_KEY_MGMT_SAE ||
-               sm->key_mgmt == WPA_KEY_MGMT_OWE) {
+               wpa_use_aes_key_wrap(sm->key_mgmt)) {
         u8 *buf;
         if (*key_data_len < 8 || *key_data_len % 8) {
             wpa_printf(MSG_DEBUG, "WPA: Unsupported "
@@ -1860,7 +1869,7 @@ int wpa_sm_rx_eapol(u8 *src_addr, u8 *buf, u32 len)
     size_t mic_len, keyhdrlen;
     u8 *key_data;
 
-    mic_len = wpa_mic_len(sm->key_mgmt);
+    mic_len = wpa_mic_len(sm->key_mgmt, sm->pmk_len);
     keyhdrlen = mic_len == 24 ? sizeof(*key192) : sizeof(*key);
 
     if (len < sizeof(*hdr) + keyhdrlen) {
@@ -1924,20 +1933,14 @@ int wpa_sm_rx_eapol(u8 *src_addr, u8 *buf, u32 len)
     if (ver != WPA_KEY_INFO_TYPE_HMAC_MD5_RC4 &&
 #ifdef CONFIG_IEEE80211W
         ver != WPA_KEY_INFO_TYPE_AES_128_CMAC &&
-#ifdef CONFIG_WPA3_SAE
-        sm->key_mgmt != WPA_KEY_MGMT_SAE &&
 #endif
-        !wpa_key_mgmt_suite_b(sm->key_mgmt) &&
-#ifdef CONFIG_OWE_STA
-        sm->key_mgmt != WPA_KEY_MGMT_OWE &&
-#endif /* CONFIG_OWE_STA */
-#endif
-        ver != WPA_KEY_INFO_TYPE_HMAC_SHA1_AES) {
+        ver != WPA_KEY_INFO_TYPE_HMAC_SHA1_AES &&
+        !wpa_use_akm_defined(sm->key_mgmt)) {
         wpa_printf(MSG_DEBUG, "WPA: Unsupported EAPOL-Key descriptor "
                "version %d.", ver);
         goto out;
     }
-    if (wpa_key_mgmt_suite_b(sm->key_mgmt) &&
+    if (wpa_use_akm_defined(sm->key_mgmt) &&
         ver != WPA_KEY_INFO_TYPE_AKM_DEFINED) {
         wpa_msg(NULL, MSG_INFO,
                 "RSN: Unsupported EAPOL-Key descriptor version %d (expected AKM defined = 0)",
@@ -1948,20 +1951,15 @@ int wpa_sm_rx_eapol(u8 *src_addr, u8 *buf, u32 len)
 #ifdef CONFIG_IEEE80211W
     if (wpa_key_mgmt_sha256(sm->key_mgmt)) {
         if (ver != WPA_KEY_INFO_TYPE_AES_128_CMAC &&
-            sm->key_mgmt != WPA_KEY_MGMT_OSEN &&
-            !wpa_key_mgmt_suite_b(sm->key_mgmt) &&
-            sm->key_mgmt != WPA_KEY_MGMT_SAE &&
-            sm->key_mgmt != WPA_KEY_MGMT_OWE) {
+            !wpa_use_akm_defined(sm->key_mgmt)) {
             goto out;
         }
     } else
 #endif
 
     if (sm->pairwise_cipher == WPA_CIPHER_CCMP &&
-        !wpa_key_mgmt_suite_b(sm->key_mgmt) &&
-        ver != WPA_KEY_INFO_TYPE_HMAC_SHA1_AES &&
-        sm->key_mgmt != WPA_KEY_MGMT_SAE &&
-        sm->key_mgmt != WPA_KEY_MGMT_OWE) {
+        !wpa_use_akm_defined(sm->key_mgmt) &&
+        ver != WPA_KEY_INFO_TYPE_HMAC_SHA1_AES ) {
         wpa_printf(MSG_DEBUG, "WPA: CCMP is used, but EAPOL-Key "
                "descriptor version (%d) is not 2.", ver);
         if (sm->group_cipher != WPA_CIPHER_CCMP &&
@@ -1982,7 +1980,7 @@ int wpa_sm_rx_eapol(u8 *src_addr, u8 *buf, u32 len)
 
 #ifdef CONFIG_GCMP
     if (sm->pairwise_cipher == WPA_CIPHER_GCMP &&
-        !wpa_key_mgmt_suite_b(sm->key_mgmt) &&
+        !wpa_use_akm_defined(sm->key_mgmt) &&
         ver != WPA_KEY_INFO_TYPE_HMAC_SHA1_AES) {
         wpa_msg(NULL, MSG_INFO,
                 "WPA: GCMP is used, but EAPOL-Key "
diff --git a/components/wpa_supplicant/src/rsn_supp/wpa_ie.c b/components/wpa_supplicant/src/rsn_supp/wpa_ie.c
index 12b40ba6c6..2fc2307d71 100644
--- a/components/wpa_supplicant/src/rsn_supp/wpa_ie.c
+++ b/components/wpa_supplicant/src/rsn_supp/wpa_ie.c
@@ -51,6 +51,7 @@ static int  wpa_gen_wpa_ie_wpa(u8 *wpa_ie, size_t wpa_ie_len,
 {
     u8 *pos;
     struct wpa_ie_hdr *hdr;
+    u32 suite;
 
     if (wpa_ie_len < sizeof(*hdr) + WPA_SELECTOR_LEN +
         2 + WPA_SELECTOR_LEN + 2 + WPA_SELECTOR_LEN)
@@ -62,34 +63,26 @@ static int  wpa_gen_wpa_ie_wpa(u8 *wpa_ie, size_t wpa_ie_len,
     WPA_PUT_LE16(hdr->version, WPA_VERSION);
     pos = (u8 *) (hdr + 1);
 
-    if (group_cipher == WPA_CIPHER_CCMP) {
-        RSN_SELECTOR_PUT(pos, WPA_CIPHER_SUITE_CCMP);
-    } else if (group_cipher == WPA_CIPHER_TKIP) {
-        RSN_SELECTOR_PUT(pos, WPA_CIPHER_SUITE_TKIP);
-    } else if (group_cipher == WPA_CIPHER_WEP104) {
-        RSN_SELECTOR_PUT(pos, WPA_CIPHER_SUITE_WEP104);
-    } else if (group_cipher == WPA_CIPHER_WEP40) {
-        RSN_SELECTOR_PUT(pos, WPA_CIPHER_SUITE_WEP40);
-    } else {
-        wpa_printf(MSG_DEBUG, "Invalid group cipher (%d).",
+    suite = wpa_cipher_to_suite(WPA_PROTO_WPA, group_cipher);
+    if (suite == 0) {
+        wpa_printf(MSG_WARNING, "Invalid group cipher (%d).",
                group_cipher);
         return -1;
     }
+    RSN_SELECTOR_PUT(pos, suite);
     pos += WPA_SELECTOR_LEN;
 
     *pos++ = 1;
     *pos++ = 0;
-    if (pairwise_cipher == WPA_CIPHER_CCMP) {
-        RSN_SELECTOR_PUT(pos, WPA_CIPHER_SUITE_CCMP);
-    } else if (pairwise_cipher == WPA_CIPHER_TKIP) {
-        RSN_SELECTOR_PUT(pos, WPA_CIPHER_SUITE_TKIP);
-    } else if (pairwise_cipher == WPA_CIPHER_NONE) {
-        RSN_SELECTOR_PUT(pos, WPA_CIPHER_SUITE_NONE);
-    } else {
-        wpa_printf(MSG_DEBUG, "Invalid pairwise cipher (%d).",
+    suite = wpa_cipher_to_suite(WPA_PROTO_WPA, pairwise_cipher);
+    if (suite == 0 ||
+        (!wpa_cipher_valid_pairwise(pairwise_cipher) &&
+         pairwise_cipher != WPA_CIPHER_NONE)) {
+        wpa_printf(MSG_WARNING, "Invalid pairwise cipher (%d).",
                pairwise_cipher);
         return -1;
     }
+    RSN_SELECTOR_PUT(pos, suite);
     pos += WPA_SELECTOR_LEN;
 
     *pos++ = 1;

From b95ace06b8427b12e61a21f61d87840edc9d7733 Mon Sep 17 00:00:00 2001
From: Jouni Malinen <quic_jouni@quicinc.com>
Date: Mon, 9 Oct 2023 08:43:18 +0530
Subject: [PATCH 2/8] feat(wifi): Pull in wpa_supplicant upstream code for SAE
 EXT key

Define new RSN AKM suite selector values

Add the new AKM suite selectors defined in IEEE P802.11-REVme/D1.3.

Signed-off-by: Jouni Malinen <quic_jouni@quicinc.com>
---
 components/wpa_supplicant/src/common/wpa_common.h | 1 +
 1 file changed, 1 insertion(+)

diff --git a/components/wpa_supplicant/src/common/wpa_common.h b/components/wpa_supplicant/src/common/wpa_common.h
index 1f9adb76a2..78148c7eb1 100644
--- a/components/wpa_supplicant/src/common/wpa_common.h
+++ b/components/wpa_supplicant/src/common/wpa_common.h
@@ -69,6 +69,7 @@
 RSN_SELECTOR(0x00, 0x0f, 0xac, 13)
 #define RSN_AUTH_KEY_MGMT_FT_SAE RSN_SELECTOR(0x00, 0x0f, 0xac, 9)
 #define RSN_AUTH_KEY_MGMT_OWE RSN_SELECTOR(0x00, 0x0f, 0xac, 18)
+#define RSN_AUTH_KEY_MGMT_SAE_EXT_KEY RSN_SELECTOR(0x00, 0x0f, 0xac, 24)
 
 #define RSN_CIPHER_SUITE_NONE RSN_SELECTOR(0x00, 0x0f, 0xac, 0)
 #define RSN_CIPHER_SUITE_WEP40 RSN_SELECTOR(0x00, 0x0f, 0xac, 1)

From 7dbfd8c6dabf4032fbbf44dcd5e7b79148ceb5a5 Mon Sep 17 00:00:00 2001
From: Jouni Malinen <quic_jouni@quicinc.com>
Date: Mon, 9 Oct 2023 08:40:41 +0530
Subject: [PATCH 3/8] feat(wifi): Pull in wpa_supplicant upstream code for SAE
 EXT key

SAE: Use wpa_key_mgmt_sae() helper

Use the existing helper function instead of maintaining multiple copies
of lists of SAE key management suites.

Signed-off-by: Jouni Malinen <quic_jouni@quicinc.com>
---
 components/wpa_supplicant/esp_supplicant/src/esp_wpas_glue.c | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/components/wpa_supplicant/esp_supplicant/src/esp_wpas_glue.c b/components/wpa_supplicant/esp_supplicant/src/esp_wpas_glue.c
index 9a74c4e6b7..d3040b168e 100644
--- a/components/wpa_supplicant/esp_supplicant/src/esp_wpas_glue.c
+++ b/components/wpa_supplicant/esp_supplicant/src/esp_wpas_glue.c
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: 2019-2021 Espressif Systems (Shanghai) CO LTD
+ * SPDX-FileCopyrightText: 2019-2023 Espressif Systems (Shanghai) CO LTD
  *
  * SPDX-License-Identifier: Apache-2.0
  */
@@ -97,7 +97,8 @@ void wpa_supplicant_transition_disable(struct wpa_sm *sm, u8 bitmap)
 {
     wpa_printf(MSG_DEBUG, "TRANSITION_DISABLE %02x", bitmap);
 
-    if (bitmap & TRANSITION_DISABLE_WPA3_PERSONAL) {
+    if  ((bitmap & TRANSITION_DISABLE_WPA3_PERSONAL) &&
+         wpa_key_mgmt_sae(sm->key_mgmt)) {
         esp_wifi_sta_disable_wpa2_authmode_internal();
     }
 }

From 21ac680eac0f60896a2713a662f352a3eae8cdeb Mon Sep 17 00:00:00 2001
From: Jouni Malinen <quic_jouni@quicinc.com>
Date: Mon, 9 Oct 2023 08:55:55 +0530
Subject: [PATCH 4/8] feat(wifi): Pull in wpa_supplicant upstream code for SAE
 EXT key

SAE: Internal WPA_KEY_MGMT_* defines for extended key AKMs

Define new WPA_KEY_MGMT_* values for the new SAE AKM suite selectors
with variable length keys. This includes updates to various mapping and
checking of the SAE key_mgmt values.

Signed-off-by: Jouni Malinen <quic_jouni@quicinc.com>
---
 components/wpa_supplicant/src/common/defs.h       | 8 ++++++++
 components/wpa_supplicant/src/common/wpa_common.c | 4 ++++
 components/wpa_supplicant/src/rsn_supp/wpa_ie.c   | 2 ++
 3 files changed, 14 insertions(+)

diff --git a/components/wpa_supplicant/src/common/defs.h b/components/wpa_supplicant/src/common/defs.h
index 6799370dec..36d37d6e6a 100644
--- a/components/wpa_supplicant/src/common/defs.h
+++ b/components/wpa_supplicant/src/common/defs.h
@@ -48,6 +48,7 @@ typedef enum { FALSE = 0, TRUE = 1 } Boolean;
 #define WPA_KEY_MGMT_IEEE8021X_SUITE_B BIT(16)
 #define WPA_KEY_MGMT_IEEE8021X_SUITE_B_192 BIT(17)
 #define WPA_KEY_MGMT_OWE BIT(22)
+#define WPA_KEY_MGMT_SAE_EXT_KEY BIT(26)
 
 static inline int wpa_key_mgmt_wpa_ieee8021x(int akm)
 {
@@ -70,6 +71,7 @@ static inline int wpa_key_mgmt_wpa_psk(int akm)
 			 WPA_KEY_MGMT_FT_PSK |
 			 WPA_KEY_MGMT_PSK_SHA256 |
 			 WPA_KEY_MGMT_SAE |
+			 WPA_KEY_MGMT_SAE_EXT_KEY |
 			 WPA_KEY_MGMT_FT_SAE));
 }
 
@@ -83,9 +85,15 @@ static inline int wpa_key_mgmt_ft(int akm)
 static inline int wpa_key_mgmt_sae(int akm)
 {
 	return !!(akm & (WPA_KEY_MGMT_SAE |
+			 WPA_KEY_MGMT_SAE_EXT_KEY |
 			 WPA_KEY_MGMT_FT_SAE));
 }
 
+static inline int wpa_key_mgmt_sae_ext_key(int akm)
+{
+	return !!(akm & (WPA_KEY_MGMT_SAE_EXT_KEY));
+}
+
 static inline int wpa_key_mgmt_sha256(int akm)
 {
 	return !!(akm & (WPA_KEY_MGMT_PSK_SHA256 |
diff --git a/components/wpa_supplicant/src/common/wpa_common.c b/components/wpa_supplicant/src/common/wpa_common.c
index cbf1fa630f..f1b302191b 100644
--- a/components/wpa_supplicant/src/common/wpa_common.c
+++ b/components/wpa_supplicant/src/common/wpa_common.c
@@ -335,6 +335,8 @@ static int rsn_key_mgmt_to_bitfield(const u8 *s)
 #ifdef CONFIG_WPA3_SAE
 	if (RSN_SELECTOR_GET(s) == RSN_AUTH_KEY_MGMT_SAE)
 		return WPA_KEY_MGMT_SAE;
+	if (RSN_SELECTOR_GET(s) == RSN_AUTH_KEY_MGMT_SAE_EXT_KEY)
+		return WPA_KEY_MGMT_SAE_EXT_KEY;
 #endif /* CONFIG_WPA3_SAE */
 #ifdef CONFIG_OWE_STA
 	if(RSN_SELECTOR_GET(s) == RSN_AUTH_KEY_MGMT_OWE)
@@ -969,6 +971,8 @@ u32 wpa_akm_to_suite(int akm)
 		return RSN_AUTH_KEY_MGMT_802_1X_SUITE_B_192;
 	if (akm & WPA_KEY_MGMT_SAE)
 		return RSN_AUTH_KEY_MGMT_SAE;
+	if (akm & WPA_KEY_MGMT_SAE_EXT_KEY)
+		return RSN_AUTH_KEY_MGMT_SAE_EXT_KEY;
 	if (akm & WPA_KEY_MGMT_FT_SAE)
 		return RSN_AUTH_KEY_MGMT_FT_SAE;
 	if (akm & WPA_KEY_MGMT_OWE)
diff --git a/components/wpa_supplicant/src/rsn_supp/wpa_ie.c b/components/wpa_supplicant/src/rsn_supp/wpa_ie.c
index 2fc2307d71..34969aee6c 100644
--- a/components/wpa_supplicant/src/rsn_supp/wpa_ie.c
+++ b/components/wpa_supplicant/src/rsn_supp/wpa_ie.c
@@ -191,6 +191,8 @@ static int wpa_gen_wpa_ie_rsn(u8 *rsn_ie, size_t rsn_ie_len,
 #ifdef CONFIG_WPA3_SAE
     } else if (key_mgmt == WPA_KEY_MGMT_SAE) {
         RSN_SELECTOR_PUT(pos, RSN_AUTH_KEY_MGMT_SAE);
+    } else if (key_mgmt == WPA_KEY_MGMT_SAE_EXT_KEY) {
+        RSN_SELECTOR_PUT(pos, RSN_AUTH_KEY_MGMT_SAE_EXT_KEY);
 #endif /* CONFIG_WPA3_SAE */
 #ifdef CONFIG_OWE_STA
     } else if (key_mgmt == WPA_KEY_MGMT_OWE) {

From 5719f7908d777943fd00f9a12ce219d0af3109e4 Mon Sep 17 00:00:00 2001
From: Jouni Malinen <quic_jouni@quicinc.com>
Date: Mon, 9 Oct 2023 09:44:55 +0530
Subject: [PATCH 5/8] feat(wifi): Pull in wpa_supplicant upstream code for SAE
 EXT key

SAE: EAPOL-Key and key/MIC length information for the new AKM suites

Update the AKM suite specific mapping of various EAPOL-Key key lengths
and algorithms to include the new SAE AKM suites with variable length
keys.

Signed-off-by: Jouni Malinen <quic_jouni@quicinc.com>
---
 .../wpa_supplicant/src/common/wpa_common.c    | 35 +++++++++++++++++++
 1 file changed, 35 insertions(+)

diff --git a/components/wpa_supplicant/src/common/wpa_common.c b/components/wpa_supplicant/src/common/wpa_common.c
index f1b302191b..2bffbbe758 100644
--- a/components/wpa_supplicant/src/common/wpa_common.c
+++ b/components/wpa_supplicant/src/common/wpa_common.c
@@ -242,6 +242,8 @@ static unsigned int wpa_kck_len(int akmp, size_t pmk_len)
 		return 24;
 	case WPA_KEY_MGMT_OWE:
 		return pmk_len / 2;
+	case WPA_KEY_MGMT_SAE_EXT_KEY:
+		return pmk_len / 2;
 	default:
 		return 16;
 	}
@@ -254,6 +256,8 @@ static unsigned int wpa_kek_len(int akmp, size_t pmk_len)
 		return 32;
 	case WPA_KEY_MGMT_OWE:
 		return pmk_len <= 32 ? 16 : 32;
+	case WPA_KEY_MGMT_SAE_EXT_KEY:
+		return pmk_len <= 32 ? 16 : 32;
 	default:
 		return 16;
 	}
@@ -265,6 +269,8 @@ unsigned int wpa_mic_len(int akmp, size_t pmk_len)
 	switch (akmp) {
 	case WPA_KEY_MGMT_IEEE8021X_SUITE_B_192:
 		return 24;
+	case WPA_KEY_MGMT_SAE_EXT_KEY:
+		return pmk_len / 2;
 	default:
 		return 16;
 	}
@@ -904,6 +910,21 @@ int wpa_eapol_key_mic(const u8 *key, size_t key_len, int akmp, int ver,
 #ifdef CONFIG_WPA3_SAE
 		case WPA_KEY_MGMT_SAE:
 			return omac1_aes_128(key, buf, len, mic);
+		case WPA_KEY_MGMT_SAE_EXT_KEY:
+			wpa_printf(MSG_DEBUG,
+				   "WPA: EAPOL-Key MIC using HMAC-SHA%u (AKM-defined - SAE-EXT-KEY)",
+				   (unsigned int) key_len * 8 * 2);
+			if (key_len == 128 / 8) {
+				if (hmac_sha256(key, key_len, buf, len, hash))
+					return -1;
+			} else {
+				wpa_printf(MSG_INFO,
+					   "SAE: Unsupported KCK length: %u",
+					   (unsigned int) key_len);
+				return -1;
+			}
+			os_memcpy(mic, hash, key_len);
+			break;
 #endif /* CONFIG_WPA3_SAE */
 #ifdef CONFIG_SUITEB
 		case WPA_KEY_MGMT_IEEE8021X_SUITE_B:
@@ -1181,6 +1202,20 @@ int wpa_pmk_to_ptk(const u8 *pmk, size_t pmk_len, const char *label,
 	if (wpa_key_mgmt_sha256(akmp)) {
 		sha256_prf(pmk, pmk_len, label, data, data_len,
 			   tmp, ptk_len);
+#ifdef CONFIG_WPA3_SAE
+	} else if (wpa_key_mgmt_sae_ext_key(akmp)) {
+		if (pmk_len == 32) {
+			wpa_printf(MSG_DEBUG,
+				   "SAE: PTK derivation using PRF(SHA256)");
+			if (sha256_prf(pmk, pmk_len, label, data, data_len,
+				       tmp, ptk_len) < 0)
+				return -1;
+		} else {
+			wpa_printf(MSG_INFO, "SAE: Unknown PMK length %u",
+				   (unsigned int) pmk_len);
+			return -1;
+		}
+#endif /* CONFIG_WPA3_SAE */
 	} else {
 		sha1_prf(pmk, pmk_len, label, data, data_len, tmp, ptk_len);
 	}

From 3ec105067c7a57aaea66e5ab187d97a41e6b87e7 Mon Sep 17 00:00:00 2001
From: Jouni Malinen <quic_jouni@quicinc.com>
Date: Mon, 9 Oct 2023 10:58:44 +0530
Subject: [PATCH 6/8] feat(wifi): Pull in wpa_supplicant upstream code for SAE
 EXT key

SAE: Indicate AKM suite selector in commit for new AKM suites

SAE authentication needs to known which AKM suite is being used to be
able to determine the correct PMK length for the new AKM suite selectors
that use variable length keys.

Signed-off-by: Jouni Malinen <quic_jouni@quicinc.com>
---
 .../src/common/ieee802_11_defs.h              |  1 +
 components/wpa_supplicant/src/common/sae.c    | 74 +++++++++++++++++++
 components/wpa_supplicant/src/common/sae.h    |  4 +
 3 files changed, 79 insertions(+)

diff --git a/components/wpa_supplicant/src/common/ieee802_11_defs.h b/components/wpa_supplicant/src/common/ieee802_11_defs.h
index 3c7eab6a44..b496690ab7 100644
--- a/components/wpa_supplicant/src/common/ieee802_11_defs.h
+++ b/components/wpa_supplicant/src/common/ieee802_11_defs.h
@@ -255,6 +255,7 @@
 #define WLAN_EID_EXT_HE_OPERATION 36
 #define WLAN_EID_EXT_REJECTED_GROUPS 92
 #define WLAN_EID_EXT_ANTI_CLOGGING_TOKEN 93
+#define WLAN_EID_EXT_AKM_SUITE_SELECTOR 114
 
 /* Extended RSN Capabilities */
 /* bits 0-3: Field length (n-1) */
diff --git a/components/wpa_supplicant/src/common/sae.c b/components/wpa_supplicant/src/common/sae.c
index bc394c80db..ee5ff1b06e 100644
--- a/components/wpa_supplicant/src/common/sae.c
+++ b/components/wpa_supplicant/src/common/sae.c
@@ -1681,6 +1681,18 @@ int sae_write_commit(struct sae_data *sae, struct wpabuf *buf,
 				"SAE: Anti-clogging token (in container)",
 				token);
 	}
+
+	if (wpa_key_mgmt_sae_ext_key(sae->akmp)) {
+		u32 suite = wpa_akm_to_suite(sae->akmp);
+
+		wpabuf_put_u8(buf, WLAN_EID_EXTENSION);
+		wpabuf_put_u8(buf, 1 + RSN_SELECTOR_LEN);
+		wpabuf_put_u8(buf, WLAN_EID_EXT_AKM_SUITE_SELECTOR);
+		RSN_SELECTOR_PUT(wpabuf_put(buf, RSN_SELECTOR_LEN), suite);
+		wpa_printf(MSG_DEBUG, "SAE: AKM Suite Selector: %08x", suite);
+		sae->own_akm_suite_selector = suite;
+	}
+
 	return ESP_OK;
 }
 
@@ -1757,6 +1769,16 @@ static int sae_is_token_container_elem(const u8 *pos, const u8 *end)
 }
 
 
+static int sae_is_akm_suite_selector_elem(const u8 *pos, const u8 *end)
+{
+	return end - pos >= 2 + 1 + RSN_SELECTOR_LEN &&
+		pos[0] == WLAN_EID_EXTENSION &&
+		pos[1] >= 1 + RSN_SELECTOR_LEN &&
+		end - pos - 2 >= pos[1] &&
+		pos[2] == WLAN_EID_EXT_AKM_SUITE_SELECTOR;
+}
+
+
 static void sae_parse_commit_token(struct sae_data *sae, const u8 **pos,
 				   const u8 *end, const u8 **token,
 				   size_t *token_len, int h2e)
@@ -2044,6 +2066,35 @@ static int sae_parse_rejected_groups(struct sae_data *sae,
 }
 
 
+static int sae_parse_akm_suite_selector(struct sae_data *sae,
+					const u8 **pos, const u8 *end)
+{
+	const u8 *epos;
+	u8 len;
+
+	wpa_hexdump(MSG_DEBUG, "SAE: Possible elements at the end of the frame",
+		    *pos, end - *pos);
+	if (!sae_is_akm_suite_selector_elem(*pos, end))
+		return WLAN_STATUS_SUCCESS;
+
+	epos = *pos;
+	epos++; /* skip IE type */
+	len = *epos++; /* IE length */
+	if (len > end - epos || len < 1)
+		return WLAN_STATUS_UNSPECIFIED_FAILURE;
+	epos++; /* skip ext ID */
+	len--;
+
+	if (len < RSN_SELECTOR_LEN)
+		return WLAN_STATUS_UNSPECIFIED_FAILURE;
+	sae->peer_akm_suite_selector = RSN_SELECTOR_GET(epos);
+	wpa_printf(MSG_DEBUG, "SAE: Received AKM Suite Selector: %08x",
+		   sae->peer_akm_suite_selector);
+	*pos = epos + len;
+	return WLAN_STATUS_SUCCESS;
+}
+
+
 u16 sae_parse_commit(struct sae_data *sae, const u8 *data, size_t len,
 		     const u8 **token, size_t *token_len, int *allowed_groups,
 		     int h2e)
@@ -2088,6 +2139,29 @@ u16 sae_parse_commit(struct sae_data *sae, const u8 *data, size_t len,
 	if (h2e)
 		sae_parse_token_container(sae, pos, end, token, token_len);
 
+	/* Conditional AKM Suite Selector element */
+	if (h2e) {
+		res = sae_parse_akm_suite_selector(sae, &pos, end);
+		if (res != WLAN_STATUS_SUCCESS)
+			return res;
+	}
+
+	if (sae->own_akm_suite_selector &&
+	    sae->own_akm_suite_selector != sae->peer_akm_suite_selector) {
+		wpa_printf(MSG_DEBUG,
+			   "SAE: AKM suite selector mismatch: own=%08x peer=%08x",
+			   sae->own_akm_suite_selector,
+			   sae->peer_akm_suite_selector);
+		return WLAN_STATUS_UNSPECIFIED_FAILURE;
+	}
+
+	if (!sae->akmp) {
+		if (sae->peer_akm_suite_selector ==
+		    RSN_AUTH_KEY_MGMT_SAE_EXT_KEY) {
+			sae->akmp = WPA_KEY_MGMT_SAE_EXT_KEY;
+		}
+	}
+
 	/*
 	 * Check whether peer-commit-scalar and PEER-COMMIT-ELEMENT are same as
 	 * the values we sent which would be evidence of a reflection attack.
diff --git a/components/wpa_supplicant/src/common/sae.h b/components/wpa_supplicant/src/common/sae.h
index 0f266bc1c2..028330268a 100644
--- a/components/wpa_supplicant/src/common/sae.h
+++ b/components/wpa_supplicant/src/common/sae.h
@@ -109,6 +109,10 @@ struct sae_data {
 	enum sae_state state;
 	u16 send_confirm;
 	u8 pmk[SAE_PMK_LEN];
+	size_t pmk_len;
+	int akmp; /* WPA_KEY_MGMT_* used in key derivation */
+	u32 own_akm_suite_selector;
+	u32 peer_akm_suite_selector;
 	u8 pmkid[SAE_PMKID_LEN];
 	struct crypto_bignum *peer_commit_scalar;
 	struct crypto_bignum *peer_commit_scalar_accepted;

From 98cc860e8602732815e6704819f06285e53d4cb6 Mon Sep 17 00:00:00 2001
From: Jouni Malinen <quic_jouni@quicinc.com>
Date: Mon, 9 Oct 2023 11:00:46 +0530
Subject: [PATCH 7/8] feat(wifi): Pull in wpa_supplicant upstream code for SAE
 EXT key

SAE: Derive a variable length PMK with the new AKM suites

Select the PMK length based on the used group (prime length) when using
the new AKM suites for SAE.

Signed-off-by: Jouni Malinen <quic_jouni@quicinc.com>
---
 components/wpa_supplicant/src/common/sae.c | 20 +++++++++++++-------
 components/wpa_supplicant/src/common/sae.h |  2 +-
 2 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/components/wpa_supplicant/src/common/sae.c b/components/wpa_supplicant/src/common/sae.c
index ee5ff1b06e..37b918230c 100644
--- a/components/wpa_supplicant/src/common/sae.c
+++ b/components/wpa_supplicant/src/common/sae.c
@@ -10,6 +10,7 @@
 
 #include "utils/includes.h"
 #include "utils/common.h"
+#include "common/wpa_common.h"
 #include "utils/const_time.h"
 #include "crypto/crypto.h"
 #include "crypto/sha256.h"
@@ -1466,10 +1467,11 @@ static int sae_derive_keys(struct sae_data *sae, const u8 *k)
 	const u8 *salt;
 	struct wpabuf *rejected_groups = NULL;
 	u8 keyseed[SAE_MAX_HASH_LEN];
-	u8 keys[2 * SAE_MAX_HASH_LEN + SAE_PMK_LEN];
+	u8 keys[2 * SAE_MAX_HASH_LEN + SAE_PMK_LEN_MAX];
 	struct crypto_bignum *tmp;
 	int ret = -1;
 	size_t hash_len, salt_len, prime_len = sae->tmp->prime_len;
+	size_t pmk_len;
 	const u8 *addr[1];
 	size_t len[1];
 
@@ -1491,6 +1493,10 @@ static int sae_derive_keys(struct sae_data *sae, const u8 *k)
 		hash_len = sae_ffc_prime_len_2_hash_len(prime_len);
 	else
 		hash_len = sae_ecc_prime_len_2_hash_len(prime_len);
+	if (wpa_key_mgmt_sae_ext_key(sae->akmp))
+		pmk_len = hash_len;
+	else
+		pmk_len = SAE_PMK_LEN;
 
 	if (sae->h2e && (sae->tmp->own_rejected_groups ||
 			 sae->tmp->peer_rejected_groups)) {
@@ -1552,27 +1558,27 @@ static int sae_derive_keys(struct sae_data *sae, const u8 *k)
 	if (sae->pk) {
 		if (sae_kdf_hash(hash_len, keyseed, "SAE-PK keys",
 				 val, sae->tmp->order_len,
-				 keys, 2 * hash_len + SAE_PMK_LEN) < 0)
+				 keys, 2 * hash_len + pmk_len) < 0)
 			goto fail;
 	} else {
 		if (sae_kdf_hash(hash_len, keyseed, "SAE KCK and PMK",
 				 val, sae->tmp->order_len,
-				 keys, hash_len + SAE_PMK_LEN) < 0)
+				 keys, hash_len + pmk_len) < 0)
 			goto fail;
 	}
 #else /* CONFIG_SAE_PK */
 	if (sae_kdf_hash(hash_len, keyseed, "SAE KCK and PMK",
 			 val, sae->tmp->order_len,
-			 keys, hash_len + SAE_PMK_LEN) < 0)
+			 keys, hash_len + pmk_len) < 0)
 		goto fail;
 #endif /* !CONFIG_SAE_PK */
 
 	forced_memzero(keyseed, sizeof(keyseed));
 	os_memcpy(sae->tmp->kck, keys, hash_len);
 	sae->tmp->kck_len = hash_len;
-	os_memcpy(sae->pmk, keys + hash_len, SAE_PMK_LEN);
+	os_memcpy(sae->pmk, keys + hash_len, pmk_len);
+	sae->pmk_len = pmk_len;
 	os_memcpy(sae->pmkid, val, SAE_PMKID_LEN);
-
 #ifdef CONFIG_SAE_PK
 	if (sae->pk) {
 		os_memcpy(sae->tmp->kek, keys + hash_len + SAE_PMK_LEN,
@@ -1585,7 +1591,7 @@ static int sae_derive_keys(struct sae_data *sae, const u8 *k)
 	forced_memzero(keys, sizeof(keys));
 	wpa_hexdump_key(MSG_DEBUG, "SAE: KCK",
 			sae->tmp->kck, sae->tmp->kck_len);
-	wpa_hexdump_key(MSG_DEBUG, "SAE: PMK", sae->pmk, SAE_PMK_LEN);
+	wpa_hexdump_key(MSG_DEBUG, "SAE: PMK", sae->pmk, sae->pmk_len);
 
 	ret = 0;
 fail:
diff --git a/components/wpa_supplicant/src/common/sae.h b/components/wpa_supplicant/src/common/sae.h
index 028330268a..ca4233159b 100644
--- a/components/wpa_supplicant/src/common/sae.h
+++ b/components/wpa_supplicant/src/common/sae.h
@@ -108,7 +108,7 @@ enum sae_state {
 struct sae_data {
 	enum sae_state state;
 	u16 send_confirm;
-	u8 pmk[SAE_PMK_LEN];
+	u8 pmk[SAE_PMK_LEN_MAX];
 	size_t pmk_len;
 	int akmp; /* WPA_KEY_MGMT_* used in key derivation */
 	u32 own_akm_suite_selector;

From ee9ce5113388123c92bd237499c540db026b7cbd Mon Sep 17 00:00:00 2001
From: Shreyas Sheth <shreyas.sheth@espressif.com>
Date: Mon, 9 Oct 2023 11:02:41 +0530
Subject: [PATCH 8/8] feat(wifi): Add SAE EXT key support for esp station

---
 components/esp_wifi/include/esp_wifi_types.h  |  3 +-
 components/esp_wifi/lib                       |  2 +-
 components/wpa_supplicant/CMakeLists.txt      |  2 +-
 .../esp_supplicant/src/esp_common.c           | 12 +--
 .../esp_supplicant/src/esp_wifi_driver.h      |  3 +-
 .../esp_supplicant/src/esp_wpa3.c             | 88 ++++++++++---------
 .../esp_supplicant/src/esp_wpas_glue.c        |  2 +-
 .../wpa_supplicant/src/common/wpa_common.c    |  3 +-
 components/wpa_supplicant/src/rsn_supp/wpa.c  | 25 +++---
 components/wpa_supplicant/src/rsn_supp/wpa.h  |  2 +-
 examples/wifi/scan/main/scan.c                |  3 +
 11 files changed, 77 insertions(+), 68 deletions(-)

diff --git a/components/esp_wifi/include/esp_wifi_types.h b/components/esp_wifi/include/esp_wifi_types.h
index cabe296cf5..ae525511ba 100644
--- a/components/esp_wifi/include/esp_wifi_types.h
+++ b/components/esp_wifi/include/esp_wifi_types.h
@@ -56,7 +56,7 @@ typedef struct {
 } wifi_country_t;
 
 /* Strength of authmodes */
-/* OPEN < WEP < WPA_PSK < OWE < WPA2_PSK = WPA_WPA2_PSK < WAPI_PSK < WPA3_PSK = WPA2_WPA3_PSK */
+/* OPEN < WEP < WPA_PSK < OWE < WPA2_PSK = WPA_WPA2_PSK < WAPI_PSK < WPA3_PSK = WPA2_WPA3_PSK < WPA3_EXT_PSK */
 typedef enum {
     WIFI_AUTH_OPEN = 0,         /**< authenticate mode : open */
     WIFI_AUTH_WEP,              /**< authenticate mode : WEP */
@@ -70,6 +70,7 @@ typedef enum {
     WIFI_AUTH_WAPI_PSK,         /**< authenticate mode : WAPI_PSK */
     WIFI_AUTH_OWE,              /**< authenticate mode : OWE */
     WIFI_AUTH_WPA3_ENT_192,     /**< authenticate mode : WPA3_ENT_SUITE_B_192_BIT */
+    WIFI_AUTH_WPA3_EXT_PSK,     /**< authenticate mode : WPA3_PSK_EXT_KEY */
     WIFI_AUTH_MAX
 } wifi_auth_mode_t;
 
diff --git a/components/esp_wifi/lib b/components/esp_wifi/lib
index 557c6fa326..e88f8cebc1 160000
--- a/components/esp_wifi/lib
+++ b/components/esp_wifi/lib
@@ -1 +1 @@
-Subproject commit 557c6fa326dc423118576a5f9d4b34ba60e0d991
+Subproject commit e88f8cebc1b5353d979ee57f22a04e40752bfd66
diff --git a/components/wpa_supplicant/CMakeLists.txt b/components/wpa_supplicant/CMakeLists.txt
index a5c31d8a73..b278467f9c 100644
--- a/components/wpa_supplicant/CMakeLists.txt
+++ b/components/wpa_supplicant/CMakeLists.txt
@@ -172,7 +172,7 @@ if(CONFIG_ESP_WIFI_ENABLE_SAE_PK)
     "src/common/sae_pk.c")
 endif()
 
-if(CONFIG_ESP_WIFI_11KV_SUPPORT OR CONFIG_ESP_WIFI_11R_SUPPORT OR CONFIG_ESP_WIFI_ENABLE_SAE_PK)
+if(CONFIG_ESP_WIFI_11KV_SUPPORT OR CONFIG_ESP_WIFI_11R_SUPPORT)
     set(srcs ${srcs}
     "src/common/bss.c"
     "src/common/scan.c"
diff --git a/components/wpa_supplicant/esp_supplicant/src/esp_common.c b/components/wpa_supplicant/esp_supplicant/src/esp_common.c
index 20c2c1187b..56399020c7 100644
--- a/components/wpa_supplicant/esp_supplicant/src/esp_common.c
+++ b/components/wpa_supplicant/esp_supplicant/src/esp_common.c
@@ -260,12 +260,12 @@ static int ieee80211_handle_rx_frm(u8 type, u8 *frame, size_t len, u8 *sender,
 	int ret = 0;
 
 	switch (type) {
-#if defined(CONFIG_IEEE80211R) || defined(CONFIG_IEEE80211KV) || defined(CONFIG_SAE_PK)
+#if defined(CONFIG_IEEE80211R) || defined(CONFIG_IEEE80211KV)
 	case WLAN_FC_STYPE_BEACON:
 	case WLAN_FC_STYPE_PROBE_RESP:
 		ret = esp_handle_beacon_probe(type, frame, len, sender, rssi, channel, current_tsf);
 		break;
-#endif /* defined(CONFIG_IEEE80211KV) || defined(CONFIG_IEEE80211R) || defined(CONFIG_SAE_PK)*/
+#endif /* defined(CONFIG_IEEE80211KV) || defined(CONFIG_IEEE80211R) */
 #ifdef CONFIG_IEEE80211R
 	case WLAN_FC_STYPE_AUTH:
 		ret = handle_auth_frame(frame, len, sender, rssi, channel);
@@ -328,7 +328,7 @@ int esp_supplicant_common_init(struct wpa_funcs *wpa_cb)
 	struct wpa_supplicant *wpa_s = &g_wpa_supp;
 	int ret = 0;
 
-#if defined(CONFIG_IEEE80211KV) || defined(CONFIG_IEEE80211R) || defined(CONFIG_SAE_PK)
+#if defined(CONFIG_IEEE80211KV) || defined(CONFIG_IEEE80211R)
 #ifdef CONFIG_SUPPLICANT_TASK
 	if (!s_supplicant_api_lock) {
 		s_supplicant_api_lock = os_recursive_mutex_create();
@@ -361,7 +361,7 @@ int esp_supplicant_common_init(struct wpa_funcs *wpa_cb)
 #endif /* CONFIG_IEEE80211KV */
 	esp_scan_init(wpa_s);
 
-#endif /* defined(CONFIG_IEEE80211KV) || defined(CONFIG_IEEE80211R) || defined(CONFIG_SAE_PK)*/
+#endif /* defined(CONFIG_IEEE80211KV) || defined(CONFIG_IEEE80211R) */
 	wpa_s->type = 0;
 	wpa_s->subtype = 0;
 	wpa_s->type |= (1 << WLAN_FC_STYPE_ASSOC_RESP) | (1 << WLAN_FC_STYPE_REASSOC_RESP) | (1 << WLAN_FC_STYPE_AUTH);
@@ -386,13 +386,13 @@ void esp_supplicant_common_deinit(void)
 {
 	struct wpa_supplicant *wpa_s = &g_wpa_supp;
 
-#if defined(CONFIG_IEEE80211KV) || defined(CONFIG_IEEE80211R) || defined(CONFIG_SAE_PK)
+#if defined(CONFIG_IEEE80211KV) || defined(CONFIG_IEEE80211R)
 	esp_scan_deinit(wpa_s);
 #ifdef CONFIG_IEEE80211KV
 	wpas_rrm_reset(wpa_s);
 	wpas_clear_beacon_rep_data(wpa_s);
 #endif /* CONFIG_IEEE80211KV */
-#endif /* defined(CONFIG_IEEE80211KV) || defined(CONFIG_IEEE80211R) || defined(CONFIG_SAE_PK)*/
+#endif /* defined(CONFIG_IEEE80211KV) || defined(CONFIG_IEEE80211R) */
 	if (wpa_s->type) {
 		wpa_s->type = 0;
 		esp_wifi_register_mgmt_frame_internal(wpa_s->type, wpa_s->subtype);
diff --git a/components/wpa_supplicant/esp_supplicant/src/esp_wifi_driver.h b/components/wpa_supplicant/esp_supplicant/src/esp_wifi_driver.h
index a6ffe6f93e..8979dfc69d 100644
--- a/components/wpa_supplicant/esp_supplicant/src/esp_wifi_driver.h
+++ b/components/wpa_supplicant/esp_supplicant/src/esp_wifi_driver.h
@@ -72,6 +72,7 @@ enum {
     WPA2_AUTH_ENT_SHA384_SUITE_B = 0x0d,
     WPA2_AUTH_FT_PSK    = 0x0e,
     WPA3_AUTH_OWE       = 0x0f,
+    WPA3_AUTH_PSK_EXT_KEY = 0x10,
     WPA2_AUTH_INVALID
 };
 
@@ -286,7 +287,6 @@ esp_err_t esp_wifi_remain_on_channel(uint8_t ifx, uint8_t type, uint8_t channel,
 bool esp_wifi_is_mbo_enabled_internal(uint8_t if_index);
 void esp_wifi_get_pmf_config_internal(wifi_pmf_config_t *pmf_cfg, uint8_t ifx);
 bool esp_wifi_is_ft_enabled_internal(uint8_t if_index);
-uint8_t esp_wifi_sta_get_use_h2e_internal(void);
 uint8_t esp_wifi_sta_get_config_sae_pk_internal(void);
 void esp_wifi_sta_disable_sae_pk_internal(void);
 void esp_wifi_sta_disable_wpa2_authmode_internal(void);
@@ -296,5 +296,6 @@ bool esp_wifi_ap_notify_node_sae_auth_done(uint8_t *mac);
 bool esp_wifi_ap_is_sta_sae_reauth_node(uint8_t *mac);
 uint8_t* esp_wifi_sta_get_sae_identifier_internal(void);
 bool esp_wifi_eb_tx_status_success_internal(void *eb);
+uint8_t* esp_wifi_sta_get_rsnxe(void);
 
 #endif /* _ESP_WIFI_DRIVER_H_ */
diff --git a/components/wpa_supplicant/esp_supplicant/src/esp_wpa3.c b/components/wpa_supplicant/esp_supplicant/src/esp_wpa3.c
index e4c9a4e03a..ab926caad4 100644
--- a/components/wpa_supplicant/esp_supplicant/src/esp_wpa3.c
+++ b/components/wpa_supplicant/esp_supplicant/src/esp_wpa3.c
@@ -17,10 +17,6 @@
 #include "esp_hostap.h"
 #include <inttypes.h>
 
-#ifdef CONFIG_SAE_PK
-#include "common/bss.h"
-extern struct wpa_supplicant g_wpa_supp;
-#endif
 
 static struct sae_pt *g_sae_pt;
 static struct sae_data g_sae_data;
@@ -33,12 +29,56 @@ static esp_err_t wpa3_build_sae_commit(u8 *bssid, size_t *sae_msg_len)
 {
     int default_group = IANA_SECP256R1;
     u32 len = 0;
+    uint8_t use_pt = 0;
     u8 own_addr[ETH_ALEN];
     const u8 *pw = (const u8 *)esp_wifi_sta_get_prof_password_internal();
     struct wifi_ssid *ssid = esp_wifi_sta_get_prof_ssid_internal();
-    uint8_t use_pt = esp_wifi_sta_get_use_h2e_internal();
+    uint8_t sae_pwe = esp_wifi_get_config_sae_pwe_h2e_internal(WIFI_IF_STA);
     char sae_pwd_id[SAE_H2E_IDENTIFIER_LEN+1] = {0};
     bool valid_pwd_id = false;
+    const u8 *rsnxe;
+    u8 rsnxe_capa = 0;
+
+    if (wpa_key_mgmt_sae_ext_key(gWpaSm.key_mgmt)) {
+        use_pt = 1;
+    }
+
+   rsnxe = esp_wifi_sta_get_rsnxe();
+   if (rsnxe && rsnxe[1] >= 1) {
+       rsnxe_capa = rsnxe[2];
+   }
+
+#ifdef CONFIG_SAE_PK
+    bool use_pk = false;
+    uint8_t sae_pk_mode = esp_wifi_sta_get_config_sae_pk_internal();
+    if ((rsnxe_capa & BIT(WLAN_RSNX_CAPAB_SAE_PK)) &&
+         sae_pk_mode != WPA3_SAE_PK_MODE_DISABLED &&
+         ((pw && sae_pk_valid_password((const char*)pw)))) {
+        use_pt = 1;
+        use_pk = true;
+    }
+
+    if (sae_pk_mode == WPA3_SAE_PK_MODE_ONLY && !use_pk) {
+        wpa_printf(MSG_DEBUG,
+           "SAE: Cannot use PK with the selected AP");
+        return ESP_FAIL;
+    }
+#endif /* CONFIG_SAE_PK */
+    if (use_pt || sae_pwe == SAE_PWE_HASH_TO_ELEMENT ||
+        sae_pwe == SAE_PWE_BOTH) {
+        use_pt = !!(rsnxe_capa & BIT(WLAN_RSNX_CAPAB_SAE_H2E));
+
+        if ((sae_pwe == SAE_PWE_HASH_TO_ELEMENT ||
+            wpa_key_mgmt_sae_ext_key(gWpaSm.key_mgmt)
+#ifdef CONFIG_SAE_PK
+            || (use_pk && sae_pk_mode == WPA3_SAE_PK_MODE_ONLY)
+#endif /* CONFIG_SAE_PK */
+            ) && !use_pt) {
+            wpa_printf(MSG_DEBUG,
+              "SAE: Cannot use H2E with the selected AP");
+            return ESP_FAIL;
+        }
+    }
 
     if (use_pt != 0) {
         memcpy(sae_pwd_id, esp_wifi_sta_get_sae_identifier_internal(), SAE_H2E_IDENTIFIER_LEN);
@@ -73,46 +113,14 @@ static esp_err_t wpa3_build_sae_commit(u8 *bssid, size_t *sae_msg_len)
         return ESP_FAIL;
     }
 
+    g_sae_data.akmp = gWpaSm.key_mgmt;
+
     esp_wifi_get_macaddr_internal(WIFI_IF_STA, own_addr);
     if (!bssid) {
         wpa_printf(MSG_ERROR, "wpa3: cannot prepare SAE commit with no BSSID!");
         return ESP_FAIL;
     }
 
-#ifdef CONFIG_SAE_PK
-    bool use_pk = false;
-    uint8_t sae_pk_mode = esp_wifi_sta_get_config_sae_pk_internal();
-    u8 rsnxe_capa = 0;
-    struct wpa_bss *bss = wpa_bss_get_bssid(&g_wpa_supp, (uint8_t *)bssid);
-    if (!bss) {
-        wpa_printf(MSG_ERROR,
-                "SAE: BSS not available, update scan result to get BSS");
-       // TODO: should we trigger scan again.
-        return ESP_FAIL;
-    }
-    if (bss) {
-        const u8 *rsnxe;
-
-        rsnxe = wpa_bss_get_ie(bss, WLAN_EID_RSNX);
-        if (rsnxe && rsnxe[1] >= 1) {
-            rsnxe_capa = rsnxe[2];
-        }
-    }
-
-    if (use_pt && (rsnxe_capa & BIT(WLAN_RSNX_CAPAB_SAE_PK)) &&
-        sae_pk_mode != WPA3_SAE_PK_MODE_DISABLED &&
-        ((pw && sae_pk_valid_password((const char*)pw)))) {
-        use_pt = 1;
-        use_pk = true;
-    }
-
-    if (sae_pk_mode == WPA3_SAE_PK_MODE_ONLY && !use_pk) {
-        wpa_printf(MSG_DEBUG,
-           "SAE: Cannot use PK with the selected AP");
-    return ESP_FAIL;
-    }
-#endif /* CONFIG_SAE_PK */
-
     if (use_pt &&
             sae_prepare_commit_pt(&g_sae_data, g_sae_pt,
                     own_addr, bssid, NULL, NULL) < 0) {
@@ -294,7 +302,7 @@ 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);
+    wpa_set_pmk(g_sae_data.pmk, g_sae_data.pmk_len, g_sae_data.pmkid, true);
 
     return ESP_OK;
 }
diff --git a/components/wpa_supplicant/esp_supplicant/src/esp_wpas_glue.c b/components/wpa_supplicant/esp_supplicant/src/esp_wpas_glue.c
index d3040b168e..16f8916f15 100644
--- a/components/wpa_supplicant/esp_supplicant/src/esp_wpas_glue.c
+++ b/components/wpa_supplicant/esp_supplicant/src/esp_wpas_glue.c
@@ -98,7 +98,7 @@ void wpa_supplicant_transition_disable(struct wpa_sm *sm, u8 bitmap)
     wpa_printf(MSG_DEBUG, "TRANSITION_DISABLE %02x", bitmap);
 
     if  ((bitmap & TRANSITION_DISABLE_WPA3_PERSONAL) &&
-         wpa_key_mgmt_sae(sm->key_mgmt)) {
+          wpa_key_mgmt_sae(sm->key_mgmt)) {
         esp_wifi_sta_disable_wpa2_authmode_internal();
     }
 }
diff --git a/components/wpa_supplicant/src/common/wpa_common.c b/components/wpa_supplicant/src/common/wpa_common.c
index 2bffbbe758..d8f0e90c0e 100644
--- a/components/wpa_supplicant/src/common/wpa_common.c
+++ b/components/wpa_supplicant/src/common/wpa_common.c
@@ -855,7 +855,8 @@ int wpa_pmk_r1_to_ptk(const u8 *pmk_r1, const u8 *snonce, const u8 *anonce,
 */
 
 int wpa_use_akm_defined(int akmp){
-	return akmp == WPA_KEY_MGMT_OWE ||
+	return akmp == WPA_KEY_MGMT_OSEN ||
+		akmp == WPA_KEY_MGMT_OWE ||
 		wpa_key_mgmt_sae(akmp) ||
 		wpa_key_mgmt_suite_b(akmp);
 }
diff --git a/components/wpa_supplicant/src/rsn_supp/wpa.c b/components/wpa_supplicant/src/rsn_supp/wpa.c
index d23da66d29..06f564e084 100644
--- a/components/wpa_supplicant/src/rsn_supp/wpa.c
+++ b/components/wpa_supplicant/src/rsn_supp/wpa.c
@@ -424,18 +424,6 @@ static int wpa_supplicant_get_pmk(struct wpa_sm *sm,
         //eapol_sm_notify_cached(sm->eapol);
 #ifdef CONFIG_IEEE80211R
         sm->xxkey_len = 0;
-#ifdef CONFIG_WPA3_SAE
-        if ((sm->key_mgmt == WPA_KEY_MGMT_FT_SAE) &&
-             sm->pmk_len == PMK_LEN) {
-            /* Need to allow FT key derivation to proceed with
-             * PMK from SAE being used as the XXKey in cases where
-             * the PMKID in msg 1/4 matches the PMKSA entry that was
-             * just added based on SAE authentication for the
-             * initial mobility domain association. */
-            os_memcpy(sm->xxkey, sm->pmk, sm->pmk_len);
-            sm->xxkey_len = sm->pmk_len;
-        }
-#endif /* CONFIG_WPA3_SAE */
 #endif /* CONFIG_IEEE80211R */
     } else if (wpa_key_mgmt_wpa_ieee8021x(sm->key_mgmt)) {
         int res = 0, pmk_len;
@@ -2298,18 +2286,22 @@ void wpa_set_profile(u32 wpa_proto, u8 auth_mode)
          sm->key_mgmt = WPA_KEY_MGMT_FT_PSK;
     } else if (auth_mode == WPA3_AUTH_OWE) {
          sm->key_mgmt = WPA_KEY_MGMT_OWE;
+    } else if (auth_mode == WPA3_AUTH_PSK_EXT_KEY) {
+         sm->key_mgmt = WPA_KEY_MGMT_SAE_EXT_KEY; /* for WPA3 PSK */
     } else {
         sm->key_mgmt = WPA_KEY_MGMT_PSK;  /* fixed to PSK for now */
     }
 }
 
-void wpa_set_pmk(uint8_t *pmk, const u8 *pmkid, bool cache_pmksa)
+void wpa_set_pmk(uint8_t *pmk, size_t pmk_length, const u8 *pmkid, bool cache_pmksa)
 {
     struct wpa_sm *sm = &gWpaSm;
     int pmk_len;
 
     if (wpa_key_mgmt_sha384(sm->key_mgmt))
         pmk_len = PMK_LEN_SUITE_B_192;
+    else if (wpa_key_mgmt_sae(sm->key_mgmt))
+        pmk_len = pmk_length;
     else
         pmk_len = PMK_LEN;
 
@@ -2335,7 +2327,8 @@ int wpa_set_bss(char *macddr, char * bssid, u8 pairwise_cipher, u8 group_cipher,
     /* Ideally we should use network_ctx for this purpose however currently network profile block
      * is part of libraries,
      * TODO Correct this in future during NVS restructuring */
-    if ((sm->key_mgmt == WPA_KEY_MGMT_SAE) &&
+    if ((sm->key_mgmt == WPA_KEY_MGMT_SAE ||
+         sm->key_mgmt == WPA_KEY_MGMT_SAE_EXT_KEY) &&
         (os_memcmp(sm->bssid, bssid, ETH_ALEN) == 0) &&
         (os_memcmp(sm->ssid, ssid, ssid_len) != 0)) {
         use_pmk_cache = false;
@@ -2474,7 +2467,9 @@ wpa_set_passphrase(char * passphrase, u8 *ssid, size_t ssid_len)
      *  Here only handle passphrase string.  Need extra step to handle 32B, 64Hex raw
      *    PMK.
      */
-    if (sm->key_mgmt == WPA_KEY_MGMT_SAE || sm->key_mgmt == WPA_KEY_MGMT_OWE)
+    if (sm->key_mgmt == WPA_KEY_MGMT_SAE ||
+        sm->key_mgmt == WPA_KEY_MGMT_OWE ||
+        sm->key_mgmt == WPA_KEY_MGMT_SAE_EXT_KEY)
         return;
 
     /* This is really SLOW, so just re cacl while reset param */
diff --git a/components/wpa_supplicant/src/rsn_supp/wpa.h b/components/wpa_supplicant/src/rsn_supp/wpa.h
index 4f7a768fc6..257735d270 100644
--- a/components/wpa_supplicant/src/rsn_supp/wpa.h
+++ b/components/wpa_supplicant/src/rsn_supp/wpa.h
@@ -38,7 +38,7 @@ struct l2_ethhdr {
 
 void wpa_sm_set_state(enum wpa_states state);
 
-void wpa_set_pmk(uint8_t *pmk, const u8 *pmkid, bool cache_pmksa);
+void wpa_set_pmk(uint8_t *pmk, size_t pmk_length, const u8 *pmkid, bool cache_pmksa);
 
 int wpa_michael_mic_failure(u16 isunicast);
 
diff --git a/examples/wifi/scan/main/scan.c b/examples/wifi/scan/main/scan.c
index cbb6cfdd44..6a53718e1b 100644
--- a/examples/wifi/scan/main/scan.c
+++ b/examples/wifi/scan/main/scan.c
@@ -55,6 +55,9 @@ static void print_auth_mode(int authmode)
     case WIFI_AUTH_WPA3_ENT_192:
         ESP_LOGI(TAG, "Authmode \tWIFI_AUTH_WPA3_ENT_192");
         break;
+    case WIFI_AUTH_WPA3_EXT_PSK:
+        ESP_LOGI(TAG, "Authmode \tWIFI_AUTH_WPA3_EXT_PSK");
+        break;
     default:
         ESP_LOGI(TAG, "Authmode \tWIFI_AUTH_UNKNOWN");
         break;