ClientModeImpl: Handle disconnection event in DisconnectingState

Since DisconnectedState no longer handles NETWORK_DISCONNECTION or
SUPPLICANT_STATE_CHANGE events, handle these events directly in
DisconnectingState. Otherwise WifiInfo state is not correctly updated.

Also, move Wifinfo state reset to exit of ConnectingState instead
of exit of ConnectModeState.

Note: This is a temporary fix for the cuttlefish test failure. Long term
goal is to remove the Disconnectingstate.

Bug: 156219024
Test: atest com.android.server.wifi
Change-Id: I8a7d64f1428761d928922d78aff998b3ad781526
diff --git a/service/java/com/android/server/wifi/ClientModeImpl.java b/service/java/com/android/server/wifi/ClientModeImpl.java
index 9a9be5e..39f5394 100644
--- a/service/java/com/android/server/wifi/ClientModeImpl.java
+++ b/service/java/com/android/server/wifi/ClientModeImpl.java
@@ -3818,9 +3818,6 @@
             if (!mWifiNative.removeAllNetworks(mInterfaceName)) {
                 loge("Failed to remove networks on exiting connect mode");
             }
-            mWifiInfo.reset();
-            mWifiInfo.setSupplicantState(SupplicantState.DISCONNECTED);
-            mWifiScoreCard.noteSupplicantStateChanged(mWifiInfo);
             mWifiHealthMonitor.setWifiEnabled(false);
             mWifiDataStall.reset();
             stopClientMode();
@@ -4406,7 +4403,10 @@
 
     class ConnectingState extends State {
         @Override
-        public void enter() {
+        public void exit() {
+            mWifiInfo.reset();
+            mWifiInfo.setSupplicantState(SupplicantState.DISCONNECTED);
+            mWifiScoreCard.noteSupplicantStateChanged(mWifiInfo);
         }
 
         @Override
@@ -5722,13 +5722,22 @@
                     break;
                 }
                 case WifiMonitor.SUPPLICANT_STATE_CHANGE_EVENT: {
-                    /**
-                     * If we get a SUPPLICANT_STATE_CHANGE_EVENT before NETWORK_DISCONNECTION_EVENT
-                     * we have missed the network disconnection, transition to mDisconnectedState
-                     * and handle the rest of the events there
-                     */
-                    mMessageHandlingStatus = MESSAGE_HANDLING_STATUS_DEFERRED;
-                    deferMessage(message);
+                    StateChangeResult stateChangeResult = (StateChangeResult) message.obj;
+                    SupplicantState state = handleSupplicantStateChange(stateChangeResult);
+                    if (state == SupplicantState.DISCONNECTED) {
+                        if (mVerboseLoggingEnabled) {
+                            log("Missed CTRL-EVENT-DISCONNECTED, disconnect");
+                        }
+                        handleNetworkDisconnect(false);
+                        transitionTo(mDisconnectedState);
+                    }
+                    break;
+                }
+                case WifiMonitor.NETWORK_DISCONNECTION_EVENT: {
+                    DisconnectEventInfo eventInfo = (DisconnectEventInfo) message.obj;
+                    if (mVerboseLoggingEnabled) {
+                        log("DisconnectingState: Network disconnection " + eventInfo);
+                    }
                     handleNetworkDisconnect(false);
                     transitionTo(mDisconnectedState);
                     break;
diff --git a/service/tests/wifitests/src/com/android/server/wifi/ClientModeImplTest.java b/service/tests/wifitests/src/com/android/server/wifi/ClientModeImplTest.java
index a003ea2..c62c0b6 100644
--- a/service/tests/wifitests/src/com/android/server/wifi/ClientModeImplTest.java
+++ b/service/tests/wifitests/src/com/android/server/wifi/ClientModeImplTest.java
@@ -5081,4 +5081,43 @@
         mLooper.dispatchAll();
         verify(mWifiNative, never()).removeNetworkCachedData(anyInt());
     }
+
+    @Test
+    public void testVerifyWifiInfoStateOnFrameworkDisconnect() throws Exception {
+        connect();
+
+        assertEquals(mCmi.getWifiInfo().getSupplicantState(), SupplicantState.COMPLETED);
+
+        // Now trigger disconnect
+        mCmi.disconnectCommand();
+        mLooper.dispatchAll();
+
+        // get disconnect event
+        mCmi.sendMessage(WifiMonitor.SUPPLICANT_STATE_CHANGE_EVENT, 0, 0,
+                new StateChangeResult(0, WifiSsid.createFromAsciiEncoded(mConnectedNetwork.SSID),
+                        sBSSID, SupplicantState.DISCONNECTED));
+        mLooper.dispatchAll();
+
+        assertEquals(mCmi.getWifiInfo().getSupplicantState(), SupplicantState.DISCONNECTED);
+    }
+
+    @Test
+    public void testVerifyWifiInfoStateOnFrameworkDisconnectButMissingDisconnectEvent()
+            throws Exception {
+        connect();
+
+        assertEquals(mCmi.getWifiInfo().getSupplicantState(), SupplicantState.COMPLETED);
+
+        // Now trigger disconnect
+        mCmi.disconnectCommand();
+        mLooper.dispatchAll();
+
+        // missing disconnect event, but got supplicant state change with disconnect state instead.
+        mCmi.sendMessage(WifiMonitor.SUPPLICANT_STATE_CHANGE_EVENT, 0, 0,
+                new StateChangeResult(0, WifiSsid.createFromAsciiEncoded(mConnectedNetwork.SSID),
+                        sBSSID, SupplicantState.DISCONNECTED));
+        mLooper.dispatchAll();
+
+        assertEquals(mCmi.getWifiInfo().getSupplicantState(), SupplicantState.DISCONNECTED);
+    }
 }