Fix snapshots for secure windows

First, also draw system bar backgrounds when drawing a fake
snapshot. For that, refactor the drawing into a separate class so
it can be reused. Also enable fake snapshots for secure windows.

Test: com.android.server.wm.TaskSnapshotControllerTest
Test: com.android.server.wm.TaskSnapshotSurfaceTest
Test: Secure activity with resuming delay, make sure system bars
are covered when reopening app.

Bug: 35710126
Change-Id: I2f0ebc7e7acb80015780a4e882f0a472599efa30
diff --git a/services/core/java/com/android/server/wm/AppWindowContainerController.java b/services/core/java/com/android/server/wm/AppWindowContainerController.java
index c625cbe..b5e194b 100644
--- a/services/core/java/com/android/server/wm/AppWindowContainerController.java
+++ b/services/core/java/com/android/server/wm/AppWindowContainerController.java
@@ -36,7 +36,6 @@
 import android.os.Debug;
 import android.os.Handler;
 import android.os.IBinder;
-import android.os.Looper;
 import android.os.Trace;
 import android.util.Slog;
 import android.view.IApplicationToken;
@@ -319,7 +318,7 @@
                         + " token: " + mToken);
                 return;
             }
-            mContainer.setDisablePreviewSnapshots(disable);
+            mContainer.setDisablePreviewScreenshots(disable);
         }
     }
 
diff --git a/services/core/java/com/android/server/wm/AppWindowToken.java b/services/core/java/com/android/server/wm/AppWindowToken.java
index 3c2dfa5..5b398f9 100644
--- a/services/core/java/com/android/server/wm/AppWindowToken.java
+++ b/services/core/java/com/android/server/wm/AppWindowToken.java
@@ -23,6 +23,7 @@
 import static android.content.pm.ActivityInfo.SCREEN_ORIENTATION_UNSET;
 import static android.view.Display.DEFAULT_DISPLAY;
 import static android.view.WindowManager.LayoutParams.FLAG_DISMISS_KEYGUARD;
+import static android.view.WindowManager.LayoutParams.FLAG_SECURE;
 import static android.view.WindowManager.LayoutParams.FLAG_SHOW_WHEN_LOCKED;
 import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION;
 import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION_STARTING;
@@ -48,6 +49,7 @@
 import static com.android.server.wm.WindowManagerService.logWithStack;
 
 import android.annotation.NonNull;
+import android.app.Activity;
 import android.content.res.Configuration;
 import android.graphics.Rect;
 import android.os.Binder;
@@ -1521,12 +1523,24 @@
         return candidate;
     }
 
