From ff99bb0730c23ce8c46785dfd59cc631f41dcf17 Mon Sep 17 00:00:00 2001
From: Andy Wilkinson <awilkinson@pivotal.io>
Date: Fri, 12 Feb 2016 16:40:02 +0000
Subject: [PATCH] Set main thread's context class loader when starting Tomcat

When an app is deployed to Tomcat, all of the application's startup
is performed with a WebAppClassLoader being the thread context
class loader. When an app is using embedded Tomcat, the
WebAppClassLoader is created as part of the application starting but
is never set as the thread context class loader. This difference
in TCCL can cause problems. For example, it breaks the use of JNDI
during application startup with embedded Tomcat.

This commit updates the embedded Tomcat servlet container to set
the TCCL to be the WebAppClassLoader once the Tomcat context has
been started. Once Tomcat is stopped, it sets the TCCL back to the
ClassLoader that loaded it.

Closes gh-2308
---
 .../tomcat/TomcatEmbeddedServletContainer.java  | 17 +++++++++++++++++
 ...catEmbeddedServletContainerFactoryTests.java | 17 +++++++++++++++++
 2 files changed, 34 insertions(+)

diff --git a/spring-boot/src/main/java/org/springframework/boot/context/embedded/tomcat/TomcatEmbeddedServletContainer.java b/spring-boot/src/main/java/org/springframework/boot/context/embedded/tomcat/TomcatEmbeddedServletContainer.java
index 3f9ad91a975..95a3eb98004 100644
--- a/spring-boot/src/main/java/org/springframework/boot/context/embedded/tomcat/TomcatEmbeddedServletContainer.java
+++ b/spring-boot/src/main/java/org/springframework/boot/context/embedded/tomcat/TomcatEmbeddedServletContainer.java
@@ -21,6 +21,7 @@ import java.util.Map;
 import java.util.concurrent.atomic.AtomicInteger;
 
 import org.apache.catalina.Container;
+import org.apache.catalina.Context;
 import org.apache.catalina.Engine;
 import org.apache.catalina.LifecycleException;
 import org.apache.catalina.LifecycleState;
@@ -91,6 +92,9 @@ public class TomcatEmbeddedServletContainer implements EmbeddedServletContainer
 			// We can re-throw failure exception directly in the main thread
 			rethrowDeferredStartupExceptions();
 
+			ClassLoader classLoader = findContext().getLoader().getClassLoader();
+			Thread.currentThread().setContextClassLoader(classLoader);
+
 			// Unlike Jetty, all Tomcat threads are daemon threads. We create a
 			// blocking non-daemon to stop immediate shutdown
 			startDaemonAwaitThread();
@@ -101,6 +105,15 @@ public class TomcatEmbeddedServletContainer implements EmbeddedServletContainer
 		}
 	}
 
+	private Context findContext() {
+		for (Container child : this.tomcat.getHost().findChildren()) {
+			if (child instanceof Context) {
+				return (Context) child;
+			}
+		}
+		throw new IllegalStateException("The host does not contain a Context");
+	}
+
 	private void addInstanceIdToEngineName() {
 		int instanceId = containerCounter.incrementAndGet();
 		if (instanceId > 0) {
@@ -245,6 +258,10 @@ public class TomcatEmbeddedServletContainer implements EmbeddedServletContainer
 					ex);
 		}
 		finally {
+			if (Thread.currentThread()
+					.getContextClassLoader() instanceof TomcatEmbeddedWebappClassLoader) {
+				Thread.currentThread().setContextClassLoader(getClass().getClassLoader());
+			}
 			containerCounter.decrementAndGet();
 		}
 	}
diff --git a/spring-boot/src/test/java/org/springframework/boot/context/embedded/tomcat/TomcatEmbeddedServletContainerFactoryTests.java b/spring-boot/src/test/java/org/springframework/boot/context/embedded/tomcat/TomcatEmbeddedServletContainerFactoryTests.java
index cad1a1be231..b2ddd7e12f0 100644
--- a/spring-boot/src/test/java/org/springframework/boot/context/embedded/tomcat/TomcatEmbeddedServletContainerFactoryTests.java
+++ b/spring-boot/src/test/java/org/springframework/boot/context/embedded/tomcat/TomcatEmbeddedServletContainerFactoryTests.java
@@ -38,6 +38,7 @@ import org.apache.catalina.connector.Connector;
 import org.apache.catalina.startup.Tomcat;
 import org.apache.catalina.valves.RemoteIpValve;
 import org.apache.coyote.http11.AbstractHttp11JsseProtocol;
+import org.junit.After;
 import org.junit.Test;
 import org.mockito.InOrder;
 
@@ -72,6 +73,11 @@ public class TomcatEmbeddedServletContainerFactoryTests
 		return new TomcatEmbeddedServletContainerFactory(0);
 	}
 
+	@After
+	public void restoreTccl() {
+		Thread.currentThread().setContextClassLoader(getClass().getClassLoader());
+	}
+
 	// JMX MBean names clash if you get more than one Engine with the same name...
 	@Test
 	public void tomcatEngineNames() throws Exception {
@@ -358,6 +364,17 @@ public class TomcatEmbeddedServletContainerFactoryTests
 		assertThat(s3.split(":")[0]).as(message).isNotEqualTo(s2.split(":")[1]);
 	}
 
+	@Test
+	public void tcclOfMainThreadIsTomcatWebAppClassLoader() {
+		Thread.currentThread().setContextClassLoader(getClass().getClassLoader());
+		TomcatEmbeddedServletContainerFactory factory = getFactory();
+		this.container = factory.getEmbeddedServletContainer();
+		this.container.start();
+		assertThat(Thread.currentThread().getContextClassLoader())
+				.isInstanceOf(TomcatEmbeddedWebappClassLoader.class);
+		this.container.stop();
+	}
+
 	@Override
 	protected Wrapper getJspServlet() {
 		Container context = ((TomcatEmbeddedServletContainer) this.container).getTomcat()