Add customizers for Reactive Web Servers

Closes gh-9572

spring-boot/src/main/java/org/springframework/boot/web/embedded/jetty/JettyReactiveWebServerFactory.java

@@ -17,6 +17,10 @@
 package org.springframework.boot.web.embedded.jetty;
 
 import java.net.InetSocketAddress;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.List;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -34,6 +38,7 @@ import org.springframework.boot.web.reactive.server.ReactiveWebServerFactory;
 import org.springframework.boot.web.server.WebServer;
 import org.springframework.http.server.reactive.HttpHandler;
 import org.springframework.http.server.reactive.JettyHttpHandlerAdapter;
+import org.springframework.util.Assert;
 
 /**
  * {@link ReactiveWebServerFactory} that can be used to create {@link JettyWebServer}s.
@@ -56,6 +61,8 @@ public class JettyReactiveWebServerFactory extends AbstractReactiveWebServerFact
 	 */
 	private int selectors = -1;
 
+	private List<JettyServerCustomizer> jettyServerCustomizers = new ArrayList<>();
+
 	private ThreadPool threadPool;
 
 	/**
@@ -90,6 +97,9 @@ public class JettyReactiveWebServerFactory extends AbstractReactiveWebServerFact
 		ServletContextHandler contextHandler = new ServletContextHandler(server, "",
 				false, false);
 		contextHandler.addServlet(servletHolder, "/");
+		for (JettyServerCustomizer customizer : getServerCustomizers()) {
+			customizer.customize(server);
+		}
 		JettyReactiveWebServerFactory.logger
 				.info("Server initialized with port: " + port);
 		return server;
@@ -134,6 +144,36 @@ public class JettyReactiveWebServerFactory extends AbstractReactiveWebServerFact
 		this.acceptors = acceptors;
 	}
 
+	/**
+	 * Sets {@link JettyServerCustomizer}s that will be applied to the {@link Server}
+	 * before it is started. Calling this method will replace any existing customizers.
+	 * @param customizers the Jetty customizers to apply
+	 */
+	public void setServerCustomizers(
+			Collection<? extends JettyServerCustomizer> customizers) {
+		Assert.notNull(customizers, "Customizers must not be null");
+		this.jettyServerCustomizers = new ArrayList<>(customizers);
+	}
+
+	/**
+	 * Returns a mutable collection of Jetty {@link JettyServerCustomizer}s that will be applied
+	 * to the {@link Server} before it is created.
+	 * @return the Jetty customizers
+	 */
+	public Collection<JettyServerCustomizer> getServerCustomizers() {
+		return this.jettyServerCustomizers;
+	}
+
+	/**
+	 * Add {@link JettyServerCustomizer}s that will be applied to the {@link Server}
+	 * before it is started.
+	 * @param customizers the customizers to add
+	 */
+	public void addServerCustomizers(JettyServerCustomizer... customizers) {
+		Assert.notNull(customizers, "Customizers must not be null");
+		this.jettyServerCustomizers.addAll(Arrays.asList(customizers));
+	}
+
 	/**
 	 * Set the number of selector threads to use.
 	 * @param selectors the number of selector threads to use

spring-boot/src/main/java/org/springframework/boot/web/embedded/tomcat/TomcatReactiveWebServerFactory.java

@@ -55,6 +55,8 @@ public class TomcatReactiveWebServerFactory extends AbstractReactiveWebServerFac
 
 	private List<TomcatContextCustomizer> tomcatContextCustomizers = new ArrayList<>();
 
+	private List<TomcatConnectorCustomizer> tomcatConnectorCustomizers = new ArrayList<>();
+
 	/**
 	 * Create a new {@link TomcatServletWebServerFactory} instance.
 	 */
@@ -132,6 +134,9 @@ public class TomcatReactiveWebServerFactory extends AbstractReactiveWebServerFac
 		// If ApplicationContext is slow to start we want Tomcat not to bind to the socket
 		// prematurely...
 		connector.setProperty("bindOnInit", "false");
+		for (TomcatConnectorCustomizer customizer : this.tomcatConnectorCustomizers) {
+			customizer.customize(connector);
+		}
 	}
 
 	private void customizeProtocol(AbstractProtocol<?> protocol) {
@@ -173,6 +178,39 @@ public class TomcatReactiveWebServerFactory extends AbstractReactiveWebServerFac
 		this.tomcatContextCustomizers.addAll(Arrays.asList(tomcatContextCustomizers));
 	}
 
