From 0660f31ba0d0c086c33376a9dc7309e30b7cbe3c Mon Sep 17 00:00:00 2001
From: Adil Saeed Musthafa <quic_adilm@quicinc.com>
Date: Wed, 1 Mar 2023 11:21:38 -0800
Subject: [PATCH] wlantest: Adjust kdk_len according to RSNX capability

This is needed to derive the PTK correct when Secure LTF support is used
and the additional KDK component needs to be taken into account.

Signed-off-by: Adil Saeed Musthafa <quic_adilm@quicinc.com>
---
 wlantest/rx_eapol.c | 11 ++++++++++-
 wlantest/rx_mgmt.c  | 17 +++++++++++++++++
 wlantest/wlantest.h |  4 ++++
 3 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/wlantest/rx_eapol.c b/wlantest/rx_eapol.c
index 67620ca14..5a38123ac 100644
--- a/wlantest/rx_eapol.c
+++ b/wlantest/rx_eapol.c
@@ -179,12 +179,21 @@ static int try_pmk(struct wlantest *wt, struct wlantest_bss *bss,
 	struct wpa_ptk ptk;
 	const u8 *sa, *aa;
 	bool mlo;
+	size_t kdk_len;
 
 	mlo = !is_zero_ether_addr(sta->mld_mac_addr) &&
 		!is_zero_ether_addr(bss->mld_mac_addr);
 	sa = mlo ? sta->mld_mac_addr : sta->addr;
 	aa = mlo ? bss->mld_mac_addr : bss->bssid;
 
+	if (ieee802_11_rsnx_capab_len(bss->rsnxe, bss->rsnxe_len,
+				      WLAN_RSNX_CAPAB_SECURE_LTF) &&
+	    ieee802_11_rsnx_capab_len(sta->rsnxe, sta->rsnxe_len,
+				      WLAN_RSNX_CAPAB_SECURE_LTF))
+		kdk_len = WPA_KDK_MAX_LEN;
+	else
+		kdk_len = 0;
+
 	if (wpa_key_mgmt_ft(sta->key_mgmt)) {
 		u8 ptk_name[WPA_PMK_NAME_LEN];
 		int use_sha384 = wpa_key_mgmt_sha384(sta->key_mgmt);
@@ -217,7 +226,7 @@ static int try_pmk(struct wlantest *wt, struct wlantest_bss *bss,
 				  "Pairwise key expansion",
 				  aa, sa, sta->anonce,
 				  sta->snonce, &ptk, sta->key_mgmt,
-				  sta->pairwise_cipher, NULL, 0, 0) < 0 ||
+				  sta->pairwise_cipher, NULL, 0, kdk_len) < 0 ||
 		   check_mic(sta, ptk.kck, ptk.kck_len, ver, data, len) < 0) {
 		return -1;
 	}
diff --git a/wlantest/rx_mgmt.c b/wlantest/rx_mgmt.c
index 569e2df78..3e210432e 100644
--- a/wlantest/rx_mgmt.c
+++ b/wlantest/rx_mgmt.c
@@ -346,6 +346,13 @@ static void rx_mgmt_beacon(struct wlantest *wt, const u8 *data, size_t len)
 		return;
 	}
 
+	if (elems.rsnxe) {
+		os_memcpy(bss->rsnxe, elems.rsnxe, elems.rsnxe_len);
+		bss->rsnxe_len = elems.rsnxe_len;
+	} else {
+		bss->rsnxe_len = 0;
+	}
+
 	if (!bss->proberesp_seen)
 		bss_update(wt, bss, &elems, 1);
 
@@ -920,6 +927,11 @@ static void rx_mgmt_assoc_req(struct wlantest *wt, const u8 *data, size_t len)
 		return;
 	}
 
+	if (elems.rsnxe) {
+		os_memcpy(sta->rsnxe, elems.rsnxe, elems.rsnxe_len);
+		sta->rsnxe_len = elems.rsnxe_len;
+	}
+
 	sta->assocreq_capab_info = le_to_host16(mgmt->u.assoc_req.capab_info);
 	sta->assocreq_listen_int =
 		le_to_host16(mgmt->u.assoc_req.listen_interval);
@@ -1162,6 +1174,11 @@ static void rx_mgmt_reassoc_req(struct wlantest *wt, const u8 *data,
 		return;
 	}
 
+	if (elems.rsnxe) {
+		os_memcpy(sta->rsnxe, elems.rsnxe, elems.rsnxe_len);
+		sta->rsnxe_len = elems.rsnxe_len;
+	}
+
 	sta->assocreq_capab_info =
 		le_to_host16(mgmt->u.reassoc_req.capab_info);
 	sta->assocreq_listen_int =
diff --git a/wlantest/wlantest.h b/wlantest/wlantest.h
index ed2d5d603..6a85cc19e 100644
--- a/wlantest/wlantest.h
+++ b/wlantest/wlantest.h
@@ -66,6 +66,8 @@ struct wlantest_sta {
 	bool ft_over_ds;
 	u16 aid;
 	u8 rsnie[257]; /* WPA/RSN IE */
+	u8 rsnxe[254]; /* RSNXE data */
+	size_t rsnxe_len;
 	u8 osenie[257]; /* OSEN IE */
 	int proto;
 	int pairwise_cipher;
@@ -151,6 +153,8 @@ struct wlantest_bss {
 	int parse_error_reported;
 	u8 wpaie[257];
 	u8 rsnie[257];
+	u8 rsnxe[254]; /* RSNXE data */
+	size_t rsnxe_len;
 	u8 osenie[257];
 	int proto;
 	int pairwise_cipher;