rt2x00: Move bssidx calculation into its own function

index a15eec6..bacb533 100644
@@ -1240,6 +1240,12 @@ static inline void rt2x00debug_dump_frame(struct rt2x00_dev *rt2x00dev,
 #endif /* CONFIG_RT2X00_LIB_DEBUGFS */
 
 /*
+ * Utility functions.
+ */
+u32 rt2x00lib_get_bssidx(struct rt2x00_dev *rt2x00dev,
+			 struct ieee80211_vif *vif);
+
+/*
  * Interrupt context handlers.
  */
 void rt2x00lib_beacondone(struct rt2x00_dev *rt2x00dev);

index 8374eb3..4db0ebe 100644
@@ -33,6 +33,22 @@
 #include "rt2x00lib.h"
 
 /*
+ * Utility functions.
+ */
+u32 rt2x00lib_get_bssidx(struct rt2x00_dev *rt2x00dev,
+			 struct ieee80211_vif *vif)
+{
+	/*
+	 * When in STA mode, bssidx is always 0 otherwise local_address[5]
+	 * contains the bss number, see BSS_ID_MASK comments for details.
+	 */
+	if (rt2x00dev->intf_sta_count)
+		return 0;
+	return vif->addr[5] & (rt2x00dev->ops->max_ap_intf - 1);
+}
+EXPORT_SYMBOL_GPL(rt2x00lib_get_bssidx);
+
+/*
  * Radio control handlers.
  */
 int rt2x00lib_enable_radio(struct rt2x00_dev *rt2x00dev)

index a9123d4..4cf1848 100644
@@ -504,15 +504,7 @@ int rt2x00mac_set_key(struct ieee80211_hw *hw, enum set_key_cmd cmd,
 
 	memset(&crypto, 0, sizeof(crypto));
 
-	/*
-	 * When in STA mode, bssidx is always 0 otherwise local_address[5]
-	 * contains the bss number, see BSS_ID_MASK comments for details.
-	 */
-	if (rt2x00dev->intf_sta_count)
-		crypto.bssidx = 0;
-	else
-		crypto.bssidx = vif->addr[5] & (rt2x00dev->ops->max_ap_intf - 1);
-
+	crypto.bssidx = rt2x00lib_get_bssidx(rt2x00dev, vif);
 	crypto.cipher = rt2x00crypto_key_to_cipher(key);
 	if (crypto.cipher == CIPHER_NONE)
 		return -EOPNOTSUPP;

index bacb533..2093a55 100644
@@ -512,6 +512,19 @@ struct rt2x00intf_conf {
 };
 
 /*
+ * Private structure for storing STA details
+ * wcid: Wireless Client ID
+ */
+struct rt2x00_sta {
+	int wcid;
+};
+
+static inline struct rt2x00_sta* sta_to_rt2x00_sta(struct ieee80211_sta *sta)
+{
+	return (struct rt2x00_sta *)sta->drv_priv;
+}
+
+/*
  * rt2x00lib callback functions.
  */
 struct rt2x00lib_ops {

@@ -621,6 +634,11 @@ struct rt2x00lib_ops {
 	void (*config) (struct rt2x00_dev *rt2x00dev,
 			struct rt2x00lib_conf *libconf,
 			const unsigned int changed_flags);
+	int (*sta_add) (struct rt2x00_dev *rt2x00dev,
+			struct ieee80211_vif *vif,
+			struct ieee80211_sta *sta);
+	int (*sta_remove) (struct rt2x00_dev *rt2x00dev,
+			   int wcid);
 };
 
 /*

@@ -1281,6 +1299,10 @@ int rt2x00mac_set_key(struct ieee80211_hw *hw, enum set_key_cmd cmd,
 #else
 #define rt2x00mac_set_key	NULL
 #endif /* CONFIG_RT2X00_LIB_CRYPTO */
+int rt2x00mac_sta_add(struct ieee80211_hw *hw, struct ieee80211_vif *vif,
+		      struct ieee80211_sta *sta);
+int rt2x00mac_sta_remove(struct ieee80211_hw *hw, struct ieee80211_vif *vif,
+			 struct ieee80211_sta *sta);
 void rt2x00mac_sw_scan_start(struct ieee80211_hw *hw);
 void rt2x00mac_sw_scan_complete(struct ieee80211_hw *hw);
 int rt2x00mac_get_stats(struct ieee80211_hw *hw,

index 4db0ebe..7e4e062 100644
@@ -932,6 +932,11 @@ static int rt2x00lib_probe_hw(struct rt2x00_dev *rt2x00dev)
 		rt2x00dev->hw->extra_tx_headroom += RT2X00_ALIGN_SIZE;
 
 	/*
+	 * Tell mac80211 about the size of our private STA structure.
+	 */
+	rt2x00dev->hw->sta_data_size = sizeof(struct rt2x00_sta);
+
+	/*
 	 * Allocate tx status FIFO for driver use.
 	 */
 	if (test_bit(REQUIRE_TXSTATUS_FIFO, &rt2x00dev->cap_flags)) {

index 4cf1848..3cd70ff 100644
@@ -552,6 +552,39 @@ int rt2x00mac_set_key(struct ieee80211_hw *hw, enum set_key_cmd cmd,
 EXPORT_SYMBOL_GPL(rt2x00mac_set_key);
 #endif /* CONFIG_RT2X00_LIB_CRYPTO */
 
+int rt2x00mac_sta_add(struct ieee80211_hw *hw, struct ieee80211_vif *vif,
+		      struct ieee80211_sta *sta)
+{
+	struct rt2x00_dev *rt2x00dev = hw->priv;
+	struct rt2x00_sta *sta_priv = sta_to_rt2x00_sta(sta);
+
+	/*
+	 * If there's no space left in the device table store
+	 * -1 as wcid but tell mac80211 everything went ok.
+	 */
+	if (rt2x00dev->ops->lib->sta_add(rt2x00dev, vif, sta))
+		sta_priv->wcid = -1;
+
+	return 0;
+}
+EXPORT_SYMBOL_GPL(rt2x00mac_sta_add);
+
+int rt2x00mac_sta_remove(struct ieee80211_hw *hw, struct ieee80211_vif *vif,
+			 struct ieee80211_sta *sta)
+{
+	struct rt2x00_dev *rt2x00dev = hw->priv;
+	struct rt2x00_sta *sta_priv = sta_to_rt2x00_sta(sta);
+
+	/*
+	 * If we never sent the STA to the device no need to clean it up.
+	 */
+	if (sta_priv->wcid < 0)
+		return 0;
+
+	return rt2x00dev->ops->lib->sta_remove(rt2x00dev, sta_priv->wcid);
+}
+EXPORT_SYMBOL_GPL(rt2x00mac_sta_remove);
+
 void rt2x00mac_sw_scan_start(struct ieee80211_hw *hw)
 {
 	struct rt2x00_dev *rt2x00dev = hw->priv;

index 2093a55..c360ab3 100644
@@ -478,6 +478,8 @@ struct rt2x00lib_crypto {
 	u8 key[16];
 	u8 tx_mic[8];
 	u8 rx_mic[8];
+
+	int wcid;
 };
 
 /*

index 3cd70ff..4d82d53 100644
@@ -494,6 +494,7 @@ int rt2x00mac_set_key(struct ieee80211_hw *hw, enum set_key_cmd cmd,
 	struct rt2x00lib_crypto crypto;
 	static const u8 bcast_addr[ETH_ALEN] =
 		{ 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, };
+	struct rt2x00_sta *sta_priv = NULL;
 
 	if (!test_bit(DEVICE_STATE_PRESENT, &rt2x00dev->flags))
 		return 0;

@@ -511,9 +512,11 @@ int rt2x00mac_set_key(struct ieee80211_hw *hw, enum set_key_cmd cmd,
 
 	crypto.cmd = cmd;
 
-	if (sta)
+	if (sta) {
 		crypto.address = sta->addr;
-	else
+		sta_priv = sta_to_rt2x00_sta(sta);
+		crypto.wcid = sta_priv->wcid;
+	} else
 		crypto.address = bcast_addr;
 
 	if (crypto.cipher == CIPHER_TKIP)

index c44adb1..ba0d7e6 100644
@@ -288,11 +288,16 @@ static void rt2x00queue_create_tx_descriptor_ht(struct rt2x00_dev *rt2x00dev,
 	struct ieee80211_tx_info *tx_info = IEEE80211_SKB_CB(skb);
 	struct ieee80211_tx_rate *txrate = &tx_info->control.rates[0];
 	struct ieee80211_hdr *hdr = (struct ieee80211_hdr *)skb->data;
+	struct rt2x00_sta *sta_priv = NULL;
 
-	if (tx_info->control.sta)
+	if (tx_info->control.sta) {
 		txdesc->u.ht.mpdu_density =
 		    tx_info->control.sta->ht_cap.ampdu_density;
 
+		sta_priv = sta_to_rt2x00_sta(tx_info->control.sta);
+		txdesc->u.ht.wcid = sta_priv->wcid;
+	}
+
 	txdesc->u.ht.ba_size = 7;	/* FIXME: What value is needed? */
 
 	/*

index f22aa86..349008d 100644
@@ -321,6 +321,7 @@ struct txentry_desc {
 			u8 ba_size;
 			u8 mpdu_density;
 			enum txop txop;
+			int wcid;
 		} ht;
 	} u;
 

index 348fa23..7edd487 100644
@@ -493,7 +493,7 @@ void rt2800_write_tx_data(struct queue_entry *entry,
 	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);

@@ -910,11 +910,51 @@ static void rt2800_init_led(struct rt2x00_dev *rt2x00dev,
 /*
  * 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, &reg);
+	rt2x00_set_field32(&reg, MAC_WCID_ATTRIBUTE_BSS_IDX, (bssidx & 0x7));
+	rt2x00_set_field32(&reg, 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;

@@ -934,14 +974,16 @@ static void rt2800_config_wcid_attr(struct rt2x00_dev *rt2x00dev,
 				   (crypto->cipher & 0x7));
 		rt2x00_set_field32(&reg, MAC_WCID_ATTRIBUTE_CIPHER_EXT,
 				   (crypto->cipher & 0x8) >> 3);
-		rt2x00_set_field32(&reg, MAC_WCID_ATTRIBUTE_BSS_IDX,
-				   (crypto->bssidx & 0x7));
-		rt2x00_set_field32(&reg, MAC_WCID_ATTRIBUTE_BSS_IDX_EXT,
-				   (crypto->bssidx & 0x8) >> 3);
 		rt2x00_set_field32(&reg, 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, &reg);
+		rt2x00_set_field32(&reg, MAC_WCID_ATTRIBUTE_KEYTAB, 0);
+		rt2x00_set_field32(&reg, MAC_WCID_ATTRIBUTE_CIPHER, 0);
+		rt2x00_set_field32(&reg, MAC_WCID_ATTRIBUTE_CIPHER_EXT, 0);
+		rt2x00_set_field32(&reg, MAC_WCID_ATTRIBUTE_RX_WIUDF, 0);
+		rt2800_register_write(rt2x00dev, offset, reg);
 	}
 
 	offset = MAC_IVEIV_ENTRY(key->hw_key_idx);

@@ -954,14 +996,6 @@ static void rt2800_config_wcid_attr(struct rt2x00_dev *rt2x00dev,
 	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,

@@ -1008,20 +1042,24 @@ 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).

@@ -1031,11 +1069,17 @@ static inline int rt2800_find_pairwise_keyslot(struct rt2x00_dev *rt2x00dev)
 	 * first 222 entries.
 	 */
 	for (idx = 33; idx <= 222; idx++) {
-		offset = MAC_WCID_ATTR_ENTRY(idx);
-		rt2800_register_read(rt2x00dev, offset, &reg);
-		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;
 }
 

@@ -1045,13 +1089,15 @@ int rt2800_config_pairwise_key(struct rt2x00_dev *rt2x00dev,
 {
 	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));

@@ -1068,12 +1114,59 @@ int rt2800_config_pairwise_key(struct rt2x00_dev *rt2x00dev,
 	/*
 	 * 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)
 {

index bef071c..7a2511f 100644
@@ -166,6 +166,9 @@ int rt2800_config_shared_key(struct rt2x00_dev *rt2x00dev,
 int rt2800_config_pairwise_key(struct rt2x00_dev *rt2x00dev,
 			       struct rt2x00lib_crypto *crypto,
 			       struct ieee80211_key_conf *key);
+int rt2800_sta_add(struct rt2x00_dev *rt2x00dev, struct ieee80211_vif *vif,
+		   struct ieee80211_sta *sta);
+int rt2800_sta_remove(struct rt2x00_dev *rt2x00dev, int wcid);
 void rt2800_config_filter(struct rt2x00_dev *rt2x00dev,
 			  const unsigned int filter_flags);
 void rt2800_config_intf(struct rt2x00_dev *rt2x00dev, struct rt2x00_intf *intf,

index e11a185..e0493fb 100644
@@ -1015,6 +1015,8 @@ static const struct ieee80211_ops rt2800pci_mac80211_ops = {
 	.get_stats		= rt2x00mac_get_stats,
 	.get_tkip_seq		= rt2800_get_tkip_seq,
 	.set_rts_threshold	= rt2800_set_rts_threshold,
+	.sta_add		= rt2x00mac_sta_add,
+	.sta_remove		= rt2x00mac_sta_remove,
 	.bss_info_changed	= rt2x00mac_bss_info_changed,
 	.conf_tx		= rt2800_conf_tx,
 	.get_tsf		= rt2800_get_tsf,

@@ -1076,6 +1078,8 @@ static const struct rt2x00lib_ops rt2800pci_rt2x00_ops = {
 	.config_erp		= rt2800_config_erp,
 	.config_ant		= rt2800_config_ant,
 	.config			= rt2800_config,
+	.sta_add		= rt2800_sta_add,
+	.sta_remove		= rt2800_sta_remove,
 };
 
 static const struct data_queue_desc rt2800pci_queue_rx = {

index 1e947c0..0015526 100644
@@ -760,6 +760,8 @@ static const struct ieee80211_ops rt2800usb_mac80211_ops = {
 	.get_stats		= rt2x00mac_get_stats,
 	.get_tkip_seq		= rt2800_get_tkip_seq,
 	.set_rts_threshold	= rt2800_set_rts_threshold,
+	.sta_add		= rt2x00mac_sta_add,
+	.sta_remove		= rt2x00mac_sta_remove,
 	.bss_info_changed	= rt2x00mac_bss_info_changed,
 	.conf_tx		= rt2800_conf_tx,
 	.get_tsf		= rt2800_get_tsf,

@@ -817,6 +819,8 @@ static const struct rt2x00lib_ops rt2800usb_rt2x00_ops = {
 	.config_erp		= rt2800_config_erp,
 	.config_ant		= rt2800_config_ant,
 	.config			= rt2800_config,
+	.sta_add		= rt2800_sta_add,
+	.sta_remove		= rt2800_sta_remove,
 };
 
 static const struct data_queue_desc rt2800usb_queue_rx = {

index 7edd487..2f4d51a 100644
@@ -4489,8 +4489,19 @@ int rt2800_ampdu_action(struct ieee80211_hw *hw, struct ieee80211_vif *vif,
 			struct ieee80211_sta *sta, u16 tid, u16 *ssn,
 			u8 buf_size)
 {
+	struct rt2x00_sta *sta_priv = (struct rt2x00_sta *)sta->drv_priv;
 	int ret = 0;
 
+	/*
+	 * Don't allow aggregation for stations the hardware isn't aware
+	 * of because tx status reports for frames to an unknown station
+	 * always contain wcid=255 and thus we can't distinguish between
+	 * multiple stations which leads to unwanted situations when the
+	 * hw reorders frames due to aggregation.
+	 */
+	if (sta_priv->wcid < 0)
+		return 1;
+
 	switch (action) {
 	case IEEE80211_AMPDU_RX_START:
 	case IEEE80211_AMPDU_RX_STOP:

index 2f4d51a..3d9f355 100644
@@ -2846,11 +2846,8 @@ static int rt2800_init_registers(struct rt2x00_dev *rt2x00dev)
 					 SHARED_KEY_MODE_ENTRY(i), 0);
 
 	for (i = 0; i < 256; i++) {
-		static const u32 wcid[2] = { 0xffffffff, 0x00ffffff };
-		rt2800_register_multiwrite(rt2x00dev, MAC_WCID_ENTRY(i),
-					      wcid, sizeof(wcid));
-
-		rt2800_register_write(rt2x00dev, MAC_WCID_ATTR_ENTRY(i), 0);
+		rt2800_config_wcid(rt2x00dev, NULL, i);
+		rt2800_delete_wcid_attr(rt2x00dev, i);
 		rt2800_register_write(rt2x00dev, MAC_IVEIV_ENTRY(i), 0);
 	}