Convert LockdownVpnTracker to NetworkCallbacks.

This will allow moving LockdownVpnTracker from the connectivity
to the VPN code. This requires moderate refactoring since it's
pretty tightly coupled to both.

In this CL:

1. Add an @hide API to tell ConnectivityService that legacy
   lockdown VPN is enabled. I chose not to use the existing
   setVpnRequiredForUids API because that method has specific
   semantics and because it will be required long term since
   it's used by non-legacy VPN types.

2. Instead of updating LockdownVpnTracker inline from the
   ConnectivityService handler thread, have it listen to
   NetworkCallbacks. This introduces an extra thread hop, but
   most of the interactions between the lockdown VPN and CS were
   via NetworkAgent, which is asynchronous anyway.

3. Add code to LegacyTypeTracker to send the extra
   CONNECTIVITY_ACTION broadcast for the underlying network type
   that is sent after the VPN connects. In order to do this, make
   Make LockdownVpnTracker specify its underlying network
   (via setUnderlyingNetworks) when it connects.

4. Reimplement LockdownVpnTracker#augmentNetworkInfo based on
   information that is available in ConnectivityService.

5. Remove the code in LockdownVpnTracker that counted errors.
   I think this code has not worked since lollipop, because
   ConnectivityService never sees NetworkInfo objects in state
   FAILED. This is because ConnectivityService only hears about
   NetworkInfo objects via NetworkAgents, and LegacyVpnRunner
   only registers its NetworkAgent when the connection succeeds.

