Add support for SAE-FT

Bug:149449646
Test: Tested on Android-R and verified on sniffer.

Signed-off-by: Mir Ali <mir-khizer.ali@broadcom.com>
Change-Id: I57d5642386dc7cbbcc162f4049721413c78a47ab
diff --git a/src/common/brcm_vendor.h b/src/common/brcm_vendor.h
index b1bf397..0a44a6c 100644
--- a/src/common/brcm_vendor.h
+++ b/src/common/brcm_vendor.h
@@ -27,6 +27,7 @@
  */
 enum brcm_nl80211_vendor_subcmds {
 	BRCM_VENDOR_SUBCMD_UNSPEC		= 0,
+	BRCM_VENDOR_SUBCMD_SET_PMK		= 4,
 	BRCM_VENDOR_SUBCMD_SET_MAC		= 6,
 	BRCM_VENDOR_SCMD_ACS			= 9,
 	BRCM_VENDOR_SCMD_MAX			= 10
@@ -65,6 +66,7 @@
 
 enum brcm_wlan_vendor_attr {
 	BRCM_ATTR_DRIVER_CMD            = 0,
+	BRCM_ATTR_DRIVER_KEY_PMK        = 1,
 	BRCM_ATTR_DRIVER_MAC_ADDR	= 3,
 	BRCM_ATTR_DRIVER_AFTER_LAST     = 5,
 	BRCM_ATTR_DRIVER_MAX            = BRCM_ATTR_DRIVER_AFTER_LAST - 1,
diff --git a/src/drivers/driver_nl80211.c b/src/drivers/driver_nl80211.c
index 4295f0c..5d1b52a 100644
--- a/src/drivers/driver_nl80211.c
+++ b/src/drivers/driver_nl80211.c
@@ -3240,6 +3240,36 @@
 #endif /* CONFIG_DRIVER_NL80211_QCA */
 
 
+#ifdef CONFIG_DRIVER_NL80211_BRCM
+static int key_mgmt_set_key(struct wpa_driver_nl80211_data *drv,
+				  const u8 *key, size_t key_len)
+{
+	struct nl_msg *msg;
+	int ret;
+	struct nlattr *params;
+
+	if (!(msg = nl80211_drv_msg(drv, 0, NL80211_CMD_VENDOR)) ||
+	    nla_put_u32(msg, NL80211_ATTR_VENDOR_ID, OUI_BRCM) ||
+	    nla_put_u32(msg, NL80211_ATTR_VENDOR_SUBCMD,
+				BRCM_VENDOR_SUBCMD_SET_PMK) ||
+          !(params = nla_nest_start(msg, NL80211_ATTR_VENDOR_DATA)) ||
+            nla_put(msg, BRCM_ATTR_DRIVER_KEY_PMK, key_len, key)) {
+                nl80211_nlmsg_clear(msg);
+                nlmsg_free(msg);
+                return -ENOBUFS;
+	}
+	nla_nest_end(msg, params);
+
+	ret = send_and_recv_msgs(drv, msg, NULL, (void *) -1, NULL, NULL);
+	if (ret) {
+		wpa_printf(MSG_DEBUG, "nl80211: Key mgmt set key failed: ret=%d (%s)",
+			ret, strerror(-ret));
+	}
+
+	return ret;
+}
+#endif /* CONFIG_DRIVER_NL80211_BRCM */
+
 static int nl80211_set_pmk(struct wpa_driver_nl80211_data *drv,
 			   const u8 *key, size_t key_len,
 			   const u8 *addr)
@@ -3327,6 +3357,12 @@
 	if (key_flag & KEY_FLAG_PMK) {
 		if (drv->capa.flags & WPA_DRIVER_FLAGS_4WAY_HANDSHAKE_8021X)
 			return nl80211_set_pmk(drv, key, key_len, addr);
+#ifdef CONFIG_DRIVER_NL80211_BRCM
+		if (drv->vendor_set_pmk) {
+			wpa_printf(MSG_INFO, "nl80211: key_mgmt_set_key with key_len %lu", key_len);
+			return key_mgmt_set_key(drv, key, key_len);
+		}
+#endif /* CONFIG_DRIVER_NL80211_BRCM */
 		/* The driver does not have any offload mechanism for PMK, so
 		 * there is no need to configure this key. */
 		return 0;
diff --git a/src/drivers/driver_nl80211.h b/src/drivers/driver_nl80211.h
index 0db7d55..4688924 100644
--- a/src/drivers/driver_nl80211.h
+++ b/src/drivers/driver_nl80211.h
@@ -220,6 +220,9 @@
 	 * (NL80211_CMD_VENDOR). 0 if no pending scan request.
 	 */
 	int last_scan_cmd;
+#ifdef CONFIG_DRIVER_NL80211_BRCM
+	unsigned int vendor_set_pmk:1; /* for legacy set_pmk method before NL80211_CMD_SET_PMK */
+#endif /* CONFIG_DRIVER_NL80211_BRCM */
 };
 
 struct nl_msg;
diff --git a/src/drivers/driver_nl80211_capa.c b/src/drivers/driver_nl80211_capa.c
index 416531f..0e038db 100644
--- a/src/drivers/driver_nl80211_capa.c
+++ b/src/drivers/driver_nl80211_capa.c
@@ -1009,6 +1009,9 @@
 						drv->capa.flags |=
 							WPA_DRIVER_FLAGS_SUPPORT_HW_MODE_ANY;
 						break;
+					case BRCM_VENDOR_SUBCMD_SET_PMK:
+						drv->vendor_set_pmk = 1;
+						break;
 					default:
 						break;
 				}
diff --git a/src/rsn_supp/wpa.c b/src/rsn_supp/wpa.c
index 6c6de80..d028e95 100644
--- a/src/rsn_supp/wpa.c
+++ b/src/rsn_supp/wpa.c
@@ -3824,6 +3824,20 @@
 	pmksa_cache_flush(sm->pmksa, network_ctx, NULL, 0);
 }
 
