fix(auth): verify public/private key match with retries in mTLS keystore loading

google-auth-library-java/oauth2_http/java/com/google/auth/mtls/MtlsKeyStoreUtils.java

@@ -64,11 +64,13 @@
  *
  * <p>Everything else matches the upstream method: key store type (JKS), alias, empty password,
  * accepted key format (PKCS#8 {@code PRIVATE KEY}), which private key is chosen, and error
- * messages. Any input that the upstream method accepted is still accepted, even though more of it
- * is now used: once a certificate and a private key have been found (the point where the upstream
- * method stopped reading), a malformed later PEM section ends the input instead of failing, and an
- * extra certificate that cannot be parsed is left out of the chain. Unlike the upstream method, the
- * input is decoded as UTF-8 rather than the platform default charset, and it is read to the end.
+ * messages. Any valid certificate-and-key pair that the upstream method accepted is still accepted,
+ * even though more of the input is now used: once a certificate and a private key have been found
+ * (the point where the upstream method stopped reading), a malformed later PEM section ends the
+ * input instead of failing, and an extra certificate that cannot be parsed is left out of the
+ * chain. Unlike the upstream method, the input is decoded as UTF-8 rather than the platform default
+ * charset, it is read to the end, and the private key is verified against the leaf certificate's
+ * public key.
  */
 @NullMarked
 final class MtlsKeyStoreUtils {
@@ -89,7 +91,8 @@ private MtlsKeyStoreUtils() {}
    * @throws IllegalArgumentException if the input has no certificate or no private key, or if a PEM
    *     section before both a certificate and a private key have been read is malformed
    * @throws GeneralSecurityException if the first certificate or the private key cannot be parsed,
-   *     or the key store cannot be created
+   *     if the private key does not match the leaf certificate's public key, or if the key store
+   *     cannot be created
    * @throws IOException if the input cannot be read
    */
   static KeyStore createMtlsKeyStore(InputStream certAndKey)
@@ -143,9 +146,12 @@ static KeyStore createMtlsKeyStore(InputStream certAndKey)
         LOGGER.log(Level.FINE, "Skipping unparsable extra certificate in mTLS chain", e);
       }
     }
+    X509Certificate leaf = chain.get(0);
+    String keyAlgorithm = leaf.getPublicKey().getAlgorithm();
     PrivateKey key =
-        KeyFactory.getInstance(chain.get(0).getPublicKey().getAlgorithm())
+        KeyFactory.getInstance("ECDSA".equals(keyAlgorithm) ? "EC" : keyAlgorithm)
             .generatePrivate(new PKCS8EncodedKeySpec(privateKey));
+    MtlsUtils.verifyKeyPair(leaf, key);
 
     KeyStore keyStore = KeyStore.getInstance("JKS");
     keyStore.load(null);

google-auth-library-java/oauth2_http/java/com/google/auth/mtls/MtlsUtils.java

@@ -42,12 +42,21 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.InterruptedIOException;
+import java.nio.charset.StandardCharsets;
 import java.nio.file.AccessDeniedException;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.nio.file.attribute.BasicFileAttributes;
+import java.security.GeneralSecurityException;
+import java.security.KeyFactory;
 import java.security.MessageDigest;
+import java.security.PrivateKey;
+import java.security.PublicKey;
+import java.security.Signature;
+import java.security.SignatureException;
+import java.security.cert.X509Certificate;
+import java.security.spec.X509EncodedKeySpec;
 import java.util.Locale;
 import org.jspecify.annotations.NullMarked;
 import org.jspecify.annotations.Nullable;
@@ -64,6 +73,9 @@
   static final String WELL_KNOWN_CERTIFICATE_CONFIG_FILE = "certificate_config.json";
   static final String CLOUDSDK_CONFIG_DIRECTORY = "gcloud";
 
