recipes-wifi/wpa-supplicant/files/patches-2.10.3/0017-hostapd-MLO-add-support-for-MLO-rekey.patch

From 5e6164cb6143d55409c08ae9bfd859efa188e383 Mon Sep 17 00:00:00 2001
From: Rameshkumar Sundaram <quic_ramess@quicinc.com>
Date: Thu, 28 Mar 2024 23:46:47 +0530
Subject: [PATCH 017/104] hostapd: MLO: add support for MLO rekey

Currently wpa group rekey is not supported for ML Stations when non-assoc
link initiates a group rekey, to support the same following changes have
been made-
  * Calculate links specific MLO GTK/IGTK and BIGTK KDE lengths based on
    corresponding cipher and key instead of taking length of one link and
    multiplying it by no of associated links.
  * For MLD, Arm group key rekey timer on one of the links and whenever it
    fires do group key rekey for all links.

Signed-off-by: Rameshkumar Sundaram <quic_ramess@quicinc.com>
Co-developed-by: Adil Saeed Musthafa <quic_adilm@quicinc.com>
Signed-off-by: Adil Saeed Musthafa <quic_adilm@quicinc.com>
Signed-off-by: Aditya Kumar Singh <quic_adisi@quicinc.com>
---
 src/ap/drv_callbacks.c                        |   2 +-
 src/ap/ieee802_11.c                           |  13 +-
 src/ap/wpa_auth.c                             | 310 +++++++++++++++---
 src/ap/wpa_auth.h                             |   9 +-
 src/ap/wpa_auth_glue.c                        |  22 ++
 src/ap/wpa_auth_i.h                           |   1 +
 src/ap/wpa_auth_ie.c                          |  12 +-
 src/common/wpa_common.h                       |   1 +
 tests/fuzzing/eapol-key-auth/eapol-key-auth.c |   2 +-
 wpa_supplicant/ibss_rsn.c                     |   2 +-
 10 files changed, 317 insertions(+), 57 deletions(-)

diff --git a/src/ap/drv_callbacks.c b/src/ap/drv_callbacks.c
index 064c7abae..dc21977ff 100644
--- a/src/ap/drv_callbacks.c
+++ b/src/ap/drv_callbacks.c
@@ -528,7 +528,7 @@ int hostapd_notif_assoc(struct hostapd_data *hapd, const u8 *addr,
 					  elems.rsnxe ? elems.rsnxe - 2 : NULL,
 					  elems.rsnxe ? elems.rsnxe_len + 2 : 0,
 					  elems.mdie, elems.mdie_len,
-					  elems.owe_dh, elems.owe_dh_len);
+					  elems.owe_dh, elems.owe_dh_len, NULL);
 		reason = WLAN_REASON_INVALID_IE;
 		status = WLAN_STATUS_INVALID_IE;
 		switch (res) {
diff --git a/src/ap/ieee802_11.c b/src/ap/ieee802_11.c
index 9d04bdf43..7ee18f4ae 100644
--- a/src/ap/ieee802_11.c
+++ b/src/ap/ieee802_11.c
@@ -1887,7 +1887,7 @@ void handle_auth_fils(struct hostapd_data *hapd, struct sta_info *sta,
 				  elems.rsn_ie - 2, elems.rsn_ie_len + 2,
 				  elems.rsnxe ? elems.rsnxe - 2 : NULL,
 				  elems.rsnxe ? elems.rsnxe_len + 2 : 0,
-				  elems.mdie, elems.mdie_len, NULL, 0);
+				  elems.mdie, elems.mdie_len, NULL, 0, NULL);
 	resp = wpa_res_to_status_code(res);
 	if (resp != WLAN_STATUS_SUCCESS)
 		goto fail;
@@ -3778,7 +3778,7 @@ u16 owe_process_rsn_ie(struct hostapd_data *hapd,
 	rsn_ie_len += 2;
 	res = wpa_validate_wpa_ie(hapd->wpa_auth, sta->wpa_sm,
 				  hapd->iface->freq, rsn_ie, rsn_ie_len,
-				  NULL, 0, NULL, 0, owe_dh, owe_dh_len);
+				  NULL, 0, NULL, 0, owe_dh, owe_dh_len, NULL);
 	status = wpa_res_to_status_code(res);
 	if (status != WLAN_STATUS_SUCCESS)
 		goto end;
