Merge "Update NetworkMonitor integration test constructor"
diff --git a/core/java/android/net/ConnectivityManager.java b/core/java/android/net/ConnectivityManager.java
index cb31404..9cf751d 100644
--- a/core/java/android/net/ConnectivityManager.java
+++ b/core/java/android/net/ConnectivityManager.java
@@ -2424,14 +2424,14 @@
     /**
      * Get the set of tethered dhcp ranges.
      *
-     * @return an array of 0 or more {@code String} of tethered dhcp ranges.
-     * @deprecated This API just return the default value which is not used in DhcpServer.
+     * @deprecated This method is not supported.
+     * TODO: remove this function when all of clients are removed.
      * {@hide}
      */
     @RequiresPermission(android.Manifest.permission.NETWORK_SETTINGS)
     @Deprecated
     public String[] getTetheredDhcpRanges() {
-        return getTetheringManager().getTetheredDhcpRanges();
+        throw new UnsupportedOperationException("getTetheredDhcpRanges is not supported");
     }
 
     /**
@@ -4716,19 +4716,19 @@
     /**
      * Returns the {@code uid} of the owner of a network connection.
      *
-     * @param protocol The protocol of the connection. Only {@code IPPROTO_TCP} and
-     * {@code IPPROTO_UDP} currently supported.
+     * @param protocol The protocol of the connection. Only {@code IPPROTO_TCP} and {@code
+     *     IPPROTO_UDP} currently supported.
      * @param local The local {@link InetSocketAddress} of a connection.
      * @param remote The remote {@link InetSocketAddress} of a connection.
-     *
      * @return {@code uid} if the connection is found and the app has permission to observe it
-     * (e.g., if it is associated with the calling VPN app's tunnel) or
-     * {@link android.os.Process#INVALID_UID} if the connection is not found.
-     * Throws {@link SecurityException} if the caller is not the active VPN for the current user.
-     * Throws {@link IllegalArgumentException} if an unsupported protocol is requested.
+     *     (e.g., if it is associated with the calling VPN app's VpnService tunnel) or {@link
+     *     android.os.Process#INVALID_UID} if the connection is not found.
+     * @throws {@link SecurityException} if the caller is not the active VpnService for the current
+     *     user.
+     * @throws {@link IllegalArgumentException} if an unsupported protocol is requested.
      */
-    public int getConnectionOwnerUid(int protocol, @NonNull InetSocketAddress local,
-            @NonNull InetSocketAddress remote) {
+    public int getConnectionOwnerUid(
+            int protocol, @NonNull InetSocketAddress local, @NonNull InetSocketAddress remote) {
         ConnectionInfo connectionInfo = new ConnectionInfo(protocol, local, remote);
         try {
             return mService.getConnectionOwnerUid(connectionInfo);
diff --git a/services/core/java/com/android/server/ConnectivityService.java b/services/core/java/com/android/server/ConnectivityService.java
index 1f027d6..52a2ca9 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);
             }
         }
     }
@@ -6940,6 +6940,15 @@
             // worry about multiple different substates of CONNECTED.
             newInfo.setDetailedState(NetworkInfo.DetailedState.SUSPENDED, info.getReason(),
                     info.getExtraInfo());
+        } else if (!suspended && info.getDetailedState() == NetworkInfo.DetailedState.SUSPENDED) {
+            // SUSPENDED state is currently only overridden from CONNECTED state. In the case the
+            // network agent is created, then goes to suspended, then goes out of suspended without
+            // ever setting connected. Check if network agent is ever connected to update the state.
+            newInfo.setDetailedState(nai.everConnected
+                    ? NetworkInfo.DetailedState.CONNECTED
+                    : NetworkInfo.DetailedState.CONNECTING,
+                    info.getReason(),
+                    info.getExtraInfo());
         }
         newInfo.setRoaming(!nai.networkCapabilities.hasCapability(NET_CAPABILITY_NOT_ROAMING));
         return newInfo;
