rt2x00_set_field32(&word, TXWI_W1_BW_WIN_SIZE, txdesc->u.ht.ba_size);
        rt2x00_set_field32(&word, TXWI_W1_WIRELESS_CLI_ID,
                           test_bit(ENTRY_TXD_ENCRYPT, &txdesc->flags) ?
-                          txdesc->key_idx : 0xff);
+                          txdesc->key_idx : txdesc->u.ht.wcid);
        rt2x00_set_field32(&word, TXWI_W1_MPDU_TOTAL_BYTE_COUNT,
                           txdesc->length);
        rt2x00_set_field32(&word, TXWI_W1_PACKETID_QUEUE, entry->queue->qid);
 /*
  * Configuration handlers.
  */
-static void rt2800_config_wcid_attr(struct rt2x00_dev *rt2x00dev,
-                                   struct rt2x00lib_crypto *crypto,
-                                   struct ieee80211_key_conf *key)
+static void rt2800_config_wcid(struct rt2x00_dev *rt2x00dev,
+                              const u8 *address,
+                              int wcid)
 {
        struct mac_wcid_entry wcid_entry;
+       u32 offset;
+
+       offset = MAC_WCID_ENTRY(wcid);
+
+       memset(&wcid_entry, 0xff, sizeof(wcid_entry));
+       if (address)
+               memcpy(wcid_entry.mac, address, ETH_ALEN);
+
+       rt2800_register_multiwrite(rt2x00dev, offset,
+                                     &wcid_entry, sizeof(wcid_entry));
+}
+
+static void rt2800_delete_wcid_attr(struct rt2x00_dev *rt2x00dev, int wcid)
+{
+       u32 offset;
+       offset = MAC_WCID_ATTR_ENTRY(wcid);
+       rt2800_register_write(rt2x00dev, offset, 0);
+}
+
+static void rt2800_config_wcid_attr_bssidx(struct rt2x00_dev *rt2x00dev,
+                                          int wcid, u32 bssidx)
+{
+       u32 offset = MAC_WCID_ATTR_ENTRY(wcid);
+       u32 reg;
+
+       /*
+        * The BSS Idx numbers is split in a main value of 3 bits,
+        * and a extended field for adding one additional bit to the value.
+        */
+       rt2800_register_read(rt2x00dev, offset, ®);
+       rt2x00_set_field32(®, MAC_WCID_ATTRIBUTE_BSS_IDX, (bssidx & 0x7));
+       rt2x00_set_field32(®, MAC_WCID_ATTRIBUTE_BSS_IDX_EXT,
+                          (bssidx & 0x8) >> 3);
+       rt2800_register_write(rt2x00dev, offset, reg);
+}
+
+static void rt2800_config_wcid_attr_cipher(struct rt2x00_dev *rt2x00dev,
+                                          struct rt2x00lib_crypto *crypto,
+                                          struct ieee80211_key_conf *key)
+{
        struct mac_iveiv_entry iveiv_entry;
        u32 offset;
        u32 reg;
                                   (crypto->cipher & 0x7));
                rt2x00_set_field32(®, MAC_WCID_ATTRIBUTE_CIPHER_EXT,
                                   (crypto->cipher & 0x8) >> 3);
-               rt2x00_set_field32(®, MAC_WCID_ATTRIBUTE_BSS_IDX,
-                                  (crypto->bssidx & 0x7));
-               rt2x00_set_field32(®, MAC_WCID_ATTRIBUTE_BSS_IDX_EXT,
-                                  (crypto->bssidx & 0x8) >> 3);
                rt2x00_set_field32(®, MAC_WCID_ATTRIBUTE_RX_WIUDF, crypto->cipher);
                rt2800_register_write(rt2x00dev, offset, reg);
        } else {
-               rt2800_register_write(rt2x00dev, offset, 0);
+               /* Delete the cipher without touching the bssidx */
+               rt2800_register_read(rt2x00dev, offset, ®);
+               rt2x00_set_field32(®, MAC_WCID_ATTRIBUTE_KEYTAB, 0);
+               rt2x00_set_field32(®, MAC_WCID_ATTRIBUTE_CIPHER, 0);
+               rt2x00_set_field32(®, MAC_WCID_ATTRIBUTE_CIPHER_EXT, 0);
+               rt2x00_set_field32(®, MAC_WCID_ATTRIBUTE_RX_WIUDF, 0);
+               rt2800_register_write(rt2x00dev, offset, reg);
        }
 
        offset = MAC_IVEIV_ENTRY(key->hw_key_idx);
        iveiv_entry.iv[3] |= key->keyidx << 6;
        rt2800_register_multiwrite(rt2x00dev, offset,
                                      &iveiv_entry, sizeof(iveiv_entry));
