Add support for flagging xml and png files

This extends the previous change that added the ability to flag resource
directories so that xml and png files are now supported.

Test: Automated
Bug: 329436914
Flag: EXEMPT Aconfig not supported on host tools
Change-Id: I9f2b6b15ba0078ea33188f1a554377784cff9786

17 files changed, 133 insertions, 32 deletions

diff --git a/core/tests/resourceflaggingtests/src/com/android/resourceflaggingtests/ResourceFlaggingTest.java b/core/tests/resourceflaggingtests/src/com/android/resourceflaggingtests/ResourceFlaggingTest.java
index ee701f674715..d9e90fa001c5 100644
--- a/core/tests/resourceflaggingtests/src/com/android/resourceflaggingtests/ResourceFlaggingTest.java
+++ b/core/tests/resourceflaggingtests/src/com/android/resourceflaggingtests/ResourceFlaggingTest.java
@@ -29,6 +29,7 @@ import android.util.DisplayMetrics;
 import android.view.LayoutInflater;
 import android.view.View;
 import android.widget.LinearLayout;
+import android.widget.TextView;
 
 import androidx.test.InstrumentationRegistry;
 import androidx.test.filters.SmallTest;
@@ -132,6 +133,24 @@ public class ResourceFlaggingTest {
         assertThat((View) ll.findViewById(R.id.text2)).isNotNull();
     }
 