-    void setDisablePreviewSnapshots(boolean disable) {
+    /**
+     * See {@link Activity#setDisablePreviewScreenshots}.
+     */
+    void setDisablePreviewScreenshots(boolean disable) {
         mDisbalePreviewScreenshots = disable;
     }
 
-    boolean shouldDisablePreviewScreenshots() {
-        return mDisbalePreviewScreenshots;
+    /**
+     * Retrieves whether we'd like to generate a snapshot that's based solely on the theme. This is
+     * the case when preview screenshots are disabled {@link #setDisablePreviewScreenshots} or when
+     * we can't take a snapshot for other reasons, for example, if we have a secure window.
+     *
+     * @return True if we need to generate an app theme snapshot, false if we'd like to take a real
+     *         screenshot.
+     */
+    boolean shouldUseAppThemeSnapshot() {
+        return mDisbalePreviewScreenshots || forAllWindows(w -> (w.mAttrs.flags & FLAG_SECURE) != 0,
+                true /* topToBottom */);
     }
 
     @Override
diff --git a/services/core/java/com/android/server/wm/TaskSnapshotController.java b/services/core/java/com/android/server/wm/TaskSnapshotController.java
index fbb826d..b79173c 100644
--- a/services/core/java/com/android/server/wm/TaskSnapshotController.java
+++ b/services/core/java/com/android/server/wm/TaskSnapshotController.java
@@ -31,11 +31,13 @@
 import android.graphics.Rect;
 import android.os.Environment;
 import android.util.ArraySet;
+import android.view.WindowManager.LayoutParams;
 import android.view.WindowManagerPolicy.StartingSurface;
 
 import com.google.android.collect.Sets;
 
 import com.android.internal.annotations.VisibleForTesting;
+import com.android.server.wm.TaskSnapshotSurface.SystemBarBackgroundPainter;
 
 import java.io.PrintWriter;
 
@@ -206,7 +208,7 @@
         final AppWindowToken topChild = task.getTopChild();
         if (StackId.isHomeOrRecentsStack(task.mStack.mStackId)) {
             return SNAPSHOT_MODE_NONE;
-        } else if (topChild != null && topChild.shouldDisablePreviewScreenshots()) {
+        } else if (topChild != null && topChild.shouldUseAppThemeSnapshot()) {
             return SNAPSHOT_MODE_APP_THEME;
         } else {
             return SNAPSHOT_MODE_REAL;
@@ -227,6 +229,8 @@
             return null;
         }
         final int color = task.getTaskDescription().getBackgroundColor();
+        final int statusBarColor = task.getTaskDescription().getStatusBarColor();
+        final int navigationBarColor = task.getTaskDescription().getNavigationBarColor();
         final GraphicBuffer buffer = GraphicBuffer.create(mainWindow.getFrameLw().width(),
                 mainWindow.getFrameLw().height(),
                 RGBA_8888, USAGE_HW_TEXTURE | USAGE_SW_WRITE_RARELY | USAGE_SW_READ_NEVER);
@@ -235,6 +239,11 @@
         }
         final Canvas c = buffer.lockCanvas();
         c.drawColor(color);
+        final LayoutParams attrs = mainWindow.getAttrs();
+        final SystemBarBackgroundPainter decorPainter = new SystemBarBackgroundPainter(attrs.flags,
+                attrs.privateFlags, attrs.systemUiVisibility, statusBarColor, navigationBarColor);
+        decorPainter.setInsets(mainWindow.mContentInsets, mainWindow.mStableInsets);
+        decorPainter.drawDecors(c, null /* statusBarExcludeFrame */);
         buffer.unlockCanvasAndPost(c);
         return new TaskSnapshot(buffer, topChild.getConfiguration().orientation,
                 mainWindow.mStableInsets, false /* reduced */, 1.0f /* scale */);
diff --git a/services/core/java/com/android/server/wm/TaskSnapshotSurface.java b/services/core/java/com/android/server/wm/TaskSnapshotSurface.java
index c816ba3..2b9e800 100644
--- a/services/core/java/com/android/server/wm/TaskSnapshotSurface.java
+++ b/services/core/java/com/android/server/wm/TaskSnapshotSurface.java
@@ -42,8 +42,11 @@
 import static com.android.server.wm.WindowManagerDebugConfig.TAG_WITH_CLASS_NAME;
 import static com.android.server.wm.WindowManagerDebugConfig.TAG_WM;
 
+import android.annotation.Nullable;
 import android.app.ActivityManager.TaskDescription;
 import android.app.ActivityManager.TaskSnapshot;
+import android.app.ActivityThread;
+import android.content.Context;
 import android.graphics.Canvas;
 import android.graphics.GraphicBuffer;
 import android.graphics.Paint;
@@ -118,13 +121,8 @@
     private final Handler mHandler;
     private boolean mSizeMismatch;
     private final Paint mBackgroundPaint = new Paint();
-    private final Paint mStatusBarPaint = new Paint();
-    private final Paint mNavigationBarPaint = new Paint();
     private final int mStatusBarColor;
-    private final int mNavigationBarColor;
-    private final int mSysUiVis;
-    private final int mWindowFlags;
-    private final int mWindowPrivateFlags;
+    @VisibleForTesting final SystemBarBackgroundPainter mSystemBarBackgroundPainter;
 
     static TaskSnapshotSurface create(WindowManagerService service, AppWindowToken token,
             TaskSnapshot snapshot) {
@@ -224,15 +222,9 @@
         mTitle = title;
         mBackgroundPaint.setColor(backgroundColor != 0 ? backgroundColor : WHITE);
         mTaskBounds = taskBounds;
-        mSysUiVis = sysUiVis;
-        mWindowFlags = windowFlags;
-        mWindowPrivateFlags = windowPrivateFlags;
-        mStatusBarColor = DecorView.calculateStatusBarColor(windowFlags,
-                service.mContext.getColor(R.color.system_bar_background_semi_transparent),
-                statusBarColor);
-        mNavigationBarColor = navigationBarColor;
-        mStatusBarPaint.setColor(mStatusBarColor);
-        mNavigationBarPaint.setColor(navigationBarColor);
+        mSystemBarBackgroundPainter = new SystemBarBackgroundPainter(windowFlags,
+                windowPrivateFlags, sysUiVis, statusBarColor, navigationBarColor);
+        mStatusBarColor = statusBarColor;
     }
 
     @Override
@@ -258,6 +250,7 @@
         mStableInsets.set(stableInsets);
         mSizeMismatch = (mFrame.width() != mSnapshot.getSnapshot().getWidth()
                 || mFrame.height() != mSnapshot.getSnapshot().getHeight());
+        mSystemBarBackgroundPainter.setInsets(contentInsets, stableInsets);
     }
 
     private void drawSnapshot() {
@@ -346,7 +339,7 @@
 
     @VisibleForTesting
     void drawBackgroundAndBars(Canvas c, Rect frame) {
-        final int statusBarHeight = getStatusBarColorViewHeight();
+        final int statusBarHeight = mSystemBarBackgroundPainter.getStatusBarColorViewHeight();
         final boolean fillHorizontally = c.getWidth() > frame.right;
         final boolean fillVertically = c.getHeight() > frame.bottom;
         if (fillHorizontally) {
@@ -359,44 +352,7 @@
         if (fillVertically) {
             c.drawRect(0, frame.bottom, c.getWidth(), c.getHeight(), mBackgroundPaint);
         }
-        drawStatusBarBackground(c, frame, statusBarHeight);
-        drawNavigationBarBackground(c);
-    }
-
-    private int getStatusBarColorViewHeight() {
-        final boolean forceStatusBarBackground =
-                (mWindowPrivateFlags & PRIVATE_FLAG_FORCE_DRAW_STATUS_BAR_BACKGROUND) != 0;
-        if (STATUS_BAR_COLOR_VIEW_ATTRIBUTES.isVisible(
-                mSysUiVis, mStatusBarColor, mWindowFlags, forceStatusBarBackground)) {
-            return getColorViewTopInset(mStableInsets.top, mContentInsets.top);
-        } else {
-            return 0;
-        }
-    }
-
-    private boolean isNavigationBarColorViewVisible() {
-        return NAVIGATION_BAR_COLOR_VIEW_ATTRIBUTES.isVisible(
-                mSysUiVis, mNavigationBarColor, mWindowFlags, false /* force */);
-    }
-
-    @VisibleForTesting
-    void drawStatusBarBackground(Canvas c, Rect frame, int statusBarHeight) {
-        if (statusBarHeight > 0 && c.getWidth() > frame.right) {
-            final int rightInset = DecorView.getColorViewRightInset(mStableInsets.right,
-                    mContentInsets.right);
-            c.drawRect(frame.right, 0, c.getWidth() - rightInset, statusBarHeight, mStatusBarPaint);
-        }
-    }
-
-    @VisibleForTesting
-    void drawNavigationBarBackground(Canvas c) {
-        final Rect navigationBarRect = new Rect();
-        getNavigationBarRect(c.getWidth(), c.getHeight(), mStableInsets, mContentInsets,
-                navigationBarRect);
-        final boolean visible = isNavigationBarColorViewVisible();
-        if (visible && !navigationBarRect.isEmpty()) {
-            c.drawRect(navigationBarRect, mNavigationBarPaint);
-        }
+        mSystemBarBackgroundPainter.drawDecors(c, frame);
     }
 
     private void reportDrawn() {
@@ -450,4 +406,84 @@
             }
         }
     }
