WifiStateMachine: Fix WPS handling

Since wpa_supplicant is no longer our config store, we need to read the
network params out of it add it to WifiConfigManager after WPS network
addition.

Bug: 34714021
Test: Manual tests. Able to connect to wps networks now.
Test: Some unit tests. It can be better, but need to figure out how to
extract the message reply from WSM.
Change-Id: I41231698c7e99ec46d34a60a70c04a6ee25ebfe4

3 files changed, 174 insertions, 6 deletions

diff --git a/service/java/com/android/server/wifi/SupplicantStaNetworkHal.java b/service/java/com/android/server/wifi/SupplicantStaNetworkHal.java
index d81a790da..46f8c41b9 100644
--- a/service/java/com/android/server/wifi/SupplicantStaNetworkHal.java
+++ b/service/java/com/android/server/wifi/SupplicantStaNetworkHal.java
@@ -230,8 +230,7 @@ public class SupplicantStaNetworkHal {
             Map<String, String> metadata = parseNetworkExtra(mIdStr);
             networkExtras.putAll(metadata);
         } else {
-            Log.e(TAG, "getIdStr failed");
-            return false;
+            Log.w(TAG, "getIdStr failed or empty");
         }
         return loadWifiEnterpriseConfig(config.SSID, config.enterpriseConfig);
     }
diff --git a/service/java/com/android/server/wifi/WifiStateMachine.java b/service/java/com/android/server/wifi/WifiStateMachine.java
index 58b4ff639..d6179e513 100644
--- a/service/java/com/android/server/wifi/WifiStateMachine.java
+++ b/service/java/com/android/server/wifi/WifiStateMachine.java
@@ -130,9 +130,11 @@ import java.net.Inet4Address;
 import java.net.InetAddress;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.LinkedList;
 import java.util.List;
+import java.util.Map;
 import java.util.Queue;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -5130,6 +5132,11 @@ public class WifiStateMachine extends StateMachine implements WifiNative.WifiRss
                 case WifiManager.START_WPS:
                     WpsInfo wpsInfo = (WpsInfo) message.obj;
                     WpsResult wpsResult = new WpsResult();
+                    // TODO(b/32898136): Not needed when we start deleting networks from supplicant
+                    // on disconnect.
+                    if (!mWifiNative.removeAllNetworks()) {
+                        loge("Failed to remove networks before WPS");
+                    }
                     switch (wpsInfo.setup) {
                         case WpsInfo.PBC:
                             if (mWifiNative.startWpsPbc(wpsInfo.BSSID)) {
@@ -6465,6 +6472,23 @@ public class WifiStateMachine extends StateMachine implements WifiNative.WifiRss
         }
     }
 
+    /**
+     * WPS connection flow:
+     * 1. WifiStateMachine receive WPS_START message from WifiManager API.
+     * 2. WifiStateMachine initiates the appropriate WPS operation using WifiNative methods:
+     * {@link WifiNative#startWpsPbc(String)}, {@link WifiNative#startWpsPinDisplay(String)}, etc.
+     * 3. WifiStateMachine then transitions to this WpsRunningState.
+     * 4a. Once WifiStateMachine receive the connected event:
+     * {@link WifiMonitor#NETWORK_CONNECTION_EVENT},
+     * 4a.1 Load the network params out of wpa_supplicant.
+     * 4a.2 Add the network with params to WifiConfigManager.
+     * 4a.3 Enable the network with |disableOthers| set to true.
+     * 4a.4 Send a response to the original source of WifiManager API using {@link #mSourceMessage}.
+     * 4b. Any failures are notified to the original source of WifiManager API
+     * using {@link #mSourceMessage}.
+     * 5. We then transition to disconnected state and let network selection reconnect to the newly
+     * added network.
+     */
     class WpsRunningState extends State {
         // Tracks the source to provide a reply
         private Message mSourceMessage;
@@ -6481,7 +6505,12 @@ public class WifiStateMachine extends StateMachine implements WifiNative.WifiRss
                     // Ignore intermediate success, wait for full connection
                     break;
                 case WifiMonitor.NETWORK_CONNECTION_EVENT:
-                    replyToMessage(mSourceMessage, WifiManager.WPS_COMPLETED);
+                    if (loadNetworksFromSupplicantAfterWps()) {
+                        replyToMessage(mSourceMessage, WifiManager.WPS_COMPLETED);
+                    } else {
+                        replyToMessage(mSourceMessage, WifiManager.WPS_FAILED,
+                                WifiManager.ERROR);
+                    }
                     mSourceMessage.recycle();
                     mSourceMessage = null;
                     deferMessage(message);
@@ -6570,9 +6599,30 @@ public class WifiStateMachine extends StateMachine implements WifiNative.WifiRss
             return HANDLED;
         }
 
