.tx_last_beacon         = rt2400pci_tx_last_beacon,
        .rfkill_poll            = rt2x00mac_rfkill_poll,
        .flush                  = rt2x00mac_flush,
+       .set_antenna            = rt2x00mac_set_antenna,
+       .get_antenna            = rt2x00mac_get_antenna,
        .get_ringparam          = rt2x00mac_get_ringparam,
 };
 
 
        .tx_last_beacon         = rt2500pci_tx_last_beacon,
        .rfkill_poll            = rt2x00mac_rfkill_poll,
        .flush                  = rt2x00mac_flush,
+       .set_antenna            = rt2x00mac_set_antenna,
+       .get_antenna            = rt2x00mac_get_antenna,
        .get_ringparam          = rt2x00mac_get_ringparam,
 };
 
 
        .conf_tx                = rt2x00mac_conf_tx,
        .rfkill_poll            = rt2x00mac_rfkill_poll,
        .flush                  = rt2x00mac_flush,
+       .set_antenna            = rt2x00mac_set_antenna,
+       .get_antenna            = rt2x00mac_get_antenna,
        .get_ringparam          = rt2x00mac_get_ringparam,
 };
 
 
                      const struct ieee80211_tx_queue_params *params);
 void rt2x00mac_rfkill_poll(struct ieee80211_hw *hw);
 void rt2x00mac_flush(struct ieee80211_hw *hw, bool drop);
+int rt2x00mac_set_antenna(struct ieee80211_hw *hw, u32 tx_ant, u32 rx_ant);
+int rt2x00mac_get_antenna(struct ieee80211_hw *hw, u32 *tx_ant, u32 *rx_ant);
 void rt2x00mac_get_ringparam(struct ieee80211_hw *hw,
                             u32 *tx, u32 *tx_max, u32 *rx, u32 *rx_max);
 
 
        rt2x00dev->ops->lib->config_erp(rt2x00dev, &erp, changed);
 }
 
