Merge "AAPT2: Increase maximum proto size"

3 files changed, 28 insertions, 10 deletions

diff --git a/tools/aapt2/LoadedApk.cpp b/tools/aapt2/LoadedApk.cpp
index 6d78990e4636..a20b9b76a703 100644
--- a/tools/aapt2/LoadedApk.cpp
+++ b/tools/aapt2/LoadedApk.cpp
@@ -69,8 +69,8 @@ std::unique_ptr<LoadedApk> LoadedApk::LoadProtoApkFromFileCollection(
       return {};
     }
 
-    io::ZeroCopyInputAdaptor adaptor(in.get());
-    if (!pb_table.ParseFromZeroCopyStream(&adaptor)) {
+    io::ProtoInputStreamReader proto_reader(in.get());
+    if (!proto_reader.ReadMessage(&pb_table)) {
       diag->Error(DiagMessage(source) << "failed to read " << kProtoResourceTablePath);
       return {};
     }
@@ -97,8 +97,8 @@ std::unique_ptr<LoadedApk> LoadedApk::LoadProtoApkFromFileCollection(
   }
 
   pb::XmlNode pb_node;
-  io::ZeroCopyInputAdaptor manifest_adaptor(manifest_in.get());
-  if (!pb_node.ParseFromZeroCopyStream(&manifest_adaptor)) {
+  io::ProtoInputStreamReader proto_reader(manifest_in.get());
+  if (!proto_reader.ReadMessage(&pb_node)) {
     diag->Error(DiagMessage(source) << "failed to read proto " << kAndroidManifestPath);
     return {};
   }
@@ -270,9 +270,9 @@ std::unique_ptr<xml::XmlResource> LoadedApk::LoadXml(const std::string& file_pat
       return nullptr;
     }
 
-    io::ZeroCopyInputAdaptor adaptor(in.get());
     pb::XmlNode pb_node;
-    if (!pb_node.ParseFromZeroCopyStream(&adaptor)) {
+    io::ProtoInputStreamReader proto_reader(in.get());
+    if (!proto_reader.ReadMessage(&pb_node)) {
       diag->Error(DiagMessage() << "failed to parse file as proto XML");
       return nullptr;
     }
@@ -317,8 +317,8 @@ ApkFormat LoadedApk::DetermineApkFormat(io::IFileCollection* apk) {
     std::unique_ptr<io::InputStream> manifest_in = manifest_file->OpenInputStream();
     if (manifest_in != nullptr) {
       pb::XmlNode pb_node;
-      io::ZeroCopyInputAdaptor manifest_adaptor(manifest_in.get());
-      if (pb_node.ParseFromZeroCopyStream(&manifest_adaptor)) {
+      io::ProtoInputStreamReader proto_reader(manifest_in.get());
+      if (!proto_reader.ReadMessage(&pb_node)) {
         return ApkFormat::kProto;
       }
     }
diff --git a/tools/aapt2/cmd/Convert.cpp b/tools/aapt2/cmd/Convert.cpp
index 86b1f4c54deb..954f1ed3fb72 100644
--- a/tools/aapt2/cmd/Convert.cpp
+++ b/tools/aapt2/cmd/Convert.cpp
@@ -170,8 +170,8 @@ class BinaryApkSerializer : public IApkSerializer {
       }
 
       pb::XmlNode pb_node;
-      io::ZeroCopyInputAdaptor adaptor(in.get());
-      if (!pb_node.ParseFromZeroCopyStream(&adaptor)) {
+      io::ProtoInputStreamReader proto_reader(in.get());
+      if (!proto_reader.ReadMessage(&pb_node)) {
         context_->GetDiagnostics()->Error(DiagMessage(source_)
                                           << "failed to parse proto XML " << *file->path);
         return false;
diff --git a/tools/aapt2/io/Util.h b/tools/aapt2/io/Util.h
index b07fb5399313..5f978a8e2c35 100644
--- a/tools/aapt2/io/Util.h
+++ b/tools/aapt2/io/Util.h
@@ -20,6 +20,7 @@
 #include <string>
 
 #include "google/protobuf/message_lite.h"
+#include "google/protobuf/io/coded_stream.h"
 
 #include "format/Archive.h"
 #include "io/File.h"
@@ -122,6 +123,23 @@ class ZeroCopyInputAdaptor : public ::google::protobuf::io::ZeroCopyInputStream
   io::InputStream* in_;
 };
 
+class ProtoInputStreamReader {
+ public:
+  explicit ProtoInputStreamReader(io::InputStream* in) : in_(in) { }
+
+  /** Deserializes a MessageLite proto from the current position in the input stream.*/
+  template <typename T> bool ReadMessage(T *message_lite) {
+    ZeroCopyInputAdaptor adapter(in_);
+    google::protobuf::io::CodedInputStream coded_stream(&adapter);
+    coded_stream.SetTotalBytesLimit(std::numeric_limits<int32_t>::max(),
+                                    coded_stream.BytesUntilTotalBytesLimit());
+    return message_lite->ParseFromCodedStream(&coded_stream);
+  }
+
+ private:
+  io::InputStream* in_;
+};
+
 }  // namespace io
 }  // namespace aapt