Add always-on VPN support for platform VPNs

This commit allows Platform VPNs to be started as part of always-on
mode.

Test: FrameworksNetTests passing, new tests added in subsequent CL
Test: Manually tested.
Change-Id: I5eda88e5b406a0e425eb7424665cf702e0979324
Merged-In: I5eda88e5b406a0e425eb7424665cf702e0979324
diff --git a/core/java/android/net/VpnManager.java b/core/java/android/net/VpnManager.java
index f19ba0f..2041cfb 100644
--- a/core/java/android/net/VpnManager.java
+++ b/core/java/android/net/VpnManager.java
@@ -126,7 +126,11 @@
         return getIntentForConfirmation();
     }
 
-    /** Delete the VPN profile configuration that was provisioned by the calling app */
+    /**
+     * Delete the VPN profile configuration that was provisioned by the calling app
+     *
+     * @throws SecurityException if this would violate user settings
+     */
     public void deleteProvisionedVpnProfile() {
         try {
             mService.deleteVpnProfile(mContext.getOpPackageName());
diff --git a/services/core/java/com/android/server/ConnectivityService.java b/services/core/java/com/android/server/ConnectivityService.java
index 5f032fc..52a2ca97 100644
--- a/services/core/java/com/android/server/ConnectivityService.java
+++ b/services/core/java/com/android/server/ConnectivityService.java
@@ -4783,7 +4783,7 @@
                 return false;
             }
 
-            return vpn.startAlwaysOnVpn();
+            return vpn.startAlwaysOnVpn(mKeyStore);
         }
     }
 
@@ -4798,7 +4798,7 @@
                 Slog.w(TAG, "User " + userId + " has no Vpn configuration");
                 return false;
             }
-            return vpn.isAlwaysOnPackageSupported(packageName);
+            return vpn.isAlwaysOnPackageSupported(packageName, mKeyStore);
         }
     }
 
@@ -4819,11 +4819,11 @@
                 Slog.w(TAG, "User " + userId + " has no Vpn configuration");
                 return false;
             }
-            if (!vpn.setAlwaysOnPackage(packageName, lockdown, lockdownWhitelist)) {
+            if (!vpn.setAlwaysOnPackage(packageName, lockdown, lockdownWhitelist, mKeyStore)) {
                 return false;
             }
             if (!startAlwaysOnVpn(userId)) {
-                vpn.setAlwaysOnPackage(null, false, null);
+                vpn.setAlwaysOnPackage(null, false, null, mKeyStore);
                 return false;
             }
         }
@@ -5009,7 +5009,7 @@
                 loge("Starting user already has a VPN");
                 return;
             }
-            userVpn = new Vpn(mHandler.getLooper(), mContext, mNMS, userId);
+            userVpn = new Vpn(mHandler.getLooper(), mContext, mNMS, userId, mKeyStore);
             mVpns.put(userId, userVpn);
             if (mUserManager.getUserInfo(userId).isPrimary() && LockdownVpnTracker.isEnabled()) {
                 updateLockdownVpn();
@@ -5080,7 +5080,7 @@
             if (TextUtils.equals(vpn.getAlwaysOnPackage(), packageName)) {
                 Slog.d(TAG, "Restarting always-on VPN package " + packageName + " for user "
                         + userId);
-                vpn.startAlwaysOnVpn();
+                vpn.startAlwaysOnVpn(mKeyStore);
             }
         }
     }
@@ -5102,7 +5102,7 @@
             if (TextUtils.equals(vpn.getAlwaysOnPackage(), packageName) && !isReplacing) {
                 Slog.d(TAG, "Removing always-on VPN package " + packageName + " for user "
                         + userId);
-                vpn.setAlwaysOnPackage(null, false, null);
+                vpn.setAlwaysOnPackage(null, false, null, mKeyStore);
             }
         }
     }
