diff --git a/core/data/build.gradle.kts b/core/data/build.gradle.kts index b27e361..8eb8fb5 100644 --- a/core/data/build.gradle.kts +++ b/core/data/build.gradle.kts @@ -30,6 +30,7 @@ kotlin { } named("androidUnitTest").dependencies { implementation(libs.koin.test) + implementation(libs.ktor.client.mock) } named("androidInstrumentedTest").dependencies { implementation(libs.room.testing) diff --git a/core/data/src/androidUnitTest/kotlin/com/garfiec/librechat/core/data/datastore/CommonTokenDataStoreConcurrencyTest.kt b/core/data/src/androidUnitTest/kotlin/com/garfiec/librechat/core/data/datastore/CommonTokenDataStoreConcurrencyTest.kt new file mode 100644 index 0000000..351bbb3 --- /dev/null +++ b/core/data/src/androidUnitTest/kotlin/com/garfiec/librechat/core/data/datastore/CommonTokenDataStoreConcurrencyTest.kt @@ -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, + 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>() + 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 = 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() + 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() + 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() + } +} diff --git a/core/data/src/commonMain/kotlin/com/garfiec/librechat/core/data/datastore/CommonTokenDataStore.kt b/core/data/src/commonMain/kotlin/com/garfiec/librechat/core/data/datastore/CommonTokenDataStore.kt index b056ffc..c7df6c9 100644 --- a/core/data/src/commonMain/kotlin/com/garfiec/librechat/core/data/datastore/CommonTokenDataStore.kt +++ b/core/data/src/commonMain/kotlin/com/garfiec/librechat/core/data/datastore/CommonTokenDataStore.kt @@ -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() }