KAFKA-21418: Cleanup DelegationTokenCache and CredentialCache

clients/src/main/java/org/apache/kafka/common/network/SaslChannelBuilder.java

@@ -44,7 +44,6 @@
 import org.apache.kafka.common.security.oauthbearer.internals.unsecured.OAuthBearerUnsecuredValidatorCallbackHandler;
 import org.apache.kafka.common.security.plain.internals.PlainSaslServer;
 import org.apache.kafka.common.security.plain.internals.PlainServerCallbackHandler;
-import org.apache.kafka.common.security.scram.ScramCredential;
 import org.apache.kafka.common.security.scram.internals.ScramMechanism;
 import org.apache.kafka.common.security.scram.internals.ScramServerCallbackHandler;
 import org.apache.kafka.common.security.ssl.SslFactory;
@@ -325,7 +324,7 @@ private void createServerCallbackHandlers(Map<String, ?> configs) {
             else if (mechanism.equals(PlainSaslServer.PLAIN_MECHANISM))
                 callbackHandler = new PlainServerCallbackHandler();
             else if (ScramMechanism.isScram(mechanism))
-                callbackHandler = new ScramServerCallbackHandler(credentialCache.cache(mechanism, ScramCredential.class), tokenCache);
+                callbackHandler = new ScramServerCallbackHandler(credentialCache.cache(mechanism), tokenCache);
             else if (mechanism.equals(OAuthBearerLoginModule.OAUTHBEARER_MECHANISM))
                 callbackHandler = new OAuthBearerUnsecuredValidatorCallbackHandler();
             else

clients/src/main/java/org/apache/kafka/common/security/authenticator/CredentialCache.java

@@ -16,53 +16,50 @@
  */
 package org.apache.kafka.common.security.authenticator;
 
+import org.apache.kafka.common.security.scram.ScramCredential;
+import org.apache.kafka.common.security.scram.internals.ScramMechanism;
+
+import java.util.Collection;
 import java.util.concurrent.ConcurrentHashMap;
 
 public class CredentialCache {
 
-    private final ConcurrentHashMap<String, Cache<?>> cacheMap = new ConcurrentHashMap<>();
+    private final ConcurrentHashMap<String, Cache> cacheMap = new ConcurrentHashMap<>();
 
-    public <C> Cache<C> createCache(String mechanism, Class<C> credentialClass) {
-        Cache<C> cache = new Cache<>(credentialClass);
-        @SuppressWarnings("unchecked")
-        Cache<C> oldCache = (Cache<C>) cacheMap.putIfAbsent(mechanism, cache);
-        return oldCache == null ? cache : oldCache;
+    public CredentialCache() {
     }
 
-    @SuppressWarnings("unchecked")
-    public <C> Cache<C> cache(String mechanism, Class<C> credentialClass) {
-        Cache<?> cache = cacheMap.get(mechanism);
-        if (cache != null) {
-            if (cache.credentialClass() != credentialClass)
-                throw new IllegalArgumentException("Invalid credential class " + credentialClass + ", expected " + cache.credentialClass());
-            return (Cache<C>) cache;
-        } else
-            return null;
+    /**
+     * Creates a cache for each SCRAM mechanism in {@code mechanisms}. Other mechanisms are ignored.
+     */
+    public CredentialCache(Collection<String> mechanisms) {
+        for (String mechanism : ScramMechanism.mechanismNames()) {
+            if (mechanisms.contains(mechanism))
+                cacheMap.put(mechanism, new Cache());
+        }
     }
 
-    public static class Cache<C> {
-        private final Class<C> credentialClass;
-        private final ConcurrentHashMap<String, C> credentials;
+    public Cache createCache(String mechanism) {
+        return cacheMap.computeIfAbsent(mechanism, ignored -> new Cache());
+    }
 
-        public Cache(Class<C> credentialClass) {
-            this.credentialClass = credentialClass;
-            this.credentials = new ConcurrentHashMap<>();
-        }
+    public Cache cache(String mechanism) {
+        return cacheMap.get(mechanism);
+    }
+
+    public static class Cache {
+        private final ConcurrentHashMap<String, ScramCredential> credentials = new ConcurrentHashMap<>();
 
-        public C get(String username) {
+        public ScramCredential get(String username) {
             return credentials.get(username);
         }
 
-        public C put(String username, C credential) {
+        public ScramCredential put(String username, ScramCredential credential) {
             return credentials.put(username, credential);
         }
 
-        public C remove(String username) {
+        public ScramCredential remove(String username) {
             return credentials.remove(username);
         }
-
-        public Class<C> credentialClass() {
-            return credentialClass;
-        }
     }
 }

clients/src/main/java/org/apache/kafka/common/security/scram/internals/ScramServerCallbackHandler.java

@@ -18,7 +18,6 @@
 
 import org.apache.kafka.common.security.auth.AuthenticateCallbackHandler;
 import org.apache.kafka.common.security.authenticator.CredentialCache;
-import org.apache.kafka.common.security.scram.ScramCredential;
 import org.apache.kafka.common.security.scram.ScramCredentialCallback;
 import org.apache.kafka.common.security.token.delegation.TokenInformation;
 import org.apache.kafka.common.security.token.delegation.internals.DelegationTokenCache;
@@ -34,11 +33,11 @@
 
 public class ScramServerCallbackHandler implements AuthenticateCallbackHandler {
 
-    private final CredentialCache.Cache<ScramCredential> credentialCache;
+    private final CredentialCache.Cache credentialCache;
     private final DelegationTokenCache tokenCache;
     private String saslMechanism;
 
-    public ScramServerCallbackHandler(CredentialCache.Cache<ScramCredential> credentialCache,
+    public ScramServerCallbackHandler(CredentialCache.Cache credentialCache,
                                       DelegationTokenCache tokenCache) {
         this.credentialCache = credentialCache;
         this.tokenCache = tokenCache;

clients/src/main/java/org/apache/kafka/common/security/token/delegation/internals/DelegationTokenCache.java

@@ -19,7 +19,6 @@
 
 import org.apache.kafka.common.security.authenticator.CredentialCache;
 import org.apache.kafka.common.security.scram.ScramCredential;
-import org.apache.kafka.common.security.scram.internals.ScramCredentialUtils;
 import org.apache.kafka.common.security.scram.internals.ScramMechanism;
 import org.apache.kafka.common.security.token.delegation.DelegationToken;
 import org.apache.kafka.common.security.token.delegation.TokenInformation;
@@ -28,90 +27,80 @@
 import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.stream.Collectors;
 
 public class DelegationTokenCache {
 
-    private final CredentialCache credentialCache = new CredentialCache();
+    private final CredentialCache credentialCache;
 
     //Cache to hold all the tokens
-    private final Map<String, TokenInformation> tokenCache = new ConcurrentHashMap<>();
+    private final Map<String, DelegationToken> tokenCache = new ConcurrentHashMap<>();
 
     //Cache to hold hmac->tokenId mapping. This is required for renew, expire requests
     private final Map<String, String> hmacTokenIdCache = new ConcurrentHashMap<>();
 
-    //Cache to hold tokenId->hmac mapping. This is required for removing entry from hmacTokenIdCache using tokenId.
-    private final Map<String, String> tokenIdHmacCache = new ConcurrentHashMap<>();
-
     public DelegationTokenCache(Collection<String> scramMechanisms) {
         //Create caches for scramMechanisms
-        ScramCredentialUtils.createCache(credentialCache, scramMechanisms);
+        this.credentialCache = new CredentialCache(scramMechanisms);
     }
 
     public ScramCredential credential(String mechanism, String tokenId) {
-        CredentialCache.Cache<ScramCredential> cache = credentialCache.cache(mechanism, ScramCredential.class);
+        CredentialCache.Cache cache = credentialCache.cache(mechanism);
         return cache == null ? null : cache.get(tokenId);
     }
 
     public String owner(String tokenId) {
-        TokenInformation tokenInfo = tokenCache.get(tokenId);
-        return tokenInfo == null ? null : tokenInfo.owner().getName();
+        DelegationToken token = tokenCache.get(tokenId);
+        return token == null ? null : token.tokenInfo().owner().getName();
     }
 
     public void updateCache(DelegationToken token, Map<String, ScramCredential> scramCredentialMap) {
-        //Update TokenCache
-        String tokenId =  token.tokenInfo().tokenId();
-        addToken(tokenId, token.tokenInfo());
-        String hmac = token.hmacAsBase64String();
+        //Update TokenCache and hmac-id cache
+        addToken(token);
         //Update Scram Credentials
-        updateCredentials(tokenId, scramCredentialMap);
-        //Update hmac-id cache
-        hmacTokenIdCache.put(hmac, tokenId);
-        tokenIdHmacCache.put(tokenId, hmac);
+        updateCredentials(token.tokenInfo().tokenId(), scramCredentialMap);
     }
 
     public void removeCache(String tokenId) {
         removeToken(tokenId);
         updateCredentials(tokenId, new HashMap<>());
     }
 
-    public String tokenIdForHmac(String base64hmac) {
-        return hmacTokenIdCache.get(base64hmac);
-    }
-
     public TokenInformation tokenForHmac(String base64hmac) {
         String tokenId = hmacTokenIdCache.get(base64hmac);
-        return tokenId == null ? null : tokenCache.get(tokenId);
+        DelegationToken token = tokenId == null ? null : tokenCache.get(tokenId);
+        return token == null ? null : token.tokenInfo();
     }
 
-    public TokenInformation addToken(String tokenId, TokenInformation tokenInfo) {
-        return tokenCache.put(tokenId, tokenInfo);
+    public void addToken(DelegationToken token) {
+        String tokenId = token.tokenInfo().tokenId();
+        tokenCache.put(tokenId, token);
+        hmacTokenIdCache.put(token.hmacAsBase64String(), tokenId);
     }
 
     public void removeToken(String tokenId) {
-        TokenInformation tokenInfo = tokenCache.remove(tokenId);
-        if (tokenInfo != null) {
-            String hmac = tokenIdHmacCache.remove(tokenInfo.tokenId());
-            if (hmac != null) {
-                hmacTokenIdCache.remove(hmac);
-            }
+        DelegationToken token = tokenCache.remove(tokenId);
+        if (token != null) {
+            hmacTokenIdCache.remove(token.hmacAsBase64String());
         }
     }
 
     public Collection<TokenInformation> tokens() {
-        return tokenCache.values();
+        return tokenCache.values().stream().map(DelegationToken::tokenInfo).collect(Collectors.toList());
     }
 
     public TokenInformation token(String tokenId) {
-        return tokenCache.get(tokenId);
+        DelegationToken token = tokenCache.get(tokenId);
+        return token == null ? null : token.tokenInfo();
     }
 
-    public CredentialCache.Cache<ScramCredential> credentialCache(String mechanism) {
-        return credentialCache.cache(mechanism, ScramCredential.class);
+    public CredentialCache.Cache credentialCache(String mechanism) {
+        return credentialCache.cache(mechanism);
     }
 
     private void updateCredentials(String tokenId, Map<String, ScramCredential> scramCredentialMap) {
         for (String mechanism : ScramMechanism.mechanismNames()) {
-            CredentialCache.Cache<ScramCredential> cache = credentialCache.cache(mechanism, ScramCredential.class);
+            CredentialCache.Cache cache = credentialCache.cache(mechanism);
             if (cache != null) {
                 ScramCredential credential = scramCredentialMap.get(mechanism);
                 if (credential == null) {
@@ -122,4 +111,4 @@ private void updateCredentials(String tokenId, Map<String, ScramCredential> scra
             }
         }
     }
-}
+}

clients/src/test/java/org/apache/kafka/common/security/scram/internals/ScramCredentialUtilsTest.java → clients/src/test/java/org/apache/kafka/common/security/authenticator/CredentialCacheTest.java

@@ -14,34 +14,49 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.apache.kafka.common.security.scram.internals;
+package org.apache.kafka.common.security.authenticator;
 
-import org.apache.kafka.common.security.authenticator.CredentialCache;
 import org.apache.kafka.common.security.scram.ScramCredential;
+import org.apache.kafka.common.security.scram.internals.ScramFormatter;
+import org.apache.kafka.common.security.scram.internals.ScramMechanism;
 
 import org.junit.jupiter.api.Test;
 
-import java.util.Arrays;
+import java.util.List;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertSame;
 
-
-public class ScramCredentialUtilsTest {
+public class CredentialCacheTest {
 
     @Test
-    public void scramCredentialCache() throws Exception {
-        CredentialCache cache = new CredentialCache();
-        ScramCredentialUtils.createCache(cache, Arrays.asList("SCRAM-SHA-512", "PLAIN"));
-        assertNotNull(cache.cache(ScramMechanism.SCRAM_SHA_512.mechanismName(), ScramCredential.class), "Cache not created for enabled mechanism");
-        assertNull(cache.cache(ScramMechanism.SCRAM_SHA_256.mechanismName(), ScramCredential.class), "Cache created for disabled mechanism");
+    public void testScramCredentialCache() throws Exception {
+        CredentialCache cache = new CredentialCache(List.of("SCRAM-SHA-512", "PLAIN"));
+        assertNotNull(cache.cache(ScramMechanism.SCRAM_SHA_512.mechanismName()), "Cache not created for enabled mechanism");
+        assertNull(cache.cache(ScramMechanism.SCRAM_SHA_256.mechanismName()), "Cache created for disabled mechanism");
+        assertNull(cache.cache("PLAIN"), "Cache created for non-SCRAM mechanism");
 
-        CredentialCache.Cache<ScramCredential> sha512Cache = cache.cache(ScramMechanism.SCRAM_SHA_512.mechanismName(), ScramCredential.class);
+        CredentialCache.Cache sha512Cache = cache.cache(ScramMechanism.SCRAM_SHA_512.mechanismName());
         ScramFormatter formatter = new ScramFormatter(ScramMechanism.SCRAM_SHA_512);
         ScramCredential credentialA = formatter.generateCredential("password", 4096);
         sha512Cache.put("userA", credentialA);
         assertEquals(credentialA, sha512Cache.get("userA"));
         assertNull(sha512Cache.get("userB"), "Invalid user credential");
+
+        assertEquals(credentialA, sha512Cache.remove("userA"));
+        assertNull(sha512Cache.get("userA"));
+    }
+
+    @Test
+    public void testCreateCacheReturnsExistingCache() {
+        CredentialCache cache = new CredentialCache();
+        String mechanism = ScramMechanism.SCRAM_SHA_256.mechanismName();
+        assertNull(cache.cache(mechanism));
+
+        CredentialCache.Cache sha256Cache = cache.createCache(mechanism);
+        assertSame(sha256Cache, cache.cache(mechanism));
+        assertSame(sha256Cache, cache.createCache(mechanism));
     }
 }