[v2,05/17] MLD STA: set MLO connection info to wpa_sm

diff --git a/src/rsn_supp/wpa.c b/src/rsn_supp/wpa.c
index af79551..eae1957 100644
--- a/src/rsn_supp/wpa.c
+++ b/src/rsn_supp/wpa.c
@@ -3006,6 +3006,8 @@  struct wpa_sm * wpa_sm_init(struct wpa_sm_ctx *ctx)
  */
 void wpa_sm_deinit(struct wpa_sm *sm)
 {
+	int i;
+
 	if (sm == NULL)
 		return;
 	pmksa_cache_deinit(sm->pmksa);
@@ -3016,6 +3018,10 @@  void wpa_sm_deinit(struct wpa_sm *sm)
 	os_free(sm->ap_wpa_ie);
 	os_free(sm->ap_rsn_ie);
 	os_free(sm->ap_rsnxe);
+	for (i = 0; i < MAX_NUM_MLD_LINKS; i++) {
+		os_free(sm->mlo.links[i].ap_rsne);
+		os_free(sm->mlo.links[i].ap_rsnxe);
+	}
 	wpa_sm_drop_sa(sm);
 	os_free(sm->ctx);
 #ifdef CONFIG_IEEE80211R
@@ -3303,6 +3309,69 @@  void wpa_sm_set_config(struct wpa_sm *sm, struct rsn_supp_config *config)
 	}
 }
 
+int wpa_sm_set_mlo_params(struct wpa_sm *sm, const struct wpa_sm_mlo *mlo)
+{
+	int i;
+	char title[50];
+	int ret;
+
+	if (!sm)
+		return -1;
+
+	os_memcpy(sm->mlo.ap_mld_addr, mlo->ap_mld_addr, ETH_ALEN);
+	sm->mlo.assoc_link_id =  mlo->assoc_link_id;
+	sm->mlo.setup_links = mlo->setup_links;
+	sm->mlo.req_links = mlo->req_links;
+	for (i = 0; i < MAX_NUM_MLD_LINKS; i++) {
+		const u8 *ie;
+		size_t len;
+
+		os_memcpy(sm->mlo.links[i].addr, mlo->links[i].addr, ETH_ALEN);
+		os_memcpy(sm->mlo.links[i].bssid, mlo->links[i].bssid,
+			  ETH_ALEN);
+
+		ie = mlo->links[i].ap_rsne;
+		len = mlo->links[i].ap_rsne_len;
+		os_free(sm->mlo.links[i].ap_rsne);
+		if (ie == NULL || len == 0) {
+			wpa_dbg(sm->ctx->msg_ctx, MSG_DEBUG,
+				"WPA: clearing MLO link[%u] AP RSNE", i);
+			sm->mlo.links[i].ap_rsne = NULL;
+			sm->mlo.links[i].ap_rsne_len = 0;
+		} else {
+			ret = os_snprintf(title, sizeof(title),
+					  "RSN: set MLO link[%u] AP RSNE", i);
+			if (!os_snprintf_error(sizeof(title), ret))
+				wpa_hexdump(MSG_DEBUG, title, ie, len);
+			sm->mlo.links[i].ap_rsne = os_memdup(ie, len);
+			if (!sm->mlo.links[i].ap_rsne)
+				return -1;
+			sm->mlo.links[i].ap_rsne_len = len;
+		}
+
+		ie = mlo->links[i].ap_rsnxe;
+		len = mlo->links[i].ap_rsnxe_len;
+		os_free(sm->mlo.links[i].ap_rsnxe);
+		if (ie == NULL || len == 0) {
+			wpa_dbg(sm->ctx->msg_ctx, MSG_DEBUG,
+				"WPA: clearing MLO link[%u] AP RSNXE", i);
+			sm->mlo.links[i].ap_rsnxe = NULL;
+			sm->mlo.links[i].ap_rsnxe_len = 0;
+		} else {
+			ret = os_snprintf(title, sizeof(title),
+					  "RSN: set MLO link[%u] AP RSNXE", i);
+			if (!os_snprintf_error(sizeof(title), ret))
+				wpa_hexdump(MSG_DEBUG, title, ie, len);
+			sm->mlo.links[i].ap_rsnxe = os_memdup(ie, len);
+			if (!sm->mlo.links[i].ap_rsnxe)
+				return -1;
+			sm->mlo.links[i].ap_rsnxe_len = len;
+		}
+	}
+
+	return 0;
+}
+
 
 /**
  * wpa_sm_set_own_addr - Set own MAC address
diff --git a/src/rsn_supp/wpa.h b/src/rsn_supp/wpa.h
index 80262a2..ad0c411 100644
--- a/src/rsn_supp/wpa.h
+++ b/src/rsn_supp/wpa.h
@@ -20,6 +20,7 @@  struct wpa_config_blob;
 struct hostapd_freq_params;
 struct wpa_channel_info;
 enum frame_encryption;
+struct wpa_sm_mlo;
 
 struct wpa_sm_ctx {
 	void *ctx; /* pointer to arbitrary upper level context */
