Update WifiInfo MAC address more frequently

Currently, we reset MAC address stored in WifiInfo to 02:00:00:00:00:00
whenever we are handling disconnect with Connected MAC Randomization
enabled. The disconnect handling can happen after we have started
connecting to a new network and generated a new randomized MAC address.

Update the MAC address to the real value once we have successfully
connected to the new network.

Bug: 78013990
Test: Manual Check, Unittest
Change-Id: I8f3d66bef353e91d3938a74c5d0f10d746903961

2 files changed, 40 insertions, 0 deletions

diff --git a/service/java/com/android/server/wifi/WifiStateMachine.java b/service/java/com/android/server/wifi/WifiStateMachine.java
index 4a0a24a21..c969d252f 100644
--- a/service/java/com/android/server/wifi/WifiStateMachine.java
+++ b/service/java/com/android/server/wifi/WifiStateMachine.java
@@ -4321,6 +4321,7 @@ public class WifiStateMachine extends StateMachine {
                     if (config != null) {
                         mWifiInfo.setBSSID(mLastBssid);
                         mWifiInfo.setNetworkId(mLastNetworkId);
+                        mWifiInfo.setMacAddress(mWifiNative.getMacAddress(mInterfaceName));
 
                         ScanDetailCache scanDetailCache =
                                 mWifiConfigManager.getScanDetailCacheForNetwork(config.networkId);
@@ -4866,6 +4867,7 @@ public class WifiStateMachine extends StateMachine {
                     mWifiInfo.setBSSID((String) message.obj);
                     mLastNetworkId = message.arg1;
                     mWifiInfo.setNetworkId(mLastNetworkId);
+                    mWifiInfo.setMacAddress(mWifiNative.getMacAddress(mInterfaceName));
                     if(!mLastBssid.equals(message.obj)) {
                         mLastBssid = (String) message.obj;
                         sendNetworkStateChangeBroadcast(mLastBssid);
diff --git a/tests/wifitests/src/com/android/server/wifi/WifiStateMachineTest.java b/tests/wifitests/src/com/android/server/wifi/WifiStateMachineTest.java
index 86692c8ce..44fdd4874 100644
--- a/tests/wifitests/src/com/android/server/wifi/WifiStateMachineTest.java
+++ b/tests/wifitests/src/com/android/server/wifi/WifiStateMachineTest.java
@@ -2338,4 +2338,42 @@ public class WifiStateMachineTest {
         mLooper.dispatchAll();
         verify(mIpClient).shutdown();
     }
+
+    /**
+     * Verify that WifiInfo's MAC address is updated when the state machine receives
+     * NETWORK_CONNECTION_EVENT while in ConnectedState.
+     */
+    @Test
+    public void verifyWifiInfoMacUpdatedWithNetworkConnectionWhileConnected() throws Exception {
+        when(mWifiNative.getMacAddress(WIFI_IFACE_NAME))
+                .thenReturn(TEST_LOCAL_MAC_ADDRESS.toString());
+        connect();
+        assertEquals("ConnectedState", getCurrentState().getName());
+        assertEquals(TEST_LOCAL_MAC_ADDRESS.toString(), mWsm.getWifiInfo().getMacAddress());
+
+        when(mWifiNative.getMacAddress(WIFI_IFACE_NAME))
+                .thenReturn(TEST_GLOBAL_MAC_ADDRESS.toString());
+        mWsm.sendMessage(WifiMonitor.NETWORK_CONNECTION_EVENT, 0, 0, sBSSID);
+        mLooper.dispatchAll();
+        assertEquals(TEST_GLOBAL_MAC_ADDRESS.toString(), mWsm.getWifiInfo().getMacAddress());
+    }
+
+    /**
+     * Verify that WifiInfo's MAC address is updated when the state machine receives
+     * NETWORK_CONNECTION_EVENT while in DisconnectedState.
+     */
+    @Test
+    public void verifyWifiInfoMacUpdatedWithNetworkConnectionWhileDisconnected() throws Exception {
+        when(mWifiNative.getMacAddress(WIFI_IFACE_NAME))
+                .thenReturn(TEST_LOCAL_MAC_ADDRESS.toString());
+        disconnect();
+        assertEquals("DisconnectedState", getCurrentState().getName());
+        assertEquals(TEST_LOCAL_MAC_ADDRESS.toString(), mWsm.getWifiInfo().getMacAddress());
+
+        when(mWifiNative.getMacAddress(WIFI_IFACE_NAME))
+                .thenReturn(TEST_GLOBAL_MAC_ADDRESS.toString());
+        mWsm.sendMessage(WifiMonitor.NETWORK_CONNECTION_EVENT, 0, 0, sBSSID);
+        mLooper.dispatchAll();
+        assertEquals(TEST_GLOBAL_MAC_ADDRESS.toString(), mWsm.getWifiInfo().getMacAddress());
+    }
 }