From b17b7a8e53f28aed746518321ecd302d0c260971 Mon Sep 17 00:00:00 2001
From: Alexander Wetzel <alexander@wetzel-home.de>
Date: Fri, 20 Mar 2020 20:04:32 +0100
Subject: [PATCH] STA: Support Extended Key ID

Support Extended Key ID in wpa_supplicant according to
IEEE Std 802.11-2016 for infrastructure (AP) associations.

Extended Key ID allows to rekey pairwise keys without the otherwise
unavoidable MPDU losses on a busy link. The standard is fully backward
compatible, allowing STAs to also connect to APs not supporting it.

Signed-off-by: Alexander Wetzel <alexander@wetzel-home.de>
---
 src/rsn_supp/wpa.c                      | 117 ++++++++++++++++++++++--
 src/rsn_supp/wpa.h                      |  14 +++
 src/rsn_supp/wpa_ft.c                   |   2 +
 src/rsn_supp/wpa_i.h                    |   4 +
 src/rsn_supp/wpa_ie.c                   |   2 +
 wpa_supplicant/ap.c                     |   1 +
 wpa_supplicant/config.c                 |   2 +
 wpa_supplicant/config.h                 |  12 +++
 wpa_supplicant/config_file.c            |   3 +
 wpa_supplicant/config_winreg.c          |   9 ++
 wpa_supplicant/ctrl_iface.c             |   3 +
 wpa_supplicant/dbus/dbus_new_handlers.c |  13 ++-
 wpa_supplicant/driver_i.h               |   9 +-
 wpa_supplicant/wpa_cli.c                |   3 +-
 wpa_supplicant/wpa_supplicant.c         |  35 ++++++-
 wpa_supplicant/wpa_supplicant.conf      |   5 +
 wpa_supplicant/wpas_glue.c              |   6 +-
 17 files changed, 222 insertions(+), 18 deletions(-)

diff --git a/src/rsn_supp/wpa.c b/src/rsn_supp/wpa.c
index ba0e5c343..14fe0846d 100644
--- a/src/rsn_supp/wpa.c
+++ b/src/rsn_supp/wpa.c
@@ -183,7 +183,7 @@ void wpa_sm_key_request(struct wpa_sm *sm, int error, int pairwise)
 	int key_info, ver;
 	u8 bssid[ETH_ALEN], *rbuf, *key_mic, *mic;
 
