Set the location checksum in vdex.

The checksum we check against at runtime isn't the dex checksum
if we are compiling a jar/apk, but the crc32 of the entry.

This is to optimize the check at startup and avoid reading the
contents of the zip file.

Because with vdex we do not want to open the zip file, (the test
has been done already in getDexOptNeeded), and cannot rely
on the .oat file for fetching it, put the location checksum in
the vdex directly.

Note for later refactorings: the dex checksum location is now
in the vdex file, but also in the oat file. This can be revisited
after we eventually cleanup OatDexFile.

Test: 629-vdex
bug: 30937355
Change-Id: I7af8ca63b889370c660d694dd4eb95e78f566a1c
diff --git a/compiler/image_test.cc b/compiler/image_test.cc
index 5629dff..9bbe595 100644
--- a/compiler/image_test.cc
+++ b/compiler/image_test.cc
@@ -29,6 +29,8 @@
 #include "elf_writer_quick.h"
 #include "gc/space/image_space.h"
 #include "image_writer.h"
+#include "linker/buffered_output_stream.h"
+#include "linker/file_output_stream.h"
 #include "linker/multi_oat_relative_patcher.h"
 #include "lock_word.h"
 #include "mirror/object-inl.h"
@@ -256,6 +258,16 @@
       bool image_space_ok = writer->PrepareImageAddressSpace();
       ASSERT_TRUE(image_space_ok);
 
+      if (kIsVdexEnabled) {
+        for (size_t i = 0, size = vdex_files.size(); i != size; ++i) {
+          std::unique_ptr<BufferedOutputStream> vdex_out(
+              MakeUnique<BufferedOutputStream>(
+                  MakeUnique<FileOutputStream>(vdex_files[i].GetFile())));
+          oat_writers[i]->WriteVerifierDeps(vdex_out.get(), nullptr);
+          oat_writers[i]->WriteChecksumsAndVdexHeader(vdex_out.get());
+        }
+      }
+
       for (size_t i = 0, size = oat_files.size(); i != size; ++i) {
         linker::MultiOatRelativePatcher patcher(driver->GetInstructionSet(),
                                                 driver->GetInstructionSetFeatures());
diff --git a/compiler/oat_test.cc b/compiler/oat_test.cc
index 9458576..0a778b0 100644
--- a/compiler/oat_test.cc
+++ b/compiler/oat_test.cc
@@ -30,6 +30,8 @@
 #include "elf_writer.h"
 #include "elf_writer_quick.h"
 #include "entrypoints/quick/quick_entrypoints.h"
+#include "linker/buffered_output_stream.h"
+#include "linker/file_output_stream.h"
 #include "linker/multi_oat_relative_patcher.h"
 #include "linker/vector_output_stream.h"
 #include "mirror/class-inl.h"
@@ -218,6 +220,17 @@
                                       oat_writer.GetBssSize(),
                                       oat_writer.GetBssRootsOffset());
 
+    if (kIsVdexEnabled) {
+      std::unique_ptr<BufferedOutputStream> vdex_out(
+            MakeUnique<BufferedOutputStream>(MakeUnique<FileOutputStream>(vdex_file)));
+      if (!oat_writer.WriteVerifierDeps(vdex_out.get(), nullptr)) {
+        return false;
+      }
+      if (!oat_writer.WriteChecksumsAndVdexHeader(vdex_out.get())) {
+        return false;
+      }
+    }
+
     if (!oat_writer.WriteRodata(oat_rodata)) {
       return false;
     }
diff --git a/compiler/oat_writer.cc b/compiler/oat_writer.cc
index 153aff4..337ea80 100644
--- a/compiler/oat_writer.cc
+++ b/compiler/oat_writer.cc
@@ -300,6 +300,7 @@
     oat_data_offset_(0u),
     oat_header_(nullptr),
     size_vdex_header_(0),
+    size_vdex_checksums_(0),
     size_dex_file_alignment_(0),
     size_executable_offset_alignment_(0),
     size_oat_header_(0),
