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

* fix: asset restriction (#2831) * fix: asset sending restriction * add test * add error messages and detekt * fist test service * detekt * chore: update gradle-ios-tests.yml for runner at that time --------- Co-authored-by: Yamil Medina <[email protected]> (cherry picked from commit ebcc9f8) * fix merge issues * fix unresolved reference (cherry picked from commit 116f276)
wireapp · Jul 2, 2024 · 5c47b90 · 5c47b90
1 parent 65fb353
commit 5c47b90
Show file tree

Hide file tree

Showing 5 changed files with 218 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
@@ -1740,6 +1740,7 @@ class UserSessionScope internal constructor(
             legalHoldHandler,
             this,
             userScopedLogger,
+            observeFileSharingStatus,
         )
     val users: UserScope
         get() = UserScope(

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.UploadedAssetId
 import com.wire.kalium.logic.data.asset.isAudioMimeType
@@ -43,6 +44,7 @@ import com.wire.kalium.logic.data.id.CurrentClientIdProvider
 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
@@ -107,12 +109,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,
@@ -123,6 +127,19 @@ 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
         }
@@ -175,7 +192,7 @@ internal class ScheduleNewAssetMessageUseCaseImpl(
                 }
             }
         }.fold({
-            ScheduleNewAssetMessageResult.Failure(it)
+            ScheduleNewAssetMessageResult.Failure.Generic(it)
         }, { (_, message) ->
             ScheduleNewAssetMessageResult.Success(message.id)
         })
@@ -346,9 +363,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
@@ -56,6 +56,8 @@ import com.wire.kalium.logic.feature.asset.UpdateAssetMessageDownloadStatusUseCa
 import com.wire.kalium.logic.feature.asset.UpdateAssetMessageDownloadStatusUseCaseImpl
 import com.wire.kalium.logic.feature.asset.UpdateAssetMessageUploadStatusUseCase
 import com.wire.kalium.logic.feature.asset.UpdateAssetMessageUploadStatusUseCaseImpl
+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.ephemeral.DeleteEphemeralMessageForSelfUserAsReceiverUseCaseImpl
 import com.wire.kalium.logic.feature.message.ephemeral.DeleteEphemeralMessageForSelfUserAsSenderUseCaseImpl
@@ -68,6 +70,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
@@ -103,8 +106,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
 ) {
@@ -140,6 +144,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)
@@ -238,6 +245,8 @@ class MessageScope internal constructor(
             userPropertyRepository,
             observeSelfDeletingMessages,
             scope,
+            observeFileSharingStatusUseCase,
+            validateAssetMimeTypeUseCase,
             dispatcher
         )
 

diff --git a/...mmonTest/kotlin/com/wire/kalium/logic/feature/asset/ScheduleNewAssetMessageUseCaseTest.kt b/...mmonTest/kotlin/com/wire/kalium/logic/feature/asset/ScheduleNewAssetMessageUseCaseTest.kt
@@ -22,6 +22,7 @@ import com.wire.kalium.cryptography.utils.SHA256Key
 import com.wire.kalium.logic.CoreFailure
 import com.wire.kalium.logic.NetworkFailure
 import com.wire.kalium.logic.StorageFailure
+import com.wire.kalium.logic.configuration.FileSharingStatus
 import com.wire.kalium.logic.data.asset.AssetRepository
 import com.wire.kalium.logic.data.asset.FakeKaliumFileSystem
 import com.wire.kalium.logic.data.asset.UploadedAssetId
@@ -40,6 +41,7 @@ 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.data.message.SelfDeletionTimer
+import com.wire.kalium.logic.feature.user.ObserveFileSharingStatusUseCase
 import com.wire.kalium.logic.framework.TestAsset.dummyUploadedAssetId
 import com.wire.kalium.logic.framework.TestAsset.mockedLongAssetData
 import com.wire.kalium.logic.functional.Either