@@ -3867,6 +3867,8 @@ static int __check_assoc_ies(struct hostapd_data *hapd, struct sta_info *sta,
 	const u8 *wpa_ie;
 	size_t wpa_ie_len;
 	const u8 *p2p_dev_addr = NULL;
+	struct hostapd_data *assoc_hapd;
+	struct sta_info *assoc_sta = NULL;
 
 	resp = check_ssid(hapd, sta, elems->ssid, elems->ssid_len);
 	if (resp != WLAN_STATUS_SUCCESS)
@@ -4041,6 +4043,10 @@ static int __check_assoc_ies(struct hostapd_data *hapd, struct sta_info *sta,
 		wpa_ie_len += 2;
 
 		if (!sta->wpa_sm) {
+			if (!link)
+				assoc_sta = hostapd_ml_get_assoc_sta(hapd, sta,
+								     &assoc_hapd);
+
 			sta->wpa_sm = wpa_auth_sta_init(hapd->wpa_auth,
 							sta->addr,
 							p2p_dev_addr);
@@ -4076,7 +4082,8 @@ static int __check_assoc_ies(struct hostapd_data *hapd, struct sta_info *sta,
 					  elems->rsnxe ? elems->rsnxe_len + 2 :
 					  0,
 					  elems->mdie, elems->mdie_len,
-					  elems->owe_dh, elems->owe_dh_len);
+					  elems->owe_dh, elems->owe_dh_len,
+					  assoc_sta ? assoc_sta->wpa_sm : NULL);
 		resp = wpa_res_to_status_code(res);
 		if (resp != WLAN_STATUS_SUCCESS)
 			return resp;
diff --git a/src/ap/wpa_auth.c b/src/ap/wpa_auth.c
index 8c1052c25..7a07dcc4c 100644
--- a/src/ap/wpa_auth.c
+++ b/src/ap/wpa_auth.c
@@ -71,6 +71,9 @@ static void wpa_group_put(struct wpa_authenticator *wpa_auth,
 			  struct wpa_group *group);
 static int ieee80211w_kde_len(struct wpa_state_machine *sm);
 static u8 * ieee80211w_kde_add(struct wpa_state_machine *sm, u8 *pos);
+static void wpa_group_update_gtk(struct wpa_authenticator *wpa_auth,
+				 struct wpa_group *group);
+
 
 static const u32 eapol_key_timeout_first = 100; /* ms */
 static const u32 eapol_key_timeout_subseq = 1000; /* ms */
@@ -102,6 +105,22 @@ static const u8 * wpa_auth_get_spa(const struct wpa_state_machine *sm)
 	return sm->addr;
 }
 
+static void wpa_update_gkeydone(struct wpa_state_machine *sm, int update)
+{
+#ifdef CONFIG_IEEE80211BE
+	int link_id;
+#endif /* CONFIG_IEEE80211BE */
+	if (!sm || !sm->wpa_auth)
+		return;
+
+	sm->wpa_auth->group->GKeyDoneStations += update;
+
+#ifdef CONFIG_IEEE80211BE
+	for_each_sm_auth(sm, link_id)
+		sm->mld_links[link_id].wpa_auth->group->GKeyDoneStations += update;
+#endif /* CONFIG_IEEE80211BE */
+}
+
 #ifdef CONFIG_IEEE80211BE
 void wpa_release_link_auth_ref(struct wpa_state_machine *sm, int release_link_id)
 {
@@ -139,10 +158,12 @@ static int wpa_get_primary_wpa_auth_cb(struct wpa_authenticator *wpa_auth, void
 	ctx->wpa_auth = wpa_auth;
 	return 1;
 }
+#endif /* CONFIG_IEEE80211BE */
 
 static struct wpa_authenticator *
 wpa_get_primary_wpa_auth(struct wpa_authenticator *wpa_auth)
 {
+#ifdef CONFIG_IEEE80211BE
 	struct wpa_get_link_auth_ctx ctx;
 
 	if (!wpa_auth || !wpa_auth->is_ml || wpa_auth->primary_auth)
@@ -153,8 +174,10 @@ wpa_get_primary_wpa_auth(struct wpa_authenticator *wpa_auth)
 	wpa_auth_for_each_auth(wpa_auth, wpa_get_primary_wpa_auth_cb, &ctx);
 
 	return ctx.wpa_auth;
-}
+#else
+	return wpa_auth;
 #endif /* CONFIG_IEEE80211BE */
+}
 
 static inline int wpa_auth_mic_failure_report(
 	struct wpa_authenticator *wpa_auth, const u8 *addr)
@@ -420,15 +443,16 @@ static void wpa_rekey_gmk(void *eloop_ctx, void *timeout_ctx)
 	}
 }
 
