const void *ptr, void *data)
 {
        struct wmi_tlv_swba_parse *swba = data;
+       struct wmi_tim_info_arg *tim_info_arg;
+       const struct wmi_tim_info *tim_info_ev = ptr;
 
        if (tag != WMI_TLV_TAG_STRUCT_TIM_INFO)
                return -EPROTO;
        if (swba->n_tim >= ARRAY_SIZE(swba->arg->tim_info))
                return -ENOBUFS;
 
-       swba->arg->tim_info[swba->n_tim++] = ptr;
+       if (__le32_to_cpu(tim_info_ev->tim_len) >
+            sizeof(tim_info_ev->tim_bitmap)) {
+               ath10k_warn(ar, "refusing to parse invalid swba structure\n");
+               return -EPROTO;
+       }
+
+       tim_info_arg = &swba->arg->tim_info[swba->n_tim];
+       tim_info_arg->tim_len = tim_info_ev->tim_len;
+       tim_info_arg->tim_mcast = tim_info_ev->tim_mcast;
+       tim_info_arg->tim_bitmap = tim_info_ev->tim_bitmap;
+       tim_info_arg->tim_changed = tim_info_ev->tim_changed;
+       tim_info_arg->tim_num_ps_pending = tim_info_ev->tim_num_ps_pending;
+
+       swba->n_tim++;
+
        return 0;
 }
 
 
 static void ath10k_wmi_update_tim(struct ath10k *ar,
                                  struct ath10k_vif *arvif,
                                  struct sk_buff *bcn,
-                                 const struct wmi_tim_info *tim_info)
+                                 const struct wmi_tim_info_arg *tim_info)
 {
        struct ieee80211_hdr *hdr = (struct ieee80211_hdr *)bcn->data;
        struct ieee80211_tim_ie *tim;
        u8 *ies, *ie;
        u8 ie_len, pvm_len;
        __le32 t;
-       u32 v;
+       u32 v, tim_len;
+
+       /* When FW reports 0 in tim_len, ensure atleast first byte
+        * in tim_bitmap is considered for pvm calculation.
+        */
+       tim_len = tim_info->tim_len ? __le32_to_cpu(tim_info->tim_len) : 1;
 
        /* if next SWBA has no tim_changed the tim_bitmap is garbage.
         * we must copy the bitmap upon change and reuse it later */
        if (__le32_to_cpu(tim_info->tim_changed)) {
                int i;
 
-               BUILD_BUG_ON(sizeof(arvif->u.ap.tim_bitmap) !=
-                            sizeof(tim_info->tim_bitmap));
+               if (sizeof(arvif->u.ap.tim_bitmap) < tim_len) {
+                       ath10k_warn(ar, "SWBA TIM field is too big (%u), truncated it to %zu",
+                                   tim_len, sizeof(arvif->u.ap.tim_bitmap));
+                       tim_len = sizeof(arvif->u.ap.tim_bitmap);
+               }
 
-               for (i = 0; i < sizeof(arvif->u.ap.tim_bitmap); i++) {
+               for (i = 0; i < tim_len; i++) {
                        t = tim_info->tim_bitmap[i / 4];
                        v = __le32_to_cpu(t);
                        arvif->u.ap.tim_bitmap[i] = (v >> ((i % 4) * 8)) & 0xFF;
                }
 
-               /* FW reports either length 0 or 16
-                * so we calculate this on our own */
+               /* FW reports either length 0 or length based on max supported
+                * station. so we calculate this on our own
+                */
                arvif->u.ap.tim_len = 0;
-               for (i = 0; i < sizeof(arvif->u.ap.tim_bitmap); i++)
+               for (i = 0; i < tim_len; i++)
                        if (arvif->u.ap.tim_bitmap[i])
                                arvif->u.ap.tim_len = i;
 
        pvm_len = ie_len - 3; /* exclude dtim count, dtim period, bmap ctl */
 
        if (pvm_len < arvif->u.ap.tim_len) {
-               int expand_size = sizeof(arvif->u.ap.tim_bitmap) - pvm_len;
+               int expand_size = tim_len - pvm_len;
                int move_size = skb_tail_pointer(bcn) - (ie + 2 + ie_len);
                void *next_ie = ie + 2 + ie_len;
 
                }
        }
 
-       if (pvm_len > sizeof(arvif->u.ap.tim_bitmap)) {
+       if (pvm_len > tim_len) {
                ath10k_warn(ar, "tim pvm length is too great (%d)\n", pvm_len);
                return;
        }
                if (WARN_ON_ONCE(i == ARRAY_SIZE(arg->tim_info)))
                        break;
 
-               arg->tim_info[i] = &ev->bcn_info[i].tim_info;
+               if (__le32_to_cpu(ev->bcn_info[i].tim_info.tim_len) >
+                    sizeof(ev->bcn_info[i].tim_info.tim_bitmap)) {
+                       ath10k_warn(ar, "refusing to parse invalid swba structure\n");
+                       return -EPROTO;
+               }
+
+               arg->tim_info[i].tim_len = ev->bcn_info[i].tim_info.tim_len;
+               arg->tim_info[i].tim_mcast = ev->bcn_info[i].tim_info.tim_mcast;
+               arg->tim_info[i].tim_bitmap =
+                               ev->bcn_info[i].tim_info.tim_bitmap;
+               arg->tim_info[i].tim_changed =
+                               ev->bcn_info[i].tim_info.tim_changed;
+               arg->tim_info[i].tim_num_ps_pending =
+                               ev->bcn_info[i].tim_info.tim_num_ps_pending;
+
                arg->noa_info[i] = &ev->bcn_info[i].p2p_noa_info;
                i++;
        }
        struct wmi_swba_ev_arg arg = {};
        u32 map;
        int i = -1;
-       const struct wmi_tim_info *tim_info;
+       const struct wmi_tim_info_arg *tim_info;
        const struct wmi_p2p_noa_info *noa_info;
        struct ath10k_vif *arvif;
        struct sk_buff *bcn;
                        break;
                }
 
-               tim_info = arg.tim_info[i];
+               tim_info = &arg.tim_info[i];
                noa_info = arg.noa_info[i];
 
                ath10k_dbg(ar, ATH10K_DBG_MGMT,
                           __le32_to_cpu(tim_info->tim_bitmap[1]),
                           __le32_to_cpu(tim_info->tim_bitmap[0]));
 
+               /* TODO: Only first 4 word from tim_bitmap is dumped.
+                * Extend debug code to dump full tim_bitmap.
+                */
+
                arvif = ath10k_get_arvif(ar, vdev_id);
                if (arvif == NULL) {
                        ath10k_warn(ar, "no vif for vdev_id %d found\n",