5 files changed, 180 insertions, 9 deletions

diff --git a/sys/dev/iwm/if_iwm.c b/sys/dev/iwm/if_iwm.c
index 965a1eb1d23c..c197d991dd37 100644
--- a/sys/dev/iwm/if_iwm.c
+++ b/sys/dev/iwm/if_iwm.c
@@ -1263,6 +1263,7 @@ iwm_stop_device(struct iwm_softc *sc)
 		iv->is_uploaded = 0;
 	}
 	sc->sc_firmware_state = 0;
+	sc->sc_flags &= ~IWM_FLAG_TE_ACTIVE;
 
 	/* device going down, Stop using ICT table */
 	sc->sc_flags &= ~IWM_FLAG_USE_ICT;
@@ -4050,8 +4051,7 @@ iwm_auth(struct ieee80211vap *vap, struct iwm_softc *sc)
 	 */
 	/* XXX duration is in units of TU, not MS */
 	duration = IWM_MVM_TE_SESSION_PROTECTION_MAX_TIME_MS;
-	iwm_mvm_protect_session(sc, iv, duration, 500 /* XXX magic number */);
-	DELAY(100);
+	iwm_mvm_protect_session(sc, iv, duration, 500 /* XXX magic number */, TRUE);
 
 	error = 0;
 out:
@@ -4349,6 +4349,15 @@ iwm_newstate(struct ieee80211vap *vap, enum ieee80211_state nstate, int arg)
 			in->in_assoc = 0;
 	}
 
