Merge pull request #163 from ouchadam/bug/various-crypto-fixes

Various crypto fixes

domains/olm/src/main/kotlin/app/dapk/st/olm/OlmWrapper.kt

@@ -170,6 +170,8 @@ class OlmWrapper(
         val inBound = OlmInboundGroupSession(roomCryptoSession.key)
         olmStore.persist(roomCryptoSession.id, inBound)
 
+        logger.crypto("Creating megolm: ${roomCryptoSession.id}")
+
         return roomCryptoSession
     }
 
@@ -181,7 +183,7 @@ class OlmWrapper(
     private suspend fun deviceCrypto(input: OlmSessionInput): DeviceCryptoSession? {
         return olmStore.readSessions(listOf(input.identity))?.let {
             DeviceCryptoSession(
-                input.deviceId, input.userId, input.identity, input.fingerprint, it
+                input.deviceId, input.userId, input.identity, input.fingerprint, it.map { it.second }
             )
         }
     }

matrix/matrix/src/main/kotlin/app/dapk/st/matrix/MatrixClient.kt

@@ -43,7 +43,11 @@ data class ServiceDependencies(
 
 interface MatrixServiceInstaller {
     fun serializers(builder: SerializersModuleBuilder.() -> Unit)
-    fun install(factory: MatrixService.Factory)
+    fun <T : MatrixService> install(factory: MatrixService.Factory): InstallExtender<T>
+}
+
+interface InstallExtender<T : MatrixService> {
+    fun proxy(proxy: (T) -> T)
 }
 
 interface MatrixServiceProvider {

matrix/matrix/src/main/kotlin/app/dapk/st/matrix/ServiceInstaller.kt

@@ -11,15 +11,22 @@ internal class ServiceInstaller {
     private val services = mutableMapOf<Any, MatrixService>()
     private val serviceInstaller = object : MatrixServiceInstaller {
 
-        val serviceCollector = mutableListOf<MatrixService.Factory>()
+        val serviceCollector = mutableListOf<Pair<MatrixService.Factory, (MatrixService) -> MatrixService>>()
         val serializers = mutableListOf<SerializersModuleBuilder.() -> Unit>()
 
         override fun serializers(builder: SerializersModuleBuilder.() -> Unit) {
             serializers.add(builder)
         }
 
-        override fun install(factory: MatrixService.Factory) {
+        override fun <T : MatrixService> install(factory: MatrixService.Factory): InstallExtender<T> {
-            serviceCollector.add(factory)
+            val mutableProxy = MutableProxy<T>()
+            return object : InstallExtender<T> {
+                override fun proxy(proxy: (T) -> T) {
+                    mutableProxy.value = proxy
+                }
+            }.also {
+                serviceCollector.add(factory to mutableProxy)
+            }
         }
     }
 
@@ -39,9 +46,9 @@ internal class ServiceInstaller {
         val serviceProvider = object : MatrixServiceProvider {
             override fun <T : MatrixService> getService(key: ServiceKey) = this@ServiceInstaller.getService<T>(key)
         }
-        serviceInstaller.serviceCollector.forEach {
+        serviceInstaller.serviceCollector.forEach { (factory, extender) ->
-            val (key, service) = it.create(ServiceDependencies(httpClient, json, serviceProvider, logger))
+            val (key, service) = factory.create(ServiceDependencies(httpClient, json, serviceProvider, logger))
-            services[key] = service
+            services[key] = extender(service)
         }
     }
 
@@ -58,3 +65,12 @@ internal class ServiceInstaller {
     }
 
 }
+
+internal class MutableProxy<T : MatrixService> : (MatrixService) -> MatrixService {
+
+    var value: (T) -> T = { it }
+
+    @Suppress("UNCHECKED_CAST")
+    override fun invoke(service: MatrixService) = value(service as T)
+
+}

matrix/services/auth/src/main/kotlin/app/dapk/st/matrix/auth/AuthService.kt

@@ -1,5 +1,6 @@
 package app.dapk.st.matrix.auth
 
+import app.dapk.st.matrix.InstallExtender
 import app.dapk.st.matrix.MatrixClient
 import app.dapk.st.matrix.MatrixService
 import app.dapk.st.matrix.MatrixServiceInstaller
@@ -25,8 +26,8 @@ interface AuthService : MatrixService {
 
 fun MatrixServiceInstaller.installAuthService(
     credentialsStore: CredentialsStore,
-) {
+): InstallExtender<AuthService> {
-    this.install { (httpClient, json) ->
+    return this.install { (httpClient, json) ->
         SERVICE_KEY to DefaultAuthService(httpClient, credentialsStore, json)
     }
 }

matrix/services/crypto/src/main/kotlin/app/dapk/st/matrix/crypto/CryptoService.kt

@@ -2,10 +2,7 @@ package app.dapk.st.matrix.crypto
 
 import app.dapk.st.core.Base64
 import app.dapk.st.core.CoroutineDispatchers
-import app.dapk.st.matrix.MatrixService
+import app.dapk.st.matrix.*
-import app.dapk.st.matrix.MatrixServiceInstaller
-import app.dapk.st.matrix.MatrixServiceProvider
-import app.dapk.st.matrix.ServiceDepFactory
 import app.dapk.st.matrix.common.*
 import app.dapk.st.matrix.crypto.internal.*
 import app.dapk.st.matrix.device.deviceService
@@ -136,8 +133,8 @@ fun MatrixServiceInstaller.installCryptoService(
     roomMembersProvider: ServiceDepFactory<RoomMembersProvider>,
     base64: Base64,
     coroutineDispatchers: CoroutineDispatchers,
-) {
+): InstallExtender<CryptoService> {
-    this.install { (_, _, services, logger) ->
+    return this.install { (_, _, services, logger) ->
         val deviceService = services.deviceService()
         val accountCryptoUseCase = FetchAccountCryptoUseCaseImpl(credentialsStore, olm, deviceService)
 

matrix/services/device/src/main/kotlin/app/dapk/st/matrix/device/DeviceService.kt

@@ -1,5 +1,6 @@
 package app.dapk.st.matrix.device
 
+import app.dapk.st.matrix.InstallExtender
 import app.dapk.st.matrix.MatrixService
 import app.dapk.st.matrix.MatrixServiceInstaller
 import app.dapk.st.matrix.MatrixServiceProvider
@@ -122,8 +123,8 @@ sealed class ToDevicePayload {
     sealed interface VerificationPayload
 }
 
-fun MatrixServiceInstaller.installEncryptionService(knownDeviceStore: KnownDeviceStore) {
+fun MatrixServiceInstaller.installEncryptionService(knownDeviceStore: KnownDeviceStore): InstallExtender<DeviceService> {
-    this.install { (httpClient, _, _, logger) ->
+    return this.install { (httpClient, _, _, logger) ->
         SERVICE_KEY to DefaultDeviceService(httpClient, logger, knownDeviceStore)
     }
 }

matrix/services/message/src/main/kotlin/app/dapk/st/matrix/message/MessageService.kt

@@ -1,10 +1,7 @@
 package app.dapk.st.matrix.message
 
 import app.dapk.st.core.Base64
-import app.dapk.st.matrix.MatrixService
+import app.dapk.st.matrix.*
-import app.dapk.st.matrix.MatrixServiceInstaller
-import app.dapk.st.matrix.MatrixServiceProvider
-import app.dapk.st.matrix.ServiceDepFactory
 import app.dapk.st.matrix.common.AlgorithmName
 import app.dapk.st.matrix.common.EventId
 import app.dapk.st.matrix.common.MessageType
@@ -132,8 +129,8 @@ fun MatrixServiceInstaller.installMessageService(
     imageContentReader: ImageContentReader,
     messageEncrypter: ServiceDepFactory<MessageEncrypter> = ServiceDepFactory { MissingMessageEncrypter },
     mediaEncrypter: ServiceDepFactory<MediaEncrypter> = ServiceDepFactory { MissingMediaEncrypter },
-) {
+): InstallExtender<MessageService> {
-    this.install { (httpClient, _, installedServices) ->
+    return this.install { (httpClient, _, installedServices) ->
         SERVICE_KEY to DefaultMessageService(
             httpClient,
             localEchoStore,

matrix/services/profile/src/main/kotlin/app/dapk/st/matrix/room/ProfileService.kt

@@ -1,6 +1,7 @@
 package app.dapk.st.matrix.room
 
 import app.dapk.st.core.SingletonFlows
+import app.dapk.st.matrix.InstallExtender
 import app.dapk.st.matrix.MatrixService
 import app.dapk.st.matrix.MatrixServiceInstaller
 import app.dapk.st.matrix.MatrixServiceProvider
@@ -29,8 +30,8 @@ fun MatrixServiceInstaller.installProfileService(
     profileStore: ProfileStore,
     singletonFlows: SingletonFlows,
     credentialsStore: CredentialsStore,
-) {
+): InstallExtender<ProfileService> {
-    this.install { (httpClient, _, _, logger) ->
+    return this.install { (httpClient, _, _, logger) ->
         SERVICE_KEY to DefaultProfileService(httpClient, logger, profileStore, singletonFlows, credentialsStore)
     }
 }

matrix/services/push/src/main/kotlin/app/dapk/st/matrix/push/PushService.kt

@@ -1,5 +1,6 @@
 package app.dapk.st.matrix.push
 
+import app.dapk.st.matrix.InstallExtender
 import app.dapk.st.matrix.MatrixClient
 import app.dapk.st.matrix.MatrixService
 import app.dapk.st.matrix.MatrixServiceInstaller
@@ -38,8 +39,8 @@ interface PushService : MatrixService {
 
 fun MatrixServiceInstaller.installPushService(
     credentialsStore: CredentialsStore,
-) {
+): InstallExtender<PushService> {
-    this.install { (httpClient, _, _, logger) ->
+    return this.install { (httpClient, _, _, logger) ->
         SERVICE_KEY to DefaultPushService(httpClient, credentialsStore, logger)
     }
 }

matrix/services/room/src/main/kotlin/app/dapk/st/matrix/room/RoomService.kt

@@ -1,9 +1,6 @@
 package app.dapk.st.matrix.room
 
-import app.dapk.st.matrix.MatrixService
+import app.dapk.st.matrix.*
-import app.dapk.st.matrix.MatrixServiceInstaller
-import app.dapk.st.matrix.MatrixServiceProvider
-import app.dapk.st.matrix.ServiceDepFactory
 import app.dapk.st.matrix.common.EventId
 import app.dapk.st.matrix.common.RoomId
 import app.dapk.st.matrix.common.RoomMember
@@ -42,8 +39,8 @@ fun MatrixServiceInstaller.installRoomService(
     memberStore: MemberStore,
     roomMessenger: ServiceDepFactory<RoomMessenger>,
     roomInviteRemover: RoomInviteRemover,
-) {
+): InstallExtender<RoomService> {
-    this.install { (httpClient, _, services, logger) ->
+    return this.install { (httpClient, _, services, logger) ->
         SERVICE_KEY to DefaultRoomService(
             httpClient,
             logger,

matrix/services/sync/src/main/kotlin/app/dapk/st/matrix/sync/SyncService.kt

@@ -2,10 +2,7 @@ package app.dapk.st.matrix.sync
 
 import app.dapk.st.core.CoroutineDispatchers
 import app.dapk.st.core.extensions.ErrorTracker
-import app.dapk.st.matrix.MatrixClient
+import app.dapk.st.matrix.*
-import app.dapk.st.matrix.MatrixService
-import app.dapk.st.matrix.MatrixServiceInstaller
-import app.dapk.st.matrix.ServiceDepFactory
 import app.dapk.st.matrix.common.*
 import app.dapk.st.matrix.sync.internal.DefaultSyncService
 import app.dapk.st.matrix.sync.internal.request.*
@@ -49,7 +46,7 @@ fun MatrixServiceInstaller.installSyncService(
     errorTracker: ErrorTracker,
     coroutineDispatchers: CoroutineDispatchers,
     syncConfig: SyncConfig = SyncConfig(),
-) {
+): InstallExtender<SyncService> {
     this.serializers {
         polymorphicDefault(ApiTimelineEvent::class) {
             ApiTimelineEvent.Ignored.serializer()
@@ -71,7 +68,7 @@ fun MatrixServiceInstaller.installSyncService(
         }
     }
 
-    this.install { (httpClient, json, services, logger) ->
+    return this.install { (httpClient, json, services, logger) ->
         SERVICE_KEY to DefaultSyncService(
             httpClient = httpClient,
             syncStore = syncStore,

matrix/services/sync/src/main/kotlin/app/dapk/st/matrix/sync/internal/sync/SyncUseCase.kt

@@ -8,7 +8,6 @@ import app.dapk.st.matrix.sync.internal.SideEffectFlowIterator
 import app.dapk.st.matrix.sync.internal.overview.ReducedSyncFilterUseCase
 import app.dapk.st.matrix.sync.internal.request.syncRequest
 import app.dapk.st.matrix.sync.internal.room.SyncSideEffects
-import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.cancellable
 import kotlinx.coroutines.flow.flow
 
@@ -25,8 +24,7 @@ internal class SyncUseCase(
     private val syncConfig: SyncConfig,
 ) {
 
-    fun sync(): Flow<Unit> {
+    private val _flow = flow {
-        return flow {
         val credentials = credentialsStore.credentials()!!
         val filterId = filterUseCase.reducedFilter(credentials.userId)
         with(flowIterator) {
@@ -37,7 +35,6 @@ internal class SyncUseCase(
             )
         }
     }.cancellable()
-    }
 
     private suspend fun onEachSyncIteration(filterId: SyncService.FilterId, credentials: UserCredentials, previousState: OverviewState?): OverviewState? {
         val syncToken = syncStore.read(key = SyncStore.SyncKey.Overview)
@@ -85,4 +82,6 @@ internal class SyncUseCase(
         )
     }
 
+    fun sync() = _flow
+
 }

test-harness/src/test/kotlin/SmokeTest.kt

@@ -123,18 +123,20 @@ class SmokeTest {
         alice.expectTextMessage(SharedState.sharedRoom, message2)
 
         // Needs investigation
-//        val aliceSecondDevice = testMatrix(SharedState.alice, isTemp = true, withLogging = true).also { it.newlogin() }
+        val aliceSecondDevice = testMatrix(SharedState.alice, isTemp = true, withLogging = true).also { it.newlogin() }
-//        aliceSecondDevice.client.syncService().startSyncing().collectAsync {
+        aliceSecondDevice.client.syncService().startSyncing().collectAsync {
-//            val message3 = "from alice to bob and alice's second device".from(SharedState.alice.roomMember)
+            aliceSecondDevice.client.proxyDeviceService().waitForOneTimeKeysToBeUploaded()
-//            alice.sendTextMessage(SharedState.sharedRoom, message3.content, isEncrypted)
+
-//            aliceSecondDevice.expectTextMessage(SharedState.sharedRoom, message3)
+            val message3 = "from alice to bob and alice's second device".from(SharedState.alice.roomMember)
-//            bob.expectTextMessage(SharedState.sharedRoom, message3)
+            alice.sendTextMessage(SharedState.sharedRoom, message3.content, isEncrypted)
-//
+            aliceSecondDevice.expectTextMessage(SharedState.sharedRoom, message3)
-//            val message4 = "from alice's second device to bob and alice's first device".from(SharedState.alice.roomMember)
+            bob.expectTextMessage(SharedState.sharedRoom, message3)
-//            aliceSecondDevice.sendTextMessage(SharedState.sharedRoom, message4.content, isEncrypted)
+
-//            alice.expectTextMessage(SharedState.sharedRoom, message4)
+            val message4 = "from alice's second device to bob and alice's first device".from(SharedState.alice.roomMember)
-//            bob.expectTextMessage(SharedState.sharedRoom, message4)
+            aliceSecondDevice.sendTextMessage(SharedState.sharedRoom, message4.content, isEncrypted)
-//        }
+            alice.expectTextMessage(SharedState.sharedRoom, message4)
+            bob.expectTextMessage(SharedState.sharedRoom, message4)
+        }
     }
 }
 

test-harness/src/test/kotlin/test/TestMatrix.kt

@@ -14,6 +14,7 @@ import app.dapk.st.matrix.crypto.RoomMembersProvider
 import app.dapk.st.matrix.crypto.Verification
 import app.dapk.st.matrix.crypto.cryptoService
 import app.dapk.st.matrix.crypto.installCryptoService
+import app.dapk.st.matrix.device.DeviceService
 import app.dapk.st.matrix.device.deviceService
 import app.dapk.st.matrix.device.installEncryptionService
 import app.dapk.st.matrix.http.ktor.KtorMatrixHttpClientFactory
@@ -39,6 +40,7 @@ import test.impl.PrintingErrorTracking
 import java.io.File
 import java.time.Clock
 import javax.imageio.ImageIO
+import kotlin.coroutines.resume
 
 object TestUsers {
 
@@ -93,7 +95,9 @@ class TestMatrix(
     ).also {
         it.install {
             installAuthService(storeModule.credentialsStore())
-            installEncryptionService(storeModule.knownDevicesStore())
+            installEncryptionService(storeModule.knownDevicesStore()).proxy {
+                ProxyDeviceService(it)
+            }
 
             val olmAccountStore = OlmPersistenceWrapper(storeModule.olmStore(), base64)
             val olm = OlmWrapper(
@@ -356,3 +360,22 @@ class JavaImageContentReader : ImageContentReader {
     override fun inputStream(uri: String) = File(uri).inputStream()
 
 }
+
+class ProxyDeviceService(private val deviceService: DeviceService) : DeviceService by deviceService {
+
+    private var oneTimeKeysContinuation: (() -> Unit)? = null
+
+    override suspend fun uploadOneTimeKeys(oneTimeKeys: DeviceService.OneTimeKeys) {
+        deviceService.uploadOneTimeKeys(oneTimeKeys)
+        oneTimeKeysContinuation?.invoke()?.also { oneTimeKeysContinuation = null }
+    }
+
+    suspend fun waitForOneTimeKeysToBeUploaded() {
+        suspendCancellableCoroutine { continuation ->
+            oneTimeKeysContinuation = { continuation.resume(Unit) }
+        }
+    }
+
+}
+
+fun MatrixClient.proxyDeviceService() = this.deviceService() as ProxyDeviceService