Make some of the UID-based VPN code reusable

By changing some member refs into arguments and having one of the
functions create the UID range instead of adding to mVpnUsers.

This will be useful for other layers of UID filtering like having
UIDs explicitly blocked from the VPN.

Deleted one broken line of code that cleared the status intent when
a restricted profile is removed. Other than that, this commit shouldn't
change any behaviour. If it does, that's a bug.

Bug: 26694104
Change-Id: Ieb656835d3282a8ba63cc3f12a80bfae166bcf44
diff --git a/core/java/android/net/UidRange.java b/core/java/android/net/UidRange.java
index 2e586b3..fd465d9 100644
--- a/core/java/android/net/UidRange.java
+++ b/core/java/android/net/UidRange.java
@@ -48,6 +48,17 @@
         return start / PER_USER_RANGE;
     }
 
+    public boolean contains(int uid) {
+        return start <= uid && uid <= stop;
+    }
+
+    /**
+     * @return {@code true} if this range contains every UID contained by the {@param other} range.
+     */
+    public boolean containsRange(UidRange other) {
+        return start <= other.start && other.stop <= stop;
+    }
+
     @Override
     public int hashCode() {
         int result = 17;
diff --git a/services/core/java/com/android/server/connectivity/Vpn.java b/services/core/java/com/android/server/connectivity/Vpn.java
index 6c1e1a7..8c4e113 100644
--- a/services/core/java/com/android/server/connectivity/Vpn.java
+++ b/services/core/java/com/android/server/connectivity/Vpn.java
@@ -22,6 +22,9 @@
 import static android.net.RouteInfo.RTN_UNREACHABLE;
 
 import android.Manifest;
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.annotation.UserIdInt;
 import android.app.AppGlobals;
 import android.app.AppOpsManager;
 import android.app.PendingIntent;
@@ -67,9 +70,11 @@
 import android.security.Credentials;
 import android.security.KeyStore;
 import android.text.TextUtils;
+import android.util.ArraySet;
 import android.util.Log;
 
 import com.android.internal.annotations.GuardedBy;
+import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.net.LegacyVpnInfo;
 import com.android.internal.net.VpnConfig;
 import com.android.internal.net.VpnInfo;
@@ -88,7 +93,10 @@
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
 import java.util.List;
+import java.util.Set;
 import java.util.SortedSet;
 import java.util.TreeSet;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -119,9 +127,18 @@
     private final Looper mLooper;
     private final NetworkCapabilities mNetworkCapabilities;
 
-    /* list of users using this VPN. */
+    /**
+     * List of UIDs that are set to use this VPN by default. Normally, every UID in the user is
+     * added to this set but that can be changed by adding allowed or disallowed applications. It
+     * is non-null iff the VPN is connected.
+     *
+     * Unless the VPN has set allowBypass=true, these UIDs are forced into the VPN.
+     *
+     * @see VpnService.Builder#addAllowedApplication(String)
+     * @see VpnService.Builder#addDisallowedApplication(String)
+     */
     @GuardedBy("this")
-    private List<UidRange> mVpnUsers = null;
+    private Set<UidRange> mVpnUsers = null;
 
     // Handle of user initiating VPN.
     private final int mUserHandle;
@@ -467,22 +484,8 @@
             Binder.restoreCallingIdentity(token);
         }
 
-        addVpnUserLocked(mUserHandle);
-        // If the user can have restricted profiles, assign all its restricted profiles to this VPN
-        if (canHaveRestrictedProfile(mUserHandle)) {
-            token = Binder.clearCallingIdentity();
-            List<UserInfo> users;
-            try {
-                users = UserManager.get(mContext).getUsers();
-            } finally {
-                Binder.restoreCallingIdentity(token);
-            }
-            for (UserInfo user : users) {
-                if (user.isRestricted() && (user.restrictedProfileParentId == mUserHandle)) {
-                    addVpnUserLocked(user.id);
-                }
-            }
-        }
+        mVpnUsers = createUserAndRestrictedProfilesRanges(mUserHandle,
+                mConfig.allowedApplications, mConfig.disallowedApplications);
         mNetworkAgent.addUidRanges(mVpnUsers.toArray(new UidRange[mVpnUsers.size()]));
 
         mNetworkInfo.setIsAvailable(true);
@@ -568,7 +571,7 @@
         Connection oldConnection = mConnection;
         NetworkAgent oldNetworkAgent = mNetworkAgent;
         mNetworkAgent = null;
-        List<UidRange> oldUsers = mVpnUsers;
+        Set<UidRange> oldUsers = mVpnUsers;
 
         // Configure the interface. Abort if any of these steps fails.
         ParcelFileDescriptor tun = ParcelFileDescriptor.adoptFd(jniCreate(config.mtu));
