diff --git a/core/java/android/net/ConnectivityManager.java b/core/java/android/net/ConnectivityManager.java
index 43d9bf3..a45e6f5 100644
--- a/core/java/android/net/ConnectivityManager.java
+++ b/core/java/android/net/ConnectivityManager.java
@@ -3362,6 +3362,10 @@
     /**
      * Determines if the calling application is subject to metered network restrictions while
      * running on background.
+     *
+     * @return {@link #RESTRICT_BACKGROUND_STATUS_DISABLED},
+     * {@link #RESTRICT_BACKGROUND_STATUS_ENABLED},
+     * or {@link #RESTRICT_BACKGROUND_STATUS_WHITELISTED}
      */
     public @RestrictBackgroundStatus int getRestrictBackgroundStatus() {
         try {
diff --git a/core/java/android/net/NetworkInfo.java b/core/java/android/net/NetworkInfo.java
index b6c5c6f..42f5feb 100644
--- a/core/java/android/net/NetworkInfo.java
+++ b/core/java/android/net/NetworkInfo.java
@@ -334,6 +334,7 @@
      *
      * @return {@code true} if large transfers should be avoided, otherwise
      *         {@code false}.
+     * @hide
      */
     public boolean isMetered() {
         synchronized (this) {
diff --git a/core/java/android/net/StaticIpConfiguration.java b/core/java/android/net/StaticIpConfiguration.java
index 7f1b179..58b1b88 100644
--- a/core/java/android/net/StaticIpConfiguration.java
+++ b/core/java/android/net/StaticIpConfiguration.java
@@ -131,7 +131,7 @@
             str.append(" ").append(dnsServer.getHostAddress());
         }
 
-        str.append(" ] Domains");
+        str.append(" ] Domains ");
         if (domains != null) str.append(domains);
         return str.toString();
     }
diff --git a/core/java/android/net/UidRange.java b/core/java/android/net/UidRange.java
index 2e586b3..fd465d9 100644
--- a/core/java/android/net/UidRange.java
+++ b/core/java/android/net/UidRange.java
@@ -48,6 +48,17 @@
         return start / PER_USER_RANGE;
     }
 
+    public boolean contains(int uid) {
+        return start <= uid && uid <= stop;
+    }
+
+    /**
+     * @return {@code true} if this range contains every UID contained by the {@param other} range.
+     */
+    public boolean containsRange(UidRange other) {
+        return start <= other.start && other.stop <= stop;
+    }
+
     @Override
     public int hashCode() {
         int result = 17;
diff --git a/services/core/java/com/android/server/ConnectivityService.java b/services/core/java/com/android/server/ConnectivityService.java
index 3adca72..5118b3f 100644
--- a/services/core/java/com/android/server/ConnectivityService.java
+++ b/services/core/java/com/android/server/ConnectivityService.java
@@ -33,6 +33,7 @@
 import static android.net.NetworkPolicyManager.MASK_METERED_NETWORKS;
 import static android.net.NetworkPolicyManager.MASK_ALL_NETWORKS;
 import static android.net.NetworkPolicyManager.RULE_NONE;
+import static android.net.NetworkPolicyManager.RULE_REJECT_ALL;
 import static android.net.NetworkPolicyManager.RULE_REJECT_METERED;
 import static android.net.NetworkPolicyManager.RULE_TEMPORARY_ALLOW_METERED;
 import static android.net.NetworkPolicyManager.uidRulesToString;
@@ -218,9 +219,6 @@
     /** Flag indicating if background data is restricted. */
     @GuardedBy("mRulesLock")
     private boolean mRestrictBackground;
-    /** Flag indicating if background data is restricted due to battery savings. */
-    @GuardedBy("mRulesLock")
-    private boolean mRestrictPower;
 
     final private Context mContext;
     private int mNetworkPreference;
@@ -669,7 +667,6 @@
         try {
             mPolicyManager.setConnectivityListener(mPolicyListener);
             mRestrictBackground = mPolicyManager.getRestrictBackground();
-            mRestrictPower = mPolicyManager.getRestrictPower();
         } catch (RemoteException e) {
             // ouch, no rules updates means some processes may never get network
             loge("unable to register INetworkPolicyListener" + e);
@@ -918,6 +915,13 @@
         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);
@@ -942,13 +946,11 @@
                         + ": " + allowed);
             }
         }
