Spring OAuth2: DuplicateKeyException when using JdbcTokenStore and DefaultTokenServices under heavy load

Question

As mentioned in the title, I'm experiencing the issue, when the same client is querying the token endpoint concurrently (two processes requesting token for the same client at the same time).

The message in the logs of the auth server looks like this:

2016-12-05 19:08:03.313  INFO 31717 --- [nio-9999-exec-5] o.s.s.o.provider.endpoint.TokenEndpoint  : Handling error: DuplicateKeyException, PreparedStatementCallback; SQL [insert into oauth_access_token (token_id, token, authentication_id, user_name, client_id, authentication, refresh_token) values (?, ?, ?, ?, ?, ?, ?)]; ERROR: duplicate key value violates unique constraint "oauth_access_token_pkey"
      Detail: Key (authentication_id)=(4148f592d600ab61affc6fa90bcbf16f) already exists.; nested exception is org.postgresql.util.PSQLException: ERROR: duplicate key value violates unique constraint "oauth_access_token_pkey"
      Detail: Key (authentication_id)=(4148f592d600ab61affc6fa90bcbf16f) already exists.

I'm using PostgreSQL with table like this:

CREATE TABLE oauth_access_token
(
  token_id character varying(256),
  token bytea,
  authentication_id character varying(256) NOT NULL,
  user_name character varying(256),
  client_id character varying(256),
  authentication bytea,
  refresh_token character varying(256),
  CONSTRAINT oauth_access_token_pkey PRIMARY KEY (authentication_id)
)

And my application looks like that:

@SpringBootApplication
public class OAuthServTest {
   public static void main (String[] args) {
      SpringApplication.run (OAuthServTest.class, args);
   }

   @Configuration
   @EnableAuthorizationServer
   @EnableTransactionManagement
   protected static class OAuth2Config extends AuthorizationServerConfigurerAdapter {

      @Autowired
      private AuthenticationManager authenticationManager;

      @Autowired
      private DataSource            dataSource;

      @Bean
      public PasswordEncoder passwordEncoder ( ) {
         return new BCryptPasswordEncoder ( );
      }

      @Override
      public void configure (AuthorizationServerEndpointsConfigurer endpoints) throws Exception {
         endpoints.authenticationManager (authenticationManager);
         endpoints.tokenServices (tokenServices ( ));
      }

      @Override
      public void configure (ClientDetailsServiceConfigurer clients) throws Exception {
         clients.jdbc (this.dataSource).passwordEncoder (passwordEncoder ( ));
      }

      @Override
      public void configure (AuthorizationServerSecurityConfigurer oauthServer) throws Exception {
         oauthServer.passwordEncoder (passwordEncoder ( ));
      }

      @Bean
      public TokenStore tokenStore ( ) {
         return new JdbcTokenStore (this.dataSource);
      }

      @Bean
      @Primary
      public AuthorizationServerTokenServices tokenServices ( ) {
         final DefaultTokenServices defaultTokenServices = new DefaultTokenServices ( );
         defaultTokenServices.setTokenStore (tokenStore ( ));
         return defaultTokenServices;
      }

   }
}

My research always lead me to this problem. But this bug was fixed a long time ago and I'm on the latest version of Spring Boot (v1.4.2).

My guess is that I'm doing something wrong and the retrieval of the token in DefaultTokenServices is not happening in the transaction?

Answer 1

The problem is with a race condition on the DefaultTokenServices.

I found a workaround. Not saying it's gorgeous but it works. The idea is to add retry logic using an AOP Advice around the TokenEndpoint:

@Aspect
@Component
public class TokenEndpointRetryInterceptor {

  private static final int MAX_RETRIES = 4;

  @Around("execution(* org.springframework.security.oauth2.provider.endpoint.TokenEndpoint.*(..))")
  public Object execute (ProceedingJoinPoint aJoinPoint) throws Throwable {
    int tts = 1000;
    for (int i=0; i<MAX_RETRIES; i++) {
      try {
        return aJoinPoint.proceed();
      } catch (DuplicateKeyException e) {
        Thread.sleep(tts);
        tts=tts*2;
      }
    }
    throw new IllegalStateException("Could not execute: " + aJoinPoint.getSignature().getName());
  }

}

Answer 2

https://github.com/spring-projects/spring-security-oauth/issues/1033

 @Bean
public TokenStore tokenStore(final DataSource dataSource) {
    final JdbcTemplate jdbcTemplate = new JdbcTemplate(dataSource);
    final AuthenticationKeyGenerator authenticationKeyGenerator = new DefaultAuthenticationKeyGenerator();
    return new JdbcTokenStore(dataSource) {

        @Override
        public void storeAccessToken(final OAuth2AccessToken token, final OAuth2Authentication authentication) {
            final String key = authenticationKeyGenerator.extractKey(authentication);
            jdbcTemplate.update("delete from oauth_access_token where authentication_id = ?", key);
            super.storeAccessToken(token, authentication);
        }

    };
}