Have the status bar listen to all macro-users VPN changes.

Bug: 73217368
Test: manual
      Using Datally on work profile. Before this, enabling the VPN
      does not show the key icon. After this it does.

Change-Id: I454eb8f3881a48af1b0187c2b14a2a399d3c2445

3 files changed, 32 insertions, 1 deletions

diff --git a/core/java/android/net/NetworkRequest.java b/core/java/android/net/NetworkRequest.java
index fdcc304e1d2e..4f92fa6a7340 100644
--- a/core/java/android/net/NetworkRequest.java
+++ b/core/java/android/net/NetworkRequest.java
@@ -24,6 +24,7 @@ import android.text.TextUtils;
 import android.util.proto.ProtoOutputStream;
 
 import java.util.Objects;
+import java.util.Set;
 
 /**
  * Defines a request for a network, made through {@link NetworkRequest.Builder} and used
@@ -205,6 +206,19 @@ public class NetworkRequest implements Parcelable {
         }
 
         /**
+         * Set the watched UIDs for this request. This will be reset and wiped out unless
+         * the calling app holds the CHANGE_NETWORK_STATE permission.
+         *
+         * @param uids The watched UIDs as a set of UidRanges, or null for everything.
+         * @return The builder to facilitate chaining.
+         * @hide
+         */
+        public Builder setUids(Set<UidRange> uids) {
+            mNetworkCapabilities.setUids(uids);
+            return this;
+        }
+
+        /**
          * Add a capability that must not exist in the requested network.
          * <p>
          * If the capability was previously added to the list of required capabilities (for
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/SecurityControllerImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/policy/SecurityControllerImpl.java
index b22ce1822acd..0adb4392a174 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/SecurityControllerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/SecurityControllerImpl.java
@@ -64,6 +64,7 @@ public class SecurityControllerImpl extends CurrentUserTracker implements Securi
             .removeCapability(NetworkCapabilities.NET_CAPABILITY_NOT_VPN)
             .removeCapability(NetworkCapabilities.NET_CAPABILITY_NOT_RESTRICTED)
             .removeCapability(NetworkCapabilities.NET_CAPABILITY_TRUSTED)
+            .setUids(null)
             .build();
     private static final int NO_NETWORK = -1;
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/SecurityControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/SecurityControllerTest.java
index 7ca9d73a230e..f76de5a03d2d 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/SecurityControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/SecurityControllerTest.java
@@ -19,11 +19,15 @@ package com.android.systemui.statusbar.policy;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
+import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.anyInt;
+import static org.mockito.Matchers.argThat;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.doNothing;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
 
 import android.app.admin.DevicePolicyManager;
 import android.content.ComponentName;
@@ -32,6 +36,9 @@ import android.content.Intent;
 import android.content.pm.StringParceledListSlice;
 import android.content.pm.UserInfo;
 import android.net.ConnectivityManager;
+import android.net.ConnectivityManager.NetworkCallback;
+import android.net.NetworkCapabilities;
+import android.net.NetworkRequest;
 import android.os.UserManager;
 import android.security.IKeyChainService;
 import android.support.test.runner.AndroidJUnit4;
@@ -61,6 +68,7 @@ public class SecurityControllerTest extends SysuiTestCase implements SecurityCon
     private final UserManager mUserManager = mock(UserManager.class);
     private SecurityControllerImpl mSecurityController;
     private CountDownLatch mStateChangedLatch;
+    private ConnectivityManager mConnectivityManager = mock(ConnectivityManager.class);
 
     // implementing SecurityControllerCallback
     @Override
@@ -72,7 +80,7 @@ public class SecurityControllerTest extends SysuiTestCase implements SecurityCon
     public void setUp() throws Exception {
         mContext.addMockSystemService(Context.DEVICE_POLICY_SERVICE, mDevicePolicyManager);
         mContext.addMockSystemService(Context.USER_SERVICE, mUserManager);
-        mContext.addMockSystemService(Context.CONNECTIVITY_SERVICE, mock(ConnectivityManager.class));
+        mContext.addMockSystemService(Context.CONNECTIVITY_SERVICE, mConnectivityManager);
 
         Intent intent = new Intent(IKeyChainService.class.getName());
         ComponentName comp = intent.resolveSystemService(mContext.getPackageManager(), 0);
@@ -176,4 +184,12 @@ public class SecurityControllerTest extends SysuiTestCase implements SecurityCon
         //assertTrue(mStateChangedLatch.await(31, TimeUnit.SECONDS));
         //assertFalse(mSecurityController.hasCACertInCurrentUser());
     }
+
+    @Test
+    public void testNetworkRequest() {
+        verify(mConnectivityManager, times(1)).registerNetworkCallback(argThat(
+                (NetworkRequest request) -> request.networkCapabilities.getUids() == null
+                        && request.networkCapabilities.getCapabilities().length == 0
+                ), any(NetworkCallback.class));
+    }
 }