diff --git a/hostapd/ctrl_iface.c b/hostapd/ctrl_iface.c
index 8ef57e5aa..5eef75bf7 100644
--- a/hostapd/ctrl_iface.c
+++ b/hostapd/ctrl_iface.c
@@ -2226,11 +2226,12 @@ static int hostapd_ctrl_set_key(struct hostapd_data *hapd, const char *cmd)
 	u8 addr[ETH_ALEN];
 	const char *pos = cmd;
 	enum wpa_alg alg;
+	enum key_flag key_flag;
 	int idx, set_tx;
 	u8 seq[6], key[WPA_TK_MAX_LEN];
 	size_t key_len;
 
-	/* parameters: alg addr idx set_tx seq key */
+	/* parameters: alg addr idx set_tx seq key key_flag */
 
 	alg = atoi(pos);
 	pos = os_strchr(pos, ' ');
@@ -2259,13 +2260,24 @@ static int hostapd_ctrl_set_key(struct hostapd_data *hapd, const char *cmd)
 	if (*pos != ' ')
 		return -1;
 	pos++;
-	key_len = os_strlen(pos) / 2;
+	if (!os_strchr(pos, ' '))
+		return -1;
+	key_len = (os_strchr(pos, ' ') - pos) / 2;
 	if (hexstr2bin(pos, key, key_len) < 0)
 		return -1;
+	pos += 2 * key_len;
+	if (*pos != ' ')
+		return -1;
+
+	pos++;
+	key_flag = atoi(pos);
+	pos = os_strchr(pos, ' ');
+	if (pos)
+		return -1;
 
 	wpa_printf(MSG_INFO, "TESTING: Set key");
 	return hostapd_drv_set_key(hapd->conf->iface, hapd, alg, addr, idx, 0,
-				   set_tx, seq, 6, key, key_len, 0);
+				   set_tx, seq, 6, key, key_len, key_flag);
 }
 
 
diff --git a/tests/hwsim/test_ap_ciphers.py b/tests/hwsim/test_ap_ciphers.py
index c44d11581..ae78e0ccc 100644
--- a/tests/hwsim/test_ap_ciphers.py
+++ b/tests/hwsim/test_ap_ciphers.py
@@ -863,7 +863,12 @@ def test_ap_wpa2_delayed_m1_m3_zero_tk(dev, apdev):
     if "OK" not in hapd.request("RESEND_M3 " + addr):
         raise Exception("RESEND_M3 failed")
 
-    if "OK" not in hapd.request("SET_KEY 3 %s %d %d %s %s" % (addr, 0, 1, 6*"00", 16*"00")):
+    KEY_FLAG_RX = 0x04
+    KEY_FLAG_TX = 0x08
+    KEY_FLAG_PAIRWISE = 0x20
+    KEY_FLAG_RX_TX = KEY_FLAG_RX | KEY_FLAG_TX
+    KEY_FLAG_PAIRWISE_RX_TX = KEY_FLAG_PAIRWISE | KEY_FLAG_RX_TX
+    if "OK" not in hapd.request("SET_KEY 3 %s %d %d %s %s %d" % (addr, 0, 1, 6*"00", 16*"00", KEY_FLAG_PAIRWISE_RX_TX)):
         raise Exception("SET_KEY failed")
     time.sleep(0.1)
     hwsim_utils.test_connectivity(dev[0], hapd, timeout=1, broadcast=False,