Bug: 173331190
Test: passes existing tests in ConnectivityServiceTest
Change-Id: I66d18512882efd468ee0ecec61f28786a195b357
diff --git a/framework/src/android/net/ConnectivityManager.java b/framework/src/android/net/ConnectivityManager.java
index 0976b75..8437798 100644
--- a/framework/src/android/net/ConnectivityManager.java
+++ b/framework/src/android/net/ConnectivityManager.java
@@ -1221,6 +1221,45 @@
     }
 
     /**
+     * Informs ConnectivityService of whether the legacy lockdown VPN, as implemented by
+     * LockdownVpnTracker, is in use. This is deprecated for new devices starting from Android 12
+     * but is still supported for backwards compatibility.
+     * <p>
+     * This type of VPN is assumed always to use the system default network, and must always declare
+     * exactly one underlying network, which is the network that was the default when the VPN
+     * connected.
+     * <p>
+     * Calling this method with {@code true} enables legacy behaviour, specifically:
+     * <ul>
+     *     <li>Any VPN that applies to userId 0 behaves specially with respect to deprecated
+     *     {@link #CONNECTIVITY_ACTION} broadcasts. Any such broadcasts will have the state in the
+     *     {@link #EXTRA_NETWORK_INFO} replaced by state of the VPN network. Also, any time the VPN
+     *     connects, a {@link #CONNECTIVITY_ACTION} broadcast will be sent for the network
+     *     underlying the VPN.</li>
+     *     <li>Deprecated APIs that return {@link NetworkInfo} objects will have their state
+     *     similarly replaced by the VPN network state.</li>
+     *     <li>Information on current network interfaces passed to NetworkStatsService will not
+     *     include any VPN interfaces.</li>
+     * </ul>
+     *
+     * @param enabled whether legacy lockdown VPN is enabled or disabled
+     *
+     * TODO: @SystemApi(client = MODULE_LIBRARIES)
+     *
+     * @hide
+     */
+    @RequiresPermission(anyOf = {
+            NetworkStack.PERMISSION_MAINLINE_NETWORK_STACK,
+            android.Manifest.permission.NETWORK_SETTINGS})
+    public void setLegacyLockdownVpnEnabled(boolean enabled) {
+        try {
+            mService.setLegacyLockdownVpnEnabled(enabled);
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
+    /**
      * Returns details about the currently active default data network
      * for a given uid.  This is for internal use only to avoid spying
      * other apps.
diff --git a/framework/src/android/net/IConnectivityManager.aidl b/framework/src/android/net/IConnectivityManager.aidl
index f909d13..ab134eb 100644
--- a/framework/src/android/net/IConnectivityManager.aidl
+++ b/framework/src/android/net/IConnectivityManager.aidl
@@ -151,6 +151,7 @@
     boolean isVpnLockdownEnabled(int userId);
     List<String> getVpnLockdownWhitelist(int userId);
     void setRequireVpnForUids(boolean requireVpn, in UidRange[] ranges);
+    void setLegacyLockdownVpnEnabled(boolean enabled);
 
     void setProvisioningNotificationVisible(boolean visible, int networkType, in String action);
 
diff --git a/services/core/java/com/android/server/ConnectivityService.java b/services/core/java/com/android/server/ConnectivityService.java
index 42b6a7f..84aaca0 100644
--- a/services/core/java/com/android/server/ConnectivityService.java
+++ b/services/core/java/com/android/server/ConnectivityService.java
@@ -318,7 +318,7 @@
     // TODO: investigate if mLockdownEnabled can be removed and replaced everywhere by
     // a direct call to LockdownVpnTracker.isEnabled().
     @GuardedBy("mVpns")
-    private boolean mLockdownEnabled;
+    private volatile boolean mLockdownEnabled;
     @GuardedBy("mVpns")
     private LockdownVpnTracker mLockdownTracker;
 
@@ -755,6 +755,27 @@
             }
         }
 
+        // When a lockdown VPN connects, send another CONNECTED broadcast for the underlying
+        // network type, to preserve previous behaviour.
+        private void maybeSendLegacyLockdownBroadcast(@NonNull NetworkAgentInfo vpnNai) {
+            if (vpnNai != mService.getLegacyLockdownNai()) return;
+
+            if (vpnNai.declaredUnderlyingNetworks == null
+                    || vpnNai.declaredUnderlyingNetworks.length != 1) {
+                Log.wtf(TAG, "Legacy lockdown VPN must have exactly one underlying network: "
+                        + Arrays.toString(vpnNai.declaredUnderlyingNetworks));
+                return;
+            }
+            final NetworkAgentInfo underlyingNai =  mService.getNetworkAgentInfoForNetwork(
+                    vpnNai.declaredUnderlyingNetworks[0]);
+            if (underlyingNai == null) return;
+
+            final int type = underlyingNai.networkInfo.getType();
+            final DetailedState state = DetailedState.CONNECTED;
+            maybeLogBroadcast(underlyingNai, state, type, true /* isDefaultNetwork */);
+            mService.sendLegacyNetworkBroadcast(underlyingNai, state, type);
+        }
+
         /** Adds the given network to the specified legacy type list. */
         public void add(int type, NetworkAgentInfo nai) {
             if (!isTypeSupported(type)) {
@@ -772,9 +793,17 @@
 
             // Send a broadcast if this is the first network of its type or if it's the default.
             final boolean isDefaultNetwork = mService.isDefaultNetwork(nai);
+
+            // If a legacy lockdown VPN is active, override the NetworkInfo state in all broadcasts
+            // to preserve previous behaviour.
+            final DetailedState state = mService.getLegacyLockdownState(DetailedState.CONNECTED);
             if ((list.size() == 1) || isDefaultNetwork) {
-                maybeLogBroadcast(nai, DetailedState.CONNECTED, type, isDefaultNetwork);
-                mService.sendLegacyNetworkBroadcast(nai, DetailedState.CONNECTED, type);
+                maybeLogBroadcast(nai, state, type, isDefaultNetwork);
+                mService.sendLegacyNetworkBroadcast(nai, state, type);
+            }
+
+            if (type == TYPE_VPN && state == DetailedState.CONNECTED) {
+                maybeSendLegacyLockdownBroadcast(nai);
             }
         }
 
@@ -1474,11 +1503,9 @@
         if (isNetworkWithCapabilitiesBlocked(nc, uid, ignoreBlocked)) {
             networkInfo.setDetailedState(DetailedState.BLOCKED, null, null);
         }
-        synchronized (mVpns) {
-            if (mLockdownTracker != null) {
-                mLockdownTracker.augmentNetworkInfo(networkInfo);
-            }
-        }
+        networkInfo.setDetailedState(
+                getLegacyLockdownState(networkInfo.getDetailedState()),
+                "" /* reason */, null /* extraInfo */);
     }
 
     /**
@@ -1537,14 +1564,6 @@
         return nai.network;
     }
 
-    // Public because it's used by mLockdownTracker.
-    public NetworkInfo getActiveNetworkInfoUnfiltered() {
-        enforceAccessPermission();
-        final int uid = mDeps.getCallingUid();
-        NetworkState state = getUnfilteredActiveNetworkState(uid);
-        return state.networkInfo;
-    }
-
     @Override
     public NetworkInfo getActiveNetworkInfoForUid(int uid, boolean ignoreBlocked) {
         NetworkStack.checkNetworkStackPermission(mContext);
@@ -2340,13 +2359,6 @@
     }
 
     private Intent makeGeneralIntent(NetworkInfo info, String bcastType) {
-        synchronized (mVpns) {
-            if (mLockdownTracker != null) {
-                info = new NetworkInfo(info);
-                mLockdownTracker.augmentNetworkInfo(info);
-            }
-        }
-
         Intent intent = new Intent(bcastType);
         intent.putExtra(ConnectivityManager.EXTRA_NETWORK_INFO, new NetworkInfo(info));
         intent.putExtra(ConnectivityManager.EXTRA_NETWORK_TYPE, info.getType());
@@ -2887,7 +2899,15 @@
                         Log.wtf(TAG, "Non-virtual networks cannot have underlying networks");
                         break;
                     }
+
                     final List<Network> underlying = (List<Network>) arg.second;
+
+                    if (isLegacyLockdownNai(nai)
+                            && (underlying == null || underlying.size() != 1)) {
+                        Log.wtf(TAG, "Legacy lockdown VPN " + nai.toShortString()
+                                + " must have exactly one underlying network: " + underlying);
+                    }
+
                     final Network[] oldUnderlying = nai.declaredUnderlyingNetworks;
                     nai.declaredUnderlyingNetworks = (underlying != null)
                             ? underlying.toArray(new Network[0]) : null;
@@ -3496,7 +3516,6 @@
                     //  incorrect) behavior.
                     mNetworkActivityTracker.updateDataActivityTracking(
                             null /* newNetwork */, nai);
