From 38431a290ed55174599be0927be32daaa1c49d56 Mon Sep 17 00:00:00 2001
From: Riya Mehta <55350838+rmehta19@users.noreply.github.com>
Date: Wed, 22 Jan 2025 13:00:26 -0800
Subject: [PATCH] feat: add support for new setAllowHardBoundTokens field.
 (#3467)

Introduce new `setAllowHardBoundTokens` field.
---
 .../InstantiatingGrpcChannelProvider.java     | 33 +++++++++++++++++++
 .../InstantiatingGrpcChannelProviderTest.java |  5 +++
 2 files changed, 38 insertions(+)

diff --git a/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProvider.java b/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProvider.java
index ae4d7f9e51..4cdeb8d74f 100644
--- a/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProvider.java
+++ b/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProvider.java
@@ -126,16 +126,35 @@ public final class InstantiatingGrpcChannelProvider implements TransportChannelP
   @Nullable private final Boolean allowNonDefaultServiceAccount;
   @VisibleForTesting final ImmutableMap<String, ?> directPathServiceConfig;
   @Nullable private final MtlsProvider mtlsProvider;
+  @Nullable private final List<HardBoundTokenTypes> allowedHardBoundTokenTypes;
   @VisibleForTesting final Map<String, String> headersWithDuplicatesRemoved = new HashMap<>();
 
   @Nullable
   private final ApiFunction<ManagedChannelBuilder, ManagedChannelBuilder> channelConfigurator;
 
+  /*
+   * Experimental feature
+   *
+   * <p>{@link HardBoundTokenTypes} specifies if hard bound tokens should be used if DirectPath
+   * or S2A is used to estabilsh a connection to Google APIs.
+   *
+   */
+  @InternalApi
+  public enum HardBoundTokenTypes {
+    // If DirectPath is used to create the channel, use hard ALTS-bound tokens for requests sent on
+    // that channel.
+    ALTS,
+    // If MTLS via S2A is used to create the channel, use hard MTLS-bound tokens for requests sent
+    // on that channel.
+    MTLS_S2A
+  }
+
   private InstantiatingGrpcChannelProvider(Builder builder) {
     this.processorCount = builder.processorCount;
     this.executor = builder.executor;
     this.headerProvider = builder.headerProvider;
     this.endpoint = builder.endpoint;
+    this.allowedHardBoundTokenTypes = builder.allowedHardBoundTokenTypes;
     this.mtlsProvider = builder.mtlsProvider;
     this.envProvider = builder.envProvider;
     this.interceptorProvider = builder.interceptorProvider;
@@ -620,6 +639,7 @@ public static final class Builder {
     @Nullable private Boolean attemptDirectPathXds;
     @Nullable private Boolean allowNonDefaultServiceAccount;
     @Nullable private ImmutableMap<String, ?> directPathServiceConfig;
+    @Nullable private List<HardBoundTokenTypes> allowedHardBoundTokenTypes;
 
     private Builder() {
       processorCount = Runtime.getRuntime().availableProcessors();
@@ -700,6 +720,19 @@ public Builder setEndpoint(String endpoint) {
       return this;
     }
 
+    /*
+     * Sets the allowed hard bound token types for this TransportChannelProvider.
+     *
+     * <p>The list of
+     * {@link HardBoundTokenTypes} indicates for which methods of connecting to Google APIs hard bound tokens should
+     * be used. This is optional; if it is not provided, bearer tokens will be used.
+     */
+    @InternalApi
+    public Builder setAllowHardBoundTokenTypes(List<HardBoundTokenTypes> allowedValues) {
+      this.allowedHardBoundTokenTypes = allowedValues;
+      return this;
+    }
+
     @VisibleForTesting
     Builder setMtlsProvider(MtlsProvider mtlsProvider) {
       this.mtlsProvider = mtlsProvider;
diff --git a/gax-java/gax-grpc/src/test/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProviderTest.java b/gax-java/gax-grpc/src/test/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProviderTest.java
index a58f9b8173..6ec9317ab9 100644
--- a/gax-java/gax-grpc/src/test/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProviderTest.java
+++ b/gax-java/gax-grpc/src/test/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProviderTest.java
@@ -225,6 +225,10 @@ void testToBuilder() {
           throw new UnsupportedOperationException();
         };
     Map<String, ?> directPathServiceConfig = ImmutableMap.of("loadbalancingConfig", "grpclb");
+    List<InstantiatingGrpcChannelProvider.HardBoundTokenTypes> hardBoundTokenTypes =
+        new ArrayList<>();
+    hardBoundTokenTypes.add(InstantiatingGrpcChannelProvider.HardBoundTokenTypes.ALTS);
+    hardBoundTokenTypes.add(InstantiatingGrpcChannelProvider.HardBoundTokenTypes.MTLS_S2A);
 
     InstantiatingGrpcChannelProvider provider =
         InstantiatingGrpcChannelProvider.newBuilder()
@@ -238,6 +242,7 @@ void testToBuilder() {
             .setChannelConfigurator(channelConfigurator)
             .setChannelsPerCpu(2.5)
             .setDirectPathServiceConfig(directPathServiceConfig)
+            .setAllowHardBoundTokenTypes(hardBoundTokenTypes)
             .build();
 
     InstantiatingGrpcChannelProvider.Builder builder = provider.toBuilder();