From fc115b364ae041dd82b1f13e3d133d436d5789a8 Mon Sep 17 00:00:00 2001
From: Andrei Otcheretianski <andrei.otcheretianski@intel.com>
Date: Mon, 25 Jan 2016 12:29:20 +0200
Subject: [PATCH 3/3] AP: Set STA assoc flag in the driver before sending Assoc
 Resp frame

Previously, stations were added to the driver only after the
(Re)Association Response frame was acked. In the time period between the
station has acked the (Re)Association Response frame and the time the
station was added to the kernel, the station can already start sending
Data frames, which will be dropped by the hardware/driver. In addition
to the data loss, the driver may ignore NDPs with PM bit set from this
STA.

Fix this by setting/adding the STA with associated flag set to the
driver before the AP sends the (Re)Association Response frame with
status success. If the (Re)Association Response frame wasn't acked,
remove the station from the driver.

TODO: Either move this marking of the driver STA entry into the driver
(or mac80211) on receiving ACK frame for status=0 (Re)Association
Request frame or explain in the commit message why this behavior is
acceptable with hostapd marking the STA association before even having
sent out the (Re)Association Response frame, never mind receiving the
ACK frame for it. The IEEE 802.11 standard describes the state change
happening on the non-AP STA acknowledging the response.

Signed-off-by: Andrei Otcheretianski <andrei.otcheretianski@intel.com>
---
 src/ap/ieee802_11.c | 150 +++++++++++++++++++++++++++++-----------------------
 1 file changed, 84 insertions(+), 66 deletions(-)

diff --git a/src/ap/ieee802_11.c b/src/ap/ieee802_11.c
index 1638d51..3b8a9db 100644
--- a/src/ap/ieee802_11.c
+++ b/src/ap/ieee802_11.c
@@ -1723,6 +1723,62 @@ static void send_deauth(struct hostapd_data *hapd, const u8 *addr,
 }
 
 
+static int add_associated_sta(struct hostapd_data *hapd,
+			      struct sta_info *sta)
+{
+	struct ieee80211_ht_capabilities ht_cap;
+	struct ieee80211_vht_capabilities vht_cap;
+
+	/*
+	 * Remove the STA entry in order to make sure the STA PS state gets
+	 * cleared and configuration gets updated.
+	 * This is relevant for cases, such as FT over the DS, where a station
+	 * re-associates back to the same AP but skips the authentication flow.
+	 * Or, if working with a driver that doesn't support full AP client
+	 * state.
+	 */
+	if (!sta->added_unassoc)
+		hostapd_drv_sta_remove(hapd, sta->addr);
+
+#ifdef CONFIG_IEEE80211N
+	if (sta->flags & WLAN_STA_HT)
+		hostapd_get_ht_capab(hapd, sta->ht_capabilities, &ht_cap);
+#endif /* CONFIG_IEEE80211N */
+#ifdef CONFIG_IEEE80211AC
+	if (sta->flags & WLAN_STA_VHT)
+		hostapd_get_vht_capab(hapd, sta->vht_capabilities, &vht_cap);
+#endif /* CONFIG_IEEE80211AC */
+
+	/*
+	 * Add the station with forced WLAN_STA_ASSOC flag. The sta->flags
+	 * will be set when the assoc. resp. ACK is processed.
+	 */
+	if (hostapd_sta_add(hapd, sta->addr, sta->aid, sta->capability,
+			    sta->supported_rates, sta->supported_rates_len,
+			    sta->listen_interval,
+			    sta->flags & WLAN_STA_HT ? &ht_cap : NULL,
+			    sta->flags & WLAN_STA_VHT ? &vht_cap : NULL,
+			    sta->flags | WLAN_STA_ASSOC, sta->qosinfo,
+			    sta->vht_opmode, sta->added_unassoc)) {
+		hostapd_logger(hapd, sta->addr,
+			       HOSTAPD_MODULE_IEEE80211, HOSTAPD_LEVEL_NOTICE,
+			       "Could not %s STA to kernel driver",
+			       sta->added_unassoc ? "set" : "add");
+
+		if (sta->added_unassoc) {
+			hostapd_drv_sta_remove(hapd, sta->addr);
+			sta->added_unassoc = 0;
+		}
+
+		return -1;
+	}
+
+	sta->added_unassoc = 0;
+
+	return 0;
+}
+
+
 static u16 send_assoc_resp(struct hostapd_data *hapd, struct sta_info *sta,
 			   u16 status_code, int reassoc, const u8 *ies,
 			   size_t ies_len)