@@ -92,6 +94,7 @@ class ScheduleNewAssetMessageUseCaseTest {
             .withObserveMessageVisibility()
             .withDeleteAssetLocally()
             .withSelfDeleteTimer(SelfDeletionTimer.Disabled)
+            .withObserveFileSharingStatusResult(FileSharingStatus.Value.EnabledAll)
             .arrange()
 
         // When
@@ -125,6 +128,7 @@ class ScheduleNewAssetMessageUseCaseTest {
                 .withSelfDeleteTimer(SelfDeletionTimer.Disabled)
                 .withDeleteAssetLocally()
                 .withObserveMessageVisibility()
+                .withObserveFileSharingStatusResult(FileSharingStatus.Value.EnabledAll)
                 .arrange()
 
             // When
@@ -158,6 +162,7 @@ class ScheduleNewAssetMessageUseCaseTest {
                 .withSelfDeleteTimer(SelfDeletionTimer.Disabled)
                 .withObserveMessageVisibility()
                 .withDeleteAssetLocally()
+                .withObserveFileSharingStatusResult(FileSharingStatus.Value.EnabledAll)
                 .arrange()
 
             // When
@@ -198,6 +203,7 @@ class ScheduleNewAssetMessageUseCaseTest {
             .withSelfDeleteTimer(SelfDeletionTimer.Disabled)
             .withObserveMessageVisibility()
             .withDeleteAssetLocally()
+            .withObserveFileSharingStatusResult(FileSharingStatus.Value.EnabledAll)
             .arrange()
 
         // When
@@ -248,6 +254,7 @@ class ScheduleNewAssetMessageUseCaseTest {
                 .withSelfDeleteTimer(SelfDeletionTimer.Disabled)
                 .withObserveMessageVisibility()
                 .withDeleteAssetLocally()
+                .withObserveFileSharingStatusResult(FileSharingStatus.Value.EnabledAll)
                 .arrange()
 
             // When
@@ -298,6 +305,7 @@ class ScheduleNewAssetMessageUseCaseTest {
                 .withSelfDeleteTimer(SelfDeletionTimer.Disabled)
                 .withObserveMessageVisibility()
                 .withDeleteAssetLocally()
+                .withObserveFileSharingStatusResult(FileSharingStatus.Value.EnabledAll)
                 .arrange()
 
             // When
@@ -338,6 +346,7 @@ class ScheduleNewAssetMessageUseCaseTest {
                 .withSelfDeleteTimer(SelfDeletionTimer.Disabled)
                 .withObserveMessageVisibility()
                 .withDeleteAssetLocally()
+                .withObserveFileSharingStatusResult(FileSharingStatus.Value.EnabledAll)
                 .arrange()
 
             // When
@@ -386,6 +395,7 @@ class ScheduleNewAssetMessageUseCaseTest {
                 .withSelfDeleteTimer(SelfDeletionTimer.Disabled)
                 .withObserveMessageVisibility()
                 .withDeleteAssetLocally()
+                .withObserveFileSharingStatusResult(FileSharingStatus.Value.EnabledAll)
                 .arrange()
 
             // When
@@ -441,6 +451,7 @@ class ScheduleNewAssetMessageUseCaseTest {
                 .withSelfDeleteTimer(SelfDeletionTimer.Disabled)
                 .withObserveMessageVisibility()
                 .withDeleteAssetLocally()
+                .withObserveFileSharingStatusResult(FileSharingStatus.Value.EnabledAll)
                 .arrange()
 
             // When
@@ -486,6 +497,7 @@ class ScheduleNewAssetMessageUseCaseTest {
             .withSelfDeleteTimer(SelfDeletionTimer.Disabled)
             .withObserveMessageVisibility()
             .withDeleteAssetLocally()
+            .withObserveFileSharingStatusResult(FileSharingStatus.Value.EnabledAll)
             .arrange()
 
         // When
@@ -529,6 +541,7 @@ class ScheduleNewAssetMessageUseCaseTest {
             .withSelfDeleteTimer(SelfDeletionTimer.Enabled(expectedDuration))
             .withObserveMessageVisibility()
             .withDeleteAssetLocally()
+            .withObserveFileSharingStatusResult(FileSharingStatus.Value.EnabledAll)
             .arrange()
 
         // When
@@ -555,6 +568,111 @@ class ScheduleNewAssetMessageUseCaseTest {
             })
     }
 
+    @Test
+    fun givenFileSendingRestrictedByTeam_whenSending_thenReturnDisabledByTeam() = runTest {
+        // Given
+        val assetToSend = mockedLongAssetData()
+        val assetName = "some-asset.txt"
+        val inputDataPath = fakeKaliumFileSystem.providePersistentAssetPath(assetName)
+        val conversationId = ConversationId("some-convo-id", "some-domain-id")
+        val (_, sendAssetUseCase) = Arrangement(this)
+            .withStoredData(assetToSend, inputDataPath)
+            .withObserveFileSharingStatusResult(FileSharingStatus.Value.Disabled)
+            .arrange()
+
+        // When
+        val result = sendAssetUseCase.invoke(
+            conversationId = conversationId,
+            assetDataPath = inputDataPath,
+            assetDataSize = assetToSend.size.toLong(),
+            assetName = assetName,
+            assetMimeType = "text/plain",
+            assetWidth = null,
+            assetHeight = null,
+            audioLengthInMs = 0
+        )
+        advanceUntilIdle()
+
+        // Then
+        assertTrue(result is ScheduleNewAssetMessageResult.Failure.DisabledByTeam)
+    }
+
+    @Test
+    fun givenAseetMimeTypeRestricted_whenSending_thenReturnRestrictedFileType() = runTest {
+        // Given
+        val assetToSend = mockedLongAssetData()
+        val assetName = "some-asset.txt"
+        val inputDataPath = fakeKaliumFileSystem.providePersistentAssetPath(assetName)
+        val conversationId = ConversationId("some-convo-id", "some-domain-id")
+        val (arrangement, sendAssetUseCase) = Arrangement(this)
+            .withStoredData(assetToSend, inputDataPath)
+            .withObserveFileSharingStatusResult(FileSharingStatus.Value.EnabledSome(listOf("png")))
+            .withValidateAsseMimeTypeResult(false)
+            .arrange()
+
+        // When
+        val result = sendAssetUseCase.invoke(
+            conversationId = conversationId,
+            assetDataPath = inputDataPath,
+            assetDataSize = assetToSend.size.toLong(),
+            assetName = assetName,
+            assetMimeType = "text/plain",
+            assetWidth = null,
+            assetHeight = null,
+            audioLengthInMs = 0
+        )
+        advanceUntilIdle()
+
+        // Then
+        assertTrue(result is ScheduleNewAssetMessageResult.Failure.RestrictedFileType)
+
+        verify(arrangement.validateAssetMimeTypeUseCase)
+            .function(arrangement.validateAssetMimeTypeUseCase::invoke)
+            .with(eq("text/plain"), eq(listOf("png")))
+            .wasInvoked(exactly = once)
+    }
+
+    @Test
+    fun givenAssetMimeTypeRestrictedAndFileAllowed_whenSending_thenReturnSendTheFile() = runTest(testDispatcher.default) {
+        // Given
+        val assetToSend = mockedLongAssetData()
+        val assetName = "some-asset.txt"
+        val inputDataPath = fakeKaliumFileSystem.providePersistentAssetPath(assetName)
+        val expectedAssetId = dummyUploadedAssetId
+        val expectedAssetSha256 = SHA256Key("some-asset-sha-256".toByteArray())
+        val conversationId = ConversationId("some-convo-id", "some-domain-id")
+        val (arrangement, sendAssetUseCase) = Arrangement(this)
+            .withStoredData(assetToSend, inputDataPath)
+            .withSuccessfulResponse(expectedAssetId, expectedAssetSha256)
+            .withObserveFileSharingStatusResult(FileSharingStatus.Value.EnabledSome(listOf("png")))
+            .withValidateAsseMimeTypeResult(true)
+            .withSelfDeleteTimer(SelfDeletionTimer.Disabled)
+            .withObserveMessageVisibility()
+            .withDeleteAssetLocally()
+            .arrange()
+
+        // When
+        val result = sendAssetUseCase.invoke(
+            conversationId = conversationId,
+            assetDataPath = inputDataPath,
+            assetDataSize = assetToSend.size.toLong(),
+            assetName = assetName,
+            assetMimeType = "image/png",
+            assetWidth = null,
+            assetHeight = null,
+            audioLengthInMs = 0
+        )
+        advanceUntilIdle()
+
+        // Then
+        assertTrue(result is ScheduleNewAssetMessageResult.Success)
+
+        verify(arrangement.validateAssetMimeTypeUseCase)
+            .function(arrangement.validateAssetMimeTypeUseCase::invoke)
+            .with(eq("image/png"), eq(listOf("png")))
+            .wasInvoked(exactly = once)
+    }
+
     private class Arrangement(val coroutineScope: CoroutineScope) {
 
         @Mock
@@ -587,6 +705,12 @@ class ScheduleNewAssetMessageUseCaseTest {
         @Mock
         private val messageRepository: MessageRepository = mock(MessageRepository::class)
 
+        @Mock
+        val validateAssetMimeTypeUseCase: ValidateAssetMimeTypeUseCase = mock(ValidateAssetMimeTypeUseCase::class)
+
+        @Mock
+        val observerFileSharingStatusUseCase: ObserveFileSharingStatusUseCase = mock(ObserveFileSharingStatusUseCase::class)
+
         val someClientId = ClientId("some-client-id")
 
         val completeStateFlow = MutableStateFlow<SlowSyncStatus>(SlowSyncStatus.Complete).asStateFlow()
@@ -596,6 +720,20 @@ class ScheduleNewAssetMessageUseCaseTest {
             withToggleReadReceiptsStatus()
         }
 
+        fun withValidateAsseMimeTypeResult(result: Boolean) = apply {
+            given(validateAssetMimeTypeUseCase)
+                .function(validateAssetMimeTypeUseCase::invoke)
+                .whenInvokedWith(any(), any())
+                .thenReturn(result)
+        }
+
+        fun withObserveFileSharingStatusResult(result: FileSharingStatus.Value) = apply {
+            given(observerFileSharingStatusUseCase)
+                .function(observerFileSharingStatusUseCase::invoke)
+                .whenInvoked()
+                .thenReturn(flowOf(FileSharingStatus(result, false)))
+        }
+
         fun withToggleReadReceiptsStatus(enabled: Boolean = false) = apply {
             given(userPropertyRepository)
                 .suspendFunction(userPropertyRepository::getReadReceiptsStatus)
@@ -785,6 +923,8 @@ class ScheduleNewAssetMessageUseCaseTest {
             userPropertyRepository,
             observeSelfDeletionTimerSettingsForConversation,
             coroutineScope,
+            observerFileSharingStatusUseCase,
+            validateAssetMimeTypeUseCase,
             testDispatcher
         )
     }