fix(android-sdk): discard in-flight token responses after sign-out

android-sdk/android/src/main/kotlin/io/logto/sdk/android/LogtoClient.kt

@@ -32,6 +32,13 @@ open class LogtoClient(
     val logtoConfig: LogtoConfig,
     application: Application,
 ) {
+    /**
+     * Guards the credential fields below: token flows that were in flight when
+     * [signOut] or [clearCredentials] dropped the credentials must not persist
+     * their (now stale) results. See [CredentialGuard].
+     */
+    private val credentialGuard = CredentialGuard()
+
     /**
      * Cached access tokens.
      */
@@ -84,6 +91,11 @@ open class LogtoClient(
 
     /**
      * Sign in
+     *
+     * If a sign-out happens while the sign-in is still in progress, the sign-in result
+     * is discarded and the completion receives a
+     * [LogtoException.Type.NOT_AUTHENTICATED] error.
+     *
      * @param[context] the activity to perform a sign-in action
      * @param[options] the sign-in options
      * @param[completion] the completion which handles the result of signing in
@@ -93,6 +105,8 @@ open class LogtoClient(
         options: SignInOptions,
         completion: EmptyCompletion<LogtoException>,
     ) {
+        val credentialStamp = credentialGuard.stamp()
+
         getOidcConfig { getOidcConfigException, oidcConfig ->
             getOidcConfigException?.let {
                 completion.onComplete(it)
@@ -119,6 +133,7 @@ open class LogtoClient(
                 )
 
                 verifyAndSaveTokenResponse(
+                    credentialStamp = credentialStamp,
                     issuer = oidcConfig.issuer,
                     responseIdToken = codeToken.idToken,
                     responseRefreshToken = codeToken.refreshToken,
@@ -162,6 +177,10 @@ open class LogtoClient(
      * re-enter the account through the existing browser session. Use [signOut] for a
      * complete sign-out.
      *
+     * Any token request that is still in flight when the credentials are cleared is
+     * discarded: its result is not persisted and its completion receives a
+     * [LogtoException.Type.NOT_AUTHENTICATED] error.
+     *
      * @param[completion] the completion invoked with any error that occurs while clearing
      * the credentials
      */
@@ -171,11 +190,7 @@ open class LogtoClient(
             return
         }
 
-        accessTokenMap.clear()
-        idToken = null
-
-        refreshToken?.let { tokenToRevoke ->
-            refreshToken = null
+        dropCredentials()?.let { tokenToRevoke ->
             getOidcConfig { getOidcConfigException, oidcConfig ->
                 getOidcConfigException?.let {
                     completion?.onComplete(it)
@@ -220,6 +235,10 @@ open class LogtoClient(
      * [LogtoException.Type.INVALID_REDIRECT_URI] without opening the browser; local
      * credentials are still cleared and the revocation is still attempted.
      *
+     * Any token request that is still in flight when the sign-out starts is discarded:
+     * its result is not persisted and its completion receives a
+     * [LogtoException.Type.NOT_AUTHENTICATED] error.
+     *
      * @param[context] the activity to perform the sign-out action
      * @param[postLogoutRedirectUri] one of the post sign-out redirect URIs of this
      * application, or `null` to let the user dismiss the browser manually after the
@@ -241,10 +260,7 @@ open class LogtoClient(
             return
         }
 
-        val tokenToRevoke = refreshToken
-        accessTokenMap.clear()
-        idToken = null
-        refreshToken = null
+        val tokenToRevoke = dropCredentials()
 
         getOidcConfig { getOidcConfigException, oidcConfig ->
             getOidcConfigException?.let {
@@ -328,6 +344,11 @@ open class LogtoClient(
         organizationId: String?,
         completion: Completion<LogtoException, AccessToken>,
     ) {
+        // The stamp must be taken before any credential is read: a sign-out that lands
+        // between the read and the stamp would otherwise go unnoticed and the refreshed
+        // tokens would be committed against the already-cleared credentials.
+        val credentialStamp = credentialGuard.stamp()
+
         if (!isAuthenticated) {
             completion.onComplete(LogtoException(LogtoException.Type.NOT_AUTHENTICATED), null)
             return
@@ -354,7 +375,11 @@ open class LogtoClient(
         }
 
         // MARK: If cannot refresh the access token, then return a NOT_AUTHENTICATED error
-        if (refreshToken == null) {
+        // Snapshot the refresh token: a concurrent sign-out can null the field while this
+        // flow is between its async hops; the flow runs on the snapshot and the credential
+        // guard arbitrates at commit time.
+        val tokenForRefresh = refreshToken
+        if (tokenForRefresh == null) {
             completion.onComplete(LogtoException(LogtoException.Type.NOT_AUTHENTICATED), null)
             return
         }
@@ -369,7 +394,7 @@ open class LogtoClient(
             Core.fetchTokenByRefreshToken(
                 tokenEndpoint = requireNotNull(oidcConfig).tokenEndpoint,
                 clientId = logtoConfig.appId,
-                refreshToken = requireNotNull(refreshToken),
+                refreshToken = tokenForRefresh,
                 resource = resource,
                 organizationId = organizationId,
                 scopes = null,
@@ -396,6 +421,7 @@ open class LogtoClient(
                 )
 
                 verifyAndSaveTokenResponse(
+                    credentialStamp = credentialStamp,
                     issuer = oidcConfig.issuer,
                     responseIdToken = refreshedToken.idToken,
                     responseRefreshToken = refreshedToken.refreshToken,
@@ -414,12 +440,14 @@ open class LogtoClient(
      * @param[completion] the completion which handles the retrieved result
      */
     fun getIdTokenClaims(completion: Completion<LogtoException, IdTokenClaims>) {
-        if (!isAuthenticated) {
+        // Snapshot the ID token: a concurrent sign-out can null the field at any point
+        val currentIdToken = idToken
+        if (!isAuthenticated || currentIdToken == null) {
             completion.onComplete(LogtoException(LogtoException.Type.NOT_AUTHENTICATED), null)
             return
         }
         try {
-            val idTokenClaims = TokenUtils.decodeIdToken(requireNotNull(idToken))
+            val idTokenClaims = TokenUtils.decodeIdToken(currentIdToken)
             completion.onComplete(null, idTokenClaims)
         } catch (exception: InvalidJwtException) {
             completion.onComplete(
@@ -489,32 +517,68 @@ open class LogtoClient(
         }
     }
 
+    /**
+     * Atomically drop the local credentials and invalidate the token flows that are
+     * still in flight, so that their responses can no longer be persisted.
+     *
+     * @return the refresh token that was current, for the caller to revoke
+     */
+    private fun dropCredentials(): String? = credentialGuard.invalidate {
+        val tokenToRevoke = refreshToken
+        accessTokenMap.clear()
+        idToken = null
+        refreshToken = null
+        tokenToRevoke
+    }
+
     private fun verifyAndSaveTokenResponse(
+        credentialStamp: Int,
         issuer: String,
         responseIdToken: String?,
         responseRefreshToken: String?,
         accessTokenKey: String,
         accessToken: AccessToken,
         completion: EmptyCompletion<LogtoException>,
     ) {
+        // Discard already-stale flows before fetching the JWKS or verifying the response
+        if (!credentialGuard.isCurrent(credentialStamp)) {
+            completion.onComplete(LogtoException(LogtoException.Type.NOT_AUTHENTICATED))
+            return
+        }
+
         getJwks { getJwksException, jwks ->
-            getJwksException?.let {
-                completion.onComplete(it)
-                return@getJwks
-            }
-            responseIdToken?.let {
-                try {
-                    TokenUtils.verifyIdToken(it, logtoConfig.appId, issuer, requireNotNull(jwks))
-                } catch (exception: InvalidJwtException) {
-                    completion.onComplete(LogtoException(LogtoException.Type.INVALID_ID_TOKEN, exception))
-                    return@getJwks
+            val verificationException = getJwksException ?: verifyIdToken(responseIdToken, issuer, jwks)
+
+            val saved = verificationException == null &&
+                credentialGuard.commit(credentialStamp) {
+                    responseIdToken?.let { idToken = it }
+                    accessTokenMap[accessTokenKey] = accessToken
+                    refreshToken = responseRefreshToken
                 }
-                idToken = it
-            }
 
-            accessTokenMap[accessTokenKey] = accessToken
-            refreshToken = responseRefreshToken
-            completion.onComplete(null)
+            completion.onComplete(
+                when {
+                    saved -> null
+                    // Stale flows always complete with NOT_AUTHENTICATED, even when the
+                    // response would also have failed verification
+                    !credentialGuard.isCurrent(credentialStamp) ->
+                        LogtoException(LogtoException.Type.NOT_AUTHENTICATED)
+                    else -> verificationException
+                },
+            )
+        }
+    }
+
+    private fun verifyIdToken(
+        responseIdToken: String?,
+        issuer: String,
+        jwks: JsonWebKeySet?,
+    ): LogtoException? = responseIdToken?.let {
+        try {
+            TokenUtils.verifyIdToken(it, logtoConfig.appId, issuer, requireNotNull(jwks))
+            null
+        } catch (exception: InvalidJwtException) {
+            LogtoException(LogtoException.Type.INVALID_ID_TOKEN, exception)
         }
     }
 
@@ -599,3 +663,37 @@ open class LogtoClient(
         accessTokenMap.putAll(tokenMap)
     }
 }
+
+/**
+ * An optimistic guard for the local credential set — the in-memory equivalent of an
+ * optimistic lock's "UPDATE ... WHERE version = ?".
+ *
+ * Async token flows take a [stamp] when they start, and [commit] applies their writes
+ * only when no [invalidate] has happened in between. This keeps a token response that
+ * lands after a sign-out from resurrecting the cleared credentials, and a response
+ * from before a sign-out from clobbering the session of a later sign-in.
+ */
+private class CredentialGuard {
+    private var version = 0
+
+    @Synchronized
+    fun stamp(): Int = version
+
+    @Synchronized
+    fun isCurrent(stamp: Int): Boolean = stamp == version
+
+    @Synchronized
+    fun <T> invalidate(block: () -> T): T {
+        version++
+        return block()
+    }
+
+    @Synchronized
+    fun commit(stamp: Int, block: () -> Unit): Boolean {
+        if (stamp != version) {
+            return false
+        }
+        block()
+        return true
+    }
+}