fix: leaking UI context GetE2EICertificateUseCase (WPB-6648) (#2717)

Co-authored-by: Alexandre Ferris <[email protected]>

app/src/main/kotlin/com/wire/android/di/GetE2EICertificateUseCaseProvider.kt

@@ -0,0 +1,48 @@
+/*
+ * Wire
+ * Copyright (C) 2024 Wire Swiss GmbH
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program. If not, see http://www.gnu.org/licenses/.
+ */
+package com.wire.android.di
+
+import android.content.Context
+import com.wire.android.feature.e2ei.GetE2EICertificateUseCase
+import com.wire.android.util.dispatchers.DispatcherProvider
+import com.wire.kalium.logic.CoreLogic
+import com.wire.kalium.logic.data.user.UserId
+import dagger.assisted.Assisted
+import dagger.assisted.AssistedFactory
+import dagger.assisted.AssistedInject
+import dagger.hilt.android.qualifiers.ApplicationContext
+
+class GetE2EICertificateUseCaseProvider @AssistedInject constructor(
+    @KaliumCoreLogic private val coreLogic: CoreLogic,
+    @ApplicationContext private val applicationContext: Context,
+    @Assisted private val userId: UserId,
+    @Assisted private val dispatcherProvider: DispatcherProvider
+) {
+
+    val useCase: GetE2EICertificateUseCase
+        get() = GetE2EICertificateUseCase(
+            enrollE2EI = coreLogic.getSessionScope(userId).enrollE2EI,
+            applicationContext = applicationContext,
+            dispatcherProvider = dispatcherProvider
+        )
+
+    @AssistedFactory
+    interface Factory {
+        fun create(userId: UserId, dispatcherProvider: DispatcherProvider): GetE2EICertificateUseCaseProvider
+    }
+}

app/src/main/kotlin/com/wire/android/feature/e2ei/GetE2EICertificateUseCase.kt

@@ -26,13 +26,15 @@ import com.wire.kalium.logic.feature.e2ei.usecase.E2EIEnrollmentResult
 import com.wire.kalium.logic.feature.e2ei.usecase.EnrollE2EIUseCase
 import com.wire.kalium.logic.functional.Either
 import com.wire.kalium.logic.functional.fold
+import dagger.hilt.android.qualifiers.ApplicationContext
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.SupervisorJob
 import kotlinx.coroutines.launch
 import javax.inject.Inject
 
 class GetE2EICertificateUseCase @Inject constructor(
     private val enrollE2EI: EnrollE2EIUseCase,
+    @ApplicationContext private val applicationContext: Context,
     val dispatcherProvider: DispatcherProvider
 ) {
 
@@ -41,7 +43,6 @@ class GetE2EICertificateUseCase @Inject constructor(
     lateinit var enrollmentResultHandler: (Either<E2EIFailure, E2EIEnrollmentResult>) -> Unit
 
     operator fun invoke(
-        context: Context,
         isNewClient: Boolean,
         enrollmentResultHandler: (Either<CoreFailure, E2EIEnrollmentResult>) -> Unit
     ) {
@@ -52,8 +53,8 @@ class GetE2EICertificateUseCase @Inject constructor(
             }, {
                 if (it is E2EIEnrollmentResult.Initialized) {
                     initialEnrollmentResult = it
-                    OAuthUseCase(context, it.target, it.oAuthClaims, it.oAuthState).launch(
-                        context.getActivity()!!.activityResultRegistry,
+                    OAuthUseCase(applicationContext, it.target, it.oAuthClaims, it.oAuthState).launch(
+                        applicationContext.getActivity()!!.activityResultRegistry,
                         ::oAuthResultHandler
                     )
                 } else enrollmentResultHandler(Either.Right(it))

app/src/main/kotlin/com/wire/android/ui/WireActivity.kt

@@ -372,7 +372,7 @@ class WireActivity : AppCompatActivity() {
                     E2EIRequiredDialog(
                         e2EIRequired = e2EIRequired,
                         isE2EILoading = isE2EILoading,
-                        getCertificate = { featureFlagNotificationViewModel.getE2EICertificate(it, context) },
+                        getCertificate = { featureFlagNotificationViewModel.getE2EICertificate(it) },
                         snoozeDialog = featureFlagNotificationViewModel::snoozeE2EIdRequiredDialog
                     )
                 }
@@ -387,7 +387,7 @@ class WireActivity : AppCompatActivity() {
                 e2EIResult?.let {
                     E2EIResultDialog(
                         result = e2EIResult,
-                        updateCertificate = { featureFlagNotificationViewModel.getE2EICertificate(it, context) },
+                        updateCertificate = { featureFlagNotificationViewModel.getE2EICertificate(it) },
                         snoozeDialog = featureFlagNotificationViewModel::snoozeE2EIdRequiredDialog,
                         openCertificateDetails = { navigate(NavigationCommand(E2eiCertificateDetailsScreenDestination(it))) },
                         dismissSuccessDialog = featureFlagNotificationViewModel::dismissSuccessE2EIdDialog,

app/src/main/kotlin/com/wire/android/ui/debug/DebugDataOptions.kt

@@ -29,7 +29,6 @@ import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.setValue
 import androidx.compose.ui.Modifier
-import androidx.compose.ui.platform.LocalContext
 import androidx.compose.ui.res.stringResource
 import androidx.hilt.navigation.compose.hiltViewModel
 import androidx.lifecycle.ViewModel
@@ -127,8 +126,8 @@ class DebugDataOptionsViewModel
         }
     }
 
-    fun enrollE2EICertificate(context: Context) {
-        e2eiCertificateUseCase(context, false) { result ->
+    fun enrollE2EICertificate() {
+        e2eiCertificateUseCase(false) { result ->
             result.fold({
                 state = state.copy(
                     certificate = (it as E2EIFailure.FailedOAuth).reason, showCertificate = true
@@ -250,7 +249,7 @@ fun DebugDataOptionsContent(
     onRestartSlowSyncForRecovery: () -> Unit,
     onForceUpdateApiVersions: () -> Unit,
     onManualMigrationPressed: () -> Unit,
-    enrollE2EICertificate: (Context) -> Unit,
+    enrollE2EICertificate: () -> Unit,
     dismissCertificateDialog: () -> Unit
 ) {
     Column {
@@ -352,9 +351,8 @@ fun DebugDataOptionsContent(
 
 @Composable
 private fun GetE2EICertificateSwitch(
-    enrollE2EI: (context: Context) -> Unit
+    enrollE2EI: () -> Unit
 ) {
-    val context = LocalContext.current
     Column {
         FolderHeader(stringResource(R.string.debug_settings_e2ei_enrollment_title))
         RowItemTemplate(modifier = Modifier.wrapContentWidth(),
@@ -369,7 +367,7 @@ private fun GetE2EICertificateSwitch(
             actions = {
                 WirePrimaryButton(
                     onClick = {
-                        enrollE2EI(context)
+                        enrollE2EI()
                     },
                     text = stringResource(R.string.label_get_e2ei_cetificate),
                     fillMaxWidth = false

app/src/main/kotlin/com/wire/android/ui/e2eiEnrollment/E2EIEnrollmentScreen.kt

@@ -29,7 +29,6 @@ import androidx.compose.material3.MaterialTheme
 import androidx.compose.runtime.Composable
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
-import androidx.compose.ui.platform.LocalContext
 import androidx.compose.ui.platform.LocalUriHandler
 import androidx.compose.ui.res.stringResource
 import androidx.compose.ui.text.SpanStyle
@@ -76,7 +75,6 @@ fun E2EIEnrollmentScreen(
     viewModel: E2EIEnrollmentViewModel = hiltViewModel(),
 ) {
     val state = viewModel.state
-    val context = LocalContext.current
 
     E2EIEnrollmentScreenContent(
         state = state,
@@ -85,7 +83,7 @@ fun E2EIEnrollmentScreen(
             viewModel.finalizeMLSClient()
         },
         dismissErrorDialog = viewModel::dismissErrorDialog,
-        enrollE2EICertificate = { viewModel.enrollE2EICertificate(context) },
+        enrollE2EICertificate = { viewModel.enrollE2EICertificate() },
         openCertificateDetails = {
             navigator.navigate(NavigationCommand(E2eiCertificateDetailsScreenDestination(state.certificate)))
         },

app/src/main/kotlin/com/wire/android/ui/e2eiEnrollment/E2EIEnrollmentViewModel.kt

@@ -17,7 +17,6 @@
  */
 package com.wire.android.ui.e2eiEnrollment
 
-import android.content.Context
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.setValue
@@ -94,9 +93,9 @@ class E2EIEnrollmentViewModel @Inject constructor(
             }
         }
     }
-    fun enrollE2EICertificate(context: Context) {
+    fun enrollE2EICertificate() {
         state = state.copy(isLoading = true)
-        e2eiCertificateUseCase(context, true) { result ->
+        e2eiCertificateUseCase(true) { result ->
             result.fold({
                 state = state.copy(
                     isLoading = false,

app/src/main/kotlin/com/wire/android/ui/home/sync/FeatureFlagNotificationViewModel.kt

@@ -18,18 +18,17 @@
 
 package com.wire.android.ui.home.sync
 
-import android.content.Context
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.setValue
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.viewModelScope
 import com.wire.android.appLogger
 import com.wire.android.datastore.GlobalDataStore
+import com.wire.android.di.GetE2EICertificateUseCaseProvider
 import com.wire.android.di.KaliumCoreLogic
 import com.wire.android.feature.AppLockSource
 import com.wire.android.feature.DisableAppLockUseCase
-import com.wire.android.feature.e2ei.GetE2EICertificateUseCase
 import com.wire.android.ui.home.FeatureFlagState
 import com.wire.android.ui.home.conversations.selfdeletion.SelfDeletionMapper.toSelfDeletionDuration
 import com.wire.android.ui.home.messagecomposer.SelfDeletionDuration
@@ -59,6 +58,7 @@ class FeatureFlagNotificationViewModel @Inject constructor(
     private val currentSessionFlow: CurrentSessionFlowUseCase,
     private val globalDataStore: GlobalDataStore,
     private val disableAppLockUseCase: DisableAppLockUseCase,
+    private val getE2EICertificateUseCaseProvider: GetE2EICertificateUseCaseProvider.Factory,
     private val dispatcherProvider: DispatcherProvider
 ) : ViewModel() {
 
@@ -288,35 +288,39 @@ class FeatureFlagNotificationViewModel @Inject constructor(
 
     fun isUserAppLockSet() = globalDataStore.isAppLockPasscodeSet()
 
-    fun getE2EICertificate(e2eiRequired: FeatureFlagState.E2EIRequired, context: Context) {
+    fun getE2EICertificate(e2eiRequired: FeatureFlagState.E2EIRequired) {
         featureFlagState = featureFlagState.copy(isE2EILoading = true)
         currentUserId?.let { userId ->
-            GetE2EICertificateUseCase(coreLogic.getSessionScope(userId).enrollE2EI, dispatcherProvider).invoke(
-                context,
-                isNewClient = false
-            ) { result ->
-                result.fold({
-                    featureFlagState = featureFlagState.copy(
-                        isE2EILoading = false,
-                        e2EIRequired = null,
-                        e2EIResult = FeatureFlagState.E2EIResult.Failure(e2eiRequired)
-                    )
-                }, {
-                    if (it is E2EIEnrollmentResult.Finalized) {
-                        featureFlagState = featureFlagState.copy(
-                            isE2EILoading = false,
-                            e2EIRequired = null,
-                            e2EIResult = FeatureFlagState.E2EIResult.Success(it.certificate)
-                        )
-                    } else if (it is E2EIEnrollmentResult.Failed) {
+            getE2EICertificateUseCaseProvider.create(
+                userId = userId,
+                dispatcherProvider = dispatcherProvider
+            )
+                .useCase
+                .invoke(
+                    isNewClient = false
+                ) { result ->
+                    result.fold({
                         featureFlagState = featureFlagState.copy(
                             isE2EILoading = false,
                             e2EIRequired = null,
                             e2EIResult = FeatureFlagState.E2EIResult.Failure(e2eiRequired)
                         )
-                    }
-                })
-            }
+                    }, {
+                        if (it is E2EIEnrollmentResult.Finalized) {
+                            featureFlagState = featureFlagState.copy(
+                                isE2EILoading = false,
+                                e2EIRequired = null,
+                                e2EIResult = FeatureFlagState.E2EIResult.Success(it.certificate)
+                            )
+                        } else if (it is E2EIEnrollmentResult.Failed) {
+                            featureFlagState = featureFlagState.copy(
+                                isE2EILoading = false,
+                                e2EIRequired = null,
+                                e2EIResult = FeatureFlagState.E2EIResult.Failure(e2eiRequired)
+                            )
+                        }
+                    })
+                }
         }
     }
 

app/src/main/kotlin/com/wire/android/ui/settings/devices/DeviceDetailsScreen.kt

@@ -17,7 +17,6 @@
  */
 package com.wire.android.ui.settings.devices
 
-import android.content.Context
 import androidx.annotation.StringRes
 import androidx.compose.foundation.background
 import androidx.compose.foundation.layout.Box
@@ -132,7 +131,7 @@ fun DeviceDetailsContent(
     onRemoveConfirm: () -> Unit = {},
     onDialogDismiss: () -> Unit = {},
     onErrorDialogDismiss: () -> Unit = {},
-    enrollE2eiCertificate: (Context) -> Unit = {},
+    enrollE2eiCertificate: () -> Unit = {},
     onUpdateClientVerification: (Boolean) -> Unit = {},
     onEnrollE2EIErrorDismiss: () -> Unit = {},
     onEnrollE2EISuccessDismiss: () -> Unit = {}
@@ -173,7 +172,6 @@ fun DeviceDetailsContent(
             }
         }
     ) { internalPadding ->
-        val context = LocalContext.current
         LazyColumn(
             modifier = Modifier
                 .fillMaxSize()
@@ -195,7 +193,7 @@ fun DeviceDetailsContent(
                         certificate = state.e2eiCertificate,
                         isCurrentDevice = state.isCurrentDevice,
                         isLoadingCertificate = state.isLoadingCertificate,
-                        enrollE2eiCertificate = { enrollE2eiCertificate(context) },
+                        enrollE2eiCertificate = { enrollE2eiCertificate() },
                         showCertificate = onNavigateToE2eiCertificateDetailsScreen
                     )
                     Divider(color = colorsScheme().background)
@@ -277,7 +275,7 @@ fun DeviceDetailsContent(
         if (state.isE2EICertificateEnrollError) {
             E2EIErrorWithDismissDialog(
                 isE2EILoading = state.isLoadingCertificate,
-                updateCertificate = { enrollE2eiCertificate(context) },
+                updateCertificate = { enrollE2eiCertificate() },
                 onDismiss = onEnrollE2EIErrorDismiss
             )
         }

app/src/main/kotlin/com/wire/android/ui/settings/devices/DeviceDetailsViewModel.kt

@@ -17,7 +17,6 @@
  */
 package com.wire.android.ui.settings.devices
 
-import android.content.Context
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.setValue
@@ -128,9 +127,9 @@ class DeviceDetailsViewModel @Inject constructor(
         }
     }
 
-    fun enrollE2eiCertificate(context: Context) {
+    fun enrollE2eiCertificate() {
         state = state.copy(isLoadingCertificate = true)
-        enrolE2EICertificateUseCase(context, false) { result ->
+        enrolE2EICertificateUseCase(false) { result ->
             result.fold({
                 state = state.copy(
                     isLoadingCertificate = false,