-        @Override
-        public void exit() {
-            mWifiConfigManager.loadFromStore();
+        /**
+         * Load network config from wpa_supplicant after WPS is complete.
+         */
+        private boolean loadNetworksFromSupplicantAfterWps() {
+            Map<String, WifiConfiguration> configs = new HashMap<>();
+            SparseArray<Map<String, String>> extras = new SparseArray<>();
+            if (!mWifiNative.migrateNetworksFromSupplicant(configs, extras)) {
+                loge("Failed to load networks from wpa_supplicant after Wps");
+                return false;
+            }
+            for (Map.Entry<String, WifiConfiguration> entry : configs.entrySet()) {
+                NetworkUpdateResult result = mWifiConfigManager.addOrUpdateNetwork(
+                        entry.getValue(), mSourceMessage.sendingUid);
+                if (!result.isSuccess()) {
+                    loge("Failed to add network after WPS: " + entry.getValue());
+                    return false;
+                }
+                if (!mWifiConfigManager.enableNetwork(
+                        result.getNetworkId(), true, mSourceMessage.sendingUid)) {
+                    loge("Failed to enable network after WPS: " + entry.getValue());
+                    return false;
+                }
+            }
+            return true;
         }
     }
 
diff --git a/tests/wifitests/src/com/android/server/wifi/WifiStateMachineTest.java b/tests/wifitests/src/com/android/server/wifi/WifiStateMachineTest.java
index 4ea233a31..4a331e3cb 100644
--- a/tests/wifitests/src/com/android/server/wifi/WifiStateMachineTest.java
+++ b/tests/wifitests/src/com/android/server/wifi/WifiStateMachineTest.java
@@ -44,6 +44,7 @@ import android.net.wifi.WifiInfo;
 import android.net.wifi.WifiManager;
 import android.net.wifi.WifiScanner;
 import android.net.wifi.WifiSsid;
+import android.net.wifi.WpsInfo;
 import android.net.wifi.hotspot2.PasspointConfiguration;
 import android.net.wifi.hotspot2.pps.HomeSp;
 import android.net.wifi.p2p.IWifiP2pManager;
@@ -70,6 +71,7 @@ import android.test.mock.MockContentProvider;
 import android.test.mock.MockContentResolver;
 import android.test.suitebuilder.annotation.SmallTest;
 import android.util.Log;
+import android.util.SparseArray;
 
 import com.android.internal.R;
 import com.android.internal.app.IBatteryStats;
@@ -96,6 +98,7 @@ import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.CountDownLatch;
 
@@ -1271,4 +1274,120 @@ public class WifiStateMachineTest {
         mLooper.stopAutoDispatch();
         verify(mPasspointManager, never()).getMatchingWifiConfig(any(ScanResult.class));
     }
