[aapt2] Implement FeatureFlagsFilter

FeatureFlagsFilter looks for `android:featureFlag` attributes in the XML
and removes elements according to the given flag values.

Bug: 297373084
Test: atest aapt2_tests
Change-Id: Icf97de2152c0c53b0b4518ae837e0d198873d1c5
diff --git a/tools/aapt2/Android.bp b/tools/aapt2/Android.bp
index ed3e1ac..40cba3e 100644
--- a/tools/aapt2/Android.bp
+++ b/tools/aapt2/Android.bp
@@ -119,6 +119,7 @@
         "io/Util.cpp",
         "io/ZipArchive.cpp",
         "link/AutoVersioner.cpp",
+        "link/FeatureFlagsFilter.cpp",
         "link/ManifestFixer.cpp",
         "link/NoDefaultResourceRemover.cpp",
         "link/PrivateAttributeMover.cpp",
diff --git a/tools/aapt2/link/FeatureFlagsFilter.cpp b/tools/aapt2/link/FeatureFlagsFilter.cpp
new file mode 100644
index 0000000..fdf3f74
--- /dev/null
+++ b/tools/aapt2/link/FeatureFlagsFilter.cpp
@@ -0,0 +1,104 @@
+/*
+ * Copyright 2023 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.
+ */
+
+#include "link/FeatureFlagsFilter.h"
+
+#include <string_view>
+
+#include "androidfw/IDiagnostics.h"
+#include "androidfw/Source.h"
+#include "util/Util.h"
+#include "xml/XmlDom.h"
+#include "xml/XmlUtil.h"
+
+using ::aapt::xml::Element;
+using ::aapt::xml::Node;
+using ::aapt::xml::NodeCast;
+
+namespace aapt {
+
+class FlagsVisitor : public xml::Visitor {
+ public:
+  explicit FlagsVisitor(android::IDiagnostics* diagnostics,
+                        const FeatureFlagValues& feature_flag_values,
+                        const FeatureFlagsFilterOptions& options)
+      : diagnostics_(diagnostics), feature_flag_values_(feature_flag_values), options_(options) {
+  }
+
+  void Visit(xml::Element* node) override {
+    std::erase_if(node->children,
+                  [this](std::unique_ptr<xml::Node>& node) { return ShouldRemove(node); });
+    VisitChildren(node);
+  }
+
+  bool HasError() const {
+    return has_error_;
+  }
+
+ private:
+  bool ShouldRemove(std::unique_ptr<xml::Node>& node) {
+    if (const auto* el = NodeCast<Element>(node.get())) {
+      auto* attr = el->FindAttribute(xml::kSchemaAndroid, "featureFlag");
+      if (attr == nullptr) {
+        return false;
+      }
+
+      bool negated = false;
+      std::string_view flag_name = util::TrimWhitespace(attr->value);
+      if (flag_name.starts_with('!')) {
+        negated = true;
+        flag_name = flag_name.substr(1);
+      }
+
+      if (auto it = feature_flag_values_.find(std::string(flag_name));
+          it != feature_flag_values_.end()) {
+        if (it->second.has_value()) {
+          if (options_.remove_disabled_elements) {
+            // Remove if flag==true && attr=="!flag" (negated) OR flag==false && attr=="flag"
+            return *it->second == negated;
+          }
+        } else if (options_.flags_must_have_value) {
+          diagnostics_->Error(android::DiagMessage(node->line_number)
+                              << "attribute 'android:featureFlag' has flag '" << flag_name
+                              << "' without a true/false value from --feature_flags parameter");
+          has_error_ = true;
+          return false;
+        }
+      } else if (options_.fail_on_unrecognized_flags) {
+        diagnostics_->Error(android::DiagMessage(node->line_number)
+                            << "attribute 'android:featureFlag' has flag '" << flag_name
+                            << "' not found in flags from --feature_flags parameter");
+        has_error_ = true;
+        return false;
+      }
+    }
+
+    return false;
+  }
+
+  android::IDiagnostics* diagnostics_;
+  const FeatureFlagValues& feature_flag_values_;
+  const FeatureFlagsFilterOptions& options_;
+  bool has_error_ = false;
+};
+
+bool FeatureFlagsFilter::Consume(IAaptContext* context, xml::XmlResource* doc) {
+  FlagsVisitor visitor(context->GetDiagnostics(), feature_flag_values_, options_);
+  doc->root->Accept(&visitor);
+  return !visitor.HasError();
+}
+
+}  // namespace aapt
diff --git a/tools/aapt2/link/FeatureFlagsFilter.h b/tools/aapt2/link/FeatureFlagsFilter.h
new file mode 100644
index 0000000..1d342a7
--- /dev/null
+++ b/tools/aapt2/link/FeatureFlagsFilter.h
@@ -0,0 +1,79 @@
+/*
+ * Copyright 2023 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.
+ */
+
+#pragma once
+
+#include <optional>
+#include <string>
+#include <unordered_map>
+#include <utility>
+
+#include "android-base/macros.h"
+#include "cmd/Util.h"
+#include "process/IResourceTableConsumer.h"
+
+namespace aapt {
+
+struct FeatureFlagsFilterOptions {
+  // If true, elements whose featureFlag values are false (i.e., disabled feature) will be removed.
+  bool remove_disabled_elements = true;
+
+  // If true, `Consume()` will return false (error) if a flag was found that is not in
+  // `feature_flag_values`.
+  bool fail_on_unrecognized_flags = true;
+
+  // If true, `Consume()` will return false (error) if a flag was found whose value in
+  // `feature_flag_values` is not defined (std::nullopt).
+  bool flags_must_have_value = true;
+};
+
+// Looks for the `android:featureFlag` attribute in each XML element, validates the flag names and
+// values, and removes elements according to the values in `feature_flag_values`. An element will be
+// removed if the flag's given value is FALSE. A "!" before the flag name in the attribute indicates
+// a boolean NOT operation, i.e., an element will be removed if the flag's given value is TRUE. For
+// example, if the XML is the following:
+//
+//   <manifest xmlns:android="http://schemas.android.com/apk/res/android" package="android">
+//     <permission android:name="FOO" android:featureFlag="!flag"
+//                 android:protectionLevel="normal" />
+//     <permission android:name="FOO" android:featureFlag="flag"
+//                 android:protectionLevel="dangerous" />
+//   </manifest>
+//
+// If `feature_flag_values` contains {"flag", true}, then the <permission> element with
+// protectionLevel="normal" will be removed, and the <permission> element with
+// protectionLevel="normal" will be kept.
+//
+// The `Consume()` function will return false if there is an invalid flag found (see
+// FeatureFlagsFilterOptions for customizing the filter's validation behavior). Do not use the XML
+// further if there are errors as there may be elements removed already.
+class FeatureFlagsFilter : public IXmlResourceConsumer {
+ public:
+  explicit FeatureFlagsFilter(FeatureFlagValues feature_flag_values,
+                              FeatureFlagsFilterOptions options)
+      : feature_flag_values_(std::move(feature_flag_values)), options_(options) {
+  }
+
+  bool Consume(IAaptContext* context, xml::XmlResource* doc) override;
+
+ private:
+  DISALLOW_COPY_AND_ASSIGN(FeatureFlagsFilter);
+
+  const FeatureFlagValues feature_flag_values_;
+  const FeatureFlagsFilterOptions options_;
+};
+
+}  // namespace aapt
diff --git a/tools/aapt2/link/FeatureFlagsFilter_test.cpp b/tools/aapt2/link/FeatureFlagsFilter_test.cpp
new file mode 100644
index 0000000..53086cc
--- /dev/null
+++ b/tools/aapt2/link/FeatureFlagsFilter_test.cpp
@@ -0,0 +1,236 @@
+/*
+ * Copyright 2023 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.
+ */
+
+#include "link/FeatureFlagsFilter.h"
+
+#include <string_view>
+
+#include "test/Test.h"
+
+using ::testing::IsNull;
+using ::testing::NotNull;
+
+namespace aapt {
+
+// Returns null if there was an error from FeatureFlagsFilter.
+std::unique_ptr<xml::XmlResource> VerifyWithOptions(std::string_view str,
+                                                    const FeatureFlagValues& feature_flag_values,
+                                                    const FeatureFlagsFilterOptions& options) {
+  std::unique_ptr<xml::XmlResource> doc = test::BuildXmlDom(str);
+  FeatureFlagsFilter filter(feature_flag_values, options);
+  if (filter.Consume(test::ContextBuilder().Build().get(), doc.get())) {
+    return doc;
+  }
+  return {};
+}
+
+// Returns null if there was an error from FeatureFlagsFilter.
+std::unique_ptr<xml::XmlResource> Verify(std::string_view str,
+                                         const FeatureFlagValues& feature_flag_values) {
+  return VerifyWithOptions(str, feature_flag_values, {});
+}
+
+TEST(FeatureFlagsFilterTest, NoFeatureFlagAttributes) {
+  auto doc = Verify(R"EOF(
+    <manifest xmlns:android="http://schemas.android.com/apk/res/android" package="android">
+      <permission android:name="FOO" />
+    </manifest>)EOF",
+                    {{"flag", false}});
+  ASSERT_THAT(doc, NotNull());
+  auto root = doc->root.get();
+  ASSERT_THAT(root, NotNull());
+  auto maybe_removed = root->FindChild({}, "permission");
+  ASSERT_THAT(maybe_removed, NotNull());
+}
+TEST(FeatureFlagsFilterTest, RemoveElementWithDisabledFlag) {
+  auto doc = Verify(R"EOF(
+    <manifest xmlns:android="http://schemas.android.com/apk/res/android" package="android">
+      <permission android:name="FOO" android:featureFlag="flag" />
+    </manifest>)EOF",
+                    {{"flag", false}});
+  ASSERT_THAT(doc, NotNull());
+  auto root = doc->root.get();
+  ASSERT_THAT(root, NotNull());
+  auto maybe_removed = root->FindChild({}, "permission");
+  ASSERT_THAT(maybe_removed, IsNull());
+}
+
+TEST(FeatureFlagsFilterTest, RemoveElementWithNegatedEnabledFlag) {
+  auto doc = Verify(R"EOF(
+    <manifest xmlns:android="http://schemas.android.com/apk/res/android" package="android">
+      <permission android:name="FOO" android:featureFlag="!flag" />
+    </manifest>)EOF",
+                    {{"flag", true}});
+  ASSERT_THAT(doc, NotNull());
+  auto root = doc->root.get();
+  ASSERT_THAT(root, NotNull());
+  auto maybe_removed = root->FindChild({}, "permission");
+  ASSERT_THAT(maybe_removed, IsNull());
+}
+
+TEST(FeatureFlagsFilterTest, KeepElementWithEnabledFlag) {
+  auto doc = Verify(R"EOF(
+    <manifest xmlns:android="http://schemas.android.com/apk/res/android" package="android">
+      <permission android:name="FOO" android:featureFlag="flag" />
+    </manifest>)EOF",
+                    {{"flag", true}});
+  ASSERT_THAT(doc, NotNull());
+  auto root = doc->root.get();
+  ASSERT_THAT(root, NotNull());
+  auto maybe_removed = root->FindChild({}, "permission");
+  ASSERT_THAT(maybe_removed, NotNull());
+}
+
+TEST(FeatureFlagsFilterTest, SideBySideEnabledAndDisabled) {
+  auto doc = Verify(R"EOF(
+    <manifest xmlns:android="http://schemas.android.com/apk/res/android" package="android">
+      <permission android:name="FOO" android:featureFlag="!flag"
+                  android:protectionLevel="normal" />
+      <permission android:name="FOO" android:featureFlag="flag"
+                  android:protectionLevel="dangerous" />
+    </manifest>)EOF",
+                    {{"flag", true}});
+  ASSERT_THAT(doc, NotNull());
+  auto root = doc->root.get();
+  ASSERT_THAT(root, NotNull());
+  auto children = root->GetChildElements();
+  ASSERT_EQ(children.size(), 1);
+  auto attr = children[0]->FindAttribute(xml::kSchemaAndroid, "protectionLevel");
+  ASSERT_THAT(attr, NotNull());
+  ASSERT_EQ(attr->value, "dangerous");
+}
+
+TEST(FeatureFlagsFilterTest, RemoveDeeplyNestedElement) {
+  auto doc = Verify(R"EOF(
+    <manifest xmlns:android="http://schemas.android.com/apk/res/android" package="android">
+      <application>
+        <provider />
+        <activity>
+          <layout android:featureFlag="!flag" />
+        </activity>
+      </application>
+    </manifest>)EOF",
+                    {{"flag", true}});
+  ASSERT_THAT(doc, NotNull());
+  auto root = doc->root.get();
+  ASSERT_THAT(root, NotNull());
+  auto application = root->FindChild({}, "application");
+  ASSERT_THAT(application, NotNull());
+  auto activity = application->FindChild({}, "activity");
+  ASSERT_THAT(activity, NotNull());
+  auto maybe_removed = activity->FindChild({}, "layout");
+  ASSERT_THAT(maybe_removed, IsNull());
+}
+
+TEST(FeatureFlagsFilterTest, KeepDeeplyNestedElement) {
+  auto doc = Verify(R"EOF(
+    <manifest xmlns:android="http://schemas.android.com/apk/res/android" package="android">
+      <application>
+        <provider />
+        <activity>
+          <layout android:featureFlag="flag" />
+        </activity>
+      </application>
+    </manifest>)EOF",
+                    {{"flag", true}});
+  ASSERT_THAT(doc, NotNull());
+  auto root = doc->root.get();
+  ASSERT_THAT(root, NotNull());
+  auto application = root->FindChild({}, "application");
+  ASSERT_THAT(application, NotNull());
+  auto activity = application->FindChild({}, "activity");
+  ASSERT_THAT(activity, NotNull());
+  auto maybe_removed = activity->FindChild({}, "layout");
+  ASSERT_THAT(maybe_removed, NotNull());
+}
+
+TEST(FeatureFlagsFilterTest, FailOnEmptyFeatureFlagAttribute) {
+  auto doc = Verify(R"EOF(
+    <manifest xmlns:android="http://schemas.android.com/apk/res/android" package="android">
+      <permission android:name="FOO" android:featureFlag=" " />
+    </manifest>)EOF",
+                    {{"flag", false}});
+  ASSERT_THAT(doc, IsNull());
+}
+
+TEST(FeatureFlagsFilterTest, FailOnFlagWithNoGivenValue) {
+  auto doc = Verify(R"EOF(
+    <manifest xmlns:android="http://schemas.android.com/apk/res/android" package="android">
+      <permission android:name="FOO" android:featureFlag="flag" />
+    </manifest>)EOF",
+                    {{"flag", std::nullopt}});
+  ASSERT_THAT(doc, IsNull());
+}
+
+TEST(FeatureFlagsFilterTest, FailOnUnrecognizedFlag) {
+  auto doc = Verify(R"EOF(
+    <manifest xmlns:android="http://schemas.android.com/apk/res/android" package="android">
+      <permission android:name="FOO" android:featureFlag="unrecognized" />
+    </manifest>)EOF",
+                    {{"flag", true}});
+  ASSERT_THAT(doc, IsNull());
+}
+
+TEST(FeatureFlagsFilterTest, FailOnMultipleValidationErrors) {
+  auto doc = Verify(R"EOF(
+    <manifest xmlns:android="http://schemas.android.com/apk/res/android" package="android">
+      <permission android:name="FOO" android:featureFlag="bar" />
+      <permission android:name="FOO" android:featureFlag="unrecognized" />
+    </manifest>)EOF",
+                    {{"flag", std::nullopt}});
+  ASSERT_THAT(doc, IsNull());
+}
+
+TEST(FeatureFlagsFilterTest, OptionRemoveDisabledElementsIsFalse) {
+  auto doc = VerifyWithOptions(R"EOF(
+    <manifest xmlns:android="http://schemas.android.com/apk/res/android" package="android">
+      <permission android:name="FOO" android:featureFlag="flag" />
+    </manifest>)EOF",
+                               {{"flag", false}}, {.remove_disabled_elements = false});
+  ASSERT_THAT(doc, NotNull());
+  auto root = doc->root.get();
+  ASSERT_THAT(root, NotNull());
+  auto maybe_removed = root->FindChild({}, "permission");
+  ASSERT_THAT(maybe_removed, NotNull());
+}
+
+TEST(FeatureFlagsFilterTest, OptionFlagsMustHaveValueIsFalse) {
+  auto doc = VerifyWithOptions(R"EOF(
+    <manifest xmlns:android="http://schemas.android.com/apk/res/android" package="android">
+      <permission android:name="FOO" android:featureFlag="flag" />
+    </manifest>)EOF",
+                               {{"flag", std::nullopt}}, {.flags_must_have_value = false});
+  ASSERT_THAT(doc, NotNull());
+  auto root = doc->root.get();
+  ASSERT_THAT(root, NotNull());
+  auto maybe_removed = root->FindChild({}, "permission");
+  ASSERT_THAT(maybe_removed, NotNull());
+}
+
+TEST(FeatureFlagsFilterTest, OptionFailOnUnrecognizedFlagsIsFalse) {
+  auto doc = VerifyWithOptions(R"EOF(
+    <manifest xmlns:android="http://schemas.android.com/apk/res/android" package="android">
+      <permission android:name="FOO" android:featureFlag="unrecognized" />
+    </manifest>)EOF",
+                               {{"flag", true}}, {.fail_on_unrecognized_flags = false});
+  ASSERT_THAT(doc, NotNull());
+  auto root = doc->root.get();
+  ASSERT_THAT(root, NotNull());
+  auto maybe_removed = root->FindChild({}, "permission");
+  ASSERT_THAT(maybe_removed, NotNull());
+}
+
+}  // namespace aapt