feat: self signed jwt support

gax/src/main/java/com/google/api/gax/core/GoogleCredentialsProvider.java

@@ -38,6 +38,8 @@
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
 import java.io.IOException;
+import java.net.URI;
+import java.util.ArrayList;
 import java.util.List;
 import javax.annotation.Nullable;
 
@@ -53,6 +55,9 @@ public abstract class GoogleCredentialsProvider implements CredentialsProvider {
 
   public abstract List<String> getScopesToApply();
 
+  @BetaApi
+  public abstract List<String> getDefaultScopes();
+
   @BetaApi
   public abstract List<String> getJwtEnabledScopes();
 
@@ -62,6 +67,16 @@ public abstract class GoogleCredentialsProvider implements CredentialsProvider {
 
   @Override
   public Credentials getCredentials() throws IOException {
+    return getCredentials(false, null);
+  }
+
+  /**
+   * If user doesn't provide custom endpoint and scopes, and the credentials is service account
+   * credentials, then we need to create a new ServiceAccountJwtAccessCredentials to use self signed
+   * jwt. See https://google.aip.dev/auth/4111.
+   */
+  public Credentials getCredentials(boolean endpointIsDefault, URI audienceForSelfSignedJwt)
+      throws IOException {
     GoogleCredentials credentials = getOAuth2Credentials();
     if (credentials == null) {
       credentials = GoogleCredentials.getApplicationDefault();
@@ -75,8 +90,12 @@ public Credentials getCredentials() throws IOException {
         break;
       }
     }
+    boolean useSelfSignedJwt = false;
+    if ((getScopesToApply().isEmpty() && endpointIsDefault) || hasJwtEnabledScope) {
+      useSelfSignedJwt = true;
+    }
     // Use JWT tokens when using a service account with an appropriate scope.
-    if (credentials instanceof ServiceAccountCredentials && hasJwtEnabledScope) {
+    if (credentials instanceof ServiceAccountCredentials && useSelfSignedJwt) {
       ServiceAccountCredentials serviceAccount = (ServiceAccountCredentials) credentials;
 
       return ServiceAccountJwtAccessCredentials.newBuilder()
@@ -85,18 +104,26 @@ public Credentials getCredentials() throws IOException {
           .setPrivateKey(serviceAccount.getPrivateKey())
           .setPrivateKeyId(serviceAccount.getPrivateKeyId())
           .setQuotaProjectId(serviceAccount.getQuotaProjectId())
+          .setDefaultAudience(audienceForSelfSignedJwt)
           .build();
     }
 
     if (credentials.createScopedRequired()) {
-      credentials = credentials.createScoped(getScopesToApply());
+      List<String> scopes = new ArrayList<String>(getDefaultScopes());
+      for (String scope : getScopesToApply()) {
+        if (!scopes.contains(scope)) {
+          scopes.add(scope);
+        }
+      }
+      credentials = credentials.createScoped(scopes);
     }
     return credentials;
   }
 
   public static Builder newBuilder() {
     return new AutoValue_GoogleCredentialsProvider.Builder()
-        .setJwtEnabledScopes(ImmutableList.<String>of());
+        .setJwtEnabledScopes(ImmutableList.<String>of())
+        .setDefaultScopes(ImmutableList.<String>of());
   }
 
   public abstract Builder toBuilder();
@@ -134,9 +161,20 @@ public abstract static class Builder {
     @BetaApi
     public abstract List<String> getJwtEnabledScopes();
 
+    /**
+     * Default scopes is for client libraries to use. Users should use setScopesToApply to set
+     * scopes, or setJwtEnabledScopes to force to use ServiceAccountJWTCredentials.
+     */
+    @BetaApi
+    public abstract Builder setDefaultScopes(List<String> val);
+
+    @BetaApi
+    public abstract List<String> getDefaultScopes();
+
     public GoogleCredentialsProvider build() {
       setScopesToApply(ImmutableList.copyOf(getScopesToApply()));
       setJwtEnabledScopes(ImmutableList.copyOf(getJwtEnabledScopes()));
+      setDefaultScopes(ImmutableList.copyOf(getDefaultScopes()));
       return autoBuild();
     }
 

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

@@ -33,17 +33,25 @@
 import com.google.api.core.BetaApi;
 import com.google.api.core.NanoClock;
 import com.google.api.gax.core.BackgroundResource;
+import com.google.api.gax.core.CredentialsProvider;
 import com.google.api.gax.core.ExecutorAsBackgroundResource;
 import com.google.api.gax.core.ExecutorProvider;
+import com.google.api.gax.core.FixedCredentialsProvider;
+import com.google.api.gax.core.GoogleCredentialsProvider;
 import com.google.api.gax.rpc.internal.QuotaProjectIdHidingCredentials;
 import com.google.api.gax.tracing.ApiTracerFactory;
 import com.google.api.gax.tracing.NoopApiTracerFactory;
 import com.google.auth.Credentials;
+import com.google.auth.oauth2.ServiceAccountCredentials;
+import com.google.auth.oauth2.ServiceAccountJwtAccessCredentials;
 import com.google.auto.value.AutoValue;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Sets;
 import java.io.IOException;
+import java.net.URI;
+import java.net.URISyntaxException;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
@@ -132,6 +140,73 @@ public static ClientContext create(ClientSettings settings) throws IOException {
     return create(settings.getStubSettings());
   }
 
+  /**
+   * If user doesn't provide custom endpoint and scopes, and the credentials is service account
+   * credentials, then we need to create a new ServiceAccountJwtAccessCredentials to use self signed
+   * jwt. See https://google.aip.dev/auth/4111.
+   */
+  @VisibleForTesting
+  static Credentials determineSelfSignedJWTCredentials(
+      CredentialsProvider provider, String endpoint, String defaultEndpoint) throws IOException {
+    if (endpoint == null || defaultEndpoint == null || !endpoint.equals(defaultEndpoint)) {
+      return provider.getCredentials();
+    }
+
+    Credentials credentials = null;
+
+    // DIREGAPIC clients may set the default endpoint like
+    // "https://compute.googleapis.com/compute/v1/projects/". For self signed jwt, we cannot use
+    // this default endpoint directly as the audience, because the correct audience should be
+    // "https://compute.googleapis.com/". Here we compute the audience.
+    URI selfSignedJwtAudience = null;
+    if (defaultEndpoint != null && defaultEndpoint.startsWith("https://")) {
+      // "https://" indicates the client is DIREGAPIC.
+      try {
+        URI uri = new URI(defaultEndpoint);
+        selfSignedJwtAudience = new URI(uri.getScheme(), uri.getHost(), "/", null);
+      } catch (URISyntaxException e) {
+        throw new IOException(
+            String.format("Default endpoint {%s} cannot be parsed.", defaultEndpoint));
+      }
+    }
+
+    // For service account credentials, if the endpoint is default endpoint, and user doesn't
+    // provide scopes, then we create a ServiceAccountJwtAccessCredentials and use this credentials
+    // instead. This credential uses self signed jwt, and it is more efficient since the network
+    // call to token endpoint is avoided. See https://google.aip.dev/auth/4111.
+    if (provider instanceof GoogleCredentialsProvider) {
+      // GoogleCredentialsProvider has this logic implemented.
+      credentials =
+          ((GoogleCredentialsProvider) provider)
+              .getCredentials(endpoint.equals(defaultEndpoint), selfSignedJwtAudience);
+    } else if (provider instanceof FixedCredentialsProvider) {
+      credentials = provider.getCredentials();
+      if (credentials instanceof ServiceAccountCredentials
+          && ((ServiceAccountCredentials) credentials).getScopes().isEmpty()
+          && endpoint.equals(defaultEndpoint)) {
+        // For FixedCredentialsProvider, we need to create a ServiceAccountJwtAccessCredentials by
+        // ourselves.
+        ServiceAccountCredentials serviceAccount = (ServiceAccountCredentials) credentials;
+
+        credentials =
+            ServiceAccountJwtAccessCredentials.newBuilder()
+                .setClientEmail(serviceAccount.getClientEmail())
+                .setClientId(serviceAccount.getClientId())
+                .setPrivateKey(serviceAccount.getPrivateKey())
+                .setPrivateKeyId(serviceAccount.getPrivateKeyId())
+                .setQuotaProjectId(serviceAccount.getQuotaProjectId())
+                .setDefaultAudience(selfSignedJwtAudience)
+                .build();
+      }
+    }
+
+    if (credentials == null) {
+      credentials = provider.getCredentials();
+    }
+
+    return credentials;
+  }
+
   /**
    * Instantiates the executor, credentials, and transport context based on the given client
    * settings.
@@ -142,7 +217,11 @@ public static ClientContext create(StubSettings settings) throws IOException {
     ExecutorProvider executorProvider = settings.getExecutorProvider();
     final ScheduledExecutorService executor = executorProvider.getExecutor();
 
-    Credentials credentials = settings.getCredentialsProvider().getCredentials();
+    Credentials credentials =
+        determineSelfSignedJWTCredentials(
+            settings.getCredentialsProvider(),
+            settings.getEndpoint(),
+            settings.getDefaultApiEndpoint());
 
     if (settings.getQuotaProjectId() != null) {
       // If the quotaProjectId is set, wrap original credentials with correct quotaProjectId as

gax/src/main/java/com/google/api/gax/rpc/StubSettings.java

@@ -70,6 +70,8 @@ public abstract class StubSettings<SettingsT extends StubSettings<SettingsT>> {
   private final TransportChannelProvider transportChannelProvider;
   private final ApiClock clock;
   private final String endpoint;
+  // defaultApiEndpoint is set by client libraries.
+  private final String defaultApiEndpoint;
   private final String quotaProjectId;
   @Nullable private final WatchdogProvider streamWatchdogProvider;
   @Nonnull private final Duration streamWatchdogCheckInterval;
@@ -84,6 +86,7 @@ protected StubSettings(Builder builder) {
     this.internalHeaderProvider = builder.internalHeaderProvider;
     this.clock = builder.clock;
     this.endpoint = builder.endpoint;
+    this.defaultApiEndpoint = builder.defaultApiEndpoint;
     this.quotaProjectId = builder.quotaProjectId;
     this.streamWatchdogProvider = builder.streamWatchdogProvider;
     this.streamWatchdogCheckInterval = builder.streamWatchdogCheckInterval;
@@ -120,6 +123,10 @@ public final String getEndpoint() {
     return endpoint;
   }
 
+  public final String getDefaultApiEndpoint() {
+    return defaultApiEndpoint;
+  }
+
   public final String getQuotaProjectId() {
     return quotaProjectId;
   }
@@ -155,6 +162,7 @@ public String toString() {
         .add("internalHeaderProvider", internalHeaderProvider)
         .add("clock", clock)
         .add("endpoint", endpoint)
+        .add("defaultApiEndpoint", defaultApiEndpoint)
         .add("quotaProjectId", quotaProjectId)
         .add("streamWatchdogProvider", streamWatchdogProvider)
         .add("streamWatchdogCheckInterval", streamWatchdogCheckInterval)
@@ -175,6 +183,7 @@ public abstract static class Builder<
     private ApiClock clock;
     private String endpoint;
     private String quotaProjectId;
+    private String defaultApiEndpoint = null;
     @Nullable private WatchdogProvider streamWatchdogProvider;
     @Nonnull private Duration streamWatchdogCheckInterval;
     @Nonnull private ApiTracerFactory tracerFactory;
@@ -188,6 +197,7 @@ protected Builder(StubSettings settings) {
       this.internalHeaderProvider = settings.internalHeaderProvider;
       this.clock = settings.clock;
       this.endpoint = settings.endpoint;
+      this.defaultApiEndpoint = settings.getDefaultApiEndpoint();
       this.quotaProjectId = settings.quotaProjectId;
       this.streamWatchdogProvider = settings.streamWatchdogProvider;
       this.streamWatchdogCheckInterval = settings.streamWatchdogCheckInterval;
@@ -337,6 +347,12 @@ public B setEndpoint(String endpoint) {
       return self();
     }
 
+    // Make it protected so only client libraries can set it in client's stubSetting class.
+    protected B setDefaultApiEndpoint(String endpoint) {
+      this.defaultApiEndpoint = endpoint;
+      return self();
+    }
+
     public B setQuotaProjectId(String quotaProjectId) {
       this.quotaProjectId = quotaProjectId;
       return self();
@@ -408,6 +424,10 @@ public String getEndpoint() {
       return endpoint;
     }
 
+    public String getDefaultApiEndpoint() {
+      return defaultApiEndpoint;
+    }
+
     /** Gets the QuotaProjectId that was previously set on this Builder. */
     public String getQuotaProjectId() {
       return quotaProjectId;
@@ -445,6 +465,7 @@ public String toString() {
           .add("internalHeaderProvider", internalHeaderProvider)
           .add("clock", clock)
           .add("endpoint", endpoint)
+          .add("defaultApiEndpoint", defaultApiEndpoint)
           .add("quotaProjectId", quotaProjectId)
           .add("streamWatchdogProvider", streamWatchdogProvider)
           .add("streamWatchdogCheckInterval", streamWatchdogCheckInterval)

gax/src/test/java/com/google/api/gax/core/GoogleCredentialsProviderTest.java

@@ -35,6 +35,7 @@
 import com.google.auth.oauth2.ServiceAccountCredentials;
 import com.google.auth.oauth2.ServiceAccountJwtAccessCredentials;
 import com.google.common.collect.ImmutableList;
+import java.net.URI;
 import java.security.PrivateKey;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -43,15 +44,18 @@
 
 @RunWith(JUnit4.class)
 public class GoogleCredentialsProviderTest {
+  public ServiceAccountCredentials createServiceAccountCredentials() {
+    return ServiceAccountCredentials.newBuilder()
+        .setClientId("fake-client-id")
+        .setClientEmail("fake@example.com")
+        .setPrivateKeyId("fake-private-key")
+        .setPrivateKey(Mockito.mock(PrivateKey.class))
+        .build();
+  }
+
   @Test
   public void serviceAccountReplacedWithJwtTokens() throws Exception {
-    ServiceAccountCredentials serviceAccountCredentials =
-        ServiceAccountCredentials.newBuilder()
-            .setClientId("fake-client-id")
-            .setClientEmail("fake@example.com")
-            .setPrivateKeyId("fake-private-key")
-            .setPrivateKey(Mockito.mock(PrivateKey.class))
-            .build();
+    ServiceAccountCredentials serviceAccountCredentials = createServiceAccountCredentials();
 
     GoogleCredentialsProvider provider =
         GoogleCredentialsProvider.newBuilder()
@@ -71,13 +75,7 @@ public void serviceAccountReplacedWithJwtTokens() throws Exception {
 
   @Test
   public void noJwtWithoutScopeMatch() throws Exception {
-    ServiceAccountCredentials serviceAccountCredentials =
-        ServiceAccountCredentials.newBuilder()
-            .setClientId("fake-client-id")
-            .setClientEmail("fake@example.com")
-            .setPrivateKeyId("fake-private-key")
-            .setPrivateKey(Mockito.mock(PrivateKey.class))
-            .build();
+    ServiceAccountCredentials serviceAccountCredentials = createServiceAccountCredentials();
 
     GoogleCredentialsProvider provider =
         GoogleCredentialsProvider.newBuilder()
@@ -100,4 +98,28 @@ public void noJwtWithoutScopeMatch() throws Exception {
         .isEqualTo(serviceAccountCredentials.getPrivateKey());
     assertThat(serviceAccountCredentials2.getScopes()).containsExactly("scope1", "scope2");
   }
+
+  @Test
+  public void selfSignedJwt() throws Exception {
+    ServiceAccountCredentials serviceAccountCredentials = createServiceAccountCredentials();
+
+    // Create a provider with default scopes only.
+    GoogleCredentialsProvider provider =
+        GoogleCredentialsProvider.newBuilder()
+            .setScopesToApply(ImmutableList.<String>of())
+            .setDefaultScopes(ImmutableList.of("scope1", "scope2"))
+            .setOAuth2Credentials(serviceAccountCredentials)
+            .build();
+
+    // By default we should get a ServiceAccountCredentials.
+    Credentials credentials = provider.getCredentials();
+    assertThat(credentials).isInstanceOf(ServiceAccountCredentials.class);
+
+    // If endpoint is default, then we should also get a ServiceAccountJwtAccessCredentials instead.
+    URI audience = new URI("https://foo.googleapis.com");
+    credentials = provider.getCredentials(true, audience);
+    assertThat(credentials).isInstanceOf(ServiceAccountJwtAccessCredentials.class);
+    assertThat(((ServiceAccountJwtAccessCredentials) credentials).toBuilder().getDefaultAudience())
+        .isEqualTo(audience);
+  }
 }