From d2b2083843911b361f365415b4e794543b594b8a Mon Sep 17 00:00:00 2001
From: Jouni Malinen <jouni@codeaurora.org>
Date: Wed, 14 Aug 2019 17:49:23 +0300
Subject: [PATCH] SAE: Allow PMKID to be added into Association Request frame
 following SAE

IEEE Std 802.11-2016 does not require this behavior from a SAE STA, but
it is not disallowed either, so it is useful to have an option to
identify the derived PMKSA in the immediately following Association
Request frames. This is disabled by default (i.e., no change to previous
behavior) and can be enabled with a global wpa_supplicant configuration
parameter sae_pmkid_in_assoc=1.

Signed-off-by: Jouni Malinen <jouni@codeaurora.org>
---
 wpa_supplicant/config.c      |  1 +
 wpa_supplicant/config.h      |  5 +++++
 wpa_supplicant/config_file.c |  4 ++++
 wpa_supplicant/sme.c         | 43 +++++++++++++++++++++++++++++-------
 4 files changed, 45 insertions(+), 8 deletions(-)

diff --git a/wpa_supplicant/config.c b/wpa_supplicant/config.c
index 7a62f96d6..fc1ed4f90 100644
--- a/wpa_supplicant/config.c
+++ b/wpa_supplicant/config.c
@@ -4828,6 +4828,7 @@ static const struct global_parse_data global_fields[] = {
 	{ INT(okc), 0 },
 	{ INT(pmf), 0 },
 	{ FUNC(sae_groups), 0 },
+	{ INT_RANGE(sae_pmkid_in_assoc, 0, 1), 0 },
 	{ INT(dtim_period), 0 },
 	{ INT(beacon_int), 0 },
 	{ FUNC(ap_vendor_elements), 0 },
diff --git a/wpa_supplicant/config.h b/wpa_supplicant/config.h
index 6a297ecfe..5b5c2fdba 100644
--- a/wpa_supplicant/config.h
+++ b/wpa_supplicant/config.h
@@ -1164,6 +1164,11 @@ struct wpa_config {
 	 */
 	int *sae_groups;
 
+	/**
+	 * sae_pmkid_in_assoc - Whether to include PMKID in SAE Assoc Req
+	 */
+	int sae_pmkid_in_assoc;
+
 	/**
 	 * dtim_period - Default DTIM period in Beacon intervals
 	 *
diff --git a/wpa_supplicant/config_file.c b/wpa_supplicant/config_file.c
index 77c326df5..91d5caa3f 100644
--- a/wpa_supplicant/config_file.c
+++ b/wpa_supplicant/config_file.c
@@ -1390,6 +1390,10 @@ static void wpa_config_write_global(FILE *f, struct wpa_config *config)
 		fprintf(f, "\n");
 	}
 
+	if (config->sae_pmkid_in_assoc)
+		fprintf(f, "sae_pmkid_in_assoc=%d\n",
+			config->sae_pmkid_in_assoc);
+
 	if (config->ap_vendor_elements) {
 		int i, len = wpabuf_len(config->ap_vendor_elements);
 		const u8 *p = wpabuf_head_u8(config->ap_vendor_elements);
diff --git a/wpa_supplicant/sme.c b/wpa_supplicant/sme.c
index dd5020179..7944368cf 100644
--- a/wpa_supplicant/sme.c
+++ b/wpa_supplicant/sme.c
@@ -1197,6 +1197,37 @@ static int sme_sae_auth(struct wpa_supplicant *wpa_s, u16 auth_transaction,
 }
 
 
+static int sme_sae_set_pmk(struct wpa_supplicant *wpa_s)
+{
+	wpa_printf(MSG_DEBUG,
+		   "SME: SAE completed - setting PMK for 4-way handshake");
+	wpa_sm_set_pmk(wpa_s->wpa, wpa_s->sme.sae.pmk, PMK_LEN,
+		       wpa_s->sme.sae.pmkid, wpa_s->pending_bssid);
+	if (wpa_s->conf->sae_pmkid_in_assoc) {
+		/* Update the own RSNE contents now that we have set the PMK
+		 * and added a PMKSA cache entry based on the successfully
+		 * completed SAE exchange. In practice, this will add the PMKID
+		 * into RSNE. */
+		if (wpa_s->sme.assoc_req_ie_len + 2 + PMKID_LEN >
+		    sizeof(wpa_s->sme.assoc_req_ie)) {
+			wpa_msg(wpa_s, MSG_WARNING,
+				"RSN: Not enough room for inserting own PMKID into RSNE");
+			return -1;
+		}
+		if (wpa_insert_pmkid(wpa_s->sme.assoc_req_ie,
+				     &wpa_s->sme.assoc_req_ie_len,
+				     wpa_s->sme.sae.pmkid) < 0)
+			return -1;
+		wpa_hexdump(MSG_DEBUG,
+			    "SME: Updated Association Request IEs",
+			    wpa_s->sme.assoc_req_ie,
+			    wpa_s->sme.assoc_req_ie_len);
+	}
+
+	return 0;
+}
+
+
 void sme_external_auth_mgmt_rx(struct wpa_supplicant *wpa_s,
 			       const u8 *auth_frame, size_t len)
 {
@@ -1230,10 +1261,8 @@ void sme_external_auth_mgmt_rx(struct wpa_supplicant *wpa_s,
 		if (res != 1)
 			return;
 
-		wpa_printf(MSG_DEBUG,
-			   "SME: SAE completed - setting PMK for 4-way handshake");
-		wpa_sm_set_pmk(wpa_s->wpa, wpa_s->sme.sae.pmk, PMK_LEN,
-			       wpa_s->sme.sae.pmkid, wpa_s->pending_bssid);
+		if (sme_sae_set_pmk(wpa_s) < 0)
+			return;
 	}
 }
 
@@ -1286,10 +1315,8 @@ void sme_event_auth(struct wpa_supplicant *wpa_s, union wpa_event_data *data)
 		if (res != 1)
 			return;
 
-		wpa_printf(MSG_DEBUG, "SME: SAE completed - setting PMK for "
-			   "4-way handshake");
-		wpa_sm_set_pmk(wpa_s->wpa, wpa_s->sme.sae.pmk, PMK_LEN,
-			       wpa_s->sme.sae.pmkid, wpa_s->pending_bssid);
+		if (sme_sae_set_pmk(wpa_s) < 0)
+			return;
 	}
 #endif /* CONFIG_SAE */