+
+    /**
+     * Helper class to draw the background of the system bars in regions the task snapshot isn't
+     * filling the window.
+     */
+    static class SystemBarBackgroundPainter {
+
+        private final Rect mContentInsets = new Rect();
+        private final Rect mStableInsets = new Rect();
+        private final Paint mStatusBarPaint = new Paint();
+        private final Paint mNavigationBarPaint = new Paint();
+        private final int mStatusBarColor;
+        private final int mNavigationBarColor;
+        private final int mWindowFlags;
+        private final int mWindowPrivateFlags;
+        private final int mSysUiVis;
+
+        SystemBarBackgroundPainter( int windowFlags, int windowPrivateFlags, int sysUiVis,
+                int statusBarColor, int navigationBarColor) {
+            mWindowFlags = windowFlags;
+            mWindowPrivateFlags = windowPrivateFlags;
+            mSysUiVis = sysUiVis;
+            final Context context = ActivityThread.currentActivityThread().getSystemUiContext();
+            mStatusBarColor = DecorView.calculateStatusBarColor(windowFlags,
+                    context.getColor(R.color.system_bar_background_semi_transparent),
+                    statusBarColor);
+            mNavigationBarColor = navigationBarColor;
+            mStatusBarPaint.setColor(mStatusBarColor);
+            mNavigationBarPaint.setColor(navigationBarColor);
+        }
+
+        void setInsets(Rect contentInsets, Rect stableInsets) {
+            mContentInsets.set(contentInsets);
+            mStableInsets.set(stableInsets);
+        }
+
+        int getStatusBarColorViewHeight() {
+            final boolean forceStatusBarBackground =
+                    (mWindowPrivateFlags & PRIVATE_FLAG_FORCE_DRAW_STATUS_BAR_BACKGROUND) != 0;
+            if (STATUS_BAR_COLOR_VIEW_ATTRIBUTES.isVisible(
+                    mSysUiVis, mStatusBarColor, mWindowFlags, forceStatusBarBackground)) {
+                return getColorViewTopInset(mStableInsets.top, mContentInsets.top);
+            } else {
+                return 0;
+            }
+        }
+
+        private boolean isNavigationBarColorViewVisible() {
+            return NAVIGATION_BAR_COLOR_VIEW_ATTRIBUTES.isVisible(
+                    mSysUiVis, mNavigationBarColor, mWindowFlags, false /* force */);
+        }
+
+        void drawDecors(Canvas c, @Nullable Rect alreadyDrawnFrame) {
+            drawStatusBarBackground(c, alreadyDrawnFrame, getStatusBarColorViewHeight());
+            drawNavigationBarBackground(c);
+        }
+
+        @VisibleForTesting
+        void drawStatusBarBackground(Canvas c, @Nullable Rect alreadyDrawnFrame,
+                int statusBarHeight) {
+            if (statusBarHeight > 0
+                    && (alreadyDrawnFrame == null || c.getWidth() > alreadyDrawnFrame.right)) {
+                final int rightInset = DecorView.getColorViewRightInset(mStableInsets.right,
+                        mContentInsets.right);
+                final int left = alreadyDrawnFrame != null ? alreadyDrawnFrame.right : 0;
+                c.drawRect(left, 0, c.getWidth() - rightInset, statusBarHeight, mStatusBarPaint);
+            }
+        }
+
+        @VisibleForTesting
+        void drawNavigationBarBackground(Canvas c) {
+            final Rect navigationBarRect = new Rect();
+            getNavigationBarRect(c.getWidth(), c.getHeight(), mStableInsets, mContentInsets,
+                    navigationBarRect);
+            final boolean visible = isNavigationBarColorViewVisible();
+            if (visible && !navigationBarRect.isEmpty()) {
+                c.drawRect(navigationBarRect, mNavigationBarPaint);
+            }
+        }
+    }
 }
