diff mbox series

wifi: mac80211: fix initialization of rx->link and rx->link_sta

Message ID 20221215190503.79904-1-nbd@nbd.name
State Superseded
Headers show
Series wifi: mac80211: fix initialization of rx->link and rx->link_sta | expand

Commit Message

Felix Fietkau Dec. 15, 2022, 7:05 p.m. UTC 
There are some codepaths that do not initialize rx->link_sta properly. This
causes a crash in places which assume that rx->link_sta is valid if rx->sta
is valid.
One known instance is triggered by __ieee80211_rx_h_amsdu being called from
fast-rx.

Since the initialization of rx->link and rx->link_sta is rather convoluted
and duplicated in many places, clean it up by using a helper function to
set it.

Fixes: ccdde7c74ffd ("wifi: mac80211: properly implement MLO key handling")
Fixes: b320d6c456ff ("wifi: mac80211: use correct rx link_sta instead of default")
Signed-off-by: Felix Fietkau <nbd@nbd.name>
---
 net/mac80211/rx.c | 218 ++++++++++++++++++++--------------------------
 1 file changed, 94 insertions(+), 124 deletions(-)

Comments

kernel test robot Dec. 15, 2022, 10:31 p.m. UTC | #1 
Hi Felix,

I love your patch! Perhaps something to improve:

