Make UI.pushConnection transient to prevent null resource after deserialization (#11809)

* PushConnection is not Serializable anymore
* AtmospherePushConnection fields are not transient
* UI.setSession calls setPushConnection(null) instead of pushConnection.disconnect()
* pushConnection.disconnect() asserts isConnected()
* If UI has a push connection, it should now always have isConnected() == true

Change-Id: I3c2e877b8e723b7cc2993cacd620920aecdef5fb

4 files changed, 37 insertions, 48 deletions

diff --git a/server/src/com/vaadin/server/communication/AtmospherePushConnection.java b/server/src/com/vaadin/server/communication/AtmospherePushConnection.java
index 9e57ccb45d..e325550c4b 100644
--- a/server/src/com/vaadin/server/communication/AtmospherePushConnection.java
+++ b/server/src/com/vaadin/server/communication/AtmospherePushConnection.java
@@ -93,12 +93,13 @@ public class AtmospherePushConnection implements PushConnection {
     }
 
     private UI ui;
-    private transient AtmosphereResource resource;
-    private transient Future<String> outgoingMessage;
-    private transient FragmentedMessage incomingMessage;
+    private AtmosphereResource resource;
+    private Future<String> outgoingMessage;
+    private FragmentedMessage incomingMessage;
 
-    public AtmospherePushConnection(UI ui) {
+    public AtmospherePushConnection(UI ui, AtmosphereResource resource) {
         this.ui = ui;
+        this.resource = resource;
     }
 
     @Override
@@ -176,21 +177,6 @@ public class AtmospherePushConnection implements PushConnection {
         }
     }
 
-    /**
-     * Associates this connection with the given AtmosphereResource. If there is
-     * a push pending, commits it.
-     * 
-     * @param resource
-     *            The AtmosphereResource representing the push channel.
-     * @throws IOException
-     */
-    protected void connect(AtmosphereResource resource) throws IOException {
-        this.resource = resource;
-    }
-
-    /**
-     * Returns whether this connection is currently open.
-     */
     @Override
     public boolean isConnected() {
         return resource != null
@@ -215,6 +201,8 @@ public class AtmospherePushConnection implements PushConnection {
 
     @Override
     public void disconnect() {
+        assert isConnected();
+
         if (outgoingMessage != null) {
             // Wait for the last message to be sent before closing the
             // connection (assumes that futures are completed in order)
diff --git a/server/src/com/vaadin/server/communication/PushConnection.java b/server/src/com/vaadin/server/communication/PushConnection.java
index bb88998402..7f78d1d48e 100644
--- a/server/src/com/vaadin/server/communication/PushConnection.java
+++ b/server/src/com/vaadin/server/communication/PushConnection.java
@@ -16,8 +16,6 @@
 
 package com.vaadin.server.communication;
 
-import java.io.Serializable;
-
 import com.vaadin.ui.UI;
 
 /**
@@ -27,18 +25,20 @@ import com.vaadin.ui.UI;
  * @author Vaadin Ltd
  * @since 7.1
  */
-public interface PushConnection extends Serializable {
+public interface PushConnection {
 
     /**
-     * Pushes pending state changes and client RPC calls to the client. It is
-     * NOT safe to invoke this method if not holding the session lock.
+     * Pushes pending state changes and client RPC calls to the client. Cannot
+     * be called if {@link #isConnected()} is false. It is NOT safe to invoke
+     * this method if not holding the session lock.
      * <p>
      * This is internal API; please use {@link UI#push()} instead.
      */
     public void push();
 
     /**
-     * Disconnects the connection.
+     * Closes the connection. Cannot be called if {@link #isConnected()} is
+     * false.
      */
     public void disconnect();
 
diff --git a/server/src/com/vaadin/server/communication/PushHandler.java b/server/src/com/vaadin/server/communication/PushHandler.java
index 7efcb8fd8c..6b853063a7 100644
--- a/server/src/com/vaadin/server/communication/PushHandler.java
+++ b/server/src/com/vaadin/server/communication/PushHandler.java
@@ -109,8 +109,7 @@ public class PushHandler implements AtmosphereHandler {
             resource.suspend();
 
             AtmospherePushConnection connection = new AtmospherePushConnection(
-                    ui);
-            connection.connect(resource);
+                    ui, resource);
 
             ui.setPushConnection(connection);
         }
@@ -305,12 +304,9 @@ public class PushHandler implements AtmosphereHandler {
     private static AtmospherePushConnection getConnectionForUI(UI ui) {
         PushConnection pushConnection = ui.getPushConnection();
         if (pushConnection instanceof AtmospherePushConnection) {
-            AtmospherePushConnection apc = (AtmospherePushConnection) pushConnection;
-            if (apc.isConnected()) {
-                return apc;
-            }
+            assert pushConnection.isConnected();
+            return (AtmospherePushConnection) pushConnection;
         }
-
         return null;
     }
 
@@ -374,11 +370,14 @@ public class PushHandler implements AtmosphereHandler {
      */
     private static void sendRefreshAndDisconnect(AtmosphereResource resource)
             throws IOException {
-        AtmospherePushConnection connection = new AtmospherePushConnection(null);
-        connection.connect(resource);
-        connection.sendMessage(VaadinService.createCriticalNotificationJSON(
-                null, null, null, null));
-        connection.disconnect();
+        AtmospherePushConnection connection = new AtmospherePushConnection(
+                null, resource);
+        try {
+            connection.sendMessage(VaadinService
+                    .createCriticalNotificationJSON(null, null, null, null));
+        } finally {
+            connection.disconnect();
+        }
     }
 
     private static final Logger getLogger() {
diff --git a/server/src/com/vaadin/ui/UI.java b/server/src/com/vaadin/ui/UI.java
index fa38666dc2..6c9551ea81 100644
--- a/server/src/com/vaadin/ui/UI.java
+++ b/server/src/com/vaadin/ui/UI.java
@@ -411,12 +411,9 @@ public abstract class UI extends AbstractSingleComponentContainer implements
         } else {
             if (session == null) {
                 detach();
-                if (pushConnection != null && pushConnection.isConnected()) {
-                    // Close the push connection when UI is detached. Otherwise
-                    // the push connection and possibly VaadinSession will live
-                    // on.
-                    pushConnection.disconnect();
-                }
+                // Close the push connection when UI is detached. Otherwise the
+                // push connection and possibly VaadinSession will live on.
+                setPushConnection(null);
             }
             this.session = session;
         }
@@ -536,7 +533,7 @@ public abstract class UI extends AbstractSingleComponentContainer implements
 
     private Navigator navigator;
 
-    private PushConnection pushConnection = null;
+    private transient PushConnection pushConnection = null;
 
     private boolean hasPendingPush = false;
 
@@ -1072,7 +1069,7 @@ public abstract class UI extends AbstractSingleComponentContainer implements
 
         boolean sessionExpired = (session == null || session.isClosing());
         getRpcProxy(UIClientRpc.class).uiClosed(sessionExpired);
-        if (getPushConnection() != null && getPushConnection().isConnected()) {
+        if (getPushConnection() != null) {
             // Push the Rpc to the client. The connection will be closed when
             // the UI is detached and cleaned up.
             getPushConnection().push();
@@ -1342,19 +1339,24 @@ public abstract class UI extends AbstractSingleComponentContainer implements
 
     /**
      * Returns the internal push connection object used by this UI. This method
-     * should only be called by the framework.
+     * should only be called by the framework. If the returned PushConnection is
+     * not null, it is guaranteed to have {@code isConnected() == true}.
      */
     public PushConnection getPushConnection() {
+        assert (pushConnection == null || pushConnection.isConnected());
         return pushConnection;
     }
 
     /**
      * Sets the internal push connection object used by this UI. This method
-     * should only be called by the framework.
+     * should only be called by the framework. If {@pushConnection} is not null,
+     * its {@code isConnected()} must be true.
      */
     public void setPushConnection(PushConnection pushConnection) {
         // If pushMode is disabled then there should never be a pushConnection
-        assert (getPushConfiguration().getPushMode().isEnabled() || pushConnection == null);
+        assert (pushConnection == null || getPushConfiguration().getPushMode()
+                .isEnabled());
+        assert (pushConnection == null || pushConnection.isConnected());
 
         if (pushConnection == this.pushConnection) {
             return;