-	if (pairwise && sm->wpa_deny_ptk0_rekey &&
+	if (pairwise && sm->wpa_deny_ptk0_rekey && !sm->use_ext_key_id &&
 	    wpa_sm_get_state(sm) == WPA_COMPLETED) {
 		wpa_msg(sm->ctx->msg_ctx, MSG_INFO,
 			"WPA: PTK0 rekey not allowed, reconnecting");
@@ -608,6 +608,51 @@ static int wpa_derive_ptk(struct wpa_sm *sm, const unsigned char *src_addr,
 }
 
 
+static int wpa_handle_ext_key_id(struct wpa_sm *sm,
+				 struct wpa_eapol_ie_parse *kde)
+{
+	if (sm->ext_key_id) {
+		u16 key_id;
+
+		if (!kde->key_id) {
+			wpa_msg(sm->ctx->msg_ctx,
+				sm->use_ext_key_id ? MSG_INFO : MSG_DEBUG,
+				"RSN: No Key ID in Extended Key ID handshake");
+			sm->keyidx_active = 0;
+			return sm->use_ext_key_id ? -1 : 0;
+		}
+
+		key_id = kde->key_id[0] & 0x03;
+		if (key_id > 1) {
+			wpa_msg(sm->ctx->msg_ctx, MSG_INFO,
+				"RSN: Invalid Extended Key ID: %d", key_id);
+			return -1;
+		}
+		wpa_dbg(sm->ctx->msg_ctx, MSG_DEBUG,
+			"RSN: Using Extended Key ID %d", key_id);
+		sm->keyidx_active = key_id;
+		sm->use_ext_key_id = 1;
+	} else {
+		if (kde->key_id && (kde->key_id[0] & 0x03)) {
+			wpa_msg(sm->ctx->msg_ctx, MSG_INFO,
+				"RSN: Non-zero Extended Key ID Key ID in PTK0 handshake");
+			return -1;
+		}
+
+		if (kde->key_id) {
+			/* This is not supposed to be included here, but ignore
+			 * the case of matching Key ID 0 just in case. */
+			wpa_msg(sm->ctx->msg_ctx, MSG_DEBUG,
+				"RSN: Extended Key ID Key ID 0 in PTK0 handshake");
+		}
+		sm->keyidx_active = 0;
+		sm->use_ext_key_id = 0;
+	}
+
+	return 0;
+}
+
+
 static void wpa_supplicant_process_1_of_4(struct wpa_sm *sm,
 					  const unsigned char *src_addr,
 					  const struct wpa_eapol_key *key,
@@ -626,7 +671,8 @@ static void wpa_supplicant_process_1_of_4(struct wpa_sm *sm,
 		return;
 	}
 
-	if (sm->wpa_deny_ptk0_rekey && wpa_sm_get_state(sm) == WPA_COMPLETED) {
+	if (sm->wpa_deny_ptk0_rekey && !sm->use_ext_key_id &&
+	    wpa_sm_get_state(sm) == WPA_COMPLETED) {
 		wpa_msg(sm->ctx->msg_ctx, MSG_INFO,
 			"WPA: PTK0 rekey not allowed, reconnecting");
 		wpa_sm_reconnect(sm);
@@ -859,13 +905,14 @@ static int wpa_supplicant_install_ptk(struct wpa_sm *sm,
 		wpa_hexdump(MSG_DEBUG, "WPA: RSC", key_rsc, rsclen);
 	}
 
-	if (wpa_sm_set_key(sm, alg, sm->bssid, 0, 1, key_rsc, rsclen,
-			   sm->ptk.tk, keylen,
+	if (wpa_sm_set_key(sm, alg, sm->bssid, sm->keyidx_active, 1, key_rsc,
+			   rsclen, sm->ptk.tk, keylen,
 			   KEY_FLAG_PAIRWISE | key_flag) < 0) {
 		wpa_msg(sm->ctx->msg_ctx, MSG_WARNING,
-			"WPA: Failed to set PTK to the "
-			"driver (alg=%d keylen=%d bssid=" MACSTR ")",
-			alg, keylen, MAC2STR(sm->bssid));
+			"WPA: Failed to set PTK to the driver (alg=%d keylen=%d bssid="
+			MACSTR " idx=%d key_flag=0x%x)",
+			alg, keylen, MAC2STR(sm->bssid),
+			sm->keyidx_active, key_flag);
 		return -1;
 	}
 
@@ -879,7 +926,23 @@ static int wpa_supplicant_install_ptk(struct wpa_sm *sm,
 		eloop_register_timeout(sm->wpa_ptk_rekey, 0, wpa_sm_rekey_ptk,
 				       sm, NULL);
 	}
+	return 0;
+}
+
+
+static int wpa_supplicant_activate_ptk(struct wpa_sm *sm)
+{
+	wpa_dbg(sm->ctx->msg_ctx, MSG_DEBUG,
+		"WPA: Activate PTK (idx=%d bssid=" MACSTR ")",
+		sm->keyidx_active, MAC2STR(sm->bssid));
 
+	if (wpa_sm_set_key(sm, 0, sm->bssid, sm->keyidx_active, 0, NULL, 0,
+			   NULL, 0, KEY_FLAG_PAIRWISE_RX_TX_MODIFY) < 0) {
+		wpa_msg(sm->ctx->msg_ctx, MSG_WARNING,
+			"WPA: Failed to activate PTK for TX (idx=%d bssid="
+			MACSTR ")", sm->keyidx_active, MAC2STR(sm->bssid));
+		return -1;
+	}
 	return 0;
 }
 
@@ -1582,6 +1645,9 @@ static void wpa_supplicant_process_3_of_4(struct wpa_sm *sm,
 	if (wpa_supplicant_validate_ie(sm, sm->bssid, &ie) < 0)
 		goto failed;
 
+	if (wpa_handle_ext_key_id(sm, &ie))
+		goto failed;
+
 	if (os_memcmp(sm->anonce, key->key_nonce, WPA_NONCE_LEN) != 0) {
 		wpa_msg(sm->ctx->msg_ctx, MSG_WARNING,
 			"WPA: ANonce from message 1 of 4-Way Handshake "
@@ -1627,6 +1693,10 @@ static void wpa_supplicant_process_3_of_4(struct wpa_sm *sm,
 	}
 #endif /* CONFIG_OCV */
 
+	if (sm->use_ext_key_id &&
+	    wpa_supplicant_install_ptk(sm, key, KEY_FLAG_RX))
+		goto failed;
+
 	if (wpa_supplicant_send_4_of_4(sm, sm->bssid, key, ver, key_info,
 				       &sm->ptk) < 0) {
 		goto failed;
@@ -1638,7 +1708,14 @@ static void wpa_supplicant_process_3_of_4(struct wpa_sm *sm,
 	sm->renew_snonce = 1;
 
 	if (key_info & WPA_KEY_INFO_INSTALL) {
-		if (wpa_supplicant_install_ptk(sm, key, KEY_FLAG_RX_TX))
+		int res;
+
+		if (sm->use_ext_key_id)
+			res = wpa_supplicant_activate_ptk(sm);
+		else
+			res = wpa_supplicant_install_ptk(sm, key,
+							 KEY_FLAG_RX_TX);
+		if (res)
 			goto failed;
 	}
 
@@ -2880,6 +2957,8 @@ void wpa_sm_notify_assoc(struct wpa_sm *sm, const u8 *bssid)
 #ifdef CONFIG_P2P
 	os_memset(sm->p2p_ip_addr, 0, sizeof(sm->p2p_ip_addr));
 #endif /* CONFIG_P2P */
+
+	sm->keyidx_active = 0;
 }
 
 
@@ -2911,6 +2990,7 @@ void wpa_sm_notify_disassoc(struct wpa_sm *sm)
 
 	/* Keys are not needed in the WPA state machine anymore */
 	wpa_sm_drop_sa(sm);
+	sm->keyidx_active = 0;
 
 	sm->msg_3_of_4_ok = 0;
 	os_memset(sm->bssid, 0, ETH_ALEN);
@@ -3164,6 +3244,12 @@ int wpa_sm_set_param(struct wpa_sm *sm, enum wpa_sm_conf_params param,
 	case WPA_PARAM_DENY_PTK0_REKEY:
 		sm->wpa_deny_ptk0_rekey = value;
 		break;
+	case WPA_PARAM_EXT_KEY_ID:
+		sm->ext_key_id = value;
+		break;
+	case WPA_PARAM_USE_EXT_KEY_ID:
+		sm->use_ext_key_id = value;
+		break;
 	default:
 		break;
 	}
@@ -3238,6 +3324,18 @@ int wpa_sm_pmf_enabled(struct wpa_sm *sm)
 }
 
 
+int wpa_sm_ext_key_id(struct wpa_sm *sm)
+{
+	return sm ? sm->ext_key_id : 0;
+}
+
+
+int wpa_sm_ext_key_id_active(struct wpa_sm *sm)
+{
+	return sm ? sm->use_ext_key_id : 0;
+}
+
+
 int wpa_sm_ocv_enabled(struct wpa_sm *sm)
 {
 	struct wpa_ie_data rsn;
@@ -4253,6 +4351,8 @@ static int fils_ft_build_assoc_req_rsne(struct wpa_sm *sm, struct wpabuf *buf)
 		capab |= WPA_CAPABILITY_MFPR;
 	if (sm->ocv)
 		capab |= WPA_CAPABILITY_OCVC;
+	if (sm->ext_key_id)
+		capab |= WPA_CAPABILITY_EXT_KEY_ID_FOR_UNICAST;
 	wpabuf_put_le16(buf, capab);
 
 	/* PMKID Count */
@@ -4680,6 +4780,7 @@ int fils_process_assoc_resp(struct wpa_sm *sm, const u8 *resp, size_t len)
 			   keylen, (long unsigned int) sm->ptk.tk_len);
 		goto fail;
 	}
+
 	rsclen = wpa_cipher_rsc_len(sm->pairwise_cipher);
 	wpa_hexdump_key(MSG_DEBUG, "FILS: Set TK to driver",
 			sm->ptk.tk, keylen);
diff --git a/src/rsn_supp/wpa.h b/src/rsn_supp/wpa.h
index 0bd14495a..02b5df883 100644
--- a/src/rsn_supp/wpa.h
+++ b/src/rsn_supp/wpa.h
@@ -102,6 +102,8 @@ enum wpa_sm_conf_params {
 	WPA_PARAM_OCV,
 	WPA_PARAM_SAE_PWE,
 	WPA_PARAM_DENY_PTK0_REKEY,
+	WPA_PARAM_EXT_KEY_ID,
+	WPA_PARAM_USE_EXT_KEY_ID,
 };
 
 struct rsn_supp_config {
@@ -154,6 +156,8 @@ int wpa_sm_set_param(struct wpa_sm *sm, enum wpa_sm_conf_params param,
 int wpa_sm_get_status(struct wpa_sm *sm, char *buf, size_t buflen,
 		      int verbose);
 int wpa_sm_pmf_enabled(struct wpa_sm *sm);
+int wpa_sm_ext_key_id(struct wpa_sm *sm);
+int wpa_sm_ext_key_id_active(struct wpa_sm *sm);
 int wpa_sm_ocv_enabled(struct wpa_sm *sm);
 
 void wpa_sm_key_request(struct wpa_sm *sm, int error, int pairwise);
@@ -300,6 +304,16 @@ static inline int wpa_sm_pmf_enabled(struct wpa_sm *sm)
 	return 0;
 }
 
+static inline int wpa_sm_ext_key_id(struct wpa_sm *sm)
+{
+	return 0;
+}
+
+static inline int wpa_sm_ext_key_id_active(struct wpa_sm *sm)
+{
+	return 0;
+}
+
 static inline int wpa_sm_ocv_enabled(struct wpa_sm *sm)
 {
 	return 0;
diff --git a/src/rsn_supp/wpa_ft.c b/src/rsn_supp/wpa_ft.c
index 6d92de5a6..bec5eb0b2 100644
--- a/src/rsn_supp/wpa_ft.c
+++ b/src/rsn_supp/wpa_ft.c
@@ -265,6 +265,8 @@ static u8 * wpa_ft_gen_req_ies(struct wpa_sm *sm, size_t *len,
 		capab |= WPA_CAPABILITY_MFPR;
 	if (sm->ocv)
 		capab |= WPA_CAPABILITY_OCVC;
+	if (sm->ext_key_id)
+		capab |= WPA_CAPABILITY_EXT_KEY_ID_FOR_UNICAST;
 	WPA_PUT_LE16(pos, capab);
 	pos += 2;
 
diff --git a/src/rsn_supp/wpa_i.h b/src/rsn_supp/wpa_i.h
index 7af678dcd..5fd70a498 100644
--- a/src/rsn_supp/wpa_i.h
+++ b/src/rsn_supp/wpa_i.h
@@ -68,6 +68,10 @@ struct wpa_sm {
 	int wpa_rsc_relaxation;
 	int owe_ptk_workaround;
 	int beacon_prot;
+	int ext_key_id; /* whether Extended Key ID is enabled */
+	int use_ext_key_id; /* whether Extended Key ID has been detected
+			     * to be used */
+	int keyidx_active; /* Key ID for the active TK */
 
 	u8 own_addr[ETH_ALEN];
 	const char *ifname;
diff --git a/src/rsn_supp/wpa_ie.c b/src/rsn_supp/wpa_ie.c
index 03c0d7e85..e8a040a2b 100644
--- a/src/rsn_supp/wpa_ie.c
+++ b/src/rsn_supp/wpa_ie.c
@@ -221,6 +221,8 @@ static int wpa_gen_wpa_ie_rsn(u8 *rsn_ie, size_t rsn_ie_len,
 		capab |= WPA_CAPABILITY_MFPR;
 	if (sm->ocv)
 		capab |= WPA_CAPABILITY_OCVC;
+	if (sm->ext_key_id)
+		capab |= WPA_CAPABILITY_EXT_KEY_ID_FOR_UNICAST;
 	WPA_PUT_LE16(pos, capab);
 	pos += 2;
 
diff --git a/wpa_supplicant/ap.c b/wpa_supplicant/ap.c
index 87573ef10..ee7c755b5 100644
--- a/wpa_supplicant/ap.c
+++ b/wpa_supplicant/ap.c
@@ -344,6 +344,7 @@ static int wpa_supplicant_conf_ap(struct wpa_supplicant *wpa_s,
 #endif /* CONFIG_IEEE80211AX */
 
 	bss->isolate = !wpa_s->conf->p2p_intra_bss;
+	bss->extended_key_id = wpa_s->conf->extended_key_id;
 	bss->force_per_enrollee_psk = wpa_s->global->p2p_per_sta_psk;
 	bss->wpa_deny_ptk0_rekey = ssid->wpa_deny_ptk0_rekey;
 
diff --git a/wpa_supplicant/config.c b/wpa_supplicant/config.c
index 4f359ede0..30801abd9 100644
--- a/wpa_supplicant/config.c
+++ b/wpa_supplicant/config.c
@@ -4293,6 +4293,7 @@ struct wpa_config * wpa_config_alloc_empty(const char *ctrl_interface,
 	config->key_mgmt_offload = DEFAULT_KEY_MGMT_OFFLOAD;
 	config->cert_in_cb = DEFAULT_CERT_IN_CB;
 	config->wpa_rsc_relaxation = DEFAULT_WPA_RSC_RELAXATION;
+	config->extended_key_id = DEFAULT_EXTENDED_KEY_ID;
 
 #ifdef CONFIG_MBO
 	config->mbo_cell_capa = DEFAULT_MBO_CELL_CAPA;
@@ -5057,6 +5058,7 @@ static const struct global_parse_data global_fields[] = {
 	{ INT_RANGE(coloc_intf_reporting, 0, 1), 0 },
 #ifdef CONFIG_WNM
 	{ INT_RANGE(disable_btm, 0, 1), CFG_CHANGED_DISABLE_BTM },
+	{ INT_RANGE(extended_key_id, 0, 1), 0 },
 #endif /* CONFIG_WNM */
 };
 
diff --git a/wpa_supplicant/config.h b/wpa_supplicant/config.h
index 486432e43..2d4cb1b8f 100644
--- a/wpa_supplicant/config.h
+++ b/wpa_supplicant/config.h
@@ -44,6 +44,7 @@
 #define DEFAULT_MBO_CELL_CAPA MBO_CELL_CAPA_NOT_SUPPORTED
 #define DEFAULT_DISASSOC_IMMINENT_RSSI_THRESHOLD -75
 #define DEFAULT_OCE_SUPPORT OCE_STA
+#define DEFAULT_EXTENDED_KEY_ID 0
 
 #include "config_ssid.h"
 #include "wps/wps.h"
@@ -1570,6 +1571,17 @@ struct wpa_config {
 	 * By default BSS transition management is enabled
 	 */
 	int disable_btm;
+
+	/**
+	 * extended_key_id - Extended Key ID support
+	 *
+	 * IEEE Std 802.11-2016 optionally allows to use Key ID 0 and 1 for PTK
+	 * keys with Extended Key ID.
+	 *
+	 * 0 = don't use Extended Key ID
+	 * 1 = use Extended Key ID when possible
+	 */
+	int extended_key_id;
 };
 
 
diff --git a/wpa_supplicant/config_file.c b/wpa_supplicant/config_file.c
index b8e56f5b2..e77cbca4f 100644
--- a/wpa_supplicant/config_file.c
+++ b/wpa_supplicant/config_file.c
@@ -1596,6 +1596,9 @@ static void wpa_config_write_global(FILE *f, struct wpa_config *config)
 			config->p2p_interface_random_mac_addr);
 	if (config->disable_btm)
 		fprintf(f, "disable_btm=1\n");
+	if (config->extended_key_id != DEFAULT_EXTENDED_KEY_ID)
+		fprintf(f, "extended_key_id=%d\n",
+			config->extended_key_id);
 }
 
 #endif /* CONFIG_NO_CONFIG_WRITE */
diff --git a/wpa_supplicant/config_winreg.c b/wpa_supplicant/config_winreg.c
index 598bc7790..1b7f96ed2 100644
--- a/wpa_supplicant/config_winreg.c
+++ b/wpa_supplicant/config_winreg.c
@@ -277,6 +277,15 @@ static int wpa_config_read_global(struct wpa_config *config, HKEY hk)
 	wpa_config_read_reg_dword(hk, TEXT("okc"), &config->okc);
 	wpa_config_read_reg_dword(hk, TEXT("pmf"), &val);
 	config->pmf = val;
+	if (wpa_config_read_reg_dword(hk, TEXT("extended_key_id"),
+				      &val) == 0) {
+		if (val < 0 || val > 1) {
+			wpa_printf(MSG_ERROR,
+				   "Invalid Extended Key ID setting (%d)", val);
+			errors++;
+		}
+		config->extended_key_id = val;
+	}
 
 	return errors ? -1 : 0;
 }
diff --git a/wpa_supplicant/ctrl_iface.c b/wpa_supplicant/ctrl_iface.c
index 3b4c50020..0c79d5257 100644
--- a/wpa_supplicant/ctrl_iface.c
+++ b/wpa_supplicant/ctrl_iface.c
@@ -5374,6 +5374,9 @@ static void wpa_supplicant_ctrl_iface_drop_sa(struct wpa_supplicant *wpa_s)
 
 	wpa_drv_set_key(wpa_s, WPA_ALG_NONE, wpa_s->bssid, 0, 0, NULL, 0, NULL,
 			0, KEY_FLAG_PAIRWISE);
+	if (wpa_sm_ext_key_id(wpa_s->wpa))
+		wpa_drv_set_key(wpa_s, WPA_ALG_NONE, wpa_s->bssid, 1, 0,
+				NULL, 0, NULL, 0, KEY_FLAG_PAIRWISE);
 	/* MLME-SETPROTECTION.request(None) */
 	wpa_drv_mlme_setprotection(wpa_s, wpa_s->bssid,
 				   MLME_SETPROTECTION_PROTECT_TYPE_NONE,
diff --git a/wpa_supplicant/dbus/dbus_new_handlers.c b/wpa_supplicant/dbus/dbus_new_handlers.c
index 01fc67b6f..4e17e31a7 100644
--- a/wpa_supplicant/dbus/dbus_new_handlers.c
+++ b/wpa_supplicant/dbus/dbus_new_handlers.c
@@ -991,20 +991,25 @@ dbus_bool_t wpas_dbus_getter_global_capabilities(
 	const struct wpa_dbus_property_desc *property_desc,
 	DBusMessageIter *iter, DBusError *error, void *user_data)
 {
-	const char *capabilities[11];
+	const char *capabilities[12];
 	size_t num_items = 0;
-#ifdef CONFIG_FILS
 	struct wpa_global *global = user_data;
 	struct wpa_supplicant *wpa_s;
+#ifdef CONFIG_FILS
 	int fils_supported = 0, fils_sk_pfs_supported = 0;
+#endif /* CONFIG_FILS */
+	int ext_key_id_supported = 0;
 
 	for (wpa_s = global->ifaces; wpa_s; wpa_s = wpa_s->next) {
+#ifdef CONFIG_FILS
 		if (wpa_is_fils_supported(wpa_s))
 			fils_supported = 1;
 		if (wpa_is_fils_sk_pfs_supported(wpa_s))
 			fils_sk_pfs_supported = 1;
-	}
 #endif /* CONFIG_FILS */
+		if (wpa_s->drv_flags & WPA_DRIVER_FLAGS_EXTENDED_KEY_ID)
+			ext_key_id_supported = 1;
+	}
 
 #ifdef CONFIG_AP
 	capabilities[num_items++] = "ap";
@@ -1037,6 +1042,8 @@ dbus_bool_t wpas_dbus_getter_global_capabilities(
 #ifdef CONFIG_OWE
 	capabilities[num_items++] = "owe";
 #endif /* CONFIG_OWE */
+	if (ext_key_id_supported)
+		capabilities[num_items++] = "extended_key_id";
 
 	return wpas_dbus_simple_array_property_getter(iter,
 						      DBUS_TYPE_STRING,
diff --git a/wpa_supplicant/driver_i.h b/wpa_supplicant/driver_i.h
index d3fb58707..6a03d8e91 100644
--- a/wpa_supplicant/driver_i.h
+++ b/wpa_supplicant/driver_i.h
@@ -165,7 +165,14 @@ static inline int wpa_drv_set_key(struct wpa_supplicant *wpa_s,
 	params.key_flag = key_flag;
 
 	if (alg != WPA_ALG_NONE) {
-		if (key_idx >= 0 && key_idx <= 6)
+		/* keyidx = 1 can be either a broadcast or--with
+		 * Extended Key ID--a unicast key. Use bit 15 for
+		 * the pairwise keyidx 1 which is hopefully high enough
+		 * to not clash with future extensions.
+		 */
+		if (key_idx == 1 && (key_flag & KEY_FLAG_PAIRWISE))
+			wpa_s->keys_cleared &= ~BIT(15);
+		else if (key_idx >= 0 && key_idx <= 5)
 			wpa_s->keys_cleared &= ~BIT(key_idx);
 		else
 			wpa_s->keys_cleared = 0;
diff --git a/wpa_supplicant/wpa_cli.c b/wpa_supplicant/wpa_cli.c
index 22885e646..730d749fe 100644
--- a/wpa_supplicant/wpa_cli.c
+++ b/wpa_supplicant/wpa_cli.c
@@ -502,6 +502,7 @@ static char ** wpa_cli_complete_set(const char *str, int pos)
 		"ignore_auth_resp",
 #endif /* CONFIG_TESTING_OPTIONS */
 		"relative_rssi", "relative_band_adjust",
+		"extended_key_id",
 	};
 	int i, num_fields = ARRAY_SIZE(fields);
 
@@ -593,7 +594,7 @@ static char ** wpa_cli_complete_get(const char *str, int pos)
 		"tdls_external_control", "osu_dir", "wowlan_triggers",
 		"p2p_search_delay", "mac_addr", "rand_addr_lifetime",
 		"preassoc_mac_addr", "key_mgmt_offload", "passive_scan",
-		"reassoc_same_bss_optim"
+		"reassoc_same_bss_optim", "extended_key_id"
 	};
 	int i, num_fields = ARRAY_SIZE(fields);
 
diff --git a/wpa_supplicant/wpa_supplicant.c b/wpa_supplicant/wpa_supplicant.c
index c638fe535..f11bac017 100644
--- a/wpa_supplicant/wpa_supplicant.c
+++ b/wpa_supplicant/wpa_supplicant.c
@@ -748,10 +748,15 @@ void wpa_clear_keys(struct wpa_supplicant *wpa_s, const u8 *addr)
 		wpa_drv_set_key(wpa_s, WPA_ALG_NONE, NULL, i, 0, NULL, 0,
 				NULL, 0, KEY_FLAG_GROUP);
 	}
-	if (!(wpa_s->keys_cleared & BIT(0)) && addr &&
+	/* Pairwise Key ID 1 for Extended Key ID is tracked in bit 15 */
+	if (!(wpa_s->keys_cleared & (BIT(0) | BIT(15))) && addr &&
 	    !is_zero_ether_addr(addr)) {
-		wpa_drv_set_key(wpa_s, WPA_ALG_NONE, addr, 0, 0, NULL, 0, NULL,
-				0, KEY_FLAG_PAIRWISE);
+		if (!(wpa_s->keys_cleared & BIT(0)))
+			wpa_drv_set_key(wpa_s, WPA_ALG_NONE, addr, 0, 0, NULL,
+					0, NULL, 0, KEY_FLAG_PAIRWISE);
+		if (!(wpa_s->keys_cleared & BIT(15)))
+			wpa_drv_set_key(wpa_s, WPA_ALG_NONE, addr, 1, 0, NULL,
+					0, NULL, 0, KEY_FLAG_PAIRWISE);
 		/* MLME-SETPROTECTION.request(None) */
 		wpa_drv_mlme_setprotection(
 			wpa_s, addr,
@@ -1635,6 +1640,30 @@ int wpa_supplicant_set_suites(struct wpa_supplicant *wpa_s,
 		sae_pwe = 1;
 	wpa_sm_set_param(wpa_s->wpa, WPA_PARAM_SAE_PWE, sae_pwe);
 
+	/* Extended Key ID is only supported in infrastructure BSS so far */
+	if (ssid->mode == WPAS_MODE_INFRA && wpa_s->conf->extended_key_id &&
+	    (ssid->proto & WPA_PROTO_RSN) &&
+	    ssid->pairwise_cipher & (WPA_CIPHER_CCMP | WPA_CIPHER_CCMP_256 |
+				     WPA_CIPHER_GCMP | WPA_CIPHER_GCMP_256) &&
+	    (wpa_s->drv_flags & WPA_DRIVER_FLAGS_EXTENDED_KEY_ID)) {
+		int use_ext_key_id = 0;
+
+		wpa_msg(wpa_s, MSG_DEBUG,
+			"WPA: Enable Extended Key ID support");
+		wpa_sm_set_param(wpa_s->wpa, WPA_PARAM_EXT_KEY_ID,
+				 wpa_s->conf->extended_key_id);
+		if (bss_rsn &&
+		    wpa_s->conf->extended_key_id &&
+		    wpa_s->pairwise_cipher != WPA_CIPHER_TKIP &&
+		    (ie.capabilities & WPA_CAPABILITY_EXT_KEY_ID_FOR_UNICAST))
+			use_ext_key_id = 1;
+		wpa_sm_set_param(wpa_s->wpa, WPA_PARAM_USE_EXT_KEY_ID,
+				 use_ext_key_id);
+	} else {
+		wpa_sm_set_param(wpa_s->wpa, WPA_PARAM_EXT_KEY_ID, 0);
+		wpa_sm_set_param(wpa_s->wpa, WPA_PARAM_USE_EXT_KEY_ID, 0);
+	}
+
 	if (wpa_sm_set_assoc_wpa_ie_default(wpa_s->wpa, wpa_ie, wpa_ie_len)) {
 		wpa_msg(wpa_s, MSG_WARNING, "WPA: Failed to generate WPA IE");
 		return -1;
diff --git a/wpa_supplicant/wpa_supplicant.conf b/wpa_supplicant/wpa_supplicant.conf
index f3a750e3c..591e1343f 100644
--- a/wpa_supplicant/wpa_supplicant.conf
+++ b/wpa_supplicant/wpa_supplicant.conf
@@ -802,6 +802,11 @@ fast_reauth=1
 # Set BIT(1) to Enable OCE in STA-CFON mode
 #oce=1
 
+# Extended Key ID support for Individually Addressed frames
+# 0 = force off: Do not use Extended Key ID (default)
+# 1 = auto: Activate Extended Key ID support if the driver supports it
+#extended_key_id=0
+
 # network block
 #
 # Each network (usually AP's sharing the same SSID) is configured as a separate
diff --git a/wpa_supplicant/wpas_glue.c b/wpa_supplicant/wpas_glue.c
index 39b05b2b9..a3049daf5 100644
--- a/wpa_supplicant/wpas_glue.c
+++ b/wpa_supplicant/wpas_glue.c
@@ -533,7 +533,8 @@ static int wpa_supplicant_set_key(void *_wpa_s, enum wpa_alg alg,
 	}
 #endif /* CONFIG_TESTING_GET_GTK */
 #ifdef CONFIG_TESTING_OPTIONS
-	if (addr && !is_broadcast_ether_addr(addr)) {
+	if (addr && !is_broadcast_ether_addr(addr) &&
+	    !(key_flag & KEY_FLAG_MODIFY)) {
 		wpa_s->last_tk_alg = alg;
 		os_memcpy(wpa_s->last_tk_addr, addr, ETH_ALEN);
 		wpa_s->last_tk_key_idx = key_idx;
@@ -1077,7 +1078,8 @@ static int wpa_supplicant_eap_auth_start_cb(void *ctx)
 {
 	struct wpa_supplicant *wpa_s = ctx;
 
-	if (!wpa_s->new_connection && wpa_s->deny_ptk0_rekey) {
+	if (!wpa_s->new_connection && wpa_s->deny_ptk0_rekey &&
+	    !wpa_sm_ext_key_id_active(wpa_s->wpa)) {
 		wpa_msg(wpa_s, MSG_INFO,
 			"WPA: PTK0 rekey not allowed, reconnecting");
 		wpa_supplicant_reconnect(wpa_s);
-- 
2.39.2