Wifi: clear cached data if a connection attempt fails

After getting PMK cache, the later flow might fails and
this PMK cache become invalid.

Bug: 152931766
Test: atest FrameworksWifiTests
Change-Id: I22fe4d31c5349155ce18d5f206beb552fc62e7e2

2 files changed, 106 insertions, 0 deletions

diff --git a/service/java/com/android/server/wifi/ClientModeImpl.java b/service/java/com/android/server/wifi/ClientModeImpl.java
index 00e02df4a..2887081c9 100644
--- a/service/java/com/android/server/wifi/ClientModeImpl.java
+++ b/service/java/com/android/server/wifi/ClientModeImpl.java
@@ -2978,6 +2978,31 @@ public class ClientModeImpl extends StateMachine {
         handleConnectionAttemptEndForDiagnostics(level2FailureCode);
     }
 
+    /* If this connection attempt fails after 802.1x stage, clear intermediate cached data. */
+    void clearNetworkCachedDataIfNeeded(WifiConfiguration config, int reason) {
+        if (config == null) return;
+
+        switch(reason) {
+            case 14: // MICHAEL_MIC_FAILURE
+            case 15: // 4WAY_HANDSHAKE_TIMEOUT
+            case 16: // GROUP_KEY_UPDATE_TIMEOUT
+            case 17: // IE_IN_4WAY_DIFFERS
+            case 18: // GROUP_CIPHER_NOT_VALID
+            case 19: // PAIRWISE_CIPHER_NOT_VALID
+            case 20: // AKMP_NOT_VALID
+            case 23: // IEEE_802_1X_AUTH_FAILED
+            case 24: // CIPHER_SUITE_REJECTED
+            case 29: // BAD_CIPHER_OR_AKM
+            case 45: // PEERKEY_MISMATCH
+            case 49: // INVALID_PMKID
+                mWifiNative.removeNetworkCachedData(config.networkId);
+                break;
+            default:
+                logi("Keep PMK cache for network disconnection reason " + reason);
+                break;
+        }
+    }
+
     /**
      * Returns the sufficient RSSI for the frequency that this network is last seen on.
      */
@@ -4211,6 +4236,7 @@ public class ClientModeImpl extends StateMachine {
                     // idempotent commands are executed twice (stopping Dhcp, enabling the SPS mode
                     // at the chip etc...
                     if (mVerboseLoggingEnabled) log("ConnectModeState: Network connection lost ");
+                    clearNetworkCachedDataIfNeeded(getTargetWifiConfiguration(), message.arg2);
                     handleNetworkDisconnect();
                     transitionTo(mDisconnectedState);
                     break;
@@ -5240,6 +5266,7 @@ public class ClientModeImpl extends StateMachine {
                                 + " BSSID=" + bssid
                                 + " target=" + target);
                     }
