Test token expiry handling

Author: simolus3

core/src/commonIntegrationTest/kotlin/com/powersync/sync/SyncIntegrationTest.kt

@@ -16,11 +16,17 @@ import com.powersync.testutils.UserRow
 import com.powersync.testutils.databaseTest
 import com.powersync.testutils.waitFor
 import com.powersync.utils.JsonUtil
+import dev.mokkery.answering.returns
+import dev.mokkery.every
+import dev.mokkery.everySuspend
 import dev.mokkery.verify
+import dev.mokkery.verifyNoMoreCalls
+import dev.mokkery.verifySuspend
 import io.kotest.matchers.collections.shouldHaveSize
 import io.kotest.matchers.shouldBe
 import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.DelicateCoroutinesApi
+import kotlinx.coroutines.launch
 import kotlinx.serialization.encodeToString
 import kotlin.test.Test
 import kotlin.test.assertEquals
@@ -526,4 +532,62 @@ class SyncIntegrationTest {
             // Meaning that the two rows are now visible
             database.expectUserCount(2)
         }
+
+    @Test
+    fun testTokenExpired() =
+        databaseTest {
+            turbineScope(timeout = 10.0.seconds) {
+                val turbine = database.currentStatus.asFlow().testIn(this)
+
+                database.connect(connector, 1000L, retryDelayMs = 5000)
+                turbine.waitFor { it.connecting }
+
+                syncLines.send(SyncLine.KeepAlive(tokenExpiresIn = 4000))
+                turbine.waitFor { it.connected }
+                verifySuspend { connector.getCredentialsCached() }
+                verifyNoMoreCalls(connector)
+
+                // Should invalidate credentials when token expires
+                syncLines.send(SyncLine.KeepAlive(tokenExpiresIn = 0))
+                turbine.waitFor { !it.connected }
+                verify { connector.invalidateCredentials() }
+
+                turbine.cancel()
+            }
+        }
+
+    @Test
+    fun testTokenPrefetch() =
+        databaseTest {
+            val prefetchCalled = CompletableDeferred<Unit>()
+            val completePrefetch = CompletableDeferred<Unit>()
+            every { connector.prefetchCredentials() } returns scope.launch {
+                prefetchCalled.complete(Unit)
+                completePrefetch.await()
+            }
+
+            turbineScope(timeout = 10.0.seconds) {
+                val turbine = database.currentStatus.asFlow().testIn(this)
+
+                database.connect(connector, 1000L, retryDelayMs = 5000)
+                turbine.waitFor { it.connecting }
+
+                syncLines.send(SyncLine.KeepAlive(tokenExpiresIn = 4000))
+                turbine.waitFor { it.connected }
+                verifySuspend { connector.getCredentialsCached() }
+                verifyNoMoreCalls(connector)
+
+                syncLines.send(SyncLine.KeepAlive(tokenExpiresIn = 10))
+                prefetchCalled.complete(Unit)
+                // Should still be connected before prefetch completes
+                database.currentStatus.connected shouldBe true
+
+                // After the prefetch completes, we should reconnect
+                completePrefetch.complete(Unit)
+                turbine.waitFor { !it.connected }
+
+                turbine.waitFor { it.connected }
+                turbine.cancel()
+            }
+        }
 }

core/src/commonMain/kotlin/com/powersync/connectors/PowerSyncBackendConnector.kt

@@ -56,18 +56,19 @@ public abstract class PowerSyncBackendConnector {
      * This may be called before the current credentials have expired.
      */
     @Throws(PowerSyncException::class, CancellationException::class)
-    public open suspend fun prefetchCredentials(): Job? {
+    public open fun prefetchCredentials(): Job {
         fetchRequest?.takeIf { it.isActive }?.let { return it }
 
-        fetchRequest =
+        val request  =
             scope.launch {
                 fetchCredentials().also { value ->
                     cachedCredentials = value
                     fetchRequest = null
                 }
             }
 
-        return fetchRequest
+        fetchRequest = request
+        return request
     }
 
     /**

core/src/commonMain/kotlin/com/powersync/sync/SyncStream.kt

@@ -47,6 +47,7 @@ import kotlinx.coroutines.NonCancellable
 import kotlinx.coroutines.cancelAndJoin
 import kotlinx.coroutines.channels.BufferOverflow
 import kotlinx.coroutines.channels.Channel
+import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.emitAll
@@ -256,21 +257,25 @@ internal class SyncStream(
     }
 
     private suspend fun streamingSyncIteration() {
-        val iteration = ActiveIteration()
-
-        try {
-            iteration.start()
-        } finally {
-            // This can't be cancelled because we need to send a stop message, which is async, to
-            // clean up resources.
-            withContext(NonCancellable) {
-                iteration.stop()
+        coroutineScope {
+            val iteration = ActiveIteration(this)
+
+            try {
+                iteration.start()
+            } finally {
+                // This can't be cancelled because we need to send a stop message, which is async, to
+                // clean up resources.
+                withContext(NonCancellable) {
+                    iteration.stop()
+                }
             }
         }
     }
 
     private inner class ActiveIteration(
+        val scope: CoroutineScope,
         var fetchLinesJob: Job? = null,
+        var credentialsInvalidation: Job? = null,
     ) {
         suspend fun start() {
             control("start", JsonUtil.json.encodeToString(params))
@@ -336,7 +341,25 @@ internal class SyncStream(
                         applyCoreChanges(instruction.status)
                     }
                 }
-                is Instruction.FetchCredentials -> TODO()
+                is Instruction.FetchCredentials -> {
+                    if (instruction.didExpire) {
+                        connector.invalidateCredentials()
+                    } else {
+                        // Token expires soon - refresh it in the background
+                        if (credentialsInvalidation == null) {
+                            val job = scope.launch {
+                                connector.prefetchCredentials().join()
+
+                                // Token has been refreshed, start another iteration
+                                stop()
+                            }
+                            job.invokeOnCompletion {
+                                credentialsInvalidation = null
+                            }
+                            credentialsInvalidation = job
+                        }
+                    }
+                }
                 Instruction.DidCompleteSync -> status.update { copy(downloadError=null) }
                 is Instruction.UnknownInstruction -> {
                     throw PowerSyncException("Unknown instruction received from core extension: ${instruction.raw}", null)