fix: check the component belongs to UI before task execution (#11726) (…

…#11788) Related-to #11599 Co-authored-by: Mikhail Shabarov <[email protected]>
vaadin · Sep 9, 2021 · 8768346 · 8768346
1 parent a500d47
commit 8768346
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 1 deletion.
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
@@ -995,9 +995,12 @@ public Router getRouter() {
      *
      * @return a registration that can be used to cancel the execution of the
      *         task
+     * @throws IllegalArgumentException
+     *             if the given component doesn't belong to this UI
      */
     public ExecutionRegistration beforeClientResponse(Component component,
-            SerializableConsumer<ExecutionContext> execution) {
+            SerializableConsumer<ExecutionContext> execution)
+            throws IllegalArgumentException {
 
         if (component == null) {
             throw new IllegalArgumentException(
@@ -1008,6 +1011,11 @@ public ExecutionRegistration beforeClientResponse(Component component,
                     "The 'execution' parameter may not be null");
         }
 
+        if (component.getUI().isPresent() && component.getUI().get() != this) {
+            throw new IllegalArgumentException(
+                    "The given component doesn't belong to the UI the task to be executed on");
+        }
+
         return internals.getStateTree().beforeClientResponse(
                 component.getElement().getNode(), execution);
     }

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
@@ -634,6 +634,31 @@ public void beforeClientResponse_withReattachedNodes() {
                 callCounter.get());
     }
 
+    @Test
+    public void beforeClientResponse_componentNotAttachedToUi_noException() {
+        UI ui = createTestUI();
+        Component component = new AttachableComponent();
+        ui.beforeClientResponse(component, context -> {
+        });
+    }
+
+    @Test()
+    public void beforeClientResponse_componentBelongsToAnotherUI_throws() {
+        UI firstUI = createTestUI();
+        UI anotherUI = createTestUI();
+        Component component = new AttachableComponent();
+        anotherUI.add(component);
+
+        IllegalArgumentException exception = Assert.assertThrows(
+                IllegalArgumentException.class,
+                () -> firstUI.beforeClientResponse(component, context -> {
+                }));
+
+        Assert.assertEquals(
+                "The given component doesn't belong to the UI the task to be executed on",
+                exception.getMessage());
+    }
+
     @ListenerPriority(5)
     private static class BeforeEnterListenerFirst
             implements BeforeEnterListener {