add set/cancel_priority to ap mode as well

ROC when the ap gets the first auth request, and CROC
after authorization (or after connection fails).

Signed-off-by: Eliad Peller <eliad@wizery.com>

7 files changed, 50 insertions, 0 deletions

diff --git a/src/ap/ap_drv_ops.h b/src/ap/ap_drv_ops.h
index 169c91b..e3b15c0 100644
--- a/src/ap/ap_drv_ops.h
+++ b/src/ap/ap_drv_ops.h
@@ -209,4 +209,19 @@ static inline void hostapd_drv_poll_client(struct hostapd_data *hapd,
 	hapd->driver->poll_client(hapd->drv_priv, own_addr, addr, qos);
 }
 
+static inline int hostapd_drv_set_priority(struct hostapd_data *hapd)
+{
+	if (hapd->driver == NULL || hapd->driver->set_priority == NULL)
+		return 0;
+	return hapd->driver->set_priority(hapd->drv_priv);
+}
+
+static inline int hostapd_drv_cancel_priority(struct hostapd_data *hapd)
+{
+	if (hapd->driver == NULL || hapd->driver->cancel_priority == NULL)
+		return 0;
+	return hapd->driver->cancel_priority(hapd->drv_priv);
+}
+
+
 #endif /* AP_DRV_OPS */
diff --git a/src/ap/drv_callbacks.c b/src/ap/drv_callbacks.c
index 27fc7d5..f80a08c 100644
--- a/src/ap/drv_callbacks.c
+++ b/src/ap/drv_callbacks.c
@@ -252,6 +252,7 @@ void hostapd_notif_disassoc(struct hostapd_data *hapd, const u8 *addr)
 	ap_sta_set_authorized(hapd, sta, 0);
 	sta->flags &= ~(WLAN_STA_AUTH | WLAN_STA_ASSOC |
 			WLAN_STA_ASSOC_REQ_OK);
+	ap_sta_cancel_priority(hapd, sta);
 	wpa_auth_sm_event(sta->wpa_sm, WPA_DISASSOC);
 	sta->acct_terminate_cause = RADIUS_ACCT_TERMINATE_CAUSE_USER_REQUEST;
 	ieee802_1x_notify_port_enabled(sta->eapol_sm, 0);
