Change SecurityProblemSupport bean after AdviceTrait created

problem-spring-web-autoconfigure/src/main/java/org/zalando/problem/spring/web/autoconfigure/security/ProblemSecurityAutoConfiguration.java

@@ -1,7 +1,7 @@
 package org.zalando.problem.spring.web.autoconfigure.security;
 
 import org.apiguardian.api.API;
-import org.springframework.beans.factory.ObjectProvider;
+import org.springframework.beans.factory.annotation.Qualifier;
 import org.springframework.boot.autoconfigure.AutoConfigureAfter;
 import org.springframework.boot.autoconfigure.AutoConfigureBefore;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
@@ -12,9 +12,9 @@
 import org.springframework.boot.autoconfigure.security.servlet.SecurityAutoConfiguration;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
-import org.springframework.context.annotation.Import;
 import org.springframework.security.config.annotation.web.WebSecurityConfigurer;
 import org.springframework.security.config.annotation.web.configuration.WebSecurityConfiguration;
+import org.springframework.web.servlet.HandlerExceptionResolver;
 import org.zalando.problem.spring.web.advice.AdviceTrait;
 import org.zalando.problem.spring.web.advice.security.SecurityProblemSupport;
 import org.zalando.problem.spring.web.autoconfigure.ProblemAutoConfiguration;
@@ -26,7 +26,6 @@
 @ConditionalOnWebApplication(type = Type.SERVLET)
 @ConditionalOnClass(WebSecurityConfigurer.class)
 @ConditionalOnBean(WebSecurityConfiguration.class)
