#include "iwl-io.h"
 #include "iwl-helpers.h"
 #include "iwl-4965.h"
+#include "iwl-sta.h"
+
+int iwl_get_free_ucode_key_index(struct iwl_priv *priv)
+{
+       int i;
+
+       for (i = 0; i < STA_KEY_MAX_NUM; i++)
+               if (test_and_set_bit(i, &priv->ucode_key_table))
+                       return i;
+
+       return -1;
+}
 
 int iwl_send_static_wepkey_cmd(struct iwl_priv *priv, u8 send_if_empty)
 {
 }
 
 int iwl_remove_default_wep_key(struct iwl_priv *priv,
-                              struct ieee80211_key_conf *key)
+                              struct ieee80211_key_conf *keyconf)
 {
        int ret;
        unsigned long flags;
 
        spin_lock_irqsave(&priv->sta_lock, flags);
+
+       if (!test_and_clear_bit(keyconf->keyidx, &priv->ucode_key_table))
+               IWL_ERROR("index %d not used in uCode key table.\n",
+                         keyconf->keyidx);
+
        priv->default_wep_key--;
-       memset(&priv->wep_keys[key->keyidx], 0, sizeof(priv->wep_keys[0]));
+       memset(&priv->wep_keys[keyconf->keyidx], 0, sizeof(priv->wep_keys[0]));
        ret = iwl_send_static_wepkey_cmd(priv, 1);
        spin_unlock_irqrestore(&priv->sta_lock, flags);
 
        spin_lock_irqsave(&priv->sta_lock, flags);
        priv->default_wep_key++;
 
+       if (test_and_set_bit(keyconf->keyidx, &priv->ucode_key_table))
+               IWL_ERROR("index %d already used in uCode key table.\n",
+                       keyconf->keyidx);
+
        priv->wep_keys[keyconf->keyidx].key_size = keyconf->keylen;
        memcpy(&priv->wep_keys[keyconf->keyidx].key, &keyconf->key,
                                                        keyconf->keylen);
        memcpy(&priv->stations[sta_id].sta.key.key[3],
                                keyconf->key, keyconf->keylen);
 
-       priv->stations[sta_id].sta.key.key_offset = sta_id % 8; /* FIXME */
+       priv->stations[sta_id].sta.key.key_offset =
+                                iwl_get_free_ucode_key_index(priv);
        priv->stations[sta_id].sta.key.key_flags = key_flags;
 
        priv->stations[sta_id].sta.sta.modify_mask = STA_MODIFY_KEY_MASK;
 
        memcpy(priv->stations[sta_id].sta.key.key, keyconf->key,
               keyconf->keylen);
 
-       priv->stations[sta_id].sta.key.key_offset
-                       = (sta_id % STA_KEY_MAX_NUM);/*FIXME*/
+       priv->stations[sta_id].sta.key.key_offset =
+                               iwl_get_free_ucode_key_index(priv);
        priv->stations[sta_id].sta.key.key_flags = key_flags;
        priv->stations[sta_id].sta.sta.modify_mask = STA_MODIFY_KEY_MASK;
        priv->stations[sta_id].sta.mode = STA_CONTROL_MODIFY_MSK;
        priv->key_mapping_key = 0;
 
        spin_lock_irqsave(&priv->sta_lock, flags);
+       if (!test_and_clear_bit(priv->stations[sta_id].sta.key.key_offset,
+               &priv->ucode_key_table))
+               IWL_ERROR("index %d not used in uCode key table.\n",
+                       priv->stations[sta_id].sta.key.key_offset);
        memset(&priv->stations[sta_id].keyinfo, 0, sizeof(struct iwl4965_hw_key));
        memset(&priv->stations[sta_id].sta.key, 0, sizeof(struct iwl4965_keyinfo));
        priv->stations[sta_id].sta.key.key_flags = STA_KEY_FLG_NO_ENC;
        spin_lock_irqsave(&priv->sta_lock, flags);
 
        priv->stations[sta_id].sta.key.key_offset =
-                                       (sta_id % STA_KEY_MAX_NUM);/* FIXME */
+                               iwl_get_free_ucode_key_index(priv);
        priv->stations[sta_id].sta.key.key_flags = key_flags;
        priv->stations[sta_id].sta.key.tkip_rx_tsc_byte2 = (u8) iv32;