Populate UI.getCurrent() before running ErrorHandler (#18127)

* chore: add test testing for correct UI.current instance in ErrorHandler * fix: Sets the current UI to this before calling ErrorHandler. Fixes #10533 * chore: star import * chore: star import * refactoring: use CurrentInstace to preserve current UI
vaadin · Nov 28, 2023 · da7ffab · da7ffab
1 parent 3d5faef
commit da7ffab
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 2 deletions.
diff --git a/flow-server/src/main/java/com/vaadin/flow/component/UI.java b/flow-server/src/main/java/com/vaadin/flow/component/UI.java
@@ -571,8 +571,14 @@ public void handleError(Exception exception) {
                         ErrorHandlingCommand errorHandlingCommand = (ErrorHandlingCommand) command;
                         errorHandlingCommand.handleError(exception);
                     } else if (getSession() != null) {
-                        getSession().getErrorHandler()
-                                .error(new ErrorEvent(exception));
+                        final Map<Class<?>, CurrentInstance> map = CurrentInstance
+                                .setCurrent(UI.this);
+                        try {
+                            getSession().getErrorHandler()
+                                    .error(new ErrorEvent(exception));
+                        } finally {
+                            CurrentInstance.restoreInstances(map);
+                        }
                     } else {
                         /*
                          * The session has expired after `ui.access` was called.

diff --git a/flow-server/src/test/java/com/vaadin/flow/component/UITest.java b/flow-server/src/test/java/com/vaadin/flow/component/UITest.java
@@ -27,6 +27,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 
@@ -78,6 +79,7 @@
 import com.vaadin.flow.router.internal.AfterNavigationHandler;
 import com.vaadin.flow.router.internal.BeforeEnterHandler;
 import com.vaadin.flow.router.internal.BeforeLeaveHandler;
+import com.vaadin.flow.server.ErrorHandler;
 import com.vaadin.flow.server.InvalidRouteConfigurationException;
 import com.vaadin.flow.server.MockVaadinContext;
 import com.vaadin.flow.server.MockVaadinServletService;
@@ -608,6 +610,27 @@ public void unsetSession_accessErrorHandlerStillWorks() throws IOException {
                 logOutputNoDebug.contains("UIDetachedException"));
     }
 
+    @Test
+    public void access_currentUIFilledInErrorHandler() {
+        UI ui = createTestUI();
+        initUI(ui, "", null);
+        final AtomicReference<UI> uiInErrorHandler = new AtomicReference<>();
+        final AtomicBoolean errorHandlerCalled = new AtomicBoolean();
+        ui.getSession().setErrorHandler((ErrorHandler) event -> {
+            errorHandlerCalled.set(true);
+            uiInErrorHandler.set(UI.getCurrent());
+        });
+        ui.access(() -> {
+            throw new RuntimeException("Simulated");
+        });
+
+        // Unlock to run pending access tasks
+        ui.getSession().unlock();
+
+        Assert.assertTrue(errorHandlerCalled.get());
+        Assert.assertEquals(ui, uiInErrorHandler.get());
+    }
+
     @Test
     public void beforeClientResponse_regularOrder() {
         UI ui = createTestUI();