diff --git a/services/tests/servicestests/src/com/android/server/wm/TaskSnapshotControllerTest.java b/services/tests/servicestests/src/com/android/server/wm/TaskSnapshotControllerTest.java
index 2b4d9fb..f253632 100644
--- a/services/tests/servicestests/src/com/android/server/wm/TaskSnapshotControllerTest.java
+++ b/services/tests/servicestests/src/com/android/server/wm/TaskSnapshotControllerTest.java
@@ -17,6 +17,7 @@
 package com.android.server.wm;
 
 import static android.view.WindowManager.LayoutParams.FIRST_APPLICATION_WINDOW;
+import static android.view.WindowManager.LayoutParams.FLAG_SECURE;
 import static com.android.server.wm.AppTransition.TRANSIT_UNSET;
 import static com.android.server.wm.TaskSnapshotController.*;
 import static junit.framework.Assert.assertEquals;
@@ -76,12 +77,19 @@
     public void testGetSnapshotMode() throws Exception {
         final WindowState disabledWindow = createWindow(null,
                 FIRST_APPLICATION_WINDOW, mDisplayContent, "disabledWindow");
-        disabledWindow.mAppToken.setDisablePreviewSnapshots(true);
+        disabledWindow.mAppToken.setDisablePreviewScreenshots(true);
         assertEquals(SNAPSHOT_MODE_APP_THEME,
                 sWm.mTaskSnapshotController.getSnapshotMode(disabledWindow.getTask()));
+
         final WindowState normalWindow = createWindow(null,
                 FIRST_APPLICATION_WINDOW, mDisplayContent, "normalWindow");
         assertEquals(SNAPSHOT_MODE_REAL,
                 sWm.mTaskSnapshotController.getSnapshotMode(normalWindow.getTask()));
+
+        final WindowState secureWindow = createWindow(null,
+                FIRST_APPLICATION_WINDOW, mDisplayContent, "secureWindow");
+        secureWindow.mAttrs.flags |= FLAG_SECURE;
+        assertEquals(SNAPSHOT_MODE_APP_THEME,
+                sWm.mTaskSnapshotController.getSnapshotMode(secureWindow.getTask()));
     }
 }
