wifi: rtw89: support WoWLAN for more chips and some fixes

index 23204b2706c9..3c666fd0f8ed 100644
@@ -13,22 +13,20 @@
 #include "ps.h"
 #include "reg.h"
 #include "util.h"
+#include "wow.h"
 
 struct rtw89_eapol_2_of_2 {
-	struct ieee80211_hdr_3addr hdr;
 	u8 gtkbody[14];
 	u8 key_des_ver;
 	u8 rsvd[92];
-} __packed __aligned(2);
+} __packed;
 
 struct rtw89_sa_query {
-	struct ieee80211_hdr_3addr hdr;
 	u8 category;
 	u8 action;
-} __packed __aligned(2);
+} __packed;
 
 struct rtw89_arp_rsp {
-	struct ieee80211_hdr_3addr addr;
 	u8 llc_hdr[sizeof(rfc1042_header)];
 	__be16 llc_type;
 	struct arphdr arp_hdr;

@@ -36,7 +34,7 @@ struct rtw89_arp_rsp {
 	__be32 sender_ip;
 	u8 target_hw[ETH_ALEN];
 	__be32 target_ip;
-} __packed __aligned(2);
+} __packed;
 
 static const u8 mss_signature[] = {0x4D, 0x53, 0x53, 0x4B, 0x50, 0x4F, 0x4F, 0x4C};
 

@@ -2197,8 +2195,10 @@ static struct sk_buff *rtw89_eapol_get(struct rtw89_dev *rtwdev,
 				     0x8E, 0x01, 0x03, 0x00, 0x5F, 0x02, 0x03};
 	struct ieee80211_vif *vif = rtwvif_to_vif(rtwvif);
 	struct ieee80211_bss_conf *bss_conf = &vif->bss_conf;
+	u8 sec_hdr_len = rtw89_wow_get_sec_hdr_len(rtwdev);
 	struct rtw89_wow_param *rtw_wow = &rtwdev->wow;
 	struct rtw89_eapol_2_of_2 *eapol_pkt;
+	struct ieee80211_hdr_3addr *hdr;
 	struct sk_buff *skb;
 	u8 key_des_ver;
 

