Fix WindowContext leak

The root causes of this issue are:
  1. WindowTokenClient#attachContext makes WindowTokenClient has strong
     reference to WindowContext, which leads to WindowContext cannot be
     GC'd.
  2. WMS#removeWindowToken needs MANAGE_APP_TOKEN permission which
     normal apps don't hold.

This CL does following things:
  1. Use weak reference instead on WindowTokenClient#mContext.
  2. Relax WMS#removeWindowToken to check callingUid if
     MANAGE_WINDOW_TOKEN permission is not held
  3. Deliver config changes to the client side in
     WMS#addWindowTokenWithOption
  4. Some minor fixes

fixes: 150812449
Bug: 150715095
Test: atest WindowContextTest
Test: atest WindowManagerServiceTests
Test: atest WindowTokenTests
Test: atest WindowManagerPermissionTests#testMANAGE_APP_TOKENS
Test: atest
WindowManagerPermissionTests#testADD_WINDOW_TOKEN_WITH_OPTIONS

Change-Id: I9f1d73af2abb78fc9844e6d9eb25e9f0293514e7
diff --git a/core/java/android/app/IWindowToken.aidl b/core/java/android/app/IWindowToken.aidl
index 8ea881f..3627b0f 100644
--- a/core/java/android/app/IWindowToken.aidl
+++ b/core/java/android/app/IWindowToken.aidl
@@ -30,4 +30,6 @@
  */
 oneway interface IWindowToken {
     void onConfigurationChanged(in Configuration newConfig, int newDisplayId);
+
+    void onWindowTokenRemoved();
 }
diff --git a/core/java/android/app/WindowContext.java b/core/java/android/app/WindowContext.java
index 878993e..3a06c9d 100644
--- a/core/java/android/app/WindowContext.java
+++ b/core/java/android/app/WindowContext.java
@@ -28,6 +28,8 @@
 import android.view.WindowManagerGlobal;
 import android.view.WindowManagerImpl;
 
+import com.android.internal.annotations.VisibleForTesting;
+
 import java.lang.ref.Reference;
 
 /**
@@ -75,8 +77,6 @@
             // config back to the client.
             result = mWms.addWindowTokenWithOptions(
                     mToken, type, getDisplayId(), options, getPackageName());
-
-            // TODO(window-context): remove token with a DeathObserver
         }  catch (RemoteException e) {
             mOwnsToken = false;
             throw e.rethrowFromSystemServer();
@@ -100,6 +100,13 @@
 
     @Override
     protected void finalize() throws Throwable {
+        release();
+        super.finalize();
+    }
+
+    /** Used for test to invoke because we can't invoke finalize directly. */
+    @VisibleForTesting
+    public void release() {
         if (mOwnsToken) {
             try {
                 mWms.removeWindowToken(mToken, getDisplayId());
@@ -108,6 +115,12 @@
                 throw e.rethrowFromSystemServer();
             }
         }
-        super.finalize();
+        destroy();
+    }
+
+    void destroy() {
+        final ContextImpl impl = (ContextImpl) getBaseContext();
+        impl.scheduleFinalCleanup(getClass().getName(), "WindowContext");
+        Reference.reachabilityFence(this);
     }
 }
diff --git a/core/java/android/app/WindowTokenClient.java b/core/java/android/app/WindowTokenClient.java
index ed0179b..301960e 100644
--- a/core/java/android/app/WindowTokenClient.java
+++ b/core/java/android/app/WindowTokenClient.java
@@ -20,6 +20,9 @@
 import android.content.res.Configuration;
 import android.os.Bundle;
 import android.os.IBinder;
+import android.view.WindowManagerGlobal;
+
+import java.lang.ref.WeakReference;
 
 /**
  * Client implementation of {@link IWindowToken}. It can receive configuration change callbacks from
@@ -31,9 +34,9 @@
 public class WindowTokenClient extends IWindowToken.Stub {
     /**
      * Attached {@link Context} for this window token to update configuration and resources.
-     * Initialized by {@link #attachContext(Context)}.
+     * Initialized by {@link #attachContext(WindowContext)}.
      */
-    private Context mContext = null;
+    private WeakReference<WindowContext> mContextRef = null;
 
     private final ResourcesManager mResourcesManager = ResourcesManager.getInstance();
 
@@ -47,30 +50,46 @@
      * @param context context to be attached
      * @throws IllegalStateException if attached context has already existed.
      */
