Fix OpenID storage selection per mechanism

Author: jungm

tomee/tomee-security/src/main/java/org/apache/tomee/security/cdi/OpenIdAuthenticationMechanism.java

@@ -113,6 +113,14 @@ private String redirectPath() {
 
     @Override
     public void cleanSubject(HttpServletRequest request, HttpServletResponse response, HttpMessageContext httpMessageContext) {
+        OpenIdStorageHandler.withDefinition(getDefinition(), () -> {
+            cleanSubjectWithSelectedDefinition(request, response, httpMessageContext);
+            return null;
+        });
+    }
+
+    private void cleanSubjectWithSelectedDefinition(final HttpServletRequest request, final HttpServletResponse response,
+                                                    final HttpMessageContext httpMessageContext) {
         String redirectTarget = buildRedirectUri();
 
         HttpSession session = request.getSession(false);
@@ -153,6 +161,13 @@ private String buildRedirectUri()
 
     @Override
     public AuthenticationStatus validateRequest(HttpServletRequest request, HttpServletResponse response, HttpMessageContext httpMessageContext) throws AuthenticationException {
+        return OpenIdStorageHandler.withDefinition(getDefinition(),
+                () -> validateRequestWithSelectedDefinition(request, response, httpMessageContext));
+    }
+
+    private AuthenticationStatus validateRequestWithSelectedDefinition(final HttpServletRequest request,
+                                                                       final HttpServletResponse response,
+                                                                       final HttpMessageContext httpMessageContext) {
         if (request.getUserPrincipal() != null) {
             AuthenticationStatus result = handleExpiredTokens(request, response, httpMessageContext);
             if (result != null) {

tomee/tomee-security/src/main/java/org/apache/tomee/security/cdi/TomEESecurityExtension.java

@@ -27,8 +27,7 @@
 import org.apache.tomee.security.cdi.openid.OpenIdIdentityStore;
 import org.apache.tomee.security.cdi.openid.TomEEOpenIdContext;
 import org.apache.tomee.security.cdi.openid.storage.OpenIdStorageHandler;
-import org.apache.tomee.security.cdi.openid.storage.impl.CookieBasedOpenIdStorageHandler;
-import org.apache.tomee.security.cdi.openid.storage.impl.SessionBasedOpenIdStorageHandler;
+import org.apache.tomee.security.cdi.openid.storage.impl.DefinitionAwareOpenIdStorageHandler;
 import org.apache.tomee.security.http.openid.OpenIdAuthenticationMechanismDefinitionDelegate;
 import org.apache.tomee.security.identitystore.TomEEDatabaseIdentityStore;
 import org.apache.tomee.security.identitystore.TomEEDefaultIdentityStore;
@@ -420,9 +419,7 @@ void registerAuthenticationMechanism(
                         OpenIdAuthenticationMechanismDefinition definition = (OpenIdAuthenticationMechanismDefinition)
                                 beanManager.getReference(definitionBean, OpenIdAuthenticationMechanismDefinition.class, creationalContext);
 
-                        return definition.useSession()
-                                ? new SessionBasedOpenIdStorageHandler()
-                                : new CookieBasedOpenIdStorageHandler();
+                        return new DefinitionAwareOpenIdStorageHandler();
                     });
 
             afterBeanDiscovery.addBean(createBean(TomEEOpenIdContext.class, beanManager));

tomee/tomee-security/src/main/java/org/apache/tomee/security/cdi/openid/storage/OpenIdStorageHandler.java

@@ -18,15 +18,40 @@
 
 import org.apache.commons.lang3.RandomStringUtils;
 
+import jakarta.security.enterprise.authentication.mechanism.http.OpenIdAuthenticationMechanismDefinition;
 import jakarta.servlet.http.HttpServletRequest;
 import jakarta.servlet.http.HttpServletResponse;
+import java.util.function.Supplier;
 
 public abstract class OpenIdStorageHandler {
     protected static final String PREFIX = "openid.";
 
     public static final String REQUEST_KEY = "REQUEST";
     public static final String STATE_KEY = "STATE";
     public static final String NONCE_KEY = "NONCE";
+    private static final ThreadLocal<OpenIdAuthenticationMechanismDefinition> CURRENT_DEFINITION = new ThreadLocal<>();
+
+    public static <T> T withDefinition(final OpenIdAuthenticationMechanismDefinition definition, final Supplier<T> task) {
+        final OpenIdAuthenticationMechanismDefinition previous = CURRENT_DEFINITION.get();
+        if (definition == null) {
+            CURRENT_DEFINITION.remove();
+        } else {
+            CURRENT_DEFINITION.set(definition);
+        }
+        try {
+            return task.get();
+        } finally {
+            if (previous == null) {
+                CURRENT_DEFINITION.remove();
+            } else {
+                CURRENT_DEFINITION.set(previous);
+            }
+        }
+    }
+
+    protected static OpenIdAuthenticationMechanismDefinition currentDefinition() {
+        return CURRENT_DEFINITION.get();
+    }
 
     public abstract String get(HttpServletRequest request, HttpServletResponse response, String key);
 

tomee/tomee-security/src/main/java/org/apache/tomee/security/cdi/openid/storage/impl/DefinitionAwareOpenIdStorageHandler.java

@@ -0,0 +1,50 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.tomee.security.cdi.openid.storage.impl;
+
+import jakarta.security.enterprise.authentication.mechanism.http.OpenIdAuthenticationMechanismDefinition;
+import jakarta.servlet.http.HttpServletRequest;
+import jakarta.servlet.http.HttpServletResponse;
+import org.apache.tomee.security.cdi.openid.storage.OpenIdStorageHandler;
+
+public class DefinitionAwareOpenIdStorageHandler extends OpenIdStorageHandler {
+    private final OpenIdStorageHandler session = new SessionBasedOpenIdStorageHandler();
+    private final OpenIdStorageHandler cookie = new CookieBasedOpenIdStorageHandler();
+
+    @Override
+    public String get(final HttpServletRequest request, final HttpServletResponse response, final String key) {
+        return delegate().get(request, response, key);
+    }
+
+    @Override
+    public void set(final HttpServletRequest request, final HttpServletResponse response, final String key, final String value) {
+        delegate().set(request, response, key, value);
+    }
+
+    @Override
+    public void delete(final HttpServletRequest request, final HttpServletResponse response, final String key) {
+        delegate().delete(request, response, key);
+    }
+
+    private OpenIdStorageHandler delegate() {
+        final OpenIdAuthenticationMechanismDefinition definition = currentDefinition();
+        if (definition == null) {
+            throw new IllegalStateException("OpenID storage access requires the selected OpenIdAuthenticationMechanismDefinition");
+        }
+        return definition.useSession() ? session : cookie;
+    }
+}

tomee/tomee-security/src/test/java/org/apache/tomee/security/cdi/openid/storage/impl/DefinitionAwareOpenIdStorageHandlerTest.java

@@ -0,0 +1,83 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.tomee.security.cdi.openid.storage.impl;
+
+import jakarta.security.enterprise.authentication.mechanism.http.OpenIdAuthenticationMechanismDefinition;
+import jakarta.servlet.http.Cookie;
+import jakarta.servlet.http.HttpServletRequest;
+import jakarta.servlet.http.HttpServletResponse;
+import jakarta.servlet.http.HttpSession;
+import org.apache.tomee.security.cdi.openid.storage.OpenIdStorageHandler;
+import org.junit.Test;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+public class DefinitionAwareOpenIdStorageHandlerTest {
+    @Test
+    public void usesCurrentDefinitionSessionStorage() {
+        final DefinitionAwareOpenIdStorageHandler handler = new DefinitionAwareOpenIdStorageHandler();
+        final OpenIdAuthenticationMechanismDefinition definition = mock(OpenIdAuthenticationMechanismDefinition.class);
+        final HttpServletRequest request = mock(HttpServletRequest.class);
+        final HttpServletResponse response = mock(HttpServletResponse.class);
+        final HttpSession session = mock(HttpSession.class);
+
+        when(definition.useSession()).thenReturn(true);
+        when(request.getSession()).thenReturn(session);
+
+        OpenIdStorageHandler.withDefinition(definition, () -> {
+            handler.set(request, response, OpenIdStorageHandler.STATE_KEY, "state");
+            return null;
+        });
+
+        verify(session).setAttribute("openid.STATE", "state");
+        verify(response, never()).addCookie(any(Cookie.class));
+    }
+
+    @Test
+    public void usesCurrentDefinitionCookieStorage() {
+        final DefinitionAwareOpenIdStorageHandler handler = new DefinitionAwareOpenIdStorageHandler();
+        final OpenIdAuthenticationMechanismDefinition definition = mock(OpenIdAuthenticationMechanismDefinition.class);
+        final HttpServletRequest request = mock(HttpServletRequest.class);
+        final HttpServletResponse response = mock(HttpServletResponse.class);
+        final HttpSession session = mock(HttpSession.class);
+
+        when(definition.useSession()).thenReturn(false);
+        when(request.getSession()).thenReturn(session);
+
+        OpenIdStorageHandler.withDefinition(definition, () -> {
+            handler.set(request, response, OpenIdStorageHandler.STATE_KEY, "state");
+            return null;
+        });
+
+        verify(response).addCookie(any(Cookie.class));
+        verify(session, never()).setAttribute(eq("openid.STATE"), eq("state"));
+    }
+
+    @Test(expected = IllegalStateException.class)
+    public void failsWithoutSelectedDefinition() {
+        final DefinitionAwareOpenIdStorageHandler handler = new DefinitionAwareOpenIdStorageHandler();
+        final HttpServletRequest request = mock(HttpServletRequest.class);
+        final HttpServletResponse response = mock(HttpServletResponse.class);
+
+        handler.set(request, response, OpenIdStorageHandler.STATE_KEY, "state");
+    }
+}