From: Eason Lai <Eason.Lai@mediatek.com>

A use-after-free issue occurs in mt76_rx_poll_complete due to a race
condition. The STA has already been removed, but the rx_status still
had a pointer to the wcid in the STA.

Use wcid_idx instead of storing the wcid pointer, and look up the wcid
via rcu_dereference() by wcid_idx.

BUG: KASAN: invalid-access in mt76_rx_poll_complete+0x280/0x470 [mt76]
Call trace:
dump_backtrace+0xec/0x128
show_stack+0x18/0x28
dump_stack_lvl+0x40/0xc8
print_report+0x1b8/0x710
kasan_report+0xe0/0x144
do_bad_area+0x120/0x260
do_tag_check_fault+0x20/0x34
do_mem_abort+0x54/0xa8
el1_abort+0x3c/0x5c
el1h_64_sync_handler+0x40/0xcc
el1h_64_sync+0x7c/0x80
mt76_rx_poll_complete+0x280/0x470
mt76_dma_rx_poll+0x114/0x51c
mt792x_poll_rx+0x60/0xf8
napi_threaded_poll_loop+0xe0/0x450
napi_threaded_poll+0x80/0x9c
kthread+0x11c/0x158
ret_from_fork+0x10/0x20

Fixes: 9c68a57bc22d ("mt76: get station pointer by wcid and pass it to mac80211")
Signed-off-by: Eason Lai <Eason.Lai@mediatek.com>
---
v2: fix mt76x02 build errors
---
 drivers/net/wireless/mediatek/mt76/agg-rx.c   | 18 ++++---
 drivers/net/wireless/mediatek/mt76/mac80211.c | 47 ++++++++-----------
 drivers/net/wireless/mediatek/mt76/mt76.h     |  9 ++--
 .../net/wireless/mediatek/mt76/mt7603/mac.c   |  6 ++-
 .../net/wireless/mediatek/mt76/mt7615/mac.c   | 19 +++++---
 .../net/wireless/mediatek/mt76/mt76_connac.h  |  3 +-
 .../wireless/mediatek/mt76/mt76_connac_mac.c  | 12 ++++-
 .../net/wireless/mediatek/mt76/mt76x02_mac.c  |  9 ++--
 .../net/wireless/mediatek/mt76/mt7915/mac.c   | 16 +++++--
 .../net/wireless/mediatek/mt76/mt7921/mac.c   | 13 +++--
 .../net/wireless/mediatek/mt76/mt7925/mac.c   | 19 +++++---
 .../net/wireless/mediatek/mt76/mt7996/mac.c   | 29 ++++++++----
 12 files changed, 119 insertions(+), 81 deletions(-)

diff --git a/drivers/net/wireless/mediatek/mt76/agg-rx.c b/drivers/net/wireless/mediatek/mt76/agg-rx.c
index bf1babac3895..b72d13975466 100644
--- a/drivers/net/wireless/mediatek/mt76/agg-rx.c
+++ b/drivers/net/wireless/mediatek/mt76/agg-rx.c
@@ -116,14 +116,15 @@ mt76_rx_aggr_reorder_work(struct work_struct *work)
 }
 
 static void
-mt76_rx_aggr_check_ctl(struct sk_buff *skb, struct sk_buff_head *frames)
+mt76_rx_aggr_check_ctl(struct mt76_dev *dev, struct sk_buff *skb,
+		       struct sk_buff_head *frames)
 {
 	struct mt76_rx_status *status = (struct mt76_rx_status *)skb->cb;
 	struct ieee80211_bar *bar = mt76_skb_get_hdr(skb);
-	struct mt76_wcid *wcid = status->wcid;
 	struct mt76_rx_tid *tid;
-	u8 tidno;
+	struct mt76_wcid *wcid;
 	u16 seqno;
+	u8 tidno;
 
 	if (!ieee80211_is_ctl(bar->frame_control))
 		return;
@@ -131,6 +132,10 @@ mt76_rx_aggr_check_ctl(struct sk_buff *skb, struct sk_buff_head *frames)
 	if (!ieee80211_is_back_req(bar->frame_control))
 		return;
 
+	wcid = __mt76_wcid_ptr(dev, status->wcid_idx);
+	if (!wcid)
+		return;
+
 	status->qos_ctl = tidno = le16_to_cpu(bar->control) >> 12;
 	seqno = IEEE80211_SEQ_TO_SN(le16_to_cpu(bar->start_seq_num));
 	tid = rcu_dereference(wcid->aggr[tidno]);
@@ -145,10 +150,11 @@ mt76_rx_aggr_check_ctl(struct sk_buff *skb, struct sk_buff_head *frames)
 	spin_unlock_bh(&tid->lock);
 }
 
