diff --git a/src/drivers/driver_nl80211.c b/src/drivers/driver_nl80211.c
index 18e4b8eef..964fbd26f 100644
--- a/src/drivers/driver_nl80211.c
+++ b/src/drivers/driver_nl80211.c
@@ -3053,7 +3053,6 @@ static int wpa_driver_nl80211_set_key(struct i802_bss *bss,
 	struct nl_msg *msg;
 	struct nl_msg *key_msg;
 	int ret;
-	int tdls = 0;
 	const char *ifname = params->ifname;
 	enum wpa_alg alg = params->alg;
 	const u8 *addr = params->addr;
@@ -3064,6 +3063,7 @@ static int wpa_driver_nl80211_set_key(struct i802_bss *bss,
 	const u8 *key = params->key;
 	size_t key_len = params->key_len;
 	int vlan_id = params->vlan_id;
+	enum key_flag key_flag = params->key_flag;
 
 	/* Ignore for P2P Device */
 	if (drv->nlmode == NL80211_IFTYPE_P2P_DEVICE)
@@ -3071,15 +3071,17 @@ static int wpa_driver_nl80211_set_key(struct i802_bss *bss,
 
 	ifindex = if_nametoindex(ifname);
 	wpa_printf(MSG_DEBUG, "%s: ifindex=%d (%s) alg=%d addr=%p key_idx=%d "
-		   "set_tx=%d seq_len=%lu key_len=%lu",
+		   "set_tx=%d seq_len=%lu key_len=%lu key_flag=0x%x",
 		   __func__, ifindex, ifname, alg, addr, key_idx, set_tx,
-		   (unsigned long) seq_len, (unsigned long) key_len);
+		   (unsigned long) seq_len, (unsigned long) key_len, key_flag);
 #ifdef CONFIG_TDLS
 	if (key_idx == -1) {
 		key_idx = 0;
-		tdls = 1;
 	}
 #endif /* CONFIG_TDLS */
+	if (key_flag & KEY_FLAG_PAIRWISE &&
+	    key_flag & (KEY_FLAG_GROUP | KEY_FLAG_DEFAULT))
+		return -EINVAL;
 
 #ifdef CONFIG_DRIVER_NL80211_QCA
 	if (alg == WPA_ALG_PMK &&
@@ -3094,10 +3096,13 @@ static int wpa_driver_nl80211_set_key(struct i802_bss *bss,
 	if (alg == WPA_ALG_PMK &&
 	    (drv->capa.flags & WPA_DRIVER_FLAGS_4WAY_HANDSHAKE_8021X))
 		return nl80211_set_pmk(drv, key, key_len, addr);
+	if (key_flag & KEY_FLAG_PMK)
+		return -EINVAL;
 
+	ret = -ENOBUFS;
 	key_msg = nlmsg_alloc();
 	if (!key_msg)
-		return -ENOBUFS;
+		return ret;
 
 	if (alg == WPA_ALG_NONE) {
 		msg = nl80211_ifindex_msg(drv, ifindex, 0, NL80211_CMD_DEL_KEY);
@@ -3107,8 +3112,10 @@ static int wpa_driver_nl80211_set_key(struct i802_bss *bss,
 		u32 suite;
 
 		suite = wpa_alg_to_cipher_suite(alg, key_len);
-		if (!suite)
+		if (!suite) {
+			ret = -EINVAL;
 			goto fail2;
+		}
 		msg = nl80211_ifindex_msg(drv, ifindex, 0, NL80211_CMD_NEW_KEY);
 		if (!msg)
 			goto fail2;
@@ -3129,22 +3136,24 @@ static int wpa_driver_nl80211_set_key(struct i802_bss *bss,
 		if (nla_put(msg, NL80211_ATTR_MAC, ETH_ALEN, addr))
 			goto fail;
 
-		if (alg != WPA_ALG_WEP && key_idx && !set_tx) {
+		if (key_flag == KEY_FLAG_GROUP_RX) {
 			wpa_printf(MSG_DEBUG, "   RSN IBSS RX GTK");
 			if (nla_put_u32(key_msg, NL80211_KEY_TYPE,
 					NL80211_KEYTYPE_GROUP))
 				goto fail;
+		} else if (key_flag & (KEY_FLAG_GROUP | KEY_FLAG_DEFAULT) ||
+			   !(key_flag & KEY_FLAG_PAIRWISE)) {
+			ret = -EINVAL;
+			goto fail;
+		} else {
+			wpa_printf(MSG_DEBUG, "   pairwise key");
 		}
-	} else if (addr && is_broadcast_ether_addr(addr)) {
-		struct nlattr *types;
-
+	} else if (key_flag & KEY_FLAG_PAIRWISE ||
+		   !(key_flag & KEY_FLAG_GROUP)) {
+		ret = -EINVAL;
+		goto fail;
+	} else {
 		wpa_printf(MSG_DEBUG, "   broadcast key");
-
-		types = nla_nest_start(key_msg, NL80211_KEY_DEFAULT_TYPES);
-		if (!types ||
-		    nla_put_flag(key_msg, NL80211_KEY_DEFAULT_TYPE_MULTICAST))
-			goto fail;
-		nla_nest_end(key_msg, types);
 	}
 	if (nla_put_u8(key_msg, NL80211_KEY_IDX, key_idx) ||
 	    nla_put_nested(msg, NL80211_ATTR_KEY, key_msg))
@@ -3167,18 +3176,18 @@ static int wpa_driver_nl80211_set_key(struct i802_bss *bss,
 			   ret, strerror(-ret));
 
 	/*
-	 * If we failed or don't need to set the default TX key (below),
+	 * If we failed or don't need to set the key as default (below),
 	 * we're done here.
 	 */
-	if (ret || !set_tx || alg == WPA_ALG_NONE || tdls)
-		return ret;
-	if (is_ap_interface(drv->nlmode) && addr &&
-	    !is_broadcast_ether_addr(addr))
+	if (ret || !(key_flag & KEY_FLAG_DEFAULT))
 		return ret;
+	if (!(key_flag & KEY_FLAG_GROUP))
+		return -EINVAL;
 
+	ret = -ENOBUFS;
 	key_msg = nlmsg_alloc();
 	if (!key_msg)
-		return -ENOBUFS;
+		return ret;
 
 	msg = nl80211_ifindex_msg(drv, ifindex, 0, NL80211_CMD_SET_KEY);
 	if (!msg)
@@ -3197,19 +3206,19 @@ static int wpa_driver_nl80211_set_key(struct i802_bss *bss,
 	if (addr && is_broadcast_ether_addr(addr)) {
 		struct nlattr *types;
 
+		wpa_printf(MSG_DEBUG, "   group key");
 		types = nla_nest_start(key_msg, NL80211_KEY_DEFAULT_TYPES);
 		if (!types ||
 		    nla_put_flag(key_msg, NL80211_KEY_DEFAULT_TYPE_MULTICAST))
 			goto fail;
 		nla_nest_end(key_msg, types);
 	} else if (addr) {
-		struct nlattr *types;
-
-		types = nla_nest_start(key_msg, NL80211_KEY_DEFAULT_TYPES);
-		if (!types ||
-		    nla_put_flag(key_msg, NL80211_KEY_DEFAULT_TYPE_UNICAST))
-			goto fail;
-		nla_nest_end(key_msg, types);
+		wpa_printf(MSG_DEBUG,
+			   "nl80211: Default group key can't use a unicast address");
+		ret = -EINVAL;
+		goto fail;
+	} else {
+		wpa_printf(MSG_DEBUG, "   WEP/WPA-NONE key");
 	}
 
 	if (nla_put_nested(msg, NL80211_ATTR_KEY, key_msg))
@@ -3226,8 +3235,6 @@ static int wpa_driver_nl80211_set_key(struct i802_bss *bss,
 	}
 
 	ret = send_and_recv_msgs(drv, msg, NULL, NULL);
-	if (ret == -ENOENT)
-		ret = 0;
 	if (ret)
 		wpa_printf(MSG_DEBUG,
 			   "nl80211: set_key default failed; err=%d %s",
@@ -3240,7 +3247,7 @@ fail:
 fail2:
 	nl80211_nlmsg_clear(key_msg);
 	nlmsg_free(key_msg);
-	return -ENOBUFS;
+	return ret;
 }
 
 
