[PATCH wireless] wifi: mac80211: include TIM bitmap control for buffered S1G mcast traffic

[Lachlan Hodges <lachlan.hodges@morsemicro.com>]

Currently when building the S1G TIM element, we only build the bitmap
control if we have buffered unicast traffic. Since AID 0 sits within
the bitmap control if we have buffered multicast traffic with no
buffered unicast traffic the bitmap control won't be emitted and
dozing stations will be unaware of buffered multicast.

To fix, only exclude the bitmap control byte when we don't have
both buffered unicast and multicast traffic.

Fixes: ee6360945483 ("wifi: mac80211: support block bitmap S1G TIM encoding")
Signed-off-by: Lachlan Hodges <lachlan.hodges@morsemicro.com>
---
 net/mac80211/tx.c | 41 ++++++++++++++++++++++-------------------
 1 file changed, 22 insertions(+), 19 deletions(-)

diff --git a/net/mac80211/tx.c b/net/mac80211/tx.c
index 3a1e2c9e1565..3896c7b2c4e5 100644
--- a/net/mac80211/tx.c
+++ b/net/mac80211/tx.c
@@ -5089,10 +5089,18 @@ static void ieee80211_beacon_add_tim_pvb(struct ps_data *ps,
  */
 static void ieee80211_s1g_beacon_add_tim_pvb(struct ps_data *ps,
 					     struct sk_buff *skb,
-					     bool mcast_traffic)
+					     bool mcast_traffic,
+					     bool ucast_traffic)
 {
 	int blk;
 
+	/*
+	 * if no unicast and multicast traffic don't emit a bitmap control
+	 * or pvb
+	 */
+	if (!mcast_traffic && !ucast_traffic)
+		return;
+
 	/*
 	 * Emit a bitmap control block with a page slice number of 31 and a
 	 * page index of 0 which indicates as per IEEE80211-2024 9.4.2.5.1
@@ -5101,6 +5109,10 @@ static void ieee80211_s1g_beacon_add_tim_pvb(struct ps_data *ps,
 	 */
 	skb_put_u8(skb, mcast_traffic | (31 << 1));
 
+	/* If there's no unicast traffic we don't need to include a PVB. */
+	if (!ucast_traffic)
+		return;
+
 	/* Emit an encoded block for each non-zero sub-block */
 	for (blk = 0; blk < IEEE80211_MAX_SUPPORTED_S1G_TIM_BLOCKS; blk++) {
 		u8 blk_bmap = 0;
@@ -5182,25 +5194,16 @@ static void __ieee80211_beacon_add_tim(struct ieee80211_sub_if_data *sdata,
 
 	ps->dtim_bc_mc = mcast_traffic;
 
-	if (have_bits) {
-		if (s1g)
-			ieee80211_s1g_beacon_add_tim_pvb(ps, skb,
-							 mcast_traffic);
-		else
-			ieee80211_beacon_add_tim_pvb(ps, skb, mcast_traffic);
+	if (s1g) {
+		ieee80211_s1g_beacon_add_tim_pvb(ps, skb, mcast_traffic,
+						 have_bits);
+	} else if (have_bits) {
+		ieee80211_beacon_add_tim_pvb(ps, skb, mcast_traffic);
 	} else {
-		/*
-		 * If there is no buffered unicast traffic for an S1G
-		 * interface, we can exclude the bitmap control. This is in
-		 * contrast to other phy types as they do include the bitmap
-		 * control and pvb even when there is no buffered traffic.
-		 */
-		if (!s1g) {
-			/* Bitmap control */
-			skb_put_u8(skb, mcast_traffic);
-			/* Part Virt Bitmap */
-			skb_put_u8(skb, 0);
-		}
+		/* Bitmap control */
+		skb_put_u8(skb, mcast_traffic);
+		/* Part Virt Bitmap */
+		skb_put_u8(skb, 0);
 	}
 
 	tim->datalen = skb_tail_pointer(skb) - tim->data;
-- 
2.43.0