-        // ...then Battery Saver Mode.
-        if (allowed && mRestrictPower) {
-            allowed = (uidRules & RULE_ALLOW_ALL) != 0;
+        // ...then power restrictions.
+        if (allowed) {
+            allowed = (uidRules & RULE_REJECT_ALL) == 0;
             if (LOGD_RULES) Log.d(TAG, "allowed status for uid " + uid + " when"
-                    + " mRestrictPower=" + mRestrictPower
-                    + ", whitelisted=" + ((uidRules & RULE_ALLOW_ALL) != 0)
-                    + ": " + allowed);
+                    + " rule is " + uidRulesToString(uidRules) + ": " + allowed);
         }
         return !allowed;
     }
@@ -1400,7 +1402,11 @@
                 final int oldRules = mUidRules.get(uid, RULE_NONE);
                 if (oldRules == uidRules) return;
 
-                mUidRules.put(uid, uidRules);
+                if (uidRules == RULE_NONE) {
+                    mUidRules.delete(uid);
+                } else {
+                    mUidRules.put(uid, uidRules);
+                }
             }
 
             // TODO: notify UID when it has requested targeted updates
@@ -1439,18 +1445,6 @@
         }
 
         @Override
-        public void onRestrictPowerChanged(boolean restrictPower) {
-            // caller is NPMS, since we only register with them
-            if (LOGD_RULES) {
-                log("onRestrictPowerChanged(restrictPower=" + restrictPower + ")");
-            }
-
-            synchronized (mRulesLock) {
-                mRestrictPower = restrictPower;
-            }
-        }
-
-        @Override
         public void onRestrictBackgroundWhitelistChanged(int uid, boolean whitelisted) {
             if (LOGD_RULES) {
                 // caller is NPMS, since we only register with them
@@ -1458,6 +1452,14 @@
                         + whitelisted + ")");
             }
         }
+        @Override
+        public void onRestrictBackgroundBlacklistChanged(int uid, boolean blacklisted) {
+            if (LOGD_RULES) {
+                // caller is NPMS, since we only register with them
+                log("onRestrictBackgroundBlacklistChanged(uid=" + uid + ", blacklisted="
+                        + blacklisted + ")");
+            }
+        }
     };
 
     /**
@@ -1891,10 +1893,6 @@
         pw.println(mRestrictBackground);
         pw.println();
 
-        pw.print("Restrict power: ");
-        pw.println(mRestrictPower);
-        pw.println();
-
         pw.println("Status for known UIDs:");
         pw.increaseIndent();
         final int size = mUidRules.size();
@@ -2247,6 +2245,11 @@
             final boolean wasDefault = isDefaultNetwork(nai);
             if (wasDefault) {
                 mDefaultInetConditionPublished = 0;
+                // Log default network disconnection before required book-keeping.
+                // Let rematchAllNetworksAndRequests() below record a new default network event
+                // if there is a fallback. Taken together, the two form a X -> 0, 0 -> Y sequence
+                // whose timestamps tell how long it takes to recover a default network.
+                logDefaultNetworkEvent(null, nai);
             }
             notifyIfacesChangedForNetworkStats();
             // TODO - we shouldn't send CALLBACK_LOST to requests that can be satisfied
@@ -2280,10 +2283,6 @@
             }
             mLegacyTypeTracker.remove(nai, wasDefault);
             rematchAllNetworksAndRequests(null, 0);
-            if (wasDefault && getDefaultNetwork() == null) {
-                // Log that we lost the default network and there is no replacement.
-                logDefaultNetworkEvent(null, nai);
-            }
             if (nai.created) {
                 // Tell netd to clean up the configuration for this network
                 // (routing rules, DNS, etc).
@@ -3371,23 +3370,42 @@
     }
 
     /**
-     * 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) {
-        final String lockdownPackage = getAlwaysOnVpnPackage(user);
-        if (lockdownPackage == null) {
-            return true;
+    private boolean startAlwaysOnVpn(int userId) {
+        final String alwaysOnPackage;
+        synchronized (mVpns) {
+            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;
         }
     }
@@ -3402,25 +3420,35 @@
             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)) {
-                vpn.setAlwaysOnPackage(null);
+            if (!startAlwaysOnVpn(userId)) {
+                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;
     }
@@ -3691,11 +3719,18 @@
             }
             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);
         }
     }
 
@@ -3706,6 +3741,7 @@
                 loge("Stopped user has no VPN");
                 return;
             }
+            userVpn.onUserStopped();
             mVpns.delete(userId);
         }
     }
@@ -3735,7 +3771,7 @@
         if (mUserManager.getUserInfo(userId).isPrimary() && LockdownVpnTracker.isEnabled()) {
             updateLockdownVpn();
         } else {
-            updateAlwaysOnVpn(userId);
+            startAlwaysOnVpn(userId);
         }
     }
 
@@ -4028,7 +4064,8 @@
             synchronized(mRulesLock) {
                 uidRules = mUidRules.get(uid, RULE_ALLOW_ALL);
             }
-            if ((uidRules & RULE_ALLOW_ALL) == 0) {
+            if (mRestrictBackground && (uidRules & RULE_ALLOW_METERED) == 0
+                    && (uidRules & RULE_TEMPORARY_ALLOW_METERED) == 0) {
                 // we could silently fail or we can filter the available nets to only give
                 // them those they have access to.  Chose the more useful option.
                 networkCapabilities.addCapability(NET_CAPABILITY_NOT_METERED);
@@ -4545,7 +4582,7 @@
         teardownUnneededNetwork(oldNetwork);
     }
 
-    private void makeDefault(NetworkAgentInfo newNetwork, NetworkAgentInfo prevNetwork) {
+    private void makeDefault(NetworkAgentInfo newNetwork) {
         if (DBG) log("Switching to new default network: " + newNetwork);
         setupDataActivityTracking(newNetwork);
         try {
@@ -4557,7 +4594,6 @@
         handleApplyDefaultProxy(newNetwork.linkProperties.getHttpProxy());
         updateTcpBufferSizes(newNetwork);
         setDefaultDnsSystemProperties(newNetwork.linkProperties.getDnsServers());
-        logDefaultNetworkEvent(newNetwork, prevNetwork);
     }
 
     // Handles a network appearing or improving its score.
@@ -4708,7 +4744,9 @@
         }
         if (isNewDefault) {
             // Notify system services that this network is up.
-            makeDefault(newNetwork, oldDefaultNetwork);
+            makeDefault(newNetwork);
+            // Log 0 -> X and Y -> X default network transitions, where X is the new default.
+            logDefaultNetworkEvent(newNetwork, oldDefaultNetwork);
             synchronized (ConnectivityService.this) {
                 // have a new default network, release the transition wakelock in
                 // a second if it's held.  The second pause is to allow apps
diff --git a/services/net/java/android/net/apf/ApfCapabilities.java b/services/net/java/android/net/apf/ApfCapabilities.java
index 0ec50c4..b0e0230 100644
--- a/services/net/java/android/net/apf/ApfCapabilities.java
+++ b/services/net/java/android/net/apf/ApfCapabilities.java
@@ -38,7 +38,8 @@
      */
     public final int apfPacketFormat;
 
