From dda03218baa91f03dff06136bfee2c6aa9c10630 Mon Sep 17 00:00:00 2001
From: Mikhail Shabarov <mikhail@vaadin.com>
Date: Thu, 2 Sep 2021 09:59:48 +0300
Subject: [PATCH 1/2] fix: check the component belongs to UI before task
 execution

Related-to #11599
---
 .../java/com/vaadin/flow/component/UI.java    |  7 ++++++
 .../com/vaadin/flow/component/UITest.java     | 25 +++++++++++++++++++
 2 files changed, 32 insertions(+)

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 5ce464aa0b0..3aca79b731f 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
@@ -995,6 +995,8 @@ 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) {
@@ -1008,6 +1010,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 72d4bd3c143..a2eb93d326e 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
@@ -635,6 +635,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 {

From 2042cde43a6aee7b5e54e42cd2e1b3737cdb06d3 Mon Sep 17 00:00:00 2001
From: Mikhail Shabarov <mikhail@vaadin.com>
Date: Mon, 6 Sep 2021 14:34:59 +0300
Subject: [PATCH 2/2] Add IAE to method signature

---
 flow-server/src/main/java/com/vaadin/flow/component/UI.java | 3 ++-
 1 file changed, 2 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 3aca79b731f..342fddc7d74 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
@@ -999,7 +999,8 @@ public Router getRouter() {
      *             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(