+	/**
+	 * Set {@link TomcatConnectorCustomizer}s that should be applied to the Tomcat
+	 * {@link Connector} . Calling this method will replace any existing customizers.
+	 * @param tomcatConnectorCustomizers the customizers to set
+	 */
+	public void setTomcatConnectorCustomizers(
+			Collection<? extends TomcatConnectorCustomizer> tomcatConnectorCustomizers) {
+		Assert.notNull(tomcatConnectorCustomizers,
+				"TomcatConnectorCustomizers must not be null");
+		this.tomcatConnectorCustomizers = new ArrayList<>(tomcatConnectorCustomizers);
+	}
+
+	/**
+	 * Add {@link TomcatConnectorCustomizer}s that should be added to the Tomcat
+	 * {@link Connector}.
+	 * @param tomcatConnectorCustomizers the customizers to add
+	 */
+	public void addConnectorCustomizers(
+			TomcatConnectorCustomizer... tomcatConnectorCustomizers) {
+		Assert.notNull(tomcatConnectorCustomizers,
+				"TomcatConnectorCustomizers must not be null");
+		this.tomcatConnectorCustomizers.addAll(Arrays.asList(tomcatConnectorCustomizers));
+	}
+
+	/**
+	 * Returns a mutable collection of the {@link TomcatConnectorCustomizer}s that will be
+	 * applied to the Tomcat {@link Connector} .
+	 * @return the customizers that will be applied
+	 */
+	public Collection<TomcatConnectorCustomizer> getTomcatConnectorCustomizers() {
+		return this.tomcatConnectorCustomizers;
+	}
+
 	/**
 	 * Factory method called to create the {@link TomcatWebServer}. Subclasses can
 	 * override this method to return a different {@link TomcatWebServer} or apply

spring-boot/src/main/java/org/springframework/boot/web/embedded/undertow/UndertowReactiveWebServerFactory.java

@@ -16,6 +16,11 @@
 
 package org.springframework.boot.web.embedded.undertow;
 
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.List;
+
 import io.undertow.Undertow;
 
 import org.springframework.boot.web.reactive.server.AbstractReactiveWebServerFactory;
@@ -23,6 +28,7 @@ import org.springframework.boot.web.reactive.server.ReactiveWebServerFactory;
 import org.springframework.boot.web.server.WebServer;
 import org.springframework.http.server.reactive.HttpHandler;
 import org.springframework.http.server.reactive.UndertowHttpHandlerAdapter;
+import org.springframework.util.Assert;
 
 /**
  * {@link ReactiveWebServerFactory} that can be used to create {@link UndertowWebServer}s.
@@ -40,6 +46,8 @@ public class UndertowReactiveWebServerFactory extends AbstractReactiveWebServerF
 
 	private Boolean directBuffers;
 
+	private List<UndertowBuilderCustomizer> builderCustomizers = new ArrayList<>();
+
 	/**
 	 * Create a new {@link UndertowReactiveWebServerFactory} instance.
 	 */
@@ -78,6 +86,9 @@ public class UndertowReactiveWebServerFactory extends AbstractReactiveWebServerF
 			builder.setDirectBuffers(this.directBuffers);
 		}
 		builder.addHttpListener(port, getListenAddress());
+		for (UndertowBuilderCustomizer customizer : this.builderCustomizers) {
+			customizer.customize(builder);
+		}
 		return builder;
 	}
 
@@ -104,4 +115,34 @@ public class UndertowReactiveWebServerFactory extends AbstractReactiveWebServerF
 		this.directBuffers = directBuffers;
 	}
 
