Improves token manager's refresh token handling

basyx.common/basyx.client/src/main/java/org/eclipse/digitaltwin/basyx/client/internal/authorization/TokenManager.java

@@ -26,30 +26,45 @@
 package org.eclipse.digitaltwin.basyx.client.internal.authorization;
 
 import java.io.IOException;
+import java.text.ParseException;
 
 import org.eclipse.digitaltwin.basyx.client.internal.authorization.grant.AccessTokenProvider;
 import org.eclipse.digitaltwin.basyx.core.exceptions.AccessTokenRetrievalException;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import com.nimbusds.oauth2.sdk.AccessTokenResponse;
 import com.nimbusds.oauth2.sdk.token.AccessToken;
 import com.nimbusds.oauth2.sdk.token.RefreshToken;
 
+import net.minidev.json.JSONObject;
+import com.nimbusds.jwt.JWT;
+import com.nimbusds.jwt.JWTParser;
+import com.nimbusds.jwt.SignedJWT;
+
+import java.util.Date;
+
 /**
  * Requests and manages the Access Tokens and Refresh Tokens.
  * 
- *  @author danish
+ * @author danish
  */
+import java.time.Instant;
+
 public class TokenManager {
-
-	private String tokenEndpoint;
-	private AccessTokenProvider accessTokenProvider;
+
+	private static final Logger LOGGER = LoggerFactory.getLogger(TokenManager.class);
+
+	private static final String EXPIRES_IN = "expires_in";
+	private static final String REFRESH_EXPIRES_IN = "refresh_expires_in";
+	private final String tokenEndpoint;
+	private final AccessTokenProvider accessTokenProvider;
 	private String accessToken;
-    private String refreshToken;
-    private long accessTokenExpiryTime;
-    private long refreshTokenExpiryTime;
-	
+	private String refreshToken;
+	private Instant accessTokenExpiryTime;
+	private Instant refreshTokenExpiryTime;
+
 	public TokenManager(String tokenEndpoint, AccessTokenProvider accessTokenProvider) {
-		super();
 		this.tokenEndpoint = tokenEndpoint;
 		this.accessTokenProvider = accessTokenProvider;
 	}
@@ -61,46 +76,139 @@ public String getTokenEndpoint() {
 	public AccessTokenProvider getAccessTokenProvider() {
 		return this.accessTokenProvider;
 	}
-
+
+	/**
+	 * Provides the access token, refreshing it if necessary.
+	 * 
+	 * @return the current valid access token
+	 * @throws IOException
+	 *             if an error occurs while retrieving the token
+	 */
+	public String getAccessToken() throws IOException {
+		Instant currentTime = Instant.now();
+
+		if (accessToken != null && currentTime.isBefore(accessTokenExpiryTime))
+			return accessToken;
+
+		synchronized (this) {
+			if (accessToken != null && currentTime.isBefore(accessTokenExpiryTime))
+				return accessToken;
+
+			if (refreshToken != null && currentTime.isBefore(refreshTokenExpiryTime))
+				return refreshAccessToken(currentTime);
+
+			return obtainNewAccessToken(currentTime);
+		}
+	}
+
 	/**
-	 * Provides access token
+	 * Updates the tokens and their expiry times.
 	 * 
-	 * @return accessToken
+	 * @param accessTokenResponse
+	 *            the response containing the new tokens
+	 * @param currentTime
+	 *            the current timestamp for consistency
+	 * @return the new access token
 	 * @throws IOException
+	 *             if an error occurs while processing the response
+	 */
+	private String updateTokens(AccessTokenResponse accessTokenResponse, Instant currentTime) throws IOException {
+		AccessToken accessTokenObj = accessTokenResponse.getTokens().getAccessToken();
+		accessToken = accessTokenObj.getValue();
+		accessTokenExpiryTime = calculateExpiryTime(accessTokenObj, currentTime);
+
+		RefreshToken refreshTokenObj = accessTokenResponse.getTokens().getRefreshToken();
+
+		if (refreshTokenObj != null) {
+			refreshToken = refreshTokenObj.getValue();
+			refreshTokenExpiryTime = calculateRefreshExpiryTime(refreshTokenObj, accessTokenResponse, currentTime);
+		}
+
+		return accessToken;
+	}
+
+	/**
+	 * Calculates the expiry time for a JWT token. First checks the 'exp' field in
+	 * the JWT, falling back to 'expires_in'.
+	 * 
+	 * @param tokenObj
+	 *            the AccessToken or RefreshToken object
+	 * @param currentTime
+	 *            the current timestamp
+	 * @return the calculated expiry time as Instant
+	 */
+	private Instant calculateExpiryTime(AccessToken tokenObj, Instant currentTime) {
+		String tokenValue = tokenObj.getValue();
+		Date expirationDate = extractExpirationTimeAsDateFromToken(tokenValue);
+
+		if (expirationDate != null)
+			return expirationDate.toInstant();
+
+		LOGGER.info("Unable to find 'exp' claim inside Access Token! Falling back to the alternative, the '{}' field.", EXPIRES_IN);
+
+		return currentTime.plusSeconds(tokenObj.getLifetime());
+	}
+
+	/**
+	 * Calculates the expiry time for a refresh token. First checks the 'exp' field
+	 * in the JWT refresh token, falling back to 'refresh_expires_in'.
+	 * 
+	 * @param refreshTokenObj
+	 *            the RefreshToken object
+	 * @param accessTokenResponse
+	 *            the response containing the refresh token
+	 * @param currentTime
+	 *            the current timestamp
+	 * @return the calculated expiry time as Instant
 	 */
-	public synchronized String getAccessToken() throws IOException {
+	private Instant calculateRefreshExpiryTime(RefreshToken refreshTokenObj, AccessTokenResponse accessTokenResponse, Instant currentTime) {
+		String tokenValue = refreshTokenObj.getValue();
+		Date expirationDate = extractExpirationTimeAsDateFromToken(tokenValue);
+
+		if (expirationDate != null)
+			return expirationDate.toInstant();
+
+		LOGGER.info("Unable to find 'exp' claim inside Refresh Token! Falling back to the alternative, the '{}' field", REFRESH_EXPIRES_IN);
 
-        if (accessToken != null && System.currentTimeMillis() < accessTokenExpiryTime)
-            return accessToken;
+		JSONObject jsonObject = accessTokenResponse.toJSONObject();
+		Number refreshExpiresInSeconds = jsonObject.getAsNumber(REFRESH_EXPIRES_IN);
 
-        if (refreshToken != null && System.currentTimeMillis() < refreshTokenExpiryTime) {
-            try {
-				return requestAccessToken(accessTokenProvider.getAccessTokenResponse(tokenEndpoint, refreshToken));
-			} catch (IOException e) {
-				throw new AccessTokenRetrievalException("Error occurred while retrieving access token" + e.getMessage());
+		if (refreshExpiresInSeconds == null)
+			return Instant.EPOCH;
+
+		return currentTime.plusSeconds(refreshExpiresInSeconds.longValue());
+	}
+
+	private Date extractExpirationTimeAsDateFromToken(String tokenValue) {
+		try {
+			JWT jwt = JWTParser.parse(tokenValue);
+
+			if (jwt instanceof SignedJWT) {
+				SignedJWT signedJwt = (SignedJWT) jwt;
+				return signedJwt.getJWTClaimsSet().getExpirationTime();
 			}
-        }
+		} catch (ParseException e) {
+			LOGGER.error("Failed to parse the token. Invalid JWT format: " + e.getMessage());
+		} catch (Exception e) {
+			LOGGER.error("Unexpected error occurred while extracting expiration time from the Token: " + e.getMessage());
+		}
+
+		return null;
+	}
+
+	private String obtainNewAccessToken(Instant currentTime) {
+		try {
+			return updateTokens(accessTokenProvider.getAccessTokenResponse(tokenEndpoint), currentTime);
+		} catch (IOException e) {
+			throw new AccessTokenRetrievalException("Error occurred while retrieving access token: " + e.getMessage());
+		}
+	}
 
-        try {
-			return requestAccessToken(accessTokenProvider.getAccessTokenResponse(tokenEndpoint));
+	private String refreshAccessToken(Instant currentTime) {
+		try {
+			return updateTokens(accessTokenProvider.getAccessTokenResponse(tokenEndpoint, refreshToken), currentTime);
 		} catch (IOException e) {
-			throw new AccessTokenRetrievalException("Error occurred while retrieving access token" + e.getMessage());
+			throw new AccessTokenRetrievalException("Error occurred while retrieving access token: " + e.getMessage());
 		}
-    }
-
-	private String requestAccessToken(AccessTokenResponse accessTokenResponse) throws IOException {        
-        AccessToken accessTokenObj = accessTokenResponse.getTokens().getAccessToken();
-        accessToken = accessTokenObj.getValue();
-        accessTokenExpiryTime = accessTokenObj.getLifetime();
-
-        RefreshToken refreshTokenObj = accessTokenResponse.getTokens().getRefreshToken();
-
-        if (refreshTokenObj != null) {
-        	refreshToken = refreshTokenObj.getValue();
-            refreshTokenExpiryTime = System.currentTimeMillis() + (30L * 24L * 60L * 60L * 1000L);
-        }
-
-        return accessToken;
-    }
-
+	}
 }

basyx.common/basyx.client/src/test/java/org/eclipse/digitaltwin/basyx/client/internal/TokenManagerTest.java

@@ -25,43 +25,106 @@
 
 package org.eclipse.digitaltwin.basyx.client.internal;
 
+import static org.junit.Assert.*;
 import org.eclipse.digitaltwin.basyx.client.internal.authorization.TokenManager;
-import org.eclipse.digitaltwin.basyx.client.internal.authorization.credential.ClientCredential;
-import org.eclipse.digitaltwin.basyx.client.internal.authorization.grant.ClientCredentialAccessTokenProvider;
 import org.junit.Before;
 import org.junit.Test;
+import org.eclipse.digitaltwin.basyx.client.internal.authorization.credential.ClientCredential;
+import org.eclipse.digitaltwin.basyx.client.internal.authorization.credential.PasswordCredential;
+import org.eclipse.digitaltwin.basyx.client.internal.authorization.grant.PasswordCredentialAccessTokenProvider;
 
 import java.io.IOException;
 
-import static org.junit.Assert.*;
-
 /**
  * Tests the behaviour of {@link TokenManager}
  * 
  * @author danish
  */
 public class TokenManagerTest {
-
+
+    private static final String TOKEN_ENDPOINT = "http://localhost:9096/realms/BaSyx/protocol/openid-connect/token";
     private TokenManager tokenManager;
 
     @Before
     public void setUp() {
-        tokenManager = new TokenManager("http://localhost:9096/realms/BaSyx/protocol/openid-connect/token", new ClientCredentialAccessTokenProvider(new ClientCredential("workstation-1", "nY0mjyECF60DGzNmQUjL81XurSl8etom")));
+        tokenManager = new TokenManager(TOKEN_ENDPOINT, new PasswordCredentialAccessTokenProvider(new PasswordCredential("basyx.reader", "basyxreader"), new ClientCredential("max-sso", "8ccc227xJflxGgtFkwwssHRZUh99nAAc")));
     }
 
     @Test
-    public void testGetAccessToken_RetrievesNewTokenAfterExpiry() throws IOException, InterruptedException {
-
+    public void retrieveNewAccessTokenWhenExpired() throws IOException, InterruptedException {        
         String initialAccessToken = tokenManager.getAccessToken();
         assertNotNull(initialAccessToken);
 
-        long tokenLifetime = 5000;
-        Thread.sleep(tokenLifetime + 1000);
+        long tokenLifetime = 1000;
+        Thread.sleep(tokenLifetime + 5);
 
         String newAccessToken = tokenManager.getAccessToken();
         assertNotNull(newAccessToken);
 
         assertNotEquals(initialAccessToken, newAccessToken);
     }
 
+    @Test
+    public void provideSameAccessTokenWhenNotExpired() throws IOException, InterruptedException {        
+        String initialAccessToken = tokenManager.getAccessToken();
+        assertNotNull(initialAccessToken);
+
+        long tokenLifetime = 350;
+        Thread.sleep(tokenLifetime);
+
+        String newAccessToken = tokenManager.getAccessToken();
+        assertNotNull(newAccessToken);
+
+        assertEquals(initialAccessToken, newAccessToken);
+    }
+
+    @Test
+    public void retrieveNewAccessTokenUsingRefreshTokenWhenAccessTokenIsExpired() throws IOException, InterruptedException {
+    	String initialAccessToken = tokenManager.getAccessToken();
+        assertNotNull(initialAccessToken);
+
+        long tokenLifetime = 1000;
+        Thread.sleep(tokenLifetime + 15);
+
+        String newAccessToken = tokenManager.getAccessToken();
+        assertNotNull(newAccessToken);
+
+        assertNotEquals(initialAccessToken, newAccessToken);
+    }
+
+    @Test
+    public void retrieveNewAccessTokenUsingRefreshTokenWhenExpired() throws IOException, InterruptedException {
+    	String initialAccessToken = tokenManager.getAccessToken();
+        assertNotNull(initialAccessToken);
+
+        long tokenLifetime = 3000;
+        Thread.sleep(tokenLifetime + 5);
+
+        String newAccessToken = tokenManager.getAccessToken();
+        assertNotNull(newAccessToken);
+
+        assertNotEquals(initialAccessToken, newAccessToken);
+    }
+
+    @Test
+    // This test is for verifying the behavior described in the issue: https://github.com/eclipse-basyx/basyx-java-server-sdk/issues/530
+    public void retrieveNewAccessTokenWhenSSOMaxReached() throws IOException, InterruptedException {
+    	String initialAccessToken = tokenManager.getAccessToken();
+        assertNotNull(initialAccessToken);
+
+        long timeLimit = 4000;
+        String intermediateAccessToken = null;
+
+        while (timeLimit > 0) {
+        	intermediateAccessToken = tokenManager.getAccessToken();
+
+        	Thread.sleep(2000);
+
+        	timeLimit -= 2000;
+        }
+
+        assertNotNull(intermediateAccessToken);
+        assertNotEquals(initialAccessToken, intermediateAccessToken);
+    }
+
 }