From 150aba2191ed0bb527a41aedce45f28faf060b1c Mon Sep 17 00:00:00 2001
From: Christoffer Sawicki <christoffer.sawicki@gmail.com>
Date: Mon, 18 Jul 2016 15:22:25 -0700
Subject: [PATCH] Add support for engine valves

Update TomcatEmbeddedServletContainerFactory to allow registration of
engine valves as well as context values. For clarity the ambiguous
`getValues()` method has been deprecated in favor of
`getContextValves()`

See gh-6311
---
 ...TomcatEmbeddedServletContainerFactory.java | 55 ++++++++++++++++++-
 1 file changed, 53 insertions(+), 2 deletions(-)

diff --git a/spring-boot/src/main/java/org/springframework/boot/context/embedded/tomcat/TomcatEmbeddedServletContainerFactory.java b/spring-boot/src/main/java/org/springframework/boot/context/embedded/tomcat/TomcatEmbeddedServletContainerFactory.java
index 5aad85009b..30e4c0b192 100644
--- a/spring-boot/src/main/java/org/springframework/boot/context/embedded/tomcat/TomcatEmbeddedServletContainerFactory.java
+++ b/spring-boot/src/main/java/org/springframework/boot/context/embedded/tomcat/TomcatEmbeddedServletContainerFactory.java
@@ -34,6 +34,7 @@ import javax.servlet.ServletContainerInitializer;
 import javax.servlet.ServletContext;
 
 import org.apache.catalina.Context;
+import org.apache.catalina.Engine;
 import org.apache.catalina.Host;
 import org.apache.catalina.Lifecycle;
 import org.apache.catalina.LifecycleEvent;
@@ -88,6 +89,7 @@ import org.springframework.util.StringUtils;
  * @author Stephane Nicoll
  * @author Andy Wilkinson
  * @author Eddú Meléndez
+ * @author Christoffer Sawicki
  * @see #setPort(int)
  * @see #setContextLifecycleListeners(Collection)
  * @see TomcatEmbeddedServletContainer
@@ -106,6 +108,8 @@ public class TomcatEmbeddedServletContainerFactory
 
 	private File baseDirectory;
 
+	private List<Valve> engineValves = new ArrayList<Valve>();
+
 	private List<Valve> contextValves = new ArrayList<Valve>();
 
 	private List<LifecycleListener> contextLifecycleListeners = new ArrayList<LifecycleListener>();
@@ -162,7 +166,7 @@ public class TomcatEmbeddedServletContainerFactory
 		customizeConnector(connector);
 		tomcat.setConnector(connector);
 		tomcat.getHost().setAutoDeploy(false);
-		tomcat.getEngine().setBackgroundProcessorDelay(-1);
+		configureEngine(tomcat.getEngine());
 		for (Connector additionalConnector : this.additionalTomcatConnectors) {
 			tomcat.getService().addConnector(additionalConnector);
 		}
@@ -170,6 +174,13 @@ public class TomcatEmbeddedServletContainerFactory
 		return getTomcatEmbeddedServletContainer(tomcat);
 	}
 
+	private void configureEngine(Engine engine) {
+		engine.setBackgroundProcessorDelay(-1);
+		for (Valve valve : this.engineValves) {
+			engine.getPipeline().addValve(valve);
+		}
+	}
+
 	protected void prepareContext(Host host, ServletContextInitializer[] initializers) {
 		File docBase = getValidDocumentRoot();
 		docBase = (docBase != null ? docBase : createTempDir("tomcat-docbase"));
@@ -527,9 +538,37 @@ public class TomcatEmbeddedServletContainerFactory
 		this.protocol = protocol;
 	}
 
+	/**
+	 * Set {@link Valve}s that should be applied to the Tomcat {@link Engine}. Calling
+	 * this method will replace any existing valves.
+	 * @param engineValves the valves to set
+	 */
+	public void setEngineValves(Collection<? extends Valve> engineValves) {
+		Assert.notNull(engineValves, "Valves must not be null");
+		this.engineValves = new ArrayList<Valve>(engineValves);
+	}
+
+	/**
+	 * Returns a mutable collection of the {@link Valve}s that will be applied to the
+	 * Tomcat {@link Engine}.
+	 * @return the engineValves the valves that will be applied
+	 */
+	public Collection<Valve> getEngineValves() {
+		return this.engineValves;
+	}
+
+	/**
+	 * Add {@link Valve}s that should be applied to the Tomcat {@link Engine}.
+	 * @param engineValves the valves to add
+	 */
+	public void addEngineValves(Valve... engineValves) {
+		Assert.notNull(engineValves, "Valves must not be null");
+		this.engineValves.addAll(Arrays.asList(engineValves));
+	}
+
 	/**
 	 * Set {@link Valve}s that should be applied to the Tomcat {@link Context}. Calling
-	 * this method will replace any existing listeners.
+	 * this method will replace any existing valves.
 	 * @param contextValves the valves to set
 	 */
 	public void setContextValves(Collection<? extends Valve> contextValves) {
@@ -541,8 +580,20 @@ public class TomcatEmbeddedServletContainerFactory
 	 * Returns a mutable collection of the {@link Valve}s that will be applied to the
 	 * Tomcat {@link Context}.
 	 * @return the contextValves the valves that will be applied
+	 * @deprecated as of 1.4 in favor of {@link #getContextValves()}
 	 */
+	@Deprecated
 	public Collection<Valve> getValves() {
+		return getContextValves();
+	}
+
+	/**
+	 * Returns a mutable collection of the {@link Valve}s that will be applied to the
+	 * Tomcat {@link Context}.
+	 * @return the contextValves the valves that will be applied
+	 * @see #getEngineValves()
+	 */
+	public Collection<Valve> getContextValves() {
 		return this.contextValves;
 	}