diff --git a/src/pae/ieee802_1x_kay.c b/src/pae/ieee802_1x_kay.c
index 74b2c843b..259c8cc71 100644
--- a/src/pae/ieee802_1x_kay.c
+++ b/src/pae/ieee802_1x_kay.c
@@ -72,7 +72,7 @@ static struct mka_alg mka_alg_tbl[] = {
 
 		.cak_trfm = ieee802_1x_cak_128bits_aes_cmac,
 		.ckn_trfm = ieee802_1x_ckn_128bits_aes_cmac,
-		.kek_trfm = ieee802_1x_kek_128bits_aes_cmac,
+		.kek_trfm = ieee802_1x_kek_aes_cmac,
 		.ick_trfm = ieee802_1x_ick_128bits_aes_cmac,
 		.icv_hash = ieee802_1x_icv_128bits_aes_cmac,
 
@@ -1556,7 +1556,7 @@ ieee802_1x_mka_encode_dist_sak_body(
 		os_memcpy(body->sak, &cs, CS_ID_LEN);
 		sak_pos = CS_ID_LEN;
 	}
-	if (aes_wrap(participant->kek.key, 16,
+	if (aes_wrap(participant->kek.key, participant->kek.len,
 		     cipher_suite_tbl[cs_index].sak_len / 8,
 		     sak->key, body->sak + sak_pos)) {
 		wpa_printf(MSG_ERROR, "KaY: AES wrap failed");
@@ -1693,8 +1693,8 @@ ieee802_1x_mka_decode_dist_sak_body(
 		wpa_printf(MSG_ERROR, "KaY-%s: Out of memory", __func__);
 		return -1;
 	}
-	if (aes_unwrap(participant->kek.key, 16, sak_len >> 3, wrap_sak,
-		       unwrap_sak)) {
+	if (aes_unwrap(participant->kek.key, participant->kek.len,
+		       sak_len >> 3, wrap_sak, unwrap_sak)) {
 		wpa_printf(MSG_ERROR, "KaY: AES unwrap failed");
 		os_free(unwrap_sak);
 		return -1;
@@ -3532,9 +3532,11 @@ ieee802_1x_kay_create_mka(struct ieee802_1x_kay *kay,
 	/* to derive KEK from CAK and CKN */
 	participant->kek.len = mka_alg_tbl[kay->mka_algindex].kek_len;
 	if (mka_alg_tbl[kay->mka_algindex].kek_trfm(participant->cak.key,
+						    participant->cak.len,
 						    participant->ckn.name,
 						    participant->ckn.len,
-						    participant->kek.key)) {
+						    participant->kek.key,
+						    participant->kek.len)) {
 		wpa_printf(MSG_ERROR, "KaY: Derived KEK failed");
 		goto fail;
 	}
diff --git a/src/pae/ieee802_1x_kay_i.h b/src/pae/ieee802_1x_kay_i.h
index ea7155ab2..024ffac73 100644
--- a/src/pae/ieee802_1x_kay_i.h
+++ b/src/pae/ieee802_1x_kay_i.h
@@ -74,7 +74,9 @@ struct mka_alg {
 	int (*cak_trfm)(const u8 *msk, const u8 *mac1, const u8 *mac2, u8 *cak);
 	int (*ckn_trfm)(const u8 *msk, const u8 *mac1, const u8 *mac2,
 			const u8 *sid, size_t sid_len, u8 *ckn);
-	int (*kek_trfm)(const u8 *cak, const u8 *ckn, size_t ckn_len, u8 *kek);
+	int (*kek_trfm)(const u8 *cak, size_t cak_bytes,
+			const u8 *ckn, size_t ckn_len,
+			u8 *kek, size_t kek_bytes);
 	int (*ick_trfm)(const u8 *cak, const u8 *ckn, size_t ckn_len, u8 *ick);
 	int (*icv_hash)(const u8 *ick, const u8 *msg, size_t msg_len, u8 *icv);
 
diff --git a/src/pae/ieee802_1x_key.c b/src/pae/ieee802_1x_key.c
index 2d6d03339..d43cb9157 100644
--- a/src/pae/ieee802_1x_key.c
+++ b/src/pae/ieee802_1x_key.c
@@ -130,13 +130,13 @@ int ieee802_1x_ckn_128bits_aes_cmac(const u8 *msk, const u8 *mac1,
 
 
 /**
- * ieee802_1x_kek_128bits_aes_cmac
+ * ieee802_1x_kek_aes_cmac
  *
  * IEEE Std 802.1X-2010, 9.3.3
  * KEK = KDF(Key, Label, Keyid, KEKLength)
  */
-int ieee802_1x_kek_128bits_aes_cmac(const u8 *cak, const u8 *ckn,
-				    size_t ckn_bytes, u8 *kek)
+int ieee802_1x_kek_aes_cmac(const u8 *cak, size_t cak_bytes, const u8 *ckn,
+			    size_t ckn_bytes, u8 *kek, size_t kek_bytes)
 {
 	u8 context[16];
 
@@ -144,8 +144,9 @@ int ieee802_1x_kek_128bits_aes_cmac(const u8 *cak, const u8 *ckn,
 	os_memset(context, 0, sizeof(context));
 	os_memcpy(context, ckn, (ckn_bytes < 16) ? ckn_bytes : 16);
 
-	return aes_kdf(cak, 128, "IEEE8021 KEK", context, sizeof(context) * 8,
-		       128, kek);
+	return aes_kdf(cak, 8 * cak_bytes, "IEEE8021 KEK",
+		       context, sizeof(context) * 8,
+		       8 * kek_bytes, kek);
 }
 
 
diff --git a/src/pae/ieee802_1x_key.h b/src/pae/ieee802_1x_key.h
index 1e464c451..e77a81662 100644
--- a/src/pae/ieee802_1x_key.h
+++ b/src/pae/ieee802_1x_key.h
@@ -14,8 +14,8 @@ int ieee802_1x_cak_128bits_aes_cmac(const u8 *msk, const u8 *mac1,
 int ieee802_1x_ckn_128bits_aes_cmac(const u8 *msk, const u8 *mac1,
 				    const u8 *mac2, const u8 *sid,
 				    size_t sid_bytes, u8 *ckn);
-int ieee802_1x_kek_128bits_aes_cmac(const u8 *cak, const u8 *ckn,
-				    size_t ckn_bytes, u8 *kek);
+int ieee802_1x_kek_aes_cmac(const u8 *cak, size_t cak_bytes, const u8 *ckn,
+			    size_t ckn_bytes, u8 *kek, size_t kek_bytes);
 int ieee802_1x_ick_128bits_aes_cmac(const u8 *cak, const u8 *ckn,
 				    size_t ckn_bytes, u8 *ick);
 int ieee802_1x_icv_128bits_aes_cmac(const u8 *ick, const u8 *msg,