crypto: Add support for key backup restoring

matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/crypto/DefaultCryptoService.kt

@@ -58,7 +58,6 @@ import org.matrix.android.sdk.api.util.JsonDict
 import org.matrix.android.sdk.internal.auth.registration.handleUIA
 import org.matrix.android.sdk.internal.crypto.crosssigning.DeviceTrustLevel
 import org.matrix.android.sdk.internal.crypto.keysbackup.RustKeyBackupService
-import org.matrix.android.sdk.internal.crypto.keysbackup.model.MegolmBackupAuthData
 import org.matrix.android.sdk.internal.crypto.keysbackup.model.rest.BackupKeysResult
 import org.matrix.android.sdk.internal.crypto.keysbackup.model.rest.CreateKeysBackupVersionBody
 import org.matrix.android.sdk.internal.crypto.keysbackup.model.rest.KeysBackupData
@@ -70,6 +69,9 @@ import org.matrix.android.sdk.internal.crypto.keysbackup.tasks.CreateKeysBackupV
 import org.matrix.android.sdk.internal.crypto.keysbackup.tasks.DeleteBackupTask
 import org.matrix.android.sdk.internal.crypto.keysbackup.tasks.GetKeysBackupLastVersionTask
 import org.matrix.android.sdk.internal.crypto.keysbackup.tasks.GetKeysBackupVersionTask
+import org.matrix.android.sdk.internal.crypto.keysbackup.tasks.GetRoomSessionDataTask
+import org.matrix.android.sdk.internal.crypto.keysbackup.tasks.GetRoomSessionsDataTask
+import org.matrix.android.sdk.internal.crypto.keysbackup.tasks.GetSessionsDataTask
 import org.matrix.android.sdk.internal.crypto.keysbackup.tasks.StoreSessionsDataTask
 import org.matrix.android.sdk.internal.crypto.keysbackup.tasks.UpdateKeysBackupVersionTask
 import org.matrix.android.sdk.internal.crypto.model.CryptoDeviceInfo
