AgentIdentityUtils.java

/*
 * Copyright 2025 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 com.google.api.client.json.GenericJson;
import com.google.api.client.json.JsonObjectParser;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Strings;
import com.google.common.collect.ImmutableList;
import com.google.common.io.BaseEncoding;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.security.GeneralSecurityException;
import java.security.MessageDigest;
import java.security.cert.CertificateFactory;
import java.security.cert.CertificateParsingException;
import java.security.cert.X509Certificate;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.regex.Pattern;
import org.slf4j.Logger;

/** Internal utility class for handling Agent Identity certificate-bound access tokens. */
final class AgentIdentityUtils {
  private static final Logger LOGGER = Slf4jUtils.getLogger(AgentIdentityUtils.class);

  static final String GOOGLE_API_CERTIFICATE_CONFIG = "GOOGLE_API_CERTIFICATE_CONFIG";
  static final String GOOGLE_API_PREVENT_AGENT_TOKEN_SHARING_FOR_GCP_SERVICES =
      "GOOGLE_API_PREVENT_AGENT_TOKEN_SHARING_FOR_GCP_SERVICES";

  /**
   * Patterns for matching Agent Identity SPIFFE trust domains.
   *
   * <p>SPIFFE (Secure Production Identity Framework for Everyone) is a set of open-source standards
   * for securely identifying software services. A SPIFFE ID is a URI structured as {@code
   * spiffe://<trust-domain>/<path>}.
   */
  private static final List<Pattern> AGENT_IDENTITY_SPIFFE_PATTERNS =
      ImmutableList.of(
          Pattern.compile("^agents\\.global\\.org-\\d+\\.system\\.id\\.goog$"),
          Pattern.compile("^agents\\.global\\.proj-\\d+\\.system\\.id\\.goog$"));

  /**
   * SAN (Subject Alternative Name) entry type for URIs.
   *
   * <p>SAN is an extension to X.509 certificates that allows additional identities, such as URIs,
   * to be bound to the subject of the certificate.
   */
  private static final int SAN_URI_TYPE = 6;

  private static final String SPIFFE_SCHEME_PREFIX = "spiffe://";

  // Polling configuration
  private static final int FAST_POLL_CYCLES = 50;
  private static final long FAST_POLL_INTERVAL_MS = 100; // 0.1 seconds
  private static final long SLOW_POLL_INTERVAL_MS = 500; // 0.5 seconds
  private static final long TOTAL_TIMEOUT_MS = 30000; // 30 seconds
  private static final List<Long> POLLING_INTERVALS;

  // Pre-calculates the sequence of polling intervals
  static {
    List<Long> intervals = new ArrayList<>();

    for (int i = 0; i < FAST_POLL_CYCLES; i++) {
      intervals.add(FAST_POLL_INTERVAL_MS);
    }

    long remainingTime = TOTAL_TIMEOUT_MS - (FAST_POLL_CYCLES * FAST_POLL_INTERVAL_MS);
    // Integer division is sufficient here as we want full cycles
    int slowPollCycles = (int) (remainingTime / SLOW_POLL_INTERVAL_MS);

    for (int i = 0; i < slowPollCycles; i++) {
      intervals.add(SLOW_POLL_INTERVAL_MS);
    }

    POLLING_INTERVALS = Collections.unmodifiableList(intervals);
  }

  // Interface to allow mocking System.getenv for tests without exposing it publicly.
  interface EnvReader {
    String getEnv(String name);
  }

  private static EnvReader envReader = System::getenv;

  /**
   * Internal interface to allow mocking time and sleep for tests. This is used to prevent tests
   * from running for long periods of time when polling is involved.
   */
  @VisibleForTesting
  interface TimeService {
    long currentTimeMillis();

    void sleep(long millis) throws InterruptedException;
  }

  private static TimeService timeService =
      new TimeService() {
        @Override
        public long currentTimeMillis() {
          return System.currentTimeMillis();
        }

        @Override
        public void sleep(long millis) throws InterruptedException {
          Thread.sleep(millis);
        }
      };

  private AgentIdentityUtils() {}

