Skip to content

Commit

Permalink
Merge pull request #45302 from sberyozkin/optimize_default_token_stat…
Browse files Browse the repository at this point in the history 
…e_manager

Minor refactoring of OIDC DefaultTokenStateManager
  • Loading branch information
@sberyozkin
sberyozkin authored Dec 28, 2024
2 parents 1ffcee2 + abed24b commit 283659c
Showing 1 changed file with 90 additions and 57 deletions.
147 changes: 90 additions & 57 deletions extensions/oidc/runtime/src/main/java/io/quarkus/oidc/runtime/DefaultTokenStateManager.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,95 +22,128 @@ public class DefaultTokenStateManager implements TokenStateManager {
public Uni<String> createTokenState(RoutingContext routingContext, OidcTenantConfig oidcConfig,
AuthorizationCodeTokens tokens, OidcRequestContext<String> requestContext) {

boolean encryptAll = !oidcConfig.tokenStateManager().splitTokens();
if (!oidcConfig.tokenStateManager().splitTokens()) {
// ID, access and refresh tokens are all represented by a single cookie.
// In this case they are all encrypted once all tokens have been added to the buffer.

StringBuilder sb = new StringBuilder();
sb.append(encryptAll ? tokens.getIdToken() : encryptToken(tokens.getIdToken(), routingContext, oidcConfig));
if (oidcConfig.tokenStateManager().strategy() == Strategy.KEEP_ALL_TOKENS) {
if (!oidcConfig.tokenStateManager().splitTokens()) {
StringBuilder sb = new StringBuilder();

// Add ID token
sb.append(tokens.getIdToken());

// By default, all three tokens are retained
if (oidcConfig.tokenStateManager().strategy() == Strategy.KEEP_ALL_TOKENS) {
// Add access and refresh tokens
sb.append(CodeAuthenticationMechanism.COOKIE_DELIM)
.append(tokens.getAccessToken())
.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
sb.append(CodeAuthenticationMechanism.COOKIE_DELIM)
.append(encryptAll ? tokens.getAccessToken()
: encryptToken(tokens.getAccessToken(), routingContext, oidcConfig))
.append("")
.append(CodeAuthenticationMechanism.COOKIE_DELIM)
.append(encryptAll ? tokens.getRefreshToken()
: encryptToken(tokens.getRefreshToken(), routingContext, oidcConfig));
} else {
.append(tokens.getRefreshToken());
}

// Now all three tokens are encrypted
String encryptedTokens = encryptToken(sb.toString(), routingContext, oidcConfig);
return Uni.createFrom().item(encryptedTokens);
} else {
// ID, access and refresh tokens are represented as individual cookies

// Encrypt ID token
String encryptedIdToken = encryptToken(tokens.getIdToken(), routingContext, oidcConfig);

// By default, all three tokens are retained
if (oidcConfig.tokenStateManager().strategy() == Strategy.KEEP_ALL_TOKENS) {

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

// Encrypt refresh token and create a `q_session_rt` cookie.
if (tokens.getRefreshToken() != null) {
CodeAuthenticationMechanism.createCookie(routingContext,
oidcConfig,
getRefreshTokenCookieName(oidcConfig),
encryptToken(tokens.getRefreshToken(), routingContext, oidcConfig),
routingContext.get(CodeAuthenticationMechanism.SESSION_MAX_AGE_PARAM), true);
}
} else if (oidcConfig.tokenStateManager().strategy() == Strategy.ID_REFRESH_TOKENS
&& tokens.getRefreshToken() != null) {
// Encrypt refresh token and create a `q_session_rt` cookie.
CodeAuthenticationMechanism.createCookie(routingContext,
oidcConfig,
getRefreshTokenCookieName(oidcConfig),
encryptToken(tokens.getRefreshToken(), routingContext, oidcConfig),
routingContext.get(CodeAuthenticationMechanism.SESSION_MAX_AGE_PARAM));
}
} else if (oidcConfig.tokenStateManager().strategy() == Strategy.ID_REFRESH_TOKENS) {
if (!oidcConfig.tokenStateManager().splitTokens()) {
sb.append(CodeAuthenticationMechanism.COOKIE_DELIM)
.append("")
.append(CodeAuthenticationMechanism.COOKIE_DELIM)
.append(encryptAll ? tokens.getRefreshToken()
: encryptToken(tokens.getRefreshToken(), routingContext, oidcConfig));
} else {
if (tokens.getRefreshToken() != null) {
CodeAuthenticationMechanism.createCookie(routingContext,
oidcConfig,
getRefreshTokenCookieName(oidcConfig),
encryptToken(tokens.getRefreshToken(), routingContext, oidcConfig),
routingContext.get(CodeAuthenticationMechanism.SESSION_MAX_AGE_PARAM));
}
}

// q_session cookie
return Uni.createFrom().item(encryptedIdToken);
}
String state = encryptAll ? encryptToken(sb.toString(), routingContext, oidcConfig) : sb.toString();
return Uni.createFrom().item(state);

}

@Override
public Uni<AuthorizationCodeTokens> getTokens(RoutingContext routingContext, OidcTenantConfig oidcConfig, String tokenState,
OidcRequestContext<AuthorizationCodeTokens> requestContext) {
boolean decryptAll = !oidcConfig.tokenStateManager().splitTokens();

tokenState = decryptAll ? decryptToken(tokenState, routingContext, oidcConfig) : tokenState;
String idToken = null;
String accessToken = null;
String refreshToken = null;

String[] tokens = CodeAuthenticationMechanism.COOKIE_PATTERN.split(tokenState);
if (!oidcConfig.tokenStateManager().splitTokens()) {
// ID, access and refresh tokens are all be represented by a single cookie.

String idToken = decryptAll ? tokens[0] : decryptToken(tokens[0], routingContext, oidcConfig);
String decryptedTokenState = decryptToken(tokenState, routingContext, oidcConfig);

String[] tokens = CodeAuthenticationMechanism.COOKIE_PATTERN.split(decryptedTokenState);

try {
idToken = tokens[0];
accessToken = null;
refreshToken = null;

if (oidcConfig.tokenStateManager().strategy() == Strategy.KEEP_ALL_TOKENS) {
accessToken = tokens[1];
refreshToken = tokens[2];
} else if (oidcConfig.tokenStateManager().strategy() == Strategy.ID_REFRESH_TOKENS) {
refreshToken = tokens[2];
}
} catch (ArrayIndexOutOfBoundsException ex) {
return Uni.createFrom().failure(new AuthenticationCompletionException("Session cookie is malformed"));
}
} else {
// Decrypt ID token from the q_session cookie
idToken = decryptToken(tokenState, routingContext, oidcConfig);
accessToken = null;
refreshToken = null;

String accessToken = null;
String refreshToken = null;
try {
if (oidcConfig.tokenStateManager().strategy() == Strategy.KEEP_ALL_TOKENS) {
if (!oidcConfig.tokenStateManager().splitTokens()) {
accessToken = decryptAll ? tokens[1] : decryptToken(tokens[1], routingContext, oidcConfig);
refreshToken = decryptAll ? tokens[2] : decryptToken(tokens[2], routingContext, oidcConfig);
} else {
Cookie atCookie = getAccessTokenCookie(routingContext, oidcConfig);
if (atCookie != null) {
accessToken = decryptToken(atCookie.getValue(), routingContext, oidcConfig);
}
Cookie rtCookie = getRefreshTokenCookie(routingContext, oidcConfig);
if (rtCookie != null) {
refreshToken = decryptToken(rtCookie.getValue(), routingContext, oidcConfig);
}
Cookie atCookie = getAccessTokenCookie(routingContext, oidcConfig);
if (atCookie != null) {
// Decrypt access token from the q_session_at cookie
accessToken = decryptToken(atCookie.getValue(), routingContext, oidcConfig);
}
Cookie rtCookie = getRefreshTokenCookie(routingContext, oidcConfig);
if (rtCookie != null) {
// Decrypt refresh token from the q_session_rt cookie
refreshToken = decryptToken(rtCookie.getValue(), routingContext, oidcConfig);
}
} else if (oidcConfig.tokenStateManager().strategy() == Strategy.ID_REFRESH_TOKENS) {
if (!oidcConfig.tokenStateManager().splitTokens()) {
refreshToken = decryptAll ? tokens[2] : decryptToken(tokens[2], routingContext, oidcConfig);
} else {
Cookie rtCookie = getRefreshTokenCookie(routingContext, oidcConfig);
if (rtCookie != null) {
refreshToken = decryptToken(rtCookie.getValue(), routingContext, oidcConfig);
}
Cookie rtCookie = getRefreshTokenCookie(routingContext, oidcConfig);
if (rtCookie != null) {
refreshToken = decryptToken(rtCookie.getValue(), routingContext, oidcConfig);
}
}
} catch (ArrayIndexOutOfBoundsException ex) {
return Uni.createFrom().failure(new AuthenticationCompletionException("Session cookie is malformed"));
}

return Uni.createFrom().item(new AuthorizationCodeTokens(idToken, accessToken, refreshToken));
}

Expand Down

0 comments on commit 283659c

Please sign in to comment.