diff --git a/net/mac80211/ieee80211_ioctl.c b/net/mac80211/ieee80211_ioctl.c
index 54ad07a..7551db3 100644
--- a/net/mac80211/ieee80211_ioctl.c
+++ b/net/mac80211/ieee80211_ioctl.c
@@ -33,8 +33,8 @@
 				    size_t key_len)
 {
 	struct ieee80211_local *local = wdev_priv(dev->ieee80211_ptr);
-	int ret = 0;
-	struct sta_info *sta;
+	int ret;
+	struct sta_info *sta = NULL;
 	struct ieee80211_key *key;
 	struct ieee80211_sub_if_data *sdata;
 
@@ -46,58 +46,64 @@
 		return -EINVAL;
 	}
 
-	if (is_broadcast_ether_addr(sta_addr)) {
-		sta = NULL;
-		key = sdata->keys[idx];
-	} else {
-		set_tx_key = 0;
-		/*
-		 * According to the standard, the key index of a pairwise
-		 * key must be zero. However, some AP are broken when it
-		 * comes to WEP key indices, so we work around this.
-		 */
-		if (idx != 0 && alg != ALG_WEP) {
-			printk(KERN_DEBUG "%s: set_encrypt - non-zero idx for "
-			       "individual key\n", dev->name);
-			return -EINVAL;
-		}
-
-		sta = sta_info_get(local, sta_addr);
-		if (!sta) {
-#ifdef CONFIG_MAC80211_VERBOSE_DEBUG
-			DECLARE_MAC_BUF(mac);
-			printk(KERN_DEBUG "%s: set_encrypt - unknown addr "
-			       "%s\n",
-			       dev->name, print_mac(mac, sta_addr));
-#endif /* CONFIG_MAC80211_VERBOSE_DEBUG */
-
-			return -ENOENT;
-		}
-
-		key = sta->key;
-	}
-
 	if (remove) {
-		ieee80211_key_free(key);
-		key = NULL;
-	} else {
-		/*
-		 * Automatically frees any old key if present.
-		 */
-		key = ieee80211_key_alloc(sdata, sta, alg, idx, key_len, _key);
-		if (!key) {
-			ret = -ENOMEM;
-			goto err_out;
+		if (is_broadcast_ether_addr(sta_addr)) {
+			key = sdata->keys[idx];
+		} else {
+			sta = sta_info_get(local, sta_addr);
+			if (!sta) {
+				ret = -ENOENT;
+				key = NULL;
+				goto err_out;
+			}
+
+			key = sta->key;
 		}
+
+		if (!key)
+			ret = -ENOENT;
+		else
+			ret = 0;
+	} else {
+		key = ieee80211_key_alloc(alg, idx, key_len, _key);
+		if (!key)
+			return -ENOMEM;
+
+		if (!is_broadcast_ether_addr(sta_addr)) {
+			set_tx_key = 0;
+			/*
+			 * According to the standard, the key index of a
+			 * pairwise key must be zero. However, some AP are
+			 * broken when it comes to WEP key indices, so we
+			 * work around this.
+			 */
+			if (idx != 0 && alg != ALG_WEP) {
+				ret = -EINVAL;
+				goto err_out;
+			}
+
+			sta = sta_info_get(local, sta_addr);
+			if (!sta) {
+				ret = -ENOENT;
+				goto err_out;
+			}
+		}
+
+		ieee80211_key_link(key, sdata, sta);
+
+		if (set_tx_key || (!sta && !sdata->default_key && key))
+			ieee80211_set_default_key(sdata, idx);
+
+		/* don't free key later */
+		key = NULL;
+
+		ret = 0;
 	}
 
-	if (set_tx_key || (!sta && !sdata->default_key && key))
-		ieee80211_set_default_key(sdata, idx);
-
-	ret = 0;
  err_out:
 	if (sta)
 		sta_info_put(sta);
+	ieee80211_key_free(key);
 	return ret;
 }