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

package com.osglab.account.features.auth

import com.osglab.account.common.errors.ExternalServiceUnavailableException
import com.osglab.account.common.errors.InvalidRequestException
import com.osglab.account.common.errors.TokenReuseException
import com.osglab.account.common.errors.UnauthorizedException
import com.osglab.account.common.security.FieldEncryptor
import com.osglab.account.common.security.IdentityFingerprint
import com.osglab.account.common.security.SecureTokenGenerator
import com.osglab.account.common.security.Sha256SecureTokenGenerator
import com.osglab.account.common.security.SessionJwt
import com.osglab.account.common.security.AccountPrincipal
import com.osglab.account.common.security.TokenHash
import com.osglab.account.config.SessionConfig
import com.osglab.account.features.integrity.AppleSignInIntegrityPayload
import com.osglab.account.features.integrity.IntegrityEvidence
import com.osglab.account.features.integrity.IntegrityService
import java.time.Clock
import java.time.Duration
import java.time.Instant
import java.util.UUID

data class SessionTokens(
    val accountId: UUID,
    val accessToken: String,
    val accessTokenExpiresAt: Instant,
    val refreshToken: String,
    val refreshTokenExpiresAt: Instant,
) {
    override fun toString(): String =
        "SessionTokens(accountId=$accountId, accessToken=[REDACTED], " +
            "accessTokenExpiresAt=$accessTokenExpiresAt, refreshToken=[REDACTED], " +
            "refreshTokenExpiresAt=$refreshTokenExpiresAt)"
}

fun interface AccountProvisioner {
    suspend fun provision(accountId: UUID, deviceCheckToken: String?, displayName: String?)
}

class SessionService(
    private val repository: AuthRepository,
    private val appleIdentityVerifier: AppleIdentityTokenVerifier,
    private val appleTokenClient: AppleTokenClient,
    private val integrityService: IntegrityService,
    private val sessionJwt: SessionJwt,
    private val fieldEncryptor: FieldEncryptor,
    private val identityFingerprint: IdentityFingerprint,
    private val sessionConfig: SessionConfig,
    private val accountProvisioner: AccountProvisioner = AccountProvisioner { _, _, _ -> },
    private val tokenGenerator: SecureTokenGenerator = Sha256SecureTokenGenerator(),
    private val clock: Clock = Clock.systemUTC(),
) {
    suspend fun signInWithApple(
        identityToken: String,
        authorizationCode: String,
        nonce: String,
        integrityEvidence: IntegrityEvidence,
        displayName: String? = null,
    ): SessionTokens {
        requireValue(identityToken, "identityToken", MAX_IDENTITY_TOKEN_LENGTH)
        requireValue(authorizationCode, "authorizationCode", MAX_AUTHORIZATION_CODE_LENGTH)
        requireValue(nonce, "nonce", MAX_NONCE_LENGTH)
        val verifiedIntegrity = integrityService.verifyAppleSignIn(
            integrityEvidence,
            AppleSignInIntegrityPayload(identityToken, authorizationCode, nonce),
        )

        val suppliedIdentity = verifyIdentityToken(identityToken, nonce)
        val exchange = exchangeCode(authorizationCode)
        val exchangedIdentity = verifyIdentityToken(exchange.identityToken, nonce)
        if (suppliedIdentity.subject != exchangedIdentity.subject) {
            throw UnauthorizedException("Apple authorization code does not match identity token")
        }

        val now = clock.instant()
        val fingerprint = identityFingerprint.ofAppleSubject(suppliedIdentity.subject)
        val account = repository.findOrCreateAccount(
            identityFingerprint = fingerprint,
            encryptedAppleSubject = fieldEncryptor.encrypt(
                suppliedIdentity.subject,
                appleSubjectContext(fingerprint),
            ),
            now = now,
        )
        integrityService.bindVerifiedKey(verifiedIntegrity.appAttestKeyId, account.id)
        repository.updateAppleRefreshToken(
            account.id,
            fieldEncryptor.encrypt(exchange.refreshToken, appleRefreshContext(account.id)),
            now,
        )
        accountProvisioner.provision(
            account.id,
            verifiedIntegrity.deviceCheckTokenForTrial.takeUnless { account.antiAbuseRestricted },
            displayName,
        )
        return createSession(account.id, now)
    }

    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(
                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 -> 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,
                    sessionId = result.sessionId,
                    refreshToken = replayedRefreshToken,
                    refreshTokenExpiresAt = result.refreshTokenExpiresAt,
                )
            }
        }
    }

    suspend fun logout(principal: AccountPrincipal) {
        repository.revokeSessionFamily(principal.userId, principal.sessionId, clock.instant())
    }

    private suspend fun createSession(accountId: UUID, now: Instant): SessionTokens {
        val refreshToken = tokenGenerator.newRefreshToken()
        val refreshExpiresAt = now.plus(Duration.ofDays(sessionConfig.refreshDays))
        val created = repository.createSession(
            accountId = accountId,
            refreshTokenHash = TokenHash.sha256(refreshToken),
            expiresAt = refreshExpiresAt,
            now = now,
        )
        val access = sessionJwt.issue(accountId, created.sessionId)
        return SessionTokens(
            accountId = accountId,
            accessToken = access.value,
            accessTokenExpiresAt = access.expiresAt,
            refreshToken = refreshToken,
            refreshTokenExpiresAt = refreshExpiresAt,
        )
    }

    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)
        } catch (exception: AppleVerificationUnavailableException) {
            throw ExternalServiceUnavailableException("Apple identity verification")
        } catch (exception: AppleTokenInvalidException) {
            throw UnauthorizedException("Apple identity token is invalid")
        }

    private suspend fun exchangeCode(code: String): AppleTokenExchange =
        try {
            appleTokenClient.exchangeAuthorizationCode(code)
        } catch (exception: AppleClientUnavailableException) {
            throw ExternalServiceUnavailableException("Apple token service")
        } catch (exception: AppleTokenEndpointException) {
            if (exception.retryable) {
                throw ExternalServiceUnavailableException("Apple token service")
            }
            throw UnauthorizedException("Apple authorization code is invalid")
        }

    private fun requireValue(value: String, name: String, maxLength: Int) {
        if (value.isBlank()) throw InvalidRequestException("$name must not be blank")
        if (value.length > maxLength) {
            throw InvalidRequestException("$name exceeds the maximum length")
        }
    }

    private companion object {
        const val MAX_IDENTITY_TOKEN_LENGTH = 16_384
        const val MAX_AUTHORIZATION_CODE_LENGTH = 2_048
        const val MAX_NONCE_LENGTH = 256
        const val MAX_REFRESH_TOKEN_LENGTH = 512
    }
}

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"