OAuth2TokenCustomizer not respecting @Primary

[MatthiasDrewsCS]

Describe the bug
I am opening a new issue, because it looks like my last comment on the already closed #2183 did not reach its audience.

My use-case is, in our authorization-server implementation we are supporting multiple issuers. We are following the suggested pattern for multi-tenancy described in https://docs.spring.io/spring-authorization-server/reference/guides/how-to-multitenancy.html. However we like to have the tenant-specific components also to have as beans, to have things like auto-wiring etc.

More specific, we want to have a dedicated implementation of OAuth2TokenCustomizer<JwtEncodingContext> per tenant.

Doing this, and registering each implementation as a bean we run into following error:

Failed to instantiate [org.springframework.security.web.SecurityFilterChain]: Factory method 'authorizationServerSecurityFilterChain' threw exception with message: No qualifying bean of type 'org.springframework.security.oauth2.server.authorization.token.OAuth2TokenCustomizer<org.springframework.security.oauth2.server.authorization.token.JwtEncodingContext>' available: expected single matching bean but found 3: clientAssertionOAuth2TokenCustomizer,customClaimsOAuth2TokenCustomizer,tenantAwareOAuth2TokenCustomizer

And the root cause:

Caused by: org.springframework.beans.factory.NoUniqueBeanDefinitionException: No qualifying bean of type 'org.springframework.security.oauth2.server.authorization.token.OAuth2TokenCustomizer<org.springframework.security.oauth2.server.authorization.token.JwtEncodingContext>' available: expected single matching bean but found 3: clientAssertionOAuth2TokenCustomizer,customClaimsOAuth2TokenCustomizer,tenantAwareOAuth2TokenCustomizer
	at app//org.springframework.security.oauth2.server.authorization.config.annotation.web.configurers.OAuth2ConfigurerUtils.getOptionalBean(OAuth2ConfigurerUtils.java:241)
	at app//org.springframework.security.oauth2.server.authorization.config.annotation.web.configurers.OAuth2ConfigurerUtils.getJwtCustomizer(OAuth2ConfigurerUtils.java:173)
	at app//org.springframework.security.oauth2.server.authorization.config.annotation.web.configurers.OAuth2ConfigurerUtils.getJwtGenerator(OAuth2ConfigurerUtils.java:131)
	at app//org.springframework.security.oauth2.server.authorization.config.annotation.web.configurers.OAuth2ConfigurerUtils.getTokenGenerator(OAuth2ConfigurerUtils.java:108)
	at app//org.springframework.security.oauth2.server.authorization.config.annotation.web.configurers.OAuth2TokenEndpointConfigurer.createDefaultAuthenticationProviders(OAuth2TokenEndpointConfigurer.java:251)
	at app//org.springframework.security.oauth2.server.authorization.config.annotation.web.configurers.OAuth2TokenEndpointConfigurer.init(OAuth2TokenEndpointConfigurer.java:194)
	at app//org.springframework.security.oauth2.server.authorization.config.annotation.web.configurers.OAuth2AuthorizationServerConfigurer.lambda$init$5(OAuth2AuthorizationServerConfigurer.java:367)
	at java.base@21.0.7/java.util.LinkedHashMap$LinkedValues.forEach(LinkedHashMap.java:833)
	at app//org.springframework.security.oauth2.server.authorization.config.annotation.web.configurers.OAuth2AuthorizationServerConfigurer.init(OAuth2AuthorizationServerConfigurer.java:366)
	at app//org.springframework.security.oauth2.server.authorization.config.annotation.web.configurers.OAuth2AuthorizationServerConfigurer.init(OAuth2AuthorizationServerConfigurer.java:86)
	at app//org.springframework.security.config.annotation.AbstractConfiguredSecurityBuilder.init(AbstractConfiguredSecurityBuilder.java:388)
	at app//org.springframework.security.config.annotation.AbstractConfiguredSecurityBuilder.doBuild(AbstractConfiguredSecurityBuilder.java:350)
	at app//org.springframework.security.config.annotation.AbstractSecurityBuilder.build(AbstractSecurityBuilder.java:38)