-@Import(SecurityProblemSupport.class)
 @AutoConfigureAfter(SecurityAutoConfiguration.class)
 @AutoConfigureBefore(ProblemAutoConfiguration.class)
 public class ProblemSecurityAutoConfiguration {
@@ -38,7 +37,13 @@ public AdviceTrait securityExceptionHandling() {
     }
 
     @Bean
-    public ProblemSecurityBeanPostProcessor problemSecurityBeanPostProcessor(final ObjectProvider<SecurityProblemSupport> securityProblemSupport) {
-        return new ProblemSecurityBeanPostProcessor(securityProblemSupport);
+    public SecurityProblemSupport securityProblemSupport(@Qualifier("handlerExceptionResolver") final HandlerExceptionResolver resolver,
+                                                         final AdviceTrait adviceTrait) {
+        return new SecurityProblemSupport(resolver);
+    }
+
+    @Bean
+    public ProblemSecurityBeanPostProcessor problemSecurityBeanPostProcessor() {
+        return new ProblemSecurityBeanPostProcessor();
     }
 }

problem-spring-web-autoconfigure/src/main/java/org/zalando/problem/spring/web/autoconfigure/security/ProblemSecurityBeanPostProcessor.java

@@ -1,34 +1,32 @@
 package org.zalando.problem.spring.web.autoconfigure.security;
 
-import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.beans.BeansException;
 import org.springframework.beans.factory.BeanCreationException;
-import org.springframework.beans.factory.ObjectProvider;
 import org.springframework.beans.factory.config.BeanPostProcessor;
 import org.springframework.security.config.annotation.web.builders.HttpSecurity;
 import org.springframework.util.ClassUtils;
-import org.zalando.problem.spring.web.advice.security.SecurityProblemSupport;
+
+import java.util.Objects;
 
 @Slf4j
-@RequiredArgsConstructor
 public class ProblemSecurityBeanPostProcessor implements BeanPostProcessor {
-    private final ObjectProvider<SecurityProblemSupport> securityProblemSupport;
 
     @Override
     public Object postProcessAfterInitialization(final Object bean, final String beanName) throws BeansException {
-        if (ClassUtils.isAssignableValue(HttpSecurity.class, bean)) {
-            securityProblemSupport.ifAvailable(support -> register((HttpSecurity) bean, beanName, support));
+        if (!ClassUtils.isAssignableValue(HttpSecurity.class, bean)) {
+            return bean;
         }
-        return bean;
-    }
-
-    private void register(final HttpSecurity http, final String beanName, final SecurityProblemSupport support) {
+        final HttpSecurity http = (HttpSecurity) bean;
         try {
-            http.exceptionHandling().authenticationEntryPoint(support).accessDeniedHandler(support);
+            final ProblemHttpConfigurer configurer = http.getConfigurer(ProblemHttpConfigurer.class);
+            if (Objects.isNull(configurer)) {
+                http.apply(new ProblemHttpConfigurer());
+            }
         } catch (final Exception cause) {
-            throw new BeanCreationException(beanName, cause);
+            throw new BeanCreationException("Fail to configure HttpSecurity's exceptionHandling", cause);
         }
-        log.info("Register HttpSecurity's exceptionHandling");
+        return bean;
     }
+
 }

problem-spring-web-autoconfigure/src/test/java/org/zalando/problem/spring/web/autoconfigure/security/ProblemSecurityBeanPostProcessorTest.java

@@ -12,68 +12,58 @@
 import org.zalando.problem.spring.web.advice.security.SecurityProblemSupport;
 
 import java.util.HashMap;
+import java.util.Map;
 import java.util.function.Consumer;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.verify;
 
 
 class ProblemSecurityBeanPostProcessorTest {
 
-    private ObjectProvider<SecurityProblemSupport> provider;
     private ProblemSecurityBeanPostProcessor processor;
+    private HttpSecurity http;
 
     @SuppressWarnings("unchecked")
     @BeforeEach
     void setUp() {
-        provider = mock(ObjectProvider.class);
-        processor = new ProblemSecurityBeanPostProcessor(provider);
+        processor = new ProblemSecurityBeanPostProcessor();
+        final ObjectPostProcessor<Object> objectPostProcessor = mock(ObjectPostProcessor.class);
+        final AuthenticationManagerBuilder authenticationManagerBuilder = mock(AuthenticationManagerBuilder.class);
+        http = spy(new HttpSecurity(objectPostProcessor, authenticationManagerBuilder, new HashMap<>()));
     }
 
-    @SuppressWarnings("unchecked")
     @Test
     void shouldBeConfigured() throws Exception {
-        final ObjectPostProcessor<Object> objectPostProcessor = mock(ObjectPostProcessor.class);
-        final AuthenticationManagerBuilder authenticationManagerBuilder = mock(AuthenticationManagerBuilder.class);
-        final SecurityProblemSupport securityProblemSupport = mock(SecurityProblemSupport.class);
-        final HttpSecurity http = new HttpSecurity(objectPostProcessor, authenticationManagerBuilder, new HashMap<>());
-
-        doAnswer(answer -> {
-            final Consumer<SecurityProblemSupport> consumer = (Consumer<SecurityProblemSupport>) answer.getArguments()[0];
-            consumer.accept(securityProblemSupport);
-            return null;
-        }).when(provider).ifAvailable(any(Consumer.class));
-
         assertThat(processor.postProcessAfterInitialization(http, "httpSecurity")).isEqualTo(http);
-        assertThat(ReflectionTestUtils.getField(http.exceptionHandling(), "authenticationEntryPoint")).isEqualTo(securityProblemSupport);
-        assertThat(ReflectionTestUtils.getField(http.exceptionHandling(), "accessDeniedHandler")).isEqualTo(securityProblemSupport);
+        verify(http).apply(any(ProblemHttpConfigurer.class));
     }
 
     @Test
-    void shouldBeNotConfiguredWhenBeanIsNotHttpSecurity() {
-        assertThat(processor.postProcessAfterInitialization("test", "notHttpSecurity")).isEqualTo("test");
+    void shouldBeNotConfiguredWhenAlreadyConfigured() throws Exception {
+        doReturn(new ProblemHttpConfigurer()).when(http).getConfigurer(ProblemHttpConfigurer.class);
+
+        assertThat(processor.postProcessAfterInitialization(http, "httpSecurity")).isEqualTo(http);
+        verify(http, never()).apply(any(ProblemHttpConfigurer.class));
     }
 
-    @SuppressWarnings("unchecked")
     @Test
-    void shouldBeThrowBeanCreationExceptionWhenRegisterExceptionHandling() throws Exception {
-        final ObjectPostProcessor<Object> objectPostProcessor = mock(ObjectPostProcessor.class);
-        final AuthenticationManagerBuilder authenticationManagerBuilder = mock(AuthenticationManagerBuilder.class);
-        final SecurityProblemSupport securityProblemSupport = mock(SecurityProblemSupport.class);
-        final HttpSecurity http = spy(new HttpSecurity(objectPostProcessor, authenticationManagerBuilder, new HashMap<>()));
+    void shouldBeFailWhenThrowException() {
+        doThrow(RuntimeException.class).when(http).getConfigurer(ProblemHttpConfigurer.class);
 
-        doAnswer(answer -> {
-            final Consumer<SecurityProblemSupport> consumer = (Consumer<SecurityProblemSupport>) answer.getArguments()[0];
-            consumer.accept(securityProblemSupport);
-            return null;
-        }).when(provider).ifAvailable(any(Consumer.class));
-        doThrow(new BeanInitializationException("test")).when(http).exceptionHandling();
+        assertThatThrownBy(() -> processor.postProcessAfterInitialization(http, "httpSecurity")).isInstanceOf(BeanCreationException.class);
+    }
 
-        assertThatThrownBy(() -> processor.postProcessAfterInitialization(http, "httpSecurity")).isExactlyInstanceOf(BeanCreationException.class);
+    @Test
+    void shouldBeNotConfiguredWhenBeanIsNotHttpSecurity() {
+        assertThat(processor.postProcessAfterInitialization("test", "notHttpSecurity")).isEqualTo("test");
     }
 }