Merge pull request #917 from synonymdev/codex-pubky-ring-callbacks-913

Author: ovitrif

app/src/main/java/to/bitkit/models/PubkyRingAuthCallback.kt

@@ -0,0 +1,70 @@
+package to.bitkit.models
+
+import android.net.Uri
+
+private const val NONCE_PARAM = "nonce"
+
+sealed interface PubkyRingAuthCallback {
+    companion object {
+        private const val BITKIT_SCHEME = "bitkit"
+        private const val PUBKY_AUTH_HOST = "pubky-auth"
+        private const val SUCCESS_PATH = "/success"
+        private const val CANCEL_PATH = "/cancel"
+        private const val ERROR_PATH = "/error"
+        private const val ERROR_MESSAGE_PARAM = "errorMessage"
+
+        fun parse(uri: Uri): PubkyRingAuthCallback? {
+            if (uri.scheme != BITKIT_SCHEME || uri.host != PUBKY_AUTH_HOST) return null
+
+            val nonce = uri.getQueryParameter(NONCE_PARAM)?.takeIf { it.isNotBlank() }
+            return when (uri.path) {
+                SUCCESS_PATH -> Success(nonce)
+                CANCEL_PATH -> Cancel(nonce)
+                ERROR_PATH -> Error(uri.getQueryParameter(ERROR_MESSAGE_PARAM), nonce)
+                else -> null
+            }
+        }
+    }
+
+    val nonce: String?
+
+    data class Success(override val nonce: String?) : PubkyRingAuthCallback
+    data class Cancel(override val nonce: String?) : PubkyRingAuthCallback
+    data class Error(val message: String?, override val nonce: String?) : PubkyRingAuthCallback
+}
+
+sealed interface PubkyRingAuthCallbackHandlingResult {
+    data object Ignored : PubkyRingAuthCallbackHandlingResult
+    data object Handled : PubkyRingAuthCallbackHandlingResult
+    data class TrustedError(val message: String?) : PubkyRingAuthCallbackHandlingResult
+}
+
+object PubkyRingAuthUrlBuilder {
+    const val SUCCESS_CALLBACK = "bitkit://pubky-auth/success"
+    const val CANCEL_CALLBACK = "bitkit://pubky-auth/cancel"
+    const val ERROR_CALLBACK = "bitkit://pubky-auth/error"
+    const val SOURCE = "Bitkit"
+
+    fun addCallbacks(authUrl: String, nonce: String? = null): String? {
+        val uri = Uri.parse(authUrl)
+        if (uri.scheme.isNullOrBlank()) return null
+
+        return uri.buildUpon()
+            .appendQueryParameter("x-success", callbackUrl(SUCCESS_CALLBACK, nonce))
+            .appendQueryParameter("x-cancel", callbackUrl(CANCEL_CALLBACK, nonce))
+            .appendQueryParameter("x-error", callbackUrl(ERROR_CALLBACK, nonce))
+            .appendQueryParameter("x-source", SOURCE)
+            .build()
+            .toString()
+    }
+
+    private fun callbackUrl(baseUrl: String, nonce: String?): String {
+        if (nonce.isNullOrBlank()) return baseUrl
+
+        return Uri.parse(baseUrl)
+            .buildUpon()
+            .appendQueryParameter(NONCE_PARAM, nonce)
+            .build()
+            .toString()
+    }
+}

app/src/main/java/to/bitkit/repositories/PubkyRepo.kt

@@ -13,9 +13,11 @@ import kotlinx.coroutines.SupervisorJob
 import kotlinx.coroutines.async
 import kotlinx.coroutines.awaitAll
 import kotlinx.coroutines.coroutineScope
+import kotlinx.coroutines.flow.MutableSharedFlow
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.asSharedFlow
 import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.map
@@ -35,24 +37,34 @@ import to.bitkit.models.PubkyProfile
 import to.bitkit.models.PubkyProfileData
 import to.bitkit.models.PubkyProfileLink
 import to.bitkit.models.PubkyPublicKeyFormat
