fix(data): mutex guard clearTokens vs refreshAccessToken race

Wraps `CommonTokenDataStore.clearTokens()` in the existing `refreshMutex`,
serializing logout against any in-flight `refreshAccessToken()` call
(which already holds the same mutex for its full HTTP duration). Without
this, a refresh completing successfully after `clearTokens()` would
silently resurrect the just-killed session by writing a fresh token to
disk. A private `clearTokensLocked()` helper holds the inner clear logic
and is called from inside `refreshAccessToken()` to avoid self-deadlock
on the non-reentrant mutex.
This commit is contained in:
Garfie 2026-04-27 18:32:06 -06:00
parent f236d97c19
commit e2e3423dcb
3 changed files with 213 additions and 2 deletions

View file

@ -30,6 +30,7 @@ kotlin {
}
named("androidUnitTest").dependencies {
implementation(libs.koin.test)
implementation(libs.ktor.client.mock)
}
named("androidInstrumentedTest").dependencies {
implementation(libs.room.testing)

View file

@ -0,0 +1,204 @@
package com.garfiec.librechat.core.data.datastore
import com.google.common.truth.Truth.assertThat
import io.ktor.client.HttpClient
import io.ktor.client.engine.mock.MockEngine
import io.ktor.client.engine.mock.respond
import io.ktor.client.plugins.contentnegotiation.ContentNegotiation
import io.ktor.client.plugins.defaultRequest
import io.ktor.client.request.url
import io.ktor.http.ContentType
import io.ktor.http.HttpStatusCode
import io.ktor.http.contentType
import io.ktor.http.headersOf
import io.ktor.serialization.kotlinx.json.json
import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.ExperimentalCoroutinesApi
import kotlinx.coroutines.Job
import kotlinx.coroutines.async
import kotlinx.coroutines.launch
import kotlinx.coroutines.test.UnconfinedTestDispatcher
import kotlinx.coroutines.test.advanceUntilIdle
import kotlinx.coroutines.test.runTest
import kotlinx.coroutines.yield
import org.junit.Test
@OptIn(ExperimentalCoroutinesApi::class)
class CommonTokenDataStoreConcurrencyTest {
/**
* Test double that records every platform storage mutation. It deliberately
* mirrors `TokenDataStore` semantics (read-through reads, write-through
* writes) without pulling in `EncryptedSharedPreferences`.
*/
private class FakeTokenDataStore(
refreshClient: Lazy<HttpClient>,
initialAccess: String? = "initial-access",
initialRefresh: String? = "initial-refresh",
) : CommonTokenDataStore(refreshClient) {
@Volatile
private var storedAccess: String? = initialAccess
@Volatile
private var storedRefresh: String? = initialRefresh
val writeLog = mutableListOf<Pair<String, String>>()
var removeCount = 0
fun persistedAccess(): String? = storedAccess
fun persistedRefresh(): String? = storedRefresh
init {
initializeTokenCache()
}
override fun readAccessToken(): String? = storedAccess
override fun readRefreshToken(): String? = storedRefresh
override fun writeTokens(accessToken: String, refreshToken: String) {
storedAccess = accessToken
storedRefresh = refreshToken
writeLog += accessToken to refreshToken
}
override fun removeTokens() {
storedAccess = null
storedRefresh = null
removeCount++
}
override fun onKeystoreCorruption() = Unit
}
private fun mockRefreshClient(engine: MockEngine): Lazy<HttpClient> = lazy {
HttpClient(engine) {
install(ContentNegotiation) { json() }
defaultRequest {
url("https://chat.example.com")
contentType(ContentType.Application.Json)
}
}
}
private fun refreshResponseBody(token: String): String =
"""{"token":"$token"}"""
private val jsonHeaders = headersOf("Content-Type", ContentType.Application.Json.toString())
@Test
fun `clearTokens blocks until in-flight refresh releases the mutex`() =
runTest(UnconfinedTestDispatcher()) {
val release = CompletableDeferred<Unit>()
val engine = MockEngine {
release.await()
respond(
content = refreshResponseBody("new-access"),
status = HttpStatusCode.OK,
headers = jsonHeaders,
)
}
val store = FakeTokenDataStore(mockRefreshClient(engine))
val refreshJob = async { store.refreshAccessToken() }
// Give the refresh coroutine a chance to acquire the mutex and
// suspend inside the MockEngine's release.await() before we start
// the concurrent clear.
yield()
advanceUntilIdle()
val clearJob: Job = launch { store.clearTokens() }
yield()
advanceUntilIdle()
// Clear must still be suspended on the mutex while refresh waits on HTTP.
assertThat(refreshJob.isCompleted).isFalse()
assertThat(clearJob.isCompleted).isFalse()
assertThat(store.removeCount).isEqualTo(0)
release.complete(Unit)
advanceUntilIdle()
// Refresh completes and transiently writes the new token, but the
// queued clearTokens immediately wipes it. Final state: logged out.
assertThat(refreshJob.await()).isTrue()
assertThat(clearJob.isCompleted).isTrue()
assertThat(store.persistedAccess()).isNull()
assertThat(store.persistedRefresh()).isNull()
assertThat(store.removeCount).isEqualTo(1)
}
@Test
fun `refresh does not resurrect a logout that raced its network call`() =
runTest(UnconfinedTestDispatcher()) {
val release = CompletableDeferred<Unit>()
val engine = MockEngine {
release.await()
respond(
content = refreshResponseBody("resurrected-access"),
status = HttpStatusCode.OK,
headers = jsonHeaders,
)
}
val store = FakeTokenDataStore(mockRefreshClient(engine))
val refreshJob = async { store.refreshAccessToken() }
yield()
advanceUntilIdle()
launch { store.clearTokens() }
yield()
advanceUntilIdle()
release.complete(Unit)
advanceUntilIdle()
// Mutex serializes: refresh writes transiently, clear wipes after.
// The invariant is the FINAL persisted state, not intermediate writes.
refreshJob.await()
assertThat(store.persistedAccess()).isNull()
assertThat(store.persistedRefresh()).isNull()
}
@Test
fun `refresh without a concurrent clear still writes the new token`() =
runTest(UnconfinedTestDispatcher()) {
val engine = MockEngine {
respond(
content = refreshResponseBody("fresh-access"),
status = HttpStatusCode.OK,
headers = jsonHeaders,
)
}
val store = FakeTokenDataStore(mockRefreshClient(engine))
val result = store.refreshAccessToken()
assertThat(result).isTrue()
assertThat(store.writeLog).hasSize(1)
assertThat(store.writeLog.single().first).isEqualTo("fresh-access")
assertThat(store.persistedAccess()).isEqualTo("fresh-access")
}
@Test
fun `sequential clearTokens and refreshAccessToken do not deadlock`() =
runTest(UnconfinedTestDispatcher()) {
val engine = MockEngine {
respond(
content = refreshResponseBody("seq-access"),
status = HttpStatusCode.OK,
headers = jsonHeaders,
)
}
val store = FakeTokenDataStore(mockRefreshClient(engine))
store.clearTokens()
assertThat(store.removeCount).isEqualTo(1)
// Without a refresh token in storage the refresh should short-circuit
// false without touching the network — and without deadlocking on a
// mutex that clearTokens already released.
val result = store.refreshAccessToken()
assertThat(result).isFalse()
}
}

View file

@ -93,13 +93,13 @@ abstract class CommonTokenDataStore(
true
} catch (e: ClientRequestException) {
Logger.w(e) { "Auth error during token refresh (status=${e.response.status})" }
clearTokens()
clearTokensLocked()
false
} catch (e: Exception) {
if (isKeystoreException(e)) {
Logger.e(e) { "Keystore corruption—clearing tokens" }
onKeystoreCorruption()
clearTokens()
clearTokensLocked()
} else {
Logger.w(e) { "Error during token refresh" }
}
@ -108,6 +108,12 @@ abstract class CommonTokenDataStore(
}
override suspend fun clearTokens() {
refreshMutex.withLock {
clearTokensLocked()
}
}
private fun clearTokensLocked() {
cachedAccessToken = null
removeTokens()
}