From 99fb6920f69d31088becad6fd5f39c7200c2c6b9 Mon Sep 17 00:00:00 2001
From: Lorenzo Colitti <lorenzo@google.com>
Date: Thu, 19 Nov 2020 18:55:12 +0900
Subject: [PATCH] Add a test for restricted profile added/removed with VPN up.

Bug: 173331190
Test: test-only change
Change-Id: Ibbf96a259a73068d110a159d54059720121117cc
---
 .../server/ConnectivityServiceTest.java       | 71 +++++++++++++++++++
 1 file changed, 71 insertions(+)

diff --git a/tests/net/java/com/android/server/ConnectivityServiceTest.java b/tests/net/java/com/android/server/ConnectivityServiceTest.java
index 6293befb2c..0595ac9032 100644
--- a/tests/net/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/net/java/com/android/server/ConnectivityServiceTest.java
@@ -18,6 +18,8 @@ package com.android.server;
 
 import static android.Manifest.permission.CHANGE_NETWORK_STATE;
 import static android.Manifest.permission.CONNECTIVITY_USE_RESTRICTED_NETWORKS;
+import static android.content.Intent.ACTION_USER_ADDED;
+import static android.content.Intent.ACTION_USER_REMOVED;
 import static android.content.pm.PackageInfo.REQUESTED_PERMISSION_GRANTED;
 import static android.content.pm.PackageManager.GET_PERMISSIONS;
 import static android.content.pm.PackageManager.MATCH_ANY_USER;
@@ -5880,6 +5882,75 @@ public class ConnectivityServiceTest {
         mMockVpn.disconnect();
     }
 
+    @Test
+    public void testVpnRestrictedUsers() throws Exception {
+        // NETWORK_SETTINGS is necessary to see the UID ranges in NetworkCapabilities.
+        mServiceContext.setPermission(Manifest.permission.NETWORK_SETTINGS,
+                PERMISSION_GRANTED);
+
+        final NetworkRequest request = new NetworkRequest.Builder()
+                .removeCapability(NET_CAPABILITY_NOT_VPN)
+                .build();
+        final TestNetworkCallback callback = new TestNetworkCallback();
+        mCm.registerNetworkCallback(request, callback);
+
+        // Bring up a VPN
+        mMockVpn.establishForMyUid();
+        callback.expectAvailableThenValidatedCallbacks(mMockVpn);
+        callback.assertNoCallback();
+
+        final int uid = Process.myUid();
+        NetworkCapabilities nc = mCm.getNetworkCapabilities(mMockVpn.getNetwork());
+        assertNotNull("nc=" + nc, nc.getUids());
+        assertEquals(nc.getUids(), uidRangesForUid(uid));
+
+        // Set an underlying network and expect to see the VPN transports change.
+        mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
+        mWiFiNetworkAgent.connect(true);
+        callback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
+        callback.expectCapabilitiesThat(mMockVpn, (caps)
+                -> caps.hasTransport(TRANSPORT_VPN)
+                && caps.hasTransport(TRANSPORT_WIFI));
+        callback.expectCapabilitiesThat(mWiFiNetworkAgent, (caps)
+                -> caps.hasCapability(NET_CAPABILITY_VALIDATED));
+
+        // Create a fake restricted profile whose parent is our user ID.
+        final int userId = UserHandle.getUserId(uid);
+        final int restrictedUserId = userId + 1;
+        final UserInfo info = new UserInfo(restrictedUserId, "user", UserInfo.FLAG_RESTRICTED);
+        info.restrictedProfileParentId = userId;
+        assertTrue(info.isRestricted());
+        when(mUserManager.getUserInfo(restrictedUserId)).thenReturn(info);
+        final Intent addedIntent = new Intent(ACTION_USER_ADDED);
+        addedIntent.putExtra(Intent.EXTRA_USER_HANDLE, restrictedUserId);
+
+        // Send a USER_ADDED broadcast for it.
+        // The BroadcastReceiver for this broadcast checks that is being run on the handler thread.
+        final Handler handler = new Handler(mCsHandlerThread.getLooper());
+        handler.post(() -> mServiceContext.sendBroadcast(addedIntent));
+
+        // Expect that the VPN UID ranges contain both |uid| and the UID range for the newly-added
+        // restricted user.
+        callback.expectCapabilitiesThat(mMockVpn, (caps)
+                -> caps.getUids().size() == 2
+                && caps.getUids().contains(new UidRange(uid, uid))
+                && caps.getUids().contains(UidRange.createForUser(restrictedUserId))
+                && caps.hasTransport(TRANSPORT_VPN)
+                && caps.hasTransport(TRANSPORT_WIFI));
+
+        // Send a USER_REMOVED broadcast and expect to lose the UID range for the restricted user.
+        final Intent removedIntent = new Intent(ACTION_USER_REMOVED);
+        removedIntent.putExtra(Intent.EXTRA_USER_HANDLE, restrictedUserId);
+        handler.post(() -> mServiceContext.sendBroadcast(removedIntent));
+
+        // Expect that the VPN gains the UID range for the restricted user.
+        callback.expectCapabilitiesThat(mMockVpn, (caps)
+                -> caps.getUids().size() == 1
+                && caps.getUids().contains(new UidRange(uid, uid))
+                && caps.hasTransport(TRANSPORT_VPN)
+                && caps.hasTransport(TRANSPORT_WIFI));
+    }
+
     @Test
     public void testIsActiveNetworkMeteredOverWifi() throws Exception {
         // Returns true by default when no network is available.