feat: Byoid metrics framework (#1232)

* feat: adding byoid metrics framework

* fix: formatting

* fix: formatting

* Addressing PR comments

* Update oauth2_http/java/com/google/auth/oauth2/MetricsUtils.java

Co-authored-by: Leo <[email protected]>

* Update oauth2_http/java/com/google/auth/oauth2/MetricsUtils.java

Co-authored-by: Leo <[email protected]>

* Update oauth2_http/java/com/google/auth/oauth2/MetricsUtils.java

Co-authored-by: Leo <[email protected]>

* fix test

* formatting

* regex fix

* lint fix

* remove builder method and fix tests

* rename handler class, fix tests

* addressing comments, credential source type

---------

Co-authored-by: Leo <[email protected]>

Author: aeitzman

oauth2_http/java/com/google/auth/oauth2/AwsCredentials.java

@@ -197,6 +197,11 @@ public GoogleCredentials createScoped(Collection<String> newScopes) {
     return new AwsCredentials((AwsCredentials.Builder) newBuilder(this).setScopes(newScopes));
   }
 
+  @Override
+  String getCredentialSourceType() {
+    return "aws";
+  }
+
   private String retrieveResource(String url, String resourceName, Map<String, Object> headers)
       throws IOException {
     return retrieveResource(url, resourceName, HttpMethods.GET, headers, /* content= */ null);

oauth2_http/java/com/google/auth/oauth2/ExternalAccountCredentials.java

@@ -33,6 +33,7 @@
 
 import static com.google.common.base.Preconditions.checkNotNull;
 
+import com.google.api.client.http.HttpHeaders;
 import com.google.api.client.json.GenericJson;
 import com.google.api.client.json.JsonObjectParser;
 import com.google.auth.RequestMetadataCallback;
@@ -90,6 +91,7 @@ abstract static class CredentialSource implements java.io.Serializable {
   private final CredentialSource credentialSource;
   private final Collection<String> scopes;
   private final ServiceAccountImpersonationOptions serviceAccountImpersonationOptions;
+  private ExternalAccountMetricsHandler metricsHandler;
 
   @Nullable private final String tokenInfoUrl;
   @Nullable private final String serviceAccountImpersonationUrl;
@@ -224,6 +226,8 @@ protected ExternalAccountCredentials(
       validateServiceAccountImpersonationInfoUrl(serviceAccountImpersonationUrl);
     }
 
+    this.metricsHandler = new ExternalAccountMetricsHandler(this);
+
     this.impersonatedCredentials = buildImpersonatedCredentials();
   }
 
@@ -274,6 +278,11 @@ protected ExternalAccountCredentials(ExternalAccountCredentials.Builder builder)
       validateServiceAccountImpersonationInfoUrl(serviceAccountImpersonationUrl);
     }
 
+    this.metricsHandler =
+        builder.metricsHandler == null
+            ? new ExternalAccountMetricsHandler(this)
+            : builder.metricsHandler;
+
     this.impersonatedCredentials = buildImpersonatedCredentials();
   }
 
@@ -505,6 +514,12 @@ protected AccessToken exchangeExternalCredentialForAccessToken(
       requestHandler.setInternalOptions(options.toString());
     }
 