-    ApfCapabilities(int apfVersionSupported, int maximumApfProgramSize, int apfPacketFormat) {
+    public ApfCapabilities(int apfVersionSupported, int maximumApfProgramSize, int apfPacketFormat)
+    {
         this.apfVersionSupported = apfVersionSupported;
         this.maximumApfProgramSize = maximumApfProgramSize;
         this.apfPacketFormat = apfPacketFormat;
diff --git a/services/tests/servicestests/src/com/android/server/connectivity/VpnTest.java b/services/tests/servicestests/src/com/android/server/connectivity/VpnTest.java
new file mode 100644
index 0000000..5d8b843
--- /dev/null
+++ b/services/tests/servicestests/src/com/android/server/connectivity/VpnTest.java
@@ -0,0 +1,310 @@
+/*
+ * Copyright (C) 2016 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.connectivity;
+
+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;
+import android.net.UidRange;
+import android.os.INetworkManagementService;
+import android.os.Looper;
+import android.os.UserHandle;
+import android.os.UserManager;
+import android.test.AndroidTestCase;
+import android.test.suitebuilder.annotation.SmallTest;
+import android.util.ArrayMap;
+import android.util.ArraySet;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Map;
+import java.util.Set;
+
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+/**
+ * Tests for {@link Vpn}.
+ *
+ * Build, install and run with:
+ *  runtest --path src/com/android/server/connectivity/VpnTest.java
+ */
+public class VpnTest extends AndroidTestCase {
+    private static final String TAG = "VpnTest";
+
+    // Mock users
+    static final UserInfo primaryUser = new UserInfo(27, "Primary", FLAG_ADMIN | FLAG_PRIMARY);
+    static final UserInfo secondaryUser = new UserInfo(15, "Secondary", FLAG_ADMIN);
+    static final UserInfo restrictedProfileA = new UserInfo(40, "RestrictedA", FLAG_RESTRICTED);
+    static final UserInfo restrictedProfileB = new UserInfo(42, "RestrictedB", FLAG_RESTRICTED);
+    static final UserInfo managedProfileA = new UserInfo(45, "ManagedA", FLAG_MANAGED_PROFILE);
+    static {
+        restrictedProfileA.restrictedProfileParentId = primaryUser.id;
+        restrictedProfileB.restrictedProfileParentId = secondaryUser.id;
+        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());
+    }
+
+    @SmallTest
+    public void testRestrictedProfilesAreAddedToVpn() {
+        setMockedUsers(primaryUser, secondaryUser, restrictedProfileA, restrictedProfileB);
+
+        final Vpn vpn = new MockVpn(primaryUser.id);
+        final Set<UidRange> ranges = vpn.createUserAndRestrictedProfilesRanges(primaryUser.id,
+                null, null);
+
+        assertEquals(new ArraySet<>(Arrays.asList(new UidRange[] {
+            UidRange.createForUser(primaryUser.id),
+            UidRange.createForUser(restrictedProfileA.id)
+        })), ranges);
+    }
+
+    @SmallTest
+    public void testManagedProfilesAreNotAddedToVpn() {
+        setMockedUsers(primaryUser, managedProfileA);
+
+        final Vpn vpn = new MockVpn(primaryUser.id);
+        final Set<UidRange> ranges = vpn.createUserAndRestrictedProfilesRanges(primaryUser.id,
+                null, null);
+
+        assertEquals(new ArraySet<>(Arrays.asList(new UidRange[] {
+            UidRange.createForUser(primaryUser.id)
+        })), ranges);
+    }
+
+    @SmallTest
+    public void testAddUserToVpnOnlyAddsOneUser() {
+        setMockedUsers(primaryUser, restrictedProfileA, managedProfileA);
+
+        final Vpn vpn = new MockVpn(primaryUser.id);
+        final Set<UidRange> ranges = new ArraySet<>();
+        vpn.addUserToRanges(ranges, primaryUser.id, null, null);
+
+        assertEquals(new ArraySet<>(Arrays.asList(new UidRange[] {
+            UidRange.createForUser(primaryUser.id)
+        })), ranges);
+    }
+
+    @SmallTest
+    public void testUidWhiteAndBlacklist() throws Exception {
+        final Vpn vpn = new MockVpn(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,
+                Arrays.asList(packages), null);
+        assertEquals(new ArraySet<>(Arrays.asList(new UidRange[] {
+            new UidRange(user.start + PKG_UIDS[0], user.start + PKG_UIDS[0]),
+            new UidRange(user.start + PKG_UIDS[1], user.start + PKG_UIDS[2])
+        })), allow);
+
+        // Blacklist
+        final Set<UidRange> disallow = vpn.createUserAndRestrictedProfilesRanges(primaryUser.id,
+                null, Arrays.asList(packages));
+        assertEquals(new ArraySet<>(Arrays.asList(new UidRange[] {
+            new UidRange(user.start, user.start + PKG_UIDS[0] - 1),
+            new UidRange(user.start + PKG_UIDS[0] + 1, user.start + PKG_UIDS[1] - 1),
+            /* 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)
+        }));
+    }
+
+    /**
+     * A subclass of {@link Vpn} with some of the fields pre-mocked.
+     */
+    private class MockVpn extends Vpn {
+        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));
+            }
+        }
+    }
+
+    /**
+     * Populate {@link #mUserManager} with a list of fake users.
+     */
+    private void setMockedUsers(UserInfo... users) {
+        final Map<Integer, UserInfo> userMap = new ArrayMap<>();
+        for (UserInfo user : users) {
+            userMap.put(user.id, user);
+        }
+
+        /**
+         * @see UserManagerService#getUsers(boolean)
+         */
+        doAnswer(invocation -> {
+            final boolean excludeDying = (boolean) invocation.getArguments()[0];
+            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];
+            return userMap.get(id);
+        }).when(mUserManager).getUserInfo(anyInt());
+
+        doAnswer(invocation -> {
+            final int id = (int) invocation.getArguments()[0];
+            return (userMap.get(id).flags & UserInfo.FLAG_ADMIN) != 0;
+        }).when(mUserManager).canHaveRestrictedProfile(anyInt());
+    }
+
+    /**
+     * Populate {@link #mPackageManager} with a fake packageName-to-UID mapping.
+     */
+    private void setMockedPackages(final Map<String, Integer> packages) {
+        try {
+            doAnswer(invocation -> {
+                final String appName = (String) invocation.getArguments()[0];
+                final int userId = (int) invocation.getArguments()[1];
+                return UserHandle.getUid(userId, packages.get(appName));
+            }).when(mPackageManager).getPackageUidAsUser(anyString(), anyInt());
+        } catch (Exception e) {
+        }
+    }
+}