From b5022108c74763d0695f5655a3a685c364477afd Mon Sep 17 00:00:00 2001
From: Rossen Stoyanchev <rstoyanchev@pivotal.io>
Date: Thu, 12 May 2016 11:21:41 -0400
Subject: [PATCH] Fix STOMP connect failure related memory leak

Normally heartbeats keep connections from hanging. However in some
cases a connection may hang before a CONNECTED frame is received
and heartbeats are put in place. This commit adds a change to enforce
a 60 limit on receiving the CONNECTED frame.

Issue: SPR-14266
---
 .../stomp/StompBrokerRelayMessageHandler.java   | 17 +++++++++++++++++
 .../tcp/reactor/Reactor2TcpClient.java          |  3 ++-
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompBrokerRelayMessageHandler.java b/spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompBrokerRelayMessageHandler.java
index 1eb2e49669..18ab77f67f 100644
--- a/spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompBrokerRelayMessageHandler.java
+++ b/spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompBrokerRelayMessageHandler.java
@@ -88,6 +88,14 @@ public class StompBrokerRelayMessageHandler extends AbstractBrokerMessageHandler
 
 	private static final Message<byte[]> HEARTBEAT_MESSAGE;
 
+	/**
+	 * A heartbeat is setup once a CONNECTED frame is received which contains
+	 * the heartbeat settings we need. If we don't receive CONNECTED within
+	 * a minute, the connection is closed proactively.
+	 */
+	private static final int MAX_TIME_TO_CONNECTED_FRAME = 60 * 1000;
+
+
 
 	static {
 		EMPTY_TASK.run();
@@ -567,6 +575,15 @@ public class StompBrokerRelayMessageHandler extends AbstractBrokerMessageHandler
 				logger.debug("TCP connection opened in session=" + getSessionId());
 			}
 			this.tcpConnection = connection;
+			this.tcpConnection.onReadInactivity(new Runnable() {
+				@Override
+				public void run() {
+					if (tcpConnection != null && !isStompConnected) {
+						handleTcpConnectionFailure("No CONNECTED frame received in " +
+								MAX_TIME_TO_CONNECTED_FRAME + " ms.", null);
+					}
+				}
+			}, MAX_TIME_TO_CONNECTED_FRAME);
 			connection.send(MessageBuilder.createMessage(EMPTY_PAYLOAD, this.connectHeaders.getMessageHeaders()));
 		}
 
diff --git a/spring-messaging/src/main/java/org/springframework/messaging/tcp/reactor/Reactor2TcpClient.java b/spring-messaging/src/main/java/org/springframework/messaging/tcp/reactor/Reactor2TcpClient.java
index 664e23c219..71541dcd79 100644
--- a/spring-messaging/src/main/java/org/springframework/messaging/tcp/reactor/Reactor2TcpClient.java
+++ b/spring-messaging/src/main/java/org/springframework/messaging/tcp/reactor/Reactor2TcpClient.java
@@ -168,7 +168,7 @@ public class Reactor2TcpClient<P> implements TcpOperations<P> {
 		Assert.notNull(connectionHandler, "TcpConnectionHandler must not be null");
 
 		final TcpClient<Message<P>, Message<P>> tcpClient;
-		Runnable cleanupTask;
+		final Runnable cleanupTask;
 		synchronized (this.tcpClients) {
 			if (this.stopping) {
 				IllegalStateException ex = new IllegalStateException("Shutting down.");
@@ -194,6 +194,7 @@ public class Reactor2TcpClient<P> implements TcpOperations<P> {
 				promise.onError(new Consumer<Throwable>() {
 					@Override
 					public void accept(Throwable ex) {
+						cleanupTask.run();
 						connectionHandler.afterConnectFailure(ex);
 					}
 				})