From 021a6fe49954768221f025d644851579b1aa6c6c Mon Sep 17 00:00:00 2001
From: Jouni Malinen <j@w1.fi>
Date: Sat, 6 Nov 2010 18:10:34 +0200
Subject: [PATCH] wlantest: Maintain a copy of WPA/RSN IE from (Re)AssocReq

---
 wlantest/rx_mgmt.c  | 22 ++++++++++++++++++++++
 wlantest/sta.c      | 24 ++++++++++++++++++++++++
 wlantest/wlantest.h |  3 +++
 3 files changed, 49 insertions(+)

diff --git a/wlantest/rx_mgmt.c b/wlantest/rx_mgmt.c
index 843c2b533..2100a8ad7 100644
--- a/wlantest/rx_mgmt.c
+++ b/wlantest/rx_mgmt.c
@@ -192,6 +192,7 @@ static void rx_mgmt_assoc_req(struct wlantest *wt, const u8 *data, size_t len)
 	const struct ieee80211_mgmt *mgmt;
 	struct wlantest_bss *bss;
 	struct wlantest_sta *sta;
+	struct ieee802_11_elems elems;
 
 	mgmt = (const struct ieee80211_mgmt *) data;
 	bss = bss_get(wt, mgmt->bssid);
@@ -212,6 +213,16 @@ static void rx_mgmt_assoc_req(struct wlantest *wt, const u8 *data, size_t len)
 		   MAC2STR(mgmt->sa), MAC2STR(mgmt->da),
 		   le_to_host16(mgmt->u.assoc_req.capab_info),
 		   le_to_host16(mgmt->u.assoc_req.listen_interval));
+
+	if (ieee802_11_parse_elems(mgmt->u.assoc_req.variable,
+				   len - (mgmt->u.assoc_req.variable - data),
+				   &elems, 0) == ParseFailed) {
+		wpa_printf(MSG_INFO, "Invalid IEs in Association Request "
+			   "frame from " MACSTR, MAC2STR(mgmt->sa));
+		return;
+	}
+
+	sta_update_assoc(sta, &elems);
 }
 
 
@@ -275,6 +286,7 @@ static void rx_mgmt_reassoc_req(struct wlantest *wt, const u8 *data,
 	const struct ieee80211_mgmt *mgmt;
 	struct wlantest_bss *bss;
 	struct wlantest_sta *sta;
+	struct ieee802_11_elems elems;
 
 	mgmt = (const struct ieee80211_mgmt *) data;
 	bss = bss_get(wt, mgmt->bssid);
@@ -296,6 +308,16 @@ static void rx_mgmt_reassoc_req(struct wlantest *wt, const u8 *data,
 		   le_to_host16(mgmt->u.reassoc_req.capab_info),
 		   le_to_host16(mgmt->u.reassoc_req.listen_interval),
 		   MAC2STR(mgmt->u.reassoc_req.current_ap));
+
+	if (ieee802_11_parse_elems(mgmt->u.reassoc_req.variable,
+				   len - (mgmt->u.reassoc_req.variable - data),
+				   &elems, 0) == ParseFailed) {
+		wpa_printf(MSG_INFO, "Invalid IEs in Reassociation Request "
+			   "frame from " MACSTR, MAC2STR(mgmt->sa));
+		return;
+	}
+
+	sta_update_assoc(sta, &elems);
 }
 
 
diff --git a/wlantest/sta.c b/wlantest/sta.c
index fc84ee1d0..447bd7a0d 100644
--- a/wlantest/sta.c
+++ b/wlantest/sta.c
@@ -15,6 +15,7 @@
 #include "utils/includes.h"
 
 #include "utils/common.h"
+#include "common/ieee802_11_common.h"
 #include "wlantest.h"
 
 
@@ -46,3 +47,26 @@ void sta_deinit(struct wlantest_sta *sta)
 	dl_list_del(&sta->list);
 	os_free(sta);
 }
+
+
+void sta_update_assoc(struct wlantest_sta *sta, struct ieee802_11_elems *elems)
+{
+	if (elems->wpa_ie && elems->rsn_ie) {
+		wpa_printf(MSG_INFO, "Both WPA IE and RSN IE included in "
+			   "Association Request frame from " MACSTR,
+			   MAC2STR(sta->addr));
+	}
+
+	if (elems->rsn_ie) {
+		wpa_hexdump(MSG_DEBUG, "RSN IE", elems->rsn_ie - 2,
+			    elems->rsn_ie_len + 2);
+		os_memcpy(sta->rsnie, elems->rsn_ie - 2,
+			  elems->rsn_ie_len + 2);
+	} else if (elems->wpa_ie) {
+		wpa_hexdump(MSG_DEBUG, "WPA IE", elems->wpa_ie - 2,
+			    elems->wpa_ie_len + 2);
+		os_memcpy(sta->rsnie, elems->wpa_ie - 2,
+			  elems->wpa_ie_len + 2);
+	} else
+		sta->rsnie[0] = 0;
+}
diff --git a/wlantest/wlantest.h b/wlantest/wlantest.h
index 909e37195..10ca61c99 100644
--- a/wlantest/wlantest.h
+++ b/wlantest/wlantest.h
@@ -29,6 +29,7 @@ struct wlantest_sta {
 		STATE3 /* associated */
 	} state;
 	u16 aid;
+	u8 rsnie[257]; /* WPA/RSN IE */
 };
 
 struct wlantest_bss {
@@ -69,5 +70,7 @@ void bss_update(struct wlantest_bss *bss, struct ieee802_11_elems *elems);
 
 struct wlantest_sta * sta_get(struct wlantest_bss *bss, const u8 *addr);
 void sta_deinit(struct wlantest_sta *sta);
+void sta_update_assoc(struct wlantest_sta *sta,
+		      struct ieee802_11_elems *elems);
 
 #endif /* WLANTEST_H */