Make ClassResource find associated Application class (#9402)

3 files changed, 20 insertions, 6 deletions

diff --git a/server/src/com/vaadin/Application.java b/server/src/com/vaadin/Application.java
index d755817896..c439d30425 100644
--- a/server/src/com/vaadin/Application.java
+++ b/server/src/com/vaadin/Application.java
@@ -79,6 +79,7 @@ public abstract class Application extends AbstractUIProvider implements
             // no current UI -> set the main window as the current UI
             UI.setCurrent(mainWindow);
         }
+        mainWindow.setApplication(this);
         this.mainWindow = mainWindow;
     }
 
@@ -99,8 +100,7 @@ public abstract class Application extends AbstractUIProvider implements
     }
 
     @Override
-    public UI createInstance(Class<? extends UI> type,
-            WrappedRequest request) {
+    public UI createInstance(Class<? extends UI> type, WrappedRequest request) {
         return getUIInstance(request);
     }
 
@@ -240,6 +240,8 @@ public abstract class Application extends AbstractUIProvider implements
             uI.setName(name);
         }
 
+        uI.setApplication(this);
+
         legacyUINames.put(uI.getName(), uI);
         uI.setSession(VaadinSession.getCurrent());
     }
diff --git a/server/src/com/vaadin/server/ClassResource.java b/server/src/com/vaadin/server/ClassResource.java
index 27643eda13..a5b654ea61 100644
--- a/server/src/com/vaadin/server/ClassResource.java
+++ b/server/src/com/vaadin/server/ClassResource.java
@@ -112,11 +112,13 @@ public class ClassResource implements ConnectorResource, Serializable {
 
     protected Class<?> getAssociatedClass() {
         if (associatedClass == null) {
-            Class<? extends UI> associatedClass = UI.getCurrent().getClass();
-            if (associatedClass == LegacyWindow.class) {
-                return VaadinSession.getCurrent().getClass();
+            UI current = UI.getCurrent();
+            if (current instanceof LegacyWindow) {
+                LegacyWindow legacyWindow = (LegacyWindow) current;
+                return legacyWindow.getApplication().getClass();
+            } else {
+                return current.getClass();
             }
-            return associatedClass;
         }
         return associatedClass;
     }
diff --git a/server/src/com/vaadin/ui/UI.java b/server/src/com/vaadin/ui/UI.java
index fc3df8d052..7ef51dfc3e 100644
--- a/server/src/com/vaadin/ui/UI.java
+++ b/server/src/com/vaadin/ui/UI.java
@@ -27,6 +27,7 @@ import java.util.Iterator;
 import java.util.LinkedHashSet;
 import java.util.Map;
 
+import com.vaadin.Application;
 import com.vaadin.event.Action;
 import com.vaadin.event.Action.Handler;
 import com.vaadin.event.ActionManager;
@@ -94,6 +95,7 @@ public abstract class UI extends AbstractComponentContainer implements
     @Deprecated
     public static class LegacyWindow extends UI {
         private String name;
+        private Application application;
 
         /**
          * Create a new legacy window
@@ -127,6 +129,14 @@ public abstract class UI extends AbstractComponentContainer implements
             // Just empty
         }
 
+        public void setApplication(Application application) {
+            this.application = application;
+        }
+
+        public Application getApplication() {
+            return application;
+        }
+
         /**
          * Gets the unique name of the window. The name of the window is used to
          * uniquely identify it.