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.AdditionalMatchers.*;
24import static org.mockito.Mockito.*;
25
26import android.annotation.UserIdInt;
27import android.app.AppOpsManager;
28import android.content.Context;
29import android.content.pm.PackageManager;
30import android.content.pm.UserInfo;
31import android.net.UidRange;
32import android.os.INetworkManagementService;
33import android.os.Looper;
34import android.os.UserHandle;
35import android.os.UserManager;
36import android.test.AndroidTestCase;
37import android.test.suitebuilder.annotation.SmallTest;
38import android.util.ArrayMap;
39import android.util.ArraySet;
40
41import java.util.ArrayList;
42import java.util.Arrays;
43import java.util.Map;
44import java.util.Set;
45
46import org.mockito.Mock;
47import org.mockito.MockitoAnnotations;
48
49/**
50 * Tests for {@link Vpn}.
51 *
52 * Build, install and run with:
53 *  runtest --path src/com/android/server/connectivity/VpnTest.java
54 */
55public class VpnTest extends AndroidTestCase {
56    private static final String TAG = "VpnTest";
57
58    // Mock users
59    static final UserInfo primaryUser = new UserInfo(27, "Primary", FLAG_ADMIN | FLAG_PRIMARY);
60    static final UserInfo secondaryUser = new UserInfo(15, "Secondary", FLAG_ADMIN);
61    static final UserInfo restrictedProfileA = new UserInfo(40, "RestrictedA", FLAG_RESTRICTED);
62    static final UserInfo restrictedProfileB = new UserInfo(42, "RestrictedB", FLAG_RESTRICTED);
63    static final UserInfo managedProfileA = new UserInfo(45, "ManagedA", FLAG_MANAGED_PROFILE);
64    static {
65        restrictedProfileA.restrictedProfileParentId = primaryUser.id;
66        restrictedProfileB.restrictedProfileParentId = secondaryUser.id;
67        managedProfileA.profileGroupId = primaryUser.id;
68    }
69
70    /**
71     * Names and UIDs for some fake packages. Important points:
72     *  - UID is ordered increasing.
73     *  - One pair of packages have consecutive UIDs.
74     */
75    static final String[] PKGS = {"com.example", "org.example", "net.example", "web.vpn"};
76    static final int[] PKG_UIDS = {66, 77, 78, 400};
77
78    // Mock packages
79    static final Map<String, Integer> mPackages = new ArrayMap<>();
80    static {
81        for (int i = 0; i < PKGS.length; i++) {
82            mPackages.put(PKGS[i], PKG_UIDS[i]);
83        }
84    }
85
86    @Mock private Context mContext;
87    @Mock private UserManager mUserManager;
88    @Mock private PackageManager mPackageManager;
89    @Mock private INetworkManagementService mNetService;
90    @Mock private AppOpsManager mAppOps;
91
92    @Override
93    public void setUp() throws Exception {
94        MockitoAnnotations.initMocks(this);
95        when(mContext.getPackageManager()).thenReturn(mPackageManager);
96        setMockedPackages(mPackages);
97        when(mContext.getSystemService(eq(Context.USER_SERVICE))).thenReturn(mUserManager);
98        when(mContext.getSystemService(eq(Context.APP_OPS_SERVICE))).thenReturn(mAppOps);
99        doNothing().when(mNetService).registerObserver(any());
100    }
101
102    @SmallTest
103    public void testRestrictedProfilesAreAddedToVpn() {
104        setMockedUsers(primaryUser, secondaryUser, restrictedProfileA, restrictedProfileB);
105
106        final Vpn vpn = new MockVpn(primaryUser.id);
107        final Set<UidRange> ranges = vpn.createUserAndRestrictedProfilesRanges(primaryUser.id,
108                null, null);
109
110        assertEquals(new ArraySet<>(Arrays.asList(new UidRange[] {
111            UidRange.createForUser(primaryUser.id),
112            UidRange.createForUser(restrictedProfileA.id)
113        })), ranges);
114    }
115
116    @SmallTest
117    public void testManagedProfilesAreNotAddedToVpn() {
118        setMockedUsers(primaryUser, managedProfileA);
119
120        final Vpn vpn = new MockVpn(primaryUser.id);
121        final Set<UidRange> ranges = vpn.createUserAndRestrictedProfilesRanges(primaryUser.id,
122                null, null);
123
124        assertEquals(new ArraySet<>(Arrays.asList(new UidRange[] {
125            UidRange.createForUser(primaryUser.id)
126        })), ranges);
127    }
128
129    @SmallTest
130    public void testAddUserToVpnOnlyAddsOneUser() {
131        setMockedUsers(primaryUser, restrictedProfileA, managedProfileA);
132
133        final Vpn vpn = new MockVpn(primaryUser.id);
134        final Set<UidRange> ranges = new ArraySet<>();
135        vpn.addUserToRanges(ranges, primaryUser.id, null, null);
136
137        assertEquals(new ArraySet<>(Arrays.asList(new UidRange[] {
138            UidRange.createForUser(primaryUser.id)
139        })), ranges);
140    }
141
142    @SmallTest
143    public void testUidWhiteAndBlacklist() throws Exception {
144        final Vpn vpn = new MockVpn(primaryUser.id);
145        final UidRange user = UidRange.createForUser(primaryUser.id);
146        final String[] packages = {PKGS[0], PKGS[1], PKGS[2]};
147
148        // Whitelist
149        final Set<UidRange> allow = vpn.createUserAndRestrictedProfilesRanges(primaryUser.id,
150                Arrays.asList(packages), null);
151        assertEquals(new ArraySet<>(Arrays.asList(new UidRange[] {
152            new UidRange(user.start + PKG_UIDS[0], user.start + PKG_UIDS[0]),
153            new UidRange(user.start + PKG_UIDS[1], user.start + PKG_UIDS[2])
154        })), allow);
155
156        // Blacklist
157        final Set<UidRange> disallow = vpn.createUserAndRestrictedProfilesRanges(primaryUser.id,
158                null, Arrays.asList(packages));
159        assertEquals(new ArraySet<>(Arrays.asList(new UidRange[] {
160            new UidRange(user.start, user.start + PKG_UIDS[0] - 1),
161            new UidRange(user.start + PKG_UIDS[0] + 1, user.start + PKG_UIDS[1] - 1),
162            /* Empty range between UIDS[1] and UIDS[2], should be excluded, */
163            new UidRange(user.start + PKG_UIDS[2] + 1, user.stop)
164        })), disallow);
165    }
166
167    @SmallTest
168    public void testLockdownChangingPackage() throws Exception {
169        final MockVpn vpn = new MockVpn(primaryUser.id);
170        final UidRange user = UidRange.createForUser(primaryUser.id);
171
172        // Default state.
173        vpn.assertUnblocked(user.start + PKG_UIDS[0], user.start + PKG_UIDS[1], user.start + PKG_UIDS[2], user.start + PKG_UIDS[3]);
174
175        // Set always-on without lockdown.
176        assertTrue(vpn.setAlwaysOnPackage(PKGS[1], false));
177        vpn.assertUnblocked(user.start + PKG_UIDS[0], user.start + PKG_UIDS[1], user.start + PKG_UIDS[2], user.start + PKG_UIDS[3]);
178
179        // Set always-on with lockdown.
180        assertTrue(vpn.setAlwaysOnPackage(PKGS[1], true));
181        verify(mNetService).setAllowOnlyVpnForUids(eq(true), aryEq(new UidRange[] {
182            new UidRange(user.start, user.start + PKG_UIDS[1] - 1),
183            new UidRange(user.start + PKG_UIDS[1] + 1, user.stop)
184        }));
185        vpn.assertBlocked(user.start + PKG_UIDS[0], user.start + PKG_UIDS[2], user.start + PKG_UIDS[3]);
186        vpn.assertUnblocked(user.start + PKG_UIDS[1]);
187
188        // Switch to another app.
189        assertTrue(vpn.setAlwaysOnPackage(PKGS[3], true));
190        verify(mNetService).setAllowOnlyVpnForUids(eq(false), aryEq(new UidRange[] {
191            new UidRange(user.start, user.start + PKG_UIDS[1] - 1),
192            new UidRange(user.start + PKG_UIDS[1] + 1, user.stop)
193        }));
194        verify(mNetService).setAllowOnlyVpnForUids(eq(true), aryEq(new UidRange[] {
195            new UidRange(user.start, user.start + PKG_UIDS[3] - 1),
196            new UidRange(user.start + PKG_UIDS[3] + 1, user.stop)
197        }));
198        vpn.assertBlocked(user.start + PKG_UIDS[0], user.start + PKG_UIDS[1], user.start + PKG_UIDS[2]);
199        vpn.assertUnblocked(user.start + PKG_UIDS[3]);
200    }
201
202    @SmallTest
203    public void testLockdownAddingAProfile() throws Exception {
204        final MockVpn vpn = new MockVpn(primaryUser.id);
205        setMockedUsers(primaryUser);
206
207        // Make a copy of the restricted profile, as we're going to mark it deleted halfway through.
208        final UserInfo tempProfile = new UserInfo(restrictedProfileA.id, restrictedProfileA.name,
209                restrictedProfileA.flags);
210        tempProfile.restrictedProfileParentId = primaryUser.id;
211
212        final UidRange user = UidRange.createForUser(primaryUser.id);
213        final UidRange profile = UidRange.createForUser(tempProfile.id);
214
215        // Set lockdown.
216        assertTrue(vpn.setAlwaysOnPackage(PKGS[3], true));
217        verify(mNetService).setAllowOnlyVpnForUids(eq(true), aryEq(new UidRange[] {
218            new UidRange(user.start, user.start + PKG_UIDS[3] - 1),
219            new UidRange(user.start + PKG_UIDS[3] + 1, user.stop)
220        }));
221
222        // Verify restricted user isn't affected at first.
223        vpn.assertUnblocked(profile.start + PKG_UIDS[0]);
224
225        // Add the restricted user.
226        setMockedUsers(primaryUser, tempProfile);
227        vpn.onUserAdded(tempProfile.id);
228        verify(mNetService).setAllowOnlyVpnForUids(eq(true), aryEq(new UidRange[] {
229            new UidRange(profile.start, profile.start + PKG_UIDS[3] - 1),
230            new UidRange(profile.start + PKG_UIDS[3] + 1, profile.stop)
231        }));
232
233        // Remove the restricted user.
234        tempProfile.partial = true;
235        vpn.onUserRemoved(tempProfile.id);
236        verify(mNetService).setAllowOnlyVpnForUids(eq(false), aryEq(new UidRange[] {
237            new UidRange(profile.start, profile.start + PKG_UIDS[3] - 1),
238            new UidRange(profile.start + PKG_UIDS[3] + 1, profile.stop)
239        }));
240    }
241
242    /**
243     * A subclass of {@link Vpn} with some of the fields pre-mocked.
244     */
245    private class MockVpn extends Vpn {
246        public MockVpn(@UserIdInt int userId) {
247            super(Looper.myLooper(), mContext, mNetService, userId);
248        }
249
250        public void assertBlocked(int... uids) {
251            for (int uid : uids) {
252                assertTrue("Uid " + uid + " should be blocked", isBlockingUid(uid));
253            }
254        }
255
256        public void assertUnblocked(int... uids) {
257            for (int uid : uids) {
258                assertFalse("Uid " + uid + " should not be blocked", isBlockingUid(uid));
259            }
260        }
261    }
262
263    /**
264     * Populate {@link #mUserManager} with a list of fake users.
265     */
266    private void setMockedUsers(UserInfo... users) {
267        final Map<Integer, UserInfo> userMap = new ArrayMap<>();
268        for (UserInfo user : users) {
269            userMap.put(user.id, user);
270        }
271
272        /**
273         * @see UserManagerService#getUsers(boolean)
274         */
275        doAnswer(invocation -> {
276            final boolean excludeDying = (boolean) invocation.getArguments()[0];
277            final ArrayList<UserInfo> result = new ArrayList<>(users.length);
278            for (UserInfo ui : users) {
279                if (!excludeDying || (ui.isEnabled() && !ui.partial)) {
280                    result.add(ui);
281                }
282            }
283            return result;
284        }).when(mUserManager).getUsers(anyBoolean());
285
286        doAnswer(invocation -> {
287            final int id = (int) invocation.getArguments()[0];
288            return userMap.get(id);
289        }).when(mUserManager).getUserInfo(anyInt());
290
291        doAnswer(invocation -> {
292            final int id = (int) invocation.getArguments()[0];
293            return (userMap.get(id).flags & UserInfo.FLAG_ADMIN) != 0;
294        }).when(mUserManager).canHaveRestrictedProfile(anyInt());
295    }
296
297    /**
298     * Populate {@link #mPackageManager} with a fake packageName-to-UID mapping.
299     */
300    private void setMockedPackages(final Map<String, Integer> packages) {
301        try {
302            doAnswer(invocation -> {
303                final String appName = (String) invocation.getArguments()[0];
304                final int userId = (int) invocation.getArguments()[1];
305                return UserHandle.getUid(userId, packages.get(appName));
306            }).when(mPackageManager).getPackageUidAsUser(anyString(), anyInt());
307        } catch (Exception e) {
308        }
309    }
310}
311