From 2599aa4dfd29627a5690460bdcd5568d1029d4dc Mon Sep 17 00:00:00 2001
From: Pierre Lecesne <lecesne@google.com>
Date: Wed, 1 Feb 2017 22:47:03 +0000
Subject: [PATCH] Add a method LoadedApk::WriteToArchive.

Test: Manually.

Change-Id: I61717204e58ca2bbfba9a52c7aecf27882a853f8
---
 tools/aapt2/LoadedApk.cpp                     | 73 ++++++++++++++++++-
 tools/aapt2/LoadedApk.h                       | 15 ++--
 tools/aapt2/strip/Strip.cpp                   |  6 ++
 .../aapt2/unflatten/BinaryResourceParser.cpp  |  3 +-
 4 files changed, 88 insertions(+), 9 deletions(-)

diff --git a/tools/aapt2/LoadedApk.cpp b/tools/aapt2/LoadedApk.cpp
index 3d466ef8ff651..76792ced86b00 100644
--- a/tools/aapt2/LoadedApk.cpp
+++ b/tools/aapt2/LoadedApk.cpp
@@ -16,10 +16,15 @@
 
 #include "LoadedApk.h"
 
+#include "ResourceValues.h"
+#include "ValueVisitor.h"
+#include "flatten/Archive.h"
+#include "flatten/TableFlattener.h"
+
 namespace aapt {
 
-std::unique_ptr<LoadedApk> LoadedApk::LoadApkFromPath(
-    IAaptContext* context, const StringPiece& path) {
+std::unique_ptr<LoadedApk> LoadedApk::LoadApkFromPath(IAaptContext* context,
+                                                      const android::StringPiece& path) {
   Source source(path);
   std::string error;
   std::unique_ptr<io::ZipFileCollection> apk =
@@ -53,4 +58,68 @@ std::unique_ptr<LoadedApk> LoadedApk::LoadApkFromPath(
   return util::make_unique<LoadedApk>(source, std::move(apk), std::move(table));
 }
 
+bool LoadedApk::WriteToArchive(IAaptContext* context, IArchiveWriter* writer) {
+  std::set<std::string> referenced_resources;
+  // List the files being referenced in the resource table.
+  for (auto& pkg : table_->packages) {
+    for (auto& type : pkg->types) {
+      for (auto& entry : type->entries) {
+        for (auto& config_value : entry->values) {
+          FileReference* file_ref = ValueCast<FileReference>(config_value->value.get());
+          if (file_ref) {
+            referenced_resources.insert(*file_ref->path);
+          }
+        }
+      }
+    }
+  }
+
+  std::unique_ptr<io::IFileCollectionIterator> iterator = apk_->Iterator();
+  while (iterator->HasNext()) {
+    io::IFile* file = iterator->Next();
+
+    std::string path = file->GetSource().path;
+    // The name of the path has the format "<zip-file-name>@<path-to-file>".
+    path = path.substr(path.find("@") + 1);
+
+    // Skip resources that are not referenced if requested.
+    if (path.find("res/") == 0 && referenced_resources.find(path) == referenced_resources.end()) {
+      if (context->IsVerbose()) {
+        context->GetDiagnostics()->Note(DiagMessage()
+                                        << "Resource '" << path << "' not referenced in "
+                                        << "resource table; removing from APK.");
+      }
+      continue;
+    }
+
+    // The resource table needs to be reserialized since it might have changed.
+    if (path == "resources.arsc") {
+      BigBuffer buffer = BigBuffer(1024);
+      TableFlattener flattener(&buffer);
+      if (!flattener.Consume(context, table_.get())) {
+        return false;
+      }
+
+      if (!writer->StartEntry(path, ArchiveEntry::kAlign) || !writer->WriteEntry(buffer) ||
+          !writer->FinishEntry()) {
+        context->GetDiagnostics()->Error(DiagMessage()
+                                         << "Error when writing file '" << path << "' in APK.");
+        return false;
+      }
+      continue;
+    }
+
+    std::unique_ptr<io::IData> data = file->OpenAsData();
+    // TODO(lecesne): Only compress the files that were compressed in the original APK.
+    if (!writer->StartEntry(path, ArchiveEntry::kCompress) ||
+        !writer->WriteEntry(data->data(), data->size()) || !writer->FinishEntry()) {
+      context->GetDiagnostics()->Error(DiagMessage()
+                                       << "Error when writing file '" << path << "' in APK.");
+      return false;
+    }
+  }
+
+  return true;
+}
+
 }  // namespace aapt
diff --git a/tools/aapt2/LoadedApk.h b/tools/aapt2/LoadedApk.h
index 0cc2d22eb7b7a..f8878d13ae35b 100644
--- a/tools/aapt2/LoadedApk.h
+++ b/tools/aapt2/LoadedApk.h
@@ -19,12 +19,11 @@
 
 #include "androidfw/StringPiece.h"
 
-#include "io/ZipArchive.h"
 #include "ResourceTable.h"
+#include "flatten/Archive.h"
+#include "io/ZipArchive.h"
 #include "unflatten/BinaryResourceParser.h"
 
-using android::StringPiece;
-
 namespace aapt {
 
 /** Info about an APK loaded in memory. */
@@ -42,8 +41,14 @@ class LoadedApk {
 
   const Source& GetSource() { return source_; }
 
-  static std::unique_ptr<LoadedApk> LoadApkFromPath(
-      IAaptContext* context, const StringPiece& path);
+  /**
+   * Writes the APK on disk at the given path, while also removing the resource
+   * files that are not referenced in the resource table.
+   */
+  bool WriteToArchive(IAaptContext* context, IArchiveWriter* writer);
+
+  static std::unique_ptr<LoadedApk> LoadApkFromPath(IAaptContext* context,
+                                                    const android::StringPiece& path);
 
  private:
   Source source_;
diff --git a/tools/aapt2/strip/Strip.cpp b/tools/aapt2/strip/Strip.cpp
index b3787ec69f083..ac3f244b090f9 100644
--- a/tools/aapt2/strip/Strip.cpp
+++ b/tools/aapt2/strip/Strip.cpp
@@ -80,6 +80,12 @@ class StripCommand {
 
     // TODO(lecesne): Implement stripping here.
 
+    std::unique_ptr<IArchiveWriter> writer =
+        CreateZipFileArchiveWriter(context_->GetDiagnostics(), options_.output_path);
+    if (!apk->WriteToArchive(context_, writer.get())) {
+      return 1;
+    }
+
     return 0;
   }
 
diff --git a/tools/aapt2/unflatten/BinaryResourceParser.cpp b/tools/aapt2/unflatten/BinaryResourceParser.cpp
index aeabcff011edf..7098fe9f7cb6e 100644
--- a/tools/aapt2/unflatten/BinaryResourceParser.cpp
+++ b/tools/aapt2/unflatten/BinaryResourceParser.cpp
@@ -370,8 +370,7 @@ bool BinaryResourceParser::ParseType(const ResourceTablePackage* package,
       return false;
     }
 
-    if (!table_->AddResourceAllowMangled(name, config, {},
-                                         std::move(resource_value),
+    if (!table_->AddResourceAllowMangled(name, res_id, config, {}, std::move(resource_value),
                                          context_->GetDiagnostics())) {
       return false;
     }