-
-static void wpa_rekey_gtk(void *eloop_ctx, void *timeout_ctx)
+static void wpa_rekey_all_groups(struct wpa_authenticator *wpa_auth)
 {
-	struct wpa_authenticator *wpa_auth = eloop_ctx;
 	struct wpa_group *group, *next;
 
 	wpa_auth_logger(wpa_auth, NULL, LOGGER_DEBUG, "rekeying GTK");
 	group = wpa_auth->group;
 	while (group) {
+		wpa_printf(MSG_DEBUG, "GTK rekey start for authenticator("
+			   MACSTR "), group vlan %d",
+			   MAC2STR(wpa_auth->addr), group->vlan_id);
 		wpa_group_get(wpa_auth, group);
 
 		group->GTKReKey = true;
@@ -441,6 +465,80 @@ static void wpa_rekey_gtk(void *eloop_ctx, void *timeout_ctx)
 		wpa_group_put(wpa_auth, group);
 		group = next;
 	}
+}
+
+#ifdef CONFIG_IEEE80211BE
+static void wpa_update_all_gtks(struct wpa_authenticator *wpa_auth)
+{
+	struct wpa_group *group, *next;
+
+	group = wpa_auth->group;
+	while (group) {
+		wpa_group_get(wpa_auth, group);
+
+		wpa_group_update_gtk(wpa_auth, group);
+		next = group->next;
+		wpa_group_put(wpa_auth, group);
+		group = next;
+	}
+}
+
+static int wpa_update_all_gtks_cb(struct wpa_authenticator *wpa_auth, void *ctx)
+{
+	u8 *mld_addr = ctx;
+
+	if (os_memcmp(wpa_auth->mld_addr, mld_addr, ETH_ALEN) != 0)
+		return 0;
+
+	wpa_update_all_gtks(wpa_auth);
+	return 0;
+}
+
+static int wpa_rekey_all_groups_cb(struct wpa_authenticator *wpa_auth,
+				   void *ctx)
+{
+	u8 *mld_addr = ctx;
+
+	if (os_memcmp(wpa_auth->mld_addr, mld_addr, ETH_ALEN) != 0)
+		return 0;
+
+	wpa_rekey_all_groups(wpa_auth);
+	return 0;
+}
+#endif /* CONFIG_IEEE80211BE */
+
+static void wpa_rekey_gtk(void *eloop_ctx, void *timeout_ctx)
+{
+	struct wpa_authenticator *wpa_auth = eloop_ctx;
+
+#ifdef CONFIG_IEEE80211BE
+	if (wpa_auth->is_ml) {
+		/* Non Primary ML authenticator eloop timer for group rekey is never
+		 * started and shouldn't fire too check and warn just in case
+		 */
+		if (!wpa_auth->primary_auth) {
+			wpa_printf(MSG_DEBUG,
+				   "WPA: Can't start GTK rekey on non-primary ML authenticator");
+			return;
+		}
+		/*
+		 * Generate all the new I/BIG/GTKs
+		 */
+		wpa_auth_for_each_auth(wpa_auth, wpa_update_all_gtks_cb,
+				       wpa_auth->mld_addr);
+
+		/*
+		 * Send all the generated I/BIG/GTKs to the respective
+		 * stations via G1 messages
+		 */
+		wpa_auth_for_each_auth(wpa_auth, wpa_rekey_all_groups_cb,
+				       wpa_auth->mld_addr);
+	} else {
+		wpa_rekey_all_groups(wpa_auth);
+	}
+#else
+	wpa_rekey_all_groups(wpa_auth);
+#endif /* CONFIG_IEEE80211BE */
 
 	if (wpa_auth->conf.wpa_group_rekey) {
 		eloop_register_timeout(wpa_auth->conf.wpa_group_rekey,
@@ -590,8 +688,19 @@ struct wpa_authenticator * wpa_init(const u8 *addr,
 	wpa_auth = os_zalloc(sizeof(struct wpa_authenticator));
 	if (!wpa_auth)
 		return NULL;
+
 	os_memcpy(wpa_auth->addr, addr, ETH_ALEN);
 	os_memcpy(&wpa_auth->conf, conf, sizeof(*conf));
+
+#ifdef CONFIG_IEEE80211BE
+	if (conf->mld_addr) {
+		wpa_auth->is_ml = true;
+		wpa_auth->link_id = conf->link_id;
+		wpa_auth->primary_auth = !conf->first_link_auth;
+		os_memcpy(wpa_auth->mld_addr, conf->mld_addr, ETH_ALEN);
+	}
+#endif /* CONFIG_IEEE80211BE */
+
 	wpa_auth->cb = cb;
 	wpa_auth->cb_ctx = cb_ctx;
 
@@ -635,7 +744,15 @@ struct wpa_authenticator * wpa_init(const u8 *addr,
 				       wpa_rekey_gmk, wpa_auth, NULL);
 	}
 
+#ifdef CONFIG_IEEE80211BE
+	/* For ML AP, run Group rekey timer only on one link(first) and whenever
+	 * it fires do rekey on all associated ML links at one shot.
+	 */
+	if ((!wpa_auth->is_ml || !conf->first_link_auth) &&
+	    wpa_auth->conf.wpa_group_rekey) {
+#else
 	if (wpa_auth->conf.wpa_group_rekey) {
+#endif /* CONFIG_IEEE80211BE */
 		eloop_register_timeout(wpa_auth->conf.wpa_group_rekey, 0,
 				       wpa_rekey_gtk, wpa_auth, NULL);
 	}
@@ -699,6 +816,10 @@ void wpa_deinit(struct wpa_authenticator *wpa_auth)
 	struct wpa_group *group, *prev;
 
 	eloop_cancel_timeout(wpa_rekey_gmk, wpa_auth, NULL);
+
+	/* TODO: assign ML Primary authenticator to next link auth and
+	 * start rekey timer.
+	 */
 	eloop_cancel_timeout(wpa_rekey_gtk, wpa_auth, NULL);
 
 	pmksa_cache_auth_deinit(wpa_auth->pmksa);
@@ -868,7 +989,7 @@ static void wpa_free_sta_sm(struct wpa_state_machine *sm)
 	}
 #endif /* CONFIG_P2P */
 	if (sm->GUpdateStationKeys) {
-		sm->group->GKeyDoneStations--;
+		wpa_update_gkeydone(sm, -1);
 		sm->GUpdateStationKeys = false;
 	}
 #ifdef CONFIG_IEEE80211R_AP
@@ -1669,12 +1790,14 @@ void wpa_receive(struct wpa_authenticator *wpa_auth,
 			wpa_auth_logger(wpa_auth, wpa_auth_get_spa(sm),
 					LOGGER_INFO,
 					"received EAPOL-Key Request for GTK rekeying");
-			eloop_cancel_timeout(wpa_rekey_gtk, wpa_auth, NULL);
+
+			eloop_cancel_timeout(wpa_rekey_gtk,
+					     wpa_get_primary_wpa_auth(wpa_auth), NULL);
 			if (wpa_auth_gtk_rekey_in_process(wpa_auth))
 				wpa_auth_logger(wpa_auth, NULL, LOGGER_DEBUG,
 						"skip new GTK rekey - already in process");
 			else
-				wpa_rekey_gtk(wpa_auth, NULL);
+				wpa_rekey_gtk(wpa_get_primary_wpa_auth(wpa_auth), NULL);
 		}
 	} else {
 		/* Do not allow the same key replay counter to be reused. */
@@ -2207,7 +2330,7 @@ int wpa_auth_sm_event(struct wpa_state_machine *sm, enum wpa_event event)
 			 * Reauthentication cancels the pending group key
 			 * update for this STA.
 			 */
-			sm->group->GKeyDoneStations--;
+			wpa_update_gkeydone(sm, -1);
 			sm->GUpdateStationKeys = false;
 			sm->PtkGroupInit = true;
 		}
@@ -2284,7 +2407,7 @@ SM_STATE(WPA_PTK, INITIALIZE)
 
 	sm->keycount = 0;
 	if (sm->GUpdateStationKeys)
-		sm->group->GKeyDoneStations--;
+		wpa_update_gkeydone(sm, -1);
 	sm->GUpdateStationKeys = false;
 	if (sm->wpa == WPA_VERSION_WPA)
 		sm->PInitAKeys = false;
@@ -4058,41 +4181,54 @@ static void wpa_auth_get_ml_key_info(struct wpa_authenticator *wpa_auth,
 	wpa_auth->cb->get_ml_key_info(wpa_auth->cb_ctx, info);
 }
 
+#define KDE_HDR_LEN (1 + 1 + RSN_SELECTOR_LEN)
 
 static size_t wpa_auth_ml_group_kdes_len(struct wpa_state_machine *sm)
 {
-	struct wpa_authenticator *wpa_auth = sm->wpa_auth;
-	struct wpa_group *gsm = sm->group;
-	size_t gtk_len = gsm->GTK_len;
-	size_t igtk_len;
-	size_t kde_len;
-	unsigned int n_links;
+	struct wpa_authenticator *wpa_auth;
+	size_t kde_len = 0;
+	int link_id;
 
 	if (sm->mld_assoc_link_id < 0)
 		return 0;
 
-	n_links = sm->n_mld_affiliated_links + 1;
+	for (link_id = 0; link_id < MAX_NUM_MLD_LINKS; link_id++) {
+		if (!sm->mld_links[link_id].valid)
+			continue;
+
+		wpa_auth = sm->mld_links[link_id].wpa_auth;
+		if (!wpa_auth || !wpa_auth->group)
+			continue;
 
-	/* MLO GTK KDE for each link */
-	kde_len = n_links * (2 + RSN_SELECTOR_LEN + 1 + 6 + gtk_len);
+		/* MLO GTK KDE
+		 * Header + Key-idx and Link-id + PN
+		 */
+		kde_len += (KDE_HDR_LEN  + 1 + WPA_MLO_GTK_KDE_PN_LEN);
+		kde_len += wpa_auth->group->GTK_len;
 
-	if (!sm->mgmt_frame_prot)
-		return kde_len;
+		if (!sm->mgmt_frame_prot)
+			continue;
 
-	/* MLO IGTK KDE for each link */
-	igtk_len = wpa_cipher_key_len(wpa_auth->conf.group_mgmt_cipher);
-	kde_len += n_links * (2 + RSN_SELECTOR_LEN + 2 + 6 + 1 + igtk_len);
+		if (wpa_auth->conf.tx_bss_auth)
+			wpa_auth = wpa_auth->conf.tx_bss_auth;
 
-	if (wpa_auth->conf.tx_bss_auth) {
-		wpa_auth = wpa_auth->conf.tx_bss_auth;
-		igtk_len = wpa_cipher_key_len(wpa_auth->conf.group_mgmt_cipher);
-	}
+		/* MLO IGTK KDE
+		 * Header + Key-idx & IPN + Link-id
+		 */
+		kde_len += (KDE_HDR_LEN + WPA_IGTK_KDE_PREFIX_LEN + 1);
+		kde_len += wpa_cipher_key_len(wpa_auth->conf.group_mgmt_cipher);
 
-	if (!wpa_auth->conf.beacon_prot)
-		return kde_len;
+		if (!wpa_auth->conf.beacon_prot)
+			continue;
+
+		/* MLO BIGTK KDE
+		 * Header + Key-idx & IPN + Link-id
+		 */
+		kde_len += (KDE_HDR_LEN + WPA_BIGTK_KDE_PREFIX_LEN + 1);
+		kde_len += wpa_cipher_key_len(wpa_auth->conf.group_mgmt_cipher);
+	}
 
-	/* MLO BIGTK KDE for each link */
-	kde_len += n_links * (2 + RSN_SELECTOR_LEN + 2 + 6 + 1 + igtk_len);
+	wpa_printf(MSG_DEBUG, "MLO Group kdes len = %zu", kde_len);
 
 	return kde_len;
 }
@@ -4102,6 +4238,7 @@ static u8 * wpa_auth_ml_group_kdes(struct wpa_state_machine *sm, u8 *pos)
 {
 	struct wpa_auth_ml_key_info ml_key_info;
 	unsigned int i, link_id;
+	u8 *start = pos;
 
 	/* First fetch the key information from all the authenticators */
 	os_memset(&ml_key_info, 0, sizeof(ml_key_info));
@@ -4153,8 +4290,10 @@ static u8 * wpa_auth_ml_group_kdes(struct wpa_state_machine *sm, u8 *pos)
 		i++;
 	}
 
-	if (!sm->mgmt_frame_prot)
+	if (!sm->mgmt_frame_prot) {
+		wpa_printf(MSG_DEBUG, "RSN: MLO Group kde len = %ld", pos - start);
 		return pos;
+	}
 
 	/* Add MLO IGTK KDEs */
 	for (i = 0, link_id = 0; link_id < MAX_NUM_MLD_LINKS; link_id++) {
@@ -4193,8 +4332,10 @@ static u8 * wpa_auth_ml_group_kdes(struct wpa_state_machine *sm, u8 *pos)
 		i++;
 	}
 
-	if (!sm->wpa_auth->conf.beacon_prot)
+	if (!sm->wpa_auth->conf.beacon_prot) {
+		wpa_printf(MSG_DEBUG, "RSN: MLO Group kde len = %ld", pos - start);
 		return pos;
+	}
 
 	/* Add MLO BIGTK KDEs */
 	for (i = 0, link_id = 0; link_id < MAX_NUM_MLD_LINKS; link_id++) {
@@ -4234,6 +4375,7 @@ static u8 * wpa_auth_ml_group_kdes(struct wpa_state_machine *sm, u8 *pos)
 		i++;
 	}
 
+	wpa_printf(MSG_DEBUG, "RSN: MLO Group kde len = %ld", pos - start);
 	return pos;
 }
 
@@ -4274,6 +4416,7 @@ static u8 * wpa_auth_ml_kdes(struct wpa_state_machine *sm, u8 *pos)
 {
 #ifdef CONFIG_IEEE80211BE
 	u8 link_id;
+	u8 *start = pos;
 
 	if (sm->mld_assoc_link_id < 0)
 		return pos;
@@ -4324,6 +4467,7 @@ static u8 * wpa_auth_ml_kdes(struct wpa_state_machine *sm, u8 *pos)
 		}
 	}
 
