Move and improve withOlmDecryption() and withOlmEncryption()

matrix-sdk-android/src/androidTest/java/im/vector/matrix/android/internal/crypto/ssss/QuadSTests.kt

@@ -35,6 +35,7 @@ import im.vector.matrix.android.common.TestMatrixCallback
 import im.vector.matrix.android.internal.crypto.SSSS_ALGORITHM_CURVE25519_AES_SHA2
 import im.vector.matrix.android.internal.crypto.crosssigning.toBase64NoPadding
 import im.vector.matrix.android.internal.crypto.secrets.DefaultSharedSecretStorageService
+import im.vector.matrix.android.internal.crypto.tools.withOlmDecryption
 import im.vector.matrix.android.internal.session.sync.model.accountdata.UserAccountDataEvent
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.GlobalScope
@@ -100,10 +101,10 @@ class QuadSTests : InstrumentedTest {
         assertNotNull("Pubkey should be defined", parsed.publicKey)
 
         val privateKeySpec = Curve25519AesSha2KeySpec.fromRecoveryKey(ssssKeyCreationInfo.recoveryKey)
-        DefaultSharedSecretStorageService.withOlmDecryption { olmPkDecryption ->
+        val pubKey = withOlmDecryption { olmPkDecryption ->
-            val pubKey = olmPkDecryption.setPrivateKey(privateKeySpec!!.privateKey)
+            olmPkDecryption.setPrivateKey(privateKeySpec!!.privateKey)
-            assertEquals("Unexpected Public Key", pubKey, parsed.publicKey)
         }
+        assertEquals("Unexpected Public Key", pubKey, parsed.publicKey)
 
         // Set as default key
         quadS.setDefaultKey(TEST_KEY_ID, object : MatrixCallback<Unit> {})

matrix-sdk-android/src/main/java/im/vector/matrix/android/internal/crypto/secrets/DefaultSharedSecretStorageService.kt

@@ -25,21 +25,21 @@ import im.vector.matrix.android.api.session.securestorage.EncryptedSecretContent
 import im.vector.matrix.android.api.session.securestorage.KeyInfo
 import im.vector.matrix.android.api.session.securestorage.KeyInfoResult
 import im.vector.matrix.android.api.session.securestorage.KeySigner
-import im.vector.matrix.android.api.session.securestorage.SsssKeyCreationInfo
 import im.vector.matrix.android.api.session.securestorage.SSSSKeySpec
 import im.vector.matrix.android.api.session.securestorage.SSSSPassphrase
 import im.vector.matrix.android.api.session.securestorage.SecretStorageKeyContent
 import im.vector.matrix.android.api.session.securestorage.SharedSecretStorageError
 import im.vector.matrix.android.api.session.securestorage.SharedSecretStorageService
+import im.vector.matrix.android.api.session.securestorage.SsssKeyCreationInfo
 import im.vector.matrix.android.internal.crypto.SSSS_ALGORITHM_CURVE25519_AES_SHA2
 import im.vector.matrix.android.internal.crypto.keysbackup.generatePrivateKeyWithPassword
 import im.vector.matrix.android.internal.crypto.keysbackup.util.computeRecoveryKey
+import im.vector.matrix.android.internal.crypto.tools.withOlmEncryption
 import im.vector.matrix.android.internal.extensions.foldToCallback
 import im.vector.matrix.android.internal.util.MatrixCoroutineDispatchers
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.launch
 import org.matrix.olm.OlmPkDecryption
-import org.matrix.olm.OlmPkEncryption
 import org.matrix.olm.OlmPkMessage
 import javax.inject.Inject
 
@@ -198,15 +198,15 @@ internal class DefaultSharedSecretStorageService @Inject constructor(
                     when (key) {
                         is KeyInfoResult.Success -> {
                             if (key.keyInfo.content.algorithm == SSSS_ALGORITHM_CURVE25519_AES_SHA2) {
-                                withOlmEncryption { olmEncrypt ->
+                                val encryptedResult = withOlmEncryption { olmEncrypt ->
                                     olmEncrypt.setRecipientKey(key.keyInfo.content.publicKey)
-                                    val encryptedResult = olmEncrypt.encrypt(secretBase64)
+                                    olmEncrypt.encrypt(secretBase64)
-                                    encryptedContents[key.keyInfo.id] = EncryptedSecretContent(
-                                            ciphertext = encryptedResult.mCipherText,
-                                            ephemeral = encryptedResult.mEphemeralKey,
-                                            mac = encryptedResult.mMac
-                                    )
                                 }
+                                encryptedContents[key.keyInfo.id] = EncryptedSecretContent(
+                                        ciphertext = encryptedResult.mCipherText,
+                                        ephemeral = encryptedResult.mEphemeralKey,
+                                        mac = encryptedResult.mMac
+                                )
                             } else {
                                 // Unknown algorithm
                                 callback.onFailure(SharedSecretStorageError.UnknownAlgorithm(key.keyInfo.content.algorithm ?: ""))
@@ -226,15 +226,15 @@ internal class DefaultSharedSecretStorageService @Inject constructor(
                         when (key) {
                             is KeyInfoResult.Success -> {
                                 if (key.keyInfo.content.algorithm == SSSS_ALGORITHM_CURVE25519_AES_SHA2) {
-                                    withOlmEncryption { olmEncrypt ->
+                                    val encryptedResult = withOlmEncryption { olmEncrypt ->
                                         olmEncrypt.setRecipientKey(key.keyInfo.content.publicKey)
-                                        val encryptedResult = olmEncrypt.encrypt(secretBase64)
+                                        olmEncrypt.encrypt(secretBase64)
-                                        encryptedContents[keyId] = EncryptedSecretContent(
-                                                ciphertext = encryptedResult.mCipherText,
-                                                ephemeral = encryptedResult.mEphemeralKey,
-                                                mac = encryptedResult.mMac
-                                        )
                                     }
+                                    encryptedContents[keyId] = EncryptedSecretContent(
+                                            ciphertext = encryptedResult.mCipherText,
+                                            ephemeral = encryptedResult.mEphemeralKey,
+                                            mac = encryptedResult.mMac
+                                    )
                                 } else {
                                     // Unknown algorithm
                                     callback.onFailure(SharedSecretStorageError.UnknownAlgorithm(key.keyInfo.content.algorithm ?: ""))
@@ -332,27 +332,5 @@ internal class DefaultSharedSecretStorageService @Inject constructor(
         const val KEY_ID_BASE = "m.secret_storage.key"
         const val ENCRYPTED = "encrypted"
         const val DEFAULT_KEY_ID = "m.secret_storage.default_key"
-
-        fun withOlmEncryption(block: (OlmPkEncryption) -> Unit) {
-            val olmPkEncryption = OlmPkEncryption()
-            try {
-                block(olmPkEncryption)
-            } catch (failure: Throwable) {
-                throw failure
-            } finally {
-                olmPkEncryption.releaseEncryption()
-            }
-        }
-
-        fun withOlmDecryption(block: (OlmPkDecryption) -> Unit) {
-            val olmPkDecryption = OlmPkDecryption()
-            try {
-                block(olmPkDecryption)
-            } catch (failure: Throwable) {
-                throw failure
-            } finally {
-                olmPkDecryption.releaseDecryption()
-            }
-        }
     }
 }

matrix-sdk-android/src/main/java/im/vector/matrix/android/internal/crypto/tools/Tools.kt

@@ -0,0 +1,38 @@
+/*
+ * Copyright (c) 2020 New Vector Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package im.vector.matrix.android.internal.crypto.tools
+
+import org.matrix.olm.OlmPkDecryption
+import org.matrix.olm.OlmPkEncryption
+
+fun <T> withOlmEncryption(block: (OlmPkEncryption) -> T): T {
+    val olmPkEncryption = OlmPkEncryption()
+    try {
+        return block(olmPkEncryption)
+    } finally {
+        olmPkEncryption.releaseEncryption()
+    }
+}
+
+fun <T> withOlmDecryption(block: (OlmPkDecryption) -> T): T {
+    val olmPkDecryption = OlmPkDecryption()
+    try {
+        return block(olmPkDecryption)
+    } finally {
+        olmPkDecryption.releaseDecryption()
+    }
+}