Expected behavior
OAuth2ConfigurerUtils should pickup the bean annotated with @Primary

[jgrandja]

@MatthiasDrewsCS Can you provide me a sample configuration of all the @Bean's you would like to register, for example, OAuth2TokenGenerator, OAuth2TokenCustomizer, OAuth2AuthorizationService, etc.

This will help me understand how you plan on auto-wiring each @Bean into the correct component.

[MatthiasDrewsCS]

OAuth2TokenCustomizers:

@Component
@RequiredArgsConstructor
public class CustomClaimsOAuth2TokenCustomizer implements OAuth2TokenCustomizer<JwtEncodingContext> {
...
}

@Component
@Slf4j
@RequiredArgsConstructor
public class ClientAssertionOAuth2TokenCustomizer implements OAuth2TokenCustomizer<JwtEncodingContext> {
...
}

@Component
@Primary
@RequiredArgsConstructor
@Slf4j
public class TenantAwareOAuth2TokenCustomizer implements OAuth2TokenCustomizer<JwtEncodingContext>, ApplicationListener<ApplicationStartedEvent> {

    private final TenantPerIssuerComponentRegistry componentRegistry;
 
...

    @Override
    public void customize(JwtEncodingContext context) {
        var delegate = componentRegistry.get(OAuth2TokenCustomizer.class);
        Assert.state(delegate != null,
            "OAuth2TokenCustomizer not found for \"requested\" issuer identifier.");
        delegate.customize(context);
    }

OAuth2AuthorizationServices:

@Component
@Primary
@RequiredArgsConstructor
public class TenantAwareOAuth2AuthorizationService implements OAuth2AuthorizationService {

    private final List<RedisOAuth2AuthorizationService> delegates;

...

    @Override
    public OAuth2Authorization findByToken(String token, OAuth2TokenType tokenType) {
        return getOAuth2AuthorizationService().findByToken(token, tokenType);
    }

    private OAuth2AuthorizationService getOAuth2AuthorizationService() {
        var context = AuthorizationServerContextHolder.getContext();
        if (context == null || context.getIssuer() == null) {
            throw new IllegalStateException("AuthorizationServerContext or issuer is not set in the current context.");
        }
        var oAuth2AuthorizationService = delegates.stream().filter(delegate -> delegate.supports(context.getIssuer()))
            .findFirst()
            .orElse(null);
        Assert.state(oAuth2AuthorizationService != null,
            "OAuth2AuthorizationService not found for \"requested\" issuer identifier.");
        return oAuth2AuthorizationService;
    }

public class RedisOAuth2AuthorizationService implements OAuth2AuthorizationService {
...
}

RegisteredClientRepository

@RequiredArgsConstructor
public class TenantAwareRegisteredClientRepository implements RegisteredClientRepository {

    private final TenantPerIssuerComponentRegistry componentRegistry;

...

    @Override
    public RegisteredClient findById(String id) {
        return getRegisteredClientRepository().findById(id);
    }

...

    private RegisteredClientRepository getRegisteredClientRepository() {
        RegisteredClientRepository registeredClientRepository = this.componentRegistry.get(RegisteredClientRepository.class);
        Assert.state(registeredClientRepository != null,
            "RegisteredClientRepository not found for \"requested\" issuer identifier.");
        return registeredClientRepository;
    }
}

PostProcessors for bean registration:

@Slf4j
@Component
public class TenantAwareBeanFactoryPostProcessor implements BeanFactoryPostProcessor, EnvironmentAware {

    private Environment environment;