@@ -601,8 +604,6 @@
             mConfig = config;
 
             // Set up forwarding and DNS rules.
-            mVpnUsers = new ArrayList<UidRange>();
-
             agentConnect();
 
             if (oldConnection != null) {
@@ -657,44 +658,93 @@
         return uids;
     }
 
-    // Note: This function adds to mVpnUsers but does not publish list to NetworkAgent.
-    private void addVpnUserLocked(int userHandle) {
-        if (mVpnUsers == null) {
-            throw new IllegalStateException("VPN is not active");
-        }
+    /**
+     * Creates a {@link Set} of non-intersecting {@link UidRange} objects including all UIDs
+     * associated with one user, and any restricted profiles attached to that user.
+     *
+     * <p>If one of {@param allowedApplications} or {@param disallowedApplications} is provided,
+     * the UID ranges will match the app whitelist or blacklist specified there. Otherwise, all UIDs
+     * in each user and profile will be included.
+     *
+     * @param userHandle The userId to create UID ranges for along with any of its restricted
+     *                   profiles.
+     * @param allowedApplications (optional) whitelist of applications to include.
+     * @param disallowedApplications (optional) blacklist of applications to exclude.
+     */
+    @VisibleForTesting
+    Set<UidRange> createUserAndRestrictedProfilesRanges(@UserIdInt int userHandle,
+            @Nullable List<String> allowedApplications,
+            @Nullable List<String> disallowedApplications) {
+        final Set<UidRange> ranges = new ArraySet<>();
 
-        if (mConfig.allowedApplications != null) {
+        // Assign the top-level user to the set of ranges
+        addUserToRanges(ranges, userHandle, allowedApplications, disallowedApplications);
+
+        // If the user can have restricted profiles, assign all its restricted profiles too
+        if (canHaveRestrictedProfile(userHandle)) {
+            final long token = Binder.clearCallingIdentity();
+            List<UserInfo> users;
+            try {
+                users = UserManager.get(mContext).getUsers();
+            } finally {
+                Binder.restoreCallingIdentity(token);
+            }
+            for (UserInfo user : users) {
+                if (user.isRestricted() && (user.restrictedProfileParentId == userHandle)) {
+                    addUserToRanges(ranges, user.id, allowedApplications, disallowedApplications);
+                }
+            }
+        }
+        return ranges;
+    }
+
+    /**
+     * Updates a {@link Set} of non-intersecting {@link UidRange} objects to include all UIDs
+     * associated with one user.
+     *
+     * <p>If one of {@param allowedApplications} or {@param disallowedApplications} is provided,
+     * the UID ranges will match the app whitelist or blacklist specified there. Otherwise, all UIDs
+     * in the user will be included.
+     *
+     * @param ranges {@link Set} of {@link UidRange}s to which to add.
+     * @param userHandle The userId to add to {@param ranges}.
+     * @param allowedApplications (optional) whitelist of applications to include.
+     * @param disallowedApplications (optional) blacklist of applications to exclude.
+     */
+    @VisibleForTesting
+    void addUserToRanges(@NonNull Set<UidRange> ranges, @UserIdInt int userHandle,
+            @Nullable List<String> allowedApplications,
+            @Nullable List<String> disallowedApplications) {
+        if (allowedApplications != null) {
             // Add ranges covering all UIDs for allowedApplications.
             int start = -1, stop = -1;
-            for (int uid : getAppsUids(mConfig.allowedApplications, userHandle)) {
+            for (int uid : getAppsUids(allowedApplications, userHandle)) {
                 if (start == -1) {
                     start = uid;
                 } else if (uid != stop + 1) {
-                    mVpnUsers.add(new UidRange(start, stop));
+                    ranges.add(new UidRange(start, stop));
                     start = uid;
                 }
                 stop = uid;
             }
-            if (start != -1) mVpnUsers.add(new UidRange(start, stop));
-        } else if (mConfig.disallowedApplications != null) {
+            if (start != -1) ranges.add(new UidRange(start, stop));
+        } else if (disallowedApplications != null) {
             // Add all ranges for user skipping UIDs for disallowedApplications.
             final UidRange userRange = UidRange.createForUser(userHandle);
             int start = userRange.start;
-            for (int uid : getAppsUids(mConfig.disallowedApplications, userHandle)) {
+            for (int uid : getAppsUids(disallowedApplications, userHandle)) {
                 if (uid == start) {
                     start++;
                 } else {
-                    mVpnUsers.add(new UidRange(start, uid - 1));
+                    ranges.add(new UidRange(start, uid - 1));
                     start = uid + 1;
                 }
             }
-            if (start <= userRange.stop) mVpnUsers.add(new UidRange(start, userRange.stop));
+            if (start <= userRange.stop) ranges.add(new UidRange(start, userRange.stop));
         } else {
             // Add all UIDs for the user.
-            mVpnUsers.add(UidRange.createForUser(userHandle));
+            ranges.add(UidRange.createForUser(userHandle));
         }
-
-        prepareStatusIntent();
     }
 
     // Returns the subset of the full list of active UID ranges the VPN applies to (mVpnUsers) that
@@ -703,7 +753,7 @@
         final UidRange userRange = UidRange.createForUser(userHandle);
         final List<UidRange> ranges = new ArrayList<UidRange>();
         for (UidRange range : mVpnUsers) {
-            if (range.start >= userRange.start && range.stop <= userRange.stop) {
+            if (userRange.containsRange(range)) {
                 ranges.add(range);
             }
         }
@@ -719,7 +769,6 @@
             mNetworkAgent.removeUidRanges(ranges.toArray(new UidRange[ranges.size()]));
         }
         mVpnUsers.removeAll(ranges);
-        mStatusIntent = null;
     }
 
     public void onUserAdded(int userHandle) {
@@ -729,7 +778,8 @@
                 && mVpnUsers != null) {
             synchronized(Vpn.this) {
                 try {
-                    addVpnUserLocked(userHandle);
+                    addUserToRanges(mVpnUsers, userHandle, mConfig.allowedApplications,
+                            mConfig.disallowedApplications);
                     if (mNetworkAgent != null) {
                         final List<UidRange> ranges = uidRangesForUser(userHandle);
                         mNetworkAgent.addUidRanges(ranges.toArray(new UidRange[ranges.size()]));
@@ -902,7 +952,7 @@
             return false;
         }
         for (UidRange uidRange : mVpnUsers) {
-            if (uidRange.start <= uid && uid <= uidRange.stop) {
+            if (uidRange.contains(uid)) {
                 return true;
             }
         }
@@ -1408,7 +1458,7 @@
 
                     // Now INetworkManagementEventObserver is watching our back.
                     mInterface = mConfig.interfaze;
-                    mVpnUsers = new ArrayList<UidRange>();
+                    prepareStatusIntent();
 
                     agentConnect();
 
diff --git a/services/tests/servicestests/src/com/android/server/connectivity/VpnTest.java b/services/tests/servicestests/src/com/android/server/connectivity/VpnTest.java
new file mode 100644
index 0000000..3295bf5
--- /dev/null
+++ b/services/tests/servicestests/src/com/android/server/connectivity/VpnTest.java
@@ -0,0 +1,198 @@
+/*
+ * Copyright (C) 2016 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.connectivity;
+
+import static android.content.pm.UserInfo.FLAG_ADMIN;
+import static android.content.pm.UserInfo.FLAG_MANAGED_PROFILE;
+import static android.content.pm.UserInfo.FLAG_PRIMARY;
+import static android.content.pm.UserInfo.FLAG_RESTRICTED;
+import static org.mockito.Mockito.*;
+
+import android.annotation.UserIdInt;
+import android.content.Context;
+import android.content.pm.PackageManager;
+import android.content.pm.UserInfo;
+import android.net.UidRange;
+import android.os.INetworkManagementService;
+import android.os.Looper;
+import android.os.UserHandle;
+import android.os.UserManager;
+import android.test.AndroidTestCase;
+import android.test.suitebuilder.annotation.SmallTest;
+import android.util.ArrayMap;
+import android.util.ArraySet;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Map;
+import java.util.Set;
+
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+/**
+ * Tests for {@link Vpn}.
+ *
+ * Build, install and run with:
+ *  runtest --path src/com/android/server/connectivity/VpnTest.java
+ */
+public class VpnTest extends AndroidTestCase {
+    private static final String TAG = "VpnTest";
+
+    // Mock users
+    static final UserInfo primaryUser = new UserInfo(27, "Primary", FLAG_ADMIN | FLAG_PRIMARY);
+    static final UserInfo secondaryUser = new UserInfo(15, "Secondary", FLAG_ADMIN);
+    static final UserInfo restrictedProfileA = new UserInfo(40, "RestrictedA", FLAG_RESTRICTED);
+    static final UserInfo restrictedProfileB = new UserInfo(42, "RestrictedB", FLAG_RESTRICTED);
+    static final UserInfo managedProfileA = new UserInfo(45, "ManagedA", FLAG_MANAGED_PROFILE);
+    static {
+        restrictedProfileA.restrictedProfileParentId = primaryUser.id;
+        restrictedProfileB.restrictedProfileParentId = secondaryUser.id;
+        managedProfileA.profileGroupId = primaryUser.id;
+    }
+
+    @Mock private Context mContext;
+    @Mock private UserManager mUserManager;
+    @Mock private PackageManager mPackageManager;
+    @Mock private INetworkManagementService mNetService;
+
+    @Override
+    public void setUp() throws Exception {
+        MockitoAnnotations.initMocks(this);
+        when(mContext.getPackageManager()).thenReturn(mPackageManager);
+        when(mContext.getSystemService(eq(Context.USER_SERVICE))).thenReturn(mUserManager);
+        doNothing().when(mNetService).registerObserver(any());
+    }
+
+    @SmallTest
+    public void testRestrictedProfilesAreAddedToVpn() {
+        setMockedUsers(primaryUser, secondaryUser, restrictedProfileA, restrictedProfileB);
+
+        final Vpn vpn = createVpn(primaryUser.id);
+        final Set<UidRange> ranges = vpn.createUserAndRestrictedProfilesRanges(primaryUser.id,
+                null, null);
+
+        assertEquals(new ArraySet<>(Arrays.asList(new UidRange[] {
+            UidRange.createForUser(primaryUser.id),
+            UidRange.createForUser(restrictedProfileA.id)
+        })), ranges);
+    }
+
+    @SmallTest
+    public void testManagedProfilesAreNotAddedToVpn() {
+        setMockedUsers(primaryUser, managedProfileA);
+
+        final Vpn vpn = createVpn(primaryUser.id);
+        final Set<UidRange> ranges = vpn.createUserAndRestrictedProfilesRanges(primaryUser.id,
+                null, null);
+
+        assertEquals(new ArraySet<>(Arrays.asList(new UidRange[] {
+            UidRange.createForUser(primaryUser.id)
+        })), ranges);
+    }
+
+    @SmallTest
+    public void testAddUserToVpnOnlyAddsOneUser() {
+        setMockedUsers(primaryUser, restrictedProfileA, managedProfileA);
+
+        final Vpn vpn = createVpn(primaryUser.id);
+        final Set<UidRange> ranges = new ArraySet<>();
+        vpn.addUserToRanges(ranges, primaryUser.id, null, null);
+
+        assertEquals(new ArraySet<>(Arrays.asList(new UidRange[] {
+            UidRange.createForUser(primaryUser.id)
+        })), ranges);
+    }
+
+    @SmallTest
+    public void testUidWhiteAndBlacklist() throws Exception {
+        final Map<String, Integer> packages = new ArrayMap<>();
+        packages.put("com.example", 66);
+        packages.put("org.example", 77);
+        packages.put("net.example", 78);
+        setMockedPackages(packages);
+
+        final Vpn vpn = createVpn(primaryUser.id);
+        final UidRange user = UidRange.createForUser(primaryUser.id);
+
+        // Whitelist
+        final Set<UidRange> allow = vpn.createUserAndRestrictedProfilesRanges(primaryUser.id,
+                new ArrayList<String>(packages.keySet()), null);
+        assertEquals(new ArraySet<>(Arrays.asList(new UidRange[] {
+            new UidRange(user.start + 66, user.start + 66),
+            new UidRange(user.start + 77, user.start + 78)
+        })), allow);
+
+        // Blacklist
+        final Set<UidRange> disallow = vpn.createUserAndRestrictedProfilesRanges(primaryUser.id,
+                null, new ArrayList<String>(packages.keySet()));
+        assertEquals(new ArraySet<>(Arrays.asList(new UidRange[] {
+            new UidRange(user.start, user.start + 65),
+            new UidRange(user.start + 67, user.start + 76),
+            new UidRange(user.start + 79, user.stop)
+        })), disallow);
+    }
+
+    /**
+     * @return A subclass of {@link Vpn} which is reliably:
+     * <ul>
+     *   <li>Associated with a specific user ID</li>
+     *   <li>Not in always-on mode</li>
+     * </ul>
+     */
+    private Vpn createVpn(@UserIdInt int userId) {
+        return new Vpn(Looper.myLooper(), mContext, mNetService, userId);
+    }
+
+    /**
+     * Populate {@link #mUserManager} with a list of fake users.
+     */
+    private void setMockedUsers(UserInfo... users) {
+        final Map<Integer, UserInfo> userMap = new ArrayMap<>();
+        for (UserInfo user : users) {
+            userMap.put(user.id, user);
+        }
+
+        doAnswer(invocation -> {
+            return new ArrayList(userMap.values());
+        }).when(mUserManager).getUsers();
+
+        doAnswer(invocation -> {
+            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(anyInt());
+    }
+
+    /**
+     * Populate {@link #mPackageManager} with a fake packageName-to-UID mapping.
+     */
+    private void setMockedPackages(final Map<String, Integer> packages) {
+        try {
+            doAnswer(invocation -> {
+                final String appName = (String) invocation.getArguments()[0];
+                final int userId = (int) invocation.getArguments()[1];
+                return UserHandle.getUid(userId, packages.get(appName));
+            }).when(mPackageManager).getPackageUidAsUser(anyString(), anyInt());
+        } catch (Exception e) {
+        }
+    }
+}