set mtls endpoint for S2A in EndpointContext.

Author: rmehta19

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

@@ -123,7 +123,6 @@ 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.
@@ -154,7 +153,6 @@ private InstantiatingGrpcChannelProvider(Builder builder) {
     this.headerProvider = builder.headerProvider;
     this.useS2A = builder.useS2A;
     this.endpoint = builder.endpoint;
-    this.mtlsEndpoint = builder.mtlsEndpoint;
     this.mtlsProvider = builder.mtlsProvider;
     this.s2aConfigProvider = builder.s2aConfigProvider;
     this.envProvider = builder.envProvider;
@@ -231,11 +229,6 @@ public boolean needsEndpoint() {
     return endpoint == null;
   }
 
-  @Override
-  public boolean needsMtlsEndpoint() {
-    return mtlsEndpoint == null;
-  }
-
   /**
    * Specify the endpoint the channel should connect to.
    *
@@ -250,21 +243,6 @@ public TransportChannelProvider withEndpoint(String endpoint) {
     return toBuilder().setEndpoint(endpoint).build();
   }
 
-  /**
-   * Specify the mtlsEndpoint the channel should connect to.
-   *
-   * <p>The value of {@code mtlsEndpoint} must be of the form {@code host:port}.
-   *
-   * @param mtlsEndpoint The mtlsEndpoint to connect to
-   * @return A new {@link InstantiatingGrpcChannelProvider} with the specified mtlsEndpoint
-   *     configured
-   */
-  @Override
-  public TransportChannelProvider withMtlsEndpoint(String mtlsEndpoint) {
-    validateEndpoint(mtlsEndpoint);
-    return toBuilder().setMtlsEndpoint(mtlsEndpoint).build();
-  }
-
   /**
    * Specify whether or not to use S2A.
    *
@@ -549,8 +527,7 @@ ChannelCredentials createPlaintextToS2AChannelCredentials(String plaintextAddres
    * returns null; in this case S2A will not be used, and a TLS connection to the service will be
    * established.
    *
-   * @return {@link ChannelCredentials} configured to use S2A to create mTLS connection to
-   *     mtlsEndpoint.
+   * @return {@link ChannelCredentials} configured to use S2A to create mTLS connection.
    */
   ChannelCredentials createS2ASecuredChannelCredentials() {
     SecureSessionAgentConfig config = s2aConfigProvider.getConfig();
@@ -648,7 +625,7 @@ private ManagedChannel createSingleChannel() throws IOException {
         }
         if (channelCredentials != null) {
           // Create the channel using S2A-secured channel credentials.
-          builder = Grpc.newChannelBuilder(mtlsEndpoint, channelCredentials);
+          builder = Grpc.newChannelBuilder(endpoint, channelCredentials);
         } else {
           // Use default if we cannot initialize channel credentials via DCA or S2A.
           builder = ManagedChannelBuilder.forAddress(serviceAddress, port);
@@ -800,7 +777,6 @@ 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();
@@ -831,7 +807,6 @@ private Builder(InstantiatingGrpcChannelProvider provider) {
       this.executor = provider.executor;
       this.headerProvider = provider.headerProvider;
       this.endpoint = provider.endpoint;
-      this.mtlsEndpoint = provider.mtlsEndpoint;
       this.useS2A = provider.useS2A;
       this.envProvider = provider.envProvider;
       this.interceptorProvider = provider.interceptorProvider;
@@ -902,22 +877,11 @@ public Builder setEndpoint(String endpoint) {
       return this;
     }
 
-    /** Sets the mtlsEndpoint used to reach the service, eg "localhost:8080". */
-    public Builder setMtlsEndpoint(String mtlsEndpoint) {
-      validateEndpoint(mtlsEndpoint);
-      this.mtlsEndpoint = mtlsEndpoint;
-      return this;
-    }
-
     Builder setUseS2A(boolean useS2A) {
       this.useS2A = useS2A;
       return this;
     }
 
-    public String getMtlsEndpoint() {
-      return mtlsEndpoint;
-    }
-
     @VisibleForTesting
     Builder setMtlsProvider(MtlsProvider mtlsProvider) {
       this.mtlsProvider = mtlsProvider;

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

@@ -103,8 +103,6 @@ void setUp() throws IOException {
     when(operationsChannelProvider.getTransportChannel()).thenReturn(transportChannel);
     when(operationsChannelProvider.withUseS2A(Mockito.any(boolean.class)))
         .thenReturn(operationsChannelProvider);
-    when(operationsChannelProvider.withMtlsEndpoint(Mockito.any(String.class)))
-        .thenReturn(operationsChannelProvider);
 
     clock = new FakeApiClock(0L);
     executor = RecordingScheduler.create(clock);

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

@@ -101,21 +101,11 @@ public boolean needsEndpoint() {
     return false;
   }
 
-  @Override
-  public boolean needsMtlsEndpoint() {
-    return false;
-  }
-
   @Override
   public TransportChannelProvider withEndpoint(String endpoint) {
     throw new UnsupportedOperationException("LocalChannelProvider doesn't need an endpoint");
   }
 
-  @Override
-  public TransportChannelProvider withMtlsEndpoint(String mtlsEndpoint) {
-    throw new UnsupportedOperationException("LocalChannelProvider doesn't need an mtlsEndpoint");
-  }
-
   @Override
   public TransportChannelProvider withUseS2A(boolean useS2A) {
     // Overriden for technical reasons. This method is a no-op for LocalChannelProvider.

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

@@ -119,21 +119,11 @@ public boolean needsEndpoint() {
     return endpoint == null;
   }
 
-  @Override
-  public boolean needsMtlsEndpoint() {
-    return false;
-  }
-
   @Override
   public TransportChannelProvider withEndpoint(String endpoint) {
     return toBuilder().setEndpoint(endpoint).build();
   }
 
-  @Override
-  public TransportChannelProvider withMtlsEndpoint(String mtlsEndpoint) {
-    return this;
-  }
-
   @Override
   public TransportChannelProvider withUseS2A(boolean useS2A) {
     return this;

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

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

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

@@ -303,6 +303,10 @@ private String determineEndpoint() throws IOException {
             "mTLS is not supported in any universe other than googleapis.com");
       }
 
+      if (shouldUseS2A()) {
+        return mtlsEndpoint();
+      }
+
       return endpoint;
     }
 

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

@@ -83,23 +83,12 @@ public boolean needsEndpoint() {
     return false;
   }
 
-  @Override
-  public boolean needsMtlsEndpoint() {
-    return false;
-  }
-
   @Override
   public TransportChannelProvider withEndpoint(String endpoint) {
     throw new UnsupportedOperationException(
         "FixedTransportChannelProvider doesn't need an endpoint");
   }
 
-  @Override
-  public TransportChannelProvider withMtlsEndpoint(String mtlsEndpoint) {
-    throw new UnsupportedOperationException(
-        "FixedTransportChannelProvider doesn't need an mtlsEndpoint");
-  }
-
   @Override
   public TransportChannelProvider withUseS2A(boolean useS2A) throws UnsupportedOperationException {
     // Overriden for technical reasons. This method is a no-op for FixedTransportChannelProvider.

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

@@ -90,23 +90,13 @@ public interface TransportChannelProvider {
   /** True if the TransportProvider has no endpoint set. */
   boolean needsEndpoint();
 
-  /** True if the TransportProvider has no mtlsEndpoint set. */
-  boolean needsMtlsEndpoint();
-
   /**
    * Sets the endpoint to use when constructing a new {@link TransportChannel}.
    *
    * <p>This method should only be called if {@link #needsEndpoint()} returns true.
    */
   TransportChannelProvider withEndpoint(String endpoint);
 
-  /**
-   * Sets the mtlsEndpoint to use when constructing a new {@link TransportChannel}.
-   *
-   * <p>This method should only be called if {@link #needsMtlsEndpoint()} returns true.
-   */
-  TransportChannelProvider withMtlsEndpoint(String mtlsEndpoint);
-
   /** Sets whether to use S2A when constructing a new {@link TransportChannel}. */
   default TransportChannelProvider withUseS2A(boolean useS2A) {
     throw new UnsupportedOperationException("S2A is not supported");

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

@@ -179,11 +179,6 @@ public boolean needsEndpoint() {
       return true;
     }
 
-    @Override
-    public boolean needsMtlsEndpoint() {
-      return false;
-    }
-
     @Override
     public String getEndpoint() {
       return endpoint;
@@ -200,17 +195,6 @@ public TransportChannelProvider withEndpoint(String endpoint) {
           endpoint);
     }
 
-    @Override
-    public TransportChannelProvider withMtlsEndpoint(String mtlsEndpoint) {
-      return new FakeTransportProvider(
-          this.transport,
-          this.executor,
-          this.shouldAutoClose,
-          this.headers,
-          this.credentials,
-          this.endpoint);
-    }
-
     @Override
     public TransportChannelProvider withUseS2A(boolean useS2A) {
       return new FakeTransportProvider(

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

@@ -373,6 +373,21 @@ void endpointContextBuild_multipleUniverseDomainConfigurations_clientSettingsHas
         .isEqualTo(clientSettingsUniverseDomain);
   }
 
+  @Test
+  void endpointContextBuild_shouldUseS2A_mtlsEndpoint() throws IOException {
+    EnvironmentProvider envProvider = Mockito.mock(EnvironmentProvider.class);
+    Mockito.when(envProvider.getenv(EndpointContext.S2A_ENV_ENABLE_USE_S2A)).thenReturn("true");
+    defaultEndpointContextBuilder =
+        defaultEndpointContextBuilder
+            .setEnvProvider(envProvider)
+            .setClientSettingsEndpoint("")
+            .setTransportChannelProviderEndpoint("")
+            .setUsingGDCH(false);
+    EndpointContext endpointContext = defaultEndpointContextBuilder.build();
+    Truth.assertThat(defaultEndpointContextBuilder.shouldUseS2A()).isTrue();
+    Truth.assertThat(endpointContext.resolvedEndpoint()).isEqualTo(DEFAULT_MTLS_ENDPOINT);
+  }
+
   @Test
   void hasValidUniverseDomain_gdchFlow_anyCredentials() throws IOException {
     Credentials noCredentials = NoCredentialsProvider.create().getCredentials();