use endpoint override.

Author: rmehta19

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

@@ -122,6 +122,7 @@ public final class InstantiatingGrpcChannelProvider implements TransportChannelP
   private final HeaderProvider headerProvider;
   private final String endpoint;
   private final String mtlsEndpoint;
+  private final String endpointOverride;
   // TODO: remove.
   // envProvider currently provides DirectPath and S2A environment variables, and is only used
   // during initial rollout for DirectPath and S2A. This provider will be removed once the
@@ -153,6 +154,7 @@ private InstantiatingGrpcChannelProvider(Builder builder) {
     this.headerProvider = builder.headerProvider;
     this.endpoint = builder.endpoint;
     this.mtlsEndpoint = builder.mtlsEndpoint;
+    this.endpointOverride = builder.endpointOverride;
     this.mtlsProvider = builder.mtlsProvider;
     this.envProvider = builder.envProvider;
     this.interceptorProvider = builder.interceptorProvider;
@@ -228,10 +230,16 @@ public boolean needsEndpoint() {
     return endpoint == null;
   }
 
+  @Override
   public boolean needsMtlsEndpoint() {
     return mtlsEndpoint == null;
   }
 
+  @Override
+  public boolean needsEndpointOverride() {
+    return endpointOverride == null;
+  }
+
   /**
    * Specify the endpoint the channel should connect to.
    *
@@ -255,11 +263,31 @@ public TransportChannelProvider withEndpoint(String endpoint) {
    * @return A new {@link InstantiatingGrpcChannelProvider} with the specified MTLS endpoint
    *     configured
    */
+  @Override
   public TransportChannelProvider withMtlsEndpoint(String mtlsEndpoint) {
-    validateEndpoint(mtlsEndpoint);
+    if (!mtlsEndpoint.isEmpty()) {
+      validateEndpoint(mtlsEndpoint);
+    }
     return toBuilder().setMtlsEndpoint(mtlsEndpoint).build();
   }
 
+  /**
+   * Specify the endpoint override.
+   *
+   * <p>The value of {@code endpointOverride} must be of the form {@code host:port}.
+   *
+   * @param endpointOverride
+   * @return A new {@link InstantiatingGrpcChannelProvider} with the specified endpoint Override
+   *     configured
+   */
+  @Override
+  public TransportChannelProvider withEndpointOverride(String endpointOverride) {
+    if (!endpointOverride.isEmpty()) {
+      validateEndpoint(endpointOverride);
+    }
+    return toBuilder().setEndpointOverride(endpointOverride).build();
+  }
+
   /** @deprecated Please modify pool settings via {@link #toBuilder()} */
   @Deprecated
   @Override
@@ -462,15 +490,13 @@ boolean shouldUseS2A() {
       return false;
     }
 
-    // If {@link mtlsEndpoint} is not set, skip S2A. S2A is also skipped when there is endpoint
-    // override. Endpoint override is respected when the {@link endpoint} is resolved via AIP#4114,
-    // see EndpointContext.java
-    if (endpoint != mtlsEndpoint) {
+    // If {@code mtlsEndpoint} is not set, or {@code endpointOverride} is specified, skip S2A.
+    if (mtlsEndpoint.isEmpty() || !endpointOverride.isEmpty()) {
       return false;
     }
 
     // mTLS via S2A is not supported in any universe other than googleapis.com.
-    if (!endpoint.contains(Credentials.GOOGLE_DEFAULT_UNIVERSE)) {
+    if (!mtlsEndpoint.contains(Credentials.GOOGLE_DEFAULT_UNIVERSE)) {
       return false;
     }
 
@@ -686,6 +712,11 @@ public String getMtlsEndpoint() {
     return mtlsEndpoint;
   }
 
+  /** The endpoint override */
+  public String getEndpointOverride() {
+    return endpointOverride;
+  }
+
   /** This method is obsolete. Use {@link #getKeepAliveTimeDuration()} instead. */
   @ObsoleteApi("Use getKeepAliveTimeDuration() instead")
   public org.threeten.bp.Duration getKeepAliveTime() {
@@ -744,6 +775,7 @@ public static final class Builder {
     private HeaderProvider headerProvider;
     private String endpoint;
     private String mtlsEndpoint;
+    private String endpointOverride;
     private EnvironmentProvider envProvider;
     private MtlsProvider mtlsProvider = new MtlsProvider();
     @Nullable private GrpcInterceptorProvider interceptorProvider;
@@ -773,6 +805,7 @@ private Builder(InstantiatingGrpcChannelProvider provider) {
       this.headerProvider = provider.headerProvider;
       this.endpoint = provider.endpoint;
       this.mtlsEndpoint = provider.mtlsEndpoint;
+      this.endpointOverride = provider.endpointOverride;
       this.envProvider = provider.envProvider;
       this.interceptorProvider = provider.interceptorProvider;
       this.maxInboundMessageSize = provider.maxInboundMessageSize;
@@ -842,11 +875,21 @@ public Builder setEndpoint(String endpoint) {
     }
 
     public Builder setMtlsEndpoint(String mtlsEndpoint) {
-      validateEndpoint(mtlsEndpoint);
+      if (!mtlsEndpoint.isEmpty()) {
+        validateEndpoint(mtlsEndpoint);
+      }
       this.mtlsEndpoint = mtlsEndpoint;
       return this;
     }
 
+    public Builder setEndpointOverride(String endpointOverride) {
+      if (!endpointOverride.isEmpty()) {
+        validateEndpoint(endpointOverride);
+      }
+      this.endpointOverride = endpointOverride;
+      return this;
+    }
+
     @VisibleForTesting
     Builder setMtlsProvider(MtlsProvider mtlsProvider) {
       this.mtlsProvider = mtlsProvider;
@@ -873,6 +916,10 @@ public String getMtlsEndpoint() {
       return mtlsEndpoint;
     }
 
+    public String getEndpointOverride() {
+      return endpointOverride;
+    }
+
     /** The maximum message size allowed to be received on the channel. */
     public Builder setMaxInboundMessageSize(Integer max) {
       this.maxInboundMessageSize = max;

gax-java/gax-grpc/src/test/java/com/google/api/gax/grpc/InstantiatingGrpcChannelProviderTest.java

@@ -86,6 +86,7 @@
 class InstantiatingGrpcChannelProviderTest extends AbstractMtlsTransportChannelTest {
   private static final String DEFAULT_ENDPOINT = "test.googleapis.com:443";
   private static final String DEFAULT_MTLS_ENDPOINT = "test.mtls.googleapis.com:443";
+  private static final String DEFAULT_ENDPOINT_OVERRIDE = "test.endpoint.com:443";
   private static final String API_KEY_HEADER_VALUE = "fake_api_key_2";
   private static final String API_KEY_AUTH_HEADER_KEY = "x-goog-api-key";
   private static String originalOSName;
@@ -161,6 +162,32 @@ void testMtlsEndpointBadPort() {
                 .setEndpoint("test.mtls.googleapis.com:abcd"));
   }
 
+  @Test
+  void testEndpointOverrideEndpoint() {
+    InstantiatingGrpcChannelProvider.Builder builder =
+        InstantiatingGrpcChannelProvider.newBuilder();
+    builder.setEndpointOverride(DEFAULT_ENDPOINT_OVERRIDE);
+    assertEquals(builder.getEndpointOverride(), DEFAULT_ENDPOINT_OVERRIDE);
+
+    InstantiatingGrpcChannelProvider provider = builder.build();
+    assertEquals(provider.getEndpointOverride(), DEFAULT_ENDPOINT_OVERRIDE);
+  }
+
+  @Test
+  void testEndpointOverrideNoPort() {
+    assertThrows(
+        IllegalArgumentException.class,
+        () ->
+            InstantiatingGrpcChannelProvider.newBuilder().setEndpointOverride("test.endpoint.com"));
+  }
+
+  @Test
+  void testEndpointOverrideBadPort() {
+    assertThrows(
+        IllegalArgumentException.class,
+        () -> InstantiatingGrpcChannelProvider.newBuilder().setEndpoint("test.endpoint.com:abcd"));
+  }
+
   @Test
   void testKeepAlive() {
     final long millis = 15;
@@ -1041,8 +1068,9 @@ void shouldUseS2A_envVarNotSet_returnsFalse() {
         .thenReturn("false");
     InstantiatingGrpcChannelProvider provider =
         InstantiatingGrpcChannelProvider.newBuilder()
-            .setEndpoint(DEFAULT_MTLS_ENDPOINT)
+            .setEndpoint(DEFAULT_ENDPOINT)
             .setMtlsEndpoint(DEFAULT_MTLS_ENDPOINT)
+            .setEndpointOverride("")
             .setEnvProvider(envProvider)
             .build();
     Truth.assertThat(provider.shouldUseS2A()).isFalse();
@@ -1056,6 +1084,8 @@ void shouldUseS2A_mtlsEndpointNotSet_returnsFalse() {
     InstantiatingGrpcChannelProvider provider =
         InstantiatingGrpcChannelProvider.newBuilder()
             .setEndpoint(DEFAULT_ENDPOINT)
+            .setMtlsEndpoint("")
+            .setEndpointOverride("")
             .setEnvProvider(envProvider)
             .build();
     Truth.assertThat(provider.shouldUseS2A()).isFalse();
@@ -1070,6 +1100,7 @@ void shouldUseS2A_endpointOverrideIsSet_returnsFalse() {
         InstantiatingGrpcChannelProvider.newBuilder()
             .setEndpoint(DEFAULT_ENDPOINT)
             .setMtlsEndpoint(DEFAULT_MTLS_ENDPOINT)
+            .setEndpointOverride(DEFAULT_ENDPOINT_OVERRIDE)
             .setEnvProvider(envProvider)
             .build();
     Truth.assertThat(provider.shouldUseS2A()).isFalse();
@@ -1082,8 +1113,9 @@ void shouldUseS2A_nonGDUUniverse_returnsFalse() {
         .thenReturn("true");
     InstantiatingGrpcChannelProvider provider =
         InstantiatingGrpcChannelProvider.newBuilder()
-            .setEndpoint("test.mtls.abcd.com:443")
+            .setEndpoint("test.abcd.com:443")
             .setMtlsEndpoint("test.mtls.abcd.com:443")
+            .setEndpointOverride("")
             .setEnvProvider(envProvider)
             .build();
     Truth.assertThat(provider.shouldUseS2A()).isFalse();
@@ -1096,8 +1128,9 @@ void shouldUseS2A_returnsTrue() {
         .thenReturn("true");
     InstantiatingGrpcChannelProvider provider =
         InstantiatingGrpcChannelProvider.newBuilder()
-            .setEndpoint(DEFAULT_MTLS_ENDPOINT)
+            .setEndpoint(DEFAULT_ENDPOINT)
             .setMtlsEndpoint(DEFAULT_MTLS_ENDPOINT)
+            .setEndpointOverride("")
             .setEnvProvider(envProvider)
             .build();
     Truth.assertThat(provider.shouldUseS2A()).isTrue();

gax-java/gax-grpc/src/test/java/com/google/api/gax/grpc/testing/LocalChannelProvider.java

@@ -106,6 +106,11 @@ public boolean needsMtlsEndpoint() {
     return false;
   }
 
+  @Override
+  public boolean needsEndpointOverride() {
+    return false;
+  }
+
   @Override
   public TransportChannelProvider withEndpoint(String endpoint) {
     throw new UnsupportedOperationException("LocalChannelProvider doesn't need an endpoint");
@@ -116,6 +121,12 @@ public TransportChannelProvider withMtlsEndpoint(String mtlsEndpoint) {
     throw new UnsupportedOperationException("LocalChannelProvider doesn't need an mtlsEndpoint");
   }
 
+  @Override
+  public TransportChannelProvider withEndpointOverride(String endpointOverride) {
+    throw new UnsupportedOperationException(
+        "LocalChannelProvider doesn't need an endpointOverride");
+  }
+
   @Override
   @BetaApi("The surface for customizing pool size is not stable yet and may change in the future.")
   public boolean acceptsPoolSize() {

gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/InstantiatingHttpJsonChannelProvider.java

@@ -124,6 +124,11 @@ public boolean needsMtlsEndpoint() {
     return false;
   }
 
+  @Override
+  public boolean needsEndpointOverride() {
+    return false;
+  }
+
   @Override
   public TransportChannelProvider withEndpoint(String endpoint) {
     return toBuilder().setEndpoint(endpoint).build();
@@ -135,6 +140,12 @@ public TransportChannelProvider withMtlsEndpoint(String mtlsEndpoint) {
         "InstantiatingHttpJsonChannelProvider doesn't need an mtlsEndpoint");
   }
 
+  @Override
+  public TransportChannelProvider withEndpointOverride(String endpointOverride) {
+    throw new UnsupportedOperationException(
+        "InstantiatingHttpJsonChannelProvider doesn't need an endpointOverride");
+  }
+
   /** @deprecated REST transport channel doesn't support channel pooling */
   @Deprecated
   @Override

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

@@ -178,7 +178,10 @@ public static ClientContext create(StubSettings settings) throws IOException {
     // A valid EndpointContext should have been created in the StubSettings
     EndpointContext endpointContext = settings.getEndpointContext();
     String endpoint = endpointContext.resolvedEndpoint();
-    String mtlsEndpoint = settings.getMtlsEndpoint();
+    String mtlsEndpoint = "";
+    mtlsEndpoint = settings.getMtlsEndpoint();
+    String endpointOverride = "";
+    endpointOverride = settings.getEndpointContext().transportChannelProviderEndpoint();
     Credentials credentials = getCredentials(settings);
     // check if need to adjust credentials/endpoint/endpointContext for GDC-H
     String settingsGdchApiAudience = settings.getGdchApiAudience();
@@ -190,6 +193,9 @@ public static ClientContext create(StubSettings settings) throws IOException {
       endpointContext = endpointContext.withGDCH();
       // Resolve the new endpoint with the GDC-H flow
       endpoint = endpointContext.resolvedEndpoint();
+      // Skip S2A in GDC-H.
+      mtlsEndpoint = "";
+      endpointOverride = "";
       // We recompute the GdchCredentials with the audience
       credentials = getGdchCredentials(settingsGdchApiAudience, endpoint, credentials);
     } else if (!Strings.isNullOrEmpty(settingsGdchApiAudience)) {
@@ -226,6 +232,9 @@ public static ClientContext create(StubSettings settings) throws IOException {
     if (transportChannelProvider.needsMtlsEndpoint()) {
       transportChannelProvider = transportChannelProvider.withMtlsEndpoint(mtlsEndpoint);
     }
+    if (transportChannelProvider.needsMtlsEndpoint()) {
+      transportChannelProvider = transportChannelProvider.withEndpointOverride(endpointOverride);
+    }
     TransportChannel transportChannel = transportChannelProvider.getTransportChannel();
 
     ApiCallContext defaultCallContext =

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

@@ -88,6 +88,11 @@ public boolean needsMtlsEndpoint() {
     return false;
   }
 
+  @Override
+  public boolean needsEndpointOverride() {
+    return false;
+  }
+
   @Override
   public TransportChannelProvider withEndpoint(String endpoint) {
     throw new UnsupportedOperationException(
@@ -100,6 +105,12 @@ public TransportChannelProvider withMtlsEndpoint(String mtlsEndpoint) {
         "FixedTransportChannelProvider doesn't need an mtlsEndpoint");
   }
 
+  @Override
+  public TransportChannelProvider withEndpointOverride(String endpointOverride) {
+    throw new UnsupportedOperationException(
+        "FixedTransportChannelProvider doesn't need an endpointOverride");
+  }
+
   /** @deprecated FixedTransportChannelProvider doesn't support ChannelPool configuration */
   @Deprecated
   @Override

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

@@ -93,6 +93,9 @@ public interface TransportChannelProvider {
   /** True if the TransportProvider has no mtlsEndpoint set. */
   boolean needsMtlsEndpoint();
 
+  /** True if the TransportProvider has no endpointOverride set */
+  boolean needsEndpointOverride();
+
   /**
    * Sets the endpoint to use when constructing a new {@link TransportChannel}.
    *
@@ -107,6 +110,13 @@ public interface TransportChannelProvider {
    */
   TransportChannelProvider withMtlsEndpoint(String mtlsEndpoint);
 
+  /**
+   * Sets the endpoint override when constructing a new {@link TransportChannel}.
+   *
+   * <p>This method should only be called if {@link #needsEndpointOverride()} returns true.
+   */
+  TransportChannelProvider withEndpointOverride(String endpointOverride);
+
   /**
    * Reports whether this provider allows pool size customization.
    *