@@ -7519,6 +7528,13 @@
      */
     public int getConnectionOwnerUid(ConnectionInfo connectionInfo) {
         final Vpn vpn = enforceActiveVpnOrNetworkStackPermission();
+
+        // Only VpnService based VPNs should be able to get this information.
+        if (vpn != null && vpn.getActiveAppVpnType() != VpnManager.TYPE_VPN_SERVICE) {
+            throw new SecurityException(
+                    "getConnectionOwnerUid() not allowed for non-VpnService VPNs");
+        }
+
         if (connectionInfo.protocol != IPPROTO_TCP && connectionInfo.protocol != IPPROTO_UDP) {
             throw new IllegalArgumentException("Unsupported protocol " + connectionInfo.protocol);
         }
diff --git a/tests/net/java/android/net/Ikev2VpnProfileTest.java b/tests/net/java/android/net/Ikev2VpnProfileTest.java
index d6a2176..2273bc6 100644
--- a/tests/net/java/android/net/Ikev2VpnProfileTest.java
+++ b/tests/net/java/android/net/Ikev2VpnProfileTest.java
@@ -22,7 +22,6 @@
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
-import static org.mockito.Mockito.mock;
 
 import android.test.mock.MockContext;
 
@@ -232,10 +231,12 @@
         builder.setAuthDigitalSignature(mUserCert, mPrivateKey, mServerRootCa);
         final VpnProfile profile = builder.build().toVpnProfile();
 
+        final String expectedSecret = Ikev2VpnProfile.PREFIX_INLINE
+                + Ikev2VpnProfile.encodeForIpsecSecret(mPrivateKey.getEncoded());
         verifyVpnProfileCommon(profile);
         assertEquals(Ikev2VpnProfile.certificateToPemString(mUserCert), profile.ipsecUserCert);
         assertEquals(
-                Ikev2VpnProfile.encodeForIpsecSecret(mPrivateKey.getEncoded()),
+                expectedSecret,
                 profile.ipsecSecret);
         assertEquals(Ikev2VpnProfile.certificateToPemString(mServerRootCa), profile.ipsecCaCert);
 
diff --git a/tests/net/java/com/android/server/ConnectivityServiceTest.java b/tests/net/java/com/android/server/ConnectivityServiceTest.java
index 64591a3..77147c8 100644
--- a/tests/net/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/net/java/com/android/server/ConnectivityServiceTest.java
@@ -78,6 +78,7 @@
 import static android.net.NetworkPolicyManager.RULE_REJECT_ALL;
 import static android.net.NetworkPolicyManager.RULE_REJECT_METERED;
 import static android.net.RouteInfo.RTN_UNREACHABLE;
+import static android.system.OsConstants.IPPROTO_TCP;
 
 import static com.android.server.ConnectivityServiceTestUtilsKt.transportToLegacyType;
 import static com.android.testutils.ConcurrentUtilsKt.await;
@@ -138,6 +139,7 @@
 import android.content.res.Resources;
 import android.location.LocationManager;
 import android.net.CaptivePortalData;
+import android.net.ConnectionInfo;
 import android.net.ConnectivityManager;
 import android.net.ConnectivityManager.NetworkCallback;
 import android.net.ConnectivityManager.PacketKeepalive;
@@ -153,6 +155,7 @@
 import android.net.INetworkPolicyListener;
 import android.net.INetworkPolicyManager;
 import android.net.INetworkStatsService;
+import android.net.InetAddresses;
 import android.net.InterfaceConfiguration;
 import android.net.IpPrefix;
 import android.net.IpSecManager;
@@ -176,6 +179,7 @@
 import android.net.SocketKeepalive;
 import android.net.UidRange;
 import android.net.Uri;
+import android.net.VpnManager;
 import android.net.metrics.IpConnectivityLog;
 import android.net.shared.NetworkMonitorUtils;
 import android.net.shared.PrivateDnsConfig;
@@ -200,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;
@@ -272,6 +277,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Predicate;
+import java.util.function.Supplier;
 
 import kotlin.reflect.KClass;
 
@@ -445,15 +451,21 @@
             return mPackageManager;
         }
 
