Multi AKM(SAE+WPA2PSK) plumb from connect command

Currently the connect command from supplicant facilitates plumbing
of single AKM to the driver. In the current change, multi-AKM plumb
for the combination of SAE and WPA2PSK AKMs is supported.

Bug:210811506
Test:Validated on Pixel phone with Android T.

Signed-off-by: Vinayak Yadawad <vinayak.yadawad@broadcom.com>
Change-Id: Ie4ace52fa54c35151c607419bb709d9aa3f67f54
diff --git a/src/common/ieee802_11_defs.h b/src/common/ieee802_11_defs.h
index 5ee691b..4300ae5 100644
--- a/src/common/ieee802_11_defs.h
+++ b/src/common/ieee802_11_defs.h
@@ -2529,4 +2529,9 @@
 #define WFA_CAPA_QM_DSCP_POLICY BIT(0)
 #define WFA_CAPA_QM_UNSOLIC_DSCP BIT(1)
 
+#ifdef CONFIG_DRIVER_NL80211_BRCM
+#define WPA_KEY_MGMT_CROSS_AKM_ROAM (WPA_KEY_MGMT_SAE | WPA_KEY_MGMT_PSK)
+#define IS_CROSS_AKM_ROAM_KEY_MGMT(key_mgmt) \
+	((key_mgmt & WPA_KEY_MGMT_CROSS_AKM_ROAM) == WPA_KEY_MGMT_CROSS_AKM_ROAM)
+#endif /* CONFIG_DRIVER_NL80211_BRCM */
 #endif /* IEEE802_11_DEFS_H */
diff --git a/src/drivers/driver_nl80211.c b/src/drivers/driver_nl80211.c
index 0783aa0..a2c4842 100644
--- a/src/drivers/driver_nl80211.c
+++ b/src/drivers/driver_nl80211.c
@@ -3243,6 +3243,24 @@
 }
 #endif /* CONFIG_DRIVER_NL80211_BRCM */
 
