Merge changes Id9f81fdf,I3428e8b3,Ifa895f71,I61cd4751,I47a25e9f into main

* changes:
  Add tests for always on lockdown VPN on system user.
  Remove MockVpn.setAlwaysOnPackage() non-lockdown.
  Mock onUserAdded() and onUserRemoved()
  Refactor helper method to return integer ranges.
  Add tests for onUserAdded and onUserRemoved
diff --git a/tests/unit/java/com/android/server/ConnectivityServiceTest.java b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
index 42b0d59..8bca4dd 100755
--- a/tests/unit/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
@@ -937,22 +937,43 @@
 
     // This function assumes the UID range for user 0 ([1, 99999])
     private static UidRangeParcel[] uidRangeParcelsExcludingUids(Integer... excludedUids) {
-        int start = 1;
-        Arrays.sort(excludedUids);
-        List<UidRangeParcel> parcels = new ArrayList<UidRangeParcel>();
+        final List<Integer> uids = Arrays.asList(excludedUids);
+        return intToUidRangeStableParcels(intRangesPrimaryExcludingUids(uids));
+    }
+
+    // Create the list of ranges for the primary user (User 0), excluding excludedUids.
+    private static List<Range<Integer>> intRangesPrimaryExcludingUids(List<Integer> excludedUids) {
+        final List<Integer> excludedUidsList = new ArrayList<>(excludedUids);
+        // Uid 0 is always excluded
+        if (!excludedUidsList.contains(0)) {
+            excludedUidsList.add(0);
+        }
+        return intRangesExcludingUids(PRIMARY_USER, excludedUidsList);
+    }
+
+    private static List<Range<Integer>> intRangesExcludingUids(int userId,
+            List<Integer> excludedAppIds) {
+        final List<Integer> excludedUids = CollectionUtils.map(excludedAppIds,
+                appId -> UserHandle.getUid(userId, appId));
+        final int userBase = userId * UserHandle.PER_USER_RANGE;
+        final int maxUid = userBase + UserHandle.PER_USER_RANGE - 1;
+
+        int start = userBase;
+        Collections.sort(excludedUids);
+        final List<Range<Integer>> ranges = new ArrayList<>();
         for (int excludedUid : excludedUids) {
             if (excludedUid == start) {
                 start++;
             } else {
-                parcels.add(new UidRangeParcel(start, excludedUid - 1));
+                ranges.add(new Range<>(start, excludedUid - 1));
                 start = excludedUid + 1;
             }
         }
-        if (start <= 99999) {
-            parcels.add(new UidRangeParcel(start, 99999));
+        if (start <= maxUid) {
+            ranges.add(new Range<>(start, maxUid));
         }
 
-        return parcels.toArray(new UidRangeParcel[0]);
+        return ranges;
     }
 
     private void waitForIdle() {
@@ -1739,6 +1760,12 @@
         return ranges.stream().map(r -> new UidRangeParcel(r, r)).toArray(UidRangeParcel[]::new);
     }
 
+    private static UidRangeParcel[] intToUidRangeStableParcels(
+            final @NonNull List<Range<Integer>> ranges) {
+        return ranges.stream().map(
+                r -> new UidRangeParcel(r.getLower(), r.getUpper())).toArray(UidRangeParcel[]::new);
+    }
+
     private void assertVpnTransportInfo(NetworkCapabilities nc, int type) {
         assertNotNull(nc);
         final TransportInfo ti = nc.getTransportInfo();
@@ -1871,6 +1898,8 @@
     private static final UserHandle TERTIARY_USER_HANDLE = new UserHandle(TERTIARY_USER);
 
     private static final int RESTRICTED_USER = 1;
+    private static final UidRange RESTRICTED_USER_UIDRANGE =
+            UidRange.createForUser(UserHandle.of(RESTRICTED_USER));
     private static final UserInfo RESTRICTED_USER_INFO = new UserInfo(RESTRICTED_USER, "",
             UserInfo.FLAG_RESTRICTED);
     static {
@@ -9426,11 +9455,11 @@
                         && c.hasTransport(TRANSPORT_WIFI));
         callback.expectCaps(mWiFiAgent, c -> c.hasCapability(NET_CAPABILITY_VALIDATED));
 
-        doReturn(UserHandle.getUid(RESTRICTED_USER, VPN_UID)).when(mPackageManager)
-                .getPackageUidAsUser(ALWAYS_ON_PACKAGE, RESTRICTED_USER);
-
-        // New user added
-        mMockVpn.onUserAdded(RESTRICTED_USER);
+        // New user added, this updates the Vpn uids, coverage in VpnTest.
+        // This is equivalent to `mMockVpn.onUserAdded(RESTRICTED_USER);`
+        final Set<UidRange> ranges = uidRangesForUids(uid);
+        ranges.add(RESTRICTED_USER_UIDRANGE);
+        mMockVpn.setUids(ranges);
 
         // Expect that the VPN UID ranges contain both |uid| and the UID range for the newly-added
         // restricted user.
@@ -9455,7 +9484,9 @@
                 && !c.hasTransport(TRANSPORT_WIFI));
 
         // User removed and expect to lose the UID range for the restricted user.
-        mMockVpn.onUserRemoved(RESTRICTED_USER);
+        // This updates the Vpn uids, coverage in VpnTest.
+        // This is equivalent to `mMockVpn.onUserRemoved(RESTRICTED_USER);`
+        mMockVpn.setUids(uidRangesForUids(uid));
 
         // Expect that the VPN gains the UID range for the restricted user, and that the capability
         // change made just before that (i.e., loss of TRANSPORT_WIFI) is preserved.
@@ -9506,28 +9537,29 @@
         assertNotNull(mCm.getActiveNetworkForUid(restrictedUid));
 
         // Start the restricted profile, and check that the UID within it loses network access.
-        doReturn(UserHandle.getUid(RESTRICTED_USER, VPN_UID)).when(mPackageManager)
-                .getPackageUidAsUser(ALWAYS_ON_PACKAGE, RESTRICTED_USER);
-        doReturn(asList(PRIMARY_USER_INFO, RESTRICTED_USER_INFO)).when(mUserManager)
-                .getAliveUsers();
         // TODO: check that VPN app within restricted profile still has access, etc.
-        mMockVpn.onUserAdded(RESTRICTED_USER);
-        final Intent addedIntent = new Intent(ACTION_USER_ADDED);
-        addedIntent.putExtra(Intent.EXTRA_USER, UserHandle.of(RESTRICTED_USER));
-        addedIntent.putExtra(Intent.EXTRA_USER_HANDLE, RESTRICTED_USER);
-        processBroadcast(addedIntent);
+        // Add a restricted user.
+        // This is equivalent to `mMockVpn.onUserAdded(RESTRICTED_USER);`, coverage in VpnTest.
+        final List<Integer> excludedUids = new ArrayList<Integer>();
+        excludedUids.add(VPN_UID);
+        if (mDeps.isAtLeastT()) {
+            // On T onwards, the corresponding SDK sandbox UID should also be excluded
+            excludedUids.add(toSdkSandboxUid(VPN_UID));
+        }
+        final List<Range<Integer>> restrictedRanges =
+                intRangesExcludingUids(RESTRICTED_USER, excludedUids);
+        mCm.setRequireVpnForUids(true, restrictedRanges);
+        waitForIdle();
+
         assertNull(mCm.getActiveNetworkForUid(uid));
         assertNull(mCm.getActiveNetworkForUid(restrictedUid));
 
         // Stop the restricted profile, and check that the UID within it has network access again.
-        doReturn(asList(PRIMARY_USER_INFO)).when(mUserManager).getAliveUsers();
+        // Remove the restricted user.
+        // This is equivalent to `mMockVpn.onUserRemoved(RESTRICTED_USER);`, coverage in VpnTest.
+        mCm.setRequireVpnForUids(false, restrictedRanges);
+        waitForIdle();
 
-        // Send a USER_REMOVED broadcast and expect to lose the UID range for the restricted user.
-        mMockVpn.onUserRemoved(RESTRICTED_USER);
-        final Intent removedIntent = new Intent(ACTION_USER_REMOVED);
-        removedIntent.putExtra(Intent.EXTRA_USER, UserHandle.of(RESTRICTED_USER));
-        removedIntent.putExtra(Intent.EXTRA_USER_HANDLE, RESTRICTED_USER);
-        processBroadcast(removedIntent);
         assertNull(mCm.getActiveNetworkForUid(uid));
         assertNotNull(mCm.getActiveNetworkForUid(restrictedUid));
 
@@ -10099,34 +10131,6 @@
         assertNetworkInfo(TYPE_MOBILE, DetailedState.DISCONNECTED);
         assertNetworkInfo(TYPE_WIFI, DetailedState.CONNECTED);
 
-        // Enable and disable an always-on VPN package without lockdown. Expect no changes.
-        reset(mMockNetd);
-        mMockVpn.setAlwaysOnPackage(ALWAYS_ON_PACKAGE, false /* lockdown */, allowList);
-        inOrder.verify(mMockNetd, never()).networkRejectNonSecureVpn(anyBoolean(), any());
-        callback.assertNoCallback();
-        defaultCallback.assertNoCallback();
-        vpnUidCallback.assertNoCallback();
-        vpnUidDefaultCallback.assertNoCallback();
-        vpnDefaultCallbackAsUid.assertNoCallback();
-        assertEquals(mWiFiAgent.getNetwork(), mCm.getActiveNetworkForUid(VPN_UID));
-        assertEquals(mWiFiAgent.getNetwork(), mCm.getActiveNetwork());
-        assertActiveNetworkInfo(TYPE_WIFI, DetailedState.CONNECTED);
-        assertNetworkInfo(TYPE_MOBILE, DetailedState.DISCONNECTED);
-        assertNetworkInfo(TYPE_WIFI, DetailedState.CONNECTED);
-
-        mMockVpn.setAlwaysOnPackage(null, false /* lockdown */, allowList);
-        inOrder.verify(mMockNetd, never()).networkRejectNonSecureVpn(anyBoolean(), any());
-        callback.assertNoCallback();
-        defaultCallback.assertNoCallback();
-        vpnUidCallback.assertNoCallback();
-        vpnUidDefaultCallback.assertNoCallback();
-        vpnDefaultCallbackAsUid.assertNoCallback();
-        assertEquals(mWiFiAgent.getNetwork(), mCm.getActiveNetworkForUid(VPN_UID));
-        assertEquals(mWiFiAgent.getNetwork(), mCm.getActiveNetwork());
-        assertActiveNetworkInfo(TYPE_WIFI, DetailedState.CONNECTED);
-        assertNetworkInfo(TYPE_MOBILE, DetailedState.DISCONNECTED);
-        assertNetworkInfo(TYPE_WIFI, DetailedState.CONNECTED);
-
         // Enable lockdown and connect a VPN. The VPN is not blocked.
         mMockVpn.setAlwaysOnPackage(ALWAYS_ON_PACKAGE, true /* lockdown */, allowList);
         defaultCallback.expect(BLOCKED_STATUS, mWiFiAgent, cb -> cb.getBlocked());
@@ -10262,7 +10266,8 @@
         // Init lockdown state to simulate LockdownVpnTracker behavior.
         mCm.setLegacyLockdownVpnEnabled(true);
         mMockVpn.setEnableTeardown(false);
-        final Set<Range<Integer>> ranges = UidRange.toIntRanges(Set.of(PRIMARY_UIDRANGE));
+        final List<Range<Integer>> ranges =
+                intRangesPrimaryExcludingUids(Collections.EMPTY_LIST /* excludedeUids */);
         mCm.setRequireVpnForUids(true /* requireVpn */, ranges);
 
         // Bring up a network.
@@ -10468,7 +10473,8 @@
 
     @Test @IgnoreUpTo(Build.VERSION_CODES.S_V2)
     public void testLockdownSetFirewallUidRule() throws Exception {
-        final Set<Range<Integer>> lockdownRange = UidRange.toIntRanges(Set.of(PRIMARY_UIDRANGE));
+        final List<Range<Integer>> lockdownRange =
+                intRangesPrimaryExcludingUids(Collections.EMPTY_LIST /* excludedeUids */);
         // Enable Lockdown
         mCm.setRequireVpnForUids(true /* requireVpn */, lockdownRange);
         waitForIdle();
diff --git a/tests/unit/java/com/android/server/connectivity/VpnTest.java b/tests/unit/java/com/android/server/connectivity/VpnTest.java
index 478ff61..48cfe77 100644
--- a/tests/unit/java/com/android/server/connectivity/VpnTest.java
+++ b/tests/unit/java/com/android/server/connectivity/VpnTest.java
@@ -579,6 +579,18 @@
     }
 
     @Test
+    public void testAlwaysOnWithoutLockdown() throws Exception {
+        final Vpn vpn = createVpn(PRIMARY_USER.id);
+        assertTrue(vpn.setAlwaysOnPackage(
+                PKGS[1], false /* lockdown */, null /* lockdownAllowlist */));
+        verify(mConnectivityManager, never()).setRequireVpnForUids(anyBoolean(), any());
+
+        assertTrue(vpn.setAlwaysOnPackage(
+                null /* packageName */, false /* lockdown */, null /* lockdownAllowlist */));
+        verify(mConnectivityManager, never()).setRequireVpnForUids(anyBoolean(), any());
+    }
+
+    @Test
     public void testLockdownChangingPackage() throws Exception {
         final Vpn vpn = createVpn(PRIMARY_USER.id);
         final Range<Integer> user = PRIMARY_USER_RANGE;
@@ -724,6 +736,37 @@
     }
 
     @Test
+    public void testLockdownSystemUser() throws Exception {
+        final Vpn vpn = createVpn(SYSTEM_USER_ID);
+
+        // Uid 0 is always excluded and PKG_UIDS[1] is the uid of the VPN.
+        final List<Integer> excludedUids = new ArrayList<>(List.of(0, PKG_UIDS[1]));
+        final List<Range<Integer>> ranges = makeVpnUidRange(SYSTEM_USER_ID, excludedUids);
+
+        // Set always-on with lockdown.
+        assertTrue(vpn.setAlwaysOnPackage(
+                PKGS[1], true /* lockdown */, null /* lockdownAllowlist */));
+        verify(mConnectivityManager).setRequireVpnForUids(true, ranges);
+
+        // Disable always-on with lockdown.
+        assertTrue(vpn.setAlwaysOnPackage(
+                null /* packageName */, false /* lockdown */, null /* lockdownAllowlist */));
+        verify(mConnectivityManager).setRequireVpnForUids(false, ranges);
+
+        // Set always-on with lockdown and allow the app PKGS[2].
+        excludedUids.add(PKG_UIDS[2]);
+        final List<Range<Integer>> ranges2 = makeVpnUidRange(SYSTEM_USER_ID, excludedUids);
+        assertTrue(vpn.setAlwaysOnPackage(
+                PKGS[1], true /* lockdown */, Collections.singletonList(PKGS[2])));
+        verify(mConnectivityManager).setRequireVpnForUids(true, ranges2);
+
+        // Disable always-on with lockdown.
+        assertTrue(vpn.setAlwaysOnPackage(
+                null /* packageName */, false /* lockdown */, null /* lockdownAllowlist */));
+        verify(mConnectivityManager).setRequireVpnForUids(false, ranges2);
+    }
+
+    @Test
     public void testLockdownRuleRepeatability() throws Exception {
         final Vpn vpn = createVpn(PRIMARY_USER.id);
         final UidRangeParcel[] primaryUserRangeParcel = new UidRangeParcel[] {
@@ -788,6 +831,101 @@
     }
 
     @Test
+    public void testOnUserAddedAndRemoved_restrictedUser() throws Exception {
+        final InOrder order = inOrder(mMockNetworkAgent);
+        final Vpn vpn = createVpn(PRIMARY_USER.id);
+        final Set<Range<Integer>> initialRange = rangeSet(PRIMARY_USER_RANGE);
+        // Note since mVpnProfile is a Ikev2VpnProfile, this starts an IkeV2VpnRunner.
+        startLegacyVpn(vpn, mVpnProfile);
+        // Set an initial Uid range and mock the network agent
+        vpn.mNetworkCapabilities.setUids(initialRange);
+        vpn.mNetworkAgent = mMockNetworkAgent;
+
+        // Add the restricted user
+        setMockedUsers(PRIMARY_USER, RESTRICTED_PROFILE_A);
+        vpn.onUserAdded(RESTRICTED_PROFILE_A.id);
+        // Expect restricted user range to be added to the NetworkCapabilities.
+        final Set<Range<Integer>> expectRestrictedRange =
+                rangeSet(PRIMARY_USER_RANGE, uidRangeForUser(RESTRICTED_PROFILE_A.id));
+        assertEquals(expectRestrictedRange, vpn.mNetworkCapabilities.getUids());
+        order.verify(mMockNetworkAgent).doSendNetworkCapabilities(
+                argThat(nc -> expectRestrictedRange.equals(nc.getUids())));
+
+        // Remove the restricted user
+        vpn.onUserRemoved(RESTRICTED_PROFILE_A.id);
+        // Expect restricted user range to be removed from the NetworkCapabilities.
+        assertEquals(initialRange, vpn.mNetworkCapabilities.getUids());
+        order.verify(mMockNetworkAgent).doSendNetworkCapabilities(
+                argThat(nc -> initialRange.equals(nc.getUids())));
+    }
+
+    @Test
+    public void testOnUserAddedAndRemoved_restrictedUserLockdown() throws Exception {
+        final UidRangeParcel[] primaryUserRangeParcel = new UidRangeParcel[] {
+                new UidRangeParcel(PRIMARY_USER_RANGE.getLower(), PRIMARY_USER_RANGE.getUpper())};
+        final Range<Integer> restrictedUserRange = uidRangeForUser(RESTRICTED_PROFILE_A.id);
+        final UidRangeParcel[] restrictedUserRangeParcel = new UidRangeParcel[] {
+                new UidRangeParcel(restrictedUserRange.getLower(), restrictedUserRange.getUpper())};
+        final Vpn vpn = createVpn(PRIMARY_USER.id);
+
+        // Set lockdown calls setRequireVpnForUids
+        vpn.setLockdown(true);
+        verify(mConnectivityManager).setRequireVpnForUids(true, toRanges(primaryUserRangeParcel));
+
+        // Add the restricted user
+        doReturn(true).when(mUserManager).canHaveRestrictedProfile();
+        setMockedUsers(PRIMARY_USER, RESTRICTED_PROFILE_A);
+        vpn.onUserAdded(RESTRICTED_PROFILE_A.id);
+
+        // Expect restricted user range to be added.
+        verify(mConnectivityManager).setRequireVpnForUids(true,
+                toRanges(restrictedUserRangeParcel));
+
+        // Mark as partial indicates that the user is removed, mUserManager.getAliveUsers() does not
+        // return the restricted user but it is still returned in mUserManager.getUserInfo().
+        RESTRICTED_PROFILE_A.partial = true;
+        // Remove the restricted user
+        vpn.onUserRemoved(RESTRICTED_PROFILE_A.id);
+        verify(mConnectivityManager).setRequireVpnForUids(false,
+                toRanges(restrictedUserRangeParcel));
+        // reset to avoid affecting other tests since RESTRICTED_PROFILE_A is static.
+        RESTRICTED_PROFILE_A.partial = false;
+    }
+
+    @Test
+    public void testOnUserAddedAndRemoved_restrictedUserAlwaysOn() throws Exception {
+        final Vpn vpn = createVpn(PRIMARY_USER.id);
+
+        // setAlwaysOnPackage() calls setRequireVpnForUids()
+        assertTrue(vpn.setAlwaysOnPackage(
+                PKGS[0], true /* lockdown */, null /* lockdownAllowlist */));
+        final List<Integer> excludedUids = List.of(PKG_UIDS[0]);
+        final List<Range<Integer>> primaryRanges =
+                makeVpnUidRange(PRIMARY_USER.id, excludedUids);
+        verify(mConnectivityManager).setRequireVpnForUids(true, primaryRanges);
+
+        // Add the restricted user
+        doReturn(true).when(mUserManager).canHaveRestrictedProfile();
+        setMockedUsers(PRIMARY_USER, RESTRICTED_PROFILE_A);
+        vpn.onUserAdded(RESTRICTED_PROFILE_A.id);
+
+        final List<Range<Integer>> restrictedRanges =
+                makeVpnUidRange(RESTRICTED_PROFILE_A.id, excludedUids);
+        // Expect restricted user range to be added.
+        verify(mConnectivityManager).setRequireVpnForUids(true, restrictedRanges);
+
+        // Mark as partial indicates that the user is removed, mUserManager.getAliveUsers() does not
+        // return the restricted user but it is still returned in mUserManager.getUserInfo().
+        RESTRICTED_PROFILE_A.partial = true;
+        // Remove the restricted user
+        vpn.onUserRemoved(RESTRICTED_PROFILE_A.id);
+        verify(mConnectivityManager).setRequireVpnForUids(false, restrictedRanges);
+
+        // reset to avoid affecting other tests since RESTRICTED_PROFILE_A is static.
+        RESTRICTED_PROFILE_A.partial = false;
+    }
+
+    @Test
     public void testPrepare_throwSecurityExceptionWhenGivenPackageDoesNotBelongToTheCaller()
             throws Exception {
         mTestDeps.mIgnoreCallingUidChecks = false;
@@ -1002,12 +1140,12 @@
 
         // List in keystore is not changed, but UID for the removed packages is no longer exempted.
         assertEquals(Arrays.asList(PKGS), vpn.getAppExclusionList(TEST_VPN_PKG));
-        assertEquals(makeVpnUidRange(PRIMARY_USER.id, newExcludedUids),
+        assertEquals(makeVpnUidRangeSet(PRIMARY_USER.id, newExcludedUids),
                 vpn.mNetworkCapabilities.getUids());
         ArgumentCaptor<NetworkCapabilities> ncCaptor =
                 ArgumentCaptor.forClass(NetworkCapabilities.class);
         verify(mMockNetworkAgent).doSendNetworkCapabilities(ncCaptor.capture());
-        assertEquals(makeVpnUidRange(PRIMARY_USER.id, newExcludedUids),
+        assertEquals(makeVpnUidRangeSet(PRIMARY_USER.id, newExcludedUids),
                 ncCaptor.getValue().getUids());
 
         reset(mMockNetworkAgent);
@@ -1019,26 +1157,28 @@
 
         // List in keystore is not changed and the uid list should be updated in the net cap.
         assertEquals(Arrays.asList(PKGS), vpn.getAppExclusionList(TEST_VPN_PKG));
-        assertEquals(makeVpnUidRange(PRIMARY_USER.id, newExcludedUids),
+        assertEquals(makeVpnUidRangeSet(PRIMARY_USER.id, newExcludedUids),
                 vpn.mNetworkCapabilities.getUids());
         verify(mMockNetworkAgent).doSendNetworkCapabilities(ncCaptor.capture());
-        assertEquals(makeVpnUidRange(PRIMARY_USER.id, newExcludedUids),
+        assertEquals(makeVpnUidRangeSet(PRIMARY_USER.id, newExcludedUids),
                 ncCaptor.getValue().getUids());
     }
 
-    private Set<Range<Integer>> makeVpnUidRange(int userId, List<Integer> excludedList) {
+    private List<Range<Integer>> makeVpnUidRange(int userId, List<Integer> excludedAppIdList) {
         final SortedSet<Integer> list = new TreeSet<>();
 
         final int userBase = userId * UserHandle.PER_USER_RANGE;
-        for (int uid : excludedList) {
-            final int applicationUid = UserHandle.getUid(userId, uid);
-            list.add(applicationUid);
-            list.add(Process.toSdkSandboxUid(applicationUid)); // Add Sdk Sandbox UID
+        for (int appId : excludedAppIdList) {
+            final int uid = UserHandle.getUid(userId, appId);
+            list.add(uid);
+            if (Process.isApplicationUid(uid)) {
+                list.add(Process.toSdkSandboxUid(uid)); // Add Sdk Sandbox UID
+            }
         }
 
         final int minUid = userBase;
         final int maxUid = userBase + UserHandle.PER_USER_RANGE - 1;
-        final Set<Range<Integer>> ranges = new ArraySet<>();
+        final List<Range<Integer>> ranges = new ArrayList<>();
 
         // Iterate the list to create the ranges between each uid.
         int start = minUid;
@@ -1059,6 +1199,10 @@
         return ranges;
     }
 
+    private Set<Range<Integer>> makeVpnUidRangeSet(int userId, List<Integer> excludedAppIdList) {
+        return new ArraySet<>(makeVpnUidRange(userId, excludedAppIdList));
+    }
+
     @Test
     public void testSetAndGetAppExclusionListRestrictedUser() throws Exception {
         final Vpn vpn = prepareVpnForVerifyAppExclusionList();