*/
        TLC_MNG_UPDATE_NOTIF = 0xF7,
 
+       /**
+        * @BEACON_FILTER_IN_NOTIF: &struct iwl_beacon_filter_notif
+        */
+       BEACON_FILTER_IN_NOTIF = 0xF8,
+
        /**
         * @STA_PM_NOTIF: &struct iwl_mvm_pm_state_notification
         */
 
 };
 
 enum iwl_rx_mpdu_mac_phy_band {
+       /* whether or not this is MAC or LINK depends on the API */
        IWL_RX_MPDU_MAC_PHY_BAND_MAC_MASK       = 0x0f,
+       IWL_RX_MPDU_MAC_PHY_BAND_LINK_MASK      = 0x0f,
        IWL_RX_MPDU_MAC_PHY_BAND_PHY_MASK       = 0x30,
        IWL_RX_MPDU_MAC_PHY_BAND_BAND_MASK      = 0xc0,
 };
         */
        __le16 phy_info;
        /**
-        * @mac_phy_band: MAC ID, PHY ID, band;
+        * @mac_phy_band: MAC/link ID, PHY ID, band;
         *      see &enum iwl_rx_mpdu_mac_phy_band
         */
        u8 mac_phy_band;
        struct iwl_rfh_queue_data data[];
 } __packed; /* RFH_QUEUE_CONFIG_API_S_VER_1 */
 
+/**
+ * struct iwl_beacon_filter_notif_v1 - beacon filter notification
+ * @average_energy: average energy for the received beacon
+ * @mac_id: MAC ID the beacon was received for
+ */
+struct iwl_beacon_filter_notif_v1 {
+       __le32 average_energy;
+       __le32 mac_id;
+} __packed; /* BEACON_FILTER_IN_NTFY_API_S_VER_1 */
+
+/**
+ * struct iwl_beacon_filter_notif - beacon filter notification
+ * @average_energy: average energy for the received beacon
+ * @link_id: link ID the beacon was received for
+ */
+struct iwl_beacon_filter_notif {
+       __le32 average_energy;
+       __le32 link_id;
+} __packed; /* BEACON_FILTER_IN_NTFY_API_S_VER_2 */
+
 #endif /* __iwl_fw_api_rx_h__ */
 
 {
        mld_link->vif = link->vif;
        mld_link->link_id = link->link_id;
+       mld_link->average_beacon_energy = 0;
 
        iwl_mld_init_internal_sta(&mld_link->bcast_sta);
        iwl_mld_init_internal_sta(&mld_link->mcast_sta);
        return grade;
 }
 EXPORT_SYMBOL_IF_IWLWIFI_KUNIT(iwl_mld_get_link_grade);
+
+void iwl_mld_handle_beacon_filter_notif(struct iwl_mld *mld,
+                                       struct iwl_rx_packet *pkt)
+{
+       const struct iwl_beacon_filter_notif *notif = (const void *)pkt->data;
+       u32 link_id = le32_to_cpu(notif->link_id);
+       struct ieee80211_bss_conf *link_conf =
+               iwl_mld_fw_id_to_link_conf(mld, link_id);
+       struct iwl_mld_link *mld_link;
+
+       if (IWL_FW_CHECK(mld, !link_conf, "invalid link ID %d\n", link_id))
+               return;
+
+       mld_link = iwl_mld_link_from_mac80211(link_conf);
+       if (WARN_ON_ONCE(!mld_link))
+               return;
+
+       mld_link->average_beacon_energy = le32_to_cpu(notif->average_energy);
+}
 
  * @mcast_sta: station used for multicast packets. Used in AP, GO and IBSS.
  * @mon_sta: station used for TX injection in monitor interface.
  * @link_id: over the air link ID
+ * @average_beacon_energy: average beacon energy for beacons received during
+ *     client connections
  * @ap_early_keys: The firmware cannot install keys before bcast/mcast STAs,
  *     but higher layers work differently, so we store the keys here for
  *     later installation.
 
        /* we can only have 2 GTK + 2 IGTK + 2 BIGTK active at a time */
        struct ieee80211_key_conf *ap_early_keys[6];
