Merge "Enable strict_updatability_linting in connectivity src"
diff --git a/Tethering/Android.bp b/Tethering/Android.bp
index 20efc3e..0b54783 100644
--- a/Tethering/Android.bp
+++ b/Tethering/Android.bp
@@ -97,7 +97,7 @@
     min_sdk_version: "30",
     header_libs: [
         "bpf_syscall_wrappers",
-        "bpf_tethering_headers",
+        "bpf_connectivity_headers",
     ],
     srcs: [
         "jni/*.cpp",
@@ -185,7 +185,11 @@
 
 android_app {
     name: "TetheringNext",
-    defaults: ["TetheringAppDefaults", "TetheringApiLevel"],
+    defaults: [
+        "TetheringAppDefaults",
+        "TetheringApiLevel",
+        "ConnectivityNextEnableDefaults",
+    ],
     static_libs: ["TetheringApiCurrentLib"],
     certificate: "networkstack",
     manifest: "AndroidManifest.xml",
diff --git a/Tethering/apex/Android.bp b/Tethering/apex/Android.bp
index a5216f7..7863572 100644
--- a/Tethering/apex/Android.bp
+++ b/Tethering/apex/Android.bp
@@ -18,8 +18,26 @@
     default_applicable_licenses: ["Android-Apache-2.0"],
 }
 
+// Defaults to enable/disable java targets which uses development APIs. "enabled" may have a
+// different value depending on the branch.
+java_defaults {
+    name: "ConnectivityNextEnableDefaults",
+    enabled: true,
+}
+apex_defaults {
+    name: "ConnectivityApexDefaults",
+    // Tethering app to include in the AOSP apex. Branches that disable the "next" targets may use
+    // a stable tethering app instead, but will generally override the AOSP apex to use updatable
+    // package names and keys, so that apex will be unused anyway.
+    apps: ["TetheringNext"], // Replace to "Tethering" if ConnectivityNextEnableDefaults is false.
+}
+// This is a placeholder comment to avoid merge conflicts
+// as the above target may have different "enabled" values
+// depending on the branch
+
 apex {
     name: "com.android.tethering",
+    defaults: ["ConnectivityApexDefaults"],
     compile_multilib: "both",
     updatable: true,
     min_sdk_version: "30",
@@ -43,7 +61,6 @@
     ],
     apps: [
         "ServiceConnectivityResources",
-        "TetheringNext",
     ],
     prebuilts: ["current_sdkinfo"],
     manifest: "manifest.json",
diff --git a/Tethering/jni/com_android_networkstack_tethering_BpfUtils.cpp b/Tethering/jni/com_android_networkstack_tethering_BpfUtils.cpp
index 2d679a8..f9e4824 100644
--- a/Tethering/jni/com_android_networkstack_tethering_BpfUtils.cpp
+++ b/Tethering/jni/com_android_networkstack_tethering_BpfUtils.cpp
@@ -26,6 +26,7 @@
 #include <net/if.h>
 #include <stdio.h>
 #include <sys/socket.h>
+#include <sys/utsname.h>
 
 // TODO: use unique_fd.
 #define BPF_FD_JUST_USE_INT
@@ -158,6 +159,37 @@
     return rv;
 }
 
+// -----------------------------------------------------------------------------
+// TODO - just use BpfUtils.h once that is available in sc-mainline-prod and has kernelVersion()
+//
+// In the mean time copying verbatim from:
+//   system/bpf/libbpf_android/include/bpf/BpfUtils.h
+// and
+//   system/bpf/libbpf_android/BpfUtils.cpp
+
+#define KVER(a, b, c) (((a) << 24) + ((b) << 16) + (c))
+
+static unsigned kernelVersion() {
+    struct utsname buf;
+    int ret = uname(&buf);
+    if (ret) return 0;
+
+    unsigned kver_major;
+    unsigned kver_minor;
+    unsigned kver_sub;
+    char discard;
+    ret = sscanf(buf.release, "%u.%u.%u%c", &kver_major, &kver_minor, &kver_sub, &discard);
+    // Check the device kernel version
+    if (ret < 3) return 0;
+
+    return KVER(kver_major, kver_minor, kver_sub);
+}
+
+static inline bool isAtLeastKernelVersion(unsigned major, unsigned minor, unsigned sub) {
+    return kernelVersion() >= KVER(major, minor, sub);
+}
+// -----------------------------------------------------------------------------
+
 static jboolean com_android_networkstack_tethering_BpfUtils_isEthernet(JNIEnv* env, jobject clazz,
                                                                        jstring iface) {
     ScopedUtfChars interface(env, iface);
@@ -170,12 +202,30 @@
         return false;
     }
 
+    // Backwards compatibility with pre-GKI kernels that use various custom
+    // ARPHRD_* for their cellular interface
+    switch (rv) {
+        // ARPHRD_PUREIP on at least some Mediatek Android kernels
+        // example: wembley with 4.19 kernel
+        case 520:
+        // in Linux 4.14+ rmnet support was upstreamed and ARHRD_RAWIP became 519,
+        // but it is 530 on at least some Qualcomm Android 4.9 kernels with rmnet
+        // example: Pixel 3 family
+        case 530:
+            // >5.4 kernels are GKI2.0 and thus upstream compatible, however 5.10
+            // shipped with Android S, so (for safety) let's limit ourselves to
+            // >5.10, ie. 5.11+ as a guarantee we're on Android T+ and thus no
+            // longer need this non-upstream compatibility logic
+            static bool is_pre_5_11_kernel = !isAtLeastKernelVersion(5, 11, 0);
+            if (is_pre_5_11_kernel) return false;
+    }
+
     switch (rv) {
         case ARPHRD_ETHER:
             return true;
         case ARPHRD_NONE:
-        case ARPHRD_RAWIP:  // in Linux 4.14+ rmnet support was upstreamed and this is 519
-        case 530:           // this is ARPHRD_RAWIP on some Android 4.9 kernels with rmnet
+        case ARPHRD_PPP:
+        case ARPHRD_RAWIP:
             return false;
         default:
             jniThrowExceptionFmt(env, "java/io/IOException",
diff --git a/Tethering/tests/privileged/Android.bp b/Tethering/tests/privileged/Android.bp
index 214b014..c890197 100644
--- a/Tethering/tests/privileged/Android.bp
+++ b/Tethering/tests/privileged/Android.bp
@@ -34,6 +34,7 @@
     name: "TetheringPrivilegedTests",
     defaults: [
         "TetheringPrivilegedTestsJniDefaults",
+        "ConnectivityNextEnableDefaults",
     ],
     srcs: [
         "src/**/*.java",
diff --git a/Tethering/tests/privileged/src/com/android/networkstack/tethering/BpfMapTest.java b/Tethering/tests/privileged/src/com/android/networkstack/tethering/BpfMapTest.java
index 646c75f..ad2faa0 100644
--- a/Tethering/tests/privileged/src/com/android/networkstack/tethering/BpfMapTest.java
+++ b/Tethering/tests/privileged/src/com/android/networkstack/tethering/BpfMapTest.java
@@ -50,7 +50,7 @@
 @RunWith(DevSdkIgnoreRunner.class)
 @IgnoreUpTo(Build.VERSION_CODES.R)
 public final class BpfMapTest {
-    // Sync from packages/modules/Connectivity/Tethering/bpf_progs/offload.c.
+    // Sync from packages/modules/Connectivity/bpf_progs/offload.c.
     private static final int TEST_MAP_SIZE = 16;
     private static final String TETHER_DOWNSTREAM6_FS_PATH =
             "/sys/fs/bpf/tethering/map_test_tether_downstream6_map";
diff --git a/Tethering/tests/unit/Android.bp b/Tethering/tests/unit/Android.bp
index 228f3fd..5150d39 100644
--- a/Tethering/tests/unit/Android.bp
+++ b/Tethering/tests/unit/Android.bp
@@ -49,7 +49,6 @@
         "src/**/*.kt",
     ],
     static_libs: [
-        "TetheringApiCurrentLib",
         "TetheringCommonTests",
         "androidx.test.rules",
         "frameworks-base-testutils",
@@ -85,6 +84,9 @@
 android_library {
     name: "TetheringTestsLatestSdkLib",
     defaults: ["TetheringTestsDefaults"],
+    static_libs: [
+        "TetheringApiStableLib",
+    ],
     target_sdk_version: "30",
     visibility: [
         "//packages/modules/Connectivity/tests:__subpackages__",
@@ -99,7 +101,13 @@
         "device-tests",
         "mts-tethering",
     ],
-    defaults: ["TetheringTestsDefaults"],
+    defaults: [
+        "TetheringTestsDefaults",
+        "ConnectivityNextEnableDefaults",
+    ],
+    static_libs: [
+        "TetheringApiCurrentLib",
+    ],
     compile_multilib: "both",
     jarjar_rules: ":TetheringTestsJarJarRules",
 }
diff --git a/Tethering/bpf_progs/Android.bp b/bpf_progs/Android.bp
similarity index 69%
rename from Tethering/bpf_progs/Android.bp
rename to bpf_progs/Android.bp
index 5b00dfe..17eebe0 100644
--- a/Tethering/bpf_progs/Android.bp
+++ b/bpf_progs/Android.bp
@@ -22,7 +22,7 @@
 }
 
 cc_library_headers {
-    name: "bpf_tethering_headers",
+    name: "bpf_connectivity_headers",
     vendor_available: false,
     host_supported: false,
     export_include_dirs: ["."],
@@ -32,9 +32,20 @@
     ],
     sdk_version: "30",
     min_sdk_version: "30",
-    apex_available: ["com.android.tethering"],
+    apex_available: [
+        "//apex_available:platform",
+        "com.android.tethering",
+        ],
     visibility: [
+        // TODO: remove it when NetworkStatsService is moved into the mainline module and no more
+        // calls to JNI in libservices.core.
+        "//frameworks/base/services/core/jni",
         "//packages/modules/Connectivity/Tethering",
+        "//packages/modules/Connectivity/tests/unit/jni",
+        // TODO: remove system/netd/* when all BPF code is moved out of Netd.
+        "//system/netd/libnetdbpf",
+        "//system/netd/server",
+        "//system/netd/tests",
     ],
 }
 
diff --git a/Tethering/bpf_progs/bpf_net_helpers.h b/bpf_progs/bpf_net_helpers.h
similarity index 100%
rename from Tethering/bpf_progs/bpf_net_helpers.h
rename to bpf_progs/bpf_net_helpers.h
diff --git a/bpf_progs/bpf_shared.h b/bpf_progs/bpf_shared.h
new file mode 100644
index 0000000..8577d9d
--- /dev/null
+++ b/bpf_progs/bpf_shared.h
@@ -0,0 +1,209 @@
+/*
+ * Copyright (C) 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <linux/if.h>
+#include <linux/if_ether.h>
+#include <linux/in.h>
+#include <linux/in6.h>
+#include <netdutils/UidConstants.h>
+
+// This header file is shared by eBPF kernel programs (C) and netd (C++) and
+// some of the maps are also accessed directly from Java mainline module code.
+//
+// Hence: explicitly pad all relevant structures and assert that their size
+// is the sum of the sizes of their fields.
+#define STRUCT_SIZE(name, size) _Static_assert(sizeof(name) == (size), "Incorrect struct size.")
+
+typedef struct {
+    uint32_t uid;
+    uint32_t tag;
+} UidTagValue;
+STRUCT_SIZE(UidTagValue, 2 * 4);  // 8
+
+typedef struct {
+    uint32_t uid;
+    uint32_t tag;
+    uint32_t counterSet;
+    uint32_t ifaceIndex;
+} StatsKey;
+STRUCT_SIZE(StatsKey, 4 * 4);  // 16
+
+typedef struct {
+    uint64_t rxPackets;
+    uint64_t rxBytes;
+    uint64_t txPackets;
+    uint64_t txBytes;
+} StatsValue;
+STRUCT_SIZE(StatsValue, 4 * 8);  // 32
+
+typedef struct {
+    char name[IFNAMSIZ];
+} IfaceValue;
+STRUCT_SIZE(IfaceValue, 16);
+
+typedef struct {
+    uint64_t rxBytes;
+    uint64_t rxPackets;
+    uint64_t txBytes;
+    uint64_t txPackets;
+    uint64_t tcpRxPackets;
+    uint64_t tcpTxPackets;
+} Stats;
+
+// Since we cannot garbage collect the stats map since device boot, we need to make these maps as
+// large as possible. The maximum size of number of map entries we can have is depend on the rlimit
+// of MEM_LOCK granted to netd. The memory space needed by each map can be calculated by the
+// following fomula:
+//      elem_size = 40 + roundup(key_size, 8) + roundup(value_size, 8)
+//      cost = roundup_pow_of_two(max_entries) * 16 + elem_size * max_entries +
+//              elem_size * number_of_CPU
+// And the cost of each map currently used is(assume the device have 8 CPUs):
+// cookie_tag_map:      key:  8 bytes, value:  8 bytes, cost:  822592 bytes    =   823Kbytes
+// uid_counter_set_map: key:  4 bytes, value:  1 bytes, cost:  145216 bytes    =   145Kbytes
+// app_uid_stats_map:   key:  4 bytes, value: 32 bytes, cost: 1062784 bytes    =  1063Kbytes
+// uid_stats_map:       key: 16 bytes, value: 32 bytes, cost: 1142848 bytes    =  1143Kbytes
+// tag_stats_map:       key: 16 bytes, value: 32 bytes, cost: 1142848 bytes    =  1143Kbytes
+// iface_index_name_map:key:  4 bytes, value: 16 bytes, cost:   80896 bytes    =    81Kbytes
+// iface_stats_map:     key:  4 bytes, value: 32 bytes, cost:   97024 bytes    =    97Kbytes
+// dozable_uid_map:     key:  4 bytes, value:  1 bytes, cost:  145216 bytes    =   145Kbytes
+// standby_uid_map:     key:  4 bytes, value:  1 bytes, cost:  145216 bytes    =   145Kbytes
+// powersave_uid_map:   key:  4 bytes, value:  1 bytes, cost:  145216 bytes    =   145Kbytes
+// total:                                                                         4930Kbytes
+// It takes maximum 4.9MB kernel memory space if all maps are full, which requires any devices
+// running this module to have a memlock rlimit to be larger then 5MB. In the old qtaguid module,
+// we don't have a total limit for data entries but only have limitation of tags each uid can have.
+// (default is 1024 in kernel);
+
+// 'static' - otherwise these constants end up in .rodata in the resulting .o post compilation
+static const int COOKIE_UID_MAP_SIZE = 10000;
+static const int UID_COUNTERSET_MAP_SIZE = 2000;
+static const int APP_STATS_MAP_SIZE = 10000;
+static const int STATS_MAP_SIZE = 5000;
+static const int IFACE_INDEX_NAME_MAP_SIZE = 1000;
+static const int IFACE_STATS_MAP_SIZE = 1000;
+static const int CONFIGURATION_MAP_SIZE = 2;
+static const int UID_OWNER_MAP_SIZE = 2000;
+
+#define BPF_PATH "/sys/fs/bpf/"
+
+#define BPF_EGRESS_PROG_PATH BPF_PATH "prog_netd_cgroupskb_egress_stats"
+#define BPF_INGRESS_PROG_PATH BPF_PATH "prog_netd_cgroupskb_ingress_stats"
+#define XT_BPF_INGRESS_PROG_PATH BPF_PATH "prog_netd_skfilter_ingress_xtbpf"
+#define XT_BPF_EGRESS_PROG_PATH BPF_PATH "prog_netd_skfilter_egress_xtbpf"
+#define XT_BPF_ALLOWLIST_PROG_PATH BPF_PATH "prog_netd_skfilter_allowlist_xtbpf"
+#define XT_BPF_DENYLIST_PROG_PATH BPF_PATH "prog_netd_skfilter_denylist_xtbpf"
+#define CGROUP_SOCKET_PROG_PATH BPF_PATH "prog_netd_cgroupsock_inet_create"
+
+#define TC_BPF_INGRESS_ACCOUNT_PROG_NAME "prog_netd_schedact_ingress_account"
+#define TC_BPF_INGRESS_ACCOUNT_PROG_PATH BPF_PATH TC_BPF_INGRESS_ACCOUNT_PROG_NAME
+
+#define COOKIE_TAG_MAP_PATH BPF_PATH "map_netd_cookie_tag_map"
+#define UID_COUNTERSET_MAP_PATH BPF_PATH "map_netd_uid_counterset_map"
+#define APP_UID_STATS_MAP_PATH BPF_PATH "map_netd_app_uid_stats_map"
+#define STATS_MAP_A_PATH BPF_PATH "map_netd_stats_map_A"
+#define STATS_MAP_B_PATH BPF_PATH "map_netd_stats_map_B"
+#define IFACE_INDEX_NAME_MAP_PATH BPF_PATH "map_netd_iface_index_name_map"
+#define IFACE_STATS_MAP_PATH BPF_PATH "map_netd_iface_stats_map"
+#define CONFIGURATION_MAP_PATH BPF_PATH "map_netd_configuration_map"
+#define UID_OWNER_MAP_PATH BPF_PATH "map_netd_uid_owner_map"
+#define UID_PERMISSION_MAP_PATH BPF_PATH "map_netd_uid_permission_map"
+
+enum UidOwnerMatchType {
+    NO_MATCH = 0,
+    HAPPY_BOX_MATCH = (1 << 0),
+    PENALTY_BOX_MATCH = (1 << 1),
+    DOZABLE_MATCH = (1 << 2),
+    STANDBY_MATCH = (1 << 3),
+    POWERSAVE_MATCH = (1 << 4),
+    RESTRICTED_MATCH = (1 << 5),
+    IIF_MATCH = (1 << 6),
+};
+
+enum BpfPermissionMatch {
+    BPF_PERMISSION_INTERNET = 1 << 2,
+    BPF_PERMISSION_UPDATE_DEVICE_STATS = 1 << 3,
+};
+// In production we use two identical stats maps to record per uid stats and
+// do swap and clean based on the configuration specified here. The statsMapType
+// value in configuration map specified which map is currently in use.
+enum StatsMapType {
+    SELECT_MAP_A,
+    SELECT_MAP_B,
+};
+
+// TODO: change the configuration object from an 8-bit bitmask to an object with clearer
+// semantics, like a struct.
+typedef uint8_t BpfConfig;
+static const BpfConfig DEFAULT_CONFIG = 0;
+
+typedef struct {
+    // Allowed interface index. Only applicable if IIF_MATCH is set in the rule bitmask above.
+    uint32_t iif;
+    // A bitmask of enum values in UidOwnerMatchType.
+    uint32_t rule;
+} UidOwnerValue;
+STRUCT_SIZE(UidOwnerValue, 2 * 4);  // 8
+
+#define UID_RULES_CONFIGURATION_KEY 1
+#define CURRENT_STATS_MAP_CONFIGURATION_KEY 2
+
+#define CLAT_INGRESS6_PROG_RAWIP_NAME "prog_clatd_schedcls_ingress6_clat_rawip"
+#define CLAT_INGRESS6_PROG_ETHER_NAME "prog_clatd_schedcls_ingress6_clat_ether"
+
+#define CLAT_INGRESS6_PROG_RAWIP_PATH BPF_PATH CLAT_INGRESS6_PROG_RAWIP_NAME
+#define CLAT_INGRESS6_PROG_ETHER_PATH BPF_PATH CLAT_INGRESS6_PROG_ETHER_NAME
+
+#define CLAT_INGRESS6_MAP_PATH BPF_PATH "map_clatd_clat_ingress6_map"
+
+typedef struct {
+    uint32_t iif;            // The input interface index
+    struct in6_addr pfx96;   // The source /96 nat64 prefix, bottom 32 bits must be 0
+    struct in6_addr local6;  // The full 128-bits of the destination IPv6 address
+} ClatIngress6Key;
+STRUCT_SIZE(ClatIngress6Key, 4 + 2 * 16);  // 36
+
+typedef struct {
+    uint32_t oif;           // The output interface to redirect to (0 means don't redirect)
+    struct in_addr local4;  // The destination IPv4 address
+} ClatIngress6Value;
+STRUCT_SIZE(ClatIngress6Value, 4 + 4);  // 8
+
+#define CLAT_EGRESS4_PROG_RAWIP_NAME "prog_clatd_schedcls_egress4_clat_rawip"
+#define CLAT_EGRESS4_PROG_ETHER_NAME "prog_clatd_schedcls_egress4_clat_ether"
+
+#define CLAT_EGRESS4_PROG_RAWIP_PATH BPF_PATH CLAT_EGRESS4_PROG_RAWIP_NAME
+#define CLAT_EGRESS4_PROG_ETHER_PATH BPF_PATH CLAT_EGRESS4_PROG_ETHER_NAME
+
+#define CLAT_EGRESS4_MAP_PATH BPF_PATH "map_clatd_clat_egress4_map"
+
+typedef struct {
+    uint32_t iif;           // The input interface index
+    struct in_addr local4;  // The source IPv4 address
+} ClatEgress4Key;
+STRUCT_SIZE(ClatEgress4Key, 4 + 4);  // 8
+
+typedef struct {
+    uint32_t oif;            // The output interface to redirect to
+    struct in6_addr local6;  // The full 128-bits of the source IPv6 address
+    struct in6_addr pfx96;   // The destination /96 nat64 prefix, bottom 32 bits must be 0
+    bool oifIsEthernet;      // Whether the output interface requires ethernet header
+    uint8_t pad[3];
+} ClatEgress4Value;
+STRUCT_SIZE(ClatEgress4Value, 4 + 2 * 16 + 1 + 3);  // 40
+
+#undef STRUCT_SIZE
diff --git a/Tethering/bpf_progs/bpf_tethering.h b/bpf_progs/bpf_tethering.h
similarity index 98%
rename from Tethering/bpf_progs/bpf_tethering.h
rename to bpf_progs/bpf_tethering.h
index 5fdf8cd..b0ec8f6 100644
--- a/Tethering/bpf_progs/bpf_tethering.h
+++ b/bpf_progs/bpf_tethering.h
@@ -24,7 +24,7 @@
 // Common definitions for BPF code in the tethering mainline module.
 // These definitions are available to:
 // - The BPF programs in Tethering/bpf_progs/
-// - JNI code that depends on the bpf_tethering_headers library.
+// - JNI code that depends on the bpf_connectivity_headers library.
 
 #define BPF_TETHER_ERRORS    \
     ERR(INVALID_IP_VERSION)  \
diff --git a/Tethering/bpf_progs/offload.c b/bpf_progs/offload.c
similarity index 100%
rename from Tethering/bpf_progs/offload.c
rename to bpf_progs/offload.c
diff --git a/Tethering/bpf_progs/test.c b/bpf_progs/test.c
similarity index 100%
rename from Tethering/bpf_progs/test.c
rename to bpf_progs/test.c
diff --git a/framework/src/android/net/Network.java b/framework/src/android/net/Network.java
index b3770ea..53f171a 100644
--- a/framework/src/android/net/Network.java
+++ b/framework/src/android/net/Network.java
@@ -382,13 +382,14 @@
         // Query a property of the underlying socket to ensure that the socket's file descriptor
         // exists, is available to bind to a network and is not closed.
         socket.getReuseAddress();
-        final ParcelFileDescriptor pfd = ParcelFileDescriptor.fromDatagramSocket(socket);
-        bindSocket(pfd.getFileDescriptor());
-        // ParcelFileDescriptor.fromSocket() creates a dup of the original fd. The original and the
-        // dup share the underlying socket in the kernel. The socket is never truly closed until the
-        // last fd pointing to the socket being closed. So close the dup one after binding the
-        // socket to control the lifetime of the dup fd.
-        pfd.close();
+
+        // ParcelFileDescriptor.fromDatagramSocket() creates a dup of the original fd. The original
+        // and the dup share the underlying socket in the kernel. The socket is never truly closed
+        // until the last fd pointing to the socket being closed. Try and eventually close the dup
+        // one after binding the socket to control the lifetime of the dup fd.
+        try (ParcelFileDescriptor pfd = ParcelFileDescriptor.fromDatagramSocket(socket)) {
+            bindSocket(pfd.getFileDescriptor());
+        }
     }
 
     /**
@@ -400,13 +401,13 @@
         // Query a property of the underlying socket to ensure that the socket's file descriptor
         // exists, is available to bind to a network and is not closed.
         socket.getReuseAddress();
-        final ParcelFileDescriptor pfd = ParcelFileDescriptor.fromSocket(socket);
-        bindSocket(pfd.getFileDescriptor());
-        // ParcelFileDescriptor.fromSocket() creates a dup of the original fd. The original and the
-        // dup share the underlying socket in the kernel. The socket is never truly closed until the
-        // last fd pointing to the socket being closed. So close the dup one after binding the
-        // socket to control the lifetime of the dup fd.
-        pfd.close();
+        // ParcelFileDescriptor.fromSocket() creates a dup of the original fd. The original and
+        // the dup share the underlying socket in the kernel. The socket is never truly closed
+        // until the last fd pointing to the socket being closed. Try and eventually close the dup
+        // one after binding the socket to control the lifetime of the dup fd.
+        try (ParcelFileDescriptor pfd = ParcelFileDescriptor.fromSocket(socket)) {
+            bindSocket(pfd.getFileDescriptor());
+        }
     }
 
     /**
diff --git a/service/Android.bp b/service/Android.bp
index d1abdec..b595ef2 100644
--- a/service/Android.bp
+++ b/service/Android.bp
@@ -68,7 +68,6 @@
         "dnsresolver_aidl_interface-V9-java",
         "modules-utils-build",
         "modules-utils-shell-command-handler",
-        "modules-utils-statemachine",
         "net-utils-device-common",
         "net-utils-device-common-netlink",
         "net-utils-framework-common",
diff --git a/service/jarjar-rules.txt b/service/jarjar-rules.txt
index 4ba6837..e5d1a88 100644
--- a/service/jarjar-rules.txt
+++ b/service/jarjar-rules.txt
@@ -8,10 +8,7 @@
 # the one in com.android.internal.util
 rule android.util.IndentingPrintWriter* com.android.connectivity.@0
 rule com.android.internal.util.IndentingPrintWriter* com.android.connectivity.@0
-rule com.android.internal.util.IState* com.android.connectivity.@0
 rule com.android.internal.util.MessageUtils* com.android.connectivity.@0
-rule com.android.internal.util.State* com.android.connectivity.@0
-rule com.android.internal.util.StateMachine* com.android.connectivity.@0
 rule com.android.internal.util.WakeupMessage* com.android.connectivity.@0
 
 rule com.android.internal.messages.** com.android.connectivity.@0
diff --git a/service/src/com/android/server/ConnectivityServiceInitializer.java b/service/src/com/android/server/ConnectivityServiceInitializer.java
index 2465479..b1a56ae 100644
--- a/service/src/com/android/server/ConnectivityServiceInitializer.java
+++ b/service/src/com/android/server/ConnectivityServiceInitializer.java
@@ -31,7 +31,6 @@
         super(context);
         // Load JNI libraries used by ConnectivityService and its dependencies
         System.loadLibrary("service-connectivity");
-        // TODO: Define formal APIs to get the needed services.
         mConnectivity = new ConnectivityService(context);
     }
 
diff --git a/service/src/com/android/server/connectivity/NetworkRanker.java b/service/src/com/android/server/connectivity/NetworkRanker.java
index d7eb9c8..43da1d0 100644
--- a/service/src/com/android/server/connectivity/NetworkRanker.java
+++ b/service/src/com/android/server/connectivity/NetworkRanker.java
@@ -63,8 +63,6 @@
         NetworkCapabilities getCapsNoCopy();
     }
 
-    private static final boolean USE_POLICY_RANKING = true;
-
     public NetworkRanker() { }
 
     /**
@@ -77,11 +75,7 @@
         final ArrayList<NetworkAgentInfo> candidates = filter(nais, nai -> nai.satisfies(request));
         if (candidates.size() == 1) return candidates.get(0); // Only one potential satisfier
         if (candidates.size() <= 0) return null; // No network can satisfy this request
-        if (USE_POLICY_RANKING) {
-            return getBestNetworkByPolicy(candidates, currentSatisfier);
-        } else {
-            return getBestNetworkByLegacyInt(candidates);
-        }
+        return getBestNetworkByPolicy(candidates, currentSatisfier);
     }
 
     // Transport preference order, if it comes down to that.
@@ -278,23 +272,6 @@
         return candidates.get(0);
     }
 
-    // TODO : switch to the policy implementation and remove
-    // Almost equivalent to Collections.max(nais), but allows returning null if no network
-    // satisfies the request.
-    private NetworkAgentInfo getBestNetworkByLegacyInt(
-            @NonNull final Collection<NetworkAgentInfo> nais) {
-        NetworkAgentInfo bestNetwork = null;
-        int bestScore = Integer.MIN_VALUE;
-        for (final NetworkAgentInfo nai : nais) {
-            final int naiScore = nai.getCurrentScore();
-            if (naiScore > bestScore) {
-                bestNetwork = nai;
-                bestScore = naiScore;
-            }
-        }
-        return bestNetwork;
-    }
-
     /**
      * Returns whether a {@link Scoreable} has a chance to beat a champion network for a request.
      *
@@ -322,30 +299,11 @@
         // If there is no satisfying network, then this network can beat, because some network
         // is always better than no network.
         if (null == champion) return true;
-        if (USE_POLICY_RANKING) {
-            // If there is no champion, the offer can always beat.
-            // Otherwise rank them.
-            final ArrayList<Scoreable> candidates = new ArrayList<>();
-            candidates.add(champion);
-            candidates.add(contestant);
-            return contestant == getBestNetworkByPolicy(candidates, champion);
-        } else {
-            return mightBeatByLegacyInt(champion.getScore(), contestant);
-        }
-    }
-
-    /**
-     * Returns whether a contestant might beat a champion according to the legacy int.
-     */
-    private boolean mightBeatByLegacyInt(@Nullable final FullScore championScore,
-            @NonNull final Scoreable contestant) {
-        final int offerIntScore;
-        if (contestant.getCapsNoCopy().hasCapability(NetworkCapabilities.NET_CAPABILITY_INTERNET)) {
-            // If the offer might have Internet access, then it might validate.
-            offerIntScore = contestant.getScore().getLegacyIntAsValidated();
-        } else {
-            offerIntScore = contestant.getScore().getLegacyInt();
-        }
-        return championScore.getLegacyInt() < offerIntScore;
+        // If there is no champion, the offer can always beat.
+        // Otherwise rank them.
+        final ArrayList<Scoreable> candidates = new ArrayList<>();
+        candidates.add(champion);
+        candidates.add(contestant);
+        return contestant == getBestNetworkByPolicy(candidates, champion);
     }
 }
diff --git a/service/src/com/android/server/connectivity/PermissionMonitor.java b/service/src/com/android/server/connectivity/PermissionMonitor.java
index da2715e..439db89 100755
--- a/service/src/com/android/server/connectivity/PermissionMonitor.java
+++ b/service/src/com/android/server/connectivity/PermissionMonitor.java
@@ -23,7 +23,6 @@
 import static android.Manifest.permission.UPDATE_DEVICE_STATS;
 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;
 import static android.net.ConnectivitySettingsManager.UIDS_ALLOWED_ON_RESTRICTED_NETWORKS;
 import static android.net.INetd.PERMISSION_INTERNET;
 import static android.net.INetd.PERMISSION_NETWORK;
@@ -60,9 +59,9 @@
 import android.os.UserManager;
 import android.provider.Settings;
 import android.system.OsConstants;
+import android.util.ArrayMap;
 import android.util.ArraySet;
 import android.util.Log;
-import android.util.SparseArray;
 import android.util.SparseIntArray;
 
 import com.android.internal.annotations.GuardedBy;
@@ -98,9 +97,9 @@
     @GuardedBy("this")
     private final Set<UserHandle> mUsers = new HashSet<>();
 
-    // Keys are appIds. Values are true for SYSTEM permission and false for NETWORK permission.
+    // Keys are uids. Values are netd network permissions.
     @GuardedBy("this")
-    private final SparseIntArray mApps = new SparseIntArray();
+    private final SparseIntArray mUidToNetworkPerm = new SparseIntArray();
 
     // Keys are active non-bypassable and fully-routed VPN's interface name, Values are uid ranges
     // for apps under the VPN
@@ -121,6 +120,11 @@
     @GuardedBy("this")
     private final Set<Integer> mUidsAllowedOnRestrictedNetworks = new ArraySet<>();
 
+    @GuardedBy("this")
+    private final Map<UserHandle, PackageManager> mUsersPackageManager = new ArrayMap<>();
+
+    private static final int SYSTEM_APPID = SYSTEM_UID;
+
     private static final int MAX_PERMISSION_UPDATE_LOGS = 40;
     private final SharedLog mPermissionUpdateLogs = new SharedLog(MAX_PERMISSION_UPDATE_LOGS, TAG);
 
@@ -212,6 +216,83 @@
         return targetPermission > currentPermission;
     }
 
