Skip to content

Remove mutex on get for JwtAuthService.getOrRefreshToken #431

@trema96

Description

@trema96

Current implementation of JwtAuthService uses a mutex both on on set, to avoid doing multiple requests when refreshing the JWT, but the mutex is also applied on read.

This means that if there are multiple parallel requests all the requests will have a small bottleneck in the getting the JWT, since all the requests will get the cached jwt sequentially. This shouldn't be an issue for most applications, but in case it becomes a problem we could in future switch to an AtomicReference<Deferred<JWT>> like in this unmerged PR

@InternalIcureApi
abstract class TokenRefresher(
initialBearer: JwtBearer?,
private var jwtRefresh: JwtRefresh,
private val refreshPadding: Duration,
protected val authApi: RawAnonymousAuthApi,
) {
companion object {
private const val MILLISECONDS_BETWEEN_ERROR_RETRIES = 10_000
}
private val scope = CoroutineScope(Dispatchers.Default + SupervisorJob())
private val status = atomic<Deferred<Status>?>(initialBearer?.let { CompletableDeferred(Status.CurrentToken(it)) })
private fun needsRefresh(status: Status, acceptExpired: Boolean): Boolean = when(status) {
is Status.CurrentToken -> !acceptExpired && isJwtExpiredOrInvalid(status.jwt.token, refreshPadding)
is Status.Failure ->
if (Clock.System.now().toEpochMilliseconds() - status.timestamp < MILLISECONDS_BETWEEN_ERROR_RETRIES) {
true
} else {
throw status.error
}
is Status.Unauthorized -> throw status.error
}
private fun Status.getToken() = if (this is Status.CurrentToken) jwt else throw IllegalStateException("Cannot get token from status")
private suspend fun regenerateAuthJwtUsingRefreshToken(refreshToken: JwtRefresh): Status = try {
authApi.refresh(refreshToken.token).successBody().token?.let {
Status.CurrentToken(JwtBearer(it))
} ?: Status.Failure(
error = InternalCardinalException("Token refresh was successful but token was null"),
timestamp = Clock.System.now().toEpochMilliseconds()
)
} catch (e: Exception) {
e.toStatus()
}
tailrec suspend fun getAuthenticationJwtOrRefresh(acceptExpired: Boolean): JwtBearer {
val deferredStatus = status.value
if (deferredStatus != null && !needsRefresh(deferredStatus.await(), acceptExpired)) {
return deferredStatus.await().getToken()
} else {
val newDeferred = scope.async(start = CoroutineStart.LAZY) {
if (!isJwtExpiredOrInvalid(jwtRefresh.token, refreshPadding)) {
regenerateAuthJwtUsingRefreshToken(jwtRefresh)
} else {
try {
val refreshResult = regenerateRefreshToken()
jwtRefresh = refreshResult.refresh
Status.CurrentToken(refreshResult.bearer)
} catch (e: Exception) {
e.toStatus()
}
}
}
return if (status.compareAndSet(deferredStatus, newDeferred)) {
newDeferred.await().getToken()
} else {
newDeferred.cancel()
getAuthenticationJwtOrRefresh(acceptExpired)
}
}
}
fun getRefreshJwt(): JwtRefresh = jwtRefresh
abstract suspend fun regenerateRefreshToken(): JwtBearerAndRefresh
sealed interface Status {
data class CurrentToken(val jwt: JwtBearer): Status
data class Unauthorized(val error: Throwable): Status
data class Failure(val error: Throwable, val timestamp: Long): Status
}
private fun Throwable.toStatus() = if (this is RequestStatusException && statusCode == 401) {
Status.Unauthorized(error = this)
} else {
Status.Failure(error = this, timestamp = Clock.System.now().toEpochMilliseconds())
}
}

Currently however AtomicFU is still in beta, it might be safer to wait for it to become stable before we use it in such a core part of the SDK. If needed a possibility would be to continue using the mutex on native, and use the atomic reference only on jvm or js

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions