From 077232f603dbbef3ac06b33cc793d346be8c5af8 Mon Sep 17 00:00:00 2001
From: Beni Lev <beni.lev@intel.com>
Date: Sun, 26 Aug 2018 20:45:04 +0300
Subject: [PATCH] OCE: Add OCE capability attribute only when associating to an
 OCE AP

Signed-off-by: Beni Lev <beni.lev@intel.com>
---
 wpa_supplicant/mbo.c              | 18 ++++++++++++++++--
 wpa_supplicant/sme.c              | 10 ++++++++--
 wpa_supplicant/wpa_supplicant.c   | 10 ++++++++--
 wpa_supplicant/wpa_supplicant_i.h |  4 +++-
 4 files changed, 35 insertions(+), 7 deletions(-)

diff --git a/wpa_supplicant/mbo.c b/wpa_supplicant/mbo.c
index 129d20545..5adf61e58 100644
--- a/wpa_supplicant/mbo.c
+++ b/wpa_supplicant/mbo.c
@@ -38,6 +38,19 @@ static int wpas_mbo_validate_non_pref_chan(u8 oper_class, u8 chan, u8 reason)
 }
 
 
+const u8 * mbo_attr_from_mbo_ie(const u8 *mbo_ie, enum mbo_attr_id attr)
+{
+	const u8 *mbo;
+	u8 ie_len = mbo_ie[1];
+
+	if (ie_len < MBO_IE_HEADER - 2)
+		return NULL;
+	mbo = mbo_ie + MBO_IE_HEADER;
+
+	return get_ie(mbo, 2 + ie_len - MBO_IE_HEADER, attr);
+}
+
+
 const u8 * wpas_mbo_get_bss_attr(struct wpa_bss *bss, enum mbo_attr_id attr)
 {
 	const u8 *mbo, *end;
@@ -149,7 +162,8 @@ static void wpas_mbo_non_pref_chan_attrs(struct wpa_supplicant *wpa_s,
 }
 
 
-int wpas_mbo_ie(struct wpa_supplicant *wpa_s, u8 *buf, size_t len)
+int wpas_mbo_ie(struct wpa_supplicant *wpa_s, u8 *buf, size_t len,
+		int add_oce_capa)
 {
 	struct wpabuf *mbo;
 	int res;
@@ -175,7 +189,7 @@ int wpas_mbo_ie(struct wpa_supplicant *wpa_s, u8 *buf, size_t len)
 	wpabuf_put_u8(mbo, wpa_s->conf->mbo_cell_capa);
 
 	/* Add OCE capability indication attribute if OCE is enabled */
-	if (wpa_s->enable_oce & OCE_STA) {
+	if ((wpa_s->enable_oce & OCE_STA) && add_oce_capa) {
 		wpabuf_put_u8(mbo, OCE_ATTR_ID_CAPA_IND);
 		wpabuf_put_u8(mbo, 1);
 		wpabuf_put_u8(mbo, OCE_RELEASE);
diff --git a/wpa_supplicant/sme.c b/wpa_supplicant/sme.c
index b68506474..708a347a8 100644
--- a/wpa_supplicant/sme.c
+++ b/wpa_supplicant/sme.c
@@ -240,6 +240,9 @@ static void sme_send_authentication(struct wpa_supplicant *wpa_s,
 	u8 ext_capab[18];
 	int ext_capab_len;
 	int skip_auth;
+#ifdef CONFIG_MBO
+	const u8 *mbo_ie;
+#endif /* CONFIG_MBO */
 
 	if (bss == NULL) {
 		wpa_msg(wpa_s, MSG_ERROR, "SME: No scan result available for "
@@ -539,13 +542,16 @@ static void sme_send_authentication(struct wpa_supplicant *wpa_s,
 	}
 
 #ifdef CONFIG_MBO
-	if (wpa_bss_get_vendor_ie(bss, MBO_IE_VENDOR_TYPE)) {
+	mbo_ie = wpa_bss_get_vendor_ie(bss, MBO_IE_VENDOR_TYPE);
+	if (mbo_ie) {
 		int len;
 
 		len = wpas_mbo_ie(wpa_s, wpa_s->sme.assoc_req_ie +
 				  wpa_s->sme.assoc_req_ie_len,
 				  sizeof(wpa_s->sme.assoc_req_ie) -
-				  wpa_s->sme.assoc_req_ie_len);
+				  wpa_s->sme.assoc_req_ie_len,
+				  !!mbo_attr_from_mbo_ie(mbo_ie,
+							 OCE_ATTR_ID_CAPA_IND));
 		if (len >= 0)
 			wpa_s->sme.assoc_req_ie_len += len;
 	}
diff --git a/wpa_supplicant/wpa_supplicant.c b/wpa_supplicant/wpa_supplicant.c
index 8bd27274f..218023832 100644
--- a/wpa_supplicant/wpa_supplicant.c
+++ b/wpa_supplicant/wpa_supplicant.c
@@ -2444,6 +2444,9 @@ static u8 * wpas_populate_assoc_ies(
 	size_t max_wpa_ie_len = 500;
 	size_t wpa_ie_len;
 	int algs = WPA_AUTH_ALG_OPEN;
+#ifdef CONFIG_MBO
+	const u8 *mbo_ie;
+#endif
 #ifdef CONFIG_FILS
 	const u8 *realm, *username, *rrk;
 	size_t realm_len, username_len, rrk_len;
@@ -2699,11 +2702,14 @@ static u8 * wpas_populate_assoc_ies(
 #endif /* CONFIG_FST */
 
 #ifdef CONFIG_MBO
-	if (bss && wpa_bss_get_vendor_ie(bss, MBO_IE_VENDOR_TYPE)) {
+	mbo_ie = bss ? wpa_bss_get_vendor_ie(bss, MBO_IE_VENDOR_TYPE) : NULL;
+	if (mbo_ie) {
 		int len;
 
 		len = wpas_mbo_ie(wpa_s, wpa_ie + wpa_ie_len,
-				  max_wpa_ie_len - wpa_ie_len);
+				  max_wpa_ie_len - wpa_ie_len,
+				  !!mbo_attr_from_mbo_ie(mbo_ie,
+							 OCE_ATTR_ID_CAPA_IND));
 		if (len >= 0)
 			wpa_ie_len += len;
 	}
diff --git a/wpa_supplicant/wpa_supplicant_i.h b/wpa_supplicant/wpa_supplicant_i.h
index c6d6d84d8..9b8d1fa1e 100644
--- a/wpa_supplicant/wpa_supplicant_i.h
+++ b/wpa_supplicant/wpa_supplicant_i.h
@@ -1359,7 +1359,9 @@ void wpas_flush_fils_hlp_req(struct wpa_supplicant *wpa_s);
 
 
 /* MBO functions */
-int wpas_mbo_ie(struct wpa_supplicant *wpa_s, u8 *buf, size_t len);
+int wpas_mbo_ie(struct wpa_supplicant *wpa_s, u8 *buf, size_t len,
+		int add_oce_capa);
+const u8 * mbo_attr_from_mbo_ie(const u8 *mbo_ie, enum mbo_attr_id attr);
 const u8 * wpas_mbo_get_bss_attr(struct wpa_bss *bss, enum mbo_attr_id attr);
 int wpas_mbo_update_non_pref_chan(struct wpa_supplicant *wpa_s,
 				  const char *non_pref_chan);