[auto build test WARNING on wireless-next/main]
[also build test WARNING on wireless/main linus/master v6.1 next-20221215]
[If your patch is applied to the wrong git tree, kindly drop us a note.
And when submitting patch, we suggest to use '--base' as documented in
https://git-scm.com/docs/git-format-patch#_base_tree_information]

url:    https://github.com/intel-lab-lkp/linux/commits/Felix-Fietkau/wifi-mac80211-fix-initialization-of-rx-link-and-rx-link_sta/20221216-030758
base:   https://git.kernel.org/pub/scm/linux/kernel/git/wireless/wireless-next.git main
patch link:    https://lore.kernel.org/r/20221215190503.79904-1-nbd%40nbd.name
patch subject: [PATCH] wifi: mac80211: fix initialization of rx->link and rx->link_sta
config: riscv-randconfig-r014-20221215
compiler: clang version 16.0.0 (https://github.com/llvm/llvm-project 98b13979fb05f3ed288a900deb843e7b27589e58)
reproduce (this is a W=1 build):
        wget https://raw.githubusercontent.com/intel/lkp-tests/master/sbin/make.cross -O ~/bin/make.cross
        chmod +x ~/bin/make.cross
        # install riscv cross compiling tool for clang build
        # apt-get install binutils-riscv64-linux-gnu
        # https://github.com/intel-lab-lkp/linux/commit/833b68b914bcc2b2804ba4673e408fbf8915b604
        git remote add linux-review https://github.com/intel-lab-lkp/linux
        git fetch --no-tags linux-review Felix-Fietkau/wifi-mac80211-fix-initialization-of-rx-link-and-rx-link_sta/20221216-030758
        git checkout 833b68b914bcc2b2804ba4673e408fbf8915b604
        # save the config file
        mkdir build_dir && cp config build_dir/.config
        COMPILER_INSTALL_PATH=$HOME/0day COMPILER=clang make.cross W=1 O=build_dir ARCH=riscv SHELL=/bin/bash net/mac80211/

If you fix the issue, kindly add following tag where applicable
| Reported-by: kernel test robot <lkp@intel.com>

All warnings (new ones prefixed by >>):

>> net/mac80211/rx.c:4841:24: warning: variable 'link' is uninitialized when used here [-Wuninitialized]
                   if (ether_addr_equal(link->conf->addr, hdr->addr1))
                                        ^~~~
   net/mac80211/rx.c:4798:34: note: initialize the variable 'link' to silence this warning
           struct ieee80211_link_data *link;
                                           ^
                                            = NULL
   1 warning generated.


vim +/link +4841 net/mac80211/rx.c

49ddf8e6e2347c Johannes Berg 2016-03-31  4785  
4406c376895608 Johannes Berg 2010-09-24  4786  /*
4406c376895608 Johannes Berg 2010-09-24  4787   * This function returns whether or not the SKB
4406c376895608 Johannes Berg 2010-09-24  4788   * was destined for RX processing or not, which,
4406c376895608 Johannes Berg 2010-09-24  4789   * if consume is true, is equivalent to whether
4406c376895608 Johannes Berg 2010-09-24  4790   * or not the skb was consumed.
4406c376895608 Johannes Berg 2010-09-24  4791   */
4406c376895608 Johannes Berg 2010-09-24  4792  static bool ieee80211_prepare_and_rx_handle(struct ieee80211_rx_data *rx,
4406c376895608 Johannes Berg 2010-09-24  4793  					    struct sk_buff *skb, bool consume)
4406c376895608 Johannes Berg 2010-09-24  4794  {
4406c376895608 Johannes Berg 2010-09-24  4795  	struct ieee80211_local *local = rx->local;
4406c376895608 Johannes Berg 2010-09-24  4796  	struct ieee80211_sub_if_data *sdata = rx->sdata;
42fb9148c07800 Johannes Berg 2022-07-14  4797  	struct ieee80211_hdr *hdr = (void *)skb->data;
42fb9148c07800 Johannes Berg 2022-07-14  4798  	struct ieee80211_link_data *link;
4406c376895608 Johannes Berg 2010-09-24  4799  
4406c376895608 Johannes Berg 2010-09-24  4800  	rx->skb = skb;
4406c376895608 Johannes Berg 2010-09-24  4801  
49ddf8e6e2347c Johannes Berg 2016-03-31  4802  	/* See if we can do fast-rx; if we have to copy we already lost,
49ddf8e6e2347c Johannes Berg 2016-03-31  4803  	 * so punt in that case. We should never have to deliver a data
49ddf8e6e2347c Johannes Berg 2016-03-31  4804  	 * frame to multiple interfaces anyway.
49ddf8e6e2347c Johannes Berg 2016-03-31  4805  	 *
49ddf8e6e2347c Johannes Berg 2016-03-31  4806  	 * We skip the ieee80211_accept_frame() call and do the necessary
49ddf8e6e2347c Johannes Berg 2016-03-31  4807  	 * checking inside ieee80211_invoke_fast_rx().
49ddf8e6e2347c Johannes Berg 2016-03-31  4808  	 */
49ddf8e6e2347c Johannes Berg 2016-03-31  4809  	if (consume && rx->sta) {
49ddf8e6e2347c Johannes Berg 2016-03-31  4810  		struct ieee80211_fast_rx *fast_rx;
49ddf8e6e2347c Johannes Berg 2016-03-31  4811  
49ddf8e6e2347c Johannes Berg 2016-03-31  4812  		fast_rx = rcu_dereference(rx->sta->fast_rx);
49ddf8e6e2347c Johannes Berg 2016-03-31  4813  		if (fast_rx && ieee80211_invoke_fast_rx(rx, fast_rx))
49ddf8e6e2347c Johannes Berg 2016-03-31  4814  			return true;
49ddf8e6e2347c Johannes Berg 2016-03-31  4815  	}
49ddf8e6e2347c Johannes Berg 2016-03-31  4816  
a58fbe1a8540d8 Johannes Berg 2015-04-22  4817  	if (!ieee80211_accept_frame(rx))
4406c376895608 Johannes Berg 2010-09-24  4818  		return false;
4406c376895608 Johannes Berg 2010-09-24  4819  
4406c376895608 Johannes Berg 2010-09-24  4820  	if (!consume) {
f9202638df3447 Avraham Stern 2022-01-26  4821  		struct skb_shared_hwtstamps *shwt;
f9202638df3447 Avraham Stern 2022-01-26  4822  
f9202638df3447 Avraham Stern 2022-01-26  4823  		rx->skb = skb_copy(skb, GFP_ATOMIC);
f9202638df3447 Avraham Stern 2022-01-26  4824  		if (!rx->skb) {
4406c376895608 Johannes Berg 2010-09-24  4825  			if (net_ratelimit())
4406c376895608 Johannes Berg 2010-09-24  4826  				wiphy_debug(local->hw.wiphy,
b305dae488193b Ben Greear    2011-01-08  4827  					"failed to copy skb for %s\n",
4406c376895608 Johannes Berg 2010-09-24  4828  					sdata->name);
4406c376895608 Johannes Berg 2010-09-24  4829  			return true;
4406c376895608 Johannes Berg 2010-09-24  4830  		}
4406c376895608 Johannes Berg 2010-09-24  4831  
f9202638df3447 Avraham Stern 2022-01-26  4832  		/* skb_copy() does not copy the hw timestamps, so copy it
f9202638df3447 Avraham Stern 2022-01-26  4833  		 * explicitly
f9202638df3447 Avraham Stern 2022-01-26  4834  		 */
f9202638df3447 Avraham Stern 2022-01-26  4835  		shwt = skb_hwtstamps(rx->skb);
f9202638df3447 Avraham Stern 2022-01-26  4836  		shwt->hwtstamp = skb_hwtstamps(skb)->hwtstamp;
4406c376895608 Johannes Berg 2010-09-24  4837  	}
4406c376895608 Johannes Berg 2010-09-24  4838  
833b68b914bcc2 Felix Fietkau 2022-12-15  4839  	if (unlikely(rx->sta && rx->sta->sta.mlo)) {
42fb9148c07800 Johannes Berg 2022-07-14  4840  		/* translate to MLD addresses */
42fb9148c07800 Johannes Berg 2022-07-14 @4841  		if (ether_addr_equal(link->conf->addr, hdr->addr1))
42fb9148c07800 Johannes Berg 2022-07-14  4842  			ether_addr_copy(hdr->addr1, rx->sdata->vif.addr);
833b68b914bcc2 Felix Fietkau 2022-12-15  4843  		if (ether_addr_equal(rx->link_sta->addr, hdr->addr2))
42fb9148c07800 Johannes Berg 2022-07-14  4844  			ether_addr_copy(hdr->addr2, rx->sta->addr);
1f6389440cebfb Johannes Berg 2022-07-18  4845  		/* translate A3 only if it's the BSSID */
1f6389440cebfb Johannes Berg 2022-07-18  4846  		if (!ieee80211_has_tods(hdr->frame_control) &&
1f6389440cebfb Johannes Berg 2022-07-18  4847  		    !ieee80211_has_fromds(hdr->frame_control)) {
833b68b914bcc2 Felix Fietkau 2022-12-15  4848  			if (ether_addr_equal(rx->link_sta->addr, hdr->addr3))
42fb9148c07800 Johannes Berg 2022-07-14  4849  				ether_addr_copy(hdr->addr3, rx->sta->addr);
42fb9148c07800 Johannes Berg 2022-07-14  4850  			else if (ether_addr_equal(link->conf->addr, hdr->addr3))
42fb9148c07800 Johannes Berg 2022-07-14  4851  				ether_addr_copy(hdr->addr3, rx->sdata->vif.addr);
1f6389440cebfb Johannes Berg 2022-07-18  4852  		}
42fb9148c07800 Johannes Berg 2022-07-14  4853  		/* not needed for A4 since it can only carry the SA */
42fb9148c07800 Johannes Berg 2022-07-14  4854  	}
42fb9148c07800 Johannes Berg 2022-07-14  4855  
4406c376895608 Johannes Berg 2010-09-24  4856  	ieee80211_invoke_rx_handlers(rx);
4406c376895608 Johannes Berg 2010-09-24  4857  	return true;
4406c376895608 Johannes Berg 2010-09-24  4858  }
4406c376895608 Johannes Berg 2010-09-24  4859
Felix Fietkau Dec. 16, 2022, 10:50 a.m. UTC | #2 
On 16.12.22 11:36, Kalle Valo wrote:
> Felix Fietkau <nbd@nbd.name> writes:
> 
>> There are some codepaths that do not initialize rx->link_sta properly. This
>> causes a crash in places which assume that rx->link_sta is valid if rx->sta
>> is valid.
>> One known instance is triggered by __ieee80211_rx_h_amsdu being called from
>> fast-rx.
> 
> An example crash log would be nice to include, it would make it easier
> to find this fix.
Sure. Will add the one that was sent to the list yesterday to v3.

- Felix
diff mbox series

Patch

diff --git a/net/mac80211/rx.c b/net/mac80211/rx.c
index 3fa7b36d4324..247dd68893af 100644
--- a/net/mac80211/rx.c
+++ b/net/mac80211/rx.c
@@ -4091,6 +4091,55 @@  static void ieee80211_invoke_rx_handlers(struct ieee80211_rx_data *rx)
 #undef CALL_RXH
 }
 
+static bool
+ieee80211_rx_is_valid_sta_link_id(struct ieee80211_sta *sta, u8 link_id)
+{
+	if (!sta->mlo)
+		return false;
+
+	return !!(sta->valid_links & BIT(link_id));
+}
+
+static bool ieee80211_rx_data_set_link(struct ieee80211_rx_data *rx,
+				       u8 link_id)
+{
+	if (!ieee80211_rx_is_valid_sta_link_id(&rx->sta->sta, link_id))
+		return false;
+
+	rx->link_id = link_id;
+	rx->link = rcu_dereference(rx->sdata->link[link_id]);
+	rx->link_sta = rcu_dereference(rx->sta->link[link_id]);
+
+	return rx->link && rx->link_sta;
+}
+
+static bool ieee80211_rx_data_set_sta(struct ieee80211_rx_data *rx,
+				      struct ieee80211_sta *pubsta,
+				      int link_id)
+{
+	struct sta_info *sta;
+
+	sta = container_of(pubsta, struct sta_info, sta);
+
+	rx->link_id = link_id;
+	rx->sta = sta;
+
+	if (sta) {
+		rx->local = sta->sdata->local;
+		rx->sdata = sta->sdata;
+		rx->link_sta = &sta->deflink;
+
+		if (link_id >= 0 &&
+		    !ieee80211_rx_data_set_link(rx, link_id))
+			return false;
+	}
+
+	if (link_id < 0)
+		rx->link = &rx->sdata->deflink;
+
+	return true;
+}
+
 /*
  * This function makes calls into the RX path, therefore
  * it has to be invoked under RCU read lock.
@@ -4099,16 +4148,19 @@  void ieee80211_release_reorder_timeout(struct sta_info *sta, int tid)
 {
 	struct sk_buff_head frames;
 	struct ieee80211_rx_data rx = {
-		.sta = sta,
-		.sdata = sta->sdata,
-		.local = sta->local,
 		/* This is OK -- must be QoS data frame */
 		.security_idx = tid,
 		.seqno_idx = tid,
-		.link_id = -1,
 	};
 	struct tid_ampdu_rx *tid_agg_rx;
