fix: plumb mtls endpoint to TransportChannelProvider (#3673)

This PR plumbs the MTLS endpoint (separately from the resolved endpoint)
from the `EndpointContext` to the `InstantiatingGrpcChannelProvider`.

Why not just set the MTLS endpoint in `EndpointContext` if S2A can be
used?
- Although we can decide whether or not to try to use S2A in
`EndpointContext`, we could fail to use S2A in
`InstantiatingGrpcChannelProvider` (if autoconfig doesn't return an
address), in which case we fall back to using a TLS connection
- DirectPath supersedes S2A in `InstantiatingGrpcChannelProvider`, and
the decision to use DirectPath is made in
`InstantiatingGrpcChannelProvider`, not `EndpointContext`. We may decide
to use S2A in `EndpointContext`, but when we go to create the channel in
`InstantiatingGrpcChannelProvider`, we may find that we should be using
DirectPath, in which case we need to use the non-MTLS endpoint

Author: rmehta19

gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProvider.java

@@ -129,6 +129,7 @@ public final class InstantiatingGrpcChannelProvider implements TransportChannelP
   private final HeaderProvider headerProvider;
   private final boolean useS2A;
   private final String endpoint;
+  private final String mtlsEndpoint;
   // TODO: remove. envProvider currently provides DirectPath environment variable, and is only used
   // during initial rollout for DirectPath. This provider will be removed once the DirectPath
   // environment is not used.
@@ -179,6 +180,7 @@ private InstantiatingGrpcChannelProvider(Builder builder) {
     this.headerProvider = builder.headerProvider;
     this.useS2A = builder.useS2A;
     this.endpoint = builder.endpoint;
+    this.mtlsEndpoint = builder.mtlsEndpoint;
     this.allowedHardBoundTokenTypes = builder.allowedHardBoundTokenTypes;
     this.mtlsProvider = builder.mtlsProvider;
     this.s2aConfigProvider = builder.s2aConfigProvider;
@@ -258,6 +260,12 @@ public boolean needsEndpoint() {
     return endpoint == null;
   }
 
+  @InternalApi("This public method is used by Gax to help configure the MTLS endpoint for S2A")
+  @Override
+  public boolean needsMtlsEndpoint() {
+    return mtlsEndpoint == null;
+  }
+
   /**
    * Specify the endpoint the channel should connect to.
    *
@@ -272,6 +280,22 @@ public TransportChannelProvider withEndpoint(String endpoint) {
     return toBuilder().setEndpoint(endpoint).build();
   }
 
+  /**
+   * Specify the mTLS endpoint the channel should connect to when using S2A.
+   *
+   * <p>The value of {@code mtlsEndpoint} must be of the form {@code host:port}.
+   *
+   * @param mtlsEndpoint The mtTLS endpoint to connect to
+   * @return A new {@link InstantiatingGrpcChannelProvider} with the specified mTLS endpoint
+   *     configured
+   */
+  @InternalApi("This public method is used by Gax to help configure the MTLS endpoint for S2A")
+  @Override
+  public TransportChannelProvider withMtlsEndpoint(String mtlsEndpoint) {
+    validateEndpoint(mtlsEndpoint);
+    return toBuilder().setMtlsEndpoint(mtlsEndpoint).build();
+  }
+
   /**
    * Specify whether or not to use S2A.
    *
@@ -666,7 +690,9 @@ private ManagedChannel createSingleChannel() throws IOException {
             channelCredentials =
                 CompositeChannelCredentials.create(channelCredentials, mtlsS2ACallCredentials);
           }
-          builder = Grpc.newChannelBuilder(endpoint, channelCredentials);
+          // Connect to the MTLS endpoint when using S2A because S2A is used to perform an MTLS
+          // handshake.
+          builder = Grpc.newChannelBuilder(mtlsEndpoint, channelCredentials);
         } else {
           // Use default if we cannot initialize channel credentials via DCA or S2A.
           builder = ManagedChannelBuilder.forAddress(serviceAddress, port);
@@ -819,6 +845,7 @@ public static final class Builder {
     private Executor executor;
     private HeaderProvider headerProvider;
     private String endpoint;
+    private String mtlsEndpoint;
     private boolean useS2A;
     private EnvironmentProvider envProvider;
     private SecureSessionAgent s2aConfigProvider = SecureSessionAgent.create();
@@ -926,6 +953,14 @@ public Builder setEndpoint(String endpoint) {
       return this;
     }
 
+    /** Sets the mTLS Endpoint used to reach the service, eg "localhost:8080". */
+    @InternalApi("This public method is used by Gax to help configure the MTLS endpoint for S2A")
+    public Builder setMtlsEndpoint(String mtlsEndpoint) {
+      validateEndpoint(mtlsEndpoint);
+      this.mtlsEndpoint = mtlsEndpoint;
+      return this;
+    }
+
     Builder setUseS2A(boolean useS2A) {
       this.useS2A = useS2A;
       return this;

gax-java/gax/clirr-ignored-differences.xml

@@ -118,4 +118,14 @@
     <className>com/google/api/gax/rpc/TransportChannelProvider</className>
     <method>* withUseS2A(*)</method>
   </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/api/gax/rpc/TransportChannelProvider</className>
+    <method>* withMtlsEndpoint(*)</method>
+  </difference>
+  <difference>
+    <differenceType>7012</differenceType>
+    <className>com/google/api/gax/rpc/TransportChannelProvider</className>
+    <method>* needsMtlsEndpoint()</method>
+  </difference>
 </differences>

gax-java/gax/src/main/java/com/google/api/gax/rpc/ClientContext.java

@@ -223,6 +223,10 @@ public static ClientContext create(StubSettings settings) throws IOException {
       transportChannelProvider = transportChannelProvider.withEndpoint(endpoint);
     }
     transportChannelProvider = transportChannelProvider.withUseS2A(endpointContext.useS2A());
+    if (transportChannelProvider.needsMtlsEndpoint()) {
+      transportChannelProvider =
+          transportChannelProvider.withMtlsEndpoint(endpointContext.mtlsEndpoint());
+    }
     TransportChannel transportChannel = transportChannelProvider.getTransportChannel();
 
     ApiCallContext defaultCallContext =

gax-java/gax/src/main/java/com/google/api/gax/rpc/EndpointContext.java

@@ -303,13 +303,6 @@ private String determineEndpoint() throws IOException {
             "mTLS is not supported in any universe other than googleapis.com");
       }
 
-      // Check if Experimental S2A feature enabled. When feature is non-experimental, remove this
-      // check from this function, and plumb MTLS endpoint to channel creation logic separately.
-      // Note that mTLS via S2A is an independent feature from mTLS via DCA (for which endpoint
-      // determined by {@code mtlsEndpointResolver} above).
-      if (shouldUseS2A()) {
-        return mtlsEndpoint();
-      }
       return endpoint;
     }
 

gax-java/gax/src/main/java/com/google/api/gax/rpc/TransportChannelProvider.java

@@ -97,6 +97,20 @@ public interface TransportChannelProvider {
    */
   TransportChannelProvider withEndpoint(String endpoint);
 
+  /** True for gRPC transport provider which has no mtlsEndpoint set. */
+  default boolean needsMtlsEndpoint() {
+    return false;
+  }
+
+  /**
+   * Sets the mTLS endpoint to use when constructing a new {@link TransportChannel} using S2A.
+   *
+   * <p>This method should only be called if {@link #needsMtlsEndpoint()} returns true.
+   */
+  default TransportChannelProvider withMtlsEndpoint(String mtlsEndpoint) {
+    return this;
+  }
+
   /** Sets whether to use S2A when constructing a new {@link TransportChannel}. */
   @BetaApi(
       "The S2A feature is not stable yet and may change in the future. https://github.com/grpc/grpc-java/issues/11533.")

gax-java/gax/src/test/java/com/google/api/gax/rpc/EndpointContextTest.java

@@ -374,7 +374,7 @@ void endpointContextBuild_multipleUniverseDomainConfigurations_clientSettingsHas
   }
 
   @Test
-  void endpointContextBuild_shouldUseS2A_mtlsEndpoint() throws IOException {
+  void endpointContextBuild_shouldUseS2A_tlsEndpoint() throws IOException {
     EnvironmentProvider envProvider = Mockito.mock(EnvironmentProvider.class);
     Mockito.when(envProvider.getenv(EndpointContext.S2A_ENV_ENABLE_USE_S2A)).thenReturn("true");
     defaultEndpointContextBuilder =
@@ -385,7 +385,7 @@ void endpointContextBuild_shouldUseS2A_mtlsEndpoint() throws IOException {
             .setUsingGDCH(false);
     EndpointContext endpointContext = defaultEndpointContextBuilder.build();
     Truth.assertThat(defaultEndpointContextBuilder.shouldUseS2A()).isTrue();
-    Truth.assertThat(endpointContext.resolvedEndpoint()).isEqualTo(DEFAULT_MTLS_ENDPOINT);
+    Truth.assertThat(endpointContext.resolvedEndpoint()).isEqualTo(DEFAULT_ENDPOINT);
   }
 
   @Test