+                    clearNetworkCachedDataIfNeeded(getTargetWifiConfiguration(), message.arg2);
                     if (bssid != null && bssid.equals(mTargetBssid)) {
                         handleNetworkDisconnect();
                         transitionTo(mDisconnectedState);
@@ -5395,6 +5422,7 @@ public class ClientModeImpl extends StateMachine {
                         mWifiDiagnostics.captureBugReportData(
                                 WifiDiagnostics.REPORT_REASON_UNEXPECTED_DISCONNECT);
                     }
+
                     boolean localGen = message.arg1 == 1;
                     if (!localGen) { // ignore disconnects initiated by wpa_supplicant.
                         mWifiScoreCard.noteNonlocalDisconnect(message.arg2);
@@ -5598,6 +5626,7 @@ public class ClientModeImpl extends StateMachine {
                                         getTargetSsid(), bssid,
                                         WifiLastResortWatchdog.FAILURE_CODE_AUTHENTICATION);
                     }
+                    clearNetworkCachedDataIfNeeded(getTargetWifiConfiguration(), message.arg2);
                     break;
                 case WifiMonitor.SUPPLICANT_STATE_CHANGE_EVENT:
                     StateChangeResult stateChangeResult = (StateChangeResult) message.obj;
diff --git a/tests/wifitests/src/com/android/server/wifi/ClientModeImplTest.java b/tests/wifitests/src/com/android/server/wifi/ClientModeImplTest.java
index 2923315d6..b9e9df4af 100644
--- a/tests/wifitests/src/com/android/server/wifi/ClientModeImplTest.java
+++ b/tests/wifitests/src/com/android/server/wifi/ClientModeImplTest.java
@@ -4989,4 +4989,81 @@ public class ClientModeImplTest extends WifiBaseTest {
 
         verifyNoMoreInteractions(mNetworkAgentHandler);
     }
+
+    /*
+     * Verify that network cached data is cleared correctly in
+     * disconnected state.
+     */
+    @Test
+    public void testNetworkCachedDataIsClearedCorrectlyInDisconnectedState() throws Exception {
+        // Setup CONNECT_MODE & a WifiConfiguration
+        initializeAndAddNetworkAndVerifySuccess();
+        mCmi.sendMessage(ClientModeImpl.CMD_START_CONNECT, 0, 0, sBSSID);
+        mLooper.dispatchAll();
+
+        // got UNSPECIFIED during this connection attempt
+        mCmi.sendMessage(WifiMonitor.NETWORK_DISCONNECTION_EVENT, 0, 1, sBSSID);
+        mLooper.dispatchAll();
+
+        assertEquals("DisconnectedState", getCurrentState().getName());
+        verify(mWifiNative, never()).removeNetworkCachedData(anyInt());
+
+        // got 4WAY_HANDSHAKE_TIMEOUT during this connection attempt
+        mCmi.sendMessage(WifiMonitor.NETWORK_DISCONNECTION_EVENT, 0, 15, sBSSID);
+        mLooper.dispatchAll();
+
+        assertEquals("DisconnectedState", getCurrentState().getName());
+        verify(mWifiNative).removeNetworkCachedData(FRAMEWORK_NETWORK_ID);
+    }
+
+    /*
+     * Verify that network cached data is cleared correctly in
+     * disconnected state.
+     */
+    @Test
+    public void testNetworkCachedDataIsClearedCorrectlyInObtainingIpState() throws Exception {
+        initializeAndAddNetworkAndVerifySuccess();
+
+        verify(mWifiNative).removeAllNetworks(WIFI_IFACE_NAME);
+
+        IActionListener connectActionListener = mock(IActionListener.class);
+        mCmi.connect(null, 0, mock(Binder.class), connectActionListener, 0, Binder.getCallingUid());
+        mLooper.dispatchAll();
+        verify(connectActionListener).onSuccess();
+
+        verify(mWifiConfigManager).enableNetwork(eq(0), eq(true), anyInt(), any());
+
+        mCmi.sendMessage(WifiMonitor.NETWORK_CONNECTION_EVENT, 0, 0, sBSSID);
+        mLooper.dispatchAll();
+
+        mCmi.sendMessage(WifiMonitor.SUPPLICANT_STATE_CHANGE_EVENT, 0, 0,
+                new StateChangeResult(0, sWifiSsid, sBSSID, SupplicantState.COMPLETED));
+        mLooper.dispatchAll();
+
+        assertEquals("ObtainingIpState", getCurrentState().getName());
+
+        // got 4WAY_HANDSHAKE_TIMEOUT during this connection attempt
+        mCmi.sendMessage(WifiMonitor.NETWORK_DISCONNECTION_EVENT, 0, 15, sBSSID);
+        mLooper.dispatchAll();
+
+        verify(mWifiNative).removeNetworkCachedData(FRAMEWORK_NETWORK_ID);
+    }
+
+    /*
+     * Verify that network cached data is NOT cleared in ConnectedState.
+     */
+    @Test
+    public void testNetworkCachedDataIsClearedIf4WayHandshakeFailure() throws Exception {
+        when(mWifiScoreCard.detectAbnormalDisconnection())
+                .thenReturn(WifiHealthMonitor.REASON_SHORT_CONNECTION_NONLOCAL);
+        InOrder inOrderWifiLockManager = inOrder(mWifiLockManager);
+        connect();
+        inOrderWifiLockManager.verify(mWifiLockManager).updateWifiClientConnected(true);
+
+        // got 4WAY_HANDSHAKE_TIMEOUT
+        mCmi.sendMessage(WifiMonitor.NETWORK_DISCONNECTION_EVENT, 0, 15, sBSSID);
+        mLooper.dispatchAll();
+        verify(mWifiNative, never()).removeNetworkCachedData(anyInt());
+    }
+
 }