Ktor server: correct way to implement a custom response authentication wrapper

Support
1

Hey there,

I’m wondering what the correct way of setting a thread-bound authentication is.

I built a custom plugin that should start the authentication on AuthenticatioChecked and end it after every response is sent (or if there was an error).

Basically, it works, but it looks like the authentication is not cleared correctly every time, that’s why I’m wondering if I’m doing something wrong. Please have a look at the code:

private val authenticatedUserKey = AttributeKey<String>("authenticatedUser")

val CustomAuthPlugin: RouteScopedPlugin<CustomAuthPluginConfig> = createRouteScopedPlugin(
    "CustomAuthPlugin",
    ::CustomAuthPluginConfig
) {
    handleAuthentication()
}

@KtorDsl
class CustomAuthPluginConfig {
    val systemAuthenticator = Authenticator() // has begin(...) and end() methods
}

private fun PluginBuilder<CustomAuthPluginConfig>.handleAuthentication() {
    on(AuthenticationChecked) { call ->
        val principal = call.principal<JWTPrincipal>() ?: error("No principal")
        val username = principal.payload.getClaim("username").asString()
        try {
            while (SecurityContextHelper.getAuthentication() != null) {
                // This should not be necessary, but it seems that the security context is not always cleared properly
                log.warn("Had to clear security context from auth: ${SecurityContextHelper.getAuthentication()}")
                pluginConfig.systemAuthenticator.end()
            }
            pluginConfig.systemAuthenticator.begin(username)
            call.attributes.put(authenticatedUserKey, username)
        } catch (e: UsernameNotFoundException) {
            log.error("Authentication failed for user $username " + e.message)
            throw e
        } catch (e: AuthenticationException) {
            log.error("Authentication failed for user $username " + e.message)
            throw e
        }
    }
    on(CallFailed) { call, error ->
        endAuthentication(call)
    }
//    on(ResponseBodyReadyForSend) { call, outgoingContent ->  // Maybe this is better ???
//        endAuthentication(call)
//    }
    on(ResponseSent) { call ->
        endAuthentication(call)
    }
}

private fun PluginBuilder<CustomAuthPluginConfig>.endAuthentication(call: ApplicationCall) {
    if (call.attributes.contains(authenticatedUserKey)) {
        pluginConfig.systemAuthenticator.end()
        call.attributes.remove(authenticatedUserKey)
    }
}

Behind the scenes I’m using a framework that in turn uses spring security to set the authentication. ( SecurityContextHelper.setAuthentication(…) )

I’m installing my custom plugin via:

fun Application.main() {
    routing {
            route("/api") {
                authenticate("auth-jvm") {
                    install(CustomAuthPlugin)
                    // authenticatedRoutes()
                }
            }

    }
}

Please help if you have any idea what the cause of the problem could be.