+        private int checkMockedPermission(String permission, Supplier<Integer> ifAbsent) {
+            final Integer granted = mMockedPermissions.get(permission);
+            return granted != null ? granted : ifAbsent.get();
+        }
+
         @Override
         public int checkPermission(String permission, int pid, int uid) {
-            final Integer granted = mMockedPermissions.get(permission);
-            if (granted == null) {
-                // All non-mocked permissions should be held by the test or unnecessary: check as
-                // normal to make sure the code does not rely on unexpected permissions.
-                return super.checkPermission(permission, pid, uid);
-            }
-            return granted;
+            return checkMockedPermission(
+                    permission, () -> super.checkPermission(permission, pid, uid));
+        }
+
+        @Override
+        public int checkCallingOrSelfPermission(String permission) {
+            return checkMockedPermission(
+                    permission, () -> super.checkCallingOrSelfPermission(permission));
         }
 
         @Override
@@ -1002,12 +1014,13 @@
         // Careful ! This is different from mNetworkAgent, because MockNetworkAgent does
         // not inherit from NetworkAgent.
         private TestNetworkAgentWrapper mMockNetworkAgent;
+        private int mVpnType = VpnManager.TYPE_VPN_SERVICE;
 
         private VpnInfo mVpnInfo;
 
         public MockVpn(int userId) {
             super(startHandlerThreadAndReturnLooper(), mServiceContext, mNetworkManagementService,
-                    userId);
+                    userId, mock(KeyStore.class));
         }
 
         public void setNetworkAgent(TestNetworkAgentWrapper agent) {
@@ -1022,6 +1035,10 @@
             updateCapabilities(null /* defaultNetwork */);
         }
 
+        public void setVpnType(int vpnType) {
+            mVpnType = vpnType;
+        }
+
         @Override
         public int getNetId() {
             if (mMockNetworkAgent == null) {
@@ -1040,6 +1057,11 @@
             return mConnected;  // Similar trickery
         }
 
+        @Override
+        public int getActiveAppVpnType() {
+            return mVpnType;
+        }
+
         private void connect(boolean isAlwaysMetered) {
             mNetworkCapabilities.set(mMockNetworkAgent.getNetworkCapabilities());
             mConnected = true;
@@ -3199,6 +3221,7 @@
                 mCellNetworkAgent);
         cellNetworkCallback.expectCallback(CallbackEntry.SUSPENDED, mCellNetworkAgent);
         cellNetworkCallback.assertNoCallback();
+        assertEquals(NetworkInfo.State.SUSPENDED, mCm.getActiveNetworkInfo().getState());
 
         // Register a garden variety default network request.
         TestNetworkCallback dfltNetworkCallback = new TestNetworkCallback();
@@ -3214,6 +3237,7 @@
                 mCellNetworkAgent);
         cellNetworkCallback.expectCallback(CallbackEntry.RESUMED, mCellNetworkAgent);
         cellNetworkCallback.assertNoCallback();
+        assertEquals(NetworkInfo.State.CONNECTED, mCm.getActiveNetworkInfo().getState());
 
         dfltNetworkCallback = new TestNetworkCallback();
         mCm.registerDefaultNetworkCallback(dfltNetworkCallback);
@@ -6427,6 +6451,90 @@
         assertEquals(Process.INVALID_UID, newNc.getOwnerUid());
     }
 
