From bbda0a9baad699abfb22e054b8dca0336cd4ec0f Mon Sep 17 00:00:00 2001
From: Shreyas Sheth <shreyas.sheth@espressif.com>
Date: Wed, 29 Jan 2025 13:49:56 +0530
Subject: [PATCH] fix(wifi): Fixed memory leak occurring in SAE PK connection

---
 components/wpa_supplicant/src/common/sae_pk.c | 39 ++++++++++++++-----
 1 file changed, 30 insertions(+), 9 deletions(-)

diff --git a/components/wpa_supplicant/src/common/sae_pk.c b/components/wpa_supplicant/src/common/sae_pk.c
index 472ef03d2c..6361de2aa0 100644
--- a/components/wpa_supplicant/src/common/sae_pk.c
+++ b/components/wpa_supplicant/src/common/sae_pk.c
@@ -620,6 +620,7 @@ int sae_check_confirm_pk(struct sae_data *sae, const u8 *ies, size_t ies_len)
 	int group;
 	struct wpa_supplicant *wpa_s = &g_wpa_supp;
 	struct sae_pk_elems elems;
+	int ret = 0;
 
 	if (!tmp) {
 		return -1;
@@ -650,7 +651,8 @@ int sae_check_confirm_pk(struct sae_data *sae, const u8 *ies, size_t ies_len)
 	if (!elems.fils_pk || !elems.fils_key_confirm || !elems.sae_pk) {
 		wpa_printf(MSG_INFO,
 				   "SAE-PK: Not all mandatory IEs included in confirm");
-		return -1;
+		ret = -1;
+		goto done;
 	}
 
 	/* TODO: Fragment reassembly */
@@ -658,7 +660,8 @@ int sae_check_confirm_pk(struct sae_data *sae, const u8 *ies, size_t ies_len)
 	if (elems.sae_pk_len < SAE_PK_M_LEN + AES_BLOCK_SIZE) {
 		wpa_printf(MSG_INFO,
 				   "SAE-PK: No room for EncryptedModifier in SAE-PK element");
-		return -1;
+		ret = -1;
+		goto done;
 	}
 
 	wpa_hexdump(MSG_DEBUG, "SAE-PK: EncryptedModifier",
@@ -669,14 +672,16 @@ int sae_check_confirm_pk(struct sae_data *sae, const u8 *ies, size_t ies_len)
 						0, NULL, NULL, m) < 0) {
 		wpa_printf(MSG_INFO,
 				   "SAE-PK: Failed to decrypt EncryptedModifier");
-		return -1;
+		ret = -1;
+		goto done;
 	}
 	wpa_hexdump_key(MSG_DEBUG, "SAE-PK: Modifier M", m, SAE_PK_M_LEN);
 
 	if (elems.fils_pk[0] != 2) {
 		wpa_printf(MSG_INFO, "SAE-PK: Unsupported public key type %u",
 				   elems.fils_pk[0]);
-		return -1;
+		ret = -1;
+		goto done;
 	}
 	k_ap_len = elems.fils_pk_len - 1;
 	k_ap = elems.fils_pk + 1;
@@ -686,13 +691,15 @@ int sae_check_confirm_pk(struct sae_data *sae, const u8 *ies, size_t ies_len)
 	key = crypto_ec_key_parse_pub(k_ap, k_ap_len);
 	if (!key) {
 		wpa_printf(MSG_INFO, "SAE-PK: Failed to parse K_AP");
-		return -1;
+		ret = -1;
+		goto done;
 	}
 	group = crypto_ec_key_group(key);
 	if (!sae_pk_valid_fingerprint(sae, m, SAE_PK_M_LEN, k_ap, k_ap_len,
 								  group)) {
 		crypto_ec_key_deinit(key);
-		return -1;
+		ret = -1;
+		goto done;
 	}
 
 	wpa_hexdump(MSG_DEBUG, "SAE-PK: Received KeyAuth",
@@ -702,7 +709,8 @@ int sae_check_confirm_pk(struct sae_data *sae, const u8 *ies, size_t ies_len)
 	if (sae_pk_hash_sig_data(sae, hash_len, false, m, SAE_PK_M_LEN,
 							 k_ap, k_ap_len, hash) < 0) {
 		crypto_ec_key_deinit(key);
-		return -1;
+		ret = -1;
+		goto done;
 	}
 
 	res = crypto_ec_key_verify_signature(key, hash, hash_len,
@@ -713,12 +721,25 @@ int sae_check_confirm_pk(struct sae_data *sae, const u8 *ies, size_t ies_len)
 	if (res != 1) {
 		wpa_printf(MSG_INFO,
 				   "SAE-PK: Invalid or incorrect signature in KeyAuth");
-		return -1;
+		ret = -1;
+		goto done;
 	}
 
 	wpa_printf(MSG_DEBUG, "SAE-PK: Valid KeyAuth signature received");
 
 	/* TODO: Store validated public key into network profile */
-	return 0;
+done:
+	if (wpa_s->sae_pk_elems.fils_pk) {
+		os_free(wpa_s->sae_pk_elems.fils_pk);
+	}
+	if (wpa_s->sae_pk_elems.sae_pk) {
+		os_free(wpa_s->sae_pk_elems.sae_pk);
+	}
+	if (wpa_s->sae_pk_elems.fils_key_confirm) {
+		os_free(wpa_s->sae_pk_elems.fils_key_confirm);
+	}
+	os_memset(&wpa_s->sae_pk_elems, 0, sizeof(wpa_s->sae_pk_elems));
+
+	return ret;
 }
 #endif /* CONFIG_SAE_PK */