feat(jans-auth-server): added configurable option to put all tokens i…

…nto cache service #7563 (#7686) #7563 Signed-off-by: YuriyZ <yzabrovarniy@gmail.com>
JanssenProject · Feb 9, 2024 · fe4c000 · fe4c000
1 parent 0f62d10
commit fe4c000
Show file tree

Hide file tree

Showing 3 changed files with 91 additions and 14 deletions.
diff --git a/jans-auth-server/model/src/main/java/io/jans/as/model/configuration/AppConfiguration.java b/jans-auth-server/model/src/main/java/io/jans/as/model/configuration/AppConfiguration.java
@@ -354,6 +354,12 @@ public class AppConfiguration implements Configuration {
     @DocProperty(description = "Boolean value specifying whether idToken filters claims based on accessToken")
     private Boolean idTokenFilterClaimsBasedOnAccessToken;
 
+    @DocProperty(description = "Boolean value specifying whether to save access_token, id_token and refresh_token in cache (with cacheKey=sha256Hex(token_code))")
+    private Boolean saveTokensInCache;
+
+    @DocProperty(description = "Boolean value specifying whether to save access_token, id_token and refresh_token in cache and skip persistence in DB at the same time (with cacheKey=sha256Hex(token_code))")
+    private Boolean saveTokensInCacheAndDontSaveInPersistence;
+
     @DocProperty(description = "The lifetime of the short lived Access Token")
     private int accessTokenLifetime;
 
@@ -2294,6 +2300,22 @@ public void setAccessTokenLifetime(int accessTokenLifetime) {
         this.accessTokenLifetime = accessTokenLifetime;
     }
 
+    public Boolean getSaveTokensInCache() {
+        return saveTokensInCache;
+    }
+
+    public void setSaveTokensInCache(Boolean saveTokensInCache) {
+        this.saveTokensInCache = saveTokensInCache;
+    }
+
+    public Boolean getSaveTokensInCacheAndDontSaveInPersistence() {
+        return saveTokensInCacheAndDontSaveInPersistence;
+    }
+
+    public void setSaveTokensInCacheAndDontSaveInPersistence(Boolean saveTokensInCacheAndDontSaveInPersistence) {
+        this.saveTokensInCacheAndDontSaveInPersistence = saveTokensInCacheAndDontSaveInPersistence;
+    }
+
     public int getUmaRptLifetime() {
         return umaRptLifetime;
     }

diff --git a/jans-auth-server/server/src/main/java/io/jans/as/server/service/GrantService.java b/jans-auth-server/server/src/main/java/io/jans/as/server/service/GrantService.java
@@ -6,18 +6,7 @@
 
 package io.jans.as.server.service;
 
-import static org.apache.commons.lang.BooleanUtils.isTrue;
-
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.List;
-import java.util.UUID;
-
-import org.apache.commons.lang.StringUtils;
-import org.slf4j.Logger;
-
 import com.google.common.collect.Lists;
-
 import io.jans.as.model.config.StaticConfiguration;
 import io.jans.as.model.configuration.AppConfiguration;
 import io.jans.as.model.configuration.LockMessageConfig;
@@ -35,6 +24,15 @@
 import jakarta.ejb.Stateless;
 import jakarta.inject.Inject;
 import jakarta.inject.Named;
+import org.apache.commons.lang.StringUtils;
+import org.slf4j.Logger;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.UUID;
+
+import static org.apache.commons.lang.BooleanUtils.isTrue;
 
 /**
  * @author Yuriy Zabrovarnyy
@@ -93,19 +91,42 @@ public void mergeSilently(TokenEntity token) {
         }
     }
 
+    public boolean shouldPersist() {
+        if (isTrue(appConfiguration.getSaveTokensInCacheAndDontSaveInPersistence())) {
+            return false;
+        }
+        return true;
+    }
+
+    public boolean shouldSaveInCache() {
+        return isTrue(appConfiguration.getSaveTokensInCache()) || isTrue(appConfiguration.getSaveTokensInCacheAndDontSaveInPersistence());
+    }
+
     public void persist(TokenEntity token) {
-        persistenceEntryManager.persist(token);
-
+        if (shouldPersist()) {
+            persistenceEntryManager.persist(token);
+        }
+
+        if (shouldSaveInCache()) {
+            saveInCache(token);
+        }
+
         if (TokenType.ID_TOKEN.getValue().equals(token.getTokenType())) {
         	publishIdTokenLockMessage(token, "add");
         }
     }
 
+    private void saveInCache(TokenEntity token) {
+        long lifeTimeAsMillis = token.getExpirationDate().getTime() - System.currentTimeMillis();
+        int lifetimeInSeconds = (int) (lifeTimeAsMillis / 1000);
+        cacheService.put(lifetimeInSeconds, TokenHashUtil.hash(token.getTokenCode()), token.getTokenCode());
+    }
+
     public void remove(TokenEntity token) {
         persistenceEntryManager.remove(token);
         log.trace("Removed token from LDAP, code: {}", token.getTokenCode());
 
-        if (TokenType.ID_TOKEN.equals(token.getTokenType())) {
+        if (TokenType.ID_TOKEN == token.getTokenTypeEnum()) {
         	publishIdTokenLockMessage(token, "del");
         }
     }
@@ -129,6 +150,9 @@ public void removeSilently(TokenEntity token) {
             if (StringUtils.isNotBlank(token.getAuthorizationCode())) {
                 cacheService.remove(CacheGrant.cacheKey(token.getAuthorizationCode(), token.getGrantId()));
             }
+            if (shouldSaveInCache()) {
+                cacheService.remove(token.getTokenCode());
+            }
         } catch (Exception e) {
             log.error(e.getMessage(), e);
         }

diff --git a/jans-auth-server/server/src/test/java/io/jans/as/server/service/GrantServiceTest.java b/jans-auth-server/server/src/test/java/io/jans/as/server/service/GrantServiceTest.java
@@ -50,6 +50,37 @@ public class GrantServiceTest {
     @Mock
     private CacheConfiguration cacheConfiguration;
 
+    @Test
+    public void shouldPersist_byDefault_shouldReturnTrue() {
+        assertTrue(grantService.shouldPersist());
+    }
+
+    @Test
+    public void shouldPersist_whenDontPersistConfIsTrue_shouldReturnFalse() {
+        Mockito.doReturn(true).when(appConfiguration).getSaveTokensInCacheAndDontSaveInPersistence();
+
+        assertFalse(grantService.shouldPersist());
+    }
+
+    @Test
+    public void shouldSaveInCache_byDefault_shoultReturnFalse() {
+        assertFalse(grantService.shouldSaveInCache());
+    }
+
+    @Test
+    public void shouldSaveInCache_whenAllowedByConfig_shoultReturnTrue() {
+        Mockito.doReturn(true).when(appConfiguration).getSaveTokensInCacheAndDontSaveInPersistence();
+
+        assertTrue(grantService.shouldSaveInCache());
+    }
+
+    @Test
+    public void shouldSaveInCache_whenAllowedByMainConfig_shoultReturnTrue() {
+        Mockito.doReturn(true).when(appConfiguration).getSaveTokensInCache();
+
+        assertTrue(grantService.shouldSaveInCache());
+    }
+
     @Test
     public void filterOutRefreshTokenFromDeletion_forTokenWithoutOnlineAccess_shouldFilterOut() {
         Mockito.doReturn(false).when(appConfiguration).getRemoveRefreshTokensForClientOnLogout();