}
 
 static void ath9k_mgd_prepare_tx(struct ieee80211_hw *hw,
-                                struct ieee80211_vif *vif)
+                                struct ieee80211_vif *vif,
+                                u16 duration)
 {
        struct ath_softc *sc = hw->priv;
        struct ath_common *common = ath9k_hw_common(sc->sc_ah);
 
 }
 
 static void iwl_mvm_mac_mgd_prepare_tx(struct ieee80211_hw *hw,
-                                     struct ieee80211_vif *vif)
+                                      struct ieee80211_vif *vif,
+                                      u16 req_duration)
 {
        struct iwl_mvm *mvm = IWL_MAC80211_GET_MVM(hw);
        u32 duration = IWL_MVM_TE_SESSION_PROTECTION_MAX_TIME_MS;
        if (iwl_mvm_ref_sync(mvm, IWL_MVM_REF_PREPARE_TX))
                return;
 
+       if (req_duration > duration)
+               duration = req_duration;
+
        mutex_lock(&mvm->mutex);
        /* Try really hard to protect the session and hear a beacon */
        iwl_mvm_protect_session(mvm, vif, duration, min_duration, 500, false);
 
  *     frame in case that no beacon was heard from the AP/P2P GO.
  *     The callback will be called before each transmission and upon return
  *     mac80211 will transmit the frame right away.
+ *      If duration is greater than zero, mac80211 hints to the driver the
+ *      duration for which the operation is requested.
  *     The callback is optional and can (should!) sleep.
  *
  * @mgd_protect_tdls_discover: Protect a TDLS discovery session. After sending
                                  u32 sset, u8 *data);
 
        void    (*mgd_prepare_tx)(struct ieee80211_hw *hw,
-                                 struct ieee80211_vif *vif);
+                                 struct ieee80211_vif *vif,
+                                 u16 duration);
 
        void    (*mgd_protect_tdls_discover)(struct ieee80211_hw *hw,
                                             struct ieee80211_vif *vif);
 
 /*
 * Portions of this file
 * Copyright(c) 2016 Intel Deutschland GmbH
+* Copyright (C) 2018 Intel Corporation
 */
 
 #ifndef __MAC80211_DRIVER_OPS
 }
 
 static inline void drv_mgd_prepare_tx(struct ieee80211_local *local,
-                                     struct ieee80211_sub_if_data *sdata)
+                                     struct ieee80211_sub_if_data *sdata,
+                                     u16 duration)
 {
        might_sleep();
 
                return;
        WARN_ON_ONCE(sdata->vif.type != NL80211_IFTYPE_STATION);
 
-       trace_drv_mgd_prepare_tx(local, sdata);
+       trace_drv_mgd_prepare_tx(local, sdata, duration);
        if (local->ops->mgd_prepare_tx)
-               local->ops->mgd_prepare_tx(&local->hw, &sdata->vif);
+               local->ops->mgd_prepare_tx(&local->hw, &sdata->vif, duration);
        trace_drv_return_void(local);
 }
 
 
                return;
        }
 
-       drv_mgd_prepare_tx(local, sdata);
+       drv_mgd_prepare_tx(local, sdata, 0);
 
        IEEE80211_SKB_CB(skb)->flags |= IEEE80211_TX_INTFL_DONT_ENCRYPT;
        if (ieee80211_hw_check(&local->hw, REPORTS_TX_ACK_STATUS))
                 */
                if (ieee80211_hw_check(&local->hw, DEAUTH_NEED_MGD_TX_PREP) &&
                    !ifmgd->have_beacon)
-                       drv_mgd_prepare_tx(sdata->local, sdata);
+                       drv_mgd_prepare_tx(sdata->local, sdata, 0);
 
                ieee80211_send_deauth_disassoc(sdata, ifmgd->bssid, stype,
                                               reason, tx, frame_buf);
        if (!elems.challenge)
                return;
        auth_data->expected_transaction = 4;
-       drv_mgd_prepare_tx(sdata->local, sdata);
+       drv_mgd_prepare_tx(sdata->local, sdata, 0);
        if (ieee80211_hw_check(&local->hw, REPORTS_TX_ACK_STATUS))
                tx_flags = IEEE80211_TX_CTL_REQ_TX_STATUS |
                           IEEE80211_TX_INTFL_MLME_CONN_TX;
        u32 tx_flags = 0;
        u16 trans = 1;
        u16 status = 0;
+       u16 prepare_tx_duration = 0;
 
        sdata_assert_lock(sdata);
 
                return -ETIMEDOUT;
        }
 
-       drv_mgd_prepare_tx(local, sdata);
+       if (auth_data->algorithm == WLAN_AUTH_SAE)
+               prepare_tx_duration =
+                       jiffies_to_msecs(IEEE80211_AUTH_TIMEOUT_SAE);
+
+       drv_mgd_prepare_tx(local, sdata, prepare_tx_duration);
 
        sdata_info(sdata, "send auth to %pM (try %d/%d)\n",
                   auth_data->bss->bssid, auth_data->tries,
                           req->bssid, req->reason_code,
                           ieee80211_get_reason_code_string(req->reason_code));
 
-               drv_mgd_prepare_tx(sdata->local, sdata);
+               drv_mgd_prepare_tx(sdata->local, sdata, 0);
                ieee80211_send_deauth_disassoc(sdata, req->bssid,
                                               IEEE80211_STYPE_DEAUTH,
                                               req->reason_code, tx,
                           req->bssid, req->reason_code,
                           ieee80211_get_reason_code_string(req->reason_code));
 
-               drv_mgd_prepare_tx(sdata->local, sdata);
+               drv_mgd_prepare_tx(sdata->local, sdata, 0);
                ieee80211_send_deauth_disassoc(sdata, req->bssid,
                                               IEEE80211_STYPE_DEAUTH,
                                               req->reason_code, tx,
 
 /*
 * Portions of this file
 * Copyright(c) 2016 Intel Deutschland GmbH
+* Copyright (C) 2018 Intel Corporation
 */
 
 #if !defined(__MAC80211_DRIVER_TRACE) || defined(TRACE_HEADER_MULTI_READ)
        TP_ARGS(local, sta, tids, num_frames, reason, more_data)
 );
 
-DEFINE_EVENT(local_sdata_evt, drv_mgd_prepare_tx,
+TRACE_EVENT(drv_mgd_prepare_tx,
        TP_PROTO(struct ieee80211_local *local,
-                struct ieee80211_sub_if_data *sdata),
+                struct ieee80211_sub_if_data *sdata,
+                u16 duration),
 
-       TP_ARGS(local, sdata)
+       TP_ARGS(local, sdata, duration),
+
+       TP_STRUCT__entry(
+               LOCAL_ENTRY
+               VIF_ENTRY
+               __field(u32, duration)
+       ),
+
+       TP_fast_assign(
+               LOCAL_ASSIGN;
+               VIF_ASSIGN;
+               __entry->duration = duration;
+       ),
+
+       TP_printk(
+               LOCAL_PR_FMT VIF_PR_FMT " duration: %u",
+               LOCAL_PR_ARG, VIF_PR_ARG, __entry->duration
+       )
 );
 
 DEFINE_EVENT(local_sdata_evt, drv_mgd_protect_tdls_discover,