+import to.bitkit.models.PubkyRingAuthCallback
+import to.bitkit.models.PubkyRingAuthCallbackHandlingResult
 import to.bitkit.models.PubkySessionBackupKind
 import to.bitkit.models.PubkySessionBackupV1
 import to.bitkit.services.PubkyService
 import to.bitkit.utils.AppError
 import to.bitkit.utils.Logger
 import java.io.ByteArrayOutputStream
+import java.util.UUID
 import javax.inject.Inject
 import javax.inject.Singleton
 import kotlin.math.min
 
 enum class PubkyAuthState { Idle, Authenticating, Authenticated }
 
+data class PubkyRingAuthRequest(
+    val authUrl: String,
+    val callbackNonce: String,
+)
+
 sealed class PubkyContactError(message: String) : AppError(message) {
     data object AlreadyExists : PubkyContactError("Contact already exists")
     data object CannotAddSelf : PubkyContactError("Cannot add your own pubky as a contact")
     data object InvalidFormat : PubkyContactError("Invalid pubky key format")
 }
 
+private class PubkyAuthAttemptInactive : AppError("Auth attempt is no longer active")
+
 @Suppress("TooManyFunctions", "LargeClass", "LongParameterList")
 @Singleton
 class PubkyRepo @Inject constructor(
@@ -80,6 +92,9 @@ class PubkyRepo @Inject constructor(
     private var isServiceInitialized = false
 
     private val _authState = MutableStateFlow(PubkyAuthState.Idle)
+    private val _activeAuthAttemptId = MutableStateFlow<String?>(null)
+    private val _authCancelEvents = MutableSharedFlow<Unit>(extraBufferCapacity = 1)
+    val authCancelEvents = _authCancelEvents.asSharedFlow()
 
     private val _profile = MutableStateFlow<PubkyProfile?>(null)
     val profile: StateFlow<PubkyProfile?> = _profile.asStateFlow()
@@ -108,7 +123,7 @@ class PubkyRepo @Inject constructor(
     private val _backupStateVersion = MutableStateFlow(0L)
     val backupStateVersion: StateFlow<Long> = _backupStateVersion.asStateFlow()
 
-    val isAuthenticated: StateFlow<Boolean> = _authState.map { it == PubkyAuthState.Authenticated }
+    val isAuthenticated: StateFlow<Boolean> = _publicKey.map { it != null }
         .stateIn(scope, SharingStarted.Eagerly, false)
 
     val displayName: StateFlow<String?> = combine(_profile, pubkyStore.data) { profile, cached ->
@@ -236,20 +251,27 @@ class PubkyRepo @Inject constructor(
 
     // region Ring auth flow
 
-    suspend fun startAuthentication(): Result<String> {
+    suspend fun startAuthentication(): Result<PubkyRingAuthRequest> {
+        val attemptId = UUID.randomUUID().toString()
+        _activeAuthAttemptId.update { attemptId }
         _authState.update { PubkyAuthState.Authenticating }
         return runCatching {
-            withContext(ioDispatcher) { pubkyService.startAuth() }
+            val authUrl = withContext(ioDispatcher) { pubkyService.startAuth() }
+            PubkyRingAuthRequest(authUrl = authUrl, callbackNonce = attemptId)
         }.onFailure {
-            _authState.update { PubkyAuthState.Idle }
+            _activeAuthAttemptId.update { null }
+            restoreAuthStateAfterAuthFlow()
         }
     }
 
     suspend fun completeAuthentication(): Result<Unit> {
+        val attemptId = _activeAuthAttemptId.value ?: return Result.failure(PubkyAuthAttemptInactive())
         return runCatching {
             withContext(ioDispatcher) {
                 val sessionSecret = pubkyService.completeAuth()
+                ensureAuthAttemptActive(attemptId)
                 val pk = pubkyService.importSession(sessionSecret).ensurePubkyPrefix()
+                ensureAuthAttemptActive(attemptId)
 
                 runCatching { keychain.delete(Keychain.Key.PUBKY_SECRET_KEY.name) }
                 keychain.upsertString(Keychain.Key.PAYKIT_SESSION.name, sessionSecret)
@@ -258,8 +280,14 @@ class PubkyRepo @Inject constructor(
                 pk
             }
         }.onFailure {
-            _authState.update { PubkyAuthState.Idle }
+            if (_activeAuthAttemptId.value == attemptId) {
+                _activeAuthAttemptId.update { null }
+            }
+            restoreAuthStateAfterAuthFlow()
         }.onSuccess { pk ->
+            if (_activeAuthAttemptId.value == attemptId) {
+                _activeAuthAttemptId.update { null }
+            }
             _publicKey.update { pk }
             _authState.update { PubkyAuthState.Authenticated }
             Logger.info("Completed pubky auth for '$pk'", context = TAG)
@@ -272,13 +300,82 @@ class PubkyRepo @Inject constructor(
         runCatching {
             withContext(ioDispatcher) { pubkyService.cancelAuth() }
         }.onFailure { Logger.warn("Failed to cancel auth", it, context = TAG) }
-        _authState.update { PubkyAuthState.Idle }
+        endAuthAttempt()
     }
 
     fun cancelAuthenticationSync() {
         scope.launch { cancelAuthentication() }
     }
 
+    suspend fun handleAuthCallback(callback: PubkyRingAuthCallback): PubkyRingAuthCallbackHandlingResult {
+        if (!isCurrentAuthCallback(callback)) {
+            return handleInvalidAuthCallback(callback)
+        }
+
+        return when (callback) {
+            is PubkyRingAuthCallback.Success -> {
+                Logger.info("Received Pubky Ring auth success callback", context = TAG)
+                PubkyRingAuthCallbackHandlingResult.Handled
+            }
+            is PubkyRingAuthCallback.Cancel -> {
+                Logger.info("Received Pubky Ring auth cancel callback", context = TAG)
+                cancelAuthentication()
+                PubkyRingAuthCallbackHandlingResult.Handled
+            }
+            is PubkyRingAuthCallback.Error -> {
+                Logger.warn("Received Pubky Ring auth error callback", context = TAG)
+                cancelAuthentication()
+                PubkyRingAuthCallbackHandlingResult.TrustedError(callback.message)
+            }
+        }
+    }
+
+    private fun handleInvalidAuthCallback(
+        callback: PubkyRingAuthCallback,
+    ): PubkyRingAuthCallbackHandlingResult {
+        if (_activeAuthAttemptId.value == null) {
+            Logger.warn("Ignoring Pubky Ring auth callback with missing or invalid nonce", context = TAG)
+            return PubkyRingAuthCallbackHandlingResult.Ignored
+        }
+
+        return when (callback) {
+            is PubkyRingAuthCallback.Success -> {
+                Logger.warn("Ignoring Pubky Ring auth success callback with missing or invalid nonce", context = TAG)
+                PubkyRingAuthCallbackHandlingResult.Ignored
+            }
+            is PubkyRingAuthCallback.Cancel -> {
+                Logger.warn("Ignoring Pubky Ring auth cancel callback with missing or invalid nonce", context = TAG)
+                PubkyRingAuthCallbackHandlingResult.Ignored
+            }
+            is PubkyRingAuthCallback.Error -> {
+                Logger.warn("Ignoring Pubky Ring auth error callback with missing or invalid nonce", context = TAG)
+                PubkyRingAuthCallbackHandlingResult.Ignored
+            }
+        }
+    }
+
+    private fun isCurrentAuthCallback(callback: PubkyRingAuthCallback): Boolean {
+        val activeAuthAttemptId = _activeAuthAttemptId.value ?: return false
+        return callback.nonce == activeAuthAttemptId
+    }
+
+    private fun ensureAuthAttemptActive(attemptId: String?) {
+        if (attemptId == null) return
+        if (_activeAuthAttemptId.value == attemptId) return
+
+        throw PubkyAuthAttemptInactive()
+    }
+
+    private fun endAuthAttempt() {
+        _activeAuthAttemptId.update { null }
+        _authCancelEvents.tryEmit(Unit)
+        restoreAuthStateAfterAuthFlow()
+    }
+
+    private fun restoreAuthStateAfterAuthFlow() {
+        _authState.update { if (_publicKey.value == null) PubkyAuthState.Idle else PubkyAuthState.Authenticated }
+    }
+
     // endregion
 
     // region Payment endpoints

app/src/main/java/to/bitkit/ui/screens/profile/PubkyChoiceViewModel.kt

@@ -17,6 +17,7 @@ import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.flow.update
 import kotlinx.coroutines.launch
 import to.bitkit.R
+import to.bitkit.models.PubkyRingAuthUrlBuilder
 import to.bitkit.models.Toast
 import to.bitkit.repositories.PubkyRepo
 import to.bitkit.ui.shared.toast.ToastEventBus
@@ -41,6 +42,16 @@ class PubkyChoiceViewModel @Inject constructor(
 
     private var approvalJob: Job? = null
 
+    init {
+        viewModelScope.launch {
+            pubkyRepo.authCancelEvents.collect {
+                approvalJob?.cancel()
+                approvalJob = null
+                _uiState.update { it.copy(isWaitingForRing = false, isLoadingAfterAuth = false) }
+            }
+        }
+    }
+
     override fun onCleared() {
         super.onCleared()
         if (_uiState.value.isWaitingForRing) {
@@ -63,8 +74,12 @@ class PubkyChoiceViewModel @Inject constructor(
             }
 
             pubkyRepo.startAuthentication()
-                .onSuccess { authUrl ->
-                    val ringIntent = createRingAuthIntent(authUrl)
+                .onSuccess { authRequest ->
+                    val callbackAuthUrl = PubkyRingAuthUrlBuilder.addCallbacks(
+                        authUrl = authRequest.authUrl,
+                        nonce = authRequest.callbackNonce,
+                    ) ?: authRequest.authUrl
+                    val ringIntent = createRingAuthIntent(callbackAuthUrl)
                     if (!canOpenWithRing(ringIntent)) {
                         cancelAuthAndShowRingDialog()
                         return@launch

app/src/main/java/to/bitkit/viewmodels/AppViewModel.kt

@@ -101,6 +101,8 @@ import to.bitkit.models.NewTransactionSheetType
 import to.bitkit.models.NodeLifecycleState
 import to.bitkit.models.PubkyProfile
 import to.bitkit.models.PubkyPublicKeyFormat
+import to.bitkit.models.PubkyRingAuthCallback
+import to.bitkit.models.PubkyRingAuthCallbackHandlingResult
 import to.bitkit.models.Suggestion
 import to.bitkit.models.Toast
 import to.bitkit.models.TransactionSpeed
@@ -2677,6 +2679,11 @@ class AppViewModel @Inject constructor(
             return@launch
         }
 
+        PubkyRingAuthCallback.parse(uri)?.let {
+            handlePubkyRingAuthCallback(it)
+            return@launch
+        }
+
         if (uri.scheme == PUBKYAUTH_SCHEME) {
             handlePubkyAuth(uri.toString())
             return@launch
@@ -2698,6 +2705,21 @@ class AppViewModel @Inject constructor(
         showSheet(Sheet.PubkyAuth(authUrl))
     }
 
+    private suspend fun handlePubkyRingAuthCallback(callback: PubkyRingAuthCallback) {
+        when (val result = pubkyRepo.handleAuthCallback(callback)) {
+            is PubkyRingAuthCallbackHandlingResult.TrustedError -> {
+                ToastEventBus.send(
+                    type = Toast.ToastType.ERROR,
+                    title = context.getString(R.string.profile__auth_error_title),
+                    description = result.message ?: context.getString(R.string.other__qr_error_text),
+                )
+            }
+            PubkyRingAuthCallbackHandlingResult.Handled,
+            PubkyRingAuthCallbackHandlingResult.Ignored,
+            -> Unit
+        }
+    }
+
     // TODO Temporary fix while these schemes can't be decoded https://github.com/synonymdev/bitkit-core/issues/70
     private fun String.removeLightningSchemes(): String = LIGHTNING_SCHEME_PATTERNS.fold(this) { acc, regex ->
         acc.replace(regex, "")