    @Override
    public void postProcessBeanFactory(@NonNull ConfigurableListableBeanFactory beanFactory) throws BeansException {
        var bindResult = Binder.get(environment).bind("acp.oauth2", OAuth2Properties.class);
        var oAuth2Properties = bindResult.get();
        oAuth2Properties.getTenants().values().forEach(tenant -> {
            var tenantId = tenant.getTenantId();
            ((BeanDefinitionRegistry) beanFactory).registerBeanDefinition(
                tenantId + "OAuth2AuthorizationService",
                BeanDefinitionBuilder.genericBeanDefinition(RedisOAuth2AuthorizationService.class)
                    .addConstructorArgReference("tenantPerIssuerComponentRegistry")
                    .addConstructorArgReference("oauth2RedisTemplate")
                    .addConstructorArgValue(tenantId)
                    .getBeanDefinition());
            log.info("registered OAuth2AuthorizationService for tenant {}", tenantId);

          ...
    }

    @Override
    public void setEnvironment(@NonNull Environment environment) {
        this.environment = environment;
    }
}

@Component
@RequiredArgsConstructor
@Slf4j
public class TenantAwareBeanPostProcessor implements BeanPostProcessor, ApplicationContextAware {

    private ApplicationContext applicationContext;

    @Override
    public Object postProcessAfterInitialization(@NonNull Object bean, @NonNull String beanName) throws BeansException {
        if (bean instanceof RegisteredClientRepository registeredClientRepository) {
            var oAuth2Properties = applicationContext.getBean(OAuth2Properties.class);
            var componentRegistry = applicationContext.getBean(TenantPerIssuerComponentRegistry.class);
            oAuth2Properties.getTenants().values().forEach(tenant -> {
                var tenantId = tenant.getTenantId();
                var tenantRegisteredClientRepository = createRegisteredClientRepository(registeredClientRepository, tenantId, tenant);
                componentRegistry.register(tenantId, RegisteredClientRepository.class, tenantRegisteredClientRepository);
                log.info("registered RegisteredClientRepository for tenant {}", tenantId);
            });
            return new TenantAwareRegisteredClientRepository(componentRegistry);
        }
        return bean;
    }

    private RegisteredClientRepository createRegisteredClientRepository(
        RegisteredClientRepository registeredClientRepository,
        String tenantId,
        Tenant tenant
    ) {
        var clientIds = tenant.getClients();
        var registeredClients = new ArrayList<RegisteredClient>();
        if (clientIds != null && !clientIds.isEmpty()) {
            clientIds.forEach(clientId -> {
                var registeredClient = registeredClientRepository.findByClientId(clientId);
                if (registeredClient != null) {
                    registeredClients.add(registeredClient);
                } else {
                    log.warn("Client with id {} not found for tenant {}", clientId, tenantId);
                }
            });
        }
        if (registeredClients.isEmpty()) {
            return new NoopRegisteredClientRepository();
        } else {
            return new InMemoryRegisteredClientRepository(registeredClients);
        }
    }

    @Override
    public void setApplicationContext(@NonNull ApplicationContext applicationContext) throws BeansException {
        this.applicationContext = applicationContext;
    }
}

Config:

@Configuration
public class OAuth2Config {

    @Bean
    @Order(Ordered.HIGHEST_PRECEDENCE)
    public SecurityFilterChain authorizationServerSecurityFilterChain(
        HttpSecurity http,
        OAuth2AuthorizationService oAuth2AuthorizationService,
        LazilyAuthenticatingOAuth2AuthorizationService lazilyAuthenticatingOAuth2AuthorizationService,
        TenantService tenantService
    ) throws Exception {
        var authorizationServerConfigurer = OAuth2AuthorizationServerConfigurer.authorizationServer();

        var requestMatcher = authorizationServerConfigurer.getEndpointsMatcher();

        http
            .securityMatcher(requestMatcher)
            .with(authorizationServerConfigurer, configurer -> configurer
                .authorizationService(oAuth2AuthorizationService)
...
            )
            .authorizeHttpRequests(authorize -> authorize.anyRequest().authenticated());
        return http.build();
    }

Note: what is happening in the TenantAwareBeanPostProcessor, you can easily designate as a hack. But unfortunately we did not find a clean option to have the conversion of Spring Boot properties into RegisteredClients for the several InMemoryRegisteredClientRepositorys. But it works for us, and does not affect the discussion about @Primary.

Note: we also have tenant specific JWK keys. However they are based on a custom implementation so I did not paste the configuration for them here.

You can see, we spent some effort into registration and of the tenant specific beans. The TenantAware... beans are able to pick the right bean instance for the requested tenant.