Exempt-From-Owner-Approval: Report non-visual Context misuse

Make obtaining a visual service from non-visual Context instance
report a strict mode violation and print the stacktrace.

Make calling getDisplay() throw an exception if called on an instance
that is not associated with a display. For existing usages introduce
a new internal method that does not perform the verification until
the usages are properly fixed.

Bug: 128338354
Test: StrictModeTest#testIncorrectContextUse_GetSystemService
Test: StrictModeTest#testIncorrectContextUse_GetDisplay
Change-Id: Id25d590eca6e10066e55d7ed6436d3bc9e433beb

26 files changed, 210 insertions, 44 deletions

diff --git a/api/test-current.txt b/api/test-current.txt
index 957794cac057..57606333bb21 100644
--- a/api/test-current.txt
+++ b/api/test-current.txt
@@ -2529,6 +2529,11 @@ package android.os {
     method public void log(android.os.StrictMode.ViolationInfo);
   }
 
+  public static final class StrictMode.VmPolicy.Builder {
+    method @NonNull public android.os.StrictMode.VmPolicy.Builder detectIncorrectContextUse();
+    method @NonNull public android.os.StrictMode.VmPolicy.Builder permitIncorrectContextUse();
+  }
+
   public class SystemConfigManager {
     method @NonNull @RequiresPermission("android.permission.READ_CARRIER_APP_INFO") public java.util.Set<java.lang.String> getDisabledUntilUsedPreinstalledCarrierApps();
     method @NonNull @RequiresPermission("android.permission.READ_CARRIER_APP_INFO") public java.util.Map<java.lang.String,java.util.List<java.lang.String>> getDisabledUntilUsedPreinstalledCarrierAssociatedApps();
diff --git a/core/java/android/app/ContextImpl.java b/core/java/android/app/ContextImpl.java
index 6b5bfda92cd0..8df26cbbb4f4 100644
--- a/core/java/android/app/ContextImpl.java
+++ b/core/java/android/app/ContextImpl.java
@@ -16,6 +16,9 @@
 
 package android.app;
 
+import static android.content.pm.PackageManager.PERMISSION_GRANTED;
+import static android.os.StrictMode.vmIncorrectContextUseEnabled;
+
 import android.annotation.IntDef;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
@@ -64,6 +67,7 @@ import android.os.IBinder;
 import android.os.Looper;
 import android.os.Process;
 import android.os.RemoteException;
+import android.os.StrictMode;
 import android.os.Trace;
 import android.os.UserHandle;
 import android.os.UserManager;
@@ -247,6 +251,7 @@ class ContextImpl extends Context {
 
     private boolean mIsSystemOrSystemUiContext;
     private boolean mIsUiContext;
+    private boolean mIsAssociatedWithDisplay;
 
     @GuardedBy("mSync")
     private File mDatabasesDir;
@@ -1891,9 +1896,20 @@ class ContextImpl extends Context {
 
     @Override
     public Object getSystemService(String name) {
-        if (isUiComponent(name) && !isUiContext()) {
-            Log.w(TAG, name + " should be accessed from Activity or other visual Context");
+        // Check incorrect Context usage.
+        if (isUiComponent(name) && !isUiContext() && vmIncorrectContextUseEnabled()) {
+            final String errorMessage = "Tried to access visual service " + name
+                    + " from a non-visual Context.";
+            final String message = "Visual services, such as WindowManager, WallpaperService or "
+                    + "LayoutInflater should be accessed from Activity or other visual Context. "
+                    + "Use an Activity or a Context created with "
+                    + "Context#createWindowContext(int, Bundle), which are adjusted to the "
+                    + "configuration and visual bounds of an area on screen.";
+            final Exception exception = new IllegalAccessException(errorMessage);
+            StrictMode.onIncorrectContextUsed(message, exception);
+            Log.e(TAG, errorMessage + message, exception);
         }
+
         return SystemServiceRegistry.getSystemService(this, name);
     }
 
@@ -1902,8 +1918,17 @@ class ContextImpl extends Context {
         return SystemServiceRegistry.getSystemServiceName(serviceClass);
     }
 
-    boolean isUiContext() {
-        return mIsSystemOrSystemUiContext || mIsUiContext;
+    private boolean isUiContext() {
+        return mIsSystemOrSystemUiContext || mIsUiContext || isSystemOrSystemUI();
+    }
+
+    /**
+     * Temporary workaround to permit incorrect usages of Context by SystemUI.
+     * TODO(b/149790106): Fix usages and remove.
+     */
+    private boolean isSystemOrSystemUI() {
+        return ActivityThread.isSystem() || checkPermission("android.permission.STATUS_BAR_SERVICE",
+                Binder.getCallingPid(), Binder.getCallingUid()) == PERMISSION_GRANTED;
     }
 
     private static boolean isUiComponent(String name) {
@@ -1925,7 +1950,7 @@ class ContextImpl extends Context {
             final int appId = UserHandle.getAppId(uid);
             if (appId == Process.ROOT_UID || appId == Process.SYSTEM_UID) {
                 Slog.w(TAG, "Missing ActivityManager; assuming " + uid + " holds " + permission);
-                return PackageManager.PERMISSION_GRANTED;
+                return PERMISSION_GRANTED;
             }
             Slog.w(TAG, "Missing ActivityManager; assuming " + uid + " does not hold "
                     + permission);
@@ -1989,7 +2014,7 @@ class ContextImpl extends Context {
     private void enforce(
             String permission, int resultOfCheck,
             boolean selfToo, int uid, String message) {
-        if (resultOfCheck != PackageManager.PERMISSION_GRANTED) {
+        if (resultOfCheck != PERMISSION_GRANTED) {
             throw new SecurityException(
                     (message != null ? (message + ": ") : "") +
                     (selfToo
@@ -2116,15 +2141,15 @@ class ContextImpl extends Context {
         if ((modeFlags&Intent.FLAG_GRANT_READ_URI_PERMISSION) != 0) {
             if (readPermission == null
                     || checkPermission(readPermission, pid, uid)
-                    == PackageManager.PERMISSION_GRANTED) {
-                return PackageManager.PERMISSION_GRANTED;
+                    == PERMISSION_GRANTED) {
+                return PERMISSION_GRANTED;
             }
         }
         if ((modeFlags&Intent.FLAG_GRANT_WRITE_URI_PERMISSION) != 0) {
             if (writePermission == null
                     || checkPermission(writePermission, pid, uid)
-                    == PackageManager.PERMISSION_GRANTED) {
-                return PackageManager.PERMISSION_GRANTED;
+                    == PERMISSION_GRANTED) {
+                return PERMISSION_GRANTED;
             }
         }
         return uri != null ? checkUriPermission(uri, pid, uid, modeFlags)
@@ -2157,7 +2182,7 @@ class ContextImpl extends Context {
     private void enforceForUri(
             int modeFlags, int resultOfCheck, boolean selfToo,
             int uid, Uri uri, String message) {
-        if (resultOfCheck != PackageManager.PERMISSION_GRANTED) {
+        if (resultOfCheck != PERMISSION_GRANTED) {
             throw new SecurityException(
                     (message != null ? (message + ": ") : "") +
                     (selfToo
@@ -2373,6 +2398,7 @@ class ContextImpl extends Context {
                 null, getDisplayAdjustments(displayId).getCompatibilityInfo(),
                 mResources.getLoaders()));
         context.mDisplay = display;
+        context.mIsAssociatedWithDisplay = true;
         return context;
     }
 
@@ -2390,6 +2416,7 @@ class ContextImpl extends Context {
         ContextImpl context = new ContextImpl(this, mMainThread, mPackageInfo, mFeatureId,
                 mSplitName, token, mUser, mFlags, mClassLoader, null);
         context.mIsUiContext = true;
+        context.mIsAssociatedWithDisplay = true;
         return context;
     }
 
@@ -2440,6 +2467,19 @@ class ContextImpl extends Context {
 
     @Override
     public Display getDisplay() {
+        if (!mIsSystemOrSystemUiContext && !mIsAssociatedWithDisplay && !isSystemOrSystemUI()) {
+            throw new UnsupportedOperationException("Tried to obtain display from a Context not "
+                    + "associated with  one. Only visual Contexts (such as Activity or one created "
+                    + "with Context#createWindowContext) or ones created with "
+                    + "Context#createDisplayContext are associated with displays. Other types of "
+                    + "Contexts are typically related to background entities and may return an "
+                    + "arbitrary display.");
+        }
+        return getDisplayNoVerify();
+    }
+
+    @Override
+    public Display getDisplayNoVerify() {
         if (mDisplay == null) {
             return mResourcesManager.getAdjustedDisplay(Display.DEFAULT_DISPLAY,
                     mResources);
@@ -2450,13 +2490,14 @@ class ContextImpl extends Context {
 
     @Override
     public int getDisplayId() {
-        final Display display = getDisplay();
+        final Display display = getDisplayNoVerify();
         return display != null ? display.getDisplayId() : Display.DEFAULT_DISPLAY;
     }
 
     @Override
     public void updateDisplay(int displayId) {
         mDisplay = mResourcesManager.getAdjustedDisplay(displayId, mResources);
+        mIsAssociatedWithDisplay = true;
     }
 
     @Override
@@ -2630,6 +2671,7 @@ class ContextImpl extends Context {
         ContextImpl context = new ContextImpl(null, mainThread, packageInfo, null,
                 activityInfo.splitName, activityToken, null, 0, classLoader, null);
         context.mIsUiContext = true;
+        context.mIsAssociatedWithDisplay = true;
 
         // Clamp display ID to DEFAULT_DISPLAY if it is INVALID_DISPLAY.
         displayId = (displayId != Display.INVALID_DISPLAY) ? displayId : Display.DEFAULT_DISPLAY;
diff --git a/core/java/android/app/TaskEmbedder.java b/core/java/android/app/TaskEmbedder.java
index 761b225a7cdc..5ebcc46aa0d8 100644
--- a/core/java/android/app/TaskEmbedder.java
+++ b/core/java/android/app/TaskEmbedder.java
@@ -597,7 +597,7 @@ public class TaskEmbedder {
         if (mTmpDisplayMetrics == null) {
             mTmpDisplayMetrics = new DisplayMetrics();
         }
-        mContext.getDisplay().getMetrics(mTmpDisplayMetrics);
+        mContext.getDisplayNoVerify().getRealMetrics(mTmpDisplayMetrics);
         return mTmpDisplayMetrics.densityDpi;
     }
 
diff --git a/core/java/android/app/WallpaperManager.java b/core/java/android/app/WallpaperManager.java
index 5f74d2e1dd57..d9405e18a162 100644
--- a/core/java/android/app/WallpaperManager.java
+++ b/core/java/android/app/WallpaperManager.java
@@ -2097,7 +2097,7 @@ public class WallpaperManager {
 
         public ColorManagementProxy(@NonNull Context context) {
             // Get a list of supported wide gamut color spaces.
-            Display display = context.getDisplay();
+            Display display = context.getDisplayNoVerify();
             if (display != null) {
                 mSupportedColorSpaces.addAll(Arrays.asList(display.getSupportedWideColorGamut()));
             }
diff --git a/core/java/android/content/Context.java b/core/java/android/content/Context.java
index ae12de027e6e..c6e84b7e46b0 100644
--- a/core/java/android/content/Context.java
+++ b/core/java/android/content/Context.java
@@ -3495,13 +3495,23 @@ public abstract class Context {
      * <dl>
      *  <dt> {@link #WINDOW_SERVICE} ("window")
      *  <dd> The top-level window manager in which you can place custom
-     *  windows.  The returned object is a {@link android.view.WindowManager}.
+     *  windows.  The returned object is a {@link android.view.WindowManager}. Must only be obtained
+     *  from a visual context such as Activity or a Context created with
+     *  {@link #createWindowContext(int, Bundle)}, which are adjusted to the configuration and
+     *  visual bounds of an area on screen.
      *  <dt> {@link #LAYOUT_INFLATER_SERVICE} ("layout_inflater")
      *  <dd> A {@link android.view.LayoutInflater} for inflating layout resources
-     *  in this context.
+     *  in this context. Must only be obtained from a visual context such as Activity or a Context
+     *  created with {@link #createWindowContext(int, Bundle)}, which are adjusted to the
+     *  configuration and visual bounds of an area on screen.
      *  <dt> {@link #ACTIVITY_SERVICE} ("activity")
      *  <dd> A {@link android.app.ActivityManager} for interacting with the
      *  global activity state of the system.
+     *  <dt> {@link #WALLPAPER_SERVICE} ("wallpaper")
+     *  <dd> A {@link android.service.wallpaper.WallpaperService} for accessing wallpapers in this
+     *  context. Must only be obtained from a visual context such as Activity or a Context created
+     *  with {@link #createWindowContext(int, Bundle)}, which are adjusted to the configuration and
+     *  visual bounds of an area on screen.
      *  <dt> {@link #POWER_SERVICE} ("power")
      *  <dd> A {@link android.os.PowerManager} for controlling power
      *  management.
@@ -5901,6 +5911,8 @@ public abstract class Context {
      * {@link #createDisplayContext(Display)} to get a display object associated with a Context, or
      * {@link android.hardware.display.DisplayManager#getDisplay} to get a display object by id.
      * @return Returns the {@link Display} object this context is associated with.
+     * @throws UnsupportedOperationException if the method is called on an instance that is not
+     *         associated with any display.
      */
     @Nullable
     public Display getDisplay() {
@@ -5908,6 +5920,17 @@ public abstract class Context {
     }
 
     /**
+     * A version of {@link #getDisplay()} that does not perform a Context misuse check to be used by
+     * legacy APIs.
+     * TODO(b/149790106): Fix usages and remove.
+     * @hide
+     */
+    @Nullable
+    public Display getDisplayNoVerify() {
+        throw new RuntimeException("Not implemented. Must override in a subclass.");
+    }
+
+    /**
      * Gets the ID of the display this context is associated with.
      *
      * @return display ID associated with this {@link Context}.
diff --git a/core/java/android/content/ContextWrapper.java b/core/java/android/content/ContextWrapper.java
index f6515e806caa..e5381ea0dd41 100644
--- a/core/java/android/content/ContextWrapper.java
+++ b/core/java/android/content/ContextWrapper.java
@@ -1003,6 +1003,12 @@ public class ContextWrapper extends Context {
         return mBase.getDisplay();
     }
 
+    /** @hide */
+    @Override
+    public @Nullable Display getDisplayNoVerify() {
+        return mBase.getDisplayNoVerify();
+    }
+
     /**
      * @hide
      */
diff --git a/core/java/android/os/RecoverySystem.java b/core/java/android/os/RecoverySystem.java
index be2245820bb6..d60820ef0f57 100644
--- a/core/java/android/os/RecoverySystem.java
+++ b/core/java/android/os/RecoverySystem.java
@@ -16,6 +16,8 @@
 
 package android.os;
 
+import static android.view.Display.DEFAULT_DISPLAY;
+
 import static java.nio.charset.StandardCharsets.UTF_8;
 
 import android.annotation.NonNull;
@@ -33,6 +35,7 @@ import android.content.Intent;
 import android.content.IntentFilter;
 import android.content.IntentSender;
 import android.content.pm.PackageManager;
+import android.hardware.display.DisplayManager;
 import android.provider.Settings;
 import android.telephony.SubscriptionInfo;
 import android.telephony.SubscriptionManager;
@@ -614,7 +617,8 @@ public class RecoverySystem {
 
             // On TV, reboot quiescently if the screen is off
             if (context.getPackageManager().hasSystemFeature(PackageManager.FEATURE_LEANBACK)) {
-                if (context.getDisplay().getState() != Display.STATE_ON) {
+                DisplayManager dm = context.getSystemService(DisplayManager.class);
+                if (dm.getDisplay(DEFAULT_DISPLAY).getState() != Display.STATE_ON) {
                     reason += ",quiescent";
                 }
             }
diff --git a/core/java/android/os/StrictMode.java b/core/java/android/os/StrictMode.java
index 3faaff73a0ea..e8af56478230 100644
--- a/core/java/android/os/StrictMode.java
+++ b/core/java/android/os/StrictMode.java
@@ -42,6 +42,7 @@ import android.os.strictmode.DiskWriteViolation;
 import android.os.strictmode.ExplicitGcViolation;
 import android.os.strictmode.FileUriExposedViolation;
 import android.os.strictmode.ImplicitDirectBootViolation;
+import android.os.strictmode.IncorrectContextUseViolation;
 import android.os.strictmode.InstanceCountViolation;
 import android.os.strictmode.IntentReceiverLeakedViolation;
 import android.os.strictmode.LeakedClosableViolation;
@@ -250,6 +251,7 @@ public final class StrictMode {
             DETECT_VM_UNTAGGED_SOCKET,
             DETECT_VM_NON_SDK_API_USAGE,
             DETECT_VM_IMPLICIT_DIRECT_BOOT,
+            DETECT_VM_INCORRECT_CONTEXT_USE,
             PENALTY_GATHER,
             PENALTY_LOG,
             PENALTY_DIALOG,
@@ -289,6 +291,8 @@ public final class StrictMode {
     private static final int DETECT_VM_IMPLICIT_DIRECT_BOOT = 1 << 10;
     /** @hide */
     private static final int DETECT_VM_CREDENTIAL_PROTECTED_WHILE_LOCKED = 1 << 11;
+    /** @hide */
+    private static final int DETECT_VM_INCORRECT_CONTEXT_USE = 1 << 12;
 
     /** @hide */
     private static final int DETECT_VM_ALL = 0x0000ffff;
@@ -871,6 +875,9 @@ public final class StrictMode {
                 if (targetSdk >= Build.VERSION_CODES.Q) {
                     detectCredentialProtectedWhileLocked();
                 }
+                if (targetSdk >= Build.VERSION_CODES.R) {
+                    detectIncorrectContextUse();
+                }
 
                 // TODO: Decide whether to detect non SDK API usage beyond a certain API level.
                 // TODO: enable detectImplicitDirectBoot() once system is less noisy
@@ -1028,6 +1035,32 @@ public final class StrictMode {
             }
 
             /**
+             * Detect attempts to invoke a method on a {@link Context} that is not suited for such
+             * operation.
+             * <p>An example of this is trying to obtain an instance of visual service (e.g.
+             * {@link android.view.WindowManager}) from a non-visual {@link Context}. This is not
+             * allowed, since a non-visual {@link Context} is not adjusted to any visual area, and
+             * therefore can report incorrect metrics or resources.
+             * @see Context#getDisplay()
+             * @see Context#getSystemService(String)
+             * @hide
+             */
+            @TestApi
+            public @NonNull Builder detectIncorrectContextUse() {
+                return enable(DETECT_VM_INCORRECT_CONTEXT_USE);
+            }
+
+            /**
+             * Disable detection of incorrect context use.
+             * TODO(b/149790106): Fix usages and remove.
+             * @hide
+             */
+            @TestApi
+            public @NonNull Builder permitIncorrectContextUse() {
+                return disable(DETECT_VM_INCORRECT_CONTEXT_USE);
+            }
+
+            /**
              * Crashes the whole process on violation. This penalty runs at the end of all enabled
              * penalties so you'll still get your logging or other violations before the process
              * dies.
@@ -2057,6 +2090,11 @@ public final class StrictMode {
     }
 
     /** @hide */
+    public static boolean vmIncorrectContextUseEnabled() {
+        return (sVmPolicy.mask & DETECT_VM_INCORRECT_CONTEXT_USE) != 0;
+    }
+
+    /** @hide */
     public static void onSqliteObjectLeaked(String message, Throwable originStack) {
         onVmPolicyViolation(new SqliteObjectLeakedViolation(message, originStack));
     }
@@ -2130,6 +2168,11 @@ public final class StrictMode {
         onVmPolicyViolation(new ImplicitDirectBootViolation());
     }
 
+    /** @hide */
+    public static void onIncorrectContextUsed(String message, Throwable originStack) {
+        onVmPolicyViolation(new IncorrectContextUseViolation(message, originStack));
+    }
+
     /** Assume locked until we hear otherwise */
     private static volatile boolean sUserKeyUnlocked = false;
 
diff --git a/core/java/android/os/strictmode/IncorrectContextUseViolation.java b/core/java/android/os/strictmode/IncorrectContextUseViolation.java
new file mode 100644
index 000000000000..647db171e080
--- /dev/null
+++ b/core/java/android/os/strictmode/IncorrectContextUseViolation.java
@@ -0,0 +1,35 @@
+/*
+ * Copyright (C) 2020 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.os.strictmode;
+
+import android.content.Context;
+
+/**
+ * Incorrect usage of {@link Context}, such as obtaining a visual service from non-visual
+ * {@link Context} instance.
+ * @see Context#getSystemService(String)
+ * @see Context#getDisplayNoVerify()
+ * @hide
+ */
+public final class IncorrectContextUseViolation extends Violation {
+
+    /** @hide */
+    public IncorrectContextUseViolation(String message, Throwable originStack) {
+        super(message);
+        initCause(originStack);
+    }
+}
diff --git a/core/java/android/view/WindowManagerImpl.java b/core/java/android/view/WindowManagerImpl.java
index 56683dd9a5d1..d40f505f87a0 100644
--- a/core/java/android/view/WindowManagerImpl.java
+++ b/core/java/android/view/WindowManagerImpl.java
@@ -101,7 +101,7 @@ public final class WindowManagerImpl implements WindowManager {
     @Override
     public void addView(@NonNull View view, @NonNull ViewGroup.LayoutParams params) {
         applyDefaultToken(params);
-        mGlobal.addView(view, params, mContext.getDisplay(), mParentWindow);
+        mGlobal.addView(view, params, mContext.getDisplayNoVerify(), mParentWindow);
         mIsViewAdded = true;
         mLastView = view;
         mLastParams = (WindowManager.LayoutParams) params;
@@ -158,7 +158,7 @@ public final class WindowManagerImpl implements WindowManager {
 
     @Override
     public Display getDefaultDisplay() {
-        return mContext.getDisplay();
+        return mContext.getDisplayNoVerify();
     }
 
     @Override
@@ -240,7 +240,7 @@ public final class WindowManagerImpl implements WindowManager {
     private Rect getMaximumBounds() {
         // TODO(b/128338354): Current maximum bound is display size, but it should be displayArea
         //  bound after displayArea feature is finished.
-        final Display display = mContext.getDisplay();
+        final Display display = mContext.getDisplayNoVerify();
         final Point displaySize = new Point();
         display.getRealSize(displaySize);
         return new Rect(0, 0, displaySize.x, displaySize.y);
diff --git a/core/java/com/android/internal/view/FloatingActionMode.java b/core/java/com/android/internal/view/FloatingActionMode.java
index 69b1609377b7..633d6848030e 100644
--- a/core/java/com/android/internal/view/FloatingActionMode.java
+++ b/core/java/com/android/internal/view/FloatingActionMode.java
@@ -210,7 +210,7 @@ public final class FloatingActionMode extends ActionMode {
     }
 
     private boolean isContentRectWithinBounds() {
-        mContext.getDisplay().getRealSize(mDisplaySize);
+        mContext.getDisplayNoVerify().getRealSize(mDisplaySize);
         mScreenRect.set(0, 0, mDisplaySize.x, mDisplaySize.y);
 
         return intersectsClosed(mContentRectOnScreen, mScreenRect)
diff --git a/core/java/com/android/internal/view/RotationPolicy.java b/core/java/com/android/internal/view/RotationPolicy.java
index 13425e5e9bec..cfb2bf9df9ed 100644
--- a/core/java/com/android/internal/view/RotationPolicy.java
+++ b/core/java/com/android/internal/view/RotationPolicy.java
@@ -77,10 +77,7 @@ public final class RotationPolicy {
             final Point size = new Point();
             final IWindowManager wm = WindowManagerGlobal.getWindowManagerService();
             try {
-                final Display display = context.getDisplay();
-                final int displayId = display != null
-                        ? display.getDisplayId()
-                        : Display.DEFAULT_DISPLAY;
+                final int displayId = context.getDisplayId();
                 wm.getInitialDisplaySize(displayId, size);
                 return size.x < size.y ?
                         Configuration.ORIENTATION_PORTRAIT : Configuration.ORIENTATION_LANDSCAPE;
diff --git a/core/tests/coretests/src/android/content/ContextTest.java b/core/tests/coretests/src/android/content/ContextTest.java
index f074233070d1..f0997a68492b 100644
--- a/core/tests/coretests/src/android/content/ContextTest.java
+++ b/core/tests/coretests/src/android/content/ContextTest.java
@@ -16,11 +16,13 @@
 
 package android.content;
 
+import static android.view.Display.DEFAULT_DISPLAY;
+
 import static org.junit.Assert.assertEquals;
 
 import android.app.ActivityThread;
+import android.hardware.display.DisplayManager;
 import android.os.UserHandle;
-import android.view.WindowManager;
 
 import androidx.test.InstrumentationRegistry;
 import androidx.test.filters.SmallTest;
@@ -45,17 +47,16 @@ public class ContextTest {
         final Context testContext =
                 InstrumentationRegistry.getInstrumentation().getTargetContext();
 
-        assertEquals(testContext.getDisplay().getDisplayId(), testContext.getDisplayId());
+        assertEquals(testContext.getDisplayNoVerify().getDisplayId(), testContext.getDisplayId());
     }
 
-    // TODO(b/128338354): Re-visit this test after introducing WindowContext
     @Test
     public void testDisplayIdForDefaultDisplayContext() {
         final Context testContext =
                 InstrumentationRegistry.getInstrumentation().getTargetContext();
-        final WindowManager wm = testContext.getSystemService(WindowManager.class);
+        final DisplayManager dm = testContext.getSystemService(DisplayManager.class);
         final Context defaultDisplayContext =
-                testContext.createDisplayContext(wm.getDefaultDisplay());
+                testContext.createDisplayContext(dm.getDisplay(DEFAULT_DISPLAY));
 
         assertEquals(defaultDisplayContext.getDisplay().getDisplayId(),
                 defaultDisplayContext.getDisplayId());
diff --git a/core/tests/coretests/src/android/view/ImeInsetsSourceConsumerTest.java b/core/tests/coretests/src/android/view/ImeInsetsSourceConsumerTest.java
index 7c78bce25b1c..7f0e0d2f54c7 100644
--- a/core/tests/coretests/src/android/view/ImeInsetsSourceConsumerTest.java
+++ b/core/tests/coretests/src/android/view/ImeInsetsSourceConsumerTest.java
@@ -61,7 +61,7 @@ public class ImeInsetsSourceConsumerTest {
                 .setName("testSurface")
                 .build();
         InstrumentationRegistry.getInstrumentation().runOnMainSync(() -> {
-            ViewRootImpl viewRootImpl = new ViewRootImpl(mContext, mContext.getDisplay());
+            ViewRootImpl viewRootImpl = new ViewRootImpl(mContext, mContext.getDisplayNoVerify());
             try {
                 viewRootImpl.setView(new TextView(mContext), new LayoutParams(), null);
             } catch (BadTokenException e) {
diff --git a/core/tests/coretests/src/android/view/InsetsControllerTest.java b/core/tests/coretests/src/android/view/InsetsControllerTest.java
index 24fe2a0a1823..7737b1a2a776 100644
--- a/core/tests/coretests/src/android/view/InsetsControllerTest.java
+++ b/core/tests/coretests/src/android/view/InsetsControllerTest.java
@@ -110,7 +110,7 @@ public class InsetsControllerTest {
         InstrumentationRegistry.getInstrumentation().runOnMainSync(() -> {
             Context context = InstrumentationRegistry.getTargetContext();
             // cannot mock ViewRootImpl since it's final.
-            mViewRoot = new ViewRootImpl(context, context.getDisplay());
+            mViewRoot = new ViewRootImpl(context, context.getDisplayNoVerify());
             try {
                 mViewRoot.setView(new TextView(context), new LayoutParams(), null);
             } catch (BadTokenException e) {
diff --git a/core/tests/coretests/src/android/view/InsetsSourceConsumerTest.java b/core/tests/coretests/src/android/view/InsetsSourceConsumerTest.java
index 5e9e2f0065ed..754c6791cade 100644
--- a/core/tests/coretests/src/android/view/InsetsSourceConsumerTest.java
+++ b/core/tests/coretests/src/android/view/InsetsSourceConsumerTest.java
@@ -77,7 +77,8 @@ public class InsetsSourceConsumerTest {
         instrumentation.runOnMainSync(() -> {
             final Context context = instrumentation.getTargetContext();
             // cannot mock ViewRootImpl since it's final.
-            final ViewRootImpl viewRootImpl = new ViewRootImpl(context, context.getDisplay());
+            final ViewRootImpl viewRootImpl = new ViewRootImpl(context,
+                    context.getDisplayNoVerify());
             try {
                 viewRootImpl.setView(new TextView(context), new LayoutParams(), null);
             } catch (BadTokenException e) {
diff --git a/core/tests/coretests/src/android/view/ViewRootImplTest.java b/core/tests/coretests/src/android/view/ViewRootImplTest.java
index df6ed8c3fe0d..e2adbcc600d7 100644
--- a/core/tests/coretests/src/android/view/ViewRootImplTest.java
+++ b/core/tests/coretests/src/android/view/ViewRootImplTest.java
@@ -61,7 +61,7 @@ public class ViewRootImplTest {
 
         InstrumentationRegistry.getInstrumentation().runOnMainSync(() -> {
             mViewRootImpl = new ViewRootImplAccessor(
-                    new ViewRootImpl(mContext, mContext.getDisplay()));
+                    new ViewRootImpl(mContext, mContext.getDisplayNoVerify()));
         });
     }
 
diff --git a/core/tests/coretests/src/com/android/internal/policy/DecorContextTest.java b/core/tests/coretests/src/com/android/internal/policy/DecorContextTest.java
index cdcf23fc10bb..3e40466e4b64 100644
--- a/core/tests/coretests/src/com/android/internal/policy/DecorContextTest.java
+++ b/core/tests/coretests/src/com/android/internal/policy/DecorContextTest.java
@@ -53,7 +53,10 @@ public final class DecorContextTest {
 
     @Test
     public void testDecorContextWithDefaultDisplay() {
-        DecorContext context = new DecorContext(mContext.getApplicationContext(), mContext);
+        Display defaultDisplay = new Display(DisplayManagerGlobal.getInstance(), DEFAULT_DISPLAY,
+                new DisplayInfo(), DisplayAdjustments.DEFAULT_DISPLAY_ADJUSTMENTS);
+        DecorContext context = new DecorContext(mContext.getApplicationContext(),
+                mContext.createDisplayContext(defaultDisplay));
 
         assertDecorContextDisplay(DEFAULT_DISPLAY, context);
     }
diff --git a/packages/SettingsLib/src/com/android/settingslib/display/DisplayDensityUtils.java b/packages/SettingsLib/src/com/android/settingslib/display/DisplayDensityUtils.java
index e0ca1ab0c07c..a38091debb64 100644
--- a/packages/SettingsLib/src/com/android/settingslib/display/DisplayDensityUtils.java
+++ b/packages/SettingsLib/src/com/android/settingslib/display/DisplayDensityUtils.java
@@ -97,7 +97,7 @@ public class DisplayDensityUtils {
 
         final Resources res = context.getResources();
         final DisplayMetrics metrics = new DisplayMetrics();
-        context.getDisplay().getRealMetrics(metrics);
+        context.getDisplayNoVerify().getRealMetrics(metrics);
 
         final int currentDensity = metrics.densityDpi;
         int currentDensityIndex = -1;
diff --git a/packages/SystemUI/src/com/android/keyguard/KeyguardStatusView.java b/packages/SystemUI/src/com/android/keyguard/KeyguardStatusView.java
index 61caf3bc5d8f..241f96e19e38 100644
--- a/packages/SystemUI/src/com/android/keyguard/KeyguardStatusView.java
+++ b/packages/SystemUI/src/com/android/keyguard/KeyguardStatusView.java
@@ -149,7 +149,7 @@ public class KeyguardStatusView extends GridLayout implements
                             new WindowlessWindowManager(context.getResources().getConfiguration(),
                                     surfaceControl, input);
                     mUniversalSmartspaceViewHost = new SurfaceControlViewHost(context,
-                            context.getDisplay(), windowlessWindowManager);
+                            context.getDisplayNoVerify(), windowlessWindowManager);
                     WindowManager.LayoutParams layoutParams =
                             new WindowManager.LayoutParams(
                                     surfaceControl.getWidth(),
diff --git a/packages/SystemUI/tests/src/com/android/keyguard/KeyguardPresentationTest.java b/packages/SystemUI/tests/src/com/android/keyguard/KeyguardPresentationTest.java
index b6ca8d8e4afb..7231b8a143d0 100644
--- a/packages/SystemUI/tests/src/com/android/keyguard/KeyguardPresentationTest.java
+++ b/packages/SystemUI/tests/src/com/android/keyguard/KeyguardPresentationTest.java
@@ -41,8 +41,8 @@ public class KeyguardPresentationTest extends SysuiTestCase {
         InjectionInflationController inflationController = new InjectionInflationController(
                 SystemUIFactory.getInstance().getRootComponent());
         Context context = getContext();
-        KeyguardPresentation keyguardPresentation =
-                new KeyguardPresentation(context, context.getDisplay(), inflationController);
+        KeyguardPresentation keyguardPresentation = new KeyguardPresentation(context,
+                context.getDisplayNoVerify(), inflationController);
         keyguardPresentation.onCreate(null /*savedInstanceState */);
     }
 }
diff --git a/packages/WallpaperCropper/src/com/android/photos/views/TiledImageRenderer.java b/packages/WallpaperCropper/src/com/android/photos/views/TiledImageRenderer.java
index fcb113e90720..210fdc662bed 100644
--- a/packages/WallpaperCropper/src/com/android/photos/views/TiledImageRenderer.java
+++ b/packages/WallpaperCropper/src/com/android/photos/views/TiledImageRenderer.java
@@ -163,7 +163,7 @@ public class TiledImageRenderer {
 
     private static boolean isHighResolution(Context context) {
         DisplayMetrics metrics = new DisplayMetrics();
-        context.getDisplay().getMetrics(metrics);
+        context.getDisplayNoVerify().getMetrics(metrics);
         return metrics.heightPixels > 2048 ||  metrics.widthPixels > 2048;
     }
 
diff --git a/services/appwidget/java/com/android/server/appwidget/AppWidgetServiceImpl.java b/services/appwidget/java/com/android/server/appwidget/AppWidgetServiceImpl.java
index 7151d2b86e83..a8a27916f56a 100644
--- a/services/appwidget/java/com/android/server/appwidget/AppWidgetServiceImpl.java
+++ b/services/appwidget/java/com/android/server/appwidget/AppWidgetServiceImpl.java
@@ -281,7 +281,7 @@ class AppWidgetServiceImpl extends IAppWidgetService.Stub implements WidgetBacku
     }
 
     private void computeMaximumWidgetBitmapMemory() {
-        Display display = mContext.getDisplay();
+        Display display = mContext.getDisplayNoVerify();
         Point size = new Point();
         display.getRealSize(size);
         // Cap memory usage at 1.5 times the size of the display
diff --git a/services/autofill/java/com/android/server/autofill/ui/CustomScrollView.java b/services/autofill/java/com/android/server/autofill/ui/CustomScrollView.java
index 813fc8d5f561..14bd7d78f3bf 100644
--- a/services/autofill/java/com/android/server/autofill/ui/CustomScrollView.java
+++ b/services/autofill/java/com/android/server/autofill/ui/CustomScrollView.java
@@ -75,7 +75,7 @@ public class CustomScrollView extends ScrollView {
         final TypedValue typedValue = new TypedValue();
         final Point point = new Point();
         final Context context = getContext();
-        context.getDisplay().getSize(point);
+        context.getDisplayNoVerify().getSize(point);
         context.getTheme().resolveAttribute(R.attr.autofillSaveCustomSubtitleMaxHeight,
                 typedValue, true);
         final View child = getChildAt(0);
diff --git a/services/autofill/java/com/android/server/autofill/ui/FillUi.java b/services/autofill/java/com/android/server/autofill/ui/FillUi.java
index 5dc43ef8ad56..344b92f43089 100644
--- a/services/autofill/java/com/android/server/autofill/ui/FillUi.java
+++ b/services/autofill/java/com/android/server/autofill/ui/FillUi.java
@@ -165,7 +165,7 @@ final class FillUi {
         // In full screen we only initialize size once assuming screen size never changes
         if (mFullScreen) {
             final Point outPoint = mTempPoint;
-            mContext.getDisplay().getSize(outPoint);
+            mContext.getDisplayNoVerify().getSize(outPoint);
             // full with of screen and half height of screen
             mContentWidth = LayoutParams.MATCH_PARENT;
             mContentHeight = outPoint.y / 2;
@@ -559,7 +559,7 @@ final class FillUi {
     }
 
     private static void resolveMaxWindowSize(Context context, Point outPoint) {
-        context.getDisplay().getSize(outPoint);
+        context.getDisplayNoVerify().getSize(outPoint);
         final TypedValue typedValue = sTempTypedValue;
         context.getTheme().resolveAttribute(R.attr.autofillDatasetPickerMaxWidth,
                 typedValue, true);
diff --git a/test-mock/src/android/test/mock/MockContext.java b/test-mock/src/android/test/mock/MockContext.java
index 359c44849634..2c6604759813 100644
--- a/test-mock/src/android/test/mock/MockContext.java
+++ b/test-mock/src/android/test/mock/MockContext.java
@@ -834,6 +834,12 @@ public class MockContext extends Context {
 
     /** @hide */
     @Override
+    public Display getDisplayNoVerify() {
+        throw new UnsupportedOperationException();
+    }
+
+    /** @hide */
+    @Override
     public int getDisplayId() {
         throw new UnsupportedOperationException();
     }