diff --git a/core/java/android/net/ConnectivityDiagnosticsManager.java b/core/java/android/net/ConnectivityDiagnosticsManager.java
index 704f31d..5234494 100644
--- a/core/java/android/net/ConnectivityDiagnosticsManager.java
+++ b/core/java/android/net/ConnectivityDiagnosticsManager.java
@@ -623,32 +623,41 @@
         /** @hide */
         @VisibleForTesting
         public void onConnectivityReportAvailable(@NonNull ConnectivityReport report) {
-            Binder.withCleanCallingIdentity(() -> {
+            final long token = Binder.clearCallingIdentity();
+            try {
                 mExecutor.execute(() -> {
                     mCb.onConnectivityReportAvailable(report);
                 });
-            });
+            } finally {
+                Binder.restoreCallingIdentity(token);
+            }
         }
 
         /** @hide */
         @VisibleForTesting
         public void onDataStallSuspected(@NonNull DataStallReport report) {
-            Binder.withCleanCallingIdentity(() -> {
+            final long token = Binder.clearCallingIdentity();
+            try {
                 mExecutor.execute(() -> {
                     mCb.onDataStallSuspected(report);
                 });
-            });
+            } finally {
+                Binder.restoreCallingIdentity(token);
+            }
         }
 
         /** @hide */
         @VisibleForTesting
         public void onNetworkConnectivityReported(
                 @NonNull Network network, boolean hasConnectivity) {
-            Binder.withCleanCallingIdentity(() -> {
+            final long token = Binder.clearCallingIdentity();
+            try {
                 mExecutor.execute(() -> {
                     mCb.onNetworkConnectivityReported(network, hasConnectivity);
                 });
-            });
+            } finally {
+                Binder.restoreCallingIdentity(token);
+            }
         }
     }
 
diff --git a/core/java/android/net/ConnectivityManager.java b/core/java/android/net/ConnectivityManager.java
index 540ea5c..06c1598 100644
--- a/core/java/android/net/ConnectivityManager.java
+++ b/core/java/android/net/ConnectivityManager.java
@@ -59,8 +59,10 @@
 import android.telephony.TelephonyManager;
 import android.util.ArrayMap;
 import android.util.Log;
+import android.util.Range;
 import android.util.SparseIntArray;
 
+import com.android.connectivity.aidl.INetworkAgent;
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.util.Preconditions;
 import com.android.internal.util.Protocol;