+    private void setupConnectionOwnerUid(int vpnOwnerUid, @VpnManager.VpnType int vpnType)
+            throws Exception {
+        final Set<UidRange> vpnRange = Collections.singleton(UidRange.createForUser(VPN_USER));
+        establishVpn(new LinkProperties(), vpnOwnerUid, vpnRange);
+        mMockVpn.setVpnType(vpnType);
+
+        final VpnInfo vpnInfo = new VpnInfo();
+        vpnInfo.ownerUid = vpnOwnerUid;
+        mMockVpn.setVpnInfo(vpnInfo);
+    }
+
+    private void setupConnectionOwnerUidAsVpnApp(int vpnOwnerUid, @VpnManager.VpnType int vpnType)
+            throws Exception {
+        setupConnectionOwnerUid(vpnOwnerUid, vpnType);
+
+        // Test as VPN app
+        mServiceContext.setPermission(android.Manifest.permission.NETWORK_STACK, PERMISSION_DENIED);
+        mServiceContext.setPermission(
+                NetworkStack.PERMISSION_MAINLINE_NETWORK_STACK, PERMISSION_DENIED);
+    }
+
+    private ConnectionInfo getTestConnectionInfo() throws Exception {
+        return new ConnectionInfo(
+                IPPROTO_TCP,
+                new InetSocketAddress(InetAddresses.parseNumericAddress("1.2.3.4"), 1234),
+                new InetSocketAddress(InetAddresses.parseNumericAddress("2.3.4.5"), 2345));
+    }
+
+    @Test
+    public void testGetConnectionOwnerUidPlatformVpn() throws Exception {
+        final int myUid = Process.myUid();
+        setupConnectionOwnerUidAsVpnApp(myUid, VpnManager.TYPE_VPN_PLATFORM);
+
+        try {
+            mService.getConnectionOwnerUid(getTestConnectionInfo());
+            fail("Expected SecurityException for non-VpnService app");
+        } catch (SecurityException expected) {
+        }
+    }
+
+    @Test
+    public void testGetConnectionOwnerUidVpnServiceWrongUser() throws Exception {
+        final int myUid = Process.myUid();
+        setupConnectionOwnerUidAsVpnApp(myUid + 1, VpnManager.TYPE_VPN_SERVICE);
+
+        try {
+            mService.getConnectionOwnerUid(getTestConnectionInfo());
+            fail("Expected SecurityException for non-VpnService app");
+        } catch (SecurityException expected) {
+        }
+    }
+
+    @Test
+    public void testGetConnectionOwnerUidVpnServiceDoesNotThrow() throws Exception {
+        final int myUid = Process.myUid();
+        setupConnectionOwnerUidAsVpnApp(myUid, VpnManager.TYPE_VPN_SERVICE);
+
+        // TODO: Test the returned UID
+        mService.getConnectionOwnerUid(getTestConnectionInfo());
+    }
+
+    @Test
+    public void testGetConnectionOwnerUidVpnServiceNetworkStackDoesNotThrow() throws Exception {
+        final int myUid = Process.myUid();
+        setupConnectionOwnerUid(myUid, VpnManager.TYPE_VPN_SERVICE);
+        mServiceContext.setPermission(
+                android.Manifest.permission.NETWORK_STACK, PERMISSION_GRANTED);
+
+        // TODO: Test the returned UID
+        mService.getConnectionOwnerUid(getTestConnectionInfo());
+    }
+
+    @Test
+    public void testGetConnectionOwnerUidVpnServiceMainlineNetworkStackDoesNotThrow()
+            throws Exception {
+        final int myUid = Process.myUid();
+        setupConnectionOwnerUid(myUid, VpnManager.TYPE_VPN_SERVICE);
+        mServiceContext.setPermission(
+                NetworkStack.PERMISSION_MAINLINE_NETWORK_STACK, PERMISSION_GRANTED);
+
+        // TODO: Test the returned UID
+        mService.getConnectionOwnerUid(getTestConnectionInfo());
+    }
+
     private TestNetworkAgentWrapper establishVpn(
             LinkProperties lp, int ownerUid, Set<UidRange> vpnRange) throws Exception {
         final TestNetworkAgentWrapper
diff --git a/tests/net/java/com/android/server/connectivity/VpnTest.java b/tests/net/java/com/android/server/connectivity/VpnTest.java
index eb78529..1994d1f 100644
--- a/tests/net/java/com/android/server/connectivity/VpnTest.java
+++ b/tests/net/java/com/android/server/connectivity/VpnTest.java
@@ -59,9 +59,15 @@
 import android.content.pm.UserInfo;
 import android.content.res.Resources;
 import android.net.ConnectivityManager;
+import android.net.Ikev2VpnProfile;
+import android.net.InetAddresses;
+import android.net.IpPrefix;
+import android.net.IpSecManager;
+import android.net.LinkProperties;
 import android.net.Network;
 import android.net.NetworkCapabilities;
 import android.net.NetworkInfo.DetailedState;
+import android.net.RouteInfo;
 import android.net.UidRange;
 import android.net.VpnManager;
 import android.net.VpnService;
@@ -72,6 +78,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;
@@ -83,6 +90,7 @@
 import com.android.internal.R;
 import com.android.internal.net.VpnConfig;
 import com.android.internal.net.VpnProfile;
+import com.android.server.IpSecService;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -92,6 +100,7 @@
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
+import java.net.Inet4Address;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -124,6 +133,9 @@
     }
 
     static final String TEST_VPN_PKG = "com.dummy.vpn";