-static inline
-enum antenna rt2x00lib_config_antenna_check(enum antenna current_ant,
-                                           enum antenna default_ant)
-{
-       if (current_ant != ANTENNA_SW_DIVERSITY)
-               return current_ant;
-       return (default_ant != ANTENNA_SW_DIVERSITY) ? default_ant : ANTENNA_B;
-}
-
 void rt2x00lib_config_antenna(struct rt2x00_dev *rt2x00dev,
                              struct antenna_setup config)
 {
        struct antenna_setup *active = &rt2x00dev->link.ant.active;
 
        /*
-        * Failsafe: Make sure we are not sending the
-        * ANTENNA_SW_DIVERSITY state to the driver.
-        * If that happens, fallback to hardware defaults,
-        * or our own default.
+        * When the caller tries to send the SW diversity,
+        * we must update the ANTENNA_RX_DIVERSITY flag to
+        * enable the antenna diversity in the link tuner.
+        *
+        * Secondly, we must guarentee we never send the
+        * software antenna diversity command to the driver.
         */
-       if (!(ant->flags & ANTENNA_RX_DIVERSITY))
-               config.rx = rt2x00lib_config_antenna_check(config.rx, def->rx);
-       else if (config.rx == ANTENNA_SW_DIVERSITY)
+       if (!(ant->flags & ANTENNA_RX_DIVERSITY)) {
+               if (config.rx == ANTENNA_SW_DIVERSITY) {
+                       ant->flags |= ANTENNA_RX_DIVERSITY;
+
+                       if (def->rx == ANTENNA_SW_DIVERSITY)
+                               config.rx = ANTENNA_B;
+                       else
+                               config.rx = def->rx;
+               }
+       } else if (config.rx == ANTENNA_SW_DIVERSITY)
                config.rx = active->rx;
 
-       if (!(ant->flags & ANTENNA_TX_DIVERSITY))
-               config.tx = rt2x00lib_config_antenna_check(config.tx, def->tx);
-       else if (config.tx == ANTENNA_SW_DIVERSITY)
+       if (!(ant->flags & ANTENNA_TX_DIVERSITY)) {
+               if (config.tx == ANTENNA_SW_DIVERSITY) {
+                       ant->flags |= ANTENNA_TX_DIVERSITY;
+
+                       if (def->tx == ANTENNA_SW_DIVERSITY)
+                               config.tx = ANTENNA_B;
+                       else
+                               config.tx = def->tx;
+               }
+       } else if (config.tx == ANTENNA_SW_DIVERSITY)
                config.tx = active->tx;
 
        /*
 
        /*
         * Determine if software diversity is enabled for
         * either the TX or RX antenna (or both).
-        * Always perform this check since within the link
-        * tuner interval the configuration might have changed.
         */
-       ant->flags &= ~ANTENNA_RX_DIVERSITY;
-       ant->flags &= ~ANTENNA_TX_DIVERSITY;
-
-       if (rt2x00dev->default_ant.rx == ANTENNA_SW_DIVERSITY)
-               ant->flags |= ANTENNA_RX_DIVERSITY;
-       if (rt2x00dev->default_ant.tx == ANTENNA_SW_DIVERSITY)
-               ant->flags |= ANTENNA_TX_DIVERSITY;
-
        if (!(ant->flags & ANTENNA_RX_DIVERSITY) &&
            !(ant->flags & ANTENNA_TX_DIVERSITY)) {
                ant->flags = 0;
 
 }
 EXPORT_SYMBOL_GPL(rt2x00mac_flush);
 
+int rt2x00mac_set_antenna(struct ieee80211_hw *hw, u32 tx_ant, u32 rx_ant)
+{
+       struct rt2x00_dev *rt2x00dev = hw->priv;
+       struct link_ant *ant = &rt2x00dev->link.ant;
+       struct antenna_setup *def = &rt2x00dev->default_ant;
+       struct antenna_setup setup;
+
+       // The antenna value is not supposed to be 0,
+       // or exceed the maximum number of antenna's.
+       if (!tx_ant || (tx_ant & ~3) || !rx_ant || (rx_ant & ~3))
+               return -EINVAL;
+
+       // When the client tried to configure the antenna to or from
+       // diversity mode, we must reset the default antenna as well
+       // as that controls the diversity switch.
+       if (ant->flags & ANTENNA_TX_DIVERSITY && tx_ant != 3)
+               ant->flags &= ~ANTENNA_TX_DIVERSITY;
+       if (ant->flags & ANTENNA_RX_DIVERSITY && rx_ant != 3)
+               ant->flags &= ~ANTENNA_RX_DIVERSITY;
+
+       // If diversity is being enabled, check if we need hardware
+       // or software diversity. In the latter case, reset the value,
+       // and make sure we update the antenna flags to have the
+       // link tuner pick up the diversity tuning.
+       if (tx_ant == 3 && def->tx == ANTENNA_SW_DIVERSITY) {
+               tx_ant = ANTENNA_SW_DIVERSITY;
+               ant->flags |= ANTENNA_TX_DIVERSITY;
+       }
+
+       if (rx_ant == 3 && def->rx == ANTENNA_SW_DIVERSITY) {
+               rx_ant = ANTENNA_SW_DIVERSITY;
+               ant->flags |= ANTENNA_RX_DIVERSITY;
+       }
+
+       setup.tx = tx_ant;
+       setup.rx = rx_ant;
+
+       rt2x00lib_config_antenna(rt2x00dev, setup);
+
+       return 0;
+}
+EXPORT_SYMBOL_GPL(rt2x00mac_set_antenna);
+
+int rt2x00mac_get_antenna(struct ieee80211_hw *hw, u32 *tx_ant, u32 *rx_ant)
+{
+       struct rt2x00_dev *rt2x00dev = hw->priv;
+       struct link_ant *ant = &rt2x00dev->link.ant;
+       struct antenna_setup *active = &rt2x00dev->link.ant.active;
+
+       // When software diversity is active, we must report this to the
+       // client and not the current active antenna state.
+       if (ant->flags & ANTENNA_TX_DIVERSITY)
+               *tx_ant = ANTENNA_HW_DIVERSITY;
+       else
+               *tx_ant = active->tx;
+
+       if (ant->flags & ANTENNA_RX_DIVERSITY)
+               *rx_ant = ANTENNA_HW_DIVERSITY;
+       else
+               *rx_ant = active->rx;
+
+       return 0;
+}
+EXPORT_SYMBOL_GPL(rt2x00mac_get_antenna);
+
 void rt2x00mac_get_ringparam(struct ieee80211_hw *hw,
                             u32 *tx, u32 *tx_max, u32 *rx, u32 *rx_max)
 {
 
        .get_tsf                = rt61pci_get_tsf,
        .rfkill_poll            = rt2x00mac_rfkill_poll,
        .flush                  = rt2x00mac_flush,
+       .set_antenna            = rt2x00mac_set_antenna,
+       .get_antenna            = rt2x00mac_get_antenna,
        .get_ringparam          = rt2x00mac_get_ringparam,
 };
 
 
        .get_tsf                = rt73usb_get_tsf,
        .rfkill_poll            = rt2x00mac_rfkill_poll,
        .flush                  = rt2x00mac_flush,
+       .set_antenna            = rt2x00mac_set_antenna,
+       .get_antenna            = rt2x00mac_get_antenna,
        .get_ringparam          = rt2x00mac_get_ringparam,
 };