+    private List<PackageInfo> getInstalledPackagesAsUser(final UserHandle user) {
+        return mPackageManager.getInstalledPackagesAsUser(GET_PERMISSIONS, user.getIdentifier());
+    }
+
+    private synchronized void updateAllApps(final List<PackageInfo> apps) {
+        for (PackageInfo app : apps) {
+            final int appId = app.applicationInfo != null
+                    ? UserHandle.getAppId(app.applicationInfo.uid) : INVALID_UID;
+            if (appId < 0) {
+                continue;
+            }
+            mAllApps.add(appId);
+        }
+    }
+
+    // Return the network permission for the passed list of apps. Note that this depends on the
+    // current settings of the device (See isUidAllowedOnRestrictedNetworks).
+    private SparseIntArray makeUidsNetworkPerm(final List<PackageInfo> apps) {
+        final SparseIntArray uidsPerm = new SparseIntArray();
+        for (PackageInfo app : apps) {
+            final int uid = app.applicationInfo != null ? app.applicationInfo.uid : INVALID_UID;
+            if (uid < 0) {
+                continue;
+            }
+            final int permission = getPackageNetdNetworkPermission(app);
+            if (isHigherNetworkPermission(permission, uidsPerm.get(uid, PERMISSION_NONE))) {
+                uidsPerm.put(uid, permission);
+            }
+        }
+        return uidsPerm;
+    }
+
+    private static SparseIntArray makeAppIdsTrafficPerm(final List<PackageInfo> apps) {
+        final SparseIntArray appIdsPerm = new SparseIntArray();
+        for (PackageInfo app : apps) {
+            final int appId = app.applicationInfo != null
+                    ? UserHandle.getAppId(app.applicationInfo.uid) : INVALID_UID;
+            if (appId < 0) {
+                continue;
+            }
+            final int otherNetdPerms = getNetdPermissionMask(app.requestedPermissions,
+                    app.requestedPermissionsFlags);
+            appIdsPerm.put(appId, appIdsPerm.get(appId) | otherNetdPerms);
+        }
+        return appIdsPerm;
+    }
+
+    private synchronized void updateUidsNetworkPermission(final SparseIntArray uids) {
+        for (int i = 0; i < uids.size(); i++) {
+            mUidToNetworkPerm.put(uids.keyAt(i), uids.valueAt(i));
+        }
+        sendUidsNetworkPermission(uids, true /* add */);
+    }
+
+    private void updateAppIdsTrafficPermission(final SparseIntArray appIds,
+            final SparseIntArray extraAppIds) {
+        for (int i = 0; i < extraAppIds.size(); i++) {
+            final int appId = extraAppIds.keyAt(i);
+            final int permission = extraAppIds.valueAt(i);
+            appIds.put(appId, appIds.get(appId) | permission);
+        }
+        sendAppIdsTrafficPermission(appIds);
+    }
+
+    private SparseIntArray getSystemTrafficPerm() {
+        final SparseIntArray appIdsPerm = new SparseIntArray();
+        for (final int uid : mSystemConfigManager.getSystemPermissionUids(INTERNET)) {
+            final int appId = UserHandle.getAppId(uid);
+            appIdsPerm.put(appId, appIdsPerm.get(appId) | PERMISSION_INTERNET);
+        }
+        for (final int uid : mSystemConfigManager.getSystemPermissionUids(UPDATE_DEVICE_STATS)) {
+            final int appId = UserHandle.getAppId(uid);
+            appIdsPerm.put(appId, appIdsPerm.get(appId) | PERMISSION_UPDATE_DEVICE_STATS);
+        }
+        return appIdsPerm;
+    }
+
     // Intended to be called only once at startup, after the system is ready. Installs a broadcast
     // receiver to monitor ongoing UID changes, so this shouldn't/needn't be called again.
     public synchronized void startMonitoring() {
@@ -252,65 +333,18 @@
         // mUidsAllowedOnRestrictedNetworks.
         updateUidsAllowedOnRestrictedNetworks(mDeps.getUidsAllowedOnRestrictedNetworks(mContext));
 
-        List<PackageInfo> apps = mPackageManager.getInstalledPackages(GET_PERMISSIONS
-                | MATCH_ANY_USER);
-        if (apps == null) {
-            loge("No apps");
-            return;
+        final List<UserHandle> usrs = mUserManager.getUserHandles(true /* excludeDying */);
+        // Update netd permissions for all users.
+        for (UserHandle user : usrs) {
+            onUserAdded(user);
         }
-
-        final SparseIntArray netdPermsAppIds = new SparseIntArray();
-
-        for (PackageInfo app : apps) {
-            int uid = app.applicationInfo != null ? app.applicationInfo.uid : INVALID_UID;
-            if (uid < 0) {
-                continue;
-            }
-            final int appId = UserHandle.getAppId(uid);
-            mAllApps.add(appId);
-
-            final int permission = getPackageNetdNetworkPermission(app);
-            if (isHigherNetworkPermission(permission, mApps.get(appId, PERMISSION_NONE))) {
-                mApps.put(appId, permission);
-            }
-
-            //TODO: unify the management of the permissions into one codepath.
-            int otherNetdPerms = getNetdPermissionMask(app.requestedPermissions,
-                    app.requestedPermissionsFlags);
-            netdPermsAppIds.put(appId, netdPermsAppIds.get(appId) | otherNetdPerms);
-        }
-
-        mUsers.addAll(mUserManager.getUserHandles(true /* excludeDying */));
-
-        final SparseArray<String> netdPermToSystemPerm = new SparseArray<>();
-        netdPermToSystemPerm.put(PERMISSION_INTERNET, INTERNET);
-        netdPermToSystemPerm.put(PERMISSION_UPDATE_DEVICE_STATS, UPDATE_DEVICE_STATS);
-        for (int i = 0; i < netdPermToSystemPerm.size(); i++) {
-            final int netdPermission = netdPermToSystemPerm.keyAt(i);
-            final String systemPermission = netdPermToSystemPerm.valueAt(i);
-            final int[] hasPermissionUids =
-                    mSystemConfigManager.getSystemPermissionUids(systemPermission);
-            for (int j = 0; j < hasPermissionUids.length; j++) {
-                final int appId = UserHandle.getAppId(hasPermissionUids[j]);
-                netdPermsAppIds.put(appId, netdPermsAppIds.get(appId) | netdPermission);
-            }
-        }
-        log("Users: " + mUsers.size() + ", Apps: " + mApps.size());
-        update(mUsers, mApps, true);
-        sendPackagePermissionsToNetd(netdPermsAppIds);
+        log("Users: " + mUsers.size() + ", UidToNetworkPerm: " + mUidToNetworkPerm.size());
     }
 
     @VisibleForTesting
     synchronized void updateUidsAllowedOnRestrictedNetworks(final Set<Integer> uids) {
         mUidsAllowedOnRestrictedNetworks.clear();
-        // This is necessary for the app id to match in isUidAllowedOnRestrictedNetworks, and will
-        // grant the permission to all uids associated with the app ID. This is safe even if the app
-        // is only installed on some users because the uid cannot match some other app – this uid is
-        // in effect not installed and can't be run.
-        // TODO (b/192431153): Change appIds back to uids.
-        for (int uid : uids) {
-            mUidsAllowedOnRestrictedNetworks.add(UserHandle.getAppId(uid));
-        }
+        mUidsAllowedOnRestrictedNetworks.addAll(uids);
     }
 
     @VisibleForTesting
@@ -324,7 +358,8 @@
         return (appInfo.targetSdkVersion < VERSION_Q && isVendorApp(appInfo))
                 // Backward compatibility for b/114245686, on devices that launched before Q daemons
                 // and apps running as the system UID are exempted from this check.
-                || (appInfo.uid == SYSTEM_UID && mDeps.getDeviceFirstSdkInt() < VERSION_Q);
+                || (UserHandle.getAppId(appInfo.uid) == SYSTEM_APPID
+                        && mDeps.getDeviceFirstSdkInt() < VERSION_Q);
     }
 
     @VisibleForTesting
@@ -332,7 +367,7 @@
         if (appInfo == null) return false;
         // Check whether package's uid is in allowed on restricted networks uid list. If so, this
         // uid can have netd system permission.
-        return mUidsAllowedOnRestrictedNetworks.contains(UserHandle.getAppId(appInfo.uid));
+        return mUidsAllowedOnRestrictedNetworks.contains(appInfo.uid);
     }
 
     @VisibleForTesting
@@ -365,33 +400,30 @@
     public synchronized boolean hasUseBackgroundNetworksPermission(final int uid) {
         // Apps with any of the CHANGE_NETWORK_STATE, NETWORK_STACK, CONNECTIVITY_INTERNAL or
         // CONNECTIVITY_USE_RESTRICTED_NETWORKS permission has the permission to use background
-        // networks. mApps contains the result of checks for both hasNetworkPermission and
-        // hasRestrictedNetworkPermission. If uid is in the mApps list that means uid has one of
-        // permissions at least.
-        return mApps.get(UserHandle.getAppId(uid), PERMISSION_NONE) != PERMISSION_NONE;
+        // networks. mUidToNetworkPerm contains the result of checks for hasNetworkPermission and
+        // hasRestrictedNetworkPermission, as well as the list of UIDs allowed on restricted
+        // networks. If uid is in the mUidToNetworkPerm list that means uid has one of permissions
+        // at least.
+        return mUidToNetworkPerm.get(uid, PERMISSION_NONE) != PERMISSION_NONE;
     }
 
     /**
      * Returns whether the given uid has permission to use restricted networks.
      */
     public synchronized boolean hasRestrictedNetworksPermission(int uid) {
-        return PERMISSION_SYSTEM == mApps.get(UserHandle.getAppId(uid), PERMISSION_NONE);
+        return PERMISSION_SYSTEM == mUidToNetworkPerm.get(uid, PERMISSION_NONE);
     }
 
