How to Implement Custom Token-Based Authentication in Spring Boot and Kotlin

When building Spring Boot REST web services, we have to deal with security.

In order to achieve our security goals, such as authorization and authentication, using a specifically designed framework like Spring Security may be the best solution.

“Spring Security is a framework that focuses on providing both authentication and authorization to Java applications. Like all Spring projects, the real power of Spring Security is found in how easily it can be extended to meet custom requirements.” – Spring Security

However, sometimes implementing a specific authentication logic to keep the application simple might be necessary.

The system we are going to present will allow us to choose whether or not to protect an API. Moreover, we assume that every valid authentication token identifies a particular user. The token is in a specific header or cookie and is used by authentication logic to extract a user whose data will be automatically passed to a protected API’s function body.

Let’s see how custom token-based authentication can be achieved in Spring Boot and Kotlin.

1. Defining a Custom Annotation

To decide whether an API should be protected by the authentication system, we are going to use a custom-defined annotation. This annotation will be used to mark a parameter of type User to define whether or not the API is protected. The instance of the particular user identified by the token is automatically retrieved and can be used inside the API function body.

Let’s see how a custom Auth annotation can be defined:

@Retention(AnnotationRetention.RUNTIME)
@Target(AnnotationTarget.VALUE_PARAMETER)
annotation class Auth

2. Defining Authentication Logic

Authentication logic should be placed in a specific component, which we are going to call AuthTokenHandler. The purpose of this class is to verify if the token is valid and extract its related user. This can be achieved in many ways. We are going to show two different verification approaches.

Using a custom DAO:

@Component
class AuthTokenHandler {

    @Autowired
    lateinit var authTokenDao: AuthTokenDao

    @Transactional(readOnly = true)
    fun getUserFromToken(token : String?) : User {    
        if (token == null)
            throw AuthenticationException()

        val authTokenOptional = authTokenDao.findByTokenNotExpired(token, Timestamp(System.currentTimeMillis()))
        
        authTokenOptional.orElseThrow { AuthenticationException() }

        return authTokenOptional.get().user
    }
}

Calling an external API:

@Component
class AuthTokenHandler {

    @Autowired
    private lateinit var tokenAPIHandler: TokenAPI

    fun getUserFromToken(token : String?) : User {
        if (token == null)
            throw AuthenticationException()

        try {
            return tokenAPIHandler.getUserByToken(token)
        } catch (e: Exception) {
            throw AuthenticationException()
        }
    }
}

In both cases, when the token is missing or not valid, a custom AuthenticationException is thrown. In this case, the protected API should respond with “401 Unauthorized.”

“The HTTP 401 Unauthorized client error status response code indicates that the request has not been applied because it lacks valid authentication credentials for the target resource.” — MDN web docs

To achieve this, a class marked with @ControllerAdvice can be used as follows:

@ControllerAdvice
class CustomExceptionsHandler {
    @ExceptionHandler(AuthenticationException::class)
    fun authenticationExceptionHandler(e: Exception): ResponseEntity<String> {
        return ResponseEntity("authToken missing or not valid!", HttpStatus.UNAUTHORIZED)
    }
}

3. Retrieving the Token

To allow Spring Boot to automatically look for the token in the headers or cookies when the custom Auth annotation is identified, an AuthTokenWebResolver implementing HandlerMethodArgumentResolver has to be defined.

Let’s assume that the authentication token can be placed in a header or cookie called authToken. The retrieving logic can be implemented as follows:

class AuthTokenWebResolver : HandlerMethodArgumentResolver {
    @Autowired
    lateinit var authTokenHandler: AuthTokenHandler

    // to register Auth annotation
    override fun supportsParameter(methodParameter: MethodParameter): Boolean {
        return methodParameter.getParameterAnnotation(Auth::class.java) != null
    }

    override fun resolveArgument(parameter: MethodParameter,
                                 mavContainer: ModelAndViewContainer?,
                                 webRequest: NativeWebRequest,
                                 binderFactory: WebDataBinderFactory?): Any? {
        if (parameter.parameterType == User::class.java) {
          // looking for the auth token in the headers
            var authToken = webRequest.getHeader("authToken")

          // looking for the auth token in the cookies
            if (authToken == null) {
                val servletRequest = webRequest.nativeRequest as HttpServletRequest

                val authTokenCookie = WebUtils.getCookie(servletRequest, "authToken")
              
                if (authTokenCookie != null)
                    authToken = authTokenCookie.value
            }

            return authTokenHandler.getUserFromToken(authToken)
        }

        return UNRESOLVED
    }
}

4. Configuring Spring Boot

Now, we have to define a custom class for the configurations. This way, Spring Boot will be able to use the custom Auth annotation as designed.

For everything to work, we need to add the previously defined AuthTokenWebResolver to the default argument resolvers. This can be achieved by harnessing the WebMvcConfigurationSupport class.

“[WebMvcConfigurationSupport] is typically imported by adding @EnableWebMvc to an application @Configuration class. An alternative more advanced option is to extend directly from this class and override methods as necessary, remembering to add @Configuration to the subclass and @Bean to overridden @Bean methods.” — Spring’s official documentation

We are going to define a @Configuration class that extends WebMvcConfigurationSupport:

@Configuration
class CustomConfig : WebMvcConfigurationSupport() {

    @Bean
    fun authWebArgumentResolverFactory() : HandlerMethodArgumentResolver {
        return AuthWebResolver()
    }

    // Addding the AuthWebResolver to the default argument resolvers
    override fun addArgumentResolvers(argumentResolvers: MutableList<HandlerMethodArgumentResolver>) {
        argumentResolvers.add(authWebArgumentResolverFactory())
    }

    // TODO: define CORS mappings
    public override fun addCorsMappings(registry: CorsRegistry) {
        registry
            .addMapping("/**")
            .allowedOrigins("*")
            .allowedMethods("GET", "DELETE", "PATCH", "POST", "PUT")
            .allowCredentials(true)
    }

    @Bean
    fun restTemplate(builder: RestTemplateBuilder): RestTemplate {
        return builder.build()
    }
}

When using WebMvcConfigurationSupport, do not forget that we have to deal with CORS configurations. Otherwise, our APIs might not be reachable as expected.

5. Putting It All Together

Now, it is time to see how Auth annotation can be used to make an API work only with authenticated users. This can be easily achieved by adding a User type parameter marked with Auth annotation to the chosen Controller API function:

@GetMapping("data/{id}")
fun getData(
  @Auth user : User,
  @PathVariable(value = "id") id: Int
) : ResponseEntity<Void> {

  // NOTE: user can be used inside the method body

  // API logic

  return ResponseEntity(HttpStatus.OK)
}

Moreover, defining an API lacking protection is possible as well:

@GetMapping("data/{id}")
fun getData(
  @PathVariable(value = "id") id: Int
) : ResponseEntity<Void> {

  // API logic

  return ResponseEntity(HttpStatus.OK)
}

Conclusion

That’s all, folks! I hope this helps you define custom token-based authentication in Spring Boot and Kotlin.