Refactor AuthViewModel to match iOS auth flow

Replace PendingAuth/RSA key passing with tempPassword pattern:
- register(): saves keys to disk, stores password temporarily
- confirmRegistration(): verifies code, then calls performLogin(email, tempPassword)
  which re-loads keys from disk - exactly as iOS does
- Extract shared performLogin() suspend fn used by both login() and post-confirm flow
- Remove PendingAuth data class entirely

Simpler, correct, matches reference iOS implementation.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

app/src/main/java/com/kecalek/chat/core/SessionManager.kt

@@ -64,9 +64,9 @@ class SessionManager @Inject constructor(
         // Re-authenticate automatically whenever the connection is (re)established.
         // During the initial login() call, lastEmail is null (cleared before connect),
         // so this handler is a no-op for the first connection.
-        connection.onConnected = {
-            val email = lastEmail ?: return@onConnected
-            val key   = lastRsaPrivateKey ?: return@onConnected
+        connection.onConnected = reconnect@{
+            val email = lastEmail ?: return@reconnect
+            val key   = lastRsaPrivateKey ?: return@reconnect
             scope.launch {
                 try {
                     val session = performAuthHandshake(email, key, lastDeviceId, "Android")
@@ -100,9 +100,13 @@ class SessionManager @Inject constructor(
         lastRsaPrivateKey = null
 
         try {
-            if (connection.state.value != ConnectionManager.State.CONNECTED) {
-                connection.connect(host, port, useTls)
+            // Always start with a fresh connection.
+            // This handles stale post-registration connections and ensures reconnect is armed.
+            if (connection.state.value == ConnectionManager.State.CONNECTED) {
+                connection.disconnect()
             }
+            connection.enableReconnect()
+            connection.connect(host, port, useTls)
 
             val session = performAuthHandshake(email, rsaPrivateKey, deviceId, deviceName)
 

app/src/main/java/com/kecalek/chat/ui/auth/AuthViewModel.kt

@@ -1,6 +1,7 @@
 package com.kecalek.chat.ui.auth
 
 import android.util.Base64
+import android.util.Log
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.viewModelScope
 import com.kecalek.chat.core.AuthException
@@ -17,16 +18,12 @@ import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.flow.update
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.withContext
-import java.security.interfaces.RSAPrivateKey
 import java.security.interfaces.RSAPublicKey
 import javax.inject.Inject
 
-/**
- * UI state for all auth screens (Login, Register, Pairing).
- */
 data class AuthUiState(
     val isLoading: Boolean = false,
-    val loadingMessage: String? = null,   // e.g. "Generating keys…", "Connecting…"
+    val loadingMessage: String? = null,
     val error: String? = null,
     val isLoggedIn: Boolean = false,
     val isRegistered: Boolean = false,
@@ -42,19 +39,6 @@ data class AuthUiState(
     val registeredEmail: String? = null,
 )
 
-/**
- * Holds the already-decrypted key material between register() and confirmRegistration()
- * so we can auto-login immediately after email confirmation without re-asking for the
- * password or repeating the expensive PBKDF2 derivation.
- *
- * Cleared as soon as it is consumed (or when the ViewModel is cleared).
- */
-private data class PendingAuth(
-    val email: String,
-    val rsaPrivate: RSAPrivateKey,
-    val identityPrivateBytes: ByteArray,  // raw Ed25519 seed — used for initLocalKey
-)
-
 @HiltViewModel
 class AuthViewModel @Inject constructor(
     private val sessionManager: SessionManager,
@@ -64,38 +48,44 @@ class AuthViewModel @Inject constructor(
     private val _uiState = MutableStateFlow(AuthUiState())
     val uiState: StateFlow<AuthUiState> = _uiState.asStateFlow()
 
-    // Holds key material between register() and confirmRegistration(). Never leaves memory.
-    private var pendingAuth: PendingAuth? = null
+    // Stored temporarily between register() and confirmRegistration() so we can
+    // auto-login after confirmation without asking the user for their password again.
+    // Cleared immediately after use (or in onCleared).
+    // Matches the iOS pattern: AuthViewModel stores self.password for this purpose.
+    private var tempPassword: String? = null
 
     init {
         _uiState.update { it.copy(hasExistingAccount = keyStorage.hasRsaKeys()) }
     }
 
+    // ───────────────────────────── LOGIN ─────────────────────────────
+
     fun login(emailOrUsername: String, password: String) {
         viewModelScope.launch {
+            performLogin(emailOrUsername, password)
+        }
+    }
+
+    /**
+     * Core login logic shared by login() and the auto-login after confirmRegistration().
+     * Matches iOS ChatClient.login() + AuthViewModel.login(appState:).
+     */
+    private suspend fun performLogin(emailOrUsername: String, password: String) {
         _uiState.update { it.copy(isLoading = true, loadingMessage = "Dešifruji klíče…", error = null) }
         try {
             if (!keyStorage.hasRsaKeys()) {
                 _uiState.update {
-                        it.copy(
-                            isLoading = false,
-                            error = "No account on this device. Register or pair first."
-                        )
+                    it.copy(isLoading = false, error = "No account on this device. Register or pair first.")
                 }
-                    return@launch
+                return
             }
 
-                // Load RSA private key (decrypted with user's password via ECP1).
-                // PBKDF2 600k iterations — must run off the main thread.
+            // Load RSA private key — PBKDF2 600k iterations, must run off the main thread
             val rsaPrivate = try {
-                    withContext(Dispatchers.Default) {
-                        keyStorage.loadRsaPrivate(password)
-                    }
+                withContext(Dispatchers.Default) { keyStorage.loadRsaPrivate(password) }
             } catch (e: Exception) {
-                    _uiState.update {
-                        it.copy(isLoading = false, error = "Wrong password or corrupted key.")
-                    }
-                    return@launch
+                _uiState.update { it.copy(isLoading = false, error = "Wrong password or corrupted key.") }
+                return
             }
 
             _uiState.update { it.copy(loadingMessage = "Přihlašuji se…") }
@@ -109,7 +99,7 @@ class AuthViewModel @Inject constructor(
                 useTls = state.useTls,
             )
 
-                // Load identity key and init local storage key (also PBKDF2 — off main thread)
+            // Init local DB encryption key from identity key (also PBKDF2, off main thread)
             if (keyStorage.hasIdentityKeys()) {
                 val identityPrivate = withContext(Dispatchers.Default) {
                     keyStorage.loadIdentityPrivate(password)
@@ -118,6 +108,7 @@ class AuthViewModel @Inject constructor(
             }
 
             _uiState.update { it.copy(isLoading = false, loadingMessage = null, isLoggedIn = true) }
+
         } catch (e: AuthException) {
             _uiState.update { it.copy(isLoading = false, loadingMessage = null, error = e.message) }
         } catch (e: Exception) {
@@ -126,60 +117,38 @@ class AuthViewModel @Inject constructor(
             }
         }
     }
-    }
+
+    // ───────────────────────────── REGISTER ─────────────────────────────
 
     fun register(username: String, email: String, password: String) {
         viewModelScope.launch {
             _uiState.update { it.copy(isLoading = true, loadingMessage = "Generuji klíče…", error = null) }
             try {
-                // Steps 1-4 are CPU-intensive (RSA-4096 keygen + 2× PBKDF2 600k iters).
-                // Run on Default dispatcher to avoid blocking the UI thread.
-                data class KeyMaterial(
-                    val rsaPublicPem: String,
-                    val identityKeyBase64: String,
-                    val rsaPrivate: RSAPrivateKey,
-                    val identityPrivateBytes: ByteArray,
-                )
-                val keys = withContext(Dispatchers.Default) {
-                    // 1. Generate RSA-4096 keypair (~5-30 seconds)
+                // CPU-intensive: RSA-4096 keygen + 2× PBKDF2 600k iterations
+                val (rsaPublicPem, identityKeyBase64) = withContext(Dispatchers.Default) {
                     val (rsaPrivate, rsaPublic) = RSACrypto.generateKeypair()
-
-                    // 2. Generate Ed25519 identity keypair (fast)
                     val (identityPrivate, identityPublic) = Ed25519Crypto.generateKeypair()
 
-                    // 3. Save keys encrypted with password (PBKDF2 600k iters each)
                     keyStorage.saveRsaKeys(rsaPrivate, rsaPublic, password)
                     keyStorage.saveIdentityKeys(identityPrivate, identityPublic, password)
 
-                    // 4. Convert to server format and capture private key material
-                    KeyMaterial(
-                        rsaPublicPem = rsaPublicKeyToPem(rsaPublic),
-                        identityKeyBase64 = Base64.encodeToString(
-                            Ed25519Crypto.serializePublic(identityPublic),
-                            Base64.NO_WRAP,
-                        ),
-                        rsaPrivate = rsaPrivate,
-                        identityPrivateBytes = Ed25519Crypto.serializePrivate(identityPrivate),
+                    Pair(
+                        rsaPublicKeyToPem(rsaPublic),
+                        Base64.encodeToString(Ed25519Crypto.serializePublic(identityPublic), Base64.NO_WRAP),
                     )
                 }
 
-                // Save decrypted keys for use in confirmRegistration (auto-login).
-                // pendingAuth is cleared after use or when this ViewModel is destroyed.
-                pendingAuth = PendingAuth(
-                    email = email,
-                    rsaPrivate = keys.rsaPrivate,
-                    identityPrivateBytes = keys.identityPrivateBytes,
-                )
+                // Save password for auto-login after email confirmation (same pattern as iOS)
+                tempPassword = password
 
                 _uiState.update { it.copy(loadingMessage = "Připojuji se k serveru…") }
 
-                // 5. Register on server (network I/O — SessionManager uses IO dispatcher internally)
                 val state = _uiState.value
                 sessionManager.register(
                     username = username,
                     email = email,
-                    rsaPublicKeyPem = keys.rsaPublicPem,
-                    identityKeyBase64 = keys.identityKeyBase64,
+                    rsaPublicKeyPem = rsaPublicPem,
+                    identityKeyBase64 = identityKeyBase64,
                     host = state.serverHost,
                     port = state.serverPort,
                     useTls = state.useTls,
@@ -196,10 +165,10 @@ class AuthViewModel @Inject constructor(
                     )
                 }
             } catch (e: AuthException) {
-                pendingAuth = null
+                tempPassword = null
                 _uiState.update { it.copy(isLoading = false, loadingMessage = null, error = e.message) }
             } catch (e: Exception) {
-                pendingAuth = null
+                tempPassword = null
                 _uiState.update {
                     it.copy(isLoading = false, loadingMessage = null, error = "Registration failed: ${e.message}")
                 }
@@ -207,38 +176,36 @@ class AuthViewModel @Inject constructor(
         }
     }
 
+    // ───────────────────────────── CONFIRM ─────────────────────────────
+
     fun confirmRegistration(email: String, code: String) {
         viewModelScope.launch {
             _uiState.update { it.copy(isLoading = true, loadingMessage = "Potvrzuji kód…", error = null) }
             try {
+                // Step 1: Verify the email code — errors here are genuinely "wrong code"
                 sessionManager.confirmRegistration(email, code)
 
-                // Auto-login immediately after confirmation using the already-decrypted
-                // key material from register(). This avoids re-asking for the password.
-                val auth = pendingAuth
-                if (auth != null && auth.email == email) {
-                    _uiState.update { it.copy(loadingMessage = "Přihlašuji se…") }
-                    val state = _uiState.value
-                    sessionManager.login(
-                        email = email,
-                        rsaPrivateKey = auth.rsaPrivate,
-                        host = state.serverHost,
-                        port = state.serverPort,
-                        useTls = state.useTls,
-                    )
-                    // Init local DB encryption key (no password needed — bytes already decrypted)
-                    keyStorage.initLocalKey(auth.identityPrivateBytes)
-                    pendingAuth = null   // consumed — clear for security
-                }
+                // Step 2: Auto-login exactly as iOS does:
+                // call login() with the stored password (keys are already on disk from register())
+                val pwd = tempPassword
+                tempPassword = null
 
+                if (pwd != null) {
+                    // performLogin updates uiState (isLoggedIn = true on success, error on failure)
+                    performLogin(email, pwd)
+                } else {
+                    // No stored password (e.g. app restarted between register and confirm).
+                    // Mark registration done and send to LoginScreen.
                     _uiState.update {
                         it.copy(
                             isLoading = false,
                             loadingMessage = null,
-                        isLoggedIn = true,
                             needsConfirmation = false,
+                            hasExistingAccount = true,
                         )
                     }
+                }
+
             } catch (e: AuthException) {
                 _uiState.update { it.copy(isLoading = false, loadingMessage = null, error = e.message) }
             } catch (e: Exception) {
@@ -249,32 +216,28 @@ class AuthViewModel @Inject constructor(
         }
     }
 
+    // ───────────────────────────── PAIRING / BIOMETRIC ─────────────────────────────
+
     fun startPairing() {
         viewModelScope.launch {
-            _uiState.update {
-                it.copy(isLoading = false, error = "Device pairing not yet implemented.")
-            }
+            _uiState.update { it.copy(isLoading = false, error = "Device pairing not yet implemented.") }
         }
     }
 
     fun cancelPairing() {
-        _uiState.update {
-            it.copy(isPairingWaiting = false, pairingCode = null, isPairingComplete = false)
-        }
+        _uiState.update { it.copy(isPairingWaiting = false, pairingCode = null, isPairingComplete = false) }
     }
 
     fun loginWithBiometric() {
         viewModelScope.launch {
-            _uiState.update {
-                it.copy(isLoading = false, error = "Biometric login not yet implemented.")
-            }
+            _uiState.update { it.copy(isLoading = false, error = "Biometric login not yet implemented.") }
         }
     }
 
+    // ───────────────────────────── MISC ─────────────────────────────
+
     fun updateServerConfig(host: String, port: Int, useTls: Boolean) {
-        _uiState.update {
-            it.copy(serverHost = host, serverPort = port, useTls = useTls)
-        }
+        _uiState.update { it.copy(serverHost = host, serverPort = port, useTls = useTls) }
     }
 
     fun clearError() {
@@ -282,16 +245,18 @@ class AuthViewModel @Inject constructor(
     }
 
     fun resetState() {
-        pendingAuth = null
+        tempPassword = null
         _uiState.value = AuthUiState(hasExistingAccount = keyStorage.hasRsaKeys())
     }
 
     override fun onCleared() {
         super.onCleared()
-        pendingAuth = null  // Ensure key material doesn't linger after ViewModel is destroyed
+        tempPassword = null  // don't let the password linger after ViewModel is destroyed
     }
 
     companion object {
+        private const val TAG = "AuthViewModel"
+
         fun rsaPublicKeyToPem(key: RSAPublicKey): String {
             val der = key.encoded
             val base64 = Base64.encodeToString(der, Base64.NO_WRAP)