Merge "Fix flaky test: testDataMigration_differentFromFallback" into main
diff --git a/service/src/com/android/server/ConnectivityService.java b/service/src/com/android/server/ConnectivityService.java
index bc703a9..a29f47f 100755
--- a/service/src/com/android/server/ConnectivityService.java
+++ b/service/src/com/android/server/ConnectivityService.java
@@ -2985,19 +2985,17 @@
     }
 
     private void handleFrozenUids(int[] uids, int[] frozenStates) {
-        final ArraySet<Range<Integer>> ranges = new ArraySet<>();
+        final ArraySet<Integer> ownerUids = new ArraySet<>();
 
         for (int i = 0; i < uids.length; i++) {
             if (frozenStates[i] == UID_FROZEN_STATE_FROZEN) {
-                Integer uidAsInteger = Integer.valueOf(uids[i]);
-                ranges.add(new Range(uidAsInteger, uidAsInteger));
+                ownerUids.add(uids[i]);
             }
         }
 
-        if (!ranges.isEmpty()) {
-            final Set<Integer> exemptUids = new ArraySet<>();
+        if (!ownerUids.isEmpty()) {
             try {
-                mDeps.destroyLiveTcpSockets(ranges, exemptUids);
+                mDeps.destroyLiveTcpSocketsByOwnerUids(ownerUids);
             } catch (Exception e) {
                 loge("Exception in socket destroy: " + e);
             }
diff --git a/service/src/com/android/server/connectivity/PermissionMonitor.java b/service/src/com/android/server/connectivity/PermissionMonitor.java
index c15f042..beaa174 100755
--- a/service/src/com/android/server/connectivity/PermissionMonitor.java
+++ b/service/src/com/android/server/connectivity/PermissionMonitor.java
@@ -1011,9 +1011,8 @@
      * @param ranges The updated UID ranges under VPN Lockdown. This function does not treat the VPN
      *               app's UID in any special way. The caller is responsible for excluding the VPN
      *               app UID from the passed-in ranges.
-     *               Ranges can have duplications and/or contain the range that is already subject
-     *               to lockdown. However, ranges can not have overlaps with other ranges including
-     *               ranges that are currently subject to lockdown.
+     *               Ranges can have duplications, overlaps, and/or contain the range that is
+     *               already subject to lockdown.
      */
     public synchronized void updateVpnLockdownUidRanges(boolean add, UidRange[] ranges) {
         final Set<UidRange> affectedUidRanges = new HashSet<>();
@@ -1045,8 +1044,10 @@
         // exclude privileged apps from the prohibit routing rules used to implement outgoing packet
         // filtering, privileged apps can still bypass outgoing packet filtering because the
         // prohibit rules observe the protected from VPN bit.
+        // If removing a UID, we ensure it is not present anywhere in the set first.
         for (final int uid: affectedUids) {
-            if (!hasRestrictedNetworksPermission(uid)) {
+            if (!hasRestrictedNetworksPermission(uid)
+                    && (add || !UidRange.containsUid(mVpnLockdownUidRanges.getSet(), uid))) {
                 updateLockdownUidRule(uid, add);
             }
         }
diff --git a/tests/unit/java/com/android/server/ConnectivityServiceTest.java b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
index 5d957b7..3688d83 100755
--- a/tests/unit/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
@@ -18547,12 +18547,7 @@
 
         waitForIdle();
 
-        final Set<Integer> exemptUids = new ArraySet();
-        final UidRange frozenUidRange = new UidRange(TEST_FROZEN_UID, TEST_FROZEN_UID);
-        final Set<UidRange> ranges = Collections.singleton(frozenUidRange);
-
-        verify(mDestroySocketsWrapper).destroyLiveTcpSockets(eq(UidRange.toIntRanges(ranges)),
-                eq(exemptUids));
+        verify(mDestroySocketsWrapper).destroyLiveTcpSocketsByOwnerUids(Set.of(TEST_FROZEN_UID));
     }
 
     @Test
diff --git a/tests/unit/java/com/android/server/connectivity/PermissionMonitorTest.java b/tests/unit/java/com/android/server/connectivity/PermissionMonitorTest.java
index cf02e3a..8dcfffa 100644
--- a/tests/unit/java/com/android/server/connectivity/PermissionMonitorTest.java
+++ b/tests/unit/java/com/android/server/connectivity/PermissionMonitorTest.java
@@ -55,6 +55,7 @@
 import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.ArgumentMatchers.intThat;
+import static org.mockito.Mockito.atLeast;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.doCallRealMethod;
 import static org.mockito.Mockito.doReturn;
@@ -126,12 +127,14 @@
     private static final int MOCK_APPID1 = 10001;
     private static final int MOCK_APPID2 = 10086;
     private static final int MOCK_APPID3 = 10110;
+    private static final int MOCK_APPID4 = 10111;
     private static final int SYSTEM_APPID1 = 1100;
     private static final int SYSTEM_APPID2 = 1108;
     private static final int VPN_APPID = 10002;
     private static final int MOCK_UID11 = MOCK_USER1.getUid(MOCK_APPID1);
     private static final int MOCK_UID12 = MOCK_USER1.getUid(MOCK_APPID2);
     private static final int MOCK_UID13 = MOCK_USER1.getUid(MOCK_APPID3);
+    private static final int MOCK_UID14 = MOCK_USER1.getUid(MOCK_APPID4);
     private static final int SYSTEM_APP_UID11 = MOCK_USER1.getUid(SYSTEM_APPID1);
     private static final int VPN_UID = MOCK_USER1.getUid(VPN_APPID);
     private static final int MOCK_UID21 = MOCK_USER2.getUid(MOCK_APPID1);
@@ -965,6 +968,66 @@
     }
 
     @Test
+    public void testLockdownUidFilteringWithLockdownEnableDisableWithMultiAddAndOverlap() {
+        doReturn(List.of(buildPackageInfo(SYSTEM_PACKAGE1, SYSTEM_APP_UID11, CHANGE_NETWORK_STATE,
+                        CONNECTIVITY_USE_RESTRICTED_NETWORKS),
+                buildPackageInfo(MOCK_PACKAGE1, MOCK_UID13),
+                buildPackageInfo(MOCK_PACKAGE2, MOCK_UID14),
+                buildPackageInfo(SYSTEM_PACKAGE2, VPN_UID)))
+                .when(mPackageManager).getInstalledPackagesAsUser(eq(GET_PERMISSIONS), anyInt());
+        startMonitoring();
+        // MOCK_UID13 is subject to the VPN.
+        final UidRange range1 = new UidRange(MOCK_UID13, MOCK_UID13);
+        final UidRange[] lockdownRange1 = {range1};
+
+        // Add Lockdown uid range at 1st time, expect a rule to be set up
+        mPermissionMonitor.updateVpnLockdownUidRanges(true /* add */, lockdownRange1);
+        verify(mBpfNetMaps).updateUidLockdownRule(anyInt(), eq(true) /* add */);
+        verify(mBpfNetMaps).updateUidLockdownRule(MOCK_UID13, true /* add */);
+
+        reset(mBpfNetMaps);
+
+        // MOCK_UID13 and MOCK_UID14 are sequential and subject to the VPN in a separate range.
+        final UidRange range2 = new UidRange(MOCK_UID13, MOCK_UID14);
+        final UidRange[] lockdownRange2 = {range2};
+
+        // Add overlapping multiple-UID range. Rule may be set again, which is functionally
+        // a no-op, so it is fine.
+        mPermissionMonitor.updateVpnLockdownUidRanges(true /* add */, lockdownRange2);
+        verify(mBpfNetMaps, atLeast(1)).updateUidLockdownRule(anyInt(), eq(true) /* add */);
+        verify(mBpfNetMaps).updateUidLockdownRule(MOCK_UID14, true /* add */);
+
+        reset(mBpfNetMaps);
+
+        // Remove the multiple-UID range. UID from first rule should not be removed.
+        mPermissionMonitor.updateVpnLockdownUidRanges(false /* false */, lockdownRange2);
+        verify(mBpfNetMaps, times(1)).updateUidLockdownRule(anyInt(), eq(false) /* add */);
+        verify(mBpfNetMaps).updateUidLockdownRule(MOCK_UID14, false /* add */);
+
+        reset(mBpfNetMaps);
+
+        // Add the multiple-UID range back again to be able to test removing the first range, too.
+        mPermissionMonitor.updateVpnLockdownUidRanges(true /* add */, lockdownRange2);
+        verify(mBpfNetMaps, atLeast(1)).updateUidLockdownRule(anyInt(), eq(true) /* add */);
+        verify(mBpfNetMaps).updateUidLockdownRule(MOCK_UID14, true /* add */);
+
+        reset(mBpfNetMaps);
+
+        // Remove the single-UID range. The rule for MOCK_UID11 should not change because it is
+        // still covered by the second, multiple-UID range rule.
+        mPermissionMonitor.updateVpnLockdownUidRanges(false /* false */, lockdownRange1);
+        verify(mBpfNetMaps, never()).updateUidLockdownRule(anyInt(),  anyBoolean());
+
+        reset(mBpfNetMaps);
+
+        // Remove the multiple-UID range. Expect both UID rules to be torn down.
+        mPermissionMonitor.updateVpnLockdownUidRanges(false /* false */, lockdownRange2);
+        verify(mBpfNetMaps, times(2)).updateUidLockdownRule(anyInt(), eq(false) /* add */);
+        verify(mBpfNetMaps).updateUidLockdownRule(MOCK_UID13, false /* add */);
+        verify(mBpfNetMaps).updateUidLockdownRule(MOCK_UID14, false /* add */);
+    }
+
+    @Test
     public void testLockdownUidFilteringWithLockdownEnableDisableWithDuplicates() {
         doReturn(List.of(
                 buildPackageInfo(SYSTEM_PACKAGE1, SYSTEM_APP_UID11, CHANGE_NETWORK_STATE,
diff --git a/tests/unit/java/com/android/server/connectivity/VpnTest.java b/tests/unit/java/com/android/server/connectivity/VpnTest.java
index ad178c0..385f831 100644
--- a/tests/unit/java/com/android/server/connectivity/VpnTest.java
+++ b/tests/unit/java/com/android/server/connectivity/VpnTest.java
@@ -182,7 +182,6 @@
 import com.android.server.VpnTestBase;
 import com.android.server.vcn.util.PersistableBundleUtils;
 import com.android.testutils.DevSdkIgnoreRule;
-import com.android.testutils.SkipPresubmit;
 
 import org.junit.Before;
 import org.junit.Rule;
@@ -1712,7 +1711,6 @@
                 errorCode);
     }
 
-    @SkipPresubmit(reason = "Out of SLO flakiness")
     @Test
     public void testStartPlatformVpnFailedWithRecoverableError() throws Exception {
         final IkeProtocolException exception = mock(IkeProtocolException.class);
@@ -1722,7 +1720,6 @@
                 VpnManager.CATEGORY_EVENT_IKE_ERROR, VpnManager.ERROR_CLASS_RECOVERABLE, errorCode);
     }
 
-    @SkipPresubmit(reason = "Out of SLO flakiness")
     @Test
     public void testStartPlatformVpnFailedWithUnknownHostException() throws Exception {
         final IkeNonProtocolException exception = mock(IkeNonProtocolException.class);
@@ -1734,7 +1731,6 @@
                 errorCode);
     }
 
-    @SkipPresubmit(reason = "Out of SLO flakiness")
     @Test
     public void testStartPlatformVpnFailedWithIkeTimeoutException() throws Exception {
         final IkeNonProtocolException exception = mock(IkeNonProtocolException.class);
@@ -1756,7 +1752,6 @@
                 VpnManager.ERROR_CODE_NETWORK_LOST);
     }
 
-    @SkipPresubmit(reason = "Out of SLO flakiness")
     @Test
     public void testStartPlatformVpnFailedWithIOException() throws Exception {
         final IkeNonProtocolException exception = mock(IkeNonProtocolException.class);
@@ -2389,7 +2384,6 @@
                 true /* areLongLivedTcpConnectionsExpensive */);
     }
 
-    @SkipPresubmit(reason = "Out of SLO flakiness")
     @Test
     public void testPreferredIpProtocolFromCarrierConfig_v4UDP() throws Exception {
         doTestReadCarrierConfig(createTestCellNc(),
@@ -2402,7 +2396,6 @@
                 false /* areLongLivedTcpConnectionsExpensive */);
     }
 
-    @SkipPresubmit(reason = "Out of SLO flakiness")
     @Test
     public void testPreferredIpProtocolFromCarrierConfig_v6ESP() throws Exception {
         doTestReadCarrierConfig(createTestCellNc(),
@@ -2415,7 +2408,6 @@
                 false /* areLongLivedTcpConnectionsExpensive */);
     }
 
-    @SkipPresubmit(reason = "Out of SLO flakiness")
     @Test
     public void testPreferredIpProtocolFromCarrierConfig_v6UDP() throws Exception {
         doTestReadCarrierConfig(createTestCellNc(),