chore(auth): clean up authentication, jwt, add some exception handling

Signed-off-by: Alan Brault <alan.brault@visus.io>

src/main/kotlin/io/visus/demos/kotlinapi/config/SecurityConfig.kt

@@ -1,13 +1,19 @@
 package io.visus.demos.kotlinapi.config
 
+import io.visus.demos.kotlinapi.domain.service.UserService
+import io.visus.demos.kotlinapi.infrastructure.security.JwtAuthenticationFilter
 import org.springframework.context.annotation.Bean
 import org.springframework.context.annotation.Configuration
+import org.springframework.security.authentication.ReactiveAuthenticationManager
+import org.springframework.security.authentication.UserDetailsRepositoryReactiveAuthenticationManager
 import org.springframework.security.config.annotation.method.configuration.EnableReactiveMethodSecurity
 import org.springframework.security.config.annotation.web.reactive.EnableWebFluxSecurity
+import org.springframework.security.config.web.server.SecurityWebFiltersOrder
 import org.springframework.security.config.web.server.ServerHttpSecurity
 import org.springframework.security.crypto.password.PasswordEncoder
 import org.springframework.security.crypto.password4j.Argon2Password4jPasswordEncoder
 import org.springframework.security.web.server.SecurityWebFilterChain
+import org.springframework.security.web.server.context.NoOpServerSecurityContextRepository
 import org.springframework.web.cors.CorsConfiguration
 import org.springframework.web.cors.reactive.CorsConfigurationSource
 import org.springframework.web.cors.reactive.UrlBasedCorsConfigurationSource
@@ -15,7 +21,10 @@ import org.springframework.web.cors.reactive.UrlBasedCorsConfigurationSource
 @Configuration
 @EnableWebFluxSecurity
 @EnableReactiveMethodSecurity
