Improve synchronization in AuthenticationService

server/src/main/java/timely/auth/AuthCache.java

@@ -50,31 +50,31 @@ public class AuthCache {
     }
 
     public static boolean containsKey(String key) {
-        return getCache().asMap().containsKey(key);
+        return getCache().getIfPresent(key) != null;
     }
 
     public static TimelyPrincipal get(String key) {
-        return getCache().asMap().get(key);
+        return getCache().getIfPresent(key);
     }
 
     public static void put(String key, TimelyPrincipal principal) {
-        getCache().asMap().put(key, principal);
+        getCache().put(key, principal);
     }
 
-    public static TimelyPrincipal remove(String key) {
-        return getCache().asMap().remove(key);
+    public static void remove(String key) {
+        getCache().invalidate(key);
     }
 
     public static void clear() {
         if (CACHE != null) {
-            CACHE.asMap().clear();
+            CACHE.invalidateAll();
         }
     }
 
     protected static Collection<Authorizations> getAuthorizations(String entityName) {
         if (!StringUtils.isEmpty(entityName)) {
             List<Authorizations> authorizationsList = new ArrayList<>();
-            TimelyPrincipal principal = CACHE.asMap().get(entityName);
+            TimelyPrincipal principal = AuthCache.get(entityName);
             if (principal != null) {
                 for (Collection<String> authCollection : principal.getAuthorizations()) {
                     authorizationsList.add(AuthorizationsUtil.toAuthorizations(authCollection));

server/src/main/java/timely/auth/AuthenticationService.java

@@ -2,11 +2,12 @@ package timely.auth;
 
 import java.security.cert.X509Certificate;
 import java.util.ArrayList;
-import java.util.HashSet;
+import java.util.Arrays;
+import java.util.Collection;
 import java.util.List;
 import java.util.ServiceConfigurationError;
-import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.stream.Collectors;
 
 import com.google.common.collect.Multimap;
 import io.netty.channel.ChannelHandlerContext;
@@ -23,7 +24,7 @@ import org.springframework.security.core.Authentication;
 import org.springframework.security.web.authentication.preauth.x509.SubjectDnX509PrincipalExtractor;
 import timely.api.request.AuthenticatedRequest;
 import timely.api.response.TimelyException;
-import timely.auth.util.AuthenticationUtils;
+import timely.auth.util.DnUtils;
 import timely.netty.http.auth.TimelyAuthenticationToken;
 
 public class AuthenticationService {
@@ -32,8 +33,8 @@ public class AuthenticationService {
     private static ApplicationContext springContext = null;
     private static AuthenticationManager authManager = null;
     private static SubjectDnX509PrincipalExtractor x509 = null;
-    private static ConcurrentHashMap<String, Set<Long>> fetchingEntity = new ConcurrentHashMap<>();
-    private static ConcurrentHashMap<String, Set<Long>> fetchingPrincipal = new ConcurrentHashMap<>();
+    private static ConcurrentHashMap<String, Object> fetchingEntity = new ConcurrentHashMap<>();
+    private static ConcurrentHashMap<String, Object> fetchingPrincipal = new ConcurrentHashMap<>();
 
     static {
         try {
@@ -51,46 +52,36 @@ public class AuthenticationService {
     }
 
     public static TimelyPrincipal authenticate(Authentication authentication, String cacheName) {
-        return authenticate(authentication, null, cacheName);
+        return authenticate(authentication, SubjectIssuerDNPair.of(cacheName), cacheName);
     }
 
-    public static Authentication authenticate(Authentication authentication) {
-        return null;
-    }
-
-    public static TimelyPrincipal authenticate(Authentication authentication, X509Certificate clientCert,
-            String entity) {
+    public static TimelyPrincipal authenticate(Authentication authentication, SubjectIssuerDNPair pair, String entity) {
 
         // only one thread should call the authenticationManager for a given entity
-        long threadId = Thread.currentThread().getId();
         TimelyPrincipal principal = AuthCache.get(entity);
         // once an entity is in the AuthCache, we can bypass all of this
         if (principal == null) {
-            Set<Long> o = fetchingEntity.getOrDefault(entity, new HashSet<>());
-            // add thread id to the set so we can track when to get rid of the Set
-            o.add(threadId);
-            synchronized (o) {
-                try {
+            Object newObj = new Object();
+            Object o = fetchingEntity.getOrDefault(entity, newObj);
+            if (o == null) {
+                o = newObj;
+            }
+            try {
+                synchronized (o) {
                     // if multiple threads are waiting on this synchronized, then the 2nd and all
                     // others will get the principal here
                     principal = AuthCache.get(entity);
                     if (principal == null) {
                         Authentication token = getAuthenticationManager().authenticate(authentication);
-                        TimelyAuthenticationToken timelyAuthenticationToken = AuthenticationUtils
-                                .getTimelyAuthentication(token, clientCert);
-                        principal = timelyAuthenticationToken.getTimelyPrincipal();
+                        principal = getTimelyPrincipal(token, pair);
                         LOG.trace("Got entity principal for {} from AuthenticationManager", entity);
                         AuthCache.put(entity, principal);
                     } else {
                         LOG.trace("Got entity principal for {} from AuthCache on 2nd attempt", entity);
                     }
-                } finally {
-                    o.remove(threadId);
-                    // if last thread fetching, remove Set from fetchingEntity
-                    if (o.isEmpty()) {
-                        fetchingEntity.remove(entity);
-                    }
                 }
+            } finally {
+                fetchingEntity.remove(entity);
             }
         } else {
             LOG.trace("Got entity principal for {} from AuthCache on 1st attempt", entity);
@@ -113,21 +104,24 @@ public class AuthenticationService {
             try {
                 TimelyAuthenticationToken token = request.getToken();
                 TimelyPrincipal requestPrincipal = token.getTimelyPrincipal();
-                long threadId = Thread.currentThread().getId();
+                String entity = requestPrincipal.getName();
                 if (!AuthCache.containsKey(requestPrincipal.getName())) {
-                    Set<Long> o = fetchingPrincipal.getOrDefault(requestPrincipal.getName(), new HashSet<>());
-                    o.add(threadId);
-                    synchronized (o) {
-                        try {
+                    Object newObj = new Object();
+                    Object o = fetchingPrincipal.getOrDefault(entity, newObj);
+                    if (o == null) {
+                        o = newObj;
+                    }
+                    try {
+                        synchronized (o) {
                             // check again if in the AuthCache
-                            if (!AuthCache.containsKey(requestPrincipal.getName())) {
+                            if (!AuthCache.containsKey(entity)) {
                                 // first thread here, so authenticate all entities in requestPrincipal
                                 List<TimelyUser> authenticatedTimelyUsers = new ArrayList<>();
                                 for (TimelyUser user : requestPrincipal.getProxiedUsers()) {
                                     SubjectIssuerDNPair p = user.getDn();
                                     TimelyAuthenticationToken entityToken = getAuthenticationToken(clientCert,
                                             p.subjectDN(), p.issuerDN());
-                                    TimelyPrincipal entityPrincipal = authenticate(entityToken, clientCert,
+                                    TimelyPrincipal entityPrincipal = authenticate(entityToken, p,
                                             entityToken.getTimelyPrincipal().getName());
                                     authenticatedTimelyUsers.addAll(entityPrincipal.getProxiedUsers());
                                 }
@@ -136,18 +130,14 @@ public class AuthenticationService {
                                 LOG.debug("Authenticated user {} for request {} with authorizations {}",
                                         principal.getName(), request.toString(), principal.getAuthorizationsString());
                             } else {
-                                LOG.trace("Verified principal {} in AuthCache on 2nd attempt",
-                                        requestPrincipal.getName());
-                            }
-                        } finally {
-                            o.remove(threadId);
-                            if (o.isEmpty()) {
-                                fetchingPrincipal.remove(requestPrincipal.getName());
+                                LOG.trace("Verified principal {} in AuthCache on 2nd attempt", entity);
                             }
                         }
+                    } finally {
+                        fetchingPrincipal.remove(entity);
                     }
                 } else {
-                    LOG.trace("Verified principal {} in AuthCache on 1st attempt", requestPrincipal.getName());
+                    LOG.trace("Verified principal {} in AuthCache on 1st attempt", entity);
                 }
             } catch (Exception e) {
                 LOG.error(e.getMessage(), e);
@@ -195,4 +185,16 @@ public class AuthenticationService {
         }
         return authenticationToken;
     }
+
+    private static TimelyPrincipal getTimelyPrincipal(Authentication authentication,
+            SubjectIssuerDNPair subjectIssuerDNPair) {
+        String subjectDn = subjectIssuerDNPair.subjectDN();
+        TimelyUser.UserType userType = DnUtils.isServerDN(subjectDn) ? TimelyUser.UserType.SERVER
+                : TimelyUser.UserType.USER;
+        Collection<String> auths = authentication.getAuthorities().stream().map(a -> a.getAuthority())
+                .collect(Collectors.toList());
+        TimelyUser timelyUser = new TimelyUser(subjectIssuerDNPair, userType, auths, null, null,
+                System.currentTimeMillis());
+        return new TimelyPrincipal(Arrays.asList(timelyUser));
+    }
 }

server/src/main/java/timely/auth/util/AuthenticationUtils.java

@@ -1,50 +0,0 @@
-package timely.auth.util;
-
-import java.security.cert.X509Certificate;
-import java.util.Collection;
-import java.util.stream.Collectors;
-
-import com.google.common.collect.HashMultimap;
-import org.springframework.security.core.Authentication;
-import org.springframework.security.core.GrantedAuthority;
-import timely.auth.SubjectIssuerDNPair;
-import timely.auth.TimelyPrincipal;
-import timely.auth.TimelyUser;
-import timely.netty.http.auth.TimelyAuthenticationToken;
-
-public class AuthenticationUtils {
-
-    public static Collection<String> getAuthCollection(Authentication authentication) {
-        Collection<? extends GrantedAuthority> authorities = authentication.getAuthorities();
-        return authorities.stream().map(a -> a.getAuthority()).collect(Collectors.toList());
-    }
-
-    public static TimelyAuthenticationToken getTimelyAuthentication(Authentication authentication,
-            X509Certificate clientCert) {
-        TimelyAuthenticationToken timelyAuthenticationToken = null;
-        if (authentication instanceof TimelyAuthenticationToken) {
-            return (TimelyAuthenticationToken) authentication;
-        } else {
-            if (clientCert == null) {
-                timelyAuthenticationToken = new TimelyAuthenticationToken(authentication.getName(), null,
-                        HashMultimap.create());
-                TimelyUser timelyUser = new TimelyUser(SubjectIssuerDNPair.of(authentication.getName()),
-                        TimelyUser.UserType.USER, AuthenticationUtils.getAuthCollection(authentication), null, null,
-                        System.currentTimeMillis());
-                timelyAuthenticationToken.setTimelyPrincipal(new TimelyPrincipal(timelyUser));
-            } else {
-                timelyAuthenticationToken = new TimelyAuthenticationToken(authentication.getName(), clientCert,
-                        HashMultimap.create());
-                TimelyUser.UserType userType = DnUtils.isServerDN(clientCert.getSubjectDN().getName())
-                        ? TimelyUser.UserType.SERVER
-                        : TimelyUser.UserType.USER;
-                TimelyUser timelyUser = new TimelyUser(
-                        SubjectIssuerDNPair.of(clientCert.getSubjectDN().getName(), clientCert.getIssuerDN().getName()),
-                        userType, AuthenticationUtils.getAuthCollection(authentication), null, null,
-                        System.currentTimeMillis());
-                timelyAuthenticationToken.setTimelyPrincipal(new TimelyPrincipal(timelyUser));
-            }
-        }
-        return timelyAuthenticationToken;
-    }
-}

server/src/main/java/timely/netty/http/auth/X509LoginRequestHandler.java

@@ -9,6 +9,7 @@ import org.slf4j.LoggerFactory;
 import timely.api.request.auth.X509LoginRequest;
 import timely.api.response.TimelyException;
 import timely.auth.AuthenticationService;
+import timely.auth.SubjectIssuerDNPair;
 import timely.auth.TimelyPrincipal;
 import timely.configuration.Http;
 import timely.configuration.Security;
@@ -35,7 +36,9 @@ public class X509LoginRequestHandler extends TimelyLoginRequestHandler<X509Login
                 // login endpoint is for direct certificate only, not for proxying
                 TimelyAuthenticationToken token = AuthenticationService.getAuthenticationToken(clientCert,
                         clientCert.getSubjectDN().getName(), clientCert.getIssuerDN().getName());
-                principal = AuthenticationService.authenticate(token, clientCert, sessionId);
+                SubjectIssuerDNPair pair = SubjectIssuerDNPair.of(clientCert.getSubjectDN().getName(),
+                        clientCert.getIssuerDN().getName());
+                principal = AuthenticationService.authenticate(token, pair, sessionId);
                 LOG.debug("Authenticated user {} with authorizations {}",
                         principal.getPrimaryUser().getDn().subjectDN(), principal.getAuthorizationsString());
             }

server/src/test/java/timely/netty/websocket/WebSocketRequestDecoderTest.java

@@ -30,7 +30,6 @@ import timely.api.request.timeseries.SuggestRequest;
 import timely.auth.AuthCache;
 import timely.auth.TimelyPrincipal;
 import timely.auth.TimelyUser;
-import timely.auth.util.AuthenticationUtils;
 import timely.configuration.Configuration;
 import timely.model.Metric;
 import timely.subscription.SubscriptionRegistry;
@@ -67,8 +66,7 @@ public class WebSocketRequestDecoderTest {
         cookie = URLEncoder.encode(UUID.randomUUID().toString(), StandardCharsets.UTF_8.name());
         AuthCache.setSessionMaxAge(config.getSecurity());
         UsernamePasswordAuthenticationToken token = new UsernamePasswordAuthenticationToken("test", "test1");
-        AuthCache.put(cookie,
-                new TimelyPrincipal(new TimelyUser("ANONYMOUS", AuthenticationUtils.getAuthCollection(token))));
+        AuthCache.put(cookie, new TimelyPrincipal(TimelyUser.ANONYMOUS_USER));
     }
 
     @AfterClass