+#ifdef CONFIG_DRIVER_NL80211_BRCM
+static int wpa_cross_akm_key_mgmt_to_suites(unsigned int key_mgmt_suites, u32 suites[],
+                        int max_suites)
+{
+    int num_suites = 0;
+
+#define __AKM_TO_SUITES_ARRAY(a, b) \
+    if (num_suites < max_suites && \
+        (key_mgmt_suites & (WPA_KEY_MGMT_ ## a))) \
+        suites[num_suites++] = (RSN_AUTH_KEY_MGMT_ ## b)
+    __AKM_TO_SUITES_ARRAY(PSK, PSK_OVER_802_1X);
+    __AKM_TO_SUITES_ARRAY(SAE, SAE);
+#undef __AKM_TO_SUITES_ARRAY
+
+    return num_suites;
+}
+#endif /* CONFIG_DRIVER_NL80211_BRCM */
+
 #ifdef CONFIG_DRIVER_NL80211_QCA
 static int issue_key_mgmt_set_key(struct wpa_driver_nl80211_data *drv,
 				  const u8 *key, size_t key_len)
@@ -6383,6 +6401,22 @@
 			return -1;
 	}
 
+#ifdef CONFIG_DRIVER_NL80211_BRCM
+	if (IS_CROSS_AKM_ROAM_KEY_MGMT(params->key_mgmt_suite)) {
+		int num_suites;
+		u32 suites[NL80211_MAX_NR_AKM_SUITES];
+
+		wpa_printf(MSG_INFO, "nl80211: key_mgmt_suites=0x%x",
+			params->key_mgmt_suite);
+		num_suites = wpa_cross_akm_key_mgmt_to_suites(params->key_mgmt_suite,
+			suites, ARRAY_SIZE(suites));
+		if (num_suites &&
+			nla_put(msg, NL80211_ATTR_AKM_SUITES, num_suites * sizeof(u32), suites)) {
+			wpa_printf(MSG_ERROR, "Updating multi akm_suite failed");
+			return -1;
+		}
+	}
+#endif /* CONFIG_DRIVER_NL80211_BRCM */
 	if (params->req_handshake_offload &&
 	    (drv->capa.flags & WPA_DRIVER_FLAGS_4WAY_HANDSHAKE_8021X)) {
 		    wpa_printf(MSG_DEBUG, "  * WANT_1X_4WAY_HS");
@@ -6445,7 +6479,12 @@
 	    nl80211_put_fils_connect_params(drv, params, msg) != 0)
 		return -1;
 
-	if ((params->key_mgmt_suite == WPA_KEY_MGMT_SAE ||
+	if ((
+#ifdef CONFIG_DRIVER_NL80211_BRCM
+	     (params->key_mgmt_suite & WPA_KEY_MGMT_SAE) ||
+#else
+	     params->key_mgmt_suite == WPA_KEY_MGMT_SAE ||
+#endif /* CONFIG_DRIVER_NL80211_BRCM */
 	     params->key_mgmt_suite == WPA_KEY_MGMT_FT_SAE) &&
 	    (!(drv->capa.flags & WPA_DRIVER_FLAGS_SME)) &&
 	    nla_put_flag(msg, NL80211_ATTR_EXTERNAL_AUTH_SUPPORT))
@@ -6496,7 +6535,12 @@
 		goto fail;
 
 #ifdef CONFIG_SAE
-	if ((params->key_mgmt_suite == WPA_KEY_MGMT_SAE ||
+	if ((
+#ifdef CONFIG_DRIVER_NL80211_BRCM
+	     (params->key_mgmt_suite & WPA_KEY_MGMT_SAE) ||
+#else
+	     params->key_mgmt_suite == WPA_KEY_MGMT_SAE ||
+#endif /* CONFIG_DRIVER_NL80211_BRCM */
 	     params->key_mgmt_suite == WPA_KEY_MGMT_FT_SAE) &&
 	    nl80211_put_sae_pwe(msg, params->sae_pwe) < 0)
 		goto fail;
@@ -6605,7 +6649,12 @@
 
 		if (wpa_driver_nl80211_set_mode(priv, nlmode) < 0)
 			return -1;
-		if (params->key_mgmt_suite == WPA_KEY_MGMT_SAE ||
+		if (
+#ifdef CONFIG_DRIVER_NL80211_BRCM
+		    (params->key_mgmt_suite & WPA_KEY_MGMT_SAE) ||
+#else
+		    params->key_mgmt_suite == WPA_KEY_MGMT_SAE ||
+#endif /* CONFIG_DRIVER_NL80211_BRCM */
 		    params->key_mgmt_suite == WPA_KEY_MGMT_FT_SAE)
 			bss->use_nl_connect = 1;
 		else
diff --git a/wpa_supplicant/events.c b/wpa_supplicant/events.c
index ae494f6..b724eec 100644
--- a/wpa_supplicant/events.c
+++ b/wpa_supplicant/events.c
@@ -1319,6 +1319,9 @@
 #ifdef CONFIG_SAE
 	if ((wpa_s->conf->sae_pwe == 1 || ssid->sae_password_id) &&
 	    wpa_s->conf->sae_pwe != 3 && wpa_key_mgmt_sae(ssid->key_mgmt) &&
+#ifdef CONFIG_DRIVER_NL80211_BRCM
+	    !(wpa_key_mgmt_wpa_psk_no_sae(ssid->key_mgmt)) &&
+#endif /* CONFIG_DRIVER_NL80211_BRCM */
 	    !(rsnxe_capa & BIT(WLAN_RSNX_CAPAB_SAE_H2E))) {
 		if (debug_print)
 			wpa_dbg(wpa_s, MSG_DEBUG,
@@ -2893,6 +2896,9 @@
 	const u8 *p;
 	u8 bssid[ETH_ALEN];
 	bool bssid_known;
+#ifdef CONFIG_DRIVER_NL80211_BRCM
+	struct wpa_ie_data ie;
+#endif /* CONFIG_DRIVER_NL80211_BRCM */
 
 	wpa_dbg(wpa_s, MSG_DEBUG, "Association info event");
 	bssid_known = wpa_drv_get_bssid(wpa_s, bssid) == 0;
@@ -3017,6 +3023,52 @@
 	if (!found_x && data->assoc_info.req_ies)
 		wpa_sm_set_assoc_rsnxe(wpa_s->wpa, NULL, 0);
 
+#ifdef CONFIG_DRIVER_NL80211_BRCM
+	/* The WPA/RSN IE has been updated at this point. Since the Firmware could have roamed
+	 * to a different security type, update the current supplicant configuration to use the AKM
+	 * and pairwise suites from the assoc IE passed by the driver.
+	 */
+	if (wpas_driver_bss_selection(wpa_s)) {
+		if (!(wpa_sm_parse_own_wpa_ie(wpa_s->wpa, &ie) < 0)) {
+			/* Check if firmware has roamed to a different security network */
+			if(wpa_s->key_mgmt != ie.key_mgmt) {
+				wpa_dbg(wpa_s, MSG_DEBUG, "Update to AKM suite 0x%x from Assoc IE",
+					ie.key_mgmt);
+				wpa_s->key_mgmt = ie.key_mgmt;
+				wpa_sm_set_param(wpa_s->wpa, WPA_PARAM_KEY_MGMT, wpa_s->key_mgmt);
+
+				if (wpa_key_mgmt_wpa_psk_no_sae(wpa_s->key_mgmt)) {
+					/* Restore PMK as it can get overwritten if the previous
+					* association was to 802.1X.
+					*/
+					if ((!(wpa_s->drv_flags &
+					    WPA_DRIVER_FLAGS_4WAY_HANDSHAKE_PSK)) &&
+					    (wpa_s->current_ssid) &&
+					    (wpa_s->current_ssid->psk_set)) {
+						if (wpa_drv_get_bssid(wpa_s, bssid) < 0) {
+							wpa_dbg(wpa_s, MSG_ERROR, "Failed to get "
+								"BSSID");
+							wpa_supplicant_deauthenticate(
+								wpa_s, WLAN_REASON_DEAUTH_LEAVING);
+							return -1;
+						}
+						wpa_sm_set_pmk(wpa_s->wpa, wpa_s->current_ssid->psk,
+							PMK_LEN, NULL, bssid);
+					}
+				}
+			}
+			if(wpa_s->pairwise_cipher != ie.pairwise_cipher) {
+				wpa_dbg(wpa_s, MSG_DEBUG, "Update to pairwise cipher suite 0x%x "
+					"from Assoc IE", ie.pairwise_cipher);
+				wpa_s->pairwise_cipher = ie.pairwise_cipher;
+				wpa_sm_set_param(wpa_s->wpa, WPA_PARAM_PAIRWISE,
+					wpa_s->pairwise_cipher);
+			}
+			// TODO: Notify the framework about security type change b/230766005
+		}
+	}
+#endif /* CONFIG_DRIVER_NL80211_BRCM */
+
 #ifdef CONFIG_FILS
 #ifdef CONFIG_SME
 	if ((wpa_s->sme.auth_alg == WPA_AUTH_ALG_FILS ||
diff --git a/wpa_supplicant/wpa_supplicant.c b/wpa_supplicant/wpa_supplicant.c
index 139907f..99af85e 100644
--- a/wpa_supplicant/wpa_supplicant.c
+++ b/wpa_supplicant/wpa_supplicant.c
@@ -1891,6 +1891,15 @@
 		wpa_sm_set_param(wpa_s->wpa, WPA_PARAM_DENY_PTK0_REKEY, 0);
 	}
 
+#ifdef CONFIG_DRIVER_NL80211_BRCM
+	if ((wpa_s->key_mgmt & WPA_KEY_MGMT_CROSS_AKM_ROAM) &&
+	    IS_CROSS_AKM_ROAM_KEY_MGMT(ssid->key_mgmt)) {
+		wpa_s->key_mgmt = WPA_KEY_MGMT_SAE | WPA_KEY_MGMT_PSK;
+		wpa_dbg(wpa_s, MSG_INFO,
+			"WPA: Updating to KEY_MGMT SAE+PSK for seamless roaming");
+	}
+#endif /* CONFIG_DRIVER_NL80211_BRCM */
+
 	return 0;
 }
 
@@ -3904,7 +3913,12 @@
 #endif /* CONFIG_WEP */
 
 	if ((wpa_s->drv_flags & WPA_DRIVER_FLAGS_4WAY_HANDSHAKE_PSK) &&
-	    (params.key_mgmt_suite == WPA_KEY_MGMT_PSK ||
+	    (
+#ifdef CONFIG_DRIVER_NL80211_BRCM
+	     (params.key_mgmt_suite & WPA_KEY_MGMT_PSK) ||
+#else
+	     params.key_mgmt_suite == WPA_KEY_MGMT_PSK ||
+#endif /* CONFIG_DRIVER_NL80211_BRCM */
 	     params.key_mgmt_suite == WPA_KEY_MGMT_FT_PSK)) {
 		params.passphrase = ssid->passphrase;
 		if (ssid->psk_set)
@@ -3929,7 +3943,12 @@
 		else
 			params.req_key_mgmt_offload = 1;
 
-		if ((params.key_mgmt_suite == WPA_KEY_MGMT_PSK ||
+		if ((
+#ifdef CONFIG_DRIVER_NL80211_BRCM
+		     (params.key_mgmt_suite & WPA_KEY_MGMT_PSK) ||
+#else
+		     params.key_mgmt_suite == WPA_KEY_MGMT_PSK ||
+#endif /* CONFIG_DRIVER_NL80211_BRCM */
 		     params.key_mgmt_suite == WPA_KEY_MGMT_PSK_SHA256 ||
 		     params.key_mgmt_suite == WPA_KEY_MGMT_FT_PSK) &&
 		    ssid->psk_set)