Can I create a plugin that adds context to the cor...

# ktor

[rocketraman]

Can I create a plugin that adds context to the coroutine started by ktor (server) to handle a request?

[Anastasios Georgousakis]

What I did is the following. I don’t know if this is the best way to do that but it worked in my server

import io.ktor.http.*
import io.ktor.server.application.*
import io.ktor.server.auth.*
import io.ktor.server.request.*
import io.ktor.server.routing.*
import io.ktor.util.*
import io.ktor.util.pipeline.*
import kotlinx.coroutines.withContext
import kotlin.coroutines.AbstractCoroutineContextElement
import kotlin.coroutines.CoroutineContext
import kotlin.coroutines.coroutineContext

/**
 * The configuration for the request context config
 */
@KtorDsl
class RequestContextRouteConfig {
    // empty
}

internal object RequestContextRouteHook : Hook<suspend (ApplicationCall, suspend () -> Unit) -> Unit> {
    internal val RequestContextPhase: PipelinePhase = PipelinePhase("RequestContextPhase")

    override fun install(
        pipeline: ApplicationCallPipeline,
        handler: suspend (ApplicationCall, suspend () -> Unit) -> Unit
    ) {
        pipeline.insertPhaseAfter(ApplicationCallPipeline.Plugins, RequestContextPhase)
        pipeline.intercept(RequestContextPhase) { handler(call, ::proceed) }
    }
}

/**
 * A plugin that process calls with request context. Usually used via the [withRequestContext] function inside routing.
 */
val RequestContextRoutePlugin: RouteScopedPlugin<RequestContextRouteConfig> =
    createRouteScopedPlugin("RequestContextRoutePlugin", ::RequestContextRouteConfig) {
        on(RequestContextRouteHook) { call, block ->
//            if (pluginConfig.requirePrincipal && call.authentication.principal == null)
//                block()
//            else
                withRequestContext(call, block)
        }
    }

fun Route.withRequestContext(
    build: Route.() -> Unit
): Route {
    val route = this
    route.install(RequestContextRoutePlugin)
    route.build()
    return route
}

internal suspend inline fun withRequestContext(
    call: ApplicationCall,
    crossinline block: suspend () -> Unit
) {
    withContext(
        call.toRequestContext()
    ) {
        block()
    }
}

fun ApplicationCall.toRequestContext(): CoroutineContext =
    RequestContext(
        locale = computeRequestLocale(request),
        acceptLanguage = request.acceptLanguage(),
        correlationId = request.headers[HttpHeaders.XCorrelationId],
        authorization = request.headers[HttpHeaders.Authorization],
        principal = principal(),
    )

class RequestContext(
    val locale: Locale,
    val acceptLanguage: String? = null,
    val correlationId: String? = null,
    val authorization: String? = null,
    val principal: Principal? = null,
) : AbstractCoroutineContextElement(RequestContext) {
    companion object Key : CoroutineContext.Key<RequestContext>

    /**
     * Custom implementation of equals to exclude principal
     */
    override fun equals(other: Any?): Boolean {
        if (this === other) return true
        if (javaClass != other?.javaClass) return false

        other as RequestContext

        if (locale != other.locale) return false
        if (acceptLanguage != other.acceptLanguage) return false
        if (authorization != other.authorization) return false
        if (correlationId != other.correlationId) return false

        return true
    }

    /**
     * Custom calculation of hash code to exclude principal
     */
    override fun hashCode(): Int {
        var result = locale.hashCode()
        result = 31 * result + (acceptLanguage?.hashCode() ?: 0)
        result = 31 * result + (authorization?.hashCode() ?: 0)
        result = 31 * result + (correlationId?.hashCode() ?: 0)
        return result
    }

    override fun toString(): String {
        return "RequestContext(" +
                "locale=$locale, " +
                "acceptLanguage=$acceptLanguage, " +
                "correlationId=$correlationId, " +
                "authorization=$authorization" +
                ")"
    }

}

[Anastasios Georgousakis]

In you routes you have to add also this to have the context in your coroutine scope

routing {
	authenticate {
		withRequestContext {
			authRoutes()
			userRoutes()
		}
	}
	
}

[CLOVIS]

Here's a custom authentication plugin which injects the user's auth in the coroutine context: https://gitlab.com/opensavvy/formulaide/-/blob/main/backend/src/main/kotlin/Authentication.kt#L55

[rocketraman]

Thanks both!

[rocketraman]

I used this information to integrate Sentry perf mon with Ktor. See https://kotlinlang.slack.com/archives/C0A974TJ9/p1680708507742939.