-    private void update(Set<UserHandle> users, SparseIntArray apps, boolean add) {
+    private void sendUidsNetworkPermission(SparseIntArray uids, boolean add) {
         List<Integer> network = new ArrayList<>();
         List<Integer> system = new ArrayList<>();
-        for (int i = 0; i < apps.size(); i++) {
-            final int permission = apps.valueAt(i);
+        for (int i = 0; i < uids.size(); i++) {
+            final int permission = uids.valueAt(i);
             if (PERMISSION_NONE == permission) {
                 continue; // Normally NONE is not stored in this map, but just in case
             }
             List<Integer> list = (PERMISSION_SYSTEM == permission) ? system : network;
-            for (UserHandle user : users) {
-                if (user == null) continue;
-
-                list.add(user.getUid(apps.keyAt(i)));
-            }
+            list.add(uids.keyAt(i));
         }
         try {
             if (add) {
@@ -415,7 +447,19 @@
      */
     public synchronized void onUserAdded(@NonNull UserHandle user) {
         mUsers.add(user);
-        update(Set.of(user), mApps, true);
+
+        final List<PackageInfo> apps = getInstalledPackagesAsUser(user);
+
+        // Save all apps
+        updateAllApps(apps);
+
+        // Uids network permissions
+        final SparseIntArray uids = makeUidsNetworkPerm(apps);
+        updateUidsNetworkPermission(uids);
+
+        // App ids traffic permission
+        final SparseIntArray appIds = makeAppIdsTrafficPerm(apps);
+        updateAppIdsTrafficPermission(appIds, getSystemTrafficPerm());
     }
 
     /**
@@ -427,44 +471,54 @@
      */
     public synchronized void onUserRemoved(@NonNull UserHandle user) {
         mUsers.remove(user);
-        update(Set.of(user), mApps, false);
+
+        final SparseIntArray removedUids = new SparseIntArray();
+        final SparseIntArray allUids = mUidToNetworkPerm.clone();
+        for (int i = 0; i < allUids.size(); i++) {
+            final int uid = allUids.keyAt(i);
+            if (user.equals(UserHandle.getUserHandleForUid(uid))) {
+                mUidToNetworkPerm.delete(uid);
+                removedUids.put(uid, allUids.valueAt(i));
+            }
+        }
+        sendUidsNetworkPermission(removedUids, false /* add */);
     }
 
     /**
      * Compare the current network permission and the given package's permission to find out highest
      * permission for the uid.
      *
+     * @param uid The target uid
      * @param currentPermission Current uid network permission
      * @param name The package has same uid that need compare its permission to update uid network
      *             permission.
      */
     @VisibleForTesting
-    protected int highestPermissionForUid(int currentPermission, String name) {
+    protected int highestPermissionForUid(int uid, int currentPermission, String name) {
+        // If multiple packages share a UID (cf: android:sharedUserId) and ask for different
+        // permissions, don't downgrade (i.e., if it's already SYSTEM, leave it as is).
         if (currentPermission == PERMISSION_SYSTEM) {
             return currentPermission;
         }
-        try {
-            final PackageInfo app = mPackageManager.getPackageInfo(name,
-                    GET_PERMISSIONS | MATCH_ANY_USER);
-            final int permission = getPackageNetdNetworkPermission(app);
-            if (isHigherNetworkPermission(permission, currentPermission)) {
-                return permission;
-            }
-        } catch (NameNotFoundException e) {
-            // App not found.
-            loge("NameNotFoundException " + name);
+        final PackageInfo app = getPackageInfoAsUser(name, UserHandle.getUserHandleForUid(uid));
+        if (app == null) return currentPermission;
+
+        final int permission = getPackageNetdNetworkPermission(app);
+        if (isHigherNetworkPermission(permission, currentPermission)) {
+            return permission;
         }
         return currentPermission;
     }
 
-    private int getPermissionForUid(final int uid) {
+    private int getTrafficPermissionForUid(final int uid) {
         int permission = PERMISSION_NONE;
         // Check all the packages for this UID. The UID has the permission if any of the
         // packages in it has the permission.
         final String[] packages = mPackageManager.getPackagesForUid(uid);
         if (packages != null && packages.length > 0) {
             for (String name : packages) {
-                final PackageInfo app = getPackageInfo(name);
+                final PackageInfo app = getPackageInfoAsUser(name,
+                        UserHandle.getUserHandleForUid(uid));
                 if (app != null && app.requestedPermissions != null) {
                     permission |= getNetdPermissionMask(app.requestedPermissions,
                             app.requestedPermissionsFlags);
@@ -524,24 +578,23 @@
      */
     public synchronized void onPackageAdded(@NonNull final String packageName, final int uid) {
         final int appId = UserHandle.getAppId(uid);
-        final int trafficPerm = getPermissionForUid(uid);
+        final int trafficPerm = getTrafficPermissionForUid(uid);
         sendPackagePermissionsForAppId(appId, trafficPerm);
 
-        // If multiple packages share a UID (cf: android:sharedUserId) and ask for different
-        // permissions, don't downgrade (i.e., if it's already SYSTEM, leave it as is).
-        final int currentPermission = mApps.get(appId, PERMISSION_NONE);
-        final int permission = highestPermissionForUid(currentPermission, packageName);
+        final int currentPermission = mUidToNetworkPerm.get(uid, PERMISSION_NONE);
+        final int permission = highestPermissionForUid(uid, currentPermission, packageName);
         if (permission != currentPermission) {
-            mApps.put(appId, permission);
+            mUidToNetworkPerm.put(uid, permission);
 
             SparseIntArray apps = new SparseIntArray();
-            apps.put(appId, permission);
-            update(mUsers, apps, true);
+            apps.put(uid, permission);
+            sendUidsNetworkPermission(apps, true /* add */);
         }
 
         // If the newly-installed package falls within some VPN's uid range, update Netd with it.
-        // This needs to happen after the mApps update above, since removeBypassingUids() in
-        // updateVpnUid() depends on mApps to check if the package can bypass VPN.
+        // This needs to happen after the mUidToNetworkPerm update above, since
+        // removeBypassingUids() in updateVpnUid() depends on mUidToNetworkPerm to check if the
+        // package can bypass VPN.
         updateVpnUid(uid, true /* add */);
         mAllApps.add(appId);
         mPermissionUpdateLogs.log("Package add: name=" + packageName + ", uid=" + uid
@@ -557,7 +610,7 @@
             for (String name : packages) {
                 // If multiple packages have the same UID, give the UID all permissions that
                 // any package in that UID has.
-                permission = highestPermissionForUid(permission, name);
+                permission = highestPermissionForUid(uid, permission, name);
                 if (permission == PERMISSION_SYSTEM) {
                     break;
                 }
@@ -576,48 +629,42 @@
      */
     public synchronized void onPackageRemoved(@NonNull final String packageName, final int uid) {
         final int appId = UserHandle.getAppId(uid);
-        final int trafficPerm = getPermissionForUid(uid);
+        final int trafficPerm = getTrafficPermissionForUid(uid);
         sendPackagePermissionsForAppId(appId, trafficPerm);
 
         // If the newly-removed package falls within some VPN's uid range, update Netd with it.
-        // This needs to happen before the mApps update below, since removeBypassingUids() in
-        // updateVpnUid() depends on mApps to check if the package can bypass VPN.
+        // This needs to happen before the mUidToNetworkPerm update below, since
+        // removeBypassingUids() in updateVpnUid() depends on mUidToNetworkPerm to check if the
+        // package can bypass VPN.
         updateVpnUid(uid, false /* add */);
         // If the package has been removed from all users on the device, clear it form mAllApps.
         if (mPackageManager.getNameForUid(uid) == null) {
             mAllApps.remove(appId);
         }
 
-        final int currentPermission = mApps.get(appId, PERMISSION_NONE);
+        final int currentPermission = mUidToNetworkPerm.get(uid, PERMISSION_NONE);
         final int permission = highestUidNetworkPermission(uid);
         mPermissionUpdateLogs.log("Package remove: name=" + packageName + ", uid=" + uid
                 + ", nPerm=(" + permissionToString(permission) + "/"
                 + permissionToString(currentPermission) + ")"
                 + ", tPerm=" + permissionToString(trafficPerm));
-        if (permission == PERMISSION_SYSTEM) {
-            // An app with this UID still has the SYSTEM permission.
-            // Therefore, this UID must already have the SYSTEM permission.
-            // Nothing to do.
-            return;
-        }
-        // If the permissions of this UID have not changed, do nothing.
-        if (permission == currentPermission) return;
-
-        final SparseIntArray apps = new SparseIntArray();
-        if (permission != PERMISSION_NONE) {
-            mApps.put(appId, permission);
-            apps.put(appId, permission);
-            update(mUsers, apps, true);
-        } else {
-            mApps.delete(appId);
-            apps.put(appId, PERMISSION_NETWORK);  // doesn't matter which permission we pick here
-            update(mUsers, apps, false);
+        if (permission != currentPermission) {
+            final SparseIntArray apps = new SparseIntArray();
+            if (permission == PERMISSION_NONE) {
+                mUidToNetworkPerm.delete(uid);
+                apps.put(uid, PERMISSION_NETWORK);  // doesn't matter which permission we pick here
+                sendUidsNetworkPermission(apps, false);
+            } else {
+                mUidToNetworkPerm.put(uid, permission);
+                apps.put(uid, permission);
+                sendUidsNetworkPermission(apps, true);
+            }
         }
     }
 
     private static int getNetdPermissionMask(String[] requestedPermissions,
                                              int[] requestedPermissionsFlags) {
-        int permissions = 0;
+        int permissions = PERMISSION_NONE;
         if (requestedPermissions == null || requestedPermissionsFlags == null) return permissions;
         for (int i = 0; i < requestedPermissions.length; i++) {
             if (requestedPermissions[i].equals(INTERNET)
@@ -632,12 +679,23 @@
         return permissions;
     }
 
-    private PackageInfo getPackageInfo(String packageName) {
+    private synchronized PackageManager getPackageManagerAsUser(UserHandle user) {
+        PackageManager pm = mUsersPackageManager.get(user);
+        if (pm == null) {
+            pm = mContext.createContextAsUser(user, 0 /* flag */).getPackageManager();
+            mUsersPackageManager.put(user, pm);
+        }
+        return pm;
+    }
+
+    private PackageInfo getPackageInfoAsUser(String packageName, UserHandle user) {
         try {
-            PackageInfo app = mPackageManager.getPackageInfo(packageName, GET_PERMISSIONS
-                    | MATCH_ANY_USER);
-            return app;
+            final PackageInfo info = getPackageManagerAsUser(user)
+                    .getPackageInfo(packageName, GET_PERMISSIONS);
+            return info;
         } catch (NameNotFoundException e) {
+            // App not found.
+            loge("NameNotFoundException " + packageName);
             return null;
         }
     }
@@ -725,8 +783,7 @@
      */
     private void removeBypassingUids(Set<Integer> uids, int vpnAppUid) {
         uids.remove(vpnAppUid);
-        uids.removeIf(uid ->
-                mApps.get(UserHandle.getAppId(uid), PERMISSION_NONE) == PERMISSION_SYSTEM);
+        uids.removeIf(uid -> mUidToNetworkPerm.get(uid, PERMISSION_NONE) == PERMISSION_SYSTEM);
     }
 
     /**
@@ -773,7 +830,7 @@
     void sendPackagePermissionsForAppId(int appId, int permissions) {
         SparseIntArray netdPermissionsAppIds = new SparseIntArray();
         netdPermissionsAppIds.put(appId, permissions);
-        sendPackagePermissionsToNetd(netdPermissionsAppIds);
+        sendAppIdsTrafficPermission(netdPermissionsAppIds);
     }
 
     /**
@@ -785,16 +842,16 @@
      * @hide
      */
     @VisibleForTesting
-    void sendPackagePermissionsToNetd(SparseIntArray netdPermissionsAppIds) {
+    void sendAppIdsTrafficPermission(SparseIntArray netdPermissionsAppIds) {
         if (mNetd == null) {
             Log.e(TAG, "Failed to get the netd service");
             return;
         }
-        ArrayList<Integer> allPermissionAppIds = new ArrayList<>();
-        ArrayList<Integer> internetPermissionAppIds = new ArrayList<>();
-        ArrayList<Integer> updateStatsPermissionAppIds = new ArrayList<>();
-        ArrayList<Integer> noPermissionAppIds = new ArrayList<>();
-        ArrayList<Integer> uninstalledAppIds = new ArrayList<>();
+        final ArrayList<Integer> allPermissionAppIds = new ArrayList<>();
+        final ArrayList<Integer> internetPermissionAppIds = new ArrayList<>();
+        final ArrayList<Integer> updateStatsPermissionAppIds = new ArrayList<>();
+        final ArrayList<Integer> noPermissionAppIds = new ArrayList<>();
+        final ArrayList<Integer> uninstalledAppIds = new ArrayList<>();
         for (int i = 0; i < netdPermissionsAppIds.size(); i++) {
             int permissions = netdPermissionsAppIds.valueAt(i);
             switch(permissions) {
@@ -865,20 +922,19 @@
         for (Integer uid : uidsToUpdate) {
             final int permission = highestUidNetworkPermission(uid);
 
-            final int appId = UserHandle.getAppId(uid);
             if (PERMISSION_NONE == permission) {
                 // Doesn't matter which permission is set here.
-                removedUids.put(appId, PERMISSION_NETWORK);
-                mApps.delete(appId);
+                removedUids.put(uid, PERMISSION_NETWORK);
+                mUidToNetworkPerm.delete(uid);
             } else {
-                updatedUids.put(appId, permission);
-                mApps.put(appId, permission);
+                updatedUids.put(uid, permission);
+                mUidToNetworkPerm.put(uid, permission);
             }
         }
 
         // Step3. Update or revoke permission for uids with netd.
-        update(mUsers, updatedUids, true /* add */);
-        update(mUsers, removedUids, false /* add */);
+        sendUidsNetworkPermission(updatedUids, true /* add */);
+        sendUidsNetworkPermission(removedUids, false /* add */);
         mPermissionUpdateLogs.log("Setting change: update=" + updatedUids
                 + ", remove=" + removedUids);
     }
@@ -890,11 +946,13 @@
         }
 
         for (String app : pkgList) {
-            final PackageInfo info = getPackageInfo(app);
-            if (info == null || info.applicationInfo == null) continue;
+            for (UserHandle user : mUsers) {
+                final PackageInfo info = getPackageInfoAsUser(app, user);
+                if (info == null || info.applicationInfo == null) continue;
 
-            final int appId = info.applicationInfo.uid;
-            onPackageAdded(app, appId); // Use onPackageAdded to add package one by one.
+                final int uid = info.applicationInfo.uid;
+                onPackageAdded(app, uid); // Use onPackageAdded to add package one by one.
+            }
         }
     }
 
diff --git a/tests/cts/hostside/app/src/com/android/cts/net/hostside/VpnTest.java b/tests/cts/hostside/app/src/com/android/cts/net/hostside/VpnTest.java
index 311b3f0..3abc4fb 100755
--- a/tests/cts/hostside/app/src/com/android/cts/net/hostside/VpnTest.java
+++ b/tests/cts/hostside/app/src/com/android/cts/net/hostside/VpnTest.java
@@ -727,11 +727,15 @@
     @Test
     public void testDefault() throws Exception {
         if (!supportedHardware()) return;
-        // If adb TCP port opened, this test may running by adb over network.
-        // All of socket would be destroyed in this test. So this test don't
-        // support adb over network, see b/119382723.
-        if (SystemProperties.getInt("persist.adb.tcp.port", -1) > -1
-                || SystemProperties.getInt("service.adb.tcp.port", -1) > -1) {
+        if (!SdkLevel.isAtLeastS() && (
+                SystemProperties.getInt("persist.adb.tcp.port", -1) > -1
+                        || SystemProperties.getInt("service.adb.tcp.port", -1) > -1)) {
+            // If adb TCP port opened, this test may running by adb over network.
+            // All of socket would be destroyed in this test. So this test don't
+            // support adb over network, see b/119382723.
+            // This is fixed in S, but still affects previous Android versions,
+            // and this test must be backwards compatible.
+            // TODO: Delete this code entirely when R is no longer supported.
             Log.i(TAG, "adb is running over the network, so skip this test");
             return;
         }
@@ -842,11 +846,16 @@
         FileDescriptor remoteFd = openSocketFdInOtherApp(TEST_HOST, 80, TIMEOUT_MS);
 
         String disallowedApps = mRemoteSocketFactoryClient.getPackageName() + "," + mPackageName;
-        // If adb TCP port opened, this test may running by adb over TCP.
-        // Add com.android.shell appllication into blacklist to exclude adb socket for VPN test,
-        // see b/119382723.
-        // Note: The test don't support running adb over network for root device
-        disallowedApps = disallowedApps + ",com.android.shell";
+        if (!SdkLevel.isAtLeastS()) {
+            // If adb TCP port opened, this test may running by adb over TCP.
+            // Add com.android.shell application into disallowedApps to exclude adb socket for VPN
+            // test, see b/119382723 (the test doesn't support adb over TCP when adb runs as root).
+            //
+            // This is fixed in S, but still affects previous Android versions,
+            // and this test must be backwards compatible.
+            // TODO: Delete this code entirely when R is no longer supported.
+            disallowedApps = disallowedApps + ",com.android.shell";
+        }
         Log.i(TAG, "Append shell app to disallowedApps: " + disallowedApps);
         startVpn(new String[] {"192.0.2.2/32", "2001:db8:1:2::ffe/128"},
                  new String[] {"192.0.2.0/24", "2001:db8::/32"},
@@ -930,11 +939,17 @@
         if (!supportedHardware()) return;
         ProxyInfo initialProxy = mCM.getDefaultProxy();
 
-        // If adb TCP port opened, this test may running by adb over TCP.
-        // Add com.android.shell appllication into blacklist to exclude adb socket for VPN test,
-        // see b/119382723.
-        // Note: The test don't support running adb over network for root device
-        String disallowedApps = mPackageName + ",com.android.shell";
+        String disallowedApps = mPackageName;
+        if (!SdkLevel.isAtLeastS()) {
+            // If adb TCP port opened, this test may running by adb over TCP.
+            // Add com.android.shell application into disallowedApps to exclude adb socket for VPN
+            // test, see b/119382723 (the test doesn't support adb over TCP when adb runs as root).
+            //
+            // This is fixed in S, but still affects previous Android versions,
+            // and this test must be backwards compatible.
+            // TODO: Delete this code entirely when R is no longer supported.
+            disallowedApps += ",com.android.shell";
+        }
         ProxyInfo testProxyInfo = ProxyInfo.buildDirectProxy("10.0.0.1", 8888);
         startVpn(new String[] {"192.0.2.2/32", "2001:db8:1:2::ffe/128"},
                 new String[] {"0.0.0.0/0", "::/0"}, "", disallowedApps,
diff --git a/tests/cts/net/src/android/net/cts/CaptivePortalTest.kt b/tests/cts/net/src/android/net/cts/CaptivePortalTest.kt
index 9f98e3f..a378aa7 100644
--- a/tests/cts/net/src/android/net/cts/CaptivePortalTest.kt
+++ b/tests/cts/net/src/android/net/cts/CaptivePortalTest.kt
@@ -40,7 +40,6 @@
 import android.net.cts.util.CtsNetUtils
 import android.net.util.NetworkStackUtils.TEST_CAPTIVE_PORTAL_HTTPS_URL
 import android.net.util.NetworkStackUtils.TEST_CAPTIVE_PORTAL_HTTP_URL
-import android.net.wifi.WifiManager
 import android.os.Build
 import android.platform.test.annotations.AppModeFull
 import android.provider.DeviceConfig
@@ -77,7 +76,7 @@
 private const val LOCALHOST_HOSTNAME = "localhost"
 
 // Re-connecting to the AP, obtaining an IP address, revalidating can take a long time
-private const val WIFI_CONNECT_TIMEOUT_MS = 120_000L
+private const val WIFI_CONNECT_TIMEOUT_MS = 40_000L
 private const val TEST_TIMEOUT_MS = 10_000L
 
 private const val TAG = "CaptivePortalTest"
@@ -94,7 +93,6 @@
 @RunWith(AndroidJUnit4::class)
 class CaptivePortalTest {
     private val context: android.content.Context by lazy { getInstrumentation().context }
-    private val wm by lazy { context.getSystemService(WifiManager::class.java) }
     private val cm by lazy { context.getSystemService(ConnectivityManager::class.java) }
     private val pm by lazy { context.packageManager }
     private val utils by lazy { CtsNetUtils(context) }
diff --git a/tests/cts/net/src/android/net/cts/NetworkValidationTest.kt b/tests/cts/net/src/android/net/cts/NetworkValidationTest.kt
index 5290f0d..8e98dba 100644
--- a/tests/cts/net/src/android/net/cts/NetworkValidationTest.kt
+++ b/tests/cts/net/src/android/net/cts/NetworkValidationTest.kt
@@ -25,7 +25,6 @@
 import android.net.InetAddresses
 import android.net.NetworkCapabilities.NET_CAPABILITY_CAPTIVE_PORTAL
 import android.net.NetworkCapabilities.NET_CAPABILITY_TRUSTED
-import android.net.NetworkCapabilities.TRANSPORT_ETHERNET
 import android.net.NetworkCapabilities.TRANSPORT_TEST
 import android.net.NetworkRequest
 import android.net.TestNetworkInterface
@@ -96,7 +95,6 @@
     private val ethRequest = NetworkRequest.Builder()
             // ETHERNET|TEST transport networks do not have NET_CAPABILITY_TRUSTED
             .removeCapability(NET_CAPABILITY_TRUSTED)
-            .addTransportType(TRANSPORT_ETHERNET)
             .addTransportType(TRANSPORT_TEST).build()
     private val ethRequestCb = TestableNetworkCallback()
 
diff --git a/tests/cts/net/src/android/net/cts/NetworkValidationTestUtil.kt b/tests/cts/net/src/android/net/cts/NetworkValidationTestUtil.kt
index dde14ac..391d03a 100644
--- a/tests/cts/net/src/android/net/cts/NetworkValidationTestUtil.kt
+++ b/tests/cts/net/src/android/net/cts/NetworkValidationTestUtil.kt
@@ -19,12 +19,20 @@
 import android.Manifest
 import android.net.util.NetworkStackUtils
 import android.provider.DeviceConfig
+import android.provider.DeviceConfig.NAMESPACE_CONNECTIVITY
+import android.util.Log
 import com.android.testutils.runAsShell
+import com.android.testutils.tryTest
+import java.util.concurrent.CompletableFuture
+import java.util.concurrent.Executor
+import java.util.concurrent.TimeUnit
 
 /**
  * Collection of utility methods for configuring network validation.
  */
 internal object NetworkValidationTestUtil {
+    val TAG = NetworkValidationTestUtil::class.simpleName
+    const val TIMEOUT_MS = 20_000L
 
     /**
      * Clear the test network validation URLs.
@@ -59,10 +67,52 @@
     @JvmStatic fun setUrlExpirationDeviceConfig(timestamp: Long?) =
             setConfig(NetworkStackUtils.TEST_URL_EXPIRATION_TIME, timestamp?.toString())
 
-    private fun setConfig(configKey: String, value: String?) {
-        runAsShell(Manifest.permission.WRITE_DEVICE_CONFIG) {
-            DeviceConfig.setProperty(
-                    DeviceConfig.NAMESPACE_CONNECTIVITY, configKey, value, false /* makeDefault */)
+    private fun setConfig(configKey: String, value: String?): String? {
+        Log.i(TAG, "Setting config \"$configKey\" to \"$value\"")
+        val readWritePermissions = arrayOf(
+                Manifest.permission.READ_DEVICE_CONFIG,
+                Manifest.permission.WRITE_DEVICE_CONFIG)
+
+        val existingValue = runAsShell(*readWritePermissions) {
+            DeviceConfig.getProperty(NAMESPACE_CONNECTIVITY, configKey)
+        }
+        if (existingValue == value) {
+            // Already the correct value. There may be a race if a change is already in flight,
+            // but if multiple threads update the config there is no way to fix that anyway.
+            Log.i(TAG, "\$configKey\" already had value \"$value\"")
+            return value
+        }
+
+        val future = CompletableFuture<String>()
+        val listener = DeviceConfig.OnPropertiesChangedListener {
+            // The listener receives updates for any change to any key, so don't react to
+            // changes that do not affect the relevant key
+            if (!it.keyset.contains(configKey)) return@OnPropertiesChangedListener
+            if (it.getString(configKey, null) == value) {
+                future.complete(value)
+            }
+        }
+
+        return tryTest {
+            runAsShell(*readWritePermissions) {
+                DeviceConfig.addOnPropertiesChangedListener(
+                        NAMESPACE_CONNECTIVITY,
+                        inlineExecutor,
+                        listener)
+                DeviceConfig.setProperty(
+                        NAMESPACE_CONNECTIVITY,
+                        configKey,
+                        value,
+                        false /* makeDefault */)
+                // Don't drop the permission until the config is applied, just in case
+                future.get(TIMEOUT_MS, TimeUnit.MILLISECONDS)
+            }.also {
+                Log.i(TAG, "Config \"$configKey\" successfully set to \"$value\"")
+            }
+        } cleanup {
+            DeviceConfig.removeOnPropertiesChangedListener(listener)
         }
     }
-}
\ No newline at end of file
+
+    private val inlineExecutor get() = Executor { r -> r.run() }
+}
diff --git a/tests/cts/net/util/java/android/net/cts/util/CtsNetUtils.java b/tests/cts/net/util/java/android/net/cts/util/CtsNetUtils.java
index e17200e..ce873f7 100644
--- a/tests/cts/net/util/java/android/net/cts/util/CtsNetUtils.java
+++ b/tests/cts/net/util/java/android/net/cts/util/CtsNetUtils.java
@@ -306,14 +306,18 @@
         }
 
         try {
+            if (wasWifiConnected) {
+                // Make sure the callback is registered before turning off WiFi.
+                callback.waitForAvailable();
+            }
             SystemUtil.runShellCommand("svc wifi disable");
             if (wasWifiConnected) {
                 // Ensure we get both an onLost callback and a CONNECTIVITY_ACTION.
                 assertNotNull("Did not receive onLost callback after disabling wifi",
                         callback.waitForLost());
-            }
-            if (wasWifiConnected && expectLegacyBroadcast) {
-                assertTrue("Wifi failed to reach DISCONNECTED state.", receiver.waitForState());
+                if (expectLegacyBroadcast) {
+                    assertTrue("Wifi failed to reach DISCONNECTED state.", receiver.waitForState());
+                }
             }
         } catch (InterruptedException ex) {
             fail("disconnectFromWifi was interrupted");
@@ -472,6 +476,7 @@
         NetworkCallback callback = new NetworkCallback() {
             @Override
             public void onLinkPropertiesChanged(Network n, LinkProperties lp) {
+                Log.i(TAG, "Link properties of network " + n + " changed to " + lp);
                 if (requiresValidatedServer && lp.getValidatedPrivateDnsServers().isEmpty()) {
                     return;
                 }
@@ -599,12 +604,14 @@
 
         @Override
         public void onAvailable(Network network) {
+            Log.i(TAG, "CtsNetUtils TestNetworkCallback onAvailable " + network);
             currentNetwork = network;
             mAvailableCv.open();
         }
 
         @Override
         public void onLost(Network network) {
+            Log.i(TAG, "CtsNetUtils TestNetworkCallback onLost " + network);
             lastLostNetwork = network;
             if (network.equals(currentNetwork)) {
                 mAvailableCv.close();
diff --git a/tests/unit/Android.bp b/tests/unit/Android.bp
index 448869d..bd30b77 100644
--- a/tests/unit/Android.bp
+++ b/tests/unit/Android.bp
@@ -93,6 +93,17 @@
     ]
 }
 
+// Subset of services-core used to by ConnectivityService tests to test VPN realistically.
+// This is stripped by jarjar (see rules below) from other unrelated classes, so tests do not
+// include most classes from services-core, which are unrelated and cause wrong code coverage
+// calculations.
+java_library {
+    name: "services.core-vpn",
+    static_libs: ["services.core"],
+    jarjar_rules: "vpn-jarjar-rules.txt",
+    visibility: ["//visibility:private"],
+}
+
 android_library {
     name: "FrameworksNetTestsLib",
     min_sdk_version: "30",
@@ -116,11 +127,11 @@
         "framework-protos",
         "mockito-target-minus-junit4",
         "net-tests-utils",
+        "net-utils-services-common",
         "platform-compat-test-rules",
         "platform-test-annotations",
         "service-connectivity-pre-jarjar",
-        "services.core",
-        "services.net",
+        "services.core-vpn",
     ],
     libs: [
         "android.net.ipsec.ike.stubs.module_lib",
@@ -144,6 +155,8 @@
     srcs: [":non-connectivity-module-test"],
     test_suites: ["device-tests"],
     static_libs: [
+        "services.core",
+        "services.net",
         "FrameworksNetTestsLib",
     ],
     libs: [
diff --git a/tests/unit/java/android/net/nsd/NsdManagerTest.java b/tests/unit/java/android/net/nsd/NsdManagerTest.java
index de77d23..30b8fcd 100644
--- a/tests/unit/java/android/net/nsd/NsdManagerTest.java
+++ b/tests/unit/java/android/net/nsd/NsdManagerTest.java
@@ -20,38 +20,32 @@
 import static libcore.junit.util.compat.CoreCompatChangeRule.EnableCompatChanges;
 
 import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.fail;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.reset;
-import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.timeout;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
 import android.compat.testing.PlatformCompatChangeRule;
 import android.content.Context;
 import android.os.Build;
-import android.os.Handler;
-import android.os.HandlerThread;
-import android.os.Looper;
-import android.os.Message;
-import android.os.Messenger;
 
 import androidx.test.filters.SmallTest;
 
-import com.android.internal.util.AsyncChannel;
 import com.android.testutils.DevSdkIgnoreRule;
 import com.android.testutils.DevSdkIgnoreRunner;
-import com.android.testutils.HandlerUtils;
+import com.android.testutils.ExceptionUtils;
 
-import org.junit.After;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.TestRule;
 import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
@@ -67,9 +61,10 @@
 
     @Mock Context mContext;
     @Mock INsdManager mService;
-    MockServiceHandler mServiceHandler;
+    @Mock INsdServiceConnector mServiceConn;
 
     NsdManager mManager;
+    INsdManagerCallback mCallback;
 
     long mTimeoutMs = 200; // non-final so that tests can adjust the value.
 
@@ -77,91 +72,85 @@
     public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
 
-        mServiceHandler = spy(MockServiceHandler.create(mContext));
-        doReturn(new Messenger(mServiceHandler)).when(mService).getMessenger();
-    }
-
-    @After
-    public void tearDown() throws Exception {
-        HandlerUtils.waitForIdle(mServiceHandler, mTimeoutMs);
-        mServiceHandler.chan.disconnect();
-        mServiceHandler.stop();
-        if (mManager != null) {
-            mManager.disconnect();
-        }
+        doReturn(mServiceConn).when(mService).connect(any());
+        mManager = new NsdManager(mContext, mService);
+        final ArgumentCaptor<INsdManagerCallback> cbCaptor = ArgumentCaptor.forClass(
+                INsdManagerCallback.class);
+        verify(mService).connect(cbCaptor.capture());
+        mCallback = cbCaptor.getValue();
     }
 
     @Test
     @EnableCompatChanges(NsdManager.RUN_NATIVE_NSD_ONLY_IF_LEGACY_APPS)
-    public void testResolveServiceS() {
-        mManager = makeNsdManagerS();
+    public void testResolveServiceS() throws Exception {
+        verify(mServiceConn, never()).startDaemon();
         doTestResolveService();
     }
 
     @Test
     @DisableCompatChanges(NsdManager.RUN_NATIVE_NSD_ONLY_IF_LEGACY_APPS)
-    public void testResolveServicePreS() {
-        mManager = makeNsdManagerPreS();
+    public void testResolveServicePreS() throws Exception {
+        verify(mServiceConn).startDaemon();
         doTestResolveService();
     }
 
     @Test
     @EnableCompatChanges(NsdManager.RUN_NATIVE_NSD_ONLY_IF_LEGACY_APPS)
-    public void testDiscoverServiceS() {
-        mManager = makeNsdManagerS();
+    public void testDiscoverServiceS() throws Exception {
+        verify(mServiceConn, never()).startDaemon();
         doTestDiscoverService();
     }
 
     @Test
     @DisableCompatChanges(NsdManager.RUN_NATIVE_NSD_ONLY_IF_LEGACY_APPS)
-    public void testDiscoverServicePreS() {
-        mManager = makeNsdManagerPreS();
+    public void testDiscoverServicePreS() throws Exception {
+        verify(mServiceConn).startDaemon();
         doTestDiscoverService();
     }
 
     @Test
     @EnableCompatChanges(NsdManager.RUN_NATIVE_NSD_ONLY_IF_LEGACY_APPS)
-    public void testParallelResolveServiceS() {
-        mManager = makeNsdManagerS();
+    public void testParallelResolveServiceS() throws Exception {
+        verify(mServiceConn, never()).startDaemon();
         doTestParallelResolveService();
     }
 
     @Test
     @DisableCompatChanges(NsdManager.RUN_NATIVE_NSD_ONLY_IF_LEGACY_APPS)
-    public void testParallelResolveServicePreS() {
-        mManager = makeNsdManagerPreS();
+    public void testParallelResolveServicePreS() throws Exception {
+        verify(mServiceConn).startDaemon();
         doTestParallelResolveService();
     }
 
     @Test
     @EnableCompatChanges(NsdManager.RUN_NATIVE_NSD_ONLY_IF_LEGACY_APPS)
-    public void testInvalidCallsS() {
-        mManager = makeNsdManagerS();
+    public void testInvalidCallsS() throws Exception {
+        verify(mServiceConn, never()).startDaemon();
         doTestInvalidCalls();
     }
 
     @Test
     @DisableCompatChanges(NsdManager.RUN_NATIVE_NSD_ONLY_IF_LEGACY_APPS)
-    public void testInvalidCallsPreS() {
-        mManager = makeNsdManagerPreS();
+    public void testInvalidCallsPreS() throws Exception {
+        verify(mServiceConn).startDaemon();
         doTestInvalidCalls();
     }
 
     @Test
     @EnableCompatChanges(NsdManager.RUN_NATIVE_NSD_ONLY_IF_LEGACY_APPS)
-    public void testRegisterServiceS() {
-        mManager = makeNsdManagerS();
+    public void testRegisterServiceS() throws Exception {
+        verify(mServiceConn, never()).startDaemon();
         doTestRegisterService();
     }
 
     @Test
     @DisableCompatChanges(NsdManager.RUN_NATIVE_NSD_ONLY_IF_LEGACY_APPS)
-    public void testRegisterServicePreS() {
-        mManager = makeNsdManagerPreS();
+    public void testRegisterServicePreS() throws Exception {
+        verify(mServiceConn).startDaemon();
         doTestRegisterService();
     }
 
-    public void doTestResolveService() {
+    private void doTestResolveService() throws Exception {
         NsdManager manager = mManager;
 
         NsdServiceInfo request = new NsdServiceInfo("a_name", "a_type");
@@ -169,18 +158,19 @@
         NsdManager.ResolveListener listener = mock(NsdManager.ResolveListener.class);
 
         manager.resolveService(request, listener);
-        int key1 = verifyRequest(NsdManager.RESOLVE_SERVICE);
+        int key1 = getRequestKey(req -> verify(mServiceConn).resolveService(req.capture(), any()));
         int err = 33;
-        sendResponse(NsdManager.RESOLVE_SERVICE_FAILED, err, key1, null);
+        mCallback.onResolveServiceFailed(key1, err);
         verify(listener, timeout(mTimeoutMs).times(1)).onResolveFailed(request, err);
 
         manager.resolveService(request, listener);
-        int key2 = verifyRequest(NsdManager.RESOLVE_SERVICE);
-        sendResponse(NsdManager.RESOLVE_SERVICE_SUCCEEDED, 0, key2, reply);
+        int key2 = getRequestKey(req ->
+                verify(mServiceConn, times(2)).resolveService(req.capture(), any()));
+        mCallback.onResolveServiceSucceeded(key2, reply);
         verify(listener, timeout(mTimeoutMs).times(1)).onServiceResolved(reply);
     }
 
-    public void doTestParallelResolveService() {
+    private void doTestParallelResolveService() throws Exception {
         NsdManager manager = mManager;
 
         NsdServiceInfo request = new NsdServiceInfo("a_name", "a_type");
@@ -190,19 +180,20 @@
         NsdManager.ResolveListener listener2 = mock(NsdManager.ResolveListener.class);
 
         manager.resolveService(request, listener1);
-        int key1 = verifyRequest(NsdManager.RESOLVE_SERVICE);
+        int key1 = getRequestKey(req -> verify(mServiceConn).resolveService(req.capture(), any()));
 
         manager.resolveService(request, listener2);
-        int key2 = verifyRequest(NsdManager.RESOLVE_SERVICE);
+        int key2 = getRequestKey(req ->
+                verify(mServiceConn, times(2)).resolveService(req.capture(), any()));
 
-        sendResponse(NsdManager.RESOLVE_SERVICE_SUCCEEDED, 0, key2, reply);
-        sendResponse(NsdManager.RESOLVE_SERVICE_SUCCEEDED, 0, key1, reply);
+        mCallback.onResolveServiceSucceeded(key2, reply);
+        mCallback.onResolveServiceSucceeded(key1, reply);
 
         verify(listener1, timeout(mTimeoutMs).times(1)).onServiceResolved(reply);
         verify(listener2, timeout(mTimeoutMs).times(1)).onServiceResolved(reply);
     }
 
-    public void doTestRegisterService() {
+    private void doTestRegisterService() throws Exception {
         NsdManager manager = mManager;
 
         NsdServiceInfo request1 = new NsdServiceInfo("a_name", "a_type");
@@ -214,40 +205,43 @@
 
         // Register two services
         manager.registerService(request1, PROTOCOL, listener1);
-        int key1 = verifyRequest(NsdManager.REGISTER_SERVICE);
+        int key1 = getRequestKey(req -> verify(mServiceConn).registerService(req.capture(), any()));
 
         manager.registerService(request2, PROTOCOL, listener2);
-        int key2 = verifyRequest(NsdManager.REGISTER_SERVICE);
+        int key2 = getRequestKey(req ->
+                verify(mServiceConn, times(2)).registerService(req.capture(), any()));
 
         // First reques fails, second request succeeds
-        sendResponse(NsdManager.REGISTER_SERVICE_SUCCEEDED, 0, key2, request2);
+        mCallback.onRegisterServiceSucceeded(key2, request2);
         verify(listener2, timeout(mTimeoutMs).times(1)).onServiceRegistered(request2);
 
         int err = 1;
-        sendResponse(NsdManager.REGISTER_SERVICE_FAILED, err, key1, request1);
+        mCallback.onRegisterServiceFailed(key1, err);
         verify(listener1, timeout(mTimeoutMs).times(1)).onRegistrationFailed(request1, err);
 
         // Client retries first request, it succeeds
         manager.registerService(request1, PROTOCOL, listener1);
-        int key3 = verifyRequest(NsdManager.REGISTER_SERVICE);
+        int key3 = getRequestKey(req ->
+                verify(mServiceConn, times(3)).registerService(req.capture(), any()));
 
-        sendResponse(NsdManager.REGISTER_SERVICE_SUCCEEDED, 0, key3, request1);
+        mCallback.onRegisterServiceSucceeded(key3, request1);
         verify(listener1, timeout(mTimeoutMs).times(1)).onServiceRegistered(request1);
 
         // First request is unregistered, it succeeds
         manager.unregisterService(listener1);
-        int key3again = verifyRequest(NsdManager.UNREGISTER_SERVICE);
+        int key3again = getRequestKey(req -> verify(mServiceConn).unregisterService(req.capture()));
         assertEquals(key3, key3again);
 
-        sendResponse(NsdManager.UNREGISTER_SERVICE_SUCCEEDED, 0, key3again, null);
+        mCallback.onUnregisterServiceSucceeded(key3again);
         verify(listener1, timeout(mTimeoutMs).times(1)).onServiceUnregistered(request1);
 
         // Second request is unregistered, it fails
         manager.unregisterService(listener2);
-        int key2again = verifyRequest(NsdManager.UNREGISTER_SERVICE);
+        int key2again = getRequestKey(req ->
+                verify(mServiceConn, times(2)).unregisterService(req.capture()));
         assertEquals(key2, key2again);
 
-        sendResponse(NsdManager.UNREGISTER_SERVICE_FAILED, err, key2again, null);
+        mCallback.onUnregisterServiceFailed(key2again, err);
         verify(listener2, timeout(mTimeoutMs).times(1)).onUnregistrationFailed(request2, err);
 
         // TODO: do not unregister listener until service is unregistered
@@ -260,7 +254,7 @@
         //verify(listener2, timeout(mTimeoutMs).times(1)).onServiceUnregistered(request2);
     }
 
-    public void doTestDiscoverService() {
+    private void doTestDiscoverService() throws Exception {
         NsdManager manager = mManager;
 
         NsdServiceInfo reply1 = new NsdServiceInfo("a_name", "a_type");
@@ -271,69 +265,73 @@
 
         // Client registers for discovery, request fails
         manager.discoverServices("a_type", PROTOCOL, listener);
-        int key1 = verifyRequest(NsdManager.DISCOVER_SERVICES);
+        int key1 = getRequestKey(req ->
+                verify(mServiceConn).discoverServices(req.capture(), any()));
 
         int err = 1;
-        sendResponse(NsdManager.DISCOVER_SERVICES_FAILED, err, key1, null);
+        mCallback.onDiscoverServicesFailed(key1, err);
         verify(listener, timeout(mTimeoutMs).times(1)).onStartDiscoveryFailed("a_type", err);
 
         // Client retries, request succeeds
         manager.discoverServices("a_type", PROTOCOL, listener);
-        int key2 = verifyRequest(NsdManager.DISCOVER_SERVICES);
+        int key2 = getRequestKey(req ->
+                verify(mServiceConn, times(2)).discoverServices(req.capture(), any()));
 
-        sendResponse(NsdManager.DISCOVER_SERVICES_STARTED, 0, key2, reply1);
+        mCallback.onDiscoverServicesStarted(key2, reply1);
         verify(listener, timeout(mTimeoutMs).times(1)).onDiscoveryStarted("a_type");
 
 
         // mdns notifies about services
-        sendResponse(NsdManager.SERVICE_FOUND, 0, key2, reply1);
+        mCallback.onServiceFound(key2, reply1);
         verify(listener, timeout(mTimeoutMs).times(1)).onServiceFound(reply1);
 
-        sendResponse(NsdManager.SERVICE_FOUND, 0, key2, reply2);
+        mCallback.onServiceFound(key2, reply2);
         verify(listener, timeout(mTimeoutMs).times(1)).onServiceFound(reply2);
 
-        sendResponse(NsdManager.SERVICE_LOST, 0, key2, reply2);
+        mCallback.onServiceLost(key2, reply2);
         verify(listener, timeout(mTimeoutMs).times(1)).onServiceLost(reply2);
 
 
         // Client unregisters its listener
         manager.stopServiceDiscovery(listener);
-        int key2again = verifyRequest(NsdManager.STOP_DISCOVERY);
+        int key2again = getRequestKey(req -> verify(mServiceConn).stopDiscovery(req.capture()));
         assertEquals(key2, key2again);
 
         // TODO: unregister listener immediately and stop notifying it about services
         // Notifications are still passed to the client's listener
-        sendResponse(NsdManager.SERVICE_LOST, 0, key2, reply1);
+        mCallback.onServiceLost(key2, reply1);
         verify(listener, timeout(mTimeoutMs).times(1)).onServiceLost(reply1);
 
         // Client is notified of complete unregistration
-        sendResponse(NsdManager.STOP_DISCOVERY_SUCCEEDED, 0, key2again, "a_type");
+        mCallback.onStopDiscoverySucceeded(key2again);
         verify(listener, timeout(mTimeoutMs).times(1)).onDiscoveryStopped("a_type");
 
         // Notifications are not passed to the client anymore
-        sendResponse(NsdManager.SERVICE_FOUND, 0, key2, reply3);
+        mCallback.onServiceFound(key2, reply3);
         verify(listener, timeout(mTimeoutMs).times(0)).onServiceLost(reply3);
 
 
         // Client registers for service discovery
         reset(listener);
         manager.discoverServices("a_type", PROTOCOL, listener);
-        int key3 = verifyRequest(NsdManager.DISCOVER_SERVICES);
+        int key3 = getRequestKey(req ->
+                verify(mServiceConn, times(3)).discoverServices(req.capture(), any()));
 
-        sendResponse(NsdManager.DISCOVER_SERVICES_STARTED, 0, key3, reply1);
+        mCallback.onDiscoverServicesStarted(key3, reply1);
         verify(listener, timeout(mTimeoutMs).times(1)).onDiscoveryStarted("a_type");
 
         // Client unregisters immediately, it fails
         manager.stopServiceDiscovery(listener);
-        int key3again = verifyRequest(NsdManager.STOP_DISCOVERY);
+        int key3again = getRequestKey(req ->
+                verify(mServiceConn, times(2)).stopDiscovery(req.capture()));
         assertEquals(key3, key3again);
 
         err = 2;
-        sendResponse(NsdManager.STOP_DISCOVERY_FAILED, err, key3again, "a_type");
+        mCallback.onStopDiscoveryFailed(key3again, err);
         verify(listener, timeout(mTimeoutMs).times(1)).onStopDiscoveryFailed("a_type", err);
 
         // New notifications are not passed to the client anymore
-        sendResponse(NsdManager.SERVICE_FOUND, 0, key3, reply1);
+        mCallback.onServiceFound(key3, reply1);
         verify(listener, timeout(mTimeoutMs).times(0)).onServiceFound(reply1);
     }
 
@@ -398,77 +396,10 @@
         }
     }
 
-    NsdManager makeNsdManagerS() {
-        // Expect we'll get 2 AsyncChannel related msgs.
-        return makeManager(2);
-    }
-
-    NsdManager makeNsdManagerPreS() {
-        // Expect we'll get 3 msgs. 2 AsyncChannel related msgs + 1 additional daemon startup msg.
-        return makeManager(3);
-    }
-
-    NsdManager makeManager(int expectedMsgCount) {
-        NsdManager manager = new NsdManager(mContext, mService);
-        // Acknowledge first two messages connecting the AsyncChannel.
-        verify(mServiceHandler, timeout(mTimeoutMs).times(expectedMsgCount)).handleMessage(any());
-
-        reset(mServiceHandler);
-        assertNotNull(mServiceHandler.chan);
-        return manager;
-    }
-
-    int verifyRequest(int expectedMessageType) {
-        HandlerUtils.waitForIdle(mServiceHandler, mTimeoutMs);
-        verify(mServiceHandler, timeout(mTimeoutMs)).handleMessage(any());
-        reset(mServiceHandler);
-        Message received = mServiceHandler.getLastMessage();
-        assertEquals(NsdManager.nameOf(expectedMessageType), NsdManager.nameOf(received.what));
-        return received.arg2;
-    }
-
-    void sendResponse(int replyType, int arg, int key, Object obj) {
-        mServiceHandler.chan.sendMessage(replyType, arg, key, obj);
-    }
-
-    // Implements the server side of AsyncChannel connection protocol
-    public static class MockServiceHandler extends Handler {
-        public final Context context;
-        public AsyncChannel chan;
-        public Message lastMessage;
-
-        MockServiceHandler(Looper l, Context c) {
-            super(l);
-            context = c;
-        }
-
-        synchronized Message getLastMessage() {
-            return lastMessage;
-        }
-
-        synchronized void setLastMessage(Message msg) {
-            lastMessage = obtainMessage();
-            lastMessage.copyFrom(msg);
-        }
-
-        @Override
-        public void handleMessage(Message msg) {
-            setLastMessage(msg);
-            if (msg.what == AsyncChannel.CMD_CHANNEL_FULL_CONNECTION) {
-                chan = new AsyncChannel();
-                chan.connect(context, this, msg.replyTo);
-                chan.sendMessage(AsyncChannel.CMD_CHANNEL_FULLY_CONNECTED);
-            }
-        }
-
-        void stop() {
-            getLooper().quitSafely();
-        }
-
-        static MockServiceHandler create(Context context) {
-            HandlerThread t = new HandlerThread("mock-service-handler");
-            t.start();
-            return new MockServiceHandler(t.getLooper(), context);
-        }
+    int getRequestKey(ExceptionUtils.ThrowingConsumer<ArgumentCaptor<Integer>> verifier)
+            throws Exception {
+        final ArgumentCaptor<Integer> captor = ArgumentCaptor.forClass(Integer.class);
+        verifier.accept(captor);
+        return captor.getValue();
     }
 }
diff --git a/tests/unit/java/com/android/server/ConnectivityServiceTest.java b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
index e0dd886..ad2f1dc 100644
--- a/tests/unit/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/unit/java/com/android/server/ConnectivityServiceTest.java
@@ -38,7 +38,6 @@
 import static android.content.pm.PackageManager.FEATURE_WIFI;
 import static android.content.pm.PackageManager.FEATURE_WIFI_DIRECT;
 import static android.content.pm.PackageManager.GET_PERMISSIONS;
-import static android.content.pm.PackageManager.MATCH_ANY_USER;
 import static android.content.pm.PackageManager.PERMISSION_DENIED;
 import static android.content.pm.PackageManager.PERMISSION_GRANTED;
 import static android.net.ConnectivityManager.ACTION_CAPTIVE_PORTAL_SIGN_IN;
@@ -249,6 +248,8 @@
 import android.net.NetworkStateSnapshot;
 import android.net.NetworkTestResultParcelable;
 import android.net.OemNetworkPreferences;
+import android.net.PacProxyManager;
+import android.net.Proxy;
 import android.net.ProxyInfo;
 import android.net.QosCallbackException;
 import android.net.QosFilter;
@@ -478,6 +479,7 @@
     private Context mContext;
     private NetworkPolicyCallback mPolicyCallback;
     private WrappedMultinetworkPolicyTracker mPolicyTracker;
+    private ProxyTracker mProxyTracker;
     private HandlerThread mAlarmManagerThread;
     private TestNetIdManager mNetIdManager;
     private QosCallbackMockHelper mQosCallbackMockHelper;
@@ -511,7 +513,7 @@
     @Mock VpnProfileStore mVpnProfileStore;
     @Mock SystemConfigManager mSystemConfigManager;
     @Mock Resources mResources;
-    @Mock ProxyTracker mProxyTracker;
+    @Mock PacProxyManager mPacProxyManager;
 
     // BatteryStatsManager is final and cannot be mocked with regular mockito, so just mock the
     // underlying binder calls.
@@ -613,6 +615,7 @@
             if (Context.SYSTEM_CONFIG_SERVICE.equals(name)) return mSystemConfigManager;
             if (Context.NETWORK_STATS_SERVICE.equals(name)) return mStatsManager;
             if (Context.BATTERY_STATS_SERVICE.equals(name)) return mBatteryStatsManager;
+            if (Context.PAC_PROXY_SERVICE.equals(name)) return mPacProxyManager;
             return super.getSystemService(name);
         }
 
@@ -1703,6 +1706,8 @@
 
         mCsHandlerThread = new HandlerThread("TestConnectivityService");
         mVMSHandlerThread = new HandlerThread("TestVpnManagerService");
+        mProxyTracker = new ProxyTracker(mServiceContext, mock(Handler.class),
+                16 /* EVENT_PROXY_HAS_CHANGED */);
 
         initMockedResources();
         final Context mockResContext = mock(Context.class);
@@ -2000,7 +2005,7 @@
                 buildPackageInfo(/* SYSTEM */ false, APP1_UID),
                 buildPackageInfo(/* SYSTEM */ false, APP2_UID),
                 buildPackageInfo(/* SYSTEM */ false, VPN_UID)
-        })).when(mPackageManager).getInstalledPackages(eq(GET_PERMISSIONS | MATCH_ANY_USER));
+        })).when(mPackageManager).getInstalledPackagesAsUser(eq(GET_PERMISSIONS), anyInt());
 
         // Create a fake always-on VPN package.
         final int userId = UserHandle.getCallingUserId();
@@ -2135,6 +2140,39 @@
         return expected;
     }
 
+    private ExpectedBroadcast expectProxyChangeAction(ProxyInfo proxy) {
+        return registerPacProxyBroadcastThat(intent -> {
+            final ProxyInfo actualProxy = (ProxyInfo) intent.getExtra(Proxy.EXTRA_PROXY_INFO,
+                    ProxyInfo.buildPacProxy(Uri.EMPTY));
+            return proxy.equals(actualProxy);
+        });
+    }
+
+    private ExpectedBroadcast registerPacProxyBroadcast() {
+        return registerPacProxyBroadcastThat(intent -> true);
+    }
+
+    private ExpectedBroadcast registerPacProxyBroadcastThat(
+            @NonNull final Predicate<Intent> filter) {
+        final IntentFilter intentFilter = new IntentFilter(Proxy.PROXY_CHANGE_ACTION);
+        // AtomicReference allows receiver to access expected even though it is constructed later.
+        final AtomicReference<ExpectedBroadcast> expectedRef = new AtomicReference<>();
+        final BroadcastReceiver receiver = new BroadcastReceiver() {
+            public void onReceive(Context context, Intent intent) {
+                final ProxyInfo proxy = (ProxyInfo) intent.getExtra(
+                            Proxy.EXTRA_PROXY_INFO, ProxyInfo.buildPacProxy(Uri.EMPTY));
+                Log.d(TAG, "Receive PROXY_CHANGE_ACTION, proxy = " + proxy);
+                if (filter.test(intent)) {
+                    expectedRef.get().complete(intent);
+                }
+            }
+        };
+        final ExpectedBroadcast expected = new ExpectedBroadcast(receiver);
+        expectedRef.set(expected);
+        mServiceContext.registerReceiver(receiver, intentFilter);
+        return expected;
+    }
+
     private boolean extraInfoInBroadcastHasExpectedNullness(NetworkInfo ni) {
         final DetailedState state = ni.getDetailedState();
         if (state == DetailedState.CONNECTED && ni.getExtraInfo() == null) return false;
@@ -3295,7 +3333,7 @@
     private void grantUsingBackgroundNetworksPermissionForUid(
             final int uid, final String packageName) throws Exception {
         doReturn(buildPackageInfo(true /* hasSystemPermission */, uid)).when(mPackageManager)
-                .getPackageInfo(eq(packageName), eq(GET_PERMISSIONS | MATCH_ANY_USER));
+                .getPackageInfo(eq(packageName), eq(GET_PERMISSIONS));
         mService.mPermissionMonitor.onPackageAdded(packageName, uid);
     }
 
@@ -9945,7 +9983,7 @@
         final ProxyInfo testProxyInfo = ProxyInfo.buildDirectProxy("test", 8888);
         mWiFiNetworkAgent = new TestNetworkAgentWrapper(TRANSPORT_WIFI);
         final Network wifiNetwork = mWiFiNetworkAgent.getNetwork();
-        doReturn(testProxyInfo).when(mService.mProxyTracker).getGlobalProxy();
+        mProxyTracker.setGlobalProxy(testProxyInfo);
         assertEquals(testProxyInfo, mService.getProxyForNetwork(wifiNetwork));
     }
 
@@ -11353,6 +11391,7 @@
         assertNull(mService.getProxyForNetwork(null));
         assertNull(mCm.getDefaultProxy());
 
+        final ExpectedBroadcast b1 = registerPacProxyBroadcast();
         final LinkProperties lp = new LinkProperties();
         lp.setInterfaceName("tun0");
         lp.addRoute(new RouteInfo(new IpPrefix(Inet4Address.ANY, 0), null));
@@ -11362,9 +11401,10 @@
         mMockVpn.establish(lp, VPN_UID, vpnRanges);
         assertVpnUidRangesUpdated(true, vpnRanges, VPN_UID);
         // VPN is connected but proxy is not set, so there is no need to send proxy broadcast.
-        verify(mProxyTracker, never()).sendProxyBroadcast();
+        b1.expectNoBroadcast(500);
 
         // Update to new range which is old range minus APP1, i.e. only APP2
+        final ExpectedBroadcast b2 = registerPacProxyBroadcast();
         final Set<UidRange> newRanges = new HashSet<>(asList(
                 new UidRange(vpnRange.start, APP1_UID - 1),
                 new UidRange(APP1_UID + 1, vpnRange.stop)));
@@ -11375,37 +11415,37 @@
         assertVpnUidRangesUpdated(false, vpnRanges, VPN_UID);
 
         // Uid has changed but proxy is not set, so there is no need to send proxy broadcast.
-        verify(mProxyTracker, never()).sendProxyBroadcast();
+        b2.expectNoBroadcast(500);
 
         final ProxyInfo testProxyInfo = ProxyInfo.buildDirectProxy("test", 8888);
+        final ExpectedBroadcast b3 = registerPacProxyBroadcast();
         lp.setHttpProxy(testProxyInfo);
         mMockVpn.sendLinkProperties(lp);
         waitForIdle();
         // Proxy is set, so send a proxy broadcast.
-        verify(mProxyTracker, times(1)).sendProxyBroadcast();
-        reset(mProxyTracker);
+        b3.expectBroadcast();
 
+        final ExpectedBroadcast b4 = registerPacProxyBroadcast();
         mMockVpn.setUids(vpnRanges);
         waitForIdle();
         // Uid has changed and proxy is already set, so send a proxy broadcast.
-        verify(mProxyTracker, times(1)).sendProxyBroadcast();
-        reset(mProxyTracker);
+        b4.expectBroadcast();
 
+        final ExpectedBroadcast b5 = registerPacProxyBroadcast();
         // Proxy is removed, send a proxy broadcast.
         lp.setHttpProxy(null);
         mMockVpn.sendLinkProperties(lp);
         waitForIdle();
-        verify(mProxyTracker, times(1)).sendProxyBroadcast();
-        reset(mProxyTracker);
+        b5.expectBroadcast();
 
         // Proxy is added in WiFi(default network), setDefaultProxy will be called.
         final LinkProperties wifiLp = mCm.getLinkProperties(mWiFiNetworkAgent.getNetwork());
         assertNotNull(wifiLp);
+        final ExpectedBroadcast b6 = expectProxyChangeAction(testProxyInfo);
         wifiLp.setHttpProxy(testProxyInfo);
         mWiFiNetworkAgent.sendLinkProperties(wifiLp);
         waitForIdle();
-        verify(mProxyTracker, times(1)).setDefaultProxy(eq(testProxyInfo));
-        reset(mProxyTracker);
+        b6.expectBroadcast();
     }
 
     @Test
@@ -11424,18 +11464,21 @@
         lp.setHttpProxy(testProxyInfo);
         final UidRange vpnRange = PRIMARY_UIDRANGE;
         final Set<UidRange> vpnRanges = Collections.singleton(vpnRange);
+        final ExpectedBroadcast b1 = registerPacProxyBroadcast();
         mMockVpn.setOwnerAndAdminUid(VPN_UID);
         mMockVpn.registerAgent(false, vpnRanges, lp);
         // In any case, the proxy broadcast won't be sent before VPN goes into CONNECTED state.
         // Otherwise, the app that calls ConnectivityManager#getDefaultProxy() when it receives the
         // proxy broadcast will get null.
-        verify(mProxyTracker, never()).sendProxyBroadcast();
+        b1.expectNoBroadcast(500);
+
+        final ExpectedBroadcast b2 = registerPacProxyBroadcast();
         mMockVpn.connect(true /* validated */, true /* hasInternet */, false /* isStrictMode */);
         waitForIdle();
         assertVpnUidRangesUpdated(true, vpnRanges, VPN_UID);
         // Vpn is connected with proxy, so the proxy broadcast will be sent to inform the apps to
         // update their proxy data.
-        verify(mProxyTracker, times(1)).sendProxyBroadcast();
+        b2.expectBroadcast();
     }
 
     @Test
@@ -11464,10 +11507,10 @@
         final LinkProperties cellularLp = new LinkProperties();
         cellularLp.setInterfaceName(MOBILE_IFNAME);
         final ProxyInfo testProxyInfo = ProxyInfo.buildDirectProxy("test", 8888);
+        final ExpectedBroadcast b = registerPacProxyBroadcast();
         cellularLp.setHttpProxy(testProxyInfo);
         mCellNetworkAgent.sendLinkProperties(cellularLp);
-        waitForIdle();
-        verify(mProxyTracker, times(1)).sendProxyBroadcast();
+        b.expectBroadcast();
     }
 
     @Test
diff --git a/tests/unit/java/com/android/server/NsdServiceTest.java b/tests/unit/java/com/android/server/NsdServiceTest.java
index 4d2970a..6d1d765 100644
--- a/tests/unit/java/com/android/server/NsdServiceTest.java
+++ b/tests/unit/java/com/android/server/NsdServiceTest.java
@@ -20,6 +20,7 @@
 import static libcore.junit.util.compat.CoreCompatChangeRule.EnableCompatChanges;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
@@ -33,14 +34,19 @@
 import android.compat.testing.PlatformCompatChangeRule;
 import android.content.ContentResolver;
 import android.content.Context;
+import android.net.nsd.INsdManagerCallback;
+import android.net.nsd.INsdServiceConnector;
 import android.net.nsd.NsdManager;
 import android.net.nsd.NsdServiceInfo;
+import android.os.Binder;
 import android.os.Build;
 import android.os.Handler;
 import android.os.HandlerThread;
+import android.os.IBinder;
 import android.os.Looper;
 import android.os.Message;
 
+import androidx.annotation.NonNull;
 import androidx.test.filters.SmallTest;
 
 import com.android.server.NsdService.DaemonConnection;
@@ -56,11 +62,15 @@
 import org.junit.Test;
 import org.junit.rules.TestRule;
 import org.junit.runner.RunWith;
+import org.mockito.AdditionalAnswers;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 import org.mockito.Spy;
 
+import java.util.LinkedList;
+import java.util.Queue;
+
 // TODOs:
 //  - test client can send requests and receive replies
 //  - test NSD_ON ENABLE/DISABLED listening
@@ -73,16 +83,30 @@
     private static final long CLEANUP_DELAY_MS = 500;
     private static final long TIMEOUT_MS = 500;
 
+    // Records INsdManagerCallback created when NsdService#connect is called.
+    // Only accessed on the test thread, since NsdService#connect is called by the NsdManager
+    // constructor called on the test thread.
+    private final Queue<INsdManagerCallback> mCreatedCallbacks = new LinkedList<>();
+
     @Rule
     public TestRule compatChangeRule = new PlatformCompatChangeRule();
     @Mock Context mContext;
     @Mock ContentResolver mResolver;
-    @Mock NsdService.NsdSettings mSettings;
     NativeCallbackReceiver mDaemonCallback;
     @Spy DaemonConnection mDaemon = new DaemonConnection(mDaemonCallback);
     HandlerThread mThread;
     TestHandler mHandler;
 
+    private static class LinkToDeathRecorder extends Binder {
+        IBinder.DeathRecipient mDr;
+
+        @Override
+        public void linkToDeath(@NonNull DeathRecipient recipient, int flags) {
+            super.linkToDeath(recipient, flags);
+            mDr = recipient;
+        }
+    }
+
     @Before
     public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
@@ -103,26 +127,29 @@
 
     @Test
     @DisableCompatChanges(NsdManager.RUN_NATIVE_NSD_ONLY_IF_LEGACY_APPS)
-    public void testPreSClients() {
-        when(mSettings.isEnabled()).thenReturn(true);
+    public void testPreSClients() throws Exception {
         NsdService service = makeService();
 
         // Pre S client connected, the daemon should be started.
-        NsdManager client1 = connectClient(service);
+        connectClient(service);
         waitForIdle();
+        final INsdManagerCallback cb1 = getCallback();
+        final IBinder.DeathRecipient deathRecipient1 = verifyLinkToDeath(cb1);
         verify(mDaemon, times(1)).maybeStart();
         verifyDaemonCommands("start-service");
 
-        NsdManager client2 = connectClient(service);
+        connectClient(service);
         waitForIdle();
+        final INsdManagerCallback cb2 = getCallback();
+        final IBinder.DeathRecipient deathRecipient2 = verifyLinkToDeath(cb2);
         verify(mDaemon, times(1)).maybeStart();
 
-        client1.disconnect();
+        deathRecipient1.binderDied();
         // Still 1 client remains, daemon shouldn't be stopped.
         waitForIdle();
         verify(mDaemon, never()).maybeStop();
 
-        client2.disconnect();
+        deathRecipient2.binderDied();
         // All clients are disconnected, the daemon should be stopped.
         verifyDelayMaybeStopDaemon(CLEANUP_DELAY_MS);
         verifyDaemonCommands("stop-service");
@@ -130,43 +157,49 @@
 
     @Test
     @EnableCompatChanges(NsdManager.RUN_NATIVE_NSD_ONLY_IF_LEGACY_APPS)
-    public void testNoDaemonStartedWhenClientsConnect() {
-        when(mSettings.isEnabled()).thenReturn(true);
-
-        NsdService service = makeService();
+    public void testNoDaemonStartedWhenClientsConnect() throws Exception {
+        final NsdService service = makeService();
 
         // Creating an NsdManager will not cause any cmds executed, which means
         // no daemon is started.
-        NsdManager client1 = connectClient(service);
+        connectClient(service);
         waitForIdle();
         verify(mDaemon, never()).execute(any());
+        final INsdManagerCallback cb1 = getCallback();
+        final IBinder.DeathRecipient deathRecipient1 = verifyLinkToDeath(cb1);
 
         // Creating another NsdManager will not cause any cmds executed.
-        NsdManager client2 = connectClient(service);
+        connectClient(service);
         waitForIdle();
         verify(mDaemon, never()).execute(any());
+        final INsdManagerCallback cb2 = getCallback();
+        final IBinder.DeathRecipient deathRecipient2 = verifyLinkToDeath(cb2);
 
         // If there is no active request, try to clean up the daemon
         // every time the client disconnects.
-        client1.disconnect();
+        deathRecipient1.binderDied();
         verifyDelayMaybeStopDaemon(CLEANUP_DELAY_MS);
         reset(mDaemon);
-        client2.disconnect();
+        deathRecipient2.binderDied();
         verifyDelayMaybeStopDaemon(CLEANUP_DELAY_MS);
+    }
 
-        client1.disconnect();
-        client2.disconnect();
+    private IBinder.DeathRecipient verifyLinkToDeath(INsdManagerCallback cb)
+            throws Exception {
+        final IBinder.DeathRecipient dr = ((LinkToDeathRecorder) cb.asBinder()).mDr;
+        assertNotNull(dr);
+        return dr;
     }
 
     @Test
     @EnableCompatChanges(NsdManager.RUN_NATIVE_NSD_ONLY_IF_LEGACY_APPS)
-    public void testClientRequestsAreGCedAtDisconnection() {
-        when(mSettings.isEnabled()).thenReturn(true);
-
+    public void testClientRequestsAreGCedAtDisconnection() throws Exception {
         NsdService service = makeService();
-        NsdManager client = connectClient(service);
 
+        NsdManager client = connectClient(service);
         waitForIdle();
+        final INsdManagerCallback cb1 = getCallback();
+        final IBinder.DeathRecipient deathRecipient = verifyLinkToDeath(cb1);
         verify(mDaemon, never()).maybeStart();
         verify(mDaemon, never()).execute(any());
 
@@ -195,20 +228,16 @@
         verifyDaemonCommand("resolve 4 a_name a_type local.");
 
         // Client disconnects, stop the daemon after CLEANUP_DELAY_MS.
-        client.disconnect();
+        deathRecipient.binderDied();
         verifyDelayMaybeStopDaemon(CLEANUP_DELAY_MS);
         // checks that request are cleaned
         verifyDaemonCommands("stop-register 2", "stop-discover 3",
                 "stop-resolve 4", "stop-service");
-
-        client.disconnect();
     }
 
     @Test
     @EnableCompatChanges(NsdManager.RUN_NATIVE_NSD_ONLY_IF_LEGACY_APPS)
-    public void testCleanupDelayNoRequestActive() {
-        when(mSettings.isEnabled()).thenReturn(true);
-
+    public void testCleanupDelayNoRequestActive() throws Exception {
         NsdService service = makeService();
         NsdManager client = connectClient(service);
 
@@ -218,6 +247,8 @@
         client.registerService(request, PROTOCOL, listener1);
         waitForIdle();
         verify(mDaemon, times(1)).maybeStart();
+        final INsdManagerCallback cb1 = getCallback();
+        final IBinder.DeathRecipient deathRecipient = verifyLinkToDeath(cb1);
         verifyDaemonCommands("start-service", "register 2 a_name a_type 2201");
 
         client.unregisterService(listener1);
@@ -226,7 +257,7 @@
         verifyDelayMaybeStopDaemon(CLEANUP_DELAY_MS);
         verifyDaemonCommand("stop-service");
         reset(mDaemon);
-        client.disconnect();
+        deathRecipient.binderDied();
         // Client disconnects, after CLEANUP_DELAY_MS, maybeStop the daemon.
         verifyDelayMaybeStopDaemon(CLEANUP_DELAY_MS);
     }
@@ -240,12 +271,28 @@
             mDaemonCallback = callback;
             return mDaemon;
         };
-        NsdService service = new NsdService(mContext, mSettings,
-                mHandler, supplier, CLEANUP_DELAY_MS);
+        final NsdService service = new NsdService(mContext, mHandler, supplier, CLEANUP_DELAY_MS) {
+            @Override
+            public INsdServiceConnector connect(INsdManagerCallback baseCb) {
+                // Wrap the callback in a transparent mock, to mock asBinder returning a
+                // LinkToDeathRecorder. This will allow recording the binder death recipient
+                // registered on the callback. Use a transparent mock and not a spy as the actual
+                // implementation class is not public and cannot be spied on by Mockito.
+                final INsdManagerCallback cb = mock(INsdManagerCallback.class,
+                        AdditionalAnswers.delegatesTo(baseCb));
+                doReturn(new LinkToDeathRecorder()).when(cb).asBinder();
+                mCreatedCallbacks.add(cb);
+                return super.connect(cb);
+            }
+        };
         verify(mDaemon, never()).execute(any(String.class));
         return service;
     }
 
+    private INsdManagerCallback getCallback() {
+        return mCreatedCallbacks.remove();
+    }
+
     NsdManager connectClient(NsdService service) {
         return new NsdManager(mContext, service);
     }
diff --git a/tests/unit/java/com/android/server/connectivity/PermissionMonitorTest.java b/tests/unit/java/com/android/server/connectivity/PermissionMonitorTest.java
index ecda338..99ef80b 100644
--- a/tests/unit/java/com/android/server/connectivity/PermissionMonitorTest.java
+++ b/tests/unit/java/com/android/server/connectivity/PermissionMonitorTest.java
@@ -60,6 +60,7 @@
 import static org.mockito.Mockito.doCallRealMethod;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.reset;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
@@ -118,10 +119,13 @@
     private static final int SYSTEM_APPID1 = 1100;
     private static final int SYSTEM_APPID2 = 1108;
     private static final int VPN_APPID = 10002;
-    private static final int MOCK_UID1 = MOCK_USER1.getUid(MOCK_APPID1);
-    private static final int MOCK_UID2 = MOCK_USER1.getUid(MOCK_APPID2);
-    private static final int SYSTEM_APP_UID1 = MOCK_USER1.getUid(SYSTEM_APPID1);
+    private static final int MOCK_UID11 = MOCK_USER1.getUid(MOCK_APPID1);
+    private static final int MOCK_UID12 = MOCK_USER1.getUid(MOCK_APPID2);
+    private static final int SYSTEM_APP_UID11 = MOCK_USER1.getUid(SYSTEM_APPID1);
     private static final int VPN_UID = MOCK_USER1.getUid(VPN_APPID);
+    private static final int MOCK_UID21 = MOCK_USER2.getUid(MOCK_APPID1);
+    private static final int MOCK_UID22 = MOCK_USER2.getUid(MOCK_APPID2);
+    private static final int SYSTEM_APP_UID21 = MOCK_USER2.getUid(SYSTEM_APPID1);
     private static final String REAL_SYSTEM_PACKAGE_NAME = "android";
     private static final String MOCK_PACKAGE1 = "appName1";
     private static final String MOCK_PACKAGE2 = "appName2";
@@ -161,9 +165,13 @@
             doCallRealMethod().when(mContext).getSystemService(SystemConfigManager.class);
         }
         when(mSystemConfigManager.getSystemPermissionUids(anyString())).thenReturn(new int[0]);
-        final Context asUserCtx = mock(Context.class, AdditionalAnswers.delegatesTo(mContext));
-        doReturn(UserHandle.ALL).when(asUserCtx).getUser();
-        when(mContext.createContextAsUser(eq(UserHandle.ALL), anyInt())).thenReturn(asUserCtx);
+        doAnswer(invocation -> {
+            final Object[] args = invocation.getArguments();
+            final Context asUserCtx = mock(Context.class, AdditionalAnswers.delegatesTo(mContext));
+            final UserHandle user = (UserHandle) args[0];
+            doReturn(user).when(asUserCtx).getUser();
+            return asUserCtx;
+        }).when(mContext).createContextAsUser(any(), anyInt());
         when(mDeps.getUidsAllowedOnRestrictedNetworks(any())).thenReturn(Set.of());
         // Set DEVICE_INITIAL_SDK_INT to Q that SYSTEM_UID won't have restricted network permission
         // by default.
@@ -172,7 +180,7 @@
         mPermissionMonitor = new PermissionMonitor(mContext, mNetdService, mDeps);
         mNetdMonitor = new NetdMonitor(mNetdService);
 
-        doReturn(List.of()).when(mPackageManager).getInstalledPackages(anyInt());
+        doReturn(List.of()).when(mPackageManager).getInstalledPackagesAsUser(anyInt(), anyInt());
     }
 
     private boolean hasRestrictedNetworkPermission(String partition, int targetSdkVersion,
@@ -330,26 +338,26 @@
     @Test
     public void testHasRestrictedNetworkPermission() {
         assertFalse(hasRestrictedNetworkPermission(
-                PARTITION_SYSTEM, VERSION_P, MOCK_PACKAGE1, MOCK_UID1));
+                PARTITION_SYSTEM, VERSION_P, MOCK_PACKAGE1, MOCK_UID11));
         assertFalse(hasRestrictedNetworkPermission(
-                PARTITION_SYSTEM, VERSION_P, MOCK_PACKAGE1, MOCK_UID1, CHANGE_NETWORK_STATE));
+                PARTITION_SYSTEM, VERSION_P, MOCK_PACKAGE1, MOCK_UID11, CHANGE_NETWORK_STATE));
         assertTrue(hasRestrictedNetworkPermission(
-                PARTITION_SYSTEM, VERSION_P, MOCK_PACKAGE1, MOCK_UID1, NETWORK_STACK));
+                PARTITION_SYSTEM, VERSION_P, MOCK_PACKAGE1, MOCK_UID11, NETWORK_STACK));
         assertFalse(hasRestrictedNetworkPermission(
-                PARTITION_SYSTEM, VERSION_P, MOCK_PACKAGE1, MOCK_UID1, CONNECTIVITY_INTERNAL));
+                PARTITION_SYSTEM, VERSION_P, MOCK_PACKAGE1, MOCK_UID11, CONNECTIVITY_INTERNAL));
         assertTrue(hasRestrictedNetworkPermission(
-                PARTITION_SYSTEM, VERSION_P, MOCK_PACKAGE1, MOCK_UID1,
+                PARTITION_SYSTEM, VERSION_P, MOCK_PACKAGE1, MOCK_UID11,
                 CONNECTIVITY_USE_RESTRICTED_NETWORKS));
         assertFalse(hasRestrictedNetworkPermission(
-                PARTITION_SYSTEM, VERSION_P, MOCK_PACKAGE1, MOCK_UID1, CHANGE_WIFI_STATE));
+                PARTITION_SYSTEM, VERSION_P, MOCK_PACKAGE1, MOCK_UID11, CHANGE_WIFI_STATE));
         assertTrue(hasRestrictedNetworkPermission(
-                PARTITION_SYSTEM, VERSION_P, MOCK_PACKAGE1, MOCK_UID1,
+                PARTITION_SYSTEM, VERSION_P, MOCK_PACKAGE1, MOCK_UID11,
                 PERMISSION_MAINLINE_NETWORK_STACK));
 
         assertFalse(hasRestrictedNetworkPermission(
-                PARTITION_SYSTEM, VERSION_Q, MOCK_PACKAGE1, MOCK_UID1));
+                PARTITION_SYSTEM, VERSION_Q, MOCK_PACKAGE1, MOCK_UID11));
         assertFalse(hasRestrictedNetworkPermission(
-                PARTITION_SYSTEM, VERSION_Q, MOCK_PACKAGE1, MOCK_UID1, CONNECTIVITY_INTERNAL));
+                PARTITION_SYSTEM, VERSION_Q, MOCK_PACKAGE1, MOCK_UID11, CONNECTIVITY_INTERNAL));
     }
 
     @Test
@@ -376,43 +384,43 @@
     @Test
     public void testHasRestrictedNetworkPermissionVendorApp() {
         assertTrue(hasRestrictedNetworkPermission(
-                PARTITION_VENDOR, VERSION_P, MOCK_PACKAGE1, MOCK_UID1));
+                PARTITION_VENDOR, VERSION_P, MOCK_PACKAGE1, MOCK_UID11));
         assertTrue(hasRestrictedNetworkPermission(
-                PARTITION_VENDOR, VERSION_P, MOCK_PACKAGE1, MOCK_UID1, CHANGE_NETWORK_STATE));
+                PARTITION_VENDOR, VERSION_P, MOCK_PACKAGE1, MOCK_UID11, CHANGE_NETWORK_STATE));
         assertTrue(hasRestrictedNetworkPermission(
-                PARTITION_VENDOR, VERSION_P, MOCK_PACKAGE1, MOCK_UID1, NETWORK_STACK));
+                PARTITION_VENDOR, VERSION_P, MOCK_PACKAGE1, MOCK_UID11, NETWORK_STACK));
         assertTrue(hasRestrictedNetworkPermission(
-                PARTITION_VENDOR, VERSION_P, MOCK_PACKAGE1, MOCK_UID1, CONNECTIVITY_INTERNAL));
+                PARTITION_VENDOR, VERSION_P, MOCK_PACKAGE1, MOCK_UID11, CONNECTIVITY_INTERNAL));
         assertTrue(hasRestrictedNetworkPermission(
-                PARTITION_VENDOR, VERSION_P, MOCK_PACKAGE1, MOCK_UID1,
+                PARTITION_VENDOR, VERSION_P, MOCK_PACKAGE1, MOCK_UID11,
                 CONNECTIVITY_USE_RESTRICTED_NETWORKS));
         assertTrue(hasRestrictedNetworkPermission(
-                PARTITION_VENDOR, VERSION_P, MOCK_PACKAGE1, MOCK_UID1, CHANGE_WIFI_STATE));
+                PARTITION_VENDOR, VERSION_P, MOCK_PACKAGE1, MOCK_UID11, CHANGE_WIFI_STATE));
 
         assertFalse(hasRestrictedNetworkPermission(
-                PARTITION_VENDOR, VERSION_Q, MOCK_PACKAGE1, MOCK_UID1));
+                PARTITION_VENDOR, VERSION_Q, MOCK_PACKAGE1, MOCK_UID11));
         assertFalse(hasRestrictedNetworkPermission(
-                PARTITION_VENDOR, VERSION_Q, MOCK_PACKAGE1, MOCK_UID1, CONNECTIVITY_INTERNAL));
+                PARTITION_VENDOR, VERSION_Q, MOCK_PACKAGE1, MOCK_UID11, CONNECTIVITY_INTERNAL));
         assertFalse(hasRestrictedNetworkPermission(
-                PARTITION_VENDOR, VERSION_Q, MOCK_PACKAGE1, MOCK_UID1, CHANGE_NETWORK_STATE));
+                PARTITION_VENDOR, VERSION_Q, MOCK_PACKAGE1, MOCK_UID11, CHANGE_NETWORK_STATE));
     }
 
     @Test
     public void testHasRestrictedNetworkPermissionUidAllowedOnRestrictedNetworks() {
-        mPermissionMonitor.updateUidsAllowedOnRestrictedNetworks(Set.of(MOCK_UID1));
+        mPermissionMonitor.updateUidsAllowedOnRestrictedNetworks(Set.of(MOCK_UID11));
         assertTrue(hasRestrictedNetworkPermission(
-                PARTITION_VENDOR, VERSION_Q, MOCK_PACKAGE1, MOCK_UID1));
+                PARTITION_VENDOR, VERSION_Q, MOCK_PACKAGE1, MOCK_UID11));
         assertTrue(hasRestrictedNetworkPermission(
-                PARTITION_VENDOR, VERSION_Q, MOCK_PACKAGE1, MOCK_UID1, CHANGE_NETWORK_STATE));
+                PARTITION_VENDOR, VERSION_Q, MOCK_PACKAGE1, MOCK_UID11, CHANGE_NETWORK_STATE));
         assertTrue(hasRestrictedNetworkPermission(
-                PARTITION_VENDOR, VERSION_Q, MOCK_PACKAGE1, MOCK_UID1, CONNECTIVITY_INTERNAL));
+                PARTITION_VENDOR, VERSION_Q, MOCK_PACKAGE1, MOCK_UID11, CONNECTIVITY_INTERNAL));
 
         assertFalse(hasRestrictedNetworkPermission(
-                PARTITION_VENDOR, VERSION_Q, MOCK_PACKAGE2, MOCK_UID2));
+                PARTITION_VENDOR, VERSION_Q, MOCK_PACKAGE2, MOCK_UID12));
         assertFalse(hasRestrictedNetworkPermission(
-                PARTITION_VENDOR, VERSION_Q, MOCK_PACKAGE2, MOCK_UID2, CHANGE_NETWORK_STATE));
+                PARTITION_VENDOR, VERSION_Q, MOCK_PACKAGE2, MOCK_UID12, CHANGE_NETWORK_STATE));
         assertFalse(hasRestrictedNetworkPermission(
-                PARTITION_VENDOR, VERSION_Q, MOCK_PACKAGE2, MOCK_UID2, CONNECTIVITY_INTERNAL));
+                PARTITION_VENDOR, VERSION_Q, MOCK_PACKAGE2, MOCK_UID12, CONNECTIVITY_INTERNAL));
 
     }
 
