Merge pull request #45327 from sberyozkin/oidc_session_cookie_access_token_expiry

sberyozkin · web-flow · commit 85d9208d4147 · 2025-01-02T11:03:24.000Z
Save access token expires_in in the session cookie
diff --git a/extensions/oidc/deployment/src/test/java/io/quarkus/oidc/test/CodeTenantReauthenticateTestCase.java b/extensions/oidc/deployment/src/test/java/io/quarkus/oidc/test/CodeTenantReauthenticateTestCase.java
@@ -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;
@@ -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;
     }
 }
diff --git a/extensions/oidc/runtime/src/main/java/io/quarkus/oidc/AuthorizationCodeTokens.java b/extensions/oidc/runtime/src/main/java/io/quarkus/oidc/AuthorizationCodeTokens.java
@@ -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.
      */
diff --git a/extensions/oidc/runtime/src/main/java/io/quarkus/oidc/runtime/DefaultTokenStateManager.java b/extensions/oidc/runtime/src/main/java/io/quarkus/oidc/runtime/DefaultTokenStateManager.java
@@ -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());
@@ -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 access token and its expires_in property
+                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.
@@ -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()) {
@@ -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"));
@@ -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) {
@@ -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
diff --git a/integration-tests/oidc-code-flow/src/test/java/io/quarkus/it/keycloak/CodeFlowTest.java b/integration-tests/oidc-code-flow/src/test/java/io/quarkus/it/keycloak/CodeFlowTest.java
@@ -352,9 +352,15 @@ public void testCodeFlowForceHttpsRedirectUriWithQueryAndPkce() throws Exception
                     "AES");
             String decryptedSessionCookieValue = OidcUtils.decryptString(sessionCookie.getValue(), key);
 
-            String encodedIdToken = decryptedSessionCookieValue.split("\\|")[0];
+            String decrypedSessionCookieValues[] = decryptedSessionCookieValue.split("\\|");
+            assertEquals(4, decrypedSessionCookieValues.length);
+
+            // ID token
+            String encodedIdToken = decrypedSessionCookieValues[0];
 
             JsonObject idToken = OidcCommonUtils.decodeJwtContent(encodedIdToken);
+            assertEquals("ID", idToken.getString("typ"));
+
             String expiresAt = idToken.getInteger("exp").toString();
             page = webClient.getPage(endpointLocationWithoutQueryUri.toURL());
             String response = page.getBody().asNormalizedText();
@@ -363,6 +369,13 @@ public void testCodeFlowForceHttpsRedirectUriWithQueryAndPkce() throws Exception
             Integer duration = Integer.valueOf(response.substring(response.length() - 1));
             assertTrue(duration > 1 && duration < 5);
 
+            // Access token and its expires_in
+            assertEquals("Bearer", OidcCommonUtils.decodeJwtContent(decrypedSessionCookieValues[1]).getString("typ"));
+            long atExpiresIn = Long.valueOf(decrypedSessionCookieValues[2]);
+            assertTrue(atExpiresIn >= 2 && atExpiresIn <= 4);
+            // Refresh token
+            assertEquals("Refresh", OidcCommonUtils.decodeJwtContent(decrypedSessionCookieValues[3]).getString("typ"));
+
             assertNull(getSessionCookie(webClient, "tenant-https"));
 
             webClient.getCookieManager().clearCookies();
@@ -1211,10 +1224,13 @@ 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"));
+            // No access token
             assertEquals("", parts[1]);
-            assertEquals("Refresh", OidcCommonUtils.decodeJwtContent(parts[2]).getString("typ"));
+            // No access token expires_in
+            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"));
@@ -1348,7 +1364,20 @@ private void checkSingleTokenCookie(Cookie tokenCookie, String type, String decr
                 SecretKey key = new SecretKeySpec(OidcUtils
                         .getSha256Digest(decryptSecret.getBytes(StandardCharsets.UTF_8)),
                         "AES");
-                token = OidcUtils.decryptString(token, key);
+                String decryptedString = OidcUtils.decryptString(token, key);
+                String[] decryptedStringParts = decryptedString.split("\\|");
+
+                // If it is an access token then an expiry date should follow the actual token
+                if ("Bearer".equals(type)) {
+                    assertEquals(2, decryptedStringParts.length);
+                    // Test access token has 3 seconds lifetime
+                    long atExpiresIn = Long.valueOf(decryptedStringParts[1]);
+                    assertTrue(atExpiresIn >= 2 && atExpiresIn <= 4);
+                } else {
+                    // For ID and referh tokens it is only a token
+                    assertEquals(1, decryptedStringParts.length);
+                }
+                token = decryptedStringParts[0];
                 tokenParts = token.split("\\.");
             } catch (Exception ex) {
                 fail("Token decryption has failed");

Original file line number	Diff line number	Diff line change
`@@ -91,7 +91,6 @@ public Long getAccessTokenExpiresIn() {`
`91`	`91`	`/**`
`92`	`92`	`* Set the access token expires_in value in seconds.`
`93`	`93`	`* It is relative to the time the access token is issued at.`
`94`		`- * This property is only checked when an authorization code flow grant completes and does not have to be persisted..`
`95`	`94`	`*`
`96`	`95`	`* @param accessTokenExpiresIn access token expires_in value in seconds.`
`97`	`96`	`*/`