+    @Test
+    public void testEnabledFlagLayoutOverrides() {
+        LinearLayout ll = (LinearLayout) getLayoutInflater().inflate(R.layout.layout3, null);
+        assertThat(ll).isNotNull();
+        assertThat((View) ll.findViewById(R.id.text1)).isNotNull();
+        assertThat(((TextView) ll.findViewById(R.id.text1)).getText()).isEqualTo("foobar");
+    }
+
+    @Test(expected = Resources.NotFoundException.class)
+    public void testDisabledLayout() {
+        getLayoutInflater().inflate(R.layout.layout2, null);
+    }
+
+    @Test(expected = Resources.NotFoundException.class)
+    public void testDisabledDrawable() {
+        mResources.getDrawable(R.drawable.removedpng);
+    }
+
     private LayoutInflater getLayoutInflater() {
         ContextWrapper c = new ContextWrapper(mContext) {
             private LayoutInflater mInflater;
diff --git a/tools/aapt2/Resource.h b/tools/aapt2/Resource.h
index 446fdd4dbc1b..0d261abd728d 100644
--- a/tools/aapt2/Resource.h
+++ b/tools/aapt2/Resource.h
@@ -243,6 +243,12 @@ struct ResourceFile {
 
   // Exported symbols
   std::vector<SourcedResourceName> exported_symbols;
+
+  // Flag status
+  FlagStatus flag_status = FlagStatus::NoFlag;
+
+  // Flag
+  std::optional<FeatureFlagAttribute> flag;
 };
 
 /**
diff --git a/tools/aapt2/ResourceParser.cpp b/tools/aapt2/ResourceParser.cpp
index a2383ac0285f..fce6aa7c80d9 100644
--- a/tools/aapt2/ResourceParser.cpp
+++ b/tools/aapt2/ResourceParser.cpp
@@ -547,7 +547,7 @@ bool ResourceParser::ParseResource(xml::XmlPullParser* parser,
   });
 
   std::string_view resource_type = parser->element_name();
-  if (auto flag = GetFlag(parser)) {
+  if (auto flag = ParseFlag(xml::FindAttribute(parser, xml::kSchemaAndroid, "featureFlag"))) {
     if (options_.flag) {
       diag_->Error(android::DiagMessage(source_.WithLine(parser->line_number()))
                    << "Resource flag are not allowed both in the path and in the file");
@@ -747,22 +747,6 @@ bool ResourceParser::ParseResource(xml::XmlPullParser* parser,
   return false;
 }
 
-std::optional<FeatureFlagAttribute> ResourceParser::GetFlag(xml::XmlPullParser* parser) {
-  auto name = xml::FindAttribute(parser, xml::kSchemaAndroid, "featureFlag");
-  if (name) {
-    FeatureFlagAttribute flag;
-    if (name->starts_with('!')) {
-      flag.negated = true;
-      flag.name = name->substr(1);
-    } else {
-      flag.name = name.value();
-    }
-    return flag;
-  } else {
-    return {};
-  }
-}
-
 bool ResourceParser::ParseItem(xml::XmlPullParser* parser,
                                ParsedResource* out_resource,
                                const uint32_t format) {
@@ -1669,7 +1653,7 @@ bool ResourceParser::ParseArrayImpl(xml::XmlPullParser* parser,
     const std::string& element_namespace = parser->element_namespace();
     const std::string& element_name = parser->element_name();
     if (element_namespace.empty() && element_name == "item") {
-      auto flag = GetFlag(parser);
+      auto flag = ParseFlag(xml::FindAttribute(parser, xml::kSchemaAndroid, "featureFlag"));
       std::unique_ptr<Item> item = ParseXml(parser, typeMask, kNoRawString);
       if (!item) {
         diag_->Error(android::DiagMessage(item_source) << "could not parse array item");
diff --git a/tools/aapt2/ResourceParser.h b/tools/aapt2/ResourceParser.h
index 4ad334c94518..90690d522ef2 100644
--- a/tools/aapt2/ResourceParser.h
+++ b/tools/aapt2/ResourceParser.h
@@ -90,8 +90,6 @@ class ResourceParser {
  private:
   DISALLOW_COPY_AND_ASSIGN(ResourceParser);
 
-  std::optional<FeatureFlagAttribute> GetFlag(xml::XmlPullParser* parser);
-
   std::optional<FlattenedXmlSubTree> CreateFlattenSubTree(xml::XmlPullParser* parser);
 
   // Parses the XML subtree as a StyleString (flattened XML representation for strings with
diff --git a/tools/aapt2/ResourcesInternal.proto b/tools/aapt2/ResourcesInternal.proto
index b0ed3da33368..f4735a2f6ce7 100644
--- a/tools/aapt2/ResourcesInternal.proto
+++ b/tools/aapt2/ResourcesInternal.proto
@@ -49,4 +49,9 @@ message CompiledFile {
 
   // Any symbols this file auto-generates/exports (eg. @+id/foo in an XML file).
   repeated Symbol exported_symbol = 5;
+
+  // The status of the flag the file is behind if any
+  uint32 flag_status = 6;
+  bool flag_negated = 7;
+  string flag_name = 8;
 }
diff --git a/tools/aapt2/cmd/Compile.cpp b/tools/aapt2/cmd/Compile.cpp
index cd8c187c1e99..52372fa38525 100644
--- a/tools/aapt2/cmd/Compile.cpp
+++ b/tools/aapt2/cmd/Compile.cpp
@@ -196,20 +196,11 @@ static bool CompileTable(IAaptContext* context, const CompileOptions& options,
     // If visibility was forced, we need to use it when creating a new resource and also error if
     // we try to parse the <public>, <public-group>, <java-symbol> or <symbol> tags.
     parser_options.visibility = options.visibility;
+    parser_options.flag = ParseFlag(path_data.flag_name);
 
-    if (!path_data.flag_name.empty()) {
-      FeatureFlagAttribute flag;
-      const auto& name = path_data.flag_name;
-      if (name.starts_with('!')) {
-        flag.negated = true;
-        flag.name = name.substr(1);
-      } else {
-        flag.name = name;
-      }
-      parser_options.flag = flag;
-
+    if (parser_options.flag) {
       std::string error;
-      auto flag_status = GetFlagStatus(flag, options.feature_flag_values, &error);
+      auto flag_status = GetFlagStatus(parser_options.flag, options.feature_flag_values, &error);
       if (flag_status) {
         parser_options.flag_status = std::move(flag_status.value());
       } else {
@@ -443,6 +434,18 @@ static bool CompileXml(IAaptContext* context, const CompileOptions& options,
   xmlres->file.config = path_data.config;
   xmlres->file.source = path_data.source;
   xmlres->file.type = ResourceFile::Type::kProtoXml;
+  xmlres->file.flag = ParseFlag(path_data.flag_name);
+
+  if (xmlres->file.flag) {
+    std::string error;
+    auto flag_status = GetFlagStatus(xmlres->file.flag, options.feature_flag_values, &error);
+    if (flag_status) {
+      xmlres->file.flag_status = flag_status.value();
+    } else {
+      context->GetDiagnostics()->Error(android::DiagMessage(path_data.source) << error);
+      return false;
+    }
+  }
 
   // Collect IDs that are defined here.
   XmlIdCollector collector;
@@ -532,6 +535,27 @@ static bool CompilePng(IAaptContext* context, const CompileOptions& options,
   res_file.source = path_data.source;
   res_file.type = ResourceFile::Type::kPng;
 
+  if (!path_data.flag_name.empty()) {
+    FeatureFlagAttribute flag;
+    auto name = path_data.flag_name;
+    if (name.starts_with('!')) {
+      flag.negated = true;
+      flag.name = name.substr(1);
+    } else {
+      flag.name = name;
+    }
+    res_file.flag = flag;
+
+    std::string error;
+    auto flag_status = GetFlagStatus(flag, options.feature_flag_values, &error);
+    if (flag_status) {
+      res_file.flag_status = flag_status.value();
+    } else {
+      context->GetDiagnostics()->Error(android::DiagMessage(path_data.source) << error);
+      return false;
+    }
+  }
+
   {
     auto data = file->OpenAsData();
     if (!data) {
diff --git a/tools/aapt2/cmd/Util.cpp b/tools/aapt2/cmd/Util.cpp
index 2177c345260c..08f8f0d85807 100644
--- a/tools/aapt2/cmd/Util.cpp
+++ b/tools/aapt2/cmd/Util.cpp
@@ -34,6 +34,20 @@ using ::android::base::StringPrintf;
 
 namespace aapt {
 
+std::optional<FeatureFlagAttribute> ParseFlag(std::optional<std::string_view> flag_text) {
+  if (!flag_text || flag_text->empty()) {
+    return {};
+  }
+  FeatureFlagAttribute flag;
+  if (flag_text->starts_with('!')) {
+    flag.negated = true;
+    flag.name = flag_text->substr(1);
+  } else {
+    flag.name = flag_text.value();
+  }
+  return flag;
+}
+
 std::optional<FlagStatus> GetFlagStatus(const std::optional<FeatureFlagAttribute>& flag,
                                         const FeatureFlagValues& feature_flag_values,
                                         std::string* out_err) {
diff --git a/tools/aapt2/cmd/Util.h b/tools/aapt2/cmd/Util.h
index f8e44b72247c..d32e532b86a8 100644
--- a/tools/aapt2/cmd/Util.h
+++ b/tools/aapt2/cmd/Util.h
@@ -49,6 +49,8 @@ struct FeatureFlagProperties {
 
 using FeatureFlagValues = std::map<std::string, FeatureFlagProperties, std::less<>>;
 
+std::optional<FeatureFlagAttribute> ParseFlag(std::optional<std::string_view> flag_text);
+
 std::optional<FlagStatus> GetFlagStatus(const std::optional<FeatureFlagAttribute>& flag,
                                         const FeatureFlagValues& feature_flag_values,
                                         std::string* out_err);
diff --git a/tools/aapt2/format/proto/ProtoDeserialize.cpp b/tools/aapt2/format/proto/ProtoDeserialize.cpp
index c6bd6dd3c4b2..8583cadff6d2 100644
--- a/tools/aapt2/format/proto/ProtoDeserialize.cpp
+++ b/tools/aapt2/format/proto/ProtoDeserialize.cpp
@@ -643,6 +643,12 @@ bool DeserializeCompiledFileFromPb(const pb::internal::CompiledFile& pb_file,
   out_file->source.path = pb_file.source_path();
   out_file->type = DeserializeFileReferenceTypeFromPb(pb_file.type());
 
+  out_file->flag_status = (FlagStatus)pb_file.flag_status();
+  if (!pb_file.flag_name().empty()) {
+    out_file->flag =
+        FeatureFlagAttribute{.name = pb_file.flag_name(), .negated = pb_file.flag_negated()};
+  }
+
   std::string config_error;
   if (!DeserializeConfigFromPb(pb_file.config(), &out_file->config, &config_error)) {
     std::ostringstream error;
diff --git a/tools/aapt2/format/proto/ProtoSerialize.cpp b/tools/aapt2/format/proto/ProtoSerialize.cpp
index 9c28780e7880..d83fe916ee95 100644
--- a/tools/aapt2/format/proto/ProtoSerialize.cpp
+++ b/tools/aapt2/format/proto/ProtoSerialize.cpp
@@ -755,6 +755,11 @@ void SerializeCompiledFileToPb(const ResourceFile& file, pb::internal::CompiledF
   out_file->set_source_path(file.source.path);
   out_file->set_type(SerializeFileReferenceTypeToPb(file.type));
   SerializeConfig(file.config, out_file->mutable_config());
+  out_file->set_flag_status((uint32_t)file.flag_status);
+  if (file.flag) {
+    out_file->set_flag_negated(file.flag->negated);
+    out_file->set_flag_name(file.flag->name);
+  }
 
   for (const SourcedResourceName& exported : file.exported_symbols) {
     pb::internal::CompiledFile_Symbol* pb_symbol = out_file->add_exported_symbol();
diff --git a/tools/aapt2/integration-tests/FlaggedResourcesTest/Android.bp b/tools/aapt2/integration-tests/FlaggedResourcesTest/Android.bp
index 432b84e5ff90..7160b35033da 100644
--- a/tools/aapt2/integration-tests/FlaggedResourcesTest/Android.bp
+++ b/tools/aapt2/integration-tests/FlaggedResourcesTest/Android.bp
@@ -31,12 +31,17 @@ genrule {
         "res/values/ints.xml",
         "res/values/strings.xml",
         "res/layout/layout1.xml",
+        "res/layout/layout3.xml",
         "res/flag(test.package.falseFlag)/values/bools.xml",
+        "res/flag(test.package.falseFlag)/layout/layout2.xml",
+        "res/flag(test.package.falseFlag)/drawable/removedpng.png",
+        "res/flag(test.package.trueFlag)/layout/layout3.xml",
         "res/values/flag(test.package.trueFlag)/bools.xml",
         "res/values/flag(!test.package.trueFlag)/bools.xml",
         "res/values/flag(!test.package.falseFlag)/bools.xml",
     ],
     out: [
+        "drawable_removedpng.(test.package.falseFlag).png.flat",
         "values_bools.arsc.flat",
         "values_bools.(test.package.falseFlag).arsc.flat",
         "values_bools.(test.package.trueFlag).arsc.flat",
@@ -46,6 +51,9 @@ genrule {
         "values_ints.arsc.flat",
         "values_strings.arsc.flat",
         "layout_layout1.xml.flat",
+        "layout_layout2.(test.package.falseFlag).xml.flat",
+        "layout_layout3.xml.flat",
+        "layout_layout3.(test.package.trueFlag).xml.flat",
     ],
     cmd: "$(location aapt2) compile $(in) -o $(genDir) " +
         "--feature-flags test.package.falseFlag:ro=false,test.package.trueFlag:ro=true",
diff --git a/tools/aapt2/integration-tests/FlaggedResourcesTest/res/flag(test.package.falseFlag)/drawable/removedpng.png b/tools/aapt2/integration-tests/FlaggedResourcesTest/res/flag(test.package.falseFlag)/drawable/removedpng.png
new file mode 100644
index 000000000000..8a9e6984be96
--- /dev/null
+++ b/tools/aapt2/integration-tests/FlaggedResourcesTest/res/flag(test.package.falseFlag)/drawable/removedpng.png
diff --git a/tools/aapt2/integration-tests/FlaggedResourcesTest/res/flag(test.package.falseFlag)/layout/layout2.xml b/tools/aapt2/integration-tests/FlaggedResourcesTest/res/flag(test.package.falseFlag)/layout/layout2.xml
new file mode 100644
index 000000000000..dec5de72925a
--- /dev/null
+++ b/tools/aapt2/integration-tests/FlaggedResourcesTest/res/flag(test.package.falseFlag)/layout/layout2.xml
@@ -0,0 +1,6 @@
+<?xml version="1.0" encoding="utf-8"?>
+<LinearLayout xmlns:android="http://schemas.android.com/apk/res/android"
+    android:layout_width="match_parent"
+    android:layout_height="match_parent"
+    android:orientation="vertical" >
+</LinearLayout>
\ No newline at end of file
diff --git a/tools/aapt2/integration-tests/FlaggedResourcesTest/res/flag(test.package.trueFlag)/layout/layout3.xml b/tools/aapt2/integration-tests/FlaggedResourcesTest/res/flag(test.package.trueFlag)/layout/layout3.xml
new file mode 100644
index 000000000000..5aeee0ee1e28
--- /dev/null
+++ b/tools/aapt2/integration-tests/FlaggedResourcesTest/res/flag(test.package.trueFlag)/layout/layout3.xml
@@ -0,0 +1,10 @@
+<?xml version="1.0" encoding="utf-8"?>
+<LinearLayout xmlns:android="http://schemas.android.com/apk/res/android"
+    android:layout_width="match_parent"
+    android:layout_height="match_parent"
+    android:orientation="vertical" >
+    <TextView android:id="@+id/text1"
+        android:layout_width="wrap_content"
+        android:layout_height="wrap_content"
+        android:text="foobar" />
+</LinearLayout>
\ No newline at end of file
diff --git a/tools/aapt2/integration-tests/FlaggedResourcesTest/res/layout/layout3.xml b/tools/aapt2/integration-tests/FlaggedResourcesTest/res/layout/layout3.xml
new file mode 100644
index 000000000000..dec5de72925a
--- /dev/null
+++ b/tools/aapt2/integration-tests/FlaggedResourcesTest/res/layout/layout3.xml
@@ -0,0 +1,6 @@
+<?xml version="1.0" encoding="utf-8"?>
+<LinearLayout xmlns:android="http://schemas.android.com/apk/res/android"
+    android:layout_width="match_parent"
+    android:layout_height="match_parent"
+    android:orientation="vertical" >
+</LinearLayout>
\ No newline at end of file
diff --git a/tools/aapt2/link/FlaggedResources_test.cpp b/tools/aapt2/link/FlaggedResources_test.cpp
index 67d0c41f6b6f..629300838bbe 100644
--- a/tools/aapt2/link/FlaggedResources_test.cpp
+++ b/tools/aapt2/link/FlaggedResources_test.cpp
@@ -76,6 +76,10 @@ TEST_F(FlaggedResourcesTest, DisabledResourcesRemovedFromTable) {
 
   std::string output;
   DumpResourceTableToString(loaded_apk.get(), &output);
+  ASSERT_EQ(output.find("bool4"), std::string::npos);
+  ASSERT_EQ(output.find("str1"), std::string::npos);
+  ASSERT_EQ(output.find("layout2"), std::string::npos);
+  ASSERT_EQ(output.find("removedpng"), std::string::npos);
 }
 
 TEST_F(FlaggedResourcesTest, DisabledResourcesRemovedFromTableChunks) {
@@ -87,6 +91,8 @@ TEST_F(FlaggedResourcesTest, DisabledResourcesRemovedFromTableChunks) {
 
   ASSERT_EQ(output.find("bool4"), std::string::npos);
   ASSERT_EQ(output.find("str1"), std::string::npos);
+  ASSERT_EQ(output.find("layout2"), std::string::npos);
+  ASSERT_EQ(output.find("removedpng"), std::string::npos);
 }
 
 TEST_F(FlaggedResourcesTest, DisabledResourcesInRJava) {
diff --git a/tools/aapt2/link/TableMerger.cpp b/tools/aapt2/link/TableMerger.cpp
index d21697938c62..1bef5f8b17f6 100644
--- a/tools/aapt2/link/TableMerger.cpp
+++ b/tools/aapt2/link/TableMerger.cpp
@@ -377,6 +377,8 @@ bool TableMerger::MergeFile(const ResourceFile& file_desc, bool overlay, io::IFi
   file_ref->SetSource(file_desc.source);
   file_ref->type = file_desc.type;
   file_ref->file = file;
+  file_ref->SetFlagStatus(file_desc.flag_status);
+  file_ref->SetFlag(file_desc.flag);
 
   ResourceTablePackage* pkg = table.FindOrCreatePackage(file_desc.name.package);
   pkg->FindOrCreateType(file_desc.name.type)