-	u8 link_id;
+	int link_id = -1;
+
+	/* FIXME: statistics won't be right with this */
+	if (sta->sta.valid_links)
+		link_id = ffs(sta->sta.valid_links) - 1;
+
+	if (!ieee80211_rx_data_set_sta(&rx, &sta->sta, link_id))
+		return;
 
 	tid_agg_rx = rcu_dereference(sta->ampdu_mlme.tid_rx[tid]);
 	if (!tid_agg_rx)
@@ -4128,10 +4180,6 @@  void ieee80211_release_reorder_timeout(struct sta_info *sta, int tid)
 		};
 		drv_event_callback(rx.local, rx.sdata, &event);
 	}
-	/* FIXME: statistics won't be right with this */
-	link_id = sta->sta.valid_links ? ffs(sta->sta.valid_links) - 1 : 0;
-	rx.link = rcu_dereference(sta->sdata->link[link_id]);
-	rx.link_sta = rcu_dereference(sta->link[link_id]);
 
 	ieee80211_rx_handlers(&rx, &frames);
 }
@@ -4147,7 +4195,6 @@  void ieee80211_mark_rx_ba_filtered_frames(struct ieee80211_sta *pubsta, u8 tid,
 		/* This is OK -- must be QoS data frame */
 		.security_idx = tid,
 		.seqno_idx = tid,
-		.link_id = -1,
 	};
 	int i, diff;
 