@@ -429,27 +437,27 @@
         doReturn(VERSION_P).when(mDeps).getDeviceFirstSdkInt();
         assertTrue(wouldBeCarryoverPackage(PARTITION_SYSTEM, VERSION_P, SYSTEM_UID));
         assertTrue(wouldBeCarryoverPackage(PARTITION_VENDOR, VERSION_P, SYSTEM_UID));
-        assertFalse(wouldBeCarryoverPackage(PARTITION_SYSTEM, VERSION_P, MOCK_UID1));
-        assertTrue(wouldBeCarryoverPackage(PARTITION_VENDOR, VERSION_P, MOCK_UID1));
+        assertFalse(wouldBeCarryoverPackage(PARTITION_SYSTEM, VERSION_P, MOCK_UID11));
+        assertTrue(wouldBeCarryoverPackage(PARTITION_VENDOR, VERSION_P, MOCK_UID11));
         assertTrue(wouldBeCarryoverPackage(PARTITION_SYSTEM, VERSION_Q, SYSTEM_UID));
         assertTrue(wouldBeCarryoverPackage(PARTITION_VENDOR, VERSION_Q, SYSTEM_UID));
-        assertFalse(wouldBeCarryoverPackage(PARTITION_SYSTEM, VERSION_Q, MOCK_UID1));
-        assertFalse(wouldBeCarryoverPackage(PARTITION_VENDOR, VERSION_Q, MOCK_UID1));
+        assertFalse(wouldBeCarryoverPackage(PARTITION_SYSTEM, VERSION_Q, MOCK_UID11));
+        assertFalse(wouldBeCarryoverPackage(PARTITION_VENDOR, VERSION_Q, MOCK_UID11));
 
         doReturn(VERSION_Q).when(mDeps).getDeviceFirstSdkInt();
         assertFalse(wouldBeCarryoverPackage(PARTITION_SYSTEM, VERSION_P, SYSTEM_UID));
         assertTrue(wouldBeCarryoverPackage(PARTITION_VENDOR, VERSION_P, SYSTEM_UID));