@@ -224,6 +225,7 @@  void wpa_sm_set_ptk_kck_kek(struct wpa_sm *sm,
 			    const u8 *ptk_kek, size_t ptk_kek_len);
 int wpa_fils_is_completed(struct wpa_sm *sm);
 void wpa_sm_pmksa_cache_reconfig(struct wpa_sm *sm);
+int wpa_sm_set_mlo_params(struct wpa_sm *sm, const struct wpa_sm_mlo *mlo);
 
 #else /* CONFIG_NO_WPA */
 
@@ -438,6 +440,12 @@  static inline void wpa_sm_pmksa_cache_reconfig(struct wpa_sm *sm)
 {
 }
 
+static inline int wpa_sm_set_mlo_params(struct wpa_sm *sm,
+					const struct wpa_sm_mlo *mlo)
+{
+	return 0;
+}
+
 #endif /* CONFIG_NO_WPA */
 
 #ifdef CONFIG_IEEE80211R
diff --git a/src/rsn_supp/wpa_i.h b/src/rsn_supp/wpa_i.h
index 3811c3b..64b0b11 100644
--- a/src/rsn_supp/wpa_i.h
+++ b/src/rsn_supp/wpa_i.h
@@ -19,6 +19,21 @@  struct pasn_ft_r1kh {
 	u8 r1kh_id[FT_R1KH_ID_LEN];
 };
 
+struct wpa_sm_link {
+	u8 addr[ETH_ALEN];
+	u8 bssid[ETH_ALEN];
+	u8 *ap_rsne, *ap_rsnxe;
+	size_t ap_rsne_len, ap_rsnxe_len;
+} links[MAX_NUM_MLD_LINKS];
+
+struct wpa_sm_mlo {
+	u8 ap_mld_addr[ETH_ALEN];
+	u8 assoc_link_id;
+	u16 setup_links; /* bitmap of accepted links */
+	u16 req_links; /* bitmap of requested links */
+	struct wpa_sm_link links[MAX_NUM_MLD_LINKS];
+};
+
 /**
  * struct wpa_sm - Internal WPA state machine data
  */
@@ -218,6 +233,7 @@  struct wpa_sm {
 	struct wpabuf *dpp_z;
 	int dpp_pfs;
 #endif /* CONFIG_DPP2 */
+	struct wpa_sm_mlo mlo;
 };
 
 
diff --git a/wpa_supplicant/events.c b/wpa_supplicant/events.c
index f3cbe97..691465a 100644
--- a/wpa_supplicant/events.c
+++ b/wpa_supplicant/events.c
@@ -50,6 +50,7 @@ 
 #include "mesh_mpm.h"
 #include "wmm_ac.h"
 #include "dpp_supplicant.h"
+#include "rsn_supp/wpa_i.h"
 
 
 #define MAX_OWE_TRANSITION_BSS_SELECT_COUNT 5
@@ -3410,6 +3411,169 @@  static int wpa_drv_get_mlo_info(struct wpa_supplicant *wpa_s)
 }
 
 
