fix(oxauth): corrected race condition during refresh token usage (4.5.5)

#1909
GluuFederation · Jul 31, 2024 · dcaa88f · dcaa88f
1 parent ca0679b
commit dcaa88f
Show file tree

Hide file tree

Showing 2 changed files with 62 additions and 4 deletions.
diff --git a/Server/src/main/java/org/gluu/oxauth/service/GrantService.java b/Server/src/main/java/org/gluu/oxauth/service/GrantService.java
@@ -30,6 +30,7 @@
 import javax.inject.Inject;
 import java.util.*;
 
+import static org.gluu.oxauth.util.ServerUtil.calculateTtl;
 import static org.gluu.oxauth.util.ServerUtil.isTrue;
 
 /**
@@ -73,13 +74,18 @@ private String tokenBaseDn() {
         return staticConfiguration.getBaseDn().getTokens();  // ou=tokens,o=gluu
     }
 
-    public void merge(TokenLdap p_token) {
-        ldapEntryManager.merge(p_token);
+    public void merge(TokenLdap token) {
+        if (shouldPutInCache(token.getTokenTypeEnum(), token.isImplicitFlow())) {
+            final int expiration = calculateTtl(new Date(), token.getExpirationDate());
+            cacheService.put(expiration, token.getTokenCode(), token);
+        } else {
+            ldapEntryManager.merge(token);
+        }
     }
 
-    public void mergeSilently(TokenLdap p_token) {
+    public void mergeSilently(TokenLdap token) {
         try {
-            ldapEntryManager.merge(p_token);
+            merge(token);
         } catch (Exception e) {
             log.error(e.getMessage(), e);
         }

diff --git a/Server/src/main/java/org/gluu/oxauth/token/ws/rs/TokenRestWebServiceImpl.java b/Server/src/main/java/org/gluu/oxauth/token/ws/rs/TokenRestWebServiceImpl.java
@@ -18,6 +18,7 @@
 import org.gluu.oxauth.model.configuration.AppConfiguration;
 import org.gluu.oxauth.model.crypto.binding.TokenBindingMessage;
 import org.gluu.oxauth.model.error.ErrorResponseFactory;
+import org.gluu.oxauth.model.ldap.TokenLdap;
 import org.gluu.oxauth.model.registration.Client;
 import org.gluu.oxauth.model.session.SessionClient;
 import org.gluu.oxauth.model.session.SessionId;
@@ -52,6 +53,9 @@
 import javax.ws.rs.core.SecurityContext;
 import java.util.Arrays;
 import java.util.Date;
+import java.util.UUID;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
 
 import static org.gluu.oxauth.util.ServerUtil.prepareForLogs;
 
@@ -65,6 +69,8 @@
 @Path("/")
 public class TokenRestWebServiceImpl implements TokenRestWebService {
 
+    private static final String NODE_ID = UUID.randomUUID().toString();
+
     @Inject
     private Logger log;
 
@@ -116,6 +122,8 @@ public class TokenRestWebServiceImpl implements TokenRestWebService {
     @Inject
     private ExternalUpdateTokenService externalUpdateTokenService;
 
+    private final ConcurrentMap<String, TokenLdap> refreshTokenLocalLock = new ConcurrentHashMap<>();
+
     @Override
     public Response requestAccessToken(String grantType, String code,
                                        String redirectUri, String username, String password, String scope,
@@ -270,6 +278,12 @@ grantType, code, redirectUri, username, refreshToken, clientId, prepareForLogs(r
                     return response(error(400, TokenErrorResponseType.INVALID_GRANT, "Unable to find refresh token or otherwise token type or client does not match."), oAuth2AuditLog);
                 }
 
+                TokenLdap lockedRefreshToken = lockRefreshToken(refreshToken);
+                if (lockedRefreshToken == null) {
+                    log.trace("Failed to lock refresh token {}", refreshToken);
+                    return response(error(400, TokenErrorResponseType.INVALID_GRANT, "Failed to lock refresh token."), oAuth2AuditLog);
+                }
+
                 checkUser(authorizationGrant, oAuth2AuditLog);
                 executionContext.setGrant(authorizationGrant);
 
@@ -553,6 +567,44 @@ grantType, code, redirectUri, username, refreshToken, clientId, prepareForLogs(r
         return response(builder, oAuth2AuditLog);
     }
 
+    private TokenLdap lockRefreshToken(String refreshTokenCode) {
+        try {
+            if (refreshTokenLocalLock.containsKey(refreshTokenCode)) {
+                log.trace("Refresh token is already used by another request. Refresh token  code: {}", refreshTokenCode);
+                return null;
+            }
+
+            for (int attempt = 1; attempt <= 3; attempt++) {
+                try {
+                    final TokenLdap token = grantService.getGrantByCode(refreshTokenCode);
+                    if (token == null) {
+                        log.trace("Refresh token is not found by code {}", refreshTokenCode);
+                        return null;
+                    }
+
+                    refreshTokenLocalLock.put(refreshTokenCode, token);
+
+                    token.getAttributes().getAttributes().put("lockKey", NODE_ID);
+                    grantService.mergeSilently(token);
+                    final TokenLdap tokenFromDb = grantService.getGrantByCode(refreshTokenCode);
+                    if (NODE_ID.equals(tokenFromDb.getAttributes().getAttributes().get("lockKey"))) {
+                        log.trace("Successfully locked refresh token {}, attempt {}", refreshTokenCode, attempt);
+                        return token;
+                    }
+
+                    log.trace("Failed to lock refresh token {}, attempt {}", refreshTokenCode, attempt);
+                    Thread.sleep(100);
+                } catch (InterruptedException e) {
+                    // ignore and make next attempt
+                    log.trace(e.getMessage(), e);
+                }
+            }
+        } finally {
+            refreshTokenLocalLock.remove(refreshTokenCode);
+        }
+        return null;
+    }
+
     private void checkUser(AuthorizationGrant authorizationGrant, OAuth2AuditLog oAuth2AuditLog) {
         if (!appConfiguration.getCheckUserPresenceOnRefreshToken()) {
             return;