From 2169bbbc9b46c74706b04d9663b6dba8b442edc4 Mon Sep 17 00:00:00 2001
From: Dave Syer <dsyer@pivotal.io>
Date: Thu, 1 Oct 2015 17:45:43 +0100
Subject: [PATCH] Check Environment matches webEnvironment after it is
 initialized

After the ApplicationEnvironmentPrepared we now check that the
webEnvironment flag and/or the Environment hasn't changed, in case
the user wanted to switch the context from a web to non-web in a
listener.

Fixes gh-2716
---
 .../boot/SpringApplication.java               | 86 +++++++++++-----
 .../boot/SpringApplicationTests.java          | 98 +++++++++++++------
 2 files changed, 127 insertions(+), 57 deletions(-)

diff --git a/spring-boot/src/main/java/org/springframework/boot/SpringApplication.java b/spring-boot/src/main/java/org/springframework/boot/SpringApplication.java
index 3e04228d4ed..3d96535f6dd 100644
--- a/spring-boot/src/main/java/org/springframework/boot/SpringApplication.java
+++ b/spring-boot/src/main/java/org/springframework/boot/SpringApplication.java
@@ -120,7 +120,8 @@ import org.springframework.web.context.support.StandardServletEnvironment;
  * your application, however, any of the following sources can also be used:
  *
  * <ul>
- * <li>{@link Class} - A Java class to be loaded by {@link AnnotatedBeanDefinitionReader}</li>
+ * <li>{@link Class} - A Java class to be loaded by {@link AnnotatedBeanDefinitionReader}
+ * </li>
  *
  * <li>{@link Resource} - An XML resource to be loaded by {@link XmlBeanDefinitionReader},
  * or a groovy script to be loaded by {@link GroovyBeanDefinitionReader}</li>
@@ -235,7 +236,8 @@ public class SpringApplication {
 			this.sources.addAll(Arrays.asList(sources));
 		}
 		this.webEnvironment = deduceWebEnvironment();
-		setInitializers((Collection) getSpringFactoriesInstances(ApplicationContextInitializer.class));
+		setInitializers((Collection) getSpringFactoriesInstances(
+				ApplicationContextInitializer.class));
 		setListeners((Collection) getSpringFactoriesInstances(ApplicationListener.class));
 		this.mainApplicationClass = deduceMainApplicationClass();
 	}
@@ -281,8 +283,8 @@ public class SpringApplication {
 			context = doRun(listeners, args);
 			stopWatch.stop();
 			if (this.logStartupInfo) {
-				new StartupInfoLogger(this.mainApplicationClass).logStarted(
-						getApplicationLog(), stopWatch);
+				new StartupInfoLogger(this.mainApplicationClass)
+						.logStarted(getApplicationLog(), stopWatch);
 			}
 			return context;
 		}
@@ -312,6 +314,7 @@ public class SpringApplication {
 			printBanner(environment);
 		}
 