  /**
   * Gets the Agent Identity certificate if certificate is available and agent token sharing is not
   * disabled.
   *
   * @return The X509Certificate if found and Agent Identities are enabled, null otherwise.
   * @throws IOException If there is an error reading the certificate file after retries.
   */
  static X509Certificate getAgentIdentityCertificate() throws IOException {
    if (isOptedOut()) {
      return null;
    }

    String certConfigPath = envReader.getEnv(GOOGLE_API_CERTIFICATE_CONFIG);
    if (Strings.isNullOrEmpty(certConfigPath)) {
      return null;
    }

    String certPath = getCertificatePathWithRetry(certConfigPath);
    return parseCertificate(certPath);
  }

  /**
   * Checks if Agent Identity token sharing is disabled via an environment variable.
   *
   * @return {@code true} if the {@link #GOOGLE_API_PREVENT_AGENT_TOKEN_SHARING_FOR_GCP_SERVICES}
   *     variable is set to {@code "false"}, otherwise returns {@code false}.
   */
  private static boolean isOptedOut() {
    String optOut = envReader.getEnv(GOOGLE_API_PREVENT_AGENT_TOKEN_SHARING_FOR_GCP_SERVICES);
    return "false".equalsIgnoreCase(optOut);
  }

  /**
   * Polls for the certificate config file and the certificate file it references, and returns the
   * certificate's path.
   *
   * <p>This method will retry for a total of {@link #TOTAL_TIMEOUT_MS} milliseconds before failing.
   *
   * @param certConfigPath The path to the certificate configuration JSON file.
   * @return The path to the certificate file extracted from the config.
   * @throws IOException If the files cannot be found after the timeout, or if the thread is
   *     interrupted while waiting.
   */
  private static String getCertificatePathWithRetry(String certConfigPath) throws IOException {
    boolean warned = false;

    // Deterministic polling loop based on pre-calculated intervals.
    for (long sleepInterval : POLLING_INTERVALS) {
      try {
        if (Files.exists(Paths.get(certConfigPath))) {
          String certPath = extractCertPathFromConfig(certConfigPath);
          if (!Strings.isNullOrEmpty(certPath) && Files.exists(Paths.get(certPath))) {
            return certPath;
          }
        }
      } catch (IOException e) {
        // Do not log here to prevent noise in the logs per iteration.
        // Fall through to the sleep logic to retry.
      }

      // If we are here, we failed to find the certificate, log a warning only once.
      if (!warned) {
        Slf4jUtils.log(
            LOGGER,
            org.slf4j.event.Level.WARN,
            Collections.emptyMap(),
            String.format(
                "Certificate config file not found at %s (from %s environment variable). "
                    + "Retrying for up to %d seconds.",
                certConfigPath, GOOGLE_API_CERTIFICATE_CONFIG, TOTAL_TIMEOUT_MS / 1000));
        warned = true;
      }

      // Sleep before the next attempt.
      try {
        timeService.sleep(sleepInterval);
      } catch (InterruptedException e) {
        Thread.currentThread().interrupt();
        throw new IOException(
            "Interrupted while waiting for Agent Identity certificate files for bound token request.",
            e);
      }
    }

    // If the loop completes without returning, we have timed out.
    throw new IOException(
        "Unable to find Agent Identity certificate config or file for bound token request after multiple retries. "
            + "Token binding protection is failing. You can turn off this protection by setting "
            + GOOGLE_API_PREVENT_AGENT_TOKEN_SHARING_FOR_GCP_SERVICES
            + " to false to fall back to unbound tokens.");
  }

  /**
   * Parses the certificate configuration JSON file and extracts the path to the certificate.
   *
   * @param certConfigPath The path to the certificate configuration JSON file.
   * @return The certificate file path, or {@code null} if not found in the config.
   * @throws IOException If the configuration file cannot be read.
   */
  @SuppressWarnings("unchecked")
  private static String extractCertPathFromConfig(String certConfigPath) throws IOException {
    try (InputStream stream = new FileInputStream(certConfigPath)) {
      JsonObjectParser parser = new JsonObjectParser(OAuth2Utils.JSON_FACTORY);
      GenericJson config = parser.parseAndClose(stream, StandardCharsets.UTF_8, GenericJson.class);
      Map<String, Object> certConfigs = (Map<String, Object>) config.get("cert_configs");
      if (certConfigs != null) {
        Map<String, Object> workload = (Map<String, Object>) certConfigs.get("workload");
        if (workload != null) {
          return (String) workload.get("cert_path");
        }
      }
    }
    return null;
  }

