Make SafeDexInstructionIterator more robust

Make operator++ set the error state when if incrementing would go
past the end of the code region.

Added test in CodeItemOverrun.

Bug: 67104794
Test: test-art-host

Change-Id: I7b27a1ddb60aaae08ae9ed4ba8762ce2a7746770

2 files changed, 15 insertions, 3 deletions

diff --git a/dexlayout/dexlayout_test.cc b/dexlayout/dexlayout_test.cc
index cdd0219ef3..38d3c6ec07 100644
--- a/dexlayout/dexlayout_test.cc
+++ b/dexlayout/dexlayout_test.cc
@@ -727,6 +727,13 @@ TEST_F(DexLayoutTest, CodeItemOverrun) {
               const_cast<Instruction*>(last_instruction.Inst())->SetOpcode(
                   Instruction::CONST_STRING_JUMBO);
               mutated_successfully = true;
+              // Test that the safe iterator doesn't go past the end.
+              SafeDexInstructionIterator it2(instructions.begin(), instructions.end());
+              while (!it2.IsErrorState()) {
+                ++it2;
+              }
+              EXPECT_TRUE(it2 == last_instruction);
+              EXPECT_TRUE(it2 < instructions.end());
             }
           }
         }
diff --git a/runtime/dex_instruction_iterator.h b/runtime/dex_instruction_iterator.h
index afa8d4aad3..9e4dea34bb 100644
--- a/runtime/dex_instruction_iterator.h
+++ b/runtime/dex_instruction_iterator.h
@@ -123,11 +123,16 @@ class SafeDexInstructionIterator : public DexInstructionIteratorBase {
     const size_t size_code_units = Inst()->CodeUnitsRequiredForSizeComputation();
     const size_t available = reinterpret_cast<const uint16_t*>(end_) -
         reinterpret_cast<const uint16_t*>(Inst());
-    if (size_code_units > available) {
+    if (UNLIKELY(size_code_units > available)) {
       error_state_ = true;
-    } else {
-      inst_ = inst_->Next();
+      return *this;
     }
+    const size_t instruction_size = inst_->SizeInCodeUnits();
+    if (UNLIKELY(instruction_size > available)) {
+      error_state_ = true;
+      return *this;
+    }
+    inst_ = inst_->RelativeAt(instruction_size);
     return *this;
   }