Add test coverage for LockdownVpnTracker.

Also add a comment to an existing test.

Bug: 173331190
Test: test-only change
Change-Id: I4211475f053f0cdf2b804132c52ce5205a772193
diff --git a/services/core/java/com/android/server/ConnectivityService.java b/services/core/java/com/android/server/ConnectivityService.java
index 1b7912e..d72a8ef 100644
--- a/services/core/java/com/android/server/ConnectivityService.java
+++ b/services/core/java/com/android/server/ConnectivityService.java
@@ -889,6 +889,13 @@
         }
 
         /**
+         * Get a reference to the system keystore.
+         */
+        public KeyStore getKeyStore() {
+            return KeyStore.getInstance();
+        }
+
+        /**
          * @see ProxyTracker
          */
         public ProxyTracker makeProxyTracker(@NonNull Context context,
@@ -990,7 +997,7 @@
         mProxyTracker = mDeps.makeProxyTracker(mContext, mHandler);
 
         mNetd = netd;
-        mKeyStore = KeyStore.getInstance();
+        mKeyStore = mDeps.getKeyStore();
         mTelephonyManager = (TelephonyManager) mContext.getSystemService(Context.TELEPHONY_SERVICE);
         mAppOpsManager = (AppOpsManager) mContext.getSystemService(Context.APP_OPS_SERVICE);
         mLocationPermissionChecker = new LocationPermissionChecker(mContext);
@@ -4977,16 +4984,21 @@
         mVpnBlockedUidRanges = newVpnBlockedUidRanges;
     }
 
