wlan-ap-Telecominfraproject/feeds/qca/hostapd/patches/r00-002-hostapd-Add-support-for-cohosted-ML-BSS.patch

From 9ed4fb3c716d4e5fc942b5ae3eb4b531746dad90 Mon Sep 17 00:00:00 2001
From: Sriram R <quic_srirrama@quicinc.com>
Date: Mon, 11 Sep 2023 09:38:50 +0530
Subject: [PATCH] hostapd: Add support for cohosted ML BSS

Currently MLO is being supported with an assumption
of only single bss per link in the hostapd conf file.
This needs to be extended when cohosted ML BSS needs
to be supported and this is required for MBSSID MLO
support as well.

Add necessary helper apis to support Multiple BSS support
for MLO to make the changes scalable.

Modify auth and assoc frames to be always sent with link
address as A1 and A3 for ease of Tx status handling.

Signed-off-by: Sriram R <quic_srirrama@quicinc.com>
---
 hostapd/main.c                     |  37 ++-----
 src/ap/beacon.c                    |  26 +++--
 src/ap/drv_callbacks.c             | 165 ++++++++++++++++-------------
 src/ap/hostapd.c                   | 125 +++++++++++++++++++---
 src/ap/hostapd.h                   |   6 ++
 src/ap/ieee802_11.c                | 159 +++++++--------------------
 src/ap/ieee802_11.h                |   2 -
 src/ap/ieee802_11_eht.c            |  23 ++--
 src/ap/wpa_auth_glue.c             |  40 ++++---
 src/drivers/driver.h               |   8 ++
 src/drivers/driver_nl80211.c       |   6 +-
 src/drivers/driver_nl80211_event.c |  25 +++--
 12 files changed, 329 insertions(+), 293 deletions(-)

