Logowanie z opcją "Pozostań zalogowany" w Springu z Kotlinem

Nie przedłużając, proszę o krytykę mojej implementacji mechanizmu autentykacji użytkownika przy użyciu access oraz refresh tokenów w aplikacji Springowej napisanej w Kotlinie. Zależy mi zarówno na weryfikacji ewentualnych błędów logicznych w działaniu mechnizmu logowania, jak i zwyczajnym sprawdzeniu kodu pod względem best practices.

1. Logowanie użytkownika. Po udanej walidacji użytkownika, zwracany jest HTTP 200 z pustym body, a same tokeny wracają jako HttpOnly cookies. Dodatkowo, jeżeli w requeście przekazany został parameter stayLoggedIn: true, to refresh token danego użytkownika zostaje zapisany w bazie danych.

@RestController
@RequestMapping("/api/auth")
internal class AuthApi(private val signInService: SignInService) {
    @PostMapping("/sign-in")
    fun signIn(@RequestBody request: SignInRequest, response: HttpServletResponse): ResponseEntity<Void> {
        val signInResponse = signInService.signIn(request)
        response.addCookie(signInResponse.accessToken.cookie)
        response.addCookie(signInResponse.refreshToken.cookie)
        return ResponseEntity(HttpStatus.OK)
    }
}

@Service
class SignInService(
    private val userRepository: UserRepository,
    private val tokenRepository: RefreshTokenRepository,
    private val tokenUtils: JwtTokenUtils,
    private val passwordEncoder: PasswordEncoder
) {

    fun signIn(request: SignInRequest): AccessTokenDto {
        return userRepository.findByEmail(request.email)
            ?.let { user -> authenticate(request, user) }
            ?: throw EmailNotFound(request.email)
    }


    private fun authenticate(request: SignInRequest, user: User): AccessTokenDto {
        when (validPassword(request.password, user.password)) {
            true -> {
                val refreshToken = tokenUtils.generateRefreshToken(user.id)
                val response = AccessTokenDto(
                    TokenCookie.accessTokenCookie(tokenUtils.generateAccessToken(user.id)),
                    TokenCookie.refreshTokenCookie(refreshToken)
                )
                if (request.stayLoggedIn) saveUserRefreshToken(user, refreshToken)
                return response
            }
            false -> throw IncorrectPassword()
        }
    }

    private fun saveUserRefreshToken(user: User, refreshToken: String) {
        tokenRepository.findByUserId(user.id)?.let {
            tokenRepository.save(it.copy(token = refreshToken))
        } ?: run {
            tokenRepository.save(RefreshToken(userId = user.id, token = refreshToken))
        }
    }

    private fun validPassword(providedPassword: String, actualPassword: String) =
        passwordEncoder.matches(providedPassword, actualPassword)
}

Użyte wyżej TokenCookie wygląda w ten sposób

enum class TokenType(val value: String) {
    ACCESS_TOKEN("access_token"), REFRESH_TOKEN("refresh_token")
}

class TokenCookie(name: String, token: String) {
    val cookie: Cookie = Cookie(name, token)

    init {
        this.cookie.isHttpOnly = true
        this.cookie.path = "/"
    }

    companion object {
        fun accessTokenCookie(token: String) = TokenCookie(ACCESS_TOKEN.value, token)
        fun refreshTokenCookie(token: String) = TokenCookie(REFRESH_TOKEN.value, token)
    }
}

2. Weryfikacja access tokenu przy kolejnych requestach. Pierwszym krokiem jest sprawdzenie poprawności access tokenu, jeżeli ten jest poprawny, to w SecurityContext zostaje ustawiony odpowiedni obiekt autentykacji a także sprawdzane jest czy w requeście był dołączony poprawny refresh token, jeżeli tak to access token zostaje wygenerowany ponownie z odświeżonym expiration date. Ma to na celu uniknięcie sytuacji gdy ktoś zaloguje się bez opcji "Pozostań zalogowany", aktywnie korzysta z serwisu przez ponad 15 minut i nagle jego token wygasa, przez co wymagane jest ponowne wpisanie nazwy użytkownika i hasła.

fun validateToken(authToken: String?, tokenType: TokenType = ACCESS_TOKEN): TokenValidationResult {
    return try {
        val claims = Jwts.parser().setSigningKey(tokenSecret).parseClaimsJws(authToken)
        if (claims.body["type"] == tokenType.value) SUCCESS else INVALID
    } catch (ex: UnsupportedJwtException) {
        INVALID
    } catch (ex: MalformedJwtException) {
        INVALID
    } catch (ex: IllegalArgumentException) {
        INVALID
    } catch (ex: SignatureException) {
        INVALID
    } catch (ex: ExpiredJwtException) {
        EXPIRED
    }
}

