How to Implement Custom Token-Based Authentication in Spring Boot and Kotlin
Tags: cybersecurity, kotlin, programming, softwaredevelopment
Published: 2024-01-19 20:55:55

When building [Spring Boot](https://spring.io/projects/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](https://spring.io/projects/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](https://spring.io/projects/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:

```kotlin
@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:

```kotlin
@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:

```kotlin
@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](https://developer.mozilla.org/en-US/docs/Web/HTTP/Status/401#)

To achieve this, a class marked with [`@ControllerAdvice`](https://docs.spring.io/spring-framework/docs/current/javadoc-api/org/springframework/web/bind/annotation/ControllerAdvice.html) can be used as follows:

```kotlin
@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`](https://docs.spring.io/spring-framework/docs/current/javadoc-api/org/springframework/web/method/support/HandlerMethodArgumentResolver.html) 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:

```kotlin
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`](https://docs.spring.io/spring-framework/docs/current/javadoc-api/org/springframework/web/servlet/config/annotation/WebMvcConfigurationSupport.html) 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](https://docs.spring.io/spring-framework/docs/current/javadoc-api/org/springframework/web/servlet/config/annotation/WebMvcConfigurationSupport.html)

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

```kotlin
@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:

```kotlin
@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:

```kotlin
@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.
  
***
_The post "[How to Implement Custom Token-Based Authentication in Spring Boot and Kotlin](https://writech.run/blog/how-to-implement-custom-token-based-authentication-in-spring-boot-and-kotlin-5b59b55c1de2/)" appeared first on [Writech](https://writech.run)._