Fix OIDC cookie related tenant id and chunk calculation issues

extensions/oidc/deployment/src/test/java/io/quarkus/oidc/test/CodeFlowDevModeTestCase.java

@@ -15,6 +15,10 @@
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.SortedMap;
+import java.util.TreeMap;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 
@@ -26,13 +30,15 @@
 import org.junit.jupiter.api.TestMethodOrder;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
+import com.gargoylesoftware.htmlunit.CookieManager;
 import com.gargoylesoftware.htmlunit.FailingHttpStatusCodeException;
 import com.gargoylesoftware.htmlunit.SilentCssErrorHandler;
 import com.gargoylesoftware.htmlunit.WebClient;
 import com.gargoylesoftware.htmlunit.WebRequest;
 import com.gargoylesoftware.htmlunit.WebResponse;
 import com.gargoylesoftware.htmlunit.html.HtmlForm;
 import com.gargoylesoftware.htmlunit.html.HtmlPage;
+import com.gargoylesoftware.htmlunit.util.Cookie;
 
 import io.quarkus.test.QuarkusDevModeTest;
 import io.quarkus.test.common.QuarkusTestResource;
@@ -114,7 +120,12 @@ public void testAccessAndRefreshTokenInjectionDevMode() throws IOException, Inte
 
             assertEquals("alice", page.getBody().asNormalizedText());
 
-            assertEquals("custom", page.getWebClient().getCookieManager().getCookie("q_session").getValue().split("\\|")[1]);
+            List<Cookie> sessionCookies = getSessionCookies(webClient, null);
+            assertEquals(2, sessionCookies.size());
+            assertEquals("q_session_chunk_1", sessionCookies.get(0).getName());
+            assertEquals("q_session_chunk_2", sessionCookies.get(1).getName());
+            String sessionCookieValue = sessionCookies.get(0).getValue() + sessionCookies.get(1).getValue();
+            assertEquals("custom", sessionCookieValue.split("\\|")[1]);
 
             webClient.getOptions().setRedirectEnabled(false);
             WebResponse webResponse = webClient
@@ -217,4 +228,16 @@ public void run() throws Throwable {
         assertTrue(checkPassed.get(), "Can not confirm Secret key for encrypting state cookie has been generated");
     }
 
+    private List<Cookie> getSessionCookies(WebClient webClient, String tenantId) {
+        String sessionCookieNameChunk = "q_session" + (tenantId == null ? "" : "_" + tenantId) + "_chunk_";
+        CookieManager cookieManager = webClient.getCookieManager();
+        SortedMap<String, Cookie> sessionCookies = new TreeMap<>();
+        for (Cookie cookie : cookieManager.getCookies()) {
+            if (cookie.getName().startsWith(sessionCookieNameChunk)) {
+                sessionCookies.put(cookie.getName(), cookie);
+            }
+        }
+
+        return sessionCookies.isEmpty() ? null : new ArrayList<Cookie>(sessionCookies.values());
+    }
 }

extensions/oidc/deployment/src/test/java/io/quarkus/oidc/test/CodeTenantReauthenticateTestCase.java