+       u32 average_beacon_energy;
        bool silent_deactivation;
        struct iwl_probe_resp_data __rcu *probe_resp_data;
 };
                               struct ieee80211_bss_conf *link_conf,
                               enum ieee80211_sta_rx_bandwidth bw);
 
+void iwl_mld_handle_beacon_filter_notif(struct iwl_mld *mld,
+                                       struct iwl_rx_packet *pkt);
+
 #endif /* __iwl_mld_link_h__ */
 
        HCMD_NAME(ESR_MODE_NOTIF),
        HCMD_NAME(MONITOR_NOTIF),
        HCMD_NAME(TLC_MNG_UPDATE_NOTIF),
+       HCMD_NAME(BEACON_FILTER_IN_NOTIF),
        HCMD_NAME(MU_GROUP_MGMT_NOTIF),
 };
 
 
 CMD_VERSIONS(omi_status_notif,
             CMD_VER_ENTRY(1, iwl_omi_send_status_notif))
 CMD_VERSIONS(ftm_resp_notif, CMD_VER_ENTRY(10, iwl_tof_range_rsp_ntfy))
+CMD_VERSIONS(beacon_filter_notif, CMD_VER_ENTRY(2, iwl_beacon_filter_notif))
 
 DEFINE_SIMPLE_CANCELLATION(session_prot, iwl_session_prot_notif, mac_link_id)
 DEFINE_SIMPLE_CANCELLATION(tlc, iwl_tlc_update_notif, sta_id)
                           mac_id)
 #define iwl_mld_cancel_omi_status_notif iwl_mld_always_cancel
 DEFINE_SIMPLE_CANCELLATION(ftm_resp, iwl_tof_range_rsp_ntfy, request_id)
+DEFINE_SIMPLE_CANCELLATION(beacon_filter, iwl_beacon_filter_notif, link_id)
 
 /**
  * DOC: Handlers for fw notifications
                             time_sync_confirm_notif, RX_HANDLER_ASYNC)
        RX_HANDLER_OF_LINK(DATA_PATH_GROUP, OMI_SEND_STATUS_NOTIF,
                           omi_status_notif)
+       RX_HANDLER_OF_LINK(DATA_PATH_GROUP, BEACON_FILTER_IN_NOTIF,
+                          beacon_filter_notif)
        RX_HANDLER_OF_FTM_REQ(LOCATION_GROUP, TOF_RANGE_RESPONSE_NOTIF,
                              ftm_resp_notif)
 };
 
 }
 EXPORT_SYMBOL_IF_IWLWIFI_KUNIT(iwl_mld_pass_packet_to_mac80211);
 
-static void iwl_mld_fill_signal(struct iwl_mld *mld,
+static bool iwl_mld_used_average_energy(struct iwl_mld *mld, int link_id,
+                                       struct ieee80211_hdr *hdr,
+                                       struct ieee80211_rx_status *rx_status)
+{
+       struct ieee80211_bss_conf *link_conf;
+       struct iwl_mld_link *mld_link;
+
+       if (unlikely(!hdr || link_id < 0))
+               return false;
+
+       if (likely(!ieee80211_is_beacon(hdr->frame_control)))
+               return false;
+
+       /*
+        * if link ID is >= valid ones then that means the RX
+        * was on the AUX link and no correction is needed
+        */
+       if (link_id >= mld->fw->ucode_capa.num_links)
+               return false;
+
+       /* for the link conf lookup */
+       guard(rcu)();
+
+       link_conf = rcu_dereference(mld->fw_id_to_bss_conf[link_id]);
+       if (!link_conf)
+               return false;
+
+       mld_link = iwl_mld_link_from_mac80211(link_conf);
+       if (!mld_link)
+               return false;
+
+       /*
+        * If we know the link by link ID then the frame was
+        * received for the link, so by filtering it means it
+        * was from the AP the link is connected to.
+        */
+
+       /* skip also in case we don't have it (yet) */
+       if (!mld_link->average_beacon_energy)
+               return false;
+
+       IWL_DEBUG_STATS(mld, "energy override by average %d\n",
+                       mld_link->average_beacon_energy);
+       rx_status->signal = -mld_link->average_beacon_energy;
+       return true;
+}
+
+static void iwl_mld_fill_signal(struct iwl_mld *mld, int link_id,
+                               struct ieee80211_hdr *hdr,
                                struct ieee80211_rx_status *rx_status,
                                struct iwl_mld_rx_phy_data *phy_data)
 {
        IWL_DEBUG_STATS(mld, "energy in A %d B %d, and max %d\n",
                        energy_a, energy_b, max_energy);
 
+       if (iwl_mld_used_average_energy(mld, link_id, hdr, rx_status))
+               return;
+
        rx_status->signal = max_energy;
-       rx_status->chains =
-           (rate_n_flags & RATE_MCS_ANT_AB_MSK) >> RATE_MCS_ANT_POS;
+       rx_status->chains = u32_get_bits(rate_n_flags, RATE_MCS_ANT_AB_MSK);
        rx_status->chain_signal[0] = energy_a;
        rx_status->chain_signal[1] = energy_b;
 }
 }
 #endif
 
