SOLR-16951: Add PKI Auth Caching for both generation and validation (#3334)

Author: HoustonPutman

solr/CHANGES.txt

@@ -245,6 +245,9 @@ Improvements
 
 * SOLR-17750: S3 File downloads now handle connection issues more gracefully (Houston Putman, Mark Miller)
 
+* SOLR-16951: Add PKI Auth Caching for both generation of the PKI Auth Tokens and validation of received tokens.
+  The default PKI Auth validity TTL has been increased from 5 seconds to 10 seconds. (Houston Putman)
+
 Optimizations
 ---------------------
 * SOLR-17578: Remove ZkController internal core supplier, for slightly faster reconnection after Zookeeper session loss. (Pierre Salagnac)

solr/core/src/java/org/apache/solr/security/PKIAuthenticationPlugin.java

@@ -19,6 +19,8 @@
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static org.apache.solr.client.solrj.SolrRequest.METHOD.GET;
 
+import com.github.benmanes.caffeine.cache.Caffeine;
+import com.github.benmanes.caffeine.cache.LoadingCache;
 import com.google.common.annotations.VisibleForTesting;
 import jakarta.servlet.FilterChain;
 import jakarta.servlet.http.HttpServletRequest;
@@ -37,6 +39,7 @@
 import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.TimeUnit;
 import org.apache.http.HttpEntity;
 import org.apache.http.HttpException;
 import org.apache.http.HttpHeaders;
@@ -94,7 +97,10 @@ public static void withServerIdentity(final boolean enabled) {
   private final Map<String, PublicKey> keyCache = new ConcurrentHashMap<>();
   private final PublicKeyHandler publicKeyHandler;
   private final CoreContainer cores;
-  private static final int MAX_VALIDITY = Integer.getInteger("pkiauth.ttl", 5000);
+  private final LoadingCache<String, PKIHeaderData> validatedHeaderCache;
+  private final LoadingCache<String, String> generatedV1TokenCache;
+  private final LoadingCache<String, String> generatedV2TokenCache;
+  private static final int MAX_VALIDITY = Integer.getInteger("pkiauth.ttl", 10000);
   private final String myNodeName;
   private final HttpHeaderClientInterceptor interceptor = new HttpHeaderClientInterceptor();
   private boolean interceptorRegistered = false;
@@ -114,6 +120,38 @@ public PKIAuthenticationPlugin(
     this.cores = cores;
     myNodeName = nodeName;
 
+    // Don't expire after read, because there is no reason to add the overhead of updating expiry
+    // information after each read. The expiration time here doesn't matter too much, because we
+    // still check the PKI Token TTL after fetching from the cache. We just want to make sure cache
+    // entries are cleaned up regularly.
+    validatedHeaderCache =
+        Caffeine.newBuilder()
+            .maximumSize(1000)
+            .expireAfterWrite(MAX_VALIDITY, TimeUnit.MILLISECONDS)
+            .build(this::decipherHeaderV2);
+    // We must expire much earlier than the max validity, because these cached Auth tokens still
+    // need to be sent to the server, which will validate the TTL. If we expire at maxValidity, the
+    // TTL check will always fail before the cache entry is expired.
+    long expireAfterTime = MAX_VALIDITY / 4;
+    // Refreshing is done asynchronously, so we want to do it before expiration. This means that
+    // requests are not synchronously blocked when generating new Auth tokens. However, the refresh
+    // will only happen when the cached header is requested. Therefore, we want to give a long-ish
+    // runway for requests to come in to trigger an asynchronous-refresh before expiry causes a
+    // synchronous-refresh.
+    long shouldRefreshTime = Math.max(1, expireAfterTime / 2);
+    generatedV1TokenCache =
+        Caffeine.newBuilder()
+            .maximumSize(100)
+            .refreshAfterWrite(shouldRefreshTime, TimeUnit.MILLISECONDS)
+            .expireAfterWrite(expireAfterTime, TimeUnit.MILLISECONDS)
+            .build(this::generateToken);
+    generatedV2TokenCache =
+        Caffeine.newBuilder()
+            .maximumSize(100)
+            .refreshAfterWrite(shouldRefreshTime, TimeUnit.MILLISECONDS)
+            .expireAfterWrite(expireAfterTime, TimeUnit.MILLISECONDS)
+            .build(this::generateTokenV2);
+
     Set<String> knownPkiVersions = Set.of("v1", "v2");
     // We always accept v2 even if it is not specified
     String[] versions = System.getProperty(ACCEPT_VERSIONS, "v2").split(",");
@@ -155,7 +193,7 @@ public boolean doAuthenticate(
         return sendError(response, true, "Could not parse node name from SolrAuthV2 header.");
       }
 
-      headerData = decipherHeaderV2(headerV2, headerV2.substring(0, nodeNameEnd));
+      headerData = validatedHeaderCache.get(headerV2);
     } else if (headerV1 != null && acceptPkiV1) {
       List<String> authInfo = StrUtils.splitWS(headerV1, false);
       if (authInfo.size() != 2) {
@@ -223,7 +261,8 @@ private PublicKey getOrFetchPublicKey(String nodeName) {
     return key;
   }
 
-  private PKIHeaderData decipherHeaderV2(String header, String nodeName) {
+  private PKIHeaderData decipherHeaderV2(String header) {
+    String nodeName = header.substring(0, header.indexOf(' '));
     PublicKey key = getOrFetchPublicKey(nodeName);
 
     int sigStart = header.lastIndexOf(' ');
@@ -414,11 +453,11 @@ public void onBegin(Request request) {
             final Optional<String> preFetchedUser = getUserFromJettyRequest(request);
             if ("v1".equals(System.getProperty(SEND_VERSION))) {
               preFetchedUser
-                  .map(PKIAuthenticationPlugin.this::generateToken)
+                  .map(generatedV1TokenCache::get)
                   .ifPresent(token -> request.headers(httpFields -> httpFields.add(HEADER, token)));
             } else {
               preFetchedUser
-                  .map(PKIAuthenticationPlugin.this::generateTokenV2)
+                  .map(generatedV2TokenCache::get)
                   .ifPresent(
                       token -> request.headers(httpFields -> httpFields.add(HEADER_V2, token)));
             }
@@ -520,10 +559,12 @@ private String generateTokenV2(String user) {
 
   void setHeader(HttpRequest httpRequest) {
     if ("v1".equals(System.getProperty(SEND_VERSION))) {
-      getUser().map(this::generateToken).ifPresent(token -> httpRequest.setHeader(HEADER, token));
+      getUser()
+          .map(generatedV1TokenCache::get)
+          .ifPresent(token -> httpRequest.setHeader(HEADER, token));
     } else {
       getUser()
-          .map(this::generateTokenV2)
+          .map(generatedV2TokenCache::get)
           .ifPresent(token -> httpRequest.setHeader(HEADER_V2, token));
     }
   }