fix: address PR review findings across mTLS rotation, retries, and auth caching

Author: macastelaz

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

@@ -52,6 +52,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.regex.Pattern;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -234,7 +235,78 @@ static CertInfo getAgentIdentityCertInfo() throws IOException {
       return null;
     }
 
-    return loadAndVerifyCredentials(paths.getCertPath(), paths.getKeyPath());
+    return getCachedOrLoadCredentials(paths.getCertPath(), paths.getKeyPath());
+  }
+
+  private static final Object certInfoCacheLock = new Object();
+  private static volatile CachedCertInfo cachedCertInfo = null;
+
+  static void clearCertInfoCache() {
+    synchronized (certInfoCacheLock) {
+      cachedCertInfo = null;
+    }
+  }
+
+  private static CertInfo getCachedOrLoadCredentials(String certPath, String keyPath)
+      throws IOException {
+    CachedCertInfo currentCache = cachedCertInfo;
+    if (currentCache != null && isCacheValid(currentCache, certPath, keyPath)) {
+      return currentCache.certInfo;
+    }
+    synchronized (certInfoCacheLock) {
+      if (cachedCertInfo != null && isCacheValid(cachedCertInfo, certPath, keyPath)) {
+        return cachedCertInfo.certInfo;
+      }
+      CertInfo info = loadAndVerifyCredentials(certPath, keyPath);
+      long certMtime = getFileMtime(certPath);
+      long keyMtime = getFileMtime(keyPath);
+      cachedCertInfo = new CachedCertInfo(certPath, keyPath, certMtime, keyMtime, info);
+      return info;
+    }
+  }
+
+  private static boolean isCacheValid(CachedCertInfo cache, String certPath, String keyPath) {
+    if (!Objects.equals(cache.certPath, certPath) || !Objects.equals(cache.keyPath, keyPath)) {
+      return false;
+    }
+    return cache.certLastModifiedTime == getFileMtime(certPath)
+        && cache.keyLastModifiedTime == getFileMtime(keyPath);
+  }
+
+  private static long getFileMtime(String path) {
+    if (path == null) {
+      return -1;
+    }
+    try {
+      java.nio.file.Path p = Paths.get(path);
+      if (Files.exists(p)) {
+        return Files.getLastModifiedTime(p).toMillis();
+      }
+    } catch (IOException e) {
+      // Ignore exception and return -1
+    }
+    return -1;
+  }
+
+  private static class CachedCertInfo {
+    final String certPath;
+    final String keyPath;
+    final long certLastModifiedTime;
+    final long keyLastModifiedTime;
+    final CertInfo certInfo;
+
+    CachedCertInfo(
+        String certPath,
+        String keyPath,
+        long certLastModifiedTime,
+        long keyLastModifiedTime,
+        CertInfo certInfo) {
+      this.certPath = certPath;
+      this.keyPath = keyPath;
+      this.certLastModifiedTime = certLastModifiedTime;
+      this.keyLastModifiedTime = keyLastModifiedTime;
+      this.certInfo = certInfo;
+    }
   }
 
   /**

google-auth-library-java/oauth2_http/javatests/com/google/auth/oauth2/AgentIdentityUtilsTest.java

@@ -34,6 +34,7 @@
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertSame;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.Mockito.mock;
@@ -75,13 +76,15 @@ class AgentIdentityUtilsTest {
 
   @BeforeEach
   void setUp() throws IOException {
+    AgentIdentityUtils.clearCertInfoCache();
     envProvider = new TestEnvironmentProvider();
     AgentIdentityUtils.setEnvReader(envProvider::getEnv);
     tempDir = Files.createTempDirectory("agent_identity_test");
   }
 
   @AfterEach
   void tearDown() throws IOException {
+    AgentIdentityUtils.clearCertInfoCache();
     AgentIdentityUtils.resetTimeService();
     AgentIdentityUtils.setWellKnownDir("/var/run/secrets/workload-spiffe-credentials/");
     AgentIdentityUtils.setEnvReader(System::getenv);
@@ -186,6 +189,32 @@ public void getAgentIdentityCertificate_happyPath_loadsCertificate() throws IOEx
     assertTrue(info.getCertificate().getIssuerDN().getName().contains("unit-tests"));
   }
 
+  @Test
+  public void testAgentIdentityCertInfoIsCachedAndReloadedWhenModified() throws IOException {
+    URL certUrl = getClass().getClassLoader().getResource("x509_leaf_certificate.pem");
+    assertNotNull(certUrl, "Test resource x509_leaf_certificate.pem not found");
+    String certPath = new File(certUrl.getFile()).getAbsolutePath();
+    File configFile = tempDir.resolve("config_cache.json").toFile();
+    String configJson =
+        "{"
+            + " \"cert_configs\": {"
+            + " \"workload\": {"
+            + " \"cert_path\": \""
+            + certPath.replace("\\", "\\\\")
+            + "\""
+            + " }"
+            + " }"
+            + "}";
+    try (FileOutputStream fos = new FileOutputStream(configFile)) {
+      fos.write(configJson.getBytes(StandardCharsets.UTF_8));
+    }
+    envProvider.setEnv("GOOGLE_API_CERTIFICATE_CONFIG", configFile.getAbsolutePath());
+    AgentIdentityUtils.CertInfo info1 = AgentIdentityUtils.getAgentIdentityCertInfo();
+    AgentIdentityUtils.CertInfo info2 = AgentIdentityUtils.getAgentIdentityCertInfo();
+    assertNotNull(info1);
+    assertSame(info1, info2);
+  }
+
   @Test
   public void getAgentIdentityCertificate_timeout_throwsIOException() {
     envProvider.setEnv(

sdk-platform-java/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/ChannelPool.java

@@ -712,7 +712,13 @@ public <RequestT, ResponseT> ClientCall<RequestT, ResponseT> newCall(
     }
   }
 
-  /** ClientCall wrapper that makes sure to decrement the outstanding RPC count on completion. */
+  /**
+   * ClientCall wrapper that makes sure to decrement the outstanding RPC count on completion.
+   *
+   * <p>Contract: Exactly one call to {@link #start(Listener, Metadata)} or explicit release via
+   * {@link #cancel(String, Throwable)} is required to balance reference counts. Early cancellation
+   * before {@code start()} safely decrements the reference count via atomic compare-and-set.
+   */
   static class ReleasingClientCall<ReqT, RespT> extends SimpleForwardingClientCall<ReqT, RespT> {
     private @Nullable CancellationException cancellationException;
     final Entry entry;

sdk-platform-java/gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/RefreshingHttpJsonChannel.java

@@ -153,7 +153,6 @@ public void refresh() {
         return;
       }
 
-      this.activeCertFingerprint = currentDiskFingerprint;
       LOG.info("mTLS certificate rotation detected. Triggering HTTP/JSON channel pool refresh.");
 
       // Prune terminated entries to prevent memory leak
@@ -162,6 +161,7 @@ public void refresh() {
       ChannelEntry newEntry = new ChannelEntry(channelFactory.get());
       allEntries.add(newEntry);
       ChannelEntry oldEntry = activeEntry.getAndSet(newEntry);
+      this.activeCertFingerprint = currentDiskFingerprint;
 
       if (oldEntry != null) {
         oldEntry.requestShutdown();
@@ -234,15 +234,22 @@ public boolean isTerminated() {
 
   @Override
   public void shutdownNow() {
-    for (ChannelEntry entry : allEntries) {
-      entry.channel.shutdownNow();
+    synchronized (refreshLock) {
+      isShuttingDown = true;
+      for (ChannelEntry entry : allEntries) {
+        entry.requestShutdown();
+        entry.channel.shutdownNow();
+      }
     }
   }
 
   @Override
   public boolean awaitTermination(long duration, TimeUnit unit) throws InterruptedException {
     long endNanos = System.nanoTime() + unit.toNanos(duration);
     for (ChannelEntry entry : allEntries) {
+      if (entry.channel.isTerminated()) {
+        continue;
+      }
       long remainingNanos = endNanos - System.nanoTime();
       if (remainingNanos <= 0) {
         return false;

sdk-platform-java/gax-java/gax-httpjson/src/test/java/com/google/api/gax/httpjson/RefreshingHttpJsonChannelTest.java

@@ -31,15 +31,21 @@
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Supplier;
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.mockito.ArgumentCaptor;
@@ -49,9 +55,14 @@ class RefreshingHttpJsonChannelTest {
   private ManagedHttpJsonChannel lastCreatedChannel;
   private String testCertPath = "/fake/path";
   private String testFingerprint = "fingerprint1";
+  private boolean shouldThrowOnFactory = false;
+  private List<RefreshingHttpJsonChannel> createdChannels;
 
   private Supplier<ManagedHttpJsonChannel> channelFactory =
       () -> {
+        if (shouldThrowOnFactory) {
+          throw new RuntimeException("Simulated factory failure");
+        }
         channelFactoryCount.incrementAndGet();
         lastCreatedChannel = mock(ManagedHttpJsonChannel.class);
         return lastCreatedChannel;
@@ -62,20 +73,32 @@ void setUp() {
     channelFactoryCount = new AtomicInteger(0);
     testCertPath = "/fake/path";
     testFingerprint = "fingerprint1";
+    shouldThrowOnFactory = false;
+    createdChannels = new ArrayList<>();
+  }
+
+  @AfterEach
+  void tearDown() {
+    for (RefreshingHttpJsonChannel channel : createdChannels) {
+      channel.shutdownNow();
+    }
   }
 
   private RefreshingHttpJsonChannel createTestChannel() {
-    return new RefreshingHttpJsonChannel(channelFactory, "fake/cert/path.json") {
-      @Override
-      protected String getWorkloadCertPath() {
-        return testCertPath;
-      }
-
-      @Override
-      protected String getCertificateFingerprint(String certPath) {
-        return testFingerprint;
-      }
-    };
+    RefreshingHttpJsonChannel ch =
+        new RefreshingHttpJsonChannel(channelFactory, "fake/cert/path.json") {
+          @Override
+          protected String getWorkloadCertPath() {
+            return testCertPath;
+          }
+
+          @Override
+          protected String getCertificateFingerprint(String certPath) {
+            return testFingerprint;
+          }
+        };
+    createdChannels.add(ch);
+    return ch;
   }
 
   @Test
@@ -194,4 +217,46 @@ void testRefreshDoesNotSpawnChannelWhenShutdown() throws InterruptedException {
     // Verify no new channel was spawned
     assertEquals(1, channelFactoryCount.get());
   }
+
+  @Test
+  void testRefreshFactoryExceptionDoesNotWedgeFingerprint() throws InterruptedException {
+    RefreshingHttpJsonChannel channel = createTestChannel();
+    assertEquals(1, channelFactoryCount.get());
+
+    shouldThrowOnFactory = true;
+    Thread.sleep(1001); // Invalidate 1-second cache
+    testFingerprint = "fingerprint2";
+
+    assertThrows(RuntimeException.class, channel::refresh);
+
+    // Because factory threw, activeCertFingerprint should NOT be updated to fingerprint2
+    // Therefore shouldRefresh() should still return true
+    assertTrue(channel.shouldRefresh());
+
+    shouldThrowOnFactory = false;
+    channel.refresh();
+    assertEquals(2, channelFactoryCount.get());
+    assertFalse(channel.shouldRefresh());
+  }
+
+  @Test
+  void testShutdownNowSetsIsShutdown() {
+    RefreshingHttpJsonChannel channel = createTestChannel();
+    assertFalse(channel.isShutdown());
+
+    channel.shutdownNow();
+
+    assertTrue(channel.isShutdown());
+  }
+
+  @Test
+  void testAwaitTerminationZeroTimeoutOnTerminatedChannelReturnsTrue() throws InterruptedException {
+    RefreshingHttpJsonChannel channel = createTestChannel();
+    ManagedHttpJsonChannel firstChannel = lastCreatedChannel;
+    when(firstChannel.isTerminated()).thenReturn(true);
+    when(firstChannel.awaitTermination(anyLong(), any())).thenReturn(true);
+
+    channel.shutdown();
+    assertTrue(channel.awaitTermination(0, TimeUnit.MILLISECONDS));
+  }
 }

sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ApiCallContext.java

@@ -66,6 +66,10 @@ public interface ApiCallContext extends RetryingContext {
   /**
    * Returns the {@link TransportChannel} associated with this call context, or {@code null} if none
    * is set.
+   *
+   * <p>Note: By default, this method returns {@code null}. If an implementation does not override
+   * this method, automatic mTLS certificate rotation and channel refreshing in retrying callables
+   * will be disabled.
    */
   default TransportChannel getTransportChannel() {
     return null;

sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ApiResultRetryAlgorithm.java

@@ -52,6 +52,8 @@ public boolean shouldRetry(Throwable previousThrowable, ResponseT previousRespon
   @Override
   public boolean shouldRetry(
       RetryingContext context, Throwable previousThrowable, ResponseT previousResponse) {
+    // Check UnauthenticatedException retryability first to ensure mTLS certificate
+    // rotation retries take precedence over static method retry codes.
     if (previousThrowable instanceof UnauthenticatedException
         && ((UnauthenticatedException) previousThrowable).isRetryable()) {
       return true;

sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/AttemptCallable.java

@@ -95,12 +95,18 @@ public ResponseT call() {
                 TransportChannel transportChannel = finalContext.getTransportChannel();
                 if (transportChannel != null && transportChannel.shouldRefresh()) {
                   transportChannel.refresh();
-                  throw new UnauthenticatedException(
-                      unauthenticatedException.getMessage(),
-                      unauthenticatedException.getCause(),
-                      unauthenticatedException.getStatusCode(),
-                      true, // isRetryable = true
-                      unauthenticatedException.getErrorDetails());
+                  UnauthenticatedException newEx =
+                      new UnauthenticatedException(
+                          unauthenticatedException.getMessage(),
+                          unauthenticatedException,
+                          unauthenticatedException.getStatusCode(),
+                          true, // isRetryable = true
+                          unauthenticatedException.getErrorDetails());
+                  newEx.setStackTrace(unauthenticatedException.getStackTrace());
+                  for (Throwable suppressed : unauthenticatedException.getSuppressed()) {
+                    newEx.addSuppressed(suppressed);
+                  }
+                  throw newEx;
                 }
                 throw unauthenticatedException;
               },