From c2eccca3bc1037be121b299ec3a3a4c47f6ee506 Mon Sep 17 00:00:00 2001
From: Eric Anderson <ejona@google.com>
Date: Tue, 13 Aug 2024 10:27:48 -0700
Subject: [PATCH] cronet: Add internal API to specify Network

cl/661194496
---
 .../io/grpc/cronet/CronetChannelBuilder.java  | 28 +++++++++++++++++--
 .../cronet/InternalCronetChannelBuilder.java  |  7 +++++
 2 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/cronet/src/main/java/io/grpc/cronet/CronetChannelBuilder.java b/cronet/src/main/java/io/grpc/cronet/CronetChannelBuilder.java
index 1c60f82846..f42dabdd55 100644
--- a/cronet/src/main/java/io/grpc/cronet/CronetChannelBuilder.java
+++ b/cronet/src/main/java/io/grpc/cronet/CronetChannelBuilder.java
@@ -20,9 +20,12 @@ import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkNotNull;
 import static io.grpc.internal.GrpcUtil.DEFAULT_MAX_MESSAGE_SIZE;
 
+import android.net.Network;
+import android.os.Build;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.util.concurrent.MoreExecutors;
+import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import com.google.errorprone.annotations.DoNotCall;
 import io.grpc.ChannelCredentials;
 import io.grpc.ChannelLogger;
@@ -105,6 +108,7 @@ public final class CronetChannelBuilder extends ForwardingChannelBuilder2<Cronet
   private int trafficStatsTag;
   private boolean trafficStatsUidSet;
   private int trafficStatsUid;
+  private Network network;
 
   private CronetChannelBuilder(String host, int port, CronetEngine cronetEngine) {
     final class CronetChannelTransportFactoryBuilder implements ClientTransportFactoryBuilder {
@@ -190,6 +194,13 @@ public final class CronetChannelBuilder extends ForwardingChannelBuilder2<Cronet
     return this;
   }
 
+  /** Sets the network ID to use for this channel traffic. */
+  @CanIgnoreReturnValue
+  CronetChannelBuilder bindToNetwork(@Nullable Network network) {
+    this.network = network;
+    return this;
+  }
+
   /**
    * Provides a custom scheduled executor service.
    *
@@ -210,7 +221,12 @@ public final class CronetChannelBuilder extends ForwardingChannelBuilder2<Cronet
   ClientTransportFactory buildTransportFactory() {
     return new CronetTransportFactory(
         new TaggingStreamFactory(
-            cronetEngine, trafficStatsTagSet, trafficStatsTag, trafficStatsUidSet, trafficStatsUid),
+            cronetEngine,
+            trafficStatsTagSet,
+            trafficStatsTag,
+            trafficStatsUidSet,
+            trafficStatsUid,
+            network),
         MoreExecutors.directExecutor(),
         scheduledExecutorService,
         maxMessageSize,
@@ -294,18 +310,21 @@ public final class CronetChannelBuilder extends ForwardingChannelBuilder2<Cronet
     private final int trafficStatsTag;
     private final boolean trafficStatsUidSet;
     private final int trafficStatsUid;
+    private final Network network;
 
     TaggingStreamFactory(
         CronetEngine cronetEngine,
         boolean trafficStatsTagSet,
         int trafficStatsTag,
         boolean trafficStatsUidSet,
-        int trafficStatsUid) {
+        int trafficStatsUid,
+        Network network) {
       this.cronetEngine = cronetEngine;
       this.trafficStatsTagSet = trafficStatsTagSet;
       this.trafficStatsTag = trafficStatsTag;
       this.trafficStatsUidSet = trafficStatsUidSet;
       this.trafficStatsUid = trafficStatsUid;
+      this.network = network;
     }
 
     @Override
@@ -320,6 +339,11 @@ public final class CronetChannelBuilder extends ForwardingChannelBuilder2<Cronet
       if (trafficStatsUidSet) {
         builder.setTrafficStatsUid(trafficStatsUid);
       }
+      if (network != null) {
+        if (Build.VERSION.SDK_INT >= 23) {
+          builder.bindToNetwork(network.getNetworkHandle());
+        }
+      }
       return builder;
     }
   }
diff --git a/cronet/src/main/java/io/grpc/cronet/InternalCronetChannelBuilder.java b/cronet/src/main/java/io/grpc/cronet/InternalCronetChannelBuilder.java
index 2954f1eee8..7e5e610ca6 100644
--- a/cronet/src/main/java/io/grpc/cronet/InternalCronetChannelBuilder.java
+++ b/cronet/src/main/java/io/grpc/cronet/InternalCronetChannelBuilder.java
@@ -16,7 +16,9 @@
 
 package io.grpc.cronet;
 
+import android.net.Network;
 import io.grpc.Internal;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 /**
  * Internal {@link CronetChannelBuilder} accessor. This is intended for usage internal to the gRPC
@@ -58,4 +60,9 @@ public final class InternalCronetChannelBuilder {
   public static void setTrafficStatsUid(CronetChannelBuilder builder, int uid) {
     builder.setTrafficStatsUid(uid);
   }
+
+  /** Sets the network {@link android.net.Network} to use when relying traffic by this channel. */
+  public static void bindToNetwork(CronetChannelBuilder builder, @Nullable Network network) {
+    builder.bindToNetwork(network);
+  }
 }