VpnTest.java revision 4d03abcd49af490dba3850d341b955dd72f24959 
1/*
2 * Copyright (C) 2016 The Android Open Source Project
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *      http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17package com.android.server.connectivity;
18
19import static android.content.pm.UserInfo.FLAG_ADMIN;
20import static android.content.pm.UserInfo.FLAG_MANAGED_PROFILE;
21import static android.content.pm.UserInfo.FLAG_PRIMARY;
22import static android.content.pm.UserInfo.FLAG_RESTRICTED;
23import static org.mockito.Mockito.*;
24
25import android.annotation.UserIdInt;
26import android.content.Context;
27import android.content.pm.PackageManager;
28import android.content.pm.UserInfo;
29import android.net.UidRange;
30import android.os.INetworkManagementService;
31import android.os.Looper;
32import android.os.UserHandle;
33import android.os.UserManager;
34import android.test.AndroidTestCase;
35import android.test.suitebuilder.annotation.SmallTest;
36import android.util.ArrayMap;
37import android.util.ArraySet;
38
39import java.util.ArrayList;
40import java.util.Arrays;
41import java.util.Map;
42import java.util.Set;
43
44import org.mockito.Mock;
45import org.mockito.MockitoAnnotations;
46
47/**
48 * Tests for {@link Vpn}.
49 *
50 * Build, install and run with:
51 *  runtest --path src/com/android/server/connectivity/VpnTest.java
52 */
53public class VpnTest extends AndroidTestCase {
54    private static final String TAG = "VpnTest";
55
56    // Mock users
57    static final UserInfo primaryUser = new UserInfo(27, "Primary", FLAG_ADMIN | FLAG_PRIMARY);
58    static final UserInfo secondaryUser = new UserInfo(15, "Secondary", FLAG_ADMIN);
59    static final UserInfo restrictedProfileA = new UserInfo(40, "RestrictedA", FLAG_RESTRICTED);
60    static final UserInfo restrictedProfileB = new UserInfo(42, "RestrictedB", FLAG_RESTRICTED);
61    static final UserInfo managedProfileA = new UserInfo(45, "ManagedA", FLAG_MANAGED_PROFILE);
62    static {
63        restrictedProfileA.restrictedProfileParentId = primaryUser.id;
64        restrictedProfileB.restrictedProfileParentId = secondaryUser.id;
65        managedProfileA.profileGroupId = primaryUser.id;
66    }
67
68    @Mock private Context mContext;
69    @Mock private UserManager mUserManager;
70    @Mock private PackageManager mPackageManager;
71    @Mock private INetworkManagementService mNetService;
72
73    @Override
74    public void setUp() throws Exception {
75        MockitoAnnotations.initMocks(this);
76        when(mContext.getPackageManager()).thenReturn(mPackageManager);
77        when(mContext.getSystemService(eq(Context.USER_SERVICE))).thenReturn(mUserManager);
78        doNothing().when(mNetService).registerObserver(any());
79    }
80
81    @SmallTest
82    public void testRestrictedProfilesAreAddedToVpn() {
83        setMockedUsers(primaryUser, secondaryUser, restrictedProfileA, restrictedProfileB);
84
85        final Vpn vpn = createVpn(primaryUser.id);
86        final Set<UidRange> ranges = vpn.createUserAndRestrictedProfilesRanges(primaryUser.id,
87                null, null);
88
89        assertEquals(new ArraySet<>(Arrays.asList(new UidRange[] {
90            UidRange.createForUser(primaryUser.id),
91            UidRange.createForUser(restrictedProfileA.id)
92        })), ranges);
93    }
94
95    @SmallTest
96    public void testManagedProfilesAreNotAddedToVpn() {
97        setMockedUsers(primaryUser, managedProfileA);
98
99        final Vpn vpn = createVpn(primaryUser.id);
100        final Set<UidRange> ranges = vpn.createUserAndRestrictedProfilesRanges(primaryUser.id,
101                null, null);
102
103        assertEquals(new ArraySet<>(Arrays.asList(new UidRange[] {
104            UidRange.createForUser(primaryUser.id)
105        })), ranges);
106    }
107
108    @SmallTest
109    public void testAddUserToVpnOnlyAddsOneUser() {
110        setMockedUsers(primaryUser, restrictedProfileA, managedProfileA);
111
112        final Vpn vpn = createVpn(primaryUser.id);
113        final Set<UidRange> ranges = new ArraySet<>();
114        vpn.addUserToRanges(ranges, primaryUser.id, null, null);
115
116        assertEquals(new ArraySet<>(Arrays.asList(new UidRange[] {
117            UidRange.createForUser(primaryUser.id)
118        })), ranges);
119    }
120
121    @SmallTest
122    public void testUidWhiteAndBlacklist() throws Exception {
123        final Map<String, Integer> packages = new ArrayMap<>();
124        packages.put("com.example", 66);
125        packages.put("org.example", 77);
126        packages.put("net.example", 78);
127        setMockedPackages(packages);
128
129        final Vpn vpn = createVpn(primaryUser.id);
130        final UidRange user = UidRange.createForUser(primaryUser.id);
131
132        // Whitelist
133        final Set<UidRange> allow = vpn.createUserAndRestrictedProfilesRanges(primaryUser.id,
134                new ArrayList<String>(packages.keySet()), null);
135        assertEquals(new ArraySet<>(Arrays.asList(new UidRange[] {
136            new UidRange(user.start + 66, user.start + 66),
137            new UidRange(user.start + 77, user.start + 78)
138        })), allow);
139
140        // Blacklist
141        final Set<UidRange> disallow = vpn.createUserAndRestrictedProfilesRanges(primaryUser.id,
142                null, new ArrayList<String>(packages.keySet()));
143        assertEquals(new ArraySet<>(Arrays.asList(new UidRange[] {
144            new UidRange(user.start, user.start + 65),
145            new UidRange(user.start + 67, user.start + 76),
146            new UidRange(user.start + 79, user.stop)
147        })), disallow);
148    }
149
150    /**
151     * @return A subclass of {@link Vpn} which is reliably:
152     * <ul>
153     *   <li>Associated with a specific user ID</li>
154     *   <li>Not in always-on mode</li>
155     * </ul>
156     */
157    private Vpn createVpn(@UserIdInt int userId) {
158        return new Vpn(Looper.myLooper(), mContext, mNetService, userId);
159    }
160
161    /**
162     * Populate {@link #mUserManager} with a list of fake users.
163     */
164    private void setMockedUsers(UserInfo... users) {
165        final Map<Integer, UserInfo> userMap = new ArrayMap<>();
166        for (UserInfo user : users) {
167            userMap.put(user.id, user);
168        }
169
170        doAnswer(invocation -> {
171            return new ArrayList(userMap.values());
172        }).when(mUserManager).getUsers();
173
174        doAnswer(invocation -> {
175            final int id = (int) invocation.getArguments()[0];
176            return userMap.get(id);
177        }).when(mUserManager).getUserInfo(anyInt());
178
179        doAnswer(invocation -> {
180            final int id = (int) invocation.getArguments()[0];
181            return (userMap.get(id).flags & UserInfo.FLAG_ADMIN) != 0;
182        }).when(mUserManager).canHaveRestrictedProfile(anyInt());
183    }
184
185    /**
186     * Populate {@link #mPackageManager} with a fake packageName-to-UID mapping.
187     */
188    private void setMockedPackages(final Map<String, Integer> packages) {
189        try {
190            doAnswer(invocation -> {
191                final String appName = (String) invocation.getArguments()[0];
192                final int userId = (int) invocation.getArguments()[1];
193                return UserHandle.getUid(userId, packages.get(appName));
194            }).when(mPackageManager).getPackageUidAsUser(anyString(), anyInt());
195        } catch (Exception e) {
196        }
197    }
198}
199