From 0497e4148197fb8fdf92b5c901ce06df1bfa548c Mon Sep 17 00:00:00 2001
From: John Crispin <john@phrozen.org>
Date: Tue, 13 Aug 2019 15:10:46 +0200
Subject: [PATCH] HE: Fix HE Capabilities element size

Set the max value of optional bytes inside the data structure. This
requires us to calculate the actually used size when copying the
HE capabilities and generating the IE.

Signed-off-by: John Crispin <john@phrozen.org>
Signed-off-by: Sven Eckelmann <seckelmann@datto.com>
---
 src/ap/ieee802_11_he.c       | 38 +++++++++++++++++++++++++++++++++++-
 src/common/ieee802_11_defs.h |  2 +-
 2 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/src/ap/ieee802_11_he.c b/src/ap/ieee802_11_he.c
index bb6083e8c..42b05d7ee 100644
--- a/src/ap/ieee802_11_he.c
+++ b/src/ap/ieee802_11_he.c
@@ -44,6 +44,41 @@ static u8 ieee80211_he_ppet_size(u8 ppe_thres_hdr, const u8 *phy_cap_info)
 }
 
 
+static u8 ieee80211_he_mcs_set_size(const u8 *phy_cap_info)
+{
+	u8 sz = 4;
+
+	if (phy_cap_info[HE_PHYCAP_CHANNEL_WIDTH_SET_IDX] &
+	    HE_PHYCAP_CHANNEL_WIDTH_SET_80PLUS80MHZ_IN_5G)
+		sz += 4;
+	if (phy_cap_info[HE_PHYCAP_CHANNEL_WIDTH_SET_IDX] &
+	    HE_PHYCAP_CHANNEL_WIDTH_SET_160MHZ_IN_5G)
+		sz += 4;
+
+	return sz;
+}
+
+
+static int ieee80211_invalid_he_cap_size(const u8 *buf, size_t len)
+{
+	struct ieee80211_he_capabilities *cap;
+	size_t cap_len;
+
+	cap = (struct ieee80211_he_capabilities *) buf;
+	cap_len = sizeof(*cap) - sizeof(cap->optional);
+	if (len < cap_len)
+		return 1;
+
+	cap_len += ieee80211_he_mcs_set_size(cap->he_phy_capab_info);
+	if (len < cap_len)
+		return 1;
+
+	cap_len += ieee80211_he_ppet_size(buf[cap_len], cap->he_phy_capab_info);
+
+	return len != cap_len;
+}
+
+
 u8 * hostapd_eid_he_capab(struct hostapd_data *hapd, u8 *eid,
 			  enum ieee80211_op_mode opmode)
 {
@@ -56,7 +91,7 @@ u8 * hostapd_eid_he_capab(struct hostapd_data *hapd, u8 *eid,
 	if (!mode)
 		return eid;
 
-	ie_size = sizeof(struct ieee80211_he_capabilities);
+	ie_size = sizeof(*cap) - sizeof(cap->optional);
 	ppet_size = ieee80211_he_ppet_size(mode->he_capab[opmode].ppet[0],
 					   mode->he_capab[opmode].phy_cap);
 
@@ -324,6 +359,7 @@ u16 copy_sta_he_capab(struct hostapd_data *hapd, struct sta_info *sta,
 {
 	if (!he_capab || !hapd->iconf->ieee80211ax ||
 	    !check_valid_he_mcs(hapd, he_capab, opmode) ||
+	    ieee80211_invalid_he_cap_size(he_capab, he_capab_len) ||
 	    he_capab_len > sizeof(struct ieee80211_he_capabilities)) {
 		sta->flags &= ~WLAN_STA_HE;
 		os_free(sta->he_capab);
diff --git a/src/common/ieee802_11_defs.h b/src/common/ieee802_11_defs.h
index b0aa913bb..214ba0e0f 100644
--- a/src/common/ieee802_11_defs.h
+++ b/src/common/ieee802_11_defs.h
@@ -2109,7 +2109,7 @@ struct ieee80211_he_capabilities {
 	u8 he_phy_capab_info[11];
 	/* Followed by 4, 8, or 12 octets of Supported HE-MCS And NSS Set field
 	* and optional variable length PPE Thresholds field. */
-	u8 optional[];
+	u8 optional[37];
 } STRUCT_PACKED;
 
 struct ieee80211_he_operation {