From 2423842fd57898277584192c8a25b819f2599eac Mon Sep 17 00:00:00 2001
From: Mikhail Shabarov <61410877+mshabarov@users.noreply.github.com>
Date: Tue, 7 Sep 2021 15:18:57 +0300
Subject: [PATCH] fix: check the component belongs to UI before task execution
 (#11726)

Related-to #11599
---
 .../java/com/vaadin/flow/component/UI.java    | 10 +++++++-
 .../com/vaadin/flow/component/UITest.java     | 25 +++++++++++++++++++
 2 files changed, 34 insertions(+), 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
index 595058fcbc4..0ea9e350b3a 100644
--- a/flow-server/src/main/java/com/vaadin/flow/component/UI.java
+++ b/flow-server/src/main/java/com/vaadin/flow/component/UI.java
@@ -1009,9 +1009,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(
@@ -1022,6 +1025,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
index bf50c463407..efa1573b35f 100644
--- a/flow-server/src/test/java/com/vaadin/flow/component/UITest.java
+++ b/flow-server/src/test/java/com/vaadin/flow/component/UITest.java
@@ -623,6 +623,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 {