@@ -2211,17 +2211,21 @@ static struct sk_buff *rtw89_eapol_get(struct rtw89_dev *rtwdev,
 	else
 		key_des_ver = 0;
 
-	skb = dev_alloc_skb(sizeof(*eapol_pkt));
+	skb = dev_alloc_skb(sizeof(*hdr) + sec_hdr_len + sizeof(*eapol_pkt));
 	if (!skb)
 		return NULL;
 
+	hdr = skb_put_zero(skb, sizeof(*hdr));
+	hdr->frame_control = cpu_to_le16(IEEE80211_FTYPE_DATA |
+					 IEEE80211_FCTL_TODS |
+					 IEEE80211_FCTL_PROTECTED);
+	ether_addr_copy(hdr->addr1, bss_conf->bssid);
+	ether_addr_copy(hdr->addr2, vif->addr);
+	ether_addr_copy(hdr->addr3, bss_conf->bssid);
+
+	skb_put_zero(skb, sec_hdr_len);
+
 	eapol_pkt = skb_put_zero(skb, sizeof(*eapol_pkt));
-	eapol_pkt->hdr.frame_control = cpu_to_le16(IEEE80211_FTYPE_DATA |
-						   IEEE80211_FCTL_TODS |
-						   IEEE80211_FCTL_PROTECTED);
-	ether_addr_copy(eapol_pkt->hdr.addr1, bss_conf->bssid);
-	ether_addr_copy(eapol_pkt->hdr.addr2, vif->addr);
-	ether_addr_copy(eapol_pkt->hdr.addr3, bss_conf->bssid);
 	memcpy(eapol_pkt->gtkbody, gtkbody, sizeof(gtkbody));
 	eapol_pkt->key_des_ver = key_des_ver;
 

@@ -2233,20 +2237,26 @@ static struct sk_buff *rtw89_sa_query_get(struct rtw89_dev *rtwdev,
 {
 	struct ieee80211_vif *vif = rtwvif_to_vif(rtwvif);
 	struct ieee80211_bss_conf *bss_conf = &vif->bss_conf;
+	u8 sec_hdr_len = rtw89_wow_get_sec_hdr_len(rtwdev);
+	struct ieee80211_hdr_3addr *hdr;
 	struct rtw89_sa_query *sa_query;
 	struct sk_buff *skb;
 
-	skb = dev_alloc_skb(sizeof(*sa_query));
+	skb = dev_alloc_skb(sizeof(*hdr) + sec_hdr_len + sizeof(*sa_query));
 	if (!skb)
 		return NULL;
 
+	hdr = skb_put_zero(skb, sizeof(*hdr));
+	hdr->frame_control = cpu_to_le16(IEEE80211_FTYPE_MGMT |
+					 IEEE80211_STYPE_ACTION |
+					 IEEE80211_FCTL_PROTECTED);
+	ether_addr_copy(hdr->addr1, bss_conf->bssid);
+	ether_addr_copy(hdr->addr2, vif->addr);
+	ether_addr_copy(hdr->addr3, bss_conf->bssid);
+
+	skb_put_zero(skb, sec_hdr_len);
+
 	sa_query = skb_put_zero(skb, sizeof(*sa_query));
-	sa_query->hdr.frame_control = cpu_to_le16(IEEE80211_FTYPE_MGMT |
-						  IEEE80211_STYPE_ACTION |
-						  IEEE80211_FCTL_PROTECTED);
-	ether_addr_copy(sa_query->hdr.addr1, bss_conf->bssid);
-	ether_addr_copy(sa_query->hdr.addr2, vif->addr);
-	ether_addr_copy(sa_query->hdr.addr3, bss_conf->bssid);
 	sa_query->category = WLAN_CATEGORY_SA_QUERY;
 	sa_query->action = WLAN_ACTION_SA_QUERY_RESPONSE;
 

@@ -2256,17 +2266,19 @@ static struct sk_buff *rtw89_sa_query_get(struct rtw89_dev *rtwdev,
 static struct sk_buff *rtw89_arp_response_get(struct rtw89_dev *rtwdev,
 					      struct rtw89_vif *rtwvif)
 {
+	u8 sec_hdr_len = rtw89_wow_get_sec_hdr_len(rtwdev);
 	struct rtw89_wow_param *rtw_wow = &rtwdev->wow;
+	struct ieee80211_hdr_3addr *hdr;
 	struct rtw89_arp_rsp *arp_skb;
 	struct arphdr *arp_hdr;
 	struct sk_buff *skb;
 	__le16 fc;
 
-	skb = dev_alloc_skb(sizeof(struct rtw89_arp_rsp));
+	skb = dev_alloc_skb(sizeof(*hdr) + sec_hdr_len + sizeof(*arp_skb));
 	if (!skb)
 		return NULL;
 
-	arp_skb = skb_put_zero(skb, sizeof(*arp_skb));
+	hdr = skb_put_zero(skb, sizeof(*hdr));
 
 	if (rtw_wow->ptk_alg)
 		fc = cpu_to_le16(IEEE80211_FTYPE_DATA | IEEE80211_FCTL_TODS |

@@ -2274,11 +2286,14 @@ static struct sk_buff *rtw89_arp_response_get(struct rtw89_dev *rtwdev,
 	else
 		fc = cpu_to_le16(IEEE80211_FTYPE_DATA | IEEE80211_FCTL_TODS);
 
-	arp_skb->addr.frame_control = fc;
-	ether_addr_copy(arp_skb->addr.addr1, rtwvif->bssid);
-	ether_addr_copy(arp_skb->addr.addr2, rtwvif->mac_addr);
-	ether_addr_copy(arp_skb->addr.addr3, rtwvif->bssid);
+	hdr->frame_control = fc;
+	ether_addr_copy(hdr->addr1, rtwvif->bssid);
+	ether_addr_copy(hdr->addr2, rtwvif->mac_addr);
+	ether_addr_copy(hdr->addr3, rtwvif->bssid);
+
+	skb_put_zero(skb, sec_hdr_len);
 
+	arp_skb = skb_put_zero(skb, sizeof(*arp_skb));
 	memcpy(arp_skb->llc_hdr, rfc1042_header, sizeof(rfc1042_header));
 	arp_skb->llc_type = htons(ETH_P_ARP);
 

index fa61484c3839..df22a53b8f41 100644
@@ -27,17 +27,23 @@ void rtw89_wow_parse_akm(struct rtw89_dev *rtwdev, struct sk_buff *skb)
 	rtw_wow->akm = rsn_ie->akm_cipher_suite.type;
 }
 
+#define RTW89_CIPHER_INFO_DEF(cipher) \
+	{WLAN_CIPHER_SUITE_ ## cipher, .fw_alg = RTW89_WOW_FW_ALG_ ## cipher, \
+	 .len = WLAN_KEY_LEN_ ## cipher}
+
 static const struct rtw89_cipher_info rtw89_cipher_info_defs[] = {
-	{WLAN_CIPHER_SUITE_WEP40,	.fw_alg = 1,	.len = WLAN_KEY_LEN_WEP40,},
-	{WLAN_CIPHER_SUITE_WEP104,	.fw_alg = 2,	.len = WLAN_KEY_LEN_WEP104,},
-	{WLAN_CIPHER_SUITE_TKIP,	.fw_alg = 3,	.len = WLAN_KEY_LEN_TKIP,},
-	{WLAN_CIPHER_SUITE_CCMP,	.fw_alg = 6,	.len = WLAN_KEY_LEN_CCMP,},
-	{WLAN_CIPHER_SUITE_GCMP,	.fw_alg = 8,	.len = WLAN_KEY_LEN_GCMP,},
-	{WLAN_CIPHER_SUITE_CCMP_256,	.fw_alg = 7,	.len = WLAN_KEY_LEN_CCMP_256,},
-	{WLAN_CIPHER_SUITE_GCMP_256,	.fw_alg = 23,	.len = WLAN_KEY_LEN_GCMP_256,},
-	{WLAN_CIPHER_SUITE_AES_CMAC,	.fw_alg = 32,	.len = WLAN_KEY_LEN_AES_CMAC,},
+	RTW89_CIPHER_INFO_DEF(WEP40),
+	RTW89_CIPHER_INFO_DEF(WEP104),
+	RTW89_CIPHER_INFO_DEF(TKIP),
+	RTW89_CIPHER_INFO_DEF(CCMP),
+	RTW89_CIPHER_INFO_DEF(GCMP),
+	RTW89_CIPHER_INFO_DEF(CCMP_256),
+	RTW89_CIPHER_INFO_DEF(GCMP_256),
+	RTW89_CIPHER_INFO_DEF(AES_CMAC),
 };
 
+#undef RTW89_CIPHER_INFO_DEF
+
 static const
 struct rtw89_cipher_info *rtw89_cipher_alg_recognize(u32 cipher)
 {

index e595aee0196d..0d90add0e88d 100644
@@ -35,6 +35,17 @@ enum rtw89_wake_reason {
 	RTW89_WOW_RSN_RX_NLO = 0x55,
 };
 
+enum rtw89_fw_alg {
+	RTW89_WOW_FW_ALG_WEP40 = 0x1,
+	RTW89_WOW_FW_ALG_WEP104 = 0x2,
+	RTW89_WOW_FW_ALG_TKIP = 0x3,
+	RTW89_WOW_FW_ALG_CCMP = 0x6,
+	RTW89_WOW_FW_ALG_CCMP_256 = 0x7,
+	RTW89_WOW_FW_ALG_GCMP = 0x8,
+	RTW89_WOW_FW_ALG_GCMP_256 = 0x9,
+	RTW89_WOW_FW_ALG_AES_CMAC = 0xa,
+};
+
 struct rtw89_cipher_suite {
 	u8 oui[3];
 	u8 type;

@@ -64,6 +75,25 @@ struct rtw89_set_key_info_iter_data {
 	bool error;
 };
 
+static inline int rtw89_wow_get_sec_hdr_len(struct rtw89_dev *rtwdev)
+{
+	struct rtw89_wow_param *rtw_wow = &rtwdev->wow;
+
+	if (!(rtwdev->chip->chip_id == RTL8852A || rtw89_is_rtl885xb(rtwdev)))
+		return 0;
+
+	switch (rtw_wow->ptk_alg) {
+	case RTW89_WOW_FW_ALG_WEP40:
+		return 4;
+	case RTW89_WOW_FW_ALG_TKIP:
+	case RTW89_WOW_FW_ALG_CCMP:
+	case RTW89_WOW_FW_ALG_GCMP_256:
+		return 8;
+	default:
+		return 0;
+	}
+}
+
 #ifdef CONFIG_PM
 int rtw89_wow_suspend(struct rtw89_dev *rtwdev, struct cfg80211_wowlan *wowlan);
 int rtw89_wow_resume(struct rtw89_dev *rtwdev);

index e8ec0246bccd..f1de975c5356 100644
@@ -4236,7 +4236,7 @@ struct rtw89_chip_info {
 	const u32 *c2h_regs;
 	struct rtw89_reg_def c2h_counter_reg;
 	const struct rtw89_page_regs *page_regs;
-	u32 wow_reason_reg;
+	const u32 *wow_reason_reg;
 	bool cfo_src_fd;
 	bool cfo_hw_comp;
 	const struct rtw89_reg_def *dcfo_comp;

@@ -4346,6 +4346,7 @@ enum rtw89_fw_feature {
 	RTW89_FW_FEATURE_NO_LPS_PG,
 	RTW89_FW_FEATURE_BEACON_FILTER,
 	RTW89_FW_FEATURE_MACID_PAUSE_SLEEP,
+	RTW89_FW_FEATURE_WOW_REASON_V1,
 };
 
 struct rtw89_fw_suit {

index 3c666fd0f8ed..a907d381ac77 100644
@@ -675,10 +675,12 @@ static const struct __fw_feat_cfg fw_feat_tbl[] = {
 	__CFG_FW_FEAT(RTL8852C, ge, 0, 27, 36, 0, SCAN_OFFLOAD),
 	__CFG_FW_FEAT(RTL8852C, ge, 0, 27, 40, 0, CRASH_TRIGGER),
 	__CFG_FW_FEAT(RTL8852C, ge, 0, 27, 56, 10, BEACON_FILTER),
+	__CFG_FW_FEAT(RTL8852C, ge, 0, 27, 80, 0, WOW_REASON_V1),
 	__CFG_FW_FEAT(RTL8922A, ge, 0, 34, 30, 0, CRASH_TRIGGER),
 	__CFG_FW_FEAT(RTL8922A, ge, 0, 34, 11, 0, MACID_PAUSE_SLEEP),
 	__CFG_FW_FEAT(RTL8922A, ge, 0, 34, 35, 0, SCAN_OFFLOAD),
 	__CFG_FW_FEAT(RTL8922A, ge, 0, 35, 12, 0, BEACON_FILTER),
+	__CFG_FW_FEAT(RTL8922A, ge, 0, 35, 22, 0, WOW_REASON_V1),
 };
 
 static void rtw89_fw_iterate_feature_cfg(struct rtw89_fw_info *fw,

index 01fea0b004d3..c3b4324c621c 100644
@@ -4659,4 +4659,10 @@ const struct rtw89_rfe_parms *
 rtw89_load_rfe_data_from_fw(struct rtw89_dev *rtwdev,
 			    const struct rtw89_rfe_parms *init);
 
+enum rtw89_wow_wakeup_ver {
+	RTW89_WOW_REASON_V0,
+	RTW89_WOW_REASON_V1,
+	RTW89_WOW_REASON_NUM,
+};
+
 #endif

index 5a59234301ca..fd81cf863bea 100644
@@ -311,6 +311,8 @@
 #define B_AX_S1_LDO2PWRCUT_F BIT(23)
 #define B_AX_S0_LDO_VSEL_F_MASK GENMASK(22, 21)
 
+#define R_AX_DBG_WOW 0x0504
+
 #define R_AX_SEC_CTRL 0x0C00
 #define B_AX_SEC_IDMEM_SIZE_CONFIG_MASK GENMASK(17, 16)
 

@@ -4315,6 +4317,8 @@
 
 #define R_BE_WLCPU_PORT_PC 0x03FC
 
+#define R_BE_DBG_WOW 0x0504
+
 #define R_BE_DCPU_PLATFORM_ENABLE 0x0888
 #define B_BE_DCPU_SYM_DPLT_MEM_MUX_EN BIT(10)
 #define B_BE_DCPU_WARM_EN BIT(9)

index 78e276f321c2..40cf84a79c46 100644
@@ -105,6 +105,10 @@ static const u32 rtw8851b_c2h_regs[RTW89_C2HREG_MAX] = {
 	R_AX_C2HREG_DATA3
 };
 
+static const u32 rtw8851b_wow_wakeup_regs[RTW89_WOW_REASON_NUM] = {
+	R_AX_C2HREG_DATA3 + 3, R_AX_C2HREG_DATA3 + 3,
+};
+
 static const struct rtw89_page_regs rtw8851b_page_regs = {
 	.hci_fc_ctrl	= R_AX_HCI_FC_CTRL,
 	.ch_page_ctrl	= R_AX_CH_PAGE_CTRL,

@@ -2509,7 +2513,7 @@ const struct rtw89_chip_info rtw8851b_chip_info = {
 	.c2h_counter_reg	= {R_AX_UDM1 + 1, B_AX_UDM1_HALMAC_C2H_ENQ_CNT_MASK >> 8},
 	.c2h_regs		= rtw8851b_c2h_regs,
 	.page_regs		= &rtw8851b_page_regs,
-	.wow_reason_reg		= R_AX_C2HREG_DATA3 + 3,
+	.wow_reason_reg		= rtw8851b_wow_wakeup_regs,
 	.cfo_src_fd		= true,
 	.cfo_hw_comp		= true,
 	.dcfo_comp		= &rtw8851b_dcfo_comp,

index aebbceea93f8..08e148328c62 100644
@@ -398,6 +398,10 @@ static const u32 rtw8852a_c2h_regs[RTW89_C2HREG_MAX] = {
 	R_AX_C2HREG_DATA3
 };
 
+static const u32 rtw8852a_wow_wakeup_regs[RTW89_WOW_REASON_NUM] = {
+	R_AX_C2HREG_DATA3 + 3, R_AX_C2HREG_DATA3 + 3,
+};
+
 static const struct rtw89_page_regs rtw8852a_page_regs = {
 	.hci_fc_ctrl	= R_AX_HCI_FC_CTRL,
 	.ch_page_ctrl	= R_AX_CH_PAGE_CTRL,

@@ -2225,7 +2229,7 @@ const struct rtw89_chip_info rtw8852a_chip_info = {
 	.c2h_regs		= rtw8852a_c2h_regs,
 	.c2h_counter_reg	= {R_AX_UDM1 + 1, B_AX_UDM1_HALMAC_C2H_ENQ_CNT_MASK >> 8},
 	.page_regs		= &rtw8852a_page_regs,
-	.wow_reason_reg		= R_AX_C2HREG_DATA3 + 3,
+	.wow_reason_reg		= rtw8852a_wow_wakeup_regs,
 	.cfo_src_fd		= false,
 	.cfo_hw_comp            = false,
 	.dcfo_comp		= &rtw8852a_dcfo_comp,

index 381c881fee69..24f9fd444ca0 100644
@@ -76,6 +76,10 @@ static const u32 rtw8852b_c2h_regs[RTW89_C2HREG_MAX] = {
 	R_AX_C2HREG_DATA3
 };
 
+static const u32 rtw8852b_wow_wakeup_regs[RTW89_WOW_REASON_NUM] = {
+	R_AX_C2HREG_DATA3 + 3, R_AX_C2HREG_DATA3 + 3,
+};
+
 static const struct rtw89_page_regs rtw8852b_page_regs = {
 	.hci_fc_ctrl	= R_AX_HCI_FC_CTRL,
 	.ch_page_ctrl	= R_AX_CH_PAGE_CTRL,

@@ -1026,7 +1030,7 @@ const struct rtw89_chip_info rtw8852b_chip_info = {
 	.c2h_counter_reg	= {R_AX_UDM1 + 1, B_AX_UDM1_HALMAC_C2H_ENQ_CNT_MASK >> 8},
 	.c2h_regs		= rtw8852b_c2h_regs,
 	.page_regs		= &rtw8852b_page_regs,
-	.wow_reason_reg		= R_AX_C2HREG_DATA3 + 3,
+	.wow_reason_reg		= rtw8852b_wow_wakeup_regs,
 	.cfo_src_fd		= true,
 	.cfo_hw_comp		= true,
 	.dcfo_comp		= &rtw8852b_dcfo_comp,

index 544ea1d3b577..193168dc7b6c 100644
@@ -73,6 +73,10 @@ static const u32 rtw8852c_c2h_regs[RTW89_H2CREG_MAX] = {
 	R_AX_C2HREG_DATA3_V1
 };
 
+static const u32 rtw8852c_wow_wakeup_regs[RTW89_WOW_REASON_NUM] = {
+	R_AX_C2HREG_DATA3_V1 + 3, R_AX_DBG_WOW,
+};
+
 static const struct rtw89_page_regs rtw8852c_page_regs = {
 	.hci_fc_ctrl	= R_AX_HCI_FC_CTRL_V1,
 	.ch_page_ctrl	= R_AX_CH_PAGE_CTRL_V1,

@@ -3007,7 +3011,7 @@ const struct rtw89_chip_info rtw8852c_chip_info = {
 	.c2h_counter_reg	= {R_AX_UDM1 + 1, B_AX_UDM1_HALMAC_C2H_ENQ_CNT_MASK >> 8},
 	.c2h_regs		= rtw8852c_c2h_regs,
 	.page_regs		= &rtw8852c_page_regs,
-	.wow_reason_reg		= R_AX_C2HREG_DATA3_V1 + 3,
+	.wow_reason_reg		= rtw8852c_wow_wakeup_regs,
 	.cfo_src_fd		= false,
 	.cfo_hw_comp            = false,
 	.dcfo_comp		= &rtw8852c_dcfo_comp,

index cc4f251caadd..2af568a3264d 100644
@@ -85,6 +85,10 @@ static const u32 rtw8922a_c2h_regs[RTW89_H2CREG_MAX] = {
 	R_BE_C2HREG_DATA3
 };
 
+static const u32 rtw8922a_wow_wakeup_regs[RTW89_WOW_REASON_NUM] = {
+	R_AX_C2HREG_DATA3_V1 + 3, R_BE_DBG_WOW,
+};
+
 static const struct rtw89_page_regs rtw8922a_page_regs = {
 	.hci_fc_ctrl	= R_BE_HCI_FC_CTRL,
 	.ch_page_ctrl	= R_BE_CH_PAGE_CTRL,

@@ -2609,7 +2613,7 @@ const struct rtw89_chip_info rtw8922a_chip_info = {
 	.c2h_counter_reg	= {R_BE_UDM1 + 1, B_BE_UDM1_HALMAC_C2H_ENQ_CNT_MASK >> 8},
 	.c2h_regs		= rtw8922a_c2h_regs,
 	.page_regs		= &rtw8922a_page_regs,
-	.wow_reason_reg		= R_AX_C2HREG_DATA3_V1 + 3,
+	.wow_reason_reg		= rtw8922a_wow_wakeup_regs,
 	.cfo_src_fd		= true,
 	.cfo_hw_comp            = true,
 	.dcfo_comp		= NULL,

index df22a53b8f41..a49265cda98f 100644
@@ -723,13 +723,18 @@ static void rtw89_wow_show_wakeup_reason(struct rtw89_dev *rtwdev)
 {
 	struct rtw89_wow_param *rtw_wow = &rtwdev->wow;
 	struct rtw89_wow_aoac_report *aoac_rpt = &rtw_wow->aoac_rpt;
-	u32 wow_reason_reg = rtwdev->chip->wow_reason_reg;
 	struct cfg80211_wowlan_nd_info nd_info;
 	struct cfg80211_wowlan_wakeup wakeup = {
 		.pattern_idx = -1,
 	};
+	u32 wow_reason_reg;
 	u8 reason;
 
+	if (RTW89_CHK_FW_FEATURE(WOW_REASON_V1, &rtwdev->fw))
+		wow_reason_reg = rtwdev->chip->wow_reason_reg[RTW89_WOW_REASON_V1];
+	else
+		wow_reason_reg = rtwdev->chip->wow_reason_reg[RTW89_WOW_REASON_V0];
+
 	reason = rtw89_read8(rtwdev, wow_reason_reg);
 	switch (reason) {
 	case RTW89_WOW_RSN_RX_DEAUTH:

index b759fa8fbca9..73462f3343e3 100644
@@ -6320,9 +6320,30 @@ int rtw89_mac_ptk_drop_by_band_and_wait(struct rtw89_dev *rtwdev,
 	return ret;
 }
 
+int rtw89_mac_cpu_io_rx(struct rtw89_dev *rtwdev, bool wow_enable)
+{
+	struct rtw89_mac_h2c_info h2c_info = {};
+	struct rtw89_mac_c2h_info c2h_info = {};
+	u32 ret;
+
+	h2c_info.id = RTW89_FWCMD_H2CREG_FUNC_WOW_CPUIO_RX_CTRL;
+	h2c_info.content_len = sizeof(h2c_info.u.hdr);
+	h2c_info.u.hdr.w0 = u32_encode_bits(wow_enable, RTW89_H2CREG_WOW_CPUIO_RX_CTRL_EN);
+
+	ret = rtw89_fw_msg_reg(rtwdev, &h2c_info, &c2h_info);
+	if (ret)
+		return ret;
+
+	if (c2h_info.id != RTW89_FWCMD_C2HREG_FUNC_WOW_CPUIO_RX_ACK)
+		ret = -EINVAL;
+
+	return ret;
+}
+
 static int rtw89_wow_config_mac_ax(struct rtw89_dev *rtwdev, bool enable_wow)
 {
 	const struct rtw89_mac_gen_def *mac = rtwdev->chip->mac_def;
+	const struct rtw89_chip_info *chip = rtwdev->chip;
 	int ret;
 
 	if (enable_wow) {

@@ -6333,12 +6354,19 @@ static int rtw89_wow_config_mac_ax(struct rtw89_dev *rtwdev, bool enable_wow)
 		}
 
 		rtw89_write32_set(rtwdev, R_AX_RX_FUNCTION_STOP, B_AX_HDR_RX_STOP);
+		rtw89_mac_cpu_io_rx(rtwdev, enable_wow);
 		rtw89_write32_clr(rtwdev, mac->rx_fltr, B_AX_SNIFFER_MODE);
 		rtw89_mac_cfg_ppdu_status(rtwdev, RTW89_MAC_0, false);
 		rtw89_write32(rtwdev, R_AX_ACTION_FWD0, 0);
 		rtw89_write32(rtwdev, R_AX_ACTION_FWD1, 0);
 		rtw89_write32(rtwdev, R_AX_TF_FWD, 0);
 		rtw89_write32(rtwdev, R_AX_HW_RPT_FWD, 0);
+
+		if (chip->chip_id == RTL8852A || rtw89_is_rtl885xb(rtwdev))
+			rtw89_write8(rtwdev, R_BE_DBG_WOW_READY, WOWLAN_NOT_READY);
+		else
+			rtw89_write32_set(rtwdev, R_AX_DBG_WOW,
+					  B_AX_DBG_WOW_CPU_IO_RX_EN);
 	} else {
 		ret = rtw89_mac_resize_ple_rx_quota(rtwdev, false);
 		if (ret) {

@@ -6346,6 +6374,7 @@ static int rtw89_wow_config_mac_ax(struct rtw89_dev *rtwdev, bool enable_wow)
 			return ret;
 		}
 
+		rtw89_mac_cpu_io_rx(rtwdev, enable_wow);
 		rtw89_write32_clr(rtwdev, R_AX_RX_FUNCTION_STOP, B_AX_HDR_RX_STOP);
 		rtw89_mac_cfg_ppdu_status(rtwdev, RTW89_MAC_0, true);
 		rtw89_write32(rtwdev, R_AX_ACTION_FWD0, TRXCFG_MPDU_PROC_ACT_FRWD);

index a580cb719233..cc9f914d440c 100644
@@ -1446,5 +1446,6 @@ int rtw89_mac_dle_quota_change(struct rtw89_dev *rtwdev, enum rtw89_qta_mode mod
 int rtw89_mac_get_dle_rsvd_qt_cfg(struct rtw89_dev *rtwdev,
 				  enum rtw89_mac_dle_rsvd_qt_type type,
 				  struct rtw89_mac_dle_rsvd_qt_cfg *cfg);
+int rtw89_mac_cpu_io_rx(struct rtw89_dev *rtwdev, bool wow_enable);
 
 #endif

index 934bdf3b398f..f212b67771d5 100644
@@ -2312,26 +2312,6 @@ static void rtw89_mac_dump_qta_lost_be(struct rtw89_dev *rtwdev)
 	dump_err_status_dispatcher_be(rtwdev);
 }
 
-static int rtw89_mac_cpu_io_rx(struct rtw89_dev *rtwdev, bool wow_enable)
-{
-	struct rtw89_mac_h2c_info h2c_info = {};
-	struct rtw89_mac_c2h_info c2h_info = {};
-	u32 ret;
-
-	h2c_info.id = RTW89_FWCMD_H2CREG_FUNC_WOW_CPUIO_RX_CTRL;
-	h2c_info.content_len = sizeof(h2c_info.u.hdr);
-	h2c_info.u.hdr.w0 = u32_encode_bits(wow_enable, RTW89_H2CREG_WOW_CPUIO_RX_CTRL_EN);
-
-	ret = rtw89_fw_msg_reg(rtwdev, &h2c_info, &c2h_info);
-	if (ret)
-		return ret;
-
-	if (c2h_info.id != RTW89_FWCMD_C2HREG_FUNC_WOW_CPUIO_RX_ACK)
-		ret = -EINVAL;
-
-	return ret;
-}
-
 static int rtw89_wow_config_mac_be(struct rtw89_dev *rtwdev, bool enable_wow)
 {
 	if (enable_wow) {

index fd81cf863bea..c9f6bdc84d03 100644
@@ -312,6 +312,7 @@
 #define B_AX_S0_LDO_VSEL_F_MASK GENMASK(22, 21)
 
 #define R_AX_DBG_WOW 0x0504
+#define B_AX_DBG_WOW_CPU_IO_RX_EN BIT(8)
 
 #define R_AX_SEC_CTRL 0x0C00
 #define B_AX_SEC_IDMEM_SIZE_CONFIG_MASK GENMASK(17, 16)

index a907d381ac77..e3fc9f168ae5 100644
@@ -6767,10 +6767,8 @@ int rtw89_fw_h2c_wow_gtk_ofld(struct rtw89_dev *rtwdev,
 	skb_put(skb, len);
 	h2c = (struct rtw89_h2c_wow_gtk_ofld *)skb->data;
 
-	if (!enable) {
-		skb_put_zero(skb, sizeof(*gtk_info));
+	if (!enable)
 		goto hdr;
-	}
 
 	ret = rtw89_fw_h2c_add_general_pkt(rtwdev, rtwvif,
 					   RTW89_PKT_OFLD_TYPE_EAPOL_KEY,

index 87c38cdc691b..4acd4301d3c2 100644
@@ -1921,7 +1921,8 @@ static void rtw89_vif_rx_stats_iter(void *data, u8 *mac,
 		return;
 
 	if (ieee80211_is_beacon(hdr->frame_control)) {
-		if (vif->type == NL80211_IFTYPE_STATION) {
+		if (vif->type == NL80211_IFTYPE_STATION &&
+		    !test_bit(RTW89_FLAG_WOWLAN, rtwdev->flags)) {
 			rtw89_vif_sync_bcn_tsf(rtwvif, hdr, skb->len);
 			rtw89_fw_h2c_rssi_offload(rtwdev, phy_ppdu);
 		}

index a49265cda98f..9882064ef68d 100644
@@ -1295,12 +1295,16 @@ static int rtw89_wow_disable_trx_pre(struct rtw89_dev *rtwdev)
 
 static int rtw89_wow_disable_trx_post(struct rtw89_dev *rtwdev)
 {
+	struct rtw89_wow_param *rtw_wow = &rtwdev->wow;
+	struct ieee80211_vif *vif = rtw_wow->wow_vif;
 	int ret;
 
 	ret = rtw89_mac_cfg_ppdu_status(rtwdev, RTW89_MAC_0, true);
 	if (ret)
 		rtw89_err(rtwdev, "cfg ppdu status\n");
 
+	rtw89_fw_h2c_set_bcn_fltr_cfg(rtwdev, vif, true);
+
 	return ret;
 }
 

index e3fc9f168ae5..fbe08c162b93 100644
@@ -2496,6 +2496,7 @@ int rtw89_fw_h2c_lps_ch_info(struct rtw89_dev *rtwdev, struct rtw89_vif *rtwvif)
 	struct rtw89_h2c_lps_ch_info *h2c;
 	u32 len = sizeof(*h2c);
 	struct sk_buff *skb;
+	u32 done;
 	int ret;
 
 	if (chip->chip_gen != RTW89_CHIP_BE)

@@ -2519,12 +2520,18 @@ int rtw89_fw_h2c_lps_ch_info(struct rtw89_dev *rtwdev, struct rtw89_vif *rtwvif)
 			      H2C_CAT_OUTSRC, H2C_CL_OUTSRC_DM,
 			      H2C_FUNC_FW_LPS_CH_INFO, 0, 0, len);
 
+	rtw89_phy_write32_mask(rtwdev, R_CHK_LPS_STAT, B_CHK_LPS_STAT, 0);
 	ret = rtw89_h2c_tx(rtwdev, skb, false);
 	if (ret) {
 		rtw89_err(rtwdev, "failed to send h2c\n");
 		goto fail;
 	}
 
+	ret = read_poll_timeout(rtw89_phy_read32_mask, done, done, 50, 5000,
+				true, rtwdev, R_CHK_LPS_STAT, B_CHK_LPS_STAT);
+	if (ret)
+		rtw89_warn(rtwdev, "h2c_lps_ch_info done polling timeout\n");
+
 	return 0;
 fail:
 	dev_kfree_skb_any(skb);

index c9f6bdc84d03..897884e9667f 100644
@@ -7818,6 +7818,8 @@
 #define B_UPD_P0_EN BIT(31)
 #define R_EMLSR 0x0044
 #define B_EMLSR_PARM GENMASK(27, 12)
+#define R_CHK_LPS_STAT 0x0058
+#define B_CHK_LPS_STAT BIT(0)
 #define R_SPOOF_CG 0x00B4
 #define B_SPOOF_CG_EN BIT(17)
 #define R_CHINFO_SEG 0x00B4