diff --git a/services/core/java/com/android/server/connectivity/Vpn.java b/services/core/java/com/android/server/connectivity/Vpn.java
index 8692436..1e04994 100644
--- a/services/core/java/com/android/server/connectivity/Vpn.java
+++ b/services/core/java/com/android/server/connectivity/Vpn.java
@@ -216,14 +216,14 @@
      * Whether to keep the connection active after rebooting, or upgrading or reinstalling. This
      * only applies to {@link VpnService} connections.
      */
-    private boolean mAlwaysOn = false;
+    @VisibleForTesting protected 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.
+     * not set. Applies to all types of VPNs.
      */
-    private boolean mLockdown = false;
+    @VisibleForTesting protected boolean mLockdown = false;
 
     /**
      * Set of packages in addition to the VPN app itself that can access the network directly when
@@ -252,14 +252,14 @@
     private final int mUserHandle;
 
     public Vpn(Looper looper, Context context, INetworkManagementService netService,
-            @UserIdInt int userHandle) {
-        this(looper, context, netService, userHandle,
+            @UserIdInt int userHandle, @NonNull KeyStore keyStore) {
+        this(looper, context, netService, userHandle, keyStore,
                 new SystemServices(context), new Ikev2SessionCreator());
     }
 
     @VisibleForTesting
     protected Vpn(Looper looper, Context context, INetworkManagementService netService,
-            int userHandle, SystemServices systemServices,
+            int userHandle, @NonNull KeyStore keyStore, SystemServices systemServices,
             Ikev2SessionCreator ikev2SessionCreator) {
         mContext = context;
         mNetd = netService;
@@ -285,7 +285,7 @@
         mNetworkCapabilities.removeCapability(NetworkCapabilities.NET_CAPABILITY_NOT_VPN);
         updateCapabilities(null /* defaultNetwork */);
 
-        loadAlwaysOnPackage();
+        loadAlwaysOnPackage(keyStore);
     }
 
     /**
@@ -437,23 +437,36 @@
     /**
      * Checks if a VPN app supports always-on mode.
      *
-     * In order to support the always-on feature, an app has to
+     * <p>In order to support the always-on feature, an app has to either have an installed
+     * PlatformVpnProfile, or:
+     *
      * <ul>
-     *     <li>target {@link VERSION_CODES#N API 24} or above, and
-     *     <li>not opt out through the {@link VpnService#SERVICE_META_DATA_SUPPORTS_ALWAYS_ON}
-     *         meta-data field.
+     *   <li>target {@link VERSION_CODES#N API 24} or above, and
+     *   <li>not opt out through the {@link VpnService#SERVICE_META_DATA_SUPPORTS_ALWAYS_ON}
+     *       meta-data field.
      * </ul>
      *
      * @param packageName the canonical package name of the VPN app
+     * @param keyStore the keystore instance to use for checking if the app has a Platform VPN
+     *     profile installed.
      * @return {@code true} if and only if the VPN app exists and supports always-on mode
      */
