From 44a8e91b9e2f6182b251999957740a858eb569c8 Mon Sep 17 00:00:00 2001
From: Stephane Nicoll <snicoll@vmware.com>
Date: Tue, 31 May 2022 12:34:19 +0200
Subject: [PATCH] Prevent main method to continue once the AOT context has been
 optimized

This commit makes sure that processing of the main method upon
completion of SpringApplication#run. Previously, any instructions in
the user's main method were invoked, which is not suitable to build-time
processing.

Closes gh-31219
---
 .../org/springframework/boot/AotProcessingHook.java | 13 +++++++++++++
 .../java/org/springframework/boot/AotProcessor.java |  9 +++++++++
 .../org/springframework/boot/AotProcessorTests.java |  6 ++++++
 3 files changed, 28 insertions(+)

diff --git a/spring-boot-project/spring-boot/src/main/java/org/springframework/boot/AotProcessingHook.java b/spring-boot-project/spring-boot/src/main/java/org/springframework/boot/AotProcessingHook.java
index 5d6feb62b6..eee7f07c51 100644
--- a/spring-boot-project/spring-boot/src/main/java/org/springframework/boot/AotProcessingHook.java
+++ b/spring-boot-project/spring-boot/src/main/java/org/springframework/boot/AotProcessingHook.java
@@ -39,8 +39,21 @@ class AotProcessingHook implements Hook {
 		return false;
 	}
 
+	@Override
+	public void postRun(SpringApplication application, ConfigurableApplicationContext context) {
+		throw new MainMethodSilentExitException();
+	}
+
 	GenericApplicationContext getApplicationContext() {
 		return this.context;
 	}
 
+	/**
+	 * Internal exception used to prevent main method to continue once
+	 * {@code SpringApplication#run} completes.
+	 */
+	static class MainMethodSilentExitException extends RuntimeException {
+
+	}
+
 }
diff --git a/spring-boot-project/spring-boot/src/main/java/org/springframework/boot/AotProcessor.java b/spring-boot-project/spring-boot/src/main/java/org/springframework/boot/AotProcessor.java
index ad8a3d4a8f..e8bc5016f0 100644
--- a/spring-boot-project/spring-boot/src/main/java/org/springframework/boot/AotProcessor.java
+++ b/spring-boot-project/spring-boot/src/main/java/org/springframework/boot/AotProcessor.java
@@ -17,6 +17,7 @@
 package org.springframework.boot;
 
 import java.io.IOException;
+import java.lang.reflect.InvocationTargetException;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
@@ -34,6 +35,7 @@ import org.springframework.aot.hint.ExecutableMode;
 import org.springframework.aot.hint.RuntimeHints;
 import org.springframework.aot.hint.TypeReference;
 import org.springframework.aot.nativex.FileNativeConfigurationWriter;
+import org.springframework.boot.AotProcessingHook.MainMethodSilentExitException;
 import org.springframework.context.aot.ApplicationContextAotGenerator;
 import org.springframework.context.support.GenericApplicationContext;
 import org.springframework.javapoet.ClassName;
@@ -124,6 +126,13 @@ public class AotProcessor {
 		try {
 			this.application.getMethod("main", String[].class).invoke(null, new Object[] { this.applicationArgs });
 		}
+		catch (InvocationTargetException ex) {
+			Throwable targetException = ex.getTargetException();
+			if (!(targetException instanceof MainMethodSilentExitException)) {
+				throw (targetException instanceof RuntimeException runtimeEx) ? runtimeEx
+						: new RuntimeException(targetException);
+			}
+		}
 		catch (Exception ex) {
 			throw new RuntimeException(ex);
 		}
diff --git a/spring-boot-project/spring-boot/src/test/java/org/springframework/boot/AotProcessorTests.java b/spring-boot-project/spring-boot/src/test/java/org/springframework/boot/AotProcessorTests.java
index 840cd6c582..cec676da1a 100644
--- a/spring-boot-project/spring-boot/src/test/java/org/springframework/boot/AotProcessorTests.java
+++ b/spring-boot-project/spring-boot/src/test/java/org/springframework/boot/AotProcessorTests.java
@@ -41,6 +41,7 @@ class AotProcessorTests {
 	@BeforeEach
 	void setup() {
 		SampleApplication.argsHolder = null;
+		SampleApplication.postRunInvoked = false;
 	}
 
 	@Test
@@ -50,6 +51,7 @@ class AotProcessorTests {
 				directory.resolve("resource"), directory.resolve("class"), "com.example", "example");
 		processor.process();
 		assertThat(SampleApplication.argsHolder).isEqualTo(arguments);
+		assertThat(SampleApplication.postRunInvoked).isFalse();
 		assertThat(directory).satisfies(hasGeneratedAssetsForSampleApplication());
 	}
 
@@ -69,6 +71,7 @@ class AotProcessorTests {
 				directory.resolve("class").toString(), "com.example", "example", "1", "2" };
 		AotProcessor.main(mainArguments);
 		assertThat(SampleApplication.argsHolder).containsExactly("1", "2");
+		assertThat(SampleApplication.postRunInvoked).isFalse();
 		assertThat(directory).satisfies(hasGeneratedAssetsForSampleApplication());
 	}
 
@@ -125,9 +128,12 @@ class AotProcessorTests {
 
 		public static String[] argsHolder;
 
+		public static boolean postRunInvoked;
+
 		public static void main(String[] args) {
 			argsHolder = args;
 			SpringApplication.run(SampleApplication.class, args);
+			postRunInvoked = true;
 		}
 
 	}