3 files changed, 15 insertions, 5 deletions

diff --git a/drivers/net/wireless/iwlwifi/mvm/mac-ctxt.c b/drivers/net/wireless/iwlwifi/mvm/mac-ctxt.c
index c08a17a3cab9..0854dc338881 100644
--- a/drivers/net/wireless/iwlwifi/mvm/mac-ctxt.c
+++ b/drivers/net/wireless/iwlwifi/mvm/mac-ctxt.c
@@ -584,7 +584,11 @@ static void iwl_mvm_mac_ctxt_cmd_fill_sta(struct iwl_mvm *mvm,
                                           struct ieee80211_vif *vif,
                                           struct iwl_mac_data_sta *ctxt_sta)
 {
-        ctxt_sta->is_assoc = cpu_to_le32(vif->bss_conf.assoc ? 1 : 0);
+        /* We need the dtim_period to set the MAC as associated */
+        if (vif->bss_conf.assoc && vif->bss_conf.dtim_period)
+                ctxt_sta->is_assoc = cpu_to_le32(1);
+        else
+                ctxt_sta->is_assoc = cpu_to_le32(0);
 
         ctxt_sta->bi = cpu_to_le32(vif->bss_conf.beacon_int);
         ctxt_sta->bi_reciprocal =
diff --git a/drivers/net/wireless/iwlwifi/mvm/mac80211.c b/drivers/net/wireless/iwlwifi/mvm/mac80211.c
index e0d79e328d0c..6bfcb3b5bfa6 100644
--- a/drivers/net/wireless/iwlwifi/mvm/mac80211.c
+++ b/drivers/net/wireless/iwlwifi/mvm/mac80211.c
@@ -114,7 +114,6 @@ int iwl_mvm_mac_setup_register(struct iwl_mvm *mvm)
                     IEEE80211_HW_QUEUE_CONTROL |
                     IEEE80211_HW_WANT_MONITOR_VIF |
                     IEEE80211_HW_SCAN_WHILE_IDLE |
-                    IEEE80211_HW_NEED_DTIM_BEFORE_ASSOC |
                     IEEE80211_HW_SUPPORTS_PS |
                     IEEE80211_HW_SUPPORTS_DYNAMIC_PS |
                     IEEE80211_HW_AMPDU_AGGREGATION;
@@ -671,8 +670,6 @@ static void iwl_mvm_bss_info_changed_station(struct iwl_mvm *mvm,
                                 IWL_ERR(mvm, "failed to update quotas\n");
                                 return;
                         }
-                        iwl_mvm_remove_time_event(mvm, mvmvif,
-                                                  &mvmvif->time_event_data);
                 } else if (mvmvif->ap_sta_id != IWL_MVM_STATION_COUNT) {
                         /* remove AP station now that the MAC is unassoc */
                         ret = iwl_mvm_rm_sta_id(mvm, vif, mvmvif->ap_sta_id);
@@ -684,6 +681,13 @@ static void iwl_mvm_bss_info_changed_station(struct iwl_mvm *mvm,
                         if (ret)
                                 IWL_ERR(mvm, "failed to update quotas\n");
                 }
+        } else if (changes & BSS_CHANGED_DTIM_PERIOD) {
+                /*
+                 * We received a beacon _after_ association so
+                 * remove the session protection.
+                 */
+                iwl_mvm_remove_time_event(mvm, mvmvif,
+                                          &mvmvif->time_event_data);
         } else if (changes & BSS_CHANGED_PS) {
                 /*
                  * TODO: remove this temporary code.
diff --git a/drivers/net/wireless/iwlwifi/mvm/time-event.c b/drivers/net/wireless/iwlwifi/mvm/time-event.c
index 4d62a5d4254d..194bfb7a7d7f 100644
--- a/drivers/net/wireless/iwlwifi/mvm/time-event.c
+++ b/drivers/net/wireless/iwlwifi/mvm/time-event.c
@@ -184,9 +184,11 @@ static void iwl_mvm_te_handle_notif(struct iwl_mvm *mvm,
                  */
                 if (te_data->vif->type == NL80211_IFTYPE_STATION &&
                     (!te_data->vif->bss_conf.assoc ||
-                     !te_data->vif->bss_conf.dtim_period))
+                     !te_data->vif->bss_conf.dtim_period)) {
                         IWL_ERR(mvm,
                                 "No assocation and the time event is over already...\n");
+                        ieee80211_connection_loss(te_data->vif);
+                }
 
                 iwl_mvm_te_clear_data(mvm, te_data);
         } else if (le32_to_cpu(notif->action) == TE_NOTIF_HOST_START) {