-    void attachContext(@NonNull Context context) {
-        if (mContext != null) {
+    void attachContext(@NonNull WindowContext context) {
+        if (mContextRef != null) {
             throw new IllegalStateException("Context is already attached.");
         }
-        mContext = context;
-        ContextImpl impl = ContextImpl.getImpl(mContext);
+        mContextRef = new WeakReference<>(context);
+        final ContextImpl impl = ContextImpl.getImpl(context);
         impl.setResources(impl.createWindowContextResources());
     }
 
     @Override
     public void onConfigurationChanged(Configuration newConfig, int newDisplayId) {
-        final int currentDisplayId = mContext.getDisplayId();
+        final Context context = mContextRef.get();
+        if (context == null) {
+            return;
+        }
+        final int currentDisplayId = context.getDisplayId();
         final boolean displayChanged = newDisplayId != currentDisplayId;
-        final Configuration config = new Configuration(mContext.getResources()
+        final Configuration config = new Configuration(context.getResources()
                 .getConfiguration());
         final boolean configChanged = config.isOtherSeqNewer(newConfig)
                 && config.updateFrom(newConfig) != 0;
         if (displayChanged || configChanged) {
             // TODO(ag/9789103): update resource manager logic to track non-activity tokens
-            mResourcesManager.updateResourcesForActivity(asBinder(), config, newDisplayId,
+            mResourcesManager.updateResourcesForActivity(this, config, newDisplayId,
                     displayChanged);
         }
         if (displayChanged) {
-            mContext.updateDisplay(newDisplayId);
+            context.updateDisplay(newDisplayId);
         }
     }
+
+    @Override
+    public void onWindowTokenRemoved() {
+        final WindowContext context = mContextRef.get();
+        if (context != null) {
+            context.destroy();
+            mContextRef.clear();
+        }
+        // If a secondary display is detached, release all views attached to this token.
+        WindowManagerGlobal.getInstance().closeAll(this, mContextRef.getClass().getName(),
+                "WindowContext");
+    }
 }
diff --git a/core/java/android/view/IWindowManager.aidl b/core/java/android/view/IWindowManager.aidl
index 84ac90b..5648adf 100644
--- a/core/java/android/view/IWindowManager.aidl
+++ b/core/java/android/view/IWindowManager.aidl
@@ -124,13 +124,20 @@
      * @param type Window type to be used with this token.
      * @param options A bundle used to pass window-related options.
      * @param displayId The ID of the display where this token should be added.
-     * @param packageName The name of package to request to add window token.
+     * @param packageName The name of package to request to add window token. Could be {@code null}
+     *                    if callers holds the MANAGE_APP_TOKENS permission.
      * @return {@link WindowManagerGlobal#ADD_OKAY} if the addition was successful, an error code
      *         otherwise.
      */
     int addWindowTokenWithOptions(IBinder token, int type, int displayId, in Bundle options,
             String packageName);
     void addWindowToken(IBinder token, int type, int displayId);
+    /**
+     * Remove window token on a specific display.
+     *
+     * @param token Token to be removed
+     * @displayId The ID of the display where this token should be removed.
+     */
     void removeWindowToken(IBinder token, int displayId);
     void prepareAppTransition(int transit, boolean alwaysKeepCurrent);
 
diff --git a/core/tests/coretests/src/android/app/WindowContextTest.java b/core/tests/coretests/src/android/app/WindowContextTest.java
new file mode 100644
index 0000000..630e16a
--- /dev/null
+++ b/core/tests/coretests/src/android/app/WindowContextTest.java
@@ -0,0 +1,74 @@
+/*
+ * 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.app;
+
+import static android.view.Display.DEFAULT_DISPLAY;
+import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION_OVERLAY;
+
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+import android.content.Context;
+import android.hardware.display.DisplayManager;
+import android.os.IBinder;
+import android.platform.test.annotations.Presubmit;
+import android.view.Display;
+import android.view.IWindowManager;
+import android.view.WindowManagerGlobal;
+
+import androidx.test.filters.FlakyTest;
+import androidx.test.filters.SmallTest;
+import androidx.test.platform.app.InstrumentationRegistry;
+import androidx.test.runner.AndroidJUnit4;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+/**
+ * Tests for {@link WindowContext}
+ *
+ * <p>Build/Install/Run:
+ *  atest FrameworksCoreTests:WindowContextTest
+ *
+ * <p>This test class is a part of Window Manager Service tests and specified in
+ * {@link com.android.server.wm.test.filters.FrameworksTestsFilter}.
+ */
+@FlakyTest(bugId = 150812449, detail = "Remove after confirmed it's stable.")
+@RunWith(AndroidJUnit4.class)
+@SmallTest
+@Presubmit
+public class WindowContextTest {
+    @Test
+    public void testWindowContextRelease_doRemoveWindowToken() throws Throwable {
+        final Context instContext = InstrumentationRegistry.getInstrumentation()
+                .getTargetContext();
+        final Display display = instContext.getSystemService(DisplayManager.class)
+                .getDisplay(DEFAULT_DISPLAY);
+        final Context context = instContext.createDisplayContext(display);
+        final WindowContext windowContext = new WindowContext(context, TYPE_APPLICATION_OVERLAY,
+                null /* options */);
+
+        final IBinder token = windowContext.getActivityToken();
+
+        final IWindowManager wms = WindowManagerGlobal.getWindowManagerService();
+        assertTrue("Token must be registered to WMS", wms.isWindowToken(token));
+
+        windowContext.release();
+
+        assertFalse("Token must be unregistered to WMS", wms.isWindowToken(token));
+    }
+}
diff --git a/core/tests/coretests/src/android/view/WindowMetricsTest.java b/core/tests/coretests/src/android/view/WindowMetricsTest.java
index 74524bf..ddc977d 100644
--- a/core/tests/coretests/src/android/view/WindowMetricsTest.java
+++ b/core/tests/coretests/src/android/view/WindowMetricsTest.java
@@ -56,17 +56,17 @@
 
     @Before
     public void setUp() {
-        final Context insetContext = InstrumentationRegistry.getInstrumentation()
+        final Context instContext = InstrumentationRegistry.getInstrumentation()
                 .getTargetContext();
-        final Display display = insetContext.getSystemService(DisplayManager.class)
+        final Display display = instContext.getSystemService(DisplayManager.class)
                 .getDisplay(DEFAULT_DISPLAY);
-        mWindowContext = insetContext.createDisplayContext(display)
+        mWindowContext = instContext.createDisplayContext(display)
                 .createWindowContext(TYPE_APPLICATION_OVERLAY, null /* options */);
         mWm = mWindowContext.getSystemService(WindowManager.class);
     }
 
     @Test
-    public void testAddViewANdRemoveView_GetMetrics_DoNotCrash() {
+    public void testAddViewAndRemoveView_GetMetrics_DoNotCrash() {
         final View view = new View(mWindowContext);
         final WindowManager.LayoutParams params =
                 new WindowManager.LayoutParams(TYPE_APPLICATION_OVERLAY);
diff --git a/data/etc/services.core.protolog.json b/data/etc/services.core.protolog.json
index a871047..72998ad 100644
--- a/data/etc/services.core.protolog.json
+++ b/data/etc/services.core.protolog.json
@@ -1525,6 +1525,18 @@
       "group": "WM_DEBUG_REMOTE_ANIMATIONS",
       "at": "com\/android\/server\/wm\/RemoteAnimationController.java"
     },
+    "838570988": {
+      "message": "Could not report token removal to the window token client.",
+      "level": "WARN",
+      "group": "WM_ERROR",
+      "at": "com\/android\/server\/wm\/WindowToken.java"
+    },
+    "845234215": {
+      "message": "App is requesting an orientation, return %d for display id=%d",
+      "level": "VERBOSE",
+      "group": "WM_DEBUG_ORIENTATION",
+      "at": "com\/android\/server\/wm\/DisplayContent.java"
+    },
     "853091290": {
       "message": "Moved stack=%s behind stack=%s",
       "level": "DEBUG",
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index b1f22f8..0674133 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -32,6 +32,7 @@
 import static android.content.pm.PackageManager.FEATURE_FREEFORM_WINDOW_MANAGEMENT;
 import static android.content.pm.PackageManager.FEATURE_PC;
 import static android.content.pm.PackageManager.PERMISSION_GRANTED;
+import static android.os.Process.INVALID_UID;
 import static android.os.Process.SYSTEM_UID;
 import static android.os.Process.myPid;
 import static android.os.Trace.TRACE_TAG_WINDOW_MANAGER;
@@ -2558,6 +2559,8 @@
             String packageName, boolean fromClientToken) {
         final boolean callerCanManageAppTokens =
                 checkCallingPermission(MANAGE_APP_TOKENS, "addWindowToken()");
+        // WindowContext users usually don't hold MANAGE_APP_TOKEN permission. Check permissions
+        // by checkAddPermission.
         if (!callerCanManageAppTokens) {
             final int res = mPolicy.checkAddPermission(type, false /* isRoundedCornerOverlay */,
                     packageName, new int[1]);
@@ -2572,7 +2575,7 @@
             synchronized (mGlobalLock) {
                 if (!callerCanManageAppTokens) {
                     if (packageName == null || !unprivilegedAppCanCreateTokenWith(
-                            null /* parentWindow */, callingUid, type, type, null /* tokenForLog */,
+                            null /* parentWindow */, callingUid, type, type, binder,
                             packageName)) {
                         throw new SecurityException("Requires MANAGE_APP_TOKENS permission");
                     }
@@ -2597,7 +2600,7 @@
                     new WallpaperWindowToken(this, binder, true, dc, callerCanManageAppTokens);
                 } else {
                     new WindowToken(this, binder, type, true, dc, callerCanManageAppTokens,
-                            false /* roundedCornerOverlay */, fromClientToken);
+                            callingUid, false /* roundedCornerOverlay */, fromClientToken);
                 }
             }
         } finally {
@@ -2620,8 +2623,25 @@
 
     @Override
     public void removeWindowToken(IBinder binder, int displayId) {
-        if (!checkCallingPermission(MANAGE_APP_TOKENS, "removeWindowToken()")) {
-            throw new SecurityException("Requires MANAGE_APP_TOKENS permission");
+        final boolean callerCanManageAppTokens =
+                checkCallingPermission(MANAGE_APP_TOKENS, "removeWindowToken()");
+        final WindowToken windowToken;
+        synchronized (mGlobalLock) {
+            windowToken = mRoot.getWindowToken(binder);
+        }
+        if (windowToken == null) {
+            ProtoLog.w(WM_ERROR,
+                    "removeWindowToken: Attempted to remove non-existing token: %s", binder);
+            return;
+        }
+        final int callingUid = Binder.getCallingUid();
+
+        // If MANAGE_APP_TOKEN permission is not held(usually from WindowContext), callers can only
+        // remove the window tokens which they added themselves.
+        if (!callerCanManageAppTokens && (windowToken.getOwnerUid() == INVALID_UID
+                || callingUid != windowToken.getOwnerUid())) {
+            throw new SecurityException("removeWindowToken: Requires MANAGE_APP_TOKENS permission"
+                    + " to remove token owned by another uid");
         }
 
         final long origId = Binder.clearCallingIdentity();
@@ -2634,14 +2654,7 @@
                     return;
                 }
 
-                final WindowToken token = dc.removeWindowToken(binder);
-                if (token == null) {
-                    ProtoLog.w(WM_ERROR,
-                            "removeWindowToken: Attempted to remove non-existing token: %s",
-                            binder);
-                    return;
-                }
-
+                dc.removeWindowToken(binder);
                 dc.getInputMonitor().updateInputWindowsLw(true /*force*/);
             }
         } finally {
diff --git a/services/core/java/com/android/server/wm/WindowToken.java b/services/core/java/com/android/server/wm/WindowToken.java
index 3c2b6ec..ccf5121 100644
--- a/services/core/java/com/android/server/wm/WindowToken.java
+++ b/services/core/java/com/android/server/wm/WindowToken.java
@@ -16,6 +16,7 @@
 
 package com.android.server.wm;
 
+import static android.os.Process.INVALID_UID;
 import static android.view.Display.INVALID_DISPLAY;
 import static android.view.WindowManager.LayoutParams.FLAG_SHOW_WALLPAPER;
 import static android.view.WindowManager.LayoutParams.TYPE_DOCK_DIVIDER;
@@ -44,6 +45,7 @@
 import android.os.Debug;
 import android.os.IBinder;
 import android.os.RemoteException;
+import android.util.Slog;
 import android.util.proto.ProtoOutputStream;
 import android.view.DisplayInfo;
 import android.view.InsetsState;
@@ -106,6 +108,11 @@
     @VisibleForTesting
     final boolean mFromClientToken;
 
+    private DeathRecipient mDeathRecipient;
+    private boolean mBinderDied = false;
+
+    private final int mOwnerUid;
+
     /**
      * Used to fix the transform of the token to be rotated to a rotation different than it's
      * display. The window frames and surfaces corresponding to this token will be layouted and
@@ -165,6 +172,30 @@
         }
     }
 
+    private class DeathRecipient implements IBinder.DeathRecipient {
+        private boolean mHasUnlinkToDeath = false;
+
+        @Override
+        public void binderDied() {
+            synchronized (mWmService.mGlobalLock) {
+                mBinderDied = true;
+                removeImmediately();
+            }
+        }
+
+        void linkToDeath() throws RemoteException {
+            token.linkToDeath(DeathRecipient.this, 0);
+        }
+
+        void unlinkToDeath() {
+            if (mHasUnlinkToDeath) {
+                return;
+            }
+            token.unlinkToDeath(DeathRecipient.this, 0);
+            mHasUnlinkToDeath = true;
+        }
+    }
+
     /**
      * Compares two child window of this token and returns -1 if the first is lesser than the
      * second in terms of z-order and 1 otherwise.
@@ -193,23 +224,35 @@
 
     WindowToken(WindowManagerService service, IBinder _token, int type, boolean persistOnEmpty,
             DisplayContent dc, boolean ownerCanManageAppTokens, boolean roundedCornerOverlay) {
-        this(service, _token, type, persistOnEmpty, dc, ownerCanManageAppTokens,
+        this(service, _token, type, persistOnEmpty, dc, ownerCanManageAppTokens, INVALID_UID,
                 roundedCornerOverlay, false /* fromClientToken */);
     }
 
     WindowToken(WindowManagerService service, IBinder _token, int type, boolean persistOnEmpty,
-            DisplayContent dc, boolean ownerCanManageAppTokens, boolean roundedCornerOverlay,
-            boolean fromClientToken) {
+            DisplayContent dc, boolean ownerCanManageAppTokens, int ownerUid,
+            boolean roundedCornerOverlay, boolean fromClientToken) {
         super(service);
         token = _token;
         windowType = type;
         mPersistOnEmpty = persistOnEmpty;
         mOwnerCanManageAppTokens = ownerCanManageAppTokens;
+        mOwnerUid = ownerUid;
         mRoundedCornerOverlay = roundedCornerOverlay;
         mFromClientToken = fromClientToken;
         if (dc != null) {
             dc.addWindowToken(token, this);
         }
+        if (shouldReportToClient()) {
+            try {
+                mDeathRecipient = new DeathRecipient();
+                mDeathRecipient.linkToDeath();
+            } catch (RemoteException e) {
+                Slog.e(TAG, "Unable to add window token with type " + windowType + " on "
+                        + "display " + dc.getDisplayId(), e);
+                mDeathRecipient = null;
+                return;
+            }
+        }
     }
 
     void removeAllWindowsIfPossible() {
@@ -222,7 +265,7 @@
     }
 
     void setExiting() {
-        if (mChildren.size() == 0) {
+        if (isEmpty()) {
             super.removeImmediately();
             return;
         }
@@ -340,6 +383,21 @@
         // Needs to occur after the token is removed from the display above to avoid attempt at
         // duplicate removal of this window container from it's parent.
         super.removeImmediately();
+
+        reportWindowTokenRemovedToClient();
+    }
+
+    private void reportWindowTokenRemovedToClient() {
+        if (!shouldReportToClient()) {
+            return;
+        }
+        mDeathRecipient.unlinkToDeath();
+        IWindowToken windowTokenClient = IWindowToken.Stub.asInterface(token);
+        try {
+            windowTokenClient.onWindowTokenRemoved();
+        } catch (RemoteException e) {
+            ProtoLog.w(WM_ERROR, "Could not report token removal to the window token client.");
+        }
     }
 
     @Override
@@ -361,17 +419,9 @@
     }
 
     void reportConfigToWindowTokenClient() {
-        if (asActivityRecord() != null) {
-            // Activities are updated through ATM callbacks.
+        if (!shouldReportToClient()) {
             return;
         }
-
-        // Unfortunately, this WindowToken is not from WindowContext so it cannot handle
-        // its own configuration changes.
-        if (!mFromClientToken) {
-            return;
-        }
-
         final Configuration config = getConfiguration();
         final int displayId = getDisplayContent().getDisplayId();
         if (config.equals(mLastReportedConfig) && displayId == mLastReportedDisplay) {
@@ -383,16 +433,26 @@
         mLastReportedDisplay = displayId;
 
         IWindowToken windowTokenClient = IWindowToken.Stub.asInterface(token);
-        if (windowTokenClient != null) {
-            try {
-                windowTokenClient.onConfigurationChanged(config, displayId);
-            } catch (RemoteException e) {
-                ProtoLog.w(WM_ERROR,
-                        "Could not report config changes to the window token client.");
-            }
+        try {
+            windowTokenClient.onConfigurationChanged(config, displayId);
+        } catch (RemoteException e) {
+            ProtoLog.w(WM_ERROR,
+                    "Could not report config changes to the window token client.");
         }
     }
 
+    /**
+     * @return {@code true} if this {@link WindowToken} is not an {@link ActivityRecord} and
+     * registered from client side.
+     */
+    private boolean shouldReportToClient() {
+        // Only report to client for WindowToken because Activities are updated through ATM
+        // callbacks.
+        return asActivityRecord() == null
+        // Report to {@link android.view.WindowTokenClient} if this token was registered from it.
+                && mFromClientToken && !mBinderDied;
+    }
+
     @Override
     void assignLayer(SurfaceControl.Transaction t, int layer) {
         if (windowType == TYPE_DOCK_DIVIDER) {
@@ -615,4 +675,8 @@
     int getWindowLayerFromType() {
         return mWmService.mPolicy.getWindowLayerFromTypeLw(windowType, mOwnerCanManageAppTokens);
     }
+
+    int getOwnerUid() {
+        return mOwnerUid;
+    }
 }
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowManagerServiceTests.java b/services/tests/wmtests/src/com/android/server/wm/WindowManagerServiceTests.java
index da4bde5..79b9ae1 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowManagerServiceTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowManagerServiceTests.java
@@ -16,13 +16,18 @@
 
 package com.android.server.wm;
 
+import static android.os.Process.INVALID_UID;
 import static android.view.WindowManager.LayoutParams.TYPE_TOAST;
 
 import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
 
+import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
+
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 import android.content.pm.PackageManager;
 import android.os.IBinder;
@@ -85,4 +90,19 @@
         assertFalse(windowToken.mRoundedCornerOverlay);
         assertTrue(windowToken.mFromClientToken);
     }
+
+    @Test(expected = SecurityException.class)
+    public void testRemoveWindowToken_ownerUidNotMatch_throwException() {
+        IBinder token = mock(IBinder.class);
+        mWm.addWindowTokenWithOptions(token, TYPE_TOAST, mDisplayContent.getDisplayId(),
+                null /* options */, null /* options */);
+
+        spyOn(mWm);
+        when(mWm.checkCallingPermission(anyString(), anyString())).thenReturn(false);
+        WindowToken windowToken = mWm.mRoot.getWindowToken(token);
+        spyOn(windowToken);
+        when(windowToken.getOwnerUid()).thenReturn(INVALID_UID);
+
+        mWm.removeWindowToken(token, mDisplayContent.getDisplayId());
+    }
 }
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowTokenTests.java b/services/tests/wmtests/src/com/android/server/wm/WindowTokenTests.java
index 7a347cb..76479cb 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowTokenTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowTokenTests.java
@@ -16,6 +16,7 @@
 
 package com.android.server.wm;
 
+import static android.os.Process.INVALID_UID;
 import static android.view.WindowManager.LayoutParams.FIRST_SUB_WINDOW;
 import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION;
 import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION_OVERLAY;
@@ -152,7 +153,7 @@
 
         token = new WindowToken(mDisplayContent.mWmService, mock(IBinder.class), TYPE_TOAST,
                 true /* persistOnEmpty */, mDisplayContent, true /* ownerCanManageAppTokens */,
-                true /* roundedCornerOverlay */, true /* fromClientToken */);
+                INVALID_UID, true /* roundedCornerOverlay */, true /* fromClientToken */);
         assertTrue(token.mRoundedCornerOverlay);
         assertTrue(token.mFromClientToken);
     }
@@ -166,7 +167,7 @@
     public void testSurfaceCreatedForWindowToken() {
         final WindowToken fromClientToken = new WindowToken(mDisplayContent.mWmService,
                 mock(IBinder.class), TYPE_APPLICATION_OVERLAY, true /* persistOnEmpty */,
-                mDisplayContent, true /* ownerCanManageAppTokens */,
+                mDisplayContent, true /* ownerCanManageAppTokens */, INVALID_UID,
                 true /* roundedCornerOverlay */, true /* fromClientToken */);
         assertNull(fromClientToken.mSurfaceControl);
 
@@ -175,7 +176,7 @@
 
         final WindowToken nonClientToken = new WindowToken(mDisplayContent.mWmService,
                 mock(IBinder.class), TYPE_TOAST, true /* persistOnEmpty */, mDisplayContent,
-                true /* ownerCanManageAppTokens */, true /* roundedCornerOverlay */,
+                true /* ownerCanManageAppTokens */, INVALID_UID, true /* roundedCornerOverlay */,
                 false /* fromClientToken */);
         assertNotNull(nonClientToken.mSurfaceControl);
     }
diff --git a/tests/permission/src/com/android/framework/permission/tests/WindowManagerPermissionTests.java b/tests/permission/src/com/android/framework/permission/tests/WindowManagerPermissionTests.java
index 4c2a984..737665f 100644
--- a/tests/permission/src/com/android/framework/permission/tests/WindowManagerPermissionTests.java
+++ b/tests/permission/src/com/android/framework/permission/tests/WindowManagerPermissionTests.java
@@ -17,6 +17,7 @@
 package com.android.framework.permission.tests;
 
 import static android.view.Display.DEFAULT_DISPLAY;
+import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION;
 
 import android.os.Binder;
 import android.os.RemoteException;
@@ -27,6 +28,8 @@
 
 import junit.framework.TestCase;
 
+import org.junit.Test;
+
 /**
  * TODO: Remove this. This is only a placeholder, need to implement this.
  */
@@ -53,7 +56,7 @@
         }
 
         try {
-            mWm.addWindowToken(null, 0, DEFAULT_DISPLAY);
+            mWm.addWindowToken(null, TYPE_APPLICATION, DEFAULT_DISPLAY);
             fail("IWindowManager.addWindowToken did not throw SecurityException as"
                     + " expected");
         } catch (SecurityException e) {
@@ -63,16 +66,6 @@
         }
 
         try {
-            mWm.removeWindowToken(null, DEFAULT_DISPLAY);
-            fail("IWindowManager.removeWindowToken did not throw SecurityException as"
-                    + " expected");
-        } catch (SecurityException e) {
-            // expected
-        } catch (RemoteException e) {
-            fail("Unexpected remote exception");
-        }
-
-        try {
             mWm.prepareAppTransition(0, false);
             fail("IWindowManager.prepareAppTransition did not throw SecurityException as"
                     + " expected");
@@ -182,4 +175,29 @@
             fail("Unexpected remote exception");
         }
     }
+
+    @Test
+    public void testADD_WINDOW_TOKEN_WITH_OPTIONS() {
+        // Verify if addWindowTokenWithOptions throw SecurityException for privileged window type.
+        try {
+            mWm.addWindowTokenWithOptions(null, TYPE_APPLICATION, DEFAULT_DISPLAY, null, "");
+            fail("IWindowManager.addWindowTokenWithOptions did not throw SecurityException as"
+                    + " expected");
+        } catch (SecurityException e) {
+            // expected
+        } catch (RemoteException e) {
+            fail("Unexpected remote exception");
+        }
+
+        // Verify if addWindowTokenWithOptions throw SecurityException for null packageName.
+        try {
+            mWm.addWindowTokenWithOptions(null, TYPE_APPLICATION, DEFAULT_DISPLAY, null, null);
+            fail("IWindowManager.addWindowTokenWithOptions did not throw SecurityException as"
+                    + " expected");
+        } catch (SecurityException e) {
+            // expected
+        } catch (RemoteException e) {
+            fail("Unexpected remote exception");
+        }
+    }
 }
diff --git a/tests/utils/testutils/java/com/android/server/wm/test/filters/FrameworksTestsFilter.java b/tests/utils/testutils/java/com/android/server/wm/test/filters/FrameworksTestsFilter.java
index 3026e0b..0dd45ba 100644
--- a/tests/utils/testutils/java/com/android/server/wm/test/filters/FrameworksTestsFilter.java
+++ b/tests/utils/testutils/java/com/android/server/wm/test/filters/FrameworksTestsFilter.java
@@ -47,7 +47,8 @@
             "android.view.InsetsSourceConsumerTest",
             "android.view.InsetsStateTest",
             "android.view.WindowMetricsTest",
-            "android.view.PendingInsetsControllerTest"
+            "android.view.PendingInsetsControllerTest",
+            "android.app.WindowContextTest"
     };
 
     public FrameworksTestsFilter(Bundle testArgs) {