+	/**
+	 * Set {@link UndertowBuilderCustomizer}s that should be applied to the Undertow
+	 * {@link Undertow.Builder}. Calling this method will replace any existing customizers.
+	 * @param customizers the customizers to set
+	 */
+	public void setBuilderCustomizers(
+			Collection<? extends UndertowBuilderCustomizer> customizers) {
+		Assert.notNull(customizers, "Customizers must not be null");
+		this.builderCustomizers = new ArrayList<>(customizers);
+	}
+
+	/**
+	 * Returns a mutable collection of the {@link UndertowBuilderCustomizer}s that will be
+	 * applied to the Undertow {@link Undertow.Builder} .
+	 * @return the customizers that will be applied
+	 */
+	public Collection<UndertowBuilderCustomizer> getBuilderCustomizers() {
+		return this.builderCustomizers;
+	}
+
+	/**
+	 * Add {@link UndertowBuilderCustomizer}s that should be used to customize the
+	 * Undertow {@link Undertow.Builder}.
+	 * @param customizers the customizers to add
+	 */
+	public void addBuilderCustomizers(UndertowBuilderCustomizer... customizers) {
+		Assert.notNull(customizers, "Customizers must not be null");
+		this.builderCustomizers.addAll(Arrays.asList(customizers));
+	}
+
 }

spring-boot/src/test/java/org/springframework/boot/web/embedded/jetty/JettyReactiveWebServerFactoryTests.java

@@ -16,20 +16,64 @@
 
 package org.springframework.boot.web.embedded.jetty;
 
-import org.springframework.boot.web.reactive.server.AbstractReactiveWebServerFactory;
+import java.util.Arrays;
+
+import org.eclipse.jetty.server.Server;
+import org.junit.Test;
+import org.mockito.InOrder;
+
 import org.springframework.boot.web.reactive.server.AbstractReactiveWebServerFactoryTests;
+import org.springframework.http.server.reactive.HttpHandler;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.inOrder;
+import static org.mockito.Mockito.mock;
 
 /**
  * Tests for {@link JettyReactiveWebServerFactory} and {@link JettyWebServer}.
  *
  * @author Brian Clozel
+ * @author Madhura Bhave
  */
 public class JettyReactiveWebServerFactoryTests
 		extends AbstractReactiveWebServerFactoryTests {
 
 	@Override
-	protected AbstractReactiveWebServerFactory getFactory() {
+	protected JettyReactiveWebServerFactory getFactory() {
 		return new JettyReactiveWebServerFactory(0);
 	}
 
+	@Test
+	public void setNullServerCustomizersShouldThrowException() {
+		JettyReactiveWebServerFactory factory = getFactory();
+		this.thrown.expect(IllegalArgumentException.class);
+		this.thrown.expectMessage("Customizers must not be null");
+		factory.setServerCustomizers(null);
+	}
+
+	@Test
+	public void addNullServerCustomizersShouldThrowException() {
+		JettyReactiveWebServerFactory factory = getFactory();
+		this.thrown.expect(IllegalArgumentException.class);
+		this.thrown.expectMessage("Customizers must not be null");
+		factory.addServerCustomizers((JettyServerCustomizer[]) null);
+	}
+
+	@Test
+	public void jettyCustomizersShouldBeInvoked() throws Exception {
+		HttpHandler handler = mock(HttpHandler.class);
+		JettyReactiveWebServerFactory factory = getFactory();
+		JettyServerCustomizer[] configurations = new JettyServerCustomizer[4];
+		for (int i = 0; i < configurations.length; i++) {
+			configurations[i] = mock(JettyServerCustomizer.class);
+		}
+		factory.setServerCustomizers(Arrays.asList(configurations[0], configurations[1]));
+		factory.addServerCustomizers(configurations[2], configurations[3]);
+		this.webServer = factory.getWebServer(handler);
+		InOrder ordered = inOrder((Object[]) configurations);
+		for (JettyServerCustomizer configuration : configurations) {
+			ordered.verify(configuration).customize(any(Server.class));
+		}
+	}
+
 }

spring-boot/src/test/java/org/springframework/boot/web/embedded/tomcat/TomcatReactiveWebServerFactoryTests.java

@@ -19,6 +19,7 @@ package org.springframework.boot.web.embedded.tomcat;
 import java.util.Arrays;
 
 import org.apache.catalina.Context;
+import org.apache.catalina.connector.Connector;
 import org.junit.Test;
 import org.mockito.InOrder;
 
@@ -33,6 +34,7 @@ import static org.mockito.Mockito.mock;
  * Tests for {@link TomcatReactiveWebServerFactory}.
  *
  * @author Brian Clozel
+ * @author Madhura Bhave
  */
 public class TomcatReactiveWebServerFactoryTests
 		extends AbstractReactiveWebServerFactoryTests {
@@ -58,4 +60,37 @@ public class TomcatReactiveWebServerFactoryTests
 		}
 	}
 