+    // Set BYOID Metrics header.
+    HttpHeaders additionalHeaders = new HttpHeaders();
+    additionalHeaders.set(
+        MetricsUtils.API_CLIENT_HEADER, this.metricsHandler.getExternalAccountMetricsHeader());
+    requestHandler.setHeaders(additionalHeaders);
+
     if (stsTokenExchangeRequest.getInternalOptions() != null) {
       // Overwrite internal options. Let subclass handle setting options.
       requestHandler.setInternalOptions(stsTokenExchangeRequest.getInternalOptions());
@@ -589,6 +604,10 @@ public ServiceAccountImpersonationOptions getServiceAccountImpersonationOptions(
     return serviceAccountImpersonationOptions;
   }
 
+  String getCredentialSourceType() {
+    return "unknown";
+  }
+
   EnvironmentProvider getEnvironmentProvider() {
     return environmentProvider;
   }
@@ -663,8 +682,11 @@ static final class ServiceAccountImpersonationOptions implements java.io.Seriali
 
     private final int lifetime;
 
+    final boolean customTokenLifetimeRequested;
+
     ServiceAccountImpersonationOptions(Map<String, Object> optionsMap) {
-      if (!optionsMap.containsKey(TOKEN_LIFETIME_SECONDS_KEY)) {
+      customTokenLifetimeRequested = optionsMap.containsKey(TOKEN_LIFETIME_SECONDS_KEY);
+      if (!customTokenLifetimeRequested) {
         lifetime = DEFAULT_TOKEN_LIFETIME_SECONDS;
         return;
       }
@@ -714,6 +736,7 @@ public abstract static class Builder extends GoogleCredentials.Builder {
     @Nullable protected String workforcePoolUserProject;
     @Nullable protected ServiceAccountImpersonationOptions serviceAccountImpersonationOptions;
     @Nullable protected String universeDomain;
+    @Nullable protected ExternalAccountMetricsHandler metricsHandler;
 
     protected Builder() {}
 
@@ -733,6 +756,7 @@ protected Builder(ExternalAccountCredentials credentials) {
       this.workforcePoolUserProject = credentials.workforcePoolUserProject;
       this.serviceAccountImpersonationOptions = credentials.serviceAccountImpersonationOptions;
       this.universeDomain = credentials.universeDomain;
+      this.metricsHandler = credentials.metricsHandler;
     }
 
     /**

oauth2_http/java/com/google/auth/oauth2/ExternalAccountMetricsHandler.java

@@ -0,0 +1,78 @@
+/*
+ * Copyright 2023 Google LLC
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are
+ * met:
+ *
+ *    * Redistributions of source code must retain the above copyright
+ * notice, this list of conditions and the following disclaimer.
+ *    * Redistributions in binary form must reproduce the above
+ * copyright notice, this list of conditions and the following disclaimer
+ * in the documentation and/or other materials provided with the
+ * distribution.
+ *
+ *    * Neither the name of Google LLC nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+ * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+ * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+ * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+ * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+ * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+ * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+ * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+
+package com.google.auth.oauth2;
+
+/**
+ * A handler for generating the x-goog-api-client header value for BYOID external account
+ * credentials.
+ */
+class ExternalAccountMetricsHandler implements java.io.Serializable {
+  private static final String SOURCE_KEY = "source";
+  private static final String IMPERSONATION_KEY = "sa-impersonation";
+  private static final String CONFIG_LIFETIME_KEY = "config-lifetime";
+  private static final String BYOID_METRICS_SECTION = "google-byoid-sdk";
+
+  private final boolean configLifetime;
+  private final boolean saImpersonation;
+  private String credentialSourceType;
+
+  /**
+   * Constructor for the external account metrics handler.
+   *
+   * @param creds the {@code ExternalAccountCredentials} object to set the external account metrics
+   *     options from.
+   */
+  ExternalAccountMetricsHandler(ExternalAccountCredentials creds) {
+    this.saImpersonation = creds.getServiceAccountImpersonationUrl() != null;
+    this.configLifetime =
+        creds.getServiceAccountImpersonationOptions().customTokenLifetimeRequested;
+    this.credentialSourceType = creds.getCredentialSourceType();
+  }
+
+  /**
+   * Gets the external account metrics header value for the x-goog-api-client header.
+   *
+   * @return the header value.
+   */
+  String getExternalAccountMetricsHeader() {
+    return String.format(
+        "%s %s %s/%s %s/%s %s/%s",
+        MetricsUtils.getLanguageAndAuthLibraryVersions(),
+        BYOID_METRICS_SECTION,
+        SOURCE_KEY,
+        this.credentialSourceType,
+        IMPERSONATION_KEY,
+        this.saImpersonation,
+        CONFIG_LIFETIME_KEY,
+        this.configLifetime);
+  }
+}

oauth2_http/java/com/google/auth/oauth2/IdentityPoolCredentials.java

@@ -38,6 +38,7 @@
 import com.google.api.client.json.GenericJson;
 import com.google.api.client.json.JsonObjectParser;
 import com.google.auth.oauth2.IdentityPoolCredentials.IdentityPoolCredentialSource.CredentialFormatType;
+import com.google.auth.oauth2.IdentityPoolCredentials.IdentityPoolCredentialSource.IdentityPoolCredentialSourceType;
 import com.google.common.io.CharStreams;
 import java.io.BufferedReader;
 import java.io.File;
@@ -192,6 +193,16 @@ public String retrieveSubjectToken() throws IOException {
     return getSubjectTokenFromMetadataServer();
   }
 
+  @Override
+  String getCredentialSourceType() {
+    if (((IdentityPoolCredentialSource) this.getCredentialSource()).credentialSourceType
+        == IdentityPoolCredentialSourceType.FILE) {
+      return "file";
+    } else {
+      return "url";
+    }
+  }
+
   private String retrieveSubjectTokenFromCredentialFile() throws IOException {
     String credentialFilePath = identityPoolCredentialSource.credentialLocation;
     if (!Files.exists(Paths.get(credentialFilePath), LinkOption.NOFOLLOW_LINKS)) {

oauth2_http/java/com/google/auth/oauth2/MetricsUtils.java

@@ -0,0 +1,70 @@
+/*
+ * Copyright 2023 Google LLC
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are
+ * met:
+ *
+ *    * Redistributions of source code must retain the above copyright
+ * notice, this list of conditions and the following disclaimer.
+ *    * Redistributions in binary form must reproduce the above
+ * copyright notice, this list of conditions and the following disclaimer
+ * in the documentation and/or other materials provided with the
+ * distribution.
+ *
+ *    * Neither the name of Google LLC nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+ * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+ * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+ * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+ * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+ * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+ * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+ * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+
+package com.google.auth.oauth2;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.util.Properties;
+
+class MetricsUtils {
+  static final String API_CLIENT_HEADER = "x-goog-api-client";
+  private static final String authLibraryVersion = getAuthLibraryVersion();
+  private static final String javaLanguageVersion = System.getProperty("java.version");
+
+  /**
+   * Gets the x-goog-api-client header value for the current Java language version and the auth
+   * library version.
+   *
+   * @return the header value.
+   */
+  static String getLanguageAndAuthLibraryVersions() {
+    return String.format("gl-java/%s auth/%s", javaLanguageVersion, authLibraryVersion);
+  }
+
+  private static String getAuthLibraryVersion() {
+    // Attempt to read the library's version from a properties file generated during the build.
+    // This value should be read and cached for later use.
+    String version = "unknown-version";
+    try (InputStream inputStream =
+        MetricsUtils.class.getResourceAsStream(
+            "/com/google/auth/oauth2/google-auth-library.properties")) {
+      if (inputStream != null) {
+        final Properties properties = new Properties();
+        properties.load(inputStream);
+        version = properties.getProperty("google-auth-library.version");
+      }
+    } catch (IOException e) {
+      // Ignore.
+    }
+    return version;
+  }
+}

oauth2_http/java/com/google/auth/oauth2/PluggableAuthCredentials.java

@@ -292,6 +292,11 @@ public PluggableAuthCredentials createScoped(Collection<String> newScopes) {
         (PluggableAuthCredentials.Builder) newBuilder(this).setScopes(newScopes));
   }
 
+  @Override
+  String getCredentialSourceType() {
+    return "executable";
+  }
+
   public static Builder newBuilder() {
     return new Builder();
   }

oauth2_http/javatests/com/google/auth/oauth2/AwsCredentialsTest.java

@@ -131,6 +131,11 @@ public void refreshAccessToken_withoutServiceAccountImpersonation() throws IOExc
     AccessToken accessToken = awsCredential.refreshAccessToken();
 
     assertEquals(transportFactory.transport.getAccessToken(), accessToken.getTokenValue());
+
+    // Validate metrics header is set correctly on the sts request.
+    Map<String, List<String>> headers =
+        transportFactory.transport.getRequests().get(3).getHeaders();
+    ExternalAccountCredentialsTest.validateMetricsHeader(headers, "aws", false, false);
   }
 
   @Test
@@ -142,18 +147,26 @@ public void refreshAccessToken_withServiceAccountImpersonation() throws IOExcept
 
     AwsCredentials awsCredential =
         (AwsCredentials)
-            AwsCredentials.newBuilder(AWS_CREDENTIAL)
+            AwsCredentials.newBuilder()
+                .setHttpTransportFactory(transportFactory)
+                .setAudience("audience")
+                .setSubjectTokenType("subjectTokenType")
                 .setTokenUrl(transportFactory.transport.getStsUrl())
+                .setTokenInfoUrl("tokenInfoUrl")
+                .setCredentialSource(buildAwsCredentialSource(transportFactory))
                 .setServiceAccountImpersonationUrl(
                     transportFactory.transport.getServiceAccountImpersonationUrl())
-                .setHttpTransportFactory(transportFactory)
-                .setCredentialSource(buildAwsCredentialSource(transportFactory))
                 .build();
 
     AccessToken accessToken = awsCredential.refreshAccessToken();
 
     assertEquals(
         transportFactory.transport.getServiceAccountAccessToken(), accessToken.getTokenValue());
+
+    // Validate metrics header is set correctly on the sts request.
+    Map<String, List<String>> headers =
+        transportFactory.transport.getRequests().get(6).getHeaders();
+    ExternalAccountCredentialsTest.validateMetricsHeader(headers, "aws", true, false);
   }
 
   @Test
@@ -165,12 +178,15 @@ public void refreshAccessToken_withServiceAccountImpersonationOptions() throws I
 
     AwsCredentials awsCredential =
         (AwsCredentials)
-            AwsCredentials.newBuilder(AWS_CREDENTIAL)
+            AwsCredentials.newBuilder()
+                .setHttpTransportFactory(transportFactory)
+                .setAudience("audience")
+                .setSubjectTokenType("subjectTokenType")
                 .setTokenUrl(transportFactory.transport.getStsUrl())
+                .setTokenInfoUrl("tokenInfoUrl")
+                .setCredentialSource(buildAwsCredentialSource(transportFactory))
                 .setServiceAccountImpersonationUrl(
                     transportFactory.transport.getServiceAccountImpersonationUrl())
-                .setHttpTransportFactory(transportFactory)
-                .setCredentialSource(buildAwsCredentialSource(transportFactory))
                 .setServiceAccountImpersonationOptions(
                     ExternalAccountCredentialsTest.buildServiceAccountImpersonationOptions(2800))
                 .build();
@@ -187,6 +203,11 @@ public void refreshAccessToken_withServiceAccountImpersonationOptions() throws I
             .parseAndClose(GenericJson.class);
 
     assertEquals("2800s", query.get("lifetime"));
+
+    // Validate metrics header is set correctly on the sts request.
+    Map<String, List<String>> headers =
+        transportFactory.transport.getRequests().get(6).getHeaders();
+    ExternalAccountCredentialsTest.validateMetricsHeader(headers, "aws", true, true);
   }
 
   @Test