Fix DexInstructionIterator overrun bug

Handle cases where the dex instructions can go past the end of the
code item in dexlayout.

Since dexlayout runs before method verification, we need to be
careful to not go past the end of the code item.

Added test.

Bug: 67104794
Test: test-art-host

Change-Id: Idf7d51344659b2c75207bdf444e39f271feb8d3a
diff --git a/dexlayout/dex_ir.cc b/dexlayout/dex_ir.cc
index 8c4ee6e..23c3a5c 100644
--- a/dexlayout/dex_ir.cc
+++ b/dexlayout/dex_ir.cc
@@ -167,10 +167,17 @@
                                std::vector<MethodId*>* method_ids,
                                std::vector<FieldId*>* field_ids) {
   bool has_id = false;
-  for (const Instruction& instruction : code->Instructions()) {
-    CHECK_GT(instruction.SizeInCodeUnits(), 0u);
+  IterationRange<DexInstructionIterator> instructions = code->Instructions();
+  SafeDexInstructionIterator it(instructions.begin(), instructions.end());
+  for (; !it.IsErrorState() && it < instructions.end(); ++it) {
+    // In case the instruction goes past the end of the code item, make sure to not process it.
+    SafeDexInstructionIterator next = it;
+    ++next;
+    if (next.IsErrorState() || next > instructions.end()) {
+      break;
+    }
     has_id |= GetIdFromInstruction(collections,
-                                   &instruction,
+                                   it.Inst(),
                                    type_ids,
                                    string_ids,
                                    method_ids,
diff --git a/dexlayout/dexlayout_test.cc b/dexlayout/dexlayout_test.cc
index f34e7ec..cdd0219 100644
--- a/dexlayout/dexlayout_test.cc
+++ b/dexlayout/dexlayout_test.cc
@@ -317,6 +317,30 @@
     return true;
   }
 
+  template <typename Mutator>
+  bool MutateDexFile(File* output_dex, const std::string& input_jar, const Mutator& mutator) {
+    std::vector<std::unique_ptr<const DexFile>> dex_files;
+    std::string error_msg;
+    CHECK(DexFileLoader::Open(input_jar.c_str(),
+                              input_jar.c_str(),
+                              /*verify*/ true,
+                              /*verify_checksum*/ true,
+                              &error_msg,
+                              &dex_files)) << error_msg;
+    EXPECT_EQ(dex_files.size(), 1u) << "Only one input dex is supported";
+    for (const std::unique_ptr<const DexFile>& dex : dex_files) {
+      CHECK(dex->EnableWrite()) << "Failed to enable write";
+      mutator(const_cast<DexFile*>(dex.get()));
+      if (!output_dex->WriteFully(dex->Begin(), dex->Size())) {
+        return false;
+      }
+    }
+    if (output_dex->Flush() != 0) {
+      PLOG(FATAL) << "Could not flush the output file.";
+    }
+    return true;
+  }
+
   // Create a profile with some subset of methods and classes.
   void CreateProfile(const std::string& input_dex,
                      const std::string& out_profile,
@@ -518,8 +542,10 @@
                      const char* dex_filename,
                      ScratchFile* profile_file,
                      std::vector<std::string>& dexlayout_exec_argv) {
-    WriteBase64ToFile(dex_filename, dex_file->GetFile());
-    EXPECT_EQ(dex_file->GetFile()->Flush(), 0);
+    if (dex_filename != nullptr) {
+      WriteBase64ToFile(dex_filename, dex_file->GetFile());
+      EXPECT_EQ(dex_file->GetFile()->Flush(), 0);
+    }
     if (profile_file != nullptr) {
       CreateProfile(dex_file->GetFilename(), profile_file->GetFilename(), dex_file->GetFilename());
     }
@@ -673,4 +699,51 @@
                             dexlayout_exec_argv));
 }
 
+// Test that instructions that go past the end of the code items don't cause crashes.
+TEST_F(DexLayoutTest, CodeItemOverrun) {
+  ScratchFile temp_dex;
+  MutateDexFile(temp_dex.GetFile(), GetTestDexFileName("ManyMethods"), [] (DexFile* dex) {
+    bool mutated_successfully = false;
+    // Change the dex instructions to make an opcode that spans past the end of the code item.
+    for (size_t i = 0; i < dex->NumClassDefs(); ++i) {
+      const DexFile::ClassDef& def = dex->GetClassDef(i);
+      const uint8_t* data = dex->GetClassData(def);
+      if (data == nullptr) {
+        continue;
+      }
+      ClassDataItemIterator it(*dex, data);
+      it.SkipAllFields();
+      while (it.HasNextDirectMethod() || it.HasNextVirtualMethod()) {
+        DexFile::CodeItem* item = const_cast<DexFile::CodeItem*>(it.GetMethodCodeItem());
+        if (item != nullptr) {
+          IterationRange<DexInstructionIterator> instructions = item->Instructions();
+          if (instructions.begin() != instructions.end()) {
+            DexInstructionIterator last_instruction = instructions.begin();
+            for (auto dex_it = instructions.begin(); dex_it != instructions.end(); ++dex_it) {
+              last_instruction = dex_it;
+            }
+            if (last_instruction->SizeInCodeUnits() == 1) {
+              // Set the opcode to something that will go past the end of the code item.
+              const_cast<Instruction*>(last_instruction.Inst())->SetOpcode(
+                  Instruction::CONST_STRING_JUMBO);
+              mutated_successfully = true;
+            }
+          }
+        }
+        it.Next();
+      }
+    }
+    CHECK(mutated_successfully)
+        << "Failed to find candidate code item with only one code unit in last instruction.";
+  });
+  std::string dexlayout = GetTestAndroidRoot() + "/bin/dexlayout";
+  EXPECT_TRUE(OS::FileExists(dexlayout.c_str())) << dexlayout << " should be a valid file path";
+  std::vector<std::string> dexlayout_exec_argv =
+      { dexlayout, "-i", "-o", "/dev/null", temp_dex.GetFilename() };
+  ASSERT_TRUE(DexLayoutExec(&temp_dex,
+                            /*dex_filename*/ nullptr,
+                            nullptr /* profile_file */,
+                            dexlayout_exec_argv));
+}
+
 }  // namespace art