From 1ea829e00300d6124d26405d2951ecca78877d46 Mon Sep 17 00:00:00 2001
From: Grigory Fadeev <grigory.fadeev@gmail.com>
Date: Sat, 19 Nov 2016 19:42:16 +0200
Subject: [PATCH] Fix connector used to configure connection timeout

Fix Tomcat customization so that the main connection is configured with
any timeout.

Closes gh-7425
---
 .../autoconfigure/web/ServerProperties.java   | 20 ++++++++++++-------
 .../web/ServerPropertiesTests.java            |  8 ++++++++
 2 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/web/ServerProperties.java b/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/web/ServerProperties.java
index ae4a553491..4579f32d04 100644
--- a/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/web/ServerProperties.java
+++ b/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/web/ServerProperties.java
@@ -809,14 +809,20 @@ public class ServerProperties
 		}
 
 		private void customizeConnectionTimeout(
-				TomcatEmbeddedServletContainerFactory factory, int connectionTimeout) {
-			for (Connector connector : factory.getAdditionalTomcatConnectors()) {
-				if (connector.getProtocolHandler() instanceof AbstractProtocol) {
-					AbstractProtocol<?> handler = (AbstractProtocol<?>) connector
-							.getProtocolHandler();
-					handler.setConnectionTimeout(connectionTimeout);
+				TomcatEmbeddedServletContainerFactory factory,
+				final int connectionTimeout) {
+			factory.addConnectorCustomizers(new TomcatConnectorCustomizer() {
+
+				@Override
+				public void customize(Connector connector) {
+					ProtocolHandler handler = connector.getProtocolHandler();
+					if (handler instanceof AbstractProtocol) {
+						AbstractProtocol<?> protocol = (AbstractProtocol<?>) handler;
+						protocol.setConnectionTimeout(connectionTimeout);
+					}
 				}
-			}
+
+			});
 		}
 
 		private void customizeRemoteIpValve(ServerProperties properties,
diff --git a/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/web/ServerPropertiesTests.java b/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/web/ServerPropertiesTests.java
index 7a76eaffea..e581c42e7a 100644
--- a/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/web/ServerPropertiesTests.java
+++ b/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/web/ServerPropertiesTests.java
@@ -111,6 +111,14 @@ public class ServerPropertiesTests {
 		assertThat(this.properties.getServerHeader()).isEqualTo("Custom Server");
 	}
 
+	@Test
+	public void testConnectionTimeout() throws Exception {
+		Map<String, String> map = new HashMap<String, String>();
+		map.put("server.connection-timeout", "60000");
+		bindProperties(map);
+		assertThat(this.properties.getConnectionTimeout()).isEqualTo(60000);
+	}
+
 	@Test
 	public void testServletPathAsMapping() throws Exception {
 		RelaxedDataBinder binder = new RelaxedDataBinder(this.properties, "server");