mt76: clear sta powersave flag after notifying driver

If the driver has configured the hardware to drop packets for a sta that was
in powersave mode, then clearing the flag too early opens up a small race
window during which packets could be dropped on wakeup
Signed-off-by: Felix Fietkau <nbd@nbd.name>

drivers/net/wireless/mediatek/mt76/mac80211.c

@@ -1163,10 +1163,12 @@ mt76_check_sta(struct mt76_dev *dev, struct sk_buff *skb)
 
 	if (ps)
 		set_bit(MT_WCID_FLAG_PS, &wcid->flags);
-	else
-		clear_bit(MT_WCID_FLAG_PS, &wcid->flags);
 
 	dev->drv->sta_ps(dev, sta, ps);
+
+	if (!ps)
+		clear_bit(MT_WCID_FLAG_PS, &wcid->flags);
+
 	ieee80211_sta_ps_transition(sta, ps);
 }