offload_assist));
 }
 
+static u32 iwl_mvm_get_tx_ant(struct iwl_mvm *mvm,
+                             struct ieee80211_tx_info *info,
+                             struct ieee80211_sta *sta, __le16 fc)
+{
+       if (info->band == NL80211_BAND_2GHZ &&
+           !iwl_mvm_bt_coex_is_shared_ant_avail(mvm))
+               return mvm->cfg->non_shared_ant << RATE_MCS_ANT_POS;
+
+       if (sta && ieee80211_is_data(fc)) {
+               struct iwl_mvm_sta *mvmsta = iwl_mvm_sta_from_mac80211(sta);
+
+               return BIT(mvmsta->tx_ant) << RATE_MCS_ANT_POS;
+       }
+
+       return BIT(mvm->mgmt_last_antenna_idx) << RATE_MCS_ANT_POS;
+}
+
 static u32 iwl_mvm_get_tx_rate(struct iwl_mvm *mvm,
                               struct ieee80211_tx_info *info,
                               struct ieee80211_sta *sta)
 {
        int rate_idx;
        u8 rate_plcp;
-       u32 rate_flags;
+       u32 rate_flags = 0;
 
        /* HT rate doesn't make sense for a non data frame */
        WARN_ONCE(info->control.rates[0].flags & IEEE80211_TX_RC_MCS,
        /* Get PLCP rate for tx_cmd->rate_n_flags */
        rate_plcp = iwl_mvm_mac80211_idx_to_hwrate(rate_idx);
 
-       if (info->band == NL80211_BAND_2GHZ &&
-           !iwl_mvm_bt_coex_is_shared_ant_avail(mvm))
-               rate_flags = mvm->cfg->non_shared_ant << RATE_MCS_ANT_POS;
-       else
-               rate_flags =
-                       BIT(mvm->mgmt_last_antenna_idx) << RATE_MCS_ANT_POS;
-
        /* Set CCK flag as needed */
        if ((rate_idx >= IWL_FIRST_CCK_RATE) && (rate_idx <= IWL_LAST_CCK_RATE))
                rate_flags |= RATE_MCS_CCK_MSK;
        return (u32)rate_plcp | rate_flags;
 }
 
+static u32 iwl_mvm_get_tx_rate_n_flags(struct iwl_mvm *mvm,
+                                      struct ieee80211_tx_info *info,
+                                      struct ieee80211_sta *sta, __le16 fc)
+{
+       return iwl_mvm_get_tx_rate(mvm, info, sta) |
+               iwl_mvm_get_tx_ant(mvm, info, sta, fc);
+}
+
 /*
  * Sets the fields in the Tx cmd that are rate related
  */
         */
 
        if (ieee80211_is_data(fc) && sta) {
-               tx_cmd->initial_rate_index = 0;
-               tx_cmd->tx_flags |= cpu_to_le32(TX_CMD_FLG_STA_RATE);
-               return;
+               struct iwl_mvm_sta *mvmsta = iwl_mvm_sta_from_mac80211(sta);
+
+               if (mvmsta->sta_state >= IEEE80211_STA_AUTHORIZED) {
+                       tx_cmd->initial_rate_index = 0;
+                       tx_cmd->tx_flags |= cpu_to_le32(TX_CMD_FLG_STA_RATE);
+                       return;
+               }
        } else if (ieee80211_is_back_req(fc)) {
                tx_cmd->tx_flags |=
                        cpu_to_le32(TX_CMD_FLG_ACK | TX_CMD_FLG_BAR);
        }
 
        /* Set the rate in the TX cmd */
-       tx_cmd->rate_n_flags = cpu_to_le32(iwl_mvm_get_tx_rate(mvm, info, sta));
+       tx_cmd->rate_n_flags =
+               cpu_to_le32(iwl_mvm_get_tx_rate_n_flags(mvm, info, sta, fc));
 }
 
 static inline void iwl_mvm_set_tx_cmd_pn(struct ieee80211_tx_info *info,
                u16 offload_assist = 0;
                u32 rate_n_flags = 0;
                u16 flags = 0;
+               struct iwl_mvm_sta *mvmsta = sta ?
+                       iwl_mvm_sta_from_mac80211(sta) : NULL;
 
                if (ieee80211_is_data_qos(hdr->frame_control)) {
                        u8 *qc = ieee80211_get_qos_ctl(hdr);
                if (!info->control.hw_key)
                        flags |= IWL_TX_FLAGS_ENCRYPT_DIS;
 
-               /* For data packets rate info comes from the fw */
-               if (!(ieee80211_is_data(hdr->frame_control) && sta)) {
+               /*
+                * For data packets rate info comes from the fw. Only
+                * set rate/antenna during connection establishment.
+                */
+               if (sta && (!ieee80211_is_data(hdr->frame_control) ||
+                           mvmsta->sta_state < IEEE80211_STA_AUTHORIZED)) {
                        flags |= IWL_TX_FLAGS_CMD_RATE;
-                       rate_n_flags = iwl_mvm_get_tx_rate(mvm, info, sta);
+                       rate_n_flags =
+                               iwl_mvm_get_tx_rate_n_flags(mvm, info, sta,
+                                                           hdr->frame_control);
                }
 
                if (mvm->trans->cfg->device_family >=
                iwl_mvm_tx_airtime(mvm, mvmsta,
                                   le16_to_cpu(tx_resp->wireless_media_time));
 
+               if ((status & TX_STATUS_MSK) != TX_STATUS_SUCCESS &&
+                   mvmsta->sta_state < IEEE80211_STA_AUTHORIZED)
+                       iwl_mvm_toggle_tx_ant(mvm, &mvmsta->tx_ant);
+
                if (sta->wme && tid != IWL_MGMT_TID) {
                        struct iwl_mvm_tid_data *tid_data =
                                &mvmsta->tid_data[tid];