1 files changed, 111 insertions, 31 deletions

diff --git a/tools/aapt2/proto/TableProtoSerializer.cpp b/tools/aapt2/proto/TableProtoSerializer.cpp
index 5d1b72b0ebbd..a5c2cbc0277b 100644
--- a/tools/aapt2/proto/TableProtoSerializer.cpp
+++ b/tools/aapt2/proto/TableProtoSerializer.cpp
@@ -22,6 +22,10 @@
 #include "proto/ProtoSerialize.h"
 #include "util/BigBuffer.h"
 
+using google::protobuf::io::CodedOutputStream;
+using google::protobuf::io::CodedInputStream;
+using google::protobuf::io::ZeroCopyOutputStream;
+
 namespace aapt {
 
 namespace {
@@ -171,7 +175,7 @@ private:
     void serializeItemCommonToPb(const Item& item, T* pbItem) {
         serializeSourceToPb(item.getSource(), mSourcePool, pbItem->mutable_source());
         if (!item.getComment().empty()) {
-            pbItem->set_comment(util::utf16ToUtf8(item.getComment()));
+            pbItem->set_comment(item.getComment());
         }
     }
 
@@ -210,7 +214,7 @@ std::unique_ptr<pb::ResourceTable> serializeTableToPb(ResourceTable* table) {
     });
     table->stringPool.prune();
 
-    std::unique_ptr<pb::ResourceTable> pbTable = util::make_unique<pb::ResourceTable>();
+    auto pbTable = util::make_unique<pb::ResourceTable>();
     serializeStringPoolToPb(table->stringPool, pbTable->mutable_string_pool());
 
     StringPool sourcePool, symbolPool;
@@ -220,28 +224,28 @@ std::unique_ptr<pb::ResourceTable> serializeTableToPb(ResourceTable* table) {
         if (package->id) {
             pbPackage->set_package_id(package->id.value());
         }
-        pbPackage->set_package_name(util::utf16ToUtf8(package->name));
+        pbPackage->set_package_name(package->name);
 
         for (auto& type : package->types) {
             pb::Type* pbType = pbPackage->add_types();
             if (type->id) {
                 pbType->set_id(type->id.value());
             }
-            pbType->set_name(util::utf16ToUtf8(toString(type->type)));
+            pbType->set_name(toString(type->type).toString());
 
             for (auto& entry : type->entries) {
                 pb::Entry* pbEntry = pbType->add_entries();
                 if (entry->id) {
                     pbEntry->set_id(entry->id.value());
                 }
-                pbEntry->set_name(util::utf16ToUtf8(entry->name));
+                pbEntry->set_name(entry->name);
 
                 // Write the SymbolStatus struct.
                 pb::SymbolStatus* pbStatus = pbEntry->mutable_symbol_status();
                 pbStatus->set_visibility(serializeVisibilityToPb(entry->symbolStatus.state));
                 serializeSourceToPb(entry->symbolStatus.source, &sourcePool,
                                     pbStatus->mutable_source());
-                pbStatus->set_comment(util::utf16ToUtf8(entry->symbolStatus.comment));
+                pbStatus->set_comment(entry->symbolStatus.comment);
 
                 for (auto& configValue : entry->values) {
                     pb::ConfigValue* pbConfigValue = pbEntry->add_config_values();
@@ -254,7 +258,7 @@ std::unique_ptr<pb::ResourceTable> serializeTableToPb(ResourceTable* table) {
                     serializeSourceToPb(configValue->value->getSource(), &sourcePool,
                                         pbValue->mutable_source());
                     if (!configValue->value->getComment().empty()) {
-                        pbValue->set_comment(util::utf16ToUtf8(configValue->value->getComment()));
+                        pbValue->set_comment(configValue->value->getComment());
                     }
 
                     if (configValue->value->isWeak()) {
@@ -274,49 +278,125 @@ std::unique_ptr<pb::ResourceTable> serializeTableToPb(ResourceTable* table) {
 }
 
 std::unique_ptr<pb::CompiledFile> serializeCompiledFileToPb(const ResourceFile& file) {
-    std::unique_ptr<pb::CompiledFile> pbFile = util::make_unique<pb::CompiledFile>();
-    pbFile->set_resource_name(util::utf16ToUtf8(file.name.toString()));
+    auto pbFile = util::make_unique<pb::CompiledFile>();
+    pbFile->set_resource_name(file.name.toString());
     pbFile->set_source_path(file.source.path);
     serializeConfig(file.config, pbFile->mutable_config());
 
     for (const SourcedResourceName& exported : file.exportedSymbols) {
         pb::CompiledFile_Symbol* pbSymbol = pbFile->add_exported_symbols();
-        pbSymbol->set_resource_name(util::utf16ToUtf8(exported.name.toString()));
+        pbSymbol->set_resource_name(exported.name.toString());
         pbSymbol->set_line_no(exported.line);
     }
     return pbFile;
 }
 
-CompiledFileOutputStream::CompiledFileOutputStream(google::protobuf::io::ZeroCopyOutputStream* out,
-                                                   pb::CompiledFile* pbFile) :
-        mOut(out), mPbFile(pbFile) {
+CompiledFileOutputStream::CompiledFileOutputStream(ZeroCopyOutputStream* out) : mOut(out) {
 }
 
-bool CompiledFileOutputStream::ensureFileWritten() {
-    if (mPbFile) {
-        const uint64_t pbSize = mPbFile->ByteSize();
-        mOut.WriteLittleEndian64(pbSize);
-        mPbFile->SerializeWithCachedSizes(&mOut);
-        const size_t padding = 4 - (pbSize & 0x03);
-        if (padding > 0) {
-            uint32_t zero = 0u;
-            mOut.WriteRaw(&zero, padding);
-        }
-        mPbFile = nullptr;
+void CompiledFileOutputStream::ensureAlignedWrite() {
+    const int padding = mOut.ByteCount() % 4;
+    if (padding > 0) {
+        uint32_t zero = 0u;
+        mOut.WriteRaw(&zero, padding);
+    }
+}
+
+void CompiledFileOutputStream::WriteLittleEndian32(uint32_t val) {
+    ensureAlignedWrite();
+    mOut.WriteLittleEndian32(val);
+}
+
+void CompiledFileOutputStream::WriteCompiledFile(const pb::CompiledFile* compiledFile) {
+    ensureAlignedWrite();
+    mOut.WriteLittleEndian64(static_cast<uint64_t>(compiledFile->ByteSize()));
+    compiledFile->SerializeWithCachedSizes(&mOut);
+}
+
+void CompiledFileOutputStream::WriteData(const BigBuffer* buffer) {
+    ensureAlignedWrite();
+    mOut.WriteLittleEndian64(static_cast<uint64_t>(buffer->size()));
+    for (const BigBuffer::Block& block : *buffer) {
+        mOut.WriteRaw(block.buffer.get(), block.size);
+    }
+}
+
+void CompiledFileOutputStream::WriteData(const void* data, size_t len) {
+    ensureAlignedWrite();
+    mOut.WriteLittleEndian64(static_cast<uint64_t>(len));
+    mOut.WriteRaw(data, len);
+}
+
+bool CompiledFileOutputStream::HadError() {
+    return mOut.HadError();
+}
+
+CompiledFileInputStream::CompiledFileInputStream(const void* data, size_t size) :
+        mIn(static_cast<const uint8_t*>(data), size) {
+}
+
+void CompiledFileInputStream::ensureAlignedRead() {
+    const int padding = mIn.CurrentPosition() % 4;
+    if (padding > 0) {
+        // Reads are always 4 byte aligned.
+        mIn.Skip(padding);
     }
-    return !mOut.HadError();
 }
 
-bool CompiledFileOutputStream::Write(const void* data, int size) {
-    if (!ensureFileWritten()) {
+bool CompiledFileInputStream::ReadLittleEndian32(uint32_t* outVal) {
+    ensureAlignedRead();
+    return mIn.ReadLittleEndian32(outVal);
+}
+
+bool CompiledFileInputStream::ReadCompiledFile(pb::CompiledFile* outVal) {
+    ensureAlignedRead();
+
+    uint64_t pbSize = 0u;
+    if (!mIn.ReadLittleEndian64(&pbSize)) {
+        return false;
+    }
+
+    CodedInputStream::Limit l = mIn.PushLimit(static_cast<int>(pbSize));
+
+    // Check that we haven't tried to read past the end.
+    if (static_cast<uint64_t>(mIn.BytesUntilLimit()) != pbSize) {
+        mIn.PopLimit(l);
+        mIn.PushLimit(0);
         return false;
     }
-    mOut.WriteRaw(data, size);
-    return !mOut.HadError();
+
+    if (!outVal->ParsePartialFromCodedStream(&mIn)) {
+        mIn.PopLimit(l);
+        mIn.PushLimit(0);
+        return false;
+    }
+
+    mIn.PopLimit(l);
+    return true;
 }
 
-bool CompiledFileOutputStream::Finish() {
-    return ensureFileWritten();
+bool CompiledFileInputStream::ReadDataMetaData(uint64_t* outOffset, uint64_t* outLen) {
+    ensureAlignedRead();
+
+    uint64_t pbSize = 0u;
+    if (!mIn.ReadLittleEndian64(&pbSize)) {
+        return false;
+    }
+
+    // Check that we aren't trying to read past the end.
+    if (pbSize > static_cast<uint64_t>(mIn.BytesUntilLimit())) {
+        mIn.PushLimit(0);
+        return false;
+    }
+
+    uint64_t offset = static_cast<uint64_t>(mIn.CurrentPosition());
+    if (!mIn.Skip(pbSize)) {
+        return false;
+    }
+
+    *outOffset = offset;
+    *outLen = pbSize;
+    return true;
 }
 
 } // namespace aapt