Support for SAE-EXT-KEY & SAE-FT-EXT AKM

If the wifi network configuration has SAE key management scheme,
check if the device supports new AKMs SAE-EXT-KEY & SAE-FT-EXT.
Based on the device support for these AKMs, add the additional
SAE-EXT-KEY & SAE-FT-EXT AKMs in wifi network configuration.

Bug: 260375081
Test: Manual - Connect/Disconnect STA multiple times
Change-Id: I6041bbfd3613bd79610e4d9da18d5dd78858d73d
diff --git a/wpa_supplicant/aidl/sta_network.cpp b/wpa_supplicant/aidl/sta_network.cpp
index 0d5d4e6..db13509 100644
--- a/wpa_supplicant/aidl/sta_network.cpp
+++ b/wpa_supplicant/aidl/sta_network.cpp
@@ -2179,10 +2179,19 @@
 	KeyMgmtMask mask)
 {
 	uint32_t key_mgmt_mask = static_cast<uint32_t>(mask);
+	struct wpa_supplicant *wpa_s = retrieveIfacePtr();
 	struct wpa_ssid *wpa_ssid = retrieveNetworkPtr();
 	if (key_mgmt_mask & ~kAllowedKeyMgmtMask) {
 		return createStatus(SupplicantStatusCode::FAILURE_ARGS_INVALID);
 	}
+#ifdef CONFIG_SAE
+	struct wpa_driver_capa capa;
+	int res = wpa_drv_get_capa(wpa_s, &capa);
+	if ((res == 0) && (key_mgmt_mask & WPA_KEY_MGMT_SAE) &&
+		(capa.key_mgmt_iftype[WPA_IF_STATION] & WPA_DRIVER_CAPA_KEY_MGMT_SAE_EXT_KEY)) {
+		key_mgmt_mask |= WPA_KEY_MGMT_SAE_EXT_KEY;
+	}
+#endif
 	setFastTransitionKeyMgmt(key_mgmt_mask);
 
 	if (key_mgmt_mask & WPA_KEY_MGMT_OWE) {
@@ -2516,6 +2525,11 @@
 			(capa.key_mgmt_iftype[WPA_IF_STATION] & WPA_DRIVER_CAPA_KEY_MGMT_FT_SAE)) {
 			key_mgmt_mask |= WPA_KEY_MGMT_FT_SAE;
 		}
+		if ((key_mgmt_mask & WPA_KEY_MGMT_SAE_EXT_KEY) &&
+			(capa.key_mgmt_iftype[WPA_IF_STATION] &
+			    WPA_DRIVER_CAPA_KEY_MGMT_FT_SAE_EXT_KEY)) {
+			key_mgmt_mask |= WPA_KEY_MGMT_FT_SAE_EXT_KEY;
+		}
 #endif
 #ifdef CONFIG_FILS
 		if ((key_mgmt_mask & WPA_KEY_MGMT_FILS_SHA256) &&