-class SecurityConfig {
+class SecurityConfig(
+    private val jwtAuthenticationFilter: JwtAuthenticationFilter,
+    private val userDetailsService: UserService,
+) {
     @Bean
     fun corsConfigurationSource(): CorsConfigurationSource {
         val configuration =
@@ -34,11 +43,19 @@ class SecurityConfig {
     @Bean
     fun passwordEncoder(): PasswordEncoder = Argon2Password4jPasswordEncoder()
 
+    @Bean
+    fun reactiveAuthenticationManager(passwordEncoder: PasswordEncoder): ReactiveAuthenticationManager {
+        val authManager = UserDetailsRepositoryReactiveAuthenticationManager(userDetailsService)
+        authManager.setPasswordEncoder(passwordEncoder)
+        return authManager
+    }
+
     @Bean
     fun securityFilterChain(http: ServerHttpSecurity): SecurityWebFilterChain {
         http
             .csrf { it.disable() }
             .cors { it.configurationSource(corsConfigurationSource()) }
+            .securityContextRepository(NoOpServerSecurityContextRepository.getInstance())
             .authorizeExchange { authorize ->
                 authorize
                     .pathMatchers(
@@ -51,10 +68,11 @@ class SecurityConfig {
                         "/api/v1/auth/refresh",
                         "/api/v1/auth/register",
                     ).permitAll()
-                    .pathMatchers("/api/v1/auth/logout").authenticated()
+                    .pathMatchers("/api/v1/auth/logout")
+                    .authenticated()
                     .anyExchange()
                     .authenticated()
-            }
+            }.addFilterAt(jwtAuthenticationFilter, SecurityWebFiltersOrder.AUTHENTICATION)
 
         return http.build()
     }

src/main/kotlin/io/visus/demos/kotlinapi/domain/exception/GlobalExceptionHandler.kt

@@ -0,0 +1,50 @@
+package io.visus.demos.kotlinapi.domain.exception
+
+import io.visus.demos.kotlinapi.api.dto.ErrorResponse
+import org.slf4j.LoggerFactory
+import org.springframework.http.HttpStatus
+import org.springframework.http.ResponseEntity
+import org.springframework.security.authentication.BadCredentialsException
+import org.springframework.security.core.userdetails.UsernameNotFoundException
+import org.springframework.web.bind.annotation.ExceptionHandler
+import org.springframework.web.bind.annotation.RestControllerAdvice
+import org.springframework.web.server.ServerWebExchange
+
+@RestControllerAdvice
+class GlobalExceptionHandler {
+    private val logger = LoggerFactory.getLogger(this::class.java)
+
+    @ExceptionHandler(BadCredentialsException::class)
+    fun handleBadCredentialsException(
+        ex: BadCredentialsException,
+        exchange: ServerWebExchange,
+    ): ResponseEntity<ErrorResponse> {
+        logger.warn("Bad credentials attempt from: {}", exchange.request.remoteAddress)
+
+        val error =
+            ErrorResponse(
+                status = HttpStatus.UNAUTHORIZED.value(),
+                message = ex.message ?: "Invalid credentials",
+                path = exchange.request.path.toString(),
+            )
+
+        return ResponseEntity.status(HttpStatus.UNAUTHORIZED).body(error)
+    }
+
+    @ExceptionHandler(UsernameNotFoundException::class)
+    fun handleUsernameNotFoundException(
+        ex: UsernameNotFoundException,
+        exchange: ServerWebExchange,
+    ): ResponseEntity<ErrorResponse> {
+        logger.warn("User not found: {}", ex.message)
+
+        val error =
+            ErrorResponse(
+                status = HttpStatus.UNAUTHORIZED.value(),
+                message = "Invalid credentials",
+                path = exchange.request.path.toString(),
+            )
+
+        return ResponseEntity.status(HttpStatus.UNAUTHORIZED).body(error)
+    }
+}

src/main/kotlin/io/visus/demos/kotlinapi/domain/service/RefreshTokenServiceImpl.kt

@@ -33,12 +33,12 @@ class RefreshTokenServiceImpl(
         userId: String,
         familyId: String,
     ): String {
-        val activeCount = refreshTokenRepository.countActive(userId)
+        val activeCount = refreshTokenRepository.countActive(ObjectId(userId))
 
         if (activeCount >= MAX_SESSIONS_PER_USER) {
             val sessions =
                 refreshTokenRepository
-                    .findActive(userId)
+                    .findActive(ObjectId(userId))
                     .sortedBy { it.createdAt }
 
             val oldest = sessions.first()
@@ -120,7 +120,7 @@ class RefreshTokenServiceImpl(
     }
 
     override suspend fun revokeAllUserTokens(userId: String) {
-        val tokens = refreshTokenRepository.findActive(userId)
+        val tokens = refreshTokenRepository.findActive(ObjectId(userId))
 
         tokens.forEach {
             it.isRevoked = true

src/main/kotlin/io/visus/demos/kotlinapi/domain/service/UserService.kt

@@ -1,5 +1,5 @@
 package io.visus.demos.kotlinapi.domain.service
 
-import org.springframework.security.core.userdetails.UserDetailsService
+import org.springframework.security.core.userdetails.ReactiveUserDetailsService
 
-interface UserService : UserDetailsService
+interface UserService : ReactiveUserDetailsService

src/main/kotlin/io/visus/demos/kotlinapi/domain/service/UserServiceImpl.kt

@@ -1,21 +1,23 @@
 package io.visus.demos.kotlinapi.domain.service
 
 import io.visus.demos.kotlinapi.infrastructure.user.UserRepository
-import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.reactor.mono
 import org.springframework.security.core.userdetails.UserDetails
 import org.springframework.security.core.userdetails.UsernameNotFoundException
 import org.springframework.stereotype.Service
+import reactor.core.publisher.Mono
 
 @Service
 class UserServiceImpl(
     private val userRepository: UserRepository,
 ) : UserService {
-    override fun loadUserByUsername(username: String): UserDetails {
-        if (username.isBlank()) throw UsernameNotFoundException("Username must not be null or empty")
+    override fun findByUsername(username: String): Mono<UserDetails> =
+        mono {
+            if (username.isBlank()) {
+                throw UsernameNotFoundException("Username must not be null or empty")
+            }
 
-        return runBlocking {
             userRepository.findByEmail(username)
                 ?: throw UsernameNotFoundException("User not found with email: $username")
         }
-    }
 }

src/main/kotlin/io/visus/demos/kotlinapi/infrastructure/security/JwtAuthenticationFilter.kt

@@ -0,0 +1,79 @@
+package io.visus.demos.kotlinapi.infrastructure.security
+
+import io.visus.demos.kotlinapi.domain.model.JwtTokenType
+import io.visus.demos.kotlinapi.domain.service.UserService
+import org.slf4j.LoggerFactory
+import org.springframework.http.HttpHeaders
+import org.springframework.security.authentication.UsernamePasswordAuthenticationToken
+import org.springframework.security.core.context.ReactiveSecurityContextHolder
+import org.springframework.stereotype.Component
+import org.springframework.web.server.ServerWebExchange
+import org.springframework.web.server.WebFilter
+import org.springframework.web.server.WebFilterChain
+import reactor.core.publisher.Mono
+
+@Component
+class JwtAuthenticationFilter(
+    private val jwtTokenProvider: JwtTokenProvider,
+    private val userService: UserService,
+) : WebFilter {
+    private val logger = LoggerFactory.getLogger(JwtAuthenticationFilter::class.java)
+
+    companion object {
+        private const val BEARER_PREFIX = "Bearer "
+    }
+
+    override fun filter(
+        exchange: ServerWebExchange,
+        chain: WebFilterChain,
+    ): Mono<Void> {
+        val jwt = extractJwtFromRequest(exchange)
+
+        return if (jwt != null && isValidAccessToken(jwt)) {
+            val username = jwtTokenProvider.extractUsername(jwt)
+
+            userService
+                .findByUsername(username)
+                .flatMap { userDetails ->
+                    val authentication =
+                        UsernamePasswordAuthenticationToken(
+                            userDetails,
+                            null,
+                            userDetails.authorities,
+                        )
+
+                    logger.debug("Set authentication for user: {}", username)
+
+                    // Set authentication in reactive security context
+                    chain
+                        .filter(exchange)
+                        .contextWrite(ReactiveSecurityContextHolder.withAuthentication(authentication))
+                }.onErrorResume { ex ->
+                    logger.error("Could not set user authentication in security context", ex)
+                    chain.filter(exchange)
+                }
+        } else {
+            // No JWT or invalid JWT, continue without authentication
+            chain.filter(exchange)
+        }
+    }
+
+    private fun isValidAccessToken(token: String): Boolean =
+        try {
+            jwtTokenProvider.validateToken(token) &&
+                jwtTokenProvider.extractTokenType(token) == JwtTokenType.ACCESS
+        } catch (ex: Exception) {
+            logger.debug("Token validation failed: {}", ex.message)
+            false
+        }
+
+    private fun extractJwtFromRequest(exchange: ServerWebExchange): String? {
+        val bearerToken = exchange.request.headers.getFirst(HttpHeaders.AUTHORIZATION)
+
+        return if (bearerToken != null && bearerToken.startsWith(BEARER_PREFIX)) {
+            bearerToken.substring(BEARER_PREFIX.length)
+        } else {
+            null
+        }
+    }
+}

src/main/kotlin/io/visus/demos/kotlinapi/infrastructure/security/JwtTokenProvider.kt

@@ -105,7 +105,7 @@ class JwtTokenProvider(
         return expiration.before(Date())
     }
 
-    private fun validateToken(token: String): Boolean {
+    fun validateToken(token: String): Boolean {
         try {
             Jwts
                 .parser()

src/main/kotlin/io/visus/demos/kotlinapi/infrastructure/user/RefreshTokenRepository.kt

@@ -1,22 +1,23 @@
 package io.visus.demos.kotlinapi.infrastructure.user
 
 import io.visus.demos.kotlinapi.domain.model.RefreshToken
+import org.bson.types.ObjectId
 import org.springframework.data.mongodb.repository.Query
 import org.springframework.data.mongodb.repository.Update
 import org.springframework.data.repository.kotlin.CoroutineCrudRepository
 
 interface RefreshTokenRepository : CoroutineCrudRepository<RefreshToken, String> {
-    suspend fun deleteByUserId(userId: String): Long
+    suspend fun deleteByUserId(userId: ObjectId): Long
 
     @Query(value = "{ 'userId': ?0, 'isRevoked': false }", count = true)
-    suspend fun countActive(userId: String): Long
+    suspend fun countActive(userId: ObjectId): Long
 
     suspend fun findByFamilyId(familyId: String): List<RefreshToken>
 
     suspend fun findByTokenHash(tokenHash: String): RefreshToken?
 
     @Query("{ 'userId': ?0, 'isRevoked': false }")
-    suspend fun findActive(userId: String): List<RefreshToken>
+    suspend fun findActive(userId: ObjectId): List<RefreshToken>
 
     @Query("{ 'familyId': ?0 }")
     @Update($$"{ '$set': { 'isRevoked': true } }")