@@ -0,0 +1,173 @@
+package io.quarkus.oidc.test;
+
+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 org.junit.jupiter.api.MethodOrderer;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.TestMethodOrder;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+import com.gargoylesoftware.htmlunit.SilentCssErrorHandler;
+import com.gargoylesoftware.htmlunit.TextPage;
+import com.gargoylesoftware.htmlunit.WebClient;
+import com.gargoylesoftware.htmlunit.html.HtmlForm;
+import com.gargoylesoftware.htmlunit.html.HtmlPage;
+import com.gargoylesoftware.htmlunit.util.Cookie;
+
+import io.quarkus.test.QuarkusUnitTest;
+import io.quarkus.test.common.QuarkusTestResource;
+import io.quarkus.test.keycloak.server.KeycloakTestResourceLifecycleManager;
+
+@TestMethodOrder(MethodOrderer.OrderAnnotation.class)
+@QuarkusTestResource(KeycloakTestResourceLifecycleManager.class)
+public class CodeTenantReauthenticateTestCase {
+    private static Class<?>[] testClasses = {
+            TenantReauthentication.class,
+            CustomTenantResolver.class,
+            CustomTenantConfigResolver.class
+    };
+
+    @RegisterExtension
+    static final QuarkusUnitTest test = new QuarkusUnitTest()
+            .withApplicationRoot((jar) -> jar
+                    .addClasses(testClasses)
+                    .addAsResource("application-tenant-reauthenticate.properties", "application.properties"));
+
+    @Test
+    public void testDefaultTenant() throws Exception {
+        try (final WebClient webClient = createWebClient()) {
+
+            callTenant(webClient, null, "/protected", "alice");
+
+            webClient.getCookieManager().clearCookies();
+        }
+    }
+
+    @Test
+    public void testTenantResolver() throws Exception {
+        try (final WebClient webClient = createWebClient()) {
+
+            callTenant(webClient, "tenant-resolver", "/protected/tenant/tenant-resolver", "tenant-resolver:alice");
+
+            webClient.getCookieManager().clearCookies();
+        }
+    }
+
+    @Test
+    public void testTenantConfigResolver() throws Exception {
+        try (final WebClient webClient = createWebClient()) {
+
+            callTenant(webClient, "tenant-config-resolver", "/protected/tenant/tenant-config-resolver",
+                    "tenant-config-resolver:alice");
+
+            webClient.getCookieManager().clearCookies();
+        }
+    }
+
+    @Test
+    public void testSwitchFromTenantResolverToDefaultTenant() throws Exception {
+        try (final WebClient webClient = createWebClient()) {
+
+            callTenant(webClient, "tenant-resolver", "/protected/tenant/tenant-resolver", "tenant-resolver:alice");
+            expectReauthentication(webClient, "/protected", "tenant-resolver");
+
+            webClient.getCookieManager().clearCookies();
+        }
+    }
+
+    @Test
+    public void testSwitchFromDefaultTenantToTenantResover() throws Exception {
+        try (final WebClient webClient = createWebClient()) {
+
+            callTenant(webClient, null, "/protected", "alice");
+            expectReauthentication(webClient, "/protected/tenant/tenant-resolver", null);
+
+            webClient.getCookieManager().clearCookies();
+        }
+    }
+
+    @Test
+    public void testSwitchFromTenantConfigResolverToDefaultTenant() throws Exception {
+        try (final WebClient webClient = createWebClient()) {
+
+            callTenant(webClient, "tenant-config-resolver", "/protected/tenant/tenant-config-resolver",
+                    "tenant-config-resolver:alice");
+            expectReauthentication(webClient, "/protected", "tenant-config-resolver");
+
+            webClient.getCookieManager().clearCookies();
+        }
+    }
+
+    @Test
+    public void testSwitchFromDefaultTenantToTenantConfigResolver() throws Exception {
+        try (final WebClient webClient = createWebClient()) {
+
+            callTenant(webClient, null, "/protected", "alice");
+            expectReauthentication(webClient, "/protected/tenant/tenant-config-resolver", null);
+
+            webClient.getCookieManager().clearCookies();
+        }
+    }
+
+    @Test
+    public void testSwitchFromTenantResolverToTenantConfigResolver() throws Exception {
+        try (final WebClient webClient = createWebClient()) {
+
+            callTenant(webClient, "tenant-resolver", "/protected/tenant/tenant-resolver", "tenant-resolver:alice");
+            expectReauthentication(webClient, "/protected/tenant/tenant-config-resolver", "tenant-resolver");
+
+            webClient.getCookieManager().clearCookies();
+        }
+    }
+
+    @Test
+    public void testSwitchFromTenantConfigResolverToTenantResolver() throws Exception {
+        try (final WebClient webClient = createWebClient()) {
+
+            callTenant(webClient, "tenant-config-resolver", "/protected/tenant/tenant-config-resolver",
+                    "tenant-config-resolver:alice");
+            expectReauthentication(webClient, "/protected/tenant/tenant-resolver", "tenant-config-resolver");
+
+            webClient.getCookieManager().clearCookies();
+        }
+    }
+
+    private static void callTenant(WebClient webClient, String tenant, String relativePath, String expectedResponse)
+            throws Exception {
+        HtmlPage page = webClient.getPage("http://localhost:8081" + relativePath);
+
+        assertEquals("Sign in to quarkus", page.getTitleText());
+
+        HtmlForm loginForm = page.getForms().get(0);
+
+        loginForm.getInputByName("username").setValueAttribute("alice");
+        loginForm.getInputByName("password").setValueAttribute("alice");
+
+        page = loginForm.getInputByName("login").click();
+
+        assertEquals(expectedResponse, page.getBody().asNormalizedText());
+        assertNotNull(getSessionCookie(webClient, tenant));
+    }
+
+    private static void expectReauthentication(WebClient webClient, String relativePath,
+            String oldTenant) throws Exception {
+        webClient.getOptions().setRedirectEnabled(false);
+        webClient.getOptions().setThrowExceptionOnFailingStatusCode(false);
+
+        TextPage textPage = webClient.getPage("http://localhost:8081" + relativePath);
+        assertEquals(302, textPage.getWebResponse().getStatusCode());
+        assertNull(getSessionCookie(webClient, oldTenant));
+    }
+
+    private static WebClient createWebClient() {
+        WebClient webClient = new WebClient();
+        webClient.setCssErrorHandler(new SilentCssErrorHandler());
+        return webClient;
+    }
+
+    private static Cookie getSessionCookie(WebClient webClient, String tenantId) {
+        return webClient.getCookieManager().getCookie("q_session" + (tenantId == null ? "" : "_" + tenantId));
+    }
+}

