Use lambda to refactor the obfuscating resource name.

There are 2 copy codes for handling the obfuscating resources names
between serializing to pb format file and writing to apks. The
obfuscator also needs to dump resources names. It means there are
3 places to handle the obfuscating resources names.

So, using C++ lambda to apply the callback mechanism refactors the
codes.

Obfuscator
* Initial a Obfuscator according to Optimizer's options
* Add Obfuscator.IsEnabled() function.
    return true either shorten_resource_paths_ is true or
    collapse_key_stringpool_ is true.

Bug: 228192695

Test: atest aapt2_test idmap2_test

Change-Id: Idd2442beecf41e9392620ff801a36fd1285e06f9
diff --git a/tools/aapt2/cmd/Optimize.cpp b/tools/aapt2/cmd/Optimize.cpp
index 042926c..ef2d70c 100644
--- a/tools/aapt2/cmd/Optimize.cpp
+++ b/tools/aapt2/cmd/Optimize.cpp
@@ -154,8 +154,8 @@
       return 1;
     }
 
-    if (options_.shorten_resource_paths) {
-      Obfuscator obfuscator(options_.table_flattener_options.shortened_path_map);
+    Obfuscator obfuscator(options_);
+    if (obfuscator.IsEnabled()) {
       if (!obfuscator.Consume(context_, apk->GetResourceTable())) {
         context_->GetDiagnostics()->Error(android::DiagMessage()
                                           << "failed shortening resource paths");
diff --git a/tools/aapt2/format/binary/TableFlattener.cpp b/tools/aapt2/format/binary/TableFlattener.cpp
index f192234..8c594ba 100644
--- a/tools/aapt2/format/binary/TableFlattener.cpp
+++ b/tools/aapt2/format/binary/TableFlattener.cpp
@@ -32,6 +32,7 @@
 #include "format/binary/ChunkWriter.h"
 #include "format/binary/ResEntryWriter.h"
 #include "format/binary/ResourceTypeExtensions.h"
+#include "optimize/Obfuscator.h"
 #include "trace/TraceBuffer.h"
 
 using namespace android;
@@ -466,9 +467,6 @@
       // table.
       std::map<ConfigDescription, std::vector<FlatEntry>> config_to_entry_list_map;
 
-      // hardcoded string uses characters which make it an invalid resource name
-      const std::string obfuscated_resource_name = "0_resource_name_obfuscated";
-
       for (const ResourceTableEntryView& entry : type.entries) {
         if (entry.staged_id) {
           aliases_.insert(std::make_pair(
@@ -477,30 +475,31 @@
         }
 
         uint32_t local_key_index;
-        ResourceName resource_name({}, type.named_type, entry.name);
-        if (!collapse_key_stringpool_ ||
-            name_collapse_exemptions_.find(resource_name) != name_collapse_exemptions_.end()) {
-          local_key_index = (uint32_t)key_pool_.MakeRef(entry.name).index();
-        } else {
-          // resource isn't exempt from collapse, add it as obfuscated value
-          if (entry.overlayable_item) {
+        auto onObfuscate = [this, &local_key_index, &entry](Obfuscator::Result obfuscatedResult,
+                                                            const ResourceName& resource_name) {
+          if (obfuscatedResult == Obfuscator::Result::Keep_ExemptionList) {
+            local_key_index = (uint32_t)key_pool_.MakeRef(entry.name).index();
+          } else if (obfuscatedResult == Obfuscator::Result::Keep_Overlayable) {
             // if the resource name of the specific entry is obfuscated and this
             // entry is in the overlayable list, the overlay can't work on this
             // overlayable at runtime because the name has been obfuscated in
             // resources.arsc during flatten operation.
             const OverlayableItem& item = entry.overlayable_item.value();
             context_->GetDiagnostics()->Warn(android::DiagMessage(item.overlayable->source)
-                                             << "The resource name of overlayable entry "
-                                             << resource_name.to_string() << "'"
-                                             << " shouldn't be obfuscated in resources.arsc");
+                                             << "The resource name of overlayable entry '"
+                                             << resource_name.to_string()
+                                             << "' shouldn't be obfuscated in resources.arsc");
 
             local_key_index = (uint32_t)key_pool_.MakeRef(entry.name).index();
           } else {
-            // TODO(b/228192695): output the entry.name and Resource id to make
-            //  de-obfuscated possible.
-            local_key_index = (uint32_t)key_pool_.MakeRef(obfuscated_resource_name).index();
+            local_key_index =
+                (uint32_t)key_pool_.MakeRef(Obfuscator::kObfuscatedResourceName).index();
           }
-        }
+        };
+
+        Obfuscator::ObfuscateResourceName(collapse_key_stringpool_, name_collapse_exemptions_,
+                                          type.named_type, entry, onObfuscate);
+
         // Group values by configuration.
         for (auto& config_value : entry.values) {
           config_to_entry_list_map[config_value->config].push_back(
diff --git a/tools/aapt2/format/binary/TableFlattener.h b/tools/aapt2/format/binary/TableFlattener.h
index 35254ba..60605d2 100644
--- a/tools/aapt2/format/binary/TableFlattener.h
+++ b/tools/aapt2/format/binary/TableFlattener.h
@@ -14,8 +14,13 @@
  * limitations under the License.
  */
 
-#ifndef AAPT_FORMAT_BINARY_TABLEFLATTENER_H
-#define AAPT_FORMAT_BINARY_TABLEFLATTENER_H
+#ifndef TOOLS_AAPT2_FORMAT_BINARY_TABLEFLATTENER_H_
+#define TOOLS_AAPT2_FORMAT_BINARY_TABLEFLATTENER_H_
+
+#include <map>
+#include <set>
+#include <string>
+#include <unordered_map>
 
 #include "Resource.h"
 #include "ResourceTable.h"
@@ -71,6 +76,9 @@
   //
   // This applies only to simple entries (entry->flags & ResTable_entry::FLAG_COMPLEX == 0).
   bool deduplicate_entry_values = false;
+
+  // Map from original resource ids to obfuscated names.
+  std::unordered_map<uint32_t, std::string> id_resource_map;
 };
 
 class TableFlattener : public IResourceTableConsumer {
@@ -82,12 +90,12 @@
   bool Consume(IAaptContext* context, ResourceTable* table) override;
 
  private:
-  DISALLOW_COPY_AND_ASSIGN(TableFlattener);
-
   TableFlattenerOptions options_;
   android::BigBuffer* buffer_;
+
+  DISALLOW_COPY_AND_ASSIGN(TableFlattener);
 };
 
 }  // namespace aapt
 
-#endif /* AAPT_FORMAT_BINARY_TABLEFLATTENER_H */
+#endif  // TOOLS_AAPT2_FORMAT_BINARY_TABLEFLATTENER_H_
diff --git a/tools/aapt2/format/proto/ProtoSerialize.cpp b/tools/aapt2/format/proto/ProtoSerialize.cpp
index a6d58fd..0e40124 100644
--- a/tools/aapt2/format/proto/ProtoSerialize.cpp
+++ b/tools/aapt2/format/proto/ProtoSerialize.cpp
@@ -18,6 +18,7 @@
 
 #include "ValueVisitor.h"
 #include "androidfw/BigBuffer.h"
+#include "optimize/Obfuscator.h"
 
 using android::ConfigDescription;
 
@@ -366,21 +367,21 @@
       }
       pb_type->set_name(type.named_type.to_string());
 
-      // hardcoded string uses characters which make it an invalid resource name
-      static const char* obfuscated_resource_name = "0_resource_name_obfuscated";
       for (const auto& entry : type.entries) {
         pb::Entry* pb_entry = pb_type->add_entry();
         if (entry.id) {
           pb_entry->mutable_entry_id()->set_id(entry.id.value());
         }
-        ResourceName resource_name({}, type.named_type, entry.name);
-        if (options.collapse_key_stringpool &&
-            options.name_collapse_exemptions.find(resource_name) ==
-            options.name_collapse_exemptions.end()) {
-          pb_entry->set_name(obfuscated_resource_name);
-        } else {
-          pb_entry->set_name(entry.name);
-        }
+        auto onObfuscate = [pb_entry, &entry](Obfuscator::Result obfuscatedResult,
+                                              const ResourceName& resource_name) {
+          pb_entry->set_name(obfuscatedResult == Obfuscator::Result::Obfuscated
+                                 ? Obfuscator::kObfuscatedResourceName
+                                 : entry.name);
+        };
+
+        Obfuscator::ObfuscateResourceName(options.collapse_key_stringpool,
+                                          options.name_collapse_exemptions, type.named_type, entry,
+                                          onObfuscate);
 
         // Write the Visibility struct.
         pb::Visibility* pb_visibility = pb_entry->mutable_visibility();
diff --git a/tools/aapt2/optimize/Obfuscator.cpp b/tools/aapt2/optimize/Obfuscator.cpp
index f704f26..b6df3d6 100644
--- a/tools/aapt2/optimize/Obfuscator.cpp
+++ b/tools/aapt2/optimize/Obfuscator.cpp
@@ -16,6 +16,7 @@
 
 #include "optimize/Obfuscator.h"
 
+#include <map>
 #include <set>
 #include <string>
 #include <unordered_set>
@@ -32,7 +33,10 @@
 
 namespace aapt {
 
-Obfuscator::Obfuscator(std::map<std::string, std::string>& path_map_out) : path_map_(path_map_out) {
+Obfuscator::Obfuscator(OptimizeOptions& optimizeOptions)
+    : options_(optimizeOptions.table_flattener_options),
+      shorten_resource_paths_(optimizeOptions.shorten_resource_paths),
+      collapse_key_stringpool_(optimizeOptions.table_flattener_options.collapse_key_stringpool) {
 }
 
 std::string ShortenFileName(const android::StringPiece& file_path, int output_length) {
@@ -77,7 +81,8 @@
   }
 };
 
-bool Obfuscator::Consume(IAaptContext* context, ResourceTable* table) {
+static bool HandleShortenFilePaths(ResourceTable* table,
+                                   std::map<std::string, std::string>& shortened_path_map) {
   // used to detect collisions
   std::unordered_set<std::string> shortened_paths;
   std::set<FileReference*, PathComparator> file_refs;
@@ -109,10 +114,94 @@
       shortened_path = GetShortenedPath(shortened_filename, extension, collision_count);
     }
     shortened_paths.insert(shortened_path);
-    path_map_.insert({*file_ref->path, shortened_path});
+    shortened_path_map.insert({*file_ref->path, shortened_path});
     file_ref->path = table->string_pool.MakeRef(shortened_path, file_ref->path.GetContext());
   }
   return true;
 }
 
+void Obfuscator::ObfuscateResourceName(
+    const bool collapse_key_stringpool, const std::set<ResourceName>& name_collapse_exemptions,
+    const ResourceNamedType& type_name, const ResourceTableEntryView& entry,
+    const android::base::function_ref<void(Result obfuscatedResult, const ResourceName&)>
+        onObfuscate) {
+  ResourceName resource_name({}, type_name, entry.name);
+  if (!collapse_key_stringpool ||
+      name_collapse_exemptions.find(resource_name) != name_collapse_exemptions.end()) {
+    onObfuscate(Result::Keep_ExemptionList, resource_name);
+  } else {
+    // resource isn't exempt from collapse, add it as obfuscated value
+    if (entry.overlayable_item) {
+      // if the resource name of the specific entry is obfuscated and this
+      // entry is in the overlayable list, the overlay can't work on this
+      // overlayable at runtime because the name has been obfuscated in
+      // resources.arsc during flatten operation.
+      onObfuscate(Result::Keep_Overlayable, resource_name);
+    } else {
+      onObfuscate(Result::Obfuscated, resource_name);
+    }
+  }
+}
+
+static bool HandleCollapseKeyStringPool(
+    const ResourceTable* table, const bool collapse_key_string_pool,
+    const std::set<ResourceName>& name_collapse_exemptions,
+    std::unordered_map<uint32_t, std::string>& id_resource_map) {
+  if (!collapse_key_string_pool) {
+    return true;
+  }
+
+  int entryResId = 0;
+  auto onObfuscate = [&entryResId, &id_resource_map](const Obfuscator::Result obfuscatedResult,
+                                                     const ResourceName& resource_name) {
+    if (obfuscatedResult == Obfuscator::Result::Obfuscated) {
+      id_resource_map.insert({entryResId, resource_name.entry});
+    }
+  };
+
+  for (auto& package : table->packages) {
+    for (auto& type : package->types) {
+      for (auto& entry : type->entries) {
+        if (!entry->id.has_value() || entry->name.empty()) {
+          continue;
+        }
+        entryResId = entry->id->id;
+        ResourceTableEntryView entry_view{
+            .name = entry->name,
+            .id = entry->id ? entry->id.value().entry_id() : (std::optional<uint16_t>)std::nullopt,
+            .visibility = entry->visibility,
+            .allow_new = entry->allow_new,
+            .overlayable_item = entry->overlayable_item,
+            .staged_id = entry->staged_id};
+
+        Obfuscator::ObfuscateResourceName(collapse_key_string_pool, name_collapse_exemptions,
+                                          type->named_type, entry_view, onObfuscate);
+      }
+    }
+  }
+
+  return true;
+}
+
+bool Obfuscator::Consume(IAaptContext* context, ResourceTable* table) {
+  HandleCollapseKeyStringPool(table, options_.collapse_key_stringpool,
+                              options_.name_collapse_exemptions, options_.id_resource_map);
+  if (shorten_resource_paths_) {
+    return HandleShortenFilePaths(table, options_.shortened_path_map);
+  }
+  return true;
+}
+
+/**
+ * Tell the optimizer whether it's needed to dump information for de-obfuscating.
+ *
+ * There are two conditions need to dump the information for de-obfuscating.
+ * * the option of shortening file paths is enabled.
+ * * the option of collapsing resource names is enabled.
+ * @return true if the information needed for de-obfuscating, otherwise false
+ */
+bool Obfuscator::IsEnabled() const {
+  return shorten_resource_paths_ || collapse_key_stringpool_;
+}
+
 }  // namespace aapt
diff --git a/tools/aapt2/optimize/Obfuscator.h b/tools/aapt2/optimize/Obfuscator.h
index 1ea32db..786bf8c 100644
--- a/tools/aapt2/optimize/Obfuscator.h
+++ b/tools/aapt2/optimize/Obfuscator.h
@@ -17,10 +17,14 @@
 #ifndef TOOLS_AAPT2_OPTIMIZE_OBFUSCATOR_H_
 #define TOOLS_AAPT2_OPTIMIZE_OBFUSCATOR_H_
 
-#include <map>
+#include <set>
 #include <string>
 
+#include "ResourceTable.h"
+#include "android-base/function_ref.h"
 #include "android-base/macros.h"
+#include "cmd/Optimize.h"
+#include "format/binary/TableFlattener.h"
 #include "process/IResourceTableConsumer.h"
 
 namespace aapt {
@@ -30,12 +34,26 @@
 // Maps resources in the apk to shortened paths.
 class Obfuscator : public IResourceTableConsumer {
  public:
-  explicit Obfuscator(std::map<std::string, std::string>& path_map_out);
+  explicit Obfuscator(OptimizeOptions& optimizeOptions);
 
   bool Consume(IAaptContext* context, ResourceTable* table) override;
 
+  bool IsEnabled() const;
+
+  enum class Result { Obfuscated, Keep_ExemptionList, Keep_Overlayable };
+
+  // hardcoded string uses characters which make it an invalid resource name
+  static constexpr char kObfuscatedResourceName[] = "0_resource_name_obfuscated";
+
+  static void ObfuscateResourceName(
+      const bool collapse_key_stringpool, const std::set<ResourceName>& name_collapse_exemptions,
+      const ResourceNamedType& type_name, const ResourceTableEntryView& entry,
+      const android::base::function_ref<void(Result, const ResourceName&)> onObfuscate);
+
  private:
-  std::map<std::string, std::string>& path_map_;
+  TableFlattenerOptions& options_;
+  const bool shorten_resource_paths_;
+  const bool collapse_key_stringpool_;
   DISALLOW_COPY_AND_ASSIGN(Obfuscator);
 };
 
diff --git a/tools/aapt2/optimize/Obfuscator_test.cpp b/tools/aapt2/optimize/Obfuscator_test.cpp
index a3339d4..17d1e52 100644
--- a/tools/aapt2/optimize/Obfuscator_test.cpp
+++ b/tools/aapt2/optimize/Obfuscator_test.cpp
@@ -16,6 +16,7 @@
 
 #include "optimize/Obfuscator.h"
 
+#include <map>
 #include <memory>
 #include <string>
 
@@ -51,8 +52,9 @@
           .AddString("android:string/string", "res/should/still/be/the/same.png")
           .Build();
 
-  std::map<std::string, std::string> path_map;
-  ASSERT_TRUE(Obfuscator(path_map).Consume(context.get(), table.get()));
+  OptimizeOptions options{.shorten_resource_paths = true};
+  std::map<std::string, std::string>& path_map = options.table_flattener_options.shortened_path_map;
+  ASSERT_TRUE(Obfuscator(options).Consume(context.get(), table.get()));
 
   // Expect that the path map is populated
   ASSERT_THAT(path_map.find("res/drawables/xmlfile.xml"), Not(Eq(path_map.end())));
@@ -87,8 +89,9 @@
                             test::ParseConfigOrDie("mdp-v21"))
           .Build();
 
-  std::map<std::string, std::string> path_map;
-  ASSERT_TRUE(Obfuscator(path_map).Consume(context.get(), table.get()));
+  OptimizeOptions options{.shorten_resource_paths = true};
+  std::map<std::string, std::string>& path_map = options.table_flattener_options.shortened_path_map;
+  ASSERT_TRUE(Obfuscator(options).Consume(context.get(), table.get()));
 
   // Expect that the path map to not contain the ColorStateList
   ASSERT_THAT(path_map.find("res/color/colorlist.xml"), Eq(path_map.end()));
@@ -107,8 +110,9 @@
           .AddFileReference("android:color/pngfile", original_png_path)
           .Build();
 
-  std::map<std::string, std::string> path_map;
-  ASSERT_TRUE(Obfuscator(path_map).Consume(context.get(), table.get()));
+  OptimizeOptions options{.shorten_resource_paths = true};
+  std::map<std::string, std::string>& path_map = options.table_flattener_options.shortened_path_map;
+  ASSERT_TRUE(Obfuscator(options).Consume(context.get(), table.get()));
 
   // Expect that the path map is populated
   ASSERT_THAT(path_map.find("res/drawable/xmlfile.xml"), Not(Eq(path_map.end())));
@@ -133,8 +137,10 @@
   test::ResourceTableBuilder builder1;
   FillTable(builder1, 0, kNumResources);
   std::unique_ptr<ResourceTable> table1 = builder1.Build();
-  std::map<std::string, std::string> expected_mapping;
-  ASSERT_TRUE(Obfuscator(expected_mapping).Consume(context.get(), table1.get()));
+  OptimizeOptions options{.shorten_resource_paths = true};
+  std::map<std::string, std::string>& expected_mapping =
+      options.table_flattener_options.shortened_path_map;
+  ASSERT_TRUE(Obfuscator(options).Consume(context.get(), table1.get()));
 
   // We are trying to ensure lack of non-determinism, it is not simple to prove
   // a negative, thus we must try the test a few times so that the test itself
@@ -153,8 +159,10 @@
     FillTable(builder2, 0, start_index);
     std::unique_ptr<ResourceTable> table2 = builder2.Build();
 
-    std::map<std::string, std::string> actual_mapping;
-    ASSERT_TRUE(Obfuscator(actual_mapping).Consume(context.get(), table2.get()));
+    OptimizeOptions actualOptimizerOptions{.shorten_resource_paths = true};
+    TableFlattenerOptions& actual_options = actualOptimizerOptions.table_flattener_options;
+    std::map<std::string, std::string>& actual_mapping = actual_options.shortened_path_map;
+    ASSERT_TRUE(Obfuscator(actualOptimizerOptions).Consume(context.get(), table2.get()));
 
     for (auto& item : actual_mapping) {
       ASSERT_THAT(expected_mapping[item.first], Eq(item.second));
@@ -162,4 +170,70 @@
   }
 }
 
+TEST(ObfuscatorTest, DumpIdResourceMap) {
+  std::unique_ptr<IAaptContext> context = test::ContextBuilder().Build();
+
+  OverlayableItem overlayable_item(std::make_shared<Overlayable>("TestName", "overlay://theme"));
+  overlayable_item.policies |= PolicyFlags::PRODUCT_PARTITION;
+  overlayable_item.policies |= PolicyFlags::SYSTEM_PARTITION;
+  overlayable_item.policies |= PolicyFlags::VENDOR_PARTITION;
+
+  std::string original_xml_path = "res/drawable/xmlfile.xml";
+  std::string original_png_path = "res/drawable/pngfile.png";
+
+  std::string name = "com.app.test:string/overlayable";
+  std::unique_ptr<ResourceTable> table =
+      test::ResourceTableBuilder()
+          .AddFileReference("android:color/xmlfile", original_xml_path)
+          .AddFileReference("android:color/pngfile", original_png_path)
+          .AddValue("com.app.test:color/mycolor", aapt::ResourceId(0x7f020000),
+                    aapt::util::make_unique<aapt::BinaryPrimitive>(
+                        uint8_t(android::Res_value::TYPE_INT_COLOR_ARGB8), 0xffaabbcc))
+          .AddString("com.app.test:string/mystring", ResourceId(0x7f030000), "hi")
+          .AddString("com.app.test:string/in_exemption", ResourceId(0x7f030001), "Hi")
+          .AddString(name, ResourceId(0x7f030002), "HI")
+          .SetOverlayable(name, overlayable_item)
+          .Build();
+
+  OptimizeOptions options{.shorten_resource_paths = true};
+  TableFlattenerOptions& flattenerOptions = options.table_flattener_options;
+  flattenerOptions.collapse_key_stringpool = true;
+  flattenerOptions.name_collapse_exemptions.insert(
+      ResourceName({}, ResourceType::kString, "in_exemption"));
+  auto& id_resource_map = flattenerOptions.id_resource_map;
+  ASSERT_TRUE(Obfuscator(options).Consume(context.get(), table.get()));
+
+  // Expect that the id resource name map is populated
+  EXPECT_THAT(id_resource_map.at(0x7f020000), Eq("mycolor"));
+  EXPECT_THAT(id_resource_map.at(0x7f030000), Eq("mystring"));
+  EXPECT_THAT(id_resource_map.find(0x7f030001), Eq(id_resource_map.end()));
+  EXPECT_THAT(id_resource_map.find(0x7f030002), Eq(id_resource_map.end()));
+}
+
+TEST(ObfuscatorTest, IsEnabledWithDefaultOption) {
+  OptimizeOptions options;
+  Obfuscator obfuscatorWithDefaultOption(options);
+  ASSERT_THAT(obfuscatorWithDefaultOption.IsEnabled(), Eq(false));
+}
+
+TEST(ObfuscatorTest, IsEnabledWithShortenPathOption) {
+  OptimizeOptions options{.shorten_resource_paths = true};
+  Obfuscator obfuscatorWithShortenPathOption(options);
+  ASSERT_THAT(obfuscatorWithShortenPathOption.IsEnabled(), Eq(true));
+}
+
+TEST(ObfuscatorTest, IsEnabledWithCollapseStringPoolOption) {
+  OptimizeOptions options;
+  options.table_flattener_options.collapse_key_stringpool = true;
+  Obfuscator obfuscatorWithCollapseStringPoolOption(options);
+  ASSERT_THAT(obfuscatorWithCollapseStringPoolOption.IsEnabled(), Eq(true));
+}
+
+TEST(ObfuscatorTest, IsEnabledWithShortenPathAndCollapseStringPoolOption) {
+  OptimizeOptions options{.shorten_resource_paths = true};
+  options.table_flattener_options.collapse_key_stringpool = true;
+  Obfuscator obfuscatorWithCollapseStringPoolOption(options);
+  ASSERT_THAT(obfuscatorWithCollapseStringPoolOption.IsEnabled(), Eq(true));
+}
+
 }  // namespace aapt