fix: jwt authentication on batch-bigtable.googleapis.com

.kokoro/nightly/integration.cfg

@@ -8,7 +8,7 @@ env_vars: {
 
 env_vars: {
     key: "INTEGRATION_TEST_ARGS"
-    value: "-P bigtable-emulator-it,bigtable-prod-it -Dbigtable.project=gcloud-devel -Dbigtable.instance=google-cloud-bigtable -Dbigtable.table=integration-tests -Dbigtable.kms_key_name=projects/gcloud-devel/locations/us-east1/keyRings/cmek-test-key-ring/cryptoKeys/cmek-test-key -Dbigtable.wait-for-cmek-key-status=true"
+    value: "-P bigtable-emulator-it,bigtable-prod-it,bigtable-prod-batch-it -Dbigtable.project=gcloud-devel -Dbigtable.instance=google-cloud-bigtable -Dbigtable.table=integration-tests -Dbigtable.kms_key_name=projects/gcloud-devel/locations/us-east1/keyRings/cmek-test-key-ring/cryptoKeys/cmek-test-key -Dbigtable.wait-for-cmek-key-status=true"
 }
 
 env_vars: {

.kokoro/presubmit/integration.cfg

@@ -8,7 +8,7 @@ env_vars: {
 
 env_vars: {
     key: "INTEGRATION_TEST_ARGS"
-    value: "-P bigtable-emulator-it,bigtable-prod-it -Dbigtable.project=gcloud-devel -Dbigtable.instance=google-cloud-bigtable -Dbigtable.table=integration-tests -Dbigtable.kms_key_name=projects/gcloud-devel/locations/us-east1/keyRings/cmek-test-key-ring/cryptoKeys/cmek-test-key -Dbigtable.wait-for-cmek-key-status=true"
+    value: "-P bigtable-emulator-it,bigtable-prod-it,bigtable-prod-batch-it -Dbigtable.project=gcloud-devel -Dbigtable.instance=google-cloud-bigtable -Dbigtable.table=integration-tests -Dbigtable.kms_key_name=projects/gcloud-devel/locations/us-east1/keyRings/cmek-test-key-ring/cryptoKeys/cmek-test-key -Dbigtable.wait-for-cmek-key-status=true"
 }
 
 env_vars: {

google-cloud-bigtable/pom.xml

@@ -323,6 +323,41 @@
       </build>
     </profile>
 
+    <profile>
+      <id>bigtable-prod-batch-it</id>
+      <build>
+        <plugins>
+          <plugin>
+            <artifactId>maven-failsafe-plugin</artifactId>
+            <executions>
+              <execution>
+                <id>prod-batch-it</id>
+                <goals>
+                  <goal>integration-test</goal>
+                  <goal>verify</goal>
+                </goals>
+                <configuration>
+                  <skip>false</skip>
+
+                  <systemPropertyVariables>
+                    <bigtable.env>cloud</bigtable.env>
+                    <bigtable.data-endpoint>batch-${bigtable.cfe-data-endpoint}</bigtable.data-endpoint>
+                    <bigtable.admin-endpoint>${bigtable.cfe-admin-endpoint}</bigtable.admin-endpoint>
+                    <bigtable.grpc-log-dir>${project.build.directory}/test-grpc-logs/prod-batch-it</bigtable.grpc-log-dir>
+                  </systemPropertyVariables>
+                  <includes>
+                    <include>com.google.cloud.bigtable.data.v2.it.*IT</include>
+                  </includes>
+                  <summaryFile>${project.build.directory}/failsafe-reports/failsafe-summary-prod-batch-it.xml</summaryFile>
+                  <reportsDirectory>${project.build.directory}/failsafe-reports/prod-batch-it</reportsDirectory>
+                </configuration>
+              </execution>
+            </executions>
+          </plugin>
+        </plugins>
+      </build>
+    </profile>
+
     <profile>
       <id>bigtable-directpath-it</id>
       <build>

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/internal/JwtCredentialsWithAudience.java

@@ -0,0 +1,79 @@
+/*
+ * Copyright 2021 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.google.cloud.bigtable.data.v2.internal;
+
+import com.google.api.core.InternalApi;
+import com.google.auth.Credentials;
+import com.google.auth.RequestMetadataCallback;
+import com.google.auth.oauth2.ServiceAccountJwtAccessCredentials;
+import java.io.IOException;
+import java.net.URI;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.Executor;
+
+/**
+ * Internal helper to fix the mapping between JWT audiences and service endpoints.
+ *
+ * <p>In most cases JWT audiences correspond to service endpoints. However, in some cases they
+ * diverge. To workaround this, this class hardcodes the audience and forces the underlying
+ * implementation to use it.
+ *
+ * <p>Internal Only - public for technical reasons
+ */
+@InternalApi
+public class JwtCredentialsWithAudience extends Credentials {
+  private final ServiceAccountJwtAccessCredentials delegate;
+
+  public JwtCredentialsWithAudience(ServiceAccountJwtAccessCredentials delegate, URI audience) {
+    this.delegate = delegate.toBuilder().setDefaultAudience(audience).build();
+  }
+
+  @Override
+  public String getAuthenticationType() {
+    return delegate.getAuthenticationType();
+  }
+
+  @Override
+  public Map<String, List<String>> getRequestMetadata() throws IOException {
+    return delegate.getRequestMetadata();
+  }
+
+  @Override
+  public void getRequestMetadata(URI ignored, Executor executor, RequestMetadataCallback callback) {
+    delegate.getRequestMetadata(null, executor, callback);
+  }
+
+  @Override
+  public Map<String, List<String>> getRequestMetadata(URI ignored) throws IOException {
+    return delegate.getRequestMetadata(null);
+  }
+
+  @Override
+  public boolean hasRequestMetadata() {
+    return delegate.hasRequestMetadata();
+  }
+
+  @Override
+  public boolean hasRequestMetadataOnly() {
+    return delegate.hasRequestMetadataOnly();
+  }
+
+  @Override
+  public void refresh() throws IOException {
+    delegate.refresh();
+  }
+}

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStub.java

@@ -21,6 +21,7 @@
 import com.google.api.gax.batching.BatcherImpl;
 import com.google.api.gax.batching.FlowController;
 import com.google.api.gax.core.BackgroundResource;
+import com.google.api.gax.core.CredentialsProvider;
 import com.google.api.gax.core.FixedCredentialsProvider;
 import com.google.api.gax.grpc.GaxGrpcProperties;
 import com.google.api.gax.grpc.GrpcCallContext;
@@ -42,6 +43,7 @@
 import com.google.api.gax.tracing.TracedServerStreamingCallable;
 import com.google.api.gax.tracing.TracedUnaryCallable;
 import com.google.auth.Credentials;
+import com.google.auth.oauth2.ServiceAccountJwtAccessCredentials;
 import com.google.bigtable.v2.BigtableGrpc;
 import com.google.bigtable.v2.CheckAndMutateRowRequest;
 import com.google.bigtable.v2.CheckAndMutateRowResponse;
@@ -56,6 +58,7 @@
 import com.google.bigtable.v2.SampleRowKeysRequest;
 import com.google.bigtable.v2.SampleRowKeysResponse;
 import com.google.cloud.bigtable.Version;
+import com.google.cloud.bigtable.data.v2.internal.JwtCredentialsWithAudience;
 import com.google.cloud.bigtable.data.v2.internal.RequestContext;
 import com.google.cloud.bigtable.data.v2.models.BulkMutation;
 import com.google.cloud.bigtable.data.v2.models.ConditionalRowMutation;
@@ -94,6 +97,8 @@
 import io.opencensus.tags.Tagger;
 import io.opencensus.tags.Tags;
 import java.io.IOException;
+import java.net.URI;
+import java.net.URISyntaxException;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
@@ -146,6 +151,9 @@ public static EnhancedBigtableStubSettings finalizeSettings(
     // TODO: this implementation is on the cusp of unwieldy, if we end up adding more features
     // consider splitting it up by feature.
 
+    // workaround JWT audience issues
+    patchCredentials(builder);
+
     // Inject channel priming
     if (settings.isRefreshingChannel()) {
       // Fix the credentials so that they can be shared
@@ -218,6 +226,41 @@ public static EnhancedBigtableStubSettings finalizeSettings(
     return builder.build();
   }
 
+  private static void patchCredentials(EnhancedBigtableStubSettings.Builder settings)
+      throws IOException {
+    int i = settings.getEndpoint().lastIndexOf(":");
+    String host = settings.getEndpoint().substring(0, i);
+    String audience = settings.getJwtAudienceMapping().get(host);
+
+    if (audience == null) {
+      return;
+    }
+    URI audienceUri = null;
+    try {
+      audienceUri = new URI(audience);
+    } catch (URISyntaxException e) {
+      throw new IllegalStateException("invalid JWT audience override", e);
+    }
+
+    CredentialsProvider credentialsProvider = settings.getCredentialsProvider();
+    if (credentialsProvider == null) {
+      return;
+    }
+
+    Credentials credentials = credentialsProvider.getCredentials();
+    if (credentials == null) {
+      return;
+    }
+
+    if (!(credentials instanceof ServiceAccountJwtAccessCredentials)) {
+      return;
+    }
+
+    ServiceAccountJwtAccessCredentials jwtCreds = (ServiceAccountJwtAccessCredentials) credentials;
+    JwtCredentialsWithAudience patchedCreds = new JwtCredentialsWithAudience(jwtCreds, audienceUri);
+    settings.setCredentialsProvider(FixedCredentialsProvider.create(patchedCreds));
+  }
+
   public EnhancedBigtableStub(EnhancedBigtableStubSettings settings, ClientContext clientContext) {
     this.settings = settings;
     this.clientContext = clientContext;

google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/data/v2/stub/EnhancedBigtableStubSettings.java

@@ -16,6 +16,7 @@
 package com.google.cloud.bigtable.data.v2.stub;
 
 import com.google.api.core.BetaApi;
+import com.google.api.core.InternalApi;
 import com.google.api.gax.batching.BatchingCallSettings;
 import com.google.api.gax.batching.BatchingSettings;
 import com.google.api.gax.batching.FlowControlSettings;
@@ -151,12 +152,20 @@ public class EnhancedBigtableStubSettings extends StubSettings<EnhancedBigtableS
           .add("https://www.googleapis.com/auth/cloud-platform")
           .build();
 
+  /**
+   * In most cases, jwt audience == service name. However in some cases, this is not the case. The
+   * following mapping is used to patch the audience in a JWT token.
+   */
+  private static final Map<String, String> DEFAULT_JWT_AUDIENCE_MAPPING =
+      ImmutableMap.of("batch-bigtable.googleapis.com", "https://bigtable.googleapis.com/");
+
   private final String projectId;
   private final String instanceId;
   private final String appProfileId;
   private final boolean isRefreshingChannel;
   private ImmutableList<String> primedTableIds;
   private HeaderTracer headerTracer;
+  private final Map<String, String> jwtAudienceMapping;
 
   private final ServerStreamingCallSettings<Query, Row> readRowsSettings;
   private final UnaryCallSettings<Query, Row> readRowSettings;
@@ -191,6 +200,7 @@ private EnhancedBigtableStubSettings(Builder builder) {
     isRefreshingChannel = builder.isRefreshingChannel;
     primedTableIds = builder.primedTableIds;
     headerTracer = builder.headerTracer;
+    jwtAudienceMapping = builder.jwtAudienceMapping;
 
     // Per method settings.
     readRowsSettings = builder.readRowsSettings.build();
@@ -240,6 +250,11 @@ HeaderTracer getHeaderTracer() {
     return headerTracer;
   }
 
+  @InternalApi("Used for internal testing")
+  public Map<String, String> getJwtAudienceMapping() {
+    return jwtAudienceMapping;
+  }
+
   /** Returns a builder for the default ChannelProvider for this service. */
   public static InstantiatingGrpcChannelProvider.Builder defaultGrpcTransportProviderBuilder() {
     return BigtableStubSettings.defaultGrpcTransportProviderBuilder()
@@ -498,6 +513,7 @@ public static class Builder extends StubSettings.Builder<EnhancedBigtableStubSet
     private boolean isRefreshingChannel;
     private ImmutableList<String> primedTableIds;
     private HeaderTracer headerTracer;
+    private Map<String, String> jwtAudienceMapping;
 
     private final ServerStreamingCallSettings.Builder<Query, Row> readRowsSettings;
     private final UnaryCallSettings.Builder<Query, Row> readRowSettings;
@@ -522,6 +538,7 @@ private Builder() {
       this.isRefreshingChannel = false;
       primedTableIds = ImmutableList.of();
       headerTracer = HeaderTracer.newBuilder().build();
+      jwtAudienceMapping = DEFAULT_JWT_AUDIENCE_MAPPING;
       setCredentialsProvider(defaultCredentialsProviderBuilder().build());
 
       // Defaults provider
@@ -629,6 +646,7 @@ private Builder(EnhancedBigtableStubSettings settings) {
       isRefreshingChannel = settings.isRefreshingChannel;
       primedTableIds = settings.primedTableIds;
       headerTracer = settings.headerTracer;
+      jwtAudienceMapping = settings.jwtAudienceMapping;
 
       // Per method settings.
       readRowsSettings = settings.readRowsSettings.toBuilder();
@@ -762,6 +780,17 @@ HeaderTracer getHeaderTracer() {
       return headerTracer;
     }
 
+    @InternalApi("Used for internal testing")
+    public Builder setJwtAudienceMapping(Map<String, String> jwtAudienceMapping) {
+      this.jwtAudienceMapping = Preconditions.checkNotNull(jwtAudienceMapping);
+      return this;
+    }
+
+    @InternalApi("Used for internal testing")
+    public Map<String, String> getJwtAudienceMapping() {
+      return jwtAudienceMapping;
+    }
+
     /** Returns the builder for the settings used for calls to readRows. */
     public ServerStreamingCallSettings.Builder<Query, Row> readRowsSettings() {
       return readRowsSettings;