From f70db167abf42caa32a7c25f192bcbc2415fad63 Mon Sep 17 00:00:00 2001
From: Jouni Malinen <quic_jouni@quicinc.com>
Date: Sun, 24 Jul 2022 22:16:40 +0300
Subject: [PATCH] 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>
---
 src/common/sae.c | 23 +++++++++++++++++------
 src/common/sae.h |  3 ++-
 2 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/src/common/sae.c b/src/common/sae.c
index 637ca29b5..33c7e5fe1 100644
--- a/src/common/sae.c
+++ b/src/common/sae.c
@@ -9,6 +9,8 @@
 #include "includes.h"
 
 #include "common.h"
+#include "common/defs.h"
+#include "common/wpa_common.h"
 #include "utils/const_time.h"
 #include "crypto/crypto.h"
 #include "crypto/sha256.h"
@@ -1520,10 +1522,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];
 
@@ -1545,6 +1548,14 @@ 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;
+	wpa_printf(MSG_DEBUG, "SAE: Derive keys - H2E=%d AKMP=0x%x = %08x (%s)",
+		   sae->h2e, sae->akmp,
+		   wpa_akm_to_suite(sae->akmp),
+		   wpa_key_mgmt_txt(sae->akmp, WPA_PROTO_RSN));
 	if (sae->h2e && (sae->tmp->own_rejected_groups ||
 			 sae->tmp->peer_rejected_groups)) {
 		struct wpabuf *own, *peer;
@@ -1602,26 +1613,26 @@ 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);
-	sae->pmk_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) {
diff --git a/src/common/sae.h b/src/common/sae.h
index 459117a71..6a6b0c824 100644
--- a/src/common/sae.h
+++ b/src/common/sae.h
@@ -11,6 +11,7 @@
 
 #define SAE_KCK_LEN 32
 #define SAE_PMK_LEN 32
+#define SAE_PMK_LEN_MAX 64
 #define SAE_PMKID_LEN 16
 #define SAE_MAX_PRIME_LEN 512
 #define SAE_MAX_ECC_PRIME_LEN 66
@@ -104,7 +105,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;