+		environment = resetEnvironment(environment);
 		// Create, load, refresh and run the ApplicationContext
 		context = createApplicationContext();
 		if (this.registerShutdownHook) {
@@ -348,11 +351,43 @@ public class SpringApplication {
 		return context;
 	}
 
+	private ConfigurableEnvironment resetEnvironment(
+			ConfigurableEnvironment environment) {
+		if (this.environment != null && environment != this.environment) {
+			environment = this.environment;
+		}
+		if (environment instanceof StandardServletEnvironment && !this.webEnvironment) {
+			Set<String> servletSources = new HashSet<>(Arrays.asList(
+					StandardServletEnvironment.SERVLET_CONTEXT_PROPERTY_SOURCE_NAME,
+					StandardServletEnvironment.SERVLET_CONFIG_PROPERTY_SOURCE_NAME,
+					StandardServletEnvironment.JNDI_PROPERTY_SOURCE_NAME));
+			ConfigurableEnvironment copy = environment;
+			environment = new StandardEnvironment();
+			environment.setActiveProfiles(copy.getActiveProfiles());
+			String current = null;
+			MutablePropertySources sources = environment.getPropertySources();
+			for (PropertySource<?> source : copy.getPropertySources()) {
+				if (!servletSources.contains(source.getName())) {
+					String name = source.getName();
+					if (sources.contains(name)) {
+						sources.replace(name, source);
+					}
+					else if (current == null) {
+						sources.addFirst(source);
+					}
+					else {
+						sources.addAfter(current, source);
+					}
+					current = name;
+				}
+			}
+		}
+		return environment;
+	}
+
 	private void configureHeadlessProperty() {
-		System.setProperty(
-				SYSTEM_PROPERTY_JAVA_AWT_HEADLESS,
-				System.getProperty(SYSTEM_PROPERTY_JAVA_AWT_HEADLESS,
-						Boolean.toString(this.headless)));
+		System.setProperty(SYSTEM_PROPERTY_JAVA_AWT_HEADLESS, System.getProperty(
+				SYSTEM_PROPERTY_JAVA_AWT_HEADLESS, Boolean.toString(this.headless)));
 	}
 
 	private SpringApplicationRunListeners getRunListeners(String[] args) {
@@ -385,8 +420,8 @@ public class SpringApplication {
 				instances.add(instance);
 			}
 			catch (Throwable ex) {
-				throw new IllegalArgumentException("Cannot instantiate " + type + " : "
-						+ name, ex);
+				throw new IllegalArgumentException(
+						"Cannot instantiate " + type + " : " + name, ex);
 			}
 		}
 
@@ -416,7 +451,8 @@ public class SpringApplication {
 	 * @see #configureProfiles(ConfigurableEnvironment, String[])
 	 * @see #configurePropertySources(ConfigurableEnvironment, String[])
 	 */
-	protected void configureEnvironment(ConfigurableEnvironment environment, String[] args) {
+	protected void configureEnvironment(ConfigurableEnvironment environment,
+			String[] args) {
 		configurePropertySources(environment, args);
 		configureProfiles(environment, args);
 	}
@@ -432,16 +468,16 @@ public class SpringApplication {
 			String[] args) {
 		MutablePropertySources sources = environment.getPropertySources();
 		if (this.defaultProperties != null && !this.defaultProperties.isEmpty()) {
-			sources.addLast(new MapPropertySource("defaultProperties",
-					this.defaultProperties));
+			sources.addLast(
+					new MapPropertySource("defaultProperties", this.defaultProperties));
 		}
 		if (this.addCommandLineProperties && args.length > 0) {
 			String name = CommandLinePropertySource.COMMAND_LINE_PROPERTY_SOURCE_NAME;
 			if (sources.contains(name)) {
 				PropertySource<?> source = sources.get(name);
 				CompositePropertySource composite = new CompositePropertySource(name);
-				composite.addPropertySource(new SimpleCommandLinePropertySource(name
-						+ "-" + args.hashCode(), args));
+				composite.addPropertySource(new SimpleCommandLinePropertySource(
+						name + "-" + args.hashCode(), args));
 				composite.addPropertySource(source);
 				sources.replace(name, composite);
 			}
@@ -508,14 +544,14 @@ public class SpringApplication {
 		Class<?> contextClass = this.applicationContextClass;
 		if (contextClass == null) {
 			try {
-				contextClass = Class
-						.forName(this.webEnvironment ? DEFAULT_WEB_CONTEXT_CLASS
-								: DEFAULT_CONTEXT_CLASS);
+				contextClass = Class.forName(this.webEnvironment
+						? DEFAULT_WEB_CONTEXT_CLASS : DEFAULT_CONTEXT_CLASS);
 			}
 			catch (ClassNotFoundException ex) {
 				throw new IllegalStateException(
 						"Unable create a default ApplicationContext, "
-								+ "please specify an ApplicationContextClass", ex);
+								+ "please specify an ApplicationContextClass",
+						ex);
 			}
 		}
 		return (ConfigurableApplicationContext) BeanUtils.instantiate(contextClass);
@@ -544,8 +580,8 @@ public class SpringApplication {
 						.setResourceLoader(this.resourceLoader);
 			}
 			if (context instanceof DefaultResourceLoader) {
-				((DefaultResourceLoader) context).setClassLoader(this.resourceLoader
-						.getClassLoader());
+				((DefaultResourceLoader) context)
+						.setClassLoader(this.resourceLoader.getClassLoader());
 			}
 		}
 	}
@@ -596,8 +632,8 @@ public class SpringApplication {
 	 */
 	protected void load(ApplicationContext context, Object[] sources) {
 		if (this.log.isDebugEnabled()) {
-			this.log.debug("Loading source "
-					+ StringUtils.arrayToCommaDelimitedString(sources));
+			this.log.debug(
+					"Loading source " + StringUtils.arrayToCommaDelimitedString(sources));
 		}
 		BeanDefinitionLoader loader = createBeanDefinitionLoader(
 				getBeanDefinitionRegistry(context), sources);
@@ -1018,8 +1054,8 @@ public class SpringApplication {
 			try {
 				List<ExitCodeGenerator> generators = new ArrayList<ExitCodeGenerator>();
 				generators.addAll(Arrays.asList(exitCodeGenerators));
-				generators.addAll(context.getBeansOfType(ExitCodeGenerator.class)
-						.values());
+				generators
+						.addAll(context.getBeansOfType(ExitCodeGenerator.class).values());
 				exitCode = getExitCode(generators);
 			}
 			finally {
diff --git a/spring-boot/src/test/java/org/springframework/boot/SpringApplicationTests.java b/spring-boot/src/test/java/org/springframework/boot/SpringApplicationTests.java
index d020166400b..b0971a90472 100644
--- a/spring-boot/src/test/java/org/springframework/boot/SpringApplicationTests.java
+++ b/spring-boot/src/test/java/org/springframework/boot/SpringApplicationTests.java
@@ -39,6 +39,7 @@ import org.springframework.boot.context.event.ApplicationEnvironmentPreparedEven
 import org.springframework.boot.context.event.ApplicationPreparedEvent;
 import org.springframework.boot.context.event.ApplicationReadyEvent;
 import org.springframework.boot.context.event.ApplicationStartedEvent;
+import org.springframework.boot.test.EnvironmentTestUtils;
 import org.springframework.boot.test.OutputCapture;
 import org.springframework.context.ApplicationContext;
 import org.springframework.context.ApplicationContextAware;
@@ -64,6 +65,7 @@ import org.springframework.core.env.StandardEnvironment;
 import org.springframework.core.io.DefaultResourceLoader;
 import org.springframework.core.io.ResourceLoader;
 import org.springframework.util.StringUtils;
+import org.springframework.web.context.support.StandardServletEnvironment;
 
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.hasItem;
@@ -214,14 +216,13 @@ public class SpringApplicationTests {
 		SpringApplication application = new SpringApplication(ExampleConfig.class);
 		application.setWebEnvironment(false);
 		final AtomicReference<ApplicationContext> reference = new AtomicReference<ApplicationContext>();
-		application
-				.setInitializers(Arrays
-						.asList(new ApplicationContextInitializer<ConfigurableApplicationContext>() {
-							@Override
-							public void initialize(ConfigurableApplicationContext context) {
-								reference.set(context);
-							}
-						}));
+		application.setInitializers(Arrays.asList(
+				new ApplicationContextInitializer<ConfigurableApplicationContext>() {
+					@Override
+					public void initialize(ConfigurableApplicationContext context) {
+						reference.set(context);
+					}
+				}));
 		this.context = application.run("--foo=bar");
 		assertThat(this.context, sameInstance(reference.get()));
 		// Custom initializers do not switch off the defaults
@@ -233,8 +234,8 @@ public class SpringApplicationTests {
 		SpringApplication application = new SpringApplication(ExampleConfig.class);
 		application.setWebEnvironment(false);
 		final AtomicReference<SpringApplication> reference = new AtomicReference<SpringApplication>();
-		class ApplicationReadyEventListener implements
-				ApplicationListener<ApplicationReadyEvent> {
+		class ApplicationReadyEventListener
+				implements ApplicationListener<ApplicationReadyEvent> {
 			@Override
 			public void onApplicationEvent(ApplicationReadyEvent event) {
 				reference.set(event.getSpringApplication());
@@ -268,8 +269,8 @@ public class SpringApplicationTests {
 		SpringApplication application = new SpringApplication(ExampleConfig.class);
 		application.setWebEnvironment(false);
 		final List<ApplicationEvent> events = new ArrayList<ApplicationEvent>();
-		class ApplicationRunningEventListener implements
-				ApplicationListener<ApplicationEvent> {
+		class ApplicationRunningEventListener
+				implements ApplicationListener<ApplicationEvent> {
 			@Override
 			public void onApplicationEvent(ApplicationEvent event) {
 				events.add((event));
@@ -305,7 +306,8 @@ public class SpringApplicationTests {
 
 	@Test
 	public void customEnvironment() throws Exception {
-		TestSpringApplication application = new TestSpringApplication(ExampleConfig.class);
+		TestSpringApplication application = new TestSpringApplication(
+				ExampleConfig.class);
 		application.setWebEnvironment(false);
 		ConfigurableEnvironment environment = new StandardEnvironment();
 		application.setEnvironment(environment);
@@ -315,7 +317,8 @@ public class SpringApplicationTests {
 
 	@Test
 	public void customResourceLoader() throws Exception {
-		TestSpringApplication application = new TestSpringApplication(ExampleConfig.class);
+		TestSpringApplication application = new TestSpringApplication(
+				ExampleConfig.class);
 		application.setWebEnvironment(false);
 		ResourceLoader resourceLoader = new DefaultResourceLoader();
 		application.setResourceLoader(resourceLoader);
@@ -362,9 +365,8 @@ public class SpringApplicationTests {
 		SpringApplication application = new SpringApplication(ExampleConfig.class);
 		application.setWebEnvironment(false);
 		ConfigurableEnvironment environment = new StandardEnvironment();
-		environment.getPropertySources().addFirst(
-				new MapPropertySource("commandLineArgs", Collections
-						.<String, Object>singletonMap("foo", "original")));
+		environment.getPropertySources().addFirst(new MapPropertySource("commandLineArgs",
+				Collections.<String, Object> singletonMap("foo", "original")));
 		application.setEnvironment(environment);
 		this.context = application.run("--foo=bar", "--bar=foo");
 		assertTrue(hasPropertySource(environment, CompositePropertySource.class,
@@ -438,8 +440,8 @@ public class SpringApplicationTests {
 		ConfigurableEnvironment environment = new StandardEnvironment();
 		application.setEnvironment(environment);
 		this.context = application.run("--foo=bar");
-		assertFalse(hasPropertySource(environment, PropertySource.class,
-				"commandLineArgs"));
+		assertFalse(
+				hasPropertySource(environment, PropertySource.class, "commandLineArgs"));
 	}
 
 	@Test
@@ -465,7 +467,8 @@ public class SpringApplicationTests {
 
 	@Test
 	public void wildcardSources() {
-		Object[] sources = { "classpath:org/springframework/boot/sample-${sample.app.test.prop}.xml" };
+		Object[] sources = {
+				"classpath:org/springframework/boot/sample-${sample.app.test.prop}.xml" };
 		TestSpringApplication application = new TestSpringApplication(sources);
 		application.setWebEnvironment(false);
 		this.context = application.run();
@@ -479,8 +482,8 @@ public class SpringApplicationTests {
 
 	@Test
 	public void runComponents() throws Exception {
-		this.context = SpringApplication.run(new Object[] { ExampleWebConfig.class,
-				Object.class }, new String[0]);
+		this.context = SpringApplication.run(
+				new Object[] { ExampleWebConfig.class, Object.class }, new String[0]);
 		assertNotNull(this.context);
 	}
 
@@ -521,7 +524,8 @@ public class SpringApplicationTests {
 
 	@Test
 	public void commandLineArgsApplyToSpringApplication() throws Exception {
-		TestSpringApplication application = new TestSpringApplication(ExampleConfig.class);
+		TestSpringApplication application = new TestSpringApplication(
+				ExampleConfig.class);
 		application.setWebEnvironment(false);
 		this.context = application.run("--spring.main.show_banner=false");
 		assertThat(application.getShowBanner(), is(false));
@@ -582,7 +586,8 @@ public class SpringApplicationTests {
 
 	@Test
 	public void headless() throws Exception {
-		TestSpringApplication application = new TestSpringApplication(ExampleConfig.class);
+		TestSpringApplication application = new TestSpringApplication(
+				ExampleConfig.class);
 		application.setWebEnvironment(false);
 		this.context = application.run();
 		assertThat(System.getProperty("java.awt.headless"), equalTo("true"));
@@ -590,7 +595,8 @@ public class SpringApplicationTests {
 
 	@Test
 	public void headlessFalse() throws Exception {
-		TestSpringApplication application = new TestSpringApplication(ExampleConfig.class);
+		TestSpringApplication application = new TestSpringApplication(
+				ExampleConfig.class);
 		application.setWebEnvironment(false);
 		application.setHeadless(false);
 		this.context = application.run();
@@ -600,7 +606,8 @@ public class SpringApplicationTests {
 	@Test
 	public void headlessSystemPropertyTakesPrecedence() throws Exception {
 		System.setProperty("java.awt.headless", "false");
-		TestSpringApplication application = new TestSpringApplication(ExampleConfig.class);
+		TestSpringApplication application = new TestSpringApplication(
+				ExampleConfig.class);
 		application.setWebEnvironment(false);
 		this.context = application.run();
 		assertThat(System.getProperty("java.awt.headless"), equalTo("false"));
@@ -608,7 +615,8 @@ public class SpringApplicationTests {
 
 	@Test
 	public void getApplicationArgumentsBean() throws Exception {
-		TestSpringApplication application = new TestSpringApplication(ExampleConfig.class);
+		TestSpringApplication application = new TestSpringApplication(
+				ExampleConfig.class);
 		application.setWebEnvironment(false);
 		this.context = application.run("--debug", "spring", "boot");
 		ApplicationArguments args = this.context.getBean(ApplicationArguments.class);
@@ -616,6 +624,31 @@ public class SpringApplicationTests {
 		assertThat(args.containsOption("debug"), equalTo(true));
 	}
 
+	@Test
+	public void webEnvironmentSwitchedOffInListener() throws Exception {
+		TestSpringApplication application = new TestSpringApplication(
+				ExampleConfig.class);
+		application.addListeners(
+				new ApplicationListener<ApplicationEnvironmentPreparedEvent>() {
+
+					@Override
+					public void onApplicationEvent(
+							ApplicationEnvironmentPreparedEvent event) {
+						assertTrue(event
+								.getEnvironment() instanceof StandardServletEnvironment);
+						EnvironmentTestUtils.addEnvironment(event.getEnvironment(),
+								"foo=bar");
+						event.getSpringApplication().setWebEnvironment(false);
+					}
+
+				});
+		this.context = application.run();
+		assertFalse(this.context.getEnvironment() instanceof StandardServletEnvironment);
+		assertEquals("bar", this.context.getEnvironment().getProperty("foo"));
+		assertEquals("test",
+				this.context.getEnvironment().getPropertySources().iterator().next());
+	}
+
 	private boolean hasPropertySource(ConfigurableEnvironment environment,
 			Class<?> propertySourceClass, String name) {
 		for (PropertySource<?> source : environment.getPropertySources()) {
@@ -637,7 +670,8 @@ public class SpringApplicationTests {
 
 	public static class SpyApplicationContext extends AnnotationConfigApplicationContext {
 
-		ConfigurableApplicationContext applicationContext = spy(new AnnotationConfigApplicationContext());
+		ConfigurableApplicationContext applicationContext = spy(
+				new AnnotationConfigApplicationContext());
 
 		@Override
 		public void registerShutdownHook() {
@@ -784,8 +818,8 @@ public class SpringApplicationTests {
 
 	}
 
-	private static class TestCommandLineRunner extends AbstractTestRunner implements
-			CommandLineRunner {
+	private static class TestCommandLineRunner extends AbstractTestRunner
+			implements CommandLineRunner {
 
 		TestCommandLineRunner(int order, String... expectedBefore) {
 			super(order, expectedBefore);
@@ -798,8 +832,8 @@ public class SpringApplicationTests {
 
 	}
 
-	private static class TestApplicationRunner extends AbstractTestRunner implements
-			ApplicationRunner {
+	private static class TestApplicationRunner extends AbstractTestRunner
+			implements ApplicationRunner {
 
 		TestApplicationRunner(int order, String... expectedBefore) {
 			super(order, expectedBefore);