Moved StreamVariable handling from AbstractCommunicationManager to

ConnectorTracker to prevent untimely unregistrations through other UIs
within the same session (#10112)

Change-Id: Id04c97970325be65b0b3c63756a2f2e731dd60d2

2 files changed, 125 insertions, 61 deletions

diff --git a/server/src/com/vaadin/server/AbstractCommunicationManager.java b/server/src/com/vaadin/server/AbstractCommunicationManager.java
index 0a3dff8120..af9118547f 100644
--- a/server/src/com/vaadin/server/AbstractCommunicationManager.java
+++ b/server/src/com/vaadin/server/AbstractCommunicationManager.java
@@ -165,10 +165,6 @@ public abstract class AbstractCommunicationManager implements Serializable {
 
     private Map<String, Class<?>> publishedFileContexts = new HashMap<String, Class<?>>();
 
-    private Map<String, Map<String, StreamVariable>> pidToNameToStreamVariable;
-
-    private Map<StreamVariable, String> streamVariableToSeckey;
-
     /**
      * TODO New constructor - document me!
      * 
@@ -644,23 +640,6 @@ public abstract class AbstractCommunicationManager implements Serializable {
         // Remove connectors that have been detached from the session during
         // handling of the request
         uI.getConnectorTracker().cleanConnectorMap();
-
-        if (pidToNameToStreamVariable != null) {
-            Iterator<String> iterator = pidToNameToStreamVariable.keySet()
-                    .iterator();
-            while (iterator.hasNext()) {
-                String connectorId = iterator.next();
-                if (uI.getConnectorTracker().getConnector(connectorId) == null) {
-                    // Owner is no longer attached to the session
-                    Map<String, StreamVariable> removed = pidToNameToStreamVariable
-                            .get(connectorId);
-                    for (String key : removed.keySet()) {
-                        streamVariableToSeckey.remove(removed.get(key));
-                    }
-                    iterator.remove();
-                }
-            }
-        }
     }
 
     protected void highlightConnector(ClientConnector highlightedConnector) {
@@ -2269,28 +2248,13 @@ public abstract class AbstractCommunicationManager implements Serializable {
          * handling post
          */
         String paintableId = owner.getConnectorId();
-        int uiId = owner.getUI().getUIId();
+        UI ui = owner.getUI();
+        int uiId = ui.getUIId();
         String key = uiId + "/" + paintableId + "/" + name;
 
-        if (pidToNameToStreamVariable == null) {
-            pidToNameToStreamVariable = new HashMap<String, Map<String, StreamVariable>>();
-        }
-        Map<String, StreamVariable> nameToStreamVariable = pidToNameToStreamVariable
-                .get(paintableId);
-        if (nameToStreamVariable == null) {
-            nameToStreamVariable = new HashMap<String, StreamVariable>();
-            pidToNameToStreamVariable.put(paintableId, nameToStreamVariable);
-        }
-        nameToStreamVariable.put(name, value);
-
-        if (streamVariableToSeckey == null) {
-            streamVariableToSeckey = new HashMap<StreamVariable, String>();
-        }
-        String seckey = streamVariableToSeckey.get(value);
-        if (seckey == null) {
-            seckey = UUID.randomUUID().toString();
-            streamVariableToSeckey.put(value, seckey);
-        }
+        ConnectorTracker connectorTracker = ui.getConnectorTracker();
+        connectorTracker.addStreamVariable(paintableId, name, value);
+        String seckey = connectorTracker.getSeckey(value);
 
         return ApplicationConstants.APP_PROTOCOL_PREFIX
                 + ServletPortletHelper.UPLOAD_URL_PREFIX + key + "/" + seckey;
@@ -2298,12 +2262,8 @@ public abstract class AbstractCommunicationManager implements Serializable {
     }
 
     public void cleanStreamVariable(ClientConnector owner, String name) {
-        Map<String, StreamVariable> nameToStreamVar = pidToNameToStreamVariable
-                .get(owner.getConnectorId());
-        nameToStreamVar.remove(name);
-        if (nameToStreamVar.isEmpty()) {
-            pidToNameToStreamVariable.remove(owner.getConnectorId());
-        }
+        owner.getUI().getConnectorTracker()
+                .cleanStreamVariable(owner.getConnectorId(), name);
     }
 
     /**
@@ -2683,9 +2643,9 @@ public abstract class AbstractCommunicationManager implements Serializable {
         UI uI = session.getUIById(Integer.parseInt(uiId));
         UI.setCurrent(uI);
 
-        StreamVariable streamVariable = getStreamVariable(connectorId,
-                variableName);
-        String secKey = streamVariableToSeckey.get(streamVariable);
+        StreamVariable streamVariable = uI.getConnectorTracker()
+                .getStreamVariable(connectorId, variableName);
+        String secKey = uI.getConnectorTracker().getSeckey(streamVariable);
         if (secKey.equals(parts[3])) {
 
             ClientConnector source = getConnector(uI, connectorId);
@@ -2741,17 +2701,6 @@ public abstract class AbstractCommunicationManager implements Serializable {
         }
     }
 
-    public StreamVariable getStreamVariable(String connectorId,
-            String variableName) {
-        Map<String, StreamVariable> map = pidToNameToStreamVariable
-                .get(connectorId);
-        if (map == null) {
-            return null;
-        }
-        StreamVariable streamVariable = map.get(variableName);
-        return streamVariable;
-    }
-
     /**
      * Stream that extracts content from another stream until the boundary
      * string is encountered.
diff --git a/server/src/com/vaadin/ui/ConnectorTracker.java b/server/src/com/vaadin/ui/ConnectorTracker.java
index 8b1a940c4b..91f9ac451c 100644
--- a/server/src/com/vaadin/ui/ConnectorTracker.java
+++ b/server/src/com/vaadin/ui/ConnectorTracker.java
@@ -23,6 +23,7 @@ import java.util.HashSet;
 import java.util.Iterator;
 import java.util.Map;
 import java.util.Set;
+import java.util.UUID;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
@@ -33,6 +34,7 @@ import com.vaadin.server.AbstractClientConnector;
 import com.vaadin.server.AbstractCommunicationManager;
 import com.vaadin.server.ClientConnector;
 import com.vaadin.server.GlobalResourceHandler;
+import com.vaadin.server.StreamVariable;
 
 /**
  * A class which takes care of book keeping of {@link ClientConnector}s for a
@@ -65,6 +67,11 @@ public class ConnectorTracker implements Serializable {
     private UI uI;
     private transient Map<ClientConnector, JSONObject> diffStates = new HashMap<ClientConnector, JSONObject>();
 
+    /** Maps connectorIds to a map of named StreamVariables */
+    private Map<String, Map<String, StreamVariable>> pidToNameToStreamVariable;
+
+    private Map<StreamVariable, String> streamVariableToSeckey;
+
     /**
      * Gets a logger for this class
      * 
@@ -259,6 +266,7 @@ public class ConnectorTracker implements Serializable {
             }
         }
 
+        cleanStreamVariables();
     }
 
     /**
@@ -497,4 +505,111 @@ public class ConnectorTracker implements Serializable {
         }
 
     }
+
+    /**
+     * Checks if the indicated connector has a StreamVariable of the given name
+     * and returns the variable if one is found.
+     * 
+     * @param connectorId
+     * @param variableName
+     * @return variable if a matching one exists, otherwise null
+     */
+    public StreamVariable getStreamVariable(String connectorId,
+            String variableName) {
+        if (pidToNameToStreamVariable == null) {
+            return null;
+        }
+        Map<String, StreamVariable> map = pidToNameToStreamVariable
+                .get(connectorId);
+        if (map == null) {
+            return null;
+        }
+        StreamVariable streamVariable = map.get(variableName);
+        return streamVariable;
+    }
+
+    /**
+     * Adds a StreamVariable of the given name to the indicated connector.
+     * 
+     * @param connectorId
+     * @param variableName
+     * @param variable
+     */
+    public void addStreamVariable(String connectorId, String variableName,
+            StreamVariable variable) {
+        if (pidToNameToStreamVariable == null) {
+            pidToNameToStreamVariable = new HashMap<String, Map<String, StreamVariable>>();
+        }
+        Map<String, StreamVariable> nameToStreamVariable = pidToNameToStreamVariable
+                .get(connectorId);
+        if (nameToStreamVariable == null) {
+            nameToStreamVariable = new HashMap<String, StreamVariable>();
+            pidToNameToStreamVariable.put(connectorId, nameToStreamVariable);
+        }
+        nameToStreamVariable.put(variableName, variable);
+
+        if (streamVariableToSeckey == null) {
+            streamVariableToSeckey = new HashMap<StreamVariable, String>();
+        }
+        String seckey = streamVariableToSeckey.get(variable);
+        if (seckey == null) {
+            seckey = UUID.randomUUID().toString();
+            streamVariableToSeckey.put(variable, seckey);
+        }
+    }
+
+    /**
+     * Removes StreamVariables that belong to connectors that are no longer
+     * attached to the session.
+     */
+    private void cleanStreamVariables() {
+        if (pidToNameToStreamVariable != null) {
+            Iterator<String> iterator = pidToNameToStreamVariable.keySet()
+                    .iterator();
+            while (iterator.hasNext()) {
+                String connectorId = iterator.next();
+                if (uI.getConnectorTracker().getConnector(connectorId) == null) {
+                    // Owner is no longer attached to the session
+                    Map<String, StreamVariable> removed = pidToNameToStreamVariable
+                            .get(connectorId);
+                    for (String key : removed.keySet()) {
+                        streamVariableToSeckey.remove(removed.get(key));
+                    }
+                    iterator.remove();
+                }
+            }
+        }
+    }
+
+    /**
+     * Removes any StreamVariable of the given name from the indicated
+     * connector.
+     * 
+     * @param connectorId
+     * @param variableName
+     */
+    public void cleanStreamVariable(String connectorId, String variableName) {
+        if (pidToNameToStreamVariable == null) {
+            return;
+        }
+        Map<String, StreamVariable> nameToStreamVar = pidToNameToStreamVariable
+                .get(connectorId);
+        nameToStreamVar.remove(variableName);
+        if (nameToStreamVar.isEmpty()) {
+            pidToNameToStreamVariable.remove(connectorId);
+        }
+    }
+
+    /**
+     * Returns the security key associated with the given StreamVariable.
+     * 
+     * @param variable
+     * @return matching security key if one exists, null otherwise
+     */
+    public String getSeckey(StreamVariable variable) {
+        if (streamVariableToSeckey == null) {
+            return null;
+        }
+        return streamVariableToSeckey.get(variable);
+    }
 }