diff --git a/src/eapol.c b/src/eapol.c
index 07120fdb..cd481317 100644
--- a/src/eapol.c
+++ b/src/eapol.c
@@ -860,6 +860,7 @@ struct eapol_sm {
 	uint8_t installed_igtk_len;
 	uint8_t installed_igtk[CRYPTO_MAX_IGTK_LEN];
 	unsigned int mic_len;
+	bool rekey : 1;
 };
 
 static void eapol_sm_destroy(void *value)
@@ -1223,6 +1224,9 @@ static void eapol_handle_ptk_1_of_4(struct eapol_sm *sm,
 			return;
 		}
 
+		if (sm->handshake->ptk_complete)
+			sm->rekey = true;
+
 		handshake_state_new_snonce(sm->handshake);
 		handshake_state_set_anonce(sm->handshake, ek->key_nonce);
 
@@ -1643,12 +1647,14 @@ static void eapol_handle_ptk_3_of_4(struct eapol_sm *sm,
 	struct handshake_state *hs = sm->handshake;
 	const uint8_t *kck;
 	const uint8_t *kek;
-	struct eapol_key *step4;
+	_auto_(l_free) struct eapol_key *step4 = NULL;
 	uint8_t mic[MIC_MAXLEN];
 	const uint8_t *gtk = NULL;
 	size_t gtk_len;
 	const uint8_t *igtk = NULL;
 	size_t igtk_len;
+	const uint8_t *key_id = NULL;
+	size_t key_id_len;
 	const uint8_t *rsne;
 	struct ie_rsn_info rsn_info;
 	const uint8_t *optional_rsne = NULL;
@@ -1856,6 +1862,44 @@ static void eapol_handle_ptk_3_of_4(struct eapol_sm *sm,
 	} else
 		igtk = NULL;
 
+	key_id = handshake_util_find_kde(HANDSHAKE_KDE_KEY_ID,
+					decrypted_key_data,
+					decrypted_key_data_size, &key_id_len);
+	if (hs->ext_key_id_capable) {
+		uint8_t idx;
+
+		if (!key_id) {
+			l_debug("No extended key KDE in frame 3/4");
+			handshake_failed(sm, MMPDU_REASON_CODE_INVALID_IE);
+			return;
+		}
+
+		if (key_id_len != 2) {
+			l_error("invalid Key ID KDE format");
+			handshake_failed(sm, MMPDU_REASON_CODE_INVALID_IE);
+			return;
+		}
+
+		idx = bit_field(key_id[0], 0, 2);
+
+		/*
+		 * IEEE 802.11-2020 - 12.7.6.4 4-way handshake message 3
+		 * "... the Authenticator assigns a new Key ID for the PTKSA in
+		 * the range of 0 to 1 that is different from the Key ID
+		 * assigned in the previous handshake"
+		 */
+		if ((idx != 0 && idx != 1) || (sm->rekey &&
+						idx == hs->active_tk_index)) {
+			l_error("invalid Key ID KDE value (%u)", idx);
+			handshake_failed(sm, MMPDU_REASON_CODE_INVALID_IE);
+			return;
+		}
+
+		hs->active_tk_index = idx;
+
+		l_debug("using Extended key ID %u", hs->active_tk_index);
+	}
+
 	if (hs->support_ip_allocation) {
 		size_t len;
 		const uint8_t *ip_alloc_kde =
@@ -1917,7 +1961,6 @@ retransmit:
 		if (!eapol_calculate_mic(hs->akm_suite, kck,
 						step4, mic, sm->mic_len)) {
 			l_debug("MIC Calculation failed");
-			l_free(step4);
 			handshake_failed(sm, MMPDU_REASON_CODE_UNSPECIFIED);
 			return;
 		}
@@ -1928,18 +1971,11 @@ retransmit:
 						handshake_state_get_kek_len(hs),
 						step4, NULL, 0)) {
 			l_debug("AES-SIV encryption failed");
-			l_free(step4);
 			handshake_failed(sm, MMPDU_REASON_CODE_UNSPECIFIED);
 			return;
 		}
 	}
 
-	eapol_sm_write(sm, (struct eapol_frame *) step4, unencrypted);
-	l_free(step4);
-
-	if (hs->ptk_complete)
-		return;
-
 	/*
 	 * For WPA1 the group handshake should be happening after we set the
 	 * ptk, this flag tells netdev to wait for the gtk/igtk before
@@ -1954,6 +1990,24 @@ retransmit:
 	if (igtk)
 		eapol_install_igtk(sm, igtk_key_index, igtk, igtk_len);
 
+	/*
+	 * Only install if this is the first 3/4 message (not retransmitting)
+	 * and a rekey. Initial associations don't need the special RX -> TX
+	 * procedure and can install the TK normally
+	 */
+	if (key_id && hs->ext_key_id_capable && sm->rekey) {
+		handshake_state_install_ext_ptk(hs, hs->active_tk_index,
+						(struct eapol_frame *) step4,
+						ETH_P_PAE, unencrypted);
+
+		return;
+	}
+
+	eapol_sm_write(sm, (struct eapol_frame *) step4, unencrypted);
+
+	if (hs->ptk_complete)
+		return;
+
 	handshake_state_install_ptk(hs);
 
 	if (rekey_offload)