3 files changed, 72 insertions, 26 deletions

diff --git a/net/mac80211/ieee80211_i.h b/net/mac80211/ieee80211_i.h
index ff7bc307827b..5bc0745368fe 100644
--- a/net/mac80211/ieee80211_i.h
+++ b/net/mac80211/ieee80211_i.h
@@ -1265,6 +1265,8 @@ void ieee80211_send_nullfunc(struct ieee80211_local *local,
                              int powersave);
 void ieee80211_sta_rx_notify(struct ieee80211_sub_if_data *sdata,
                              struct ieee80211_hdr *hdr);
+void ieee80211_sta_tx_notify(struct ieee80211_sub_if_data *sdata,
+                             struct ieee80211_hdr *hdr);
 void ieee80211_beacon_connection_loss_work(struct work_struct *work);
 
 void ieee80211_wake_queues_by_reason(struct ieee80211_hw *hw,
diff --git a/net/mac80211/mlme.c b/net/mac80211/mlme.c
index 849d9c639add..33ffce3ec605 100644
--- a/net/mac80211/mlme.c
+++ b/net/mac80211/mlme.c
@@ -1034,6 +1034,51 @@ void ieee80211_sta_rx_notify(struct ieee80211_sub_if_data *sdata,
         ieee80211_sta_reset_conn_monitor(sdata);
 }
 
+static void ieee80211_reset_ap_probe(struct ieee80211_sub_if_data *sdata)
+{
+        struct ieee80211_if_managed *ifmgd = &sdata->u.mgd;
+
+        if (!(ifmgd->flags & (IEEE80211_STA_BEACON_POLL |
+                              IEEE80211_STA_CONNECTION_POLL)))
+            return;
+
+        ifmgd->flags &= ~(IEEE80211_STA_CONNECTION_POLL |
+                          IEEE80211_STA_BEACON_POLL);
+        mutex_lock(&sdata->local->iflist_mtx);
+        ieee80211_recalc_ps(sdata->local, -1);
+        mutex_unlock(&sdata->local->iflist_mtx);
+
+        if (sdata->local->hw.flags & IEEE80211_HW_CONNECTION_MONITOR)
+                return;
+
+        /*
+         * We've received a probe response, but are not sure whether
+         * we have or will be receiving any beacons or data, so let's
+         * schedule the timers again, just in case.
+         */
+        ieee80211_sta_reset_beacon_monitor(sdata);
+
+        mod_timer(&ifmgd->conn_mon_timer,
+                  round_jiffies_up(jiffies +
+                                   IEEE80211_CONNECTION_IDLE_TIME));
+}
+
+void ieee80211_sta_tx_notify(struct ieee80211_sub_if_data *sdata,
+                             struct ieee80211_hdr *hdr)
+{
+        if (!ieee80211_is_data(hdr->frame_control) &&
+            !ieee80211_is_nullfunc(hdr->frame_control))
+            return;
+
+        ieee80211_sta_reset_conn_monitor(sdata);
+
+        if (ieee80211_is_nullfunc(hdr->frame_control) &&
+            sdata->u.mgd.probe_send_count > 0) {
+                sdata->u.mgd.probe_send_count = 0;
+                ieee80211_queue_work(&sdata->local->hw, &sdata->work);
+        }
+}
+
 static void ieee80211_mgd_probe_ap_send(struct ieee80211_sub_if_data *sdata)
 {
         struct ieee80211_if_managed *ifmgd = &sdata->u.mgd;
@@ -1049,8 +1094,19 @@ static void ieee80211_mgd_probe_ap_send(struct ieee80211_sub_if_data *sdata)
         if (ifmgd->probe_send_count >= unicast_limit)
                 dst = NULL;
 
-        ssid = ieee80211_bss_get_ie(ifmgd->associated, WLAN_EID_SSID);
-        ieee80211_send_probe_req(sdata, dst, ssid + 2, ssid[1], NULL, 0);
+        /*
+         * When the hardware reports an accurate Tx ACK status, it's
+         * better to send a nullfunc frame instead of a probe request,
+         * as it will kick us off the AP quickly if we aren't associated
+         * anymore. The timeout will be reset if the frame is ACKed by
+         * the AP.
+         */
+        if (sdata->local->hw.flags & IEEE80211_HW_REPORTS_TX_ACK_STATUS)
+                ieee80211_send_nullfunc(sdata->local, sdata, 0);
+        else {
+                ssid = ieee80211_bss_get_ie(ifmgd->associated, WLAN_EID_SSID);
+                ieee80211_send_probe_req(sdata, dst, ssid + 2, ssid[1], NULL, 0);
+        }
 
         ifmgd->probe_send_count++;
         ifmgd->probe_timeout = jiffies + IEEE80211_PROBE_WAIT;
@@ -1517,29 +1573,8 @@ static void ieee80211_rx_mgmt_probe_resp(struct ieee80211_sub_if_data *sdata,
         ieee80211_rx_bss_info(sdata, mgmt, len, rx_status, &elems, false);
 
         if (ifmgd->associated &&
-            memcmp(mgmt->bssid, ifmgd->associated->bssid, ETH_ALEN) == 0 &&
-            ifmgd->flags & (IEEE80211_STA_BEACON_POLL |
-                            IEEE80211_STA_CONNECTION_POLL)) {
-                ifmgd->flags &= ~(IEEE80211_STA_CONNECTION_POLL |
-                                  IEEE80211_STA_BEACON_POLL);
-                mutex_lock(&sdata->local->iflist_mtx);
-                ieee80211_recalc_ps(sdata->local, -1);
-                mutex_unlock(&sdata->local->iflist_mtx);
-
-                if (sdata->local->hw.flags & IEEE80211_HW_CONNECTION_MONITOR)
-                        return;
-
-                /*
-                 * We've received a probe response, but are not sure whether
-                 * we have or will be receiving any beacons or data, so let's
-                 * schedule the timers again, just in case.
-                 */
-                ieee80211_sta_reset_beacon_monitor(sdata);
-
-                mod_timer(&ifmgd->conn_mon_timer,
-                          round_jiffies_up(jiffies +
-                                           IEEE80211_CONNECTION_IDLE_TIME));
-        }
+            memcmp(mgmt->bssid, ifmgd->associated->bssid, ETH_ALEN) == 0)
+                ieee80211_reset_ap_probe(sdata);
 }
 
 /*
@@ -1891,7 +1926,12 @@ void ieee80211_sta_work(struct ieee80211_sub_if_data *sdata)
                 u8 bssid[ETH_ALEN];
 
                 memcpy(bssid, ifmgd->associated->bssid, ETH_ALEN);
-                if (time_is_after_jiffies(ifmgd->probe_timeout))
+
+                /* ACK received for nullfunc probing frame */
+                if (!ifmgd->probe_send_count)
+                        ieee80211_reset_ap_probe(sdata);
+
+                else if (time_is_after_jiffies(ifmgd->probe_timeout))
                         run_again(ifmgd, ifmgd->probe_timeout);
 
                 else if (ifmgd->probe_send_count < IEEE80211_MAX_PROBE_TRIES) {
diff --git a/net/mac80211/status.c b/net/mac80211/status.c
index 3153c19893b8..8695cd11dfd9 100644
--- a/net/mac80211/status.c
+++ b/net/mac80211/status.c
@@ -155,6 +155,10 @@ static void ieee80211_frame_acked(struct sta_info *sta, struct sk_buff *skb)
 
                 ieee80211_queue_work(&local->hw, &local->recalc_smps);
         }
+
+        if ((sdata->vif.type == NL80211_IFTYPE_STATION) &&
+            (local->hw.flags & IEEE80211_HW_REPORTS_TX_ACK_STATUS))
+                ieee80211_sta_tx_notify(sdata, (void *) skb->data);
 }
 
 void ieee80211_tx_status(struct ieee80211_hw *hw, struct sk_buff *skb)