+#ifdef CONFIG_DRIVER_NL80211_BRCM
+void wpa_sm_install_pmk(struct wpa_sm *sm)
+{
+	/* In case the driver wants to handle re-assocs, pass it down the PMK. */
+	if (wpa_sm_set_key(sm, wpa_cipher_to_alg(sm->pairwise_cipher), NULL, 0, 0, NULL, 0,
+		(u8*)sm->pmk, sm->pmk_len, KEY_FLAG_PMK) < 0) {
+		wpa_hexdump(MSG_DEBUG, "PSK: Install PMK to the driver for driver reassociations",
+			(u8*)sm->pmk, sm->pmk_len);
+		/* No harm if the driver doesn't support. */
+		wpa_msg(sm->ctx->msg_ctx, MSG_DEBUG,
+			"WPA: Failed to set PMK to the driver");
+	}
+}
+#endif /* CONFIG_DRIVER_NL80211_BRCM */
 
 #ifdef CONFIG_WNM
 int wpa_wnmsleep_install_key(struct wpa_sm *sm, u8 subelem_id, u8 *buf)
diff --git a/src/rsn_supp/wpa.h b/src/rsn_supp/wpa.h
index 2142772..5fc7ed6 100644
--- a/src/rsn_supp/wpa.h
+++ b/src/rsn_supp/wpa.h
@@ -198,6 +198,10 @@
 
 int wpa_sm_get_p2p_ip_addr(struct wpa_sm *sm, u8 *buf);
 
+#ifdef CONFIG_DRIVER_NL80211_BRCM
+void wpa_sm_install_pmk(struct wpa_sm *sm);
+#endif /* CONFIG_DRIVER_NL80211_BRCM */
+
 void wpa_sm_set_rx_replay_ctr(struct wpa_sm *sm, const u8 *rx_replay_counter);
 void wpa_sm_set_ptk_kck_kek(struct wpa_sm *sm,
 			    const u8 *ptk_kck, size_t ptk_kck_len,
@@ -410,6 +414,7 @@
 			    const u8 *ric_ies, size_t ric_ies_len);
 int wpa_ft_is_completed(struct wpa_sm *sm);
 void wpa_reset_ft_completed(struct wpa_sm *sm);
+void wpa_set_ft_completed(struct wpa_sm *sm);
 int wpa_ft_validate_reassoc_resp(struct wpa_sm *sm, const u8 *ies,
 				 size_t ies_len, const u8 *src_addr);
 int wpa_ft_start_over_ds(struct wpa_sm *sm, const u8 *target_ap,
@@ -451,6 +456,10 @@
 {
 }
 
+static inline void wpa_set_ft_completed(struct wpa_sm *sm)
+{
+}
+
 static inline int
 wpa_ft_validate_reassoc_resp(struct wpa_sm *sm, const u8 *ies, size_t ies_len,
 			     const u8 *src_addr)
diff --git a/src/rsn_supp/wpa_ft.c b/src/rsn_supp/wpa_ft.c
index bf73376..d7e3685 100644
--- a/src/rsn_supp/wpa_ft.c
+++ b/src/rsn_supp/wpa_ft.c
@@ -718,6 +718,12 @@
 }
 
 