+struct links_info {
+	/* bitmap of link IDs in Per-STA profile subelements*/
+	u16 non_assoc_links;
+	u8 addr[MAX_NUM_MLD_LINKS][ETH_ALEN];
+};
+
+
+static void wpas_get_basic_mle_links_info(const u8 *mle, size_t mle_len,
+					  struct links_info *info)
+{
+	size_t rem_len;
+	const u8 *pos;
+#define ML_CTRL_FIELD_LEN 2
+	if (mle_len < (ML_CTRL_FIELD_LEN + 1) ||
+	    (mle_len - ML_CTRL_FIELD_LEN < mle[ML_CTRL_FIELD_LEN]))
+		return;
+
+	// Skip Common Info
+	pos = mle + ML_CTRL_FIELD_LEN + mle[ML_CTRL_FIELD_LEN];
+	rem_len = mle_len - ML_CTRL_FIELD_LEN - mle[ML_CTRL_FIELD_LEN];
+
+	// Parse Subelements
+	while (rem_len > 2) {
+		int ie_len = 2 + pos[1];
+
+		if (rem_len < ie_len)
+			return;
+
+		if (pos[0] == MULTI_LINK_SUB_ELEM_ID_PER_STA_PROFILE) {
+			u8 link_id;
+			const u8 *sta_profile;
+
+#define BASIC_ML_STA_INFO_STA_MAC_IDX \
+		(2 + /* STA Control field */ \
+		 1) /* STA Info Length field (Basic) */
+			if (pos[1] < (BASIC_ML_STA_INFO_STA_MAC_IDX + ETH_ALEN))
+				goto next_subelem;
+
+			sta_profile = &pos[2];
+			link_id = sta_profile[0] &
+				  BASIC_ML_STA_CTRL0_LINK_ID_MASK;
+			if (link_id >= MAX_NUM_MLD_LINKS)
+				goto next_subelem;
+
+			if (!(sta_profile[0] & BASIC_ML_STA_CTRL0_PRES_STA_MAC))
+				goto next_subelem;
+
+			info->non_assoc_links |= BIT(link_id);
+			os_memcpy(info->addr[link_id],
+				  &sta_profile[BASIC_ML_STA_INFO_STA_MAC_IDX],
+				  ETH_ALEN);
+		}
+next_subelem:
+		pos += ie_len;
+		rem_len -= ie_len;
+	}
+}
+
+
+static int wpa_sm_set_ml_info(struct wpa_supplicant *wpa_s,
+			      union wpa_event_data *data)
+{
+	int i;
+	struct wpabuf *mle;
+	struct ieee802_11_elems req_elems, resp_elems;
+	struct links_info req_links, resp_links;
+	struct wpa_sm_mlo mlo;
+	const u8 *bss_rsn = NULL, *bss_rsnx = NULL;
+
+	os_memset(&mlo, 0, sizeof(mlo));
+	if (!wpa_s->valid_links)
+		return wpa_sm_set_mlo_params(wpa_s->wpa, &mlo);
+
+	if (!data || !data->assoc_info.req_ies || !data->assoc_info.resp_ies)
+		return -1;
+
+	if (ieee802_11_parse_elems(data->assoc_info.resp_ies,
+				   data->assoc_info.resp_ies_len, &resp_elems,
+				   0) == ParseFailed ||
+	    ieee802_11_parse_elems(data->assoc_info.req_ies,
+				   data->assoc_info.req_ies_len, &req_elems,
+				   0) == ParseFailed) {
+		wpa_dbg(wpa_s, MSG_ERROR,
+			"MLO: Failed to parse Association request/response IEs");
+		return -1;
+	}
+
+	mle = ieee802_11_defrag_mle(&req_elems, MULTI_LINK_CONTROL_TYPE_BASIC);
+	if (!mle) {
+		wpa_dbg(wpa_s, MSG_ERROR,
+			"MLO: Basic Multi-Link element not found in Association request");
+		return -1;
+	}
+	os_memset(&req_links, 0, sizeof(req_links));
+	wpas_get_basic_mle_links_info((const u8 *) wpabuf_head(mle),
+				      wpabuf_len(mle), &req_links);
+	wpabuf_free(mle);
+
+	mle = ieee802_11_defrag_mle(&resp_elems, MULTI_LINK_CONTROL_TYPE_BASIC);
+	if (!mle) {
+		wpa_dbg(wpa_s, MSG_ERROR,
+			"MLO: Basic Multi-Link element not found in Association response");
+		return -1;
+	}
+	os_memset(&resp_links, 0, sizeof(resp_links));
+	wpas_get_basic_mle_links_info((const u8 *) wpabuf_head(mle),
+				      wpabuf_len(mle), &resp_links);
+	wpabuf_free(mle);
+
+	if (req_links.non_assoc_links != resp_links.non_assoc_links) {
+		wpa_dbg(wpa_s, MSG_ERROR,
+			"MLO: Association request and response links bitmap not equal");
+		return -1;
+	}
+
+	mlo.assoc_link_id = wpa_s->mlo_assoc_link_id;
+	mlo.setup_links = wpa_s->valid_links;
+	mlo.req_links = req_links.non_assoc_links | BIT(mlo.assoc_link_id);
+	os_memcpy(mlo.ap_mld_addr, wpa_s->ap_mld_addr, ETH_ALEN);
+	for (i = 0; i < MAX_NUM_MLD_LINKS; i++) {
+		struct wpa_bss *bss;
+
+		if (!(mlo.req_links & BIT(i)))
+			continue;
+
+		if (mlo.setup_links & BIT(i)) {
+			bss = wpa_s->links[i].bss;
+		} else {
+			bss = wpa_supplicant_get_new_bss(wpa_s,
+							 resp_links.addr[i]);
+			if (!bss) {
+				wpa_supplicant_update_scan_results(wpa_s);
+				bss = wpa_supplicant_get_new_bss(
+					wpa_s, resp_links.addr[i]);
+			}
+		}
+
+		if (!bss)
+			return -1;
+
+		bss_rsn = wpa_bss_get_ie(bss, WLAN_EID_RSN);
+		if (!bss_rsn)
+			return -1;
+		bss_rsnx = wpa_bss_get_ie(bss, WLAN_EID_RSNX);
+
+		mlo.links[i].ap_rsne = (u8 *) bss_rsn;
+		mlo.links[i].ap_rsne_len = 2 + bss_rsn[1];
+		mlo.links[i].ap_rsnxe = (u8 *) bss_rsnx;
+		mlo.links[i].ap_rsnxe_len = bss_rsnx ? 2 + bss_rsnx[1] : 0;
+
+		os_memcpy(mlo.links[i].bssid, bss->bssid, ETH_ALEN);
+		if (mlo.setup_links & BIT(i))
+			os_memcpy(mlo.links[i].addr, wpa_s->links[i].addr,
+				  ETH_ALEN);
+		else
+			os_memcpy(mlo.links[i].addr, req_links.addr[i],
+				  ETH_ALEN);
+	}
+
+	return wpa_sm_set_mlo_params(wpa_s->wpa, &mlo);
+}
+
+
 static void wpa_supplicant_event_assoc(struct wpa_supplicant *wpa_s,
 				       union wpa_event_data *data)
 {
@@ -3534,6 +3698,15 @@  static void wpa_supplicant_event_assoc(struct wpa_supplicant *wpa_s,
 		wpa_supplicant_scard_init(wpa_s, wpa_s->current_ssid);
 	}
 	wpa_sm_notify_assoc(wpa_s->wpa, bssid);
+
+	if (wpa_sm_set_ml_info(wpa_s, data)) {
+		wpa_dbg(wpa_s, MSG_ERROR,
+			"Failed to set MLO connection info to wpa_sm");
+		wpa_supplicant_deauthenticate(wpa_s,
+					      WLAN_REASON_DEAUTH_LEAVING);
+		return;
+	}
+
 	if (wpa_s->l2)
 		l2_packet_notify_auth_start(wpa_s->l2);
 
diff --git a/wpa_supplicant/wpa_supplicant.c b/wpa_supplicant/wpa_supplicant.c
index 7610172..5c3d7dc 100644
--- a/wpa_supplicant/wpa_supplicant.c
+++ b/wpa_supplicant/wpa_supplicant.c
@@ -69,6 +69,7 @@ 
 #include "ap/ap_config.h"
 #include "ap/hostapd.h"
 #endif /* CONFIG_MESH */
+#include "rsn_supp/wpa_i.h"
 
 const char *const wpa_supplicant_version =
 "wpa_supplicant v" VERSION_STR "\n"
@@ -403,6 +404,7 @@  void wpa_supplicant_set_non_wpa_policy(struct wpa_supplicant *wpa_s,
 #ifdef CONFIG_WEP
 	int i;
 #endif /* CONFIG_WEP */
+	struct wpa_sm_mlo mlo;
 
 	if (ssid->key_mgmt & WPA_KEY_MGMT_WPS)
 		wpa_s->key_mgmt = WPA_KEY_MGMT_WPS;
@@ -443,6 +445,8 @@  void wpa_supplicant_set_non_wpa_policy(struct wpa_supplicant *wpa_s,
 			 wpa_s->mgmt_group_cipher);
 
 	pmksa_cache_clear_current(wpa_s->wpa);
+	os_memset(&mlo, 0, sizeof(mlo));
+	wpa_sm_set_mlo_params(wpa_s->wpa, &mlo);
 }