-static void iwl_mld_rx_fill_status(struct iwl_mld *mld, struct sk_buff *skb,
+/* Note: hdr can be NULL */
+static void iwl_mld_rx_fill_status(struct iwl_mld *mld, int link_id,
+                                  struct ieee80211_hdr *hdr,
+                                  struct sk_buff *skb,
                                   struct iwl_mld_rx_phy_data *phy_data,
                                   int queue)
 {
            phy_data->phy_info & IWL_RX_MPDU_PHY_SHORT_PREAMBLE)
                rx_status->enc_flags |= RX_ENC_FLAG_SHORTPRE;
 
-       iwl_mld_fill_signal(mld, rx_status, phy_data);
+       iwl_mld_fill_signal(mld, link_id, hdr, rx_status, phy_data);
 
        /* This may be overridden by iwl_mld_rx_he() to HE_RU */
        switch (rate_n_flags & RATE_MCS_CHAN_WIDTH_MSK) {
        struct sk_buff *skb;
        size_t mpdu_desc_size = sizeof(*mpdu_desc);
        bool drop = false;
-       u8 crypto_len = 0, band;
+       u8 crypto_len = 0, band, link_id;
        u32 pkt_len = iwl_rx_packet_payload_len(pkt);
        u32 mpdu_len;
        enum iwl_mld_reorder_result reorder_res;
                                SCHED_SCAN_PASS_ALL_STATE_FOUND;
        }
 
-       iwl_mld_rx_fill_status(mld, skb, &phy_data, queue);
+       link_id = u8_get_bits(mpdu_desc->mac_phy_band,
+                             IWL_RX_MPDU_MAC_PHY_BAND_LINK_MASK);
+
+       iwl_mld_rx_fill_status(mld, link_id, hdr, skb, &phy_data, queue);
 
        if (iwl_mld_rx_crypto(mld, sta, hdr, rx_status, mpdu_desc, queue,
                              le32_to_cpu(pkt->len_n_flags), &crypto_len))
        rx_status->freq = ieee80211_channel_to_frequency(channel,
                                                         rx_status->band);
 
-       iwl_mld_rx_fill_status(mld, skb, &phy_data, queue);
+       /* link ID is ignored for NULL header */
+       iwl_mld_rx_fill_status(mld, -1, NULL, skb, &phy_data, queue);
 
        /* No more radiotap info should be added after this point.
         * Mark it as mac header for upper layers to know where
 
        if (test_bit(IWL_MVM_STATUS_IN_HW_RESTART, &mvm->status))
                return;
 
+       mvmvif->deflink.average_beacon_energy = 0;
+
        INIT_DELAYED_WORK(&mvmvif->csa_work,
                          iwl_mvm_channel_switch_disconnect_wk);
 
 
  * @mcast_sta: multicast station
  * @phy_ctxt: phy context allocated to this link, if any
  * @bf_data: beacon filtering data
+ * @average_beacon_energy: average beacon energy for beacons received during
+ *     client connections
  */
 struct iwl_mvm_vif_link_info {
        u8 bssid[ETH_ALEN];
        u16 mgmt_queue;
 
        struct iwl_mvm_link_bf_data bf_data;
+       u32 average_beacon_energy;
 };
 
 /**
                                        struct iwl_rx_cmd_buffer *rxb);
 void iwl_mvm_channel_switch_error_notif(struct iwl_mvm *mvm,
                                        struct iwl_rx_cmd_buffer *rxb);
+void iwl_mvm_rx_beacon_filter_notif(struct iwl_mvm *mvm,
+                                   struct iwl_rx_cmd_buffer *rxb);
+
 /* Bindings */
 int iwl_mvm_binding_add_vif(struct iwl_mvm *mvm, struct ieee80211_vif *vif);
 int iwl_mvm_binding_remove_vif(struct iwl_mvm *mvm, struct ieee80211_vif *vif);
 
        RX_HANDLER_GRP(SCAN_GROUP, CHANNEL_SURVEY_NOTIF,
                       iwl_mvm_rx_channel_survey_notif, RX_HANDLER_ASYNC_LOCKED,
                       struct iwl_umac_scan_channel_survey_notif),
+       RX_HANDLER_GRP(DATA_PATH_GROUP, BEACON_FILTER_IN_NOTIF,
+                      iwl_mvm_rx_beacon_filter_notif,
+                      RX_HANDLER_ASYNC_LOCKED,
+                      /* same size as v1 */
+                      struct iwl_beacon_filter_notif),
 };
 #undef RX_HANDLER
 #undef RX_HANDLER_GRP
        HCMD_NAME(ESR_MODE_NOTIF),
        HCMD_NAME(MONITOR_NOTIF),
        HCMD_NAME(THERMAL_DUAL_CHAIN_REQUEST),