+	wpa_printf(MSG_DEBUG, "RSN: MLO Link kde len = %ld", pos - start);
 	pos = wpa_auth_ml_group_kdes(sm, pos);
 #endif /* CONFIG_IEEE80211BE */
 
@@ -5106,7 +5250,7 @@ SM_STATE(WPA_PTK_GROUP, REKEYESTABLISHED)
 #endif /* CONFIG_OCV */
 
 	if (sm->GUpdateStationKeys)
-		sm->group->GKeyDoneStations--;
+		wpa_update_gkeydone(sm, -1);
 	sm->GUpdateStationKeys = false;
 	sm->GTimeoutCtr = 0;
 	/* FIX: MLME.SetProtection.Request(TA, Tx_Rx) */
@@ -5121,7 +5265,7 @@ SM_STATE(WPA_PTK_GROUP, KEYERROR)
 {
 	SM_ENTRY_MA(WPA_PTK_GROUP, KEYERROR, wpa_ptk_group);
 	if (sm->GUpdateStationKeys)
-		sm->group->GKeyDoneStations--;
+		wpa_update_gkeydone(sm, -1);
 	sm->GUpdateStationKeys = false;
 	sm->Disconnect = true;
 	sm->disconnect_reason = WLAN_REASON_GROUP_KEY_UPDATE_TIMEOUT;
@@ -5415,18 +5559,11 @@ int wpa_wnmsleep_bigtk_subelem(struct wpa_state_machine *sm, u8 *pos)
 
 #endif /* CONFIG_WNM_AP */
 
-
-static void wpa_group_setkeys(struct wpa_authenticator *wpa_auth,
-			      struct wpa_group *group)
+static void wpa_group_update_gtk(struct wpa_authenticator *wpa_auth,
+				 struct wpa_group *group)
 {
 	int tmp;
 
-	wpa_printf(MSG_DEBUG,
-		   "WPA: group state machine entering state SETKEYS (VLAN-ID %d)",
-		   group->vlan_id);
-	group->changed = true;
-	group->wpa_group_state = WPA_GROUP_SETKEYS;
-	group->GTKReKey = false;
 	tmp = group->GM;
 	group->GM = group->GN;
 	group->GN = tmp;
@@ -5440,6 +5577,24 @@ static void wpa_group_setkeys(struct wpa_authenticator *wpa_auth,
 	 * counting the STAs that are marked with GUpdateStationKeys instead of
 	 * including all STAs that could be in not-yet-completed state. */
 	wpa_gtk_update(wpa_auth, group);
+}
+
+static void wpa_group_setkeys(struct wpa_authenticator *wpa_auth,
+			      struct wpa_group *group)
+{
+	wpa_printf(MSG_DEBUG,
+		   "WPA: group state machine entering state SETKEYS (VLAN-ID %d)",
+		   group->vlan_id);
+	group->changed = true;
+	group->wpa_group_state = WPA_GROUP_SETKEYS;
+	group->GTKReKey = false;
+
+#ifdef CONFIG_IEEE80211BE
+	if (wpa_auth->is_ml)
+		goto skip_update;
+#endif /* CONFIG_IEEE80211BE */
+
+	wpa_group_update_gtk(wpa_auth, group);
 
 	if (group->GKeyDoneStations) {
 		wpa_printf(MSG_DEBUG,
@@ -5447,6 +5602,10 @@ static void wpa_group_setkeys(struct wpa_authenticator *wpa_auth,
 			   group->GKeyDoneStations);
 		group->GKeyDoneStations = 0;
 	}
+
+#ifdef CONFIG_IEEE80211BE
+skip_update:
+#endif /* CONFIG_IEEE80211BE */
 	wpa_auth_for_each_sta(wpa_auth, wpa_group_update_sta, group);
 	wpa_printf(MSG_DEBUG, "wpa_group_setkeys: GKeyDoneStations=%d",
 		   group->GKeyDoneStations);
@@ -5564,6 +5723,57 @@ static void wpa_group_sm_step(struct wpa_authenticator *wpa_auth,
 	}
 }
 
+static void wpa_mark_group_change(struct wpa_state_machine *sm, bool change)
+{
+#ifdef CONFIG_IEEE80211BE
+	int link_id;
+#endif /* CONFIG_IEEE80211BE */
+
+	if (!sm || !sm->wpa_auth)
+		return;
+	sm->wpa_auth->group->changed = change;
+
+#ifdef CONFIG_IEEE80211BE
+	for_each_sm_auth(sm, link_id)
+		sm->mld_links[link_id].wpa_auth->group->changed = change;
+#endif /* CONFIG_IEEE80211BE */
+}
+
+static void wpa_group_sm_step_links(struct wpa_state_machine *sm)
+{
+#ifdef CONFIG_IEEE80211BE
+	int link_id;
+#endif /* CONFIG_IEEE80211BE */
+
+	if (!sm || !sm->wpa_auth)
+		return;
+	wpa_group_sm_step(sm->wpa_auth, sm->wpa_auth->group);
+
+#ifdef CONFIG_IEEE80211BE
+	for_each_sm_auth(sm, link_id)
+		wpa_group_sm_step(sm->mld_links[link_id].wpa_auth,
+				  sm->mld_links[link_id].wpa_auth->group);
+#endif /* CONFIG_IEEE80211BE */
+}
+
+static bool wpa_group_sm_changed(struct wpa_state_machine *sm)
+{
+#ifdef CONFIG_IEEE80211BE
+	int link_id;
+#endif /* CONFIG_IEEE80211BE */
+	bool changed;
+
+	if (!sm || !sm->wpa_auth)
+		return false;
+	changed = sm->wpa_auth->group->changed;
+
+#ifdef CONFIG_IEEE80211BE
+	for_each_sm_auth(sm, link_id)
+		changed |= sm->mld_links[link_id].wpa_auth->group->changed;
+#endif /* CONFIG_IEEE80211BE */
+
+	return changed;
+}
 
 static int wpa_sm_step(struct wpa_state_machine *sm)
 {
@@ -5584,7 +5794,7 @@ static int wpa_sm_step(struct wpa_state_machine *sm)
 			break;
 
 		sm->changed = false;
-		sm->wpa_auth->group->changed = false;
+		wpa_mark_group_change(sm, false);
 
 		SM_STEP_RUN(WPA_PTK);
 		if (sm->pending_deinit)
@@ -5592,8 +5802,8 @@ static int wpa_sm_step(struct wpa_state_machine *sm)
 		SM_STEP_RUN(WPA_PTK_GROUP);
 		if (sm->pending_deinit)
 			break;
-		wpa_group_sm_step(sm->wpa_auth, sm->group);
-	} while (sm->changed || sm->wpa_auth->group->changed);
+		wpa_group_sm_step_links(sm);
+	} while (sm->changed || wpa_group_sm_changed(sm));
 	sm->in_step_loop = 0;
 
 	if (sm->pending_deinit) {
@@ -6807,8 +7017,10 @@ int wpa_auth_rekey_gtk(struct wpa_authenticator *wpa_auth)
 {
 	if (!wpa_auth)
 		return -1;
-	eloop_cancel_timeout(wpa_rekey_gtk, wpa_auth, NULL);
-	return eloop_register_timeout(0, 0, wpa_rekey_gtk, wpa_auth, NULL);
+	eloop_cancel_timeout(wpa_rekey_gtk,
+			     wpa_get_primary_wpa_auth(wpa_auth), NULL);
+	return eloop_register_timeout(0, 0, wpa_rekey_gtk,
+				      wpa_get_primary_wpa_auth(wpa_auth), NULL);
 }
 
 
diff --git a/src/ap/wpa_auth.h b/src/ap/wpa_auth.h
index 1446872f3..331d217b5 100644
--- a/src/ap/wpa_auth.h
+++ b/src/ap/wpa_auth.h
@@ -285,6 +285,12 @@ struct wpa_auth_config {
 	 * Set only in nontransmitted BSSs, i.e., is NULL for transmitted BSS
 	 * and in BSSs that are not part of a Multi-BSSID set. */
 	struct wpa_authenticator *tx_bss_auth;
+
+#ifdef CONFIG_IEEE80211BE
+	u8 *mld_addr;
+	int link_id;
+	struct wpa_authenticator *first_link_auth;
+#endif /* CONFIG_IEEE80211BE */
 };
 
 typedef enum {
@@ -429,7 +435,8 @@ wpa_validate_wpa_ie(struct wpa_authenticator *wpa_auth,
 		    const u8 *wpa_ie, size_t wpa_ie_len,
 		    const u8 *rsnxe, size_t rsnxe_len,
 		    const u8 *mdie, size_t mdie_len,
-		    const u8 *owe_dh, size_t owe_dh_len);
+		    const u8 *owe_dh, size_t owe_dh_len,
+		    struct wpa_state_machine *assoc_sm);
 int wpa_validate_osen(struct wpa_authenticator *wpa_auth,
 		      struct wpa_state_machine *sm,
 		      const u8 *osen_ie, size_t osen_ie_len);
diff --git a/src/ap/wpa_auth_glue.c b/src/ap/wpa_auth_glue.c
index d3cd44695..1726c7201 100644
--- a/src/ap/wpa_auth_glue.c
+++ b/src/ap/wpa_auth_glue.c
@@ -1713,6 +1713,7 @@ int hostapd_setup_wpa(struct hostapd_data *hapd)
 
 	hostapd_wpa_auth_conf(hapd->conf, hapd->iconf, &_conf);
 	_conf.msg_ctx = hapd->msg_ctx;
+
 	tx_bss = hostapd_mbssid_get_tx_bss(hapd);
 	if (tx_bss != hapd)
 		_conf.tx_bss_auth = tx_bss->wpa_auth;
@@ -1753,6 +1754,27 @@ int hostapd_setup_wpa(struct hostapd_data *hapd)
 		!!(hapd->iface->drv_flags2 &
 		   WPA_DRIVER_FLAGS2_PROT_RANGE_NEG_AP);
 
+#ifdef CONFIG_IEEE80211BE
+	_conf.mld_addr = NULL;
+	_conf.link_id = -1;
+	_conf.first_link_auth = NULL;
+
+	if (hapd->conf->mld_ap) {
+		struct hostapd_data *lhapd;
+
+		_conf.mld_addr = hapd->mld->mld_addr;
+		_conf.link_id = hapd->mld_link_id;
+
+		for_each_mld_link(lhapd, hapd) {
+			if (lhapd == hapd)
+				continue;
+
+			if (lhapd->wpa_auth)
+				_conf.first_link_auth = lhapd->wpa_auth;
+		}
+	}
+#endif /* CONFIG_IEEE80211BE */
+
 	hapd->wpa_auth = wpa_init(hapd->own_addr, &_conf, &cb, hapd);
 	if (hapd->wpa_auth == NULL) {
 		wpa_printf(MSG_ERROR, "WPA initialization failed.");
diff --git a/src/ap/wpa_auth_i.h b/src/ap/wpa_auth_i.h
index 9ba90749d..29bb66733 100644
--- a/src/ap/wpa_auth_i.h
+++ b/src/ap/wpa_auth_i.h
@@ -176,6 +176,7 @@ struct wpa_state_machine {
 	u8 peer_mld_addr[ETH_ALEN];
 	s8 mld_assoc_link_id;
 	u8 n_mld_affiliated_links;
+	u16 valid_links;
 
 	struct mld_link {
 		bool valid;
diff --git a/src/ap/wpa_auth_ie.c b/src/ap/wpa_auth_ie.c
index a5f2861c9..bf2303e4f 100644
--- a/src/ap/wpa_auth_ie.c
+++ b/src/ap/wpa_auth_ie.c
@@ -608,7 +608,8 @@ wpa_validate_wpa_ie(struct wpa_authenticator *wpa_auth,
 		    const u8 *wpa_ie, size_t wpa_ie_len,
 		    const u8 *rsnxe, size_t rsnxe_len,
 		    const u8 *mdie, size_t mdie_len,
-		    const u8 *owe_dh, size_t owe_dh_len)
+		    const u8 *owe_dh, size_t owe_dh_len,
+		    struct wpa_state_machine *assoc_sm)
 {
 	struct wpa_auth_config *conf = &wpa_auth->conf;
 	struct wpa_ie_data data;
@@ -956,6 +957,15 @@ wpa_validate_wpa_ie(struct wpa_authenticator *wpa_auth,
 	else
 		sm->wpa = WPA_VERSION_WPA;
 
+	if (assoc_sm) {
+		/* For ML Association Link STA cannot choose a different
+		 * akm or pairwise cipher from assoc STA
+		 */
+		if (sm->wpa_key_mgmt != assoc_sm->wpa_key_mgmt)
+			return WPA_INVALID_AKMP;
+		if (sm->pairwise != assoc_sm->pairwise)
+			return WPA_INVALID_PAIRWISE;
+	}
 #if defined(CONFIG_IEEE80211R_AP) && defined(CONFIG_FILS)
 	if ((sm->wpa_key_mgmt == WPA_KEY_MGMT_FT_FILS_SHA256 ||
 	     sm->wpa_key_mgmt == WPA_KEY_MGMT_FT_FILS_SHA384) &&
diff --git a/src/common/wpa_common.h b/src/common/wpa_common.h
index 01efeea3a..24ceed600 100644
--- a/src/common/wpa_common.h
+++ b/src/common/wpa_common.h
@@ -24,6 +24,7 @@
 #define WPA_PASN_PMK_LEN 32
 #define WPA_PASN_MAX_MIC_LEN 24
 #define WPA_MAX_RSNXE_LEN 4
+#define WPA_MLO_GTK_KDE_PN_LEN 6
 
 #define OWE_DH_GROUP 19
 
diff --git a/tests/fuzzing/eapol-key-auth/eapol-key-auth.c b/tests/fuzzing/eapol-key-auth/eapol-key-auth.c
index bb46422c6..17f69fd76 100644
--- a/tests/fuzzing/eapol-key-auth/eapol-key-auth.c
+++ b/tests/fuzzing/eapol-key-auth/eapol-key-auth.c
@@ -262,7 +262,7 @@ static int auth_init(struct wpa *wpa)
 	}
 
 	if (wpa_validate_wpa_ie(wpa->auth_group, wpa->auth, 2412, supp_ie,
-				supp_ie_len, NULL, 0, NULL, 0, NULL, 0) !=
+				supp_ie_len, NULL, 0, NULL, 0, NULL, 0, NULL) !=
 	    WPA_IE_OK) {
 		wpa_printf(MSG_DEBUG, "AUTH: wpa_validate_wpa_ie() failed");
 		return -1;
diff --git a/wpa_supplicant/ibss_rsn.c b/wpa_supplicant/ibss_rsn.c
index 554268a47..2d06f1a6a 100644
--- a/wpa_supplicant/ibss_rsn.c
+++ b/wpa_supplicant/ibss_rsn.c
@@ -484,7 +484,7 @@ static int ibss_rsn_auth_init(struct ibss_rsn *ibss_rsn,
 				"\x00\x0f\xac\x04"
 				"\x01\x00\x00\x0f\xac\x04"
 				"\x01\x00\x00\x0f\xac\x02"
-				"\x00\x00", 22, NULL, 0, NULL, 0, NULL, 0) !=
+				"\x00\x00", 22, NULL, 0, NULL, 0, NULL, 0, NULL) !=
 	    WPA_IE_OK) {
 		wpa_printf(MSG_DEBUG, "AUTH: wpa_validate_wpa_ie() failed");
 		return -1;
-- 
2.39.2