-void mt76_rx_aggr_reorder(struct sk_buff *skb, struct sk_buff_head *frames)
+void mt76_rx_aggr_reorder(struct mt76_dev *dev, struct sk_buff *skb,
+			  struct sk_buff_head *frames)
 {
 	struct mt76_rx_status *status = (struct mt76_rx_status *)skb->cb;
-	struct mt76_wcid *wcid = status->wcid;
+	struct mt76_wcid *wcid = __mt76_wcid_ptr(dev, status->wcid_idx);
 	struct ieee80211_sta *sta;
 	struct mt76_rx_tid *tid;
 	bool sn_less;
@@ -164,7 +170,7 @@ void mt76_rx_aggr_reorder(struct sk_buff *skb, struct sk_buff_head *frames)
 
 	if (!status->aggr) {
 		if (!(status->flag & RX_FLAG_8023))
-			mt76_rx_aggr_check_ctl(skb, frames);
+			mt76_rx_aggr_check_ctl(dev, skb, frames);
 		return;
 	}
 
diff --git a/drivers/net/wireless/mediatek/mt76/mac80211.c b/drivers/net/wireless/mediatek/mt76/mac80211.c
index 4ae5e4715a9c..2933b29782d3 100644
--- a/drivers/net/wireless/mediatek/mt76/mac80211.c
+++ b/drivers/net/wireless/mediatek/mt76/mac80211.c
@@ -1283,21 +1283,24 @@ mt76_rx_convert(struct mt76_dev *dev, struct sk_buff *skb,
 	memcpy(status->chain_signal, mstat.chain_signal,
 	       sizeof(mstat.chain_signal));
 
-	if (mstat.wcid) {
-		status->link_valid = mstat.wcid->link_valid;
-		status->link_id = mstat.wcid->link_id;
-	}
+	if (mstat.wcid_idx != MT76_WCID_IDX_INVALID) {
+		struct mt76_wcid *wcid = __mt76_wcid_ptr(dev, mstat.wcid_idx);
 
-	*sta = wcid_to_sta(mstat.wcid);
+		if (wcid) {
+			status->link_valid = wcid->link_valid;
+			status->link_id = wcid->link_id;
+			*sta = wcid_to_sta(wcid);
+		}
+	}
 	*hw = mt76_phy_hw(dev, mstat.phy_idx);
 }
 
 static void
-mt76_check_ccmp_pn(struct sk_buff *skb)
+mt76_check_ccmp_pn(struct mt76_dev *dev, struct sk_buff *skb)
 {
 	struct mt76_rx_status *status = (struct mt76_rx_status *)skb->cb;
-	struct mt76_wcid *wcid = status->wcid;
 	struct ieee80211_hdr *hdr;
+	struct mt76_wcid *wcid;
 	int security_idx;
 	int ret;
 
@@ -1307,6 +1310,7 @@ mt76_check_ccmp_pn(struct sk_buff *skb)
 	if (status->flag & RX_FLAG_ONLY_MONITOR)
 		return;
 
+	wcid = __mt76_wcid_ptr(dev, status->wcid_idx);
 	if (!wcid || !wcid->rx_check_pn)
 		return;
 
@@ -1354,7 +1358,7 @@ static void
 mt76_airtime_report(struct mt76_dev *dev, struct mt76_rx_status *status,
 		    int len)
 {
-	struct mt76_wcid *wcid = status->wcid;
+	struct mt76_wcid *wcid = __mt76_wcid_ptr(dev, status->wcid_idx);
 	struct ieee80211_rx_status info = {
 		.enc_flags = status->enc_flags,
 		.rate_idx = status->rate_idx,
@@ -1382,19 +1386,9 @@ mt76_airtime_report(struct mt76_dev *dev, struct mt76_rx_status *status,
 static void
 mt76_airtime_flush_ampdu(struct mt76_dev *dev)
 {
-	struct mt76_wcid *wcid;
-	int wcid_idx;
-
 	if (!dev->rx_ampdu_len)
 		return;
 
-	wcid_idx = dev->rx_ampdu_status.wcid_idx;
-	if (wcid_idx < ARRAY_SIZE(dev->wcid))
-		wcid = rcu_dereference(dev->wcid[wcid_idx]);
-	else
-		wcid = NULL;
-	dev->rx_ampdu_status.wcid = wcid;
-
 	mt76_airtime_report(dev, &dev->rx_ampdu_status, dev->rx_ampdu_len);
 
 	dev->rx_ampdu_len = 0;
@@ -1405,7 +1399,7 @@ static void
 mt76_airtime_check(struct mt76_dev *dev, struct sk_buff *skb)
 {
 	struct mt76_rx_status *status = (struct mt76_rx_status *)skb->cb;
-	struct mt76_wcid *wcid = status->wcid;
+	struct mt76_wcid *wcid = __mt76_wcid_ptr(dev, status->wcid_idx);
 
 	if (!(dev->drv->drv_flags & MT_DRV_SW_RX_AIRTIME))
 		return;
@@ -1418,8 +1412,6 @@ mt76_airtime_check(struct mt76_dev *dev, struct sk_buff *skb)
 
 		if (!ether_addr_equal(hdr->addr1, dev->phy.macaddr))
 			return;
-
-		wcid = NULL;
 	}
 
 	if (!(status->flag & RX_FLAG_AMPDU_DETAILS) ||
@@ -1430,7 +1422,6 @@ mt76_airtime_check(struct mt76_dev *dev, struct sk_buff *skb)
 		if (!dev->rx_ampdu_len ||
 		    status->ampdu_ref != dev->rx_ampdu_ref) {
 			dev->rx_ampdu_status = *status;
-			dev->rx_ampdu_status.wcid_idx = wcid ? wcid->idx : 0xff;
 			dev->rx_ampdu_ref = status->ampdu_ref;
 		}
 
@@ -1448,7 +1439,7 @@ mt76_check_sta(struct mt76_dev *dev, struct sk_buff *skb)
 	struct ieee80211_hdr *hdr = mt76_skb_get_hdr(skb);
 	struct ieee80211_sta *sta;
 	struct ieee80211_hw *hw;
-	struct mt76_wcid *wcid = status->wcid;
+	struct mt76_wcid *wcid = __mt76_wcid_ptr(dev, status->wcid_idx);
 	u8 tidno = status->qos_ctl & IEEE80211_QOS_CTL_TID_MASK;
 	bool ps;
 
@@ -1456,8 +1447,10 @@ mt76_check_sta(struct mt76_dev *dev, struct sk_buff *skb)
 	if (ieee80211_is_pspoll(hdr->frame_control) && !wcid &&
 	    !(status->flag & RX_FLAG_8023)) {
 		sta = ieee80211_find_sta_by_ifaddr(hw, hdr->addr2, NULL);
-		if (sta)
-			wcid = status->wcid = (struct mt76_wcid *)sta->drv_priv;
+		if (sta) {
+			wcid = (struct mt76_wcid *)sta->drv_priv;
+			status->wcid_idx = wcid->idx;
+		}
 	}
 
 	mt76_airtime_check(dev, skb);
@@ -1521,7 +1514,7 @@ void mt76_rx_complete(struct mt76_dev *dev, struct sk_buff_head *frames,
 	while ((skb = __skb_dequeue(frames)) != NULL) {
 		struct sk_buff *nskb = skb_shinfo(skb)->frag_list;
 
-		mt76_check_ccmp_pn(skb);
+		mt76_check_ccmp_pn(dev, skb);
 		skb_shinfo(skb)->frag_list = NULL;
 		mt76_rx_convert(dev, skb, &hw, &sta);
 		ieee80211_rx_list(hw, sta, skb, &list);
@@ -1563,7 +1556,7 @@ void mt76_rx_poll_complete(struct mt76_dev *dev, enum mt76_rxq_id q,
 		    mt76_npu_device_active(dev))
 			__skb_queue_tail(&frames, skb);
 		else
-			mt76_rx_aggr_reorder(skb, &frames);
+			mt76_rx_aggr_reorder(dev, skb, &frames);
 	}
 
 	mt76_rx_complete(dev, &frames, napi);
diff --git a/drivers/net/wireless/mediatek/mt76/mt76.h b/drivers/net/wireless/mediatek/mt76/mt76.h
index 527bef97e122..1e1bd410a7a7 100644
--- a/drivers/net/wireless/mediatek/mt76/mt76.h
+++ b/drivers/net/wireless/mediatek/mt76/mt76.h
@@ -364,6 +364,7 @@ enum mt76_wcid_flags {
 };
 
 #define MT76_N_WCIDS 1088
+#define MT76_WCID_IDX_INVALID		0xffff
 #define MT76_BEACON_MON_MAX_MISS	7
 
 /* stored in ieee80211_tx_info::hw_queue */
@@ -728,10 +729,7 @@ struct mt76_mmio {
 };
 
 struct mt76_rx_status {
-	union {
-		struct mt76_wcid *wcid;
-		u16 wcid_idx;
-	};
+	u16 wcid_idx;
 
 	u32 reorder_time;
 
@@ -1793,7 +1791,8 @@ void mt76_rx_complete(struct mt76_dev *dev, struct sk_buff_head *frames,
 		      struct napi_struct *napi);
 void mt76_rx_poll_complete(struct mt76_dev *dev, enum mt76_rxq_id q,
 			   struct napi_struct *napi);
-void mt76_rx_aggr_reorder(struct sk_buff *skb, struct sk_buff_head *frames);
+void mt76_rx_aggr_reorder(struct mt76_dev *dev, struct sk_buff *skb,
+			  struct sk_buff_head *frames);
 void mt76_testmode_tx_pending(struct mt76_phy *phy);
 void mt76_queue_tx_complete(struct mt76_dev *dev, struct mt76_queue *q,
 			    struct mt76_queue_entry *e);
diff --git a/drivers/net/wireless/mediatek/mt76/mt7603/mac.c b/drivers/net/wireless/mediatek/mt76/mt7603/mac.c
index d3110eeb45d7..aeb890df808e 100644
--- a/drivers/net/wireless/mediatek/mt76/mt7603/mac.c
+++ b/drivers/net/wireless/mediatek/mt76/mt7603/mac.c
@@ -513,6 +513,7 @@ mt7603_mac_fill_rx(struct mt7603_dev *dev, struct sk_buff *skb)
 	u32 rxd2 = le32_to_cpu(rxd[2]);
 	bool unicast = rxd1 & MT_RXD1_NORMAL_U2M;
 	bool insert_ccmp_hdr = false;
+	struct mt76_wcid *wcid;
 	bool remove_pad;
 	int idx;
 	int i;
@@ -524,7 +525,8 @@ mt7603_mac_fill_rx(struct mt7603_dev *dev, struct sk_buff *skb)
 	i >>= 1;
 
 	idx = FIELD_GET(MT_RXD2_NORMAL_WLAN_IDX, rxd2);
-	status->wcid = mt7603_rx_get_wcid(dev, idx, unicast);
+	wcid = mt7603_rx_get_wcid(dev, idx, unicast);
+	status->wcid_idx = wcid ? wcid->idx : MT76_WCID_IDX_INVALID;
 
 	status->band = sband->band;
 	if (i < sband->n_channels)
@@ -672,7 +674,7 @@ mt7603_mac_fill_rx(struct mt7603_dev *dev, struct sk_buff *skb)
 	}
 
 	hdr = (struct ieee80211_hdr *)skb->data;
-	if (!status->wcid || !ieee80211_is_data_qos(hdr->frame_control))
+	if (!wcid || !ieee80211_is_data_qos(hdr->frame_control))
 		return 0;
 
 	status->aggr = unicast &&
diff --git a/drivers/net/wireless/mediatek/mt76/mt7615/mac.c b/drivers/net/wireless/mediatek/mt76/mt7615/mac.c
index ce0051468501..84be00a13930 100644
--- a/drivers/net/wireless/mediatek/mt76/mt7615/mac.c
+++ b/drivers/net/wireless/mediatek/mt76/mt7615/mac.c
@@ -246,17 +246,20 @@ static void mt7615_mac_fill_tm_rx(struct mt7615_phy *phy, __le32 *rxv)
 }
 
 /* The HW does not translate the mac header to 802.3 for mesh point */
-static int mt7615_reverse_frag0_hdr_trans(struct sk_buff *skb, u16 hdr_gap)
+static int mt7615_reverse_frag0_hdr_trans(struct mt7615_dev *dev,
+					  struct sk_buff *skb, u16 hdr_gap)
 {
 	struct mt76_rx_status *status = (struct mt76_rx_status *)skb->cb;
 	struct ethhdr *eth_hdr = (struct ethhdr *)(skb->data + hdr_gap);
-	struct mt7615_sta *msta = (struct mt7615_sta *)status->wcid;
 	__le32 *rxd = (__le32 *)skb->data;
 	struct ieee80211_sta *sta;
 	struct ieee80211_vif *vif;
 	struct ieee80211_hdr hdr;
+	struct mt7615_sta *msta;
 	u16 frame_control;
 
+	msta = (struct mt7615_sta *)mt76_wcid_ptr(dev, status->wcid_idx);
+
 	if (le32_get_bits(rxd[1], MT_RXD1_NORMAL_ADDR_TYPE) !=
 	    MT_RXD1_NORMAL_U2M)
 		return -EINVAL;
@@ -335,6 +338,7 @@ static int mt7615_mac_fill_rx(struct mt7615_dev *dev, struct sk_buff *skb)
 	struct ieee80211_supported_band *sband;
 	struct ieee80211_hdr *hdr;
 	struct mt7615_phy *phy2;
+	struct mt76_wcid *wcid;
 	__le32 *rxd = (__le32 *)skb->data;
 	u32 rxd0 = le32_to_cpu(rxd[0]);
 	u32 rxd1 = le32_to_cpu(rxd[1]);
@@ -378,12 +382,13 @@ static int mt7615_mac_fill_rx(struct mt7615_dev *dev, struct sk_buff *skb)
 
 	unicast = (rxd1 & MT_RXD1_NORMAL_ADDR_TYPE) == MT_RXD1_NORMAL_U2M;
 	idx = FIELD_GET(MT_RXD2_NORMAL_WLAN_IDX, rxd2);
-	status->wcid = mt7615_rx_get_wcid(dev, idx, unicast);
+	wcid = mt7615_rx_get_wcid(dev, idx, unicast);
+	status->wcid_idx = wcid ? wcid->idx : MT76_WCID_IDX_INVALID;
 
-	if (status->wcid) {
+	if (wcid) {
 		struct mt7615_sta *msta;
 
-		msta = container_of(status->wcid, struct mt7615_sta, wcid);
+		msta = container_of(wcid, struct mt7615_sta, wcid);
 		mt76_wcid_add_poll(&dev->mt76, &msta->wcid);
 	}
 
@@ -590,7 +595,7 @@ static int mt7615_mac_fill_rx(struct mt7615_dev *dev, struct sk_buff *skb)
 
 	hdr_gap = (u8 *)rxd - skb->data + 2 * remove_pad;
 	if (hdr_trans && ieee80211_has_morefrags(fc)) {
-		if (mt7615_reverse_frag0_hdr_trans(skb, hdr_gap))
+		if (mt7615_reverse_frag0_hdr_trans(dev, skb, hdr_gap))
 			return -EINVAL;
 		hdr_trans = false;
 	} else {
@@ -638,7 +643,7 @@ static int mt7615_mac_fill_rx(struct mt7615_dev *dev, struct sk_buff *skb)
 		status->flag |= RX_FLAG_8023;
 	}
 
-	if (!status->wcid || !ieee80211_is_data_qos(fc))
+	if (!wcid || !ieee80211_is_data_qos(fc))
 		return 0;
 
 	status->aggr = unicast &&
diff --git a/drivers/net/wireless/mediatek/mt76/mt76_connac.h b/drivers/net/wireless/mediatek/mt76/mt76_connac.h
index 51423c7740bd..d9ae8d204442 100644
--- a/drivers/net/wireless/mediatek/mt76/mt76_connac.h
+++ b/drivers/net/wireless/mediatek/mt76/mt76_connac.h
@@ -448,7 +448,8 @@ bool mt76_connac2_mac_add_txs_skb(struct mt76_dev *dev, struct mt76_wcid *wcid,
 void mt76_connac2_mac_decode_he_radiotap(struct mt76_dev *dev,
 					 struct sk_buff *skb,
 					 __le32 *rxv, u32 mode);
-int mt76_connac2_reverse_frag0_hdr_trans(struct ieee80211_vif *vif,
+int mt76_connac2_reverse_frag0_hdr_trans(struct mt76_dev *dev,
+					 struct ieee80211_vif *vif,
 					 struct sk_buff *skb, u16 hdr_offset);
 int mt76_connac2_mac_fill_rx_rate(struct mt76_dev *dev,
 				  struct mt76_rx_status *status,
diff --git a/drivers/net/wireless/mediatek/mt76/mt76_connac_mac.c b/drivers/net/wireless/mediatek/mt76/mt76_connac_mac.c
index 0339e2e7ab60..376d59d517cc 100644
--- a/drivers/net/wireless/mediatek/mt76/mt76_connac_mac.c
+++ b/drivers/net/wireless/mediatek/mt76/mt76_connac_mac.c
@@ -952,7 +952,8 @@ void mt76_connac2_mac_decode_he_radiotap(struct mt76_dev *dev,
 EXPORT_SYMBOL_GPL(mt76_connac2_mac_decode_he_radiotap);
 
 /* The HW does not translate the mac header to 802.3 for mesh point */
-int mt76_connac2_reverse_frag0_hdr_trans(struct ieee80211_vif *vif,
+int mt76_connac2_reverse_frag0_hdr_trans(struct mt76_dev *dev,
+					 struct ieee80211_vif *vif,
 					 struct sk_buff *skb, u16 hdr_offset)
 {
 	struct mt76_rx_status *status = (struct mt76_rx_status *)skb->cb;
@@ -960,6 +961,7 @@ int mt76_connac2_reverse_frag0_hdr_trans(struct ieee80211_vif *vif,
 	__le32 *rxd = (__le32 *)skb->data;
 	struct ieee80211_sta *sta;
 	struct ieee80211_hdr hdr;
+	struct mt76_wcid *wcid;
 	u16 frame_control;
 
 	if (le32_get_bits(rxd[3], MT_RXD3_NORMAL_ADDR_TYPE) !=
@@ -969,7 +971,13 @@ int mt76_connac2_reverse_frag0_hdr_trans(struct ieee80211_vif *vif,
 	if (!(le32_to_cpu(rxd[1]) & MT_RXD1_NORMAL_GROUP_4))
 		return -EINVAL;
 
-	sta = container_of((void *)status->wcid, struct ieee80211_sta, drv_priv);
+	wcid = __mt76_wcid_ptr(dev, status->wcid_idx);
+	if (!wcid)
+		return -EINVAL;
+
+	sta = container_of((void *)wcid, struct ieee80211_sta, drv_priv);
+	if (!sta)
+		return -EINVAL;
 
 	/* store the info from RXD and ethhdr to avoid being overridden */
 	frame_control = le32_get_bits(rxd[6], MT_RXD6_FRAME_CONTROL);
diff --git a/drivers/net/wireless/mediatek/mt76/mt76x02_mac.c b/drivers/net/wireless/mediatek/mt76/mt76x02_mac.c
index 14ee5b3b94d3..ce5d3f8f019a 100644
--- a/drivers/net/wireless/mediatek/mt76/mt76x02_mac.c
+++ b/drivers/net/wireless/mediatek/mt76/mt76x02_mac.c
@@ -779,11 +779,10 @@ int mt76x02_mac_process_rx(struct mt76x02_dev *dev, struct sk_buff *skb,
 	u32 ctl = le32_to_cpu(rxwi->ctl);
 	u16 rate = le16_to_cpu(rxwi->rate);
 	u16 tid_sn = le16_to_cpu(rxwi->tid_sn);
-	bool unicast = rxwi->rxinfo & cpu_to_le32(MT_RXINFO_UNICAST);
 	int pad_len = 0, nstreams = dev->mphy.chainmask & 0xf;
+	u8 wcid_idx;
 	s8 signal;
 	u8 pn_len;
-	u8 wcid;
 	int len;
 
 	if (!test_bit(MT76_STATE_RUNNING, &dev->mphy.state))
@@ -799,9 +798,9 @@ int mt76x02_mac_process_rx(struct mt76x02_dev *dev, struct sk_buff *skb,
 		status->flag |= RX_FLAG_IV_STRIPPED;
 	}
 
-	wcid = FIELD_GET(MT_RXWI_CTL_WCID, ctl);
-	sta = mt76x02_rx_get_sta(&dev->mt76, wcid);
-	status->wcid = mt76x02_rx_get_sta_wcid(sta, unicast);
+	wcid_idx = FIELD_GET(MT_RXWI_CTL_WCID, ctl);
+	sta = mt76x02_rx_get_sta(&dev->mt76, wcid_idx);
+	status->wcid_idx = wcid_idx;
 
 	len = FIELD_GET(MT_RXWI_CTL_MPDU_LEN, ctl);
 	pn_len = FIELD_GET(MT_RXINFO_PN_LEN, rxinfo);
diff --git a/drivers/net/wireless/mediatek/mt76/mt7915/mac.c b/drivers/net/wireless/mediatek/mt76/mt7915/mac.c
index cec2c4208255..30672713323f 100644
--- a/drivers/net/wireless/mediatek/mt76/mt7915/mac.c
+++ b/drivers/net/wireless/mediatek/mt76/mt7915/mac.c
@@ -280,6 +280,7 @@ mt7915_mac_fill_rx(struct mt7915_dev *dev, struct sk_buff *skb,
 	u8 mode = 0, qos_ctl = 0;
 	struct mt7915_sta *msta = NULL;
 	u32 csum_status = *(u32 *)skb->cb;
+	struct mt76_wcid *wcid;
 	bool hdr_trans;
 	u16 hdr_gap;
 	u16 seq_ctrl = 0;
@@ -313,10 +314,13 @@ mt7915_mac_fill_rx(struct mt7915_dev *dev, struct sk_buff *skb,
 
 	unicast = FIELD_GET(MT_RXD3_NORMAL_ADDR_TYPE, rxd3) == MT_RXD3_NORMAL_U2M;
 	idx = FIELD_GET(MT_RXD1_NORMAL_WLAN_IDX, rxd1);
-	status->wcid = mt7915_rx_get_wcid(dev, idx, unicast);
 
-	if (status->wcid) {
-		msta = container_of(status->wcid, struct mt7915_sta, wcid);
+	wcid = mt7915_rx_get_wcid(dev, idx, unicast);
+
+	status->wcid_idx = wcid ? wcid->idx : MT76_WCID_IDX_INVALID;
+
+	if (wcid) {
+		msta = container_of(wcid, struct mt7915_sta, wcid);
 		mt76_wcid_add_poll(&dev->mt76, &msta->wcid);
 	}
 
@@ -475,7 +479,8 @@ mt7915_mac_fill_rx(struct mt7915_dev *dev, struct sk_buff *skb,
 
 		vif = container_of((void *)msta->vif, struct ieee80211_vif,
 				   drv_priv);
-		err = mt76_connac2_reverse_frag0_hdr_trans(vif, skb, hdr_gap);
+		err = mt76_connac2_reverse_frag0_hdr_trans(&dev->mt76, vif,
+							   skb, hdr_gap);
 		if (err)
 			return err;
 
@@ -532,7 +537,8 @@ mt7915_mac_fill_rx(struct mt7915_dev *dev, struct sk_buff *skb,
 	if (rxv && mode >= MT_PHY_TYPE_HE_SU && !(status->flag & RX_FLAG_8023))
 		mt76_connac2_mac_decode_he_radiotap(&dev->mt76, skb, rxv, mode);
 
-	if (!status->wcid || !ieee80211_is_data_qos(fc))
+	if (status->wcid_idx == MT76_WCID_IDX_INVALID ||
+	    !ieee80211_is_data_qos(fc))
 		return 0;
 
 	status->aggr = unicast &&
diff --git a/drivers/net/wireless/mediatek/mt76/mt7921/mac.c b/drivers/net/wireless/mediatek/mt76/mt7921/mac.c
index 03b4960db73f..59c64b951678 100644
--- a/drivers/net/wireless/mediatek/mt76/mt7921/mac.c
+++ b/drivers/net/wireless/mediatek/mt76/mt7921/mac.c
@@ -184,6 +184,7 @@ mt7921_mac_fill_rx(struct mt792x_dev *dev, struct sk_buff *skb)
 	u32 rxd4 = le32_to_cpu(rxd[4]);
 	struct mt792x_sta *msta = NULL;
 	struct mt792x_link_sta *mlink;
+	struct mt76_wcid *wcid;
 	u16 seq_ctrl = 0;
 	__le16 fc = 0;
 	u8 mode = 0;
@@ -211,10 +212,11 @@ mt7921_mac_fill_rx(struct mt792x_dev *dev, struct sk_buff *skb)
 	chfreq = FIELD_GET(MT_RXD3_NORMAL_CH_FREQ, rxd3);
 	unicast = FIELD_GET(MT_RXD3_NORMAL_ADDR_TYPE, rxd3) == MT_RXD3_NORMAL_U2M;
 	idx = FIELD_GET(MT_RXD1_NORMAL_WLAN_IDX, rxd1);
-	status->wcid = mt792x_rx_get_wcid(dev, idx, unicast);
+	wcid = mt792x_rx_get_wcid(dev, idx, unicast);
+	status->wcid_idx = wcid ? wcid->idx : MT76_WCID_IDX_INVALID;
 
-	if (status->wcid) {
-		mlink = container_of(status->wcid, struct mt792x_link_sta, wcid);
+	if (wcid) {
+		mlink = container_of(wcid, struct mt792x_link_sta, wcid);
 		msta = container_of(mlink, struct mt792x_sta, deflink);
 		mt76_wcid_add_poll(&dev->mt76, &mlink->wcid);
 	}
@@ -395,7 +397,8 @@ mt7921_mac_fill_rx(struct mt792x_dev *dev, struct sk_buff *skb)
 
 		vif = container_of((void *)msta->vif, struct ieee80211_vif,
 				   drv_priv);
-		err = mt76_connac2_reverse_frag0_hdr_trans(vif, skb, hdr_gap);
+		err = mt76_connac2_reverse_frag0_hdr_trans(&dev->mt76, vif,
+							   skb, hdr_gap);
 		if (err)
 			return err;
 
@@ -433,7 +436,7 @@ mt7921_mac_fill_rx(struct mt792x_dev *dev, struct sk_buff *skb)
 	if (rxv && mode >= MT_PHY_TYPE_HE_SU && !(status->flag & RX_FLAG_8023))
 		mt76_connac2_mac_decode_he_radiotap(&dev->mt76, skb, rxv, mode);
 
-	if (!status->wcid || !ieee80211_is_data_qos(fc))
+	if (!wcid || !ieee80211_is_data_qos(fc))
 		return 0;
 
 	status->aggr = unicast && !ieee80211_is_qos_nullfunc(fc);
diff --git a/drivers/net/wireless/mediatek/mt76/mt7925/mac.c b/drivers/net/wireless/mediatek/mt76/mt7925/mac.c
index c47bd812b66b..e429f78596e2 100644
--- a/drivers/net/wireless/mediatek/mt76/mt7925/mac.c
+++ b/drivers/net/wireless/mediatek/mt76/mt7925/mac.c
@@ -166,17 +166,20 @@ void mt7925_mac_set_fixed_rate_table(struct mt792x_dev *dev,
 }
 
 /* The HW does not translate the mac header to 802.3 for mesh point */
-static int mt7925_reverse_frag0_hdr_trans(struct sk_buff *skb, u16 hdr_gap)
+static int mt7925_reverse_frag0_hdr_trans(struct mt792x_dev *dev,
+					  struct sk_buff *skb, u16 hdr_gap)
 {
 	struct mt76_rx_status *status = (struct mt76_rx_status *)skb->cb;
 	struct ethhdr *eth_hdr = (struct ethhdr *)(skb->data + hdr_gap);
-	struct mt792x_sta *msta = (struct mt792x_sta *)status->wcid;
 	__le32 *rxd = (__le32 *)skb->data;
 	struct ieee80211_sta *sta;
 	struct ieee80211_vif *vif;
 	struct ieee80211_hdr hdr;
+	struct mt792x_sta *msta;
 	u16 frame_control;
 
+	msta = (struct mt792x_sta *)mt76_wcid_ptr(dev, status->wcid_idx);
+
 	if (le32_get_bits(rxd[3], MT_RXD3_NORMAL_ADDR_TYPE) !=
 	    MT_RXD3_NORMAL_U2M)
 		return -EINVAL;
@@ -368,6 +371,7 @@ mt7925_mac_fill_rx(struct mt792x_dev *dev, struct sk_buff *skb)
 	u32 rxd3 = le32_to_cpu(rxd[3]);
 	u32 rxd4 = le32_to_cpu(rxd[4]);
 	struct mt792x_link_sta *mlink;
+	struct mt76_wcid *wcid;
 	u8 mode = 0; /* , band_idx; */
 	u16 seq_ctrl = 0;
 	__le16 fc = 0;
@@ -392,10 +396,11 @@ mt7925_mac_fill_rx(struct mt792x_dev *dev, struct sk_buff *skb)
 	chfreq = FIELD_GET(MT_RXD3_NORMAL_CH_FREQ, rxd3);
 	unicast = FIELD_GET(MT_RXD3_NORMAL_ADDR_TYPE, rxd3) == MT_RXD3_NORMAL_U2M;
 	idx = FIELD_GET(MT_RXD1_NORMAL_WLAN_IDX, rxd1);
-	status->wcid = mt792x_rx_get_wcid(dev, idx, unicast);
+	wcid = mt792x_rx_get_wcid(dev, idx, unicast);
+	status->wcid_idx = wcid ? wcid->idx : MT76_WCID_IDX_INVALID;
 
-	if (status->wcid) {
-		mlink = container_of(status->wcid, struct mt792x_link_sta, wcid);
+	if (wcid) {
+		mlink = container_of(wcid, struct mt792x_link_sta, wcid);
 		mt76_wcid_add_poll(&dev->mt76, &mlink->wcid);
 	}
 
@@ -545,7 +550,7 @@ mt7925_mac_fill_rx(struct mt792x_dev *dev, struct sk_buff *skb)
 
 	hdr_gap = (u8 *)rxd - skb->data + 2 * remove_pad;
 	if (hdr_trans && ieee80211_has_morefrags(fc)) {
-		if (mt7925_reverse_frag0_hdr_trans(skb, hdr_gap))
+		if (mt7925_reverse_frag0_hdr_trans(dev, skb, hdr_gap))
 			return -EINVAL;
 		hdr_trans = false;
 	} else {
@@ -608,7 +613,7 @@ mt7925_mac_fill_rx(struct mt792x_dev *dev, struct sk_buff *skb)
 		}
 	}
 
-	if (!status->wcid || !ieee80211_is_data_qos(fc))
+	if (!wcid || !ieee80211_is_data_qos(fc))
 		return 0;
 
 	status->aggr = unicast && !ieee80211_is_qos_nullfunc(fc);
diff --git a/drivers/net/wireless/mediatek/mt76/mt7996/mac.c b/drivers/net/wireless/mediatek/mt76/mt7996/mac.c
index e2a83da3a09c..98c30957dd76 100644
--- a/drivers/net/wireless/mediatek/mt76/mt7996/mac.c
+++ b/drivers/net/wireless/mediatek/mt76/mt7996/mac.c
@@ -186,17 +186,19 @@ static void mt7996_mac_sta_poll(struct mt7996_dev *dev)
 }
 
 /* The HW does not translate the mac header to 802.3 for mesh point */
-static int mt7996_reverse_frag0_hdr_trans(struct sk_buff *skb, u16 hdr_gap)
+static int mt7996_reverse_frag0_hdr_trans(struct mt7996_dev *dev,
+					  struct sk_buff *skb, u16 hdr_gap)
 {
 	struct mt76_rx_status *status = (struct mt76_rx_status *)skb->cb;
 	struct ethhdr *eth_hdr = (struct ethhdr *)(skb->data + hdr_gap);
-	struct mt7996_sta_link *msta_link = (void *)status->wcid;
-	struct mt7996_sta *msta = msta_link->sta;
+	struct mt76_wcid *wcid = mt76_wcid_ptr(dev, status->wcid_idx);
+	struct mt7996_sta_link *msta_link;
 	struct ieee80211_bss_conf *link_conf;
 	__le32 *rxd = (__le32 *)skb->data;
 	struct ieee80211_sta *sta;
 	struct ieee80211_vif *vif;
 	struct ieee80211_hdr hdr;
+	struct mt7996_sta *msta;
 	u16 frame_control;
 
 	if (le32_get_bits(rxd[3], MT_RXD3_NORMAL_ADDR_TYPE) !=
@@ -206,10 +208,16 @@ static int mt7996_reverse_frag0_hdr_trans(struct sk_buff *skb, u16 hdr_gap)
 	if (!(le32_to_cpu(rxd[1]) & MT_RXD1_NORMAL_GROUP_4))
 		return -EINVAL;
 
+	if (!wcid)
+		return -EINVAL;
+
+	msta_link = container_of(wcid, struct mt7996_sta_link, wcid);
+	msta = msta_link->sta;
+
 	if (!msta || !msta->vif)
 		return -EINVAL;
 
-	sta = wcid_to_sta(status->wcid);
+	sta = wcid_to_sta(wcid);
 	vif = container_of((void *)msta->vif, struct ieee80211_vif, drv_priv);
 	link_conf = rcu_dereference(vif->link_conf[msta_link->wcid.link_id]);
 	if (!link_conf)
@@ -429,6 +437,7 @@ mt7996_mac_fill_rx(struct mt7996_dev *dev, enum mt76_rxq_id q,
 	bool unicast, insert_ccmp_hdr = false;
 	u8 remove_pad, amsdu_info, band_idx;
 	u8 mode = 0, qos_ctl = 0;
+	struct mt76_wcid *wcid;
 	bool hdr_trans;
 	u16 hdr_gap;
 	u16 seq_ctrl = 0;
@@ -461,12 +470,14 @@ mt7996_mac_fill_rx(struct mt7996_dev *dev, enum mt76_rxq_id q,
 
 	unicast = FIELD_GET(MT_RXD3_NORMAL_ADDR_TYPE, rxd3) == MT_RXD3_NORMAL_U2M;
 	idx = FIELD_GET(MT_RXD1_NORMAL_WLAN_IDX, rxd1);
-	status->wcid = mt7996_rx_get_wcid(dev, idx, band_idx);
 
-	if (status->wcid) {
+	wcid = mt7996_rx_get_wcid(dev, idx, band_idx);
+	status->wcid_idx = wcid ? wcid->idx : MT76_WCID_IDX_INVALID;
+
+	if (wcid) {
 		struct mt7996_sta_link *msta_link;
 
-		msta_link = container_of(status->wcid, struct mt7996_sta_link,
+		msta_link = container_of(wcid, struct mt7996_sta_link,
 					 wcid);
 		msta = msta_link->sta;
 		mt76_wcid_add_poll(&dev->mt76, &msta_link->wcid);
@@ -620,7 +631,7 @@ mt7996_mac_fill_rx(struct mt7996_dev *dev, enum mt76_rxq_id q,
 
 	hdr_gap = (u8 *)rxd - skb->data + 2 * remove_pad;
 	if (hdr_trans && ieee80211_has_morefrags(fc)) {
-		if (mt7996_reverse_frag0_hdr_trans(skb, hdr_gap))
+		if (mt7996_reverse_frag0_hdr_trans(dev, skb, hdr_gap))
 			return -EINVAL;
 		hdr_trans = false;
 	} else {
@@ -697,7 +708,7 @@ mt7996_mac_fill_rx(struct mt7996_dev *dev, enum mt76_rxq_id q,
 		}
 	}
 
-	if (!status->wcid || !ieee80211_is_data_qos(fc) || hw_aggr)
+	if (!wcid || !ieee80211_is_data_qos(fc) || hw_aggr)
 		return 0;
 
 	status->aggr = unicast &&
-- 
2.45.2