+       HCMD_NAME(BEACON_FILTER_IN_NOTIF),
        HCMD_NAME(STA_PM_NOTIF),
        HCMD_NAME(MU_GROUP_MGMT_NOTIF),
        HCMD_NAME(RX_QUEUES_NOTIFICATION),
 
        ieee80211_rx_napi(mvm->hw, sta, skb, napi);
 }
 
+static bool iwl_mvm_used_average_energy(struct iwl_mvm *mvm,
+                                       struct iwl_rx_mpdu_desc *desc,
+                                       struct ieee80211_hdr *hdr,
+                                       struct ieee80211_rx_status *rx_status)
+{
+       struct iwl_mvm_vif *mvm_vif;
+       struct ieee80211_vif *vif;
+       u32 id;
+
+       if (unlikely(!hdr || !desc))
+               return false;
+
+       if (likely(!ieee80211_is_beacon(hdr->frame_control)))
+               return false;
+
+       /* for the link conf lookup */
+       guard(rcu)();
+
+       /* MAC or link ID depending on FW, but driver has them equal */
+       id = u8_get_bits(desc->mac_phy_band,
+                        IWL_RX_MPDU_MAC_PHY_BAND_MAC_MASK);
+
+       /* >= means AUX MAC/link ID, no energy correction needed then */
+       if (id >= ARRAY_SIZE(mvm->vif_id_to_mac))
+               return false;
+
+       vif = iwl_mvm_rcu_dereference_vif_id(mvm, id, true);
+       if (!vif)
+               return false;
+
+       mvm_vif = iwl_mvm_vif_from_mac80211(vif);
+
+       /*
+        * If we know the MAC by MAC or link ID then the frame was
+        * received for the link, so by filtering it means it was
+        * from the AP the link is connected to.
+        */
+
+       /* skip also in case we don't have it (yet) */
+       if (!mvm_vif->deflink.average_beacon_energy)
+               return false;
+
+       IWL_DEBUG_STATS(mvm, "energy override by average %d\n",
+                       mvm_vif->deflink.average_beacon_energy);
+       rx_status->signal = -mvm_vif->deflink.average_beacon_energy;
+       return true;
+}
+
 static void iwl_mvm_get_signal_strength(struct iwl_mvm *mvm,
+                                       struct iwl_rx_mpdu_desc *desc,
+                                       struct ieee80211_hdr *hdr,
                                        struct ieee80211_rx_status *rx_status,
                                        u32 rate_n_flags, int energy_a,
                                        int energy_b)
 {
        int max_energy;
-       u32 rate_flags = rate_n_flags;
 
        energy_a = energy_a ? -energy_a : S8_MIN;
        energy_b = energy_b ? -energy_b : S8_MIN;
        IWL_DEBUG_STATS(mvm, "energy In A %d B %d, and max %d\n",
                        energy_a, energy_b, max_energy);
 
+       if (iwl_mvm_used_average_energy(mvm, desc, hdr, rx_status))
+               return;
+
        rx_status->signal = max_energy;
-       rx_status->chains =
-               (rate_flags & RATE_MCS_ANT_AB_MSK) >> RATE_MCS_ANT_POS;
+       rx_status->chains = u32_get_bits(rate_n_flags, RATE_MCS_ANT_AB_MSK);
        rx_status->chain_signal[0] = energy_a;
        rx_status->chain_signal[1] = energy_b;
 }
 /*
  * Note: requires also rx_status->band to be prefilled, as well
  * as phy_data (apart from phy_data->info_type)
+ * Note: desc/hdr may be NULL
  */
 static void iwl_mvm_rx_fill_status(struct iwl_mvm *mvm,
+                                  struct iwl_rx_mpdu_desc *desc,
+                                  struct ieee80211_hdr *hdr,
                                   struct sk_buff *skb,
                                   struct iwl_mvm_rx_phy_data *phy_data,
                                   int queue)
 
        rx_status->freq = ieee80211_channel_to_frequency(phy_data->channel,
                                                         rx_status->band);