+    private boolean isLockdownVpnEnabled() {
+        return mKeyStore.contains(Credentials.LOCKDOWN_VPN);
+    }
+
     @Override
     public boolean updateLockdownVpn() {
-        if (mDeps.getCallingUid() != Process.SYSTEM_UID) {
-            logw("Lockdown VPN only available to AID_SYSTEM");
+        if (mDeps.getCallingUid() != Process.SYSTEM_UID
+                && Binder.getCallingPid() != Process.myPid()) {
+            logw("Lockdown VPN only available to system process or AID_SYSTEM");
             return false;
         }
 
         synchronized (mVpns) {
             // Tear down existing lockdown if profile was removed
-            mLockdownEnabled = LockdownVpnTracker.isEnabled();
+            mLockdownEnabled = isLockdownVpnEnabled();
             if (mLockdownEnabled) {
                 byte[] profileTag = mKeyStore.get(Credentials.LOCKDOWN_VPN);
                 if (profileTag == null) {
@@ -5007,7 +5019,8 @@
                     logw("VPN for user " + user + " not ready yet. Skipping lockdown");
                     return false;
                 }
-                setLockdownTracker(new LockdownVpnTracker(mContext, this, mHandler, vpn, profile));
+                setLockdownTracker(
+                        new LockdownVpnTracker(mContext, this, mHandler, mKeyStore, vpn,  profile));
             } else {
                 setLockdownTracker(null);
             }
@@ -5095,7 +5108,7 @@
 
         synchronized (mVpns) {
             // Can't set always-on VPN if legacy VPN is already in lockdown mode.
-            if (LockdownVpnTracker.isEnabled()) {
+            if (isLockdownVpnEnabled()) {
                 return false;
             }
 
@@ -5201,7 +5214,7 @@
             }
             userVpn = new Vpn(mHandler.getLooper(), mContext, mNMS, mNetd, userId, mKeyStore);
             mVpns.put(userId, userVpn);
-            if (mUserManager.getUserInfo(userId).isPrimary() && LockdownVpnTracker.isEnabled()) {
+            if (mUserManager.getUserInfo(userId).isPrimary() && isLockdownVpnEnabled()) {
                 updateLockdownVpn();
             }
         }
@@ -5285,7 +5298,7 @@
     private void onUserUnlocked(int userId) {
         synchronized (mVpns) {
             // User present may be sent because of an unlock, which might mean an unlocked keystore.
-            if (mUserManager.getUserInfo(userId).isPrimary() && LockdownVpnTracker.isEnabled()) {
+            if (mUserManager.getUserInfo(userId).isPrimary() && isLockdownVpnEnabled()) {
                 updateLockdownVpn();
             } else {
                 startAlwaysOnVpn(userId);
diff --git a/tests/net/java/com/android/server/ConnectivityServiceTest.java b/tests/net/java/com/android/server/ConnectivityServiceTest.java
index 2a0414f..0223fdd 100644
--- a/tests/net/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/net/java/com/android/server/ConnectivityServiceTest.java
@@ -21,6 +21,7 @@
 import static android.app.PendingIntent.FLAG_IMMUTABLE;
 import static android.content.Intent.ACTION_USER_ADDED;
 import static android.content.Intent.ACTION_USER_REMOVED;
+import static android.content.Intent.ACTION_USER_UNLOCKED;
 import static android.content.pm.PackageInfo.REQUESTED_PERMISSION_GRANTED;
 import static android.content.pm.PackageManager.GET_PERMISSIONS;
 import static android.content.pm.PackageManager.MATCH_ANY_USER;
@@ -220,6 +221,7 @@
 import android.os.UserHandle;
 import android.os.UserManager;
 import android.provider.Settings;
+import android.security.Credentials;
 import android.security.KeyStore;
 import android.system.Os;
 import android.telephony.TelephonyManager;
@@ -236,6 +238,7 @@
 import com.android.internal.app.IBatteryStats;
 import com.android.internal.net.VpnConfig;
 import com.android.internal.net.VpnInfo;
+import com.android.internal.net.VpnProfile;
 import com.android.internal.util.ArrayUtils;
 import com.android.internal.util.WakeupMessage;
 import com.android.internal.util.test.BroadcastInterceptingContext;
@@ -280,6 +283,7 @@
 import java.net.InetAddress;
 import java.net.InetSocketAddress;
 import java.net.Socket;
+import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
@@ -387,6 +391,7 @@
     @Mock MockableSystemProperties mSystemProperties;
     @Mock EthernetManager mEthernetManager;
     @Mock NetworkPolicyManager mNetworkPolicyManager;
+    @Mock KeyStore mKeyStore;
 
     private ArgumentCaptor<ResolverParamsParcel> mResolverParamsParcelCaptor =
             ArgumentCaptor.forClass(ResolverParamsParcel.class);
@@ -1065,6 +1070,15 @@
         private int mVpnType = VpnManager.TYPE_VPN_SERVICE;
         private VpnInfo mVpnInfo;
 
+        // These ConditionVariables allow tests to wait for LegacyVpnRunner to be stopped/started.
+        // TODO: this scheme is ad-hoc and error-prone because it does not fail if, for example, the
+        // test expects two starts in a row, or even if the production code calls start twice in a
+        // row. find a better solution. Simply putting a method to create a LegacyVpnRunner into
+        // Vpn.Dependencies doesn't work because LegacyVpnRunner is not a static class and has
+        // extensive access into the internals of Vpn.
+        private ConditionVariable mStartLegacyVpnCv = new ConditionVariable();
+        private ConditionVariable mStopVpnRunnerCv = new ConditionVariable();
+
         public MockVpn(int userId) {
             super(startHandlerThreadAndReturnLooper(), mServiceContext,
                     new Dependencies() {
@@ -1078,7 +1092,7 @@
                             return mDeviceIdleInternal;
                         }
                     },
-                    mNetworkManagementService, mMockNetd, userId, mock(KeyStore.class));
+                    mNetworkManagementService, mMockNetd, userId, mKeyStore);
         }
 
         public void setUids(Set<UidRange> uids) {
@@ -1190,10 +1204,44 @@
             }
             mAgentRegistered = false;
             setUids(null);
+            // Remove NET_CAPABILITY_INTERNET or MockNetworkAgent will refuse to connect later on.
+            mNetworkCapabilities.removeCapability(NET_CAPABILITY_INTERNET);
             mInterface = null;
         }
 
         @Override
+        public void startLegacyVpnRunner() {
+            mStartLegacyVpnCv.open();
+        }
+
+        public void expectStartLegacyVpnRunner() {
+            assertTrue("startLegacyVpnRunner not called after " + TIMEOUT_MS + " ms",
+                    mStartLegacyVpnCv.block(TIMEOUT_MS));
+
+            // startLegacyVpn calls stopVpnRunnerPrivileged, which will open mStopVpnRunnerCv, just
+            // before calling startLegacyVpnRunner. Restore mStopVpnRunnerCv, so the test can expect
+            // that the VpnRunner is stopped and immediately restarted by calling
+            // expectStartLegacyVpnRunner() and expectStopVpnRunnerPrivileged() back-to-back.
+            mStopVpnRunnerCv = new ConditionVariable();
+        }
+
+        @Override
+        public void stopVpnRunnerPrivileged() {
+            if (mVpnRunner != null) {
+                super.stopVpnRunnerPrivileged();
+                disconnect();
+                mStartLegacyVpnCv = new ConditionVariable();
+            }
+            mVpnRunner = null;
+            mStopVpnRunnerCv.open();
+        }
+
+        public void expectStopVpnRunnerPrivileged() {
+            assertTrue("stopVpnRunnerPrivileged not called after " + TIMEOUT_MS + " ms",
+                    mStopVpnRunnerCv.block(TIMEOUT_MS));
+        }
+
+        @Override
         public synchronized VpnInfo getVpnInfo() {
             if (mVpnInfo != null) return mVpnInfo;
 
@@ -1274,10 +1322,12 @@
         }
     }
 
-    private static final int VPN_USER = 0;
-    private static final int APP1_UID = UserHandle.getUid(VPN_USER, 10100);
-    private static final int APP2_UID = UserHandle.getUid(VPN_USER, 10101);
-    private static final int VPN_UID = UserHandle.getUid(VPN_USER, 10043);
+    private static final int PRIMARY_USER = 0;
+    private static final int APP1_UID = UserHandle.getUid(PRIMARY_USER, 10100);
+    private static final int APP2_UID = UserHandle.getUid(PRIMARY_USER, 10101);
+    private static final int VPN_UID = UserHandle.getUid(PRIMARY_USER, 10043);
+    private static final UserInfo PRIMARY_USER_INFO = new UserInfo(PRIMARY_USER, "",
+            UserInfo.FLAG_PRIMARY);
 
     @Before
     public void setUp() throws Exception {
@@ -1288,10 +1338,9 @@
         MockitoAnnotations.initMocks(this);
         when(mMetricsService.defaultNetworkMetrics()).thenReturn(mDefaultNetworkMetrics);
 
-        when(mUserManager.getAliveUsers()).thenReturn(
-                Arrays.asList(new UserInfo[] {
-                        new UserInfo(VPN_USER, "", 0),
-                }));
+        when(mUserManager.getAliveUsers()).thenReturn(Arrays.asList(PRIMARY_USER_INFO));
+        when(mUserManager.getUserInfo(PRIMARY_USER)).thenReturn(PRIMARY_USER_INFO);
+
         final ApplicationInfo applicationInfo = new ApplicationInfo();
         applicationInfo.targetSdkVersion = Build.VERSION_CODES.Q;
         when(mPackageManager.getApplicationInfoAsUser(anyString(), anyInt(), any()))
@@ -1333,6 +1382,9 @@
         verify(mNetworkPolicyManager).registerListener(policyListenerCaptor.capture());
         mPolicyListener = policyListenerCaptor.getValue();
 
+        mServiceContext.setPermission(
+                Manifest.permission.CONTROL_VPN, PERMISSION_GRANTED);
+
         // Create local CM before sending system ready so that we can answer
         // getSystemService() correctly.
         mCm = new WrappedConnectivityManager(InstrumentationRegistry.getContext(), mService);
@@ -1364,6 +1416,7 @@
         doReturn(mMetricsService).when(deps).getMetricsLogger();
         doReturn(true).when(deps).queryUserAccess(anyInt(), anyInt());
         doReturn(mBatteryStatsService).when(deps).getBatteryStatsService();
+        doReturn(mKeyStore).when(deps).getKeyStore();
         doAnswer(inv -> {
             mPolicyTracker = new WrappedMultinetworkPolicyTracker(
                     inv.getArgument(0), inv.getArgument(1), inv.getArgument(2));
@@ -1529,10 +1582,9 @@
         }
 
         public void expectNoBroadcast(int timeoutMs) throws Exception {
-            waitForIdle();
             try {
                 final Intent intent = get(timeoutMs, TimeUnit.MILLISECONDS);
-                fail("Unexpected broadcast: " + intent.getAction());
+                fail("Unexpected broadcast: " + intent.getAction() + " " + intent.getExtras());
             } catch (TimeoutException expected) {
             } finally {
                 mServiceContext.unregisterReceiver(mReceiver);
@@ -6401,8 +6453,6 @@
         mServiceContext.setPermission(
                 Manifest.permission.CONTROL_ALWAYS_ON_VPN, PERMISSION_GRANTED);
         mServiceContext.setPermission(
-                Manifest.permission.CONTROL_VPN, PERMISSION_GRANTED);
-        mServiceContext.setPermission(
                 Manifest.permission.NETWORK_SETTINGS, PERMISSION_GRANTED);
 
         // Connect wifi and check that UIDs in the main and restricted profiles have network access.
@@ -6418,14 +6468,13 @@
         mService.setAlwaysOnVpnPackage(userId, ALWAYS_ON_PACKAGE, true /* lockdown */, allowList);
         waitForIdle();
         assertNull(mCm.getActiveNetworkForUid(uid));
+        // This is arguably overspecified: a UID that is not running doesn't have an active network.
+        // But it's useful to check that non-default users do not lose network access, and to prove
+        // that the loss of connectivity below is indeed due to the restricted profile coming up.
         assertNotNull(mCm.getActiveNetworkForUid(restrictedUid));
 
         // Start the restricted profile, and check that the UID within it loses network access.
-        when(mUserManager.getAliveUsers()).thenReturn(
-                Arrays.asList(new UserInfo[] {
-                        new UserInfo(userId, "", 0),
-                        info
-                }));
+        when(mUserManager.getAliveUsers()).thenReturn(Arrays.asList(PRIMARY_USER_INFO, info));
         // TODO: check that VPN app within restricted profile still has access, etc.
         handler.post(() -> mServiceContext.sendBroadcast(addedIntent));
         waitForIdle();
@@ -6433,10 +6482,7 @@
         assertNull(mCm.getActiveNetworkForUid(restrictedUid));
 
         // Stop the restricted profile, and check that the UID within it has network access again.
-        when(mUserManager.getAliveUsers()).thenReturn(
-                Arrays.asList(new UserInfo[] {
-                        new UserInfo(userId, "", 0),
-                }));
+        when(mUserManager.getAliveUsers()).thenReturn(Arrays.asList(PRIMARY_USER_INFO));
         handler.post(() -> mServiceContext.sendBroadcast(removedIntent));
         waitForIdle();
         assertNull(mCm.getActiveNetworkForUid(uid));
@@ -6930,6 +6976,194 @@
         mCm.unregisterNetworkCallback(vpnUidCallback);
     }
 
+    private void setupLegacyLockdownVpn() {
+        final String profileName = "testVpnProfile";
+        final byte[] profileTag = profileName.getBytes(StandardCharsets.UTF_8);
+        when(mKeyStore.contains(Credentials.LOCKDOWN_VPN)).thenReturn(true);
+        when(mKeyStore.get(Credentials.LOCKDOWN_VPN)).thenReturn(profileTag);
+
+        final VpnProfile profile = new VpnProfile(profileName);
+        profile.name = "My VPN";
+        profile.server = "192.0.2.1";
+        profile.dnsServers = "8.8.8.8";
+        profile.type = VpnProfile.TYPE_IPSEC_XAUTH_PSK;
+        final byte[] encodedProfile = profile.encode();
+        when(mKeyStore.get(Credentials.VPN + profileName)).thenReturn(encodedProfile);
+    }
+
+    @Test
+    public void testLegacyLockdownVpn() throws Exception {
+        final NetworkRequest request = new NetworkRequest.Builder().clearCapabilities().build();
+        final TestNetworkCallback callback = new TestNetworkCallback();
+        mCm.registerNetworkCallback(request, callback);
+
+        final TestNetworkCallback defaultCallback = new TestNetworkCallback();
+        mCm.registerDefaultNetworkCallback(defaultCallback);
+
+        // Pretend lockdown VPN was configured.
+        setupLegacyLockdownVpn();
+
+        // LockdownVpnTracker disables the Vpn teardown code and enables lockdown.
+        // Check the VPN's state before it does so.
+        assertTrue(mMockVpn.getEnableTeardown());
+        assertFalse(mMockVpn.getLockdown());
+
+        // Send a USER_UNLOCKED broadcast so CS starts LockdownVpnTracker.
+        final int userId = UserHandle.getUserId(Process.myUid());
+        final Intent addedIntent = new Intent(ACTION_USER_UNLOCKED);
+        addedIntent.putExtra(Intent.EXTRA_USER_HANDLE, userId);
+        final Handler handler = new Handler(mCsHandlerThread.getLooper());
+        handler.post(() -> mServiceContext.sendBroadcast(addedIntent));
+        waitForIdle();
+
+        // Lockdown VPN disables teardown and enables lockdown.
+        assertFalse(mMockVpn.getEnableTeardown());
+        assertTrue(mMockVpn.getLockdown());
+
+        // Bring up a network.
+        // Expect nothing to happen because the network does not have an IPv4 default route: legacy
+        // VPN only supports IPv4.
+        final LinkProperties cellLp = new LinkProperties();
+        cellLp.setInterfaceName("rmnet0");
+        cellLp.addLinkAddress(new LinkAddress("2001:db8::1/64"));
+        cellLp.addRoute(new RouteInfo(new IpPrefix("::/0"), null, "rmnet0"));
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR, cellLp);
+        mCellNetworkAgent.connect(false /* validated */);
+        callback.expectAvailableCallbacksUnvalidatedAndBlocked(mCellNetworkAgent);
+        defaultCallback.expectAvailableCallbacksUnvalidatedAndBlocked(mCellNetworkAgent);
+        waitForIdle();
+        assertNull(mMockVpn.getAgent());
+
+        // Add an IPv4 address. Ideally the VPN should start, but it doesn't because nothing calls
+        // LockdownVpnTracker#handleStateChangedLocked. This is a bug.
+        // TODO: consider fixing this.
+        cellLp.addLinkAddress(new LinkAddress("192.0.2.2/25"));
+        cellLp.addRoute(new RouteInfo(new IpPrefix("0.0.0.0/0"), null, "rmnet0"));
+        mCellNetworkAgent.sendLinkProperties(cellLp);
+        callback.expectCallback(CallbackEntry.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
+        defaultCallback.expectCallback(CallbackEntry.LINK_PROPERTIES_CHANGED, mCellNetworkAgent);
+        waitForIdle();
+        assertNull(mMockVpn.getAgent());
+
+        // Disconnect, then try again with a network that supports IPv4 at connection time.
+        // Expect lockdown VPN to come up.
+        ExpectedBroadcast b1 = expectConnectivityAction(TYPE_MOBILE, DetailedState.DISCONNECTED);
+        mCellNetworkAgent.disconnect();
+        callback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
+        defaultCallback.expectCallback(CallbackEntry.LOST, mCellNetworkAgent);
+        b1.expectBroadcast();
+
+        // When lockdown VPN is active, the NetworkInfo state in CONNECTIVITY_ACTION is overwritten
+        // with the state of the VPN network. So expect a CONNECTING broadcast.
+        b1 = expectConnectivityAction(TYPE_MOBILE, DetailedState.CONNECTING);
+        mCellNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_CELLULAR, cellLp);
+        mCellNetworkAgent.connect(false /* validated */);
+        callback.expectAvailableCallbacksUnvalidatedAndBlocked(mCellNetworkAgent);
+        defaultCallback.expectAvailableCallbacksUnvalidatedAndBlocked(mCellNetworkAgent);
+        b1.expectBroadcast();
+        assertActiveNetworkInfo(TYPE_MOBILE, DetailedState.BLOCKED);
+        assertNetworkInfo(TYPE_MOBILE, DetailedState.BLOCKED);
+        assertNetworkInfo(TYPE_WIFI, DetailedState.BLOCKED);
+        assertNetworkInfo(TYPE_VPN, DetailedState.BLOCKED);
+
+        // TODO: it would be nice if we could simply rely on the production code here, and have
+        // LockdownVpnTracker start the VPN, have the VPN code register its NetworkAgent with
+        // ConnectivityService, etc. That would require duplicating a fair bit of code from the
+        // Vpn tests around how to mock out LegacyVpnRunner. But even if we did that, this does not
+        // work for at least two reasons:
+        // 1. In this test, calling registerNetworkAgent does not actually result in an agent being
+        //    registered. This is because nothing calls onNetworkMonitorCreated, which is what
+        //    actually ends up causing handleRegisterNetworkAgent to be called. Code in this test
+        //    that wants to register an agent must use TestNetworkAgentWrapper.
+        // 2. Even if we exposed Vpn#agentConnect to the test, and made MockVpn#agentConnect call
+        //    the TestNetworkAgentWrapper code, this would deadlock because the
+        //    TestNetworkAgentWrapper code cannot be called on the handler thread since it calls
+        //    waitForIdle().
+        mMockVpn.expectStartLegacyVpnRunner();
+        b1 = expectConnectivityAction(TYPE_VPN, DetailedState.CONNECTED);
+        ExpectedBroadcast b2 = expectConnectivityAction(TYPE_MOBILE, DetailedState.CONNECTED);
+        mMockVpn.establishForMyUid();
+        callback.expectAvailableThenValidatedCallbacks(mMockVpn);
+        defaultCallback.expectAvailableThenValidatedCallbacks(mMockVpn);
+        b1.expectBroadcast();
+        b2.expectBroadcast();
+        assertActiveNetworkInfo(TYPE_MOBILE, DetailedState.CONNECTED);
+        assertNetworkInfo(TYPE_MOBILE, DetailedState.CONNECTED);
+        assertNetworkInfo(TYPE_WIFI, DetailedState.DISCONNECTED);
+        assertNetworkInfo(TYPE_VPN, DetailedState.CONNECTED);
+
+        // Switch default network from cell to wifi. Expect VPN to disconnect and reconnect.
+        final LinkProperties wifiLp = new LinkProperties();
+        wifiLp.setInterfaceName("wlan0");
+        wifiLp.addLinkAddress(new LinkAddress("192.0.2.163/25"));
+        wifiLp.addRoute(new RouteInfo(new IpPrefix("0.0.0.0/0"), null, "wlan0"));
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI, wifiLp);
+
+        b1 = expectConnectivityAction(TYPE_MOBILE, DetailedState.DISCONNECTED);
+        // Wifi is CONNECTING because the VPN isn't up yet.
+        b2 = expectConnectivityAction(TYPE_WIFI, DetailedState.CONNECTING);
+        ExpectedBroadcast b3 = expectConnectivityAction(TYPE_VPN, DetailedState.DISCONNECTED);
+        mWiFiNetworkAgent.connect(false /* validated */);
+        b1.expectBroadcast();
+        b2.expectBroadcast();
+        b3.expectBroadcast();
+        mMockVpn.expectStopVpnRunnerPrivileged();
+        mMockVpn.expectStartLegacyVpnRunner();
+
+        // TODO: why is wifi not blocked? Is this because something calls prepare()?
+        callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
+        callback.expectCapabilitiesThat(mMockVpn, (nc) -> nc.hasTransport(TRANSPORT_WIFI));
+        defaultCallback.expectCapabilitiesThat(mMockVpn, (nc) -> nc.hasTransport(TRANSPORT_WIFI));
+        callback.expectCallback(CallbackEntry.LOST, mMockVpn);
+        defaultCallback.expectCallback(CallbackEntry.LOST, mMockVpn);
+        defaultCallback.expectAvailableCallbacksUnvalidatedAndBlocked(mWiFiNetworkAgent);
+
+        // While the VPN is reconnecting on the new network, everything is blocked.
+        assertActiveNetworkInfo(TYPE_WIFI, DetailedState.BLOCKED);
+        assertNetworkInfo(TYPE_MOBILE, DetailedState.BLOCKED);
+        assertNetworkInfo(TYPE_WIFI, DetailedState.BLOCKED);
+        assertNetworkInfo(TYPE_VPN, DetailedState.BLOCKED);
+
+        // The VPN comes up again on wifi.
+        b1 = expectConnectivityAction(TYPE_VPN, DetailedState.CONNECTED);
+        b2 = expectConnectivityAction(TYPE_WIFI, DetailedState.CONNECTED);
+        mMockVpn.establishForMyUid();
+        callback.expectAvailableThenValidatedCallbacks(mMockVpn);
+        defaultCallback.expectAvailableThenValidatedCallbacks(mMockVpn);
+        b1.expectBroadcast();
+        b2.expectBroadcast();
+
+        assertActiveNetworkInfo(TYPE_WIFI, DetailedState.CONNECTED);
+        assertNetworkInfo(TYPE_MOBILE, DetailedState.DISCONNECTED);
+        assertNetworkInfo(TYPE_WIFI, DetailedState.CONNECTED);
+        assertNetworkInfo(TYPE_VPN, DetailedState.CONNECTED);
+
+        // 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();
+
+        assertActiveNetworkInfo(TYPE_WIFI, DetailedState.CONNECTED);
+        assertNetworkInfo(TYPE_MOBILE, DetailedState.DISCONNECTED);
+        assertNetworkInfo(TYPE_WIFI, DetailedState.CONNECTED);
+        assertNetworkInfo(TYPE_VPN, DetailedState.CONNECTED);
+
+        b1 = expectConnectivityAction(TYPE_WIFI, DetailedState.DISCONNECTED);
+        mWiFiNetworkAgent.disconnect();
+        callback.expectCallback(CallbackEntry.LOST, mWiFiNetworkAgent);
+        b1.expectBroadcast();
+        callback.expectCapabilitiesThat(mMockVpn, (nc) -> !nc.hasTransport(TRANSPORT_WIFI));
+        b2 = expectConnectivityAction(TYPE_VPN, DetailedState.DISCONNECTED);
+        mMockVpn.expectStopVpnRunnerPrivileged();
+        callback.expectCallback(CallbackEntry.LOST, mMockVpn);
+        b2.expectBroadcast();
+    }
+
     @Test
     public final void testLoseTrusted() throws Exception {
         final NetworkRequest trustedRequest = new NetworkRequest.Builder()
@@ -7604,7 +7838,7 @@
         lp.addRoute(new RouteInfo(new IpPrefix(Inet4Address.ANY, 0), null));
         lp.addRoute(new RouteInfo(new IpPrefix(Inet6Address.ANY, 0), RTN_UNREACHABLE));
         // The uid range needs to cover the test app so the network is visible to it.
-        final Set<UidRange> vpnRange = Collections.singleton(UidRange.createForUser(VPN_USER));
+        final Set<UidRange> vpnRange = Collections.singleton(UidRange.createForUser(PRIMARY_USER));
         mMockVpn.establish(lp, VPN_UID, vpnRange);
         assertVpnUidRangesUpdated(true, vpnRange, VPN_UID);
 
@@ -7632,7 +7866,7 @@
         lp.addRoute(new RouteInfo(new IpPrefix(Inet6Address.ANY, 0), null));
         lp.addRoute(new RouteInfo(new IpPrefix(Inet4Address.ANY, 0), null));
         // The uid range needs to cover the test app so the network is visible to it.
-        final Set<UidRange> vpnRange = Collections.singleton(UidRange.createForUser(VPN_USER));
+        final Set<UidRange> vpnRange = Collections.singleton(UidRange.createForUser(PRIMARY_USER));
         mMockVpn.establish(lp, Process.SYSTEM_UID, vpnRange);
         assertVpnUidRangesUpdated(true, vpnRange, Process.SYSTEM_UID);
 
@@ -7648,7 +7882,7 @@
         lp.addRoute(new RouteInfo(new IpPrefix("192.0.2.0/24"), null, "tun0"));
         lp.addRoute(new RouteInfo(new IpPrefix(Inet6Address.ANY, 0), RTN_UNREACHABLE));
         // The uid range needs to cover the test app so the network is visible to it.
-        final Set<UidRange> vpnRange = Collections.singleton(UidRange.createForUser(VPN_USER));
+        final Set<UidRange> vpnRange = Collections.singleton(UidRange.createForUser(PRIMARY_USER));
         mMockVpn.establish(lp, Process.SYSTEM_UID, vpnRange);
         assertVpnUidRangesUpdated(true, vpnRange, Process.SYSTEM_UID);
 
@@ -7663,7 +7897,7 @@
         lp.addRoute(new RouteInfo(new IpPrefix(Inet4Address.ANY, 0), null));
         lp.addRoute(new RouteInfo(new IpPrefix(Inet6Address.ANY, 0), null));
         // The uid range needs to cover the test app so the network is visible to it.
-        final Set<UidRange> vpnRange = Collections.singleton(UidRange.createForUser(VPN_USER));
+        final Set<UidRange> vpnRange = Collections.singleton(UidRange.createForUser(PRIMARY_USER));
         mMockVpn.establish(lp, VPN_UID, vpnRange);
         assertVpnUidRangesUpdated(true, vpnRange, VPN_UID);
 
@@ -7715,7 +7949,7 @@
         lp.addRoute(new RouteInfo(new IpPrefix(Inet4Address.ANY, 0), RTN_UNREACHABLE));
         lp.addRoute(new RouteInfo(new IpPrefix(Inet6Address.ANY, 0), null));
         // The uid range needs to cover the test app so the network is visible to it.
-        final UidRange vpnRange = UidRange.createForUser(VPN_USER);
+        final UidRange vpnRange = UidRange.createForUser(PRIMARY_USER);
         final Set<UidRange> vpnRanges = Collections.singleton(vpnRange);
         mMockVpn.establish(lp, VPN_UID, vpnRanges);
         assertVpnUidRangesUpdated(true, vpnRanges, VPN_UID);
@@ -7867,7 +8101,7 @@
 
     private void setupConnectionOwnerUid(int vpnOwnerUid, @VpnManager.VpnType int vpnType)
             throws Exception {
-        final Set<UidRange> vpnRange = Collections.singleton(UidRange.createForUser(VPN_USER));
+        final Set<UidRange> vpnRange = Collections.singleton(UidRange.createForUser(PRIMARY_USER));
         mMockVpn.establish(new LinkProperties(), vpnOwnerUid, vpnRange);
         assertVpnUidRangesUpdated(true, vpnRange, vpnOwnerUid);
         mMockVpn.setVpnType(vpnType);
@@ -8441,7 +8675,7 @@
         lp.setInterfaceName("tun0");
         lp.addRoute(new RouteInfo(new IpPrefix(Inet4Address.ANY, 0), null));
         lp.addRoute(new RouteInfo(new IpPrefix(Inet6Address.ANY, 0), null));
-        final UidRange vpnRange = UidRange.createForUser(VPN_USER);
+        final UidRange vpnRange = UidRange.createForUser(PRIMARY_USER);
         Set<UidRange> vpnRanges = Collections.singleton(vpnRange);
         mMockVpn.establish(lp, VPN_UID, vpnRanges);
         assertVpnUidRangesUpdated(true, vpnRanges, VPN_UID);