From 4371ed02ac30cb56ec4ddbf60c87e85c183d860b Mon Sep 17 00:00:00 2001
From: iximeow <me@iximeow.net>
Date: Thu, 22 Jul 2021 00:31:02 -0700
Subject: fix incorrect decodes with scas and 67-prefixes

---
 src/protected_mode/mod.rs | 56 +++++++++++++++++++++++++++++++++++++++--------
 1 file changed, 47 insertions(+), 9 deletions(-)

(limited to 'src/protected_mode')

diff --git a/src/protected_mode/mod.rs b/src/protected_mode/mod.rs
index 6e503a0..19fe881 100644
--- a/src/protected_mode/mod.rs
+++ b/src/protected_mode/mod.rs
@@ -525,6 +525,8 @@ impl OperandSpec {
             OperandSpec::DispU16 |
             OperandSpec::DispU32 |
             OperandSpec::Deref |
+            OperandSpec::Deref_si |
+            OperandSpec::Deref_di |
             OperandSpec::Deref_esi |
             OperandSpec::Deref_edi |
             OperandSpec::RegDisp |
@@ -704,6 +706,12 @@ impl Operand {
             OperandSpec::Deref => {
                 Operand::RegDeref(inst.regs[1])
             }
+            OperandSpec::Deref_si => {
+                Operand::RegDeref(RegSpec::si())
+            }
+            OperandSpec::Deref_di => {
+                Operand::RegDeref(RegSpec::di())
+            }
             OperandSpec::Deref_esi => {
                 Operand::RegDeref(RegSpec::esi())
             }
@@ -2720,6 +2728,8 @@ enum OperandSpec {
     DispU16,
     DispU32,
     Deref,
+    Deref_si,
+    Deref_di,
     Deref_esi,
     Deref_edi,
     RegDisp,
@@ -4357,7 +4367,8 @@ impl Instruction {
     /// prefixes.
     pub fn segment_override_for_op(&self, op: u8) -> Option<Segment> {
         match self.opcode {
-            Opcode::STOS => {
+            Opcode::STOS |
+            Opcode::SCAS => {
                 if op == 0 {
                     Some(Segment::ES)
                 } else {
@@ -5748,6 +5759,7 @@ fn read_M_16bit<T: Reader<<Arch as yaxpeax_arch::Arch>::Address, <Arch as yaxpea
     let modbits = modrm >> 6;
     let mmm = modrm & 7;
     if modbits == 0b00 && mmm == 0b110 {
+        instr.disp = read_num(words, 2)? as u16 as u32;
         return Ok(OperandSpec::DispU16);
     }
     match mmm {
@@ -8975,21 +8987,34 @@ fn unlikely_operands<T: Reader<<Arch as yaxpeax_arch::Arch>::Address, <Arch as y
         // sure hope these aren't backwards huh
         OperandCode::AL_Xb => {
             instruction.regs[0] = RegSpec::al();
-            instruction.regs[1] = RegSpec::esi();
+            if instruction.prefixes.address_size() {
+                instruction.regs[1] = RegSpec::si();
+            } else {
+                instruction.regs[1] = RegSpec::esi();
+            }
             instruction.operands[0] = OperandSpec::RegRRR;
             instruction.operands[1] = OperandSpec::Deref;
             instruction.mem_size = 1;
             instruction.operand_count = 2;
         }
         OperandCode::Yb_Xb => {
-            instruction.operands[0] = OperandSpec::Deref_edi;
-            instruction.operands[1] = OperandSpec::Deref_esi;
+            if instruction.prefixes.address_size() {
+                instruction.operands[0] = OperandSpec::Deref_di;
+                instruction.operands[1] = OperandSpec::Deref_si;
+            } else {
+                instruction.operands[0] = OperandSpec::Deref_edi;
+                instruction.operands[1] = OperandSpec::Deref_esi;
+            }
             instruction.mem_size = 1;
             instruction.operand_count = 2;
         }
         OperandCode::Yb_AL => {
             instruction.regs[0] = RegSpec::al();
-            instruction.regs[1] = RegSpec::esi();
+            if instruction.prefixes.address_size() {
+                instruction.regs[1] = RegSpec::di();
+            } else {
+                instruction.regs[1] = RegSpec::edi();
+            }
             instruction.operands[0] = OperandSpec::Deref;
             instruction.operands[1] = OperandSpec::RegRRR;
             instruction.mem_size = 1;
@@ -9003,7 +9028,11 @@ fn unlikely_operands<T: Reader<<Arch as yaxpeax_arch::Arch>::Address, <Arch as y
                 instruction.mem_size = 4;
                 RegSpec::eax()
             };
-            instruction.regs[1] = RegSpec::esi();
+            if instruction.prefixes.address_size() {
+                instruction.regs[1] = RegSpec::si();
+            } else {
+                instruction.regs[1] = RegSpec::esi();
+            }
             instruction.operands[1] = OperandSpec::Deref;
         }
         OperandCode::Yv_AX => {
@@ -9014,7 +9043,11 @@ fn unlikely_operands<T: Reader<<Arch as yaxpeax_arch::Arch>::Address, <Arch as y
                 instruction.mem_size = 4;
                 RegSpec::eax()
             };
-            instruction.regs[1] = RegSpec::edi();
+            if instruction.prefixes.address_size() {
+                instruction.regs[1] = RegSpec::di();
+            } else {
+                instruction.regs[1] = RegSpec::edi();
+            }
             instruction.operands[0] = OperandSpec::Deref;
             instruction.operands[1] = OperandSpec::RegRRR;
         }
@@ -9024,8 +9057,13 @@ fn unlikely_operands<T: Reader<<Arch as yaxpeax_arch::Arch>::Address, <Arch as y
             } else {
                 4
             };
-            instruction.operands[0] = OperandSpec::Deref_edi;
-            instruction.operands[1] = OperandSpec::Deref_esi;
+            if instruction.prefixes.address_size() {
+                instruction.operands[0] = OperandSpec::Deref_di;
+                instruction.operands[1] = OperandSpec::Deref_si;
+            } else {
+                instruction.operands[0] = OperandSpec::Deref_edi;
+                instruction.operands[1] = OperandSpec::Deref_esi;
+            }
         }
         OperandCode::ModRM_0x0f12 => {
             instruction.regs[0].bank = RegisterBank::X;
-- 
cgit v1.1