-    public boolean isAlwaysOnPackageSupported(String packageName) {
+    public boolean isAlwaysOnPackageSupported(String packageName, @NonNull KeyStore keyStore) {
         enforceSettingsPermission();
 
         if (packageName == null) {
             return false;
         }
 
+        final long oldId = Binder.clearCallingIdentity();
+        try {
+            if (getVpnProfilePrivileged(packageName, keyStore) != null) {
+                return true;
+            }
+        } finally {
+            Binder.restoreCallingIdentity(oldId);
+        }
+
         PackageManager pm = mContext.getPackageManager();
         ApplicationInfo appInfo = null;
         try {
@@ -485,27 +498,31 @@
     }
 
     /**
-     * Configures an always-on VPN connection through a specific application.
-     * This connection is automatically granted and persisted after a reboot.
+     * Configures an always-on VPN connection through a specific application. This connection is
+     * automatically granted and persisted after a reboot.
      *
-     * <p>The designated package should exist and declare a {@link VpnService} in its
-     *    manifest guarded by {@link android.Manifest.permission.BIND_VPN_SERVICE},
-     *    otherwise the call will fail.
+     * <p>The designated package should either have a PlatformVpnProfile installed, or declare a
+     * {@link VpnService} in its manifest guarded by {@link
+     * android.Manifest.permission.BIND_VPN_SERVICE}, otherwise the call will fail.
      *
      * <p>Note that this method does not check if the VPN app supports always-on mode. The check is
-     *    delayed to {@link #startAlwaysOnVpn()}, which is always called immediately after this
-     *    method in {@link android.net.IConnectivityManager#setAlwaysOnVpnPackage}.
+     * delayed to {@link #startAlwaysOnVpn()}, which is always called immediately after this method
+     * in {@link android.net.IConnectivityManager#setAlwaysOnVpnPackage}.
      *
      * @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.
      * @param lockdownWhitelist packages to be whitelisted from lockdown.
+     * @param keyStore the Keystore instance to use for checking of PlatformVpnProfile(s)
      * @return {@code true} if the package has been set as always-on, {@code false} otherwise.
      */
     public synchronized boolean setAlwaysOnPackage(
-            String packageName, boolean lockdown, List<String> lockdownWhitelist) {
+            @Nullable String packageName,
+            boolean lockdown,
+            @Nullable List<String> lockdownWhitelist,
+            @NonNull KeyStore keyStore) {
         enforceControlPermissionOrInternalCaller();
 
-        if (setAlwaysOnPackageInternal(packageName, lockdown, lockdownWhitelist)) {
+        if (setAlwaysOnPackageInternal(packageName, lockdown, lockdownWhitelist, keyStore)) {
             saveAlwaysOnPackage();
             return true;
         }
@@ -513,20 +530,22 @@
     }
 
     /**
-     * Configures an always-on VPN connection through a specific application, the same as
-     * {@link #setAlwaysOnPackage}.
+     * Configures an always-on VPN connection through a specific application, the same as {@link
+     * #setAlwaysOnPackage}.
      *
-     * Does not perform permission checks. Does not persist any of the changes to storage.
+     * <p>Does not perform permission checks. Does not persist any of the changes to storage.
      *
      * @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.
      * @param lockdownWhitelist packages to be whitelisted from lockdown. This is only used if
-     *        {@code lockdown} is {@code true}. Packages must not contain commas.
+     *     {@code lockdown} is {@code true}. Packages must not contain commas.
+     * @param keyStore the system keystore instance to check for profiles
      * @return {@code true} if the package has been set as always-on, {@code false} otherwise.
      */
     @GuardedBy("this")
     private boolean setAlwaysOnPackageInternal(
-            String packageName, boolean lockdown, List<String> lockdownWhitelist) {
+            @Nullable String packageName, boolean lockdown,
+            @Nullable List<String> lockdownWhitelist, @NonNull KeyStore keyStore) {
         if (VpnConfig.LEGACY_VPN.equals(packageName)) {
             Log.w(TAG, "Not setting legacy VPN \"" + packageName + "\" as always-on.");
             return false;
@@ -542,11 +561,18 @@
         }
 
         if (packageName != null) {
-            // TODO: Give the minimum permission possible; if there is a Platform VPN profile, only
-            // grant ACTIVATE_PLATFORM_VPN.
-            // Pre-authorize new always-on VPN package. Grant the full ACTIVATE_VPN appop, allowing
-            // both VpnService and Platform VPNs.
-            if (!setPackageAuthorization(packageName, VpnManager.TYPE_VPN_SERVICE)) {
+            final VpnProfile profile;
+            final long oldId = Binder.clearCallingIdentity();
+            try {
+                profile = getVpnProfilePrivileged(packageName, keyStore);
+            } finally {
+                Binder.restoreCallingIdentity(oldId);
+            }
+
+            // Pre-authorize new always-on VPN package.
+            final int grantType =
+                    (profile == null) ? VpnManager.TYPE_VPN_SERVICE : VpnManager.TYPE_VPN_PLATFORM;
+            if (!setPackageAuthorization(packageName, grantType)) {
                 return false;
             }
             mAlwaysOn = true;
@@ -611,11 +637,9 @@
         }
     }
 
-    /**
-     * Load the always-on package and lockdown config from Settings.Secure
-     */
+    /** Load the always-on package and lockdown config from Settings. */
     @GuardedBy("this")
-    private void loadAlwaysOnPackage() {
+    private void loadAlwaysOnPackage(@NonNull KeyStore keyStore) {
         final long token = Binder.clearCallingIdentity();
         try {
             final String alwaysOnPackage = mSystemServices.settingsSecureGetStringForUser(
@@ -626,17 +650,21 @@
                     Settings.Secure.ALWAYS_ON_VPN_LOCKDOWN_WHITELIST, mUserHandle);
             final List<String> whitelistedPackages = TextUtils.isEmpty(whitelistString)
                     ? Collections.emptyList() : Arrays.asList(whitelistString.split(","));
-            setAlwaysOnPackageInternal(alwaysOnPackage, alwaysOnLockdown, whitelistedPackages);
+            setAlwaysOnPackageInternal(
+                    alwaysOnPackage, alwaysOnLockdown, whitelistedPackages, keyStore);
         } finally {
             Binder.restoreCallingIdentity(token);
         }
     }
 
     /**
+     * Starts the currently selected always-on VPN
+     *
+     * @param keyStore the keyStore instance for looking up PlatformVpnProfile(s)
      * @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.
+     *     was no always-on VPN to start. {@code false} otherwise.
      */
-    public boolean startAlwaysOnVpn() {
+    public boolean startAlwaysOnVpn(@NonNull KeyStore keyStore) {
         final String alwaysOnPackage;
         synchronized (this) {
             alwaysOnPackage = getAlwaysOnPackage();
@@ -645,8 +673,8 @@
                 return true;
             }
             // Remove always-on VPN if it's not supported.
-            if (!isAlwaysOnPackageSupported(alwaysOnPackage)) {
-                setAlwaysOnPackage(null, false, null);
+            if (!isAlwaysOnPackageSupported(alwaysOnPackage, keyStore)) {
+                setAlwaysOnPackage(null, false, null, keyStore);
                 return false;
             }
             // Skip if the service is already established. This isn't bulletproof: it's not bound
@@ -657,10 +685,24 @@
             }
         }
 
-        // Tell the OS that background services in this app need to be allowed for
-        // a short time, so we can bootstrap the VPN service.
         final long oldId = Binder.clearCallingIdentity();
         try {
+            // Prefer VPN profiles, if any exist.
+            VpnProfile profile = getVpnProfilePrivileged(alwaysOnPackage, keyStore);
+            if (profile != null) {
+                startVpnProfilePrivileged(profile, alwaysOnPackage);
+
+                // If the above startVpnProfilePrivileged() call returns, the Ikev2VpnProfile was
+                // correctly parsed, and the VPN has started running in a different thread. The only
+                // other possibility is that the above call threw an exception, which will be
+                // caught below, and returns false (clearing the always-on VPN). Once started, the
+                // Platform VPN cannot permanantly fail, and is resiliant to temporary failures. It
+                // will continue retrying until shut down by the user, or always-on is toggled off.
+                return true;
+            }
+
+            // Tell the OS that background services in this app need to be allowed for
+            // a short time, so we can bootstrap the VPN service.
             DeviceIdleController.LocalService idleController =
                     LocalServices.getService(DeviceIdleController.LocalService.class);
             idleController.addPowerSaveTempWhitelistApp(Process.myUid(), alwaysOnPackage,
@@ -675,6 +717,9 @@
                 Log.e(TAG, "VpnService " + serviceIntent + " failed to start", e);
                 return false;
             }
+        } catch (Exception e) {
+            Log.e(TAG, "Error starting always-on VPN", e);
+            return false;
         } finally {
             Binder.restoreCallingIdentity(oldId);
         }
@@ -2816,6 +2861,10 @@
         return isVpnProfilePreConsented(mContext, packageName);
     }
 
+    private boolean isCurrentIkev2VpnLocked(@NonNull String packageName) {
+        return isCurrentPreparedPackage(packageName) && mVpnRunner instanceof IkeV2VpnRunner;
+    }
+
     /**
      * Deletes an app-provisioned VPN profile.
      *
@@ -2832,6 +2881,17 @@
 
         Binder.withCleanCallingIdentity(
                 () -> {
+                    // If this profile is providing the current VPN, turn it off, disabling
+                    // always-on as well if enabled.
+                    if (isCurrentIkev2VpnLocked(packageName)) {
+                        if (mAlwaysOn) {
+                            // Will transitively call prepareInternal(VpnConfig.LEGACY_VPN).
+                            setAlwaysOnPackage(null, false, null, keyStore);
+                        } else {
+                            prepareInternal(VpnConfig.LEGACY_VPN);
+                        }
+                    }
+
                     keyStore.delete(getProfileNameForPackage(packageName), Process.SYSTEM_UID);
                 });
     }
@@ -2942,11 +3002,9 @@
 
         // To stop the VPN profile, the caller must be the current prepared package and must be
         // running an Ikev2VpnProfile.
-        if (!isCurrentPreparedPackage(packageName) && mVpnRunner instanceof IkeV2VpnRunner) {
-            return;
+        if (isCurrentIkev2VpnLocked(packageName)) {
+            prepareInternal(VpnConfig.LEGACY_VPN);
         }
-
-        prepareInternal(VpnConfig.LEGACY_VPN);
     }
 
     /**
diff --git a/tests/net/java/com/android/server/ConnectivityServiceTest.java b/tests/net/java/com/android/server/ConnectivityServiceTest.java
index 8ed497b..77147c8 100644
--- a/tests/net/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/net/java/com/android/server/ConnectivityServiceTest.java
@@ -204,6 +204,7 @@
 import android.os.UserHandle;
 import android.os.UserManager;
 import android.provider.Settings;
+import android.security.KeyStore;
 import android.system.Os;
 import android.test.mock.MockContentResolver;
 import android.text.TextUtils;
@@ -1019,7 +1020,7 @@
 
         public MockVpn(int userId) {
             super(startHandlerThreadAndReturnLooper(), mServiceContext, mNetworkManagementService,
-                    userId);
+                    userId, mock(KeyStore.class));
         }
 
         public void setNetworkAgent(TestNetworkAgentWrapper agent) {
diff --git a/tests/net/java/com/android/server/connectivity/VpnTest.java b/tests/net/java/com/android/server/connectivity/VpnTest.java
index ac1c518..0e3b797 100644
--- a/tests/net/java/com/android/server/connectivity/VpnTest.java
+++ b/tests/net/java/com/android/server/connectivity/VpnTest.java
@@ -72,6 +72,7 @@
 import android.os.Process;
 import android.os.UserHandle;
 import android.os.UserManager;
+import android.provider.Settings;
 import android.security.Credentials;
 import android.security.KeyStore;
 import android.util.ArrayMap;
@@ -260,17 +261,17 @@
         assertFalse(vpn.getLockdown());
 
         // Set always-on without lockdown.
-        assertTrue(vpn.setAlwaysOnPackage(PKGS[1], false, Collections.emptyList()));
+        assertTrue(vpn.setAlwaysOnPackage(PKGS[1], false, Collections.emptyList(), mKeyStore));
         assertTrue(vpn.getAlwaysOn());
         assertFalse(vpn.getLockdown());
 
         // Set always-on with lockdown.
-        assertTrue(vpn.setAlwaysOnPackage(PKGS[1], true, Collections.emptyList()));
+        assertTrue(vpn.setAlwaysOnPackage(PKGS[1], true, Collections.emptyList(), mKeyStore));
         assertTrue(vpn.getAlwaysOn());
         assertTrue(vpn.getLockdown());
 
         // Remove always-on configuration.
-        assertTrue(vpn.setAlwaysOnPackage(null, false, Collections.emptyList()));
+        assertTrue(vpn.setAlwaysOnPackage(null, false, Collections.emptyList(), mKeyStore));
         assertFalse(vpn.getAlwaysOn());
         assertFalse(vpn.getLockdown());
     }
@@ -284,11 +285,11 @@
         assertUnblocked(vpn, 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, null));
+        assertTrue(vpn.setAlwaysOnPackage(PKGS[1], false, null, mKeyStore));
         assertUnblocked(vpn, 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, null));
+        assertTrue(vpn.setAlwaysOnPackage(PKGS[1], true, null, mKeyStore));
         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)
@@ -297,7 +298,7 @@
         assertUnblocked(vpn, user.start + PKG_UIDS[1]);
 
         // Switch to another app.
-        assertTrue(vpn.setAlwaysOnPackage(PKGS[3], true, null));
+        assertTrue(vpn.setAlwaysOnPackage(PKGS[3], true, null, mKeyStore));
         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)
@@ -316,7 +317,8 @@
         final UidRange user = UidRange.createForUser(primaryUser.id);
 
         // Set always-on with lockdown and whitelist app PKGS[2] from lockdown.
-        assertTrue(vpn.setAlwaysOnPackage(PKGS[1], true, Collections.singletonList(PKGS[2])));
+        assertTrue(vpn.setAlwaysOnPackage(
+                PKGS[1], true, Collections.singletonList(PKGS[2]), mKeyStore));
         verify(mNetService).setAllowOnlyVpnForUids(eq(true), aryEq(new UidRange[] {
                 new UidRange(user.start, user.start + PKG_UIDS[1] - 1),
                 new UidRange(user.start + PKG_UIDS[2] + 1, user.stop)
@@ -325,7 +327,8 @@
         assertUnblocked(vpn, user.start + PKG_UIDS[1], user.start + PKG_UIDS[2]);
 
         // Change whitelisted app to PKGS[3].
-        assertTrue(vpn.setAlwaysOnPackage(PKGS[1], true, Collections.singletonList(PKGS[3])));
+        assertTrue(vpn.setAlwaysOnPackage(
+                PKGS[1], true, Collections.singletonList(PKGS[3]), mKeyStore));
         verify(mNetService).setAllowOnlyVpnForUids(eq(false), aryEq(new UidRange[] {
                 new UidRange(user.start + PKG_UIDS[2] + 1, user.stop)
         }));
@@ -337,7 +340,8 @@
         assertUnblocked(vpn, user.start + PKG_UIDS[1], user.start + PKG_UIDS[3]);
 
         // Change the VPN app.
-        assertTrue(vpn.setAlwaysOnPackage(PKGS[0], true, Collections.singletonList(PKGS[3])));
+        assertTrue(vpn.setAlwaysOnPackage(
+                PKGS[0], true, Collections.singletonList(PKGS[3]), mKeyStore));
         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.start + PKG_UIDS[3] - 1)
@@ -350,7 +354,7 @@
         assertUnblocked(vpn, user.start + PKG_UIDS[0], user.start + PKG_UIDS[3]);
 
         // Remove the whitelist.
-        assertTrue(vpn.setAlwaysOnPackage(PKGS[0], true, null));
+        assertTrue(vpn.setAlwaysOnPackage(PKGS[0], true, null, mKeyStore));
         verify(mNetService).setAllowOnlyVpnForUids(eq(false), aryEq(new UidRange[] {
                 new UidRange(user.start + PKG_UIDS[0] + 1, user.start + PKG_UIDS[3] - 1),
                 new UidRange(user.start + PKG_UIDS[3] + 1, user.stop)
@@ -363,7 +367,8 @@
         assertUnblocked(vpn, user.start + PKG_UIDS[0]);
 
         // Add the whitelist.
-        assertTrue(vpn.setAlwaysOnPackage(PKGS[0], true, Collections.singletonList(PKGS[1])));
+        assertTrue(vpn.setAlwaysOnPackage(
+                PKGS[0], true, Collections.singletonList(PKGS[1]), mKeyStore));
         verify(mNetService).setAllowOnlyVpnForUids(eq(false), aryEq(new UidRange[] {
                 new UidRange(user.start + PKG_UIDS[0] + 1, user.stop)
         }));
@@ -375,12 +380,13 @@
         assertUnblocked(vpn, user.start + PKG_UIDS[0], user.start + PKG_UIDS[1]);
 
         // Try whitelisting a package with a comma, should be rejected.
-        assertFalse(vpn.setAlwaysOnPackage(PKGS[0], true, Collections.singletonList("a.b,c.d")));
+        assertFalse(vpn.setAlwaysOnPackage(
+                PKGS[0], true, Collections.singletonList("a.b,c.d"), mKeyStore));
 
         // Pass a non-existent packages in the whitelist, they (and only they) should be ignored.
         // Whitelisted package should change from PGKS[1] to PKGS[2].
-        assertTrue(vpn.setAlwaysOnPackage(PKGS[0], true,
-                Arrays.asList("com.foo.app", PKGS[2], "com.bar.app")));
+        assertTrue(vpn.setAlwaysOnPackage(
+                PKGS[0], true, Arrays.asList("com.foo.app", PKGS[2], "com.bar.app"), mKeyStore));
         verify(mNetService).setAllowOnlyVpnForUids(eq(false), aryEq(new UidRange[]{
                 new UidRange(user.start + PKG_UIDS[0] + 1, user.start + PKG_UIDS[1] - 1),
                 new UidRange(user.start + PKG_UIDS[1] + 1, user.stop)
@@ -405,7 +411,7 @@
         final UidRange profile = UidRange.createForUser(tempProfile.id);
 
         // Set lockdown.
-        assertTrue(vpn.setAlwaysOnPackage(PKGS[3], true, null));
+        assertTrue(vpn.setAlwaysOnPackage(PKGS[3], true, null, mKeyStore));
         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)
@@ -499,22 +505,22 @@
                 .thenReturn(Collections.singletonList(resInfo));
 
         // null package name should return false
-        assertFalse(vpn.isAlwaysOnPackageSupported(null));
+        assertFalse(vpn.isAlwaysOnPackageSupported(null, mKeyStore));
 
         // Pre-N apps are not supported
         appInfo.targetSdkVersion = VERSION_CODES.M;
-        assertFalse(vpn.isAlwaysOnPackageSupported(PKGS[0]));
+        assertFalse(vpn.isAlwaysOnPackageSupported(PKGS[0], mKeyStore));
 
         // N+ apps are supported by default
         appInfo.targetSdkVersion = VERSION_CODES.N;
-        assertTrue(vpn.isAlwaysOnPackageSupported(PKGS[0]));
+        assertTrue(vpn.isAlwaysOnPackageSupported(PKGS[0], mKeyStore));
 
         // Apps that opt out explicitly are not supported
         appInfo.targetSdkVersion = VERSION_CODES.CUR_DEVELOPMENT;
         Bundle metaData = new Bundle();
         metaData.putBoolean(VpnService.SERVICE_META_DATA_SUPPORTS_ALWAYS_ON, false);
         svcInfo.metaData = metaData;
-        assertFalse(vpn.isAlwaysOnPackageSupported(PKGS[0]));
+        assertFalse(vpn.isAlwaysOnPackageSupported(PKGS[0], mKeyStore));
     }
 
     @Test
@@ -531,7 +537,7 @@
                 .cancelAsUser(anyString(), anyInt(), eq(userHandle));
 
         // Start showing a notification for disconnected once always-on.
-        vpn.setAlwaysOnPackage(PKGS[0], false, null);
+        vpn.setAlwaysOnPackage(PKGS[0], false, null, mKeyStore);
         order.verify(mNotificationManager)
                 .notifyAsUser(anyString(), anyInt(), any(), eq(userHandle));
 
@@ -545,7 +551,7 @@
                 .notifyAsUser(anyString(), anyInt(), any(), eq(userHandle));
 
         // Notification should be cleared after unsetting always-on package.
-        vpn.setAlwaysOnPackage(null, false, null);
+        vpn.setAlwaysOnPackage(null, false, null, mKeyStore);
         order.verify(mNotificationManager).cancelAsUser(anyString(), anyInt(), eq(userHandle));
     }
 
@@ -920,12 +926,48 @@
                         eq(AppOpsManager.MODE_IGNORED));
     }
 
+    private void setAndVerifyAlwaysOnPackage(Vpn vpn, int uid, boolean lockdownEnabled) {
+        assertTrue(vpn.setAlwaysOnPackage(TEST_VPN_PKG, lockdownEnabled, null, mKeyStore));
+
+        verify(mKeyStore).get(eq(vpn.getProfileNameForPackage(TEST_VPN_PKG)));
+        verify(mAppOps).setMode(
+                eq(AppOpsManager.OP_ACTIVATE_PLATFORM_VPN), eq(uid), eq(TEST_VPN_PKG),
+                eq(AppOpsManager.MODE_ALLOWED));
+
+        verify(mSystemServices).settingsSecurePutStringForUser(
+                eq(Settings.Secure.ALWAYS_ON_VPN_APP), eq(TEST_VPN_PKG), eq(primaryUser.id));
+        verify(mSystemServices).settingsSecurePutIntForUser(
+                eq(Settings.Secure.ALWAYS_ON_VPN_LOCKDOWN), eq(lockdownEnabled ? 1 : 0),
+                eq(primaryUser.id));
+        verify(mSystemServices).settingsSecurePutStringForUser(
+                eq(Settings.Secure.ALWAYS_ON_VPN_LOCKDOWN_WHITELIST), eq(""), eq(primaryUser.id));
+    }
+
+    @Test
+    public void testSetAndStartAlwaysOnVpn() throws Exception {
+        final Vpn vpn = createVpn(primaryUser.id);
+        setMockedUsers(primaryUser);
+
+        // UID checks must return a different UID; otherwise it'll be treated as already prepared.
+        final int uid = Process.myUid() + 1;
+        when(mPackageManager.getPackageUidAsUser(eq(TEST_VPN_PKG), anyInt()))
+                .thenReturn(uid);
+        when(mKeyStore.get(vpn.getProfileNameForPackage(TEST_VPN_PKG)))
+                .thenReturn(mVpnProfile.encode());
+
+        setAndVerifyAlwaysOnPackage(vpn, uid, false);
+        assertTrue(vpn.startAlwaysOnVpn(mKeyStore));
+
+        // TODO: Test the Ikev2VpnRunner started up properly. Relies on utility methods added in
+        // a subsequent CL.
+    }
+
     /**
      * Mock some methods of vpn object.
      */
     private Vpn createVpn(@UserIdInt int userId) {
         return new Vpn(Looper.myLooper(), mContext, mNetService,
-                userId, mSystemServices, mIkev2SessionCreator);
+                userId, mKeyStore, mSystemServices, mIkev2SessionCreator);
     }
 
     private static void assertBlocked(Vpn vpn, int... uids) {