Migrate secure screen logic

5 files changed, 45 insertions, 38 deletions

diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/dialog/RedeemVoucherDialog.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/dialog/RedeemVoucherDialog.kt
index 1f86eedc64..93734e537d 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/dialog/RedeemVoucherDialog.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/dialog/RedeemVoucherDialog.kt
@@ -22,6 +22,9 @@ import androidx.compose.ui.res.stringResource
 import androidx.compose.ui.text.input.KeyboardType
 import androidx.compose.ui.tooling.preview.Devices
 import androidx.compose.ui.tooling.preview.Preview
+import androidx.compose.ui.window.DialogProperties
+import androidx.compose.ui.window.SecureFlagPolicy
+import net.mullvad.mullvadvpn.BuildConfig
 import net.mullvad.mullvadvpn.R
 import net.mullvad.mullvadvpn.compose.button.PrimaryButton
 import net.mullvad.mullvadvpn.compose.button.VariantButton
@@ -171,7 +174,12 @@ fun RedeemVoucherDialog(
         titleContentColor = MaterialTheme.colorScheme.onBackground,
         onDismissRequest = {
             onDismiss(uiState.voucherViewModelState is VoucherDialogState.Success)
-        }
+        },
+        properties =
+            DialogProperties(
+                securePolicy =
+                    if (BuildConfig.DEBUG) SecureFlagPolicy.Inherit else SecureFlagPolicy.SecureOn
+            )
     )
 }
 
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/AccountScreen.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/AccountScreen.kt
index 683c57987b..b19d432549 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/AccountScreen.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/screen/AccountScreen.kt
@@ -37,6 +37,7 @@ import net.mullvad.mullvadvpn.compose.component.MissingPolicy
 import net.mullvad.mullvadvpn.compose.component.NavigateBackIconButton
 import net.mullvad.mullvadvpn.compose.component.ScaffoldWithMediumTopBar
 import net.mullvad.mullvadvpn.compose.dialog.DeviceNameInfoDialog
+import net.mullvad.mullvadvpn.compose.util.SecureScreenWhileInView
 import net.mullvad.mullvadvpn.constant.IS_PLAY_BUILD
 import net.mullvad.mullvadvpn.lib.common.util.openAccountPageInBrowser
 import net.mullvad.mullvadvpn.lib.theme.AppTheme
