ath9k: Fix rx of mcast/bcast frames in PS mode with auto sleep

The functionality to keep the device awake until it is done with
the rx of any mcast/bcast frames which are pending on AP should
also be added to the hardwares which support auto sleep feature.
This patch fixes frequent failures in ARP resolution when it is
initiated by the other end. Currently auto sleep is enabled only
for ar9003 in ath9k.
Signed-off-by: NVasanthakumar Thiagarajan <vasanth@atheros.com>
Signed-off-by: NJohn W. Linville <linville@tuxdriver.com>

drivers/net/wireless/ath/ath9k/recv.c

@@ -19,6 +19,12 @@
 
 #define SKB_CB_ATHBUF(__skb)	(*((struct ath_buf **)__skb->cb))
 
+static inline bool ath9k_check_auto_sleep(struct ath_softc *sc)
+{
+	return sc->ps_enabled &&
+	       (sc->sc_ah->caps.hw_caps & ATH9K_HW_CAP_AUTOSLEEP);
+}
+
 static struct ieee80211_hw * ath_get_virt_hw(struct ath_softc *sc,
 					     struct ieee80211_hdr *hdr)
 {
@@ -616,8 +622,8 @@ static void ath_rx_ps(struct ath_softc *sc, struct sk_buff *skb)
 	hdr = (struct ieee80211_hdr *)skb->data;
 
 	/* Process Beacon and CAB receive in PS state */
-	if ((sc->ps_flags & PS_WAIT_FOR_BEACON) &&
-	    ieee80211_is_beacon(hdr->frame_control))
+	if (((sc->ps_flags & PS_WAIT_FOR_BEACON) || ath9k_check_auto_sleep(sc))
+	    && ieee80211_is_beacon(hdr->frame_control))
 		ath_rx_ps_beacon(sc, skb);
 	else if ((sc->ps_flags & PS_WAIT_FOR_CAB) &&
 		 (ieee80211_is_data(hdr->frame_control) ||
@@ -932,9 +938,10 @@ int ath_rx_tasklet(struct ath_softc *sc, int flush, bool hp)
 			sc->rx.rxotherant = 0;
 		}
 
-		if (unlikely(sc->ps_flags & (PS_WAIT_FOR_BEACON |
+		if (unlikely(ath9k_check_auto_sleep(sc) ||
+			     (sc->ps_flags & (PS_WAIT_FOR_BEACON |
 					      PS_WAIT_FOR_CAB |
-					     PS_WAIT_FOR_PSPOLL_DATA)))
+					      PS_WAIT_FOR_PSPOLL_DATA))))
 			ath_rx_ps(sc, skb);
 
 		ath_rx_send_to_mac80211(hw, sc, skb, rxs);