diff --git a/src/ap/hostapd.h b/src/ap/hostapd.h
index 4fb3c51..0465844 100644
--- a/src/ap/hostapd.h
+++ b/src/ap/hostapd.h
@@ -61,6 +61,7 @@ struct hostapd_data {
 	u8 own_addr[ETH_ALEN];
 
 	int num_sta; /* number of entries in sta_list */
+	int num_priorities; /* num of stations with set_priorities */
 	struct sta_info *sta_list; /* STA info list head */
 #define STA_HASH_SIZE 256
 #define STA_HASH(sta) (sta[5])
diff --git a/src/ap/ieee802_11.c b/src/ap/ieee802_11.c
index 6821205..ed1661a 100644
--- a/src/ap/ieee802_11.c
+++ b/src/ap/ieee802_11.c
@@ -481,6 +481,8 @@ static void handle_auth(struct hostapd_data *hapd,
 #endif /* CONFIG_IEEE80211R */
 	}
 
+	ap_sta_set_priority(hapd, sta);
+
  fail:
 	send_auth_reply(hapd, mgmt->sa, mgmt->bssid, auth_alg,
 			auth_transaction + 1, resp, resp_ies, resp_ies_len);
@@ -1221,6 +1223,7 @@ static void handle_deauth(struct hostapd_data *hapd,
 	ap_sta_set_authorized(hapd, sta, 0);
 	sta->flags &= ~(WLAN_STA_AUTH | WLAN_STA_ASSOC |
 			WLAN_STA_ASSOC_REQ_OK);
+	ap_sta_cancel_priority(hapd, sta);
 	wpa_auth_sm_event(sta->wpa_sm, WPA_DEAUTH);
 	hostapd_logger(hapd, sta->addr, HOSTAPD_MODULE_IEEE80211,
 		       HOSTAPD_LEVEL_DEBUG, "deauthenticated");
diff --git a/src/ap/sta_info.c b/src/ap/sta_info.c
index 0e52690..05b78c0 100644
--- a/src/ap/sta_info.c
+++ b/src/ap/sta_info.c
@@ -604,6 +604,7 @@ void ap_sta_deauthenticate(struct hostapd_data *hapd, struct sta_info *sta,
 	wpa_printf(MSG_DEBUG, "%s: deauthenticate STA " MACSTR,
 		   hapd->conf->iface, MAC2STR(sta->addr));
 	sta->flags &= ~(WLAN_STA_AUTH | WLAN_STA_ASSOC | WLAN_STA_ASSOC_REQ_OK);
+	ap_sta_cancel_priority(hapd, sta);
 	ap_sta_set_authorized(hapd, sta, 0);
 	sta->timeout_next = STA_REMOVE;
 	wpa_printf(MSG_DEBUG, "%s: reschedule ap_handle_timer timeout "
@@ -881,6 +882,7 @@ void ap_sta_set_authorized(struct hostapd_data *hapd, struct sta_info *sta,
 				AP_STA_CONNECTED MACSTR, MAC2STR(sta->addr));
 
 		sta->flags |= WLAN_STA_AUTHORIZED;
+		ap_sta_cancel_priority(hapd, sta);
 	} else {
 		if (dev_addr)
 			wpa_msg(hapd->msg_ctx, MSG_INFO, AP_STA_DISCONNECTED
@@ -924,6 +926,7 @@ void ap_sta_disconnect(struct hostapd_data *hapd, struct sta_info *sta,
 	wpa_auth_sm_event(sta->wpa_sm, WPA_DEAUTH);
 	ieee802_1x_notify_port_enabled(sta->eapol_sm, 0);
 	sta->flags &= ~(WLAN_STA_AUTH | WLAN_STA_ASSOC | WLAN_STA_ASSOC_REQ_OK);
+	ap_sta_cancel_priority(hapd, sta);
 	wpa_printf(MSG_DEBUG, "%s: reschedule ap_handle_timer timeout "
 		   "for " MACSTR " (%d seconds - "
 		   "AP_MAX_INACTIVITY_AFTER_DEAUTH)",
@@ -965,3 +968,26 @@ void ap_sta_disassoc_cb(struct hostapd_data *hapd, struct sta_info *sta)
 	eloop_cancel_timeout(ap_sta_disassoc_cb_timeout, hapd, sta);
 	ap_sta_disassoc_cb_timeout(hapd, sta);
 }
+
+void ap_sta_set_priority(struct hostapd_data *hapd, struct sta_info *sta)
+{
+	if (sta->priority_set)
+		return;
+
+	if (sta->flags & WLAN_STA_AUTHORIZED)
+		return;
+
+	sta->priority_set = 1;
+	if (!hapd->num_priorities++)
+		hostapd_drv_set_priority(hapd);
+}
+
+void ap_sta_cancel_priority(struct hostapd_data *hapd, struct sta_info *sta)
+{
+	if (!sta->priority_set)
+		return;
+
+	sta->priority_set = 0;
+	if (!--hapd->num_priorities)
+		hostapd_drv_cancel_priority(hapd);
+}
diff --git a/src/ap/sta_info.h b/src/ap/sta_info.h
index cef428d..42e384b 100644
--- a/src/ap/sta_info.h
+++ b/src/ap/sta_info.h
@@ -56,6 +56,7 @@ struct sta_info {
 	unsigned int no_ht_set:1;
 	unsigned int ht_20mhz_set:1;
 	unsigned int no_p2p_set:1;
+	unsigned int priority_set:1;
 
 	u16 auth_alg;
 	u8 previous_ap[6];
@@ -176,5 +177,7 @@ static inline int ap_sta_is_authorized(struct sta_info *sta)
 
 void ap_sta_deauth_cb(struct hostapd_data *hapd, struct sta_info *sta);
 void ap_sta_disassoc_cb(struct hostapd_data *hapd, struct sta_info *sta);
+void ap_sta_set_priority(struct hostapd_data *hapd, struct sta_info *sta);
+void ap_sta_cancel_priority(struct hostapd_data *hapd, struct sta_info *sta);
 
 #endif /* STA_INFO_H */
diff --git a/src/ap/tkip_countermeasures.c b/src/ap/tkip_countermeasures.c
index dd5aa68..bb6f9ef 100644
--- a/src/ap/tkip_countermeasures.c
+++ b/src/ap/tkip_countermeasures.c
@@ -55,6 +55,7 @@ static void ieee80211_tkip_countermeasures_start(struct hostapd_data *hapd)
 		}
 		hostapd_drv_sta_deauth(hapd, sta->addr,
 				       WLAN_REASON_MICHAEL_MIC_FAILURE);
+		ap_sta_cancel_priority(hapd, sta);
 		ap_free_sta(hapd, sta);
 	}
 }