Merge "Lock down networking when waiting for always-on" into nyc-dev

core/java/android/os/INetworkManagementService.aidl

@@ -436,4 +436,6 @@ interface INetworkManagementService
 
     void addInterfaceToLocalNetwork(String iface, in List<RouteInfo> routes);
     void removeInterfaceFromLocalNetwork(String iface);
+
+    void setAllowOnlyVpnForUids(boolean enable, in UidRange[] uidRanges);
 }

core/java/android/provider/Settings.java

@@ -4701,6 +4701,14 @@ public final class Settings {
          */
         public static final String ALWAYS_ON_VPN_APP = "always_on_vpn_app";
 
+        /**
+         * Whether to block networking outside of VPN connections while always-on is set.
+         * @see #ALWAYS_ON_VPN_APP
+         *
+         * @hide
+         */
+        public static final String ALWAYS_ON_VPN_LOCKDOWN = "always_on_vpn_lockdown";
+
         /**
          * Whether applications can be installed for this user via the system's
          * {@link Intent#ACTION_INSTALL_PACKAGE} mechanism.

services/core/java/com/android/server/ConnectivityService.java

@@ -915,6 +915,13 @@ public class ConnectivityService extends IConnectivityManager.Stub
         final boolean networkMetered;
         final int uidRules;
 
+        synchronized (mVpns) {
+            final Vpn vpn = mVpns.get(UserHandle.getUserId(uid));
+            if (vpn != null && vpn.isBlockingUid(uid)) {
+                return true;
+            }
+        }
+
         final String iface = (lp == null ? "" : lp.getInterfaceName());
         synchronized (mRulesLock) {
             networkMetered = mMeteredIfaces.contains(iface);
@@ -3365,23 +3372,42 @@ public class ConnectivityService extends IConnectivityManager.Stub
     }
 
     /**
-     * Sets up or tears down the always-on VPN for user {@param user} as appropriate.
+     * Starts the always-on VPN {@link VpnService} for user {@param userId}, which should perform
+     * some setup and then call {@code establish()} to connect.
      *
-     * @return {@code false} in case of errors; {@code true} otherwise.
+     * @return {@code true} if the service was started, the service was already connected, or there
+     *         was no always-on VPN to start. {@code false} otherwise.
      */
-    private boolean updateAlwaysOnVpn(int user) {
+    private boolean startAlwaysOnVpn(int userId) {
-        final String lockdownPackage = getAlwaysOnVpnPackage(user);
+        final String alwaysOnPackage;
-        if (lockdownPackage == null) {
+        synchronized (mVpns) {
-            return true;
+            Vpn vpn = mVpns.get(userId);
+            if (vpn == null) {
+                // Shouldn't happen as all codepaths that point here should have checked the Vpn
+                // exists already.
+                Slog.wtf(TAG, "User " + userId + " has no Vpn configuration");
+                return false;
+            }
+            alwaysOnPackage = vpn.getAlwaysOnPackage();
+            // Skip if there is no service to start.
+            if (alwaysOnPackage == null) {
+                return true;
+            }
+            // Skip if the service is already established. This isn't bulletproof: it's not bound
+            // until after establish(), so if it's mid-setup onStartCommand will be sent twice,
+            // which may restart the connection.
+            if (vpn.getNetworkInfo().isConnected()) {
+                return true;
+            }
         }
 
-        // Create an intent to start the VPN service declared in the app's manifest.
+        // Start the VPN service declared in the app's manifest.
         Intent serviceIntent = new Intent(VpnConfig.SERVICE_INTERFACE);
-        serviceIntent.setPackage(lockdownPackage);
+        serviceIntent.setPackage(alwaysOnPackage);
-
         try {
-            return mContext.startServiceAsUser(serviceIntent, UserHandle.of(user)) != null;
+            return mContext.startServiceAsUser(serviceIntent, UserHandle.of(userId)) != null;
         } catch (RuntimeException e) {
+            Slog.w(TAG, "VpnService " + serviceIntent + " failed to start", e);
             return false;
         }
     }
@@ -3396,25 +3422,35 @@ public class ConnectivityService extends IConnectivityManager.Stub
             return false;
         }
 
-        // If the current VPN package is the same as the new one, this is a no-op
-        final String oldPackage = getAlwaysOnVpnPackage(userId);
-        if (TextUtils.equals(oldPackage, packageName)) {
-            return true;
-        }
-
         synchronized (mVpns) {
             Vpn vpn = mVpns.get(userId);
             if (vpn == null) {
                 Slog.w(TAG, "User " + userId + " has no Vpn configuration");
                 return false;
             }
-            if (!vpn.setAlwaysOnPackage(packageName)) {
+            // If the current VPN package is the same as the new one, this is a no-op
+            if (TextUtils.equals(packageName, vpn.getAlwaysOnPackage())) {
+                return true;
+            }
+            if (!vpn.setAlwaysOnPackage(packageName, lockdown)) {
                 return false;
             }
-            if (!updateAlwaysOnVpn(userId)) {
+            if (!startAlwaysOnVpn(userId)) {
-                vpn.setAlwaysOnPackage(null);
+                vpn.setAlwaysOnPackage(null, false);
                 return false;
             }
+
+            // Save the configuration
+            final long token = Binder.clearCallingIdentity();
+            try {
+                final ContentResolver cr = mContext.getContentResolver();
+                Settings.Secure.putStringForUser(cr, Settings.Secure.ALWAYS_ON_VPN_APP,
+                        packageName, userId);
+                Settings.Secure.putIntForUser(cr, Settings.Secure.ALWAYS_ON_VPN_LOCKDOWN,
+                        (lockdown ? 1 : 0), userId);
+            } finally {
+                Binder.restoreCallingIdentity(token);
+            }
         }
         return true;
     }
@@ -3685,11 +3721,18 @@ public class ConnectivityService extends IConnectivityManager.Stub
             }
             userVpn = new Vpn(mHandler.getLooper(), mContext, mNetd, userId);
             mVpns.put(userId, userVpn);
+
+            final ContentResolver cr = mContext.getContentResolver();
+            String alwaysOnPackage = Settings.Secure.getStringForUser(cr,
+                    Settings.Secure.ALWAYS_ON_VPN_APP, userId);
+            final boolean alwaysOnLockdown = Settings.Secure.getIntForUser(cr,
+                    Settings.Secure.ALWAYS_ON_VPN_LOCKDOWN, /* default */ 0, userId) != 0;
+            if (alwaysOnPackage != null) {
+                userVpn.setAlwaysOnPackage(alwaysOnPackage, alwaysOnLockdown);
+            }
         }
         if (mUserManager.getUserInfo(userId).isPrimary() && LockdownVpnTracker.isEnabled()) {
             updateLockdownVpn();
-        } else {
-            updateAlwaysOnVpn(userId);
         }
     }
 
