This is an automated email from the ASF dual-hosted git repository.

rzo1 pushed a commit to branch tomee-10.x
in repository https://gitbox.apache.org/repos/asf/tomee.git

commit 8875e362601d9a1e42cfe98989bd913b4904a122
Author: Richard Zowalla <[email protected]>
AuthorDate: Tue Oct 13 08:49:18 2026 +0200

    TOMEE-4729 - Restore MP-JWT security state for asynchronous requests (#3080)
    
    State pushed by lazy token validation after the request went async was
    never restored; restore it on the pushing thread when the request ends.
    
    (cherry picked from commit 907f4bb6e4481cb329dd6d588779e8898fd4810f)
---
 .../apache/tomee/microprofile/jwt/MPJWTFilter.java |  85 +++++--
 .../jwt/MPJWTFilterExitWebAppTest.java             | 260 ++++++++++++++++++++-
 2 files changed, 326 insertions(+), 19 deletions(-)

diff --git 
a/mp-jwt/src/main/java/org/apache/tomee/microprofile/jwt/MPJWTFilter.java 
b/mp-jwt/src/main/java/org/apache/tomee/microprofile/jwt/MPJWTFilter.java
index 170bcabcef..4b679e2288 100644
--- a/mp-jwt/src/main/java/org/apache/tomee/microprofile/jwt/MPJWTFilter.java
+++ b/mp-jwt/src/main/java/org/apache/tomee/microprofile/jwt/MPJWTFilter.java
@@ -19,6 +19,8 @@ package org.apache.tomee.microprofile.jwt;
 import jakarta.enterprise.inject.Instance;
 import jakarta.enterprise.inject.spi.DeploymentException;
 import jakarta.inject.Inject;
+import jakarta.servlet.AsyncEvent;
+import jakarta.servlet.AsyncListener;
 import jakarta.servlet.Filter;
 import jakarta.servlet.FilterChain;
 import jakarta.servlet.FilterConfig;
@@ -69,6 +71,7 @@ import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.Callable;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Function;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
@@ -117,21 +120,80 @@ public class MPJWTFilter implements Filter {
             // the current thread; restore it on every exit path so a failing 
request does not
             // leak its security context to the next request served by this 
pooled thread
             exitWebApp(request);
+
+            // the token may still be validated lazily once the request went 
async,
+            // so restore that state when the async request is over
+            if (request.isAsyncStarted()) {
+                try {
+                    request.getAsyncContext().addListener(new 
AsyncExitWebAppListener(request));
+                } catch (final IllegalStateException e) {
+                    // the async request already completed concurrently
+                }
+            }
         }
     }
 
     private static void exitWebApp(final ServletRequest request) {
         final Object state = request.getAttribute(PRE_LOGIN_STATE);
-        if (state == null) {
-            return;
+        if (state instanceof PreLoginState preLoginState && 
preLoginState.restore()) {
+            request.removeAttribute(PRE_LOGIN_STATE);
+        }
+    }
+
+    /**
+     * Security state pushed by the token validation. {@link 
TomcatSecurityService} keeps the
+     * identity and the run-as stack in thread locals, so it can only be 
restored on the thread
+     * which pushed it, and at most once.
+     */
+    static final class PreLoginState {
+
+        private final TomcatSecurityService securityService;
+        private final Object state;
+        private final Thread thread = Thread.currentThread();
+        private final AtomicBoolean restored = new AtomicBoolean();
+
+        PreLoginState(final TomcatSecurityService securityService, final 
Object state) {
+            this.securityService = securityService;
+            this.state = state;
+        }
+
+        boolean restore() {
+            if (thread != Thread.currentThread() || 
!restored.compareAndSet(false, true)) {
+                return false;
+            }
+
+            securityService.exitWebApp(state);
+            return true;
+        }
+    }
+
+    private static final class AsyncExitWebAppListener implements 
AsyncListener {
+
+        private final ServletRequest request;
+
+        private AsyncExitWebAppListener(final ServletRequest request) {
+            this.request = request;
         }
 
-        // remove it first so the state is never restored twice for the same 
request
-        request.removeAttribute(PRE_LOGIN_STATE);
+        @Override
+        public void onComplete(final AsyncEvent event) {
+            exitWebApp(request);
+        }
 
-        final SecurityService securityService = 
SystemInstance.get().getComponent(SecurityService.class);
-        if (TomcatSecurityService.class.isInstance(securityService)) {
-            
TomcatSecurityService.class.cast(securityService).exitWebApp(state);
+        @Override
+        public void onTimeout(final AsyncEvent event) {
+            exitWebApp(request);
+        }
+
+        @Override
+        public void onError(final AsyncEvent event) {
+            exitWebApp(request);
+        }
+
+        @Override
+        public void onStartAsync(final AsyncEvent event) {
+            // listeners are dropped when the request goes async again
+            event.getAsyncContext().addListener(this);
         }
     }
 
@@ -369,19 +431,16 @@ public class MPJWTFilter implements Filter {
                 throw new InvalidTokenException(token, e);
             }
 
-            // TODO - do the login here, save the state to the request so we 
can recover it later.
-
             final SecurityService securityService = 
SystemInstance.get().getComponent(SecurityService.class);
             if (TomcatSecurityService.class.isInstance(securityService)) {
                 TomcatSecurityService tomcatSecurityService = 
TomcatSecurityService.class.cast(securityService);
                 final org.apache.catalina.connector.Request req = 
OpenEJBSecurityListener.requests.get();
-                Object state = 
tomcatSecurityService.enterWebApp(req.getWrapper().getRealm(), jsonWebToken, 
req.getWrapper().getRunAs());
+                final Object state = 
tomcatSecurityService.enterWebApp(req.getWrapper().getRealm(), jsonWebToken, 
req.getWrapper().getRunAs());
 
-                request.setAttribute(PRE_LOGIN_STATE, state);
+                // restored by the filter, or by its async listener for async 
requests
+                request.setAttribute(PRE_LOGIN_STATE, new 
PreLoginState(tomcatSecurityService, state));
             }
 
-            // TODO Also check if it is an async request and add a listener to 
close off the state
-
             return jsonWebToken;
 
         }
diff --git 
a/mp-jwt/src/test/java/org/apache/tomee/microprofile/jwt/MPJWTFilterExitWebAppTest.java
 
b/mp-jwt/src/test/java/org/apache/tomee/microprofile/jwt/MPJWTFilterExitWebAppTest.java
index 675372e1cb..28dab8844b 100644
--- 
a/mp-jwt/src/test/java/org/apache/tomee/microprofile/jwt/MPJWTFilterExitWebAppTest.java
+++ 
b/mp-jwt/src/test/java/org/apache/tomee/microprofile/jwt/MPJWTFilterExitWebAppTest.java
@@ -17,8 +17,14 @@
 package org.apache.tomee.microprofile.jwt;
 
 import jakarta.enterprise.inject.Instance;
+import jakarta.servlet.AsyncContext;
+import jakarta.servlet.AsyncEvent;
+import jakarta.servlet.AsyncListener;
 import jakarta.servlet.FilterChain;
+import jakarta.servlet.ServletContext;
 import jakarta.servlet.ServletException;
+import jakarta.servlet.ServletRequest;
+import jakarta.servlet.ServletResponse;
 import jakarta.servlet.http.HttpServletRequest;
 import jakarta.servlet.http.HttpServletResponse;
 import org.apache.catalina.core.StandardServer;
@@ -34,12 +40,18 @@ import org.junit.Test;
 import java.lang.reflect.Field;
 import java.lang.reflect.Proxy;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.Callable;
+import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertTrue;
@@ -48,6 +60,7 @@ import static org.junit.Assert.fail;
 public class MPJWTFilterExitWebAppTest {
 
     private RecordingSecurityService securityService;
+    private ExecutorService asyncThread;
 
     @Before
     public void setUp() {
@@ -55,10 +68,12 @@ public class MPJWTFilterExitWebAppTest {
         TomcatHelper.setServer(new StandardServer());
         securityService = new RecordingSecurityService();
         SystemInstance.get().setComponent(SecurityService.class, 
securityService);
+        asyncThread = Executors.newSingleThreadExecutor();
     }
 
     @After
     public void tearDown() {
+        asyncThread.shutdownNow();
         securityService.clearRunAsStack();
         SystemInstance.reset();
     }
@@ -114,10 +129,133 @@ public class MPJWTFilterExitWebAppTest {
         assertExited(attributes);
     }
 
-    private void enterWebApp(final jakarta.servlet.ServletRequest req) {
+    @Test
+    public void asyncRequestValidatedOnFilterThreadIsRestoredOnce() throws 
Exception {
+        final Map<String, Object> attributes = new HashMap<>();
+        final FakeAsyncContext asyncContext = new FakeAsyncContext();
+        final HttpServletRequest request = request(attributes, asyncContext);
+
+        filter().doFilter(request, response(), (req, res) -> enterWebApp(req));
+
+        assertExited(attributes);
+        assertEquals(1, asyncContext.listeners.size());
+
+        asyncContext.fireOnComplete();
+        assertEquals(1, securityService.exited.size());
+    }
+
+    @Test
+    public void asyncRequestValidatedLaterIsRestoredOnComplete() throws 
Exception {
+        final Map<String, Object> attributes = new HashMap<>();
+        final FakeAsyncContext asyncContext = new FakeAsyncContext();
+        final HttpServletRequest request = request(attributes, asyncContext);
+
+        filter().doFilter(request, response(), (req, res) -> {
+        });
+
+        assertTrue(securityService.exited.isEmpty());
+        assertEquals(1, asyncContext.listeners.size());
+
+        // the token is validated lazily on the thread handling the async 
request
+        assertEquals(Integer.valueOf(0), onAsyncThread(() -> {
+            enterWebApp(request);
+            asyncContext.fireOnComplete();
+            return securityService.runAsStackSize();
+        }));
+
+        assertEquals(1, securityService.exited.size());
+        assertNull(attributes.get(MPJWTFilter.PRE_LOGIN_STATE));
+    }
+
+    @Test
+    public void asyncRequestValidatedLaterIsRestoredOnErrorAndTimeout() throws 
Exception {
+        for (final String event : new String[]{"error", "timeout"}) {
+            securityService.exited.clear();
+
+            final Map<String, Object> attributes = new HashMap<>();
+            final FakeAsyncContext asyncContext = new FakeAsyncContext();
+            final HttpServletRequest request = request(attributes, 
asyncContext);
+
+            filter().doFilter(request, response(), (req, res) -> {
+            });
+
+            assertEquals(Integer.valueOf(0), onAsyncThread(() -> {
+                enterWebApp(request);
+                if ("error".equals(event)) {
+                    asyncContext.fireOnError();
+                } else {
+                    asyncContext.fireOnTimeout();
+                }
+                return securityService.runAsStackSize();
+            }));
+
+            assertEquals(event, 1, securityService.exited.size());
+            assertNull(event, attributes.get(MPJWTFilter.PRE_LOGIN_STATE));
+        }
+    }
+
+    @Test
+    public void stateIsNotRestoredOnAnotherThread() throws Exception {
+        final Map<String, Object> attributes = new HashMap<>();
+        final FakeAsyncContext asyncContext = new FakeAsyncContext();
+        final HttpServletRequest request = request(attributes, asyncContext);
+
+        // validated on another thread while the filter is still running
+        filter().doFilter(request, response(), (req, res) -> {
+            try {
+                onAsyncThread(() -> {
+                    enterWebApp(req);
+                    return null;
+                });
+            } catch (final Exception e) {
+                throw new ServletException(e);
+            }
+        });
+
+        assertTrue(securityService.exited.isEmpty());
+        assertNotNull(attributes.get(MPJWTFilter.PRE_LOGIN_STATE));
+
+        // completion notified on a thread which did not push the state
+        asyncContext.fireOnComplete();
+        assertTrue(securityService.exited.isEmpty());
+        assertEquals(0, securityService.runAsStackSize());
+
+        assertEquals(Integer.valueOf(0), onAsyncThread(() -> {
+            asyncContext.fireOnComplete();
+            return securityService.runAsStackSize();
+        }));
+        assertEquals(1, securityService.exited.size());
+        assertNull(attributes.get(MPJWTFilter.PRE_LOGIN_STATE));
+    }
+
+    @Test
+    public void listenerIsKeptWhenAsyncIsStartedAgain() throws Exception {
+        final Map<String, Object> attributes = new HashMap<>();
+        final FakeAsyncContext asyncContext = new FakeAsyncContext();
+        final HttpServletRequest request = request(attributes, asyncContext);
+
+        filter().doFilter(request, response(), (req, res) -> {
+        });
+
+        asyncContext.fireOnStartAsync();
+        assertEquals(1, asyncContext.listeners.size());
+
+        assertEquals(Integer.valueOf(0), onAsyncThread(() -> {
+            enterWebApp(request);
+            asyncContext.fireOnComplete();
+            return securityService.runAsStackSize();
+        }));
+        assertEquals(1, securityService.exited.size());
+    }
+
+    private <T> T onAsyncThread(final Callable<T> task) throws Exception {
+        return asyncThread.submit(task).get();
+    }
+
+    private void enterWebApp(final ServletRequest req) {
         // simulates what ValidateJSonWebToken does once the token has been 
validated
         final Object state = securityService.enterWebApp(null, null, "admin");
-        req.setAttribute(MPJWTFilter.PRE_LOGIN_STATE, state);
+        req.setAttribute(MPJWTFilter.PRE_LOGIN_STATE, new 
MPJWTFilter.PreLoginState(securityService, state));
         assertEquals(1, securityService.runAsStackSize());
     }
 
@@ -147,12 +285,24 @@ public class MPJWTFilterExitWebAppTest {
     }
 
     private static HttpServletRequest request(final Map<String, Object> 
attributes) {
+        return request(attributes, null);
+    }
+
+    private static HttpServletRequest request(final Map<String, Object> 
attributes, final AsyncContext asyncContext) {
+        final Map<String, Object> synchronizedAttributes = 
Collections.synchronizedMap(attributes);
         return (HttpServletRequest) Proxy.newProxyInstance(
                 MPJWTFilterExitWebAppTest.class.getClassLoader(), new 
Class<?>[]{HttpServletRequest.class},
                 (proxy, method, args) -> switch (method.getName()) {
-                    case "getAttribute" -> attributes.get((String) args[0]);
-                    case "setAttribute" -> attributes.put((String) args[0], 
args[1]);
-                    case "removeAttribute" -> attributes.remove((String) 
args[0]);
+                    case "getAttribute" -> synchronizedAttributes.get((String) 
args[0]);
+                    case "setAttribute" -> synchronizedAttributes.put((String) 
args[0], args[1]);
+                    case "removeAttribute" -> 
synchronizedAttributes.remove((String) args[0]);
+                    case "isAsyncStarted" -> asyncContext != null;
+                    case "getAsyncContext" -> {
+                        if (asyncContext == null) {
+                            throw new IllegalStateException("not async");
+                        }
+                        yield asyncContext;
+                    }
                     default -> throw new 
UnsupportedOperationException(method.getName());
                 });
     }
@@ -165,9 +315,107 @@ public class MPJWTFilterExitWebAppTest {
                 });
     }
 
+    private static class FakeAsyncContext implements AsyncContext {
+
+        private final List<AsyncListener> listeners = new ArrayList<>();
+
+        void fireOnComplete() throws Exception {
+            for (final AsyncListener listener : new ArrayList<>(listeners)) {
+                listener.onComplete(new AsyncEvent(this));
+            }
+        }
+
+        void fireOnError() throws Exception {
+            for (final AsyncListener listener : new ArrayList<>(listeners)) {
+                listener.onError(new AsyncEvent(this));
+            }
+        }
+
+        void fireOnTimeout() throws Exception {
+            for (final AsyncListener listener : new ArrayList<>(listeners)) {
+                listener.onTimeout(new AsyncEvent(this));
+            }
+        }
+
+        void fireOnStartAsync() throws Exception {
+            // the container drops the listeners before notifying them
+            final List<AsyncListener> current = new ArrayList<>(listeners);
+            listeners.clear();
+            for (final AsyncListener listener : current) {
+                listener.onStartAsync(new AsyncEvent(this));
+            }
+        }
+
+        @Override
+        public void addListener(final AsyncListener listener) {
+            listeners.add(listener);
+        }
+
+        @Override
+        public void addListener(final AsyncListener listener, final 
ServletRequest request,
+                                final ServletResponse response) {
+            listeners.add(listener);
+        }
+
+        @Override
+        public ServletRequest getRequest() {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public ServletResponse getResponse() {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public boolean hasOriginalRequestAndResponse() {
+            return true;
+        }
+
+        @Override
+        public void dispatch() {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public void dispatch(final String path) {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public void dispatch(final ServletContext context, final String path) {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public void complete() {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public void start(final Runnable run) {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public <T extends AsyncListener> T createListener(final Class<T> 
clazz) {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public void setTimeout(final long timeout) {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public long getTimeout() {
+            return 0;
+        }
+    }
+
     private static class RecordingSecurityService extends 
TomcatSecurityService {
 
-        private final List<Object> exited = new ArrayList<>();
+        private final List<Object> exited = new CopyOnWriteArrayList<>();
 
         @Override
         public void exitWebApp(final Object state) {

Reply via email to