Skip to content

Commit

Permalink
Save access token expires_in in the session cookie
Browse files Browse the repository at this point in the history
  • Loading branch information
sberyozkin committed Dec 31, 2024
1 parent 3b22087 commit 48078f6
Show file tree
Hide file tree
Showing 4 changed files with 53 additions and 11 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,9 @@
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.fail;

import java.util.List;

import org.htmlunit.SilentCssErrorHandler;
import org.htmlunit.TextPage;
Expand Down Expand Up @@ -166,7 +169,26 @@ private static WebClient createWebClient() {
return webClient;
}

private static Cookie getSessionCookie(WebClient webClient, String tenantId) {
return webClient.getCookieManager().getCookie("q_session" + (tenantId == null ? "" : "_" + tenantId));
private static List<Cookie> getSessionCookie(WebClient webClient, String tenantId) {
Cookie sessionCookieChunk1 = null;
Cookie sessionCookieChunk2 = null;

String sessionCookieSuffix = "q_session" + (tenantId == null ? "" : "_" + tenantId);
String sessionCookieChunk1Name = sessionCookieSuffix + "_chunk_1";
String sessionCookieChunk2Name = sessionCookieSuffix + "_chunk_2";
for (Cookie c : webClient.getCookieManager().getCookies()) {
if (c.getName().startsWith(sessionCookieSuffix)) {
if (c.getName().equals(sessionCookieChunk1Name)) {
sessionCookieChunk1 = c;
} else if (c.getName().equals(sessionCookieChunk2Name)) {
sessionCookieChunk2 = c;
} else {
fail("Unexpected session cookie chunk: " + c.getName());
}
}
}
return sessionCookieChunk1 != null && sessionCookieChunk2 != null
? List.of(sessionCookieChunk1, sessionCookieChunk2)
: null;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,6 @@ public Long getAccessTokenExpiresIn() {
/**
* Set the access token expires_in value in seconds.
* It is relative to the time the access token is issued at.
* This property is only checked when an authorization code flow grant completes and does not have to be persisted..
*
* @param accessTokenExpiresIn access token expires_in value in seconds.
*/
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -37,12 +37,16 @@ public Uni<String> createTokenState(RoutingContext routingContext, OidcTenantCon
sb.append(CodeAuthenticationMechanism.COOKIE_DELIM)
.append(tokens.getAccessToken())
.append(CodeAuthenticationMechanism.COOKIE_DELIM)
.append(tokens.getAccessTokenExpiresIn() != null ? tokens.getAccessTokenExpiresIn() : "")
.append(CodeAuthenticationMechanism.COOKIE_DELIM)
.append(tokens.getRefreshToken());
} else if (oidcConfig.tokenStateManager().strategy() == Strategy.ID_REFRESH_TOKENS) {
// But sometimes the access token is not required.
// For example, when the Quarkus endpoint does not need to use it to access another service.
// Skip access token, add refresh token
// Skip access token and access token expiry, add refresh token
sb.append(CodeAuthenticationMechanism.COOKIE_DELIM)
.append("")
.append(CodeAuthenticationMechanism.COOKIE_DELIM)
.append("")
.append(CodeAuthenticationMechanism.COOKIE_DELIM)
.append(tokens.getRefreshToken());
Expand All @@ -60,11 +64,18 @@ public Uni<String> createTokenState(RoutingContext routingContext, OidcTenantCon
// By default, all three tokens are retained
if (oidcConfig.tokenStateManager().strategy() == Strategy.KEEP_ALL_TOKENS) {

StringBuilder sb = new StringBuilder();

// Add ID token
sb.append(tokens.getAccessToken())
.append(CodeAuthenticationMechanism.COOKIE_DELIM)
.append(tokens.getAccessTokenExpiresIn() != null ? tokens.getAccessTokenExpiresIn() : "");

// Encrypt access token and create a `q_session_at` cookie.
CodeAuthenticationMechanism.createCookie(routingContext,
oidcConfig,
getAccessTokenCookieName(oidcConfig),
encryptToken(tokens.getAccessToken(), routingContext, oidcConfig),
encryptToken(sb.toString(), routingContext, oidcConfig),
routingContext.get(CodeAuthenticationMechanism.SESSION_MAX_AGE_PARAM), true);

// Encrypt refresh token and create a `q_session_rt` cookie.
Expand Down Expand Up @@ -97,6 +108,7 @@ public Uni<AuthorizationCodeTokens> getTokens(RoutingContext routingContext, Oid

String idToken = null;
String accessToken = null;
Long accessTokenExpiresIn = null;
String refreshToken = null;

if (!oidcConfig.tokenStateManager().splitTokens()) {
Expand All @@ -113,9 +125,10 @@ public Uni<AuthorizationCodeTokens> getTokens(RoutingContext routingContext, Oid

if (oidcConfig.tokenStateManager().strategy() == Strategy.KEEP_ALL_TOKENS) {
accessToken = tokens[1];
refreshToken = tokens[2];
accessTokenExpiresIn = tokens[2].isEmpty() ? null : Long.valueOf(tokens[2]);
refreshToken = tokens[3];
} else if (oidcConfig.tokenStateManager().strategy() == Strategy.ID_REFRESH_TOKENS) {
refreshToken = tokens[2];
refreshToken = tokens[3];
}
} catch (ArrayIndexOutOfBoundsException ex) {
return Uni.createFrom().failure(new AuthenticationCompletionException("Session cookie is malformed"));
Expand All @@ -130,7 +143,14 @@ public Uni<AuthorizationCodeTokens> getTokens(RoutingContext routingContext, Oid
Cookie atCookie = getAccessTokenCookie(routingContext, oidcConfig);
if (atCookie != null) {
// Decrypt access token from the q_session_at cookie
accessToken = decryptToken(atCookie.getValue(), routingContext, oidcConfig);
String accessTokenState = decryptToken(atCookie.getValue(), routingContext, oidcConfig);
String[] accessTokenData = CodeAuthenticationMechanism.COOKIE_PATTERN.split(accessTokenState);
accessToken = accessTokenData[0];
try {
accessTokenExpiresIn = accessTokenData[1].isEmpty() ? null : Long.valueOf(accessTokenData[1]);
} catch (ArrayIndexOutOfBoundsException ex) {
return Uni.createFrom().failure(new AuthenticationCompletionException("Session cookie is malformed"));
}
}
Cookie rtCookie = getRefreshTokenCookie(routingContext, oidcConfig);
if (rtCookie != null) {
Expand All @@ -144,7 +164,7 @@ public Uni<AuthorizationCodeTokens> getTokens(RoutingContext routingContext, Oid
}
}
}
return Uni.createFrom().item(new AuthorizationCodeTokens(idToken, accessToken, refreshToken));
return Uni.createFrom().item(new AuthorizationCodeTokens(idToken, accessToken, refreshToken, accessTokenExpiresIn));
}

@Override
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -1211,10 +1211,11 @@ public void testDefaultSessionManagerIdRefreshTokens() throws Exception {
String sessionCookieValue = OidcUtils.decryptString(sessionCookie.getValue(), key);

String[] parts = sessionCookieValue.split("\\|");
assertEquals(3, parts.length);
assertEquals(4, parts.length);
assertEquals("ID", OidcCommonUtils.decodeJwtContent(parts[0]).getString("typ"));
assertEquals("", parts[1]);
assertEquals("Refresh", OidcCommonUtils.decodeJwtContent(parts[2]).getString("typ"));
assertEquals("", parts[2]);
assertEquals("Refresh", OidcCommonUtils.decodeJwtContent(parts[3]).getString("typ"));

assertNull(getSessionAtCookie(webClient, "tenant-id-refresh-token"));
assertNull(getSessionRtCookie(webClient, "tenant-id-refresh-token"));
Expand Down

0 comments on commit 48078f6

Please sign in to comment.