[Cronet] Add TrafficStats tagging API to CronetHttpURLConnection

This just plumbs it down to the underlying UrlRequest.

Cq-Include-Trybots: master.tryserver.chromium.android:android_cronet_tester
Change-Id: Ifed23818a66e9dffe6317b2ecc0d8fefcedced52
Reviewed-on: https://chromium-review.googlesource.com/1100351
Commit-Queue: Paul Jensen <pauljensen@chromium.org>
Reviewed-by: Helen Li <xunjieli@chromium.org>
Cr-Commit-Position: refs/heads/master@{#567755}

components/cronet/android/java/src/org/chromium/net/urlconnection/CronetHttpURLConnection.java

@@ -10,6 +10,7 @@ import android.util.Pair;
 
 import org.chromium.net.CronetEngine;
 import org.chromium.net.CronetException;
+import org.chromium.net.ExperimentalUrlRequest;
 import org.chromium.net.UrlRequest;
 import org.chromium.net.UrlResponseInfo;
 
@@ -42,6 +43,10 @@ public class CronetHttpURLConnection extends HttpURLConnection {
     private final MessageLoop mMessageLoop;
     private UrlRequest mRequest;
     private final List<Pair<String, String>> mRequestHeaders;
+    private boolean mTrafficStatsTagSet;
+    private int mTrafficStatsTag;
+    private boolean mTrafficStatsUidSet;
+    private int mTrafficStatsUid;
 
     private CronetInputStream mInputStream;
     private CronetOutputStream mOutputStream;
@@ -256,8 +261,9 @@ public class CronetHttpURLConnection extends HttpURLConnection {
         if (connected) {
             return;
         }
-        final UrlRequest.Builder requestBuilder = mCronetEngine.newUrlRequestBuilder(
-                getURL().toString(), new CronetUrlRequestCallback(), mMessageLoop);
+        final ExperimentalUrlRequest.Builder requestBuilder =
+                (ExperimentalUrlRequest.Builder) mCronetEngine.newUrlRequestBuilder(
+                        getURL().toString(), new CronetUrlRequestCallback(), mMessageLoop);
         if (doOutput) {
             if (method.equals("GET")) {
                 method = "POST";
@@ -291,6 +297,12 @@ public class CronetHttpURLConnection extends HttpURLConnection {
         }
         // Set HTTP method.
         requestBuilder.setHttpMethod(method);
+        if (mTrafficStatsTagSet) {
+            requestBuilder.setTrafficStatsTag(mTrafficStatsTag);
+        }
+        if (mTrafficStatsUidSet) {
+            requestBuilder.setTrafficStatsUid(mTrafficStatsUid);
+        }
 
         connected = true;
         mRequest = requestBuilder.build();
@@ -422,6 +434,54 @@ public class CronetHttpURLConnection extends HttpURLConnection {
         mMessageLoop.loop(getReadTimeout());
     }
 
+    /**
+     * Sets {@link android.net.TrafficStats} tag to use when accounting socket traffic caused by
+     * this request. See {@link android.net.TrafficStats} for more information. If no tag is
+     * set (e.g. this method isn't called), then Android accounts for the socket traffic caused
+     * by this request as if the tag value were set to 0.
+     * <p>
+     * <b>NOTE:</b>Setting a tag disallows sharing of sockets with requests
+     * with other tags, which may adversely effect performance by prohibiting
+     * connection sharing. In other words use of multiplexed sockets (e.g. HTTP/2
+     * and QUIC) will only be allowed if all requests have the same socket tag.
+     *
+     * @param tag the tag value used to when accounting for socket traffic caused by this
+     *            request. Tags between 0xFFFFFF00 and 0xFFFFFFFF are reserved and used
+     *            internally by system services like {@link android.app.DownloadManager} when
+     *            performing traffic on behalf of an application.
+     */
+    public void setTrafficStatsTag(int tag) {
+        if (connected) {
+            throw new IllegalStateException(
+                    "Cannot modify traffic stats tag after connection is made.");
+        }
+        mTrafficStatsTagSet = true;
+        mTrafficStatsTag = tag;
+    }
+
+    /**
+     * Sets specific UID to use when accounting socket traffic caused by this request. See
+     * {@link android.net.TrafficStats} for more information. Designed for use when performing
+     * an operation on behalf of another application. Caller must hold
+     * {@link android.Manifest.permission#MODIFY_NETWORK_ACCOUNTING} permission. By default
+     * traffic is attributed to UID of caller.
+     * <p>
+     * <b>NOTE:</b>Setting a UID disallows sharing of sockets with requests
+     * with other UIDs, which may adversely effect performance by prohibiting
+     * connection sharing. In other words use of multiplexed sockets (e.g. HTTP/2
+     * and QUIC) will only be allowed if all requests have the same UID set.
+     *
+     * @param uid the UID to attribute socket traffic caused by this request.
+     */
+    public void setTrafficStatsUid(int uid) {
+        if (connected) {
+            throw new IllegalStateException(
+                    "Cannot modify traffic stats UID after connection is made.");
+        }
+        mTrafficStatsUidSet = true;
+        mTrafficStatsUid = uid;
+    }
+
     /**
      * Returns the index of request header in {@link #mRequestHeaders} or
      * -1 if not found.

components/cronet/android/test/javatests/src/org/chromium/net/urlconnection/CronetHttpURLConnectionTest.java

@@ -13,6 +13,7 @@ import static org.junit.Assert.fail;
 import static org.chromium.net.CronetTestRule.getContext;
 
 import android.os.Build;
+import android.os.Process;
 import android.support.test.filters.SmallTest;
 
 import org.junit.After;
@@ -28,6 +29,8 @@ import org.chromium.net.CronetException;
 import org.chromium.net.CronetTestRule;
 import org.chromium.net.CronetTestRule.CompareDefaultWithCronet;
 import org.chromium.net.CronetTestRule.OnlyRunCronetHttpURLConnection;
+import org.chromium.net.CronetTestRule.RequiresMinApi;
+import org.chromium.net.CronetTestUtil;
 import org.chromium.net.MockUrlRequestJobFactory;
 import org.chromium.net.NativeTestServer;
 
@@ -1348,4 +1351,55 @@ public class CronetHttpURLConnectionTest {
         }
         return headerValues;
     }
+
+    @Test
+    @SmallTest
+    @Feature({"Cronet"})
+    @RequiresMinApi(9) // Tagging support added in API level 9: crrev.com/c/chromium/src/+/930086
+    public void testTagging() throws Exception {
+        URL url = new URL(NativeTestServer.getEchoMethodURL());
+
+        // Test untagged requests are given tag 0.
+        int tag = 0;
+        long priorBytes = CronetTestUtil.nativeGetTaggedBytes(tag);
+        CronetHttpURLConnection urlConnection = (CronetHttpURLConnection) url.openConnection();
+        assertEquals(200, urlConnection.getResponseCode());
+        urlConnection.disconnect();
+        assertTrue(CronetTestUtil.nativeGetTaggedBytes(tag) > priorBytes);
+
+        // Test explicit tagging.
+        tag = 0x12345678;
+        priorBytes = CronetTestUtil.nativeGetTaggedBytes(tag);
+        urlConnection = (CronetHttpURLConnection) url.openConnection();
+        urlConnection.setTrafficStatsTag(tag);
+        assertEquals(200, urlConnection.getResponseCode());
+        urlConnection.disconnect();
+        assertTrue(CronetTestUtil.nativeGetTaggedBytes(tag) > priorBytes);
+
+        // Test a different tag value.
+        tag = 0x87654321;
+        priorBytes = CronetTestUtil.nativeGetTaggedBytes(tag);
+        urlConnection = (CronetHttpURLConnection) url.openConnection();
+        urlConnection.setTrafficStatsTag(tag);
+        assertEquals(200, urlConnection.getResponseCode());
+        urlConnection.disconnect();
+        assertTrue(CronetTestUtil.nativeGetTaggedBytes(tag) > priorBytes);
+
+        // Test tagging with our UID.
+        // NOTE(pauljensen): Explicitly setting the UID to the current UID isn't a particularly
+        // thorough test of this API but at least provides coverage of the underlying code, and
+        // verifies that traffic is still properly attributed.
+        // The code path for UID is parallel to that for the tag, which we do have more thorough
+        // testing for.  More thorough testing of setting the UID would require running tests with
+        // a rare permission which isn't realistic for most apps.  Apps are allowed to set the UID
+        // to their own UID as per this logic in the tagging kernel module:
+        // https://android.googlesource.com/kernel/common/+/21dd5d7/net/netfilter/xt_qtaguid.c#154
+        tag = 0;
+        priorBytes = CronetTestUtil.nativeGetTaggedBytes(tag);
+        urlConnection = (CronetHttpURLConnection) url.openConnection();
+        urlConnection.setTrafficStatsUid(Process.myUid());
+        assertEquals(200, urlConnection.getResponseCode());
+        urlConnection.disconnect();
+        assertTrue(CronetTestUtil.nativeGetTaggedBytes(tag) > priorBytes);
+    }
 }