diff mbox series

[3/8] wifi: ath12k: Refactor sta state machine

Message ID 20241023133004.2253830-4-kvalo@kernel.org (mailing list archive)
State Changes Requested
Delegated to: Jeff Johnson
Headers show
Series wifi: ath12k: MLO support part 2 | expand

Commit Message

Kalle Valo Oct. 23, 2024, 1:29 p.m. UTC
From: Sriram R <quic_srirrama@quicinc.com>

Refactor ath12k_mac_op_sta_state(), with generic wrappers which can be used for
both multi link stations and non-ML stations.

Tested-on: QCN9274 hw2.0 PCI WLAN.WBE.1.0.1-00029-QCAHKSWPL_SILICONZ-1
Tested-on: WCN7850 hw2.0 PCI WLAN.HMT.1.0.c5-00481-QCAHMTSWPL_V1.0_V2.0_SILICONZ-3

Signed-off-by: Sriram R <quic_srirrama@quicinc.com>
Signed-off-by: Harshitha Prem <quic_hprem@quicinc.com>
Signed-off-by: Kalle Valo <quic_kvalo@quicinc.com>
---
 drivers/net/wireless/ath/ath12k/core.h |   3 +
 drivers/net/wireless/ath/ath12k/mac.c  | 341 +++++++++++++++++--------
 2 files changed, 242 insertions(+), 102 deletions(-)

Comments