-        assertFalse(wouldBeCarryoverPackage(PARTITION_SYSTEM, VERSION_P, MOCK_UID1));
-        assertTrue(wouldBeCarryoverPackage(PARTITION_VENDOR, VERSION_P, MOCK_UID1));
+        assertFalse(wouldBeCarryoverPackage(PARTITION_SYSTEM, VERSION_P, MOCK_UID11));
+        assertTrue(wouldBeCarryoverPackage(PARTITION_VENDOR, VERSION_P, MOCK_UID11));
         assertFalse(wouldBeCarryoverPackage(PARTITION_SYSTEM, VERSION_Q, SYSTEM_UID));
         assertFalse(wouldBeCarryoverPackage(PARTITION_VENDOR, VERSION_Q, SYSTEM_UID));
-        assertFalse(wouldBeCarryoverPackage(PARTITION_SYSTEM, VERSION_Q, MOCK_UID1));
-        assertFalse(wouldBeCarryoverPackage(PARTITION_VENDOR, VERSION_Q, MOCK_UID1));
+        assertFalse(wouldBeCarryoverPackage(PARTITION_SYSTEM, VERSION_Q, MOCK_UID11));
+        assertFalse(wouldBeCarryoverPackage(PARTITION_VENDOR, VERSION_Q, MOCK_UID11));
 
         assertFalse(wouldBeCarryoverPackage(PARTITION_OEM, VERSION_Q, SYSTEM_UID));
         assertFalse(wouldBeCarryoverPackage(PARTITION_PRODUCT, VERSION_Q, SYSTEM_UID));
