Make session refresh retries idempotent

Preserve the successor session for legitimate refresh retries so transient failures no longer revoke the user's session family.

src/main/kotlin/com/osglab/account/config/AppConfig.kt

@@ -56,6 +56,10 @@ data class AppConfig(
                 hmacSecret = config.secret("app.session.secret", production).toByteArray(),
                 accessMinutes = config.positiveLong("app.session.accessMinutes"),
                 refreshDays = config.positiveLong("app.session.refreshDays"),
+                legacyRefreshReplaySeconds = config.positiveLong(
+                    "app.session.legacyRefreshReplaySeconds",
+                    30,
+                ),
                 gatewayGrantDays = config.positiveLong("app.session.gatewayGrantDays", 30),
             )
             val encryption = EncryptionConfig(
@@ -244,6 +248,9 @@ data class AppConfig(
             require(session.refreshDays in 1..365) {
                 "app.session.refreshDays must be between 1 and 365"
             }
+            require(session.legacyRefreshReplaySeconds in 5..120) {
+                "app.session.legacyRefreshReplaySeconds must be between 5 and 120"
+            }
             require(!production || providers.volcengine.credentialsAvailable) {
                 "Production Volcengine credentials are missing"
             }
@@ -380,6 +387,7 @@ data class SessionConfig(
     val hmacSecret: ByteArray,
     val accessMinutes: Long,
     val refreshDays: Long,
+    val legacyRefreshReplaySeconds: Long = 30,
     val gatewayGrantDays: Long = 30,
 )
 

src/main/kotlin/com/osglab/account/features/auth/AuthRepository.kt

@@ -6,7 +6,9 @@ import org.jetbrains.exposed.v1.core.Table
 import org.jetbrains.exposed.v1.core.and
 import org.jetbrains.exposed.v1.core.eq
 import org.jetbrains.exposed.v1.core.greater
+import org.jetbrains.exposed.v1.core.isNotNull
 import org.jetbrains.exposed.v1.core.isNull
+import org.jetbrains.exposed.v1.core.lessEq
 import org.jetbrains.exposed.v1.javatime.timestamp
 import org.jetbrains.exposed.v1.jdbc.insert
 import org.jetbrains.exposed.v1.jdbc.insertIgnore
@@ -47,6 +49,10 @@ internal object SessionsTable : Table("sessions") {
     val familyId = varchar("family_id", 36).index()
     val refreshTokenHash = varchar("refresh_token_hash", 64).uniqueIndex()
     val replacedById = varchar("replaced_by_id", 36).nullable()
+    val refreshOperationId = varchar("refresh_operation_id", 36).nullable()
+    val encryptedReplacementRefreshToken =
+        varchar("encrypted_replacement_refresh_token", 255).nullable()
+    val refreshReplayUntil = timestamp("refresh_replay_until").nullable().index()
     val createdAt = timestamp("created_at")
     val expiresAt = timestamp("expires_at")
     val revokedAt = timestamp("revoked_at").nullable()
@@ -66,6 +72,16 @@ data class CreatedSession(
     val familyId: UUID,
 )
 
+data class RefreshRotationAttempt(
+    val currentTokenHash: String,
+    val newTokenHash: String,
+    val encryptedNewToken: String,
+    val newExpiresAt: Instant,
+    val operationId: UUID?,
+    val replayUntil: Instant,
+    val now: Instant,
+)
+
 sealed interface RefreshRotationResult {
     data class Rotated(
         val accountId: UUID,
@@ -73,12 +89,21 @@ sealed interface RefreshRotationResult {
         val familyId: UUID,
     ) : RefreshRotationResult
 
+    data class Replayed(
+        val accountId: UUID,
+        val sessionId: UUID,
+        val familyId: UUID,
+        val encryptedRefreshToken: String,
+        val refreshTokenExpiresAt: Instant,
+    ) : RefreshRotationResult
+
     data object Invalid : RefreshRotationResult
     data object ReuseDetected : RefreshRotationResult
 }
 
 internal enum class RefreshRotationDecision {
     ROTATE,
+    REPLAY_ROTATION,
     REVOKE_EXPIRED,
     REVOKE_REUSED_FAMILY,
 }
@@ -91,9 +116,11 @@ internal object RefreshRotationPolicy {
     fun decide(
         revoked: Boolean,
         replaced: Boolean,
+        replayable: Boolean,
         expiresAt: Instant,
         now: Instant,
     ): RefreshRotationDecision = when {
+        replayable -> RefreshRotationDecision.REPLAY_ROTATION
         revoked || replaced -> RefreshRotationDecision.REVOKE_REUSED_FAMILY
         !expiresAt.isAfter(now) -> RefreshRotationDecision.REVOKE_EXPIRED
         else -> RefreshRotationDecision.ROTATE
@@ -114,12 +141,7 @@ interface AuthRepository {
         now: Instant,
     ): CreatedSession
 
-    suspend fun rotateRefreshToken(
-        currentTokenHash: String,
-        newTokenHash: String,
-        newExpiresAt: Instant,
-        now: Instant,
-    ): RefreshRotationResult
+    suspend fun rotateRefreshToken(attempt: RefreshRotationAttempt): RefreshRotationResult
 
     suspend fun revokeSessionFamily(accountId: UUID, sessionId: UUID, now: Instant): Boolean
     suspend fun isSessionActive(accountId: UUID, sessionId: UUID, now: Instant): Boolean
@@ -215,37 +237,78 @@ class ExposedAuthRepository(
     }
 
     override suspend fun rotateRefreshToken(
-        currentTokenHash: String,
-        newTokenHash: String,
-        newExpiresAt: Instant,
-        now: Instant,
+        attempt: RefreshRotationAttempt,
     ): RefreshRotationResult = databaseFactory.query {
+        SessionsTable.update({
+            SessionsTable.refreshReplayUntil.isNotNull() and
+                (SessionsTable.refreshReplayUntil lessEq attempt.now)
+        }) {
+            it[refreshOperationId] = null
+            it[encryptedReplacementRefreshToken] = null
+            it[refreshReplayUntil] = null
+        }
         val current = SessionsTable.selectAll()
-            .where { SessionsTable.refreshTokenHash eq currentTokenHash }
+            .where { SessionsTable.refreshTokenHash eq attempt.currentTokenHash }
             .forUpdate()
             .singleOrNull()
             ?: return@query RefreshRotationResult.Invalid
         val familyId = current[SessionsTable.familyId]
+        val replacement = current[SessionsTable.replacedById]?.takeIf {
+            current[SessionsTable.refreshOperationId] == attempt.operationId?.toString() &&
+                current[SessionsTable.refreshReplayUntil]?.isAfter(attempt.now) == true &&
+                current[SessionsTable.encryptedReplacementRefreshToken] != null
+        }?.let { replacementId ->
+            SessionsTable.selectAll()
+                .where { SessionsTable.id eq replacementId }
+                .forUpdate()
+                .singleOrNull()
+                ?.takeIf {
+                    it[SessionsTable.accountId] == current[SessionsTable.accountId] &&
+                        it[SessionsTable.familyId] == familyId &&
+                        it[SessionsTable.revokedAt] == null &&
+                        it[SessionsTable.replacedById] == null &&
+                        it[SessionsTable.expiresAt].isAfter(attempt.now)
+                }
+        }
         when (
             RefreshRotationPolicy.decide(
                 revoked = current[SessionsTable.revokedAt] != null,
                 replaced = current[SessionsTable.replacedById] != null,
+                replayable = replacement != null,
                 expiresAt = current[SessionsTable.expiresAt],
-                now = now,
+                now = attempt.now,
             )
         ) {
+            RefreshRotationDecision.REPLAY_ROTATION -> {
+                val replayed = requireNotNull(replacement)
+                return@query RefreshRotationResult.Replayed(
+                    accountId = UUID.fromString(replayed[SessionsTable.accountId]),
+                    sessionId = UUID.fromString(replayed[SessionsTable.id]),
+                    familyId = UUID.fromString(replayed[SessionsTable.familyId]),
+                    encryptedRefreshToken = requireNotNull(
+                        current[SessionsTable.encryptedReplacementRefreshToken],
+                    ),
+                    refreshTokenExpiresAt = replayed[SessionsTable.expiresAt],
+                )
+            }
             RefreshRotationDecision.REVOKE_REUSED_FAMILY -> {
                 SessionsTable.update({ SessionsTable.familyId eq familyId }) {
-                    it[SessionsTable.revokedAt] = now
+                    it[SessionsTable.revokedAt] = attempt.now
+                    it[refreshOperationId] = null
+                    it[encryptedReplacementRefreshToken] = null
+                    it[refreshReplayUntil] = null
                 }
                 SessionsTable.update({ SessionsTable.id eq current[SessionsTable.id] }) {
-                    it[SessionsTable.reuseDetectedAt] = now
+                    it[SessionsTable.reuseDetectedAt] = attempt.now
                 }
                 return@query RefreshRotationResult.ReuseDetected
             }
             RefreshRotationDecision.REVOKE_EXPIRED -> {
                 SessionsTable.update({ SessionsTable.familyId eq familyId }) {
-                    it[SessionsTable.revokedAt] = now
+                    it[SessionsTable.revokedAt] = attempt.now
+                    it[refreshOperationId] = null
+                    it[encryptedReplacementRefreshToken] = null
+                    it[refreshReplayUntil] = null
                 }
                 return@query RefreshRotationResult.Invalid
             }
@@ -257,13 +320,16 @@ class ExposedAuthRepository(
             it[SessionsTable.id] = newSessionId.toString()
             it[SessionsTable.accountId] = current[SessionsTable.accountId]
             it[SessionsTable.familyId] = familyId
-            it[SessionsTable.refreshTokenHash] = newTokenHash
-            it[SessionsTable.createdAt] = now
-            it[SessionsTable.expiresAt] = newExpiresAt
+            it[SessionsTable.refreshTokenHash] = attempt.newTokenHash
+            it[SessionsTable.createdAt] = attempt.now
+            it[SessionsTable.expiresAt] = attempt.newExpiresAt
         }
         SessionsTable.update({ SessionsTable.id eq current[SessionsTable.id] }) {
             it[SessionsTable.replacedById] = newSessionId.toString()
-            it[SessionsTable.revokedAt] = now
+            it[SessionsTable.revokedAt] = attempt.now
+            it[SessionsTable.refreshOperationId] = attempt.operationId?.toString()
+            it[SessionsTable.encryptedReplacementRefreshToken] = attempt.encryptedNewToken
+            it[SessionsTable.refreshReplayUntil] = attempt.replayUntil
         }
         RefreshRotationResult.Rotated(
             accountId = UUID.fromString(current[SessionsTable.accountId]),
@@ -290,6 +356,9 @@ class ExposedAuthRepository(
                 (SessionsTable.familyId eq session[SessionsTable.familyId])
         }) {
             it[SessionsTable.revokedAt] = now
+            it[refreshOperationId] = null
+            it[encryptedReplacementRefreshToken] = null
+            it[refreshReplayUntil] = null
         } > 0
     }
 

src/main/kotlin/com/osglab/account/features/auth/AuthRoutes.kt

@@ -1,6 +1,7 @@
 package com.osglab.account.features.auth
 
 import com.osglab.account.common.api.ApiResponse
+import com.osglab.account.common.errors.InvalidRequestException
 import com.osglab.account.common.errors.UnauthorizedException
 import com.osglab.account.common.security.AccountPrincipal
 import com.osglab.account.common.security.SESSION_AUTH_NAME
@@ -15,6 +16,7 @@ import io.ktor.server.routing.Route
 import io.ktor.server.routing.post
 import io.ktor.server.routing.route
 import kotlinx.serialization.Serializable
+import java.util.UUID
 
 @Serializable
 data class AppleSignInRequest(
@@ -44,8 +46,12 @@ data class AppAttestRequest(
 }
 
 @Serializable
-data class RefreshSessionRequest(val refreshToken: String) {
-    override fun toString(): String = "RefreshSessionRequest(refreshToken=[REDACTED])"
+data class RefreshSessionRequest(
+    val refreshToken: String,
+    val refreshOperationId: String? = null,
+) {
+    override fun toString(): String =
+        "RefreshSessionRequest(refreshToken=[REDACTED], refreshOperationId=$refreshOperationId)"
 }
 
 @Serializable
@@ -92,8 +98,11 @@ class AuthRoutes(
                 }
                 post("/refresh") {
                     val request = call.receive<RefreshSessionRequest>()
+                    val operationId = request.refreshOperationId?.let(::parseRefreshOperationId)
                     call.respond(
-                        ApiResponse(data = sessionService.refresh(request.refreshToken).toResponse()),
+                        ApiResponse(
+                            data = sessionService.refresh(request.refreshToken, operationId).toResponse(),
+                        ),
                     )
                 }
                 authenticate(SESSION_AUTH_NAME) {
@@ -112,6 +121,10 @@ class AuthRoutes(
 fun Route.authRoutes(sessionService: SessionService) =
     AuthRoutes(sessionService).register(this)
 
+private fun parseRefreshOperationId(value: String): UUID =
+    runCatching { UUID.fromString(value) }
+        .getOrElse { throw InvalidRequestException("refreshOperationId must be a UUID") }
+
 private fun SessionTokens.toResponse(): SessionTokenResponse = SessionTokenResponse(
     accountId = accountId.toString(),
     accessToken = accessToken,

src/main/kotlin/com/osglab/account/features/auth/SessionService.kt

@@ -96,29 +96,52 @@ class SessionService(
         return createSession(account.id, now)
     }
 
-    suspend fun refresh(refreshToken: String): SessionTokens {
+    suspend fun refresh(refreshToken: String, operationId: UUID? = null): SessionTokens {
         requireValue(refreshToken, "refreshToken", MAX_REFRESH_TOKEN_LENGTH)
         val now = clock.instant()
+        val currentTokenHash = TokenHash.sha256(refreshToken)
         val replacement = tokenGenerator.newRefreshToken()
         val replacementExpiresAt = now.plus(Duration.ofDays(sessionConfig.refreshDays))
+        val replayUntil = if (operationId == null) {
+            now.plusSeconds(sessionConfig.legacyRefreshReplaySeconds)
+        } else {
+            // A stable operation ID lets a crashed client recover until the successor expires.
+            replacementExpiresAt
+        }
         return when (
             val result = repository.rotateRefreshToken(
-                currentTokenHash = TokenHash.sha256(refreshToken),
-                newTokenHash = TokenHash.sha256(replacement),
-                newExpiresAt = replacementExpiresAt,
-                now = now,
+                RefreshRotationAttempt(
+                    currentTokenHash = currentTokenHash,
+                    newTokenHash = TokenHash.sha256(replacement),
+                    encryptedNewToken = fieldEncryptor.encrypt(
+                        replacement,
+                        refreshReplayContext(currentTokenHash),
+                    ),
+                    newExpiresAt = replacementExpiresAt,
+                    operationId = operationId,
+                    replayUntil = replayUntil,
+                    now = now,
+                ),
             )
         ) {
             RefreshRotationResult.Invalid -> throw UnauthorizedException("Refresh token is invalid or expired")
             RefreshRotationResult.ReuseDetected -> throw TokenReuseException()
-            is RefreshRotationResult.Rotated -> {
-                val access = sessionJwt.issue(result.accountId, result.sessionId)
-                SessionTokens(
+            is RefreshRotationResult.Rotated -> issueSessionTokens(
+                accountId = result.accountId,
+                sessionId = result.sessionId,
+                refreshToken = replacement,
+                refreshTokenExpiresAt = replacementExpiresAt,
+            )
+            is RefreshRotationResult.Replayed -> {
+                val replayedRefreshToken = fieldEncryptor.decrypt(
+                    result.encryptedRefreshToken,
+                    refreshReplayContext(currentTokenHash),
+                )
+                issueSessionTokens(
                     accountId = result.accountId,
-                    accessToken = access.value,
-                    accessTokenExpiresAt = access.expiresAt,
-                    refreshToken = replacement,
-                    refreshTokenExpiresAt = replacementExpiresAt,
+                    sessionId = result.sessionId,
+                    refreshToken = replayedRefreshToken,
+                    refreshTokenExpiresAt = result.refreshTokenExpiresAt,
                 )
             }
         }
@@ -147,6 +170,22 @@ class SessionService(
         )
     }
 
+    private fun issueSessionTokens(
+        accountId: UUID,
+        sessionId: UUID,
+        refreshToken: String,
+        refreshTokenExpiresAt: Instant,
+    ): SessionTokens {
+        val access = sessionJwt.issue(accountId, sessionId)
+        return SessionTokens(
+            accountId = accountId,
+            accessToken = access.value,
+            accessTokenExpiresAt = access.expiresAt,
+            refreshToken = refreshToken,
+            refreshTokenExpiresAt = refreshTokenExpiresAt,
+        )
+    }
+
     private suspend fun verifyIdentityToken(token: String, nonce: String): AppleIdentity =
         try {
             appleIdentityVerifier.verify(token, nonce)
@@ -185,3 +224,4 @@ class SessionService(
 
 fun appleRefreshContext(accountId: UUID): String = "apple-refresh-token:$accountId"
 fun appleSubjectContext(identityFingerprint: String): String = "apple-subject:$identityFingerprint"
+fun refreshReplayContext(currentTokenHash: String): String = "session-refresh-replay:$currentTokenHash"