@@ -142,6 +144,9 @@ internal class RequestSender @Inject constructor(
         private val createKeysBackupVersionTask: CreateKeysBackupVersionTask,
         private val backupRoomKeysTask: StoreSessionsDataTask,
         private val updateKeysBackupVersionTask: UpdateKeysBackupVersionTask,
+        private val getSessionsDataTask: GetSessionsDataTask,
+        private val getRoomSessionsDataTask: GetRoomSessionsDataTask,
+        private val getRoomSessionDataTask: GetRoomSessionDataTask,
         ) {
     companion object {
         const val REQUEST_RETRY_COUNT = 3
@@ -321,6 +326,26 @@ internal class RequestSender @Inject constructor(
         val params = UpdateKeysBackupVersionTask.Params(keysBackupVersion.version, body)
         updateKeysBackupVersionTask.executeRetry(params, REQUEST_RETRY_COUNT)
     }
+
+    suspend fun downloadBackedUpKeys(version: String, roomId: String, sessionId: String): KeysBackupData {
+        val data = getRoomSessionDataTask.execute(GetRoomSessionDataTask.Params(roomId, sessionId, version))
+
+        return KeysBackupData(mutableMapOf(
+                roomId to RoomKeysBackupData(mutableMapOf(
+                        sessionId to data
+                ))
+        ))
+    }
+
+    suspend fun downloadBackedUpKeys(version: String, roomId: String): KeysBackupData {
+        val data = getRoomSessionsDataTask.execute(GetRoomSessionsDataTask.Params(roomId, version))
+        // Convert to KeysBackupData
+        return KeysBackupData(mutableMapOf(roomId to data))
+    }
+
+    suspend fun downloadBackedUpKeys(version: String): KeysBackupData {
+        return getSessionsDataTask.execute(GetSessionsDataTask.Params(version))
+    }
 }
 
 /**

matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/crypto/OlmMachine.kt

@@ -18,10 +18,6 @@ package org.matrix.android.sdk.internal.crypto
 
 import androidx.lifecycle.LiveData
 import androidx.lifecycle.MutableLiveData
-import com.squareup.moshi.Types
-import java.io.File
-import java.nio.charset.Charset
-import java.util.concurrent.ConcurrentHashMap
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.withContext
@@ -39,7 +35,6 @@ import org.matrix.android.sdk.api.util.toOptional
 import org.matrix.android.sdk.internal.crypto.crosssigning.DeviceTrustLevel
 import org.matrix.android.sdk.internal.crypto.crosssigning.UserTrustResult
 import org.matrix.android.sdk.internal.crypto.keysbackup.model.MegolmBackupAuthData
-import org.matrix.android.sdk.internal.crypto.keysbackup.model.rest.RoomKeysBackupData
 import org.matrix.android.sdk.internal.crypto.model.CryptoDeviceInfo
 import org.matrix.android.sdk.internal.crypto.model.ImportRoomKeysResult
 import org.matrix.android.sdk.internal.crypto.model.MXUsersDevicesMap
@@ -51,7 +46,6 @@ import org.matrix.android.sdk.internal.network.parsing.CheckNumberType
 import org.matrix.android.sdk.internal.session.sync.model.DeviceListResponse
 import org.matrix.android.sdk.internal.session.sync.model.DeviceOneTimeKeysCountSyncResponse
 import org.matrix.android.sdk.internal.session.sync.model.ToDeviceSyncResponse
-import org.matrix.android.sdk.internal.util.JsonCanonicalizer
 import timber.log.Timber
 import uniffi.olm.BackupKey
 import uniffi.olm.BackupKeys
@@ -62,13 +56,16 @@ import uniffi.olm.DecryptionException
 import uniffi.olm.DeviceLists
 import uniffi.olm.KeyRequestPair
 import uniffi.olm.Logger
-import uniffi.olm.OlmMachine as InnerMachine
-import uniffi.olm.ProgressListener as RustProgressListener
 import uniffi.olm.Request
 import uniffi.olm.RequestType
 import uniffi.olm.RoomKeyCounts
-import uniffi.olm.UserIdentity as RustUserIdentity
 import uniffi.olm.setLogger
+import java.io.File
+import java.nio.charset.Charset
+import java.util.concurrent.ConcurrentHashMap
+import uniffi.olm.OlmMachine as InnerMachine
+import uniffi.olm.ProgressListener as RustProgressListener
+import uniffi.olm.UserIdentity as RustUserIdentity
 
 class CryptoLogger : Logger {
     override fun log(logLine: String) {
@@ -506,6 +503,22 @@ internal class OlmMachine(
                 ImportRoomKeysResult(result.total, result.imported)
             }
 
+    @Throws(CryptoStoreException::class)
+    suspend fun importDecryptedKeys(
+            keys: List<MegolmSessionData>,
+            listener: ProgressListener?
+    ): ImportRoomKeysResult =
+            withContext(Dispatchers.IO) {
+                val adapter = MoshiProvider.providesMoshi().adapter(List::class.java)
+                val encodedKeys = adapter.toJson(keys)
+
+                val rustListener = CryptoProgressListener(listener)
+
+                val result = inner.importDecryptedKeys(encodedKeys, rustListener)
+
+                ImportRoomKeysResult(result.total, result.imported)
+            }
+
     @Throws(CryptoStoreException::class)
     suspend fun getIdentity(userId: String): UserIdentities? {
         val identity = withContext(Dispatchers.IO) {

matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/crypto/keysbackup/RustKeyBackupService.kt

@@ -19,6 +19,8 @@ package org.matrix.android.sdk.internal.crypto.keysbackup
 import android.os.Handler
 import android.os.Looper
 import androidx.annotation.UiThread
+import androidx.annotation.VisibleForTesting
+import androidx.annotation.WorkerThread
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.delay
@@ -33,6 +35,7 @@ import org.matrix.android.sdk.api.session.crypto.keysbackup.KeysBackupService
 import org.matrix.android.sdk.api.session.crypto.keysbackup.KeysBackupState
 import org.matrix.android.sdk.api.session.crypto.keysbackup.KeysBackupStateListener
 import org.matrix.android.sdk.internal.crypto.MXCRYPTO_ALGORITHM_MEGOLM_BACKUP
+import org.matrix.android.sdk.internal.crypto.MegolmSessionData
 import org.matrix.android.sdk.internal.crypto.OlmMachine
 import org.matrix.android.sdk.internal.crypto.RequestSender
 import org.matrix.android.sdk.internal.crypto.keysbackup.model.KeysBackupVersionTrust
@@ -40,11 +43,14 @@ import org.matrix.android.sdk.internal.crypto.keysbackup.model.MegolmBackupAuthD
 import org.matrix.android.sdk.internal.crypto.keysbackup.model.MegolmBackupCreationInfo
 import org.matrix.android.sdk.internal.crypto.keysbackup.model.SignalableMegolmBackupAuthData
 import org.matrix.android.sdk.internal.crypto.keysbackup.model.rest.CreateKeysBackupVersionBody
+import org.matrix.android.sdk.internal.crypto.keysbackup.model.rest.KeyBackupData
+import org.matrix.android.sdk.internal.crypto.keysbackup.model.rest.KeysBackupData
 import org.matrix.android.sdk.internal.crypto.keysbackup.model.rest.KeysVersion
 import org.matrix.android.sdk.internal.crypto.keysbackup.model.rest.KeysVersionResult
 import org.matrix.android.sdk.internal.crypto.keysbackup.model.rest.UpdateKeysBackupVersionBody
 import org.matrix.android.sdk.internal.crypto.model.ImportRoomKeysResult
 import org.matrix.android.sdk.internal.crypto.store.SavedKeyBackupKeyInfo
+import org.matrix.android.sdk.internal.di.MoshiProvider
 import org.matrix.android.sdk.internal.extensions.foldToCallback
 import org.matrix.android.sdk.internal.session.SessionScope
 import org.matrix.android.sdk.internal.util.JsonCanonicalizer
@@ -54,6 +60,7 @@ import timber.log.Timber
 import uniffi.olm.BackupRecoveryKey
 import uniffi.olm.Request
 import uniffi.olm.RequestType
+import java.security.InvalidParameterException
 import javax.inject.Inject
 import kotlin.random.Random
 
@@ -330,9 +337,9 @@ internal class RustKeyBackupService @Inject constructor(
 
                     @Suppress("UNCHECKED_CAST")
                     UpdateKeysBackupVersionBody(
-                        algorithm = keysBackupVersion.algorithm,
+                            algorithm = keysBackupVersion.algorithm,
-                        authData = newAuthData.copy(signatures = newSignatures).toJsonDict(),
+                            authData = newAuthData.copy(signatures = newSignatures).toJsonDict(),
-                        version = keysBackupVersion.version)
+                            version = keysBackupVersion.version)
                 }
                 try {
                     sender.updateBackup(keysBackupVersion, body)
@@ -364,7 +371,6 @@ internal class RustKeyBackupService @Inject constructor(
             authData == null                          -> {
                 Timber.w("isValidRecoveryKeyForKeysBackupVersion: Key backup is missing required data")
                 throw IllegalArgumentException("Missing element")
-
             }
             backupKey.publicKey != authData.publicKey -> {
                 Timber.w("isValidRecoveryKeyForKeysBackupVersion: Public keys mismatch")
@@ -390,7 +396,6 @@ internal class RustKeyBackupService @Inject constructor(
                 val key = BackupRecoveryKey.fromBase58(recoveryKey)
                 checkRecoveryKey(key, keysBackupVersion)
                 trustKeysBackupVersion(keysBackupVersion, true, callback)
-
             } catch (exception: Throwable) {
                 callback.onFailure(exception)
             }
@@ -423,14 +428,141 @@ internal class RustKeyBackupService @Inject constructor(
         progressListener.onProgress(backedUpKeys, total)
     }
 
+    /**
+     * Same method as [RoomKeysRestClient.getRoomKey] except that it accepts nullable
+     * parameters and always returns a KeysBackupData object through the Callback
+     */
+    private suspend fun getKeys(sessionId: String?, roomId: String?, version: String): KeysBackupData {
+        return when {
+            roomId != null && sessionId != null -> {
+                sender.downloadBackedUpKeys(version, roomId, sessionId)
+            }
+            roomId != null                      -> {
+                sender.downloadBackedUpKeys(version, roomId)
+            }
+            else                                -> {
+                sender.downloadBackedUpKeys(version)
+            }
+        }
+    }
+
+    @VisibleForTesting
+    @WorkerThread
+    fun decryptKeyBackupData(keyBackupData: KeyBackupData, sessionId: String, roomId: String, key: BackupRecoveryKey): MegolmSessionData? {
+        var sessionBackupData: MegolmSessionData? = null
+
+        val jsonObject = keyBackupData.sessionData
+
+        val ciphertext = jsonObject["ciphertext"]?.toString()
+        val mac = jsonObject["mac"]?.toString()
+        val ephemeralKey = jsonObject["ephemeral"]?.toString()
+
+        if (ciphertext != null && mac != null && ephemeralKey != null) {
+            try {
+                val decrypted = key.decrypt(ephemeralKey, mac, ciphertext)
+
+                val moshi = MoshiProvider.providesMoshi()
+                val adapter = moshi.adapter(MegolmSessionData::class.java)
+
+                sessionBackupData = adapter.fromJson(decrypted)
+            } catch (e: Throwable) {
+                Timber.e(e, "OlmException")
+            }
+
+            if (sessionBackupData != null) {
+                sessionBackupData = sessionBackupData.copy(
+                        sessionId = sessionId,
+                        roomId = roomId
+                )
+            }
+        }
+
+        return sessionBackupData
+    }
+
+    private suspend fun restoreBackup(
+            keysVersionResult: KeysVersionResult,
+            recoveryKey: BackupRecoveryKey,
+            roomId: String?,
+            sessionId: String?,
+            stepProgressListener: StepProgressListener?,
+    ): ImportRoomKeysResult {
+        withContext(coroutineDispatchers.crypto) {
+            // Check if the recovery is valid before going any further
+            if (!isValidRecoveryKey(recoveryKey, keysVersionResult)) {
+                Timber.e("restoreKeysWithRecoveryKey: Invalid recovery key for this keys version")
+                throw InvalidParameterException("Invalid recovery key")
+            }
+        }
+
+        stepProgressListener?.onStepProgress(StepProgressListener.Step.DownloadingKey)
+
+        // Get backed up keys from the homeserver
+        val data = getKeys(sessionId, roomId, keysVersionResult.version)
+
+        return withContext(coroutineDispatchers.computation) {
+            val sessionsData = ArrayList<MegolmSessionData>()
+            // Restore that data
+            var sessionsFromHsCount = 0
+            for ((roomIdLoop, backupData) in data.roomIdToRoomKeysBackupData) {
+                for ((sessionIdLoop, keyBackupData) in backupData.sessionIdToKeyBackupData) {
+                    sessionsFromHsCount++
+
+                    val sessionData = decryptKeyBackupData(keyBackupData, sessionIdLoop, roomIdLoop, recoveryKey)
+
+                    sessionData?.let {
+                        sessionsData.add(it)
+                    }
+                }
+            }
+            Timber.v("restoreKeysWithRecoveryKey: Decrypted ${sessionsData.size} keys out" +
+                    " of $sessionsFromHsCount from the backup store on the homeserver")
+
+            // Do not trigger a backup for them if they come from the backup version we are using
+            val backUp = keysVersionResult.version != keysBackupVersion?.version
+            if (backUp) {
+                Timber.v("restoreKeysWithRecoveryKey: Those keys will be backed up" +
+                        " to backup version: ${keysBackupVersion?.version}")
+            }
+
+            // Import them into the crypto store
+            val progressListener = if (stepProgressListener != null) {
+                object : ProgressListener {
+                    override fun onProgress(progress: Int, total: Int) {
+                        // Note: no need to post to UI thread, importMegolmSessionsData() will do it
+                        stepProgressListener.onStepProgress(StepProgressListener.Step.ImportingKey(progress, total))
+                    }
+                }
+            } else {
+                null
+            }
+
+            val result = olmMachine.importDecryptedKeys(sessionsData, progressListener)
+
+            // Do not back up the key if it comes from a backup recovery
+            if (backUp) {
+                maybeBackupKeys()
+            }
+
+            // Save for next time and for gossiping
+            saveBackupRecoveryKey(recoveryKey.toBase64(), keysVersionResult.version)
+            result
+        }
+    }
+
     override fun restoreKeysWithRecoveryKey(keysVersionResult: KeysVersionResult,
                                             recoveryKey: String,
                                             roomId: String?,
                                             sessionId: String?,
                                             stepProgressListener: StepProgressListener?,
                                             callback: MatrixCallback<ImportRoomKeysResult>) {
-        // TODO
         Timber.v("restoreKeysWithRecoveryKey: From backup version: ${keysVersionResult.version}")
+        cryptoCoroutineScope.launch(coroutineDispatchers.main) {
+            runCatching {
+                val key = BackupRecoveryKey.fromBase58(recoveryKey)
+                restoreBackup(keysVersionResult, key, roomId, sessionId, stepProgressListener)
+            }.foldToCallback(callback)
+        }
     }
 
     override fun restoreKeyBackupWithPassword(keysBackupVersion: KeysVersionResult,
@@ -439,8 +571,16 @@ internal class RustKeyBackupService @Inject constructor(
                                               sessionId: String?,
                                               stepProgressListener: StepProgressListener?,
                                               callback: MatrixCallback<ImportRoomKeysResult>) {
-        // TODO
         Timber.v("[MXKeyBackup] restoreKeyBackup with password: From backup version: ${keysBackupVersion.version}")
+        cryptoCoroutineScope.launch(coroutineDispatchers.main) {
+            runCatching {
+                val recoveryKey = withContext(coroutineDispatchers.crypto) {
+                    BackupRecoveryKey.fromPassphrase(password)
+                }
+
+                restoreBackup(keysBackupVersion, recoveryKey, roomId, sessionId, stepProgressListener)
+            }.foldToCallback(callback)
+        }
     }
 
     override fun getVersion(version: String, callback: MatrixCallback<KeysVersionResult?>) {
@@ -573,15 +713,18 @@ internal class RustKeyBackupService @Inject constructor(
         }
     }
 
+    private fun isValidRecoveryKey(recoveryKey: BackupRecoveryKey, version: KeysVersionResult): Boolean {
+        val publicKey = recoveryKey.publicKey().publicKey
+        val authData = getMegolmBackupAuthData(version) ?: return false
+        return authData.publicKey == publicKey
+    }
+
     override fun isValidRecoveryKeyForCurrentVersion(recoveryKey: String, callback: MatrixCallback<Boolean>) {
         val keysBackupVersion = keysBackupVersion ?: return Unit.also { callback.onSuccess(false) }
 
         try {
             val key = BackupRecoveryKey.fromBase64(recoveryKey)
-            val publicKey = key.publicKey().publicKey
+            callback.onSuccess(isValidRecoveryKey(key, keysBackupVersion))
-            val authData = getMegolmBackupAuthData(keysBackupVersion) ?: return Unit.also { callback.onSuccess(false) }
-
-            callback.onSuccess(authData.publicKey == publicKey)
         } catch (error: Throwable) {
             callback.onFailure(error)
         }

rust-sdk/src/backup_recovery_key.rs

@@ -3,8 +3,9 @@ use pbkdf2::pbkdf2;
 use rand::{distributions::Alphanumeric, thread_rng, Rng};
 use sha2::Sha512;
 use std::{collections::HashMap, iter};
+use thiserror::Error;
 
-use matrix_sdk_crypto::backups::RecoveryKey;
+use matrix_sdk_crypto::backups::{RecoveryKey, OlmPkDecryptionError};
 
 /// TODO
 pub struct BackupRecoveryKey {
@@ -12,6 +13,14 @@ pub struct BackupRecoveryKey {
     passphrase_info: Option<PassphraseInfo>,
 }
 
+/// TODO
+#[derive(Debug, Error)]
+pub enum PkDecryptionError {
+    /// TODO
+    #[error("Error decryption a PkMessage {0}")]
+    Olm(#[from] OlmPkDecryptionError),
+}
+
 /// TODO
 #[derive(Debug, Clone)]
 pub struct PassphraseInfo {
@@ -114,4 +123,20 @@ impl BackupRecoveryKey {
     pub fn to_base58(&self) -> String {
         self.inner.to_base58()
     }
+
+    /// TODO
+    pub fn to_base64(&self) -> String {
+        self.inner.to_base64()
+    }
+
+    pub fn decrypt(
+        &self,
+        ephemeral_key: String,
+        mac: String,
+        ciphertext: String,
+    ) -> Result<String, PkDecryptionError> {
+        self.inner
+            .decrypt(ephemeral_key, mac, ciphertext)
+            .map_err(|e| e.into())
+    }
 }

rust-sdk/src/lib.rs

@@ -17,7 +17,7 @@ mod responses;
 mod users;
 mod verification;
 
-pub use backup_recovery_key::{BackupKey, BackupRecoveryKey, PassphraseInfo};
+pub use backup_recovery_key::{BackupKey, BackupRecoveryKey, PassphraseInfo, PkDecryptionError};
 pub use device::Device;
 pub use error::{
     CryptoStoreError, DecryptionError, KeyImportError, SecretImportError, SignatureError,

rust-sdk/src/machine.rs

@@ -35,6 +35,7 @@ use matrix_sdk_crypto::{
     backups::{MegolmV1BackupKey, RecoveryKey},
     decrypt_key_export, encrypt_key_export,
     matrix_qrcode::QrVerificationData,
+    olm::ExportedRoomKey,
     EncryptionSettings, LocalTrust, OlmMachine as InnerMachine, UserIdentities,
     Verification as RustVerification,
 };
@@ -603,6 +604,25 @@ impl OlmMachine {
         Ok(encrypted)
     }
 
+    fn impor_keys_helper(
+        &self,
+        keys: Vec<ExportedRoomKey>,
+        progress_listener: Box<dyn ProgressListener>,
+    ) -> Result<KeysImportResult, KeyImportError> {
+        let listener = |progress: usize, total: usize| {
+            progress_listener.on_progress(progress as i32, total as i32)
+        };
+
+        let result = self
+            .runtime
+            .block_on(self.inner.import_keys(keys, listener))?;
+
+        Ok(KeysImportResult {
+            total: result.1 as i32,
+            imported: result.0 as i32,
+        })
+    }
+
     /// Import room keys from the given serialized key export.
     ///
     /// # Arguments
@@ -621,19 +641,17 @@ impl OlmMachine {
     ) -> Result<KeysImportResult, KeyImportError> {
         let keys = Cursor::new(keys);
         let keys = decrypt_key_export(keys, passphrase)?;
+        self.impor_keys_helper(keys, progress_listener)
+    }
 
-        let listener = |progress: usize, total: usize| {
+    /// TODO
-            progress_listener.on_progress(progress as i32, total as i32)
+    pub fn import_decrypted_keys(
-        };
+        &self,
-
+        keys: &str,
-        let result = self
+        progress_listener: Box<dyn ProgressListener>,
-            .runtime
+    ) -> Result<KeysImportResult, KeyImportError> {
-            .block_on(self.inner.import_keys(keys, listener))?;
+        let keys: Vec<ExportedRoomKey> = serde_json::from_str(keys).unwrap();
-
+        self.impor_keys_helper(keys, progress_listener)
-        Ok(KeysImportResult {
-            total: result.1 as i32,
-            imported: result.0 as i32,
-        })
     }
 
     /// Discard the currently active room key for the given room if there is

rust-sdk/src/olm.udl

@@ -10,6 +10,11 @@ callback interface ProgressListener {
     void on_progress(i32 progress, i32 total);
 };
 
+[Error]
+enum PkDecryptionError {
+    "Olm",
+};
+
 [Error]
 enum KeyImportError {
     "Export",
@@ -334,6 +339,11 @@ interface OlmMachine {
         [ByRef] string passphrase,
         ProgressListener progress_listener
     );
+    [Throws=KeyImportError]
+    KeysImportResult import_decrypted_keys(
+        [ByRef] string keys,
+        ProgressListener progress_listener
+    );
     [Throws=CryptoStoreError]
     void discard_room_key([ByRef] string room_id);
 
@@ -394,5 +404,8 @@ interface BackupRecoveryKey {
     [Name=from_base58]
     constructor(string key);
     string to_base58();
+    string to_base64();
     BackupKey public_key();
+    [Throws=PkDecryptionError]
+    string decrypt(string ephemeral_key, string mac, string ciphertext);
 };