From 62ef82fb3942c43f14c0f878e80f21822f3f46ea Mon Sep 17 00:00:00 2001
From: yinxu <yinxu@google.com>
Date: Wed, 27 Jan 2021 13:07:21 -0800
Subject: [PATCH] Only remove the relevant network

WifiStatusTracker only monitors WiFi network or VCN network over WiFi,
so whenever NetworkCallback.onLost is triggered, we should only handle
those relevant networks.

Bug: 178611654
Test: Manual tests
Change-Id: I6cef52a4373e8fc3d4acb47d9237f750efa4a7f2
---
 .../settingslib/wifi/WifiStatusTracker.java   | 27 ++++++++++++++++---
 .../policy/NetworkControllerBaseTest.java     |  5 ++++
 2 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/packages/SettingsLib/src/com/android/settingslib/wifi/WifiStatusTracker.java b/packages/SettingsLib/src/com/android/settingslib/wifi/WifiStatusTracker.java
index 5580afe70faf0..78282fb14f5f3 100644
--- a/packages/SettingsLib/src/com/android/settingslib/wifi/WifiStatusTracker.java
+++ b/packages/SettingsLib/src/com/android/settingslib/wifi/WifiStatusTracker.java
@@ -19,6 +19,7 @@ import android.content.Context;
 import android.content.Intent;
 import android.net.ConnectivityManager;
 import android.net.ConnectivityManager.NetworkCallback;
+import android.net.LinkProperties;
 import android.net.Network;
 import android.net.NetworkCapabilities;
 import android.net.NetworkInfo;
@@ -38,7 +39,9 @@ import android.util.FeatureFlagUtils;
 import com.android.settingslib.R;
 import com.android.settingslib.Utils;
 
+import java.util.HashSet;
 import java.util.List;
+import java.util.Set;
 
 /**
  * Track status of Wi-Fi for the Sys UI.
@@ -50,6 +53,7 @@ public class WifiStatusTracker {
     private final NetworkScoreManager mNetworkScoreManager;
     private final ConnectivityManager mConnectivityManager;
     private final Handler mHandler = new Handler(Looper.getMainLooper());
+    private final Set<Integer> mNetworks = new HashSet<>();
     private final WifiNetworkScoreCache.CacheListener mCacheListener =
             new WifiNetworkScoreCache.CacheListener(mHandler) {
                 @Override
@@ -64,6 +68,20 @@ public class WifiStatusTracker {
             .addTransportType(NetworkCapabilities.TRANSPORT_WIFI)
             .addTransportType(NetworkCapabilities.TRANSPORT_CELLULAR).build();
     private final NetworkCallback mNetworkCallback = new NetworkCallback() {
+        @Override
+        public void onAvailable(
+                Network network, NetworkCapabilities networkCapabilities,
+                LinkProperties linkProperties, boolean blocked) {
+            boolean isVcnOverWifi =
+                    networkCapabilities.hasTransport(NetworkCapabilities.TRANSPORT_CELLULAR)
+                            && (Utils.tryGetWifiInfoForVcn(networkCapabilities) != null);
+            boolean isWifi =
+                    networkCapabilities.hasTransport(NetworkCapabilities.TRANSPORT_WIFI);
+            if (isVcnOverWifi || isWifi) {
+                mNetworks.add(network.getNetId());
+            }
+        }
+
         // Note: onCapabilitiesChanged is guaranteed to be called "immediately" after onAvailable
         // and onLinkPropertiesChanged.
         @Override
@@ -84,9 +102,12 @@ public class WifiStatusTracker {
 
         @Override
         public void onLost(Network network) {
-            updateWifiInfo(null);
-            updateStatusLabel();
-            mCallback.run();
+            if (mNetworks.contains(network.getNetId())) {
+                mNetworks.remove(network.getNetId());
+                updateWifiInfo(null);
+                updateStatusLabel();
+                mCallback.run();
+            }
         }
     };
     private final NetworkCallback mDefaultNetworkCallback = new NetworkCallback() {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/NetworkControllerBaseTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/NetworkControllerBaseTest.java
index deabcbe520b61..c2a030f3f6ae4 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/NetworkControllerBaseTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/NetworkControllerBaseTest.java
@@ -37,6 +37,7 @@ import static org.mockito.Mockito.when;
 import android.app.Instrumentation;
 import android.content.Intent;
 import android.net.ConnectivityManager;
+import android.net.LinkProperties;
 import android.net.Network;
 import android.net.NetworkCapabilities;
 import android.net.NetworkRequest;
@@ -343,6 +344,8 @@ public class NetworkControllerBaseTest extends SysuiTestCase {
         setConnectivityCommon(networkType, validated, isConnected);
         if (networkType == NetworkCapabilities.TRANSPORT_WIFI) {
             if (isConnected) {
+                mNetworkCallback.onAvailable(mock(Network.class),
+                        new NetworkCapabilities(mNetCapabilities), new LinkProperties(), false);
                 mNetworkCallback.onCapabilitiesChanged(
                         mock(Network.class), new NetworkCapabilities(mNetCapabilities));
             } else {
@@ -357,6 +360,8 @@ public class NetworkControllerBaseTest extends SysuiTestCase {
         setConnectivityCommon(networkType, validated, isConnected);
         if (networkType == NetworkCapabilities.TRANSPORT_CELLULAR) {
             if (isConnected) {
+                mNetworkCallback.onAvailable(mock(Network.class),
+                        new NetworkCapabilities(mNetCapabilities), new LinkProperties(), false);
                 mNetworkCallback.onCapabilitiesChanged(
                         mock(Network.class), new NetworkCapabilities(mNetCapabilities));
             } else {