@Component
internal class JwtTokenFilter(
    private val tokenUtils: JwtTokenUtils,
    private val expiredTokenHandler: ExpiredTokenHandler
) : OncePerRequestFilter() {
    override fun doFilterInternal(
        request: HttpServletRequest,
        response: HttpServletResponse,
        filterChain: FilterChain
    ) {
        if (!request.requestURI.startsWith("/api/auth")) {
            validateTokens(request, response)
        }
        filterChain.doFilter(request, response)
    }

    private fun validateTokens(
        request: HttpServletRequest,
        response: HttpServletResponse
    ) {
        getTokenFromCookies(request, ACCESS_TOKEN)?.let {
            val tokenValidationResult = tokenUtils.validateToken(it.value)
            val userId = tokenUtils.getId(it.value)
            if (tokenValidationResult == SUCCESS) {
                updateSecurityContext(userId)
                extendAccessToken(request, response, userId)
            } else if (tokenValidationResult == EXPIRED) {
                val refreshToken = getTokenFromCookies(request, REFRESH_TOKEN)?.value
                if (expiredTokenHandler.checkRefreshEligibility(response, it.value, refreshToken)) {
                    updateSecurityContext(userId)
                }
            }
        }
    }

    private fun updateSecurityContext(userId: Int) {
        SecurityContextHolder.getContext().authentication =
            UsernamePasswordAuthenticationToken(userId, null, Collections.emptyList())
    }

    private fun getTokenFromCookies(request: HttpServletRequest, tokenType: TokenType) =
        request.cookies?.find { cookie -> cookie.name == tokenType.value }

    private fun extendAccessToken(request: HttpServletRequest, response: HttpServletResponse, userId: Int) {
        getTokenFromCookies(request, REFRESH_TOKEN)?.let {
            if (tokenUtils.validateToken(it.value, REFRESH_TOKEN) == SUCCESS) {
                response.addCookie(TokenCookie.accessTokenCookie(tokenUtils.generateAccessToken(userId)).cookie)
            }
        }
    }
}

3. Obsługa refresh tokenu. Jak widać w powyższym kodzie, są 3 rezultaty walidacji tokenu. SUCCESS - wszysto ok, INVALID - bliżej nieokreślony błąd oraz EXPIRED, mówiący o tym ze token stracił ważność. W takiej sytuacji sprawdzany jest refresh token. Pierwszym krokiem jest sprawdzenie czy w bazie danych istnieje refresh token dla danego użytkownika (czyli czy zaznaczył on opcję "Pozostań zalogowany" przy logowaniu. Następnie sam token jest walidowany z tokenem który przyszedł w requeście i jeżeli wszystko jest ok, to generowana jest nowa para tokenów.

@Service
class ExpiredTokenHandler(private val tokenRepository: RefreshTokenRepository, private val tokenUtils: JwtTokenUtils) {

    fun checkRefreshEligibility(response: HttpServletResponse, expiredToken: String, refreshToken: String?): Boolean {
        val userId = tokenUtils.getId(expiredToken)
        val savedRefreshToken = tokenRepository.findByUserId(userId)
        return if (validateRefreshToken(refreshToken, savedRefreshToken)) {
            val (newAccessToken, newRefreshToken) = generateNewTokens(userId, savedRefreshToken!!)
            updateCookies(newAccessToken, newRefreshToken, response)
            true
        } else {
            false
        }
    }

    private fun validateRefreshToken(
        refreshToken: String?,
        savedRefreshToken: RefreshToken?
    ): Boolean {
        return refreshToken != null &&
                savedRefreshToken != null &&
                refreshToken == savedRefreshToken.token &&
                tokenUtils.validateToken(refreshToken, REFRESH_TOKEN) == SUCCESS
    }

    private fun updateCookies(accessToken: String, refreshToken: String, response: HttpServletResponse): Boolean {
        response.addCookie(accessTokenCookie(accessToken).cookie)
        response.addCookie(refreshTokenCookie(refreshToken).cookie)
        return true
    }

    private fun generateNewTokens(userId: Int, savedRefreshToken: RefreshToken): Pair<String, String> {
        val newAccessToken = tokenUtils.generateAccessToken(userId)
        val newRefreshToken = tokenUtils.generateRefreshToken(userId)
        tokenRepository.save(savedRefreshToken.copy(token = newRefreshToken))
        return Pair(newAccessToken, newRefreshToken)
    }
}