-                    notifyLockdownVpn(nai);
                     ensureNetworkTransitionWakelock(nai.toShortString());
                 }
             }
@@ -5071,10 +5090,59 @@
         mVpnBlockedUidRanges = newVpnBlockedUidRanges;
     }
 
+    @Override
+    public void setLegacyLockdownVpnEnabled(boolean enabled) {
+        enforceSettingsPermission();
+        mHandler.post(() -> mLockdownEnabled = enabled);
+    }
+
+    // TODO: remove when the VPN code moves out.
     private boolean isLockdownVpnEnabled() {
         return mKeyStore.contains(Credentials.LOCKDOWN_VPN);
     }
 
+    private boolean isLegacyLockdownNai(NetworkAgentInfo nai) {
+        return mLockdownEnabled
+                && getVpnType(nai) == VpnManager.TYPE_VPN_LEGACY
+                && nai.networkCapabilities.appliesToUid(Process.FIRST_APPLICATION_UID);
+    }
+
+    private NetworkAgentInfo getLegacyLockdownNai() {
+        if (!mLockdownEnabled) {
+            return null;
+        }
+        // The legacy lockdown VPN always only applies to UID 0.
+        final NetworkAgentInfo nai = getVpnForUid(Process.FIRST_APPLICATION_UID);
+        if (nai == null || !isLegacyLockdownNai(nai)) return null;
+
+        // The legacy lockdown VPN must always have exactly one underlying network.
+        if (nai.declaredUnderlyingNetworks == null ||  nai.declaredUnderlyingNetworks.length != 1) {
+            return null;
+        }
+
+        // The legacy lockdown VPN always uses the default network.
+        // If the VPN's underlying network is no longer the current default network, it means that
+        // the default network has just switched, and the VPN is about to disconnect.
+        // Report that the VPN is not connected, so when the state of NetworkInfo objects
+        // overwritten by getLegacyLockdownState will be set to CONNECTING and not CONNECTED.
+        final NetworkAgentInfo defaultNetwork = getDefaultNetwork();
+        if (defaultNetwork == null
+                || !defaultNetwork.network.equals(nai.declaredUnderlyingNetworks[0])) {
+            return null;
+        }
+
+        return nai;
+    };
+
+    private DetailedState getLegacyLockdownState(DetailedState origState) {
+        if (origState != DetailedState.CONNECTED) {
+            return origState;
+        }
+        return (mLockdownEnabled && getLegacyLockdownNai() == null)
+                ? DetailedState.CONNECTING
+                : DetailedState.CONNECTED;
+    }
+
     @Override
     public boolean updateLockdownVpn() {
         // Allow the system UID for the system server and for Settings.
@@ -5087,32 +5155,32 @@
 
         synchronized (mVpns) {
             // Tear down existing lockdown if profile was removed
-            mLockdownEnabled = isLockdownVpnEnabled();
-            if (mLockdownEnabled) {
-                byte[] profileTag = mKeyStore.get(Credentials.LOCKDOWN_VPN);
-                if (profileTag == null) {
-                    loge("Lockdown VPN configured but cannot be read from keystore");
-                    return false;
-                }
-                String profileName = new String(profileTag);
-                final VpnProfile profile = VpnProfile.decode(
-                        profileName, mKeyStore.get(Credentials.VPN + profileName));
-                if (profile == null) {
-                    loge("Lockdown VPN configured invalid profile " + profileName);
-                    setLockdownTracker(null);
-                    return true;
-                }
-                int user = UserHandle.getUserId(mDeps.getCallingUid());
-                Vpn vpn = mVpns.get(user);
-                if (vpn == null) {
-                    logw("VPN for user " + user + " not ready yet. Skipping lockdown");
-                    return false;
-                }
-                setLockdownTracker(
-                        new LockdownVpnTracker(mContext, this, mHandler, mKeyStore, vpn,  profile));
-            } else {
+            if (!isLockdownVpnEnabled()) {
                 setLockdownTracker(null);
+                return true;
             }
+
+            byte[] profileTag = mKeyStore.get(Credentials.LOCKDOWN_VPN);
+            if (profileTag == null) {
+                loge("Lockdown VPN configured but cannot be read from keystore");
+                return false;
+            }
+            String profileName = new String(profileTag);
+            final VpnProfile profile = VpnProfile.decode(
+                    profileName, mKeyStore.get(Credentials.VPN + profileName));
+            if (profile == null) {
+                loge("Lockdown VPN configured invalid profile " + profileName);
+                setLockdownTracker(null);
+                return true;
+            }
+            int user = UserHandle.getUserId(mDeps.getCallingUid());
+            Vpn vpn = mVpns.get(user);
+            if (vpn == null) {
+                logw("VPN for user " + user + " not ready yet. Skipping lockdown");
+                return false;
+            }
+            setLockdownTracker(
+                    new LockdownVpnTracker(mContext, mHandler, mKeyStore, vpn,  profile));
         }
 
         return true;
@@ -7341,7 +7409,6 @@
             mLingerMonitor.noteLingerDefaultNetwork(oldDefaultNetwork, newDefaultNetwork);
         }
         mNetworkActivityTracker.updateDataActivityTracking(newDefaultNetwork, oldDefaultNetwork);
-        notifyLockdownVpn(newDefaultNetwork);
         handleApplyDefaultProxy(null != newDefaultNetwork
                 ? newDefaultNetwork.linkProperties.getHttpProxy() : null);
         updateTcpBufferSizes(null != newDefaultNetwork
@@ -7799,12 +7866,6 @@
                 mDefaultInetConditionPublished = newDefaultNetwork.lastValidated ? 100 : 0;
                 mLegacyTypeTracker.add(
                         newDefaultNetwork.networkInfo.getType(), newDefaultNetwork);
-                // If the legacy VPN is connected, notifyLockdownVpn may end up sending a broadcast
-                // to reflect the NetworkInfo of this new network. This broadcast has to be sent
-                // after the disconnect broadcasts above, but before the broadcasts sent by the
-                // legacy type tracker below.
-                // TODO : refactor this, it's too complex
-                notifyLockdownVpn(newDefaultNetwork);
             }
         }
 
