From e1ef5f55eca96f1e9b1d80cf6d9e4d90d250b31a Mon Sep 17 00:00:00 2001
From: ekager <ekager@mozilla.com>
Date: Wed, 15 Jul 2020 16:31:51 -0400
Subject: [PATCH] For #12453 - Sets secure flags on private tab of tabs tray

---
 .../fenix/tabtray/TabTrayDialogFragment.kt    | 21 +++++++++++++++++--
 .../org/mozilla/fenix/tabtray/TabTrayView.kt  | 11 ++--------
 2 files changed, 21 insertions(+), 11 deletions(-)

diff --git a/app/src/main/java/org/mozilla/fenix/tabtray/TabTrayDialogFragment.kt b/app/src/main/java/org/mozilla/fenix/tabtray/TabTrayDialogFragment.kt
index 62f5cacf7..b7340d1d4 100644
--- a/app/src/main/java/org/mozilla/fenix/tabtray/TabTrayDialogFragment.kt
+++ b/app/src/main/java/org/mozilla/fenix/tabtray/TabTrayDialogFragment.kt
@@ -9,6 +9,7 @@ import android.os.Bundle
 import android.view.LayoutInflater
 import android.view.View
 import android.view.ViewGroup
+import android.view.WindowManager
 import androidx.appcompat.app.AppCompatDialogFragment
 import androidx.core.view.isVisible
 import androidx.core.view.updatePadding
@@ -25,6 +26,7 @@ import mozilla.components.browser.session.SessionManager
 import mozilla.components.browser.state.selector.normalTabs
 import mozilla.components.browser.state.selector.privateTabs
 import mozilla.components.browser.state.state.BrowserState
+import mozilla.components.browser.state.state.TabSessionState
 import mozilla.components.feature.tab.collections.TabCollection
 import mozilla.components.feature.tabs.TabsUseCases
 import mozilla.components.feature.tabs.tabstray.TabsFeature
@@ -37,6 +39,7 @@ import org.mozilla.fenix.components.TabCollectionStorage
 import org.mozilla.fenix.components.metrics.Event
 import org.mozilla.fenix.ext.components
 import org.mozilla.fenix.ext.requireComponents
+import org.mozilla.fenix.ext.settings
 import org.mozilla.fenix.utils.allowUndo
 
 @SuppressWarnings("TooManyFunctions", "LargeClass")
@@ -48,7 +51,7 @@ class TabTrayDialogFragment : AppCompatDialogFragment() {
 
     private val snackbarAnchor: View?
         get() = if (tabTrayView.fabView.new_tab_button.isVisible) tabTrayView.fabView.new_tab_button
-                else null
+        else null
 
     private val collectionStorageObserver = object : TabCollectionStorage.Observer {
         override fun onCollectionCreated(title: String, sessions: List<Session>) {
@@ -131,7 +134,13 @@ class TabTrayDialogFragment : AppCompatDialogFragment() {
             startingInLandscape = requireContext().resources.configuration.orientation ==
                     Configuration.ORIENTATION_LANDSCAPE,
             lifecycleScope = viewLifecycleOwner.lifecycleScope
-        ) { tabsFeature.get()?.filterTabs(it) }
+        ) { private ->
+            val filter: (TabSessionState) -> Boolean = { state -> private == state.content.private }
+
+            tabsFeature.get()?.filterTabs(filter)
+
+            setSecureFlagsIfNeeded(private)
+        }
 
         tabsFeature.set(
             TabsFeature(
@@ -171,6 +180,14 @@ class TabTrayDialogFragment : AppCompatDialogFragment() {
         }
     }
 
+    private fun setSecureFlagsIfNeeded(private: Boolean) {
+        if (private && context?.settings()?.allowScreenshotsInPrivateMode == false) {
+            dialog?.window?.addFlags(WindowManager.LayoutParams.FLAG_SECURE)
+        } else if (!(activity as HomeActivity).browsingModeManager.mode.isPrivate) {
+            dialog?.window?.clearFlags(WindowManager.LayoutParams.FLAG_SECURE)
+        }
+    }
+
     private fun showUndoSnackbarForTab(sessionId: String) {
         val sessionManager = view?.context?.components?.core?.sessionManager
         val snapshot = sessionManager
diff --git a/app/src/main/java/org/mozilla/fenix/tabtray/TabTrayView.kt b/app/src/main/java/org/mozilla/fenix/tabtray/TabTrayView.kt
index f82e5f923..b9acc4858 100644
--- a/app/src/main/java/org/mozilla/fenix/tabtray/TabTrayView.kt
+++ b/app/src/main/java/org/mozilla/fenix/tabtray/TabTrayView.kt
@@ -27,7 +27,6 @@ import mozilla.components.browser.menu.item.SimpleBrowserMenuItem
 import mozilla.components.browser.state.selector.normalTabs
 import mozilla.components.browser.state.selector.privateTabs
 import mozilla.components.browser.state.state.BrowserState
-import mozilla.components.browser.state.state.TabSessionState
 import mozilla.components.browser.tabstray.BrowserTabsTray
 import org.mozilla.fenix.R
 import org.mozilla.fenix.components.metrics.Event
@@ -44,7 +43,7 @@ class TabTrayView(
     isPrivate: Boolean,
     startingInLandscape: Boolean,
     lifecycleScope: LifecycleCoroutineScope,
-    private val filterTabs: ((TabSessionState) -> Boolean) -> Unit
+    private val filterTabs: (Boolean) -> Unit
 ) : LayoutContainer, TabLayout.OnTabSelectedListener {
     val fabView = LayoutInflater.from(container.context)
         .inflate(R.layout.component_tabstray_fab, container, true)
@@ -204,14 +203,8 @@ class TabTrayView(
     }
 
     override fun onTabSelected(tab: TabLayout.Tab?) {
-        // We need a better way to determine which tab was selected.
-        val filter: (TabSessionState) -> Boolean = when (tab?.position) {
-            1 -> { state -> state.content.private }
-            else -> { state -> !state.content.private }
-        }
-
         toggleFabText(isPrivateModeSelected)
-        filterTabs.invoke(filter)
+        filterTabs.invoke(isPrivateModeSelected)
 
         updateState(view.context.components.core.store.state)
         scrollToTab(view.context.components.core.store.state.selectedTabId)