[ipsec-qtaguid] Tag sockets upon creation of encap sockets

Added calls to tag encap sockets to that of the UID for which the encap
socket is being created on behalf of. This ensures that all data
accounting generated for the UDP-encap-ESP socket is correctly billed to
the right UID.

Bug: 62994731
Test: New tests added to IpSecServiceTest.java, passing
Change-Id: I15365ea9c982fd7b4e3cdeff314ddfba2289c86e

core/java/android/net/TrafficStats.java

@@ -19,6 +19,7 @@ package android.net;
 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 @@ public class TrafficStats {
 
     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 @@ public class TrafficStats {
         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

services/core/java/com/android/server/IpSecService.java

@@ -34,6 +34,7 @@ import android.net.IpSecTransform;
 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 @@ public class IpSecService extends IIpSecService.Stub {
     }
 
     private final IpSecServiceConfiguration mSrvConfig;
+    final UidFdTagger mUidFdTagger;
 
     /**
      * Interface for user-reference and kernel-resource cleanup.
@@ -762,8 +764,23 @@ public class IpSecService extends IIpSecService.Stub {
     /** @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() {
@@ -924,6 +941,26 @@ public class IpSecService extends IIpSecService.Stub {
         throw new IOException("Failed " + MAX_PORT_BIND_ATTEMPTS + " attempts to bind to a port");
     }
 
+    /**
+     * 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
@@ -939,7 +976,8 @@ public class IpSecService extends IIpSecService.Stub {
         }
         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 @@ public class IpSecService extends IIpSecService.Stub {
             }
 
             sockFd = Os.socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
+            mUidFdTagger.tag(sockFd, callingUid);
 
             if (port != 0) {
                 Log.v(TAG, "Binding to port " + port);