3 files changed, 114 insertions, 7 deletions

diff --git a/drivers/net/wireless/ath/ath9k/ath9k.h b/drivers/net/wireless/ath/ath9k/ath9k.h
index 10ffc9442859..34256621d417 100644
--- a/drivers/net/wireless/ath/ath9k/ath9k.h
+++ b/drivers/net/wireless/ath/ath9k/ath9k.h
@@ -515,6 +515,7 @@ struct ath_rfkill {
 #define SC_OP_LED_ON            BIT(13)
 #define SC_OP_SCANNING          BIT(14)
 #define SC_OP_TSF_RESET         BIT(15)
+#define SC_OP_WAIT_FOR_CAB      BIT(16)
 
 struct ath_bus_ops {
 	void		(*read_cachesize)(struct ath_softc *sc, int *csz);
diff --git a/drivers/net/wireless/ath/ath9k/debug.h b/drivers/net/wireless/ath/ath9k/debug.h
index 23298b90b52b..db845cf960c9 100644
--- a/drivers/net/wireless/ath/ath9k/debug.h
+++ b/drivers/net/wireless/ath/ath9k/debug.h
@@ -29,6 +29,7 @@ enum ATH_DEBUG {
 	ATH_DBG_BEACON		= 0x00000100,
 	ATH_DBG_CONFIG		= 0x00000200,
 	ATH_DBG_FATAL		= 0x00000400,
+	ATH_DBG_PS		= 0x00000800,
 	ATH_DBG_ANY		= 0xffffffff
 };
 
diff --git a/drivers/net/wireless/ath/ath9k/recv.c b/drivers/net/wireless/ath/ath9k/recv.c
index 58bb26c72771..72e9283bcf7b 100644
--- a/drivers/net/wireless/ath/ath9k/recv.c
+++ b/drivers/net/wireless/ath/ath9k/recv.c
@@ -473,6 +473,112 @@ void ath_flushrecv(struct ath_softc *sc)
 	spin_unlock_bh(&sc->rx.rxflushlock);
 }
 
+static bool ath_beacon_dtim_pending_cab(struct sk_buff *skb)
+{
+	/* Check whether the Beacon frame has DTIM indicating buffered bc/mc */
+	struct ieee80211_mgmt *mgmt;
+	u8 *pos, *end, id, elen;
+	struct ieee80211_tim_ie *tim;
+
+	mgmt = (struct ieee80211_mgmt *)skb->data;
+	pos = mgmt->u.beacon.variable;
+	end = skb->data + skb->len;
+
+	while (pos + 2 < end) {
+		id = *pos++;
+		elen = *pos++;
+		if (pos + elen > end)
+			break;
+
+		if (id == WLAN_EID_TIM) {
+			if (elen < sizeof(*tim))
+				break;
+			tim = (struct ieee80211_tim_ie *) pos;
+			if (tim->dtim_count != 0)
+				break;
+			return tim->bitmap_ctrl & 0x01;
+		}
+
+		pos += elen;
+	}
+
+	return false;
+}
+
+static void ath_rx_ps_back_to_sleep(struct ath_softc *sc)
+{
+	sc->sc_flags &= ~(SC_OP_WAIT_FOR_BEACON | SC_OP_WAIT_FOR_CAB);
+	if (sc->hw->conf.flags & IEEE80211_CONF_PS)
+		ath9k_hw_setpower(sc->sc_ah, ATH9K_PM_NETWORK_SLEEP);
+}
+
+static void ath_rx_ps_beacon(struct ath_softc *sc, struct sk_buff *skb)
+{
+	struct ieee80211_mgmt *mgmt;
+
+	if (skb->len < 24 + 8 + 2 + 2)
+		return;
+
+	mgmt = (struct ieee80211_mgmt *)skb->data;
+	if (memcmp(sc->curbssid, mgmt->bssid, ETH_ALEN) != 0)
+		return; /* not from our current AP */
+
+	if (!(sc->hw->conf.flags & IEEE80211_CONF_PS)) {
+		/* We are not in PS mode anymore; remain awake */
+		DPRINTF(sc, ATH_DBG_PS, "Not in PS mode anymore, remain "
+			"awake\n");
+		sc->sc_flags &= ~(SC_OP_WAIT_FOR_BEACON | SC_OP_WAIT_FOR_CAB);
+		return;
+	}
+
+	if (ath_beacon_dtim_pending_cab(skb)) {
+		/*
+		 * Remain awake waiting for buffered broadcast/multicast
+		 * frames.
+		 */
+		DPRINTF(sc, ATH_DBG_PS, "Received DTIM beacon indicating "
+			"buffered broadcast/multicast frame(s)\n");
+		sc->sc_flags |= SC_OP_WAIT_FOR_CAB;
+		return;
+	}
+
+	if (sc->sc_flags & SC_OP_WAIT_FOR_CAB) {
+		/*
+		 * This can happen if a broadcast frame is dropped or the AP
+		 * fails to send a frame indicating that all CAB frames have
+		 * been delivered.
+		 */
+		DPRINTF(sc, ATH_DBG_PS, "PS wait for CAB frames timed out\n");
+	}
+
+	/* No more broadcast/multicast frames to be received at this point. */
+	ath_rx_ps_back_to_sleep(sc);
+}
+
+static void ath_rx_ps(struct ath_softc *sc, struct sk_buff *skb)
+{
+	struct ieee80211_hdr *hdr;
+
+	hdr = (struct ieee80211_hdr *)skb->data;
+
+	/* Process Beacon and CAB receive in PS state */
+	if (ieee80211_is_beacon(hdr->frame_control))
+		ath_rx_ps_beacon(sc, skb);
+	else if ((sc->sc_flags & SC_OP_WAIT_FOR_CAB) &&
+		 (ieee80211_is_data(hdr->frame_control) ||
+		  ieee80211_is_action(hdr->frame_control)) &&
+		 is_multicast_ether_addr(hdr->addr1) &&
+		 !ieee80211_has_moredata(hdr->frame_control)) {
+		DPRINTF(sc, ATH_DBG_PS, "All PS CAB frames received, back to "
+			"sleep\n");
+		/*
+		 * No more broadcast/multicast frames to be received at this
+		 * point.
+		 */
+		ath_rx_ps_back_to_sleep(sc);
+	}
+}
+
 static void ath_rx_send_to_mac80211(struct ath_softc *sc, struct sk_buff *skb,
 				    struct ieee80211_rx_status *rx_status)
 {
@@ -667,8 +773,6 @@ int ath_rx_tasklet(struct ath_softc *sc, int flush)
 			rx_status.flag &= ~RX_FLAG_DECRYPTED;
 		}
 
-		ath_rx_send_to_mac80211(sc, skb, &rx_status);
-
 		/* We will now give hardware our shiny new allocated skb */
 		bf->bf_mpdu = requeue_skb;
 		bf->bf_buf_addr = dma_map_single(sc->dev, requeue_skb->data,
@@ -680,6 +784,7 @@ int ath_rx_tasklet(struct ath_softc *sc, int flush)
 			bf->bf_mpdu = NULL;
 			DPRINTF(sc, ATH_DBG_FATAL,
 				"dma_mapping_error() on RX\n");
+			ath_rx_send_to_mac80211(sc, skb, &rx_status);
 			break;
 		}
 		bf->bf_dmacontext = bf->bf_buf_addr;
@@ -695,11 +800,11 @@ int ath_rx_tasklet(struct ath_softc *sc, int flush)
 			sc->rx.rxotherant = 0;
 		}
 
-		if (ieee80211_is_beacon(fc) &&
-				(sc->sc_flags & SC_OP_WAIT_FOR_BEACON)) {
-			sc->sc_flags &= ~SC_OP_WAIT_FOR_BEACON;
-			ath9k_hw_setpower(sc->sc_ah, ATH9K_PM_NETWORK_SLEEP);
-		}
+		if (unlikely(sc->sc_flags & SC_OP_WAIT_FOR_BEACON))
+			ath_rx_ps(sc, skb);
+
+		ath_rx_send_to_mac80211(sc, skb, &rx_status);
+
 requeue:
 		list_move_tail(&bf->list, &sc->rx.rxbuf);
 		ath_rx_buf_link(sc, bf);