@@ -4158,10 +4205,8 @@  void ieee80211_mark_rx_ba_filtered_frames(struct ieee80211_sta *pubsta, u8 tid,
 
 	sta = container_of(pubsta, struct sta_info, sta);
 
-	rx.sta = sta;
-	rx.sdata = sta->sdata;
-	rx.link = &rx.sdata->deflink;
-	rx.local = sta->local;
+	if (!ieee80211_rx_data_set_sta(&rx, pubsta, -1))
+		return;
 
 	rcu_read_lock();
 	tid_agg_rx = rcu_dereference(sta->ampdu_mlme.tid_rx[tid]);
@@ -4548,15 +4593,6 @@  void ieee80211_check_fast_rx_iface(struct ieee80211_sub_if_data *sdata)
 	mutex_unlock(&local->sta_mtx);
 }
 
-static bool
-ieee80211_rx_is_valid_sta_link_id(struct ieee80211_sta *sta, u8 link_id)
-{
-	if (!sta->mlo)
-		return false;
-
-	return !!(sta->valid_links & BIT(link_id));
-}
-
 static void ieee80211_rx_8023(struct ieee80211_rx_data *rx,
 			      struct ieee80211_fast_rx *fast_rx,
 			      int orig_len)
@@ -4667,7 +4703,6 @@  static bool ieee80211_invoke_fast_rx(struct ieee80211_rx_data *rx,
 	struct sk_buff *skb = rx->skb;
 	struct ieee80211_hdr *hdr = (void *)skb->data;
 	struct ieee80211_rx_status *status = IEEE80211_SKB_RXCB(skb);
-	struct sta_info *sta = rx->sta;
 	int orig_len = skb->len;
 	int hdrlen = ieee80211_hdrlen(hdr->frame_control);
 	int snap_offs = hdrlen;
@@ -4679,7 +4714,6 @@  static bool ieee80211_invoke_fast_rx(struct ieee80211_rx_data *rx,
 		u8 da[ETH_ALEN];
 		u8 sa[ETH_ALEN];
 	} addrs __aligned(2);
-	struct link_sta_info *link_sta;
 	struct ieee80211_sta_rx_stats *stats;
 
 	/* for parallel-rx, we need to have DUP_VALIDATED, otherwise we write
@@ -4782,18 +4816,10 @@  static bool ieee80211_invoke_fast_rx(struct ieee80211_rx_data *rx,
  drop:
 	dev_kfree_skb(skb);
 
-	if (rx->link_id >= 0) {
-		link_sta = rcu_dereference(sta->link[rx->link_id]);
-		if (!link_sta)
-			return true;
-	} else {
-		link_sta = &sta->deflink;
-	}
-
 	if (fast_rx->uses_rss)
-		stats = this_cpu_ptr(link_sta->pcpu_rx_stats);
+		stats = this_cpu_ptr(rx->link_sta->pcpu_rx_stats);
 	else
-		stats = &link_sta->rx_stats;
+		stats = &rx->link_sta->rx_stats;
 
 	stats->dropped++;
 	return true;
@@ -4811,7 +4837,6 @@  static bool ieee80211_prepare_and_rx_handle(struct ieee80211_rx_data *rx,
 	struct ieee80211_local *local = rx->local;
 	struct ieee80211_sub_if_data *sdata = rx->sdata;
 	struct ieee80211_hdr *hdr = (void *)skb->data;
-	struct link_sta_info *link_sta = NULL;
 	struct ieee80211_link_data *link;
 
 	rx->skb = skb;
@@ -4834,35 +4859,6 @@  static bool ieee80211_prepare_and_rx_handle(struct ieee80211_rx_data *rx,
 	if (!ieee80211_accept_frame(rx))
 		return false;
 
-	if (rx->link_id >= 0) {
-		link = rcu_dereference(rx->sdata->link[rx->link_id]);
-
-		/* we might race link removal */
-		if (!link)
-			return true;
-		rx->link = link;
-
-		if (rx->sta) {
-			rx->link_sta =
-				rcu_dereference(rx->sta->link[rx->link_id]);
-			if (!rx->link_sta)
-				return true;
-		}
-	} else {
-		if (rx->sta)
-			rx->link_sta = &rx->sta->deflink;
-
-		rx->link = &sdata->deflink;
-	}
-
-	if (unlikely(!is_multicast_ether_addr(hdr->addr1) &&
-		     rx->link_id >= 0 && rx->sta && rx->sta->sta.mlo)) {
-		link_sta = rcu_dereference(rx->sta->link[rx->link_id]);
-
-		if (WARN_ON_ONCE(!link_sta))
-			return true;
-	}
-
 	if (!consume) {
 		struct skb_shared_hwtstamps *shwt;
 
@@ -4882,16 +4878,16 @@  static bool ieee80211_prepare_and_rx_handle(struct ieee80211_rx_data *rx,
 		shwt->hwtstamp = skb_hwtstamps(skb)->hwtstamp;
 	}
 
-	if (unlikely(link_sta)) {
+	if (unlikely(rx->sta && rx->sta->sta.mlo)) {
 		/* translate to MLD addresses */
 		if (ether_addr_equal(link->conf->addr, hdr->addr1))
 			ether_addr_copy(hdr->addr1, rx->sdata->vif.addr);
-		if (ether_addr_equal(link_sta->addr, hdr->addr2))
+		if (ether_addr_equal(rx->link_sta->addr, hdr->addr2))
 			ether_addr_copy(hdr->addr2, rx->sta->addr);
 		/* translate A3 only if it's the BSSID */
 		if (!ieee80211_has_tods(hdr->frame_control) &&
 		    !ieee80211_has_fromds(hdr->frame_control)) {
-			if (ether_addr_equal(link_sta->addr, hdr->addr3))
+			if (ether_addr_equal(rx->link_sta->addr, hdr->addr3))
 				ether_addr_copy(hdr->addr3, rx->sta->addr);
 			else if (ether_addr_equal(link->conf->addr, hdr->addr3))
 				ether_addr_copy(hdr->addr3, rx->sdata->vif.addr);
@@ -4912,6 +4908,7 @@  static void __ieee80211_rx_handle_8023(struct ieee80211_hw *hw,
 	struct ieee80211_rx_status *status = IEEE80211_SKB_RXCB(skb);
 	struct ieee80211_fast_rx *fast_rx;
 	struct ieee80211_rx_data rx;
+	int link_id = -1;
 
 	memset(&rx, 0, sizeof(rx));
 	rx.skb = skb;
@@ -4928,12 +4925,8 @@  static void __ieee80211_rx_handle_8023(struct ieee80211_hw *hw,
 	if (!pubsta)
 		goto drop;
 
-	rx.sta = container_of(pubsta, struct sta_info, sta);
-	rx.sdata = rx.sta->sdata;
-
-	if (status->link_valid &&
-	    !ieee80211_rx_is_valid_sta_link_id(pubsta, status->link_id))
-		goto drop;
+	if (status->link_valid)
+		link_id = status->link_id;
 
 	/*
 	 * TODO: Should the frame be dropped if the right link_id is not
@@ -4942,19 +4935,8 @@  static void __ieee80211_rx_handle_8023(struct ieee80211_hw *hw,
 	 * link_id is used only for stats purpose and updating the stats on
 	 * the deflink is fine?
 	 */
-	if (status->link_valid)
-		rx.link_id = status->link_id;
-
-	if (rx.link_id >= 0) {
-		struct ieee80211_link_data *link;
-
-		link =  rcu_dereference(rx.sdata->link[rx.link_id]);
-		if (!link)
-			goto drop;
-		rx.link = link;
-	} else {
-		rx.link = &rx.sdata->deflink;
-	}
+	if (!ieee80211_rx_data_set_sta(&rx, pubsta, link_id))
+		goto drop;
 
 	fast_rx = rcu_dereference(rx.sta->fast_rx);
 	if (!fast_rx)
@@ -4972,6 +4954,8 @@  static bool ieee80211_rx_for_interface(struct ieee80211_rx_data *rx,
 {
 	struct link_sta_info *link_sta;
 	struct ieee80211_hdr *hdr = (void *)skb->data;
+	struct sta_info *sta;
+	int link_id = -1;
 
 	/*
 	 * Look up link station first, in case there's a
@@ -4981,24 +4965,19 @@  static bool ieee80211_rx_for_interface(struct ieee80211_rx_data *rx,
 	 */
 	link_sta = link_sta_info_get_bss(rx->sdata, hdr->addr2);
 	if (link_sta) {
-		rx->sta = link_sta->sta;
-		rx->link_id = link_sta->link_id;
+		sta = link_sta->sta;
+		link_id = link_sta->link_id;
 	} else {
 		struct ieee80211_rx_status *status = IEEE80211_SKB_RXCB(skb);
 
-		rx->sta = sta_info_get_bss(rx->sdata, hdr->addr2);
-		if (rx->sta) {
-			if (status->link_valid &&
-			    !ieee80211_rx_is_valid_sta_link_id(&rx->sta->sta,
-							       status->link_id))
-				return false;
-
-			rx->link_id = status->link_valid ? status->link_id : -1;
-		} else {
-			rx->link_id = -1;
-		}
+		sta = sta_info_get_bss(rx->sdata, hdr->addr2);
+		if (status->link_valid)
+			link_id = status->link_id;
 	}
 
+	if (!ieee80211_rx_data_set_sta(rx, &sta->sta, link_id))
+		return false;
+
 	return ieee80211_prepare_and_rx_handle(rx, skb, consume);
 }
 
@@ -5057,19 +5036,15 @@  static void __ieee80211_rx_handle_packet(struct ieee80211_hw *hw,
 
 	if (ieee80211_is_data(fc)) {
 		struct sta_info *sta, *prev_sta;
-		u8 link_id = status->link_id;
+		int link_id = -1;
 
-		if (pubsta) {
-			rx.sta = container_of(pubsta, struct sta_info, sta);
-			rx.sdata = rx.sta->sdata;
+		if (status->link_valid)
+			link_id = status->link_id;
 
-			if (status->link_valid &&
-			    !ieee80211_rx_is_valid_sta_link_id(pubsta, link_id))
+		if (pubsta) {
+			if (!ieee80211_rx_data_set_sta(&rx, pubsta, link_id))
 				goto out;
 
-			if (status->link_valid)
-				rx.link_id = status->link_id;
-
 			/*
 			 * In MLO connection, fetch the link_id using addr2
 			 * when the driver does not pass link_id in status.
@@ -5087,7 +5062,7 @@  static void __ieee80211_rx_handle_packet(struct ieee80211_hw *hw,
 				if (!link_sta)
 					goto out;
 
-				rx.link_id = link_sta->link_id;
+				ieee80211_rx_data_set_link(&rx, link_sta->link_id);
 			}
 
 			if (ieee80211_prepare_and_rx_handle(&rx, skb, true))
@@ -5103,30 +5078,25 @@  static void __ieee80211_rx_handle_packet(struct ieee80211_hw *hw,
 				continue;
 			}
 
-			if ((status->link_valid &&
-			     !ieee80211_rx_is_valid_sta_link_id(&prev_sta->sta,
-								link_id)) ||
-			    (!status->link_valid && prev_sta->sta.mlo))
+			if (!ieee80211_rx_data_set_sta(&rx, &prev_sta->sta,
+						       link_id))
+				goto out;
+
+			if (!status->link_valid && prev_sta->sta.mlo)
 				continue;
 
-			rx.link_id = status->link_valid ? link_id : -1;
-			rx.sta = prev_sta;
-			rx.sdata = prev_sta->sdata;
 			ieee80211_prepare_and_rx_handle(&rx, skb, false);
 
 			prev_sta = sta;
 		}
 
 		if (prev_sta) {
-			if ((status->link_valid &&
-			     !ieee80211_rx_is_valid_sta_link_id(&prev_sta->sta,
-								link_id)) ||
-			    (!status->link_valid && prev_sta->sta.mlo))
+			if (!ieee80211_rx_data_set_sta(&rx, &prev_sta->sta,
+						       link_id))
 				goto out;
 
-			rx.link_id = status->link_valid ? link_id : -1;
-			rx.sta = prev_sta;
-			rx.sdata = prev_sta->sdata;
+			if (!status->link_valid && prev_sta->sta.mlo)
+				goto out;
 
 			if (ieee80211_prepare_and_rx_handle(&rx, skb, true))
 				return;