-       iwl_mvm_get_signal_strength(mvm, rx_status, rate_n_flags,
+       iwl_mvm_get_signal_strength(mvm, desc, hdr, rx_status, rate_n_flags,
                                    phy_data->energy_a, phy_data->energy_b);
 
        /* using TLV format and must be after all fixed len fields */
                goto out;
        }
 
-       iwl_mvm_rx_fill_status(mvm, skb, &phy_data, queue);
+       iwl_mvm_rx_fill_status(mvm, desc, hdr, skb, &phy_data, queue);
 
        if (sta) {
                struct iwl_mvm_sta *mvmsta = iwl_mvm_sta_from_mac80211(sta);
        rx_status->band = phy_data.channel > 14 ? NL80211_BAND_5GHZ :
                NL80211_BAND_2GHZ;
 
-       iwl_mvm_rx_fill_status(mvm, skb, &phy_data, queue);
+       iwl_mvm_rx_fill_status(mvm, NULL, NULL, skb, &phy_data, queue);
 
        /* no more radio tap info should be put after this point.
         *
 out:
        rcu_read_unlock();
 }
+
+void iwl_mvm_rx_beacon_filter_notif(struct iwl_mvm *mvm,
+                                   struct iwl_rx_cmd_buffer *rxb)
+{
+       struct iwl_rx_packet *pkt = rxb_addr(rxb);
+       /* MAC or link ID in v1/v2, but driver has the IDs equal */
+       struct iwl_beacon_filter_notif *notif = (void *)pkt->data;
+       u32 id = le32_to_cpu(notif->link_id);
+       struct iwl_mvm_vif *mvm_vif;
+       struct ieee80211_vif *vif;
+
+       /* >= means AUX MAC/link ID, no energy correction needed then */
+       if (IWL_FW_CHECK(mvm, id >= ARRAY_SIZE(mvm->vif_id_to_mac),
+                        "invalid link ID %d\n", id))
+               return;
+
+       vif = iwl_mvm_rcu_dereference_vif_id(mvm, id, false);
+       if (!vif)
+               return;
+
+       mvm_vif = iwl_mvm_vif_from_mac80211(vif);
+
+       mvm_vif->deflink.average_beacon_energy =
+               le32_to_cpu(notif->average_energy);
+}