Handle hints for CGLIB proxies consistently

This commit makes sure that hints are registered for CGLIB proxies even if the proxy itself is not created. This typically happens when AOT runs on an existing classpath, and a previous run already created the proxy. Closes gh-29295
spring-projects · Oct 10, 2022 · 4eca87b · 4eca87b
1 parent e8ce86a
commit 4eca87b
Show file tree

Hide file tree

Showing 7 changed files with 142 additions and 52 deletions.
diff --git a/...context/src/main/java/org/springframework/context/aot/ApplicationContextAotGenerator.java b/...context/src/main/java/org/springframework/context/aot/ApplicationContextAotGenerator.java
@@ -50,7 +50,7 @@ public class ApplicationContextAotGenerator {
 	 */
 	public ClassName processAheadOfTime(GenericApplicationContext applicationContext,
 			GenerationContext generationContext) {
-		return withGeneratedClassHandler(new GeneratedClassHandler(generationContext), () -> {
+		return withCglibClassHandler(new CglibClassHandler(generationContext), () -> {
 			applicationContext.refreshForAotProcessing(generationContext.getRuntimeHints());
 			DefaultListableBeanFactory beanFactory = applicationContext.getDefaultListableBeanFactory();
 			ApplicationContextInitializationCodeGenerator codeGenerator =
@@ -60,12 +60,14 @@ public ClassName processAheadOfTime(GenericApplicationContext applicationContext
 		});
 	}
 
-	private <T> T withGeneratedClassHandler(GeneratedClassHandler generatedClassHandler, Supplier<T> task) {
+	private <T> T withCglibClassHandler(CglibClassHandler cglibClassHandler, Supplier<T> task) {
 		try {
-			ReflectUtils.setGeneratedClassHandler(generatedClassHandler);
+			ReflectUtils.setLoadedClassHandler(cglibClassHandler::handleLoadedClass);
+			ReflectUtils.setGeneratedClassHandler(cglibClassHandler::handleGeneratedClass);
 			return task.get();
 		}
 		finally {
+			ReflectUtils.setLoadedClassHandler(null);
 			ReflectUtils.setGeneratedClassHandler(null);
 		}
 	}

diff --git a/...rk/context/aot/GeneratedClassHandler.java → ...mework/context/aot/CglibClassHandler.java b/...rk/context/aot/GeneratedClassHandler.java → ...mework/context/aot/CglibClassHandler.java
@@ -30,46 +30,48 @@
 import org.springframework.core.io.ByteArrayResource;
 
 /**
- * Handle generated classes by adding them to a {@link GenerationContext},
+ * Handle CGLIB classes by adding them to a {@link GenerationContext},
  * and register the necessary hints so that they can be instantiated.
  *
  * @author Stephane Nicoll
  * @see ReflectUtils#setGeneratedClassHandler(BiConsumer)
+ * @see ReflectUtils#setLoadedClassHandler(Consumer)
  */
-class GeneratedClassHandler implements BiConsumer<String, byte[]> {
+class CglibClassHandler {
 
-	private static final Consumer<Builder> asCglibProxy = hint ->
-			hint.withMembers(MemberCategory.INVOKE_DECLARED_CONSTRUCTORS,
-					MemberCategory.INVOKE_DECLARED_METHODS,
-					MemberCategory.DECLARED_FIELDS);
-
-	private static final Consumer<Builder> asCglibProxyTargetType = hint ->
-			hint.withMembers(MemberCategory.INTROSPECT_DECLARED_CONSTRUCTORS,
-					MemberCategory.INVOKE_DECLARED_METHODS);
+	private static final Consumer<Builder> instantiateCglibProxy = hint ->
+			hint.withMembers(MemberCategory.INVOKE_DECLARED_CONSTRUCTORS);
 
 	private final RuntimeHints runtimeHints;
 
 	private final GeneratedFiles generatedFiles;
 
-	GeneratedClassHandler(GenerationContext generationContext) {
+	CglibClassHandler(GenerationContext generationContext) {
 		this.runtimeHints = generationContext.getRuntimeHints();
 		this.generatedFiles = generationContext.getGeneratedFiles();
 	}
 
-	@Override
-	public void accept(String className, byte[] content) {
-		this.runtimeHints.reflection().registerType(TypeReference.of(className), asCglibProxy)
-				.registerType(TypeReference.of(getTargetTypeClassName(className)), asCglibProxyTargetType);
-		String path = className.replace(".", "/") + ".class";
+	/**
+	 * Handle the specified generated CGLIB class.
+	 * @param cglibClassName the name of the generated class
+	 * @param content the bytecode of the generated class
+	 */
+	public void handleGeneratedClass(String cglibClassName, byte[] content) {
+		registerHints(TypeReference.of(cglibClassName));
+		String path = cglibClassName.replace(".", "/") + ".class";
 		this.generatedFiles.addFile(Kind.CLASS, path, new ByteArrayResource(content));
 	}
 
-	private String getTargetTypeClassName(String proxyClassName) {
-		int index = proxyClassName.indexOf("$$SpringCGLIB$$");
-		if (index == -1) {
-			throw new IllegalArgumentException("Failed to extract target type from " + proxyClassName);
-		}
-		return proxyClassName.substring(0, index);
+	/**
+	 * Handle the specified loaded CGLIB class.
+	 * @param cglibClass a cglib class that has been loaded
+	 */
+	public void handleLoadedClass(Class<?> cglibClass) {
+		registerHints(TypeReference.of(cglibClass));
+	}
+
+	private void registerHints(TypeReference cglibTypeReference) {
+		this.runtimeHints.reflection().registerType(cglibTypeReference, instantiateCglibProxy);
 	}
 
 }
diff --git a/...-context/src/main/java/org/springframework/context/support/GenericApplicationContext.java b/...-context/src/main/java/org/springframework/context/support/GenericApplicationContext.java
@@ -21,9 +21,12 @@
 import java.lang.reflect.Proxy;
 import java.util.List;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Consumer;
 import java.util.function.Supplier;
 
+import org.springframework.aot.hint.MemberCategory;
 import org.springframework.aot.hint.RuntimeHints;
+import org.springframework.aot.hint.TypeHint.Builder;
 import org.springframework.beans.BeanUtils;
 import org.springframework.beans.BeansException;
 import org.springframework.beans.factory.BeanDefinitionStoreException;
@@ -46,6 +49,7 @@
 import org.springframework.core.metrics.ApplicationStartup;
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
+import org.springframework.util.ClassUtils;
 
 /**
  * Generic ApplicationContext implementation that holds a single internal
@@ -107,6 +111,16 @@
  */
 public class GenericApplicationContext extends AbstractApplicationContext implements BeanDefinitionRegistry {
 
+	private static final Consumer<Builder> asCglibProxy = hint ->
+			hint.withMembers(MemberCategory.INVOKE_DECLARED_CONSTRUCTORS,
+					MemberCategory.INVOKE_DECLARED_METHODS,
+					MemberCategory.DECLARED_FIELDS);
+
+	private static final Consumer<Builder> asCglibProxyTarget = hint ->
+			hint.withMembers(MemberCategory.INTROSPECT_DECLARED_CONSTRUCTORS,
+					MemberCategory.INVOKE_DECLARED_METHODS);
+
+
 	private final DefaultListableBeanFactory beanFactory;
 
 	@Nullable
@@ -433,6 +447,14 @@ private void preDetermineBeanTypes(RuntimeHints runtimeHints) {
 					if (Proxy.isProxyClass(beanType)) {
 						runtimeHints.proxies().registerJdkProxy(beanType.getInterfaces());
 					}
+					else {
+						Class<?> userClass = ClassUtils.getUserClass(beanType);
+						if (userClass != beanType) {
+							runtimeHints.reflection()
+									.registerType(beanType, asCglibProxy)
+									.registerType(userClass, asCglibProxyTarget);
+						}
+					}
 				}
 			}
 		}

diff --git a/.../java/org/springframework/context/annotation/AnnotationConfigApplicationContextTests.java b/.../java/org/springframework/context/annotation/AnnotationConfigApplicationContextTests.java
@@ -21,7 +21,10 @@
 
 import org.junit.jupiter.api.Test;
 
+import org.springframework.aot.hint.MemberCategory;
 import org.springframework.aot.hint.RuntimeHints;
+import org.springframework.aot.hint.TypeReference;
+import org.springframework.aot.hint.predicate.RuntimeHintsPredicates;
 import org.springframework.beans.factory.FactoryBean;
 import org.springframework.beans.factory.NoSuchBeanDefinitionException;
 import org.springframework.beans.factory.annotation.Autowired;
@@ -31,6 +34,8 @@
 import org.springframework.context.annotation6.ComponentForScanning;
 import org.springframework.context.annotation6.ConfigForScanning;
 import org.springframework.context.annotation6.Jsr330NamedForScanning;
+import org.springframework.context.testfixture.context.annotation.CglibConfiguration;
+import org.springframework.context.testfixture.context.annotation.LambdaBeanConfiguration;
 import org.springframework.core.ResolvableType;
 import org.springframework.util.ObjectUtils;
 
@@ -450,6 +455,40 @@ void refreshForAotCanInstantiateBeanWithFieldAutowiredApplicationContext() {
 		assertThat(bean.applicationContext).isSameAs(context);
 	}
 
+	@Test
+	void refreshForAotRegisterHintsForCglibProxy() {
+		AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext();
+		context.register(CglibConfiguration.class);
+		RuntimeHints runtimeHints = new RuntimeHints();
+		context.refreshForAotProcessing(runtimeHints);
+		TypeReference cglibType = TypeReference.of(CglibConfiguration.class.getName() + "$$SpringCGLIB$$0");
+		assertThat(RuntimeHintsPredicates.reflection().onType(cglibType)
+				.withMemberCategories(MemberCategory.INVOKE_DECLARED_CONSTRUCTORS,
+						MemberCategory.INVOKE_DECLARED_METHODS, MemberCategory.DECLARED_FIELDS))
+				.accepts(runtimeHints);
+	}
+
+	@Test
+	void refreshForAotRegisterHintsForTargetOfCglibProxy() {
+		AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext();
+		context.register(CglibConfiguration.class);
+		RuntimeHints runtimeHints = new RuntimeHints();
+		context.refreshForAotProcessing(runtimeHints);
+		assertThat(RuntimeHintsPredicates.reflection().onType(TypeReference.of(CglibConfiguration.class))
+				.withMemberCategories(MemberCategory.INTROSPECT_DECLARED_CONSTRUCTORS,
+						MemberCategory.INVOKE_DECLARED_METHODS))
+				.accepts(runtimeHints);
+	}
+
+	@Test
+	void refreshForAotRegisterDoesNotConsiderLambdaBeanAsCglibProxy() {
+		AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext();
+		context.register(LambdaBeanConfiguration.class);
+		RuntimeHints runtimeHints = new RuntimeHints();
+		context.refreshForAotProcessing(runtimeHints);
+		assertThat(runtimeHints.reflection().typeHints()).isEmpty();
+	}
+
 
 	@Configuration
 	static class Config {

diff --git a/...xt/src/test/java/org/springframework/context/aot/ApplicationContextAotGeneratorTests.java b/...xt/src/test/java/org/springframework/context/aot/ApplicationContextAotGeneratorTests.java
@@ -297,10 +297,15 @@ void processAheadOfTimeWhenHasCglibProxyWriteProxyAndGenerateReflectionHints() t
 			GenericApplicationContext applicationContext = new AnnotationConfigApplicationContext();
 			applicationContext.registerBean(CglibConfiguration.class);
 			TestGenerationContext context = processAheadOfTime(applicationContext);
-			String proxyClassName = CglibConfiguration.class.getName() + "$$SpringCGLIB$$0";
+			isRegisteredCglibClass(context, CglibConfiguration.class.getName() + "$$SpringCGLIB$$0");
+			isRegisteredCglibClass(context, CglibConfiguration.class.getName() + "$$SpringCGLIB$$1");
+			isRegisteredCglibClass(context, CglibConfiguration.class.getName() + "$$SpringCGLIB$$2");
+		}
+
+		private void isRegisteredCglibClass(TestGenerationContext context, String cglibClassName) throws IOException {
 			assertThat(context.getGeneratedFiles()
-					.getGeneratedFileContent(Kind.CLASS, proxyClassName.replace('.', '/') + ".class")).isNotNull();
-			assertThat(RuntimeHintsPredicates.reflection().onType(TypeReference.of(proxyClassName))
+					.getGeneratedFileContent(Kind.CLASS, cglibClassName.replace('.', '/') + ".class")).isNotNull();
+			assertThat(RuntimeHintsPredicates.reflection().onType(TypeReference.of(cglibClassName))
 					.withMemberCategory(MemberCategory.INVOKE_DECLARED_CONSTRUCTORS)).accepts(context.getRuntimeHints());
 		}
 

diff --git a/...ntext/aot/GeneratedClassHandlerTests.java → ...k/context/aot/CglibClassHandlerTests.java b/...ntext/aot/GeneratedClassHandlerTests.java → ...k/context/aot/CglibClassHandlerTests.java
@@ -30,57 +30,46 @@
 import org.springframework.core.io.InputStreamSource;
 
 import static org.assertj.core.api.Assertions.assertThat;
-import static org.assertj.core.api.Assertions.assertThatIllegalArgumentException;
 
 /**
- * Tests for {@link GeneratedClassHandler}.
+ * Tests for {@link CglibClassHandler}.
  *
  * @author Stephane Nicoll
  */
-class GeneratedClassHandlerTests {
+class CglibClassHandlerTests {
 
 	private static final byte[] TEST_CONTENT = new byte[] { 'a' };
 
 	private final TestGenerationContext generationContext;
 
-	private final GeneratedClassHandler handler;
+	private final CglibClassHandler handler;
 
-	public GeneratedClassHandlerTests() {
+	public CglibClassHandlerTests() {
 		this.generationContext = new TestGenerationContext();
-		this.handler = new GeneratedClassHandler(this.generationContext);
+		this.handler = new CglibClassHandler(this.generationContext);
 	}
 
 	@Test
-	void handlerGenerateRuntimeHintsForProxy() {
+	void handlerGeneratedClassCreatesRuntimeHintsForProxy() {
 		String className = "com.example.Test$$SpringCGLIB$$0";
-		this.handler.accept(className, TEST_CONTENT);
+		this.handler.handleGeneratedClass(className, TEST_CONTENT);
 		assertThat(RuntimeHintsPredicates.reflection().onType(TypeReference.of(className))
-				.withMemberCategories(MemberCategory.INVOKE_DECLARED_CONSTRUCTORS,
-						MemberCategory.INVOKE_DECLARED_METHODS, MemberCategory.DECLARED_FIELDS))
+				.withMemberCategories(MemberCategory.INVOKE_DECLARED_CONSTRUCTORS))
 				.accepts(this.generationContext.getRuntimeHints());
 	}
 
 	@Test
-	void handlerGenerateRuntimeHintsForTargetType() {
-		String className = "com.example.Test$$SpringCGLIB$$0";
-		this.handler.accept(className, TEST_CONTENT);
-		assertThat(RuntimeHintsPredicates.reflection().onType(TypeReference.of("com.example.Test"))
-				.withMemberCategories(MemberCategory.INTROSPECT_DECLARED_CONSTRUCTORS,
-						MemberCategory.INVOKE_DECLARED_METHODS))
+	void handlerLoadedClassCreatesRuntimeHintsForProxy() {
+		this.handler.handleLoadedClass(CglibClassHandler.class);
+		assertThat(RuntimeHintsPredicates.reflection().onType(CglibClassHandler.class)
+				.withMemberCategories(MemberCategory.INVOKE_DECLARED_CONSTRUCTORS))
 				.accepts(this.generationContext.getRuntimeHints());
 	}
 
-	@Test
-	void handlerFailsWithInvalidProxyClassName() {
-		String className = "com.example.Test$$AnotherProxy$$0";
-		assertThatIllegalArgumentException().isThrownBy(() -> this.handler.accept(className, TEST_CONTENT))
-				.withMessageContaining("Failed to extract target type");
-	}
-
 	@Test
 	void handlerRegisterGeneratedClass() throws IOException {
 		String className = "com.example.Test$$SpringCGLIB$$0";
-		this.handler.accept(className, TEST_CONTENT);
+		this.handler.handleGeneratedClass(className, TEST_CONTENT);
 		InMemoryGeneratedFiles generatedFiles = this.generationContext.getGeneratedFiles();
 		assertThat(generatedFiles.getGeneratedFiles(Kind.SOURCE)).isEmpty();
 		assertThat(generatedFiles.getGeneratedFiles(Kind.RESOURCE)).isEmpty();

diff --git a/...a/org/springframework/context/testfixture/context/annotation/LambdaBeanConfiguration.java b/...a/org/springframework/context/testfixture/context/annotation/LambdaBeanConfiguration.java
@@ -0,0 +1,31 @@
+/*
+ * Copyright 2002-2022 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.context.testfixture.context.annotation;
+
+import org.springframework.beans.factory.config.BeanFactoryPostProcessor;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+
+@Configuration(proxyBeanMethods = false)
+public class LambdaBeanConfiguration {
+
+	@Bean
+	public static BeanFactoryPostProcessor lambaBean() {
+		return beanFactory -> {};
+	}
+
+}