@@ -7862,18 +7923,6 @@
         sendInetConditionBroadcast(nai.networkInfo);
     }
 
-    private void notifyLockdownVpn(NetworkAgentInfo nai) {
-        synchronized (mVpns) {
-            if (mLockdownTracker != null) {
-                if (nai != null && nai.isVPN()) {
-                    mLockdownTracker.onVpnStateChanged(nai.networkInfo);
-                } else {
-                    mLockdownTracker.onNetworkInfoChanged();
-                }
-            }
-        }
-    }
-
     @NonNull
     private NetworkInfo mixInInfo(@NonNull final NetworkAgentInfo nai, @NonNull NetworkInfo info) {
         final NetworkInfo newInfo = new NetworkInfo(info);
@@ -7912,7 +7961,6 @@
             oldInfo = networkAgent.networkInfo;
             networkAgent.networkInfo = newInfo;
         }
-        notifyLockdownVpn(networkAgent);
 
         if (DBG) {
             log(networkAgent.toShortString() + " EVENT_NETWORK_INFO_CHANGED, going from "
diff --git a/tests/net/java/com/android/server/ConnectivityServiceTest.java b/tests/net/java/com/android/server/ConnectivityServiceTest.java
index 7905f57..1da4e31 100644
--- a/tests/net/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/net/java/com/android/server/ConnectivityServiceTest.java
@@ -7339,11 +7339,14 @@
         when(mKeyStore.get(Credentials.VPN + profileName)).thenReturn(encodedProfile);
     }
 
-    private void establishLegacyLockdownVpn() throws Exception {
+    private void establishLegacyLockdownVpn(Network underlying) throws Exception {
+        // The legacy lockdown VPN only supports userId 0, and must have an underlying network.
+        assertNotNull(underlying);
         mMockVpn.setVpnType(VpnManager.TYPE_VPN_LEGACY);
         // The legacy lockdown VPN only supports userId 0.
         final Set<UidRange> ranges = Collections.singleton(UidRange.createForUser(PRIMARY_USER));
         mMockVpn.registerAgent(ranges);
+        mMockVpn.setUnderlyingNetworks(new Network[]{underlying});
         mMockVpn.connect(true);
     }
 
@@ -7351,6 +7354,9 @@
     public void testLegacyLockdownVpn() throws Exception {
         mServiceContext.setPermission(
                 Manifest.permission.CONTROL_VPN, PERMISSION_GRANTED);
+        // For LockdownVpnTracker to call registerSystemDefaultNetworkCallback.
+        mServiceContext.setPermission(
+                Manifest.permission.NETWORK_SETTINGS, PERMISSION_GRANTED);
 
         final NetworkRequest request = new NetworkRequest.Builder().clearCapabilities().build();
         final TestNetworkCallback callback = new TestNetworkCallback();
@@ -7359,6 +7365,10 @@
         final TestNetworkCallback defaultCallback = new TestNetworkCallback();
         mCm.registerDefaultNetworkCallback(defaultCallback);
 
+        final TestNetworkCallback systemDefaultCallback = new TestNetworkCallback();
+        mCm.registerSystemDefaultNetworkCallback(systemDefaultCallback,
+                new Handler(ConnectivityThread.getInstanceLooper()));
+
         // Pretend lockdown VPN was configured.
         setupLegacyLockdownVpn();
 
@@ -7388,6 +7398,7 @@
         mCellNetworkAgent.connect(false /* validated */);
         callback.expectAvailableCallbacksUnvalidatedAndBlocked(mCellNetworkAgent);
         defaultCallback.expectAvailableCallbacksUnvalidatedAndBlocked(mCellNetworkAgent);
+        systemDefaultCallback.expectAvailableCallbacksUnvalidatedAndBlocked(mCellNetworkAgent);
         waitForIdle();
         assertNull(mMockVpn.getAgent());
 
@@ -7399,6 +7410,8 @@
         mCellNetworkAgent.sendLinkProperties(cellLp);
         callback.expectCallback(CallbackEntry.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
         defaultCallback.expectCallback(CallbackEntry.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
+        systemDefaultCallback.expectCallback(CallbackEntry.LINK_PROPERTIES_CHANGED,
+                mCellNetworkAgent);
         waitForIdle();
         assertNull(mMockVpn.getAgent());
 
@@ -7408,6 +7421,7 @@
         mCellNetworkAgent.disconnect();
         callback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
         defaultCallback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
+        systemDefaultCallback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
         b1.expectBroadcast();
 
         // When lockdown VPN is active, the NetworkInfo state in CONNECTIVITY_ACTION is overwritten
@@ -7417,6 +7431,7 @@
         mCellNetworkAgent.connect(false /* validated */);
         callback.expectAvailableCallbacksUnvalidatedAndBlocked(mCellNetworkAgent);
         defaultCallback.expectAvailableCallbacksUnvalidatedAndBlocked(mCellNetworkAgent);
+        systemDefaultCallback.expectAvailableCallbacksUnvalidatedAndBlocked(mCellNetworkAgent);
         b1.expectBroadcast();
         assertActiveNetworkInfo(TYPE_MOBILE, DetailedState.BLOCKED);
         assertNetworkInfo(TYPE_MOBILE, DetailedState.BLOCKED);
@@ -7439,9 +7454,10 @@
         mMockVpn.expectStartLegacyVpnRunner();
         b1 = expectConnectivityAction(TYPE_VPN, DetailedState.CONNECTED);
         ExpectedBroadcast b2 = expectConnectivityAction(TYPE_MOBILE, DetailedState.CONNECTED);
-        establishLegacyLockdownVpn();
+        establishLegacyLockdownVpn(mCellNetworkAgent.getNetwork());
         callback.expectAvailableThenValidatedCallbacks(mMockVpn);
         defaultCallback.expectAvailableThenValidatedCallbacks(mMockVpn);
+        systemDefaultCallback.assertNoCallback();
         NetworkCapabilities vpnNc = mCm.getNetworkCapabilities(mMockVpn.getNetwork());
         b1.expectBroadcast();
         b2.expectBroadcast();
@@ -7453,9 +7469,7 @@
         assertTrue(vpnNc.hasTransport(TRANSPORT_CELLULAR));
         assertFalse(vpnNc.hasTransport(TRANSPORT_WIFI));
         assertFalse(vpnNc.hasCapability(NET_CAPABILITY_NOT_METERED));
-        VpnTransportInfo ti = (VpnTransportInfo) vpnNc.getTransportInfo();
-        assertNotNull(ti);
-        assertEquals(VpnManager.TYPE_VPN_LEGACY, ti.type);
+        assertVpnTransportInfo(vpnNc, VpnManager.TYPE_VPN_LEGACY);
 
         // Switch default network from cell to wifi. Expect VPN to disconnect and reconnect.
         final LinkProperties wifiLp = new LinkProperties();
@@ -7483,11 +7497,10 @@
         // fact that a VPN is connected should only result in the VPN itself being unblocked, not
         // any other network. Bug in isUidBlockedByVpn?
         callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
-        callback.expectCapabilitiesThat(mMockVpn, nc -> nc.hasTransport(TRANSPORT_WIFI));
         callback.expectCallback(CallbackEntry.LOST, mMockVpn);
-        defaultCallback.expectCapabilitiesThat(mMockVpn, nc -> nc.hasTransport(TRANSPORT_WIFI));
         defaultCallback.expectCallback(CallbackEntry.LOST, mMockVpn);
         defaultCallback.expectAvailableCallbacksUnvalidatedAndBlocked(mWiFiNetworkAgent);
+        systemDefaultCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
 
         // While the VPN is reconnecting on the new network, everything is blocked.
         assertActiveNetworkInfo(TYPE_WIFI, DetailedState.BLOCKED);
@@ -7498,9 +7511,10 @@
         // The VPN comes up again on wifi.
         b1 = expectConnectivityAction(TYPE_VPN, DetailedState.CONNECTED);
         b2 = expectConnectivityAction(TYPE_WIFI, DetailedState.CONNECTED);
-        establishLegacyLockdownVpn();
+        establishLegacyLockdownVpn(mWiFiNetworkAgent.getNetwork());
         callback.expectAvailableThenValidatedCallbacks(mMockVpn);
         defaultCallback.expectAvailableThenValidatedCallbacks(mMockVpn);
+        systemDefaultCallback.assertNoCallback();
         b1.expectBroadcast();
         b2.expectBroadcast();
         assertActiveNetworkInfo(TYPE_WIFI, DetailedState.CONNECTED);
@@ -7514,14 +7528,10 @@
         assertTrue(vpnNc.hasCapability(NET_CAPABILITY_NOT_METERED));
 
         // Disconnect cell. Nothing much happens since it's not the default network.
-        // Whenever LockdownVpnTracker is connected, it will send a connected broadcast any time any
-        // NetworkInfo is updated. This is probably a bug.
-        // TODO: consider fixing this.
-        b1 = expectConnectivityAction(TYPE_WIFI, DetailedState.CONNECTED);
         mCellNetworkAgent.disconnect();
-        b1.expectBroadcast();
         callback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
         defaultCallback.assertNoCallback();
+        systemDefaultCallback.assertNoCallback();
 
         assertActiveNetworkInfo(TYPE_WIFI, DetailedState.CONNECTED);
         assertNetworkInfo(TYPE_MOBILE, DetailedState.DISCONNECTED);
@@ -7531,6 +7541,7 @@
         b1 = expectConnectivityAction(TYPE_WIFI, DetailedState.DISCONNECTED);
         mWiFiNetworkAgent.disconnect();
         callback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        systemDefaultCallback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
         b1.expectBroadcast();
         callback.expectCapabilitiesThat(mMockVpn, nc -> !nc.hasTransport(TRANSPORT_WIFI));
         b2 = expectConnectivityAction(TYPE_VPN, DetailedState.DISCONNECTED);