+
+    /**
+     * Verify successful Wps PBC network connection.
+     */
+    @Test
+    public void wpsPbcConnectSuccess() throws Exception {
+        loadComponentsInStaMode();
+        mWsm.setOperationalMode(WifiStateMachine.CONNECT_MODE);
+        mLooper.dispatchAll();
+
+        when(mWifiNative.startWpsPbc(anyString())).thenReturn(true);
+        WpsInfo wpsInfo = new WpsInfo();
+        wpsInfo.setup = WpsInfo.PBC;
+
+        mLooper.startAutoDispatch();
+        mWsm.sendMessage(WifiManager.START_WPS, 0, 0, wpsInfo);
+        mLooper.stopAutoDispatch();
+        verify(mWifiNative).startWpsPbc(anyString());
+
+        assertEquals("WpsRunningState", getCurrentState().getName());
+
+        setupMocksForWpsNetworkMigration();
+
+        mLooper.startAutoDispatch();
+        mWsm.sendMessage(WifiMonitor.NETWORK_CONNECTION_EVENT, 0, 0, null);
+        mLooper.stopAutoDispatch();
+
+        assertEquals("DisconnectedState", getCurrentState().getName());
+    }
+
+    /**
+     * Verify failure in starting Wps PBC network connection.
+     */
+    @Test
+    public void wpsPbcConnectFailure() throws Exception {
+        loadComponentsInStaMode();
+        mWsm.setOperationalMode(WifiStateMachine.CONNECT_MODE);
+        mLooper.dispatchAll();
+
+        when(mWifiNative.startWpsPbc(anyString())).thenReturn(false);
+        WpsInfo wpsInfo = new WpsInfo();
+        wpsInfo.setup = WpsInfo.PBC;
+
+        mLooper.startAutoDispatch();
+        mWsm.sendMessage(WifiManager.START_WPS, 0, 0, wpsInfo);
+        mLooper.stopAutoDispatch();
+        verify(mWifiNative).startWpsPbc(anyString());
+
+        assertFalse("WpsRunningState".equals(getCurrentState().getName()));
+    }
+
+    /**
+     * Verify successful Wps Pin Display network connection.
+     */
+    @Test
+    public void wpsPinDisplayConnectSuccess() throws Exception {
+        loadComponentsInStaMode();
+        mWsm.setOperationalMode(WifiStateMachine.CONNECT_MODE);
+        mLooper.dispatchAll();
+
+        when(mWifiNative.startWpsPinDisplay(anyString())).thenReturn("34545434");
+        WpsInfo wpsInfo = new WpsInfo();
+        wpsInfo.setup = WpsInfo.DISPLAY;
+
+        mLooper.startAutoDispatch();
+        mWsm.sendMessage(WifiManager.START_WPS, 0, 0, wpsInfo);
+        mLooper.stopAutoDispatch();
+        verify(mWifiNative).startWpsPinDisplay(anyString());
+
+        assertEquals("WpsRunningState", getCurrentState().getName());
+
+        setupMocksForWpsNetworkMigration();
+
+        mLooper.startAutoDispatch();
+        mWsm.sendMessage(WifiMonitor.NETWORK_CONNECTION_EVENT, 0, 0, null);
+        mLooper.stopAutoDispatch();
+
+        assertEquals("DisconnectedState", getCurrentState().getName());
+    }
+
+    /**
+     * Verify failure in Wps Pin Display network connection.
+     */
+    @Test
+    public void wpsPinDisplayConnectFailure() throws Exception {
+        loadComponentsInStaMode();
+        mWsm.setOperationalMode(WifiStateMachine.CONNECT_MODE);
+        mLooper.dispatchAll();
+
+        when(mWifiNative.startWpsPinDisplay(anyString())).thenReturn(null);
+        WpsInfo wpsInfo = new WpsInfo();
+        wpsInfo.setup = WpsInfo.DISPLAY;
+
+        mLooper.startAutoDispatch();
+        mWsm.sendMessage(WifiManager.START_WPS, 0, 0, wpsInfo);
+        mLooper.stopAutoDispatch();
+        verify(mWifiNative).startWpsPinDisplay(anyString());
+
+        assertFalse("WpsRunningState".equals(getCurrentState().getName()));
+    }
+
+    private void setupMocksForWpsNetworkMigration() {
+        int newNetworkId = 5;
+        // Now trigger the network connection event for adding the WPS network.
+        doAnswer(new AnswerWithArguments() {
+            public boolean answer(Map<String, WifiConfiguration> configs,
+                                  SparseArray<Map<String, String>> networkExtras) throws Exception {
+                configs.put("dummy", new WifiConfiguration());
+                return true;
+            }
+        }).when(mWifiNative).migrateNetworksFromSupplicant(any(Map.class), any(SparseArray.class));
+        when(mWifiConfigManager.addOrUpdateNetwork(any(WifiConfiguration.class), anyInt()))
+                .thenReturn(new NetworkUpdateResult(newNetworkId));
+        when(mWifiConfigManager.enableNetwork(eq(newNetworkId), anyBoolean(), anyInt()))
+                .thenReturn(true);
+    }
 }