@@ -74,6 +75,9 @@ fun AccountScreen(
     onLogoutClick: () -> Unit = {},
     onBackClick: () -> Unit = {}
 ) {
+    // This will enable SECURE_FLAG while this screen is visible to preview screenshot
+    SecureScreenWhileInView()
+
     val context = LocalContext.current
     val backgroundColor = MaterialTheme.colorScheme.background
     val systemUiController = rememberSystemUiController()
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/util/SecureScreen.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/util/SecureScreen.kt
new file mode 100644
index 0000000000..2680f5b4b8
--- /dev/null
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/compose/util/SecureScreen.kt
@@ -0,0 +1,30 @@
+package net.mullvad.mullvadvpn.compose.util
+
+import android.app.Activity
+import android.view.WindowManager
+import androidx.compose.runtime.Composable
+import androidx.compose.runtime.DisposableEffect
+import androidx.compose.runtime.saveable.rememberSaveable
+import androidx.compose.ui.platform.LocalContext
+import net.mullvad.mullvadvpn.BuildConfig
+
+@Composable
+fun SecureScreenWhileInView() {
+    if (BuildConfig.DEBUG) {
+        return
+    }
+    val context = LocalContext.current
+    val window = (context as Activity).window
+    val secureScreenWasEnabled = rememberSaveable {
+        window.attributes.flags and WindowManager.LayoutParams.FLAG_SECURE != 0
+    }
+
+    DisposableEffect(Unit) {
+        window.addFlags(WindowManager.LayoutParams.FLAG_SECURE)
+        onDispose {
+            if (!secureScreenWasEnabled) {
+                window.clearFlags(WindowManager.LayoutParams.FLAG_SECURE)
+            }
+        }
+    }
+}
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/MainActivity.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/MainActivity.kt
index 2ff152707d..f839472538 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/MainActivity.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/MainActivity.kt
@@ -9,13 +9,11 @@ import android.content.res.Configuration
 import android.net.VpnService
 import android.os.Bundle
 import android.util.Log
-import android.view.WindowManager
 import androidx.activity.result.ActivityResultLauncher
 import androidx.activity.result.contract.ActivityResultContracts
 import androidx.compose.runtime.collectAsState
 import androidx.compose.ui.platform.ComposeView
 import androidx.compose.ui.platform.ViewCompositionStrategy
-import androidx.fragment.app.Fragment
 import androidx.fragment.app.FragmentActivity
 import androidx.fragment.app.FragmentManager
 import androidx.lifecycle.lifecycleScope
@@ -65,8 +63,6 @@ open class MainActivity : FragmentActivity() {
             // handling the callback value.
         }
 
-    private var visibleSecureScreens = HashSet<Fragment>()
-
     private val deviceIsTv by lazy {
         val uiModeManager = getSystemService(UI_MODE_SERVICE) as UiModeManager
 
@@ -149,26 +145,6 @@ open class MainActivity : FragmentActivity() {
         super.onDestroy()
     }
 
-    fun enterSecureScreen(screen: Fragment) {
-        synchronized(this) {
-            visibleSecureScreens.add(screen)
-
-            if (!BuildConfig.DEBUG) {
-                window?.addFlags(WindowManager.LayoutParams.FLAG_SECURE)
-            }
-        }
-    }
-
-    fun leaveSecureScreen(screen: Fragment) {
-        synchronized(this) {
-            visibleSecureScreens.remove(screen)
-
-            if (!BuildConfig.DEBUG && visibleSecureScreens.isEmpty()) {
-                window?.clearFlags(WindowManager.LayoutParams.FLAG_SECURE)
-            }
-        }
-    }
-
     fun openAccount() {
         supportFragmentManager.beginTransaction().apply {
             setCustomAnimations(
@@ -307,7 +283,8 @@ open class MainActivity : FragmentActivity() {
                 .filter { it is AccountExpiry.Available }
                 .map { it.date()?.isBeforeNow }
                 .first()
-        } ?: false
+        }
+            ?: false
     }
 
     private fun openLoginView() {
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/fragment/AccountFragment.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/fragment/AccountFragment.kt
index 3bd46d9f52..efdc0783a3 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/fragment/AccountFragment.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/fragment/AccountFragment.kt
@@ -1,6 +1,5 @@
 package net.mullvad.mullvadvpn.ui.fragment
 
-import android.app.Activity
 import android.os.Bundle
 import android.view.LayoutInflater
 import android.view.View
@@ -11,7 +10,6 @@ import androidx.compose.ui.platform.ComposeView
 import net.mullvad.mullvadvpn.R
 import net.mullvad.mullvadvpn.compose.screen.AccountScreen
 import net.mullvad.mullvadvpn.lib.theme.AppTheme
-import net.mullvad.mullvadvpn.ui.extension.requireMainActivity
 import net.mullvad.mullvadvpn.viewmodel.AccountViewModel
 import org.koin.androidx.viewmodel.ext.android.viewModel
 
@@ -42,16 +40,6 @@ class AccountFragment : BaseFragment() {
         }
     }
 
-    override fun onAttach(activity: Activity) {
-        super.onAttach(activity)
-        requireMainActivity().enterSecureScreen(this)
-    }
-
-    override fun onDetach() {
-        super.onDetach()
-        requireMainActivity().leaveSecureScreen(this)
-    }
-
     override fun onEnterTransitionAnimationEnd() {
         vm.onTransitionAnimationEnd()
     }