maintain the previous key after rotation

Author: zhangweikop

util/src/main/java/io/grpc/util/AdvancedTlsX509KeyManager.java

@@ -42,10 +42,10 @@
  * AdvancedTlsX509KeyManager is an {@code X509ExtendedKeyManager} that allows users to configure
  * advanced TLS features, such as private key and certificate chain reloading.
  *
- * <p>The key material alias increments on every credential load (e.g. {@code "key-1"},
- * {@code "key-2"}, ...), ensuring the same alias always maps to the same key material. This is
- * required by Netty's {@code OpenSslCachingX509KeyManagerFactory} to correctly cache key
- * material and create a new cache entry on cert reload.
+ * <p>The alias increments on every credential load (e.g. {@code "key-1"}, {@code "key-2"}, ...),
+ * so the same alias always maps to the same key material. The previous alias is retained for one
+ * rotation to allow in-progress handshakes to complete. This is required by Netty's
+ * {@code OpenSslCachingX509KeyManagerFactory} to cache key material across cert reloads.
  *
  * <p>When using {@code SslProvider.OPENSSL}, wrap this key manager in Netty's
  * {@code OpenSslCachingX509KeyManagerFactory} to avoid per-handshake key material encoding
@@ -61,44 +61,39 @@ public final class AdvancedTlsX509KeyManager extends X509ExtendedKeyManager {
   static final String ALIAS_PREFIX = "key-";
 
   private final AtomicInteger revision = new AtomicInteger(0);
-  private final int revisionWarningThreshold;
-  // The credential information to be sent to peers to prove our identity.
-  private volatile KeyInfo keyInfo;
+  // Snapshot of current and previous KeyInfo; previous is retained for in-progress handshakes
+  // after one rotation.
+  private volatile KeyInfoSnapshot snapshot = new KeyInfoSnapshot(null, null);
 
-  public AdvancedTlsX509KeyManager() {
-    // Netty's default OpenSslCachingX509KeyManagerFactory maxCachedEntries.
-    this(1024);
-  }
-
-  /**
-   * Creates a key manager with a custom revision warning threshold.
-   * @param revisionWarningThreshold the number of credential loads after which a warning is logged.
-   *     Only relevant when using {@code SslProvider.OPENSSL} with
-   *     {@code OpenSslCachingX509KeyManagerFactory}.
-   */
-  public AdvancedTlsX509KeyManager(int revisionWarningThreshold) {
-    this.revisionWarningThreshold = revisionWarningThreshold;
-  }
+  public AdvancedTlsX509KeyManager() {}
 
   private String alias() {
-    KeyInfo info = this.keyInfo;
-    if (info == null) {
-      return null;
-    }
-    return info.alias;
+    KeyInfo curr = this.snapshot.current;
+    return curr != null ? curr.alias : null;
   }
 
   @Override
   public PrivateKey getPrivateKey(String alias) {
-    KeyInfo info = this.keyInfo;
-    return info != null && info.alias.equals(alias) ? info.key : null;
+    KeyInfoSnapshot snap = this.snapshot;
+    if (snap.current != null && snap.current.alias.equals(alias)) {
+      return snap.current.key;
+    }
+    if (snap.previous != null && snap.previous.alias.equals(alias)) {
+      return snap.previous.key;
+    }
+    return null;
   }
 
   @Override
   public X509Certificate[] getCertificateChain(String alias) {
-    KeyInfo info = this.keyInfo;
-    return info != null && info.alias.equals(alias)
-        ? Arrays.copyOf(info.certs, info.certs.length) : null;
+    KeyInfoSnapshot snap = this.snapshot;
+    if (snap.current != null && snap.current.alias.equals(alias)) {
+      return Arrays.copyOf(snap.current.certs, snap.current.certs.length);
+    }
+    if (snap.previous != null && snap.previous.alias.equals(alias)) {
+      return Arrays.copyOf(snap.previous.certs, snap.previous.certs.length);
+    }
+    return null;
   }
 
   @Override
@@ -155,15 +150,9 @@ public void updateIdentityCredentials(PrivateKey key, X509Certificate[] certs) {
    * @param key  the private key that is going to be used
    */
   public void updateIdentityCredentials(X509Certificate[] certs, PrivateKey key) {
-    // When using SslProvider.OPENSSL with OpenSslCachingX509KeyManagerFactory, its cache stops
-    // accepting new aliases once maxCachedEntries is reached (default: 1024). Beyond this,
-    // handshakes still succeed but per-handshake re-encoding overhead resumes.
-    if (revision.get() >= revisionWarningThreshold) {
-      log.warning("AdvancedTlsX509KeyManager: revision counter has reached "
-          + revisionWarningThreshold);
-    }
-    this.keyInfo = new KeyInfo(checkNotNull(certs, "certs"), checkNotNull(key, "key"),
+    KeyInfo newInfo = new KeyInfo(checkNotNull(certs, "certs"), checkNotNull(key, "key"),
         ALIAS_PREFIX + revision.incrementAndGet());
+    this.snapshot = new KeyInfoSnapshot(newInfo, this.snapshot.current);
   }
 
   /**
@@ -262,7 +251,6 @@ public Closeable updateIdentityCredentialsFromFile(File keyFile, File certFile,
   }
 
   private static class KeyInfo {
-    // The private key and the cert chain we will use to send to peers to prove our identity.
     final X509Certificate[] certs;
     final PrivateKey key;
     final String alias;
@@ -274,6 +262,16 @@ public KeyInfo(X509Certificate[] certs, PrivateKey key, String alias) {
     }
   }
 
+  private static class KeyInfoSnapshot {
+    final KeyInfo current;
+    final KeyInfo previous;
+
+    KeyInfoSnapshot(KeyInfo current, KeyInfo previous) {
+      this.current = current;
+      this.previous = previous;
+    }
+  }
+
   private class LoadFilePathExecution implements Runnable {
     File keyFile;
     File certFile;

util/src/test/java/io/grpc/util/AdvancedTlsX509KeyManagerTest.java

@@ -18,7 +18,6 @@
 
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
@@ -92,14 +91,17 @@ public void updateTrustCredentials_replacesIssuers() throws Exception {
     assertEquals(AdvancedTlsX509KeyManager.ALIAS_PREFIX + "2", alias2);
     assertEquals(clientKey0, serverKeyManager.getPrivateKey(alias2));
     assertArrayEquals(clientCert0, serverKeyManager.getCertificateChain(alias2));
-    // Old alias no longer resolves — ensures alias stability contract is enforced.
-    assertNull(serverKeyManager.getPrivateKey(alias1));
+    // Previous alias still resolves — retained to allow in-progress handshakes to complete.
+    assertEquals(serverKey0, serverKeyManager.getPrivateKey(alias1));
+    assertArrayEquals(serverCert0, serverKeyManager.getCertificateChain(alias1));
 
     serverKeyManager.updateIdentityCredentials(serverCert0File, serverKey0File, 1,
         TimeUnit.MINUTES, executor);
     String alias3 = serverKeyManager.chooseEngineServerAlias(null, null, null);
     assertEquals(serverKey0, serverKeyManager.getPrivateKey(alias3));
     assertArrayEquals(serverCert0, serverKeyManager.getCertificateChain(alias3));
+    // alias1 is now two rotations back — no longer retained.
+    assertNull(serverKeyManager.getPrivateKey(alias1));
 
     serverKeyManager.updateIdentityCredentials(serverCert0, serverKey0);
     String alias4 = serverKeyManager.chooseEngineServerAlias(null, null, null);
@@ -135,46 +137,6 @@ public void allAliasMethods_agreeAfterCredentialLoad() throws Exception {
     assertArrayEquals(new String[]{expectedAlias}, keyManager.getServerAliases(null, null));
   }
 
-  @Test
-  public void revisionWarningThreshold_logsWarningAtThreshold() throws Exception {
-    Logger log = Logger.getLogger(AdvancedTlsX509KeyManager.class.getName());
-    TestHandler handler = new TestHandler();
-    log.addHandler(handler);
-    log.setUseParentHandlers(false);
-    log.setLevel(Level.ALL);
-
-    try {
-      // Custom threshold: warning when revision reaches threshold.
-      int threshold = 3;
-      AdvancedTlsX509KeyManager customKeyManager = new AdvancedTlsX509KeyManager(threshold);
-      for (int i = 0; i < threshold; i++) {
-        customKeyManager.updateIdentityCredentials(serverCert0, serverKey0);
-      }
-      assertFalse(hasRevisionWarning(handler));
-      customKeyManager.updateIdentityCredentials(serverCert0, serverKey0);
-      assertTrue(hasRevisionWarning(handler));
-
-      // Key manager must still provide credentials correctly after soft threshold is exceeded.
-      String alias = customKeyManager.chooseEngineServerAlias(null, null, null);
-      assertEquals(serverKey0, customKeyManager.getPrivateKey(alias));
-      assertArrayEquals(serverCert0, customKeyManager.getCertificateChain(alias));
-
-      // Further credential updates must also work.
-      customKeyManager.updateIdentityCredentials(clientCert0File, clientKey0File);
-      String newAlias = customKeyManager.chooseEngineServerAlias(null, null, null);
-      assertEquals(clientKey0, customKeyManager.getPrivateKey(newAlias));
-      assertArrayEquals(clientCert0, customKeyManager.getCertificateChain(newAlias));
-    } finally {
-      log.removeHandler(handler);
-    }
-  }
-
-  private static boolean hasRevisionWarning(TestHandler handler) {
-    return handler.getRecords().stream()
-        .anyMatch(r -> Level.WARNING.equals(r.getLevel())
-            && r.getMessage().contains("revision counter has reached"));
-  }
-
   @Test
   public void credentialSettingParameterValidity() throws Exception {
     // Checking edge cases of public API parameter setting.