@@ -409,10 +410,11 @@
                                       CreateTypeLookupTable create_type_lookup_table) {
   DCHECK(write_state_ == WriteState::kAddingDexFileSources);
   const uint8_t* current_dex_data = nullptr;
-  for (size_t i = 0; ; ++i) {
+  for (size_t i = 0; i < vdex_file.GetHeader().GetNumberOfDexFiles(); ++i) {
     current_dex_data = vdex_file.GetNextDexFileData(current_dex_data);
     if (current_dex_data == nullptr) {
-      break;
+      LOG(ERROR) << "Unexpected number of dex files in vdex " << location;
+      return false;
     }
     if (!DexFile::IsMagicValid(current_dex_data)) {
       LOG(ERROR) << "Invalid magic in vdex file created from " << location;
@@ -424,7 +426,14 @@
     oat_dex_files_.emplace_back(full_location,
                                 DexFileSource(current_dex_data),
                                 create_type_lookup_table);
+    oat_dex_files_.back().dex_file_location_checksum_ = vdex_file.GetLocationChecksum(i);
   }
+
+  if (vdex_file.GetNextDexFileData(current_dex_data) != nullptr) {
+    LOG(ERROR) << "Unexpected number of dex files in vdex " << location;
+    return false;
+  }
+
   if (oat_dex_files_.empty()) {
     LOG(ERROR) << "No dex files in vdex file created from " << location;
     return false;
@@ -488,8 +497,8 @@
 
   // Initialize VDEX and OAT headers.
   if (kIsVdexEnabled) {
-    size_vdex_header_ = sizeof(VdexFile::Header);
-    vdex_size_ = size_vdex_header_;
+    // Reserve space for Vdex header and checksums.
+    vdex_size_ = sizeof(VdexFile::Header) + oat_dex_files_.size() * sizeof(VdexFile::VdexChecksum);
   }
   size_t oat_data_offset = InitOatHeader(instruction_set,
                                         instruction_set_features,
@@ -1837,6 +1846,7 @@
       size_total += (x);
 
     DO_STAT(size_vdex_header_);
+    DO_STAT(size_vdex_checksums_);
     DO_STAT(size_dex_file_alignment_);
     DO_STAT(size_executable_offset_alignment_);
     DO_STAT(size_oat_header_);
@@ -2383,6 +2393,7 @@
 
   // Update dex file size and resize class offsets in the OatDexFile.
   // Note: For raw data, the checksum is passed directly to AddRawDexFileSource().
+  // Note: For vdex, the checksum is copied from the existing vdex file.
   oat_dex_file->dex_file_size_ = header->file_size_;
   oat_dex_file->class_offsets_.resize(header->class_defs_size_);
   return true;
@@ -2592,11 +2603,31 @@
   return true;
 }
 
-bool OatWriter::WriteVdexHeader(OutputStream* vdex_out) {
+bool OatWriter::WriteChecksumsAndVdexHeader(OutputStream* vdex_out) {
   if (!kIsVdexEnabled) {
     return true;
   }
-  off_t actual_offset = vdex_out->Seek(0, kSeekSet);
+  // Write checksums
+  off_t actual_offset = vdex_out->Seek(sizeof(VdexFile::Header), kSeekSet);
+  if (actual_offset != sizeof(VdexFile::Header)) {
+    PLOG(ERROR) << "Failed to seek to the checksum location of vdex file. Actual: " << actual_offset
+                << " File: " << vdex_out->GetLocation();
+    return false;
+  }
+
+  for (size_t i = 0, size = oat_dex_files_.size(); i != size; ++i) {
+    OatDexFile* oat_dex_file = &oat_dex_files_[i];
+    if (!vdex_out->WriteFully(
+            &oat_dex_file->dex_file_location_checksum_, sizeof(VdexFile::VdexChecksum))) {
+      PLOG(ERROR) << "Failed to write dex file location checksum. File: "
+                  << vdex_out->GetLocation();
+      return false;
+    }
+    size_vdex_checksums_ += sizeof(VdexFile::VdexChecksum);
+  }
+
+  // Write header.
+  actual_offset = vdex_out->Seek(0, kSeekSet);
   if (actual_offset != 0) {
     PLOG(ERROR) << "Failed to seek to the beginning of vdex file. Actual: " << actual_offset
                 << " File: " << vdex_out->GetLocation();
@@ -2610,12 +2641,15 @@
   size_t verifier_deps_section_size = vdex_quickening_info_offset_ - vdex_verifier_deps_offset_;
   size_t quickening_info_section_size = vdex_size_ - vdex_quickening_info_offset_;
 
-  VdexFile::Header vdex_header(
-      dex_section_size, verifier_deps_section_size, quickening_info_section_size);
+  VdexFile::Header vdex_header(oat_dex_files_.size(),
+                               dex_section_size,
+                               verifier_deps_section_size,
+                               quickening_info_section_size);
   if (!vdex_out->WriteFully(&vdex_header, sizeof(VdexFile::Header))) {
     PLOG(ERROR) << "Failed to write vdex header. File: " << vdex_out->GetLocation();
     return false;
   }
+  size_vdex_header_ = sizeof(VdexFile::Header);
 
   if (!vdex_out->Flush()) {
     PLOG(ERROR) << "Failed to flush stream after writing to vdex file."
diff --git a/compiler/oat_writer.h b/compiler/oat_writer.h
index 0dcf79e..da221d6 100644
--- a/compiler/oat_writer.h
+++ b/compiler/oat_writer.h
@@ -124,7 +124,7 @@
   //   - Initialize()
   //   - WriteVerifierDeps()
   //   - WriteQuickeningInfo()
-  //   - WriteVdexHeader()
+  //   - WriteChecksumsAndVdexHeader()
   //   - PrepareLayout(),
   //   - WriteRodata(),
   //   - WriteCode(),
@@ -168,7 +168,7 @@
                             /*out*/ std::vector<std::unique_ptr<const DexFile>>* opened_dex_files);
   bool WriteQuickeningInfo(OutputStream* vdex_out);
   bool WriteVerifierDeps(OutputStream* vdex_out, verifier::VerifierDeps* verifier_deps);
-  bool WriteVdexHeader(OutputStream* vdex_out);
+  bool WriteChecksumsAndVdexHeader(OutputStream* vdex_out);
   // Initialize the writer with the given parameters.
   void Initialize(const CompilerDriver* compiler,
                   ImageWriter* image_writer,
@@ -387,6 +387,7 @@
 
   // output stats
   uint32_t size_vdex_header_;
+  uint32_t size_vdex_checksums_;
   uint32_t size_dex_file_alignment_;
   uint32_t size_executable_offset_alignment_;
   uint32_t size_oat_header_;
diff --git a/dex2oat/dex2oat.cc b/dex2oat/dex2oat.cc
index 9e6032f..57d2efe 100644
--- a/dex2oat/dex2oat.cc
+++ b/dex2oat/dex2oat.cc
@@ -1855,8 +1855,8 @@
           return false;
         }
 
-        // VDEX finalized, seek back to the beginning and write the header.
-        if (!oat_writers_[i]->WriteVdexHeader(vdex_out.get())) {
+        // VDEX finalized, seek back to the beginning and write checksums and the header.
+        if (!oat_writers_[i]->WriteChecksumsAndVdexHeader(vdex_out.get())) {
           LOG(ERROR) << "Failed to write vdex header into VDEX " << vdex_file->GetPath();
           return false;
         }
diff --git a/runtime/oat_file_assistant.cc b/runtime/oat_file_assistant.cc
index 4d1e1ea..6a62a16 100644
--- a/runtime/oat_file_assistant.cc
+++ b/runtime/oat_file_assistant.cc
@@ -314,7 +314,7 @@
   const OatFile::OatDexFile* oat_dex_file = file.GetOatDexFile(
       dex_location_.c_str(), dex_checksum_pointer, &error_msg);
   if (oat_dex_file == nullptr) {
-    VLOG(oat) << error_msg;
+    LOG(ERROR) << error_msg;
     return kOatDexOutOfDate;
   }
 
diff --git a/runtime/vdex_file.cc b/runtime/vdex_file.cc
index 843be92..dabf8c8 100644
--- a/runtime/vdex_file.cc
+++ b/runtime/vdex_file.cc
@@ -35,10 +35,12 @@
   return (memcmp(version_, kVdexVersion, sizeof(kVdexVersion)) == 0);
 }
 
-VdexFile::Header::Header(uint32_t dex_size,
+VdexFile::Header::Header(uint32_t number_of_dex_files,
+                         uint32_t dex_size,
                          uint32_t verifier_deps_size,
                          uint32_t quickening_info_size)
-    : dex_size_(dex_size),
+    : number_of_dex_files_(number_of_dex_files),
+      dex_size_(dex_size),
       verifier_deps_size_(verifier_deps_size),
       quickening_info_size_(quickening_info_size) {
   memcpy(magic_, kVdexMagic, sizeof(kVdexMagic));
diff --git a/runtime/vdex_file.h b/runtime/vdex_file.h
index 75a0d5e..3b114a9 100644
--- a/runtime/vdex_file.h
+++ b/runtime/vdex_file.h
@@ -43,7 +43,10 @@
  public:
   struct Header {
    public:
-    Header(uint32_t dex_size, uint32_t verifier_deps_size, uint32_t quickening_info_size);
+    Header(uint32_t number_of_dex_files_,
+           uint32_t dex_size,
+           uint32_t verifier_deps_size,
+           uint32_t quickening_info_size);
 
     const char* GetMagic() const { return reinterpret_cast<const char*>(magic_); }
     const char* GetVersion() const { return reinterpret_cast<const char*>(version_); }
@@ -54,18 +57,22 @@
     uint32_t GetDexSize() const { return dex_size_; }
     uint32_t GetVerifierDepsSize() const { return verifier_deps_size_; }
     uint32_t GetQuickeningInfoSize() const { return quickening_info_size_; }
+    uint32_t GetNumberOfDexFiles() const { return number_of_dex_files_; }
 
    private:
     static constexpr uint8_t kVdexMagic[] = { 'v', 'd', 'e', 'x' };
-    static constexpr uint8_t kVdexVersion[] = { '0', '0', '0', '\0' };
+    static constexpr uint8_t kVdexVersion[] = { '0', '0', '1', '\0' };
 
     uint8_t magic_[4];
     uint8_t version_[4];
+    uint32_t number_of_dex_files_;
     uint32_t dex_size_;
     uint32_t verifier_deps_size_;
     uint32_t quickening_info_size_;
   };
 
+  typedef uint32_t VdexChecksum;
+
   static VdexFile* Open(const std::string& vdex_filename,
                         bool writable,
                         bool low_4gb,
@@ -88,7 +95,7 @@
 
   ArrayRef<const uint8_t> GetVerifierDepsData() const {
     return ArrayRef<const uint8_t>(
-        Begin() + sizeof(Header) + GetHeader().GetDexSize(), GetHeader().GetVerifierDepsSize());
+        DexBegin() + GetHeader().GetDexSize(), GetHeader().GetVerifierDepsSize());
   }
 
   ArrayRef<const uint8_t> GetQuickeningInfo() const {
@@ -107,6 +114,12 @@
   // is none.
   const uint8_t* GetNextDexFileData(const uint8_t* cursor) const;
 
+  // Get the location checksum of the dex file number `dex_file_index`.
+  uint32_t GetLocationChecksum(uint32_t dex_file_index) const {
+    DCHECK_LT(dex_file_index, GetHeader().GetNumberOfDexFiles());
+    return reinterpret_cast<const uint32_t*>(Begin() + sizeof(Header))[dex_file_index];
+  }
+
  private:
   explicit VdexFile(MemMap* mmap) : mmap_(mmap) {}
 
@@ -115,11 +128,15 @@
   }
 
   const uint8_t* DexBegin() const {
-    return Begin() + sizeof(Header);
+    return Begin() + sizeof(Header) + GetSizeOfChecksumsSection();
   }
 
   const uint8_t* DexEnd() const {
-    return Begin() + sizeof(Header) + GetHeader().GetDexSize();
+    return DexBegin() + GetHeader().GetDexSize();
+  }
+
+  size_t GetSizeOfChecksumsSection() const {
+    return sizeof(VdexChecksum) * GetHeader().GetNumberOfDexFiles();
   }
 
   std::unique_ptr<MemMap> mmap_;