+	if ((vap->iv_state == IEEE80211_S_AUTH ||
+	     vap->iv_state == IEEE80211_S_ASSOC ||
+	     vap->iv_state == IEEE80211_S_RUN) &&
+	    (nstate == IEEE80211_S_INIT ||
+	     nstate == IEEE80211_S_SCAN ||
+	     nstate == IEEE80211_S_AUTH)) {
+		iwm_mvm_stop_session_protection(sc, ivp);
+	}
+
 	if ((vap->iv_state == IEEE80211_S_RUN ||
 	     vap->iv_state == IEEE80211_S_ASSOC) &&
 	    nstate == IEEE80211_S_INIT) {
diff --git a/sys/dev/iwm/if_iwm_debug.h b/sys/dev/iwm/if_iwm_debug.h
index 2a2ac731288d..35d624436e08 100644
--- a/sys/dev/iwm/if_iwm_debug.h
+++ b/sys/dev/iwm/if_iwm_debug.h
@@ -44,6 +44,7 @@ enum {
 	IWM_DEBUG_TEMP		= 0x00100000,	/* Thermal Sensor handling */
 	IWM_DEBUG_FW		= 0x00200000,	/* Firmware management */
 	IWM_DEBUG_LAR		= 0x00400000,	/* Location Aware Regulatory */
+	IWM_DEBUG_TE		= 0x00800000,	/* Time Event handling */
 	IWM_DEBUG_REGISTER	= 0x20000000,	/* print chipset register */
 	IWM_DEBUG_TRACE		= 0x40000000,	/* Print begin and start driver function */
 	IWM_DEBUG_FATAL		= 0x80000000,	/* fatal errors */
diff --git a/sys/dev/iwm/if_iwm_time_event.c b/sys/dev/iwm/if_iwm_time_event.c
index db24ea986fe4..340d19d27246 100644
--- a/sys/dev/iwm/if_iwm_time_event.c
+++ b/sys/dev/iwm/if_iwm_time_event.c
@@ -155,6 +155,7 @@ __FBSDID("$FreeBSD$");
 #include <dev/iwm/if_iwmvar.h>
 #include <dev/iwm/if_iwm_debug.h>
 #include <dev/iwm/if_iwm_util.h>
+#include <dev/iwm/if_iwm_notif_wait.h>
 #include <dev/iwm/if_iwm_pcie_trans.h>
 #include <dev/iwm/if_iwm_time_event.h>
 
@@ -166,31 +167,132 @@ __FBSDID("$FreeBSD$");
 #define IWM_MVM_ROC_TE_TYPE_MGMT_TX IWM_TE_P2P_CLIENT_ASSOC
 
 static int
+iwm_mvm_te_notif(struct iwm_softc *sc, struct iwm_rx_packet *pkt,
+    void *data)
+{
+	struct iwm_time_event_notif *resp;
+	int resp_len = iwm_rx_packet_payload_len(pkt);
+
+	if (pkt->hdr.code != IWM_TIME_EVENT_NOTIFICATION ||
+	    resp_len != sizeof(*resp)) {
+		IWM_DPRINTF(sc, IWM_DEBUG_TE,
+		    "Invalid TIME_EVENT_NOTIFICATION response\n");
+		return 1;
+	}
+
+	resp = (void *)pkt->data;
+
+	/* te_data->uid is already set in the TIME_EVENT_CMD response */
+	if (le32toh(resp->unique_id) != sc->sc_time_event_uid)
+		return false;
+
+	IWM_DPRINTF(sc, IWM_DEBUG_TE,
+	    "TIME_EVENT_NOTIFICATION response - UID = 0x%x\n",
+	    sc->sc_time_event_uid);
+	if (!resp->status) {
+		IWM_DPRINTF(sc, IWM_DEBUG_TE,
+		    "TIME_EVENT_NOTIFICATION received but not executed\n");
+	}
+
+	return 1;
+}
+
+static int
+iwm_mvm_time_event_response(struct iwm_softc *sc, struct iwm_rx_packet *pkt,
+    void *data)
+{
+	struct iwm_time_event_resp *resp;
+	int resp_len = iwm_rx_packet_payload_len(pkt);
+
+	if (pkt->hdr.code != IWM_TIME_EVENT_CMD ||
+	    resp_len != sizeof(*resp)) {
+		IWM_DPRINTF(sc, IWM_DEBUG_TE,
+		    "Invalid TIME_EVENT_CMD response\n");
+		return 1;
+	}
+
+	resp = (void *)pkt->data;
+
+	/* we should never get a response to another TIME_EVENT_CMD here */
+	if (le32toh(resp->id) != IWM_TE_BSS_STA_AGGRESSIVE_ASSOC) {
+		IWM_DPRINTF(sc, IWM_DEBUG_TE,
+		    "Got TIME_EVENT_CMD response with wrong id: %d\n",
+		    le32toh(resp->id));
+		return 0;
+	}
+
+	sc->sc_time_event_uid = le32toh(resp->unique_id);
+	IWM_DPRINTF(sc, IWM_DEBUG_TE,
+	    "TIME_EVENT_CMD response - UID = 0x%x\n", sc->sc_time_event_uid);
+	return 1;
+}
+
+
+/* XXX Use the te_data function argument properly, like in iwlwifi's code. */
+
+static int
 iwm_mvm_time_event_send_add(struct iwm_softc *sc, struct iwm_vap *ivp,
 	void *te_data, struct iwm_time_event_cmd *te_cmd)
 {
+	static const uint16_t time_event_response[] = { IWM_TIME_EVENT_CMD };
+	struct iwm_notification_wait wait_time_event;
 	int ret;
 
-	IWM_DPRINTF(sc, IWM_DEBUG_CMD | IWM_DEBUG_RESET,
+	IWM_DPRINTF(sc, IWM_DEBUG_TE,
 	    "Add new TE, duration %d TU\n", le32toh(te_cmd->duration));
 
-	ret = iwm_mvm_send_cmd_pdu(sc, IWM_TIME_EVENT_CMD, IWM_CMD_SYNC,
-	    sizeof(*te_cmd), te_cmd);
+	/*
+	 * Use a notification wait, which really just processes the
+	 * command response and doesn't wait for anything, in order
+	 * to be able to process the response and get the UID inside
+	 * the RX path. Using CMD_WANT_SKB doesn't work because it
+	 * stores the buffer and then wakes up this thread, by which
+	 * time another notification (that the time event started)
+	 * might already be processed unsuccessfully.
+	 */
+	iwm_init_notification_wait(sc->sc_notif_wait, &wait_time_event,
+				   time_event_response,
+				   nitems(time_event_response),
+				   iwm_mvm_time_event_response, /*te_data*/NULL);
+
+	ret = iwm_mvm_send_cmd_pdu(sc, IWM_TIME_EVENT_CMD, 0, sizeof(*te_cmd),
+	    te_cmd);
 	if (ret) {
-		IWM_DPRINTF(sc, IWM_DEBUG_CMD | IWM_DEBUG_RESET,
+		IWM_DPRINTF(sc, IWM_DEBUG_TE,
 		    "%s: Couldn't send IWM_TIME_EVENT_CMD: %d\n",
 		    __func__, ret);
+		iwm_remove_notification(sc->sc_notif_wait, &wait_time_event);
+		return ret;
+	}
+
+	/* No need to wait for anything, so just pass 1 (0 isn't valid) */
+	IWM_UNLOCK(sc);
+	ret = iwm_wait_notification(sc->sc_notif_wait, &wait_time_event, 1);
+	IWM_LOCK(sc);
+	/* should never fail */
+	if (ret) {
+		IWM_DPRINTF(sc, IWM_DEBUG_TE,
+		    "%s: Failed to get response for IWM_TIME_EVENT_CMD: %d\n",
+		    __func__, ret);
 	}
 
 	return ret;
 }
 
+#define TU_TO_HZ(tu)	(((uint64_t)(tu) * 1024 * hz) / 1000000)
+
 void
 iwm_mvm_protect_session(struct iwm_softc *sc, struct iwm_vap *ivp,
-	uint32_t duration, uint32_t max_delay)
+	uint32_t duration, uint32_t max_delay, boolean_t wait_for_notif)
 {
+	const uint16_t te_notif_response[] = { IWM_TIME_EVENT_NOTIFICATION };
+	struct iwm_notification_wait wait_te_notif;
 	struct iwm_time_event_cmd time_cmd = {};
 
+	/* Do nothing if a time event is already scheduled. */
+	if (sc->sc_flags & IWM_FLAG_TE_ACTIVE)
+		return;
+
 	time_cmd.action = htole32(IWM_FW_CTXT_ACTION_ADD);
 	time_cmd.id_and_color =
 	    htole32(IWM_FW_CMD_ID_AND_COLOR(ivp->id, ivp->color));
@@ -209,5 +311,58 @@ iwm_mvm_protect_session(struct iwm_softc *sc, struct iwm_vap *ivp,
 	        IWM_TE_V2_NOTIF_HOST_EVENT_END |
 		IWM_T2_V2_START_IMMEDIATELY);
 
-	iwm_mvm_time_event_send_add(sc, ivp, /*te_data*/NULL, &time_cmd);
+	if (!wait_for_notif) {
+		iwm_mvm_time_event_send_add(sc, ivp, /*te_data*/NULL, &time_cmd);
+		DELAY(100);
+		sc->sc_flags |= IWM_FLAG_TE_ACTIVE;
+		return;
+	}
+
+	/*
+	 * Create notification_wait for the TIME_EVENT_NOTIFICATION to use
+	 * right after we send the time event
+	 */
+	iwm_init_notification_wait(sc->sc_notif_wait, &wait_te_notif,
+	    te_notif_response, nitems(te_notif_response),
+	    iwm_mvm_te_notif, /*te_data*/NULL);
+
+	/* If TE was sent OK - wait for the notification that started */
+	if (iwm_mvm_time_event_send_add(sc, ivp, /*te_data*/NULL, &time_cmd)) {
+		IWM_DPRINTF(sc, IWM_DEBUG_TE,
+		    "%s: Failed to add TE to protect session\n", __func__);
+		iwm_remove_notification(sc->sc_notif_wait, &wait_te_notif);
+	} else {
+		sc->sc_flags |= IWM_FLAG_TE_ACTIVE;
+		IWM_UNLOCK(sc);
+		if (iwm_wait_notification(sc->sc_notif_wait, &wait_te_notif,
+		    TU_TO_HZ(max_delay))) {
+			IWM_DPRINTF(sc, IWM_DEBUG_TE,
+			    "%s: Failed to protect session until TE\n",
+			    __func__);
+		}
+		IWM_LOCK(sc);
+	}
+}
+
+void
+iwm_mvm_stop_session_protection(struct iwm_softc *sc, struct iwm_vap *ivp)
+{
+	struct iwm_time_event_cmd time_cmd = {};
+
+	/* Do nothing if the time event has already ended. */
+	if ((sc->sc_flags & IWM_FLAG_TE_ACTIVE) == 0)
+		return;
+
+	time_cmd.id = htole32(sc->sc_time_event_uid);
+	time_cmd.action = htole32(IWM_FW_CTXT_ACTION_REMOVE);
+	time_cmd.id_and_color =
+	    htole32(IWM_FW_CMD_ID_AND_COLOR(ivp->id, ivp->color));
+
+	IWM_DPRINTF(sc, IWM_DEBUG_TE,
+	    "%s: Removing TE 0x%x\n", __func__, le32toh(time_cmd.id));
+	if (iwm_mvm_send_cmd_pdu(sc, IWM_TIME_EVENT_CMD, 0, sizeof(time_cmd),
+	    &time_cmd) == 0)
+		sc->sc_flags &= ~IWM_FLAG_TE_ACTIVE;
+
+	DELAY(100);
 }
diff --git a/sys/dev/iwm/if_iwm_time_event.h b/sys/dev/iwm/if_iwm_time_event.h
index 4c741428e1ea..0e6dd8887997 100644
--- a/sys/dev/iwm/if_iwm_time_event.h
+++ b/sys/dev/iwm/if_iwm_time_event.h
@@ -108,6 +108,8 @@
 #define	__IF_IWM_TIME_EVENT_H__
 
 extern	void iwm_mvm_protect_session(struct iwm_softc *sc, struct iwm_vap *ivp,
-	    uint32_t duration, uint32_t max_delay);
+	    uint32_t duration, uint32_t max_delay, boolean_t wait_for_notif);
+extern	void iwm_mvm_stop_session_protection(struct iwm_softc *sc,
+	    struct iwm_vap *ivp);
 
 #endif	/* __IF_IWM_TIME_EVENT_H__ */
diff --git a/sys/dev/iwm/if_iwmvar.h b/sys/dev/iwm/if_iwmvar.h
index ac7b35d56bd9..536027493a1a 100644
--- a/sys/dev/iwm/if_iwmvar.h
+++ b/sys/dev/iwm/if_iwmvar.h
@@ -434,6 +434,7 @@ struct iwm_softc {
 #define IWM_FLAG_BUSY		(1 << 4)
 #define IWM_FLAG_SCANNING	(1 << 5)
 #define IWM_FLAG_SCAN_RUNNING	(1 << 6)
+#define IWM_FLAG_TE_ACTIVE	(1 << 7)
 
 	struct intr_config_hook sc_preinit_hook;
 	struct callout		sc_watchdog_to;
@@ -562,6 +563,9 @@ struct iwm_softc {
 
 	/* Track firmware state for STA association. */
 	int			sc_firmware_state;
+
+	/* Unique ID (assigned by the firmware) of the current Time Event. */
+	uint32_t		sc_time_event_uid;
 };
 
 #define IWM_LOCK_INIT(_sc) \