Ensure oat checksum determinism across hosts and devices.

For Cloud Compilation, we need to enable hosts to generate a boot image
that has exactly the same checksum as generated on device. In other
words, we need to make the oat checksum deterministic across hosts and
devices.

To achieve that, when writing the oat header, the non-deterministic
fields are padded to their length limits and excluded from the oat
checksum computation.

Bug: 372868052
Test: m test-art-host-gtest
Test: Generate a boot image on host. See the checksum being identical to
  the one on device.
Change-Id: I9f73a28b349b673c25909b61a0c7ae8cf883c014

5 files changed, 153 insertions, 38 deletions

diff --git a/dex2oat/dex2oat.cc b/dex2oat/dex2oat.cc
index 7564bf0b68..39b0d826b4 100644
--- a/dex2oat/dex2oat.cc
+++ b/dex2oat/dex2oat.cc
@@ -500,15 +500,6 @@ class ThreadLocalHashOverride {
   Handle<mirror::Object> old_field_value_;
 };
 
-class OatKeyValueStore : public SafeMap<std::string, std::string> {
- public:
-  using SafeMap::Put;
-
-  iterator Put(const std::string& k, bool v) {
-    return SafeMap::Put(k, v ? OatHeader::kTrueValue : OatHeader::kFalseValue);
-  }
-};
-
 class Dex2Oat final {
  public:
   explicit Dex2Oat(TimingLogger* timings)
@@ -893,7 +884,7 @@ class Dex2Oat final {
     }
 
     // Fill some values into the key-value store for the oat header.
-    key_value_store_.reset(new OatKeyValueStore());
+    key_value_store_.reset(new linker::OatKeyValueStore());
 
     // Automatically force determinism for the boot image and boot image extensions in a host build.
     if (!kIsTargetBuild && (IsBootImage() || IsBootImageExtension())) {
@@ -978,7 +969,8 @@ class Dex2Oat final {
         }
         oss << argv[i];
       }
-      key_value_store_->Put(OatHeader::kDex2OatCmdLineKey, oss.str());
+      key_value_store_->PutNonDeterministic(
+          OatHeader::kDex2OatCmdLineKey, oss.str(), /*allow_truncation=*/true);
     }
     key_value_store_->Put(OatHeader::kDebuggableKey, compiler_options_->debuggable_);
     key_value_store_->Put(OatHeader::kNativeDebuggableKey,
@@ -1696,7 +1688,10 @@ class Dex2Oat final {
         CompilerFilter::DependsOnImageChecksum(original_compiler_filter)) {
       std::string versions =
           apex_versions_argument_.empty() ? runtime->GetApexVersions() : apex_versions_argument_;
-      key_value_store_->Put(OatHeader::kApexVersionsKey, versions);
+      if (!key_value_store_->PutNonDeterministic(OatHeader::kApexVersionsKey, versions)) {
+        LOG(ERROR) << "Cannot store apex versions string because it's too long";
+        return dex2oat::ReturnCode::kOther;
+      }
     }
 
     // Now that we have adjusted whether we generate an image, encode it in the
@@ -2915,7 +2910,7 @@ class Dex2Oat final {
 
   std::unique_ptr<CompilerOptions> compiler_options_;
 
-  std::unique_ptr<OatKeyValueStore> key_value_store_;
+  std::unique_ptr<linker::OatKeyValueStore> key_value_store_;
 
   std::unique_ptr<VerificationResults> verification_results_;
 
diff --git a/dex2oat/linker/image_test.h b/dex2oat/linker/image_test.h
index 3706b685fa..349462e098 100644
--- a/dex2oat/linker/image_test.h
+++ b/dex2oat/linker/image_test.h
@@ -231,13 +231,12 @@ inline void ImageTest::DoCompile(ImageHeader::StorageMode storage_mode,
       CompileAll(class_loader, class_path, &timings);
 
       TimingLogger::ScopedTiming t("WriteElf", &timings);
-      SafeMap<std::string, std::string> key_value_store;
+      OatKeyValueStore key_value_store;
       key_value_store.Put(OatHeader::kBootClassPathKey,
                           android::base::Join(out_helper.dex_file_locations, ':'));
-      key_value_store.Put(OatHeader::kApexVersionsKey, Runtime::Current()->GetApexVersions());
-      key_value_store.Put(
-          OatHeader::kConcurrentCopying,
-          compiler_options_->EmitReadBarrier() ? OatHeader::kTrueValue : OatHeader::kFalseValue);
+      key_value_store.PutNonDeterministic(OatHeader::kApexVersionsKey,
+                                          Runtime::Current()->GetApexVersions());
+      key_value_store.Put(OatHeader::kConcurrentCopying, compiler_options_->EmitReadBarrier());
 
       std::vector<std::unique_ptr<ElfWriter>> elf_writers;
       std::vector<std::unique_ptr<OatWriter>> oat_writers;
diff --git a/dex2oat/linker/oat_writer.cc b/dex2oat/linker/oat_writer.cc
index 20787ddc8f..6cdfe7bb40 100644
--- a/dex2oat/linker/oat_writer.cc
+++ b/dex2oat/linker/oat_writer.cc
@@ -111,6 +111,33 @@ inline uint32_t CodeAlignmentSize(uint32_t header_offset, const CompiledMethod&
 
 }  // anonymous namespace
 
+bool OatKeyValueStore::PutNonDeterministic(const std::string& k,
+                                           const std::string& v,
+                                           bool allow_truncation) {
+  size_t length = OatHeader::GetNonDeterministicFieldLength(k);
+  DCHECK_GT(length, 0u);
+  if (v.length() <= length) {
+    map_.Put(k, v);
+    return true;
+  }
+  if (allow_truncation) {
+    LOG(WARNING) << "Key value store field " << k << "too long. Truncating";
+    map_.Put(k, v.substr(length));
+    return true;
+  }
+  return false;
+}
+
+void OatKeyValueStore::Put(const std::string& k, const std::string& v) {
+  DCHECK(OatHeader::IsDeterministicField(k));
+  map_.Put(k, v);
+}
+
+void OatKeyValueStore::Put(const std::string& k, bool v) {
+  DCHECK(OatHeader::IsDeterministicField(k));
+  map_.Put(k, v ? OatHeader::kTrueValue : OatHeader::kFalseValue);
+}
+
 // .bss mapping offsets used for BCP DexFiles.
 struct OatWriter::BssMappingInfo {
   // Offsets set in PrepareLayout.
@@ -550,7 +577,7 @@ bool OatWriter::WriteAndOpenDexFiles(
 
 bool OatWriter::StartRoData(const std::vector<const DexFile*>& dex_files,
                             OutputStream* oat_rodata,
-                            SafeMap<std::string, std::string>* key_value_store) {
+                            OatKeyValueStore* key_value_store) {
   CHECK(write_state_ == WriteState::kStartRoData);
 
   // Record the ELF rodata section offset, i.e. the beginning of the OAT data.
@@ -1972,9 +1999,18 @@ bool OatWriter::VisitDexMethods(DexMethodVisitor* visitor) {
   return true;
 }
 
-size_t OatWriter::InitOatHeader(uint32_t num_dex_files,
-                                SafeMap<std::string, std::string>* key_value_store) {
+size_t OatWriter::InitOatHeader(uint32_t num_dex_files, OatKeyValueStore* key_value_store) {
   TimingLogger::ScopedTiming split("InitOatHeader", timings_);
+
+  // `key_value_store` only exists in the first oat file in a multi-image boot image.
+  if (key_value_store != nullptr) {
+    // Add non-deterministic fields if they don't exist. These fields should always exist with fixed
+    // lengths.
+    for (auto [field, length] : OatHeader::kNonDeterministicFieldsAndLengths) {
+      key_value_store->map_.FindOrAdd(std::string(field));
+    }
+  }
+
   // Check that oat version when runtime was compiled matches the oat version
   // when dex2oat was compiled. We have seen cases where they got out of sync.
   constexpr std::array<uint8_t, 4> dex2oat_oat_version = OatHeader::kOatVersion;
@@ -1982,7 +2018,7 @@ size_t OatWriter::InitOatHeader(uint32_t num_dex_files,
   oat_header_ = OatHeader::Create(GetCompilerOptions().GetInstructionSet(),
                                   GetCompilerOptions().GetInstructionSetFeatures(),
                                   num_dex_files,
-                                  key_value_store,
+                                  key_value_store != nullptr ? &key_value_store->map_ : nullptr,
                                   oat_data_offset_);
   size_oat_header_ += sizeof(OatHeader);
   size_oat_header_key_value_store_ += oat_header_->GetHeaderSize() - sizeof(OatHeader);
@@ -2751,10 +2787,7 @@ bool OatWriter::WriteHeader(OutputStream* out) {
 
   // Update checksum with header data.
   DCHECK_EQ(oat_header_->GetChecksum(), 0u);  // For checksum calculation.
-  const uint8_t* header_begin = reinterpret_cast<const uint8_t*>(oat_header_);
-  const uint8_t* header_end = oat_header_->GetKeyValueStore() + oat_header_->GetKeyValueStoreSize();
-  uint32_t old_checksum = oat_checksum_;
-  oat_checksum_ = adler32(old_checksum, header_begin, header_end - header_begin);
+  oat_header_->ComputeChecksum(&oat_checksum_);
   oat_header_->SetChecksum(oat_checksum_);
 
   const size_t file_offset = oat_data_offset_;
diff --git a/dex2oat/linker/oat_writer.h b/dex2oat/linker/oat_writer.h
index 5a775fa517..4f50b1a39c 100644
--- a/dex2oat/linker/oat_writer.h
+++ b/dex2oat/linker/oat_writer.h
@@ -70,6 +70,29 @@ enum class CopyOption {
   kOnlyIfCompressed
 };
 
+class OatKeyValueStore {
+ public:
+  // Puts a key value pair whose key is in `OatHeader::kNonDeterministicFieldsAndLengths`.
+  bool PutNonDeterministic(const std::string& k,
+                           const std::string& v,
+                           bool allow_truncation = false);
+
+  // Puts a key value pair whose key is in `OatHeader::kDeterministicFields`.
+  void Put(const std::string& k, const std::string& v);
+
+  // Puts a key value pair whose key is in `OatHeader::kDeterministicFields`.
+  void Put(const std::string& k, bool v);
+
+  // Makes sure calls with `const char*` falls into the overload for `std::string`, not the one for
+  // `bool`.
+  void Put(const std::string& k, const char* v) { Put(k, std::string(v)); }
+
+ private:
+  SafeMap<std::string, std::string> map_;
+
+  friend class OatWriter;
+};
+
 // OatHeader         variable length with count of D OatDexFiles
 //
 // TypeLookupTable[0] one descriptor to class def index hash table for each OatDexFile.
@@ -167,7 +190,7 @@ class OatWriter {
   // Start writing .rodata, including supporting data structures for dex files.
   bool StartRoData(const std::vector<const DexFile*>& dex_files,
                    OutputStream* oat_rodata,
-                   SafeMap<std::string, std::string>* key_value_store);
+                   OatKeyValueStore* key_value_store);
   // Initialize the writer with the given parameters.
   void Initialize(const CompilerDriver* compiler_driver,
                   const VerificationResults* verification_results,
@@ -291,7 +314,7 @@ class OatWriter {
   void WriteVerifierDeps(verifier::VerifierDeps* verifier_deps,
                          /*out*/std::vector<uint8_t>* buffer);
 
-  size_t InitOatHeader(uint32_t num_dex_files, SafeMap<std::string, std::string>* key_value_store);
+  size_t InitOatHeader(uint32_t num_dex_files, OatKeyValueStore* key_value_store);
   size_t InitClassOffsets(size_t offset);
   size_t InitOatClasses(size_t offset);
   size_t InitOatMaps(size_t offset);
diff --git a/dex2oat/linker/oat_writer_test.cc b/dex2oat/linker/oat_writer_test.cc
index f3b598fe77..538c1abc9a 100644
--- a/dex2oat/linker/oat_writer_test.cc
+++ b/dex2oat/linker/oat_writer_test.cc
@@ -14,8 +14,11 @@
  * limitations under the License.
  */
 
-#include "android-base/stringprintf.h"
+#include "oat_writer.h"
 
+#include <cstdint>
+
+#include "android-base/stringprintf.h"
 #include "arch/instruction_set_features.h"
 #include "art_method-inl.h"
 #include "base/file_utils.h"
@@ -35,6 +38,7 @@
 #include "driver/compiler_driver.h"
 #include "driver/compiler_options.h"
 #include "entrypoints/quick/quick_entrypoints.h"
+#include "gtest/gtest.h"
 #include "linker/elf_writer.h"
 #include "linker/elf_writer_quick.h"
 #include "linker/multi_oat_relative_patcher.h"
@@ -104,7 +108,7 @@ class OatTest : public CommonCompilerDriverTest {
   bool WriteElf(File* vdex_file,
                 File* oat_file,
                 const std::vector<const DexFile*>& dex_files,
-                SafeMap<std::string, std::string>& key_value_store,
+                OatKeyValueStore& key_value_store,
                 bool verify) {
     TimingLogger timings("WriteElf", false, false);
     ClearBootImageOption();
@@ -124,7 +128,7 @@ class OatTest : public CommonCompilerDriverTest {
   bool WriteElf(File* vdex_file,
                 File* oat_file,
                 const std::vector<const char*>& dex_filenames,
-                SafeMap<std::string, std::string>& key_value_store,
+                OatKeyValueStore& key_value_store,
                 bool verify,
                 CopyOption copy,
                 ProfileCompilationInfo* profile_compilation_info) {
@@ -143,7 +147,7 @@ class OatTest : public CommonCompilerDriverTest {
                 File* oat_file,
                 File&& dex_file_fd,
                 const char* location,
-                SafeMap<std::string, std::string>& key_value_store,
+                OatKeyValueStore& key_value_store,
                 bool verify,
                 CopyOption copy,
                 ProfileCompilationInfo* profile_compilation_info = nullptr) {
@@ -159,7 +163,7 @@ class OatTest : public CommonCompilerDriverTest {
   bool DoWriteElf(File* vdex_file,
                   File* oat_file,
                   OatWriter& oat_writer,
-                  SafeMap<std::string, std::string>& key_value_store,
+                  OatKeyValueStore& key_value_store,
                   bool verify,
                   CopyOption copy) {
     std::unique_ptr<ElfWriter> elf_writer = CreateElfWriterQuick(
@@ -426,7 +430,7 @@ TEST_F(OatTest, WriteRead) {
   }
 
   ScratchFile tmp_base, tmp_oat(tmp_base, ".oat"), tmp_vdex(tmp_base, ".vdex");
-  SafeMap<std::string, std::string> key_value_store;
+  OatKeyValueStore key_value_store;
   key_value_store.Put(OatHeader::kBootClassPathChecksumsKey, "testkey");
   bool success = WriteElf(tmp_vdex.GetFile(),
                           tmp_oat.GetFile(),
@@ -494,6 +498,67 @@ TEST_F(OatTest, WriteRead) {
   }
 }
 
+TEST_F(OatTest, ChecksumDeterminism) {
+  ClassLinker* class_linker = Runtime::Current()->GetClassLinker();
+  SetupCompiler(/*compiler_options=*/{});
+
+  if (kCompile) {
+    TimingLogger timings("OatTest::ChecksumDeterminism", /*precise=*/false, /*verbose=*/false);
+    CompileAll(/*class_loader=*/nullptr, class_linker->GetBootClassPath(), &timings);
+  }
+
+  auto write_elf_and_get_checksum = [&](OatKeyValueStore& key_value_store,
+                                        /*out*/ uint32_t* checksum) {
+    ScratchFile tmp_base, tmp_oat(tmp_base, ".oat"), tmp_vdex(tmp_base, ".vdex");
+
+    bool success = WriteElf(tmp_vdex.GetFile(),
+                            tmp_oat.GetFile(),
+                            class_linker->GetBootClassPath(),
+                            key_value_store,
+                            /*verify=*/false);
+    ASSERT_TRUE(success);
+
+    std::string error_msg;
+    std::unique_ptr<OatFile> oat_file(OatFile::Open(/*zip_fd=*/-1,
+                                                    tmp_oat.GetFilename(),
+                                                    tmp_oat.GetFilename(),
+                                                    /*executable=*/false,
+                                                    /*low_4gb=*/true,
+                                                    &error_msg));
+    ASSERT_TRUE(oat_file.get() != nullptr) << error_msg;
+    const OatHeader& oat_header = oat_file->GetOatHeader();
+    ASSERT_TRUE(oat_header.IsValid());
+    *checksum = oat_header.GetChecksum();
+  };
+
+  uint32_t checksum_1, checksum_2, checksum_3;
+
+  {
+    OatKeyValueStore key_value_store;
+    key_value_store.Put(OatHeader::kBootClassPathChecksumsKey, "testkey");
+    ASSERT_NO_FATAL_FAILURE(write_elf_and_get_checksum(key_value_store, &checksum_1));
+  }
+
+  {
+    // Put non-deterministic fields. This should not affect the checksum.
+    OatKeyValueStore key_value_store;
+    key_value_store.Put(OatHeader::kBootClassPathChecksumsKey, "testkey");
+    key_value_store.PutNonDeterministic(OatHeader::kDex2OatCmdLineKey, "cmdline");
+    key_value_store.PutNonDeterministic(OatHeader::kApexVersionsKey, "apex-versions");
+    ASSERT_NO_FATAL_FAILURE(write_elf_and_get_checksum(key_value_store, &checksum_2));
+    EXPECT_EQ(checksum_1, checksum_2);
+  }
+
+  {
+    // Put deterministic fields. This should affect the checksum.
+    OatKeyValueStore key_value_store;
+    key_value_store.Put(OatHeader::kBootClassPathChecksumsKey, "testkey");
+    key_value_store.Put(OatHeader::kClassPathKey, "classpath");
+    ASSERT_NO_FATAL_FAILURE(write_elf_and_get_checksum(key_value_store, &checksum_3));
+    EXPECT_NE(checksum_1, checksum_3);
+  }
+}
+
 TEST_F(OatTest, OatHeaderSizeCheck) {
   // If this test is failing and you have to update these constants,
   // it is time to update OatHeader::kOatVersion
@@ -548,7 +613,7 @@ TEST_F(OatTest, EmptyTextSection) {
   CompileAll(class_loader, dex_files, &timings);
 
   ScratchFile tmp_base, tmp_oat(tmp_base, ".oat"), tmp_vdex(tmp_base, ".vdex");
-  SafeMap<std::string, std::string> key_value_store;
+  OatKeyValueStore key_value_store;
   bool success = WriteElf(tmp_vdex.GetFile(),
                           tmp_oat.GetFile(),
                           dex_files,
@@ -617,7 +682,7 @@ void OatTest::TestDexFileInput(bool verify, bool low_4gb, bool use_profile) {
   input_dexfiles.push_back(std::move(dex_file2_data));
   scratch_files.push_back(&dex_file2);
 
-  SafeMap<std::string, std::string> key_value_store;
+  OatKeyValueStore key_value_store;
   {
     // Test using the AddDexFileSource() interface with the dex files.
     ScratchFile tmp_base, tmp_oat(tmp_base, ".oat"), tmp_vdex(tmp_base, ".vdex");
@@ -746,7 +811,7 @@ void OatTest::TestZipFileInput(bool verify, CopyOption copy) {
   success = zip_builder.Finish();
   ASSERT_TRUE(success) << strerror(errno);
 
-  SafeMap<std::string, std::string> key_value_store;
+  OatKeyValueStore key_value_store;
   {
     // Test using the AddDexFileSource() interface with the zip file.
     std::vector<const char*> input_filenames = { zip_file.GetFilename().c_str() };
@@ -865,7 +930,7 @@ void OatTest::TestZipFileInputWithEmptyDex() {
   success = zip_builder.Finish();
   ASSERT_TRUE(success) << strerror(errno);
 
-  SafeMap<std::string, std::string> key_value_store;
+  OatKeyValueStore key_value_store;
   std::vector<const char*> input_filenames = { zip_file.GetFilename().c_str() };
   ScratchFile oat_file, vdex_file(oat_file, ".vdex");
   std::unique_ptr<ProfileCompilationInfo> profile_compilation_info(new ProfileCompilationInfo());