+void wpa_set_ft_completed(struct wpa_sm *sm)
+{
+	if (sm != NULL)
+		sm->ft_completed = 1;
+}
+
 static int wpa_ft_process_gtk_subelem(struct wpa_sm *sm, const u8 *gtk_elem,
 				      size_t gtk_elem_len)
 {
diff --git a/wpa_supplicant/events.c b/wpa_supplicant/events.c
index fd4511d..85c7190 100644
--- a/wpa_supplicant/events.c
+++ b/wpa_supplicant/events.c
@@ -49,6 +49,7 @@
 #include "mesh_mpm.h"
 #include "wmm_ac.h"
 #include "dpp_supplicant.h"
+#include "rsn_supp/wpa_i.h"
 
 
 #define MAX_OWE_TRANSITION_BSS_SELECT_COUNT 5
@@ -2916,6 +2917,12 @@
 		p += len;
 	}
 #endif /* CONFIG_SME */
+#ifdef CONFIG_DRIVER_NL80211_BRCM
+	if ((wpa_s->key_mgmt == WPA_KEY_MGMT_FT_SAE) &&
+		wpa_ft_is_completed(wpa_s->wpa)) {
+		return 0;
+	}
+#endif /* CONFIG_DRIVER_NL80211_BRCM */
 
 	/* Process FT when SME is in the driver */
 	if (!(wpa_s->drv_flags & WPA_DRIVER_FLAGS_SME) &&
@@ -3075,6 +3082,9 @@
 #if defined(CONFIG_FILS) || defined(CONFIG_MBO)
 	struct wpa_bss *bss;
 #endif /* CONFIG_FILS || CONFIG_MBO */
+#ifdef CONFIG_DRIVER_NL80211_BRCM
+	struct wpa_ie_data ie;
+#endif /* CONFIG_DRIVER_NL80211_BRCM */
 
 #ifdef CONFIG_AP
 	if (wpa_s->ap_iface) {
@@ -3092,6 +3102,27 @@
 	eloop_cancel_timeout(wpas_network_reenabled, wpa_s, NULL);
 	wpa_s->own_reconnect_req = 0;
 
+#ifdef CONFIG_DRIVER_NL80211_BRCM
+	if (!(wpa_sm_parse_own_wpa_ie(wpa_s->wpa, &ie) < 0)) {
+		struct wpa_ft_ies parse;
+		/* Check for FT reassociation is done by the driver */
+#ifdef CONFIG_IEEE80211R
+		int use_sha384 = wpa_key_mgmt_sha384(wpa_s->wpa->key_mgmt);
+		if ((wpa_s->key_mgmt == WPA_KEY_MGMT_FT_SAE) && (wpa_s->key_mgmt == ie.key_mgmt)) {
+			if (wpa_ft_parse_ies(data->assoc_info.resp_ies,
+				data->assoc_info.resp_ies_len, &parse, use_sha384) < 0) {
+				wpa_printf(MSG_DEBUG, "Failed to parse FT IEs");
+				return;
+			}
+			if (parse.rsn_pmkid != NULL) {
+				wpa_set_ft_completed(wpa_s->wpa);
+				wpa_dbg(wpa_s, MSG_DEBUG, "Assume FT reassoc completed by the driver");
+			}
+		}
+#endif  /* CONFIG_IEEE80211R */
+	}
+#endif /* CONFIG_DRIVER_NL80211_BRCM */
+
 	ft_completed = wpa_ft_is_completed(wpa_s->wpa);
 	if (data && wpa_supplicant_event_associnfo(wpa_s, data) < 0)
 		return;
@@ -3110,6 +3141,15 @@
 		return;
 	}
 
+#ifdef CONFIG_DRIVER_NL80211_BRCM
+	/* For driver based roaming, insert PSK during the initial association */
+	if (is_zero_ether_addr(wpa_s->bssid) &&
+		wpa_key_mgmt_wpa_psk(wpa_s->key_mgmt)) {
+		/* In case the driver wants to handle re-assocs, pass it down the PMK. */
+		wpa_dbg(wpa_s, MSG_DEBUG, "Pass the PMK to the driver");
+		wpa_sm_install_pmk(wpa_s->wpa);
+	}
+#endif /* CONFIG_DRIVER_NL80211_BRCM */
 	wpa_supplicant_set_state(wpa_s, WPA_ASSOCIATED);
 	if (os_memcmp(bssid, wpa_s->bssid, ETH_ALEN) != 0) {
 		if (os_reltime_initialized(&wpa_s->session_start)) {