fix: asset restriction [WPB-9947] (#2831) (#2856) 🍒 (#2862)

* Commit with unresolved merge conflicts * asset restriction * fix tests --------- Co-authored-by: Mohamad Jaara <[email protected]>
wireapp · Jul 8, 2024 · e935501 · e935501
1 parent e497778
commit e935501
Show file tree

Hide file tree

Showing 5 changed files with 258 additions and 28 deletions.
diff --git a/logic/src/commonMain/kotlin/com/wire/kalium/logic/feature/UserSessionScope.kt b/logic/src/commonMain/kotlin/com/wire/kalium/logic/feature/UserSessionScope.kt
@@ -1767,6 +1767,7 @@ class UserSessionScope internal constructor(
             messageMetadataRepository,
             staleEpochVerifier,
             legalHoldHandler,
+            observeFileSharingStatus,
             this,
             userScopedLogger,
         )

diff --git a/...c/commonMain/kotlin/com/wire/kalium/logic/feature/asset/ScheduleNewAssetMessageUseCase.kt b/...c/commonMain/kotlin/com/wire/kalium/logic/feature/asset/ScheduleNewAssetMessageUseCase.kt
@@ -24,6 +24,7 @@ import com.wire.kalium.cryptography.utils.AES256Key
 import com.wire.kalium.cryptography.utils.SHA256Key
 import com.wire.kalium.cryptography.utils.generateRandomAES256Key
 import com.wire.kalium.logic.CoreFailure
+import com.wire.kalium.logic.configuration.FileSharingStatus
 import com.wire.kalium.logic.data.asset.AssetRepository
 import com.wire.kalium.logic.data.asset.AssetTransferStatus
 import com.wire.kalium.logic.data.asset.UploadedAssetId
@@ -44,6 +45,7 @@ import com.wire.kalium.logic.data.user.UserId
 import com.wire.kalium.logic.feature.message.MessageSendFailureHandler
 import com.wire.kalium.logic.feature.message.MessageSender
 import com.wire.kalium.logic.feature.selfDeletingMessages.ObserveSelfDeletionTimerSettingsForConversationUseCase
+import com.wire.kalium.logic.feature.user.ObserveFileSharingStatusUseCase
 import com.wire.kalium.logic.functional.Either
 import com.wire.kalium.logic.functional.flatMap
 import com.wire.kalium.logic.functional.fold
@@ -108,12 +110,14 @@ internal class ScheduleNewAssetMessageUseCaseImpl(
     private val userPropertyRepository: UserPropertyRepository,
     private val selfDeleteTimer: ObserveSelfDeletionTimerSettingsForConversationUseCase,
     private val scope: CoroutineScope,
+    private val observeFileSharingStatus: ObserveFileSharingStatusUseCase,
+    private val validateAssetMimeTypeUseCase: ValidateAssetMimeTypeUseCase,
     private val dispatcher: KaliumDispatcher,
 ) : ScheduleNewAssetMessageUseCase {
 
     private var outGoingAssetUploadJob: Job? = null
 
-    @Suppress("LongMethod")
+    @Suppress("LongMethod", "ReturnCount")
     override suspend fun invoke(
         conversationId: ConversationId,
         assetDataPath: Path,
@@ -124,6 +128,18 @@ internal class ScheduleNewAssetMessageUseCaseImpl(
         assetHeight: Int?,
         audioLengthInMs: Long
     ): ScheduleNewAssetMessageResult {
+        observeFileSharingStatus().first().also {
+            when (it.state) {
+                FileSharingStatus.Value.Disabled -> return ScheduleNewAssetMessageResult.Failure.DisabledByTeam
+                FileSharingStatus.Value.EnabledAll -> { /* no-op*/ }
+
+                is FileSharingStatus.Value.EnabledSome -> if (!validateAssetMimeTypeUseCase(assetMimeType, it.state.allowedType)) {
+                    kaliumLogger.e("The asset message trying to be processed has invalid content data")
+                    return ScheduleNewAssetMessageResult.Failure.RestrictedFileType
+                }
+            }
+        }
+
         slowSyncRepository.slowSyncStatus.first {
             it is SlowSyncStatus.Complete
         }
@@ -165,7 +181,7 @@ internal class ScheduleNewAssetMessageUseCaseImpl(
                 }
             }
         }.fold({
-            ScheduleNewAssetMessageResult.Failure(it)
+            ScheduleNewAssetMessageResult.Failure.Generic(it)
         }, { (_, message) ->
             ScheduleNewAssetMessageResult.Success(message.id)
         })
@@ -345,9 +361,13 @@ internal class ScheduleNewAssetMessageUseCaseImpl(
     }
 }
 
-sealed class ScheduleNewAssetMessageResult {
-    class Success(val messageId: String) : ScheduleNewAssetMessageResult()
-    class Failure(val coreFailure: CoreFailure) : ScheduleNewAssetMessageResult()
+sealed interface ScheduleNewAssetMessageResult {
+    data class Success(val messageId: String) : ScheduleNewAssetMessageResult
+    sealed interface Failure : ScheduleNewAssetMessageResult {
+        data class Generic(val coreFailure: CoreFailure) : Failure
+        data object DisabledByTeam : Failure
+        data object RestrictedFileType : Failure
+    }
 }
 
 private data class AssetMessageMetadata(

diff --git a/logic/src/commonMain/kotlin/com/wire/kalium/logic/feature/message/MessageScope.kt b/logic/src/commonMain/kotlin/com/wire/kalium/logic/feature/message/MessageScope.kt
@@ -61,6 +61,8 @@ import com.wire.kalium.logic.feature.asset.ScheduleNewAssetMessageUseCaseImpl
 import com.wire.kalium.logic.feature.asset.UpdateAssetMessageTransferStatusUseCase
 import com.wire.kalium.logic.feature.asset.UpdateAssetMessageTransferStatusUseCaseImpl
 import com.wire.kalium.logic.feature.message.composite.SendButtonActionConfirmationMessageUseCase
+import com.wire.kalium.logic.feature.asset.ValidateAssetMimeTypeUseCase
+import com.wire.kalium.logic.feature.asset.ValidateAssetMimeTypeUseCaseImpl
 import com.wire.kalium.logic.feature.message.composite.SendButtonActionMessageUseCase
 import com.wire.kalium.logic.feature.message.composite.SendButtonMessageUseCase
 import com.wire.kalium.logic.feature.message.confirmation.ConfirmationDeliveryHandler
@@ -82,6 +84,7 @@ import com.wire.kalium.logic.feature.message.ephemeral.EphemeralMessageDeletionH
 import com.wire.kalium.logic.feature.selfDeletingMessages.ObserveSelfDeletionTimerSettingsForConversationUseCase
 import com.wire.kalium.logic.feature.sessionreset.ResetSessionUseCase
 import com.wire.kalium.logic.feature.sessionreset.ResetSessionUseCaseImpl
+import com.wire.kalium.logic.feature.user.ObserveFileSharingStatusUseCase
 import com.wire.kalium.logic.sync.SyncManager
 import com.wire.kalium.logic.sync.receiver.handler.legalhold.LegalHoldHandler
 import com.wire.kalium.logic.util.MessageContentEncoder
@@ -118,8 +121,9 @@ class MessageScope internal constructor(
     private val messageMetadataRepository: MessageMetadataRepository,
     private val staleEpochVerifier: StaleEpochVerifier,
     private val legalHoldHandler: LegalHoldHandler,
+    private val observeFileSharingStatusUseCase: ObserveFileSharingStatusUseCase,
     private val scope: CoroutineScope,
-    private val kaliumLogger: KaliumLogger,
+    kaliumLogger: KaliumLogger,
     internal val dispatcher: KaliumDispatcher = KaliumDispatcherImpl,
     private val legalHoldStatusMapper: LegalHoldStatusMapper = LegalHoldStatusMapperImpl
 ) {
@@ -155,6 +159,9 @@ class MessageScope internal constructor(
             protoContentMapper = protoContentMapper
         )
 
+    private val validateAssetMimeTypeUseCase: ValidateAssetMimeTypeUseCase
+        get() = ValidateAssetMimeTypeUseCaseImpl()
+
     private val messageContentEncoder = MessageContentEncoder()
     private val messageSendingInterceptor: MessageSendingInterceptor
         get() = MessageSendingInterceptorImpl(messageContentEncoder, messageRepository)
@@ -269,6 +276,8 @@ class MessageScope internal constructor(
             userPropertyRepository,
             observeSelfDeletingMessages,
             scope,
+            observeFileSharingStatusUseCase,
+            validateAssetMimeTypeUseCase,
             dispatcher
         )