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