-
-       offset = MAC_WCID_ENTRY(key->hw_key_idx);
-
-       memset(&wcid_entry, 0, sizeof(wcid_entry));
-       if (crypto->cmd == SET_KEY)
-               memcpy(wcid_entry.mac, crypto->address, ETH_ALEN);
-       rt2800_register_multiwrite(rt2x00dev, offset,
-                                     &wcid_entry, sizeof(wcid_entry));
 }
 
 int rt2800_config_shared_key(struct rt2x00_dev *rt2x00dev,
        /*
         * Update WCID information
         */
-       rt2800_config_wcid_attr(rt2x00dev, crypto, key);
+       rt2800_config_wcid(rt2x00dev, crypto->address, key->hw_key_idx);
+       rt2800_config_wcid_attr_bssidx(rt2x00dev, key->hw_key_idx,
+                                      crypto->bssidx);
+       rt2800_config_wcid_attr_cipher(rt2x00dev, crypto, key);
 
        return 0;
 }
 EXPORT_SYMBOL_GPL(rt2800_config_shared_key);
 
-static inline int rt2800_find_pairwise_keyslot(struct rt2x00_dev *rt2x00dev)
+static inline int rt2800_find_wcid(struct rt2x00_dev *rt2x00dev)
 {
+       struct mac_wcid_entry wcid_entry;
        int idx;
-       u32 offset, reg;
+       u32 offset;
 
        /*
-        * Search for the first free pairwise key entry and return the
-        * corresponding index.
+        * Search for the first free WCID entry and return the corresponding
+        * index.
         *
         * Make sure the WCID starts _after_ the last possible shared key
         * entry (>32).
         * first 222 entries.
         */
        for (idx = 33; idx <= 222; idx++) {
-               offset = MAC_WCID_ATTR_ENTRY(idx);
-               rt2800_register_read(rt2x00dev, offset, ®);
-               if (!reg)
+               offset = MAC_WCID_ENTRY(idx);
+               rt2800_register_multiread(rt2x00dev, offset, &wcid_entry,
+                                         sizeof(wcid_entry));
+               if (is_broadcast_ether_addr(wcid_entry.mac))
                        return idx;
        }
+
+       /*
+        * Use -1 to indicate that we don't have any more space in the WCID
+        * table.
+        */
        return -1;
 }
 
 {
        struct hw_key_entry key_entry;
        u32 offset;
-       int idx;
 
        if (crypto->cmd == SET_KEY) {
-               idx = rt2800_find_pairwise_keyslot(rt2x00dev);
-               if (idx < 0)
+               /*
+                * Allow key configuration only for STAs that are
+                * known by the hw.
+                */
+               if (crypto->wcid < 0)
                        return -ENOSPC;
-               key->hw_key_idx = idx;
+               key->hw_key_idx = crypto->wcid;
 
                memcpy(key_entry.key, crypto->key,
                       sizeof(key_entry.key));
        /*
         * Update WCID information
         */
-       rt2800_config_wcid_attr(rt2x00dev, crypto, key);
+       rt2800_config_wcid_attr_cipher(rt2x00dev, crypto, key);
 
        return 0;
 }
 EXPORT_SYMBOL_GPL(rt2800_config_pairwise_key);
 
+int rt2800_sta_add(struct rt2x00_dev *rt2x00dev, struct ieee80211_vif *vif,
+                  struct ieee80211_sta *sta)
+{
+       int wcid;
+       struct rt2x00_sta *sta_priv = sta_to_rt2x00_sta(sta);
+
+       /*
+        * Find next free WCID.
+        */
+       wcid = rt2800_find_wcid(rt2x00dev);
+
+       /*
+        * Store selected wcid even if it is invalid so that we can
+        * later decide if the STA is uploaded into the hw.
+        */
+       sta_priv->wcid = wcid;
+
+       /*
+        * No space left in the device, however, we can still communicate
+        * with the STA -> No error.
+        */
+       if (wcid < 0)
+               return 0;
+
+       /*
+        * Clean up WCID attributes and write STA address to the device.
+        */
+       rt2800_delete_wcid_attr(rt2x00dev, wcid);
+       rt2800_config_wcid(rt2x00dev, sta->addr, wcid);
+       rt2800_config_wcid_attr_bssidx(rt2x00dev, wcid,
+                                      rt2x00lib_get_bssidx(rt2x00dev, vif));
+       return 0;
+}
+EXPORT_SYMBOL_GPL(rt2800_sta_add);
+
+int rt2800_sta_remove(struct rt2x00_dev *rt2x00dev, int wcid)
+{
+       /*
+        * Remove WCID entry, no need to clean the attributes as they will
+        * get renewed when the WCID is reused.
+        */
+       rt2800_config_wcid(rt2x00dev, NULL, wcid);
+
+       return 0;
+}
+EXPORT_SYMBOL_GPL(rt2800_sta_remove);
+
 void rt2800_config_filter(struct rt2x00_dev *rt2x00dev,
                          const unsigned int filter_flags)
 {