+    private static final String TEST_VPN_SERVER = "1.2.3.4";
+    private static final String TEST_VPN_IDENTITY = "identity";
+    private static final byte[] TEST_VPN_PSK = "psk".getBytes();
 
     /**
      * Names and UIDs for some fake packages. Important points:
@@ -150,23 +162,39 @@
     @Mock private Vpn.SystemServices mSystemServices;
     @Mock private Vpn.Ikev2SessionCreator mIkev2SessionCreator;
     @Mock private ConnectivityManager mConnectivityManager;
+    @Mock private IpSecService mIpSecService;
     @Mock private KeyStore mKeyStore;
-    private final VpnProfile mVpnProfile = new VpnProfile("key");
+    private final VpnProfile mVpnProfile;
+
+    private IpSecManager mIpSecManager;
+
+    public VpnTest() throws Exception {
+        // Build an actual VPN profile that is capable of being converted to and from an
+        // Ikev2VpnProfile
+        final Ikev2VpnProfile.Builder builder =
+                new Ikev2VpnProfile.Builder(TEST_VPN_SERVER, TEST_VPN_IDENTITY);
+        builder.setAuthPsk(TEST_VPN_PSK);
+        mVpnProfile = builder.build().toVpnProfile();
+    }
 
     @Before
     public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
 
+        mIpSecManager = new IpSecManager(mContext, mIpSecService);
+
         when(mContext.getPackageManager()).thenReturn(mPackageManager);
         setMockedPackages(mPackages);
 
-        when(mContext.getPackageName()).thenReturn(Vpn.class.getPackage().getName());
+        when(mContext.getPackageName()).thenReturn(TEST_VPN_PKG);
+        when(mContext.getOpPackageName()).thenReturn(TEST_VPN_PKG);
         when(mContext.getSystemService(eq(Context.USER_SERVICE))).thenReturn(mUserManager);
         when(mContext.getSystemService(eq(Context.APP_OPS_SERVICE))).thenReturn(mAppOps);
         when(mContext.getSystemService(eq(Context.NOTIFICATION_SERVICE)))
                 .thenReturn(mNotificationManager);
         when(mContext.getSystemService(eq(Context.CONNECTIVITY_SERVICE)))
                 .thenReturn(mConnectivityManager);
+        when(mContext.getSystemService(eq(Context.IPSEC_SERVICE))).thenReturn(mIpSecManager);
         when(mContext.getString(R.string.config_customVpnAlwaysOnDisconnectedDialogComponent))
                 .thenReturn(Resources.getSystem().getString(
                         R.string.config_customVpnAlwaysOnDisconnectedDialogComponent));
@@ -260,17 +288,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 +312,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 +325,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 +344,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 +354,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 +367,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 +381,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 +394,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 +407,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 +438,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 +532,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 +564,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 +578,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));
     }
 
@@ -656,8 +689,12 @@
     }
 
     private Vpn createVpnAndSetupUidChecks(int... grantedOps) throws Exception {
-        final Vpn vpn = createVpn(primaryUser.id);
-        setMockedUsers(primaryUser);
+        return createVpnAndSetupUidChecks(primaryUser, grantedOps);
+    }
+
+    private Vpn createVpnAndSetupUidChecks(UserInfo user, int... grantedOps) throws Exception {
+        final Vpn vpn = createVpn(user.id);
+        setMockedUsers(user);
 
         when(mPackageManager.getPackageUidAsUser(eq(TEST_VPN_PKG), anyInt()))
                 .thenReturn(Process.myUid());
@@ -726,6 +763,19 @@
     }
 
     @Test
+    public void testProvisionVpnProfileRestrictedUser() throws Exception {
+        final Vpn vpn =
+                createVpnAndSetupUidChecks(
+                        restrictedProfileA, AppOpsManager.OP_ACTIVATE_PLATFORM_VPN);
+
+        try {
+            vpn.provisionVpnProfile(TEST_VPN_PKG, mVpnProfile, mKeyStore);
+            fail("Expected SecurityException due to restricted user");
+        } catch (SecurityException expected) {
+        }
+    }
+
+    @Test
     public void testDeleteVpnProfile() throws Exception {
         final Vpn vpn = createVpnAndSetupUidChecks();
 
@@ -736,6 +786,19 @@
     }
 
     @Test
+    public void testDeleteVpnProfileRestrictedUser() throws Exception {
+        final Vpn vpn =
+                createVpnAndSetupUidChecks(
+                        restrictedProfileA, AppOpsManager.OP_ACTIVATE_PLATFORM_VPN);
+
+        try {
+            vpn.deleteVpnProfile(TEST_VPN_PKG, mKeyStore);
+            fail("Expected SecurityException due to restricted user");
+        } catch (SecurityException expected) {
+        }
+    }
+
+    @Test
     public void testGetVpnProfilePrivileged() throws Exception {
         final Vpn vpn = createVpnAndSetupUidChecks();
 
@@ -820,6 +883,32 @@
     }
 
     @Test
+    public void testStartVpnProfileRestrictedUser() throws Exception {
+        final Vpn vpn =
+                createVpnAndSetupUidChecks(
+                        restrictedProfileA, AppOpsManager.OP_ACTIVATE_PLATFORM_VPN);
+
+        try {
+            vpn.startVpnProfile(TEST_VPN_PKG, mKeyStore);
+            fail("Expected SecurityException due to restricted user");
+        } catch (SecurityException expected) {
+        }
+    }
+
+    @Test
+    public void testStopVpnProfileRestrictedUser() throws Exception {
+        final Vpn vpn =
+                createVpnAndSetupUidChecks(
+                        restrictedProfileA, AppOpsManager.OP_ACTIVATE_PLATFORM_VPN);
+
+        try {
+            vpn.stopVpnProfile(TEST_VPN_PKG);
+            fail("Expected SecurityException due to restricted user");
+        } catch (SecurityException expected) {
+        }
+    }
+
+    @Test
     public void testSetPackageAuthorizationVpnService() throws Exception {
         final Vpn vpn = createVpnAndSetupUidChecks();
 
@@ -864,12 +953,68 @@
                         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.
+    }
+
+    @Test
+    public void testStartLegacyVpn() throws Exception {
+        final Vpn vpn = createVpn(primaryUser.id);
+        setMockedUsers(primaryUser);
+
+        // Dummy egress interface
+        final String egressIface = "DUMMY0";
+        final LinkProperties lp = new LinkProperties();
+        lp.setInterfaceName(egressIface);
+
+        final RouteInfo defaultRoute = new RouteInfo(new IpPrefix(Inet4Address.ANY, 0),
+                        InetAddresses.parseNumericAddress("192.0.2.0"), egressIface);
+        lp.addRoute(defaultRoute);
+
+        vpn.startLegacyVpn(mVpnProfile, mKeyStore, lp);
+
+        // 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) {