@@ -2060,9 +2116,23 @@ static void handle_assoc(struct hostapd_data *hapd,
 	sta->timeout_next = STA_NULLFUNC;
 
  fail:
+	/*
+	 * In case of successful response add the station to the driver,
+	 * otherwise the kernel may ignore data frames before we process the
+	 * ack. In case of failure, this station will be removed
+	 */
+	if (resp == WLAN_STATUS_SUCCESS && add_associated_sta(hapd, sta))
+		resp = WLAN_STATUS_AP_UNABLE_TO_HANDLE_NEW_STA;
+
 	reply_res = send_assoc_resp(hapd, sta, resp, reassoc, pos, left);
-	if (sta->added_unassoc && (resp != WLAN_STATUS_SUCCESS ||
-				   reply_res != WLAN_STATUS_SUCCESS)) {
+
+	/*
+	 * Remove the station in case of failure tx a successful response (it
+	 * was added associated to the driver) or if the station was previously
+	 * added unassociated.
+	 */
+	if ((reply_res != WLAN_STATUS_SUCCESS &&
+	     resp == WLAN_STATUS_SUCCESS) || sta->added_unassoc) {
 		hostapd_drv_sta_remove(hapd, sta->addr);
 		sta->added_unassoc = 0;
 	}
@@ -2558,8 +2628,6 @@ static void handle_assoc_cb(struct hostapd_data *hapd,
 	u16 status;
 	struct sta_info *sta;
 	int new_assoc = 1;
-	struct ieee80211_ht_capabilities ht_cap;
-	struct ieee80211_vht_capabilities vht_cap;
 
 	sta = ap_get_sta(hapd, mgmt->da);
 	if (!sta) {
@@ -2573,21 +2641,26 @@ static void handle_assoc_cb(struct hostapd_data *hapd,
 		wpa_printf(MSG_INFO,
 			   "handle_assoc_cb(reassoc=%d) - too short payload (len=%lu)",
 			   reassoc, (unsigned long) len);
-		goto remove_sta;
+		hostapd_drv_sta_remove(hapd, sta->addr);
+		return;
 	}
 
+	if (reassoc)
+		status = le_to_host16(mgmt->u.reassoc_resp.status_code);
+	else
+		status = le_to_host16(mgmt->u.assoc_resp.status_code);
+
 	if (!ok) {
 		hostapd_logger(hapd, mgmt->da, HOSTAPD_MODULE_IEEE80211,
 			       HOSTAPD_LEVEL_DEBUG,
 			       "did not acknowledge association response");
 		sta->flags &= ~WLAN_STA_ASSOC_REQ_OK;
-		goto remove_sta;
-	}
+		/* The STA is added only in case of SUCCESS */
+		if (status == WLAN_STATUS_SUCCESS)
+			hostapd_drv_sta_remove(hapd, sta->addr);
 
-	if (reassoc)
-		status = le_to_host16(mgmt->u.reassoc_resp.status_code);
-	else
-		status = le_to_host16(mgmt->u.assoc_resp.status_code);
+		return;
+	}
 
 	if (status != WLAN_STATUS_SUCCESS)
 		return;
@@ -2623,54 +2696,6 @@ static void handle_assoc_cb(struct hostapd_data *hapd,
 	sta->sa_query_timed_out = 0;
 #endif /* CONFIG_IEEE80211W */
 
-	/*
-	 * Remove the STA entry in order to make sure the STA PS state gets
-	 * cleared and configuration gets updated in case of reassociation back
-	 * to the same AP.
-	 *
-	 * This is relevant for cases, such as FT over the DS, where a station
-	 * reassociates back to the same AP but skips the authentication flow
-	 * and if working with a driver that doesn't support full AP client
-	 * state.
-	 */
-	if (!sta->added_unassoc)
-		hostapd_drv_sta_remove(hapd, sta->addr);
-
-#ifdef CONFIG_IEEE80211N
-	if (sta->flags & WLAN_STA_HT)
-		hostapd_get_ht_capab(hapd, sta->ht_capabilities, &ht_cap);
-#endif /* CONFIG_IEEE80211N */
-#ifdef CONFIG_IEEE80211AC
-	if (sta->flags & WLAN_STA_VHT)
-		hostapd_get_vht_capab(hapd, sta->vht_capabilities, &vht_cap);
-#endif /* CONFIG_IEEE80211AC */
-
-	if (hostapd_sta_add(hapd, sta->addr, sta->aid, sta->capability,
-			    sta->supported_rates, sta->supported_rates_len,
-			    sta->listen_interval,
-			    sta->flags & WLAN_STA_HT ? &ht_cap : NULL,
-			    sta->flags & WLAN_STA_VHT ? &vht_cap : NULL,
-			    sta->flags, sta->qosinfo, sta->vht_opmode,
-			    sta->added_unassoc)) {
-		hostapd_logger(hapd, sta->addr, HOSTAPD_MODULE_IEEE80211,
-			       HOSTAPD_LEVEL_NOTICE,
-			       "Could not %s STA to kernel driver",
-			       sta->added_unassoc ? "set" : "add");
-		ap_sta_disconnect(hapd, sta, sta->addr,
-				  WLAN_REASON_DISASSOC_AP_BUSY);
-		if (sta->added_unassoc)
-			goto remove_sta;
-		return;
-	}
-
-	/*
-	 * added_unassoc flag is set for a station that was added to the driver
-	 * in unassociated state. Clear this flag once the station has completed
-	 * association, to make sure the STA entry will be cleared from the
-	 * driver in case of reassocition back to the same AP.
-	 */
-	sta->added_unassoc = 0;
-
 	if (sta->flags & WLAN_STA_WDS) {
 		int ret;
 		char ifname_wds[IFNAMSIZ + 1];
@@ -2702,14 +2727,7 @@ static void handle_assoc_cb(struct hostapd_data *hapd,
 	else
 		wpa_auth_sm_event(sta->wpa_sm, WPA_ASSOC);
 	hapd->new_assoc_sta_cb(hapd, sta, !new_assoc);
-
 	ieee802_1x_notify_port_enabled(sta->eapol_sm, 1);
-
-remove_sta:
-	if (sta->added_unassoc) {
-		hostapd_drv_sta_remove(hapd, sta->addr);
-		sta->added_unassoc = 0;
-	}
 }
 
 
-- 
1.9.1