@@ -72,10 +74,12 @@
 import java.io.UncheckedIOException;
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
+import java.net.DatagramSocket;
 import java.net.InetAddress;
 import java.net.InetSocketAddress;
 import java.net.Socket;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -1162,6 +1166,55 @@
     }
 
     /**
+     * Adds or removes a requirement for given UID ranges to use the VPN.
+     *
+     * If set to {@code true}, informs the system that the UIDs in the specified ranges must not
+     * have any connectivity except if a VPN is connected and applies to the UIDs, or if the UIDs
+     * otherwise have permission to bypass the VPN (e.g., because they have the
+     * {@link android.Manifest.permission.CONNECTIVITY_USE_RESTRICTED_NETWORKS} permission, or when
+     * using a socket protected by a method such as {@link VpnService#protect(DatagramSocket)}. If
+     * set to {@code false}, a previously-added restriction is removed.
+     * <p>
+     * Each of the UID ranges specified by this method is added and removed as is, and no processing
+     * is performed on the ranges to de-duplicate, merge, split, or intersect them. In order to
+     * remove a previously-added range, the exact range must be removed as is.
+     * <p>
+     * The changes are applied asynchronously and may not have been applied by the time the method
+     * returns. Apps will be notified about any changes that apply to them via
+     * {@link NetworkCallback#onBlockedStatusChanged} callbacks called after the changes take
+     * effect.
+     * <p>
+     * This method should be called only by the VPN code.
+     *
+     * @param ranges the UID ranges to restrict
+     * @param requireVpn whether the specified UID ranges must use a VPN
+     *
+     * TODO: expose as @SystemApi.
+     * @hide
+     */
+    @RequiresPermission(anyOf = {
+            NetworkStack.PERMISSION_MAINLINE_NETWORK_STACK,
+            android.Manifest.permission.NETWORK_STACK})
+    public void setRequireVpnForUids(boolean requireVpn,
+            @NonNull Collection<Range<Integer>> ranges) {
+        Objects.requireNonNull(ranges);
+        // The Range class is not parcelable. Convert to UidRange, which is what is used internally.
+        // This method is not necessarily expected to be used outside the system server, so
+        // parceling may not be necessary, but it could be used out-of-process, e.g., by the network
+        // stack process, or by tests.
+        UidRange[] rangesArray = new UidRange[ranges.size()];
+        int index = 0;
+        for (Range<Integer> range : ranges) {
+            rangesArray[index++] = new UidRange(range.getLower(), range.getUpper());
+        }
+        try {
+            mService.setRequireVpnForUids(requireVpn, rangesArray);
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
+    /**
      * Returns details about the currently active default data network
      * for a given uid.  This is for internal use only to avoid spying
      * other apps.
@@ -1832,30 +1885,42 @@
             mCallback = new ISocketKeepaliveCallback.Stub() {
                 @Override
                 public void onStarted(int slot) {
-                    Binder.withCleanCallingIdentity(() ->
-                            mExecutor.execute(() -> {
-                                mSlot = slot;
-                                callback.onStarted();
-                            }));
+                    final long token = Binder.clearCallingIdentity();
+                    try {
+                        mExecutor.execute(() -> {
+                            mSlot = slot;
+                            callback.onStarted();
+                        });
+                    } finally {
+                        Binder.restoreCallingIdentity(token);
+                    }
                 }
 
                 @Override
                 public void onStopped() {
-                    Binder.withCleanCallingIdentity(() ->
-                            mExecutor.execute(() -> {
-                                mSlot = null;
-                                callback.onStopped();
-                            }));
+                    final long token = Binder.clearCallingIdentity();
+                    try {
+                        mExecutor.execute(() -> {
+                            mSlot = null;
+                            callback.onStopped();
+                        });
+                    } finally {
+                        Binder.restoreCallingIdentity(token);
+                    }
                     mExecutor.shutdown();
                 }
 
                 @Override
                 public void onError(int error) {
-                    Binder.withCleanCallingIdentity(() ->
-                            mExecutor.execute(() -> {
-                                mSlot = null;
-                                callback.onError(error);
-                            }));
+                    final long token = Binder.clearCallingIdentity();
+                    try {
+                        mExecutor.execute(() -> {
+                            mSlot = null;
+                            callback.onError(error);
+                        });
+                    } finally {
+                        Binder.restoreCallingIdentity(token);
+                    }
                     mExecutor.shutdown();
                 }
 
@@ -3119,39 +3184,6 @@
     }
 
     /**
-     * Check mobile provisioning.
-     *
-     * @param suggestedTimeOutMs, timeout in milliseconds
-     *
-     * @return time out that will be used, maybe less that suggestedTimeOutMs
-     * -1 if an error.
-     *
-     * {@hide}
-     */
-    public int checkMobileProvisioning(int suggestedTimeOutMs) {
-        int timeOutMs = -1;
-        try {
-            timeOutMs = mService.checkMobileProvisioning(suggestedTimeOutMs);
-        } catch (RemoteException e) {
-            throw e.rethrowFromSystemServer();
-        }
-        return timeOutMs;
-    }
-
-    /**
-     * Get the mobile provisioning url.
-     * {@hide}
-     */
-    @RequiresPermission(android.Manifest.permission.NETWORK_SETTINGS)
-    public String getMobileProvisioningUrl() {
-        try {
-            return mService.getMobileProvisioningUrl();
-        } catch (RemoteException e) {
-            throw e.rethrowFromSystemServer();
-        }
-    }
-
-    /**
      * Set sign in error notification to visible or invisible
      *
      * @hide
@@ -3287,9 +3319,9 @@
     @RequiresPermission(anyOf = {
             NetworkStack.PERMISSION_MAINLINE_NETWORK_STACK,
             android.Manifest.permission.NETWORK_FACTORY})
-    public Network registerNetworkAgent(Messenger messenger, NetworkInfo ni, LinkProperties lp,
+    public Network registerNetworkAgent(INetworkAgent na, NetworkInfo ni, LinkProperties lp,
             NetworkCapabilities nc, int score, NetworkAgentConfig config) {
-        return registerNetworkAgent(messenger, ni, lp, nc, score, config, NetworkProvider.ID_NONE);
+        return registerNetworkAgent(na, ni, lp, nc, score, config, NetworkProvider.ID_NONE);
     }
 
     /**
@@ -3300,10 +3332,10 @@
     @RequiresPermission(anyOf = {
             NetworkStack.PERMISSION_MAINLINE_NETWORK_STACK,
             android.Manifest.permission.NETWORK_FACTORY})
-    public Network registerNetworkAgent(Messenger messenger, NetworkInfo ni, LinkProperties lp,
+    public Network registerNetworkAgent(INetworkAgent na, NetworkInfo ni, LinkProperties lp,
             NetworkCapabilities nc, int score, NetworkAgentConfig config, int providerId) {
         try {
-            return mService.registerNetworkAgent(messenger, ni, lp, nc, score, config, providerId);
+            return mService.registerNetworkAgent(na, ni, lp, nc, score, config, providerId);
         } catch (RemoteException e) {
             throw e.rethrowFromSystemServer();
         }
diff --git a/core/java/android/net/IConnectivityManager.aidl b/core/java/android/net/IConnectivityManager.aidl
index 95a2f2e..b32c98b 100644
--- a/core/java/android/net/IConnectivityManager.aidl
+++ b/core/java/android/net/IConnectivityManager.aidl
@@ -29,6 +29,7 @@
 import android.net.NetworkState;
 import android.net.ISocketKeepaliveCallback;
 import android.net.ProxyInfo;
+import android.net.UidRange;
 import android.os.Bundle;
 import android.os.IBinder;
 import android.os.INetworkActivityListener;
@@ -37,6 +38,7 @@
 import android.os.PersistableBundle;
 import android.os.ResultReceiver;
 
+import com.android.connectivity.aidl.INetworkAgent;
 import com.android.internal.net.LegacyVpnInfo;
 import com.android.internal.net.VpnConfig;
 import com.android.internal.net.VpnInfo;
@@ -145,10 +147,7 @@
     String getAlwaysOnVpnPackage(int userId);
     boolean isVpnLockdownEnabled(int userId);
     List<String> getVpnLockdownWhitelist(int userId);
-
-    int checkMobileProvisioning(int suggestedTimeOutMs);
-
-    String getMobileProvisioningUrl();
+    void setRequireVpnForUids(boolean requireVpn, in UidRange[] ranges);
 
     void setProvisioningNotificationVisible(boolean visible, int networkType, in String action);
 
@@ -164,7 +163,7 @@
 
     void declareNetworkRequestUnfulfillable(in NetworkRequest request);
 
-    Network registerNetworkAgent(in Messenger messenger, in NetworkInfo ni, in LinkProperties lp,
+    Network registerNetworkAgent(in INetworkAgent na, in NetworkInfo ni, in LinkProperties lp,
             in NetworkCapabilities nc, int score, in NetworkAgentConfig config,
             in int factorySerialNumber);
 
diff --git a/core/java/android/net/NetworkAgent.java b/core/java/android/net/NetworkAgent.java
index 6780167..4166c2c 100644
--- a/core/java/android/net/NetworkAgent.java
+++ b/core/java/android/net/NetworkAgent.java
@@ -29,11 +29,12 @@
 import android.os.Handler;
 import android.os.Looper;
 import android.os.Message;
-import android.os.Messenger;
+import android.os.RemoteException;
 import android.util.Log;
 
+import com.android.connectivity.aidl.INetworkAgent;
+import com.android.connectivity.aidl.INetworkAgentRegistry;
 import com.android.internal.annotations.VisibleForTesting;
-import com.android.internal.util.AsyncChannel;
 import com.android.internal.util.Protocol;
 
 import java.lang.annotation.Retention;
@@ -94,12 +95,18 @@
     @Nullable
     private volatile Network mNetwork;
 
+    @Nullable
+    private volatile INetworkAgentRegistry mRegistry;
+
+    private interface RegistryAction {
+        void execute(@NonNull INetworkAgentRegistry registry) throws RemoteException;
+    }
+
     private final Handler mHandler;
-    private volatile AsyncChannel mAsyncChannel;
     private final String LOG_TAG;
     private static final boolean DBG = true;
     private static final boolean VDBG = false;
-    private final ArrayList<Message> mPreConnectedQueue = new ArrayList<Message>();
+    private final ArrayList<RegistryAction> mPreConnectedQueue = new ArrayList<>();
     private volatile long mLastBwRefreshTime = 0;
     private static final long BW_REFRESH_MIN_WIN_MS = 500;
     private boolean mBandwidthUpdateScheduled = false;
@@ -329,6 +336,17 @@
      */
     public static final int CMD_REMOVE_KEEPALIVE_PACKET_FILTER = BASE + 17;
 
+    /**
+     * Sent by ConnectivityService to the NetworkAgent to complete the bidirectional connection.
+     * obj = INetworkAgentRegistry
+     */
+    private static final int EVENT_AGENT_CONNECTED = BASE + 18;
+
+    /**
+     * Sent by ConnectivityService to the NetworkAgent to inform the agent that it was disconnected.
+     */
+    private static final int EVENT_AGENT_DISCONNECTED = BASE + 19;
+
     private static NetworkInfo getLegacyNetworkInfo(final NetworkAgentConfig config) {
         // The subtype can be changed with (TODO) setLegacySubtype, but it starts
         // with 0 (TelephonyManager.NETWORK_TYPE_UNKNOWN) and an empty description.
@@ -402,36 +420,33 @@
         @Override
         public void handleMessage(Message msg) {
             switch (msg.what) {
-                case AsyncChannel.CMD_CHANNEL_FULL_CONNECTION: {
-                    if (mAsyncChannel != null) {
+                case EVENT_AGENT_CONNECTED: {
+                    if (mRegistry != null) {
                         log("Received new connection while already connected!");
                     } else {
                         if (VDBG) log("NetworkAgent fully connected");
-                        AsyncChannel ac = new AsyncChannel();
-                        ac.connected(null, this, msg.replyTo);
-                        ac.replyToMessage(msg, AsyncChannel.CMD_CHANNEL_FULLY_CONNECTED,
-                                AsyncChannel.STATUS_SUCCESSFUL);
                         synchronized (mPreConnectedQueue) {
-                            mAsyncChannel = ac;
-                            for (Message m : mPreConnectedQueue) {
-                                ac.sendMessage(m);
+                            final INetworkAgentRegistry registry = (INetworkAgentRegistry) msg.obj;
+                            mRegistry = registry;
+                            for (RegistryAction a : mPreConnectedQueue) {
+                                try {
+                                    a.execute(registry);
+                                } catch (RemoteException e) {
+                                    Log.wtf(LOG_TAG, "Communication error with registry", e);
+                                    // Fall through
+                                }
                             }
                             mPreConnectedQueue.clear();
                         }
                     }
                     break;
                 }
-                case AsyncChannel.CMD_CHANNEL_DISCONNECT: {
-                    if (VDBG) log("CMD_CHANNEL_DISCONNECT");
-                    if (mAsyncChannel != null) mAsyncChannel.disconnect();
-                    break;
-                }
-                case AsyncChannel.CMD_CHANNEL_DISCONNECTED: {
+                case EVENT_AGENT_DISCONNECTED: {
                     if (DBG) log("NetworkAgent channel lost");
                     // let the client know CS is done with us.
                     onNetworkUnwanted();
                     synchronized (mPreConnectedQueue) {
-                        mAsyncChannel = null;
+                        mRegistry = null;
                     }
                     break;
                 }
@@ -494,15 +509,7 @@
                 }
 
                 case CMD_SET_SIGNAL_STRENGTH_THRESHOLDS: {
-                    ArrayList<Integer> thresholds =
-                            ((Bundle) msg.obj).getIntegerArrayList("thresholds");
-                    // TODO: Change signal strength thresholds API to use an ArrayList<Integer>
-                    // rather than convert to int[].
-                    int[] intThresholds = new int[(thresholds != null) ? thresholds.size() : 0];
-                    for (int i = 0; i < intThresholds.length; i++) {
-                        intThresholds[i] = thresholds.get(i);
-                    }
-                    onSignalStrengthThresholdsUpdated(intThresholds);
+                    onSignalStrengthThresholdsUpdated((int[]) msg.obj);
                     break;
                 }
                 case CMD_PREVENT_AUTOMATIC_RECONNECT: {
@@ -541,7 +548,7 @@
             }
             final ConnectivityManager cm = (ConnectivityManager) mInitialConfiguration.context
                     .getSystemService(Context.CONNECTIVITY_SERVICE);
-            mNetwork = cm.registerNetworkAgent(new Messenger(mHandler),
+            mNetwork = cm.registerNetworkAgent(new NetworkAgentBinder(mHandler),
                     new NetworkInfo(mInitialConfiguration.info),
                     mInitialConfiguration.properties, mInitialConfiguration.capabilities,
                     mInitialConfiguration.score, mInitialConfiguration.config, providerId);
@@ -550,6 +557,95 @@
         return mNetwork;
     }
 
+    private static class NetworkAgentBinder extends INetworkAgent.Stub {
+        private final Handler mHandler;
+
+        private NetworkAgentBinder(Handler handler) {
+            mHandler = handler;
+        }
+
+        @Override
+        public void onRegistered(@NonNull INetworkAgentRegistry registry) {
+            mHandler.sendMessage(mHandler.obtainMessage(EVENT_AGENT_CONNECTED, registry));
+        }
+
+        @Override
+        public void onDisconnected() {
+            mHandler.sendMessage(mHandler.obtainMessage(EVENT_AGENT_DISCONNECTED));
+        }
+
+        @Override
+        public void onBandwidthUpdateRequested() {
+            mHandler.sendMessage(mHandler.obtainMessage(CMD_REQUEST_BANDWIDTH_UPDATE));
+        }
+
+        @Override
+        public void onValidationStatusChanged(
+                int validationStatus, @Nullable String captivePortalUrl) {
+            // TODO: consider using a parcelable as argument when the interface is structured
+            Bundle redirectUrlBundle = new Bundle();
+            redirectUrlBundle.putString(NetworkAgent.REDIRECT_URL_KEY, captivePortalUrl);
+            mHandler.sendMessage(mHandler.obtainMessage(CMD_REPORT_NETWORK_STATUS,
+                    validationStatus, 0, redirectUrlBundle));
+        }
+
+        @Override
+        public void onSaveAcceptUnvalidated(boolean acceptUnvalidated) {
+            mHandler.sendMessage(mHandler.obtainMessage(CMD_SAVE_ACCEPT_UNVALIDATED,
+                    acceptUnvalidated ? 1 : 0, 0));
+        }
+
+        @Override
+        public void onStartNattSocketKeepalive(int slot, int intervalDurationMs,
+                @NonNull NattKeepalivePacketData packetData) {
+            mHandler.sendMessage(mHandler.obtainMessage(CMD_START_SOCKET_KEEPALIVE,
+                    slot, intervalDurationMs, packetData));
+        }
+
+        @Override
+        public void onStartTcpSocketKeepalive(int slot, int intervalDurationMs,
+                @NonNull TcpKeepalivePacketData packetData) {
+            mHandler.sendMessage(mHandler.obtainMessage(CMD_START_SOCKET_KEEPALIVE,
+                    slot, intervalDurationMs, packetData));
+        }
+
+        @Override
+        public void onStopSocketKeepalive(int slot) {
+            mHandler.sendMessage(mHandler.obtainMessage(CMD_STOP_SOCKET_KEEPALIVE, slot, 0));
+        }
+
+        @Override
+        public void onSignalStrengthThresholdsUpdated(@NonNull int[] thresholds) {
+            mHandler.sendMessage(mHandler.obtainMessage(
+                    CMD_SET_SIGNAL_STRENGTH_THRESHOLDS, thresholds));
+        }
+
+        @Override
+        public void onPreventAutomaticReconnect() {
+            mHandler.sendMessage(mHandler.obtainMessage(CMD_PREVENT_AUTOMATIC_RECONNECT));
+        }
+
+        @Override
+        public void onAddNattKeepalivePacketFilter(int slot,
+                @NonNull NattKeepalivePacketData packetData) {
+            mHandler.sendMessage(mHandler.obtainMessage(CMD_ADD_KEEPALIVE_PACKET_FILTER,
+                    slot, 0, packetData));
+        }
+
+        @Override
+        public void onAddTcpKeepalivePacketFilter(int slot,
+                @NonNull TcpKeepalivePacketData packetData) {
+            mHandler.sendMessage(mHandler.obtainMessage(CMD_ADD_KEEPALIVE_PACKET_FILTER,
+                    slot, 0, packetData));
+        }
+
+        @Override
+        public void onRemoveKeepalivePacketFilter(int slot) {
+            mHandler.sendMessage(mHandler.obtainMessage(CMD_REMOVE_KEEPALIVE_PACKET_FILTER,
+                    slot, 0));
+        }
+    }
+
     /**
      * Register this network agent with a testing harness.
      *
@@ -559,13 +655,13 @@
      *
      * @hide
      */
-    public Messenger registerForTest(final Network network) {
+    public INetworkAgent registerForTest(final Network network) {
         log("Registering NetworkAgent for test");
         synchronized (mRegisterLock) {
             mNetwork = network;
             mInitialConfiguration = null;
         }
-        return new Messenger(mHandler);
+        return new NetworkAgentBinder(mHandler);
     }
 
     /**
@@ -589,29 +685,17 @@
         return mNetwork;
     }
 
-    private void queueOrSendMessage(int what, Object obj) {
-        queueOrSendMessage(what, 0, 0, obj);
-    }
-
-    private void queueOrSendMessage(int what, int arg1, int arg2) {
-        queueOrSendMessage(what, arg1, arg2, null);
-    }
-
-    private void queueOrSendMessage(int what, int arg1, int arg2, Object obj) {
-        Message msg = Message.obtain();
-        msg.what = what;
-        msg.arg1 = arg1;
-        msg.arg2 = arg2;
-        msg.obj = obj;
-        queueOrSendMessage(msg);
-    }
-
-    private void queueOrSendMessage(Message msg) {
+    private void queueOrSendMessage(@NonNull RegistryAction action) {
         synchronized (mPreConnectedQueue) {
-            if (mAsyncChannel != null) {
-                mAsyncChannel.sendMessage(msg);
+            if (mRegistry != null) {
+                try {
+                    action.execute(mRegistry);
+                } catch (RemoteException e) {
+                    Log.wtf(LOG_TAG, "Error executing registry action", e);
+                    // Fall through: the channel is asynchronous and does not report errors back
+                }
             } else {
-                mPreConnectedQueue.add(msg);
+                mPreConnectedQueue.add(action);
             }
         }
     }
@@ -622,7 +706,8 @@
      */
     public final void sendLinkProperties(@NonNull LinkProperties linkProperties) {
         Objects.requireNonNull(linkProperties);
-        queueOrSendMessage(EVENT_NETWORK_PROPERTIES_CHANGED, new LinkProperties(linkProperties));
+        final LinkProperties lp = new LinkProperties(linkProperties);
+        queueOrSendMessage(reg -> reg.sendLinkProperties(lp));
     }
 
     /**
@@ -647,9 +732,7 @@
     public final void setUnderlyingNetworks(@Nullable List<Network> underlyingNetworks) {
         final ArrayList<Network> underlyingArray = (underlyingNetworks != null)
                 ? new ArrayList<>(underlyingNetworks) : null;
-        final Bundle bundle = new Bundle();
-        bundle.putParcelableArrayList(UNDERLYING_NETWORKS_KEY, underlyingArray);
-        queueOrSendMessage(EVENT_UNDERLYING_NETWORKS_CHANGED, bundle);
+        queueOrSendMessage(reg -> reg.sendUnderlyingNetworks(underlyingArray));
     }
 
     /**
@@ -659,7 +742,7 @@
     public void markConnected() {
         mNetworkInfo.setDetailedState(NetworkInfo.DetailedState.CONNECTED, null /* reason */,
                 mNetworkInfo.getExtraInfo());
-        queueOrSendMessage(EVENT_NETWORK_INFO_CHANGED, mNetworkInfo);
+        queueOrSendNetworkInfo(mNetworkInfo);
     }
 
     /**
@@ -672,7 +755,7 @@
         // When unregistering an agent nobody should use the extrainfo (or reason) any more.
         mNetworkInfo.setDetailedState(NetworkInfo.DetailedState.DISCONNECTED, null /* reason */,
                 null /* extraInfo */);
-        queueOrSendMessage(EVENT_NETWORK_INFO_CHANGED, mNetworkInfo);
+        queueOrSendNetworkInfo(mNetworkInfo);
     }
 
     /**
@@ -689,7 +772,7 @@
     @Deprecated
     public void setLegacySubtype(final int legacySubtype, @NonNull final String legacySubtypeName) {
         mNetworkInfo.setSubtype(legacySubtype, legacySubtypeName);
-        queueOrSendMessage(EVENT_NETWORK_INFO_CHANGED, mNetworkInfo);
+        queueOrSendNetworkInfo(mNetworkInfo);
     }
 
     /**
@@ -711,7 +794,7 @@
     @Deprecated
     public void setLegacyExtraInfo(@Nullable final String extraInfo) {
         mNetworkInfo.setExtraInfo(extraInfo);
-        queueOrSendMessage(EVENT_NETWORK_INFO_CHANGED, mNetworkInfo);
+        queueOrSendNetworkInfo(mNetworkInfo);
     }
 
     /**
@@ -720,7 +803,11 @@
      */
     @UnsupportedAppUsage(maxTargetSdk = Build.VERSION_CODES.P, trackingBug = 115609023)
     public final void sendNetworkInfo(NetworkInfo networkInfo) {
-        queueOrSendMessage(EVENT_NETWORK_INFO_CHANGED, new NetworkInfo(networkInfo));
+        queueOrSendNetworkInfo(new NetworkInfo(networkInfo));
+    }
+
+    private void queueOrSendNetworkInfo(NetworkInfo networkInfo) {
+        queueOrSendMessage(reg -> reg.sendNetworkInfo(networkInfo));
     }
 
     /**
@@ -731,8 +818,8 @@
         Objects.requireNonNull(networkCapabilities);
         mBandwidthUpdatePending.set(false);
         mLastBwRefreshTime = System.currentTimeMillis();
-        queueOrSendMessage(EVENT_NETWORK_CAPABILITIES_CHANGED,
-                new NetworkCapabilities(networkCapabilities));
+        final NetworkCapabilities nc = new NetworkCapabilities(networkCapabilities);
+        queueOrSendMessage(reg -> reg.sendNetworkCapabilities(nc));
     }
 
     /**
@@ -744,7 +831,7 @@
         if (score < 0) {
             throw new IllegalArgumentException("Score must be >= 0");
         }
-        queueOrSendMessage(EVENT_NETWORK_SCORE_CHANGED, score, 0);
+        queueOrSendMessage(reg -> reg.sendScore(score));
     }
 
     /**
@@ -784,9 +871,8 @@
      * @hide should move to NetworkAgentConfig.
      */
     public void explicitlySelected(boolean explicitlySelected, boolean acceptUnvalidated) {
-        queueOrSendMessage(EVENT_SET_EXPLICITLY_SELECTED,
-                explicitlySelected ? 1 : 0,
-                acceptUnvalidated ? 1 : 0);
+        queueOrSendMessage(reg -> reg.sendExplicitlySelected(
+                explicitlySelected, acceptUnvalidated));
     }
 
     /**
@@ -909,7 +995,7 @@
      */
     public final void sendSocketKeepaliveEvent(int slot,
             @SocketKeepalive.KeepaliveEvent int event) {
-        queueOrSendMessage(EVENT_SOCKET_KEEPALIVE, slot, event);
+        queueOrSendMessage(reg -> reg.sendSocketKeepaliveEvent(slot, event));
     }
     /** @hide TODO delete once callers have moved to sendSocketKeepaliveEvent */
     public void onSocketKeepaliveEvent(int slot, int reason) {
diff --git a/core/java/android/net/SocketKeepalive.java b/core/java/android/net/SocketKeepalive.java
index a7dce18..d007a95 100644
--- a/core/java/android/net/SocketKeepalive.java
+++ b/core/java/android/net/SocketKeepalive.java
@@ -187,38 +187,54 @@
         mCallback = new ISocketKeepaliveCallback.Stub() {
             @Override
             public void onStarted(int slot) {
-                Binder.withCleanCallingIdentity(() ->
-                        mExecutor.execute(() -> {
-                            mSlot = slot;
-                            callback.onStarted();
-                        }));
+                final long token = Binder.clearCallingIdentity();
+                try {
+                    mExecutor.execute(() -> {
+                        mSlot = slot;
+                        callback.onStarted();
+                    });
+                } finally {
+                    Binder.restoreCallingIdentity(token);
+                }
             }
 
             @Override
             public void onStopped() {
-                Binder.withCleanCallingIdentity(() ->
-                        executor.execute(() -> {
-                            mSlot = null;
-                            callback.onStopped();
-                        }));
+                final long token = Binder.clearCallingIdentity();
+                try {
+                    executor.execute(() -> {
+                        mSlot = null;
+                        callback.onStopped();
+                    });
+                } finally {
+                    Binder.restoreCallingIdentity(token);
+                }
             }
 
             @Override
             public void onError(int error) {
-                Binder.withCleanCallingIdentity(() ->
-                        executor.execute(() -> {
-                            mSlot = null;
-                            callback.onError(error);
-                        }));
+                final long token = Binder.clearCallingIdentity();
+                try {
+                    executor.execute(() -> {
+                        mSlot = null;
+                        callback.onError(error);
+                    });
+                } finally {
+                    Binder.restoreCallingIdentity(token);
+                }
             }
 
             @Override
             public void onDataReceived() {
-                Binder.withCleanCallingIdentity(() ->
-                        executor.execute(() -> {
-                            mSlot = null;
-                            callback.onDataReceived();
-                        }));
+                final long token = Binder.clearCallingIdentity();
+                try {
+                    executor.execute(() -> {
+                        mSlot = null;
+                        callback.onDataReceived();
+                    });
+                } finally {
+                    Binder.restoreCallingIdentity(token);
+                }
             }
         };
     }
diff --git a/services/core/java/com/android/server/ConnectivityService.java b/services/core/java/com/android/server/ConnectivityService.java
index 26c0e59..397eeb2 100644
--- a/services/core/java/com/android/server/ConnectivityService.java
+++ b/services/core/java/com/android/server/ConnectivityService.java
@@ -74,7 +74,6 @@
 import android.content.Intent;
 import android.content.IntentFilter;
 import android.content.pm.PackageManager;
-import android.content.res.Configuration;
 import android.database.ContentObserver;
 import android.net.CaptivePortal;
 import android.net.CaptivePortalData;
@@ -90,12 +89,10 @@
 import android.net.IDnsResolver;
 import android.net.IIpConnectivityMetrics;
 import android.net.INetd;
-import android.net.INetdEventCallback;
 import android.net.INetworkManagementEventObserver;
 import android.net.INetworkMonitor;
 import android.net.INetworkMonitorCallbacks;
 import android.net.INetworkPolicyListener;
-import android.net.INetworkPolicyManager;
 import android.net.INetworkStatsService;
 import android.net.ISocketKeepaliveCallback;
 import android.net.InetAddresses;
@@ -133,6 +130,7 @@
 import android.net.Uri;
 import android.net.VpnManager;
 import android.net.VpnService;
+import android.net.metrics.INetdEventListener;
 import android.net.metrics.IpConnectivityLog;
 import android.net.metrics.NetworkEvent;
 import android.net.netlink.InetDiagMessage;
@@ -173,8 +171,8 @@
 import android.util.Pair;
 import android.util.SparseArray;
 import android.util.SparseIntArray;
-import android.util.Xml;
 
+import com.android.connectivity.aidl.INetworkAgent;
 import com.android.internal.R;
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.annotations.VisibleForTesting;
@@ -189,7 +187,6 @@
 import com.android.internal.util.IndentingPrintWriter;
 import com.android.internal.util.LocationPermissionChecker;
 import com.android.internal.util.MessageUtils;
-import com.android.internal.util.XmlUtils;
 import com.android.modules.utils.BasicShellCommandHandler;
 import com.android.net.module.util.LinkPropertiesUtils.CompareOrUpdateResult;
 import com.android.net.module.util.LinkPropertiesUtils.CompareResult;
@@ -210,7 +207,6 @@
 import com.android.server.connectivity.PermissionMonitor;
 import com.android.server.connectivity.ProxyTracker;
 import com.android.server.connectivity.Vpn;
-import com.android.server.net.BaseNetdEventCallback;
 import com.android.server.net.BaseNetworkObserver;
 import com.android.server.net.LockdownVpnTracker;
 import com.android.server.net.NetworkPolicyManagerInternal;
@@ -220,14 +216,7 @@
 
 import libcore.io.IoUtils;
 
-import org.xmlpull.v1.XmlPullParser;
-import org.xmlpull.v1.XmlPullParserException;
-
-import java.io.File;
 import java.io.FileDescriptor;
-import java.io.FileNotFoundException;
-import java.io.FileReader;
-import java.io.IOException;
 import java.io.PrintWriter;
 import java.net.Inet4Address;
 import java.net.InetAddress;
@@ -336,7 +325,7 @@
     @VisibleForTesting
     protected INetd mNetd;
     private INetworkStatsService mStatsService;
-    private INetworkPolicyManager mPolicyManager;
+    private NetworkPolicyManager mPolicyManager;
     private NetworkPolicyManagerInternal mPolicyManagerInternal;
 
     /**
@@ -557,6 +546,13 @@
     private static final int EVENT_CAPPORT_DATA_CHANGED = 46;
 
     /**
+     * Used by setRequireVpnForUids.
+     * arg1 = whether the specified UID ranges are required to use a VPN.
+     * obj  = Array of UidRange objects.
+     */
+    private static final int EVENT_SET_REQUIRE_VPN_FOR_UIDS = 47;
+
+    /**
      * Argument for {@link #EVENT_PROVISIONING_NOTIFICATION} to indicate that the notification
      * should be shown.
      */
@@ -945,15 +941,15 @@
     }
 
     public ConnectivityService(Context context, INetworkManagementService netManager,
-            INetworkStatsService statsService, INetworkPolicyManager policyManager) {
-        this(context, netManager, statsService, policyManager, getDnsResolver(context),
-                new IpConnectivityLog(), NetdService.getInstance(), new Dependencies());
+            INetworkStatsService statsService) {
+        this(context, netManager, statsService, getDnsResolver(context), new IpConnectivityLog(),
+                NetdService.getInstance(), new Dependencies());
     }
 
     @VisibleForTesting
     protected ConnectivityService(Context context, INetworkManagementService netManager,
-            INetworkStatsService statsService, INetworkPolicyManager policyManager,
-            IDnsResolver dnsresolver, IpConnectivityLog logger, INetd netd, Dependencies deps) {
+            INetworkStatsService statsService, IDnsResolver dnsresolver, IpConnectivityLog logger,
+            INetd netd, Dependencies deps) {
         if (DBG) log("ConnectivityService starting up");
 
         mDeps = Objects.requireNonNull(deps, "missing Dependencies");
@@ -991,7 +987,7 @@
 
         mNMS = Objects.requireNonNull(netManager, "missing INetworkManagementService");
         mStatsService = Objects.requireNonNull(statsService, "missing INetworkStatsService");
-        mPolicyManager = Objects.requireNonNull(policyManager, "missing INetworkPolicyManager");
+        mPolicyManager = mContext.getSystemService(NetworkPolicyManager.class);
         mPolicyManagerInternal = Objects.requireNonNull(
                 LocalServices.getService(NetworkPolicyManagerInternal.class),
                 "missing NetworkPolicyManagerInternal");
@@ -1007,12 +1003,7 @@
         // To ensure uid rules are synchronized with Network Policy, register for
         // NetworkPolicyManagerService events must happen prior to NetworkPolicyManagerService
         // reading existing policy from disk.
-        try {
-            mPolicyManager.registerListener(mPolicyListener);
-        } catch (RemoteException e) {
-            // ouch, no rules updates means some processes may never get network
-            loge("unable to register INetworkPolicyListener" + e);
-        }
+        mPolicyManager.registerListener(mPolicyListener);
 
         final PowerManager powerManager = (PowerManager) context.getSystemService(
                 Context.POWER_SERVICE);
@@ -1289,19 +1280,28 @@
         }
     }
 
-    private Network[] getVpnUnderlyingNetworks(int uid) {
-        synchronized (mVpns) {
-            if (!mLockdownEnabled) {
-                int user = UserHandle.getUserId(uid);
-                Vpn vpn = mVpns.get(user);
-                if (vpn != null && vpn.appliesToUid(uid)) {
-                    return vpn.getUnderlyingNetworks();
+    // TODO: determine what to do when more than one VPN applies to |uid|.
+    private NetworkAgentInfo getVpnForUid(int uid) {
+        synchronized (mNetworkForNetId) {
+            for (int i = 0; i < mNetworkForNetId.size(); i++) {
+                final NetworkAgentInfo nai = mNetworkForNetId.valueAt(i);
+                if (nai.isVPN() && nai.everConnected && nai.networkCapabilities.appliesToUid(uid)) {
+                    return nai;
                 }
             }
         }
         return null;
     }
 
+    private Network[] getVpnUnderlyingNetworks(int uid) {
+        synchronized (mVpns) {
+            if (mLockdownEnabled) return null;
+        }
+        final NetworkAgentInfo nai = getVpnForUid(uid);
+        if (nai != null) return nai.declaredUnderlyingNetworks;
+        return null;
+    }
+
     private NetworkState getUnfilteredActiveNetworkState(int uid) {
         NetworkAgentInfo nai = getDefaultNetwork();
 
@@ -1327,7 +1327,7 @@
     }
 
     /**
-     * Check if UID should be blocked from using the network with the given LinkProperties.
+     * Check if UID should be blocked from using the specified network.
      */
     private boolean isNetworkWithLinkPropertiesBlocked(LinkProperties lp, int uid,
             boolean ignoreBlocked) {
@@ -1335,12 +1335,7 @@
         if (ignoreBlocked) {
             return false;
         }
-        synchronized (mVpns) {
-            final Vpn vpn = mVpns.get(UserHandle.getUserId(uid));
-            if (vpn != null && vpn.getLockdown() && vpn.isBlockingUid(uid)) {
-                return true;
-            }
-        }
+        if (isUidBlockedByVpn(uid, mVpnBlockedUidRanges)) return true;
         final String iface = (lp == null ? "" : lp.getInterfaceName());
         return mPolicyManagerInternal.isUidNetworkingBlocked(uid, iface);
     }
@@ -1566,22 +1561,14 @@
                             nc, mDeps.getCallingUid(), callingPackageName));
         }
 
-        synchronized (mVpns) {
-            if (!mLockdownEnabled) {
-                Vpn vpn = mVpns.get(userId);
-                if (vpn != null) {
-                    Network[] networks = vpn.getUnderlyingNetworks();
-                    if (networks != null) {
-                        for (Network network : networks) {
-                            nc = getNetworkCapabilitiesInternal(network);
-                            if (nc != null) {
-                                result.put(
-                                        network,
-                                        maybeSanitizeLocationInfoForCaller(
-                                                nc, mDeps.getCallingUid(), callingPackageName));
-                            }
-                        }
-                    }
+        // No need to check mLockdownEnabled. If it's true, getVpnUnderlyingNetworks returns null.
+        final Network[] networks = getVpnUnderlyingNetworks(Binder.getCallingUid());
+        if (networks != null) {
+            for (Network network : networks) {
+                nc = getNetworkCapabilitiesInternal(network);
+                if (nc != null) {
+                    result.put(network, maybeSanitizeLocationInfoForCaller(
+                                    nc, mDeps.getCallingUid(), callingPackageName));
                 }
             }
         }
@@ -1919,8 +1906,7 @@
         return true;
     }
 
-    @VisibleForTesting
-    protected final INetdEventCallback mNetdEventCallback = new BaseNetdEventCallback() {
+    private class NetdEventCallback extends INetdEventListener.Stub {
         @Override
         public void onPrivateDnsValidationEvent(int netId, String ipAddress,
                 String hostname, boolean validated) {
@@ -1936,8 +1922,8 @@
         }
 
         @Override
-        public void onDnsEvent(int netId, int eventType, int returnCode, String hostname,
-                String[] ipAddresses, int ipAddressesCount, long timestamp, int uid) {
+        public void onDnsEvent(int netId, int eventType, int returnCode, int latencyMs,
+                String hostname,  String[] ipAddresses, int ipAddressesCount, int uid) {
             NetworkAgentInfo nai = getNetworkAgentInfoForNetId(netId);
             // Netd event only allow registrants from system. Each NetworkMonitor thread is under
             // the caller thread of registerNetworkAgent. Thus, it's not allowed to register netd
@@ -1956,21 +1942,42 @@
                                        String prefixString, int prefixLength) {
             mHandler.post(() -> handleNat64PrefixEvent(netId, added, prefixString, prefixLength));
         }
-    };
 
-    private void registerNetdEventCallback() {
-        final IIpConnectivityMetrics ipConnectivityMetrics = mDeps.getIpConnectivityMetrics();
-        if (ipConnectivityMetrics == null) {
-            Log.wtf(TAG, "Missing IIpConnectivityMetrics");
-            return;
+        @Override
+        public void onConnectEvent(int netId, int error, int latencyMs, String ipAddr, int port,
+                int uid) {
         }
 
+        @Override
+        public void onWakeupEvent(String prefix, int uid, int ethertype, int ipNextHeader,
+                byte[] dstHw, String srcIp, String dstIp, int srcPort, int dstPort,
+                long timestampNs) {
+        }
+
+        @Override
+        public void onTcpSocketStatsEvent(int[] networkIds, int[] sentPackets, int[] lostPackets,
+                int[] rttsUs, int[] sentAckDiffsMs) {
+        }
+
+        @Override
+        public int getInterfaceVersion() throws RemoteException {
+            return this.VERSION;
+        }
+
+        @Override
+        public String getInterfaceHash() {
+            return this.HASH;
+        }
+    };
+
+    @VisibleForTesting
+    protected final INetdEventListener mNetdEventCallback = new NetdEventCallback();
+
+    private void registerNetdEventCallback() {
         try {
-            ipConnectivityMetrics.addNetdEventCallback(
-                    INetdEventCallback.CALLBACK_CALLER_CONNECTIVITY_SERVICE,
-                    mNetdEventCallback);
+            mDnsResolver.registerEventListener(mNetdEventCallback);
         } catch (Exception e) {
-            loge("Error registering netd callback: " + e);
+            loge("Error registering DnsResolver callback: " + e);
         }
     }
 
@@ -2007,29 +2014,18 @@
     void handleRestrictBackgroundChanged(boolean restrictBackground) {
         if (mRestrictBackground == restrictBackground) return;
 
-        for (final NetworkAgentInfo nai : mNetworkAgentInfos.values()) {
+        final List<UidRange> blockedRanges = mVpnBlockedUidRanges;
+        for (final NetworkAgentInfo nai : mNetworkAgentInfos) {
             final boolean curMetered = nai.networkCapabilities.isMetered();
             maybeNotifyNetworkBlocked(nai, curMetered, curMetered, mRestrictBackground,
-                    restrictBackground);
+                    restrictBackground, blockedRanges, blockedRanges);
         }
 
         mRestrictBackground = restrictBackground;
     }
 
-    private boolean isUidNetworkingWithVpnBlocked(int uid, int uidRules, boolean isNetworkMetered,
+    private boolean isUidBlockedByRules(int uid, int uidRules, boolean isNetworkMetered,
             boolean isBackgroundRestricted) {
-        synchronized (mVpns) {
-            final Vpn vpn = mVpns.get(UserHandle.getUserId(uid));
-            // Because the return value of this function depends on the list of UIDs the
-            // always-on VPN blocks when in lockdown mode, when the always-on VPN changes that
-            // list all state depending on the return value of this function has to be recomputed.
-            // TODO: add a trigger when the always-on VPN sets its blocked UIDs to reevaluate and
-            // send the necessary onBlockedStatusChanged callbacks.
-            if (vpn != null && vpn.getLockdown() && vpn.isBlockingUid(uid)) {
-                return true;
-            }
-        }
-
         return NetworkPolicyManagerInternal.isUidNetworkingBlocked(uid, uidRules,
                 isNetworkMetered, isBackgroundRestricted);
     }
@@ -2715,7 +2711,7 @@
      */
     private NetworkAgentInfo[] networksSortedById() {
         NetworkAgentInfo[] networks = new NetworkAgentInfo[0];
-        networks = mNetworkAgentInfos.values().toArray(networks);
+        networks = mNetworkAgentInfos.toArray(networks);
         Arrays.sort(networks, Comparator.comparingInt(nai -> nai.network.getNetId()));
         return networks;
     }
@@ -2761,11 +2757,6 @@
                     handleAsyncChannelHalfConnect(msg);
                     break;
                 }
-                case AsyncChannel.CMD_CHANNEL_DISCONNECT: {
-                    NetworkAgentInfo nai = mNetworkAgentInfos.get(msg.replyTo);
-                    if (nai != null) nai.asyncChannel.disconnect();
-                    break;
-                }
                 case AsyncChannel.CMD_CHANNEL_DISCONNECTED: {
                     handleAsyncChannelDisconnected(msg);
                     break;
@@ -2775,8 +2766,9 @@
         }
 
         private void maybeHandleNetworkAgentMessage(Message msg) {
-            NetworkAgentInfo nai = mNetworkAgentInfos.get(msg.replyTo);
-            if (nai == null) {
+            final Pair<NetworkAgentInfo, Object> arg = (Pair<NetworkAgentInfo, Object>) msg.obj;
+            final NetworkAgentInfo nai = arg.first;
+            if (!mNetworkAgentInfos.contains(nai)) {
                 if (VDBG) {
                     log(String.format("%s from unknown NetworkAgent", eventName(msg.what)));
                 }
@@ -2785,7 +2777,7 @@
 
             switch (msg.what) {
                 case NetworkAgent.EVENT_NETWORK_CAPABILITIES_CHANGED: {
-                    NetworkCapabilities networkCapabilities = (NetworkCapabilities) msg.obj;
+                    NetworkCapabilities networkCapabilities = (NetworkCapabilities) arg.second;
                     if (networkCapabilities.hasConnectivityManagedCapability()) {
                         Log.wtf(TAG, "BUG: " + nai + " has CS-managed capability.");
                     }
@@ -2802,13 +2794,13 @@
                     break;
                 }
                 case NetworkAgent.EVENT_NETWORK_PROPERTIES_CHANGED: {
-                    LinkProperties newLp = (LinkProperties) msg.obj;
+                    LinkProperties newLp = (LinkProperties) arg.second;
                     processLinkPropertiesFromAgent(nai, newLp);
                     handleUpdateLinkProperties(nai, newLp);
                     break;
                 }
                 case NetworkAgent.EVENT_NETWORK_INFO_CHANGED: {
-                    NetworkInfo info = (NetworkInfo) msg.obj;
+                    NetworkInfo info = (NetworkInfo) arg.second;
                     updateNetworkInfo(nai, info);
                     break;
                 }
@@ -2833,7 +2825,7 @@
                     break;
                 }
                 case NetworkAgent.EVENT_SOCKET_KEEPALIVE: {
-                    mKeepaliveTracker.handleEventSocketKeepalive(nai, msg);
+                    mKeepaliveTracker.handleEventSocketKeepalive(nai, msg.arg1, msg.arg2);
                     break;
                 }
                 case NetworkAgent.EVENT_UNDERLYING_NETWORKS_CHANGED: {
@@ -2844,7 +2836,7 @@
                     }
                     final ArrayList<Network> underlying;
                     try {
-                        underlying = ((Bundle) msg.obj).getParcelableArrayList(
+                        underlying = ((Bundle) arg.second).getParcelableArrayList(
                                 NetworkAgent.UNDERLYING_NETWORKS_KEY);
                     } catch (NullPointerException | ClassCastException e) {
                         break;
@@ -2923,8 +2915,7 @@
                         if (nai.lastCaptivePortalDetected &&
                             Settings.Global.CAPTIVE_PORTAL_MODE_AVOID == getCaptivePortalMode()) {
                             if (DBG) log("Avoiding captive portal network: " + nai.toShortString());
-                            nai.asyncChannel.sendMessage(
-                                    NetworkAgent.CMD_PREVENT_AUTOMATIC_RECONNECT);
+                            nai.onPreventAutomaticReconnect();
                             teardownUnneededNetwork(nai);
                             break;
                         }
@@ -3016,13 +3007,10 @@
             }
             updateInetCondition(nai);
             // Let the NetworkAgent know the state of its network
-            Bundle redirectUrlBundle = new Bundle();
-            redirectUrlBundle.putString(NetworkAgent.REDIRECT_URL_KEY, redirectUrl);
             // TODO: Evaluate to update partial connectivity to status to NetworkAgent.
-            nai.asyncChannel.sendMessage(
-                    NetworkAgent.CMD_REPORT_NETWORK_STATUS,
-                    (valid ? NetworkAgent.VALID_NETWORK : NetworkAgent.INVALID_NETWORK),
-                    0, redirectUrlBundle);
+            nai.onValidationStatusChanged(
+                    valid ? NetworkAgent.VALID_NETWORK : NetworkAgent.INVALID_NETWORK,
+                    redirectUrl);
 
             // If NetworkMonitor detects partial connectivity before
             // EVENT_PROMPT_UNVALIDATED arrives, show the partial connectivity notification
@@ -3056,6 +3044,14 @@
                     }
                     break;
                 }
+                case NetworkAgentInfo.EVENT_AGENT_REGISTERED: {
+                    handleNetworkAgentRegistered(msg);
+                    break;
+                }
+                case NetworkAgentInfo.EVENT_AGENT_DISCONNECTED: {
+                    handleNetworkAgentDisconnected(msg);
+                    break;
+                }
             }
             return true;
         }
@@ -3232,7 +3228,7 @@
     private void handlePrivateDnsSettingsChanged() {
         final PrivateDnsConfig cfg = mDnsManager.getPrivateDnsConfig();
 
-        for (NetworkAgentInfo nai : mNetworkAgentInfos.values()) {
+        for (NetworkAgentInfo nai : mNetworkAgentInfos) {
             handlePerNetworkPrivateDnsConfig(nai, cfg);
             if (networkRequiresPrivateDnsValidation(nai)) {
                 handleUpdateLinkProperties(nai, new LinkProperties(nai.linkProperties));
@@ -3330,7 +3326,6 @@
 
     private void handleAsyncChannelHalfConnect(Message msg) {
         ensureRunningOnConnectivityServiceThread();
-        final AsyncChannel ac = (AsyncChannel) msg.obj;
         if (mNetworkProviderInfos.containsKey(msg.replyTo)) {
             if (msg.arg1 == AsyncChannel.STATUS_SUCCESSFUL) {
                 if (VDBG) log("NetworkFactory connected");
@@ -3342,39 +3337,45 @@
                 loge("Error connecting NetworkFactory");
                 mNetworkProviderInfos.remove(msg.obj);
             }
-        } else if (mNetworkAgentInfos.containsKey(msg.replyTo)) {
-            if (msg.arg1 == AsyncChannel.STATUS_SUCCESSFUL) {
-                if (VDBG) log("NetworkAgent connected");
-                // A network agent has requested a connection.  Establish the connection.
-                mNetworkAgentInfos.get(msg.replyTo).asyncChannel.
-                        sendMessage(AsyncChannel.CMD_CHANNEL_FULL_CONNECTION);
-            } else {
-                loge("Error connecting NetworkAgent");
-                NetworkAgentInfo nai = mNetworkAgentInfos.remove(msg.replyTo);
-                if (nai != null) {
-                    final boolean wasDefault = isDefaultNetwork(nai);
-                    synchronized (mNetworkForNetId) {
-                        mNetworkForNetId.remove(nai.network.getNetId());
-                    }
-                    mNetIdManager.releaseNetId(nai.network.getNetId());
-                    // Just in case.
-                    mLegacyTypeTracker.remove(nai, wasDefault);
+        }
+    }
+
+    private void handleNetworkAgentRegistered(Message msg) {
+        final NetworkAgentInfo nai = (NetworkAgentInfo) msg.obj;
+        if (!mNetworkAgentInfos.contains(nai)) {
+            return;
+        }
+
+        if (msg.arg1 == NetworkAgentInfo.ARG_AGENT_SUCCESS) {
+            if (VDBG) log("NetworkAgent registered");
+        } else {
+            loge("Error connecting NetworkAgent");
+            mNetworkAgentInfos.remove(nai);
+            if (nai != null) {
+                final boolean wasDefault = isDefaultNetwork(nai);
+                synchronized (mNetworkForNetId) {
+                    mNetworkForNetId.remove(nai.network.getNetId());
                 }
+                mNetIdManager.releaseNetId(nai.network.getNetId());
+                // Just in case.
+                mLegacyTypeTracker.remove(nai, wasDefault);
             }
         }
     }
 
-    // This is a no-op if it's called with a message designating a network that has
+    private void handleNetworkAgentDisconnected(Message msg) {
+        NetworkAgentInfo nai = (NetworkAgentInfo) msg.obj;
+        if (mNetworkAgentInfos.contains(nai)) {
+            disconnectAndDestroyNetwork(nai);
+        }
+    }
+
+    // This is a no-op if it's called with a message designating a provider that has
     // already been destroyed, because its reference will not be found in the relevant
     // maps.
     private void handleAsyncChannelDisconnected(Message msg) {
-        NetworkAgentInfo nai = mNetworkAgentInfos.get(msg.replyTo);
-        if (nai != null) {
-            disconnectAndDestroyNetwork(nai);
-        } else {
-            NetworkProviderInfo npi = mNetworkProviderInfos.remove(msg.replyTo);
-            if (DBG && npi != null) log("unregisterNetworkFactory for " + npi.name);
-        }
+        NetworkProviderInfo npi = mNetworkProviderInfos.remove(msg.replyTo);
+        if (DBG && npi != null) log("unregisterNetworkFactory for " + npi.name);
     }
 
     // Destroys a network, remove references to it from the internal state managed by
@@ -3418,7 +3419,7 @@
             wakeupModifyInterface(iface, nai.networkCapabilities, false);
         }
         nai.networkMonitor().notifyNetworkDisconnected();
-        mNetworkAgentInfos.remove(nai.messenger);
+        mNetworkAgentInfos.remove(nai);
         nai.clatd.update();
         synchronized (mNetworkForNetId) {
             // Remove the NetworkAgent, but don't mark the netId as
@@ -3526,7 +3527,7 @@
         mNetworkRequests.put(nri.request, nri);
         mNetworkRequestInfoLogs.log("REGISTER " + nri);
         if (nri.request.isListen()) {
-            for (NetworkAgentInfo network : mNetworkAgentInfos.values()) {
+            for (NetworkAgentInfo network : mNetworkAgentInfos) {
                 if (nri.request.networkCapabilities.hasSignalStrength() &&
                         network.satisfiesImmutableCapabilitiesOf(nri.request)) {
                     updateSignalStrengthThresholds(network, "REGISTER", nri.request);
@@ -3591,8 +3592,8 @@
     private boolean isNetworkPotentialSatisfier(
             @NonNull final NetworkAgentInfo candidate, @NonNull final NetworkRequestInfo nri) {
         // listen requests won't keep up a network satisfying it. If this is not a multilayer
-        // request, we can return immediately. For multilayer requests, we have to check to see if
-        // any of the multilayer requests may have a potential satisfier.
+        // request, return immediately. For multilayer requests, check to see if any of the
+        // multilayer requests may have a potential satisfier.
         if (!nri.isMultilayerRequest() && nri.mRequests.get(0).isListen()) {
             return false;
         }
@@ -3742,7 +3743,7 @@
         } else {
             // listens don't have a singular affectedNetwork.  Check all networks to see
             // if this listen request applies and remove it.
-            for (NetworkAgentInfo nai : mNetworkAgentInfos.values()) {
+            for (NetworkAgentInfo nai : mNetworkAgentInfos) {
                 nai.removeRequest(nri.request.requestId);
                 if (nri.request.networkCapabilities.hasSignalStrength() &&
                         nai.satisfiesImmutableCapabilitiesOf(nri.request)) {
@@ -3815,13 +3816,12 @@
         }
 
         if (always) {
-            nai.asyncChannel.sendMessage(
-                    NetworkAgent.CMD_SAVE_ACCEPT_UNVALIDATED, encodeBool(accept));
+            nai.onSaveAcceptUnvalidated(accept);
         }
 
         if (!accept) {
             // Tell the NetworkAgent to not automatically reconnect to the network.
-            nai.asyncChannel.sendMessage(NetworkAgent.CMD_PREVENT_AUTOMATIC_RECONNECT);
+            nai.onPreventAutomaticReconnect();
             // Teardown the network.
             teardownUnneededNetwork(nai);
         }
@@ -3852,13 +3852,12 @@
 
         // TODO: Use the current design or save the user choice into IpMemoryStore.
         if (always) {
-            nai.asyncChannel.sendMessage(
-                    NetworkAgent.CMD_SAVE_ACCEPT_UNVALIDATED, encodeBool(accept));
+            nai.onSaveAcceptUnvalidated(accept);
         }
 
         if (!accept) {
             // Tell the NetworkAgent to not automatically reconnect to the network.
-            nai.asyncChannel.sendMessage(NetworkAgent.CMD_PREVENT_AUTOMATIC_RECONNECT);
+            nai.onPreventAutomaticReconnect();
             // Tear down the network.
             teardownUnneededNetwork(nai);
         } else {
@@ -3996,7 +3995,7 @@
 
     private void rematchForAvoidBadWifiUpdate() {
         rematchAllNetworksAndRequests();
-        for (NetworkAgentInfo nai: mNetworkAgentInfos.values()) {
+        for (NetworkAgentInfo nai: mNetworkAgentInfos) {
             if (nai.networkCapabilities.hasTransport(NetworkCapabilities.TRANSPORT_WIFI)) {
                 sendUpdatedScoreToFactories(nai);
             }
@@ -4139,7 +4138,7 @@
         // to a network that provides no or limited connectivity is not useful, because the user
         // cannot use that network except through the notification shown by this method, and the
         // notification is only shown if the network is explicitly selected by the user.
-        nai.asyncChannel.sendMessage(NetworkAgent.CMD_PREVENT_AUTOMATIC_RECONNECT);
+        nai.onPreventAutomaticReconnect();
 
         // TODO: Evaluate if it's needed to wait 8 seconds for triggering notification when
         // NetworkMonitor detects the network is partial connectivity. Need to change the design to
@@ -4301,6 +4300,9 @@
                 case EVENT_DATA_SAVER_CHANGED:
                     handleRestrictBackgroundChanged(toBool(msg.arg1));
                     break;
+                case EVENT_SET_REQUIRE_VPN_FOR_UIDS:
+                    handleSetRequireVpnForUids(toBool(msg.arg1), (UidRange[]) msg.obj);
+                    break;
             }
         }
     }
@@ -4469,8 +4471,7 @@
         if (!nai.everConnected) {
             return;
         }
-        LinkProperties lp = getLinkProperties(nai);
-        if (isNetworkWithLinkPropertiesBlocked(lp, uid, false)) {
+        if (isNetworkWithLinkPropertiesBlocked(nai.linkProperties, uid, false)) {
             return;
         }
         nai.networkMonitor().forceReevaluation(uid);
@@ -4789,7 +4790,7 @@
                 return new VpnInfo[0];
             }
             List<VpnInfo> infoList = new ArrayList<>();
-            for (NetworkAgentInfo nai : mNetworkAgentInfos.values()) {
+            for (NetworkAgentInfo nai : mNetworkAgentInfos) {
                 VpnInfo info = createVpnInfo(nai);
                 if (info != null) {
                     infoList.add(info);
@@ -4890,13 +4891,63 @@
      */
     private void propagateUnderlyingNetworkCapabilities(Network updatedNetwork) {
         ensureRunningOnConnectivityServiceThread();
-        for (NetworkAgentInfo nai : mNetworkAgentInfos.values()) {
+        for (NetworkAgentInfo nai : mNetworkAgentInfos) {
             if (updatedNetwork == null || hasUnderlyingNetwork(nai, updatedNetwork)) {
                 updateCapabilitiesForNetwork(nai);
             }
         }
     }
 
+    private boolean isUidBlockedByVpn(int uid, List<UidRange> blockedUidRanges) {
+        // Determine whether this UID is blocked because of always-on VPN lockdown. If a VPN applies
+        // to the UID, then the UID is not blocked because always-on VPN lockdown applies only when
+        // a VPN is not up.
+        final NetworkAgentInfo vpnNai = getVpnForUid(uid);
+        if (vpnNai != null && !vpnNai.networkAgentConfig.allowBypass) return false;
+        for (UidRange range : blockedUidRanges) {
+            if (range.contains(uid)) return true;
+        }
+        return false;
+    }
+
+    @Override
+    public void setRequireVpnForUids(boolean requireVpn, UidRange[] ranges) {
+        NetworkStack.checkNetworkStackPermission(mContext);
+        mHandler.sendMessage(mHandler.obtainMessage(EVENT_SET_REQUIRE_VPN_FOR_UIDS,
+                encodeBool(requireVpn), 0 /* arg2 */, ranges));
+    }
+
+    private void handleSetRequireVpnForUids(boolean requireVpn, UidRange[] ranges) {
+        if (DBG) {
+            Log.d(TAG, "Setting VPN " + (requireVpn ? "" : "not ") + "required for UIDs: "
+                    + Arrays.toString(ranges));
+        }
+        // Cannot use a Set since the list of UID ranges might contain duplicates.
+        final List<UidRange> newVpnBlockedUidRanges = new ArrayList(mVpnBlockedUidRanges);
+        for (int i = 0; i < ranges.length; i++) {
+            if (requireVpn) {
+                newVpnBlockedUidRanges.add(ranges[i]);
+            } else {
+                newVpnBlockedUidRanges.remove(ranges[i]);
+            }
+        }
+
+        try {
+            mNetd.networkRejectNonSecureVpn(requireVpn, toUidRangeStableParcels(ranges));
+        } catch (RemoteException | ServiceSpecificException e) {
+            Log.e(TAG, "setRequireVpnForUids(" + requireVpn + ", "
+                    + Arrays.toString(ranges) + "): netd command failed: " + e);
+        }
+
+        for (final NetworkAgentInfo nai : mNetworkAgentInfos) {
+            final boolean curMetered = nai.networkCapabilities.isMetered();
+            maybeNotifyNetworkBlocked(nai, curMetered, curMetered, mRestrictBackground,
+                    mRestrictBackground, mVpnBlockedUidRanges, newVpnBlockedUidRanges);
+        }
+
+        mVpnBlockedUidRanges = newVpnBlockedUidRanges;
+    }
+
     @Override
     public boolean updateLockdownVpn() {
         if (mDeps.getCallingUid() != Process.SYSTEM_UID) {
@@ -5081,101 +5132,6 @@
     }
 
     @Override
-    public int checkMobileProvisioning(int suggestedTimeOutMs) {
-        // TODO: Remove?  Any reason to trigger a provisioning check?
-        return -1;
-    }
-
-    /** Location to an updatable file listing carrier provisioning urls.
-     *  An example:
-     *
-     * <?xml version="1.0" encoding="utf-8"?>
-     *  <provisioningUrls>
-     *   <provisioningUrl mcc="310" mnc="4">http://myserver.com/foo?mdn=%3$s&amp;iccid=%1$s&amp;imei=%2$s</provisioningUrl>
-     *  </provisioningUrls>
-     */
-    private static final String PROVISIONING_URL_PATH =
-            "/data/misc/radio/provisioning_urls.xml";
-    private final File mProvisioningUrlFile = new File(PROVISIONING_URL_PATH);
-
-    /** XML tag for root element. */
-    private static final String TAG_PROVISIONING_URLS = "provisioningUrls";
-    /** XML tag for individual url */
-    private static final String TAG_PROVISIONING_URL = "provisioningUrl";
-    /** XML attribute for mcc */
-    private static final String ATTR_MCC = "mcc";
-    /** XML attribute for mnc */
-    private static final String ATTR_MNC = "mnc";
-
-    private String getProvisioningUrlBaseFromFile() {
-        XmlPullParser parser;
-        Configuration config = mContext.getResources().getConfiguration();
-
-        try (FileReader fileReader = new FileReader(mProvisioningUrlFile)) {
-            parser = Xml.newPullParser();
-            parser.setInput(fileReader);
-            XmlUtils.beginDocument(parser, TAG_PROVISIONING_URLS);
-
-            while (true) {
-                XmlUtils.nextElement(parser);
-
-                String element = parser.getName();
-                if (element == null) break;
-
-                if (element.equals(TAG_PROVISIONING_URL)) {
-                    String mcc = parser.getAttributeValue(null, ATTR_MCC);
-                    try {
-                        if (mcc != null && Integer.parseInt(mcc) == config.mcc) {
-                            String mnc = parser.getAttributeValue(null, ATTR_MNC);
-                            if (mnc != null && Integer.parseInt(mnc) == config.mnc) {
-                                parser.next();
-                                if (parser.getEventType() == XmlPullParser.TEXT) {
-                                    return parser.getText();
-                                }
-                            }
-                        }
-                    } catch (NumberFormatException e) {
-                        loge("NumberFormatException in getProvisioningUrlBaseFromFile: " + e);
-                    }
-                }
-            }
-            return null;
-        } catch (FileNotFoundException e) {
-            loge("Carrier Provisioning Urls file not found");
-        } catch (XmlPullParserException e) {
-            loge("Xml parser exception reading Carrier Provisioning Urls file: " + e);
-        } catch (IOException e) {
-            loge("I/O exception reading Carrier Provisioning Urls file: " + e);
-        }
-        return null;
-    }
-
-    @Override
-    public String getMobileProvisioningUrl() {
-        enforceSettingsPermission();
-        String url = getProvisioningUrlBaseFromFile();
-        if (TextUtils.isEmpty(url)) {
-            url = mContext.getResources().getString(R.string.mobile_provisioning_url);
-            log("getMobileProvisioningUrl: mobile_provisioining_url from resource =" + url);
-        } else {
-            log("getMobileProvisioningUrl: mobile_provisioning_url from File =" + url);
-        }
-        // populate the iccid, imei and phone number in the provisioning url.
-        if (!TextUtils.isEmpty(url)) {
-            String phoneNumber = mTelephonyManager.getLine1Number();
-            if (TextUtils.isEmpty(phoneNumber)) {
-                phoneNumber = "0000000000";
-            }
-            url = String.format(url,
-                    mTelephonyManager.getSimSerialNumber() /* ICCID */,
-                    mTelephonyManager.getDeviceId() /* IMEI */,
-                    phoneNumber /* Phone number */);
-        }
-
-        return url;
-    }
-
-    @Override
     public void setProvisioningNotificationVisible(boolean visible, int networkType,
             String action) {
         enforceSettingsPermission();
@@ -5591,7 +5547,7 @@
         mAppOpsManager.checkPackage(callerUid, callerPackageName);
     }
 
-    private ArrayList<Integer> getSignalStrengthThresholds(@NonNull final NetworkAgentInfo nai) {
+    private int[] getSignalStrengthThresholds(@NonNull final NetworkAgentInfo nai) {
         final SortedSet<Integer> thresholds = new TreeSet<>();
         synchronized (nai) {
             for (final NetworkRequestInfo nri : mNetworkRequests.values()) {
@@ -5603,14 +5559,13 @@
                 }
             }
         }
-        return new ArrayList<>(thresholds);
+        // TODO: use NetworkStackUtils.convertToIntArray after moving it
+        return ArrayUtils.convertToIntArray(new ArrayList<>(thresholds));
     }
 
     private void updateSignalStrengthThresholds(
             NetworkAgentInfo nai, String reason, NetworkRequest request) {
-        ArrayList<Integer> thresholdsArray = getSignalStrengthThresholds(nai);
-        Bundle thresholds = new Bundle();
-        thresholds.putIntegerArrayList("thresholds", thresholdsArray);
+        final int[] thresholdsArray = getSignalStrengthThresholds(nai);
 
         if (VDBG || (DBG && !"CONNECT".equals(reason))) {
             String detail;
@@ -5620,12 +5575,10 @@
                 detail = reason;
             }
             log(String.format("updateSignalStrengthThresholds: %s, sending %s to %s",
-                    detail, Arrays.toString(thresholdsArray.toArray()), nai.toShortString()));
+                    detail, Arrays.toString(thresholdsArray), nai.toShortString()));
         }
 
-        nai.asyncChannel.sendMessage(
-                android.net.NetworkAgent.CMD_SET_SIGNAL_STRENGTH_THRESHOLDS,
-                0, 0, thresholds);
+        nai.onSignalStrengthThresholdsUpdated(thresholdsArray);
     }
 
     private void ensureValidNetworkSpecifier(NetworkCapabilities nc) {
@@ -5735,7 +5688,7 @@
             nai = mNetworkForNetId.get(network.getNetId());
         }
         if (nai != null) {
-            nai.asyncChannel.sendMessage(android.net.NetworkAgent.CMD_REQUEST_BANDWIDTH_UPDATE);
+            nai.onBandwidthUpdateRequested();
             synchronized (mBandwidthRequests) {
                 final int uid = mDeps.getCallingUid();
                 Integer uidReqs = mBandwidthRequests.get(uid);
@@ -5978,7 +5931,13 @@
     // NetworkAgentInfo keyed off its connecting messenger
     // TODO - eval if we can reduce the number of lists/hashmaps/sparsearrays
     // NOTE: Only should be accessed on ConnectivityServiceThread, except dump().
-    private final HashMap<Messenger, NetworkAgentInfo> mNetworkAgentInfos = new HashMap<>();
+    private final ArraySet<NetworkAgentInfo> mNetworkAgentInfos = new ArraySet<>();
+
+    // UID ranges for users that are currently blocked by VPNs.
+    // This array is accessed and iterated on multiple threads without holding locks, so its
+    // contents must never be mutated. When the ranges change, the array is replaced with a new one
+    // (on the handler thread).
+    private volatile List<UidRange> mVpnBlockedUidRanges = new ArrayList<>();
 
     @GuardedBy("mBlockedAppUids")
     private final HashSet<Integer> mBlockedAppUids = new HashSet<>();
@@ -6026,17 +5985,17 @@
     /**
      * Register a new agent. {@see #registerNetworkAgent} below.
      */
-    public Network registerNetworkAgent(Messenger messenger, NetworkInfo networkInfo,
+    public Network registerNetworkAgent(INetworkAgent na, NetworkInfo networkInfo,
             LinkProperties linkProperties, NetworkCapabilities networkCapabilities,
             int currentScore, NetworkAgentConfig networkAgentConfig) {
-        return registerNetworkAgent(messenger, networkInfo, linkProperties, networkCapabilities,
+        return registerNetworkAgent(na, networkInfo, linkProperties, networkCapabilities,
                 currentScore, networkAgentConfig, NetworkProvider.ID_NONE);
     }
 
     /**
      * Register a new agent with ConnectivityService to handle a network.
      *
-     * @param messenger a messenger for ConnectivityService to contact the agent asynchronously.
+     * @param na a reference for ConnectivityService to contact the agent asynchronously.
      * @param networkInfo the initial info associated with this network. It can be updated later :
      *         see {@link #updateNetworkInfo}.
      * @param linkProperties the initial link properties of this network. They can be updated
@@ -6049,7 +6008,7 @@
      * @param providerId the ID of the provider owning this NetworkAgent.
      * @return the network created for this agent.
      */
-    public Network registerNetworkAgent(Messenger messenger, NetworkInfo networkInfo,
+    public Network registerNetworkAgent(INetworkAgent na, NetworkInfo networkInfo,
             LinkProperties linkProperties, NetworkCapabilities networkCapabilities,
             int currentScore, NetworkAgentConfig networkAgentConfig, int providerId) {
         if (networkCapabilities.hasTransport(TRANSPORT_TEST)) {
@@ -6061,14 +6020,14 @@
         final int uid = mDeps.getCallingUid();
         final long token = Binder.clearCallingIdentity();
         try {
-            return registerNetworkAgentInternal(messenger, networkInfo, linkProperties,
+            return registerNetworkAgentInternal(na, networkInfo, linkProperties,
                     networkCapabilities, currentScore, networkAgentConfig, providerId, uid);
         } finally {
             Binder.restoreCallingIdentity(token);
         }
     }
 
-    private Network registerNetworkAgentInternal(Messenger messenger, NetworkInfo networkInfo,
+    private Network registerNetworkAgentInternal(INetworkAgent na, NetworkInfo networkInfo,
             LinkProperties linkProperties, NetworkCapabilities networkCapabilities,
             int currentScore, NetworkAgentConfig networkAgentConfig, int providerId, int uid) {
         if (networkCapabilities.hasTransport(TRANSPORT_TEST)) {
@@ -6084,7 +6043,7 @@
         // TODO: Instead of passing mDefaultRequest, provide an API to determine whether a Network
         // satisfies mDefaultRequest.
         final NetworkCapabilities nc = new NetworkCapabilities(networkCapabilities);
-        final NetworkAgentInfo nai = new NetworkAgentInfo(messenger, new AsyncChannel(),
+        final NetworkAgentInfo nai = new NetworkAgentInfo(na,
                 new Network(mNetIdManager.reserveNetId()), new NetworkInfo(networkInfo), lp, nc,
                 currentScore, mContext, mTrackerHandler, new NetworkAgentConfig(networkAgentConfig),
                 this, mNetd, mDnsResolver, mNMS, providerId, uid);
@@ -6102,7 +6061,7 @@
                 nai.network, name, new NetworkMonitorCallbacks(nai));
         // NetworkAgentInfo registration will finish when the NetworkMonitor is created.
         // If the network disconnects or sends any other event before that, messages are deferred by
-        // NetworkAgent until nai.asyncChannel.connect(), which will be called when finalizing the
+        // NetworkAgent until nai.connect(), which will be called when finalizing the
         // registration.
         return nai.network;
     }
@@ -6110,7 +6069,7 @@
     private void handleRegisterNetworkAgent(NetworkAgentInfo nai, INetworkMonitor networkMonitor) {
         nai.onNetworkMonitorCreated(networkMonitor);
         if (VDBG) log("Got NetworkAgent Messenger");
-        mNetworkAgentInfos.put(nai.messenger, nai);
+        mNetworkAgentInfos.add(nai);
         synchronized (mNetworkForNetId) {
             mNetworkForNetId.put(nai.network.getNetId(), nai);
         }
@@ -6120,7 +6079,7 @@
         } catch (RemoteException e) {
             e.rethrowAsRuntimeException();
         }
-        nai.asyncChannel.connect(mContext, mTrackerHandler, nai.messenger);
+        nai.notifyRegistered();
         NetworkInfo networkInfo = nai.networkInfo;
         updateNetworkInfo(nai, networkInfo);
         updateUids(nai, null, nai.networkCapabilities);
@@ -6634,7 +6593,7 @@
 
             if (meteredChanged) {
                 maybeNotifyNetworkBlocked(nai, oldMetered, newMetered, mRestrictBackground,
-                        mRestrictBackground);
+                        mRestrictBackground, mVpnBlockedUidRanges, mVpnBlockedUidRanges);
             }
 
             final boolean roamingChanged = prevNc.hasCapability(NET_CAPABILITY_NOT_ROAMING) !=
@@ -6699,6 +6658,15 @@
         return stableRanges;
     }
 
+    private static UidRangeParcel[] toUidRangeStableParcels(UidRange[] ranges) {
+        final UidRangeParcel[] stableRanges = new UidRangeParcel[ranges.length];
+        for (int i = 0; i < ranges.length; i++) {
+            stableRanges[i] = new UidRangeParcel(ranges[i].start, ranges[i].stop);
+        }
+        return stableRanges;
+    }
+
+
     private void updateUids(NetworkAgentInfo nai, NetworkCapabilities prevNc,
             NetworkCapabilities newNc) {
         Set<UidRange> prevRanges = null == prevNc ? null : prevNc.getUids();
@@ -6933,7 +6901,7 @@
                 break;
             }
         }
-        nai.asyncChannel.disconnect();
+        nai.disconnect();
     }
 
     private void handleLingerComplete(NetworkAgentInfo oldNetwork) {
@@ -7123,7 +7091,7 @@
 
         // Gather the list of all relevant agents and sort them by score.
         final ArrayList<NetworkAgentInfo> nais = new ArrayList<>();
-        for (final NetworkAgentInfo nai : mNetworkAgentInfos.values()) {
+        for (final NetworkAgentInfo nai : mNetworkAgentInfos) {
             if (!nai.everConnected) continue;
             nais.add(nai);
         }
@@ -7158,7 +7126,7 @@
 
     private void applyNetworkReassignment(@NonNull final NetworkReassignment changes,
             final long now) {
-        final Collection<NetworkAgentInfo> nais = mNetworkAgentInfos.values();
+        final Collection<NetworkAgentInfo> nais = mNetworkAgentInfos;
 
         // Since most of the time there are only 0 or 1 background networks, it would probably
         // be more efficient to just use an ArrayList here. TODO : measure performance
@@ -7251,7 +7219,7 @@
         updateLegacyTypeTrackerAndVpnLockdownForRematch(oldDefaultNetwork, newDefaultNetwork, nais);
 
         // Tear down all unneeded networks.
-        for (NetworkAgentInfo nai : mNetworkAgentInfos.values()) {
+        for (NetworkAgentInfo nai : mNetworkAgentInfos) {
             if (unneeded(nai, UnneededFor.TEARDOWN)) {
                 if (nai.getLingerExpiry() > 0) {
                     // This network has active linger timers and no requests, but is not
@@ -7488,7 +7456,7 @@
             // This has to happen after matching the requests, because callbacks are just requests.
             notifyNetworkCallbacks(networkAgent, ConnectivityManager.CALLBACK_PRECHECK);
         } else if (state == NetworkInfo.State.DISCONNECTED) {
-            networkAgent.asyncChannel.disconnect();
+            networkAgent.disconnect();
             if (networkAgent.isVPN()) {
                 updateUids(networkAgent, networkAgent.networkCapabilities, null);
             }
@@ -7526,7 +7494,9 @@
         }
 
         final boolean metered = nai.networkCapabilities.isMetered();
-        final boolean blocked = isUidNetworkingWithVpnBlocked(nri.mUid, mUidRules.get(nri.mUid),
+        boolean blocked;
+        blocked = isUidBlockedByVpn(nri.mUid, mVpnBlockedUidRanges);
+        blocked |= isUidBlockedByRules(nri.mUid, mUidRules.get(nri.mUid),
                 metered, mRestrictBackground);
         callCallbackForRequest(nri, nai, ConnectivityManager.CALLBACK_AVAILABLE, blocked ? 1 : 0);
     }
@@ -7548,21 +7518,25 @@
      * @param newRestrictBackground True if data saver is enabled.
      */
     private void maybeNotifyNetworkBlocked(NetworkAgentInfo nai, boolean oldMetered,
-            boolean newMetered, boolean oldRestrictBackground, boolean newRestrictBackground) {
+            boolean newMetered, boolean oldRestrictBackground, boolean newRestrictBackground,
+            List<UidRange> oldBlockedUidRanges, List<UidRange> newBlockedUidRanges) {
 
         for (int i = 0; i < nai.numNetworkRequests(); i++) {
             NetworkRequest nr = nai.requestAt(i);
             NetworkRequestInfo nri = mNetworkRequests.get(nr);
             final int uidRules = mUidRules.get(nri.mUid);
-            final boolean oldBlocked, newBlocked;
-            // mVpns lock needs to be hold here to ensure that the active VPN cannot be changed
-            // between these two calls.
-            synchronized (mVpns) {
-                oldBlocked = isUidNetworkingWithVpnBlocked(nri.mUid, uidRules, oldMetered,
-                        oldRestrictBackground);
-                newBlocked = isUidNetworkingWithVpnBlocked(nri.mUid, uidRules, newMetered,
-                        newRestrictBackground);
-            }
+            final boolean oldBlocked, newBlocked, oldVpnBlocked, newVpnBlocked;
+
+            oldVpnBlocked = isUidBlockedByVpn(nri.mUid, oldBlockedUidRanges);
+            newVpnBlocked = (oldBlockedUidRanges != newBlockedUidRanges)
+                    ? isUidBlockedByVpn(nri.mUid, newBlockedUidRanges)
+                    : oldVpnBlocked;
+
+            oldBlocked = oldVpnBlocked || isUidBlockedByRules(nri.mUid, uidRules, oldMetered,
+                    oldRestrictBackground);
+            newBlocked = newVpnBlocked || isUidBlockedByRules(nri.mUid, uidRules, newMetered,
+                    newRestrictBackground);
+
             if (oldBlocked != newBlocked) {
                 callCallbackForRequest(nri, nai, ConnectivityManager.CALLBACK_BLK_CHANGED,
                         encodeBool(newBlocked));
@@ -7576,19 +7550,14 @@
      * @param newRules The new rules to apply.
      */
     private void maybeNotifyNetworkBlockedForNewUidRules(int uid, int newRules) {
-        for (final NetworkAgentInfo nai : mNetworkAgentInfos.values()) {
+        for (final NetworkAgentInfo nai : mNetworkAgentInfos) {
             final boolean metered = nai.networkCapabilities.isMetered();
+            final boolean vpnBlocked = isUidBlockedByVpn(uid, mVpnBlockedUidRanges);
             final boolean oldBlocked, newBlocked;
-            // TODO: Consider that doze mode or turn on/off battery saver would deliver lots of uid
-            // rules changed event. And this function actually loop through all connected nai and
-            // its requests. It seems that mVpns lock will be grabbed frequently in this case.
-            // Reduce the number of locking or optimize the use of lock are likely needed in future.
-            synchronized (mVpns) {
-                oldBlocked = isUidNetworkingWithVpnBlocked(
-                        uid, mUidRules.get(uid), metered, mRestrictBackground);
-                newBlocked = isUidNetworkingWithVpnBlocked(
-                        uid, newRules, metered, mRestrictBackground);
-            }
+            oldBlocked = vpnBlocked || isUidBlockedByRules(
+                    uid, mUidRules.get(uid), metered, mRestrictBackground);
+            newBlocked = vpnBlocked || isUidBlockedByRules(
+                    uid, newRules, metered, mRestrictBackground);
             if (oldBlocked == newBlocked) {
                 continue;
             }
@@ -7682,7 +7651,7 @@
         ensureRunningOnConnectivityServiceThread();
         ArrayList<Network> defaultNetworks = new ArrayList<>();
         NetworkAgentInfo defaultNetwork = getDefaultNetwork();
-        for (NetworkAgentInfo nai : mNetworkAgentInfos.values()) {
+        for (NetworkAgentInfo nai : mNetworkAgentInfos) {
             if (nai.everConnected && (nai == defaultNetwork || nai.isVPN())) {
                 defaultNetworks.add(nai.network);
             }
@@ -8233,6 +8202,13 @@
         final IBinder iCb = cb.asBinder();
         final NetworkRequestInfo nri = cbInfo.mRequestInfo;
 
+        // Connectivity Diagnostics are meant to be used with a single network request. It would be
+        // confusing for these networks to change when an NRI is satisfied in another layer.
+        if (nri.isMultilayerRequest()) {
+            throw new IllegalArgumentException("Connectivity Diagnostics do not support multilayer "
+                + "network requests.");
+        }
+
         // This means that the client registered the same callback multiple times. Do
         // not override the previous entry, and exit silently.
         if (mConnectivityDiagnosticsCallbacks.containsKey(iCb)) {
@@ -8259,7 +8235,8 @@
         synchronized (mNetworkForNetId) {
             for (int i = 0; i < mNetworkForNetId.size(); i++) {
                 final NetworkAgentInfo nai = mNetworkForNetId.valueAt(i);
-                if (nai.satisfies(nri.request)) {
+                // Connectivity Diagnostics rejects multilayer requests at registration hence get(0)
+                if (nai.satisfies(nri.mRequests.get(0))) {
                     matchingNetworks.add(nai);
                 }
             }
@@ -8387,7 +8364,8 @@
                 mConnectivityDiagnosticsCallbacks.entrySet()) {
             final ConnectivityDiagnosticsCallbackInfo cbInfo = entry.getValue();
             final NetworkRequestInfo nri = cbInfo.mRequestInfo;
-            if (nai.satisfies(nri.request)) {
+            // Connectivity Diagnostics rejects multilayer requests at registration hence get(0).
+            if (nai.satisfies(nri.mRequests.get(0))) {
                 if (checkConnectivityDiagnosticsPermissions(
                         nri.mPid, nri.mUid, nai, cbInfo.mCallingPackageName)) {
                     results.add(entry.getValue().mCb);
@@ -8416,7 +8394,7 @@
             return false;
         }
 
-        for (NetworkAgentInfo virtual : mNetworkAgentInfos.values()) {
+        for (NetworkAgentInfo virtual : mNetworkAgentInfos) {
             if (virtual.supportsUnderlyingNetworks()
                     && virtual.networkCapabilities.getOwnerUid() == callbackUid
                     && ArrayUtils.contains(virtual.declaredUnderlyingNetworks, nai.network)) {
diff --git a/services/core/java/com/android/server/ConnectivityServiceInitializer.java b/services/core/java/com/android/server/ConnectivityServiceInitializer.java
index f276329..0779f71 100644
--- a/services/core/java/com/android/server/ConnectivityServiceInitializer.java
+++ b/services/core/java/com/android/server/ConnectivityServiceInitializer.java
@@ -20,7 +20,6 @@
 import static android.os.IServiceManager.DUMP_FLAG_PRIORITY_NORMAL;
 
 import android.content.Context;
-import android.net.INetworkPolicyManager;
 import android.net.INetworkStatsService;
 import android.os.INetworkManagementService;
 import android.os.ServiceManager;
@@ -40,7 +39,7 @@
         System.loadLibrary("service-connectivity");
         // TODO: Define formal APIs to get the needed services.
         mConnectivity = new ConnectivityService(context, getNetworkManagementService(),
-                getNetworkStatsService(), getNetworkPolicyManager());
+                getNetworkStatsService());
     }
 
     @Override
@@ -59,10 +58,4 @@
         return INetworkStatsService.Stub.asInterface(
                 ServiceManager.getService(Context.NETWORK_STATS_SERVICE));
     }
-
-    private INetworkPolicyManager getNetworkPolicyManager() {
-        return INetworkPolicyManager.Stub.asInterface(
-                ServiceManager.getService(Context.NETWORK_POLICY_SERVICE));
-    }
-
 }
diff --git a/services/core/java/com/android/server/connectivity/KeepaliveTracker.java b/services/core/java/com/android/server/connectivity/KeepaliveTracker.java
index 96cbfde..34d9ccc 100644
--- a/services/core/java/com/android/server/connectivity/KeepaliveTracker.java
+++ b/services/core/java/com/android/server/connectivity/KeepaliveTracker.java
@@ -18,10 +18,7 @@
 
 import static android.content.pm.PackageManager.PERMISSION_GRANTED;
 import static android.net.NattSocketKeepalive.NATT_PORT;
-import static android.net.NetworkAgent.CMD_ADD_KEEPALIVE_PACKET_FILTER;
-import static android.net.NetworkAgent.CMD_REMOVE_KEEPALIVE_PACKET_FILTER;
 import static android.net.NetworkAgent.CMD_START_SOCKET_KEEPALIVE;
-import static android.net.NetworkAgent.CMD_STOP_SOCKET_KEEPALIVE;
 import static android.net.SocketKeepalive.BINDER_DIED;
 import static android.net.SocketKeepalive.DATA_RECEIVED;
 import static android.net.SocketKeepalive.ERROR_INSUFFICIENT_RESOURCES;
@@ -330,10 +327,9 @@
                 Log.d(TAG, "Starting keepalive " + mSlot + " on " + mNai.toShortString());
                 switch (mType) {
                     case TYPE_NATT:
-                        mNai.asyncChannel.sendMessage(
-                                CMD_ADD_KEEPALIVE_PACKET_FILTER, slot, 0 /* Unused */, mPacket);
-                        mNai.asyncChannel
-                                .sendMessage(CMD_START_SOCKET_KEEPALIVE, slot, mInterval, mPacket);
+                        final NattKeepalivePacketData nattData = (NattKeepalivePacketData) mPacket;
+                        mNai.onAddNattKeepalivePacketFilter(slot, nattData);
+                        mNai.onStartNattSocketKeepalive(slot, mInterval, nattData);
                         break;
                     case TYPE_TCP:
                         try {
@@ -342,11 +338,10 @@
                             handleStopKeepalive(mNai, mSlot, ERROR_INVALID_SOCKET);
                             return;
                         }
-                        mNai.asyncChannel.sendMessage(
-                                CMD_ADD_KEEPALIVE_PACKET_FILTER, slot, 0 /* Unused */, mPacket);
+                        final TcpKeepalivePacketData tcpData = (TcpKeepalivePacketData) mPacket;
+                        mNai.onAddTcpKeepalivePacketFilter(slot, tcpData);
                         // TODO: check result from apf and notify of failure as needed.
-                        mNai.asyncChannel
-                                .sendMessage(CMD_START_SOCKET_KEEPALIVE, slot, mInterval, mPacket);
+                        mNai.onStartTcpSocketKeepalive(slot, mInterval, tcpData);
                         break;
                     default:
                         Log.wtf(TAG, "Starting keepalive with unknown type: " + mType);
@@ -394,9 +389,8 @@
                             mTcpController.stopSocketMonitor(mSlot);
                             // fall through
                         case TYPE_NATT:
-                            mNai.asyncChannel.sendMessage(CMD_STOP_SOCKET_KEEPALIVE, mSlot);
-                            mNai.asyncChannel.sendMessage(CMD_REMOVE_KEEPALIVE_PACKET_FILTER,
-                                    mSlot);
+                            mNai.onStopSocketKeepalive(mSlot);
+                            mNai.onRemoveKeepalivePacketFilter(mSlot);
                             break;
                         default:
                             Log.wtf(TAG, "Stopping keepalive with unknown type: " + mType);
@@ -548,17 +542,13 @@
     }
 
     /** Handle keepalive events from lower layer. */
-    public void handleEventSocketKeepalive(@NonNull NetworkAgentInfo nai,
-            @NonNull Message message) {
-        int slot = message.arg1;
-        int reason = message.arg2;
-
+    public void handleEventSocketKeepalive(@NonNull NetworkAgentInfo nai, int slot, int reason) {
         KeepaliveInfo ki = null;
         try {
             ki = mKeepalives.get(nai).get(slot);
         } catch(NullPointerException e) {}
         if (ki == null) {
-            Log.e(TAG, "Event " + message.what + "," + slot + "," + reason
+            Log.e(TAG, "Event " + NetworkAgent.EVENT_SOCKET_KEEPALIVE + "," + slot + "," + reason
                     + " for unknown keepalive " + slot + " on " + nai.toShortString());
             return;
         }
@@ -601,7 +591,7 @@
             ki.mStartedState = KeepaliveInfo.NOT_STARTED;
             cleanupStoppedKeepalive(nai, slot);
         } else {
-            Log.wtf(TAG, "Event " + message.what + "," + slot + "," + reason
+            Log.wtf(TAG, "Event " + NetworkAgent.EVENT_SOCKET_KEEPALIVE + "," + slot + "," + reason
                     + " for keepalive in wrong state: " + ki.toString());
         }
     }
diff --git a/services/core/java/com/android/server/connectivity/NetworkAgentInfo.java b/services/core/java/com/android/server/connectivity/NetworkAgentInfo.java
index 52b9f5c..7bde4d5 100644
--- a/services/core/java/com/android/server/connectivity/NetworkAgentInfo.java
+++ b/services/core/java/com/android/server/connectivity/NetworkAgentInfo.java
@@ -27,25 +27,35 @@
 import android.net.INetd;
 import android.net.INetworkMonitor;
 import android.net.LinkProperties;
+import android.net.NattKeepalivePacketData;
 import android.net.Network;
+import android.net.NetworkAgent;
 import android.net.NetworkAgentConfig;
 import android.net.NetworkCapabilities;
 import android.net.NetworkInfo;
 import android.net.NetworkMonitorManager;
 import android.net.NetworkRequest;
 import android.net.NetworkState;
+import android.net.TcpKeepalivePacketData;
+import android.os.Bundle;
 import android.os.Handler;
+import android.os.IBinder;
 import android.os.INetworkManagementService;
-import android.os.Messenger;
+import android.os.RemoteException;
 import android.os.SystemClock;
 import android.util.Log;
+import android.util.Pair;
 import android.util.SparseArray;
 
-import com.android.internal.util.AsyncChannel;
+import com.android.connectivity.aidl.INetworkAgent;
+import com.android.connectivity.aidl.INetworkAgentRegistry;
 import com.android.internal.util.WakeupMessage;
 import com.android.server.ConnectivityService;
 
 import java.io.PrintWriter;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.NoSuchElementException;
 import java.util.Objects;
 import java.util.SortedSet;
 import java.util.TreeSet;
@@ -136,7 +146,11 @@
     // Underlying networks declared by the agent. Only set if supportsUnderlyingNetworks is true.
     // The networks in this list might be declared by a VPN app using setUnderlyingNetworks and are
     // not guaranteed to be current or correct, or even to exist.
-    public @Nullable Network[] declaredUnderlyingNetworks;
+    //
+    // This array is read and iterated on multiple threads with no locking so its contents must
+    // never be modified. When the list of networks changes, replace with a new array, on the
+    // handler thread.
+    public @Nullable volatile Network[] declaredUnderlyingNetworks;
 
     // The capabilities originally announced by the NetworkAgent, regardless of any capabilities
     // that were added or removed due to this network's underlying networks.
@@ -221,6 +235,31 @@
      */
     public static final int EVENT_NETWORK_LINGER_COMPLETE = 1001;
 
+    /**
+     * Inform ConnectivityService that the agent is half-connected.
+     * arg1 = ARG_AGENT_SUCCESS or ARG_AGENT_FAILURE
+     * obj = NetworkAgentInfo
+     * @hide
+     */
+    public static final int EVENT_AGENT_REGISTERED = 1002;
+
+    /**
+     * Inform ConnectivityService that the agent was disconnected.
+     * obj = NetworkAgentInfo
+     * @hide
+     */
+    public static final int EVENT_AGENT_DISCONNECTED = 1003;
+
+    /**
+     * Argument for EVENT_AGENT_HALF_CONNECTED indicating failure.
+     */
+    public static final int ARG_AGENT_FAILURE = 0;
+
+    /**
+     * Argument for EVENT_AGENT_HALF_CONNECTED indicating success.
+     */
+    public static final int ARG_AGENT_SUCCESS = 1;
+
     // All linger timers for this network, sorted by expiry time. A linger timer is added whenever
     // a request is moved to a network with a better score, regardless of whether the network is or
     // was lingering or not.
@@ -262,8 +301,9 @@
     // report is generated. Once non-null, it will never be null again.
     @Nullable private ConnectivityReport mConnectivityReport;
 
-    public final Messenger messenger;
-    public final AsyncChannel asyncChannel;
+    public final INetworkAgent networkAgent;
+    // Only accessed from ConnectivityService handler thread
+    private final AgentDeathMonitor mDeathMonitor = new AgentDeathMonitor();
 
     public final int factorySerialNumber;
 
@@ -279,13 +319,12 @@
     private final Context mContext;
     private final Handler mHandler;
 
-    public NetworkAgentInfo(Messenger messenger, AsyncChannel ac, Network net, NetworkInfo info,
+    public NetworkAgentInfo(INetworkAgent na, Network net, NetworkInfo info,
             LinkProperties lp, NetworkCapabilities nc, int score, Context context,
             Handler handler, NetworkAgentConfig config, ConnectivityService connService, INetd netd,
             IDnsResolver dnsResolver, INetworkManagementService nms, int factorySerialNumber,
             int creatorUid) {
-        this.messenger = messenger;
-        asyncChannel = ac;
+        networkAgent = na;
         network = net;
         networkInfo = info;
         linkProperties = lp;
@@ -300,6 +339,249 @@
         this.creatorUid = creatorUid;
     }
 
+    private class AgentDeathMonitor implements IBinder.DeathRecipient {
+        @Override
+        public void binderDied() {
+            notifyDisconnected();
+        }
+    }
+
+    /**
+     * Notify the NetworkAgent that it was registered, and should be unregistered if it dies.
+     *
+     * Must be called from the ConnectivityService handler thread. A NetworkAgent can only be
+     * registered once.
+     */
+    public void notifyRegistered() {
+        try {
+            networkAgent.asBinder().linkToDeath(mDeathMonitor, 0);
+            networkAgent.onRegistered(new NetworkAgentMessageHandler(mHandler));
+        } catch (RemoteException e) {
+            Log.e(TAG, "Error registering NetworkAgent", e);
+            maybeUnlinkDeathMonitor();
+            mHandler.obtainMessage(EVENT_AGENT_REGISTERED, ARG_AGENT_FAILURE, 0, this)
+                    .sendToTarget();
+            return;
+        }
+
+        mHandler.obtainMessage(EVENT_AGENT_REGISTERED, ARG_AGENT_SUCCESS, 0, this).sendToTarget();
+    }
+
+    /**
+     * Disconnect the NetworkAgent. Must be called from the ConnectivityService handler thread.
+     */
+    public void disconnect() {
+        try {
+            networkAgent.onDisconnected();
+        } catch (RemoteException e) {
+            Log.i(TAG, "Error disconnecting NetworkAgent", e);
+            // Fall through: it's fine if the remote has died
+        }
+
+        notifyDisconnected();
+        maybeUnlinkDeathMonitor();
+    }
+
+    private void maybeUnlinkDeathMonitor() {
+        try {
+            networkAgent.asBinder().unlinkToDeath(mDeathMonitor, 0);
+        } catch (NoSuchElementException e) {
+            // Was not linked: ignore
+        }
+    }
+
+    private void notifyDisconnected() {
+        // Note this may be called multiple times if ConnectivityService disconnects while the
+        // NetworkAgent also dies. ConnectivityService ignores disconnects of already disconnected
+        // agents.
+        mHandler.obtainMessage(EVENT_AGENT_DISCONNECTED, this).sendToTarget();
+    }
+
+    /**
+     * Notify the NetworkAgent that bandwidth update was requested.
+     */
+    public void onBandwidthUpdateRequested() {
+        try {
+            networkAgent.onBandwidthUpdateRequested();
+        } catch (RemoteException e) {
+            Log.e(TAG, "Error sending bandwidth update request event", e);
+        }
+    }
+
+    /**
+     * Notify the NetworkAgent that validation status has changed.
+     */
+    public void onValidationStatusChanged(int validationStatus, @Nullable String captivePortalUrl) {
+        try {
+            networkAgent.onValidationStatusChanged(validationStatus, captivePortalUrl);
+        } catch (RemoteException e) {
+            Log.e(TAG, "Error sending validation status change event", e);
+        }
+    }
+
+    /**
+     * Notify the NetworkAgent that the acceptUnvalidated setting should be saved.
+     */
+    public void onSaveAcceptUnvalidated(boolean acceptUnvalidated) {
+        try {
+            networkAgent.onSaveAcceptUnvalidated(acceptUnvalidated);
+        } catch (RemoteException e) {
+            Log.e(TAG, "Error sending accept unvalidated event", e);
+        }
+    }
+
+    /**
+     * Notify the NetworkAgent that NATT socket keepalive should be started.
+     */
+    public void onStartNattSocketKeepalive(int slot, int intervalDurationMs,
+            @NonNull NattKeepalivePacketData packetData) {
+        try {
+            networkAgent.onStartNattSocketKeepalive(slot, intervalDurationMs, packetData);
+        } catch (RemoteException e) {
+            Log.e(TAG, "Error sending NATT socket keepalive start event", e);
+        }
+    }
+
+    /**
+     * Notify the NetworkAgent that TCP socket keepalive should be started.
+     */
+    public void onStartTcpSocketKeepalive(int slot, int intervalDurationMs,
+            @NonNull TcpKeepalivePacketData packetData) {
+        try {
+            networkAgent.onStartTcpSocketKeepalive(slot, intervalDurationMs, packetData);
+        } catch (RemoteException e) {
+            Log.e(TAG, "Error sending TCP socket keepalive start event", e);
+        }
+    }
+
+    /**
+     * Notify the NetworkAgent that socket keepalive should be stopped.
+     */
+    public void onStopSocketKeepalive(int slot) {
+        try {
+            networkAgent.onStopSocketKeepalive(slot);
+        } catch (RemoteException e) {
+            Log.e(TAG, "Error sending TCP socket keepalive stop event", e);
+        }
+    }
+
+    /**
+     * Notify the NetworkAgent that signal strength thresholds should be updated.
+     */
+    public void onSignalStrengthThresholdsUpdated(@NonNull int[] thresholds) {
+        try {
+            networkAgent.onSignalStrengthThresholdsUpdated(thresholds);
+        } catch (RemoteException e) {
+            Log.e(TAG, "Error sending signal strength thresholds event", e);
+        }
+    }
+
+    /**
+     * Notify the NetworkAgent that automatic reconnect should be prevented.
+     */
+    public void onPreventAutomaticReconnect() {
+        try {
+            networkAgent.onPreventAutomaticReconnect();
+        } catch (RemoteException e) {
+            Log.e(TAG, "Error sending prevent automatic reconnect event", e);
+        }
+    }
+
+    /**
+     * Notify the NetworkAgent that a NATT keepalive packet filter should be added.
+     */
+    public void onAddNattKeepalivePacketFilter(int slot,
+            @NonNull NattKeepalivePacketData packetData) {
+        try {
+            networkAgent.onAddNattKeepalivePacketFilter(slot, packetData);
+        } catch (RemoteException e) {
+            Log.e(TAG, "Error sending add NATT keepalive packet filter event", e);
+        }
+    }
+
+    /**
+     * Notify the NetworkAgent that a TCP keepalive packet filter should be added.
+     */
+    public void onAddTcpKeepalivePacketFilter(int slot,
+            @NonNull TcpKeepalivePacketData packetData) {
+        try {
+            networkAgent.onAddTcpKeepalivePacketFilter(slot, packetData);
+        } catch (RemoteException e) {
+            Log.e(TAG, "Error sending add TCP keepalive packet filter event", e);
+        }
+    }
+
+    /**
+     * Notify the NetworkAgent that a keepalive packet filter should be removed.
+     */
+    public void onRemoveKeepalivePacketFilter(int slot) {
+        try {
+            networkAgent.onRemoveKeepalivePacketFilter(slot);
+        } catch (RemoteException e) {
+            Log.e(TAG, "Error sending remove keepalive packet filter event", e);
+        }
+    }
+
+    // TODO: consider moving out of NetworkAgentInfo into its own class
+    private class NetworkAgentMessageHandler extends INetworkAgentRegistry.Stub {
+        private final Handler mHandler;
+
+        private NetworkAgentMessageHandler(Handler handler) {
+            mHandler = handler;
+        }
+
+        @Override
+        public void sendNetworkCapabilities(NetworkCapabilities nc) {
+            mHandler.obtainMessage(NetworkAgent.EVENT_NETWORK_CAPABILITIES_CHANGED,
+                    new Pair<>(NetworkAgentInfo.this, nc)).sendToTarget();
+        }
+
+        @Override
+        public void sendLinkProperties(LinkProperties lp) {
+            mHandler.obtainMessage(NetworkAgent.EVENT_NETWORK_PROPERTIES_CHANGED,
+                    new Pair<>(NetworkAgentInfo.this, lp)).sendToTarget();
+        }
+
+        @Override
+        public void sendNetworkInfo(NetworkInfo info) {
+            mHandler.obtainMessage(NetworkAgent.EVENT_NETWORK_INFO_CHANGED,
+                    new Pair<>(NetworkAgentInfo.this, info)).sendToTarget();
+        }
+
+        @Override
+        public void sendScore(int score) {
+            mHandler.obtainMessage(NetworkAgent.EVENT_NETWORK_SCORE_CHANGED, score, 0,
+                    new Pair<>(NetworkAgentInfo.this, null)).sendToTarget();
+        }
+
+        @Override
+        public void sendExplicitlySelected(boolean explicitlySelected, boolean acceptPartial) {
+            mHandler.obtainMessage(NetworkAgent.EVENT_SET_EXPLICITLY_SELECTED,
+                    explicitlySelected ? 1 : 0, acceptPartial ? 1 : 0,
+                    new Pair<>(NetworkAgentInfo.this, null)).sendToTarget();
+        }
+
+        @Override
+        public void sendSocketKeepaliveEvent(int slot, int reason) {
+            mHandler.obtainMessage(NetworkAgent.EVENT_SOCKET_KEEPALIVE,
+                    slot, reason, new Pair<>(NetworkAgentInfo.this, null)).sendToTarget();
+        }
+
+        @Override
+        public void sendUnderlyingNetworks(@Nullable List<Network> networks) {
+            final Bundle args = new Bundle();
+            if (networks instanceof ArrayList<?>) {
+                args.putParcelableArrayList(NetworkAgent.UNDERLYING_NETWORKS_KEY,
+                        (ArrayList<Network>) networks);
+            } else {
+                args.putParcelableArrayList(NetworkAgent.UNDERLYING_NETWORKS_KEY,
+                        networks == null ? null : new ArrayList<>(networks));
+            }
+            mHandler.obtainMessage(NetworkAgent.EVENT_UNDERLYING_NETWORKS_CHANGED,
+                    new Pair<>(NetworkAgentInfo.this, args)).sendToTarget();
+        }
+    }
+
     /**
      * Inform NetworkAgentInfo that a new NetworkMonitor was created.
      */
diff --git a/tests/net/integration/src/com/android/server/net/integrationtests/ConnectivityServiceIntegrationTest.kt b/tests/net/integration/src/com/android/server/net/integrationtests/ConnectivityServiceIntegrationTest.kt
index 70f6386..8e18751 100644
--- a/tests/net/integration/src/com/android/server/net/integrationtests/ConnectivityServiceIntegrationTest.kt
+++ b/tests/net/integration/src/com/android/server/net/integrationtests/ConnectivityServiceIntegrationTest.kt
@@ -25,7 +25,6 @@
 import android.net.ConnectivityManager
 import android.net.IDnsResolver
 import android.net.INetd
-import android.net.INetworkPolicyManager
 import android.net.INetworkStatsService
 import android.net.LinkProperties
 import android.net.NetworkCapabilities.NET_CAPABILITY_CAPTIVE_PORTAL
@@ -88,8 +87,6 @@
     @Mock
     private lateinit var statsService: INetworkStatsService
     @Mock
-    private lateinit var policyManager: INetworkPolicyManager
-    @Mock
     private lateinit var log: IpConnectivityLog
     @Mock
     private lateinit var netd: INetd
@@ -171,7 +168,7 @@
     }
 
     private inner class TestConnectivityService(deps: Dependencies) : ConnectivityService(
-            context, netManager, statsService, policyManager, dnsResolver, log, netd, deps)
+            context, netManager, statsService, dnsResolver, log, netd, deps)
 
     private fun makeDependencies(): ConnectivityService.Dependencies {
         val deps = spy(ConnectivityService.Dependencies())
diff --git a/tests/net/java/com/android/server/ConnectivityServiceTest.java b/tests/net/java/com/android/server/ConnectivityServiceTest.java
index a613e5e..9421acd 100644
--- a/tests/net/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/net/java/com/android/server/ConnectivityServiceTest.java
@@ -166,7 +166,6 @@
 import android.net.INetworkMonitor;
 import android.net.INetworkMonitorCallbacks;
 import android.net.INetworkPolicyListener;
-import android.net.INetworkPolicyManager;
 import android.net.INetworkStatsService;
 import android.net.InetAddresses;
 import android.net.InterfaceConfigurationParcel;
@@ -183,6 +182,7 @@
 import android.net.NetworkFactory;
 import android.net.NetworkInfo;
 import android.net.NetworkInfo.DetailedState;
+import android.net.NetworkPolicyManager;
 import android.net.NetworkRequest;
 import android.net.NetworkSpecifier;
 import android.net.NetworkStack;
@@ -299,6 +299,7 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Predicate;
 import java.util.function.Supplier;
+import java.util.stream.Collectors;
 
 import kotlin.reflect.KClass;
 
@@ -365,7 +366,6 @@
     @Mock INetworkManagementService mNetworkManagementService;
     @Mock INetworkStatsService mStatsService;
     @Mock IBatteryStats mBatteryStatsService;
-    @Mock INetworkPolicyManager mNpm;
     @Mock IDnsResolver mMockDnsResolver;
     @Mock INetd mMockNetd;
     @Mock NetworkStackClient mNetworkStack;
@@ -380,6 +380,7 @@
     @Mock TelephonyManager mTelephonyManager;
     @Mock MockableSystemProperties mSystemProperties;
     @Mock EthernetManager mEthernetManager;
+    @Mock NetworkPolicyManager mNetworkPolicyManager;
 
     private ArgumentCaptor<ResolverParamsParcel> mResolverParamsParcelCaptor =
             ArgumentCaptor.forClass(ResolverParamsParcel.class);
@@ -412,6 +413,7 @@
 
         @Spy private Resources mResources;
         private final LinkedBlockingQueue<Intent> mStartedActivities = new LinkedBlockingQueue<>();
+
         // Map of permission name -> PermissionManager.Permission_{GRANTED|DENIED} constant
         private final HashMap<String, Integer> mMockedPermissions = new HashMap<>();
 
@@ -477,6 +479,7 @@
             if (Context.APP_OPS_SERVICE.equals(name)) return mAppOpsManager;
             if (Context.TELEPHONY_SERVICE.equals(name)) return mTelephonyManager;
             if (Context.ETHERNET_SERVICE.equals(name)) return mEthernetManager;
+            if (Context.NETWORK_POLICY_SERVICE.equals(name)) return mNetworkPolicyManager;
             return super.getSystemService(name);
         }
 
@@ -1326,7 +1329,6 @@
         mService = new ConnectivityService(mServiceContext,
                 mNetworkManagementService,
                 mStatsService,
-                mNpm,
                 mMockDnsResolver,
                 mock(IpConnectivityLog.class),
                 mMockNetd,
@@ -1336,7 +1338,7 @@
 
         final ArgumentCaptor<INetworkPolicyListener> policyListenerCaptor =
                 ArgumentCaptor.forClass(INetworkPolicyListener.class);
-        verify(mNpm).registerListener(policyListenerCaptor.capture());
+        verify(mNetworkPolicyManager).registerListener(policyListenerCaptor.capture());
         mPolicyListener = policyListenerCaptor.getValue();
 
         // Create local CM before sending system ready so that we can answer
@@ -6509,6 +6511,26 @@
         checkNetworkInfo(mCm.getNetworkInfo(type), type, state);
     }
 
+    // Checks that each of the |agents| receive a blocked status change callback with the specified
+    // |blocked| value, in any order. This is needed because when an event affects multiple
+    // networks, ConnectivityService does not guarantee the order in which callbacks are fired.
+    private void assertBlockedCallbackInAnyOrder(TestNetworkCallback callback, boolean blocked,
+            TestNetworkAgentWrapper... agents) {
+        final List<Network> expectedNetworks = Arrays.asList(agents).stream()
+                .map((agent) -> agent.getNetwork())
+                .collect(Collectors.toList());
+
+        // Expect exactly one blocked callback for each agent.
+        for (int i = 0; i < agents.length; i++) {
+            CallbackEntry e = callback.expectCallbackThat(TIMEOUT_MS, (c) ->
+                    c instanceof CallbackEntry.BlockedStatus
+                            && ((CallbackEntry.BlockedStatus) c).getBlocked() == blocked);
+            Network network = e.getNetwork();
+            assertTrue("Received unexpected blocked callback for network " + network,
+                    expectedNetworks.remove(network));
+        }
+    }
+
     @Test
     public void testNetworkBlockedStatusAlwaysOnVpn() throws Exception {
         mServiceContext.setPermission(
@@ -6555,9 +6577,10 @@
         assertNetworkInfo(TYPE_WIFI, DetailedState.BLOCKED);
 
         // Disable lockdown, expect to see the network unblocked.
-        // There are no callbacks because they are not implemented yet.
         mService.setAlwaysOnVpnPackage(userId, null, false /* lockdown */, allowList);
         expectNetworkRejectNonSecureVpn(inOrder, false, firstHalf, secondHalf);
+        callback.expectBlockedStatusCallback(false, mWiFiNetworkAgent);
+        defaultCallback.expectBlockedStatusCallback(false, mWiFiNetworkAgent);
         vpnUidCallback.assertNoCallback();
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetworkForUid(VPN_UID));
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetwork());
@@ -6605,6 +6628,8 @@
         allowList.clear();
         mService.setAlwaysOnVpnPackage(userId, ALWAYS_ON_PACKAGE, true /* lockdown */, allowList);
         expectNetworkRejectNonSecureVpn(inOrder, true, firstHalf, secondHalf);
+        defaultCallback.expectBlockedStatusCallback(true, mWiFiNetworkAgent);
+        assertBlockedCallbackInAnyOrder(callback, true, mWiFiNetworkAgent, mCellNetworkAgent);
         vpnUidCallback.assertNoCallback();
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetworkForUid(VPN_UID));
         assertNull(mCm.getActiveNetwork());
@@ -6614,6 +6639,8 @@
 
         // Disable lockdown. Everything is unblocked.
         mService.setAlwaysOnVpnPackage(userId, null, false /* lockdown */, allowList);
+        defaultCallback.expectBlockedStatusCallback(false, mWiFiNetworkAgent);
+        assertBlockedCallbackInAnyOrder(callback, false, mWiFiNetworkAgent, mCellNetworkAgent);
         vpnUidCallback.assertNoCallback();
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetworkForUid(VPN_UID));
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetwork());
@@ -6647,6 +6674,8 @@
 
         // Enable lockdown and connect a VPN. The VPN is not blocked.
         mService.setAlwaysOnVpnPackage(userId, ALWAYS_ON_PACKAGE, true /* lockdown */, allowList);
+        defaultCallback.expectBlockedStatusCallback(true, mWiFiNetworkAgent);
+        assertBlockedCallbackInAnyOrder(callback, true, mWiFiNetworkAgent, mCellNetworkAgent);
         vpnUidCallback.assertNoCallback();
         assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetworkForUid(VPN_UID));
         assertNull(mCm.getActiveNetwork());
@@ -6658,7 +6687,7 @@
         defaultCallback.expectAvailableThenValidatedCallbacks(mMockVpn);
         vpnUidCallback.assertNoCallback();  // vpnUidCallback has NOT_VPN capability.
         assertEquals(mMockVpn.getNetwork(), mCm.getActiveNetwork());
-        assertEquals(null, mCm.getActiveNetworkForUid(VPN_UID));  // BUG?
+        assertEquals(mWiFiNetworkAgent.getNetwork(), mCm.getActiveNetworkForUid(VPN_UID));
         assertActiveNetworkInfo(TYPE_WIFI, DetailedState.CONNECTED);
         assertNetworkInfo(TYPE_MOBILE, DetailedState.DISCONNECTED);
         assertNetworkInfo(TYPE_VPN, DetailedState.CONNECTED);
@@ -7810,8 +7839,7 @@
     @Test
     public void testCheckConnectivityDiagnosticsPermissionsNetworkStack() throws Exception {
         final NetworkAgentInfo naiWithoutUid =
-                new NetworkAgentInfo(
-                        null, null, null, null, null, new NetworkCapabilities(), 0,
+                new NetworkAgentInfo(null, null, null, null, new NetworkCapabilities(), 0,
                         mServiceContext, null, null, mService, null, null, null, 0, INVALID_UID);
 
         mServiceContext.setPermission(
@@ -7826,8 +7854,7 @@
     @Test
     public void testCheckConnectivityDiagnosticsPermissionsWrongUidPackageName() throws Exception {
         final NetworkAgentInfo naiWithoutUid =
-                new NetworkAgentInfo(
-                        null, null, null, null, null, new NetworkCapabilities(), 0,
+                new NetworkAgentInfo(null, null, null, null, new NetworkCapabilities(), 0,
                         mServiceContext, null, null, mService, null, null, null, 0, INVALID_UID);
 
         mServiceContext.setPermission(android.Manifest.permission.NETWORK_STACK, PERMISSION_DENIED);
@@ -7842,8 +7869,7 @@
     @Test
     public void testCheckConnectivityDiagnosticsPermissionsNoLocationPermission() throws Exception {
         final NetworkAgentInfo naiWithoutUid =
-                new NetworkAgentInfo(
-                        null, null, null, null, null, new NetworkCapabilities(), 0,
+                new NetworkAgentInfo(null, null, null, null, new NetworkCapabilities(), 0,
                         mServiceContext, null, null, mService, null, null, null, 0, INVALID_UID);
 
         mServiceContext.setPermission(android.Manifest.permission.NETWORK_STACK, PERMISSION_DENIED);
@@ -7859,8 +7885,7 @@
     public void testCheckConnectivityDiagnosticsPermissionsActiveVpn() throws Exception {
         final Network network = new Network(NET_ID);
         final NetworkAgentInfo naiWithoutUid =
-                new NetworkAgentInfo(
-                        null, null, network, null, null, new NetworkCapabilities(), 0,
+                new NetworkAgentInfo(null, network, null, null, new NetworkCapabilities(), 0,
                         mServiceContext, null, null, mService, null, null, null, 0, INVALID_UID);
 
         setupLocationPermissions(Build.VERSION_CODES.Q, true, AppOpsManager.OPSTR_FINE_LOCATION,
@@ -7894,8 +7919,7 @@
         final NetworkCapabilities nc = new NetworkCapabilities();
         nc.setAdministratorUids(new int[] {Process.myUid()});
         final NetworkAgentInfo naiWithUid =
-                new NetworkAgentInfo(
-                        null, null, null, null, null, nc, 0, mServiceContext, null, null,
+                new NetworkAgentInfo(null, null, null, null, nc, 0, mServiceContext, null, null,
                         mService, null, null, null, 0, INVALID_UID);
 
         setupLocationPermissions(Build.VERSION_CODES.Q, true, AppOpsManager.OPSTR_FINE_LOCATION,
@@ -7914,8 +7938,7 @@
         nc.setOwnerUid(Process.myUid());
         nc.setAdministratorUids(new int[] {Process.myUid()});
         final NetworkAgentInfo naiWithUid =
-                new NetworkAgentInfo(
-                        null, null, null, null, null, nc, 0, mServiceContext, null, null,
+                new NetworkAgentInfo(null, null, null, null, nc, 0, mServiceContext, null, null,
                         mService, null, null, null, 0, INVALID_UID);
 
         setupLocationPermissions(Build.VERSION_CODES.Q, true, AppOpsManager.OPSTR_FINE_LOCATION,
diff --git a/tests/net/java/com/android/server/connectivity/LingerMonitorTest.java b/tests/net/java/com/android/server/connectivity/LingerMonitorTest.java
index aafa18a..96c56e3 100644
--- a/tests/net/java/com/android/server/connectivity/LingerMonitorTest.java
+++ b/tests/net/java/com/android/server/connectivity/LingerMonitorTest.java
@@ -353,7 +353,7 @@
         NetworkCapabilities caps = new NetworkCapabilities();
         caps.addCapability(0);
         caps.addTransportType(transport);
-        NetworkAgentInfo nai = new NetworkAgentInfo(null, null, new Network(netId), info, null,
+        NetworkAgentInfo nai = new NetworkAgentInfo(null, new Network(netId), info, null,
                 caps, 50, mCtx, null, null /* config */, mConnService, mNetd, mDnsResolver, mNMS,
                 NetworkProvider.ID_NONE, Binder.getCallingUid());
         nai.everValidated = true;
diff --git a/tests/net/java/com/android/server/connectivity/VpnTest.java b/tests/net/java/com/android/server/connectivity/VpnTest.java
index cc47317..3648c4d 100644
--- a/tests/net/java/com/android/server/connectivity/VpnTest.java
+++ b/tests/net/java/com/android/server/connectivity/VpnTest.java
@@ -27,7 +27,6 @@
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
-import static org.mockito.AdditionalMatchers.aryEq;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyInt;
@@ -89,6 +88,7 @@
 import android.security.KeyStore;
 import android.util.ArrayMap;
 import android.util.ArraySet;
+import android.util.Range;
 
 import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
@@ -350,7 +350,7 @@
 
         // Set always-on with lockdown.
         assertTrue(vpn.setAlwaysOnPackage(PKGS[1], true, null, mKeyStore));
-        verify(mNetd).networkRejectNonSecureVpn(eq(true), aryEq(new UidRangeParcel[] {
+        verify(mConnectivityManager).setRequireVpnForUids(true, toRanges(new UidRangeParcel[] {
                 new UidRangeParcel(user.start, user.start + PKG_UIDS[1] - 1),
                 new UidRangeParcel(user.start + PKG_UIDS[1] + 1, user.stop)
         }));
@@ -361,12 +361,11 @@
 
         // Switch to another app.
         assertTrue(vpn.setAlwaysOnPackage(PKGS[3], true, null, mKeyStore));
-        verify(mNetd).networkRejectNonSecureVpn(eq(false), aryEq(new UidRangeParcel[] {
+        verify(mConnectivityManager).setRequireVpnForUids(false, toRanges(new UidRangeParcel[] {
                 new UidRangeParcel(user.start, user.start + PKG_UIDS[1] - 1),
                 new UidRangeParcel(user.start + PKG_UIDS[1] + 1, user.stop)
         }));
-
-        verify(mNetd).networkRejectNonSecureVpn(eq(true), aryEq(new UidRangeParcel[] {
+        verify(mConnectivityManager).setRequireVpnForUids(true, toRanges(new UidRangeParcel[] {
                 new UidRangeParcel(user.start, user.start + PKG_UIDS[3] - 1),
                 new UidRangeParcel(user.start + PKG_UIDS[3] + 1, user.stop)
         }));
@@ -383,7 +382,7 @@
         // Set always-on with lockdown and allow app PKGS[2] from lockdown.
         assertTrue(vpn.setAlwaysOnPackage(
                 PKGS[1], true, Collections.singletonList(PKGS[2]), mKeyStore));
-        verify(mNetd).networkRejectNonSecureVpn(eq(true), aryEq(new UidRangeParcel[] {
+        verify(mConnectivityManager).setRequireVpnForUids(true, toRanges(new UidRangeParcel[] {
                 new UidRangeParcel(user.start, user.start + PKG_UIDS[1] - 1),
                 new UidRangeParcel(user.start + PKG_UIDS[2] + 1, user.stop)
         }));
@@ -392,10 +391,10 @@
         // Change allowed app list to PKGS[3].
         assertTrue(vpn.setAlwaysOnPackage(
                 PKGS[1], true, Collections.singletonList(PKGS[3]), mKeyStore));
-        verify(mNetd).networkRejectNonSecureVpn(eq(false), aryEq(new UidRangeParcel[] {
+        verify(mConnectivityManager).setRequireVpnForUids(false, toRanges(new UidRangeParcel[] {
                 new UidRangeParcel(user.start + PKG_UIDS[2] + 1, user.stop)
         }));
-        verify(mNetd).networkRejectNonSecureVpn(eq(true), aryEq(new UidRangeParcel[] {
+        verify(mConnectivityManager).setRequireVpnForUids(true, toRanges(new UidRangeParcel[] {
                 new UidRangeParcel(user.start + PKG_UIDS[1] + 1, user.start + PKG_UIDS[3] - 1),
                 new UidRangeParcel(user.start + PKG_UIDS[3] + 1, user.stop)
         }));
@@ -405,11 +404,11 @@
         // Change the VPN app.
         assertTrue(vpn.setAlwaysOnPackage(
                 PKGS[0], true, Collections.singletonList(PKGS[3]), mKeyStore));
-        verify(mNetd).networkRejectNonSecureVpn(eq(false), aryEq(new UidRangeParcel[] {
+        verify(mConnectivityManager).setRequireVpnForUids(false, toRanges(new UidRangeParcel[] {
                 new UidRangeParcel(user.start, user.start + PKG_UIDS[1] - 1),
                 new UidRangeParcel(user.start + PKG_UIDS[1] + 1, user.start + PKG_UIDS[3] - 1)
         }));
-        verify(mNetd).networkRejectNonSecureVpn(eq(true), aryEq(new UidRangeParcel[] {
+        verify(mConnectivityManager).setRequireVpnForUids(true, toRanges(new UidRangeParcel[] {
                 new UidRangeParcel(user.start, user.start + PKG_UIDS[0] - 1),
                 new UidRangeParcel(user.start + PKG_UIDS[0] + 1, user.start + PKG_UIDS[3] - 1)
         }));
@@ -418,11 +417,11 @@
 
         // Remove the list of allowed packages.
         assertTrue(vpn.setAlwaysOnPackage(PKGS[0], true, null, mKeyStore));
-        verify(mNetd).networkRejectNonSecureVpn(eq(false), aryEq(new UidRangeParcel[] {
+        verify(mConnectivityManager).setRequireVpnForUids(false, toRanges(new UidRangeParcel[] {
                 new UidRangeParcel(user.start + PKG_UIDS[0] + 1, user.start + PKG_UIDS[3] - 1),
                 new UidRangeParcel(user.start + PKG_UIDS[3] + 1, user.stop)
         }));
-        verify(mNetd).networkRejectNonSecureVpn(eq(true), aryEq(new UidRangeParcel[] {
+        verify(mConnectivityManager).setRequireVpnForUids(true, toRanges(new UidRangeParcel[] {
                 new UidRangeParcel(user.start + PKG_UIDS[0] + 1, user.stop),
         }));
         assertBlocked(vpn, user.start + PKG_UIDS[1], user.start + PKG_UIDS[2],
@@ -432,10 +431,10 @@
         // Add the list of allowed packages.
         assertTrue(vpn.setAlwaysOnPackage(
                 PKGS[0], true, Collections.singletonList(PKGS[1]), mKeyStore));
-        verify(mNetd).networkRejectNonSecureVpn(eq(false), aryEq(new UidRangeParcel[] {
+        verify(mConnectivityManager).setRequireVpnForUids(false, toRanges(new UidRangeParcel[] {
                 new UidRangeParcel(user.start + PKG_UIDS[0] + 1, user.stop)
         }));
-        verify(mNetd).networkRejectNonSecureVpn(eq(true), aryEq(new UidRangeParcel[] {
+        verify(mConnectivityManager).setRequireVpnForUids(true, toRanges(new UidRangeParcel[] {
                 new UidRangeParcel(user.start + PKG_UIDS[0] + 1, user.start + PKG_UIDS[1] - 1),
                 new UidRangeParcel(user.start + PKG_UIDS[1] + 1, user.stop)
         }));
@@ -450,11 +449,11 @@
         // allowed package should change from PGKS[1] to PKGS[2].
         assertTrue(vpn.setAlwaysOnPackage(
                 PKGS[0], true, Arrays.asList("com.foo.app", PKGS[2], "com.bar.app"), mKeyStore));
-        verify(mNetd).networkRejectNonSecureVpn(eq(false), aryEq(new UidRangeParcel[]{
+        verify(mConnectivityManager).setRequireVpnForUids(false, toRanges(new UidRangeParcel[] {
                 new UidRangeParcel(user.start + PKG_UIDS[0] + 1, user.start + PKG_UIDS[1] - 1),
                 new UidRangeParcel(user.start + PKG_UIDS[1] + 1, user.stop)
         }));
-        verify(mNetd).networkRejectNonSecureVpn(eq(true), aryEq(new UidRangeParcel[]{
+        verify(mConnectivityManager).setRequireVpnForUids(true, toRanges(new UidRangeParcel[] {
                 new UidRangeParcel(user.start + PKG_UIDS[0] + 1, user.start + PKG_UIDS[2] - 1),
                 new UidRangeParcel(user.start + PKG_UIDS[2] + 1, user.stop)
         }));
@@ -475,7 +474,7 @@
 
         // Set lockdown.
         assertTrue(vpn.setAlwaysOnPackage(PKGS[3], true, null, mKeyStore));
-        verify(mNetd).networkRejectNonSecureVpn(eq(true), aryEq(new UidRangeParcel[] {
+        verify(mConnectivityManager).setRequireVpnForUids(true, toRanges(new UidRangeParcel[] {
                 new UidRangeParcel(user.start, user.start + PKG_UIDS[3] - 1),
                 new UidRangeParcel(user.start + PKG_UIDS[3] + 1, user.stop)
         }));
@@ -485,7 +484,7 @@
         // Add the restricted user.
         setMockedUsers(primaryUser, tempProfile);
         vpn.onUserAdded(tempProfile.id);
-        verify(mNetd).networkRejectNonSecureVpn(eq(true), aryEq(new UidRangeParcel[] {
+        verify(mConnectivityManager).setRequireVpnForUids(true, toRanges(new UidRangeParcel[] {
                 new UidRangeParcel(profile.start, profile.start + PKG_UIDS[3] - 1),
                 new UidRangeParcel(profile.start + PKG_UIDS[3] + 1, profile.stop)
         }));
@@ -493,7 +492,7 @@
         // Remove the restricted user.
         tempProfile.partial = true;
         vpn.onUserRemoved(tempProfile.id);
-        verify(mNetd).networkRejectNonSecureVpn(eq(false), aryEq(new UidRangeParcel[] {
+        verify(mConnectivityManager).setRequireVpnForUids(false, toRanges(new UidRangeParcel[] {
                 new UidRangeParcel(profile.start, profile.start + PKG_UIDS[3] - 1),
                 new UidRangeParcel(profile.start + PKG_UIDS[3] + 1, profile.stop)
         }));
@@ -506,22 +505,29 @@
                 new UidRangeParcel(PRI_USER_RANGE.start, PRI_USER_RANGE.stop)};
         // Given legacy lockdown is already enabled,
         vpn.setLockdown(true);
-
-        verify(mNetd).networkRejectNonSecureVpn(eq(true), aryEq(primaryUserRangeParcel));
+        verify(mConnectivityManager, times(1)).setRequireVpnForUids(true,
+                toRanges(primaryUserRangeParcel));
 
         // Enabling legacy lockdown twice should do nothing.
         vpn.setLockdown(true);
-        verify(mNetd, times(1))
-                .networkRejectNonSecureVpn(anyBoolean(), any(UidRangeParcel[].class));
+        verify(mConnectivityManager, times(1)).setRequireVpnForUids(anyBoolean(), any());
 
         // And disabling should remove the rules exactly once.
         vpn.setLockdown(false);
-        verify(mNetd).networkRejectNonSecureVpn(eq(false), aryEq(primaryUserRangeParcel));
+        verify(mConnectivityManager, times(1)).setRequireVpnForUids(false,
+                toRanges(primaryUserRangeParcel));
 
         // Removing the lockdown again should have no effect.
         vpn.setLockdown(false);
-        verify(mNetd, times(2)).networkRejectNonSecureVpn(
-                anyBoolean(), any(UidRangeParcel[].class));
+        verify(mConnectivityManager, times(2)).setRequireVpnForUids(anyBoolean(), any());
+    }
+
+    private ArrayList<Range<Integer>> toRanges(UidRangeParcel[] ranges) {
+        ArrayList<Range<Integer>> rangesArray = new ArrayList<>(ranges.length);
+        for (int i = 0; i < ranges.length; i++) {
+            rangesArray.add(new Range<>(ranges[i].start, ranges[i].stop));
+        }
+        return rangesArray;
     }
 
     @Test
@@ -535,21 +541,21 @@
             new UidRangeParcel(entireUser[0].start + PKG_UIDS[0] + 1, entireUser[0].stop)
         };
 
-        final InOrder order = inOrder(mNetd);
+        final InOrder order = inOrder(mConnectivityManager);
 
         // Given lockdown is enabled with no package (legacy VPN),
         vpn.setLockdown(true);
-        order.verify(mNetd).networkRejectNonSecureVpn(eq(true), aryEq(entireUser));
+        order.verify(mConnectivityManager).setRequireVpnForUids(true, toRanges(entireUser));
 
         // When a new VPN package is set the rules should change to cover that package.
         vpn.prepare(null, PKGS[0], VpnManager.TYPE_VPN_SERVICE);
-        order.verify(mNetd).networkRejectNonSecureVpn(eq(false), aryEq(entireUser));
-        order.verify(mNetd).networkRejectNonSecureVpn(eq(true), aryEq(exceptPkg0));
+        order.verify(mConnectivityManager).setRequireVpnForUids(false, toRanges(entireUser));
+        order.verify(mConnectivityManager).setRequireVpnForUids(true, toRanges(exceptPkg0));
 
         // When that VPN package is unset, everything should be undone again in reverse.
         vpn.prepare(null, VpnConfig.LEGACY_VPN, VpnManager.TYPE_VPN_SERVICE);
-        order.verify(mNetd).networkRejectNonSecureVpn(eq(false), aryEq(exceptPkg0));
-        order.verify(mNetd).networkRejectNonSecureVpn(eq(true), aryEq(entireUser));
+        order.verify(mConnectivityManager).setRequireVpnForUids(false, toRanges(exceptPkg0));
+        order.verify(mConnectivityManager).setRequireVpnForUids(true, toRanges(entireUser));
     }
 
     @Test
diff --git a/tests/net/java/com/android/server/net/NetworkStatsCollectionTest.java b/tests/net/java/com/android/server/net/NetworkStatsCollectionTest.java
index 89146f9..435c3c0 100644
--- a/tests/net/java/com/android/server/net/NetworkStatsCollectionTest.java
+++ b/tests/net/java/com/android/server/net/NetworkStatsCollectionTest.java
@@ -64,7 +64,6 @@
 
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
-import java.io.DataOutputStream;
 import java.io.File;
 import java.io.FileOutputStream;
 import java.io.InputStream;
@@ -124,7 +123,7 @@
 
         // now export into a unified format
         final ByteArrayOutputStream bos = new ByteArrayOutputStream();
-        collection.write(new DataOutputStream(bos));
+        collection.write(bos);
 
         // clear structure completely
         collection.reset();
@@ -152,7 +151,7 @@
 
         // now export into a unified format
         final ByteArrayOutputStream bos = new ByteArrayOutputStream();
-        collection.write(new DataOutputStream(bos));
+        collection.write(bos);
 
         // clear structure completely
         collection.reset();
@@ -180,7 +179,7 @@
 
         // now export into a unified format
         final ByteArrayOutputStream bos = new ByteArrayOutputStream();
-        collection.write(new DataOutputStream(bos));
+        collection.write(bos);
 
         // clear structure completely
         collection.reset();