Merge "Add uid information in PackageListObserver"
diff --git a/core/java/android/net/ConnectivityManager.java b/core/java/android/net/ConnectivityManager.java
index d08379f..e5802c2 100644
--- a/core/java/android/net/ConnectivityManager.java
+++ b/core/java/android/net/ConnectivityManager.java
@@ -2816,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
@@ -2878,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/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/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;
     }