@@ -3700,6 +3743,7 @@ public class ConnectivityService extends IConnectivityManager.Stub
                 loge("Stopped user has no VPN");
                 return;
             }
+            userVpn.onUserStopped();
             mVpns.delete(userId);
         }
     }
@@ -3729,7 +3773,7 @@ public class ConnectivityService extends IConnectivityManager.Stub
         if (mUserManager.getUserInfo(userId).isPrimary() && LockdownVpnTracker.isEnabled()) {
             updateLockdownVpn();
         } else {
-            updateAlwaysOnVpn(userId);
+            startAlwaysOnVpn(userId);
         }
     }
 

services/core/java/com/android/server/NetworkManagementService.java

@@ -1849,6 +1849,22 @@ public class NetworkManagementService extends INetworkManagementService.Stub
         }
     }
 
+    @Override
+    public void setAllowOnlyVpnForUids(boolean add, UidRange[] uidRanges)
+            throws ServiceSpecificException {
+        try {
+            mNetdService.networkRejectNonSecureVpn(add, uidRanges);
+        } catch (ServiceSpecificException e) {
+            Log.w(TAG, "setAllowOnlyVpnForUids(" + add + ", " + Arrays.toString(uidRanges) + ")"
+                    + ": netd command failed", e);
+            throw e;
+        } catch (RemoteException e) {
+            Log.w(TAG, "setAllowOnlyVpnForUids(" + add + ", " + Arrays.toString(uidRanges) + ")"
+                    + ": netd command failed", e);
+            throw e.rethrowAsRuntimeException();
+        }
+    }
+
     @Override
     public void setUidCleartextNetworkPolicy(int uid, int policy) {
         if (Binder.getCallingUid() != uid) {

services/core/java/com/android/server/connectivity/Vpn.java

@@ -66,7 +66,6 @@ import android.os.SystemClock;
 import android.os.SystemService;
 import android.os.UserHandle;
 import android.os.UserManager;
-import android.provider.Settings;
 import android.security.Credentials;
 import android.security.KeyStore;
 import android.text.TextUtils;
@@ -127,6 +126,19 @@ public class Vpn {
     private final Looper mLooper;
     private final NetworkCapabilities mNetworkCapabilities;
 
+    /**
+     * Whether to keep the connection active after rebooting, or upgrading or reinstalling. This
+     * only applies to {@link VpnService} connections.
+     */
+    private boolean mAlwaysOn = false;
+
+    /**
+     * Whether to disable traffic outside of this VPN even when the VPN is not connected. System
+     * apps can still bypass by choosing explicit networks. Has no effect if {@link mAlwaysOn} is
+     * not set.
+     */
+    private boolean mLockdown = false;
+
     /**
      * List of UIDs that are set to use this VPN by default. Normally, every UID in the user is
      * added to this set but that can be changed by adding allowed or disallowed applications. It
@@ -140,6 +152,14 @@ public class Vpn {
     @GuardedBy("this")
     private Set<UidRange> mVpnUsers = null;
 
+    /**
+     * List of UIDs for which networking should be blocked until VPN is ready, during brief periods
+     * when VPN is not running. For example, during system startup or after a crash.
+     * @see mLockdown
+     */
+    @GuardedBy("this")
+    private Set<UidRange> mBlockedUsers = new ArraySet<>();
+
     // Handle of user initiating VPN.
     private final int mUserHandle;
 
@@ -194,9 +214,10 @@ public class Vpn {
      *    manifest guarded by {@link android.Manifest.permission.BIND_VPN_SERVICE},
      *    otherwise the call will fail.
      *
-     * @param newPackage the package to designate as always-on VPN supplier.
+     * @param packageName the package to designate as always-on VPN supplier.
+     * @param lockdown whether to prevent traffic outside of a VPN, for example while connecting.
      */
-    public synchronized boolean setAlwaysOnPackage(String packageName) {
+    public synchronized boolean setAlwaysOnPackage(String packageName, boolean lockdown) {
         enforceControlPermissionOrInternalCaller();
 
         // Disconnect current VPN.
@@ -210,14 +231,9 @@ public class Vpn {
             prepareInternal(packageName);
         }
 
-        // Save the new package name in Settings.Secure.
+        mAlwaysOn = (packageName != null);
-        final long token = Binder.clearCallingIdentity();
+        mLockdown = (mAlwaysOn && lockdown);
-        try {
+        setVpnForcedLocked(mLockdown);
-            Settings.Secure.putStringForUser(mContext.getContentResolver(),
-                    Settings.Secure.ALWAYS_ON_VPN_APP, packageName, mUserHandle);
-        } finally {
-            Binder.restoreCallingIdentity(token);
-        }
         return true;
     }
 
@@ -229,14 +245,7 @@ public class Vpn {
      */
     public synchronized String getAlwaysOnPackage() {
         enforceControlPermissionOrInternalCaller();
-
+        return (mAlwaysOn ? mPackage : null);
-        final long token = Binder.clearCallingIdentity();
-        try {
-            return Settings.Secure.getStringForUser(mContext.getContentResolver(),
-                    Settings.Secure.ALWAYS_ON_VPN_APP, mUserHandle);
-        } finally {
-            Binder.restoreCallingIdentity(token);
-        }
     }
 
     /**
@@ -258,6 +267,11 @@ public class Vpn {
      * @return true if the operation is succeeded.
      */
     public synchronized boolean prepare(String oldPackage, String newPackage) {
+        // Stop an existing always-on VPN from being dethroned by other apps.
+        if (mAlwaysOn && !TextUtils.equals(mPackage, newPackage)) {
+            return false;
+        }
+
         if (oldPackage != null) {
             if (getAppUid(oldPackage, mUserHandle) != mOwnerUID) {
                 // The package doesn't match. We return false (to obtain user consent) unless the
@@ -281,11 +295,6 @@ public class Vpn {
             return true;
         }
 
-        // Stop an existing always-on VPN from being dethroned by other apps.
-        if (getAlwaysOnPackage() != null) {
-            return false;
-        }
-
         // Check that the caller is authorized.
         enforceControlPermission();
 
@@ -469,7 +478,7 @@ public class Vpn {
         mNetworkInfo.setDetailedState(DetailedState.CONNECTING, null, null);
 
         NetworkMisc networkMisc = new NetworkMisc();
-        networkMisc.allowBypass = mConfig.allowBypass;
+        networkMisc.allowBypass = mConfig.allowBypass && !mLockdown;
 
         long token = Binder.clearCallingIdentity();
         try {
@@ -685,7 +694,7 @@ public class Vpn {
             final long token = Binder.clearCallingIdentity();
             List<UserInfo> users;
             try {
-                users = UserManager.get(mContext).getUsers();
+                users = UserManager.get(mContext).getUsers(true);
             } finally {
                 Binder.restoreCallingIdentity(token);
             }
@@ -774,18 +783,22 @@ public class Vpn {
     public void onUserAdded(int userHandle) {
         // If the user is restricted tie them to the parent user's VPN
         UserInfo user = UserManager.get(mContext).getUserInfo(userHandle);
-        if (user.isRestricted() && user.restrictedProfileParentId == mUserHandle
+        if (user.isRestricted() && user.restrictedProfileParentId == mUserHandle) {
-                && mVpnUsers != null) {
             synchronized(Vpn.this) {
-                try {
+                if (mVpnUsers != null) {
-                    addUserToRanges(mVpnUsers, userHandle, mConfig.allowedApplications,
+                    try {
-                            mConfig.disallowedApplications);
+                        addUserToRanges(mVpnUsers, userHandle, mConfig.allowedApplications,
-                    if (mNetworkAgent != null) {
+                                mConfig.disallowedApplications);
-                        final List<UidRange> ranges = uidRangesForUser(userHandle);
+                        if (mNetworkAgent != null) {
-                        mNetworkAgent.addUidRanges(ranges.toArray(new UidRange[ranges.size()]));
+                            final List<UidRange> ranges = uidRangesForUser(userHandle);
+                            mNetworkAgent.addUidRanges(ranges.toArray(new UidRange[ranges.size()]));
+                        }
+                    } catch (Exception e) {
+                        Log.wtf(TAG, "Failed to add restricted user to owner", e);
                     }
-                } catch (Exception e) {
+                }
-                    Log.wtf(TAG, "Failed to add restricted user to owner", e);
+                if (mAlwaysOn) {
+                    setVpnForcedLocked(mLockdown);
                 }
             }
         }
@@ -794,18 +807,100 @@ public class Vpn {
     public void onUserRemoved(int userHandle) {
         // clean up if restricted
         UserInfo user = UserManager.get(mContext).getUserInfo(userHandle);
-        if (user.isRestricted() && user.restrictedProfileParentId == mUserHandle
+        if (user.isRestricted() && user.restrictedProfileParentId == mUserHandle) {
-                && mVpnUsers != null) {
             synchronized(Vpn.this) {
-                try {
+                if (mVpnUsers != null) {
-                    removeVpnUserLocked(userHandle);
+                    try {
-                } catch (Exception e) {
+                        removeVpnUserLocked(userHandle);
-                    Log.wtf(TAG, "Failed to remove restricted user to owner", e);
+                    } catch (Exception e) {
+                        Log.wtf(TAG, "Failed to remove restricted user to owner", e);
+                    }
+                }
+                if (mAlwaysOn) {
+                    setVpnForcedLocked(mLockdown);
                 }
             }
         }
     }
 
+    /**
+     * Called when the user associated with this VPN has just been stopped.
+     */
+    public synchronized void onUserStopped() {
+        // Switch off networking lockdown (if it was enabled)
+        setVpnForcedLocked(false);
+        mAlwaysOn = false;
+
+        // Quit any active connections
+        agentDisconnect();
+    }
+
+    /**
+     * Restrict network access from all UIDs affected by this {@link Vpn}, apart from the VPN
+     * service app itself, to only sockets that have had {@code protect()} called on them. All
+     * non-VPN traffic is blocked via a {@code PROHIBIT} response from the kernel.
+     *
+     * The exception for the VPN UID isn't technically necessary -- setup should use protected
+     * sockets -- but in practice it saves apps that don't protect their sockets from breaking.
+     *
+     * Calling multiple times with {@param enforce} = {@code true} will recreate the set of UIDs to
+     * block every time, and if anything has changed update using {@link #setAllowOnlyVpnForUids}.
+     *
+     * @param enforce {@code true} to require that all traffic under the jurisdiction of this
+     *                {@link Vpn} goes through a VPN connection or is blocked until one is
+     *                available, {@code false} to lift the requirement.
+     *
+     * @see #mBlockedUsers
+     */
+    @GuardedBy("this")
+    private void setVpnForcedLocked(boolean enforce) {
+        final Set<UidRange> removedRanges = new ArraySet<>(mBlockedUsers);
+        if (enforce) {
+            final Set<UidRange> addedRanges = createUserAndRestrictedProfilesRanges(mUserHandle,
+                    /* allowedApplications */ null,
+                    /* disallowedApplications */ Collections.singletonList(mPackage));
+
+            removedRanges.removeAll(addedRanges);
+            addedRanges.removeAll(mBlockedUsers);
+
+            setAllowOnlyVpnForUids(false, removedRanges);
+            setAllowOnlyVpnForUids(true, addedRanges);
+        } else {
+            setAllowOnlyVpnForUids(false, removedRanges);
+        }
+    }
+
+    /**
+     * Either add or remove a list of {@link UidRange}s to the list of UIDs that are only allowed
+     * to make connections through sockets that have had {@code protect()} called on them.
+     *
+     * @param enforce {@code true} to add to the blacklist, {@code false} to remove.
+     * @param ranges {@link Collection} of {@link UidRange}s to add (if {@param enforce} is
+     *               {@code true}) or to remove.
+     * @return {@code true} if all of the UIDs were added/removed. {@code false} otherwise,
+     *         including added ranges that already existed or removed ones that didn't.
+     */
+    @GuardedBy("this")
+    private boolean setAllowOnlyVpnForUids(boolean enforce, Collection<UidRange> ranges) {
+        if (ranges.size() == 0) {
+            return true;
+        }
+        final UidRange[] rangesArray = ranges.toArray(new UidRange[ranges.size()]);
+        try {
+            mNetd.setAllowOnlyVpnForUids(enforce, rangesArray);
+        } catch (RemoteException | RuntimeException e) {
+            Log.e(TAG, "Updating blocked=" + enforce
+                    + " for UIDs " + Arrays.toString(ranges.toArray()) + " failed", e);
+            return false;
+        }
+        if (enforce) {
+            mBlockedUsers.addAll(ranges);
+        } else {
+            mBlockedUsers.removeAll(ranges);
+        }
+        return true;
+    }
+
     /**
      * Return the configuration of the currently running VPN.
      */
@@ -959,6 +1054,21 @@ public class Vpn {
         return false;
     }
 
+    /**
+     * @return {@code true} if the set of users blocked whilst waiting for VPN to connect includes
+     *         the UID {@param uid}, {@code false} otherwise.
+     *
+     * @see #mBlockedUsers
+     */
+    public synchronized boolean isBlockingUid(int uid) {
+        for (UidRange uidRange : mBlockedUsers) {
+            if (uidRange.contains(uid)) {
+                return true;
+            }
+        }
+        return false;
+    }
+
     private native int jniCreate(int mtu);
     private native String jniGetName(int tun);
     private native int jniSetAddresses(String interfaze, String addresses);

services/tests/servicestests/src/com/android/server/connectivity/VpnTest.java

@@ -20,9 +20,11 @@ import static android.content.pm.UserInfo.FLAG_ADMIN;
 import static android.content.pm.UserInfo.FLAG_MANAGED_PROFILE;
 import static android.content.pm.UserInfo.FLAG_PRIMARY;
 import static android.content.pm.UserInfo.FLAG_RESTRICTED;
+import static org.mockito.AdditionalMatchers.*;
 import static org.mockito.Mockito.*;
 
 import android.annotation.UserIdInt;
+import android.app.AppOpsManager;
 import android.content.Context;
 import android.content.pm.PackageManager;
 import android.content.pm.UserInfo;
@@ -65,16 +67,35 @@ public class VpnTest extends AndroidTestCase {
         managedProfileA.profileGroupId = primaryUser.id;
     }
 
+    /**
+     * Names and UIDs for some fake packages. Important points:
+     *  - UID is ordered increasing.
+     *  - One pair of packages have consecutive UIDs.
+     */
+    static final String[] PKGS = {"com.example", "org.example", "net.example", "web.vpn"};
+    static final int[] PKG_UIDS = {66, 77, 78, 400};
+
+    // Mock packages
+    static final Map<String, Integer> mPackages = new ArrayMap<>();
+    static {
+        for (int i = 0; i < PKGS.length; i++) {
+            mPackages.put(PKGS[i], PKG_UIDS[i]);
+        }
+    }
+
     @Mock private Context mContext;
     @Mock private UserManager mUserManager;
     @Mock private PackageManager mPackageManager;
     @Mock private INetworkManagementService mNetService;
+    @Mock private AppOpsManager mAppOps;
 
     @Override
     public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
         when(mContext.getPackageManager()).thenReturn(mPackageManager);
+        setMockedPackages(mPackages);
         when(mContext.getSystemService(eq(Context.USER_SERVICE))).thenReturn(mUserManager);
+        when(mContext.getSystemService(eq(Context.APP_OPS_SERVICE))).thenReturn(mAppOps);
         doNothing().when(mNetService).registerObserver(any());
     }
 
@@ -82,7 +103,7 @@ public class VpnTest extends AndroidTestCase {
     public void testRestrictedProfilesAreAddedToVpn() {
         setMockedUsers(primaryUser, secondaryUser, restrictedProfileA, restrictedProfileB);
 
-        final Vpn vpn = createVpn(primaryUser.id);
+        final Vpn vpn = new MockVpn(primaryUser.id);
         final Set<UidRange> ranges = vpn.createUserAndRestrictedProfilesRanges(primaryUser.id,
                 null, null);
 
@@ -96,7 +117,7 @@ public class VpnTest extends AndroidTestCase {
     public void testManagedProfilesAreNotAddedToVpn() {
         setMockedUsers(primaryUser, managedProfileA);
 
-        final Vpn vpn = createVpn(primaryUser.id);
+        final Vpn vpn = new MockVpn(primaryUser.id);
         final Set<UidRange> ranges = vpn.createUserAndRestrictedProfilesRanges(primaryUser.id,
                 null, null);
 
@@ -109,7 +130,7 @@ public class VpnTest extends AndroidTestCase {
     public void testAddUserToVpnOnlyAddsOneUser() {
         setMockedUsers(primaryUser, restrictedProfileA, managedProfileA);
 
-        final Vpn vpn = createVpn(primaryUser.id);
+        final Vpn vpn = new MockVpn(primaryUser.id);
         final Set<UidRange> ranges = new ArraySet<>();
         vpn.addUserToRanges(ranges, primaryUser.id, null, null);
 
@@ -120,42 +141,123 @@ public class VpnTest extends AndroidTestCase {
 
     @SmallTest
     public void testUidWhiteAndBlacklist() throws Exception {
-        final Map<String, Integer> packages = new ArrayMap<>();
+        final Vpn vpn = new MockVpn(primaryUser.id);
-        packages.put("com.example", 66);
-        packages.put("org.example", 77);
-        packages.put("net.example", 78);
-        setMockedPackages(packages);
-
-        final Vpn vpn = createVpn(primaryUser.id);
         final UidRange user = UidRange.createForUser(primaryUser.id);
+        final String[] packages = {PKGS[0], PKGS[1], PKGS[2]};
 
         // Whitelist
         final Set<UidRange> allow = vpn.createUserAndRestrictedProfilesRanges(primaryUser.id,
-                new ArrayList<String>(packages.keySet()), null);
+                Arrays.asList(packages), null);
         assertEquals(new ArraySet<>(Arrays.asList(new UidRange[] {
-            new UidRange(user.start + 66, user.start + 66),
+            new UidRange(user.start + PKG_UIDS[0], user.start + PKG_UIDS[0]),
-            new UidRange(user.start + 77, user.start + 78)
+            new UidRange(user.start + PKG_UIDS[1], user.start + PKG_UIDS[2])
         })), allow);
 
         // Blacklist
         final Set<UidRange> disallow = vpn.createUserAndRestrictedProfilesRanges(primaryUser.id,
-                null, new ArrayList<String>(packages.keySet()));
+                null, Arrays.asList(packages));
         assertEquals(new ArraySet<>(Arrays.asList(new UidRange[] {
-            new UidRange(user.start, user.start + 65),
+            new UidRange(user.start, user.start + PKG_UIDS[0] - 1),
-            new UidRange(user.start + 67, user.start + 76),
+            new UidRange(user.start + PKG_UIDS[0] + 1, user.start + PKG_UIDS[1] - 1),
-            new UidRange(user.start + 79, user.stop)
+            /* Empty range between UIDS[1] and UIDS[2], should be excluded, */
+            new UidRange(user.start + PKG_UIDS[2] + 1, user.stop)
         })), disallow);
     }
 
+    @SmallTest
+    public void testLockdownChangingPackage() throws Exception {
+        final MockVpn vpn = new MockVpn(primaryUser.id);
+        final UidRange user = UidRange.createForUser(primaryUser.id);
+
+        // Default state.
+        vpn.assertUnblocked(user.start + PKG_UIDS[0], user.start + PKG_UIDS[1], user.start + PKG_UIDS[2], user.start + PKG_UIDS[3]);
+
+        // Set always-on without lockdown.
+        assertTrue(vpn.setAlwaysOnPackage(PKGS[1], false));
+        vpn.assertUnblocked(user.start + PKG_UIDS[0], user.start + PKG_UIDS[1], user.start + PKG_UIDS[2], user.start + PKG_UIDS[3]);
+
+        // Set always-on with lockdown.
+        assertTrue(vpn.setAlwaysOnPackage(PKGS[1], true));
+        verify(mNetService).setAllowOnlyVpnForUids(eq(true), aryEq(new UidRange[] {
+            new UidRange(user.start, user.start + PKG_UIDS[1] - 1),
+            new UidRange(user.start + PKG_UIDS[1] + 1, user.stop)
+        }));
+        vpn.assertBlocked(user.start + PKG_UIDS[0], user.start + PKG_UIDS[2], user.start + PKG_UIDS[3]);
+        vpn.assertUnblocked(user.start + PKG_UIDS[1]);
+
+        // Switch to another app.
+        assertTrue(vpn.setAlwaysOnPackage(PKGS[3], true));
+        verify(mNetService).setAllowOnlyVpnForUids(eq(false), aryEq(new UidRange[] {
+            new UidRange(user.start, user.start + PKG_UIDS[1] - 1),
+            new UidRange(user.start + PKG_UIDS[1] + 1, user.stop)
+        }));
+        verify(mNetService).setAllowOnlyVpnForUids(eq(true), aryEq(new UidRange[] {
+            new UidRange(user.start, user.start + PKG_UIDS[3] - 1),
+            new UidRange(user.start + PKG_UIDS[3] + 1, user.stop)
+        }));
+        vpn.assertBlocked(user.start + PKG_UIDS[0], user.start + PKG_UIDS[1], user.start + PKG_UIDS[2]);
+        vpn.assertUnblocked(user.start + PKG_UIDS[3]);
+    }
+
+    @SmallTest
+    public void testLockdownAddingAProfile() throws Exception {
+        final MockVpn vpn = new MockVpn(primaryUser.id);
+        setMockedUsers(primaryUser);
+
+        // Make a copy of the restricted profile, as we're going to mark it deleted halfway through.
+        final UserInfo tempProfile = new UserInfo(restrictedProfileA.id, restrictedProfileA.name,
+                restrictedProfileA.flags);
+        tempProfile.restrictedProfileParentId = primaryUser.id;
+
+        final UidRange user = UidRange.createForUser(primaryUser.id);
+        final UidRange profile = UidRange.createForUser(tempProfile.id);
+
+        // Set lockdown.
+        assertTrue(vpn.setAlwaysOnPackage(PKGS[3], true));
+        verify(mNetService).setAllowOnlyVpnForUids(eq(true), aryEq(new UidRange[] {
+            new UidRange(user.start, user.start + PKG_UIDS[3] - 1),
+            new UidRange(user.start + PKG_UIDS[3] + 1, user.stop)
+        }));
+
+        // Verify restricted user isn't affected at first.
+        vpn.assertUnblocked(profile.start + PKG_UIDS[0]);
+
+        // Add the restricted user.
+        setMockedUsers(primaryUser, tempProfile);
+        vpn.onUserAdded(tempProfile.id);
+        verify(mNetService).setAllowOnlyVpnForUids(eq(true), aryEq(new UidRange[] {
+            new UidRange(profile.start, profile.start + PKG_UIDS[3] - 1),
+            new UidRange(profile.start + PKG_UIDS[3] + 1, profile.stop)
+        }));
+
+        // Remove the restricted user.
+        tempProfile.partial = true;
+        vpn.onUserRemoved(tempProfile.id);
+        verify(mNetService).setAllowOnlyVpnForUids(eq(false), aryEq(new UidRange[] {
+            new UidRange(profile.start, profile.start + PKG_UIDS[3] - 1),
+            new UidRange(profile.start + PKG_UIDS[3] + 1, profile.stop)
+        }));
+    }
+
     /**
-     * @return A subclass of {@link Vpn} which is reliably:
+     * A subclass of {@link Vpn} with some of the fields pre-mocked.
-     * <ul>
-     *   <li>Associated with a specific user ID</li>
-     *   <li>Not in always-on mode</li>
-     * </ul>
      */
-    private Vpn createVpn(@UserIdInt int userId) {
+    private class MockVpn extends Vpn {
-        return new Vpn(Looper.myLooper(), mContext, mNetService, userId);
+        public MockVpn(@UserIdInt int userId) {
+            super(Looper.myLooper(), mContext, mNetService, userId);
+        }
+
+        public void assertBlocked(int... uids) {
+            for (int uid : uids) {
+                assertTrue("Uid " + uid + " should be blocked", isBlockingUid(uid));
+            }
+        }
+
+        public void assertUnblocked(int... uids) {
+            for (int uid : uids) {
+                assertFalse("Uid " + uid + " should not be blocked", isBlockingUid(uid));
+            }
+        }
     }
 
     /**
@@ -167,9 +269,19 @@ public class VpnTest extends AndroidTestCase {
             userMap.put(user.id, user);
         }
 
+        /**
+         * @see UserManagerService#getUsers(boolean)
+         */
         doAnswer(invocation -> {
-            return new ArrayList(userMap.values());
+            final boolean excludeDying = (boolean) invocation.getArguments()[0];
-        }).when(mUserManager).getUsers();
+            final ArrayList<UserInfo> result = new ArrayList<>(users.length);
+            for (UserInfo ui : users) {
+                if (!excludeDying || (ui.isEnabled() && !ui.partial)) {
+                    result.add(ui);
+                }
+            }
+            return result;
+        }).when(mUserManager).getUsers(anyBoolean());
 
         doAnswer(invocation -> {
             final int id = (int) invocation.getArguments()[0];