1 files changed, 19 insertions, 8 deletions

diff --git a/net/mac80211/cfg.c b/net/mac80211/cfg.c
index b0c41a0cee79..e7535ffc8e1c 100644
--- a/net/mac80211/cfg.c
+++ b/net/mac80211/cfg.c
@@ -123,6 +123,7 @@ static int ieee80211_add_key(struct wiphy *wiphy, struct net_device *dev,
         struct sta_info *sta = NULL;
         enum ieee80211_key_alg alg;
         int ret;
+        struct ieee80211_key *key;
 
         sdata = IEEE80211_DEV_TO_SUB_IF(dev);
 
@@ -141,16 +142,21 @@ static int ieee80211_add_key(struct wiphy *wiphy, struct net_device *dev,
                 return -EINVAL;
         }
 
+        key = ieee80211_key_alloc(alg, key_idx, params->key_len, params->key);
+        if (!key)
+                return -ENOMEM;
+
         if (mac_addr) {
                 sta = sta_info_get(sdata->local, mac_addr);
-                if (!sta)
+                if (!sta) {
+                        ieee80211_key_free(key);
                         return -ENOENT;
+                }
         }
 
+        ieee80211_key_link(key, sdata, sta);
+
         ret = 0;
-        if (!ieee80211_key_alloc(sdata, sta, alg, key_idx,
-                                 params->key_len, params->key))
-                ret = -ENOMEM;
 
         if (sta)
                 sta_info_put(sta);
@@ -164,6 +170,7 @@ static int ieee80211_del_key(struct wiphy *wiphy, struct net_device *dev,
         struct ieee80211_sub_if_data *sdata;
         struct sta_info *sta;
         int ret;
+        struct ieee80211_key *key;
 
         sdata = IEEE80211_DEV_TO_SUB_IF(dev);
 
@@ -173,9 +180,11 @@ static int ieee80211_del_key(struct wiphy *wiphy, struct net_device *dev,
                         return -ENOENT;
 
                 ret = 0;
-                if (sta->key)
-                        ieee80211_key_free(sta->key);
-                else
+                if (sta->key) {
+                        key = sta->key;
+                        ieee80211_key_free(key);
+                        WARN_ON(sta->key);
+                } else
                         ret = -ENOENT;
 
                 sta_info_put(sta);
@@ -185,7 +194,9 @@ static int ieee80211_del_key(struct wiphy *wiphy, struct net_device *dev,
         if (!sdata->keys[key_idx])
                 return -ENOENT;
 
-        ieee80211_key_free(sdata->keys[key_idx]);
+        key = sdata->keys[key_idx];
+        ieee80211_key_free(key);
+        WARN_ON(sdata->keys[key_idx]);
 
         return 0;
 }