+  private static final byte[] VERIFICATION_DATA =
+      "verification-data".getBytes(StandardCharsets.UTF_8);
+
   /**
    * Well-known directory where workload SPIFFE credentials and certificate configurations are
    * delivered at container startup.
@@ -151,13 +163,69 @@
     timeService = service != null ? service : Thread::sleep;
   }
 
+  /**
+   * Verifies that {@code privateKey} corresponds to the public key in {@code cert} by performing a
+   * test signature and verification ({@code SHA256withRSA} for RSA keys, {@code SHA256withECDSA}
+   * for EC/ECDSA keys).
+   *
+   * <p>Certificates with key algorithms other than RSA, EC, or ECDSA return without error so
+   * callers that accept any JCA key algorithm (such as {@link MtlsKeyStoreUtils}) pass through
+   * unchanged.
+   *
+   * @param cert the X.509 certificate whose public key is checked
+   * @param privateKey the private key to verify against {@code cert}
+   * @throws GeneralSecurityException if {@code privateKey} does not match {@code cert}'s public
+   *     key, or if signature initialization, signing, or verification fails
+   */
+  @InternalApi
+  public static void verifyKeyPair(X509Certificate cert, PrivateKey privateKey)
+      throws GeneralSecurityException {
+    PublicKey publicKey = cert.getPublicKey();
+    String keyAlgorithm = publicKey.getAlgorithm();
+    String sigAlg;
+    if ("RSA".equals(keyAlgorithm)) {
+      sigAlg = "SHA256withRSA";
+    } else if ("EC".equals(keyAlgorithm) || "ECDSA".equals(keyAlgorithm)) {
+      sigAlg = "SHA256withECDSA";
+    } else {
+      return;
+    }
+
+    if ("ECDSA".equals(keyAlgorithm) && publicKey.getEncoded() != null) {
+      // SunEC rejects keys whose algorithm is "ECDSA" rather than "EC"; normalize the key.
+      publicKey =
+          KeyFactory.getInstance("EC")
+              .generatePublic(new X509EncodedKeySpec(publicKey.getEncoded()));
+    }
+
+    Signature signer = Signature.getInstance(sigAlg);
+    signer.initSign(privateKey);
+    signer.update(VERIFICATION_DATA);
+    byte[] signature = signer.sign();
+
+    Signature verifier = Signature.getInstance(sigAlg);
+    verifier.initVerify(publicKey);
+    verifier.update(VERIFICATION_DATA);
+
+    boolean matches;
+    try {
+      matches = verifier.verify(signature);
+    } catch (SignatureException e) {
+      throw new GeneralSecurityException("Certificate and private key do not match", e);
+    }
+
+    if (!matches) {
+      throw new GeneralSecurityException("Certificate and private key do not match");
+    }
+  }
+
   /**
    * Checks whether {@code pathStr} resides within {@code targetWellKnownDir} (or either of the
    * default well-known SPIFFE directories {@link #WELL_KNOWN_SPIFFE_DIR} and {@link
    * #WELL_KNOWN_RUN_SPIFFE_DIR} when {@code targetWellKnownDir} is {@link #WELL_KNOWN_SPIFFE_DIR}).
    */
   @InternalApi
   public static boolean isPathInWellKnownDir(@Nullable String pathStr, String targetWellKnownDir) {
     if (Strings.isNullOrEmpty(pathStr) || Strings.isNullOrEmpty(targetWellKnownDir)) {
       return false;
     }
@@ -207,7 +275,7 @@
    * @throws IOException if the thread is interrupted while sleeping during startup polling
    */
   @InternalApi
   public static boolean waitForWorkloadConfigStartupIfNeeded(
       File configFile,
       String targetWellKnownDir,
       TimeService sleepService,

google-auth-library-java/oauth2_http/java/com/google/auth/mtls/X509Provider.java

@@ -35,13 +35,17 @@
 import com.google.auth.oauth2.PropertyProvider;
 import com.google.auth.oauth2.SystemEnvironmentProvider;
 import com.google.auth.oauth2.SystemPropertyProvider;
+import com.google.common.annotations.VisibleForTesting;
 import java.io.File;
 import java.io.FileInputStream;
 import java.io.FileNotFoundException;
 import java.io.IOException;
 import java.io.InputStream;
+import java.io.InterruptedIOException;
 import java.io.SequenceInputStream;
 import java.security.KeyStore;
+import java.util.logging.Level;
+import java.util.logging.Logger;
 import org.jspecify.annotations.NullMarked;
 import org.jspecify.annotations.Nullable;
 
@@ -52,6 +56,25 @@
 @NullMarked
 @InternalApi
 public class X509Provider implements MtlsProvider {
+  private static final Logger LOGGER = Logger.getLogger(X509Provider.class.getName());
+
+  // Retries for verifying certificate and private key matching during atomic key rotation.
+  static final int CERT_KEY_MATCH_RETRIES = 3;
+  static final long CERT_KEY_MATCH_RETRY_INTERVAL_MS = 100;
+
+  @VisibleForTesting
+  @FunctionalInterface
+  interface TimeService {
+    void sleep(long millis) throws InterruptedException;
+  }
+
+  private static volatile TimeService timeService = Thread::sleep;
+
+  @VisibleForTesting
+  static void setTimeServiceForTesting(@Nullable TimeService service) {
+    timeService = service != null ? service : Thread::sleep;
+  }
+
   private final EnvironmentProvider envProvider;
   private final PropertyProvider propProvider;
   private final @Nullable String certConfigPathOverride;
@@ -193,22 +216,48 @@ public KeyStore getKeyStore() throws CertificateSourceUnavailableException, IOEx
       throw configUnavailable;
     }
 
-    // Read the certificate and private key file paths into streams.
-    try (InputStream certStream = new FileInputStream(new File(workloadCertConfig.getCertPath()));
-        InputStream privateKeyStream =
-            new FileInputStream(new File(workloadCertConfig.getPrivateKeyPath()));
-        SequenceInputStream certAndPrivateKeyStream =
-            new SequenceInputStream(certStream, privateKeyStream)) {
+    // Read the certificate and private key file paths into streams, retrying to handle transient
+    // mid-rotation mismatches or file replacements when separate cert_path and key_path files are
+    // updated non-atomically relative to each other.
+    Exception lastException = null;
+    for (int attempt = 0; attempt < CERT_KEY_MATCH_RETRIES; attempt++) {
+      try (InputStream certStream =
+              new FileInputStream(new File(workloadCertConfig.getCertPath()));
+          InputStream privateKeyStream =
+              new FileInputStream(new File(workloadCertConfig.getPrivateKeyPath()));
+          SequenceInputStream certAndPrivateKeyStream =
+              new SequenceInputStream(certStream, privateKeyStream)) {
 
-      // Build a key store using the combined stream.
-      return MtlsKeyStoreUtils.createMtlsKeyStore(certAndPrivateKeyStream);
-    } catch (CertificateSourceUnavailableException e) {
-      // Throw the CertificateSourceUnavailableException without wrapping.
-      throw e;
-    } catch (Exception e) {
-      // Wrap all other exception types to an IOException.
-      throw new IOException("X509Provider: Unexpected IOException:", e);
+        // Build a key store using the combined stream.
+        return MtlsKeyStoreUtils.createMtlsKeyStore(certAndPrivateKeyStream);
+      } catch (CertificateSourceUnavailableException e) {
+        // Throw the CertificateSourceUnavailableException without wrapping.
+        throw e;
+      } catch (Exception e) {
+        lastException = e;
+        if (attempt + 1 < CERT_KEY_MATCH_RETRIES) {
+          LOGGER.log(
+              Level.FINE,
+              "X509Provider: Failed to load or verify certificate and private key on attempt "
+                  + (attempt + 1)
+                  + ", retrying",
+              e);
+          try {
+            timeService.sleep(CERT_KEY_MATCH_RETRY_INTERVAL_MS);
+          } catch (InterruptedException interrupted) {
+            Thread.currentThread().interrupt();
+            InterruptedIOException ioException =
+                new InterruptedIOException(
+                    "X509Provider: Interrupted while waiting for certificate and private key"
+                        + " match.");
+            ioException.initCause(interrupted);
+            throw ioException;
+          }
+        }
+      }
     }
+    // Wrap all other exception types to an IOException.
+    throw new IOException("X509Provider: Unexpected IOException:", lastException);
   }
 
   /**

google-auth-library-java/oauth2_http/java/com/google/auth/oauth2/AgentIdentityCertificateValidationUtils.java

@@ -30,6 +30,7 @@
  */
 package com.google.auth.oauth2;
 
+import com.google.auth.mtls.MtlsUtils;
 import com.google.common.collect.ImmutableList;
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
@@ -38,14 +39,10 @@
 import java.nio.file.Files;
 import java.nio.file.Paths;
 import java.security.GeneralSecurityException;
-import java.security.KeyFactory;
 import java.security.PrivateKey;
-import java.security.PublicKey;
-import java.security.Signature;
 import java.security.cert.CertificateFactory;
 import java.security.cert.CertificateParsingException;
 import java.security.cert.X509Certificate;
-import java.security.spec.X509EncodedKeySpec;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
@@ -79,9 +76,6 @@ final class AgentIdentityCertificateValidationUtils {
 
   private static final String SPIFFE_SCHEME_PREFIX = "spiffe://";
 
-  private static final byte[] VERIFICATION_DATA =
-      "verification-data".getBytes(StandardCharsets.UTF_8);
-
   private AgentIdentityCertificateValidationUtils() {}
 
   /** Thrown for key algorithms that can never be verified, so callers fail without retrying. */
@@ -159,35 +153,16 @@ static PrivateKey readPrivateKey(final String keyPath, final @Nullable String al
    */
   static boolean verifyKeyPair(final X509Certificate cert, final PrivateKey privateKey)
       throws UnsupportedKeyAlgorithmException {
-    PublicKey publicKey = cert.getPublicKey();
-    String keyAlgorithm = publicKey.getAlgorithm();
-    String sigAlg;
-    if ("RSA".equals(keyAlgorithm)) {
-      sigAlg = "SHA256withRSA";
-    } else if ("EC".equals(keyAlgorithm) || "ECDSA".equals(keyAlgorithm)) {
-      sigAlg = "SHA256withECDSA";
-    } else {
+    String keyAlgorithm = cert.getPublicKey().getAlgorithm();
+    if (!"RSA".equals(keyAlgorithm)
+        && !"EC".equals(keyAlgorithm)
+        && !"ECDSA".equals(keyAlgorithm)) {
       throw new UnsupportedKeyAlgorithmException(keyAlgorithm);
     }
 
     try {
-      if ("ECDSA".equals(keyAlgorithm) && publicKey.getEncoded() != null) {
-        // SunEC rejects keys whose algorithm is "ECDSA" rather than "EC"; normalize the key.
-        publicKey =
-            KeyFactory.getInstance("EC")
-                .generatePublic(new X509EncodedKeySpec(publicKey.getEncoded()));
-      }
-
-      Signature signer = Signature.getInstance(sigAlg);
-      signer.initSign(privateKey);
-      signer.update(VERIFICATION_DATA);
-      byte[] signature = signer.sign();
-
-      Signature verifier = Signature.getInstance(sigAlg);
-      verifier.initVerify(publicKey);
-      verifier.update(VERIFICATION_DATA);
-
-      return verifier.verify(signature);
+      MtlsUtils.verifyKeyPair(cert, privateKey);
+      return true;
     } catch (Exception e) {
       LoggingUtils.log(
           LOGGER_PROVIDER,