-        assertFalse(wouldBeCarryoverPackage(PARTITION_OEM, VERSION_Q, MOCK_UID1));
-        assertFalse(wouldBeCarryoverPackage(PARTITION_PRODUCT, VERSION_Q, MOCK_UID1));
+        assertFalse(wouldBeCarryoverPackage(PARTITION_OEM, VERSION_Q, MOCK_UID11));
+        assertFalse(wouldBeCarryoverPackage(PARTITION_PRODUCT, VERSION_Q, MOCK_UID11));
     }
 
     private boolean wouldBeUidAllowedOnRestrictedNetworks(int uid) {
@@ -461,20 +469,20 @@
     @Test
     public void testIsAppAllowedOnRestrictedNetworks() {
         mPermissionMonitor.updateUidsAllowedOnRestrictedNetworks(Set.of());
-        assertFalse(wouldBeUidAllowedOnRestrictedNetworks(MOCK_UID1));
-        assertFalse(wouldBeUidAllowedOnRestrictedNetworks(MOCK_UID2));
+        assertFalse(wouldBeUidAllowedOnRestrictedNetworks(MOCK_UID11));
+        assertFalse(wouldBeUidAllowedOnRestrictedNetworks(MOCK_UID12));
 
-        mPermissionMonitor.updateUidsAllowedOnRestrictedNetworks(Set.of(MOCK_UID1));
-        assertTrue(wouldBeUidAllowedOnRestrictedNetworks(MOCK_UID1));
-        assertFalse(wouldBeUidAllowedOnRestrictedNetworks(MOCK_UID2));
+        mPermissionMonitor.updateUidsAllowedOnRestrictedNetworks(Set.of(MOCK_UID11));
+        assertTrue(wouldBeUidAllowedOnRestrictedNetworks(MOCK_UID11));
+        assertFalse(wouldBeUidAllowedOnRestrictedNetworks(MOCK_UID12));
 
-        mPermissionMonitor.updateUidsAllowedOnRestrictedNetworks(Set.of(MOCK_UID2));
-        assertFalse(wouldBeUidAllowedOnRestrictedNetworks(MOCK_UID1));
-        assertTrue(wouldBeUidAllowedOnRestrictedNetworks(MOCK_UID2));
+        mPermissionMonitor.updateUidsAllowedOnRestrictedNetworks(Set.of(MOCK_UID12));
+        assertFalse(wouldBeUidAllowedOnRestrictedNetworks(MOCK_UID11));
+        assertTrue(wouldBeUidAllowedOnRestrictedNetworks(MOCK_UID12));
 
         mPermissionMonitor.updateUidsAllowedOnRestrictedNetworks(Set.of(123));
-        assertFalse(wouldBeUidAllowedOnRestrictedNetworks(MOCK_UID1));
-        assertFalse(wouldBeUidAllowedOnRestrictedNetworks(MOCK_UID2));
+        assertFalse(wouldBeUidAllowedOnRestrictedNetworks(MOCK_UID11));
+        assertFalse(wouldBeUidAllowedOnRestrictedNetworks(MOCK_UID12));
     }
 
     private void assertBackgroundPermission(boolean hasPermission, String name, int uid,
@@ -491,16 +499,16 @@
         assertBackgroundPermission(true, SYSTEM_PACKAGE1, SYSTEM_UID, CHANGE_NETWORK_STATE);
         assertBackgroundPermission(true, SYSTEM_PACKAGE1, SYSTEM_UID, NETWORK_STACK);
 
-        assertFalse(mPermissionMonitor.hasUseBackgroundNetworksPermission(MOCK_UID1));
-        assertBackgroundPermission(false, MOCK_PACKAGE1, MOCK_UID1);
-        assertBackgroundPermission(true, MOCK_PACKAGE1, MOCK_UID1,
+        assertFalse(mPermissionMonitor.hasUseBackgroundNetworksPermission(MOCK_UID11));
+        assertBackgroundPermission(false, MOCK_PACKAGE1, MOCK_UID11);
+        assertBackgroundPermission(true, MOCK_PACKAGE1, MOCK_UID11,
                 CONNECTIVITY_USE_RESTRICTED_NETWORKS);
 
-        assertFalse(mPermissionMonitor.hasUseBackgroundNetworksPermission(MOCK_UID2));
-        assertBackgroundPermission(false, MOCK_PACKAGE2, MOCK_UID2);
-        assertBackgroundPermission(false, MOCK_PACKAGE2, MOCK_UID2,
+        assertFalse(mPermissionMonitor.hasUseBackgroundNetworksPermission(MOCK_UID12));
+        assertBackgroundPermission(false, MOCK_PACKAGE2, MOCK_UID12);
+        assertBackgroundPermission(false, MOCK_PACKAGE2, MOCK_UID12,
                 CONNECTIVITY_INTERNAL);
-        assertBackgroundPermission(true, MOCK_PACKAGE2, MOCK_UID2, NETWORK_STACK);
+        assertBackgroundPermission(true, MOCK_PACKAGE2, MOCK_UID12, NETWORK_STACK);
     }
 
     private class NetdMonitor {
@@ -589,13 +597,13 @@
 
     @Test
     public void testUserAndPackageAddRemove() throws Exception {
-        // MOCK_UID1: MOCK_PACKAGE1 only has network permission.
-        // SYSTEM_APP_UID1: SYSTEM_PACKAGE1 has system permission.
-        // SYSTEM_APP_UID1: SYSTEM_PACKAGE2 only has network permission.
-        buildAndMockPackageInfoWithPermissions(MOCK_PACKAGE1, MOCK_UID1, CHANGE_NETWORK_STATE);
-        buildAndMockPackageInfoWithPermissions(SYSTEM_PACKAGE1, SYSTEM_APP_UID1,
+        // MOCK_UID11: MOCK_PACKAGE1 only has network permission.
+        // SYSTEM_APP_UID11: SYSTEM_PACKAGE1 has system permission.
+        // SYSTEM_APP_UID11: SYSTEM_PACKAGE2 only has network permission.
+        buildAndMockPackageInfoWithPermissions(MOCK_PACKAGE1, MOCK_UID11, CHANGE_NETWORK_STATE);
+        buildAndMockPackageInfoWithPermissions(SYSTEM_PACKAGE1, SYSTEM_APP_UID11,
                 CONNECTIVITY_USE_RESTRICTED_NETWORKS);
-        buildAndMockPackageInfoWithPermissions(SYSTEM_PACKAGE2, SYSTEM_APP_UID1,
+        buildAndMockPackageInfoWithPermissions(SYSTEM_PACKAGE2, SYSTEM_APP_UID11,
                 CHANGE_NETWORK_STATE);
 
         // Add user MOCK_USER1.
@@ -610,16 +618,19 @@
         mNetdMonitor.expectNetworkPerm(PERMISSION_SYSTEM, new UserHandle[]{MOCK_USER1},
                 SYSTEM_APPID1);
 
+        final List<PackageInfo> pkgs = List.of(buildPackageInfo(SYSTEM_PACKAGE1, SYSTEM_APP_UID21,
+                        CONNECTIVITY_USE_RESTRICTED_NETWORKS),
+                buildPackageInfo(SYSTEM_PACKAGE2, SYSTEM_APP_UID21, CHANGE_NETWORK_STATE));
+        doReturn(pkgs).when(mPackageManager).getInstalledPackagesAsUser(eq(GET_PERMISSIONS),
+                eq(MOCK_USER_ID2));
         // Add user MOCK_USER2.
         mPermissionMonitor.onUserAdded(MOCK_USER2);
         mNetdMonitor.expectNetworkPerm(PERMISSION_SYSTEM, new UserHandle[]{MOCK_USER1, MOCK_USER2},
                 SYSTEM_APPID1);
 
         // Remove SYSTEM_PACKAGE2, expect keep system permission.
-        when(mPackageManager.getPackagesForUid(SYSTEM_APP_UID1))
-                .thenReturn(new String[]{SYSTEM_PACKAGE1});
-        when(mPackageManager.getPackagesForUid(MOCK_USER2.getUid(SYSTEM_APPID1)))
-                .thenReturn(new String[]{SYSTEM_PACKAGE1});
+        doReturn(new String[]{SYSTEM_PACKAGE1}).when(mPackageManager)
+                .getPackagesForUid(intThat(uid -> UserHandle.getAppId(uid) == SYSTEM_APPID1));
         removePackageForUsers(new UserHandle[]{MOCK_USER1, MOCK_USER2},
                 SYSTEM_PACKAGE2, SYSTEM_APPID1);
         mNetdMonitor.expectNetworkPerm(PERMISSION_SYSTEM, new UserHandle[]{MOCK_USER1, MOCK_USER2},
@@ -631,6 +642,8 @@
         mNetdMonitor.expectNetworkPerm(PERMISSION_SYSTEM, new UserHandle[]{MOCK_USER1, MOCK_USER2},
                 SYSTEM_APPID1);
 
+        // Add MOCK_PACKAGE1
+        buildAndMockPackageInfoWithPermissions(MOCK_PACKAGE1, MOCK_UID21, CHANGE_NETWORK_STATE);
         addPackageForUsers(new UserHandle[]{MOCK_USER1, MOCK_USER2}, MOCK_PACKAGE1, MOCK_APPID1);
         mNetdMonitor.expectNetworkPerm(PERMISSION_SYSTEM, new UserHandle[]{MOCK_USER1, MOCK_USER2},
                 SYSTEM_APPID1);
@@ -638,9 +651,8 @@
                 MOCK_APPID1);
 
         // Remove MOCK_PACKAGE1, expect no permission left for all user.
-        when(mPackageManager.getPackagesForUid(MOCK_UID1)).thenReturn(new String[]{});
-        when(mPackageManager.getPackagesForUid(MOCK_USER2.getUid(MOCK_APPID1)))
-                .thenReturn(new String[]{});
+        doReturn(new String[]{}).when(mPackageManager)
+                .getPackagesForUid(intThat(uid -> UserHandle.getAppId(uid) == MOCK_APPID1));
         removePackageForUsers(new UserHandle[]{MOCK_USER1, MOCK_USER2}, MOCK_PACKAGE1, MOCK_APPID1);
         mNetdMonitor.expectNoNetworkPerm(new UserHandle[]{MOCK_USER1, MOCK_USER2}, MOCK_APPID1);
 
@@ -674,68 +686,71 @@
 
     @Test
     public void testUidFilteringDuringVpnConnectDisconnectAndUidUpdates() throws Exception {
-        when(mPackageManager.getInstalledPackages(eq(GET_PERMISSIONS | MATCH_ANY_USER))).thenReturn(
-                List.of(buildPackageInfo(SYSTEM_PACKAGE1, SYSTEM_APP_UID1, CHANGE_NETWORK_STATE,
-                                CONNECTIVITY_USE_RESTRICTED_NETWORKS),
-                        buildPackageInfo(MOCK_PACKAGE1, MOCK_UID1),
-                        buildPackageInfo(MOCK_PACKAGE2, MOCK_UID2),
-                        buildPackageInfo(SYSTEM_PACKAGE2, VPN_UID)));
-        buildAndMockPackageInfoWithPermissions(MOCK_PACKAGE1, MOCK_UID1);
+        doReturn(List.of(buildPackageInfo(SYSTEM_PACKAGE1, SYSTEM_APP_UID11, CHANGE_NETWORK_STATE,
+                        CONNECTIVITY_USE_RESTRICTED_NETWORKS),
+                buildPackageInfo(MOCK_PACKAGE1, MOCK_UID11),
+                buildPackageInfo(MOCK_PACKAGE2, MOCK_UID12),
+                buildPackageInfo(SYSTEM_PACKAGE2, VPN_UID)))
+                .when(mPackageManager).getInstalledPackagesAsUser(eq(GET_PERMISSIONS), anyInt());
+        buildAndMockPackageInfoWithPermissions(MOCK_PACKAGE1, MOCK_UID11);
         mPermissionMonitor.startMonitoring();
-        // Every app on user 0 except MOCK_UID2 are under VPN.
+        // Every app on user 0 except MOCK_UID12 are under VPN.
         final Set<UidRange> vpnRange1 = Set.of(
-                new UidRange(0, MOCK_UID2 - 1),
-                new UidRange(MOCK_UID2 + 1, UserHandle.PER_USER_RANGE - 1));
-        final Set<UidRange> vpnRange2 = Set.of(new UidRange(MOCK_UID2, MOCK_UID2));
+                new UidRange(0, MOCK_UID12 - 1),
+                new UidRange(MOCK_UID12 + 1, UserHandle.PER_USER_RANGE - 1));
+        final Set<UidRange> vpnRange2 = Set.of(new UidRange(MOCK_UID12, MOCK_UID12));
 
-        // When VPN is connected, expect a rule to be set up for user app MOCK_UID1
+        // When VPN is connected, expect a rule to be set up for user app MOCK_UID11
         mPermissionMonitor.onVpnUidRangesAdded("tun0", vpnRange1, VPN_UID);
-        verify(mNetdService).firewallAddUidInterfaceRules(eq("tun0"), aryEq(new int[]{MOCK_UID1}));
+        verify(mNetdService).firewallAddUidInterfaceRules(eq("tun0"), aryEq(new int[]{MOCK_UID11}));
 
         reset(mNetdService);
 
-        // When MOCK_UID1 package is uninstalled and reinstalled, expect Netd to be updated
-        mPermissionMonitor.onPackageRemoved(MOCK_PACKAGE1, MOCK_UID1);
-        verify(mNetdService).firewallRemoveUidInterfaceRules(aryEq(new int[]{MOCK_UID1}));
-        mPermissionMonitor.onPackageAdded(MOCK_PACKAGE1, MOCK_UID1);
-        verify(mNetdService).firewallAddUidInterfaceRules(eq("tun0"), aryEq(new int[]{MOCK_UID1}));
+        // When MOCK_UID11 package is uninstalled and reinstalled, expect Netd to be updated
+        mPermissionMonitor.onPackageRemoved(MOCK_PACKAGE1, MOCK_UID11);
+        verify(mNetdService).firewallRemoveUidInterfaceRules(aryEq(new int[]{MOCK_UID11}));
+        mPermissionMonitor.onPackageAdded(MOCK_PACKAGE1, MOCK_UID11);
+        verify(mNetdService).firewallAddUidInterfaceRules(eq("tun0"), aryEq(new int[]{MOCK_UID11}));
 
         reset(mNetdService);
 
         // During VPN uid update (vpnRange1 -> vpnRange2), ConnectivityService first deletes the
         // old UID rules then adds the new ones. Expect netd to be updated
         mPermissionMonitor.onVpnUidRangesRemoved("tun0", vpnRange1, VPN_UID);
-        verify(mNetdService).firewallRemoveUidInterfaceRules(aryEq(new int[] {MOCK_UID1}));
+        verify(mNetdService).firewallRemoveUidInterfaceRules(aryEq(new int[] {MOCK_UID11}));
         mPermissionMonitor.onVpnUidRangesAdded("tun0", vpnRange2, VPN_UID);
-        verify(mNetdService).firewallAddUidInterfaceRules(eq("tun0"), aryEq(new int[]{MOCK_UID2}));
+        verify(mNetdService).firewallAddUidInterfaceRules(eq("tun0"), aryEq(new int[]{MOCK_UID12}));
 
         reset(mNetdService);
 
         // When VPN is disconnected, expect rules to be torn down
         mPermissionMonitor.onVpnUidRangesRemoved("tun0", vpnRange2, VPN_UID);
-        verify(mNetdService).firewallRemoveUidInterfaceRules(aryEq(new int[] {MOCK_UID2}));
+        verify(mNetdService).firewallRemoveUidInterfaceRules(aryEq(new int[] {MOCK_UID12}));
         assertNull(mPermissionMonitor.getVpnUidRanges("tun0"));
     }
 
     @Test
     public void testUidFilteringDuringPackageInstallAndUninstall() throws Exception {
-        when(mPackageManager.getInstalledPackages(eq(GET_PERMISSIONS | MATCH_ANY_USER))).thenReturn(
-                List.of(buildPackageInfo(SYSTEM_PACKAGE1, SYSTEM_APP_UID1, CHANGE_NETWORK_STATE,
-                                NETWORK_STACK, CONNECTIVITY_USE_RESTRICTED_NETWORKS),
-                        buildPackageInfo(SYSTEM_PACKAGE2, VPN_UID)));
-        buildAndMockPackageInfoWithPermissions(MOCK_PACKAGE1, MOCK_UID1);
+        doReturn(List.of(buildPackageInfo(SYSTEM_PACKAGE1, SYSTEM_APP_UID11, CHANGE_NETWORK_STATE,
+                        NETWORK_STACK, CONNECTIVITY_USE_RESTRICTED_NETWORKS),
+                buildPackageInfo(SYSTEM_PACKAGE2, VPN_UID)))
+                .when(mPackageManager).getInstalledPackagesAsUser(eq(GET_PERMISSIONS), anyInt());
+        buildAndMockPackageInfoWithPermissions(MOCK_PACKAGE1, MOCK_UID11);
 
         mPermissionMonitor.startMonitoring();
-        final Set<UidRange> vpnRange = Set.of(UidRange.createForUser(MOCK_USER1));
+        final Set<UidRange> vpnRange = Set.of(UidRange.createForUser(MOCK_USER1),
+                UidRange.createForUser(MOCK_USER2));
         mPermissionMonitor.onVpnUidRangesAdded("tun0", vpnRange, VPN_UID);
 
         // Newly-installed package should have uid rules added
-        mPermissionMonitor.onPackageAdded(MOCK_PACKAGE1, MOCK_UID1);
-        verify(mNetdService).firewallAddUidInterfaceRules(eq("tun0"), aryEq(new int[]{MOCK_UID1}));
+        addPackageForUsers(new UserHandle[]{MOCK_USER1, MOCK_USER2}, MOCK_PACKAGE1, MOCK_APPID1);
+        verify(mNetdService).firewallAddUidInterfaceRules(eq("tun0"), aryEq(new int[]{MOCK_UID11}));
+        verify(mNetdService).firewallAddUidInterfaceRules(eq("tun0"), aryEq(new int[]{MOCK_UID21}));
 
         // Removed package should have its uid rules removed
-        mPermissionMonitor.onPackageRemoved(MOCK_PACKAGE1, MOCK_UID1);
-        verify(mNetdService).firewallRemoveUidInterfaceRules(aryEq(new int[] {MOCK_UID1}));
+        mPermissionMonitor.onPackageRemoved(MOCK_PACKAGE1, MOCK_UID11);
+        verify(mNetdService).firewallRemoveUidInterfaceRules(aryEq(new int[]{MOCK_UID11}));
+        verify(mNetdService, never()).firewallRemoveUidInterfaceRules(aryEq(new int[]{MOCK_UID21}));
     }
 
 
@@ -767,7 +782,7 @@
         netdPermissionsAppIds.put(SYSTEM_APPID2, PERMISSION_UPDATE_DEVICE_STATS);
 
         // Send the permission information to netd, expect permission updated.
-        mPermissionMonitor.sendPackagePermissionsToNetd(netdPermissionsAppIds);
+        mPermissionMonitor.sendAppIdsTrafficPermission(netdPermissionsAppIds);
 
         mNetdMonitor.expectTrafficPerm(PERMISSION_INTERNET, MOCK_APPID1);
         mNetdMonitor.expectTrafficPerm(PERMISSION_NONE, MOCK_APPID2);
@@ -790,69 +805,69 @@
 
     @Test
     public void testPackageInstall() throws Exception {
-        addPackage(MOCK_PACKAGE1, MOCK_UID1, INTERNET, UPDATE_DEVICE_STATS);
+        addPackage(MOCK_PACKAGE1, MOCK_UID11, INTERNET, UPDATE_DEVICE_STATS);
         mNetdMonitor.expectTrafficPerm(PERMISSION_TRAFFIC_ALL, MOCK_APPID1);
 
-        addPackage(MOCK_PACKAGE2, MOCK_UID2, INTERNET);
+        addPackage(MOCK_PACKAGE2, MOCK_UID12, INTERNET);
         mNetdMonitor.expectTrafficPerm(PERMISSION_INTERNET, MOCK_APPID2);
     }
 
     @Test
     public void testPackageInstallSharedUid() throws Exception {
-        addPackage(MOCK_PACKAGE1, MOCK_UID1, INTERNET, UPDATE_DEVICE_STATS);
+        addPackage(MOCK_PACKAGE1, MOCK_UID11, INTERNET, UPDATE_DEVICE_STATS);
         mNetdMonitor.expectTrafficPerm(PERMISSION_TRAFFIC_ALL, MOCK_APPID1);
 
         // Install another package with the same uid and no permissions should not cause the app id
         // to lose permissions.
-        addPackage(MOCK_PACKAGE2, MOCK_UID1);
+        addPackage(MOCK_PACKAGE2, MOCK_UID11);
         mNetdMonitor.expectTrafficPerm(PERMISSION_TRAFFIC_ALL, MOCK_APPID1);
     }
 
     @Test
     public void testPackageUninstallBasic() throws Exception {
-        addPackage(MOCK_PACKAGE1, MOCK_UID1, INTERNET, UPDATE_DEVICE_STATS);
+        addPackage(MOCK_PACKAGE1, MOCK_UID11, INTERNET, UPDATE_DEVICE_STATS);
         mNetdMonitor.expectTrafficPerm(PERMISSION_TRAFFIC_ALL, MOCK_APPID1);
 
-        when(mPackageManager.getPackagesForUid(MOCK_UID1)).thenReturn(new String[]{});
-        mPermissionMonitor.onPackageRemoved(MOCK_PACKAGE1, MOCK_UID1);
+        when(mPackageManager.getPackagesForUid(MOCK_UID11)).thenReturn(new String[]{});
+        mPermissionMonitor.onPackageRemoved(MOCK_PACKAGE1, MOCK_UID11);
         mNetdMonitor.expectTrafficPerm(PERMISSION_UNINSTALLED, MOCK_APPID1);
     }
 
     @Test
     public void testPackageRemoveThenAdd() throws Exception {
-        addPackage(MOCK_PACKAGE1, MOCK_UID1, INTERNET, UPDATE_DEVICE_STATS);
+        addPackage(MOCK_PACKAGE1, MOCK_UID11, INTERNET, UPDATE_DEVICE_STATS);
         mNetdMonitor.expectTrafficPerm(PERMISSION_TRAFFIC_ALL, MOCK_APPID1);
 
-        when(mPackageManager.getPackagesForUid(MOCK_UID1)).thenReturn(new String[]{});
-        mPermissionMonitor.onPackageRemoved(MOCK_PACKAGE1, MOCK_UID1);
+        when(mPackageManager.getPackagesForUid(MOCK_UID11)).thenReturn(new String[]{});
+        mPermissionMonitor.onPackageRemoved(MOCK_PACKAGE1, MOCK_UID11);
         mNetdMonitor.expectTrafficPerm(PERMISSION_UNINSTALLED, MOCK_APPID1);
 
-        addPackage(MOCK_PACKAGE1, MOCK_UID1, INTERNET);
+        addPackage(MOCK_PACKAGE1, MOCK_UID11, INTERNET);
         mNetdMonitor.expectTrafficPerm(PERMISSION_INTERNET, MOCK_APPID1);
     }
 
     @Test
     public void testPackageUpdate() throws Exception {
-        addPackage(MOCK_PACKAGE1, MOCK_UID1);
+        addPackage(MOCK_PACKAGE1, MOCK_UID11);
         mNetdMonitor.expectTrafficPerm(PERMISSION_NONE, MOCK_APPID1);
 
-        addPackage(MOCK_PACKAGE1, MOCK_UID1, INTERNET);
+        addPackage(MOCK_PACKAGE1, MOCK_UID11, INTERNET);
         mNetdMonitor.expectTrafficPerm(PERMISSION_INTERNET, MOCK_APPID1);
     }
 
     @Test
     public void testPackageUninstallWithMultiplePackages() throws Exception {
-        addPackage(MOCK_PACKAGE1, MOCK_UID1, INTERNET, UPDATE_DEVICE_STATS);
+        addPackage(MOCK_PACKAGE1, MOCK_UID11, INTERNET, UPDATE_DEVICE_STATS);
         mNetdMonitor.expectTrafficPerm(PERMISSION_TRAFFIC_ALL, MOCK_APPID1);
 
         // Install another package with the same uid but different permissions.
-        addPackage(MOCK_PACKAGE2, MOCK_UID1, INTERNET);
-        mNetdMonitor.expectTrafficPerm(PERMISSION_TRAFFIC_ALL, MOCK_UID1);
+        addPackage(MOCK_PACKAGE2, MOCK_UID11, INTERNET);
+        mNetdMonitor.expectTrafficPerm(PERMISSION_TRAFFIC_ALL, MOCK_UID11);
 
         // Uninstall MOCK_PACKAGE1 and expect only INTERNET permission left.
-        when(mPackageManager.getPackagesForUid(eq(MOCK_UID1)))
+        when(mPackageManager.getPackagesForUid(eq(MOCK_UID11)))
                 .thenReturn(new String[]{MOCK_PACKAGE2});
-        mPermissionMonitor.onPackageRemoved(MOCK_PACKAGE1, MOCK_UID1);
+        mPermissionMonitor.onPackageRemoved(MOCK_PACKAGE1, MOCK_UID11);
         mNetdMonitor.expectTrafficPerm(PERMISSION_INTERNET, MOCK_APPID1);
     }
 
@@ -871,9 +886,9 @@
     @Test
     public void testUpdateUidPermissionsFromSystemConfig() throws Exception {
         when(mSystemConfigManager.getSystemPermissionUids(eq(INTERNET)))
-                .thenReturn(new int[]{ MOCK_UID1, MOCK_UID2 });
+                .thenReturn(new int[]{ MOCK_UID11, MOCK_UID12 });
         when(mSystemConfigManager.getSystemPermissionUids(eq(UPDATE_DEVICE_STATS)))
-                .thenReturn(new int[]{ MOCK_UID2 });
+                .thenReturn(new int[]{ MOCK_UID12 });
 
         mPermissionMonitor.startMonitoring();
         mNetdMonitor.expectTrafficPerm(PERMISSION_INTERNET, MOCK_APPID1);
@@ -904,17 +919,17 @@
         // Verify receiving PACKAGE_ADDED intent.
         final Intent addedIntent = new Intent(Intent.ACTION_PACKAGE_ADDED,
                 Uri.fromParts("package", MOCK_PACKAGE1, null /* fragment */));
-        addedIntent.putExtra(Intent.EXTRA_UID, MOCK_UID1);
-        buildAndMockPackageInfoWithPermissions(MOCK_PACKAGE1, MOCK_UID1, INTERNET,
+        addedIntent.putExtra(Intent.EXTRA_UID, MOCK_UID11);
+        buildAndMockPackageInfoWithPermissions(MOCK_PACKAGE1, MOCK_UID11, INTERNET,
                 UPDATE_DEVICE_STATS);
         receiver.onReceive(mContext, addedIntent);
         mNetdMonitor.expectTrafficPerm(PERMISSION_TRAFFIC_ALL, MOCK_APPID1);
 
         // Verify receiving PACKAGE_REMOVED intent.
-        when(mPackageManager.getPackagesForUid(MOCK_UID1)).thenReturn(new String[]{});
+        when(mPackageManager.getPackagesForUid(MOCK_UID11)).thenReturn(new String[]{});
         final Intent removedIntent = new Intent(Intent.ACTION_PACKAGE_REMOVED,
                 Uri.fromParts("package", MOCK_PACKAGE1, null /* fragment */));
-        removedIntent.putExtra(Intent.EXTRA_UID, MOCK_UID1);
+        removedIntent.putExtra(Intent.EXTRA_UID, MOCK_UID11);
         receiver.onReceive(mContext, removedIntent);
         mNetdMonitor.expectTrafficPerm(PERMISSION_UNINSTALLED, MOCK_APPID1);
     }
@@ -934,20 +949,20 @@
                 Settings.Global.getUriFor(UIDS_ALLOWED_ON_RESTRICTED_NETWORKS));
 
         // Prepare PackageInfo for MOCK_PACKAGE1 and MOCK_PACKAGE2
-        buildAndMockPackageInfoWithPermissions(MOCK_PACKAGE1, MOCK_UID1);
-        buildAndMockPackageInfoWithPermissions(MOCK_PACKAGE2, MOCK_UID2);
+        buildAndMockPackageInfoWithPermissions(MOCK_PACKAGE1, MOCK_UID11);
+        buildAndMockPackageInfoWithPermissions(MOCK_PACKAGE2, MOCK_UID12);
 
-        // MOCK_UID1 is listed in setting that allow to use restricted networks, MOCK_UID1
+        // MOCK_UID11 is listed in setting that allow to use restricted networks, MOCK_UID11
         // should have SYSTEM permission.
-        when(mDeps.getUidsAllowedOnRestrictedNetworks(any())).thenReturn(Set.of(MOCK_UID1));
+        when(mDeps.getUidsAllowedOnRestrictedNetworks(any())).thenReturn(Set.of(MOCK_UID11));
         contentObserver.onChange(true /* selfChange */);
         mNetdMonitor.expectNetworkPerm(PERMISSION_SYSTEM, new UserHandle[]{MOCK_USER1},
                 MOCK_APPID1);
         mNetdMonitor.expectNoNetworkPerm(new UserHandle[]{MOCK_USER1}, MOCK_APPID2);
 
-        // MOCK_UID2 is listed in setting that allow to use restricted networks, MOCK_UID2
-        // should have SYSTEM permission but MOCK_UID1 should revoke permission.
-        when(mDeps.getUidsAllowedOnRestrictedNetworks(any())).thenReturn(Set.of(MOCK_UID2));
+        // MOCK_UID12 is listed in setting that allow to use restricted networks, MOCK_UID12
+        // should have SYSTEM permission but MOCK_UID11 should revoke permission.
+        when(mDeps.getUidsAllowedOnRestrictedNetworks(any())).thenReturn(Set.of(MOCK_UID12));
         contentObserver.onChange(true /* selfChange */);
         mNetdMonitor.expectNetworkPerm(PERMISSION_SYSTEM, new UserHandle[]{MOCK_USER1},
                 MOCK_APPID2);
@@ -965,29 +980,29 @@
         final ContentObserver contentObserver = expectRegisterContentObserver(
                 Settings.Global.getUriFor(UIDS_ALLOWED_ON_RESTRICTED_NETWORKS));
 
-        buildAndMockPackageInfoWithPermissions(MOCK_PACKAGE1, MOCK_UID1, CHANGE_NETWORK_STATE);
-        buildAndMockPackageInfoWithPermissions(MOCK_PACKAGE2, MOCK_UID1);
+        buildAndMockPackageInfoWithPermissions(MOCK_PACKAGE1, MOCK_UID11, CHANGE_NETWORK_STATE);
+        buildAndMockPackageInfoWithPermissions(MOCK_PACKAGE2, MOCK_UID11);
 
-        // MOCK_PACKAGE1 have CHANGE_NETWORK_STATE, MOCK_UID1 should have NETWORK permission.
+        // MOCK_PACKAGE1 have CHANGE_NETWORK_STATE, MOCK_UID11 should have NETWORK permission.
         addPackageForUsers(new UserHandle[]{MOCK_USER1}, MOCK_PACKAGE1, MOCK_APPID1);
         mNetdMonitor.expectNetworkPerm(PERMISSION_NETWORK, new UserHandle[]{MOCK_USER1},
                 MOCK_APPID1);
 
-        // MOCK_UID1 is listed in setting that allow to use restricted networks, MOCK_UID1
+        // MOCK_UID11 is listed in setting that allow to use restricted networks, MOCK_UID11
         // should upgrade to SYSTEM permission.
-        when(mDeps.getUidsAllowedOnRestrictedNetworks(any())).thenReturn(Set.of(MOCK_UID1));
+        when(mDeps.getUidsAllowedOnRestrictedNetworks(any())).thenReturn(Set.of(MOCK_UID11));
         contentObserver.onChange(true /* selfChange */);
         mNetdMonitor.expectNetworkPerm(PERMISSION_SYSTEM, new UserHandle[]{MOCK_USER1},
                 MOCK_APPID1);
 
-        // No app lists in setting, MOCK_UID1 should downgrade to NETWORK permission.
+        // No app lists in setting, MOCK_UID11 should downgrade to NETWORK permission.
         when(mDeps.getUidsAllowedOnRestrictedNetworks(any())).thenReturn(Set.of());
         contentObserver.onChange(true /* selfChange */);
         mNetdMonitor.expectNetworkPerm(PERMISSION_NETWORK, new UserHandle[]{MOCK_USER1},
                 MOCK_APPID1);
 
-        // MOCK_PACKAGE1 removed, should revoke permission from MOCK_UID1.
-        when(mPackageManager.getPackagesForUid(MOCK_UID1)).thenReturn(new String[]{MOCK_PACKAGE2});
+        // MOCK_PACKAGE1 removed, should revoke permission from MOCK_UID11.
+        when(mPackageManager.getPackagesForUid(MOCK_UID11)).thenReturn(new String[]{MOCK_PACKAGE2});
         removePackageForUsers(new UserHandle[]{MOCK_USER1}, MOCK_PACKAGE1, MOCK_APPID1);
         mNetdMonitor.expectNoNetworkPerm(new UserHandle[]{MOCK_USER1}, MOCK_APPID1);
     }
@@ -998,33 +1013,39 @@
         final ContentObserver contentObserver = expectRegisterContentObserver(
                 Settings.Global.getUriFor(UIDS_ALLOWED_ON_RESTRICTED_NETWORKS));
 
-        // Prepare PackageInfo for MOCK_PACKAGE1 and MOCK_PACKAGE2.
-        buildAndMockPackageInfoWithPermissions(MOCK_PACKAGE1, MOCK_UID1);
-        buildAndMockPackageInfoWithPermissions(MOCK_PACKAGE2, MOCK_UID2);
+        // Prepare PackageInfo for MOCK_APPID1 and MOCK_APPID2 in MOCK_USER1.
+        buildAndMockPackageInfoWithPermissions(MOCK_PACKAGE1, MOCK_UID11);
+        buildAndMockPackageInfoWithPermissions(MOCK_PACKAGE2, MOCK_UID12);
 
-        // MOCK_UID1 is listed in setting that allow to use restricted networks, MOCK_UID1
-        // should have SYSTEM permission and MOCK_UID2 has no permissions.
-        when(mDeps.getUidsAllowedOnRestrictedNetworks(any())).thenReturn(Set.of(MOCK_UID1));
+        // MOCK_UID11 is listed in setting that allow to use restricted networks, MOCK_UID11 should
+        // have SYSTEM permission and MOCK_UID12 has no permissions.
+        when(mDeps.getUidsAllowedOnRestrictedNetworks(any())).thenReturn(Set.of(MOCK_UID11));
         contentObserver.onChange(true /* selfChange */);
         mNetdMonitor.expectNetworkPerm(PERMISSION_SYSTEM, new UserHandle[]{MOCK_USER1},
                 MOCK_APPID1);
         mNetdMonitor.expectNoNetworkPerm(new UserHandle[]{MOCK_USER1}, MOCK_APPID2);
 
         // Add user MOCK_USER2.
+        final List<PackageInfo> pkgs = List.of(buildPackageInfo(MOCK_PACKAGE1, MOCK_UID21));
+        buildAndMockPackageInfoWithPermissions(MOCK_PACKAGE2, MOCK_UID22);
+        doReturn(pkgs).when(mPackageManager)
+                .getInstalledPackagesAsUser(eq(GET_PERMISSIONS), eq(MOCK_USER_ID2));
         mPermissionMonitor.onUserAdded(MOCK_USER2);
-        // MOCK_APPID1 in both users should all have SYSTEM permission and MOCK_APPID2 has no
-        // permissions in either user.
-        mNetdMonitor.expectNetworkPerm(PERMISSION_SYSTEM, new UserHandle[]{MOCK_USER1, MOCK_USER2},
+        // MOCK_APPID1 in MOCK_USER1 should have SYSTEM permission but in MOCK_USER2 should have no
+        // permissions. And MOCK_APPID2 has no permissions in either users.
+        mNetdMonitor.expectNetworkPerm(PERMISSION_SYSTEM, new UserHandle[]{MOCK_USER1},
                 MOCK_APPID1);
+        mNetdMonitor.expectNoNetworkPerm(new UserHandle[]{MOCK_USER2}, MOCK_APPID1);
         mNetdMonitor.expectNoNetworkPerm(new UserHandle[]{MOCK_USER1, MOCK_USER2}, MOCK_APPID2);
 
-        // MOCK_UID2 is listed in setting that allow to use restricted networks, MOCK_APPID2
-        // in both users should have SYSTEM permission and MOCK_APPID1 has no permissions in either
-        // user.
-        when(mDeps.getUidsAllowedOnRestrictedNetworks(any())).thenReturn(Set.of(MOCK_UID2));
+        // MOCK_UID22 is listed in setting that allow to use restricted networks,
+        // MOCK_APPID2 in MOCK_USER2 should have SYSTEM permission but in MOCK_USER1 should have no
+        // permissions. And MOCK_APPID1 has no permissions in either users.
+        doReturn(Set.of(MOCK_UID22)).when(mDeps).getUidsAllowedOnRestrictedNetworks(any());
         contentObserver.onChange(true /* selfChange */);
-        mNetdMonitor.expectNetworkPerm(PERMISSION_SYSTEM, new UserHandle[]{MOCK_USER1, MOCK_USER2},
+        mNetdMonitor.expectNetworkPerm(PERMISSION_SYSTEM, new UserHandle[]{MOCK_USER2},
                 MOCK_APPID2);
+        mNetdMonitor.expectNoNetworkPerm(new UserHandle[]{MOCK_USER1}, MOCK_APPID2);
         mNetdMonitor.expectNoNetworkPerm(new UserHandle[]{MOCK_USER1, MOCK_USER2}, MOCK_APPID1);
 
         // Remove user MOCK_USER1
@@ -1044,9 +1065,9 @@
     public void testOnExternalApplicationsAvailable() throws Exception {
         // Initial the permission state. MOCK_PACKAGE1 and MOCK_PACKAGE2 are installed on external
         // and have different uids. There has no permission for both uids.
-        when(mPackageManager.getInstalledPackages(eq(GET_PERMISSIONS | MATCH_ANY_USER))).thenReturn(
-                List.of(buildPackageInfo(MOCK_PACKAGE1, MOCK_UID1),
-                        buildPackageInfo(MOCK_PACKAGE2, MOCK_UID2)));
+        doReturn(List.of(buildPackageInfo(MOCK_PACKAGE1, MOCK_UID11),
+                buildPackageInfo(MOCK_PACKAGE2, MOCK_UID12)))
+                .when(mPackageManager).getInstalledPackagesAsUser(eq(GET_PERMISSIONS), anyInt());
         mPermissionMonitor.startMonitoring();
         mNetdMonitor.expectNoNetworkPerm(new UserHandle[]{MOCK_USER1}, MOCK_APPID1, MOCK_APPID2);
         mNetdMonitor.expectTrafficPerm(PERMISSION_NONE, MOCK_APPID1, MOCK_APPID2);
@@ -1057,9 +1078,9 @@
         final Intent externalIntent = new Intent(Intent.ACTION_EXTERNAL_APPLICATIONS_AVAILABLE);
         externalIntent.putExtra(Intent.EXTRA_CHANGED_PACKAGE_LIST,
                 new String[] { MOCK_PACKAGE1 , MOCK_PACKAGE2});
-        buildAndMockPackageInfoWithPermissions(MOCK_PACKAGE1, MOCK_UID1,
+        buildAndMockPackageInfoWithPermissions(MOCK_PACKAGE1, MOCK_UID11,
                 CONNECTIVITY_USE_RESTRICTED_NETWORKS, INTERNET);
-        buildAndMockPackageInfoWithPermissions(MOCK_PACKAGE2, MOCK_UID2, CHANGE_NETWORK_STATE,
+        buildAndMockPackageInfoWithPermissions(MOCK_PACKAGE2, MOCK_UID12, CHANGE_NETWORK_STATE,
                 UPDATE_DEVICE_STATS);
         receiver.onReceive(mContext, externalIntent);
         mNetdMonitor.expectNetworkPerm(PERMISSION_SYSTEM, new UserHandle[]{MOCK_USER1},
@@ -1079,9 +1100,9 @@
 
         // Initial the permission state. MOCK_PACKAGE1 and MOCK_PACKAGE2 are installed on external
         // and have different uids. There has no permission for both uids.
-        buildAndMockPackageInfoWithPermissions(MOCK_PACKAGE1, MOCK_UID1,
+        buildAndMockPackageInfoWithPermissions(MOCK_PACKAGE1, MOCK_UID11,
                 CONNECTIVITY_USE_RESTRICTED_NETWORKS, INTERNET);
-        buildAndMockPackageInfoWithPermissions(MOCK_PACKAGE2, MOCK_UID2, CHANGE_NETWORK_STATE,
+        buildAndMockPackageInfoWithPermissions(MOCK_PACKAGE2, MOCK_UID12, CHANGE_NETWORK_STATE,
                 UPDATE_DEVICE_STATS);
 
         // Verify receiving EXTERNAL_APPLICATIONS_AVAILABLE intent and update permission to netd.
@@ -1101,10 +1122,10 @@
     public void testOnExternalApplicationsAvailableWithSharedUid()
             throws Exception {
         // Initial the permission state. MOCK_PACKAGE1 and MOCK_PACKAGE2 are installed on external
-        // storage and shared on MOCK_UID1. There has no permission for MOCK_UID1.
-        when(mPackageManager.getInstalledPackages(eq(GET_PERMISSIONS | MATCH_ANY_USER))).thenReturn(
-                List.of(buildPackageInfo(MOCK_PACKAGE1, MOCK_UID1),
-                        buildPackageInfo(MOCK_PACKAGE2, MOCK_UID1)));
+        // storage and shared on MOCK_UID11. There has no permission for MOCK_UID11.
+        doReturn(List.of(buildPackageInfo(MOCK_PACKAGE1, MOCK_UID11),
+                buildPackageInfo(MOCK_PACKAGE2, MOCK_UID11)))
+                .when(mPackageManager).getInstalledPackagesAsUser(eq(GET_PERMISSIONS), anyInt());
         mPermissionMonitor.startMonitoring();
         mNetdMonitor.expectNoNetworkPerm(new UserHandle[]{MOCK_USER1}, MOCK_APPID1);
         mNetdMonitor.expectTrafficPerm(PERMISSION_NONE, MOCK_APPID1);
@@ -1114,8 +1135,8 @@
         // Verify receiving EXTERNAL_APPLICATIONS_AVAILABLE intent and update permission to netd.
         final Intent externalIntent = new Intent(Intent.ACTION_EXTERNAL_APPLICATIONS_AVAILABLE);
         externalIntent.putExtra(Intent.EXTRA_CHANGED_PACKAGE_LIST, new String[] {MOCK_PACKAGE1});
-        buildAndMockPackageInfoWithPermissions(MOCK_PACKAGE1, MOCK_UID1, CHANGE_NETWORK_STATE);
-        buildAndMockPackageInfoWithPermissions(MOCK_PACKAGE2, MOCK_UID1, UPDATE_DEVICE_STATS);
+        buildAndMockPackageInfoWithPermissions(MOCK_PACKAGE1, MOCK_UID11, CHANGE_NETWORK_STATE);
+        buildAndMockPackageInfoWithPermissions(MOCK_PACKAGE2, MOCK_UID11, UPDATE_DEVICE_STATS);
         receiver.onReceive(mContext, externalIntent);
         mNetdMonitor.expectNetworkPerm(PERMISSION_NETWORK, new UserHandle[]{MOCK_USER1},
                 MOCK_APPID1);
@@ -1126,12 +1147,11 @@
     public void testOnExternalApplicationsAvailableWithSharedUid_DifferentStorage()
             throws Exception {
         // Initial the permission state. MOCK_PACKAGE1 is installed on external storage and
-        // MOCK_PACKAGE2 is installed on device. These two packages are shared on MOCK_UID1.
-        // MOCK_UID1 has NETWORK and INTERNET permissions.
-        when(mPackageManager.getInstalledPackages(eq(GET_PERMISSIONS | MATCH_ANY_USER))).thenReturn(
-                List.of(buildPackageInfo(MOCK_PACKAGE1, MOCK_UID1),
-                        buildPackageInfo(MOCK_PACKAGE2, MOCK_UID1, CHANGE_NETWORK_STATE,
-                                INTERNET)));
+        // MOCK_PACKAGE2 is installed on device. These two packages are shared on MOCK_UID11.
+        // MOCK_UID11 has NETWORK and INTERNET permissions.
+        doReturn(List.of(buildPackageInfo(MOCK_PACKAGE1, MOCK_UID11),
+                buildPackageInfo(MOCK_PACKAGE2, MOCK_UID11, CHANGE_NETWORK_STATE, INTERNET)))
+                .when(mPackageManager).getInstalledPackagesAsUser(eq(GET_PERMISSIONS), anyInt());
         mPermissionMonitor.startMonitoring();
         mNetdMonitor.expectNetworkPerm(PERMISSION_NETWORK, new UserHandle[]{MOCK_USER1},
                 MOCK_APPID1);
@@ -1142,9 +1162,9 @@
         // Verify receiving EXTERNAL_APPLICATIONS_AVAILABLE intent and update permission to netd.
         final Intent externalIntent = new Intent(Intent.ACTION_EXTERNAL_APPLICATIONS_AVAILABLE);
         externalIntent.putExtra(Intent.EXTRA_CHANGED_PACKAGE_LIST, new String[] {MOCK_PACKAGE1});
-        buildAndMockPackageInfoWithPermissions(MOCK_PACKAGE1, MOCK_UID1,
+        buildAndMockPackageInfoWithPermissions(MOCK_PACKAGE1, MOCK_UID11,
                 CONNECTIVITY_USE_RESTRICTED_NETWORKS, UPDATE_DEVICE_STATS);
-        buildAndMockPackageInfoWithPermissions(MOCK_PACKAGE2, MOCK_UID1, CHANGE_NETWORK_STATE,
+        buildAndMockPackageInfoWithPermissions(MOCK_PACKAGE2, MOCK_UID11, CHANGE_NETWORK_STATE,
                 INTERNET);
         receiver.onReceive(mContext, externalIntent);
         mNetdMonitor.expectNetworkPerm(PERMISSION_SYSTEM, new UserHandle[]{MOCK_USER1},
diff --git a/tests/unit/jni/Android.bp b/tests/unit/jni/Android.bp
index 1c1ba9e..fe971e7 100644
--- a/tests/unit/jni/Android.bp
+++ b/tests/unit/jni/Android.bp
@@ -13,6 +13,8 @@
         "-Wthread-safety",
     ],
 
+    header_libs: ["bpf_connectivity_headers"],
+
     srcs: [
         ":lib_networkStatsFactory_native",
         "test_onload.cpp",
diff --git a/tests/unit/vpn-jarjar-rules.txt b/tests/unit/vpn-jarjar-rules.txt
new file mode 100644
index 0000000..16661b9
--- /dev/null
+++ b/tests/unit/vpn-jarjar-rules.txt
@@ -0,0 +1,4 @@
+# Only keep classes imported by ConnectivityServiceTest
+keep com.android.server.VpnManagerService
+keep com.android.server.connectivity.Vpn
+keep com.android.server.connectivity.VpnProfileStore
\ No newline at end of file