From d1110fae6acc2d3edd7c09b86f3eccf09ea8fc2c Mon Sep 17 00:00:00 2001
From: Alexis Engelke <alexis.engelke@googlemail.com>
Date: Sun, 13 Jan 2019 14:27:04 +0100
Subject: [PATCH] Handle address size overrides

---
 decode.c | 34 ++++++++++++++++------------------
 decode.h |  2 ++
 2 files changed, 18 insertions(+), 18 deletions(-)

diff --git a/decode.c b/decode.c
index efa52e7..96627a2 100644
--- a/decode.c
+++ b/decode.c
@@ -610,6 +610,13 @@ decode(const uint8_t* buffer, int len, DecodeMode mode, Instr* instr)
         vec_size = 32;
     }
 
+    // Compute address size.
+    uint8_t addr_size_log = mode == DECODE_64 ? 4 : 3;
+    if (prefixes & PREFIX_ADDRSZ)
+        addr_size_log -= 1;
+    instr->addr_size = addr_size_log;
+    uint8_t addr_size = (1 << addr_size_log) >> 1;
+
     uint8_t operand_sizes[4] = {
         0, 1 << desc->gp_fixed_operand_size, op_size, vec_size
     };
@@ -680,29 +687,20 @@ decode(const uint8_t* buffer, int len, DecodeMode mode, Instr* instr)
         operand->reg = REG_NONE;
         operand->size = op_size;
         instr->scale = 0;
-        // TODO: Address size overrides
+
+        if (UNLIKELY(off + addr_size > len))
+            return -1;
 #if defined(ARCH_386)
-        if (mode == DECODE_32)
-        {
-            if (UNLIKELY(off + 4 > len))
-            {
-                return -1;
-            }
+        if (addr_size == 2)
+            instr->disp = LOAD_LE_2(&buffer[off]);
+#endif
+        if (addr_size == 4)
             instr->disp = LOAD_LE_4(&buffer[off]);
-            off += 4;
-        }
-#endif
 #if defined(ARCH_X86_64)
-        if (mode == DECODE_64)
-        {
-            if (UNLIKELY(off + 8 > len))
-            {
-                return -1;
-            }
+        if (addr_size == 8)
             instr->disp = LOAD_LE_8(&buffer[off]);
-            off += 8;
-        }
 #endif
+        off += addr_size;
     }
     else if (UNLIKELY(imm_control != 0))
     {
diff --git a/decode.h b/decode.h
index 43ac1bf..d24264d 100644
--- a/decode.h
+++ b/decode.h
@@ -118,6 +118,7 @@ struct Instr
     struct Operand operands[4];
     uint8_t segment : 3;
     uint8_t width : 3;
+    uint8_t addr_size : 2;
 
     /**
      * Encoded as 1 << (scale - 1)  **or**  no scaled register at all if zero.
@@ -137,6 +138,7 @@ typedef struct Instr Instr;
 
 #define INSTR_SEGMENT(instr) ((instr)->segment)
 #define INSTR_WIDTH(instr) ((1 << (instr)->width) >> 1)
+#define INSTR_ADDRSZ(instr) ((1 << (instr)->addr_size) >> 1)
 #define INSTR_HAS_REP(instr) ((instr)->prefixes & PREFIX_REP)
 #define INSTR_HAS_REPNZ(instr) ((instr)->prefixes & PREFIX_REPNZ)
 #define INSTR_HAS_LOCK(instr) ((instr)->prefixes & PREFIX_LOCK)