diff --git a/core/java/android/net/CaptivePortal.java b/core/java/android/net/CaptivePortal.java
index 3ab35e1..ba7323d 100644
--- a/core/java/android/net/CaptivePortal.java
+++ b/core/java/android/net/CaptivePortal.java
@@ -29,15 +29,33 @@
  * {@code ACTION_CAPTIVE_PORTAL_SIGN_IN} activity.
  */
 public class CaptivePortal implements Parcelable {
-    /** @hide */
+    /**
+     * Response code from the captive portal application, indicating that the portal was dismissed
+     * and the network should be re-validated.
+     * @see ICaptivePortal#appResponse(int)
+     * @see android.net.INetworkMonitor#notifyCaptivePortalAppFinished(int)
+     * @hide
+     */
     @SystemApi
     @TestApi
     public static final int APP_RETURN_DISMISSED    = 0;
-    /** @hide */
+    /**
+     * Response code from the captive portal application, indicating that the user did not login and
+     * does not want to use the captive portal network.
+     * @see ICaptivePortal#appResponse(int)
+     * @see android.net.INetworkMonitor#notifyCaptivePortalAppFinished(int)
+     * @hide
+     */
     @SystemApi
     @TestApi
     public static final int APP_RETURN_UNWANTED     = 1;
-    /** @hide */
+    /**
+     * Response code from the captive portal application, indicating that the user does not wish to
+     * login but wants to use the captive portal network as-is.
+     * @see ICaptivePortal#appResponse(int)
+     * @see android.net.INetworkMonitor#notifyCaptivePortalAppFinished(int)
+     * @hide
+     */
     @SystemApi
     @TestApi
     public static final int APP_RETURN_WANTED_AS_IS = 2;
diff --git a/core/java/android/net/ConnectivityManager.java b/core/java/android/net/ConnectivityManager.java
index 2a357ff..e5802c2 100644
--- a/core/java/android/net/ConnectivityManager.java
+++ b/core/java/android/net/ConnectivityManager.java
@@ -38,7 +38,6 @@
 import android.os.Build.VERSION_CODES;
 import android.os.Bundle;
 import android.os.Handler;
-import android.os.HandlerThread;
 import android.os.IBinder;
 import android.os.INetworkActivityListener;
 import android.os.INetworkManagementService;
@@ -75,6 +74,9 @@
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.Executor;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.RejectedExecutionException;
 
 /**
  * Class that answers queries about the state of network connectivity. It also
@@ -1813,23 +1815,26 @@
         public static final int MIN_INTERVAL = 10;
 
         private final Network mNetwork;
-        private final PacketKeepaliveCallback mCallback;
-        private final Looper mLooper;
-        private final Messenger mMessenger;
+        private final ISocketKeepaliveCallback mCallback;
+        private final ExecutorService mExecutor;
 
         private volatile Integer mSlot;
 
-        void stopLooper() {
-            mLooper.quit();
-        }
-
         @UnsupportedAppUsage
         public void stop() {
             try {
-                mService.stopKeepalive(mNetwork, mSlot);
-            } catch (RemoteException e) {
-                Log.e(TAG, "Error stopping packet keepalive: ", e);
-                stopLooper();
+                mExecutor.execute(() -> {
+                    try {
+                        if (mSlot != null) {
+                            mService.stopKeepalive(mNetwork, mSlot);
+                        }
+                    } catch (RemoteException e) {
+                        Log.e(TAG, "Error stopping packet keepalive: ", e);
+                        throw e.rethrowFromSystemServer();
+                    }
+                });
+            } catch (RejectedExecutionException e) {
+                // The internal executor has already stopped due to previous event.
             }
         }
 
@@ -1837,40 +1842,43 @@
             Preconditions.checkNotNull(network, "network cannot be null");
             Preconditions.checkNotNull(callback, "callback cannot be null");
             mNetwork = network;
-            mCallback = callback;
-            HandlerThread thread = new HandlerThread(TAG);
-            thread.start();
-            mLooper = thread.getLooper();
-            mMessenger = new Messenger(new Handler(mLooper) {
+            mExecutor = Executors.newSingleThreadExecutor();
+            mCallback = new ISocketKeepaliveCallback.Stub() {
                 @Override
-                public void handleMessage(Message message) {
-                    switch (message.what) {
-                        case NetworkAgent.EVENT_SOCKET_KEEPALIVE:
-                            int error = message.arg2;
-                            try {
-                                if (error == SUCCESS) {
-                                    if (mSlot == null) {
-                                        mSlot = message.arg1;
-                                        mCallback.onStarted();
-                                    } else {
-                                        mSlot = null;
-                                        stopLooper();
-                                        mCallback.onStopped();
-                                    }
-                                } else {
-                                    stopLooper();
-                                    mCallback.onError(error);
-                                }
-                            } catch (Exception e) {
-                                Log.e(TAG, "Exception in keepalive callback(" + error + ")", e);
-                            }
-                            break;
-                        default:
-                            Log.e(TAG, "Unhandled message " + Integer.toHexString(message.what));
-                            break;
-                    }
+                public void onStarted(int slot) {
+                    Binder.withCleanCallingIdentity(() ->
+                            mExecutor.execute(() -> {
+                                mSlot = slot;
+                                callback.onStarted();
+                            }));
                 }
-            });
+
+                @Override
+                public void onStopped() {
+                    Binder.withCleanCallingIdentity(() ->
+                            mExecutor.execute(() -> {
+                                mSlot = null;
+                                callback.onStopped();
+                            }));
+                    mExecutor.shutdown();
+                }
+
+                @Override
+                public void onError(int error) {
+                    Binder.withCleanCallingIdentity(() ->
+                            mExecutor.execute(() -> {
+                                mSlot = null;
+                                callback.onError(error);
+                            }));
+                    mExecutor.shutdown();
+                }
+
+                @Override
+                public void onDataReceived() {
+                    // PacketKeepalive is only used for Nat-T keepalive and as such does not invoke
+                    // this callback when data is received.
+                }
+            };
         }
     }
 
@@ -1887,12 +1895,11 @@
             InetAddress srcAddr, int srcPort, InetAddress dstAddr) {
         final PacketKeepalive k = new PacketKeepalive(network, callback);
         try {
-            mService.startNattKeepalive(network, intervalSeconds, k.mMessenger, new Binder(),
+            mService.startNattKeepalive(network, intervalSeconds, k.mCallback,
                     srcAddr.getHostAddress(), srcPort, dstAddr.getHostAddress());
         } catch (RemoteException e) {
             Log.e(TAG, "Error starting packet keepalive: ", e);
-            k.stopLooper();
-            return null;
+            throw e.rethrowFromSystemServer();
         }
         return k;
     }
@@ -2809,23 +2816,6 @@
     }
 
     /**
-     * @removed
-     * @deprecated This API would be removed when all of caller has been updated.
-     * */
-    @Deprecated
-    public abstract static class TetheringEntitlementValueListener  {
-        /**
-         * Called to notify entitlement result.
-         *
-         * @param resultCode a int value of entitlement result. It may be one of
-         *         {@link #TETHER_ERROR_NO_ERROR},
-         *         {@link #TETHER_ERROR_PROVISION_FAILED}, or
-         *         {@link #TETHER_ERROR_ENTITLEMENT_UNKONWN}.
-         */
-        public void onEntitlementResult(int resultCode) {}
-    }
-
-    /**
      * Get the last value of the entitlement check on this downstream. If the cached value is
      * {@link #TETHER_ERROR_NO_ERROR} or showEntitlementUi argument is false, it just return the
      * cached value. Otherwise, a UI-based entitlement check would be performed. It is not
@@ -2871,31 +2861,6 @@
     }
 
     /**
-     * @removed
-     * @deprecated This API would be removed when all of caller has been updated.
-     * */
-    @Deprecated
-    public void getLatestTetheringEntitlementValue(int type, boolean showEntitlementUi,
-            @NonNull final TetheringEntitlementValueListener listener, @Nullable Handler handler) {
-        Preconditions.checkNotNull(listener, "TetheringEntitlementValueListener cannot be null.");
-        ResultReceiver wrappedListener = new ResultReceiver(handler) {
-            @Override
-            protected void onReceiveResult(int resultCode, Bundle resultData) {
-                listener.onEntitlementResult(resultCode);
-            }
-        };
-
-        try {
-            String pkgName = mContext.getOpPackageName();
-            Log.i(TAG, "getLatestTetheringEntitlementValue:" + pkgName);
-            mService.getLatestTetheringEntitlementResult(type, wrappedListener,
-                    showEntitlementUi, pkgName);
-        } catch (RemoteException e) {
-            throw e.rethrowFromSystemServer();
-        }
-    }
-
-    /**
      * Report network connectivity status.  This is currently used only
      * to alter status bar UI.
      * <p>This method requires the caller to hold the permission
diff --git a/core/java/android/net/DnsResolver.java b/core/java/android/net/DnsResolver.java
index d3bc3e6..93b8cf8 100644
--- a/core/java/android/net/DnsResolver.java
+++ b/core/java/android/net/DnsResolver.java
@@ -22,11 +22,11 @@
 import static android.os.MessageQueue.OnFileDescriptorEventListener.EVENT_ERROR;
 import static android.os.MessageQueue.OnFileDescriptorEventListener.EVENT_INPUT;
 
+import android.annotation.CallbackExecutor;
 import android.annotation.IntDef;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
-import android.os.Handler;
-import android.os.MessageQueue;
+import android.os.Looper;
 import android.system.ErrnoException;
 import android.util.Log;
 
@@ -37,8 +37,7 @@
 import java.net.UnknownHostException;
 import java.util.ArrayList;
 import java.util.List;
-import java.util.function.Consumer;
-
+import java.util.concurrent.Executor;
 
 /**
  * Dns resolver class for asynchronous dns querying
@@ -81,66 +80,137 @@
     public static final int FLAG_NO_CACHE_STORE = 1 << 1;
     public static final int FLAG_NO_CACHE_LOOKUP = 1 << 2;
 
-    private static final int DNS_RAW_RESPONSE = 1;
-
     private static final int NETID_UNSET = 0;
 
     private static final DnsResolver sInstance = new DnsResolver();
 
     /**
-     * listener for receiving raw answers
-     */
-    public interface RawAnswerListener {
-        /**
-         * {@code byte[]} is {@code null} if query timed out
-         */
-        void onAnswer(@Nullable byte[] answer);
-    }
-
-    /**
-     * listener for receiving parsed answers
-     */
-    public interface InetAddressAnswerListener {
-        /**
-         * Will be called exactly once with all the answers to the query.
-         * size of addresses will be zero if no available answer could be parsed.
-         */
-        void onAnswer(@NonNull List<InetAddress> addresses);
-    }
-
-    /**
      * Get instance for DnsResolver
      */
-    public static DnsResolver getInstance() {
+    public static @NonNull DnsResolver getInstance() {
         return sInstance;
     }
 
     private DnsResolver() {}
 
     /**
-     * Pass in a blob and corresponding setting,
-     * get a blob back asynchronously with the entire raw answer.
+     * Answer parser for parsing raw answers
+     *
+     * @param <T> The type of the parsed answer
+     */
+    public interface AnswerParser<T> {
+        /**
+         * Creates a <T> answer by parsing the given raw answer.
+         *
+         * @param rawAnswer the raw answer to be parsed
+         * @return a parsed <T> answer
+         * @throws ParseException if parsing failed
+         */
+        @NonNull T parse(@NonNull byte[] rawAnswer) throws ParseException;
+    }
+
+    /**
+     * Base class for answer callbacks
+     *
+     * @param <T> The type of the parsed answer
+     */
+    public abstract static class AnswerCallback<T> {
+        /** @hide */
+        public final AnswerParser<T> parser;
+
+        public AnswerCallback(@NonNull AnswerParser<T> parser) {
+            this.parser = parser;
+        };
+
+        /**
+         * Success response to
+         * {@link android.net.DnsResolver#query query()}.
+         *
+         * Invoked when the answer to a query was successfully parsed.
+         *
+         * @param answer parsed answer to the query.
+         *
+         * {@see android.net.DnsResolver#query query()}
+         */
+        public abstract void onAnswer(@NonNull T answer);
+
+        /**
+         * Error response to
+         * {@link android.net.DnsResolver#query query()}.
+         *
+         * Invoked when there is no valid answer to
+         * {@link android.net.DnsResolver#query query()}
+         *
+         * @param exception a {@link ParseException} object with additional
+         *    detail regarding the failure
+         */
+        public abstract void onParseException(@NonNull ParseException exception);
+
+        /**
+         * Error response to
+         * {@link android.net.DnsResolver#query query()}.
+         *
+         * Invoked if an error happens when
+         * issuing the DNS query or receiving the result.
+         * {@link android.net.DnsResolver#query query()}
+         *
+         * @param exception an {@link ErrnoException} object with additional detail
+         *    regarding the failure
+         */
+        public abstract void onQueryException(@NonNull ErrnoException exception);
+    }
+
+    /**
+     * Callback for receiving raw answers
+     */
+    public abstract static class RawAnswerCallback extends AnswerCallback<byte[]> {
+        public RawAnswerCallback() {
+            super(rawAnswer -> rawAnswer);
+        }
+    }
+
+    /**
+     * Callback for receiving parsed {@link InetAddress} answers
+     *
+     * Note that if the answer does not contain any IP addresses,
+     * onAnswer will be called with an empty list.
+     */
+    public abstract static class InetAddressAnswerCallback
+            extends AnswerCallback<List<InetAddress>> {
+        public InetAddressAnswerCallback() {
+            super(rawAnswer -> new DnsAddressAnswer(rawAnswer).getAddresses());
+        }
+    }
+
+    /**
+     * Send a raw DNS query.
+     * The answer will be provided asynchronously through the provided {@link AnswerCallback}.
      *
      * @param network {@link Network} specifying which network for querying.
      *         {@code null} for query on default network.
      * @param query blob message
      * @param flags flags as a combination of the FLAGS_* constants
-     * @param handler {@link Handler} to specify the thread
-     *         upon which the {@link RawAnswerListener} will be invoked.
-     * @param listener a {@link RawAnswerListener} which will be called to notify the caller
+     * @param executor The {@link Executor} that the callback should be executed on.
+     * @param callback an {@link AnswerCallback} which will be called to notify the caller
      *         of the result of dns query.
      */
-    public void query(@Nullable Network network, @NonNull byte[] query, @QueryFlag int flags,
-            @NonNull Handler handler, @NonNull RawAnswerListener listener) throws ErrnoException {
-        final FileDescriptor queryfd = resNetworkSend((network != null
+    public <T> void query(@Nullable Network network, @NonNull byte[] query, @QueryFlag int flags,
+            @NonNull @CallbackExecutor Executor executor, @NonNull AnswerCallback<T> callback) {
+        final FileDescriptor queryfd;
+        try {
+            queryfd = resNetworkSend((network != null
                 ? network.netId : NETID_UNSET), query, query.length, flags);
-        registerFDListener(handler.getLooper().getQueue(), queryfd,
-                answerbuf -> listener.onAnswer(answerbuf));
+        } catch (ErrnoException e) {
+            callback.onQueryException(e);
+            return;
+        }
+
+        registerFDListener(executor, queryfd, callback);
     }
 
     /**
-     * Pass in a domain name and corresponding setting,
-     * get a blob back asynchronously with the entire raw answer.
+     * Send a DNS query with the specified name, class and query type.
+     * The answer will be provided asynchronously through the provided {@link AnswerCallback}.
      *
      * @param network {@link Network} specifying which network for querying.
      *         {@code null} for query on default network.
@@ -148,74 +218,53 @@
      * @param nsClass dns class as one of the CLASS_* constants
      * @param nsType dns resource record (RR) type as one of the TYPE_* constants
      * @param flags flags as a combination of the FLAGS_* constants
-     * @param handler {@link Handler} to specify the thread
-     *         upon which the {@link RawAnswerListener} will be invoked.
-     * @param listener a {@link RawAnswerListener} which will be called to notify the caller
+     * @param executor The {@link Executor} that the callback should be executed on.
+     * @param callback an {@link AnswerCallback} which will be called to notify the caller
      *         of the result of dns query.
      */
-    public void query(@Nullable Network network, @NonNull String domain, @QueryClass int nsClass,
-            @QueryType int nsType, @QueryFlag int flags,
-            @NonNull Handler handler, @NonNull RawAnswerListener listener) throws ErrnoException {
-        final FileDescriptor queryfd = resNetworkQuery((network != null
-                ? network.netId : NETID_UNSET), domain, nsClass, nsType, flags);
-        registerFDListener(handler.getLooper().getQueue(), queryfd,
-                answerbuf -> listener.onAnswer(answerbuf));
+    public <T> void query(@Nullable Network network, @NonNull String domain,
+            @QueryClass int nsClass, @QueryType int nsType, @QueryFlag int flags,
+            @NonNull @CallbackExecutor Executor executor, @NonNull AnswerCallback<T> callback) {
+        final FileDescriptor queryfd;
+        try {
+            queryfd = resNetworkQuery((network != null
+                    ? network.netId : NETID_UNSET), domain, nsClass, nsType, flags);
+        } catch (ErrnoException e) {
+            callback.onQueryException(e);
+            return;
+        }
+        registerFDListener(executor, queryfd, callback);
     }
 
-    /**
-     * Pass in a domain name and corresponding setting,
-     * get back a set of InetAddresses asynchronously.
-     *
-     * @param network {@link Network} specifying which network for querying.
-     *         {@code null} for query on default network.
-     * @param domain domain name for querying
-     * @param flags flags as a combination of the FLAGS_* constants
-     * @param handler {@link Handler} to specify the thread
-     *         upon which the {@link InetAddressAnswerListener} will be invoked.
-     * @param listener an {@link InetAddressAnswerListener} which will be called to
-     *         notify the caller of the result of dns query.
-     *
-     */
-    public void query(@Nullable Network network, @NonNull String domain, @QueryFlag int flags,
-            @NonNull Handler handler, @NonNull InetAddressAnswerListener listener)
-            throws ErrnoException {
-        final FileDescriptor v4fd = resNetworkQuery((network != null
-                ? network.netId : NETID_UNSET), domain, CLASS_IN, TYPE_A, flags);
-        final FileDescriptor v6fd = resNetworkQuery((network != null
-                ? network.netId : NETID_UNSET), domain, CLASS_IN, TYPE_AAAA, flags);
-
-        final InetAddressAnswerAccumulator accmulator =
-                new InetAddressAnswerAccumulator(2, listener);
-        final Consumer<byte[]> consumer = answerbuf ->
-                accmulator.accumulate(parseAnswers(answerbuf));
-
-        registerFDListener(handler.getLooper().getQueue(), v4fd, consumer);
-        registerFDListener(handler.getLooper().getQueue(), v6fd, consumer);
-    }
-
-    private void registerFDListener(@NonNull MessageQueue queue,
-            @NonNull FileDescriptor queryfd, @NonNull Consumer<byte[]> answerConsumer) {
-        queue.addOnFileDescriptorEventListener(
+    private <T> void registerFDListener(@NonNull Executor executor,
+            @NonNull FileDescriptor queryfd, @NonNull AnswerCallback<T> answerCallback) {
+        Looper.getMainLooper().getQueue().addOnFileDescriptorEventListener(
                 queryfd,
                 FD_EVENTS,
                 (fd, events) -> {
-                    byte[] answerbuf = null;
-                    try {
-                    // TODO: Implement result function in Java side instead of using JNI
-                    //       Because JNI method close fd prior than unregistering fd on
-                    //       event listener.
-                        answerbuf = resNetworkResult(fd);
-                    } catch (ErrnoException e) {
-                        Log.e(TAG, "resNetworkResult:" + e.toString());
-                    }
-                    answerConsumer.accept(answerbuf);
+                    executor.execute(() -> {
+                        byte[] answerbuf = null;
+                        try {
+                            answerbuf = resNetworkResult(fd);
+                        } catch (ErrnoException e) {
+                            Log.e(TAG, "resNetworkResult:" + e.toString());
+                            answerCallback.onQueryException(e);
+                            return;
+                        }
 
+                        try {
+                            answerCallback.onAnswer(
+                                    answerCallback.parser.parse(answerbuf));
+                        } catch (ParseException e) {
+                            answerCallback.onParseException(e);
+                        }
+                    });
                     // Unregister this fd listener
                     return 0;
                 });
     }
 
-    private class DnsAddressAnswer extends DnsPacket {
+    private static class DnsAddressAnswer extends DnsPacket {
         private static final String TAG = "DnsResolver.DnsAddressAnswer";
         private static final boolean DBG = false;
 
@@ -226,12 +275,6 @@
             if ((mHeader.flags & (1 << 15)) == 0) {
                 throw new ParseException("Not an answer packet");
             }
-            if (mHeader.rcode != 0) {
-                throw new ParseException("Response error, rcode:" + mHeader.rcode);
-            }
-            if (mHeader.getRecordCount(ANSECTION) == 0) {
-                throw new ParseException("No available answer");
-            }
             if (mHeader.getRecordCount(QDSECTION) == 0) {
                 throw new ParseException("No question found");
             }
@@ -241,6 +284,8 @@
 
         public @NonNull List<InetAddress> getAddresses() {
             final List<InetAddress> results = new ArrayList<InetAddress>();
+            if (mHeader.getRecordCount(ANSECTION) == 0) return results;
+
             for (final DnsRecord ansSec : mRecords[ANSECTION]) {
                 // Only support A and AAAA, also ignore answers if query type != answer type.
                 int nsType = ansSec.nsType;
@@ -259,34 +304,4 @@
         }
     }
 
-    private @Nullable List<InetAddress> parseAnswers(@Nullable byte[] data) {
-        try {
-            return (data == null) ? null : new DnsAddressAnswer(data).getAddresses();
-        } catch (DnsPacket.ParseException e) {
-            Log.e(TAG, "Parse answer fail " + e.getMessage());
-            return null;
-        }
-    }
-
-    private class InetAddressAnswerAccumulator {
-        private final List<InetAddress> mAllAnswers;
-        private final InetAddressAnswerListener mAnswerListener;
-        private final int mTargetAnswerCount;
-        private int mReceivedAnswerCount = 0;
-
-        InetAddressAnswerAccumulator(int size, @NonNull InetAddressAnswerListener listener) {
-            mTargetAnswerCount = size;
-            mAllAnswers = new ArrayList<>();
-            mAnswerListener = listener;
-        }
-
-        public void accumulate(@Nullable List<InetAddress> answer) {
-            if (null != answer) {
-                mAllAnswers.addAll(answer);
-            }
-            if (++mReceivedAnswerCount == mTargetAnswerCount) {
-                mAnswerListener.onAnswer(mAllAnswers);
-            }
-        }
-    }
 }
diff --git a/core/java/android/net/IConnectivityManager.aidl b/core/java/android/net/IConnectivityManager.aidl
index 403b44d..24e6a85 100644
--- a/core/java/android/net/IConnectivityManager.aidl
+++ b/core/java/android/net/IConnectivityManager.aidl
@@ -27,6 +27,7 @@
 import android.net.NetworkQuotaInfo;
 import android.net.NetworkRequest;
 import android.net.NetworkState;
+import android.net.ISocketKeepaliveCallback;
 import android.net.ProxyInfo;
 import android.os.Bundle;
 import android.os.IBinder;
@@ -194,15 +195,15 @@
 
     void factoryReset();
 
-    void startNattKeepalive(in Network network, int intervalSeconds, in Messenger messenger,
-            in IBinder binder, String srcAddr, int srcPort, String dstAddr);
+    void startNattKeepalive(in Network network, int intervalSeconds,
+            in ISocketKeepaliveCallback cb, String srcAddr, int srcPort, String dstAddr);
 
     void startNattKeepaliveWithFd(in Network network, in FileDescriptor fd, int resourceId,
-            int intervalSeconds, in Messenger messenger, in IBinder binder, String srcAddr,
+            int intervalSeconds, in ISocketKeepaliveCallback cb, String srcAddr,
             String dstAddr);
 
     void startTcpKeepalive(in Network network, in FileDescriptor fd, int intervalSeconds,
-            in Messenger messenger, in IBinder binder);
+            in ISocketKeepaliveCallback cb);
 
     void stopKeepalive(in Network network, int slot);
 
@@ -219,4 +220,6 @@
 
     void registerTetheringEventCallback(ITetheringEventCallback callback, String callerPkg);
     void unregisterTetheringEventCallback(ITetheringEventCallback callback, String callerPkg);
+
+    IBinder startOrGetTestNetworkService();
 }
diff --git a/core/java/android/net/NattSocketKeepalive.java b/core/java/android/net/NattSocketKeepalive.java
index 88631ae..84da294 100644
--- a/core/java/android/net/NattSocketKeepalive.java
+++ b/core/java/android/net/NattSocketKeepalive.java
@@ -17,7 +17,6 @@
 package android.net;
 
 import android.annotation.NonNull;
-import android.os.Binder;
 import android.os.RemoteException;
 import android.util.Log;
 
@@ -52,24 +51,30 @@
 
     @Override
     void startImpl(int intervalSec) {
-        try {
-            mService.startNattKeepaliveWithFd(mNetwork, mFd, mResourceId, intervalSec, mMessenger,
-                    new Binder(), mSource.getHostAddress(), mDestination.getHostAddress());
-        } catch (RemoteException e) {
-            Log.e(TAG, "Error starting packet keepalive: ", e);
-            stopLooper();
-        }
+        mExecutor.execute(() -> {
+            try {
+                mService.startNattKeepaliveWithFd(mNetwork, mFd, mResourceId, intervalSec,
+                        mCallback,
+                        mSource.getHostAddress(), mDestination.getHostAddress());
+            } catch (RemoteException e) {
+                Log.e(TAG, "Error starting socket keepalive: ", e);
+                throw e.rethrowFromSystemServer();
+            }
+        });
     }
 
     @Override
     void stopImpl() {
-        try {
-            if (mSlot != null) {
-                mService.stopKeepalive(mNetwork, mSlot);
+        mExecutor.execute(() -> {
+            try {
+                if (mSlot != null) {
+                    mService.stopKeepalive(mNetwork, mSlot);
+                }
+            } catch (RemoteException e) {
+                Log.e(TAG, "Error stopping socket keepalive: ", e);
+                throw e.rethrowFromSystemServer();
             }
-        } catch (RemoteException e) {
-            Log.e(TAG, "Error stopping packet keepalive: ", e);
-            stopLooper();
-        }
+        });
+
     }
 }
diff --git a/core/java/android/net/SocketKeepalive.java b/core/java/android/net/SocketKeepalive.java
index 07728be..0e768df 100644
--- a/core/java/android/net/SocketKeepalive.java
+++ b/core/java/android/net/SocketKeepalive.java
@@ -20,13 +20,8 @@
 import android.annotation.IntRange;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
-import android.os.Handler;
-import android.os.HandlerThread;
-import android.os.Looper;
-import android.os.Message;
-import android.os.Messenger;
-import android.os.Process;
-import android.util.Log;
+import android.os.Binder;
+import android.os.RemoteException;
 
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
@@ -152,10 +147,9 @@
 
     @NonNull final IConnectivityManager mService;
     @NonNull final Network mNetwork;
-    @NonNull private final Executor mExecutor;
-    @NonNull private final SocketKeepalive.Callback mCallback;
-    @NonNull private final Looper mLooper;
-    @NonNull final Messenger mMessenger;
+    @NonNull final Executor mExecutor;
+    @NonNull final ISocketKeepaliveCallback mCallback;
+    // TODO: remove slot since mCallback could be used to identify which keepalive to stop.
     @Nullable Integer mSlot;
 
     SocketKeepalive(@NonNull IConnectivityManager service, @NonNull Network network,
@@ -163,53 +157,53 @@
         mService = service;
         mNetwork = network;
         mExecutor = executor;
-        mCallback = callback;
-        // TODO: 1. Use other thread modeling instead of create one thread for every instance to
-        //          reduce the memory cost.
-        //       2. support restart.
-        //       3. Fix race condition which caused by rapidly start and stop.
-        HandlerThread thread = new HandlerThread(TAG, Process.THREAD_PRIORITY_BACKGROUND
-                + Process.THREAD_PRIORITY_LESS_FAVORABLE);
-        thread.start();
-        mLooper = thread.getLooper();
-        mMessenger = new Messenger(new Handler(mLooper) {
+        mCallback = new ISocketKeepaliveCallback.Stub() {
             @Override
-            public void handleMessage(Message message) {
-                switch (message.what) {
-                    case NetworkAgent.EVENT_SOCKET_KEEPALIVE:
-                        final int status = message.arg2;
-                        try {
-                            if (status == SUCCESS) {
-                                if (mSlot == null) {
-                                    mSlot = message.arg1;
-                                    mExecutor.execute(() -> mCallback.onStarted());
-                                } else {
-                                    mSlot = null;
-                                    stopLooper();
-                                    mExecutor.execute(() -> mCallback.onStopped());
-                                }
-                            } else if (status == DATA_RECEIVED) {
-                                stopLooper();
-                                mExecutor.execute(() -> mCallback.onDataReceived());
-                            } else {
-                                stopLooper();
-                                mExecutor.execute(() -> mCallback.onError(status));
-                            }
-                        } catch (Exception e) {
-                            Log.e(TAG, "Exception in keepalive callback(" + status + ")", e);
-                        }
-                        break;
-                    default:
-                        Log.e(TAG, "Unhandled message " + Integer.toHexString(message.what));
-                        break;
-                }
+            public void onStarted(int slot) {
+                Binder.withCleanCallingIdentity(() ->
+                        mExecutor.execute(() -> {
+                            mSlot = slot;
+                            callback.onStarted();
+                        }));
             }
-        });
+
+            @Override
+            public void onStopped() {
+                Binder.withCleanCallingIdentity(() ->
+                        executor.execute(() -> {
+                            mSlot = null;
+                            callback.onStopped();
+                        }));
+            }
+
+            @Override
+            public void onError(int error) {
+                Binder.withCleanCallingIdentity(() ->
+                        executor.execute(() -> {
+                            mSlot = null;
+                            callback.onError(error);
+                        }));
+            }
+
+            @Override
+            public void onDataReceived() {
+                Binder.withCleanCallingIdentity(() ->
+                        executor.execute(() -> {
+                            mSlot = null;
+                            callback.onDataReceived();
+                        }));
+            }
+        };
     }
 
     /**
      * Request that keepalive be started with the given {@code intervalSec}. See
-     * {@link SocketKeepalive}.
+     * {@link SocketKeepalive}. If the remote binder dies, or the binder call throws an exception
+     * when invoking start or stop of the {@link SocketKeepalive}, a {@link RemoteException} will be
+     * thrown into the {@code executor}. This is typically not important to catch because the remote
+     * party is the system, so if it is not in shape to communicate through binder the system is
+     * probably going down anyway. If the caller cares regardless, it can use a custom
+     * {@link Executor} to catch the {@link RemoteException}.
      *
      * @param intervalSec The target interval in seconds between keepalive packet transmissions.
      *                    The interval should be between 10 seconds and 3600 seconds, otherwise
@@ -222,12 +216,6 @@
 
     abstract void startImpl(int intervalSec);
 
-    /** @hide */
-    protected void stopLooper() {
-        // TODO: remove this after changing thread modeling.
-        mLooper.quit();
-    }
-
     /**
      * Requests that keepalive be stopped. The application must wait for {@link Callback#onStopped}
      * before using the object. See {@link SocketKeepalive}.
@@ -245,7 +233,6 @@
     @Override
     public final void close() {
         stop();
-        stopLooper();
     }
 
     /**
@@ -259,7 +246,8 @@
         public void onStopped() {}
         /** An error occurred. */
         public void onError(@ErrorCode int error) {}
-        /** The keepalive on a TCP socket was stopped because the socket received data. */
+        /** The keepalive on a TCP socket was stopped because the socket received data. This is
+         * never called for UDP sockets. */
         public void onDataReceived() {}
     }
 }
diff --git a/core/java/android/net/TcpKeepalivePacketData.java b/core/java/android/net/TcpKeepalivePacketData.java
deleted file mode 100644
index f07dfb6..0000000
--- a/core/java/android/net/TcpKeepalivePacketData.java
+++ /dev/null
@@ -1,234 +0,0 @@
-/*
- * Copyright (C) 2019 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-package android.net;
-
-import static android.net.SocketKeepalive.ERROR_INVALID_IP_ADDRESS;
-
-import android.annotation.NonNull;
-import android.annotation.Nullable;
-import android.net.SocketKeepalive.InvalidPacketException;
-import android.net.util.IpUtils;
-import android.os.Parcel;
-import android.os.Parcelable;
-import android.system.OsConstants;
-
-import java.net.Inet4Address;
-import java.net.InetAddress;
-import java.nio.ByteBuffer;
-import java.nio.ByteOrder;
-import java.util.Objects;
-
-/**
- * Represents the actual tcp keep alive packets which will be used for hardware offload.
- * @hide
- */
-public class TcpKeepalivePacketData extends KeepalivePacketData implements Parcelable {
-    private static final String TAG = "TcpKeepalivePacketData";
-
-     /** TCP sequence number. */
-    public final int tcpSeq;
-
-    /** TCP ACK number. */
-    public final int tcpAck;
-
-    /** TCP RCV window. */
-    public final int tcpWnd;
-
-    /** TCP RCV window scale. */
-    public final int tcpWndScale;
-
-    private static final int IPV4_HEADER_LENGTH = 20;
-    private static final int IPV6_HEADER_LENGTH = 40;
-    private static final int TCP_HEADER_LENGTH = 20;
-
-    // This should only be constructed via static factory methods, such as
-    // tcpKeepalivePacket.
-    private TcpKeepalivePacketData(TcpSocketInfo tcpDetails, byte[] data)
-            throws InvalidPacketException {
-        super(tcpDetails.srcAddress, tcpDetails.srcPort, tcpDetails.dstAddress,
-                tcpDetails.dstPort, data);
-        tcpSeq = tcpDetails.seq;
-        tcpAck = tcpDetails.ack;
-        // In the packet, the window is shifted right by the window scale.
-        tcpWnd = tcpDetails.rcvWnd;
-        tcpWndScale = tcpDetails.rcvWndScale;
-    }
-
-    /**
-     * Factory method to create tcp keepalive packet structure.
-     */
-    public static TcpKeepalivePacketData tcpKeepalivePacket(
-            TcpSocketInfo tcpDetails) throws InvalidPacketException {
-        final byte[] packet;
-        if ((tcpDetails.srcAddress instanceof Inet4Address)
-                && (tcpDetails.dstAddress instanceof Inet4Address)) {
-            packet = buildV4Packet(tcpDetails);
-        } else {
-            // TODO: support ipv6
-            throw new InvalidPacketException(ERROR_INVALID_IP_ADDRESS);
-        }
-
-        return new TcpKeepalivePacketData(tcpDetails, packet);
-    }
-
-    /**
-     * Build ipv4 tcp keepalive packet, not including the link-layer header.
-     */
-    // TODO : if this code is ever moved to the network stack, factorize constants with the ones
-    // over there.
-    private static byte[] buildV4Packet(TcpSocketInfo tcpDetails) {
-        final int length = IPV4_HEADER_LENGTH + TCP_HEADER_LENGTH;
-        ByteBuffer buf = ByteBuffer.allocate(length);
-        buf.order(ByteOrder.BIG_ENDIAN);
-        // IP version and TOS. TODO : fetch this from getsockopt(SOL_IP, IP_TOS)
-        buf.putShort((short) 0x4500);
-        buf.putShort((short) length);
-        buf.putInt(0x4000);                         // ID, flags=DF, offset
-        // TODO : fetch TTL from getsockopt(SOL_IP, IP_TTL)
-        buf.put((byte) 64);
-        buf.put((byte) OsConstants.IPPROTO_TCP);
-        final int ipChecksumOffset = buf.position();
-        buf.putShort((short) 0);                    // IP checksum
-        buf.put(tcpDetails.srcAddress.getAddress());
-        buf.put(tcpDetails.dstAddress.getAddress());
-        buf.putShort((short) tcpDetails.srcPort);
-        buf.putShort((short) tcpDetails.dstPort);
-        buf.putInt(tcpDetails.seq);                 // Sequence Number
-        buf.putInt(tcpDetails.ack);                 // ACK
-        buf.putShort((short) 0x5010);               // TCP length=5, flags=ACK
-        buf.putShort((short) (tcpDetails.rcvWnd >> tcpDetails.rcvWndScale));   // Window size
-        final int tcpChecksumOffset = buf.position();
-        buf.putShort((short) 0);                    // TCP checksum
-        // URG is not set therefore the urgent pointer is not included
-        buf.putShort(ipChecksumOffset, IpUtils.ipChecksum(buf, 0));
-        buf.putShort(tcpChecksumOffset, IpUtils.tcpChecksum(
-                buf, 0, IPV4_HEADER_LENGTH, TCP_HEADER_LENGTH));
-
-        return buf.array();
-    }
-
-    // TODO: add buildV6Packet.
-
-    /** Represents tcp/ip information. */
-    // TODO: Replace TcpSocketInfo with TcpKeepalivePacketDataParcelable.
-    public static class TcpSocketInfo {
-        public final InetAddress srcAddress;
-        public final InetAddress dstAddress;
-        public final int srcPort;
-        public final int dstPort;
-        public final int seq;
-        public final int ack;
-        public final int rcvWnd;
-        public final int rcvWndScale;
-
-        public TcpSocketInfo(InetAddress sAddr, int sPort, InetAddress dAddr,
-                int dPort, int writeSeq, int readSeq, int rWnd, int rWndScale) {
-            srcAddress = sAddr;
-            dstAddress = dAddr;
-            srcPort = sPort;
-            dstPort = dPort;
-            seq = writeSeq;
-            ack = readSeq;
-            rcvWnd = rWnd;
-            rcvWndScale = rWndScale;
-        }
-    }
-
-    @Override
-    public boolean equals(@Nullable final Object o) {
-        if (!(o instanceof TcpKeepalivePacketData)) return false;
-        final TcpKeepalivePacketData other = (TcpKeepalivePacketData) o;
-        return this.srcAddress.equals(other.srcAddress)
-                && this.dstAddress.equals(other.dstAddress)
-                && this.srcPort == other.srcPort
-                && this.dstPort == other.dstPort
-                && this.tcpAck == other.tcpAck
-                && this.tcpSeq == other.tcpSeq
-                && this.tcpWnd == other.tcpWnd
-                && this.tcpWndScale == other.tcpWndScale;
-    }
-
-    @Override
-    public int hashCode() {
-        return Objects.hash(srcAddress, dstAddress, srcPort, dstPort, tcpAck, tcpSeq, tcpWnd,
-                tcpWndScale);
-    }
-
-    /* Parcelable Implementation. */
-    /* Note that this object implements parcelable (and needs to keep doing this as it inherits
-     * from a class that does), but should usually be parceled as a stable parcelable using
-     * the toStableParcelable() and fromStableParcelable() methods.
-     */
-    public int describeContents() {
-        return 0;
-    }
-
-    /** Write to parcel. */
-    public void writeToParcel(Parcel out, int flags) {
-        super.writeToParcel(out, flags);
-        out.writeInt(tcpSeq);
-        out.writeInt(tcpAck);
-        out.writeInt(tcpWnd);
-        out.writeInt(tcpWndScale);
-    }
-
-    private TcpKeepalivePacketData(Parcel in) {
-        super(in);
-        tcpSeq = in.readInt();
-        tcpAck = in.readInt();
-        tcpWnd = in.readInt();
-        tcpWndScale = in.readInt();
-    }
-
-    /** Parcelable Creator. */
-    public static final Parcelable.Creator<TcpKeepalivePacketData> CREATOR =
-            new Parcelable.Creator<TcpKeepalivePacketData>() {
-                public TcpKeepalivePacketData createFromParcel(Parcel in) {
-                    return new TcpKeepalivePacketData(in);
-                }
-
-                public TcpKeepalivePacketData[] newArray(int size) {
-                    return new TcpKeepalivePacketData[size];
-                }
-            };
-
-    /**
-     * Convert this TcpKeepalivePacketData to a TcpKeepalivePacketDataParcelable.
-     */
-    @NonNull
-    public TcpKeepalivePacketDataParcelable toStableParcelable() {
-        final TcpKeepalivePacketDataParcelable parcel = new TcpKeepalivePacketDataParcelable();
-        parcel.srcAddress = srcAddress.getAddress();
-        parcel.srcPort = srcPort;
-        parcel.dstAddress = dstAddress.getAddress();
-        parcel.dstPort = dstPort;
-        parcel.seq = tcpSeq;
-        parcel.ack = tcpAck;
-        return parcel;
-    }
-
-    @Override
-    public String toString() {
-        return "saddr: " + srcAddress
-                + " daddr: " + dstAddress
-                + " sport: " + srcPort
-                + " dport: " + dstPort
-                + " seq: " + tcpSeq
-                + " ack: " + tcpAck
-                + " wnd: " + tcpWnd
-                + " wndScale: " + tcpWndScale;
-    }
-}
diff --git a/core/java/android/net/TcpSocketKeepalive.java b/core/java/android/net/TcpSocketKeepalive.java
index f691a0d..26cc8ff 100644
--- a/core/java/android/net/TcpSocketKeepalive.java
+++ b/core/java/android/net/TcpSocketKeepalive.java
@@ -17,7 +17,6 @@
 package android.net;
 
 import android.annotation.NonNull;
-import android.os.Binder;
 import android.os.RemoteException;
 import android.util.Log;
 
@@ -56,24 +55,28 @@
      */
     @Override
     void startImpl(int intervalSec) {
-        try {
-            final FileDescriptor fd = mSocket.getFileDescriptor$();
-            mService.startTcpKeepalive(mNetwork, fd, intervalSec, mMessenger, new Binder());
-        } catch (RemoteException e) {
-            Log.e(TAG, "Error starting packet keepalive: ", e);
-            stopLooper();
-        }
+        mExecutor.execute(() -> {
+            try {
+                final FileDescriptor fd = mSocket.getFileDescriptor$();
+                mService.startTcpKeepalive(mNetwork, fd, intervalSec, mCallback);
+            } catch (RemoteException e) {
+                Log.e(TAG, "Error starting packet keepalive: ", e);
+                throw e.rethrowFromSystemServer();
+            }
+        });
     }
 
     @Override
     void stopImpl() {
-        try {
-            if (mSlot != null) {
-                mService.stopKeepalive(mNetwork, mSlot);
+        mExecutor.execute(() -> {
+            try {
+                if (mSlot != null) {
+                    mService.stopKeepalive(mNetwork, mSlot);
+                }
+            } catch (RemoteException e) {
+                Log.e(TAG, "Error stopping packet keepalive: ", e);
+                throw e.rethrowFromSystemServer();
             }
-        } catch (RemoteException e) {
-            Log.e(TAG, "Error stopping packet keepalive: ", e);
-            stopLooper();
-        }
+        });
     }
 }
diff --git a/core/java/android/net/UidRange.java b/core/java/android/net/UidRange.java
index 793c82d..fa0eeb9 100644
--- a/core/java/android/net/UidRange.java
+++ b/core/java/android/net/UidRange.java
@@ -19,14 +19,17 @@
 import static android.os.UserHandle.PER_USER_RANGE;
 
 import android.os.Parcel;
+import android.os.Parcelable;
 
 /**
  * An inclusive range of UIDs.
  *
  * @hide
  */
-public final class UidRange extends UidRangeParcel {
-    private UidRange() {}
+public final class UidRange implements Parcelable {
+    public final int start;
+    public final int stop;
+
     public UidRange(int startUid, int stopUid) {
         if (startUid < 0) throw new IllegalArgumentException("Invalid start UID.");
         if (stopUid < 0) throw new IllegalArgumentException("Invalid stop UID.");
@@ -86,18 +89,28 @@
         return start + "-" + stop;
     }
 
-    /**
-     * DO NOT override "writeToParcel" and "readFromParcel" in this class.
-     * The parceling code is autogenerated by the superclass.
-     */
+    // Implement the Parcelable interface
+    // TODO: Consider making this class no longer parcelable, since all users are likely in the
+    // system server.
+    @Override
+    public int describeContents() {
+        return 0;
+    }
+
+    @Override
+    public void writeToParcel(Parcel dest, int flags) {
+        dest.writeInt(start);
+        dest.writeInt(stop);
+    }
 
     public static final Creator<UidRange> CREATOR =
         new Creator<UidRange>() {
             @Override
             public UidRange createFromParcel(Parcel in) {
-                UidRange obj = new UidRange();
-                obj.readFromParcel(in);
-                return obj;
+                int start = in.readInt();
+                int stop = in.readInt();
+
+                return new UidRange(start, stop);
             }
             @Override
             public UidRange[] newArray(int size) {
diff --git a/core/java/android/net/apf/ApfCapabilities.java b/core/java/android/net/apf/ApfCapabilities.java
index e09fa8f..d6023d7 100644
--- a/core/java/android/net/apf/ApfCapabilities.java
+++ b/core/java/android/net/apf/ApfCapabilities.java
@@ -19,17 +19,20 @@
 import android.annotation.SystemApi;
 import android.annotation.TestApi;
 import android.content.Context;
+import android.os.Parcel;
+import android.os.Parcelable;
 
 import com.android.internal.R;
 
 /**
  * APF program support capabilities.
  *
+ * This class is immutable.
  * @hide
  */
 @SystemApi
 @TestApi
-public class ApfCapabilities {
+public final class ApfCapabilities implements Parcelable {
     /**
      * Version of APF instruction set supported for packet filtering. 0 indicates no support for
      * packet filtering using APF programs.
@@ -53,6 +56,37 @@
         this.apfPacketFormat = apfPacketFormat;
     }
 
+    private ApfCapabilities(Parcel in) {
+        apfVersionSupported = in.readInt();
+        maximumApfProgramSize = in.readInt();
+        apfPacketFormat = in.readInt();
+    }
+
+
+    @Override
+    public int describeContents() {
+        return 0;
+    }
+
+    @Override
+    public void writeToParcel(Parcel dest, int flags) {
+        dest.writeInt(apfVersionSupported);
+        dest.writeInt(maximumApfProgramSize);
+        dest.writeInt(apfPacketFormat);
+    }
+
+    public static final Creator<ApfCapabilities> CREATOR = new Creator<ApfCapabilities>() {
+        @Override
+        public ApfCapabilities createFromParcel(Parcel in) {
+            return new ApfCapabilities(in);
+        }
+
+        @Override
+        public ApfCapabilities[] newArray(int size) {
+            return new ApfCapabilities[size];
+        }
+    };
+
     @Override
     public String toString() {
         return String.format("%s{version: %d, maxSize: %d, format: %d}", getClass().getSimpleName(),
diff --git a/services/core/java/com/android/server/ConnectivityService.java b/services/core/java/com/android/server/ConnectivityService.java
index 8384ca0..733b8a0 100644
--- a/services/core/java/com/android/server/ConnectivityService.java
+++ b/services/core/java/com/android/server/ConnectivityService.java
@@ -41,7 +41,6 @@
 import static android.net.NetworkPolicyManager.RULE_NONE;
 import static android.net.NetworkPolicyManager.uidRulesToString;
 import static android.net.shared.NetworkMonitorUtils.isValidationRequired;
-import static android.net.shared.NetworkParcelableUtil.toStableParcelable;
 import static android.os.Process.INVALID_UID;
 import static android.system.OsConstants.IPPROTO_TCP;
 import static android.system.OsConstants.IPPROTO_UDP;
@@ -73,6 +72,7 @@
 import android.net.INetworkPolicyListener;
 import android.net.INetworkPolicyManager;
 import android.net.INetworkStatsService;
+import android.net.ISocketKeepaliveCallback;
 import android.net.ITetheringEventCallback;
 import android.net.InetAddresses;
 import android.net.IpPrefix;
@@ -299,6 +299,15 @@
     private INetworkPolicyManager mPolicyManager;
     private NetworkPolicyManagerInternal mPolicyManagerInternal;
 
+    /**
+     * TestNetworkService (lazily) created upon first usage. Locked to prevent creation of multiple
+     * instances.
+     */
+    @GuardedBy("mTNSLock")
+    private TestNetworkService mTNS;
+
+    private final Object mTNSLock = new Object();
+
     private String mCurrentTcpBufferSizes;
 
     private static final SparseArray<String> sMagicDecoderRing = MessageUtils.findMessageNames(
@@ -5367,7 +5376,7 @@
         final long token = Binder.clearCallingIdentity();
         try {
             getNetworkStack().makeNetworkMonitor(
-                    toStableParcelable(nai.network), name, new NetworkMonitorCallbacks(nai));
+                    nai.network, name, new NetworkMonitorCallbacks(nai));
         } finally {
             Binder.restoreCallingIdentity(token);
         }
@@ -6680,32 +6689,32 @@
     }
 
     @Override
-    public void startNattKeepalive(Network network, int intervalSeconds, Messenger messenger,
-            IBinder binder, String srcAddr, int srcPort, String dstAddr) {
+    public void startNattKeepalive(Network network, int intervalSeconds,
+            ISocketKeepaliveCallback cb, String srcAddr, int srcPort, String dstAddr) {
         enforceKeepalivePermission();
         mKeepaliveTracker.startNattKeepalive(
                 getNetworkAgentInfoForNetwork(network),
-                intervalSeconds, messenger, binder,
+                intervalSeconds, cb,
                 srcAddr, srcPort, dstAddr, NattSocketKeepalive.NATT_PORT);
     }
 
     @Override
     public void startNattKeepaliveWithFd(Network network, FileDescriptor fd, int resourceId,
-            int intervalSeconds, Messenger messenger, IBinder binder, String srcAddr,
+            int intervalSeconds, ISocketKeepaliveCallback cb, String srcAddr,
             String dstAddr) {
         enforceKeepalivePermission();
         mKeepaliveTracker.startNattKeepalive(
                 getNetworkAgentInfoForNetwork(network), fd, resourceId,
-                intervalSeconds, messenger, binder,
+                intervalSeconds, cb,
                 srcAddr, dstAddr, NattSocketKeepalive.NATT_PORT);
     }
 
     @Override
     public void startTcpKeepalive(Network network, FileDescriptor fd, int intervalSeconds,
-            Messenger messenger, IBinder binder) {
+            ISocketKeepaliveCallback cb) {
         enforceKeepalivePermission();
         mKeepaliveTracker.startTcpKeepalive(
-                getNetworkAgentInfoForNetwork(network), fd, intervalSeconds, messenger, binder);
+                getNetworkAgentInfoForNetwork(network), fd, intervalSeconds, cb);
     }
 
     @Override
@@ -6948,4 +6957,22 @@
             return vpn != null && vpn.getLockdown();
         }
     }
+
+    /**
+     * Returns a IBinder to a TestNetworkService. Will be lazily created as needed.
+     *
+     * <p>The TestNetworkService must be run in the system server due to TUN creation.
+     */
+    @Override
+    public IBinder startOrGetTestNetworkService() {
+        synchronized (mTNSLock) {
+            TestNetworkService.enforceTestNetworkPermissions(mContext);
+
+            if (mTNS == null) {
+                mTNS = new TestNetworkService(mContext, mNMS);
+            }
+
+            return mTNS;
+        }
+    }
 }
diff --git a/services/core/java/com/android/server/connectivity/KeepaliveTracker.java b/services/core/java/com/android/server/connectivity/KeepaliveTracker.java
index cc4c173..35d6860 100644
--- a/services/core/java/com/android/server/connectivity/KeepaliveTracker.java
+++ b/services/core/java/com/android/server/connectivity/KeepaliveTracker.java
@@ -21,8 +21,8 @@
 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.NetworkAgent.EVENT_SOCKET_KEEPALIVE;
 import static android.net.SocketKeepalive.BINDER_DIED;
+import static android.net.SocketKeepalive.DATA_RECEIVED;
 import static android.net.SocketKeepalive.ERROR_INVALID_INTERVAL;
 import static android.net.SocketKeepalive.ERROR_INVALID_IP_ADDRESS;
 import static android.net.SocketKeepalive.ERROR_INVALID_NETWORK;
@@ -34,6 +34,7 @@
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.net.ISocketKeepaliveCallback;
 import android.net.KeepalivePacketData;
 import android.net.NattKeepalivePacketData;
 import android.net.NetworkAgent;
@@ -47,7 +48,6 @@
 import android.os.Handler;
 import android.os.IBinder;
 import android.os.Message;
-import android.os.Messenger;
 import android.os.Process;
 import android.os.RemoteException;
 import android.system.ErrnoException;
@@ -99,8 +99,7 @@
      */
     class KeepaliveInfo implements IBinder.DeathRecipient {
         // Bookkeeping data.
-        private final Messenger mMessenger;
-        private final IBinder mBinder;
+        private final ISocketKeepaliveCallback mCallback;
         private final int mUid;
         private final int mPid;
         private final NetworkAgentInfo mNai;
@@ -124,15 +123,13 @@
         private static final int STARTED = 3;
         private int mStartedState = NOT_STARTED;
 
-        KeepaliveInfo(@NonNull Messenger messenger,
-                @NonNull IBinder binder,
+        KeepaliveInfo(@NonNull ISocketKeepaliveCallback callback,
                 @NonNull NetworkAgentInfo nai,
                 @NonNull KeepalivePacketData packet,
                 int interval,
                 int type,
                 @NonNull FileDescriptor fd) {
-            mMessenger = messenger;
-            mBinder = binder;
+            mCallback = callback;
             mPid = Binder.getCallingPid();
             mUid = Binder.getCallingUid();
 
@@ -143,7 +140,7 @@
             mFd = fd;
 
             try {
-                mBinder.linkToDeath(this, 0);
+                mCallback.asBinder().linkToDeath(this, 0);
             } catch (RemoteException e) {
                 binderDied();
             }
@@ -176,22 +173,14 @@
                     + " ]";
         }
 
-        /** Sends a message back to the application via its SocketKeepalive.Callback. */
-        void notifyMessenger(int slot, int err) {
-            if (DBG) {
-                Log.d(TAG, "notify keepalive " + mSlot + " on " + mNai.network + " for " + err);
-            }
-            KeepaliveTracker.this.notifyMessenger(mMessenger, slot, err);
-        }
-
         /** Called when the application process is killed. */
         public void binderDied() {
             stop(BINDER_DIED);
         }
 
         void unlinkDeathRecipient() {
-            if (mBinder != null) {
-                mBinder.unlinkToDeath(this, 0);
+            if (mCallback != null) {
+                mCallback.asBinder().unlinkToDeath(this, 0);
             }
         }
 
@@ -283,9 +272,23 @@
                     Log.wtf(TAG, "Stopping keepalive with unknown type: " + mType);
                 }
             }
-            // TODO: at the moment we unconditionally return failure here. In cases where the
-            // NetworkAgent is alive, should we ask it to reply, so it can return failure?
-            notifyMessenger(mSlot, reason);
+
+            if (reason == SUCCESS) {
+                try {
+                    mCallback.onStopped();
+                } catch (RemoteException e) {
+                    Log.w(TAG, "Discarded onStop callback: " + reason);
+                }
+            } else if (reason == DATA_RECEIVED) {
+                try {
+                    mCallback.onDataReceived();
+                } catch (RemoteException e) {
+                    Log.w(TAG, "Discarded onDataReceived callback: " + reason);
+                }
+            } else {
+                notifyErrorCallback(mCallback, reason);
+            }
+
             unlinkDeathRecipient();
         }
 
@@ -294,16 +297,12 @@
         }
     }
 
-    void notifyMessenger(Messenger messenger, int slot, int err) {
-        Message message = Message.obtain();
-        message.what = EVENT_SOCKET_KEEPALIVE;
-        message.arg1 = slot;
-        message.arg2 = err;
-        message.obj = null;
+    void notifyErrorCallback(ISocketKeepaliveCallback cb, int error) {
+        if (DBG) Log.w(TAG, "Sending onError(" + error + ") callback");
         try {
-            messenger.send(message);
+            cb.onError(error);
         } catch (RemoteException e) {
-            // Process died?
+            Log.w(TAG, "Discarded onError(" + error + ") callback");
         }
     }
 
@@ -414,7 +413,11 @@
             // Keepalive successfully started.
             if (DBG) Log.d(TAG, "Started keepalive " + slot + " on " + nai.name());
             ki.mStartedState = KeepaliveInfo.STARTED;
-            ki.notifyMessenger(slot, reason);
+            try {
+                ki.mCallback.onStarted(slot);
+            } catch (RemoteException e) {
+                Log.w(TAG, "Discarded onStarted(" + slot + ") callback");
+            }
         } else {
             // Keepalive successfully stopped, or error.
             ki.mStartedState = KeepaliveInfo.NOT_STARTED;
@@ -436,14 +439,13 @@
      **/
     public void startNattKeepalive(@Nullable NetworkAgentInfo nai,
             int intervalSeconds,
-            @NonNull Messenger messenger,
-            @NonNull IBinder binder,
+            @NonNull ISocketKeepaliveCallback cb,
             @NonNull String srcAddrString,
             int srcPort,
             @NonNull String dstAddrString,
             int dstPort) {
         if (nai == null) {
-            notifyMessenger(messenger, NO_KEEPALIVE, ERROR_INVALID_NETWORK);
+            notifyErrorCallback(cb, ERROR_INVALID_NETWORK);
             return;
         }
 
@@ -452,7 +454,7 @@
             srcAddress = NetworkUtils.numericToInetAddress(srcAddrString);
             dstAddress = NetworkUtils.numericToInetAddress(dstAddrString);
         } catch (IllegalArgumentException e) {
-            notifyMessenger(messenger, NO_KEEPALIVE, ERROR_INVALID_IP_ADDRESS);
+            notifyErrorCallback(cb, ERROR_INVALID_IP_ADDRESS);
             return;
         }
 
@@ -461,11 +463,12 @@
             packet = NattKeepalivePacketData.nattKeepalivePacket(
                     srcAddress, srcPort, dstAddress, NATT_PORT);
         } catch (InvalidPacketException e) {
-            notifyMessenger(messenger, NO_KEEPALIVE, e.error);
+            notifyErrorCallback(cb, e.error);
             return;
         }
-        KeepaliveInfo ki = new KeepaliveInfo(messenger, binder, nai, packet, intervalSeconds,
+        KeepaliveInfo ki = new KeepaliveInfo(cb, nai, packet, intervalSeconds,
                 KeepaliveInfo.TYPE_NATT, null);
+        Log.d(TAG, "Created keepalive: " + ki.toString());
         mConnectivityServiceHandler.obtainMessage(
                 NetworkAgent.CMD_START_SOCKET_KEEPALIVE, ki).sendToTarget();
     }
@@ -483,10 +486,9 @@
     public void startTcpKeepalive(@Nullable NetworkAgentInfo nai,
             @NonNull FileDescriptor fd,
             int intervalSeconds,
-            @NonNull Messenger messenger,
-            @NonNull IBinder binder) {
+            @NonNull ISocketKeepaliveCallback cb) {
         if (nai == null) {
-            notifyMessenger(messenger, NO_KEEPALIVE, ERROR_INVALID_NETWORK);
+            notifyErrorCallback(cb, ERROR_INVALID_NETWORK);
             return;
         }
 
@@ -500,10 +502,10 @@
             } catch (ErrnoException e1) {
                 Log.e(TAG, "Couldn't move fd out of repair mode after failure to start keepalive");
             }
-            notifyMessenger(messenger, NO_KEEPALIVE, e.error);
+            notifyErrorCallback(cb, e.error);
             return;
         }
-        KeepaliveInfo ki = new KeepaliveInfo(messenger, binder, nai, packet, intervalSeconds,
+        KeepaliveInfo ki = new KeepaliveInfo(cb, nai, packet, intervalSeconds,
                 KeepaliveInfo.TYPE_TCP, fd);
         Log.d(TAG, "Created keepalive: " + ki.toString());
         mConnectivityServiceHandler.obtainMessage(CMD_START_SOCKET_KEEPALIVE, ki).sendToTarget();
@@ -520,14 +522,13 @@
             @Nullable FileDescriptor fd,
             int resourceId,
             int intervalSeconds,
-            @NonNull Messenger messenger,
-            @NonNull IBinder binder,
+            @NonNull ISocketKeepaliveCallback cb,
             @NonNull String srcAddrString,
             @NonNull String dstAddrString,
             int dstPort) {
         // Ensure that the socket is created by IpSecService.
         if (!isNattKeepaliveSocketValid(fd, resourceId)) {
-            notifyMessenger(messenger, NO_KEEPALIVE, ERROR_INVALID_SOCKET);
+            notifyErrorCallback(cb, ERROR_INVALID_SOCKET);
         }
 
         // Get src port to adopt old API.
@@ -536,11 +537,11 @@
             final SocketAddress srcSockAddr = Os.getsockname(fd);
             srcPort = ((InetSocketAddress) srcSockAddr).getPort();
         } catch (ErrnoException e) {
-            notifyMessenger(messenger, NO_KEEPALIVE, ERROR_INVALID_SOCKET);
+            notifyErrorCallback(cb, ERROR_INVALID_SOCKET);
         }
 
         // Forward request to old API.
-        startNattKeepalive(nai, intervalSeconds, messenger, binder, srcAddrString, srcPort,
+        startNattKeepalive(nai, intervalSeconds, cb, srcAddrString, srcPort,
                 dstAddrString, dstPort);
     }
 
diff --git a/services/core/java/com/android/server/connectivity/NetworkNotificationManager.java b/services/core/java/com/android/server/connectivity/NetworkNotificationManager.java
index 053da0d..828a1e5 100644
--- a/services/core/java/com/android/server/connectivity/NetworkNotificationManager.java
+++ b/services/core/java/com/android/server/connectivity/NetworkNotificationManager.java
@@ -28,6 +28,7 @@
 import android.content.res.Resources;
 import android.net.wifi.WifiInfo;
 import android.os.UserHandle;
+import android.telephony.AccessNetworkConstants.TransportType;
 import android.telephony.TelephonyManager;
 import android.text.TextUtils;
 import android.util.Slog;
@@ -92,7 +93,7 @@
         return -1;
     }
 
-    private static String getTransportName(int transportType) {
+    private static String getTransportName(@TransportType int transportType) {
         Resources r = Resources.getSystem();
         String[] networkTypes = r.getStringArray(R.array.network_switch_type_name);
         try {
diff --git a/services/core/java/com/android/server/connectivity/PermissionMonitor.java b/services/core/java/com/android/server/connectivity/PermissionMonitor.java
index d84a4d2..123564e 100644
--- a/services/core/java/com/android/server/connectivity/PermissionMonitor.java
+++ b/services/core/java/com/android/server/connectivity/PermissionMonitor.java
@@ -46,13 +46,11 @@
 import android.util.Slog;
 import android.util.SparseIntArray;
 
-import com.android.internal.annotations.GuardedBy;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.util.ArrayUtils;
 import com.android.server.LocalServices;
 
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
@@ -84,20 +82,14 @@
     // Keys are App IDs. Values are true for SYSTEM permission and false for NETWORK permission.
     private final Map<Integer, Boolean> mApps = new HashMap<>();
 
-    // Keys are App packageNames, Values are app uids. . We need to keep track of this information
-    // because PackageListObserver#onPackageRemoved does not pass the UID.
-    @GuardedBy("mPackageNameUidMap")
-    private final Map<String, Integer> mPackageNameUidMap = new HashMap<>();
-
     private class PackageListObserver implements PackageManagerInternal.PackageListObserver {
         @Override
-        public void onPackageAdded(String packageName) {
+        public void onPackageAdded(String packageName, int uid) {
             final PackageInfo app = getPackageInfo(packageName);
             if (app == null) {
                 Slog.wtf(TAG, "Failed to get information of installed package: " + packageName);
                 return;
             }
-            int uid = (app.applicationInfo != null) ? app.applicationInfo.uid : INVALID_UID;
             if (uid == INVALID_UID) {
                 Slog.wtf(TAG, "Failed to get the uid of installed package: " + packageName
                         + "uid: " + uid);
@@ -107,29 +99,21 @@
                 return;
             }
             sendPackagePermissionsForUid(uid,
-                    filterPermission(Arrays.asList(app.requestedPermissions)));
-            synchronized (mPackageNameUidMap) {
-                mPackageNameUidMap.put(packageName, uid);
-            }
+                    getNetdPermissionMask(app.requestedPermissions));
         }
 
         @Override
-        public void onPackageRemoved(String packageName) {
-            int uid;
-            synchronized (mPackageNameUidMap) {
-                if (!mPackageNameUidMap.containsKey(packageName)) {
-                    return;
-                }
-                uid = mPackageNameUidMap.get(packageName);
-                mPackageNameUidMap.remove(packageName);
-            }
+        public void onPackageRemoved(String packageName, int uid) {
             int permission = 0;
+            // If there are still packages remain under the same uid, check the permission of the
+            // remaining packages. We only remove the permission for a given uid when all packages
+            // for that uid no longer have that permission.
             String[] packages = mPackageManager.getPackagesForUid(uid);
             if (packages != null && packages.length > 0) {
                 for (String name : packages) {
                     final PackageInfo app = getPackageInfo(name);
                     if (app != null && app.requestedPermissions != null) {
-                        permission |= filterPermission(Arrays.asList(app.requestedPermissions));
+                        permission |= getNetdPermissionMask(app.requestedPermissions);
                     }
                 }
             }
@@ -184,12 +168,9 @@
 
             //TODO: unify the management of the permissions into one codepath.
             if (app.requestedPermissions != null) {
-                int otherNetdPerms = filterPermission(Arrays.asList(app.requestedPermissions));
+                int otherNetdPerms = getNetdPermissionMask(app.requestedPermissions);
                 if (otherNetdPerms != 0) {
                     netdPermsUids.put(uid, netdPermsUids.get(uid) | otherNetdPerms);
-                    synchronized (mPackageNameUidMap) {
-                        mPackageNameUidMap.put(app.applicationInfo.packageName, uid);
-                    }
                 }
             }
         }
@@ -422,13 +403,15 @@
         }
     }
 
-    private static int filterPermission(List<String> requestedPermissions) {
+    private static int getNetdPermissionMask(String[] requestedPermissions) {
         int permissions = 0;
-        if (requestedPermissions.contains(INTERNET)) {
-            permissions |= INetd.PERMISSION_INTERNET;
-        }
-        if (requestedPermissions.contains(UPDATE_DEVICE_STATS)) {
-            permissions |= INetd.PERMISSION_UPDATE_DEVICE_STATS;
+        for (String permissionName : requestedPermissions) {
+            if (permissionName.equals(INTERNET)) {
+                permissions |= INetd.PERMISSION_INTERNET;
+            }
+            if (permissionName.equals(UPDATE_DEVICE_STATS)) {
+                permissions |= INetd.PERMISSION_UPDATE_DEVICE_STATS;
+            }
         }
         return permissions;
     }
@@ -439,8 +422,6 @@
                     | MATCH_ANY_USER);
             return app;
         } catch (NameNotFoundException e) {
-            // App not found.
-            loge("NameNotFoundException " + packageName);
             return null;
         }
     }
diff --git a/tests/net/Android.bp b/tests/net/Android.bp
index 2539c0f..c62d85e 100644
--- a/tests/net/Android.bp
+++ b/tests/net/Android.bp
@@ -13,7 +13,6 @@
         "mockito-target-minus-junit4",
         "platform-test-annotations",
         "services.core",
-        "services.ipmemorystore",
         "services.net",
     ],
     libs: [
diff --git a/tests/net/java/android/net/IpMemoryStoreTest.java b/tests/net/java/android/net/IpMemoryStoreTest.java
index 57ecc8f..18c6768 100644
--- a/tests/net/java/android/net/IpMemoryStoreTest.java
+++ b/tests/net/java/android/net/IpMemoryStoreTest.java
@@ -16,6 +16,9 @@
 
 package android.net;
 
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
+
 import android.content.Context;
 
 import androidx.test.filters.SmallTest;
@@ -33,13 +36,25 @@
     @Mock
     Context mMockContext;
     @Mock
+    NetworkStackClient mNetworkStackClient;
+    @Mock
     IIpMemoryStore mMockService;
     IpMemoryStore mStore;
 
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
-        mStore = new IpMemoryStore(mMockContext, mMockService);
+        doAnswer(invocation -> {
+            ((IIpMemoryStoreCallbacks) invocation.getArgument(0))
+                    .onIpMemoryStoreFetched(mMockService);
+            return null;
+        }).when(mNetworkStackClient).fetchIpMemoryStore(any());
+        mStore = new IpMemoryStore(mMockContext) {
+            @Override
+            protected NetworkStackClient getNetworkStackClient() {
+                return mNetworkStackClient;
+            }
+        };
     }
 
     @Test
diff --git a/tests/net/java/android/net/apf/ApfCapabilitiesTest.java b/tests/net/java/android/net/apf/ApfCapabilitiesTest.java
new file mode 100644
index 0000000..75752c3
--- /dev/null
+++ b/tests/net/java/android/net/apf/ApfCapabilitiesTest.java
@@ -0,0 +1,50 @@
+/*
+ * Copyright (C) 2019 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.net.apf;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotEquals;
+
+import android.net.shared.ParcelableTestUtil;
+
+import androidx.test.filters.SmallTest;
+import androidx.test.runner.AndroidJUnit4;
+
+import com.android.internal.util.TestUtils;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+@RunWith(AndroidJUnit4.class)
+@SmallTest
+public class ApfCapabilitiesTest {
+    @Test
+    public void testParcelUnparcel() {
+        final ApfCapabilities caps = new ApfCapabilities(123, 456, 789);
+        ParcelableTestUtil.assertFieldCountEquals(3, ApfCapabilities.class);
+
+        TestUtils.assertParcelingIsLossless(caps, ApfCapabilities.CREATOR);
+    }
+
+    @Test
+    public void testEquals() {
+        assertEquals(new ApfCapabilities(1, 2, 3), new ApfCapabilities(1, 2, 3));
+        assertNotEquals(new ApfCapabilities(2, 2, 3), new ApfCapabilities(1, 2, 3));
+        assertNotEquals(new ApfCapabilities(1, 3, 3), new ApfCapabilities(1, 2, 3));
+        assertNotEquals(new ApfCapabilities(1, 2, 4), new ApfCapabilities(1, 2, 3));
+    }
+}
diff --git a/tests/net/java/com/android/server/ConnectivityServiceTest.java b/tests/net/java/com/android/server/ConnectivityServiceTest.java
index c14d0af..952c0da 100644
--- a/tests/net/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/net/java/com/android/server/ConnectivityServiceTest.java
@@ -60,10 +60,10 @@
 import static android.net.NetworkPolicyManager.RULE_NONE;
 import static android.net.NetworkPolicyManager.RULE_REJECT_ALL;
 import static android.net.NetworkPolicyManager.RULE_REJECT_METERED;
-import static android.net.shared.NetworkParcelableUtil.fromStableParcelable;
 
 import static com.android.internal.util.TestUtils.waitForIdleHandler;
 import static com.android.internal.util.TestUtils.waitForIdleLooper;
+import static com.android.internal.util.TestUtils.waitForIdleSerialExecutor;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
@@ -88,6 +88,7 @@
 import static org.mockito.Mockito.verifyNoMoreInteractions;
 import static org.mockito.Mockito.when;
 
+import android.annotation.NonNull;
 import android.app.NotificationManager;
 import android.app.PendingIntent;
 import android.content.BroadcastReceiver;
@@ -123,7 +124,6 @@
 import android.net.NetworkInfo;
 import android.net.NetworkInfo.DetailedState;
 import android.net.NetworkMisc;
-import android.net.NetworkParcelable;
 import android.net.NetworkRequest;
 import android.net.NetworkSpecifier;
 import android.net.NetworkStackClient;
@@ -500,8 +500,7 @@
                 fail(e.getMessage());
             }
 
-            final ArgumentCaptor<NetworkParcelable> nmNetworkCaptor =
-                    ArgumentCaptor.forClass(NetworkParcelable.class);
+            final ArgumentCaptor<Network> nmNetworkCaptor = ArgumentCaptor.forClass(Network.class);
             final ArgumentCaptor<INetworkMonitorCallbacks> nmCbCaptor =
                     ArgumentCaptor.forClass(INetworkMonitorCallbacks.class);
             doNothing().when(mNetworkStack).makeNetworkMonitor(
@@ -541,8 +540,7 @@
                 }
             };
 
-            assertEquals(
-                    mNetworkAgent.netId, fromStableParcelable(nmNetworkCaptor.getValue()).netId);
+            assertEquals(mNetworkAgent.netId, nmNetworkCaptor.getValue().netId);
             mNmCallbacks = nmCbCaptor.getValue();
 
             try {
@@ -3761,7 +3759,7 @@
             }
         }
 
-        private LinkedBlockingQueue<CallbackValue> mCallbacks = new LinkedBlockingQueue<>();
+        private final LinkedBlockingQueue<CallbackValue> mCallbacks = new LinkedBlockingQueue<>();
 
         @Override
         public void onStarted() {
@@ -3836,6 +3834,11 @@
         }
 
         private LinkedBlockingQueue<CallbackValue> mCallbacks = new LinkedBlockingQueue<>();
+        private final Executor mExecutor;
+
+        TestSocketKeepaliveCallback(@NonNull Executor executor) {
+            mExecutor = executor;
+        }
 
         @Override
         public void onStarted() {
@@ -3873,6 +3876,12 @@
         public void expectError(int error) {
             expectCallback(new CallbackValue(CallbackType.ON_ERROR, error));
         }
+
+        public void assertNoCallback() {
+            waitForIdleSerialExecutor(mExecutor, TIMEOUT_MS);
+            CallbackValue cv = mCallbacks.peek();
+            assertNull("Unexpected callback: " + cv, cv);
+        }
     }
 
     private Network connectKeepaliveNetwork(LinkProperties lp) {
@@ -3979,19 +3988,6 @@
         myNet = connectKeepaliveNetwork(lp);
         mWiFiNetworkAgent.setStartKeepaliveError(PacketKeepalive.SUCCESS);
 
-        // Check things work as expected when the keepalive is stopped and the network disconnects.
-        ka = mCm.startNattKeepalive(myNet, validKaInterval, callback, myIPv4, 12345, dstIPv4);
-        callback.expectStarted();
-        ka.stop();
-        mWiFiNetworkAgent.disconnect();
-        waitFor(mWiFiNetworkAgent.getDisconnectedCV());
-        waitForIdle();
-        callback.expectStopped();
-
-        // Reconnect.
-        myNet = connectKeepaliveNetwork(lp);
-        mWiFiNetworkAgent.setStartKeepaliveError(PacketKeepalive.SUCCESS);
-
         // Check that keepalive slots start from 1 and increment. The first one gets slot 1.
         mWiFiNetworkAgent.setExpectedKeepaliveSlot(1);
         ka = mCm.startNattKeepalive(myNet, validKaInterval, callback, myIPv4, 12345, dstIPv4);
@@ -4067,7 +4063,7 @@
         Network notMyNet = new Network(61234);
         Network myNet = connectKeepaliveNetwork(lp);
 
-        TestSocketKeepaliveCallback callback = new TestSocketKeepaliveCallback();
+        TestSocketKeepaliveCallback callback = new TestSocketKeepaliveCallback(executor);
         SocketKeepalive ka;
 
         // Attempt to start keepalives with invalid parameters and check for errors.
@@ -4110,6 +4106,22 @@
         ka.stop();
         callback.expectStopped();
 
+        // Check that keepalive could be restarted.
+        ka.start(validKaInterval);
+        callback.expectStarted();
+        ka.stop();
+        callback.expectStopped();
+
+        // Check that keepalive can be restarted without waiting for callback.
+        ka.start(validKaInterval);
+        callback.expectStarted();
+        ka.stop();
+        ka.start(validKaInterval);
+        callback.expectStopped();
+        callback.expectStarted();
+        ka.stop();
+        callback.expectStopped();
+
         // Check that deleting the IP address stops the keepalive.
         LinkProperties bogusLp = new LinkProperties(lp);
         ka = mCm.createSocketKeepalive(myNet, testSocket, myIPv4, dstIPv4, executor, callback);
@@ -4134,20 +4146,7 @@
         final Network myNetAlias = myNet;
         assertNull(mCm.getNetworkCapabilities(myNetAlias));
         ka.stop();
-
-        // Reconnect.
-        myNet = connectKeepaliveNetwork(lp);
-        mWiFiNetworkAgent.setStartKeepaliveError(SocketKeepalive.SUCCESS);
-
-        // Check things work as expected when the keepalive is stopped and the network disconnects.
-        ka = mCm.createSocketKeepalive(myNet, testSocket, myIPv4, dstIPv4, executor, callback);
-        ka.start(validKaInterval);
-        callback.expectStarted();
-        ka.stop();
-        mWiFiNetworkAgent.disconnect();
-        waitFor(mWiFiNetworkAgent.getDisconnectedCV());
-        waitForIdle();
-        callback.expectStopped();
+        callback.assertNoCallback();
 
         // Reconnect.
         myNet = connectKeepaliveNetwork(lp);
@@ -4162,7 +4161,7 @@
         // The second one gets slot 2.
         mWiFiNetworkAgent.setExpectedKeepaliveSlot(2);
         final UdpEncapsulationSocket testSocket2 = mIpSec.openUdpEncapsulationSocket(6789);
-        TestSocketKeepaliveCallback callback2 = new TestSocketKeepaliveCallback();
+        TestSocketKeepaliveCallback callback2 = new TestSocketKeepaliveCallback(executor);
         SocketKeepalive ka2 =
                 mCm.createSocketKeepalive(myNet, testSocket2, myIPv4, dstIPv4, executor, callback2);
         ka2.start(validKaInterval);
@@ -4215,7 +4214,7 @@
         final Socket testSocketV4 = new Socket();
         final Socket testSocketV6 = new Socket();
 
-        TestSocketKeepaliveCallback callback = new TestSocketKeepaliveCallback();
+        TestSocketKeepaliveCallback callback = new TestSocketKeepaliveCallback(executor);
         SocketKeepalive ka;
 
         // Attempt to start Tcp keepalives with invalid parameters and check for errors.
diff --git a/tests/net/java/com/android/server/net/NetworkStatsCollectionTest.java b/tests/net/java/com/android/server/net/NetworkStatsCollectionTest.java
index 830c928..9b4f49c 100644
--- a/tests/net/java/com/android/server/net/NetworkStatsCollectionTest.java
+++ b/tests/net/java/com/android/server/net/NetworkStatsCollectionTest.java
@@ -101,6 +101,7 @@
     @After
     public void tearDown() throws Exception {
         RecurrenceRule.sClock = sOriginalClock;
+        NetworkTemplate.resetForceAllNetworkTypes();
     }
 
     private void setClock(Instant instant) {
diff --git a/tests/net/java/com/android/server/net/NetworkStatsServiceTest.java b/tests/net/java/com/android/server/net/NetworkStatsServiceTest.java
index 598448b..bce526d 100644
--- a/tests/net/java/com/android/server/net/NetworkStatsServiceTest.java
+++ b/tests/net/java/com/android/server/net/NetworkStatsServiceTest.java
@@ -19,6 +19,7 @@
 import static android.content.Intent.ACTION_UID_REMOVED;
 import static android.content.Intent.EXTRA_UID;
 import static android.net.ConnectivityManager.TYPE_MOBILE;
+import static android.net.ConnectivityManager.TYPE_VPN;
 import static android.net.ConnectivityManager.TYPE_WIFI;
 import static android.net.ConnectivityManager.TYPE_WIMAX;
 import static android.net.NetworkStats.DEFAULT_NETWORK_ALL;
@@ -41,6 +42,7 @@
 import static android.net.NetworkStats.UID_ALL;
 import static android.net.NetworkStatsHistory.FIELD_ALL;
 import static android.net.NetworkTemplate.buildTemplateMobileAll;
+import static android.net.NetworkTemplate.buildTemplateMobileWildcard;
 import static android.net.NetworkTemplate.buildTemplateWifiWildcard;
 import static android.net.TrafficStats.MB_IN_BYTES;
 import static android.net.TrafficStats.UID_REMOVED;
@@ -132,6 +134,8 @@
 
     private static final String TEST_IFACE = "test0";
     private static final String TEST_IFACE2 = "test1";
+    private static final String TUN_IFACE = "test_nss_tun0";
+
     private static final long TEST_START = 1194220800000L;
 
     private static final String IMSI_1 = "310004";
@@ -145,10 +149,12 @@
     private static final int UID_RED = 1001;
     private static final int UID_BLUE = 1002;
     private static final int UID_GREEN = 1003;
-
+    private static final int UID_VPN = 1004;
 
     private static final Network WIFI_NETWORK =  new Network(100);
     private static final Network MOBILE_NETWORK =  new Network(101);
+    private static final Network VPN_NETWORK = new Network(102);
+
     private static final Network[] NETWORKS_WIFI = new Network[]{ WIFI_NETWORK };
     private static final Network[] NETWORKS_MOBILE = new Network[]{ MOBILE_NETWORK };
 
@@ -914,7 +920,113 @@
         assertNetworkTotal(sTemplateImsi1, 2048L, 16L, 512L, 4L, 0);
         assertUidTotal(sTemplateImsi1, UID_RED, 128L, 2L, 128L, 2L, 0);
         assertUidTotal(sTemplateImsi1, UID_TETHERING, 1920L, 14L, 384L, 2L, 0);
+    }
 
+    @Test
+    public void vpnWithOneUnderlyingIface() throws Exception {
+        // WiFi network is connected and VPN is using WiFi (which has TEST_IFACE).
+        expectDefaultSettings();
+        NetworkState[] networkStates = new NetworkState[] {buildWifiState(), buildVpnState()};
+        VpnInfo[] vpnInfos = new VpnInfo[] {createVpnInfo(TEST_IFACE)};
+        expectNetworkStatsUidDetail(buildEmptyStats());
+        expectBandwidthControlCheck();
+
+        mService.forceUpdateIfaces(
+                new Network[] {WIFI_NETWORK, VPN_NETWORK},
+                vpnInfos,
+                networkStates,
+                getActiveIface(networkStates));
+        // create some traffic (assume 10 bytes of MTU for VPN interface and 1 byte encryption
+        // overhead per packet):
+        // 1000 bytes (100 packets) were sent/received by UID_RED over VPN.
+        // 500 bytes (50 packets) were sent/received by UID_BLUE over VPN.
+        // VPN sent/received 1650 bytes (150 packets) over WiFi.
+        // Of 1650 bytes over WiFi, expect 1000 bytes attributed to UID_RED, 500 bytes attributed to
+        // UID_BLUE, and 150 bytes attributed to UID_VPN for both rx/tx traffic.
+        incrementCurrentTime(HOUR_IN_MILLIS);
+        expectNetworkStatsUidDetail(new NetworkStats(getElapsedRealtime(), 3)
+                .addValues(TUN_IFACE, UID_RED, SET_DEFAULT, TAG_NONE, 1000L, 100L, 1000L, 100L, 1L)
+                .addValues(TUN_IFACE, UID_BLUE, SET_DEFAULT, TAG_NONE, 500L, 50L, 500L, 50L, 1L)
+                .addValues(
+                    TEST_IFACE, UID_VPN, SET_DEFAULT, TAG_NONE, 1650L, 150L, 1650L, 150L, 2L));
+
+        forcePollAndWaitForIdle();
+
+        assertUidTotal(sTemplateWifi, UID_RED, 1000L, 100L, 1000L, 100L, 1);
+        assertUidTotal(sTemplateWifi, UID_BLUE, 500L, 50L, 500L, 50L, 1);
+        assertUidTotal(sTemplateWifi, UID_VPN, 150L, 0L, 150L, 0L, 2);
+    }
+
+    @Test
+    public void vpnWithOneUnderlyingIface_withCompression() throws Exception {
+        // WiFi network is connected and VPN is using WiFi (which has TEST_IFACE).
+        expectDefaultSettings();
+        NetworkState[] networkStates = new NetworkState[] {buildWifiState(), buildVpnState()};
+        VpnInfo[] vpnInfos = new VpnInfo[] {createVpnInfo(TEST_IFACE)};
+        expectNetworkStatsUidDetail(buildEmptyStats());
+        expectBandwidthControlCheck();
+
+        mService.forceUpdateIfaces(
+                new Network[] {WIFI_NETWORK, VPN_NETWORK},
+                vpnInfos,
+                networkStates,
+                getActiveIface(networkStates));
+        // create some traffic (assume 10 bytes of MTU for VPN interface and 1 byte encryption
+        // overhead per packet):
+        // 1000 bytes (100 packets) were sent/received by UID_RED over VPN.
+        // 3000 bytes (300 packets) were sent/received by UID_BLUE over VPN.
+        // VPN sent/received 1000 bytes (100 packets) over WiFi.
+        // Of 1000 bytes over WiFi, expect 250 bytes attributed UID_RED and 750 bytes to UID_BLUE,
+        // with nothing attributed to UID_VPN for both rx/tx traffic.
+        incrementCurrentTime(HOUR_IN_MILLIS);
+        expectNetworkStatsUidDetail(new NetworkStats(getElapsedRealtime(), 3)
+                .addValues(TUN_IFACE, UID_RED, SET_DEFAULT, TAG_NONE, 1000L, 100L, 1000L, 100L, 1L)
+                .addValues(TUN_IFACE, UID_BLUE, SET_DEFAULT, TAG_NONE, 3000L, 300L, 3000L, 300L, 1L)
+                .addValues(
+                    TEST_IFACE, UID_VPN, SET_DEFAULT, TAG_NONE, 1000L, 100L, 1000L, 100L, 0L));
+
+        forcePollAndWaitForIdle();
+
+        assertUidTotal(sTemplateWifi, UID_RED, 250L, 25L, 250L, 25L, 0);
+        assertUidTotal(sTemplateWifi, UID_BLUE, 750L, 75L, 750L, 75L, 0);
+        assertUidTotal(sTemplateWifi, UID_VPN, 0L, 0L, 0L, 0L, 0);
+    }
+
+    @Test
+    public void vpnWithIncorrectUnderlyingIface() throws Exception {
+        // WiFi and Cell networks are connected and VPN is using Cell (which has TEST_IFACE2),
+        // but has declared only WiFi (TEST_IFACE) in its underlying network set.
+        expectDefaultSettings();
+        NetworkState[] networkStates =
+                new NetworkState[] {
+                    buildWifiState(), buildMobile4gState(TEST_IFACE2), buildVpnState()
+                };
+        VpnInfo[] vpnInfos = new VpnInfo[] {createVpnInfo(TEST_IFACE)};
+        expectNetworkStatsUidDetail(buildEmptyStats());
+        expectBandwidthControlCheck();
+
+        mService.forceUpdateIfaces(
+                new Network[] {WIFI_NETWORK, VPN_NETWORK},
+                vpnInfos,
+                networkStates,
+                getActiveIface(networkStates));
+        // create some traffic (assume 10 bytes of MTU for VPN interface and 1 byte encryption
+        // overhead per packet):
+        // 1000 bytes (100 packets) were sent/received by UID_RED over VPN.
+        // VPN sent/received 1100 bytes (100 packets) over Cell.
+        // Of 1100 bytes over Cell, expect all of it attributed to UID_VPN for both rx/tx traffic.
+        incrementCurrentTime(HOUR_IN_MILLIS);
+        expectNetworkStatsUidDetail(new NetworkStats(getElapsedRealtime(), 2)
+                .addValues(TUN_IFACE, UID_RED, SET_DEFAULT, TAG_NONE, 1000L, 100L, 1000L, 100L, 1L)
+                .addValues(
+                    TEST_IFACE2, UID_VPN, SET_DEFAULT, TAG_NONE, 1100L, 100L, 1100L, 100L, 1L));
+
+        forcePollAndWaitForIdle();
+
+        assertUidTotal(sTemplateWifi, UID_RED, 0L, 0L, 0L, 0L, 0);
+        assertUidTotal(sTemplateWifi, UID_VPN, 0L, 0L, 0L, 0L, 0);
+        assertUidTotal(buildTemplateMobileWildcard(), UID_RED, 0L, 0L, 0L, 0L, 0);
+        assertUidTotal(buildTemplateMobileWildcard(), UID_VPN, 1100L, 100L, 1100L, 100L, 1);
     }
 
     @Test
@@ -1262,6 +1374,22 @@
         return new NetworkStats(getElapsedRealtime(), 0);
     }
 
+    private static NetworkState buildVpnState() {
+        final NetworkInfo info = new NetworkInfo(TYPE_VPN, 0, null, null);
+        info.setDetailedState(DetailedState.CONNECTED, null, null);
+        final LinkProperties prop = new LinkProperties();
+        prop.setInterfaceName(TUN_IFACE);
+        return new NetworkState(info, prop, new NetworkCapabilities(), VPN_NETWORK, null, null);
+    }
+
+    private static VpnInfo createVpnInfo(String underlyingIface) {
+        VpnInfo info = new VpnInfo();
+        info.ownerUid = UID_VPN;
+        info.vpnIface = TUN_IFACE;
+        info.primaryUnderlyingIface = underlyingIface;
+        return info;
+    }
+
     private long getElapsedRealtime() {
         return mElapsedRealtime;
     }