Restrict size of custom views on SysUI side as well

Bug: 270553691
Flag: com.android.server.notification.notification_custom_view_uri_restriction
Test: newly added unit tests
Change-Id: I8e222355305098f0a466b41ac51a5252e2a4094a

13 files changed, 707 insertions, 8 deletions

diff --git a/Android.bp b/Android.bp
index 9d3b64d7335b..303fa2cd18da 100644
--- a/Android.bp
+++ b/Android.bp
@@ -583,6 +583,7 @@ java_library {
         "documents-ui-compat-config",
         "calendar-provider-compat-config",
         "contacts-provider-platform-compat-config",
+        "SystemUI-core-compat-config",
     ] + select(soong_config_variable("ANDROID", "release_crashrecovery_module"), {
         "true": [],
         default: [
diff --git a/packages/SystemUI/Android.bp b/packages/SystemUI/Android.bp
index 744388f47d0e..9c487be961dd 100644
--- a/packages/SystemUI/Android.bp
+++ b/packages/SystemUI/Android.bp
@@ -207,6 +207,8 @@ filegroup {
         "tests/src/**/systemui/statusbar/notification/row/NotificationConversationInfoTest.java",
         "tests/src/**/systemui/statusbar/notification/row/NotificationGutsManagerWithScenesTest.kt",
         "tests/src/**/systemui/statusbar/notification/row/wrapper/NotificationTemplateViewWrapperTest.kt",
+        "tests/src/**/systemui/statusbar/notification/row/NotificationCustomContentMemoryVerifierTest.java",
+        "tests/src/**/systemui/statusbar/notification/row/NotificationCustomContentMemoryVerifierDisabledTest.java",
         "tests/src/**/systemui/statusbar/notification/stack/NotificationStackScrollLayoutTest.java",
         "tests/src/**/systemui/statusbar/phone/CentralSurfacesImplTest.java",
         "tests/src/**/systemui/statusbar/phone/fragment/CollapsedStatusBarFragmentTest.java",
@@ -552,6 +554,11 @@ android_library {
     },
 }
 
+platform_compat_config {
+    name: "SystemUI-core-compat-config",
+    src: ":SystemUI-core",
+}
+
 filegroup {
     name: "AAA-src",
     srcs: ["tests/src/com/android/AAAPlusPlusVerifySysuiRequiredTestPropertiesTest.java"],
@@ -754,6 +761,7 @@ android_library {
         "kosmos",
         "testables",
         "androidx.test.rules",
+        "platform-compat-test-rules",
     ],
     libs: [
         "android.test.runner.stubs.system",
@@ -888,6 +896,7 @@ android_robolectric_test {
     static_libs: [
         "RoboTestLibraries",
         "androidx.compose.runtime_runtime",
+        "platform-compat-test-rules",
     ],
     libs: [
         "android.test.runner.impl",
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/NotificationEntry.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/NotificationEntry.java
index 09cc3f23032e..9dc651ed507a 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/NotificationEntry.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/NotificationEntry.java
@@ -643,6 +643,10 @@ public final class NotificationEntry extends ListEntry {
         return row.isMediaRow();
     }
 
+    public boolean containsCustomViews() {
+        return getSbn().getNotification().containsCustomViews();
+    }
+
     public void resetUserExpansion() {
         if (row != null) row.resetUserExpansion();
     }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryViewWalker.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryViewWalker.kt
index 6491223e6e10..f9e9bee4d809 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryViewWalker.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationMemoryViewWalker.kt
@@ -12,7 +12,7 @@ import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow
 import com.android.systemui.util.children
 
 /** Walks view hiearchy of a given notification to estimate its memory use. */
-internal object NotificationMemoryViewWalker {
+object NotificationMemoryViewWalker {
 
     private const val TAG = "NotificationMemory"
 
@@ -26,9 +26,13 @@ internal object NotificationMemoryViewWalker {
         private var softwareBitmaps = 0
 
         fun addSmallIcon(smallIconUse: Int) = apply { smallIcon += smallIconUse }
+
         fun addLargeIcon(largeIconUse: Int) = apply { largeIcon += largeIconUse }
+
         fun addSystem(systemIconUse: Int) = apply { systemIcons += systemIconUse }
+
         fun addStyle(styleUse: Int) = apply { style += styleUse }
+
         fun addSoftwareBitmapPenalty(softwareBitmapUse: Int) = apply {
             softwareBitmaps += softwareBitmapUse
         }
@@ -67,14 +71,14 @@ internal object NotificationMemoryViewWalker {
                     getViewUsage(ViewType.PRIVATE_EXPANDED_VIEW, row.privateLayout?.expandedChild),
                     getViewUsage(
                         ViewType.PRIVATE_CONTRACTED_VIEW,
-                        row.privateLayout?.contractedChild
+                        row.privateLayout?.contractedChild,
                     ),
                     getViewUsage(ViewType.PRIVATE_HEADS_UP_VIEW, row.privateLayout?.headsUpChild),
                     getViewUsage(
                         ViewType.PUBLIC_VIEW,
                         row.publicLayout?.expandedChild,
                         row.publicLayout?.contractedChild,
-                        row.publicLayout?.headsUpChild
+                        row.publicLayout?.headsUpChild,
                     ),
                 )
                 .filterNotNull()
@@ -107,14 +111,14 @@ internal object NotificationMemoryViewWalker {
             row.publicLayout?.expandedChild,
             row.publicLayout?.contractedChild,
             row.publicLayout?.headsUpChild,
-            seenObjects = seenObjects
+            seenObjects = seenObjects,
         )
     }
 
     private fun getViewUsage(
         type: ViewType,
         vararg rootViews: View?,
-        seenObjects: HashSet<Int> = hashSetOf()
+        seenObjects: HashSet<Int> = hashSetOf(),
     ): NotificationViewUsage? {
         val usageBuilder = lazy { UsageBuilder() }
         rootViews.forEach { rootView ->
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationContentInflater.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationContentInflater.java
index c7e15fdb98c7..73e8246907aa 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationContentInflater.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationContentInflater.java
@@ -901,6 +901,13 @@ public class NotificationContentInflater implements NotificationRowContentBinder
         if (!satisfiesMinHeightRequirement(view, entry, resources)) {
             return "inflated notification does not meet minimum height requirement";
         }
+
+        if (NotificationCustomContentMemoryVerifier.requiresImageViewMemorySizeCheck(entry)) {
+            if (!NotificationCustomContentMemoryVerifier.satisfiesMemoryLimits(view, entry)) {
+                return "inflated notification does not meet maximum memory size requirement";
+            }
+        }
+
         return null;
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationCustomContentCompat.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationCustomContentCompat.java
new file mode 100644
index 000000000000..c55cb6725e45
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationCustomContentCompat.java
@@ -0,0 +1,36 @@
+/*
+ * Copyright (C) 2025 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 com.android.systemui.statusbar.notification.row;
+
+import android.compat.annotation.ChangeId;
+import android.compat.annotation.EnabledAfter;
+import android.os.Build;
+
+/**
+ * Holds compat {@link ChangeId} for {@link NotificationCustomContentMemoryVerifier}.
+ */
+final class NotificationCustomContentCompat {
+    /**
+     * Enables memory size checking of custom views included in notifications to ensure that
+     * they conform to the size limit set in `config_notificationStripRemoteViewSizeBytes`
+     * config.xml parameter.
+     * Notifications exceeding the size will be rejected.
+     */
+    @ChangeId
+    @EnabledAfter(targetSdkVersion = Build.VERSION_CODES.BAKLAVA)
+    public static final long CHECK_SIZE_OF_INFLATED_CUSTOM_VIEWS = 270553691L;
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationCustomContentMemoryVerifier.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationCustomContentMemoryVerifier.kt
new file mode 100644
index 000000000000..a3e6a5cddc94
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationCustomContentMemoryVerifier.kt
@@ -0,0 +1,175 @@
+/*
+ * Copyright (C) 2024 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 com.android.systemui.statusbar.notification.row
+
+import android.app.compat.CompatChanges
+import android.content.Context
+import android.graphics.drawable.AdaptiveIconDrawable
+import android.graphics.drawable.BitmapDrawable
+import android.graphics.drawable.Drawable
+import android.os.Build
+import android.util.Log
+import android.view.View
+import android.view.ViewGroup
+import android.widget.ImageView
+import androidx.annotation.VisibleForTesting
+import com.android.app.tracing.traceSection
+import com.android.systemui.statusbar.notification.collection.NotificationEntry
+
+/** Checks whether Notifications with Custom content views conform to configured memory limits. */
+object NotificationCustomContentMemoryVerifier {
+
+    private const val NOTIFICATION_SERVICE_TAG = "NotificationService"
+
+    /** Notifications with custom views need to conform to maximum memory consumption. */
+    @JvmStatic
+    fun requiresImageViewMemorySizeCheck(entry: NotificationEntry): Boolean {
+        if (!com.android.server.notification.Flags.notificationCustomViewUriRestriction()) {
+            return false
+        }
+
+        return entry.containsCustomViews()
+    }
+
+    /**
+     * This walks the custom view hierarchy contained in the passed Notification view and determines
+     * if the total memory consumption of all image views satisfies the limit set by
+     * [getStripViewSizeLimit]. It will also log to logcat if the limit exceeds
+     * [getWarnViewSizeLimit].
+     *
+     * @return true if the Notification conforms to the view size limits.
+     */
+    @JvmStatic
+    fun satisfiesMemoryLimits(view: View, entry: NotificationEntry): Boolean {
+        val mainColumnView =
+            view.findViewById<View>(com.android.internal.R.id.notification_main_column)
+        if (mainColumnView == null) {
+            Log.wtf(
+                NOTIFICATION_SERVICE_TAG,
+                "R.id.notification_main_column view should not be null!",
+            )
+            return true
+        }
+
+        val memorySize =
+            traceSection("computeViewHiearchyImageViewSize") {
+                computeViewHierarchyImageViewSize(view)
+            }
+
+        if (memorySize > getStripViewSizeLimit(view.context)) {
+            val stripOversizedView = isCompatChangeEnabledForUid(entry.sbn.uid)
+            if (stripOversizedView) {
+                Log.w(
+                    NOTIFICATION_SERVICE_TAG,
+                    "Dropped notification due to too large RemoteViews ($memorySize bytes) on " +
+                        "pkg: ${entry.sbn.packageName} tag: ${entry.sbn.tag} id: ${entry.sbn.id}",
+                )
+            } else {
+                Log.w(
+                    NOTIFICATION_SERVICE_TAG,
+                    "RemoteViews too large on pkg: ${entry.sbn.packageName} " +
+                        "tag: ${entry.sbn.tag} id: ${entry.sbn.id} " +
+                        "this WILL notification WILL be dropped when targetSdk " +
+                        "is set to ${Build.VERSION_CODES.BAKLAVA}!",
+                )
+            }
+
+            // We still warn for size, but return "satisfies = ok" if the target SDK
+            // is too low.
+            return !stripOversizedView
+        }
+
+        if (memorySize > getWarnViewSizeLimit(view.context)) {
+            // We emit the same warning as NotificationManagerService does to keep some consistency
+            // for developers.
+            Log.w(
+                NOTIFICATION_SERVICE_TAG,
+                "RemoteViews too large on pkg: ${entry.sbn.packageName} " +
+                    "tag: ${entry.sbn.tag} id: ${entry.sbn.id} " +
+                    "this notifications might be dropped in a future release",
+            )
+        }
+        return true
+    }
+
+    private fun isCompatChangeEnabledForUid(uid: Int): Boolean =
+        try {
+            CompatChanges.isChangeEnabled(
+                NotificationCustomContentCompat.CHECK_SIZE_OF_INFLATED_CUSTOM_VIEWS,
+                uid,
+            )
+        } catch (e: RuntimeException) {
+            Log.wtf(NOTIFICATION_SERVICE_TAG, "Failed to contact system_server for compat change.")
+            false
+        }
+
+    @VisibleForTesting
+    @JvmStatic
+    fun computeViewHierarchyImageViewSize(view: View): Int =
+        when (view) {
+            is ViewGroup -> {
+                var use = 0
+                for (i in 0 until view.childCount) {
+                    use += computeViewHierarchyImageViewSize(view.getChildAt(i))
+                }
+                use
+            }
+            is ImageView -> computeImageViewSize(view)
+            else -> 0
+        }
+
+    /**
+     * Returns the memory size of a Bitmap contained in a passed [ImageView] in bytes. If the view
+     * contains any other kind of drawable, the memory size is estimated from its intrinsic
+     * dimensions.
+     *
+     * @return Bitmap size in bytes or 0 if no drawable is set.
+     */
+    private fun computeImageViewSize(view: ImageView): Int {
+        val drawable = view.drawable
+        return computeDrawableSize(drawable)
+    }
+
+    private fun computeDrawableSize(drawable: Drawable?): Int {
+        return when (drawable) {
+            null -> 0
+            is AdaptiveIconDrawable ->
+                computeDrawableSize(drawable.foreground) +
+                    computeDrawableSize(drawable.background) +
+                    computeDrawableSize(drawable.monochrome)
+            is BitmapDrawable -> drawable.bitmap.allocationByteCount
+            // People can sneak large drawables into those custom memory views via resources -
+            // we use the intrisic size as a proxy for how much memory rendering those will
+            // take.
+            else -> drawable.intrinsicWidth * drawable.intrinsicHeight * 4
+        }
+    }
+
+    /** @return Size of remote views after which a size warning is logged. */
+    @VisibleForTesting
+    fun getWarnViewSizeLimit(context: Context): Int =
+        context.resources.getInteger(
+            com.android.internal.R.integer.config_notificationWarnRemoteViewSizeBytes
+        )
+
+    /** @return Size of remote views after which the notification is dropped. */
+    @VisibleForTesting
+    fun getStripViewSizeLimit(context: Context): Int =
+        context.resources.getInteger(
+            com.android.internal.R.integer.config_notificationStripRemoteViewSizeBytes
+        )
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationRowContentBinderImpl.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationRowContentBinderImpl.kt
index bc3653a34fca..fda5e740ce52 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationRowContentBinderImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationRowContentBinderImpl.kt
@@ -1393,9 +1393,17 @@ constructor(
          */
         @VisibleForTesting
         fun isValidView(view: View, entry: NotificationEntry, resources: Resources): String? {
-            return if (!satisfiesMinHeightRequirement(view, entry, resources)) {
-                "inflated notification does not meet minimum height requirement"
-            } else null
+            if (!satisfiesMinHeightRequirement(view, entry, resources)) {
+                return "inflated notification does not meet minimum height requirement"
+            }
+
+            if (NotificationCustomContentMemoryVerifier.requiresImageViewMemorySizeCheck(entry)) {
+                if (!NotificationCustomContentMemoryVerifier.satisfiesMemoryLimits(view, entry)) {
+                    return "inflated notification does not meet maximum memory size requirement"
+                }
+            }
+
+            return null
         }
 
         private fun satisfiesMinHeightRequirement(
diff --git a/packages/SystemUI/tests/res/layout/custom_view_flipper.xml b/packages/SystemUI/tests/res/layout/custom_view_flipper.xml
new file mode 100644
index 000000000000..eb3ba82b043b
--- /dev/null
+++ b/packages/SystemUI/tests/res/layout/custom_view_flipper.xml
@@ -0,0 +1,14 @@
+<?xml version="1.0" encoding="utf-8"?>
+<FrameLayout xmlns:android="http://schemas.android.com/apk/res/android"
+    android:orientation="vertical"
+    android:layout_width="match_parent"
+    android:layout_height="match_parent">
+
+    <ViewFlipper
+        android:id="@+id/flipper"
+        android:layout_width="match_parent"
+        android:layout_height="400dp"
+        android:flipInterval="1000"
+        />
+
+</FrameLayout>
\ No newline at end of file
diff --git a/packages/SystemUI/tests/res/layout/custom_view_flipper_image.xml b/packages/SystemUI/tests/res/layout/custom_view_flipper_image.xml
new file mode 100644
index 000000000000..e2a00bd845cd
--- /dev/null
+++ b/packages/SystemUI/tests/res/layout/custom_view_flipper_image.xml
@@ -0,0 +1,5 @@
+<?xml version="1.0" encoding="utf-8"?>
+<ImageView xmlns:android="http://schemas.android.com/apk/res/android"
+    android:id="@+id/imageview"
+    android:layout_width="match_parent"
+    android:layout_height="400dp" />
\ No newline at end of file
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationCustomContentMemoryVerifierFlagDisabledTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationCustomContentMemoryVerifierFlagDisabledTest.java
new file mode 100644
index 000000000000..09fa3871f6e3
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationCustomContentMemoryVerifierFlagDisabledTest.java
@@ -0,0 +1,57 @@
+/*
+ * Copyright (C) 2025 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 com.android.systemui.statusbar.notification.row;
+
+import static com.android.systemui.statusbar.notification.row.NotificationCustomContentNotificationBuilder.buildAcceptableNotificationEntry;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import android.compat.testing.PlatformCompatChangeRule;
+import android.platform.test.annotations.DisableFlags;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
+import com.android.server.notification.Flags;
+import com.android.systemui.SysuiTestCase;
+import com.android.systemui.statusbar.notification.collection.NotificationEntry;
+
+import libcore.junit.util.compat.CoreCompatChangeRule.EnableCompatChanges;
+
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+@SmallTest
+@RunWith(AndroidJUnit4.class)
+public class NotificationCustomContentMemoryVerifierFlagDisabledTest extends SysuiTestCase {
+
+    @Rule
+    public PlatformCompatChangeRule mCompatChangeRule = new PlatformCompatChangeRule();
+
+    @Test
+    @DisableFlags(Flags.FLAG_NOTIFICATION_CUSTOM_VIEW_URI_RESTRICTION)
+    @EnableCompatChanges({
+            NotificationCustomContentCompat.CHECK_SIZE_OF_INFLATED_CUSTOM_VIEWS
+    })
+    public void requiresImageViewMemorySizeCheck_flagDisabled_returnsFalse() {
+        NotificationEntry entry = buildAcceptableNotificationEntry(mContext);
+        assertThat(NotificationCustomContentMemoryVerifier.requiresImageViewMemorySizeCheck(entry))
+                .isFalse();
+    }
+
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationCustomContentMemoryVerifierTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationCustomContentMemoryVerifierTest.java
new file mode 100644
index 000000000000..1cadb3c0a909
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationCustomContentMemoryVerifierTest.java
@@ -0,0 +1,285 @@
+/*
+ * Copyright (C) 2025 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 com.android.systemui.statusbar.notification.row;
+
+import static com.android.systemui.statusbar.notification.row.NotificationCustomContentNotificationBuilder.buildAcceptableNotification;
+import static com.android.systemui.statusbar.notification.row.NotificationCustomContentNotificationBuilder.buildAcceptableNotificationEntry;
+import static com.android.systemui.statusbar.notification.row.NotificationCustomContentNotificationBuilder.buildOversizedNotification;
+import static com.android.systemui.statusbar.notification.row.NotificationCustomContentNotificationBuilder.buildWarningSizedNotification;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import android.app.Notification;
+import android.compat.testing.PlatformCompatChangeRule;
+import android.content.ContentProvider;
+import android.content.ContentValues;
+import android.content.Context;
+import android.content.pm.ProviderInfo;
+import android.content.res.AssetFileDescriptor;
+import android.database.Cursor;
+import android.net.Uri;
+import android.os.Bundle;
+import android.os.CancellationSignal;
+import android.os.ParcelFileDescriptor;
+import android.os.Process;
+import android.platform.test.annotations.EnableFlags;
+import android.view.View;
+import android.widget.FrameLayout;
+import android.widget.RemoteViews;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
+import com.android.server.notification.Flags;
+import com.android.systemui.SysuiTestCase;
+import com.android.systemui.statusbar.notification.collection.NotificationEntry;
+import com.android.systemui.statusbar.notification.collection.NotificationEntryBuilder;
+
+import libcore.junit.util.compat.CoreCompatChangeRule.DisableCompatChanges;
+import libcore.junit.util.compat.CoreCompatChangeRule.EnableCompatChanges;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.io.FileNotFoundException;
+
+@SmallTest
+@RunWith(AndroidJUnit4.class)
+@EnableFlags(Flags.FLAG_NOTIFICATION_CUSTOM_VIEW_URI_RESTRICTION)
+public class NotificationCustomContentMemoryVerifierTest extends SysuiTestCase {
+
+    private static final String AUTHORITY = "notification.memory.test.authority";
+    private static final Uri TEST_URI = new Uri.Builder()
+            .scheme("content")
+            .authority(AUTHORITY)
+            .path("path")
+            .build();
+
+    @Rule
+    public PlatformCompatChangeRule mCompatChangeRule = new PlatformCompatChangeRule();
+
+    @Before
+    public void setUp() {
+        TestImageContentProvider provider = new TestImageContentProvider(mContext);
+        mContext.getContentResolver().addProvider(AUTHORITY, provider);
+        provider.onCreate();
+    }
+
+    @Test
+    @EnableCompatChanges({
+            NotificationCustomContentCompat.CHECK_SIZE_OF_INFLATED_CUSTOM_VIEWS})
+    public void requiresImageViewMemorySizeCheck_customViewNotification_returnsTrue() {
+        NotificationEntry entry =
+                buildAcceptableNotificationEntry(
+                        mContext);
+        assertThat(NotificationCustomContentMemoryVerifier.requiresImageViewMemorySizeCheck(entry))
+                .isTrue();
+    }
+
+    @Test
+    @EnableCompatChanges({
+            NotificationCustomContentCompat.CHECK_SIZE_OF_INFLATED_CUSTOM_VIEWS})
+    public void requiresImageViewMemorySizeCheck_plainNotification_returnsFalse() {
+        Notification notification =
+                new Notification.Builder(mContext, "ChannelId")
+                        .setContentTitle("Just a notification")
+                        .setContentText("Yep")
+                        .build();
+        NotificationEntry entry = new NotificationEntryBuilder().setNotification(
+                notification).build();
+        assertThat(NotificationCustomContentMemoryVerifier.requiresImageViewMemorySizeCheck(entry))
+                .isFalse();
+    }
+
+
+    @Test
+    @EnableCompatChanges({
+            NotificationCustomContentCompat.CHECK_SIZE_OF_INFLATED_CUSTOM_VIEWS})
+    public void satisfiesMemoryLimits_smallNotification_returnsTrue() {
+        Notification.Builder notification =
+                buildAcceptableNotification(mContext,
+                        TEST_URI);
+        NotificationEntry entry = toEntry(notification);
+        View inflatedView = inflateNotification(notification);
+        assertThat(
+                NotificationCustomContentMemoryVerifier.satisfiesMemoryLimits(inflatedView, entry)
+        )
+                .isTrue();
+    }
+
+    @Test
+    @EnableCompatChanges({
+            NotificationCustomContentCompat.CHECK_SIZE_OF_INFLATED_CUSTOM_VIEWS})
+    public void satisfiesMemoryLimits_oversizedNotification_returnsFalse() {
+        Notification.Builder notification =
+                buildOversizedNotification(mContext,
+                        TEST_URI);
+        NotificationEntry entry = toEntry(notification);
+        View inflatedView = inflateNotification(notification);
+        assertThat(
+                NotificationCustomContentMemoryVerifier.satisfiesMemoryLimits(inflatedView, entry)
+        ).isFalse();
+    }
+
+    @Test
+    @DisableCompatChanges(
+            {NotificationCustomContentCompat.CHECK_SIZE_OF_INFLATED_CUSTOM_VIEWS}
+    )
+    public void satisfiesMemoryLimits_oversizedNotification_compatDisabled_returnsTrue() {
+        Notification.Builder notification =
+                buildOversizedNotification(mContext,
+                        TEST_URI);
+        NotificationEntry entry = toEntry(notification);
+        View inflatedView = inflateNotification(notification);
+        assertThat(
+                NotificationCustomContentMemoryVerifier.satisfiesMemoryLimits(inflatedView, entry)
+        ).isTrue();
+    }
+
+    @Test
+    @EnableCompatChanges({
+            NotificationCustomContentCompat.CHECK_SIZE_OF_INFLATED_CUSTOM_VIEWS})
+    public void satisfiesMemoryLimits_warningSizedNotification_returnsTrue() {
+        Notification.Builder notification =
+                buildWarningSizedNotification(mContext,
+                        TEST_URI);
+        NotificationEntry entry = toEntry(notification);
+        View inflatedView = inflateNotification(notification);
+        assertThat(
+                NotificationCustomContentMemoryVerifier.satisfiesMemoryLimits(inflatedView, entry)
+        )
+                .isTrue();
+    }
+
+    @Test
+    @EnableCompatChanges({
+            NotificationCustomContentCompat.CHECK_SIZE_OF_INFLATED_CUSTOM_VIEWS})
+    public void satisfiesMemoryLimits_viewWithoutCustomNotificationRoot_returnsTrue() {
+        NotificationEntry entry = new NotificationEntryBuilder().build();
+        View view = new FrameLayout(mContext);
+        assertThat(NotificationCustomContentMemoryVerifier.satisfiesMemoryLimits(view, entry))
+                .isTrue();
+    }
+
+    @Test
+    @EnableCompatChanges({
+            NotificationCustomContentCompat.CHECK_SIZE_OF_INFLATED_CUSTOM_VIEWS})
+    public void computeViewHierarchyImageViewSize_smallNotification_returnsSensibleValue() {
+        Notification.Builder notification =
+                buildAcceptableNotification(mContext,
+                        TEST_URI);
+        // This should have a size of a single image
+        View inflatedView = inflateNotification(notification);
+        assertThat(
+                NotificationCustomContentMemoryVerifier.computeViewHierarchyImageViewSize(
+                        inflatedView))
+                .isGreaterThan(170000);
+    }
+
+    private View inflateNotification(Notification.Builder builder) {
+        RemoteViews remoteViews = builder.createBigContentView();
+        return remoteViews.apply(mContext, new FrameLayout(mContext));
+    }
+
+    private NotificationEntry toEntry(Notification.Builder builder) {
+        return new NotificationEntryBuilder().setNotification(builder.build())
+                .setUid(Process.myUid()).build();
+    }
+
+
+    /** This provider serves the images for inflation. */
+    class TestImageContentProvider extends ContentProvider {
+
+        TestImageContentProvider(Context context) {
+            ProviderInfo info = new ProviderInfo();
+            info.authority = AUTHORITY;
+            info.exported = true;
+            attachInfoForTesting(context, info);
+            setAuthorities(AUTHORITY);
+        }
+
+        @Override
+        public boolean onCreate() {
+            return true;
+        }
+
+        @Override
+        public ParcelFileDescriptor openFile(Uri uri, String mode) {
+            return getContext().getResources().openRawResourceFd(
+                    NotificationCustomContentNotificationBuilder.getDRAWABLE_IMAGE_RESOURCE())
+                        .getParcelFileDescriptor();
+        }
+
+        @Override
+        public AssetFileDescriptor openTypedAssetFile(Uri uri, String mimeTypeFilter, Bundle opts) {
+            return getContext().getResources().openRawResourceFd(
+                    NotificationCustomContentNotificationBuilder.getDRAWABLE_IMAGE_RESOURCE());
+        }
+
+        @Override
+        public AssetFileDescriptor openTypedAssetFile(Uri uri, String mimeTypeFilter, Bundle opts,
+                CancellationSignal signal) throws FileNotFoundException {
+            return openTypedAssetFile(uri, mimeTypeFilter, opts);
+        }
+
+        @Override
+        public int delete(Uri uri, Bundle extras) {
+            return 0;
+        }
+
+        @Override
+        public int delete(Uri uri, String selection, String[] selectionArgs) {
+            return 0;
+        }
+
+        @Override
+        public String getType(Uri uri) {
+            return "image/png";
+        }
+
+        @Override
+        public Uri insert(Uri uri, ContentValues values) {
+            return null;
+        }
+
+        @Override
+        public Uri insert(Uri uri, ContentValues values, Bundle extras) {
+            return super.insert(uri, values, extras);
+        }
+
+        @Override
+        public Cursor query(Uri uri, String[] projection, Bundle queryArgs,
+                CancellationSignal cancellationSignal) {
+            return super.query(uri, projection, queryArgs, cancellationSignal);
+        }
+
+        @Override
+        public Cursor query(Uri uri, String[] projection, String selection, String[] selectionArgs,
+                String sortOrder) {
+            return null;
+        }
+
+        @Override
+        public int update(Uri uri, ContentValues values, String selection, String[] selectionArgs) {
+            return 0;
+        }
+    }
+
+
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationCustomContentNotificationBuilder.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationCustomContentNotificationBuilder.kt
new file mode 100644
index 000000000000..ca4f24da3c08
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationCustomContentNotificationBuilder.kt
@@ -0,0 +1,94 @@
+/*
+ * Copyright (C) 2024 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
+ */
+
+@file:JvmName("NotificationCustomContentNotificationBuilder")
+
+package com.android.systemui.statusbar.notification.row
+
+import android.app.Notification
+import android.app.Notification.DecoratedCustomViewStyle
+import android.content.Context
+import android.graphics.drawable.BitmapDrawable
+import android.net.Uri
+import android.os.Process
+import android.widget.RemoteViews
+import com.android.systemui.statusbar.notification.collection.NotificationEntry
+import com.android.systemui.statusbar.notification.collection.NotificationEntryBuilder
+import com.android.systemui.tests.R
+import org.hamcrest.Matchers.lessThan
+import org.junit.Assume.assumeThat
+
+public val DRAWABLE_IMAGE_RESOURCE = R.drawable.romainguy_rockaway
+
+fun buildAcceptableNotificationEntry(context: Context): NotificationEntry {
+    return NotificationEntryBuilder()
+        .setNotification(buildAcceptableNotification(context, null).build())
+        .setUid(Process.myUid())
+        .build()
+}
+
+fun buildAcceptableNotification(context: Context, uri: Uri?): Notification.Builder =
+    buildNotification(context, uri, 1)
+
+fun buildOversizedNotification(context: Context, uri: Uri): Notification.Builder {
+    val numImagesForOversize =
+        (NotificationCustomContentMemoryVerifier.getStripViewSizeLimit(context) /
+            drawableSizeOnDevice(context)) + 2
+    return buildNotification(context, uri, numImagesForOversize)
+}
+
+fun buildWarningSizedNotification(context: Context, uri: Uri): Notification.Builder {
+    val numImagesForOversize =
+        (NotificationCustomContentMemoryVerifier.getWarnViewSizeLimit(context) /
+            drawableSizeOnDevice(context)) + 1
+    // The size needs to be smaller than outright stripping size.
+    assumeThat(
+        numImagesForOversize * drawableSizeOnDevice(context),
+        lessThan(NotificationCustomContentMemoryVerifier.getStripViewSizeLimit(context)),
+    )
+    return buildNotification(context, uri, numImagesForOversize)
+}
+
+fun buildNotification(context: Context, uri: Uri?, numImages: Int): Notification.Builder {
+    val remoteViews = RemoteViews(context.packageName, R.layout.custom_view_flipper)
+    repeat(numImages) { i ->
+        val remoteViewFlipperImageView =
+            RemoteViews(context.packageName, R.layout.custom_view_flipper_image)
+
+        if (uri == null) {
+            remoteViewFlipperImageView.setImageViewResource(
+                R.id.imageview,
+                R.drawable.romainguy_rockaway,
+            )
+        } else {
+            val imageUri = uri.buildUpon().appendPath(i.toString()).build()
+            remoteViewFlipperImageView.setImageViewUri(R.id.imageview, imageUri)
+        }
+        remoteViews.addView(R.id.flipper, remoteViewFlipperImageView)
+    }
+
+    return Notification.Builder(context, "ChannelId")
+        .setSmallIcon(android.R.drawable.ic_info)
+        .setStyle(DecoratedCustomViewStyle())
+        .setCustomContentView(remoteViews)
+        .setCustomBigContentView(remoteViews)
+        .setContentTitle("This is a remote view!")
+}
+
+fun drawableSizeOnDevice(context: Context): Int {
+    val drawable = context.resources.getDrawable(DRAWABLE_IMAGE_RESOURCE)
+    return (drawable as BitmapDrawable).bitmap.allocationByteCount
+}