diff --git a/services/tests/servicestests/src/com/android/server/wm/TaskSnapshotSurfaceTest.java b/services/tests/servicestests/src/com/android/server/wm/TaskSnapshotSurfaceTest.java
index 717ddf2..e2868d7 100644
--- a/services/tests/servicestests/src/com/android/server/wm/TaskSnapshotSurfaceTest.java
+++ b/services/tests/servicestests/src/com/android/server/wm/TaskSnapshotSurfaceTest.java
@@ -171,11 +171,25 @@
         final Canvas mockCanvas = mock(Canvas.class);
         when(mockCanvas.getWidth()).thenReturn(100);
         when(mockCanvas.getHeight()).thenReturn(100);
-        mSurface.drawStatusBarBackground(mockCanvas, new Rect(0, 0, 50, 100), 10);
+        mSurface.mSystemBarBackgroundPainter.drawStatusBarBackground(
+                mockCanvas, new Rect(0, 0, 50, 100), 10);
         verify(mockCanvas).drawRect(eq(50.0f), eq(0.0f), eq(90.0f), eq(10.0f), any());
     }
 
     @Test
+    public void testDrawStatusBarBackground_nullFrame() {
+        setupSurface(100, 100);
+        final Rect insets = new Rect(0, 10, 10, 0);
+        mSurface.setFrames(new Rect(0, 0, 100, 100), insets, insets);
+        final Canvas mockCanvas = mock(Canvas.class);
+        when(mockCanvas.getWidth()).thenReturn(100);
+        when(mockCanvas.getHeight()).thenReturn(100);
+        mSurface.mSystemBarBackgroundPainter.drawStatusBarBackground(
+                mockCanvas, null, 10);
+        verify(mockCanvas).drawRect(eq(0.0f), eq(0.0f), eq(90.0f), eq(10.0f), any());
+    }
+
+    @Test
     public void testDrawStatusBarBackground_nope() {
         setupSurface(100, 100);
         final Rect insets = new Rect(0, 10, 10, 0);
@@ -183,7 +197,8 @@
         final Canvas mockCanvas = mock(Canvas.class);
         when(mockCanvas.getWidth()).thenReturn(100);
         when(mockCanvas.getHeight()).thenReturn(100);
-        mSurface.drawStatusBarBackground(mockCanvas, new Rect(0, 0, 100, 100), 10);
+        mSurface.mSystemBarBackgroundPainter.drawStatusBarBackground(
+                mockCanvas, new Rect(0, 0, 100, 100), 10);
         verify(mockCanvas, never()).drawRect(anyInt(), anyInt(), anyInt(), anyInt(), any());
     }
 
@@ -196,7 +211,7 @@
         final Canvas mockCanvas = mock(Canvas.class);
         when(mockCanvas.getWidth()).thenReturn(100);
         when(mockCanvas.getHeight()).thenReturn(100);
-        mSurface.drawNavigationBarBackground(mockCanvas);
+        mSurface.mSystemBarBackgroundPainter.drawNavigationBarBackground(mockCanvas);
         verify(mockCanvas).drawRect(eq(new Rect(0, 90, 100, 100)), any());
     }
 
@@ -209,7 +224,7 @@
         final Canvas mockCanvas = mock(Canvas.class);
         when(mockCanvas.getWidth()).thenReturn(100);
         when(mockCanvas.getHeight()).thenReturn(100);
-        mSurface.drawNavigationBarBackground(mockCanvas);
+        mSurface.mSystemBarBackgroundPainter.drawNavigationBarBackground(mockCanvas);
         verify(mockCanvas).drawRect(eq(new Rect(0, 0, 10, 100)), any());
     }
 
@@ -222,7 +237,7 @@
         final Canvas mockCanvas = mock(Canvas.class);
         when(mockCanvas.getWidth()).thenReturn(100);
         when(mockCanvas.getHeight()).thenReturn(100);
-        mSurface.drawNavigationBarBackground(mockCanvas);
+        mSurface.mSystemBarBackgroundPainter.drawNavigationBarBackground(mockCanvas);
         verify(mockCanvas).drawRect(eq(new Rect(90, 0, 100, 100)), any());
     }
 }