extensions/oidc/deployment/src/test/java/io/quarkus/oidc/test/CustomTenantConfigResolver.java

@@ -2,10 +2,13 @@
 
 import jakarta.enterprise.context.ApplicationScoped;
 
+import org.eclipse.microprofile.config.ConfigProvider;
+
 import io.quarkus.oidc.OidcRequestContext;
 import io.quarkus.oidc.OidcTenantConfig;
 import io.quarkus.oidc.OidcTenantConfig.ApplicationType;
 import io.quarkus.oidc.TenantConfigResolver;
+import io.quarkus.oidc.runtime.OidcUtils;
 import io.smallrye.mutiny.Uni;
 import io.vertx.ext.web.RoutingContext;
 
@@ -22,13 +25,14 @@ public Uni<OidcTenantConfig> resolve(RoutingContext context, OidcRequestContext<
             config.setApplicationType(ApplicationType.WEB_APP);
             return Uni.createFrom().item(config);
         }
+        context.remove(OidcUtils.TENANT_ID_ATTRIBUTE);
         if (context.request().path().endsWith("/null-tenant")) {
             return null;
         }
         return Uni.createFrom().nullItem();
     }
 
     private String getIssuerUrl() {
-        return System.getProperty("keycloak.url");
+        return ConfigProvider.getConfig().getValue("keycloak.url", String.class);
     }
 }

extensions/oidc/deployment/src/test/java/io/quarkus/oidc/test/CustomTenantResolver.java

@@ -0,0 +1,19 @@
+package io.quarkus.oidc.test;
+
+import jakarta.enterprise.context.ApplicationScoped;
+
+import io.quarkus.oidc.TenantResolver;
+import io.quarkus.oidc.runtime.OidcUtils;
+import io.vertx.ext.web.RoutingContext;
+
+@ApplicationScoped
+public class CustomTenantResolver implements TenantResolver {
+    @Override
+    public String resolve(RoutingContext context) {
+        if (context.request().path().endsWith("/tenant-resolver")) {
+            return "tenant-resolver";
+        }
+        context.put(OidcUtils.TENANT_ID_ATTRIBUTE, OidcUtils.DEFAULT_TENANT_ID);
+        return null;
+    }
+}

extensions/oidc/deployment/src/test/java/io/quarkus/oidc/test/TenantReauthentication.java

@@ -0,0 +1,35 @@
+package io.quarkus.oidc.test;
+
+import jakarta.inject.Inject;
+import jakarta.ws.rs.GET;
+import jakarta.ws.rs.Path;
+import jakarta.ws.rs.PathParam;
+
+import org.eclipse.microprofile.jwt.JsonWebToken;
+
+import io.quarkus.oidc.IdToken;
+import io.quarkus.oidc.OidcSession;
+import io.quarkus.security.Authenticated;
+
+@Path("/protected")
+@Authenticated
+public class TenantReauthentication {
+
+    @Inject
+    @IdToken
+    JsonWebToken idToken;
+
+    @Inject
+    OidcSession session;
+
+    @GET
+    public String getName() {
+        return idToken.getName();
+    }
+
+    @GET
+    @Path("tenant/{id}")
+    public String getTenantName(@PathParam("id") String tenantId) {
+        return tenantId + ":" + idToken.getName();
+    }
+}