  /**
   * Parses an X.509 certificate from the given file path.
   *
   * @param certPath The path to the certificate file.
   * @return The parsed {@link X509Certificate}.
   * @throws IOException If the certificate file cannot be read or parsed.
   */
  private static X509Certificate parseCertificate(String certPath) throws IOException {
    try (InputStream stream = new FileInputStream(certPath)) {
      CertificateFactory cf = CertificateFactory.getInstance("X.509");
      return (X509Certificate) cf.generateCertificate(stream);
    } catch (GeneralSecurityException e) {
      throw new IOException(
          "Failed to parse Agent Identity certificate for bound token request.", e);
    }
  }

  /**
   * Determines whether the provided certificate belongs to an "Agent Identity," indicating that the
   * SDK should request a certificate-bound access token.
   *
   * <p>Agent Identities are distinguished from other workloads by inspecting the <b>Subject
   * Alternative Name (SAN)</b> extension of the X.509 certificate for a <b>SPIFFE (Secure
   * Production Identity Framework for Everyone)</b> ID.
   *
   * <p>Specifically, this method iterates through the URI entries in the SAN and checks if any
   * {@code spiffe://} URI has a trust domain that matches one of the following Agent Identity
   * patterns:
   *
   * <ul>
   *   <li>{@code agents.global.org-[0-9]+.system.id.goog}
   *   <li>{@code agents.global.proj-[0-9]+.system.id.goog}
   * </ul>
   *
   * <p>If a matching SPIFFE ID is found, the certificate is identified as an Agent Identity, and
   * this method returns {@code true}.
   *
   * @param cert the X.509 certificate to inspect
   * @return {@code true} if the certificate represents an Agent Identity, {@code false} otherwise
   */
  static boolean shouldRequestBoundToken(X509Certificate cert) {
    try {
      Collection<List<?>> sans = cert.getSubjectAlternativeNames();
      if (sans == null) {
        return false;
      }
      for (List<?> san : sans) {
        // SAN entry is a list where first element is the type (Integer) and second is value (mostly
        // String)
        if (san.size() >= 2
            && san.get(0) instanceof Integer
            && (Integer) san.get(0) == SAN_URI_TYPE) {
          Object value = san.get(1);
          if (value instanceof String) {
            String uri = (String) value;
            if (uri.startsWith(SPIFFE_SCHEME_PREFIX)) {
              // Extract trust domain: spiffe://<trust_domain>/...
              String withoutScheme = uri.substring(SPIFFE_SCHEME_PREFIX.length());
              int slashIndex = withoutScheme.indexOf('/');
              String trustDomain =
                  (slashIndex == -1) ? withoutScheme : withoutScheme.substring(0, slashIndex);

              for (Pattern pattern : AGENT_IDENTITY_SPIFFE_PATTERNS) {
                if (pattern.matcher(trustDomain).matches()) {
                  return true;
                }
              }
            }
          }
        }
      }
    } catch (CertificateParsingException e) {
      LOGGER.warn("Failed to parse Subject Alternative Names from certificate", e);
    }
    return false;
  }

  /** Calculates the SHA-256 fingerprint of the certificate, Base64Url encoded without padding. */
  static String calculateCertificateFingerprint(X509Certificate cert) throws IOException {
    try {
      MessageDigest md = MessageDigest.getInstance("SHA-256");
      byte[] der = cert.getEncoded();
      md.update(der);
      byte[] digest = md.digest();
      return BaseEncoding.base64Url().omitPadding().encode(digest);
    } catch (GeneralSecurityException e) {
      throw new IOException("Failed to calculate fingerprint for Agent Identity certificate.", e);
    }
  }

  @VisibleForTesting
  static void setEnvReader(EnvReader reader) {
    envReader = reader;
  }

  @VisibleForTesting
  static void setTimeService(TimeService service) {
    timeService = service;
  }

  @VisibleForTesting
  static void resetTimeService() {
    timeService =
        new TimeService() {
          @Override
          public long currentTimeMillis() {
            return System.currentTimeMillis();
          }

          @Override
          public void sleep(long millis) throws InterruptedException {
            Thread.sleep(millis);
          }
        };
  }
}