feat(gax): append cred-type header for auth metrics (#3186)

This is POC change in gax-java for auth metrics requirements on token
usage. See go/googleapis-auth-metric-design for context.


[Credentials](https://github.com/googleapis/google-auth-library-java/blob/main/credentials/java/com/google/auth/Credentials.java)
will expose `getMetricsCredentialType()` method, this change appends it
to existing `x-goog-api-client` header

Note: Currently implement in gax at client level. There are 2 edge cases
not covered and will create followups for: if handwritten library
overrides credentials at rpc level; If handwritten library does not
build on gax. (ref: b/370039645, b/370038458)

related change in
`google-auth-library`googleapis/google-auth-library-java#1503
included in 1.28.0.

Author: zhumin8

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

@@ -44,6 +44,7 @@
 import com.google.api.gax.tracing.ApiTracerFactory;
 import com.google.api.gax.tracing.BaseApiTracerFactory;
 import com.google.auth.ApiKeyCredentials;
+import com.google.auth.CredentialTypeForMetrics;
 import com.google.auth.Credentials;
 import com.google.auth.oauth2.GdchCredentials;
 import com.google.auto.value.AutoValue;
@@ -210,7 +211,8 @@ public static ClientContext create(StubSettings settings) throws IOException {
     if (transportChannelProvider.needsExecutor() && settings.getExecutorProvider() != null) {
       transportChannelProvider = transportChannelProvider.withExecutor(backgroundExecutor);
     }
-    Map<String, String> headers = getHeadersFromSettings(settings);
+
+    Map<String, String> headers = getHeaders(settings, credentials);
     if (transportChannelProvider.needsHeaders()) {
       transportChannelProvider = transportChannelProvider.withHeaders(headers);
     }
@@ -318,8 +320,11 @@ static GdchCredentials getGdchCredentials(
   /**
    * Getting a header map from HeaderProvider and InternalHeaderProvider from settings with Quota
    * Project Id.
+   *
+   * <p>Then if credentials is present and its type for metrics is not {@code
+   * CredentialTypeForMetrics.DO_NOT_SEND}, append this type info to x-goog-api-client header.
    */
-  private static Map<String, String> getHeadersFromSettings(StubSettings settings) {
+  private static Map<String, String> getHeaders(StubSettings settings, Credentials credentials) {
     // Resolve conflicts when merging headers from multiple sources
     Map<String, String> userHeaders = settings.getHeaderProvider().getHeaders();
     Map<String, String> internalHeaders = settings.getInternalHeaderProvider().getHeaders();
@@ -346,6 +351,20 @@ private static Map<String, String> getHeadersFromSettings(StubSettings settings)
     effectiveHeaders.putAll(userHeaders);
     effectiveHeaders.putAll(conflictResolution);
 
+    return appendCredentialTypeToHeaderIfPresent(effectiveHeaders, credentials);
+  }
+
+  private static Map<String, String> appendCredentialTypeToHeaderIfPresent(
+      Map<String, String> effectiveHeaders, Credentials credentials) {
+    CredentialTypeForMetrics credentialTypeForMetrics =
+        credentials == null
+            ? CredentialTypeForMetrics.DO_NOT_SEND
+            : credentials.getMetricsCredentialType();
+    if (credentialTypeForMetrics != CredentialTypeForMetrics.DO_NOT_SEND) {
+      effectiveHeaders.computeIfPresent(
+          ApiClientHeaderProvider.getDefaultApiClientHeaderKey(),
+          (key, value) -> value + " cred-type/" + credentialTypeForMetrics.getLabel());
+    }
     return ImmutableMap.copyOf(effectiveHeaders);
   }
 

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

@@ -39,6 +39,7 @@
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 import com.google.api.core.ApiClock;
 import com.google.api.gax.core.BackgroundResource;
@@ -53,6 +54,7 @@
 import com.google.api.gax.rpc.testing.FakeStubSettings;
 import com.google.api.gax.rpc.testing.FakeTransportChannel;
 import com.google.auth.ApiKeyCredentials;
+import com.google.auth.CredentialTypeForMetrics;
 import com.google.auth.Credentials;
 import com.google.auth.oauth2.ComputeEngineCredentials;
 import com.google.auth.oauth2.GdchCredentials;
@@ -677,6 +679,86 @@ void testUserAgentConcat() throws Exception {
         .containsEntry("user-agent", "user-supplied-agent internal-agent");
   }
 
+  @Test
+  void testApiClientHeaderAppendsCredType() throws Exception {
+    GoogleCredentials googleCredentials = Mockito.mock(GoogleCredentials.class);
+    when(googleCredentials.getMetricsCredentialType())
+        .thenReturn(CredentialTypeForMetrics.USER_CREDENTIALS);
+
+    Map<String, String> headers =
+        setupTestForCredentialTokenUsageMetricsAndGetTransportChannelHeaders(
+            FixedCredentialsProvider.create(googleCredentials),
+            FixedHeaderProvider.create("x-goog-api-client", "internal-agent"));
+
+    assertThat(headers).containsEntry("x-goog-api-client", "internal-agent cred-type/u");
+  }
+
+  @Test
+  void testApiClientHeaderDoNotAppendsCredType_whenNoApiClientHeader() throws Exception {
+    GoogleCredentials googleCredentials = Mockito.mock(GoogleCredentials.class);
+    when(googleCredentials.getMetricsCredentialType())
+        .thenReturn(CredentialTypeForMetrics.USER_CREDENTIALS);
+
+    Map<String, String> headers =
+        setupTestForCredentialTokenUsageMetricsAndGetTransportChannelHeaders(
+            FixedCredentialsProvider.create(googleCredentials),
+            FixedHeaderProvider.create("some-other-header", "internal-agent"));
+
+    assertThat(headers).doesNotContainKey("x-goog-api-client");
+    assertThat(headers).containsEntry("some-other-header", "internal-agent");
+  }
+
+  @Test
+  void testApiClientHeaderDoNotAppendsCredType_whenNullCredentials() throws IOException {
+    Map<String, String> headers =
+        setupTestForCredentialTokenUsageMetricsAndGetTransportChannelHeaders(
+            NoCredentialsProvider.create(),
+            FixedHeaderProvider.create("x-goog-api-client", "internal-agent"));
+
+    assertThat(headers).containsKey("x-goog-api-client");
+    assertThat(headers.get("x-goog-api-client")).doesNotContain("cred-type/");
+  }
+
+  @Test
+  void testApiClientHeaderDoNotAppendsCredType_whenCredTypeDoNotSend() throws Exception {
+    GoogleCredentials googleCredentials = Mockito.mock(GoogleCredentials.class);
+    when(googleCredentials.getMetricsCredentialType())
+        .thenReturn(CredentialTypeForMetrics.DO_NOT_SEND);
+
+    Map<String, String> headers =
+        setupTestForCredentialTokenUsageMetricsAndGetTransportChannelHeaders(
+            FixedCredentialsProvider.create(googleCredentials),
+            FixedHeaderProvider.create("x-goog-api-client", "internal-agent"));
+
+    assertThat(headers).containsKey("x-goog-api-client");
+    assertThat(headers.get("x-goog-api-client")).doesNotContain("cred-type/");
+  }
+
+  private Map<String, String> setupTestForCredentialTokenUsageMetricsAndGetTransportChannelHeaders(
+      CredentialsProvider credentialsProvider, HeaderProvider headerProvider) throws IOException {
+    TransportChannelProvider transportChannelProvider =
+        new FakeTransportProvider(
+            FakeTransportChannel.create(new FakeChannel()),
+            null,
+            true,
+            null,
+            null,
+            DEFAULT_ENDPOINT);
+
+    ClientSettings.Builder builder =
+        new FakeClientSettings.Builder()
+            .setExecutorProvider(
+                FixedExecutorProvider.create(Mockito.mock(ScheduledExecutorService.class)))
+            .setTransportChannelProvider(transportChannelProvider)
+            .setCredentialsProvider(credentialsProvider)
+            .setInternalHeaderProvider(headerProvider);
+
+    ClientContext clientContext = ClientContext.create(builder.build());
+    FakeTransportChannel transportChannel =
+        (FakeTransportChannel) clientContext.getTransportChannel();
+    return transportChannel.getHeaders();
+  }
+
   private static String endpoint = "https://foo.googleapis.com";
   private static String mtlsEndpoint = "https://foo.mtls.googleapis.com";