package dev.inmo.postssystem.features.auth.client import dev.inmo.postssystem.features.auth.common.* import dev.inmo.postssystem.features.users.common.User import dev.inmo.micro_utils.common.* import dev.inmo.micro_utils.coroutines.launchSafelyWithoutExceptions import io.ktor.client.HttpClientConfig import io.ktor.client.features.cookies.* import io.ktor.client.features.expectSuccess import io.ktor.client.request.* import io.ktor.client.statement.HttpReceivePipeline import io.ktor.client.statement.HttpResponse import io.ktor.http.* import kotlinx.coroutines.* import kotlinx.coroutines.sync.Mutex import kotlinx.coroutines.sync.withLock object AuthUnavailableException : Exception() fun HttpClientConfig<*>.installClientAuthenticator( baseUrl: String, scope: CoroutineScope, initialAuthKey: Either, onAuthKeyUpdated: suspend (AuthTokenInfo) -> Unit, onUserRetrieved: suspend (User?) -> Unit, onAuthKeyInvalidated: suspend () -> Unit ) { // install(Logging) { // logger = Logger.DEFAULT // level = LogLevel.HEADERS // } install(HttpCookies) { // Will keep an in-memory map with all the cookies from previous requests. storage = AcceptAllCookiesStorage() } val authMutex = Mutex() var currentRefreshToken: RefreshToken? = null initialAuthKey.onFirst { currentRefreshToken = it as? RefreshToken }.onSecond { currentRefreshToken = it.refresh } val creds = initialAuthKey.t1 as? AuthCreds var userRefreshJob: Job? = null install("Auth Token Refresher") { val clientAuthFeature = ClientAuthFeature(this, baseUrl) fun refreshUser(newTokenInfo: AuthTokenInfo) { userRefreshJob ?.cancel() userRefreshJob = scope.launchSafelyWithoutExceptions { onUserRetrieved(clientAuthFeature.getMe(newTokenInfo.token)) } } initialAuthKey.onSecond { refreshUser(it) } suspend fun refreshToken() { val capturedRefresh = currentRefreshToken runCatching { when { capturedRefresh == null && creds == null -> throw AuthUnavailableException capturedRefresh != null -> { currentRefreshToken = null val newTokenInfo = clientAuthFeature.refresh(capturedRefresh) currentRefreshToken = newTokenInfo ?.refresh if (newTokenInfo == null) { refreshToken() } else { onAuthKeyUpdated(newTokenInfo) refreshUser(newTokenInfo) } } creds != null -> { val newAuthTokenInfo = clientAuthFeature.auth(creds) if (newAuthTokenInfo != null) { onAuthKeyUpdated(newAuthTokenInfo) refreshUser(newAuthTokenInfo) currentRefreshToken = newAuthTokenInfo.refresh } } } }.onFailure { onAuthKeyInvalidated() } } sendPipeline.intercept(HttpSendPipeline.State) { if (!context.url.buildString().startsWith(baseUrl) || clientAuthFeature.isAuthRequest(context)) { return@intercept } context.expectSuccess = false if (authMutex.isLocked) { authMutex.withLock { /* do nothing, just wait while mutex will be freed */ } } } receivePipeline.intercept(HttpReceivePipeline.Before) { if ( context.request.url.toString().startsWith(baseUrl) && context.response.status == HttpStatusCode.Unauthorized ) { authMutex.withLock { refreshToken() } val newResponse = context.client ?.request{ takeFrom(context.request) } ?: return@intercept proceedWith(newResponse) } } } }