Merge pull request #64 from garfiec/fix/clearTokens-refresh-mutex
fix(data): mutex guard clearTokens vs refreshAccessToken race
This commit is contained in:
Garfie Chiu
GitHub
commit
cb7e7941a5
3 changed files with 213 additions and 2 deletions
|
|
@ -30,6 +30,7 @@ kotlin {
|
|||
}
|
||||
named("androidUnitTest").dependencies {
|
||||
implementation(libs.koin.test)
|
||||
implementation(libs.ktor.client.mock)
|
||||
}
|
||||
named("androidInstrumentedTest").dependencies {
|
||||
implementation(libs.room.testing)
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
}
|
||||
}
|
||||
|
|
@ -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()
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue