Filter overlays based on constraints during resource resolution

Keep AssetManager updated about displayId and deviceId changes so
that it can filter resource overlays based on displayId and deviceId
constraints.

Test: atest libandroidfw_tests
Test: atest FrameworkServicesTests
Test: atest CtsResourcesTestCases
Bug: 371801644
Flag: android.content.res.rro_constraints
Change-Id: I6163114b62501cb5f77531443e8995a83168701b
NO_IFTTT=Introduce overlay constraint type values in libandroidfw, which are already same as the ones in OverlayConstraint.java

10 files changed, 374 insertions, 68 deletions

diff --git a/cmds/idmap2/include/idmap2/Idmap.h b/cmds/idmap2/include/idmap2/Idmap.h
index 1f15daf1ba47..335dc97bf964 100644
--- a/cmds/idmap2/include/idmap2/Idmap.h
+++ b/cmds/idmap2/include/idmap2/Idmap.h
@@ -178,8 +178,8 @@ class IdmapHeader {
 };
 
 struct IdmapConstraint {
-  // Constraint type can be TYPE_DISPLAY_ID or TYP_DEVICE_ID, please refer
-  // to ConstraintType in OverlayConstraint.java
+  // Constraint type can be android::kOverlayConstraintTypeDisplayId or
+  // android::kOverlayConstraintTypeDeviceId
   uint32_t constraint_type;
   uint32_t constraint_value;
 
diff --git a/core/java/android/app/ContextImpl.java b/core/java/android/app/ContextImpl.java
index d8aa8b3df622..f191b80d2330 100644
--- a/core/java/android/app/ContextImpl.java
+++ b/core/java/android/app/ContextImpl.java
@@ -2975,6 +2975,13 @@ class ContextImpl extends Context {
         if (display != null) {
             updateDeviceIdIfChanged(display.getDisplayId());
         }
+        updateResourceOverlayConstraints();
+    }
+
+    private void updateResourceOverlayConstraints() {
+        if (mResources != null) {
+            mResources.getAssets().setOverlayConstraints(getDisplayId(), getDeviceId());
+        }
     }
 
     @Override
@@ -2987,9 +2994,11 @@ class ContextImpl extends Context {
             }
         }
 
-        return new ContextImpl(this, mMainThread, mPackageInfo, mParams,
+        final ContextImpl context = new ContextImpl(this, mMainThread, mPackageInfo, mParams,
                 mAttributionSource.getAttributionTag(), mAttributionSource.getNext(), mSplitName,
                 mToken, mUser, mFlags, mClassLoader, null, deviceId, true);
+        context.updateResourceOverlayConstraints();
+        return context;
     }
 
     @NonNull
@@ -3284,6 +3293,7 @@ class ContextImpl extends Context {
             mDeviceId = updatedDeviceId;
             mAttributionSource = createAttributionSourceWithDeviceId(mAttributionSource, mDeviceId);
             notifyOnDeviceChangedListeners(updatedDeviceId);
+            updateResourceOverlayConstraints();
         }
     }
 
@@ -3686,6 +3696,7 @@ class ContextImpl extends Context {
                 mResourcesManager.setLocaleConfig(lc);
             }
         }
+        updateResourceOverlayConstraints();
     }
 
     void installSystemApplicationInfo(ApplicationInfo info, ClassLoader classLoader) {
diff --git a/core/java/android/content/res/AssetManager.java b/core/java/android/content/res/AssetManager.java
index 7cd2d31ac974..bcb50881d327 100644
--- a/core/java/android/content/res/AssetManager.java
+++ b/core/java/android/content/res/AssetManager.java
@@ -148,8 +148,8 @@ public final class AssetManager implements AutoCloseable {
      * @hide
      */
     public static class Builder {
-        private ArrayList<ApkAssets> mUserApkAssets = new ArrayList<>();
-        private ArrayList<ResourcesLoader> mLoaders = new ArrayList<>();
+        private final ArrayList<ApkAssets> mUserApkAssets = new ArrayList<>();
+        private final ArrayList<ResourcesLoader> mLoaders = new ArrayList<>();
 
         private boolean mNoInit = false;
 
@@ -1625,6 +1625,23 @@ public final class AssetManager implements AutoCloseable {
     }
 
     /**
+     * Passes the display id and device id to AssetManager, to filter out overlays based on
+     * any {@link android.content.om.OverlayConstraint}.
+     *
+     * @hide
+     */
+    public void setOverlayConstraints(int displayId, int deviceId) {
+        if (!Flags.rroConstraints()) {
+            return;
+        }
+
+        synchronized (this) {
+            ensureValidLocked();
+            nativeSetOverlayConstraints(mObject, displayId, deviceId);
+        }
+    }
+
+    /**
      * @hide
      */
     @UnsupportedAppUsage
@@ -1717,6 +1734,7 @@ public final class AssetManager implements AutoCloseable {
             int screenWidth, int screenHeight, int smallestScreenWidthDp, int screenWidthDp,
             int screenHeightDp, int screenLayout, int uiMode, int colorMode, int grammaticalGender,
             int majorVersion, boolean forceRefresh);
+    private static native void nativeSetOverlayConstraints(long ptr, int displayId, int deviceId);
     private static native @NonNull SparseArray<String> nativeGetAssignedPackageIdentifiers(
             long ptr, boolean includeOverlays, boolean includeLoaders);
 
diff --git a/core/jni/android_util_AssetManager.cpp b/core/jni/android_util_AssetManager.cpp
index b2649a471b8a..a73ff421acf7 100644
--- a/core/jni/android_util_AssetManager.cpp
+++ b/core/jni/android_util_AssetManager.cpp
@@ -424,6 +424,15 @@ static void NativeSetConfiguration(JNIEnv* env, jclass /*clazz*/, jlong ptr, jin
   assetmanager->SetDefaultLocale(default_locale_int);
 }
 
+static void NativeSetOverlayConstraints(JNIEnv* /*env*/, jclass /*clazz*/, jlong ptr,
+                                        jint displayId, jint deviceId) {
+    ATRACE_NAME("AssetManager::SetDisplayIdAndDeviceId");
+
+    auto assetmanager = LockAndStartAssetManager(ptr);
+    assetmanager->SetOverlayConstraints(static_cast<int32_t>(displayId),
+                                        static_cast<int32_t>(deviceId));
+}
+
 static jobject NativeGetAssignedPackageIdentifiers(JNIEnv* env, jclass /*clazz*/, jlong ptr,
                                                    jboolean includeOverlays,
                                                    jboolean includeLoaders) {
@@ -1554,6 +1563,7 @@ static const JNINativeMethod gAssetManagerMethods[] = {
         {"nativeSetApkAssets", "(J[Landroid/content/res/ApkAssets;ZZ)V", (void*)NativeSetApkAssets},
         {"nativeSetConfiguration", "(JIILjava/lang/String;[Ljava/lang/String;IIIIIIIIIIIIIIIIZ)V",
          (void*)NativeSetConfiguration},
+        {"nativeSetOverlayConstraints", "(JII)V", (void*)NativeSetOverlayConstraints},
         {"nativeGetAssignedPackageIdentifiers", "(JZZ)Landroid/util/SparseArray;",
          (void*)NativeGetAssignedPackageIdentifiers},
 
diff --git a/libs/androidfw/AssetManager2.cpp b/libs/androidfw/AssetManager2.cpp
index 0fa31c7a832e..12c298ed66bc 100644
--- a/libs/androidfw/AssetManager2.cpp
+++ b/libs/androidfw/AssetManager2.cpp
@@ -44,6 +44,9 @@ namespace android {
 
 namespace {
 
+constexpr int32_t kDefaultDisplayId = 0;
+constexpr int32_t kDefaultDeviceId = 0;
+
 using EntryValue = std::variant<Res_value, incfs::verified_map_ptr<ResTable_map_entry>>;
 
 /* NOTE: table_entry has been verified in LoadedPackage::GetEntryFromOffset(),
@@ -61,7 +64,7 @@ base::expected<EntryValue, IOError> GetEntryValue(
   return table_entry->value();
 }
 
-} // namespace
+}  // namespace
 
 struct FindEntryResult {
   // The cookie representing the ApkAssets in which the value resides.
@@ -99,14 +102,15 @@ struct Theme::Entry {
   Res_value value;
 };
 
-AssetManager2::AssetManager2(ApkAssetsList apk_assets, const ResTable_config& configuration) {
+AssetManager2::AssetManager2(ApkAssetsList apk_assets, const ResTable_config& configuration)
+  : display_id_(kDefaultDisplayId), device_id_(kDefaultDeviceId) {
   configurations_.push_back(configuration);
 
   // Don't invalidate caches here as there's nothing cached yet.
   SetApkAssets(apk_assets, false);
 }
 
-AssetManager2::AssetManager2() {
+AssetManager2::AssetManager2() : display_id_(kDefaultDisplayId), device_id_(kDefaultDeviceId) {
   configurations_.emplace_back();
 }
 
@@ -172,8 +176,7 @@ void AssetManager2::BuildDynamicRefTable(ApkAssetsList apk_assets) {
       // to take effect.
       auto iter = target_assets_package_ids.find(loaded_idmap->TargetApkPath());
       if (iter == target_assets_package_ids.end()) {
-         LOG(INFO) << "failed to find target package for overlay "
-                   << loaded_idmap->OverlayApkPath();
+        LOG(INFO) << "failed to find target package for overlay " << loaded_idmap->OverlayApkPath();
       } else {
         uint8_t target_package_id = iter->second;
 
@@ -189,10 +192,11 @@ void AssetManager2::BuildDynamicRefTable(ApkAssetsList apk_assets) {
                                   << " assigned package group";
 
         PackageGroup& target_package_group = package_groups_[target_idx];
-        target_package_group.overlays_.push_back(
-            ConfiguredOverlay{loaded_idmap->GetTargetResourcesMap(target_package_id,
-                                                                  overlay_ref_table.get()),
-                              apk_assets_cookies[apk_assets]});
+        target_package_group.overlays_.push_back(ConfiguredOverlay{
+                  loaded_idmap->GetTargetResourcesMap(target_package_id, overlay_ref_table.get()),
+                  apk_assets_cookies[apk_assets],
+                  IsAnyOverlayConstraintSatisfied(loaded_idmap->GetConstraints())
+        });
       }
     }
 
@@ -291,7 +295,7 @@ void AssetManager2::DumpToLog() const {
   }
   LOG(INFO) << "Package ID map: " << list;
 
-  for (const auto& package_group: package_groups_) {
+  for (const auto& package_group : package_groups_) {
     list = "";
     for (const auto& package : package_group.packages_) {
       const LoadedPackage* loaded_package = package.loaded_package_;
@@ -347,7 +351,6 @@ std::shared_ptr<const DynamicRefTable> AssetManager2::GetDynamicRefTableForCooki
 
 const std::unordered_map<std::string, std::string>*
   AssetManager2::GetOverlayableMapForPackage(uint32_t package_id) const {
-
   if (package_id >= package_ids_.size()) {
     return nullptr;
   }
@@ -462,6 +465,28 @@ void AssetManager2::SetConfigurations(std::span<const ResTable_config> configura
   }
 }
 
+void AssetManager2::SetOverlayConstraints(int32_t display_id, int32_t device_id) {
+  bool changed = false;
+  if (display_id_ != display_id) {
+    display_id_ = display_id;
+    changed = true;
+  }
+  if (device_id_ != device_id) {
+    device_id_ = device_id;
+    changed = true;
+  }
+  if (changed) {
+    // Enable/disable overlays based on current constraints
+    for (PackageGroup& group : package_groups_) {
+      for (auto &overlay: group.overlays_) {
+        overlay.enabled = IsAnyOverlayConstraintSatisfied(
+                overlay.overlay_res_maps_.GetConstraints());
+      }
+    }
+    InvalidateCaches(static_cast<uint32_t>(-1));
+  }
+}
+
 std::set<AssetManager2::ApkAssetsPtr> AssetManager2::GetNonSystemOverlays() const {
   std::set<ApkAssetsPtr> non_system_overlays;
   for (const PackageGroup& package_group : package_groups_) {
@@ -475,6 +500,8 @@ std::set<AssetManager2::ApkAssetsPtr> AssetManager2::GetNonSystemOverlays() cons
 
     if (!found_system_package) {
       auto op = StartOperation();
+      // Return all overlays, including the disabled ones as this is used for static info
+      // collection only.
       for (const ConfiguredOverlay& overlay : package_group.overlays_) {
         if (const auto& asset = GetApkAssets(overlay.cookie)) {
           non_system_overlays.insert(std::move(asset));
@@ -651,7 +678,6 @@ base::expected<FindEntryResult, NullOrIOError> AssetManager2::FindEntry(
 
   auto op = StartOperation();
 
-
   // Retrieve the package group from the package id of the resource id.
   if (UNLIKELY(!is_valid_resid(resid))) {
     LOG(ERROR) << base::StringPrintf("Invalid resource ID 0x%08x.", resid);
@@ -672,7 +698,7 @@ base::expected<FindEntryResult, NullOrIOError> AssetManager2::FindEntry(
   std::optional<FindEntryResult> final_result;
   bool final_has_locale = false;
   bool final_overlaid = false;
-  for (auto & config : configurations_) {
+  for (auto& config : configurations_) {
     // Might use this if density_override != 0.
     ResTable_config density_override_config;
 
@@ -698,7 +724,8 @@ base::expected<FindEntryResult, NullOrIOError> AssetManager2::FindEntry(
       }
       if (!assets->IsLoader()) {
         for (const auto& id_map : package_group.overlays_) {
-          auto overlay_entry = id_map.overlay_res_maps_.Lookup(resid);
+          auto overlay_entry = id_map.enabled ?
+                  id_map.overlay_res_maps_.Lookup(resid) : IdmapResMap::Result();
           if (!overlay_entry) {
             // No id map entry exists for this target resource.
             continue;
@@ -708,7 +735,7 @@ base::expected<FindEntryResult, NullOrIOError> AssetManager2::FindEntry(
             ConfigDescription best_frro_config;
             Res_value best_frro_value;
             bool frro_found = false;
-            for( const auto& [config, value] : overlay_entry.GetInlineValue()) {
+            for (const auto& [config, value] : overlay_entry.GetInlineValue()) {
               if ((!frro_found || config.isBetterThan(best_frro_config, desired_config))
                   && config.match(*desired_config)) {
                 frro_found = true;
@@ -1011,7 +1038,7 @@ std::string AssetManager2::GetLastResourceResolution() const {
                                      resid, resource_name_string.c_str(), conf.toString().c_str());
     char str[40];
     str[0] = '\0';
-    for(auto iter = configurations_.begin(); iter < configurations_.end(); iter++) {
+    for (auto iter = configurations_.begin(); iter < configurations_.end(); iter++) {
       iter->getBcp47Locale(str);
       log_stream << base::StringPrintf(" %s%s", str, iter < configurations_.end() ? "," : "");
     }
@@ -1480,7 +1507,7 @@ base::expected<uint32_t, NullOrIOError> AssetManager2::GetResourceId(
       base::expected<uint32_t, NullOrIOError> resid = package->FindEntryByName(type16, entry16);
       if (UNLIKELY(IsIOError(resid))) {
          return base::unexpected(resid.error());
-       }
+      }
 
       if (!resid.has_value() && kAttr16 == type16) {
         // Private attributes in libraries (such as the framework) are sometimes encoded
@@ -1505,7 +1532,7 @@ void AssetManager2::RebuildFilterList() {
       package.loaded_package_->ForEachTypeSpec([&](const TypeSpec& type_spec, uint8_t type_id) {
         FilteredConfigGroup* group = nullptr;
         for (const auto& type_entry : type_spec.type_entries) {
-          for (auto & config : configurations_) {
+          for (auto& config : configurations_) {
             if (type_entry.config.match(config)) {
               if (!group) {
                 group = &package.filtered_configs_.editItemAt(type_id - 1);
@@ -1522,6 +1549,27 @@ void AssetManager2::RebuildFilterList() {
   }
 }
 
+bool AssetManager2::IsAnyOverlayConstraintSatisfied(const Idmap_constraints& constraints) const {
+  if (constraints.constraint_count == 0) {
+    // There are no constraints, return true.
+    return true;
+  }
+
+  for (uint32_t i = 0; i < constraints.constraint_count; i++) {
+    auto constraint = constraints.constraint_entries[i];
+    if (constraint.constraint_type == kOverlayConstraintTypeDisplayId &&
+        constraint.constraint_value == display_id_) {
+      return true;
+    }
+    if (constraint.constraint_type == kOverlayConstraintTypeDeviceId &&
+        constraint.constraint_value == device_id_) {
+      return true;
+    }
+  }
+
+  return false;
+}
+
 void AssetManager2::InvalidateCaches(uint32_t diff) {
   cached_resolved_values_.clear();
 
diff --git a/libs/androidfw/Idmap.cpp b/libs/androidfw/Idmap.cpp
index f0ef97e5bdcc..8d1de1af56d2 100644
--- a/libs/androidfw/Idmap.cpp
+++ b/libs/androidfw/Idmap.cpp
@@ -56,13 +56,6 @@ struct Idmap_header {
   // without having to read/store each header entry separately.
 };
 
-struct Idmap_constraint {
-  // Constraint type can be TYPE_DISPLAY_ID or TYP_DEVICE_ID, please refer
-  // to ConstraintType in OverlayConstraint.java
-  uint32_t constraint_type;
-  uint32_t constraint_value;
-};
-
 struct Idmap_data_header {
   uint32_t target_entry_count;
   uint32_t target_inline_entry_count;
@@ -148,12 +141,13 @@ status_t OverlayDynamicRefTable::lookupResourceIdNoRewrite(uint32_t* resId) cons
   return DynamicRefTable::lookupResourceId(resId);
 }
 
-IdmapResMap::IdmapResMap(const Idmap_data_header* data_header, Idmap_target_entries entries,
-                         Idmap_target_inline_entries inline_entries,
+IdmapResMap::IdmapResMap(const Idmap_data_header* data_header, const Idmap_constraints& constraints,
+                         Idmap_target_entries entries, Idmap_target_inline_entries inline_entries,
                          const Idmap_target_entry_inline_value* inline_entry_values,
                          const ConfigDescription* configs, uint8_t target_assigned_package_id,
                          const OverlayDynamicRefTable* overlay_ref_table)
     : data_header_(data_header),
+      constraints_(constraints),
       entries_(entries),
       inline_entries_(inline_entries),
       inline_entry_values_(inline_entry_values),
@@ -254,7 +248,7 @@ std::optional<std::string_view> ReadString(const uint8_t** in_out_data_ptr, size
   }
   return std::string_view(data, *len);
 }
-} // namespace
+}  // namespace
 
 // O_PATH is a lightweight way of creating an FD, only exists on Linux
 #ifndef O_PATH
@@ -262,9 +256,7 @@ std::optional<std::string_view> ReadString(const uint8_t** in_out_data_ptr, size
 #endif
 
 LoadedIdmap::LoadedIdmap(const std::string& idmap_path, const Idmap_header* header,
-                         const Idmap_constraint* constraints,
-                         uint32_t constraints_count,
-                         const Idmap_data_header* data_header,
+                         const Idmap_data_header* data_header, const Idmap_constraints& constraints,
                          Idmap_target_entries target_entries,
                          Idmap_target_inline_entries target_inline_entries,
                          const Idmap_target_entry_inline_value* inline_entry_values,
@@ -272,9 +264,8 @@ LoadedIdmap::LoadedIdmap(const std::string& idmap_path, const Idmap_header* head
                          std::unique_ptr<ResStringPool>&& string_pool,
                          std::string_view overlay_apk_path, std::string_view target_apk_path)
     : header_(header),
-      constraints_(constraints),
-      constraints_count_(constraints_count),
       data_header_(data_header),
+      constraints_(constraints),
       target_entries_(target_entries),
       target_inline_entries_(target_inline_entries),
       inline_entry_values_(inline_entry_values),
@@ -328,16 +319,20 @@ std::unique_ptr<LoadedIdmap> LoadedIdmap::Load(StringPiece idmap_path, StringPie
     return {};
   }
 
-  auto constraints_count = ReadType<uint32_t>(&data_ptr, &data_size, "constraints count");
-  if (!constraints_count) {
+  auto constraint_count = ReadType<uint32_t>(&data_ptr, &data_size, "constraint count");
+  if (!constraint_count) {
+    LOG(ERROR) << "idmap doesn't have constraint count";
     return {};
   }
-  auto constraints = *constraints_count > 0 ?
-      ReadType<Idmap_constraint>(&data_ptr, &data_size, "constraints", *constraints_count)
+  auto constraint_entries = *constraint_count > 0 ?
+      ReadType<Idmap_constraint>(&data_ptr, &data_size, "constraints", dtohl(*constraint_count))
       : nullptr;
-  if (*constraints_count > 0 && !constraints) {
+  if (*constraint_count > 0 && !constraint_entries) {
+    LOG(ERROR) << "no constraint entries in idmap with non-zero constraints";
     return {};
   }
+  Idmap_constraints constraints{.constraint_count = *constraint_count,
+                                .constraint_entries = constraint_entries};
 
   // Parse the idmap data blocks. Currently idmap2 can only generate one data block.
   auto data_header = ReadType<Idmap_data_header>(&data_ptr, &data_size, "data header");
@@ -405,10 +400,9 @@ std::unique_ptr<LoadedIdmap> LoadedIdmap::Load(StringPiece idmap_path, StringPie
 
   // Can't use make_unique because LoadedIdmap constructor is private.
   return std::unique_ptr<LoadedIdmap>(
-      new LoadedIdmap(std::string(idmap_path), header, constraints, *constraints_count,
-                      data_header, target_entries, target_inline_entries,
-                      target_inline_entry_values,configurations, overlay_entries,
-                      std::move(idmap_string_pool),*overlay_path, *target_path));
+      new LoadedIdmap(std::string(idmap_path), header, data_header, constraints, target_entries,
+                      target_inline_entries, target_inline_entry_values, configurations,
+                      overlay_entries, std::move(idmap_string_pool), *overlay_path, *target_path));
 }
 
 UpToDate LoadedIdmap::IsUpToDate() const {
diff --git a/libs/androidfw/include/androidfw/AssetManager2.h b/libs/androidfw/include/androidfw/AssetManager2.h
index 0fdeefa09e26..a47fe6a7f50d 100644
--- a/libs/androidfw/include/androidfw/AssetManager2.h
+++ b/libs/androidfw/include/androidfw/AssetManager2.h
@@ -171,6 +171,8 @@ class AssetManager2 {
     default_locale_ = default_locale;
   }
 
+  void SetOverlayConstraints(int32_t display_id, int32_t device_id);
+
   // Returns all configurations for which there are resources defined, or an I/O error if reading
   // resource data failed.
   //
@@ -389,6 +391,9 @@ class AssetManager2 {
 
       // The cookie of the overlay assets.
       ApkAssetsCookie cookie;
+
+      // Enable/disable status of the overlay based on current constraints of AssetManager.
+      bool enabled;
   };
 
   // Represents a logical package, which can be made up of many individual packages. Each package
@@ -457,6 +462,8 @@ class AssetManager2 {
   // promoted apk assets when the last operation ends.
   void FinishOperation() const;
 
+  bool IsAnyOverlayConstraintSatisfied(const Idmap_constraints& constraints) const;
+
   // The ordered list of ApkAssets to search. These are not owned by the AssetManager, and must
   // have a longer lifetime.
   // The second pair element is the promoted version of the assets, that is held for the duration
@@ -480,6 +487,9 @@ class AssetManager2 {
   // may need to be purged.
   ftl::SmallVector<ResTable_config, 1> configurations_;
 
+  int32_t display_id_;
+  int32_t device_id_;
+
   // Cached set of bags. These are cached because they can inherit keys from parent bags,
   // which involves some calculation.
   mutable std::unordered_map<uint32_t, util::unique_cptr<ResolvedBag>> cached_bags_;
diff --git a/libs/androidfw/include/androidfw/Idmap.h b/libs/androidfw/include/androidfw/Idmap.h
index 0c0856315d8f..939b62462560 100644
--- a/libs/androidfw/include/androidfw/Idmap.h
+++ b/libs/androidfw/include/androidfw/Idmap.h
@@ -59,13 +59,25 @@ inline UpToDate fromBool(bool value) {
 class LoadedIdmap;
 class IdmapResMap;
 struct Idmap_header;
-struct Idmap_constraint;
+struct Idmap_constraints;
 struct Idmap_data_header;
-struct Idmap_target_entry;
 struct Idmap_target_entry_inline;
 struct Idmap_target_entry_inline_value;
-struct Idmap_overlay_entry;
 
+// LINT.IfChange
+constexpr int32_t kOverlayConstraintTypeDisplayId = 0;
+constexpr int32_t kOverlayConstraintTypeDeviceId = 1;
+// LINT.ThenChange(../../../../core/java/android/content/om/OverlayConstraint.java)
+
+struct Idmap_constraint {
+  // Constraint type can be kOverlayConstraintTypeDisplayId or kOverlayConstraintTypeDeviceId
+  const uint32_t constraint_type;
+  const uint32_t constraint_value;
+};
+struct Idmap_constraints {
+  const uint32_t constraint_count = 0;
+  const Idmap_constraint* constraint_entries = nullptr;
+};
 struct Idmap_target_entries {
   const uint32_t* target_id = nullptr;
   const uint32_t* overlay_id = nullptr;
@@ -169,14 +181,19 @@ class IdmapResMap {
     return overlay_ref_table_;
   }
 
+  inline Idmap_constraints GetConstraints() const {
+    return constraints_;
+  }
+
  private:
-  explicit IdmapResMap(const Idmap_data_header* data_header, Idmap_target_entries entries,
-                       Idmap_target_inline_entries inline_entries,
+  explicit IdmapResMap(const Idmap_data_header* data_header, const Idmap_constraints& constraints,
+                       Idmap_target_entries entries, Idmap_target_inline_entries inline_entries,
                        const Idmap_target_entry_inline_value* inline_entry_values,
                        const ConfigDescription* configs, uint8_t target_assigned_package_id,
                        const OverlayDynamicRefTable* overlay_ref_table);
 
   const Idmap_data_header* data_header_;
+  Idmap_constraints constraints_;
   Idmap_target_entries entries_;
   Idmap_target_inline_entries inline_entries_;
   const Idmap_target_entry_inline_value* inline_entry_values_;
@@ -210,8 +227,9 @@ class LoadedIdmap {
   // Returns a mapping from target resource ids to overlay values.
   IdmapResMap GetTargetResourcesMap(uint8_t target_assigned_package_id,
                                     const OverlayDynamicRefTable* overlay_ref_table) const {
-    return IdmapResMap(data_header_, target_entries_, target_inline_entries_, inline_entry_values_,
-                       configurations_, target_assigned_package_id, overlay_ref_table);
+    return IdmapResMap(data_header_, constraints_, target_entries_, target_inline_entries_,
+                       inline_entry_values_, configurations_, target_assigned_package_id,
+                       overlay_ref_table);
   }
 
   // Returns a dynamic reference table for a loaded overlay package.
@@ -223,14 +241,17 @@ class LoadedIdmap {
   // LoadedIdmap.
   UpToDate IsUpToDate() const;
 
+  inline const Idmap_constraints GetConstraints() const {
+    return constraints_;
+  }
+
  protected:
   // Exposed as protected so that tests can subclass and mock this class out.
   LoadedIdmap() = default;
 
   const Idmap_header* header_;
-  const Idmap_constraint* constraints_;
-  uint32_t constraints_count_;
   const Idmap_data_header* data_header_;
+  Idmap_constraints constraints_;
   Idmap_target_entries target_entries_;
   Idmap_target_inline_entries target_inline_entries_;
   const Idmap_target_entry_inline_value* inline_entry_values_;
@@ -247,9 +268,7 @@ class LoadedIdmap {
   DISALLOW_COPY_AND_ASSIGN(LoadedIdmap);
 
   explicit LoadedIdmap(const std::string& idmap_path, const Idmap_header* header,
-                       const Idmap_constraint* constraints,
-                       uint32_t constraints_count,
-                       const Idmap_data_header* data_header,
+                       const Idmap_data_header* data_header, const Idmap_constraints& constraints,
                        Idmap_target_entries target_entries,
                        Idmap_target_inline_entries target_inline_entries,
                        const Idmap_target_entry_inline_value* inline_entry_values_,
diff --git a/services/tests/servicestests/AndroidManifest.xml b/services/tests/servicestests/AndroidManifest.xml
index 834fea46e505..4531b3948495 100644
--- a/services/tests/servicestests/AndroidManifest.xml
+++ b/services/tests/servicestests/AndroidManifest.xml
@@ -199,6 +199,10 @@
         <service android:name="com.android.server.job.MockBiasJobService"
              android:permission="android.permission.BIND_JOB_SERVICE"/>
 
+        <activity
+            android:name="android.app.Activity"
+            android:exported="false" />
+
         <activity android:name="com.android.server.pm.BaseShortcutManagerTest$ShortcutActivity"/>
         <activity android:name="com.android.server.pm.BaseShortcutManagerTest$ShortcutActivity2"/>
         <activity android:name="com.android.server.pm.BaseShortcutManagerTest$ShortcutActivity3"/>
diff --git a/services/tests/servicestests/src/com/android/server/om/OverlayConstraintsTests.java b/services/tests/servicestests/src/com/android/server/om/OverlayConstraintsTests.java
index b2e296a36b93..2912a0762761 100644
--- a/services/tests/servicestests/src/com/android/server/om/OverlayConstraintsTests.java
+++ b/services/tests/servicestests/src/com/android/server/om/OverlayConstraintsTests.java
@@ -19,6 +19,7 @@ package com.android.server.om;
 import static android.content.Context.DEVICE_ID_DEFAULT;
 import static android.content.om.OverlayConstraint.TYPE_DEVICE_ID;
 import static android.content.om.OverlayConstraint.TYPE_DISPLAY_ID;
+import static android.util.TypedValue.TYPE_STRING;
 import static android.view.Display.DEFAULT_DISPLAY;
 
 import static androidx.test.core.app.ApplicationProvider.getApplicationContext;
@@ -28,6 +29,9 @@ import static junit.framework.Assert.assertNotNull;
 
 import static org.testng.Assert.assertThrows;
 
+import android.app.Activity;
+import android.companion.virtual.VirtualDeviceManager;
+import android.content.Context;
 import android.content.om.FabricatedOverlay;
 import android.content.om.OverlayConstraint;
 import android.content.om.OverlayIdentifier;
@@ -35,12 +39,12 @@ import android.content.om.OverlayInfo;
 import android.content.om.OverlayManager;
 import android.content.om.OverlayManagerTransaction;
 import android.content.res.Flags;
+import android.content.res.Resources;
 import android.os.UserHandle;
 import android.platform.test.annotations.RequiresFlagsDisabled;
 import android.platform.test.annotations.RequiresFlagsEnabled;
-import android.platform.test.flag.junit.CheckFlagsRule;
-import android.platform.test.flag.junit.DeviceFlagsValueProvider;
-import android.util.TypedValue;
+import android.view.Display;
+import android.virtualdevice.cts.common.VirtualDeviceRule;
 
 import junitparams.JUnitParamsRunner;
 import junitparams.Parameters;
@@ -53,20 +57,28 @@ import org.junit.runner.RunWith;
 
 import java.util.Collections;
 import java.util.List;
+import java.util.Objects;
+import java.util.concurrent.TimeoutException;
 
 @RunWith(JUnitParamsRunner.class)
 public class OverlayConstraintsTests {
+    private static final String RESOURCE_NAME = "string/module_2_name";
+    private static final String RESOURCE_DEFAULT_VALUE = "module_2_name";
+    private static final String RESOURCE_OVERLAID_VALUE = "hello";
+    private static final long TIMEOUT_MILLIS = 2000L;
 
     @Rule
-    public final CheckFlagsRule mCheckFlagsRule = DeviceFlagsValueProvider.createCheckFlagsRule();
+    public final VirtualDeviceRule mVirtualDeviceRule = VirtualDeviceRule.createDefault();
 
     private OverlayManager mOverlayManager;
     private UserHandle mUserHandle;
     private OverlayIdentifier mOverlayIdentifier = null;
+    private final String mPackageName = getApplicationContext().getPackageName();
 
     @Before
     public void setUp() throws Exception {
-        mOverlayManager = getApplicationContext().getSystemService(OverlayManager.class);
+        final Context context = getApplicationContext();
+        mOverlayManager = context.getSystemService(OverlayManager.class);
         mUserHandle = UserHandle.of(UserHandle.myUserId());
     }
 
@@ -79,6 +91,7 @@ public class OverlayConstraintsTests {
                             .build();
             mOverlayManager.commit(transaction);
             mOverlayIdentifier = null;
+            waitForResourceValue(RESOURCE_DEFAULT_VALUE, getApplicationContext());
         }
     }
 
@@ -161,13 +174,161 @@ public class OverlayConstraintsTests {
                 List.of(new OverlayConstraint(TYPE_DISPLAY_ID, DEFAULT_DISPLAY))));
     }
 
+    @Test
+    @RequiresFlagsEnabled(Flags.FLAG_RRO_CONSTRAINTS)
+    public void enableOverlayWithoutConstraints_appliesOverlayWithoutConstraints()
+            throws Exception {
+        enableOverlay(Collections.emptyList());
+
+        // Assert than the overlay is applied for both default device context and virtual
+        // device context.
+        final Context context = getApplicationContext();
+        waitForResourceValue(RESOURCE_OVERLAID_VALUE, context);
+        VirtualDeviceManager.VirtualDevice virtualDevice =
+                mVirtualDeviceRule.createManagedVirtualDevice();
+        final Context deviceContext = context.createDeviceContext(virtualDevice.getDeviceId());
+        waitForResourceValue(RESOURCE_OVERLAID_VALUE, deviceContext);
+    }
+
+    @Test
+    @RequiresFlagsEnabled(Flags.FLAG_RRO_CONSTRAINTS)
+    public void enableOverlayWithConstraints_withTypeDeviceId_appliesOverlayWithConstraints()
+            throws Exception {
+        final int deviceId1 = mVirtualDeviceRule.createManagedVirtualDevice().getDeviceId();
+        final int deviceId2 = mVirtualDeviceRule.createManagedVirtualDevice().getDeviceId();
+        enableOverlay(List.of(new OverlayConstraint(TYPE_DEVICE_ID, deviceId1),
+                new OverlayConstraint(TYPE_DEVICE_ID, deviceId2)));
+
+        // Assert than the overlay is not applied for contexts not associated with the above
+        // devices.
+        final Context context = getApplicationContext();
+        ensureResourceValueStaysAt(RESOURCE_DEFAULT_VALUE, context);
+        final int deviceId3 = mVirtualDeviceRule.createManagedVirtualDevice().getDeviceId();
+        ensureResourceValueStaysAt(RESOURCE_DEFAULT_VALUE, context.createDeviceContext(deviceId3));
+
+        // Assert than the overlay is applied for contexts associated with the above devices.
+        waitForResourceValue(RESOURCE_OVERLAID_VALUE, context.createDeviceContext(deviceId1));
+        waitForResourceValue(RESOURCE_OVERLAID_VALUE, context.createDeviceContext(deviceId2));
+    }
+
+    @Test
+    @RequiresFlagsEnabled(Flags.FLAG_RRO_CONSTRAINTS)
+    public void enableOverlayWithConstraints_withTypeDisplayId_appliesOverlayWithConstraints()
+            throws Exception {
+        final Display display1 =
+                mVirtualDeviceRule.createManagedUnownedVirtualDisplay().getDisplay();
+        final Display display2 =
+                mVirtualDeviceRule.createManagedUnownedVirtualDisplay().getDisplay();
+        enableOverlay(List.of(new OverlayConstraint(TYPE_DISPLAY_ID, display1.getDisplayId()),
+                new OverlayConstraint(TYPE_DISPLAY_ID, display2.getDisplayId())));
+
+        // Assert than the overlay is not applied for contexts not associated with the above
+        // displays.
+        final Context context = getApplicationContext();
+        ensureResourceValueStaysAt(RESOURCE_DEFAULT_VALUE, context);
+        final Display display3 =
+                mVirtualDeviceRule.createManagedUnownedVirtualDisplay().getDisplay();
+        ensureResourceValueStaysAt(RESOURCE_DEFAULT_VALUE, context.createDisplayContext(display3));
+
+        // Assert than the overlay is applied for contexts associated with the above displays.
+        waitForResourceValue(RESOURCE_OVERLAID_VALUE, context.createDisplayContext(display1));
+        waitForResourceValue(RESOURCE_OVERLAID_VALUE, context.createDisplayContext(display2));
+    }
+
+    @Test
+    @RequiresFlagsEnabled(Flags.FLAG_RRO_CONSTRAINTS)
+    public void enableOverlayWithConstraints_withTypesDisplayIdAndDeviceId_appliesOverlayWithConstraints()
+            throws Exception {
+        final Display display1 =
+                mVirtualDeviceRule.createManagedUnownedVirtualDisplay().getDisplay();
+        final int deviceId1 = mVirtualDeviceRule.createManagedVirtualDevice().getDeviceId();
+        enableOverlay(List.of(new OverlayConstraint(TYPE_DISPLAY_ID, display1.getDisplayId()),
+                new OverlayConstraint(TYPE_DEVICE_ID, deviceId1)));
+
+        // Assert than the overlay is not applied for contexts not associated with the above
+        // display or device.
+        final Context context = getApplicationContext();
+        ensureResourceValueStaysAt(RESOURCE_DEFAULT_VALUE, context);
+        final Display display2 =
+                mVirtualDeviceRule.createManagedUnownedVirtualDisplay().getDisplay();
+        ensureResourceValueStaysAt(RESOURCE_DEFAULT_VALUE, context.createDisplayContext(display2));
+        final int deviceId2 = mVirtualDeviceRule.createManagedVirtualDevice().getDeviceId();
+        ensureResourceValueStaysAt(RESOURCE_DEFAULT_VALUE, context.createDeviceContext(deviceId2));
+
+        // Assert than the overlay is applied for contexts associated with the above display or
+        // device.
+        waitForResourceValue(RESOURCE_OVERLAID_VALUE, context.createDisplayContext(display1));
+        waitForResourceValue(RESOURCE_OVERLAID_VALUE, context.createDeviceContext(deviceId1));
+    }
+
+    @Test
+    @RequiresFlagsEnabled(Flags.FLAG_RRO_CONSTRAINTS)
+    public void enableOverlayWithConstraints_withTypeDisplayId_appliesForActivityOnDisplay()
+            throws Exception {
+        final Display display =
+                mVirtualDeviceRule.createManagedUnownedVirtualDisplay(
+                                VirtualDeviceRule.createTrustedVirtualDisplayConfigBuilder())
+                        .getDisplay();
+        final Activity activityOnDefaultDisplay = mVirtualDeviceRule.startActivityOnDisplaySync(
+                DEFAULT_DISPLAY, Activity.class);
+        final Activity activityOnVirtualDisplay = mVirtualDeviceRule.startActivityOnDisplaySync(
+                display.getDisplayId(), Activity.class);
+
+        enableOverlay(List.of(new OverlayConstraint(TYPE_DISPLAY_ID, display.getDisplayId())));
+
+        // Assert than the overlay is not applied for any existing activity on the default display.
+        ensureResourceValueStaysAt(RESOURCE_DEFAULT_VALUE, activityOnDefaultDisplay);
+        // Assert than the overlay is applied for any existing activity on the virtual display.
+        waitForResourceValue(RESOURCE_OVERLAID_VALUE, activityOnVirtualDisplay);
+
+        // Assert than the overlay is not applied for any new activity on the default display.
+        final Activity newActivityOnDefaultDisplay = mVirtualDeviceRule.startActivityOnDisplaySync(
+                DEFAULT_DISPLAY, Activity.class);
+        ensureResourceValueStaysAt(RESOURCE_DEFAULT_VALUE, newActivityOnDefaultDisplay);
+        // Assert than the overlay is applied for any new activity on the virtual display.
+        final Activity newActivityOnVirtualDisplay = mVirtualDeviceRule.startActivityOnDisplaySync(
+                display.getDisplayId(), Activity.class);
+        waitForResourceValue(RESOURCE_OVERLAID_VALUE, newActivityOnVirtualDisplay);
+    }
+
+    @Test
+    @RequiresFlagsEnabled(Flags.FLAG_RRO_CONSTRAINTS)
+    public void enableOverlayWithConstraints_withTypeDeviceId_appliesForActivityOnDevice()
+            throws Exception {
+        final VirtualDeviceManager.VirtualDevice device =
+                mVirtualDeviceRule.createManagedVirtualDevice();
+        final Display display =
+                mVirtualDeviceRule.createManagedVirtualDisplay(device,
+                                VirtualDeviceRule.createTrustedVirtualDisplayConfigBuilder())
+                        .getDisplay();
+        final Activity activityOnDefaultDevice = mVirtualDeviceRule.startActivityOnDisplaySync(
+                DEFAULT_DISPLAY, Activity.class);
+        final Activity activityOnVirtualDevice = mVirtualDeviceRule.startActivityOnDisplaySync(
+                display.getDisplayId(), Activity.class);
+
+        enableOverlay(List.of(new OverlayConstraint(TYPE_DEVICE_ID, device.getDeviceId())));
+
+        // Assert than the overlay is not applied for any existing activity on the default device.
+        ensureResourceValueStaysAt(RESOURCE_DEFAULT_VALUE, activityOnDefaultDevice);
+        // Assert than the overlay is applied for any existing activity on the virtual device.
+        waitForResourceValue(RESOURCE_OVERLAID_VALUE, activityOnVirtualDevice);
+
+        // Assert than the overlay is not applied for any new activity on the default device.
+        final Activity newActivityOnDefaultDevice = mVirtualDeviceRule.startActivityOnDisplaySync(
+                DEFAULT_DISPLAY, Activity.class);
+        ensureResourceValueStaysAt(RESOURCE_DEFAULT_VALUE, newActivityOnDefaultDevice);
+        // Assert than the overlay is applied for any new activity on the virtual device.
+        final Activity newActivityOnVirtualDevice = mVirtualDeviceRule.startActivityOnDisplaySync(
+                display.getDisplayId(), Activity.class);
+        waitForResourceValue(RESOURCE_OVERLAID_VALUE, newActivityOnVirtualDevice);
+    }
+
     private FabricatedOverlay createFabricatedOverlay() {
-        String packageName = getApplicationContext().getPackageName();
         FabricatedOverlay fabricatedOverlay = new FabricatedOverlay.Builder(
-                packageName, "testOverlay" /* name */, packageName)
+                mPackageName, "testOverlay" /* name */, mPackageName)
                 .build();
-        fabricatedOverlay.setResourceValue("string/module_2_name" /* resourceName */,
-                TypedValue.TYPE_STRING, "hello" /* value */, null /* configuration */);
+        fabricatedOverlay.setResourceValue(RESOURCE_NAME, TYPE_STRING, RESOURCE_OVERLAID_VALUE,
+                null /* configuration */);
         return fabricatedOverlay;
     }
 
@@ -183,6 +344,37 @@ public class OverlayConstraintsTests {
         mOverlayIdentifier = fabricatedOverlay.getIdentifier();
     }
 
+    private static void waitForResourceValue(final String expectedValue, Context context)
+            throws TimeoutException {
+        final long endTime = System.currentTimeMillis() + TIMEOUT_MILLIS;
+        final Resources resources = context.getResources();
+        final int resourceId = getResourceId(context);
+        String resourceValue = null;
+        while (System.currentTimeMillis() < endTime) {
+            resourceValue = resources.getString(resourceId);
+            if (Objects.equals(resourceValue, expectedValue)) {
+                return;
+            }
+        }
+        throw new TimeoutException("Timed out waiting for '" + RESOURCE_NAME + "' value to equal '"
+                + expectedValue + "': current value is '" + resourceValue + "'");
+    }
+
+    private static void ensureResourceValueStaysAt(final String expectedValue, Context context) {
+        final long endTime = System.currentTimeMillis() + TIMEOUT_MILLIS;
+        final Resources resources = context.getResources();
+        final int resourceId = getResourceId(context);
+        String resourceValue;
+        while (System.currentTimeMillis() < endTime) {
+            resourceValue = resources.getString(resourceId);
+            assertEquals(expectedValue, resourceValue);
+        }
+    }
+
+    private static int getResourceId(Context context) {
+        return context.getResources().getIdentifier(RESOURCE_NAME, "", context.getPackageName());
+    }
+
     private static List<OverlayConstraint>[] getAllConstraintLists() {
         return new List[]{
                 Collections.emptyList(),