diff --git a/core/java/android/net/TrafficStats.java b/core/java/android/net/TrafficStats.java
index 954e59c..d701550 100644
--- a/core/java/android/net/TrafficStats.java
+++ b/core/java/android/net/TrafficStats.java
@@ -19,6 +19,7 @@
 import android.annotation.RequiresPermission;
 import android.annotation.SuppressLint;
 import android.annotation.SystemApi;
+import android.annotation.TestApi;
 import android.app.DownloadManager;
 import android.app.backup.BackupManager;
 import android.app.usage.NetworkStatsManager;
@@ -154,6 +155,8 @@
 
     private static Object sProfilingLock = new Object();
 
+    private static final String LOOPBACK_IFACE = "lo";
+
     /**
      * Set active tag to use when accounting {@link Socket} traffic originating
      * from the current thread. Only one active tag per thread is supported.
@@ -542,6 +545,30 @@
         return nativeGetIfaceStat(iface, TYPE_RX_BYTES);
     }
 
+    /** {@hide} */
+    @TestApi
+    public static long getLoopbackTxPackets() {
+        return nativeGetIfaceStat(LOOPBACK_IFACE, TYPE_TX_PACKETS);
+    }
+
+    /** {@hide} */
+    @TestApi
+    public static long getLoopbackRxPackets() {
+        return nativeGetIfaceStat(LOOPBACK_IFACE, TYPE_RX_PACKETS);
+    }
+
+    /** {@hide} */
+    @TestApi
+    public static long getLoopbackTxBytes() {
+        return nativeGetIfaceStat(LOOPBACK_IFACE, TYPE_TX_BYTES);
+    }
+
+    /** {@hide} */
+    @TestApi
+    public static long getLoopbackRxBytes() {
+        return nativeGetIfaceStat(LOOPBACK_IFACE, TYPE_RX_BYTES);
+    }
+
     /**
      * Return number of packets transmitted since device boot. Counts packets
      * across all network interfaces, and always increases monotonically since
diff --git a/services/core/java/com/android/server/IpSecService.java b/services/core/java/com/android/server/IpSecService.java
index e9eb3b3..a764808 100644
--- a/services/core/java/com/android/server/IpSecService.java
+++ b/services/core/java/com/android/server/IpSecService.java
@@ -34,6 +34,7 @@
 import android.net.IpSecTransformResponse;
 import android.net.IpSecUdpEncapResponse;
 import android.net.NetworkUtils;
+import android.net.TrafficStats;
 import android.net.util.NetdService;
 import android.os.Binder;
 import android.os.IBinder;
@@ -120,6 +121,7 @@
     }
 
     private final IpSecServiceConfiguration mSrvConfig;
+    final UidFdTagger mUidFdTagger;
 
     /**
      * Interface for user-reference and kernel-resource cleanup.
@@ -762,8 +764,23 @@
     /** @hide */
     @VisibleForTesting
     public IpSecService(Context context, IpSecServiceConfiguration config) {
+        this(context, config, (fd, uid) ->  {
+            try{
+                TrafficStats.setThreadStatsUid(uid);
+                TrafficStats.tagFileDescriptor(fd);
+            } finally {
+                TrafficStats.clearThreadStatsUid();
+            }
+        });
+    }
+
+    /** @hide */
+    @VisibleForTesting
+    public IpSecService(
+            Context context, IpSecServiceConfiguration config, UidFdTagger uidFdTagger) {
         mContext = context;
         mSrvConfig = config;
+        mUidFdTagger = uidFdTagger;
     }
 
     public void systemReady() {
@@ -925,6 +942,26 @@
     }
 
     /**
+     * Functional interface to do traffic tagging of given sockets to UIDs.
+     *
+     * <p>Specifically used by openUdpEncapsulationSocket to ensure data usage on the UDP encap
+     * sockets are billed to the UID that the UDP encap socket was created on behalf of.
+     *
+     * <p>Separate class so that the socket tagging logic can be mocked; TrafficStats uses static
+     * methods that cannot be easily mocked/tested.
+     */
+    @VisibleForTesting
+    public interface UidFdTagger {
+        /**
+         * Sets socket tag to assign all traffic to the provided UID.
+         *
+         * <p>Since the socket is created on behalf of an unprivileged application, all traffic
+         * should be accounted to the UID of the unprivileged application.
+         */
+        public void tag(FileDescriptor fd, int uid) throws IOException;
+    }
+
+    /**
      * Open a socket via the system server and bind it to the specified port (random if port=0).
      * This will return a PFD to the user that represent a bound UDP socket. The system server will
      * cache the socket and a record of its owner so that it can and must be freed when no longer
@@ -939,7 +976,8 @@
         }
         checkNotNull(binder, "Null Binder passed to openUdpEncapsulationSocket");
 
-        UserRecord userRecord = mUserResourceTracker.getUserRecord(Binder.getCallingUid());
+        int callingUid = Binder.getCallingUid();
+        UserRecord userRecord = mUserResourceTracker.getUserRecord(callingUid);
         int resourceId = mNextResourceId.getAndIncrement();
         FileDescriptor sockFd = null;
         try {
@@ -948,6 +986,7 @@
             }
 
             sockFd = Os.socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
+            mUidFdTagger.tag(sockFd, callingUid);
 
             if (port != 0) {
                 Log.v(TAG, "Binding to port " + port);