extensions/oidc/deployment/src/test/resources/application-tenant-reauthenticate.properties

@@ -0,0 +1,11 @@
+quarkus.oidc.auth-server-url=${keycloak.url}/realms/quarkus
+quarkus.oidc.client-id=quarkus-web-app
+quarkus.oidc.credentials.secret=secret
+quarkus.oidc.application-type=web-app
+
+quarkus.oidc.tenant-resolver.auth-server-url=${keycloak.url}/realms/quarkus
+quarkus.oidc.tenant-resolver.client-id=quarkus-web-app
+quarkus.oidc.tenant-resolver.credentials.secret=secret
+quarkus.oidc.tenant-resolver.application-type=web-app
+
+quarkus.log.category."com.gargoylesoftware.htmlunit".level=ERROR

extensions/oidc/runtime/src/main/java/io/quarkus/oidc/runtime/CodeAuthenticationMechanism.java

@@ -560,8 +560,23 @@ public Uni<ChallengeData> apply(TenantConfigContext tenantContext) {
     }
 
     public Uni<ChallengeData> getChallengeInternal(RoutingContext context, TenantConfigContext configContext) {
-        LOG.debug("Starting an authentication challenge");
-        return removeSessionCookie(context, configContext.oidcConfig)
+        LOG.debugf("Starting an authentication challenge for tenant %s", configContext.oidcConfig.tenantId.get());
+
+        OidcTenantConfig sessionCookieConfig = configContext.oidcConfig;
+        String sessionTenantIdSetByCookie = context.get(OidcUtils.TENANT_ID_SET_BY_SESSION_COOKIE);
+
+        if (sessionTenantIdSetByCookie != null
+                && !sessionTenantIdSetByCookie.equals(sessionCookieConfig.tenantId.orElse(OidcUtils.DEFAULT_TENANT_ID))) {
+            // New tenant id has been chosen during the tenant resolution process
+            // Get the already resolved configuration, avoiding calling the tenant resolvers
+            OidcTenantConfig previousTenantConfig = resolver.getResolvedConfig(sessionTenantIdSetByCookie);
+            if (previousTenantConfig != null) {
+                sessionCookieConfig = previousTenantConfig;
+                LOG.debugf("Removing the session cookie for the previous tenant id: %s", sessionTenantIdSetByCookie);
+                OidcUtils.getSessionCookie(context, sessionCookieConfig);
+            }
+        }
+        return removeSessionCookie(context, sessionCookieConfig)
                 .chain(new Function<Void, Uni<? extends ChallengeData>>() {
 
                     @Override
@@ -963,6 +978,8 @@ public Void apply(String cookieValue) {
                                                 String nextValue = cookieValue.substring(currentPos, nextValueUpperPos);
                                                 // q_session_session_chunk_1, etc
                                                 String nextName = sessionName + OidcUtils.SESSION_COOKIE_CHUNK + sessionIndex;
+                                                LOG.debugf("Creating the %s session cookie chunk, size: %d", nextName,
+                                                        nextValue.length());
                                                 createCookie(context, configContext.oidcConfig, nextName, nextValue,
                                                         sessionMaxAge, true);
                                                 currentPos = nextPos;

extensions/oidc/runtime/src/main/java/io/quarkus/oidc/runtime/DefaultTenantConfigResolver.java

@@ -377,4 +377,19 @@ private static ImmutablePathMatcher.ImmutablePathMatcherBuilder<String> addPath(
         }
     }
 
+    public OidcTenantConfig getResolvedConfig(String sessionTenantId) {
+        if (OidcUtils.DEFAULT_TENANT_ID.equals(sessionTenantId)) {
+            return tenantConfigBean.getDefaultTenant().getOidcTenantConfig();
+        }
+
+        if (tenantConfigBean.getStaticTenantsConfig().containsKey(sessionTenantId)) {
+            return tenantConfigBean.getStaticTenantsConfig().get(sessionTenantId).getOidcTenantConfig();
+        }
+
+        if (tenantConfigBean.getDynamicTenantsConfig().containsKey(sessionTenantId)) {
+            return tenantConfigBean.getDynamicTenantsConfig().get(sessionTenantId).getOidcTenantConfig();
+        }
+        return null;
+    }
+
 }