Merge "Update VpnTest for new restricted API" am: 2954e8af0f am: 69dbb598f9 am: 5704beb453

Original change: https://android-review.googlesource.com/c/platform/frameworks/base/+/1554177

MUST ONLY BE SUBMITTED BY AUTOMERGER

Change-Id: I9188685fe0bb5ee7dfe2c0538fea6644b2cd642e

2 files changed, 17 insertions, 23 deletions

diff --git a/services/core/java/com/android/server/connectivity/Vpn.java b/services/core/java/com/android/server/connectivity/Vpn.java
index 5956fe1abdd1..e8062ae0eb57 100644
--- a/services/core/java/com/android/server/connectivity/Vpn.java
+++ b/services/core/java/com/android/server/connectivity/Vpn.java
@@ -203,6 +203,7 @@ public class Vpn {
     protected final NetworkCapabilities mNetworkCapabilities;
     private final SystemServices mSystemServices;
     private final Ikev2SessionCreator mIkev2SessionCreator;
+    private final UserManager mUserManager;
 
     /**
      * Whether to keep the connection active after rebooting, or upgrading or reinstalling. This
@@ -409,6 +410,7 @@ public class Vpn {
         mLooper = looper;
         mSystemServices = systemServices;
         mIkev2SessionCreator = ikev2SessionCreator;
+        mUserManager = mContext.getSystemService(UserManager.class);
 
         mPackage = VpnConfig.LEGACY_VPN;
         mOwnerUID = getAppUid(mPackage, mUserId);
@@ -1435,7 +1437,7 @@ public class Vpn {
             final long token = Binder.clearCallingIdentity();
             List<UserInfo> users;
             try {
-                users = UserManager.get(mContext).getAliveUsers();
+                users = mUserManager.getAliveUsers();
             } finally {
                 Binder.restoreCallingIdentity(token);
             }
@@ -1519,7 +1521,7 @@ public class Vpn {
      */
     public void onUserAdded(int userId) {
         // If the user is restricted tie them to the parent user's VPN
-        UserInfo user = UserManager.get(mContext).getUserInfo(userId);
+        UserInfo user = mUserManager.getUserInfo(userId);
         if (user.isRestricted() && user.restrictedProfileParentId == mUserId) {
             synchronized(Vpn.this) {
                 final Set<UidRange> existingRanges = mNetworkCapabilities.getUids();
@@ -1547,7 +1549,7 @@ public class Vpn {
      */
     public void onUserRemoved(int userId) {
         // clean up if restricted
-        UserInfo user = UserManager.get(mContext).getUserInfo(userId);
+        UserInfo user = mUserManager.getUserInfo(userId);
         if (user.isRestricted() && user.restrictedProfileParentId == mUserId) {
             synchronized(Vpn.this) {
                 final Set<UidRange> existingRanges = mNetworkCapabilities.getUids();
@@ -1972,8 +1974,7 @@ public class Vpn {
 
     private void enforceNotRestrictedUser() {
         Binder.withCleanCallingIdentity(() -> {
-            final UserManager mgr = UserManager.get(mContext);
-            final UserInfo user = mgr.getUserInfo(mUserId);
+            final UserInfo user = mUserManager.getUserInfo(mUserId);
 
             if (user.isRestricted()) {
                 throw new SecurityException("Restricted users cannot configure VPNs");
@@ -2008,9 +2009,8 @@ public class Vpn {
      */
     public void startLegacyVpnPrivileged(VpnProfile profile, KeyStore keyStore,
             @Nullable Network underlying, @NonNull LinkProperties egress) {
-        UserManager mgr = UserManager.get(mContext);
-        UserInfo user = mgr.getUserInfo(mUserId);
-        if (user.isRestricted() || mgr.hasUserRestriction(UserManager.DISALLOW_CONFIG_VPN,
+        UserInfo user = mUserManager.getUserInfo(mUserId);
+        if (user.isRestricted() || mUserManager.hasUserRestriction(UserManager.DISALLOW_CONFIG_VPN,
                     new UserHandle(mUserId))) {
             throw new SecurityException("Restricted users cannot establish VPNs");
         }
diff --git a/tests/net/java/com/android/server/connectivity/VpnTest.java b/tests/net/java/com/android/server/connectivity/VpnTest.java
index 3a93c5b10540..32c6a75bd904 100644
--- a/tests/net/java/com/android/server/connectivity/VpnTest.java
+++ b/tests/net/java/com/android/server/connectivity/VpnTest.java
@@ -257,12 +257,14 @@ public class VpnTest {
 
     @Test
     public void testRestrictedProfilesAreAddedToVpn() {
-        if (true) return; // TODO(b/175883995): Test disabled until updated for new UserManager API.
         setMockedUsers(primaryUser, secondaryUser, restrictedProfileA, restrictedProfileB);
 
         final Vpn vpn = createVpn(primaryUser.id);
-        final Set<UidRange> ranges = vpn.createUserAndRestrictedProfilesRanges(primaryUser.id,
-                null, null);
+
+        // Assume the user can have restricted profiles.
+        doReturn(true).when(mUserManager).canHaveRestrictedProfile();
+        final Set<UidRange> ranges =
+                vpn.createUserAndRestrictedProfilesRanges(primaryUser.id, null, null);
 
         assertEquals(new ArraySet<>(Arrays.asList(new UidRange[] {
                 PRI_USER_RANGE, UidRange.createForUser(restrictedProfileA.id)
@@ -271,7 +273,6 @@ public class VpnTest {
 
     @Test
     public void testManagedProfilesAreNotAddedToVpn() {
-        if (true) return; // TODO(b/175883995): Test disabled until updated for new UserManager API.
         setMockedUsers(primaryUser, managedProfileA);
 
         final Vpn vpn = createVpn(primaryUser.id);
@@ -294,7 +295,6 @@ public class VpnTest {
 
     @Test
     public void testUidAllowAndDenylist() throws Exception {
-        if (true) return; // TODO(b/175883995): Test disabled until updated for new UserManager API.
         final Vpn vpn = createVpn(primaryUser.id);
         final UidRange user = PRI_USER_RANGE;
         final String[] packages = {PKGS[0], PKGS[1], PKGS[2]};
@@ -320,7 +320,6 @@ public class VpnTest {
 
     @Test
     public void testGetAlwaysAndOnGetLockDown() throws Exception {
-        if (true) return; // TODO(b/175883995): Test disabled until updated for new UserManager API.
         final Vpn vpn = createVpn(primaryUser.id);
 
         // Default state.
@@ -345,7 +344,6 @@ public class VpnTest {
 
     @Test
     public void testLockdownChangingPackage() throws Exception {
-        if (true) return; // TODO(b/175883995): Test disabled until updated for new UserManager API.
         final Vpn vpn = createVpn(primaryUser.id);
         final UidRange user = PRI_USER_RANGE;
 
@@ -373,7 +371,6 @@ public class VpnTest {
 
     @Test
     public void testLockdownAllowlist() throws Exception {
-        if (true) return; // TODO(b/175883995): Test disabled until updated for new UserManager API.
         final Vpn vpn = createVpn(primaryUser.id);
         final UidRange user = PRI_USER_RANGE;
 
@@ -448,7 +445,6 @@ public class VpnTest {
 
     @Test
     public void testLockdownRuleRepeatability() throws Exception {
-        if (true) return; // TODO(b/175883995): Test disabled until updated for new UserManager API.
         final Vpn vpn = createVpn(primaryUser.id);
         final UidRangeParcel[] primaryUserRangeParcel = new UidRangeParcel[] {
                 new UidRangeParcel(PRI_USER_RANGE.start, PRI_USER_RANGE.stop)};
@@ -481,7 +477,6 @@ public class VpnTest {
 
     @Test
     public void testLockdownRuleReversibility() throws Exception {
-        if (true) return; // TODO(b/175883995): Test disabled until updated for new UserManager API.
         final Vpn vpn = createVpn(primaryUser.id);
         final UidRangeParcel[] entireUser = {
             new UidRangeParcel(PRI_USER_RANGE.start, PRI_USER_RANGE.stop)
@@ -1164,6 +1159,10 @@ public class VpnTest {
         doReturn(UserHandle.of(userId)).when(asUserContext).getUser();
         when(mContext.createContextAsUser(eq(UserHandle.of(userId)), anyInt()))
                 .thenReturn(asUserContext);
+        when(asUserContext.getSystemServiceName(UserManager.class))
+                .thenReturn(Context.USER_SERVICE);
+        when(asUserContext.getSystemService(UserManager.class))
+                .thenReturn(mUserManager);
         final TestLooper testLooper = new TestLooper();
         final Vpn vpn = new Vpn(testLooper.getLooper(), mContext, new TestDeps(), mNetService,
                 mNetd, userId, mKeyStore, mSystemServices, mIkev2SessionCreator);
@@ -1199,11 +1198,6 @@ public class VpnTest {
             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();
     }
 
     /**