Update NETWORK_CONNECTION_EVENT netId after WPS

Updates the netId argument of a supplicants NETWORK_CONNECTION_EVENT
message inside of WifiStateMachine's WpsRunningState, after successfully
connecting to a supplicant network, and learning the networkId from
supplicant. This allows WifiStateMachine to proceed with IP
provisioning, and fixes a bug where it would immediately disconnect from
the WPS network, then reconnect after a delay.

Bug: 37834109
Test: Manually tested WPS connection
Test: Added unit test
Change-Id: Ie261d0db8c386049329baee8319e014ed5302547

1 files changed, 17 insertions, 7 deletions

diff --git a/service/java/com/android/server/wifi/WifiStateMachine.java b/service/java/com/android/server/wifi/WifiStateMachine.java
index ed2bc9dc0..39c82d576 100644
--- a/service/java/com/android/server/wifi/WifiStateMachine.java
+++ b/service/java/com/android/server/wifi/WifiStateMachine.java
@@ -6751,7 +6751,11 @@ public class WifiStateMachine extends StateMachine implements WifiNative.WifiRss
                     // Ignore intermediate success, wait for full connection
                     break;
                 case WifiMonitor.NETWORK_CONNECTION_EVENT:
-                    if (loadNetworksFromSupplicantAfterWps()) {
+                    Pair<Boolean, Integer> loadResult = loadNetworksFromSupplicantAfterWps();
+                    boolean success = loadResult.first;
+                    int netId = loadResult.second;
+                    if (success) {
+                        message.arg1 = netId;
                         replyToMessage(mSourceMessage, WifiManager.WPS_COMPLETED);
                     } else {
                         replyToMessage(mSourceMessage, WifiManager.WPS_FAILED,
@@ -6847,13 +6851,17 @@ public class WifiStateMachine extends StateMachine implements WifiNative.WifiRss
 
         /**
          * Load network config from wpa_supplicant after WPS is complete.
+         * @return a boolean (true if load was successful) and integer (Network Id of currently
+         *          connected network, can be {@link WifiConfiguration#INVALID_NETWORK_ID} despite
+         *          successful loading, if multiple networks in supplicant) pair.
          */
-        private boolean loadNetworksFromSupplicantAfterWps() {
+        private Pair<Boolean, Integer> loadNetworksFromSupplicantAfterWps() {
             Map<String, WifiConfiguration> configs = new HashMap<>();
             SparseArray<Map<String, String>> extras = new SparseArray<>();
+            int netId = WifiConfiguration.INVALID_NETWORK_ID;
             if (!mWifiNative.migrateNetworksFromSupplicant(configs, extras)) {
                 loge("Failed to load networks from wpa_supplicant after Wps");
-                return false;
+                return Pair.create(false, WifiConfiguration.INVALID_NETWORK_ID);
             }
             for (Map.Entry<String, WifiConfiguration> entry : configs.entrySet()) {
                 WifiConfiguration config = entry.getValue();
@@ -6864,15 +6872,17 @@ public class WifiStateMachine extends StateMachine implements WifiNative.WifiRss
                         config, mSourceMessage.sendingUid);
                 if (!result.isSuccess()) {
                     loge("Failed to add network after WPS: " + entry.getValue());
-                    return false;
+                    return Pair.create(false, WifiConfiguration.INVALID_NETWORK_ID);
                 }
                 if (!mWifiConfigManager.enableNetwork(
                         result.getNetworkId(), true, mSourceMessage.sendingUid)) {
-                    loge("Failed to enable network after WPS: " + entry.getValue());
-                    return false;
+                    Log.wtf(TAG, "Failed to enable network after WPS: " + entry.getValue());
+                    return Pair.create(false, WifiConfiguration.INVALID_NETWORK_ID);
                 }
+                netId = result.getNetworkId();
             }
-            return true;
+            return Pair.create(true,
+                    configs.size() == 1 ? netId : WifiConfiguration.INVALID_NETWORK_ID);
         }
     }