Jeff Johnson Oct. 23, 2024, 3:38 p.m. UTC | #1
On 10/23/2024 6:29 AM, Kalle Valo wrote:
> From: Sriram R <quic_srirrama@quicinc.com>
> 
> Refactor ath12k_mac_op_sta_state(), with generic wrappers which can be used for
> both multi link stations and non-ML stations.
> 
> Tested-on: QCN9274 hw2.0 PCI WLAN.WBE.1.0.1-00029-QCAHKSWPL_SILICONZ-1
> Tested-on: WCN7850 hw2.0 PCI WLAN.HMT.1.0.c5-00481-QCAHMTSWPL_V1.0_V2.0_SILICONZ-3
> 
> Signed-off-by: Sriram R <quic_srirrama@quicinc.com>
> Signed-off-by: Harshitha Prem <quic_hprem@quicinc.com>
> Signed-off-by: Kalle Valo <quic_kvalo@quicinc.com>
> ---
>  drivers/net/wireless/ath/ath12k/core.h |   3 +
>  drivers/net/wireless/ath/ath12k/mac.c  | 341 +++++++++++++++++--------
>  2 files changed, 242 insertions(+), 102 deletions(-)
> 
> diff --git a/drivers/net/wireless/ath/ath12k/core.h b/drivers/net/wireless/ath/ath12k/core.h
> index 06b637ba8b8f..6faa46b9adc9 100644
> --- a/drivers/net/wireless/ath/ath12k/core.h
> +++ b/drivers/net/wireless/ath/ath12k/core.h
> @@ -461,6 +461,9 @@ struct ath12k_link_sta {
>  	struct ath12k_link_vif *arvif;
>  	struct ath12k_sta *ahsta;
>  
> +	/* link address similar to ieee80211_link_sta */
> +	u8 addr[ETH_ALEN];
> +
>  	/* the following are protected by ar->data_lock */
>  	u32 changed; /* IEEE80211_RC_* */
>  	u32 bw;
> diff --git a/drivers/net/wireless/ath/ath12k/mac.c b/drivers/net/wireless/ath/ath12k/mac.c
> index d4aa4540c8e6..3de6d605cd74 100644
> --- a/drivers/net/wireless/ath/ath12k/mac.c
> +++ b/drivers/net/wireless/ath/ath12k/mac.c
> @@ -4519,10 +4519,10 @@ ath12k_mac_set_peer_vht_fixed_rate(struct ath12k_link_vif *arvif,
>  	return ret;
>  }
>  
> -static int ath12k_station_assoc(struct ath12k *ar,
> -				struct ath12k_link_vif *arvif,
> -				struct ath12k_link_sta *arsta,
> -				bool reassoc)
> +static int ath12k_mac_station_assoc(struct ath12k *ar,
> +				    struct ath12k_link_vif *arvif,
> +				    struct ath12k_link_sta *arsta,
> +				    bool reassoc)
>  {
>  	struct ieee80211_vif *vif = ath12k_ahvif_to_vif(arvif->ahvif);
>  	struct ieee80211_sta *sta = ath12k_ahsta_to_sta(arsta->ahsta);
> @@ -4609,28 +4609,19 @@ static int ath12k_station_assoc(struct ath12k *ar,
>  	return 0;
>  }
>  
> -static int ath12k_station_disassoc(struct ath12k *ar,
> -				   struct ath12k_link_vif *arvif,
> -				   struct ath12k_link_sta *arsta)
> +static int ath12k_mac_station_disassoc(struct ath12k *ar,
> +				       struct ath12k_link_vif *arvif,
> +				       struct ath12k_link_sta *arsta)
>  {
>  	struct ieee80211_sta *sta = ath12k_ahsta_to_sta(arsta->ahsta);
> -	int ret;
>  
>  	lockdep_assert_wiphy(ath12k_ar_to_hw(ar)->wiphy);
>  
>  	if (!sta->wme) {
>  		arvif->num_legacy_stations--;
> -		ret = ath12k_recalc_rtscts_prot(arvif);
> -		if (ret)
> -			return ret;
> +		return ath12k_recalc_rtscts_prot(arvif);
>  	}
>  
> -	ret = ath12k_clear_peer_keys(arvif, sta->addr);
> -	if (ret) {
> -		ath12k_warn(ar->ab, "failed to clear all peer keys for vdev %i: %d\n",
> -			    arvif->vdev_id, ret);
> -		return ret;
> -	}
>  	return 0;
>  }
>  
> @@ -4826,6 +4817,145 @@ static void ath12k_mac_dec_num_stations(struct ath12k_link_vif *arvif,
>  	ar->num_stations--;
>  }
>  
> +static void ath12k_mac_station_post_remove(struct ath12k *ar,
> +					   struct ath12k_link_vif *arvif,
> +					   struct ath12k_link_sta *arsta)
> +{
> +	struct ieee80211_vif *vif = ath12k_ahvif_to_vif(arvif->ahvif);
> +	struct ieee80211_sta *sta = ath12k_ahsta_to_sta(arsta->ahsta);
> +	struct ath12k_sta *ahsta = arsta->ahsta;
> +	struct ath12k_peer *peer;
> +
> +	lockdep_assert_wiphy(ath12k_ar_to_hw(ar)->wiphy);
> +
> +	ath12k_mac_dec_num_stations(arvif, arsta);
> +
> +	spin_lock_bh(&ar->ab->base_lock);
> +
> +	peer = ath12k_peer_find(ar->ab, arvif->vdev_id, sta->addr);
> +	if (peer && peer->sta == sta) {
> +		ath12k_warn(ar->ab, "Found peer entry %pM n vdev %i after it was supposedly removed\n",
> +			    vif->addr, arvif->vdev_id);
> +		peer->sta = NULL;
> +		list_del(&peer->list);
> +		kfree(peer);
> +		ar->num_peers--;
> +	}
> +
> +	spin_unlock_bh(&ar->ab->base_lock);
> +
> +	kfree(arsta->rx_stats);
> +	arsta->rx_stats = NULL;
> +
> +	if (arsta->link_id < IEEE80211_MLD_MAX_NUM_LINKS) {
> +		rcu_assign_pointer(ahsta->link[arsta->link_id], NULL);
> +		synchronize_rcu();

I've mentioned this in the past in some internal discussion and seems now is a
good time to bring this to light.

It concerns me that this happens so late in the process. In theory another
thread could already have a valid arsta pointer and could be trying to
dereference that pointer while the code above is destroying underlying data
(i.e. arsta->rx_stats).

Should we set this to NULL and synchronize RCU at the beginning of the process
so that we know all access to the struct has finished before we start
destroying the data?

Or can this not actually happen in practice due to other synchronization
mechansims? And if so, should we document that somewhere?


> +		ahsta->links_map &= ~(BIT(arsta->link_id));
> +		arsta->link_id = ATH12K_INVALID_LINK_ID;
> +		arsta->ahsta = NULL;
> +	}
> +}
> +
> +static int ath12k_mac_station_unauthorize(struct ath12k *ar,
> +					  struct ath12k_link_vif *arvif,
> +					  struct ath12k_link_sta *arsta)
> +{
> +	struct ath12k_peer *peer;
> +	int ret;
> +
> +	lockdep_assert_wiphy(ath12k_ar_to_hw(ar)->wiphy);
> +
> +	spin_lock_bh(&ar->ab->base_lock);
> +
> +	peer = ath12k_peer_find(ar->ab, arvif->vdev_id, arsta->addr);
> +	if (peer)
> +		peer->is_authorized = false;
> +
> +	spin_unlock_bh(&ar->ab->base_lock);
> +
> +	/* Driver should clear the peer keys during mac80211's ref ptr
> +	 * gets cleared in __sta_info_destroy_part2 (trans from
> +	 * IEEE80211_STA_AUTHORIZED to IEEE80211_STA_ASSOC)

I'm unable to understand this comment

> +	 */
> +	ret = ath12k_clear_peer_keys(arvif, arsta->addr);
> +	if (ret) {
> +		ath12k_warn(ar->ab, "failed to clear all peer keys for vdev %i: %d\n",
> +			    arvif->vdev_id, ret);
> +		return ret;
> +	}
> +
> +	return 0;
> +}
> +
> +static int ath12k_mac_station_authorize(struct ath12k *ar,
> +					struct ath12k_link_vif *arvif,
> +					struct ath12k_link_sta *arsta)
> +{
> +	struct ath12k_peer *peer;
> +	struct ieee80211_vif *vif = ath12k_ahvif_to_vif(arvif->ahvif);
> +	struct ieee80211_sta *sta = ath12k_ahsta_to_sta(arsta->ahsta);
> +	int ret;
> +
> +	lockdep_assert_wiphy(ath12k_ar_to_hw(ar)->wiphy);
> +
> +	spin_lock_bh(&ar->ab->base_lock);
> +
> +	peer = ath12k_peer_find(ar->ab, arvif->vdev_id, sta->addr);
> +	if (peer)
> +		peer->is_authorized = true;
> +
> +	spin_unlock_bh(&ar->ab->base_lock);
> +
> +	if (vif->type == NL80211_IFTYPE_STATION && arvif->is_up) {
> +		ret = ath12k_wmi_set_peer_param(ar, sta->addr,
> +						arvif->vdev_id,
> +						WMI_PEER_AUTHORIZE,
> +						1);
> +		if (ret) {
> +			ath12k_warn(ar->ab, "Unable to authorize peer %pM vdev %d: %d\n",
> +				    sta->addr, arvif->vdev_id, ret);
> +			return ret;
> +		}
> +	}
> +
> +	return 0;
> +}
> +
> +static int ath12k_mac_station_remove(struct ath12k *ar,
> +				     struct ath12k_link_vif *arvif,
> +				     struct ath12k_link_sta *arsta)
> +{
> +	struct ieee80211_sta *sta = ath12k_ahsta_to_sta(arsta->ahsta);
> +	struct ath12k_vif *ahvif = arvif->ahvif;
> +	int ret;
> +
> +	lockdep_assert_wiphy(ath12k_ar_to_hw(ar)->wiphy);
> +
> +	wiphy_work_cancel(ar->ah->hw->wiphy, &arsta->update_wk);
> +
> +	if (ahvif->vdev_type == WMI_VDEV_TYPE_STA) {
> +		ath12k_bss_disassoc(ar, arvif);
> +		ret = ath12k_mac_vdev_stop(arvif);
> +		if (ret)
> +			ath12k_warn(ar->ab, "failed to stop vdev %i: %d\n",
> +				    arvif->vdev_id, ret);
> +	}
> +
> +	ath12k_dp_peer_cleanup(ar, arvif->vdev_id, sta->addr);
> +
> +	ret = ath12k_peer_delete(ar, arvif->vdev_id, sta->addr);
> +	if (ret)
> +		ath12k_warn(ar->ab, "Failed to delete peer: %pM for VDEV: %d\n",
> +			    sta->addr, arvif->vdev_id);
> +	else
> +		ath12k_dbg(ar->ab, ATH12K_DBG_MAC, "Removed peer: %pM for VDEV: %d\n",
> +			   sta->addr, arvif->vdev_id);
> +
> +	ath12k_mac_station_post_remove(ar, arvif, arsta);
> +
> +	return ret;
> +}
> +
>  static int ath12k_mac_station_add(struct ath12k *ar,
>  				  struct ath12k_link_vif *arvif,
>  				  struct ath12k_link_sta *arsta)
> @@ -4933,31 +5063,37 @@ static u32 ath12k_mac_ieee80211_sta_bw_to_wmi(struct ath12k *ar,
>  	return bw;
>  }
>  
> -static int ath12k_mac_op_sta_state(struct ieee80211_hw *hw,
> -				   struct ieee80211_vif *vif,
> -				   struct ieee80211_sta *sta,
> -				   enum ieee80211_sta_state old_state,
> -				   enum ieee80211_sta_state new_state)
> +static int ath12k_mac_handle_link_sta_state(struct ieee80211_hw *hw,
> +					    struct ath12k_link_vif *arvif,
> +					    struct ath12k_link_sta *arsta,
> +					    enum ieee80211_sta_state old_state,
> +					    enum ieee80211_sta_state new_state)
>  {
> -	struct ath12k_vif *ahvif = ath12k_vif_to_ahvif(vif);
> -	struct ath12k_sta *ahsta = ath12k_sta_to_ahsta(sta);
> -	struct ath12k *ar;
> -	struct ath12k_link_vif *arvif;
> -	struct ath12k_link_sta *arsta;
> -	struct ath12k_peer *peer;
> +	struct ath12k *ar = arvif->ar;
> +	struct ieee80211_vif *vif = ath12k_ahvif_to_vif(arvif->ahvif);
> +	struct ieee80211_sta *sta = ath12k_ahsta_to_sta(arsta->ahsta);
> +	struct ath12k_sta *ahsta = arsta->ahsta;
>  	int ret = 0;
>  
>  	lockdep_assert_wiphy(hw->wiphy);
>  
> -	arvif = &ahvif->deflink;
> -	arsta = &ahsta->deflink;
> +	/* IEEE80211_STA_NONE -> IEEE80211_STA_NOTEXIST: Remove the station
> +	 * from driver
> +	 */
> +	if ((old_state == IEEE80211_STA_NONE &&
> +	     new_state == IEEE80211_STA_NOTEXIST)) {
> +		/* ML sta needs separate handling */
> +		if (sta->mlo)
> +			return 0;
>  
> -	ar = ath12k_get_ar_by_vif(hw, vif);
> -	if (!ar) {
> -		WARN_ON_ONCE(1);
> -		return -EINVAL;
> +		ret = ath12k_mac_station_remove(ar, arvif, arsta);
> +		if (ret) {
> +			ath12k_warn(ar->ab, "Failed to remove station: %pM for VDEV: %d\n",
> +				    arsta->addr, arvif->vdev_id);
> +		}
>  	}
>  
> +	/* IEEE80211_STA_NOTEXIST -> IEEE80211_STA_NONE: Add new station to driver */
>  	if (old_state == IEEE80211_STA_NOTEXIST &&
>  	    new_state == IEEE80211_STA_NONE) {
>  		memset(arsta, 0, sizeof(*arsta));
> @@ -4975,56 +5111,16 @@ static int ath12k_mac_op_sta_state(struct ieee80211_hw *hw,
>  		if (ret)
>  			ath12k_warn(ar->ab, "Failed to add station: %pM for VDEV: %d\n",
>  				    sta->addr, arvif->vdev_id);
> -	} else if ((old_state == IEEE80211_STA_NONE &&
> -		    new_state == IEEE80211_STA_NOTEXIST)) {
> -		wiphy_work_cancel(hw->wiphy, &arsta->update_wk);
>  
> -		if (ahvif->vdev_type == WMI_VDEV_TYPE_STA) {
> -			ath12k_bss_disassoc(ar, arvif);
> -			ret = ath12k_mac_vdev_stop(arvif);
> -			if (ret)
> -				ath12k_warn(ar->ab, "failed to stop vdev %i: %d\n",
> -					    arvif->vdev_id, ret);
> -		}
> -		ath12k_dp_peer_cleanup(ar, arvif->vdev_id, sta->addr);
> -
> -		ret = ath12k_peer_delete(ar, arvif->vdev_id, sta->addr);
> -		if (ret)
> -			ath12k_warn(ar->ab, "Failed to delete peer: %pM for VDEV: %d\n",
> -				    sta->addr, arvif->vdev_id);
> -		else
> -			ath12k_dbg(ar->ab, ATH12K_DBG_MAC, "Removed peer: %pM for VDEV: %d\n",
> -				   sta->addr, arvif->vdev_id);
> -
> -		ath12k_mac_dec_num_stations(arvif, arsta);
> -		spin_lock_bh(&ar->ab->base_lock);
> -		peer = ath12k_peer_find(ar->ab, arvif->vdev_id, sta->addr);
> -		if (peer && peer->sta == sta) {
> -			ath12k_warn(ar->ab, "Found peer entry %pM n vdev %i after it was supposedly removed\n",
> -				    vif->addr, arvif->vdev_id);
> -			peer->sta = NULL;
> -			list_del(&peer->list);
> -			kfree(peer);
> -			ar->num_peers--;
> -		}
> -		spin_unlock_bh(&ar->ab->base_lock);
> -
> -		kfree(arsta->rx_stats);
> -		arsta->rx_stats = NULL;
> -
> -		if (arsta->link_id < IEEE80211_MLD_MAX_NUM_LINKS) {
> -			rcu_assign_pointer(ahsta->link[arsta->link_id], NULL);
> -			synchronize_rcu();
> -			ahsta->links_map &= ~(BIT(arsta->link_id));
> -			arsta->link_id = ATH12K_INVALID_LINK_ID;
> -			arsta->ahsta = NULL;
> -		}
> +	/* IEEE80211_STA_AUTH -> IEEE80211_STA_ASSOC: Send station assoc command for
> +	 * peer associated to AP/Mesh/ADHOC vif type.
> +	 */
>  	} else if (old_state == IEEE80211_STA_AUTH &&
>  		   new_state == IEEE80211_STA_ASSOC &&
>  		   (vif->type == NL80211_IFTYPE_AP ||
>  		    vif->type == NL80211_IFTYPE_MESH_POINT ||
>  		    vif->type == NL80211_IFTYPE_ADHOC)) {
> -		ret = ath12k_station_assoc(ar, arvif, arsta, false);
> +		ret = ath12k_mac_station_assoc(ar, arvif, arsta, false);
>  		if (ret)
>  			ath12k_warn(ar->ab, "Failed to associate station: %pM\n",
>  				    sta->addr);
> @@ -5035,40 +5131,32 @@ static int ath12k_mac_op_sta_state(struct ieee80211_hw *hw,
>  		arsta->bw_prev = sta->deflink.bandwidth;
>  
>  		spin_unlock_bh(&ar->data_lock);
> +
> +	/* IEEE80211_STA_ASSOC -> IEEE80211_STA_AUTHORIZED: set peer status as
> +	 * authorized
> +	 */
>  	} else if (old_state == IEEE80211_STA_ASSOC &&
>  		   new_state == IEEE80211_STA_AUTHORIZED) {
> -		spin_lock_bh(&ar->ab->base_lock);
> +		ret = ath12k_mac_station_authorize(ar, arvif, arsta);
> +		if (ret)
> +			ath12k_warn(ar->ab, "Failed to authorize station: %pM\n",
> +				    sta->addr);
>  
> -		peer = ath12k_peer_find(ar->ab, arvif->vdev_id, sta->addr);
> -		if (peer)
> -			peer->is_authorized = true;
> -
> -		spin_unlock_bh(&ar->ab->base_lock);
> -
> -		if (vif->type == NL80211_IFTYPE_STATION && arvif->is_up) {
> -			ret = ath12k_wmi_set_peer_param(ar, sta->addr,
> -							arvif->vdev_id,
> -							WMI_PEER_AUTHORIZE,
> -							1);
> -			if (ret)
> -				ath12k_warn(ar->ab, "Unable to authorize peer %pM vdev %d: %d\n",
> -					    sta->addr, arvif->vdev_id, ret);
> -		}
> +	/* IEEE80211_STA_AUTHORIZED -> IEEE80211_STA_ASSOC: station may be in removal,
> +	 * deauthorize it.
> +	 */
>  	} else if (old_state == IEEE80211_STA_AUTHORIZED &&
>  		   new_state == IEEE80211_STA_ASSOC) {
> -		spin_lock_bh(&ar->ab->base_lock);
> -
> -		peer = ath12k_peer_find(ar->ab, arvif->vdev_id, sta->addr);
> -		if (peer)
> -			peer->is_authorized = false;
> -
> -		spin_unlock_bh(&ar->ab->base_lock);
> +		ath12k_mac_station_unauthorize(ar, arvif, arsta);
> +	/* IEEE80211_STA_ASSOC -> IEEE80211_STA_AUTH: disassoc peer connected to
> +	 * AP/mesh/ADHOC vif type.
> +	 */
>  	} else if (old_state == IEEE80211_STA_ASSOC &&
>  		   new_state == IEEE80211_STA_AUTH &&
>  		   (vif->type == NL80211_IFTYPE_AP ||
>  		    vif->type == NL80211_IFTYPE_MESH_POINT ||
>  		    vif->type == NL80211_IFTYPE_ADHOC)) {
> -		ret = ath12k_station_disassoc(ar, arvif, arsta);
> +		ret = ath12k_mac_station_disassoc(ar, arvif, arsta);
>  		if (ret)
>  			ath12k_warn(ar->ab, "Failed to disassociate station: %pM\n",
>  				    sta->addr);
> @@ -5077,6 +5165,55 @@ static int ath12k_mac_op_sta_state(struct ieee80211_hw *hw,
>  	return ret;
>  }
>  
> +static int ath12k_mac_op_sta_state(struct ieee80211_hw *hw,
> +				   struct ieee80211_vif *vif,
> +				   struct ieee80211_sta *sta,
> +				   enum ieee80211_sta_state old_state,
> +				   enum ieee80211_sta_state new_state)
> +{
> +	struct ath12k_vif *ahvif = ath12k_vif_to_ahvif(vif);
> +	struct ath12k_sta *ahsta = ath12k_sta_to_ahsta(sta);
> +	struct ath12k_link_vif *arvif;
> +	struct ath12k_link_sta *arsta;
> +	int ret;
> +	u8 link_id = 0;
> +
> +	lockdep_assert_wiphy(hw->wiphy);
> +
> +	if (ieee80211_vif_is_mld(vif) && sta->valid_links) {
> +		WARN_ON(!sta->mlo && hweight16(sta->valid_links) != 1);
> +		link_id = ffs(sta->valid_links) - 1;
> +	}
> +
> +	/* Handle for non-ML station */
> +	if (!sta->mlo) {
> +		arvif = wiphy_dereference(hw->wiphy, ahvif->link[link_id]);
> +		arsta = &ahsta->deflink;
> +		arsta->ahsta = ahsta;
> +
> +		if (WARN_ON(!arvif || !arsta)) {
> +			ret = -EINVAL;
> +			goto exit;
> +		}
> +
> +		/* vdev might be in deleted */
> +		if (WARN_ON(!arvif->ar)) {
> +			ret = -EINVAL;
> +			goto exit;
> +		}
> +
> +		ret = ath12k_mac_handle_link_sta_state(hw, arvif, arsta,
> +						       old_state, new_state);
> +		if (ret)
> +			goto exit;
> +	}
> +
> +	ret = 0;
> +
> +exit:
> +	return ret;
> +}
> +
>  static int ath12k_mac_op_sta_set_txpwr(struct ieee80211_hw *hw,
>  				       struct ieee80211_vif *vif,
>  				       struct ieee80211_sta *sta)
Kalle Valo Oct. 26, 2024, 9:08 a.m. UTC | #2
Baochen Qiang <quic_bqiang@quicinc.com> writes:

> On 10/23/2024 9:29 PM, Kalle Valo wrote:
>> +static void ath12k_mac_station_post_remove(struct ath12k *ar,
>> +					   struct ath12k_link_vif *arvif,
>> +					   struct ath12k_link_sta *arsta)
>> +{
>> +	struct ieee80211_vif *vif = ath12k_ahvif_to_vif(arvif->ahvif);
>> +	struct ieee80211_sta *sta = ath12k_ahsta_to_sta(arsta->ahsta);
>> +	struct ath12k_sta *ahsta = arsta->ahsta;
>> +	struct ath12k_peer *peer;
>> +
>> +	lockdep_assert_wiphy(ath12k_ar_to_hw(ar)->wiphy);
>> +
>> +	ath12k_mac_dec_num_stations(arvif, arsta);
>> +
>> +	spin_lock_bh(&ar->ab->base_lock);
>> +
>> +	peer = ath12k_peer_find(ar->ab, arvif->vdev_id, sta->addr);
>> +	if (peer && peer->sta == sta) {
>> +		ath12k_warn(ar->ab, "Found peer entry %pM n vdev %i after it was supposedly removed\n",
>> +			    vif->addr, arvif->vdev_id);
>> +		peer->sta = NULL;
>> +		list_del(&peer->list);
>> +		kfree(peer);
>> +		ar->num_peers--;
>> +	}
>> +
>> +	spin_unlock_bh(&ar->ab->base_lock);
>> +
>> +	kfree(arsta->rx_stats);
>> +	arsta->rx_stats = NULL;
>> +
>> +	if (arsta->link_id < IEEE80211_MLD_MAX_NUM_LINKS) {
>> +		rcu_assign_pointer(ahsta->link[arsta->link_id], NULL);
>> +		synchronize_rcu();
>> +		ahsta->links_map &= ~(BIT(arsta->link_id));
>
> should we put this ahead of rcu_assign_pointer()?

I agree, I'll do that in v2.
Kalle Valo Oct. 29, 2024, 3:29 p.m. UTC | #3
Jeff Johnson <quic_jjohnson@quicinc.com> writes:

> On 10/23/2024 6:29 AM, Kalle Valo wrote:
>
>> +static void ath12k_mac_station_post_remove(struct ath12k *ar,
>> +					   struct ath12k_link_vif *arvif,
>> +					   struct ath12k_link_sta *arsta)
>> +{
>> +	struct ieee80211_vif *vif = ath12k_ahvif_to_vif(arvif->ahvif);
>> +	struct ieee80211_sta *sta = ath12k_ahsta_to_sta(arsta->ahsta);
>> +	struct ath12k_sta *ahsta = arsta->ahsta;
>> +	struct ath12k_peer *peer;
>> +
>> +	lockdep_assert_wiphy(ath12k_ar_to_hw(ar)->wiphy);
>> +
>> +	ath12k_mac_dec_num_stations(arvif, arsta);
>> +
>> +	spin_lock_bh(&ar->ab->base_lock);
>> +
>> +	peer = ath12k_peer_find(ar->ab, arvif->vdev_id, sta->addr);
>> +	if (peer && peer->sta == sta) {
>> +		ath12k_warn(ar->ab, "Found peer entry %pM n vdev %i after it was supposedly removed\n",
>> +			    vif->addr, arvif->vdev_id);
>> +		peer->sta = NULL;
>> +		list_del(&peer->list);
>> +		kfree(peer);
>> +		ar->num_peers--;
>> +	}
>> +
>> +	spin_unlock_bh(&ar->ab->base_lock);
>> +
>> +	kfree(arsta->rx_stats);
>> +	arsta->rx_stats = NULL;
>> +
>> +	if (arsta->link_id < IEEE80211_MLD_MAX_NUM_LINKS) {
>> +		rcu_assign_pointer(ahsta->link[arsta->link_id], NULL);
>> +		synchronize_rcu();
>
> I've mentioned this in the past in some internal discussion and seems now is a
> good time to bring this to light.
>
> It concerns me that this happens so late in the process. In theory another
> thread could already have a valid arsta pointer and could be trying to
> dereference that pointer while the code above is destroying underlying data
> (i.e. arsta->rx_stats).
>
> Should we set this to NULL and synchronize RCU at the beginning of the process
> so that we know all access to the struct has finished before we start
> destroying the data?
>
> Or can this not actually happen in practice due to other synchronization
> mechansims? And if so, should we document that somewhere?

I think you are correct, AFAICS the kfree(arsta->rx_stats) should be
after synchronize_rcu(). But this race was already in the code before
this patch so we need to fix in a separate patch. I have added this to
my todo list.
Jeff Johnson Oct. 29, 2024, 3:35 p.m. UTC | #4
On 10/29/2024 8:29 AM, Kalle Valo wrote:
> Jeff Johnson <quic_jjohnson@quicinc.com> writes:
> 
>> On 10/23/2024 6:29 AM, Kalle Valo wrote:
>>
>>> +static void ath12k_mac_station_post_remove(struct ath12k *ar,
>>> +					   struct ath12k_link_vif *arvif,
>>> +					   struct ath12k_link_sta *arsta)
>>> +{
>>> +	struct ieee80211_vif *vif = ath12k_ahvif_to_vif(arvif->ahvif);
>>> +	struct ieee80211_sta *sta = ath12k_ahsta_to_sta(arsta->ahsta);
>>> +	struct ath12k_sta *ahsta = arsta->ahsta;
>>> +	struct ath12k_peer *peer;
>>> +
>>> +	lockdep_assert_wiphy(ath12k_ar_to_hw(ar)->wiphy);
>>> +
>>> +	ath12k_mac_dec_num_stations(arvif, arsta);
>>> +
>>> +	spin_lock_bh(&ar->ab->base_lock);
>>> +
>>> +	peer = ath12k_peer_find(ar->ab, arvif->vdev_id, sta->addr);
>>> +	if (peer && peer->sta == sta) {
>>> +		ath12k_warn(ar->ab, "Found peer entry %pM n vdev %i after it was supposedly removed\n",
>>> +			    vif->addr, arvif->vdev_id);
>>> +		peer->sta = NULL;
>>> +		list_del(&peer->list);
>>> +		kfree(peer);
>>> +		ar->num_peers--;
>>> +	}
>>> +
>>> +	spin_unlock_bh(&ar->ab->base_lock);
>>> +
>>> +	kfree(arsta->rx_stats);
>>> +	arsta->rx_stats = NULL;
>>> +
>>> +	if (arsta->link_id < IEEE80211_MLD_MAX_NUM_LINKS) {
>>> +		rcu_assign_pointer(ahsta->link[arsta->link_id], NULL);
>>> +		synchronize_rcu();
>>
>> I've mentioned this in the past in some internal discussion and seems now is a
>> good time to bring this to light.
>>
>> It concerns me that this happens so late in the process. In theory another
>> thread could already have a valid arsta pointer and could be trying to
>> dereference that pointer while the code above is destroying underlying data
>> (i.e. arsta->rx_stats).
>>
>> Should we set this to NULL and synchronize RCU at the beginning of the process
>> so that we know all access to the struct has finished before we start
>> destroying the data?
>>
>> Or can this not actually happen in practice due to other synchronization
>> mechansims? And if so, should we document that somewhere?
> 
> I think you are correct, AFAICS the kfree(arsta->rx_stats) should be
> after synchronize_rcu(). But this race was already in the code before
> this patch so we need to fix in a separate patch. I have added this to
> my todo list.
> 

Sounds reasonable to me
Kalle Valo Oct. 29, 2024, 3:38 p.m. UTC | #5
+ aditya

Jeff Johnson <quic_jjohnson@quicinc.com> writes:

>> +static int ath12k_mac_station_unauthorize(struct ath12k *ar,
>> +					  struct ath12k_link_vif *arvif,
>> +					  struct ath12k_link_sta *arsta)
>> +{
>> +	struct ath12k_peer *peer;
>> +	int ret;
>> +
>> +	lockdep_assert_wiphy(ath12k_ar_to_hw(ar)->wiphy);
>> +
>> +	spin_lock_bh(&ar->ab->base_lock);
>> +
>> +	peer = ath12k_peer_find(ar->ab, arvif->vdev_id, arsta->addr);
>> +	if (peer)
>> +		peer->is_authorized = false;
>> +
>> +	spin_unlock_bh(&ar->ab->base_lock);
>> +
>> +	/* Driver should clear the peer keys during mac80211's ref ptr
>> +	 * gets cleared in __sta_info_destroy_part2 (trans from
>> +	 * IEEE80211_STA_AUTHORIZED to IEEE80211_STA_ASSOC)
>
> I'm unable to understand this comment

Indeed, that's weird. Aditya, do you have any idea what the comment is
trying to say?
Aditya Kumar Singh Oct. 30, 2024, 4:05 a.m. UTC | #6
On 10/29/24 21:08, Kalle Valo wrote:
> + aditya
> 
> Jeff Johnson <quic_jjohnson@quicinc.com> writes:
> 
>>> +static int ath12k_mac_station_unauthorize(struct ath12k *ar,
>>> +					  struct ath12k_link_vif *arvif,
>>> +					  struct ath12k_link_sta *arsta)
>>> +{
>>> +	struct ath12k_peer *peer;
>>> +	int ret;
>>> +
>>> +	lockdep_assert_wiphy(ath12k_ar_to_hw(ar)->wiphy);
>>> +
>>> +	spin_lock_bh(&ar->ab->base_lock);
>>> +
>>> +	peer = ath12k_peer_find(ar->ab, arvif->vdev_id, arsta->addr);
>>> +	if (peer)
>>> +		peer->is_authorized = false;
>>> +
>>> +	spin_unlock_bh(&ar->ab->base_lock);
>>> +
>>> +	/* Driver should clear the peer keys during mac80211's ref ptr
>>> +	 * gets cleared in __sta_info_destroy_part2 (trans from
>>> +	 * IEEE80211_STA_AUTHORIZED to IEEE80211_STA_ASSOC)
>>
>> I'm unable to understand this comment
> 
> Indeed, that's weird. Aditya, do you have any idea what the comment is
> trying to say?
> 

At present, ath12k clear the keys in ath12k_station_disassoc() which 
gets executed in state change from IEEE80211_STA_ASSOC to 
IEEE80211_STA_AUTH.

However, in mac80211, once the station moves from 
IEEE80211_STA_AUTHORIZED to IEEE80211_STA_ASSOC itself, the keys are 
deleted. Please see - __sta_info_destroy_part2() -> 
ieee80211_free_sta_keys().

Now, ath12k peer object (struct ath12k_peer) holds the key reference 
from mac80211 (see ath12k_peer::keys[]). Hence, once mac80211 deletes 
the key, driver should not keep a reference to it or else it could lead 
to issues.

Therefore, it is important that the driver should clear the peer keys 
during transition from IEEE80211_STA_AUTHORIZED to IEEE80211_STA_ASSOC 
it self since we know that once we return from here, mac80211 is going 
to remove the keys.

ath12k_mac_station_unauthorize() gets called when station moves from 
state IEEE80211_STA_AUTHORIZED to IEEE80211_STA_ASSOC hence call to 
ath12k_clear_peer_keys() is moved from ath12k_station_disassoc() to 
ath12k_mac_station_unauthorize().

Is this clear now?

May be the comment in the code could be re-written as below?

/* Driver must clear the keys during the state change from
  * IEEE80211_STA_AUTHORIZED to IEEE80211_STA_ASSOC, since after
  * returning from here, mac80211 is going to delete the keys
  * in __sta_info_destroy_part2(). This will ensure that the driver does
  * not retain stale key references after mac80211 deletes the keys.
  */
Kalle Valo Oct. 30, 2024, 6:28 p.m. UTC | #7
Aditya Kumar Singh <quic_adisi@quicinc.com> writes:

> On 10/29/24 21:08, Kalle Valo wrote:
>> + aditya
>> Jeff Johnson <quic_jjohnson@quicinc.com> writes:
>> 
>>>> +static int ath12k_mac_station_unauthorize(struct ath12k *ar,
>>>> +					  struct ath12k_link_vif *arvif,
>>>> +					  struct ath12k_link_sta *arsta)
>>>> +{
>>>> +	struct ath12k_peer *peer;
>>>> +	int ret;
>>>> +
>>>> +	lockdep_assert_wiphy(ath12k_ar_to_hw(ar)->wiphy);
>>>> +
>>>> +	spin_lock_bh(&ar->ab->base_lock);
>>>> +
>>>> +	peer = ath12k_peer_find(ar->ab, arvif->vdev_id, arsta->addr);
>>>> +	if (peer)
>>>> +		peer->is_authorized = false;
>>>> +
>>>> +	spin_unlock_bh(&ar->ab->base_lock);
>>>> +
>>>> +	/* Driver should clear the peer keys during mac80211's ref ptr
>>>> +	 * gets cleared in __sta_info_destroy_part2 (trans from
>>>> +	 * IEEE80211_STA_AUTHORIZED to IEEE80211_STA_ASSOC)
>>>
>>> I'm unable to understand this comment
>> Indeed, that's weird. Aditya, do you have any idea what the comment
>> is
>> trying to say?
>> 
>
> At present, ath12k clear the keys in ath12k_station_disassoc() which
> gets executed in state change from IEEE80211_STA_ASSOC to
> IEEE80211_STA_AUTH.
>
> However, in mac80211, once the station moves from
> IEEE80211_STA_AUTHORIZED to IEEE80211_STA_ASSOC itself, the keys are
> deleted. Please see - __sta_info_destroy_part2() ->
> ieee80211_free_sta_keys().
>
> Now, ath12k peer object (struct ath12k_peer) holds the key reference
> from mac80211 (see ath12k_peer::keys[]). Hence, once mac80211 deletes
> the key, driver should not keep a reference to it or else it could
> lead to issues.
>
> Therefore, it is important that the driver should clear the peer keys
> during transition from IEEE80211_STA_AUTHORIZED to IEEE80211_STA_ASSOC
> it self since we know that once we return from here, mac80211 is going
> to remove the keys.
>
> ath12k_mac_station_unauthorize() gets called when station moves from
> state IEEE80211_STA_AUTHORIZED to IEEE80211_STA_ASSOC hence call to
> ath12k_clear_peer_keys() is moved from ath12k_station_disassoc() to
> ath12k_mac_station_unauthorize().
>
> Is this clear now?

Super clear :) Thanks!

> May be the comment in the code could be re-written as below?
>
> /* Driver must clear the keys during the state change from
>  * IEEE80211_STA_AUTHORIZED to IEEE80211_STA_ASSOC, since after
>  * returning from here, mac80211 is going to delete the keys
>  * in __sta_info_destroy_part2(). This will ensure that the driver does
>  * not retain stale key references after mac80211 deletes the keys.
>  */

Looks good to me, I'll add that if it's ok for Jeff as well.
Jeff Johnson Oct. 30, 2024, 6:39 p.m. UTC | #8
On 10/30/2024 11:28 AM, Kalle Valo wrote:
> Aditya Kumar Singh <quic_adisi@quicinc.com> writes:
>> May be the comment in the code could be re-written as below?
>>
>> /* Driver must clear the keys during the state change from
>>  * IEEE80211_STA_AUTHORIZED to IEEE80211_STA_ASSOC, since after
>>  * returning from here, mac80211 is going to delete the keys
>>  * in __sta_info_destroy_part2(). This will ensure that the driver does
>>  * not retain stale key references after mac80211 deletes the keys.
>>  */
> 
> Looks good to me, I'll add that if it's ok for Jeff as well.
> 

Definitely ok, thanks for the clarification
diff mbox series

Patch

diff --git a/drivers/net/wireless/ath/ath12k/core.h b/drivers/net/wireless/ath/ath12k/core.h
index 06b637ba8b8f..6faa46b9adc9 100644
--- a/drivers/net/wireless/ath/ath12k/core.h
+++ b/drivers/net/wireless/ath/ath12k/core.h
@@ -461,6 +461,9 @@  struct ath12k_link_sta {
 	struct ath12k_link_vif *arvif;
 	struct ath12k_sta *ahsta;
 
+	/* link address similar to ieee80211_link_sta */
+	u8 addr[ETH_ALEN];
+
 	/* the following are protected by ar->data_lock */
 	u32 changed; /* IEEE80211_RC_* */
 	u32 bw;
diff --git a/drivers/net/wireless/ath/ath12k/mac.c b/drivers/net/wireless/ath/ath12k/mac.c
index d4aa4540c8e6..3de6d605cd74 100644
--- a/drivers/net/wireless/ath/ath12k/mac.c
+++ b/drivers/net/wireless/ath/ath12k/mac.c
@@ -4519,10 +4519,10 @@  ath12k_mac_set_peer_vht_fixed_rate(struct ath12k_link_vif *arvif,
 	return ret;
 }
 
-static int ath12k_station_assoc(struct ath12k *ar,
-				struct ath12k_link_vif *arvif,
-				struct ath12k_link_sta *arsta,
-				bool reassoc)
+static int ath12k_mac_station_assoc(struct ath12k *ar,
+				    struct ath12k_link_vif *arvif,
+				    struct ath12k_link_sta *arsta,
+				    bool reassoc)
 {
 	struct ieee80211_vif *vif = ath12k_ahvif_to_vif(arvif->ahvif);
 	struct ieee80211_sta *sta = ath12k_ahsta_to_sta(arsta->ahsta);
@@ -4609,28 +4609,19 @@  static int ath12k_station_assoc(struct ath12k *ar,
 	return 0;
 }
 
-static int ath12k_station_disassoc(struct ath12k *ar,
-				   struct ath12k_link_vif *arvif,
-				   struct ath12k_link_sta *arsta)
+static int ath12k_mac_station_disassoc(struct ath12k *ar,
+				       struct ath12k_link_vif *arvif,
+				       struct ath12k_link_sta *arsta)
 {
 	struct ieee80211_sta *sta = ath12k_ahsta_to_sta(arsta->ahsta);
-	int ret;
 
 	lockdep_assert_wiphy(ath12k_ar_to_hw(ar)->wiphy);
 
 	if (!sta->wme) {
 		arvif->num_legacy_stations--;
-		ret = ath12k_recalc_rtscts_prot(arvif);
-		if (ret)
-			return ret;
+		return ath12k_recalc_rtscts_prot(arvif);
 	}
 
-	ret = ath12k_clear_peer_keys(arvif, sta->addr);
-	if (ret) {
-		ath12k_warn(ar->ab, "failed to clear all peer keys for vdev %i: %d\n",
-			    arvif->vdev_id, ret);
-		return ret;
-	}
 	return 0;
 }
 
@@ -4826,6 +4817,145 @@  static void ath12k_mac_dec_num_stations(struct ath12k_link_vif *arvif,
 	ar->num_stations--;
 }
 
+static void ath12k_mac_station_post_remove(struct ath12k *ar,
+					   struct ath12k_link_vif *arvif,
+					   struct ath12k_link_sta *arsta)
+{
+	struct ieee80211_vif *vif = ath12k_ahvif_to_vif(arvif->ahvif);
+	struct ieee80211_sta *sta = ath12k_ahsta_to_sta(arsta->ahsta);
+	struct ath12k_sta *ahsta = arsta->ahsta;
+	struct ath12k_peer *peer;
+
+	lockdep_assert_wiphy(ath12k_ar_to_hw(ar)->wiphy);
+
+	ath12k_mac_dec_num_stations(arvif, arsta);
+
+	spin_lock_bh(&ar->ab->base_lock);
+
+	peer = ath12k_peer_find(ar->ab, arvif->vdev_id, sta->addr);
+	if (peer && peer->sta == sta) {
+		ath12k_warn(ar->ab, "Found peer entry %pM n vdev %i after it was supposedly removed\n",
+			    vif->addr, arvif->vdev_id);
+		peer->sta = NULL;
+		list_del(&peer->list);
+		kfree(peer);
+		ar->num_peers--;
+	}
+
+	spin_unlock_bh(&ar->ab->base_lock);
+
+	kfree(arsta->rx_stats);
+	arsta->rx_stats = NULL;
+
+	if (arsta->link_id < IEEE80211_MLD_MAX_NUM_LINKS) {
+		rcu_assign_pointer(ahsta->link[arsta->link_id], NULL);
+		synchronize_rcu();
+		ahsta->links_map &= ~(BIT(arsta->link_id));
+		arsta->link_id = ATH12K_INVALID_LINK_ID;
+		arsta->ahsta = NULL;
+	}
+}
+
+static int ath12k_mac_station_unauthorize(struct ath12k *ar,
+					  struct ath12k_link_vif *arvif,
+					  struct ath12k_link_sta *arsta)
+{
+	struct ath12k_peer *peer;
+	int ret;
+
+	lockdep_assert_wiphy(ath12k_ar_to_hw(ar)->wiphy);
+
+	spin_lock_bh(&ar->ab->base_lock);
+
+	peer = ath12k_peer_find(ar->ab, arvif->vdev_id, arsta->addr);
+	if (peer)
+		peer->is_authorized = false;
+
+	spin_unlock_bh(&ar->ab->base_lock);
+
+	/* Driver should clear the peer keys during mac80211's ref ptr
+	 * gets cleared in __sta_info_destroy_part2 (trans from
+	 * IEEE80211_STA_AUTHORIZED to IEEE80211_STA_ASSOC)
+	 */
+	ret = ath12k_clear_peer_keys(arvif, arsta->addr);
+	if (ret) {
+		ath12k_warn(ar->ab, "failed to clear all peer keys for vdev %i: %d\n",
+			    arvif->vdev_id, ret);
+		return ret;
+	}
+
+	return 0;
+}
+
+static int ath12k_mac_station_authorize(struct ath12k *ar,
+					struct ath12k_link_vif *arvif,
+					struct ath12k_link_sta *arsta)
+{
+	struct ath12k_peer *peer;
+	struct ieee80211_vif *vif = ath12k_ahvif_to_vif(arvif->ahvif);
+	struct ieee80211_sta *sta = ath12k_ahsta_to_sta(arsta->ahsta);
+	int ret;
+
+	lockdep_assert_wiphy(ath12k_ar_to_hw(ar)->wiphy);
+
+	spin_lock_bh(&ar->ab->base_lock);
+
+	peer = ath12k_peer_find(ar->ab, arvif->vdev_id, sta->addr);
+	if (peer)
+		peer->is_authorized = true;
+
+	spin_unlock_bh(&ar->ab->base_lock);
+
+	if (vif->type == NL80211_IFTYPE_STATION && arvif->is_up) {
+		ret = ath12k_wmi_set_peer_param(ar, sta->addr,
+						arvif->vdev_id,
+						WMI_PEER_AUTHORIZE,
+						1);
+		if (ret) {
+			ath12k_warn(ar->ab, "Unable to authorize peer %pM vdev %d: %d\n",
+				    sta->addr, arvif->vdev_id, ret);
+			return ret;
+		}
+	}
+
+	return 0;
+}
+
+static int ath12k_mac_station_remove(struct ath12k *ar,
+				     struct ath12k_link_vif *arvif,
+				     struct ath12k_link_sta *arsta)
+{
+	struct ieee80211_sta *sta = ath12k_ahsta_to_sta(arsta->ahsta);
+	struct ath12k_vif *ahvif = arvif->ahvif;
+	int ret;
+
+	lockdep_assert_wiphy(ath12k_ar_to_hw(ar)->wiphy);
+
+	wiphy_work_cancel(ar->ah->hw->wiphy, &arsta->update_wk);
+
+	if (ahvif->vdev_type == WMI_VDEV_TYPE_STA) {
+		ath12k_bss_disassoc(ar, arvif);
+		ret = ath12k_mac_vdev_stop(arvif);
+		if (ret)
+			ath12k_warn(ar->ab, "failed to stop vdev %i: %d\n",
+				    arvif->vdev_id, ret);
+	}
+
+	ath12k_dp_peer_cleanup(ar, arvif->vdev_id, sta->addr);
+
+	ret = ath12k_peer_delete(ar, arvif->vdev_id, sta->addr);
+	if (ret)
+		ath12k_warn(ar->ab, "Failed to delete peer: %pM for VDEV: %d\n",
+			    sta->addr, arvif->vdev_id);
+	else
+		ath12k_dbg(ar->ab, ATH12K_DBG_MAC, "Removed peer: %pM for VDEV: %d\n",
+			   sta->addr, arvif->vdev_id);
+
+	ath12k_mac_station_post_remove(ar, arvif, arsta);
+
+	return ret;
+}
+
 static int ath12k_mac_station_add(struct ath12k *ar,
 				  struct ath12k_link_vif *arvif,
 				  struct ath12k_link_sta *arsta)
@@ -4933,31 +5063,37 @@  static u32 ath12k_mac_ieee80211_sta_bw_to_wmi(struct ath12k *ar,
 	return bw;
 }
 
-static int ath12k_mac_op_sta_state(struct ieee80211_hw *hw,
-				   struct ieee80211_vif *vif,
-				   struct ieee80211_sta *sta,
-				   enum ieee80211_sta_state old_state,
-				   enum ieee80211_sta_state new_state)
+static int ath12k_mac_handle_link_sta_state(struct ieee80211_hw *hw,
+					    struct ath12k_link_vif *arvif,
+					    struct ath12k_link_sta *arsta,
+					    enum ieee80211_sta_state old_state,
+					    enum ieee80211_sta_state new_state)
 {
-	struct ath12k_vif *ahvif = ath12k_vif_to_ahvif(vif);
-	struct ath12k_sta *ahsta = ath12k_sta_to_ahsta(sta);
-	struct ath12k *ar;
-	struct ath12k_link_vif *arvif;
-	struct ath12k_link_sta *arsta;
-	struct ath12k_peer *peer;
+	struct ath12k *ar = arvif->ar;
+	struct ieee80211_vif *vif = ath12k_ahvif_to_vif(arvif->ahvif);
+	struct ieee80211_sta *sta = ath12k_ahsta_to_sta(arsta->ahsta);
+	struct ath12k_sta *ahsta = arsta->ahsta;
 	int ret = 0;
 
 	lockdep_assert_wiphy(hw->wiphy);
 
-	arvif = &ahvif->deflink;
-	arsta = &ahsta->deflink;
+	/* IEEE80211_STA_NONE -> IEEE80211_STA_NOTEXIST: Remove the station
+	 * from driver
+	 */
+	if ((old_state == IEEE80211_STA_NONE &&
+	     new_state == IEEE80211_STA_NOTEXIST)) {
+		/* ML sta needs separate handling */
+		if (sta->mlo)
+			return 0;
 
-	ar = ath12k_get_ar_by_vif(hw, vif);
-	if (!ar) {
-		WARN_ON_ONCE(1);
-		return -EINVAL;
+		ret = ath12k_mac_station_remove(ar, arvif, arsta);
+		if (ret) {
+			ath12k_warn(ar->ab, "Failed to remove station: %pM for VDEV: %d\n",
+				    arsta->addr, arvif->vdev_id);
+		}
 	}
 
+	/* IEEE80211_STA_NOTEXIST -> IEEE80211_STA_NONE: Add new station to driver */
 	if (old_state == IEEE80211_STA_NOTEXIST &&
 	    new_state == IEEE80211_STA_NONE) {
 		memset(arsta, 0, sizeof(*arsta));
@@ -4975,56 +5111,16 @@  static int ath12k_mac_op_sta_state(struct ieee80211_hw *hw,
 		if (ret)
 			ath12k_warn(ar->ab, "Failed to add station: %pM for VDEV: %d\n",
 				    sta->addr, arvif->vdev_id);
-	} else if ((old_state == IEEE80211_STA_NONE &&
-		    new_state == IEEE80211_STA_NOTEXIST)) {
-		wiphy_work_cancel(hw->wiphy, &arsta->update_wk);
 
-		if (ahvif->vdev_type == WMI_VDEV_TYPE_STA) {
-			ath12k_bss_disassoc(ar, arvif);
-			ret = ath12k_mac_vdev_stop(arvif);
-			if (ret)
-				ath12k_warn(ar->ab, "failed to stop vdev %i: %d\n",
-					    arvif->vdev_id, ret);
-		}
-		ath12k_dp_peer_cleanup(ar, arvif->vdev_id, sta->addr);
-
-		ret = ath12k_peer_delete(ar, arvif->vdev_id, sta->addr);
-		if (ret)
-			ath12k_warn(ar->ab, "Failed to delete peer: %pM for VDEV: %d\n",
-				    sta->addr, arvif->vdev_id);
-		else
-			ath12k_dbg(ar->ab, ATH12K_DBG_MAC, "Removed peer: %pM for VDEV: %d\n",
-				   sta->addr, arvif->vdev_id);
-
-		ath12k_mac_dec_num_stations(arvif, arsta);
-		spin_lock_bh(&ar->ab->base_lock);
-		peer = ath12k_peer_find(ar->ab, arvif->vdev_id, sta->addr);
-		if (peer && peer->sta == sta) {
-			ath12k_warn(ar->ab, "Found peer entry %pM n vdev %i after it was supposedly removed\n",
-				    vif->addr, arvif->vdev_id);
-			peer->sta = NULL;
-			list_del(&peer->list);
-			kfree(peer);
-			ar->num_peers--;
-		}
-		spin_unlock_bh(&ar->ab->base_lock);
-
-		kfree(arsta->rx_stats);
-		arsta->rx_stats = NULL;
-
-		if (arsta->link_id < IEEE80211_MLD_MAX_NUM_LINKS) {
-			rcu_assign_pointer(ahsta->link[arsta->link_id], NULL);
-			synchronize_rcu();
-			ahsta->links_map &= ~(BIT(arsta->link_id));
-			arsta->link_id = ATH12K_INVALID_LINK_ID;
-			arsta->ahsta = NULL;
-		}
+	/* IEEE80211_STA_AUTH -> IEEE80211_STA_ASSOC: Send station assoc command for
+	 * peer associated to AP/Mesh/ADHOC vif type.
+	 */
 	} else if (old_state == IEEE80211_STA_AUTH &&
 		   new_state == IEEE80211_STA_ASSOC &&
 		   (vif->type == NL80211_IFTYPE_AP ||
 		    vif->type == NL80211_IFTYPE_MESH_POINT ||
 		    vif->type == NL80211_IFTYPE_ADHOC)) {
-		ret = ath12k_station_assoc(ar, arvif, arsta, false);
+		ret = ath12k_mac_station_assoc(ar, arvif, arsta, false);
 		if (ret)
 			ath12k_warn(ar->ab, "Failed to associate station: %pM\n",
 				    sta->addr);
@@ -5035,40 +5131,32 @@  static int ath12k_mac_op_sta_state(struct ieee80211_hw *hw,
 		arsta->bw_prev = sta->deflink.bandwidth;
 
 		spin_unlock_bh(&ar->data_lock);
+
+	/* IEEE80211_STA_ASSOC -> IEEE80211_STA_AUTHORIZED: set peer status as
+	 * authorized
+	 */
 	} else if (old_state == IEEE80211_STA_ASSOC &&
 		   new_state == IEEE80211_STA_AUTHORIZED) {
-		spin_lock_bh(&ar->ab->base_lock);
+		ret = ath12k_mac_station_authorize(ar, arvif, arsta);
+		if (ret)
+			ath12k_warn(ar->ab, "Failed to authorize station: %pM\n",
+				    sta->addr);
 
-		peer = ath12k_peer_find(ar->ab, arvif->vdev_id, sta->addr);
-		if (peer)
-			peer->is_authorized = true;
-
-		spin_unlock_bh(&ar->ab->base_lock);
-
-		if (vif->type == NL80211_IFTYPE_STATION && arvif->is_up) {
-			ret = ath12k_wmi_set_peer_param(ar, sta->addr,
-							arvif->vdev_id,
-							WMI_PEER_AUTHORIZE,
-							1);
-			if (ret)
-				ath12k_warn(ar->ab, "Unable to authorize peer %pM vdev %d: %d\n",
-					    sta->addr, arvif->vdev_id, ret);
-		}
+	/* IEEE80211_STA_AUTHORIZED -> IEEE80211_STA_ASSOC: station may be in removal,
+	 * deauthorize it.
+	 */
 	} else if (old_state == IEEE80211_STA_AUTHORIZED &&
 		   new_state == IEEE80211_STA_ASSOC) {
-		spin_lock_bh(&ar->ab->base_lock);
-
-		peer = ath12k_peer_find(ar->ab, arvif->vdev_id, sta->addr);
-		if (peer)
-			peer->is_authorized = false;
-
-		spin_unlock_bh(&ar->ab->base_lock);
+		ath12k_mac_station_unauthorize(ar, arvif, arsta);
+	/* IEEE80211_STA_ASSOC -> IEEE80211_STA_AUTH: disassoc peer connected to
+	 * AP/mesh/ADHOC vif type.
+	 */
 	} else if (old_state == IEEE80211_STA_ASSOC &&
 		   new_state == IEEE80211_STA_AUTH &&
 		   (vif->type == NL80211_IFTYPE_AP ||
 		    vif->type == NL80211_IFTYPE_MESH_POINT ||
 		    vif->type == NL80211_IFTYPE_ADHOC)) {
-		ret = ath12k_station_disassoc(ar, arvif, arsta);
+		ret = ath12k_mac_station_disassoc(ar, arvif, arsta);
 		if (ret)
 			ath12k_warn(ar->ab, "Failed to disassociate station: %pM\n",
 				    sta->addr);
@@ -5077,6 +5165,55 @@  static int ath12k_mac_op_sta_state(struct ieee80211_hw *hw,
 	return ret;
 }
 
+static int ath12k_mac_op_sta_state(struct ieee80211_hw *hw,
+				   struct ieee80211_vif *vif,
+				   struct ieee80211_sta *sta,
+				   enum ieee80211_sta_state old_state,
+				   enum ieee80211_sta_state new_state)
+{
+	struct ath12k_vif *ahvif = ath12k_vif_to_ahvif(vif);
+	struct ath12k_sta *ahsta = ath12k_sta_to_ahsta(sta);
+	struct ath12k_link_vif *arvif;
+	struct ath12k_link_sta *arsta;
+	int ret;
+	u8 link_id = 0;
+
+	lockdep_assert_wiphy(hw->wiphy);
+
+	if (ieee80211_vif_is_mld(vif) && sta->valid_links) {
+		WARN_ON(!sta->mlo && hweight16(sta->valid_links) != 1);
+		link_id = ffs(sta->valid_links) - 1;
+	}
+
+	/* Handle for non-ML station */
+	if (!sta->mlo) {
+		arvif = wiphy_dereference(hw->wiphy, ahvif->link[link_id]);
+		arsta = &ahsta->deflink;
+		arsta->ahsta = ahsta;
+
+		if (WARN_ON(!arvif || !arsta)) {
+			ret = -EINVAL;
+			goto exit;
+		}
+
+		/* vdev might be in deleted */
+		if (WARN_ON(!arvif->ar)) {
+			ret = -EINVAL;
+			goto exit;
+		}
+
+		ret = ath12k_mac_handle_link_sta_state(hw, arvif, arsta,
+						       old_state, new_state);
+		if (ret)
+			goto exit;
+	}
+
+	ret = 0;
+
+exit:
+	return ret;
+}
+
 static int ath12k_mac_op_sta_set_txpwr(struct ieee80211_hw *hw,
 				       struct ieee80211_vif *vif,
 				       struct ieee80211_sta *sta)