WIP: Add TrafficStats

2 files changed, 7 insertions, 0 deletions

diff --git a/tunnel/src/main/java/com/wireguard/android/backend/GoBackend.java b/tunnel/src/main/java/com/wireguard/android/backend/GoBackend.java
index 36525c59..c5dd2044 100644
--- a/tunnel/src/main/java/com/wireguard/android/backend/GoBackend.java
+++ b/tunnel/src/main/java/com/wireguard/android/backend/GoBackend.java
@@ -15,6 +15,7 @@ import android.net.Network;
 import android.net.NetworkCapabilities;
 import android.net.NetworkRequest;
 import android.net.ProxyInfo;
+import android.net.TrafficStats;
 import android.net.Uri;
 import android.os.Build;
 import android.os.ParcelFileDescriptor;
@@ -103,6 +104,7 @@ import androidx.collection.ArraySet;
 public final class GoBackend implements Backend {
     private static final int DNS_RESOLUTION_RETRIES = 10;
     private static final String TAG = "WireGuard/GoBackend";
+    private static final int STATS_TAG = 2;
     @Nullable private static AlwaysOnCallback alwaysOnCallback;
     private static GhettoCompletableFuture<VpnService> vpnService = new GhettoCompletableFuture<>();
     private final Context context;
@@ -317,6 +319,7 @@ public final class GoBackend implements Backend {
         StringBuffer buf = new StringBuffer();
         try {
             URL url = new URL(pacFileUrl.toString());
+            TrafficStats.setThreadStatsTag(STATS_TAG);
             urlConnection = (HttpURLConnection) network.openConnection(url);
 
             InputStream in = urlConnection.getInputStream();
diff --git a/tunnel/src/main/java/com/wireguard/util/Resolver.java b/tunnel/src/main/java/com/wireguard/util/Resolver.java
index 301e03e0..654e01f5 100644
--- a/tunnel/src/main/java/com/wireguard/util/Resolver.java
+++ b/tunnel/src/main/java/com/wireguard/util/Resolver.java
@@ -17,6 +17,7 @@ import java.net.UnknownHostException;
 import android.net.IpPrefix;
 import android.net.LinkProperties;
 import android.net.Network;
+import android.net.TrafficStats;
 import android.util.Log;
 
 import androidx.annotation.Nullable;
@@ -24,6 +25,7 @@ import androidx.annotation.Nullable;
 @NonNullForAll
 public class Resolver {
     private static final String TAG = "WireGuard/Resolver";
+    private static final int STATS_TAG = 3; // FIXME
     @Nullable private final Network network;
     @Nullable private final LinkProperties linkProps;
     @Nullable private IpPrefix nat64Prefix;
@@ -82,6 +84,7 @@ public class Resolver {
     }
 
     public InetAddress resolve(String host) throws UnknownHostException {
+        TrafficStats.setThreadStatsTag(STATS_TAG);
         final InetAddress[] candidates = network != null ? network.getAllByName(host) : InetAddress.getAllByName(host);
         InetAddress address = candidates[0];
         for (final InetAddress candidate : candidates) {
@@ -89,6 +92,7 @@ public class Resolver {
 
             try {
                 sock = new DatagramSocket();
+                TrafficStats.tagDatagramSocket(sock);
                 if (network != null) {
                     network.bindSocket(sock);
                 }