+	@Test
+	public void setNullConnectorCustomizersShouldThrowException() {
+		TomcatReactiveWebServerFactory factory = getFactory();
+		this.thrown.expect(IllegalArgumentException.class);
+		this.thrown.expectMessage("Customizers must not be null");
+		factory.setTomcatConnectorCustomizers(null);
+	}
+
+	@Test
+	public void addNullAddConnectorCustomizersShouldThrowException() {
+		TomcatReactiveWebServerFactory factory = getFactory();
+		this.thrown.expect(IllegalArgumentException.class);
+		this.thrown.expectMessage("Customizers must not be null");
+		factory.addConnectorCustomizers((TomcatConnectorCustomizer[]) null);
+	}
+
+	@Test
+	public void tomcatConnectorCustomizersShouldBeInvoked() throws Exception {
+		TomcatReactiveWebServerFactory factory = getFactory();
+		HttpHandler handler = mock(HttpHandler.class);
+		TomcatConnectorCustomizer[] listeners = new TomcatConnectorCustomizer[4];
+		for (int i = 0; i < listeners.length; i++) {
+			listeners[i] = mock(TomcatConnectorCustomizer.class);
+		}
+		factory.setTomcatConnectorCustomizers(Arrays.asList(listeners[0], listeners[1]));
+		factory.addConnectorCustomizers(listeners[2], listeners[3]);
+		this.webServer = factory.getWebServer(handler);
+		InOrder ordered = inOrder((Object[]) listeners);
+		for (TomcatConnectorCustomizer listener : listeners) {
+			ordered.verify(listener).customize(any(Connector.class));
+		}
+	}
+
 }

spring-boot/src/test/java/org/springframework/boot/web/embedded/undertow/UndertowReactiveWebServerFactoryTests.java

@@ -16,12 +16,24 @@
 
 package org.springframework.boot.web.embedded.undertow;
 
+import java.util.Arrays;
+
+import io.undertow.Undertow;
+import org.junit.Test;
+import org.mockito.InOrder;
+
 import org.springframework.boot.web.reactive.server.AbstractReactiveWebServerFactoryTests;
+import org.springframework.http.server.reactive.HttpHandler;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.inOrder;
+import static org.mockito.Mockito.mock;
 
 /**
  * Tests for {@link UndertowReactiveWebServerFactory} and {@link UndertowWebServer}.
  *
  * @author Brian Clozel
+ * @author Madhura Bhave
  */
 public class UndertowReactiveWebServerFactoryTests
 		extends AbstractReactiveWebServerFactoryTests {
@@ -31,4 +43,37 @@ public class UndertowReactiveWebServerFactoryTests
 		return new UndertowReactiveWebServerFactory(0);
 	}
 
+	@Test
+	public void setNullBuilderCustomizersShouldThrowException() {
+		UndertowReactiveWebServerFactory factory = getFactory();
+		this.thrown.expect(IllegalArgumentException.class);
+		this.thrown.expectMessage("Customizers must not be null");
+		factory.setBuilderCustomizers(null);
+	}
+
+	@Test
+	public void addNullBuilderCustomizersShouldThrowException() {
+		UndertowReactiveWebServerFactory factory = getFactory();
+		this.thrown.expect(IllegalArgumentException.class);
+		this.thrown.expectMessage("Customizers must not be null");
+		factory.addBuilderCustomizers((UndertowBuilderCustomizer[]) null);
+	}
+
+	@Test
+	public void builderCustomizersShouldBeInvoked() throws Exception {
+		UndertowReactiveWebServerFactory factory = getFactory();
+		HttpHandler handler = mock(HttpHandler.class);
+		UndertowBuilderCustomizer[] customizers = new UndertowBuilderCustomizer[4];
+		for (int i = 0; i < customizers.length; i++) {
+			customizers[i] = mock(UndertowBuilderCustomizer.class);
+		}
+		factory.setBuilderCustomizers(Arrays.asList(customizers[0], customizers[1]));
+		factory.addBuilderCustomizers(customizers[2], customizers[3]);
+		this.webServer = factory.getWebServer(handler);
+		InOrder ordered = inOrder((Object[]) customizers);
+		for (UndertowBuilderCustomizer customizer : customizers) {
+			ordered.verify(customizer).customize(any(Undertow.Builder.class));
+		}
+	}
+
 }