--- a/hostapd/main.c
+++ b/hostapd/main.c
@@ -168,6 +168,7 @@ static int hostapd_driver_init(struct ho
 	struct hostapd_bss_config *conf = hapd->conf;
 	u8 *b = conf->bssid;
 	struct wpa_driver_capa capa;
+	struct hostapd_data *h_hapd = NULL;
 
 	if (hapd->driver == NULL || hapd->driver->hapd_init == NULL) {
 		wpa_printf(MSG_ERROR, "No hostapd driver wrapper available");
@@ -175,35 +176,10 @@ static int hostapd_driver_init(struct ho
 	}
 
 #ifdef CONFIG_IEEE80211BE
-	for (i = 0; conf->mld_ap && i < iface->interfaces->count; i++) {
-		struct hostapd_iface *h = iface->interfaces->iface[i];
-		struct hostapd_data *h_hapd = h->bss[0];
-		struct hostapd_bss_config *hconf = h_hapd->conf;
-
-		if (h == iface) {
-			wpa_printf(MSG_DEBUG, "MLD: Skip own interface");
-			continue;
-		}
-
-		if (!hconf->mld_ap) {
-			wpa_printf(MSG_DEBUG,
-				   "MLD: Skip non MLD");
-			continue;
-		}
-
-		if (!hostapd_is_ml_partner(hapd, h_hapd)) {
-			wpa_printf(MSG_DEBUG,
-				   "MLD: Skip non matching mld vif name");
-			continue;
-		}
-
-		wpa_printf(MSG_DEBUG, "MLD: Found matching MLD interface");
-		if (!h_hapd->drv_priv) {
-			wpa_printf(MSG_DEBUG,
-				   "MLD: Matching MLD BSS not initialized yet");
-			continue;
-		}
+	if (conf->mld_ap)
+		h_hapd = hostapd_mld_get_first_bss(hapd);
 
+	if (h_hapd) {
 		hapd->drv_priv = h_hapd->drv_priv;
 
 		/*
@@ -223,6 +199,8 @@ static int hostapd_driver_init(struct ho
 
 		hapd->mld_link_id = hapd->mld->next_link_id++;
 		hostapd_mld_add_link(hapd);
+		wpa_printf(MSG_DEBUG, "Setup of non first link BSS(link id %d) of MLD %s",
+			   hapd->mld_link_id, hapd->conf->iface);
 
 		goto setup_mld;
 	}
@@ -293,6 +271,8 @@ static int hostapd_driver_init(struct ho
 		random_mac_addr_keep_oui(hapd->own_addr);
 		hapd->mld_link_id = hapd->mld->next_link_id++;
 		hostapd_mld_add_link(hapd);
+		wpa_printf(MSG_DEBUG, "Setup of first link (link %d) BSS of MLD %s",
+			   hapd->mld_link_id, hapd->conf->iface);
 	}
 
 setup_mld:
--- a/src/ap/beacon.c
+++ b/src/ap/beacon.c
@@ -2306,7 +2306,7 @@ int ieee802_11_set_beacon(struct hostapd
 	struct hostapd_iface *iface = hapd->iface;
 	int ret;
 	size_t i, j;
-	bool is_6g;
+	bool is_6g, hapd_mld = false;
 
 	ret = __ieee802_11_set_beacon(hapd);
 	if (ret != 0)
@@ -2315,27 +2315,35 @@ int ieee802_11_set_beacon(struct hostapd
 	if (!iface->interfaces || iface->interfaces->count <= 1)
 		return 0;
 
+#ifdef CONFIG_IEEE80211BE
+	hapd_mld = hapd->conf->mld_ap;
+#endif
 	/* Update Beacon frames in case of 6 GHz colocation or AP MLD */
 	is_6g = is_6ghz_op_class(iface->conf->op_class);
 	for (j = 0; j < iface->interfaces->count; j++) {
 		struct hostapd_iface *other;
-		bool mld_ap = false;
+		bool other_iface_6g;
 
 		other = iface->interfaces->iface[j];
 		if (other == iface || !other || !other->conf)
 			continue;
 
-#ifdef CONFIG_IEEE80211BE
-		if (hapd->conf->mld_ap && other->bss[0]->conf->mld_ap &&
-		    hostapd_is_ml_partner(hapd, other->bss[0]))
-			mld_ap = true;
-#endif /* CONFIG_IEEE80211BE */
+		other_iface_6g = is_6ghz_op_class(other->conf->op_class);
 
-		if (is_6g == is_6ghz_op_class(other->conf->op_class) &&
-		    !mld_ap)
+		if (is_6g == other_iface_6g && !hapd_mld)
 			continue;
 
 		for (i = 0; i < other->num_bss; i++) {
+#ifdef CONFIG_IEEE80211BE
+			bool mld_ap = false;
+			if (hapd_mld && other->bss[i]->conf->mld_ap &&
+			    hostapd_is_ml_partner(hapd, other->bss[i]))
+				mld_ap = true;
+
+			if (is_6g == other_iface_6g && !mld_ap)
+				continue;
+#endif /* CONFIG_IEEE80211BE */
+
 			if (other->bss[i] && other->bss[i]->started)
 				__ieee802_11_set_beacon(other->bss[i]);
 		}
--- a/src/ap/drv_callbacks.c
+++ b/src/ap/drv_callbacks.c
@@ -1444,13 +1444,13 @@ switch_link_hapd(struct hostapd_data *ha
 	return hapd;
 }
 
-
 #define HAPD_BROADCAST ((struct hostapd_data *) -1)
 
 static struct hostapd_data * get_hapd_bssid(struct hostapd_iface *iface,
-					    const u8 *bssid)
+					    const u8 *bssid, int link_id)
 {
-	size_t i;
+	size_t i, j;
+	struct hostapd_data *hapd, *p_hapd;
 
 	if (bssid == NULL)
 		return NULL;
@@ -1459,8 +1459,26 @@ static struct hostapd_data * get_hapd_bs
 		return HAPD_BROADCAST;
 
 	for (i = 0; i < iface->num_bss; i++) {
-		if (os_memcmp(bssid, iface->bss[i]->own_addr, ETH_ALEN) == 0)
-			return iface->bss[i];
+		hapd = iface->bss[i];
+#ifdef CONFIG_IEEE80211BE
+		if (os_memcmp(bssid, hapd->own_addr, ETH_ALEN) == 0 ||
+		    (hapd->conf->mld_ap &&
+		     os_memcmp(bssid, hapd->mld->mld_addr, ETH_ALEN) == 0 &&
+		     link_id == hapd->mld_link_id))
+			return hapd;
+		else if (hapd->conf->mld_ap)
+			for_each_partner_bss(p_hapd, hapd) {
+				if (p_hapd == hapd)
+					continue;
+
+				if (os_memcmp(bssid, p_hapd->own_addr, ETH_ALEN) == 0 ||
+				    (os_memcmp(bssid, p_hapd->mld->mld_addr, ETH_ALEN) == 0 && link_id == p_hapd->mld_link_id))
+				    	return p_hapd;
+			}
+#else
+		if (os_memcmp(bssid, hapd->own_addr, ETH_ALEN) == 0)
+			return hapd;
+#endif /*CONFIG_IEEE80211BE */
 	}
 
 	return NULL;
@@ -1471,7 +1489,7 @@ static void hostapd_rx_from_unknown_sta(
 					const u8 *bssid, const u8 *addr,
 					int wds)
 {
-	hapd = get_hapd_bssid(hapd->iface, bssid);
+	hapd = get_hapd_bssid(hapd->iface, bssid, -1);
 	if (hapd == NULL || hapd == HAPD_BROADCAST)
 		return;
 
@@ -1486,8 +1504,8 @@ static int hostapd_mgmt_rx(struct hostap
 	const u8 *bssid;
 	struct hostapd_frame_info fi;
 	int ret;
-	bool is_mld = false;
 
+	hapd = rx_mgmt->ctx ? rx_mgmt->ctx : hapd;
 	hapd = switch_link_hapd(hapd, rx_mgmt->link_id);
 	iface = hapd->iface;
 
@@ -1511,14 +1529,7 @@ static int hostapd_mgmt_rx(struct hostap
 	if (bssid == NULL)
 		return 0;
 
-#ifdef CONFIG_IEEE80211BE
-	if (hapd->conf->mld_ap &&
-	    os_memcmp(hapd->mld->mld_addr, bssid, ETH_ALEN) == 0)
-		is_mld = true;
-#endif /* CONFIG_IEEE80211BE */
-
-	if (!is_mld)
-		hapd = get_hapd_bssid(iface, bssid);
+	hapd = get_hapd_bssid(iface, bssid, rx_mgmt->link_id);
 
 	if (!hapd) {
 		u16 fc = le_to_host16(hdr->frame_control);
@@ -1570,34 +1581,21 @@ static void hostapd_mgmt_tx_cb(struct ho
 	struct ieee80211_hdr *hdr;
 	struct hostapd_data *orig_hapd, *tmp_hapd;
 
-#ifdef CONFIG_IEEE80211BE
-	if (hapd->conf->mld_ap && link_id != -1) {
-		tmp_hapd = hostapd_mld_get_link_bss(hapd, link_id);
-		if (tmp_hapd)
-			hapd = tmp_hapd;
-	}
-#endif /* CONFIG_IEEE80211BE */
 	orig_hapd = hapd;
 
+	hapd = switch_link_hapd(hapd, link_id);
 	hdr = (struct ieee80211_hdr *) buf;
-	tmp_hapd = get_hapd_bssid(hapd->iface, get_hdr_bssid(hdr, len));
-	if (tmp_hapd) {
+	tmp_hapd = get_hapd_bssid(hapd->iface, get_hdr_bssid(hdr, len), link_id);
+	if (tmp_hapd)
 		hapd = tmp_hapd;
-#ifdef CONFIG_IEEE80211BE
-	} else if (hapd->conf->mld_ap &&
-		   os_memcmp(hapd->mld->mld_addr, get_hdr_bssid(hdr, len),
-			     ETH_ALEN) == 0) {
-		/* AP MLD address match - use hapd pointer as-is */
-#endif /* CONFIG_IEEE80211BE */
-	} else {
+	else
 		return;
-	}
 
 	if (hapd == HAPD_BROADCAST) {
 		if (stype != WLAN_FC_STYPE_ACTION || len <= 25 ||
 		    buf[24] != WLAN_ACTION_PUBLIC)
 			return;
-		hapd = get_hapd_bssid(orig_hapd->iface, hdr->addr2);
+		hapd = get_hapd_bssid(orig_hapd->iface, hdr->addr2, link_id);
 		if (!hapd || hapd == HAPD_BROADCAST)
 			return;
 		/*
@@ -1633,16 +1631,56 @@ static int hostapd_event_new_sta(struct
 }
 
 
+static struct hostapd_data * hostapd_find_sta(struct hostapd_iface *iface, const u8 *src, struct sta_info **sta_ret)
+{
+	int j;
+	struct hostapd_data *hapd;
+	struct sta_info *sta;
+
+	for (j = 0; j < iface->num_bss; j++) {
+		hapd = iface->bss[j];
+		sta = ap_get_sta(hapd, src);
+		if (sta && sta->flags & WLAN_STA_ASSOC && sta->wpa_sm) {
+			if (sta_ret)
+				*sta_ret = sta;
+			return hapd;
+		}
+	}
+	return NULL;
+}
+
 static struct hostapd_data * hostapd_find_by_sta(struct hostapd_iface *iface,
-						 const u8 *src)
+						 const u8 *src, struct sta_info **sta_ret)
 {
 	struct sta_info *sta;
+	struct hostapd_data *hapd, *p_hapd, *tmp_hapd;
 	unsigned int j;
 
+	if (sta_ret)
+		*sta_ret = NULL;
 	for (j = 0; j < iface->num_bss; j++) {
-		sta = ap_get_sta(iface->bss[j], src);
-		if (sta && sta->flags & WLAN_STA_ASSOC)
-			return iface->bss[j];
+		hapd = iface->bss[j];
+		sta = ap_get_sta(hapd, src);
+		if (sta && sta->flags & WLAN_STA_ASSOC) {
+			if (sta_ret)
+				*sta_ret = sta;
+			return hapd;
+		}
+#ifdef CONFIG_IEEE80211BE
+		else if (hapd->conf->mld_ap) {
+			for_each_partner_bss(p_hapd, hapd) {
+				if (p_hapd == hapd)
+					continue;
+
+				sta = ap_get_sta(p_hapd, src);
+				if (sta && sta->flags & WLAN_STA_ASSOC && sta->wpa_sm) {
+					if (sta_ret)
+						*sta_ret = sta;
+					return p_hapd;
+				}
+			}
+		}
+#endif /* CONFIG_IEEE80211BE */
 	}
 
 	return NULL;
@@ -1656,42 +1694,8 @@ static void hostapd_event_eapol_rx(struc
 {
 	struct hostapd_data *orig_hapd = hapd;
 
-#ifdef CONFIG_IEEE80211BE
-	if (link_id != -1) {
-		struct hostapd_data *h_hapd;
-
-		hapd = switch_link_hapd(hapd, link_id);
-		h_hapd = hostapd_find_by_sta(hapd->iface, src);
-		if (!h_hapd)
-			h_hapd = hostapd_find_by_sta(orig_hapd->iface, src);
-		if (h_hapd)
-			hapd = h_hapd;
-	} else if (hapd->conf->mld_ap) {
-		unsigned int i;
-
-		/* Search for STA on other MLO BSSs */
-		for (i = 0; i < hapd->iface->interfaces->count; i++) {
-			struct hostapd_iface *h =
-				hapd->iface->interfaces->iface[i];
-			struct hostapd_data *h_hapd = h->bss[0];
-			struct hostapd_bss_config *hconf = h_hapd->conf;
-
-			if (!hconf->mld_ap ||
-			    !hostapd_is_ml_partner(hapd, h_hapd))
-				continue;
-
-			h_hapd = hostapd_find_by_sta(h, src);
-			if (h_hapd) {
-				hapd = h_hapd;
-				break;
-			}
-		}
-	} else {
-		hapd = hostapd_find_by_sta(hapd->iface, src);
-	}
-#else /* CONFIG_IEEE80211BE */
-	hapd = hostapd_find_by_sta(hapd->iface, src);
-#endif /* CONFIG_IEEE80211BE */
+	hapd = switch_link_hapd(hapd, link_id);
+	hapd = hostapd_find_by_sta(hapd->iface, src, NULL);
 
 	if (!hapd) {
 		/* WLAN cases need to have an existing association, but non-WLAN
@@ -2067,6 +2071,26 @@ int hostapd_mld_update_bridge_conn(char
 	return 0;
 }
 
+static void hostapd_eapol_tx_status(struct hostapd_data *hapd, const u8 *dst,
+			     const u8 *data, size_t len, int ack,
+			     int link_id)
+{
+	struct sta_info *sta = NULL;
+	struct hostapd_iface *iface = hapd->iface;
+
+	hapd = switch_link_hapd(hapd, link_id);
+	hapd = hostapd_find_by_sta(hapd->iface, dst, &sta);
+
+	if (sta == NULL) {
+		wpa_printf(MSG_DEBUG, "Ignore TX status for Data frame to STA "
+			   MACSTR " that is not currently associated",
+			   MAC2STR(dst));
+		return;
+	}
+
+	ieee802_1x_eapol_tx_status(hapd, sta, data, len, ack);
+}
+
 void hostapd_wpa_event(void *ctx, enum wpa_event_type event,
 		       union wpa_event_data *data)
 {
@@ -2125,11 +2149,11 @@ void hostapd_wpa_event(void *ctx, enum w
 		}
 		break;
 	case EVENT_EAPOL_TX_STATUS:
-		hapd = switch_link_hapd(hapd, data->eapol_tx_status.link_id);
 		hostapd_eapol_tx_status(hapd, data->eapol_tx_status.dst,
 					data->eapol_tx_status.data,
 					data->eapol_tx_status.data_len,
-					data->eapol_tx_status.ack);
+					data->eapol_tx_status.ack,
+					data->eapol_tx_status.link_id);
 		break;
 	case EVENT_DRIVER_CLIENT_POLL_OK:
 		hostapd_client_poll_ok(hapd, data->client_poll.addr);
--- a/src/ap/hostapd.c
+++ b/src/ap/hostapd.c
@@ -477,17 +477,19 @@ static int hostapd_broadcast_wep_set(str
 static void hostapd_clear_drv_priv(struct hostapd_data *hapd)
 {
 #ifdef CONFIG_IEEE80211BE
-	unsigned int i;
+	unsigned int i, j;
 
-	for (i = 0; i < hapd->iface->interfaces->count; i++) {
+	for (i = 0; i < hapd->iface->interfaces->count && hapd->conf->mld_ap; i++) {
 		struct hostapd_iface *iface = hapd->iface->interfaces->iface[i];
 
-		if (hapd->iface == iface)
+		if (!iface || hapd->iface == iface)
 			continue;
 
-		if (iface->bss && iface->bss[0] &&
-		    hostapd_mld_get_first_bss(iface->bss[0]) == hapd)
-			iface->bss[0]->drv_priv = NULL;
+		for (j = 0; j < iface->num_bss; j++) {
+			if (iface->bss && iface->bss[j] &&
+			    hostapd_mld_get_first_bss(iface->bss[j]) == hapd)
+				iface->bss[j]->drv_priv = NULL;
+		}
 	}
 #endif /* CONFIG_IEEE80211BE */
 
@@ -1306,6 +1308,7 @@ static int hostapd_setup_bss(struct host
 	char force_ifname[IFNAMSIZ];
 	u8 if_addr[ETH_ALEN];
 	int flush_old_stations = 1;
+	struct hostapd_data *h_hapd = NULL;
 
 #ifdef CONFIG_IEEE80211BE
 	if (!hostapd_mld_is_first_bss(hapd))
@@ -1354,6 +1357,19 @@ static int hostapd_setup_bss(struct host
 			} while (mac_in_conf(hapd->iconf, hapd->own_addr));
 		}
 
+		if (conf->mld_ap) {
+			h_hapd = hostapd_mld_get_first_bss(hapd);
+
+			if (h_hapd) {
+				hapd->drv_priv = h_hapd->drv_priv;
+				hapd->mld_link_id = hapd->mld->next_link_id++;
+				hostapd_mld_add_link(hapd);
+				wpa_printf(MSG_DEBUG, "TSetup of non first link BSS(link id %d) of MLD %s",
+					   hapd->mld_link_id, hapd->conf->iface);
+				goto setup_mld;
+			}
+		}
+
 		hapd->interface_added = 1;
 		if (hostapd_if_add(hapd->iface->bss[0], WPA_IF_AP_BSS,
 				   conf->iface, addr, hapd,
@@ -1368,8 +1384,18 @@ static int hostapd_setup_bss(struct host
 
 		if (!addr)
 			os_memcpy(hapd->own_addr, if_addr, ETH_ALEN);
+
+		if (hapd->conf->mld_ap) {
+			wpa_printf(MSG_DEBUG, "TSetup of first link BSS of MLD %s",
+				   hapd->conf->iface);
+			os_memcpy(hapd->mld->mld_addr, hapd->own_addr, ETH_ALEN);
+			random_mac_addr_keep_oui(hapd->own_addr);
+			hapd->mld_link_id = hapd->mld->next_link_id++;
+			hostapd_mld_add_link(hapd);
+		}
 	}
 
+setup_mld:
 	 if (is_zero_ether_addr(conf->bssid))
 		 os_memcpy(conf->bssid, hapd->own_addr, ETH_ALEN);
 
@@ -1610,6 +1636,19 @@ static int hostapd_setup_bss(struct host
 		return -1;
 	}
 
+#ifdef CONFIG_IEEE80211BE
+	if (hapd->conf->mld_ap && !first) {
+		wpa_printf(MSG_DEBUG,
+			   "MLD: Set link_id=%u, mld_addr=" MACSTR
+			   ", own_addr=" MACSTR,
+			   hapd->mld_link_id, MAC2STR(hapd->mld->mld_addr),
+			   MAC2STR(hapd->own_addr));
+
+		hostapd_drv_link_add(hapd, hapd->mld_link_id,
+				     hapd->own_addr);
+	}
+#endif /* CONFIG_IEEE80211BE */
+
 	if (start_beacon && hostapd_start_beacon(hapd, flush_old_stations) < 0)
 		return -1;
 
@@ -2812,6 +2851,10 @@ static void hostapd_bss_deinit(struct ho
 		hapd->rad_attr_db = NULL;
 	}
 #endif /* CONFIG_SQLITE */
+#ifdef CONFIG_IEEE80211BE
+	if (hapd->conf->mld_ap)
+		hostapd_mld_remove_link(hapd);
+#endif
 	hostapd_cleanup(hapd);
 }
 
@@ -4671,18 +4714,30 @@ hostapd_interface_update_fils_ubpr(struc
 struct hostapd_data * hostapd_mld_get_link_bss(struct hostapd_data *hapd,
 					       u8 link_id)
 {
-	unsigned int i;
+	unsigned int i, j;
+	struct hostapd_iface *iface;
+	struct hostapd_data *bss;
+	struct hostapd_bss_config *conf;
 
 	for (i = 0; i < hapd->iface->interfaces->count; i++) {
-		struct hostapd_iface *h = hapd->iface->interfaces->iface[i];
-		struct hostapd_data *h_hapd = h->bss[0];
-		struct hostapd_bss_config *hconf = h_hapd->conf;
+		iface = hapd->iface->interfaces->iface[i];
 
-		if (!hconf->mld_ap || !hostapd_is_ml_partner(hapd, h_hapd))
+		if (!iface)
 			continue;
 
-		if (h_hapd->mld_link_id == link_id)
-			return h_hapd;
+		for (j = 0; j < iface->num_bss; j++) {
+			bss = iface->bss[j];
+			conf = bss->conf;
+
+			if (!conf->mld_ap || !hostapd_is_ml_partner(hapd, bss))
+				continue;
+
+			if (!bss->drv_priv)
+				continue;
+
+			if (bss->mld_link_id == link_id)
+				return bss;
+		}
 	}
 
 	return NULL;
@@ -4801,4 +4856,5 @@ struct hostapd_data * hostapd_mld_get_fi
 
 	return mld->fbss;
 }
+
 #endif /* CONFIG_IEEE80211BE */
--- a/src/ap/hostapd.h
+++ b/src/ap/hostapd.h
@@ -720,6 +720,13 @@ struct hostapd_iface {
 	bool is_no_ir;
 };
 
+/* Iterate over all affiliated links in MLD to which @self belongs.
+ * Caller should ensure to have a check in order to ignore the link
+ * which called this.
+ */
+#define for_each_partner_bss(partner, self) \
+	dl_list_for_each(partner, &self->mld->links, struct hostapd_data, link)
+
 /* hostapd.c */
 int hostapd_for_each_interface(struct hapd_interfaces *interfaces,
 			       int (*cb)(struct hostapd_iface *iface,
--- a/src/ap/ieee802_11.c
+++ b/src/ap/ieee802_11.c
@@ -410,14 +410,7 @@ static int send_auth_reply(struct hostap
 	struct wpabuf *ml_resp = NULL;
 
 #ifdef CONFIG_IEEE80211BE
-	/*
-	 * Once a non-AP MLD is added to the driver, the addressing should use
-	 * the MLD MAC address. Thus, use the MLD address instead of translating
-	 * the addresses.
-	 */
 	if (hapd->conf->mld_ap && sta && sta->mld_info.mld_sta) {
-		sa = hapd->mld->mld_addr;
-
 		ml_resp = hostapd_ml_auth_resp(hapd);
 		if (!ml_resp)
 			return -1;
@@ -438,7 +431,7 @@ static int send_auth_reply(struct hostap
 					    WLAN_FC_STYPE_AUTH);
 	os_memcpy(reply->da, dst, ETH_ALEN);
 	os_memcpy(reply->sa, sa, ETH_ALEN);
-	os_memcpy(reply->bssid, bssid, ETH_ALEN);
+	os_memcpy(reply->bssid, sa, ETH_ALEN);
 
 	reply->u.auth.auth_alg = host_to_le16(auth_alg);
 	reply->u.auth.auth_transaction = host_to_le16(auth_transaction);
@@ -3233,14 +3226,9 @@ static void handle_auth(struct hostapd_d
 	bssid = mgmt->bssid;
 
 #ifdef CONFIG_IEEE80211BE
-	 /*
-	  * Once a non-AP MLD is added to the driver, the addressing should use
-	  * the MLD MAC address. It is the responsibility of the driver to
-	  * handle the translations.
-	  */
 	if (hapd->conf->mld_ap && sta && sta->mld_info.mld_sta) {
 		dst = sta->addr;
-		bssid = hapd->mld->mld_addr;
+		bssid = hapd->own_addr;
 	}
 #endif /* CONFIG_IEEE80211BE */
 
@@ -4527,24 +4515,27 @@ static void hostapd_process_assoc_ml_inf
 
 	for (i = 0; i < MAX_NUM_MLD_LINKS; i++) {
 		struct hostapd_iface *iface = NULL;
+		struct hostapd_data *bss = NULL;
 		struct mld_link_info *link = &sta->mld_info.links[i];
+		bool link_bss_found = false;
 
 		if (!link->valid)
 			continue;
 
-		for (j = 0; j < hapd->iface->interfaces->count; j++) {
-			iface = hapd->iface->interfaces->iface[j];
+		for_each_partner_bss(bss, hapd) {
+			if (bss == hapd)
+				continue;
 
-			if (hapd->iface == iface)
+			if (bss->mld_link_id != i)
 				continue;
 
-			if (iface->bss[0]->conf->mld_ap &&
-			    hostapd_is_ml_partner(hapd, iface->bss[0]) &&
-			    i == iface->bss[0]->mld_link_id)
-				break;
+			ieee80211_ml_process_link(bss, sta, link,
+						  ies, ies_len, reassoc);
+			link_bss_found = true;
+			break;
 		}
 
-		if (!iface || j == hapd->iface->interfaces->count) {
+		if (!link_bss_found) {
 			wpa_printf(MSG_DEBUG,
 				   "MLD: No link match for link_id=%u", i);
 
@@ -4554,9 +4545,6 @@ static void hostapd_process_assoc_ml_inf
 					hapd, link->status,
 					link->resp_sta_profile,
 					sizeof(link->resp_sta_profile));
-		} else {
-			ieee80211_ml_process_link(iface->bss[0], sta, link,
-						  ies, ies_len, reassoc);
 		}
 	}
 #endif /* CONFIG_IEEE80211BE */
@@ -4773,15 +4761,6 @@ static u16 send_assoc_resp(struct hostap
 			     (reassoc ? WLAN_FC_STYPE_REASSOC_RESP :
 			      WLAN_FC_STYPE_ASSOC_RESP));
 
-#ifdef CONFIG_IEEE80211BE
-	/*
-	 * Once a non-AP MLD is added to the driver, the addressing should use
-	 * MLD MAC address.
-	 */
-	if (hapd->conf->mld_ap && sta && sta->mld_info.mld_sta)
-		sa = hapd->mld->mld_addr;
-#endif /* CONFIG_IEEE80211BE */
-
 	os_memcpy(reply->da, addr, ETH_ALEN);
 	os_memcpy(reply->sa, sa, ETH_ALEN);
 	os_memcpy(reply->bssid, sa, ETH_ALEN);
@@ -5715,6 +5694,7 @@ static bool hostapd_ml_handle_disconnect
 	struct hostapd_data *assoc_hapd, *tmp_hapd;
 	struct sta_info *assoc_sta;
 	unsigned int i, link_id;
+	struct sta_info *tmp_sta;
 
 	if (!hostapd_is_mld_ap(hapd))
 		return false;
@@ -5725,46 +5705,28 @@ static bool hostapd_ml_handle_disconnect
 	 */
 	assoc_sta = hostapd_ml_get_assoc_sta(hapd, sta, &assoc_hapd);
 
-	for (link_id = 0; link_id < MAX_NUM_MLD_LINKS; link_id++) {
-		for (i = 0; i < assoc_hapd->iface->interfaces->count; i++) {
-			struct sta_info *tmp_sta;
+	for_each_partner_bss(tmp_hapd, assoc_hapd) {
+		if (tmp_hapd == assoc_hapd)
+			continue;
 
-			if (!assoc_sta->mld_info.links[link_id].valid)
-				continue;
+		if (!assoc_sta->mld_info.links[tmp_hapd->mld_link_id].valid)
+			continue;
 
-			tmp_hapd =
-				assoc_hapd->iface->interfaces->iface[i]->bss[0];
+		for (tmp_sta = tmp_hapd->sta_list; tmp_sta;
+		     tmp_sta = tmp_sta->next) {
 
-			if (!tmp_hapd->conf->mld_ap ||
-			    hostapd_is_ml_partner(assoc_hapd, tmp_hapd))
+			if (tmp_sta->mld_assoc_link_id !=
+			    assoc_sta->mld_assoc_link_id ||
+			    tmp_sta->aid != assoc_sta->aid)
 				continue;
 
-			for (tmp_sta = tmp_hapd->sta_list; tmp_sta;
-			     tmp_sta = tmp_sta->next) {
-				/*
-				 * Remove the station on which the association
-				 * was done only after all other link stations
-				 * are removed. Since there is only a single
-				 * station per struct hostapd_hapd with the
-				 * same association link simply break out from
-				 * the loop.
-				 */
-				if (tmp_sta == assoc_sta)
-					break;
-
-				if (tmp_sta->mld_assoc_link_id !=
-				    assoc_sta->mld_assoc_link_id ||
-				    tmp_sta->aid != assoc_sta->aid)
-					continue;
-
-				if (!disassoc)
-					hostapd_deauth_sta(tmp_hapd, tmp_sta,
-							   mgmt);
-				else
-					hostapd_disassoc_sta(tmp_hapd, tmp_sta,
-							     mgmt);
-				break;
-			}
+			if (!disassoc)
+				hostapd_deauth_sta(tmp_hapd, tmp_sta,
+						   mgmt);
+			else
+				hostapd_disassoc_sta(tmp_hapd, tmp_sta,
+						     mgmt);
+			break;
 		}
 	}
 
@@ -6359,37 +6321,31 @@ static void hostapd_ml_handle_assoc_cb(s
 {
 #ifdef CONFIG_IEEE80211BE
 	unsigned int i, link_id;
+	struct hostapd_data *tmp_hapd;
 
 	if (!hostapd_is_mld_ap(hapd))
 		return;
 
-	for (link_id = 0; link_id < MAX_NUM_MLD_LINKS; link_id++) {
-		struct mld_link_info *link = &sta->mld_info.links[link_id];
+	for_each_partner_bss(tmp_hapd, hapd) {
+		struct mld_link_info *link = &sta->mld_info.links[tmp_hapd->mld_link_id];
+		struct sta_info *tmp_sta;
 
-		if (!link->valid)
+		if (tmp_hapd == hapd)
 			continue;
 
-		for (i = 0; i < hapd->iface->interfaces->count; i++) {
-			struct sta_info *tmp_sta;
-			struct hostapd_data *tmp_hapd =
-				hapd->iface->interfaces->iface[i]->bss[0];
+		if (!link->valid)
+			continue;
 
-			if (!hostapd_is_ml_partner(tmp_hapd, hapd))
+		for (tmp_sta = tmp_hapd->sta_list; tmp_sta;
+		     tmp_sta = tmp_sta->next) {
+			if (tmp_sta == sta ||
+			    tmp_sta->mld_assoc_link_id !=
+			    sta->mld_assoc_link_id ||
+			    tmp_sta->aid != sta->aid)
 				continue;
 
-			for (tmp_sta = tmp_hapd->sta_list; tmp_sta;
-			     tmp_sta = tmp_sta->next) {
-				if (tmp_sta == sta ||
-				    tmp_sta->mld_assoc_link_id !=
-				    sta->mld_assoc_link_id ||
-				    tmp_sta->aid != sta->aid)
-					continue;
-
-				ieee80211_ml_link_sta_assoc_cb(tmp_hapd,
-							       tmp_sta, link,
-							       ok);
-				break;
-			}
+			ieee80211_ml_link_sta_assoc_cb(tmp_hapd, tmp_sta, link, ok);
+			break;
 		}
 	}
 #endif /* CONFIG_IEEE80211BE */
@@ -6799,34 +6755,6 @@ void hostapd_tx_status(struct hostapd_da
 	ieee802_1x_tx_status(hapd, sta, buf, len, ack);
 }
 
-
-void hostapd_eapol_tx_status(struct hostapd_data *hapd, const u8 *dst,
-			     const u8 *data, size_t len, int ack)
-{
-	struct sta_info *sta;
-	struct hostapd_iface *iface = hapd->iface;
-
-	sta = ap_get_sta(hapd, dst);
-	if (sta == NULL && iface->num_bss > 1) {
-		size_t j;
-		for (j = 0; j < iface->num_bss; j++) {
-			hapd = iface->bss[j];
-			sta = ap_get_sta(hapd, dst);
-			if (sta)
-				break;
-		}
-	}
-	if (sta == NULL || !(sta->flags & WLAN_STA_ASSOC)) {
-		wpa_printf(MSG_DEBUG, "Ignore TX status for Data frame to STA "
-			   MACSTR " that is not currently associated",
-			   MAC2STR(dst));
-		return;
-	}
-
-	ieee802_1x_eapol_tx_status(hapd, sta, data, len, ack);
-}
-
-
 void hostapd_client_poll_ok(struct hostapd_data *hapd, const u8 *addr)
 {
 	struct sta_info *sta;
--- a/src/ap/ieee802_11.h
+++ b/src/ap/ieee802_11.h
@@ -135,8 +135,6 @@ int hostapd_get_he_twt_responder(struct
 u8 * hostapd_eid_cca(struct hostapd_data *hapd, u8 *eid);
 void hostapd_tx_status(struct hostapd_data *hapd, const u8 *addr,
 		       const u8 *buf, size_t len, int ack);
-void hostapd_eapol_tx_status(struct hostapd_data *hapd, const u8 *dst,
-			     const u8 *data, size_t len, int ack);
 void ieee802_11_rx_from_unknown(struct hostapd_data *hapd, const u8 *src,
 				int wds);
 u8 * hostapd_eid_assoc_comeback_time(struct hostapd_data *hapd,
--- a/src/ap/ieee802_11_eht.c
+++ b/src/ap/ieee802_11_eht.c
@@ -977,31 +977,20 @@ static int hostapd_mld_validate_assoc_in
 	for (link_id = 0; link_id < MAX_NUM_MLD_LINKS; link_id++) {
 		struct hostapd_data *other_hapd;
 
-		if (!info->links[link_id].valid)
+		if (!info->links[link_id].valid || link_id == hapd->mld_link_id)
 			continue;
 
-		for (i = 0; i < hapd->iface->interfaces->count; i++) {
-			other_hapd = hapd->iface->interfaces->iface[i]->bss[0];
+		other_hapd = hostapd_mld_get_link_bss(hapd, link_id);
 
-			if (hapd == other_hapd)
-				continue;
-
-			if (hostapd_is_ml_partner(hapd, other_hapd) &&
-			    link_id == other_hapd->mld_link_id)
-				break;
-		}
-
-		if (i == hapd->iface->interfaces->count &&
-		    link_id != hapd->mld_link_id) {
+		if (!other_hapd) {
 			wpa_printf(MSG_DEBUG, "MLD: Invalid link ID=%u",
 				   link_id);
 			return -1;
 		}
 
-		if (i < hapd->iface->interfaces->count)
-			os_memcpy(info->links[link_id].local_addr,
-				  other_hapd->own_addr,
-				  ETH_ALEN);
+		os_memcpy(info->links[link_id].local_addr,
+			  other_hapd->own_addr,
+			  ETH_ALEN);
 	}
 
 	return 0;
--- a/src/ap/wpa_auth_glue.c
+++ b/src/ap/wpa_auth_glue.c
@@ -1511,25 +1511,33 @@ static int hostapd_wpa_auth_get_ml_rsn_i
 
 	for (i = 0; i < info->n_mld_links; i++) {
 		unsigned int link_id = info->links[i].link_id;
+		struct hostapd_data *bss = NULL;
+		bool link_bss_found = false;
 
 		wpa_printf(MSG_DEBUG,
 			   "WPA_AUTH: MLD: Get link RSN CB: link_id=%u",
 			   link_id);
 
-		for (j = 0; j < hapd->iface->interfaces->count; j++) {
-			struct hostapd_iface *iface =
-				hapd->iface->interfaces->iface[j];
+		if (hapd->mld_link_id == link_id) {
+			wpa_auth_ml_get_rsn_info(hapd->wpa_auth,
+						 &info->links[i]);
+			continue;
+		}
 
-			if (!hostapd_is_ml_partner(hapd, iface->bss[0]) ||
-			    link_id != iface->bss[0]->mld_link_id)
+		for_each_partner_bss(bss, hapd) {
+			if (bss == hapd)
 				continue;
 
-			wpa_auth_ml_get_rsn_info(iface->bss[0]->wpa_auth,
+			if (bss->mld_link_id != link_id)
+				continue;
+
+			wpa_auth_ml_get_rsn_info(bss->wpa_auth,
 						 &info->links[i]);
+			link_bss_found = true;
 			break;
 		}
 
-		if (j == hapd->iface->interfaces->count)
+		if (!link_bss_found)
 			wpa_printf(MSG_DEBUG,
 				   "WPA_AUTH: MLD: link=%u not found", link_id);
 	}
@@ -1552,27 +1560,37 @@ static int hostapd_wpa_auth_get_ml_key_i
 
 	for (i = 0; i < info->n_mld_links; i++) {
 		u8 link_id = info->links[i].link_id;
+		struct hostapd_data *bss = NULL;
+		bool link_bss_found = false;
 
 		wpa_printf(MSG_DEBUG,
 			   "WPA_AUTH: MLD: Get link info CB: link_id=%u",
 			   link_id);
 
-		for (j = 0; j < hapd->iface->interfaces->count; j++) {
-			struct hostapd_iface *iface =
-				hapd->iface->interfaces->iface[j];
+		if (hapd->mld_link_id == link_id) {
+			wpa_auth_ml_get_key_info(hapd->wpa_auth,
+						 &info->links[i],
+						 info->mgmt_frame_prot,
+						 info->beacon_prot);
+			continue;
+		}
+
+		for_each_partner_bss(bss, hapd) {
+			if (bss == hapd)
+				continue;
 
-			if (!hostapd_is_ml_partner(hapd, iface->bss[0]) ||
-			    link_id != iface->bss[0]->mld_link_id)
+			if (bss->mld_link_id != link_id)
 				continue;
 
-			wpa_auth_ml_get_key_info(iface->bss[0]->wpa_auth,
+			wpa_auth_ml_get_key_info(bss->wpa_auth,
 						 &info->links[i],
 						 info->mgmt_frame_prot,
 						 info->beacon_prot);
+			link_bss_found = true;
 			break;
 		}
 
-		if (j == hapd->iface->interfaces->count)
+		if (!link_bss_found)
 			wpa_printf(MSG_DEBUG,
 				   "WPA_AUTH: MLD: link=%u not found", link_id);
 	}
--- a/src/drivers/driver.h
+++ b/src/drivers/driver.h
@@ -6337,6 +6337,14 @@ union wpa_event_data {
 		void *drv_priv;
 
 		/**
+		 * ctx - Pointer to store ctx of private BSS information
+		 *
+		 * If not set to NULL, this is used for forwarding the packet to right
+		 * link BSS of ML BSS.
+		 */
+		 void *ctx;
+
+		/**
 		 * freq - Frequency (in MHz) on which the frame was received
 		 */
 		int freq;
--- a/src/drivers/driver_nl80211.c
+++ b/src/drivers/driver_nl80211.c
@@ -8218,7 +8218,7 @@ static int i802_sta_deauth(void *priv, c
 	return wpa_driver_nl80211_send_mlme(bss, (u8 *) &mgmt,
 					    IEEE80211_HDRLEN +
 					    sizeof(mgmt.u.deauth), 0, 0, 0, 0,
-					    0, NULL, 0, 0, -1);
+					    0, NULL, 0, 0, link_id);
 }
 
 
@@ -13668,7 +13668,7 @@ static int nl80211_link_add(void *priv,
 		idx = bss->n_links;
 	}
 
-	msg = nl80211_drv_msg(drv, 0, NL80211_CMD_ADD_LINK);
+	msg = nl80211_bss_msg(bss, 0, NL80211_CMD_ADD_LINK);
 	if (!msg ||
 	    nla_put_u8(msg, NL80211_ATTR_MLO_LINK_ID, link_id) ||
 	    nla_put(msg, NL80211_ATTR_MAC, ETH_ALEN, addr)) {
@@ -13688,7 +13688,7 @@ static int nl80211_link_add(void *priv,
 
 	bss->n_links = idx + 1;
 
-	wpa_printf(MSG_DEBUG, "nl80211: MLD: n_links=%zu", bss->n_links);
+	wpa_printf(MSG_DEBUG, "nl80211: MLD: n_links=%zu ifidx %d", bss->n_links, drv->ifindex);
 	return 0;
 }
 
--- a/src/drivers/driver_nl80211_event.c
+++ b/src/drivers/driver_nl80211_event.c
@@ -22,7 +22,7 @@
 
 
 static void
-nl80211_control_port_frame_tx_status(struct wpa_driver_nl80211_data *drv,
+nl80211_control_port_frame_tx_status(struct i802_bss *bss,
 				     const u8 *frame, size_t len,
 				     struct nlattr *ack, struct nlattr *cookie);
 
@@ -1373,13 +1373,14 @@ static void mlme_event_mgmt(struct i802_
 	event.rx_mgmt.frame_len = len;
 	event.rx_mgmt.ssi_signal = ssi_signal;
 	event.rx_mgmt.drv_priv = bss;
+	event.rx_mgmt.ctx = bss->ctx;
 	event.rx_mgmt.link_id = link_id;
 
 	wpa_supplicant_event(drv->ctx, EVENT_RX_MGMT, &event);
 }
 
 
-static void mlme_event_mgmt_tx_status(struct wpa_driver_nl80211_data *drv,
+static void mlme_event_mgmt_tx_status(struct i802_bss *bss,
 				      struct nlattr *cookie, const u8 *frame,
 				      size_t len, struct nlattr *ack)
 {
@@ -1387,6 +1388,7 @@ static void mlme_event_mgmt_tx_status(st
 	const struct ieee80211_hdr *hdr = (const struct ieee80211_hdr *) frame;
 	u16 fc = le_to_host16(hdr->frame_control);
 	u64 cookie_val = 0;
+	struct wpa_driver_nl80211_data *drv = bss->drv;
 
 	if (cookie)
 		cookie_val = nla_get_u64(cookie);
@@ -1403,7 +1405,7 @@ static void mlme_event_mgmt_tx_status(st
 	    WPA_GET_BE16(frame + 2 * ETH_ALEN) == ETH_P_PAE) {
 		wpa_printf(MSG_DEBUG,
 			   "nl80211: Work around misdelivered control port TX status for EAPOL");
-		nl80211_control_port_frame_tx_status(drv, frame, len, ack,
+		nl80211_control_port_frame_tx_status(bss, frame, len, ack,
 						     cookie);
 		return;
 	}
@@ -1439,7 +1441,7 @@ static void mlme_event_mgmt_tx_status(st
 	event.tx_status.ack = ack != NULL;
 	event.tx_status.link_id = cookie_val == drv->send_frame_cookie ?
 		drv->send_frame_link_id : NL80211_DRV_LINK_ID_NA;
-	wpa_supplicant_event(drv->ctx, EVENT_TX_STATUS, &event);
+	wpa_supplicant_event(bss->ctx, EVENT_TX_STATUS, &event);
 }
 
 
@@ -1750,7 +1752,7 @@ static void mlme_event(struct i802_bss *
 				nla_len(frame), link_id);
 		break;
 	case NL80211_CMD_FRAME_TX_STATUS:
-		mlme_event_mgmt_tx_status(drv, cookie, nla_data(frame),
+		mlme_event_mgmt_tx_status(bss, cookie, nla_data(frame),
 					  nla_len(frame), ack);
 		break;
 	case NL80211_CMD_UNPROT_DEAUTHENTICATE:
@@ -3607,7 +3609,7 @@ static void nl80211_update_muedca_params
 	wpa_supplicant_event(drv->ctx, EVENT_UPDATE_MUEDCA_PARAMS, &ed);
 }
 
-static void nl80211_control_port_frame(struct wpa_driver_nl80211_data *drv,
+static void nl80211_control_port_frame(struct i802_bss *bss,
 				       struct nlattr **tb)
 {
 	u8 *src_addr;
@@ -3637,7 +3639,7 @@ static void nl80211_control_port_frame(s
 			   MAC2STR(src_addr));
 		break;
 	case ETH_P_PAE:
-		drv_event_eapol_rx2(drv->ctx, src_addr,
+		drv_event_eapol_rx2(bss->ctx, src_addr,
 				    nla_data(tb[NL80211_ATTR_FRAME]),
 				    nla_len(tb[NL80211_ATTR_FRAME]),
 				    encrypted, link_id);
@@ -3653,11 +3655,12 @@ static void nl80211_control_port_frame(s
 
 
 static void
-nl80211_control_port_frame_tx_status(struct wpa_driver_nl80211_data *drv,
+nl80211_control_port_frame_tx_status(struct i802_bss *bss,
 				     const u8 *frame, size_t len,
 				     struct nlattr *ack, struct nlattr *cookie)
 {
 	union wpa_event_data event;
+	struct wpa_driver_nl80211_data *drv = bss->drv;
 
 	if (!cookie || len < ETH_HLEN)
 		return;
@@ -3675,7 +3678,7 @@ nl80211_control_port_frame_tx_status(str
 		nla_get_u64(cookie) == drv->eapol_tx_cookie ?
 		drv->eapol_tx_link_id : NL80211_DRV_LINK_ID_NA;
 
-	wpa_supplicant_event(drv->ctx, EVENT_EAPOL_TX_STATUS, &event);
+	wpa_supplicant_event(bss->ctx, EVENT_EAPOL_TX_STATUS, &event);
 }
 
 
@@ -4028,7 +4031,7 @@ static void do_process_drv_event(struct
 	case NL80211_CMD_CONTROL_PORT_FRAME_TX_STATUS:
 		if (!frame)
 			break;
-		nl80211_control_port_frame_tx_status(drv,
+		nl80211_control_port_frame_tx_status(bss,
 						     nla_data(frame),
 						     nla_len(frame),
 						     tb[NL80211_ATTR_ACK],
@@ -4152,7 +4155,7 @@ int process_bss_event(struct nl_msg *msg
 		nl80211_external_auth(bss->drv, tb);
 		break;
 	case NL80211_CMD_CONTROL_PORT_FRAME:
-		nl80211_control_port_frame(bss->drv, tb);
+		nl80211_control_port_frame(bss, tb);
 		break;
 	default:
 		wpa_printf(MSG_DEBUG, "nl80211: Ignored unknown event "