Merge pull request #558 from serpilliere/fix_x86_segm

Fix x86 segm

3 files changed, 269 insertions, 50 deletions

diff --git a/miasm2/arch/x86/jit.py b/miasm2/arch/x86/jit.py
index 6d9be8ac..9ffd8a82 100644
--- a/miasm2/arch/x86/jit.py
+++ b/miasm2/arch/x86/jit.py
@@ -80,6 +80,15 @@ class jitter_x86_32(jitter):
     def ir_archbloc_fix_regs_for_mode(self, irblock, attrib=64):
         self.orig_irbloc_fix_regs_for_mode(irblock, 64)
 
+    def push_uint16_t(self, value):
+        self.cpu.ESP -= self.ir_arch.sp.size / 8
+        self.vm.set_mem(self.cpu.ESP, pck16(value))
+
+    def pop_uint16_t(self):
+        value = upck16(self.vm.get_mem(self.cpu.ESP, self.ir_arch.sp.size / 8))
+        self.cpu.ESP += self.ir_arch.sp.size / 8
+        return value
+
     def push_uint32_t(self, value):
         self.cpu.ESP -= self.ir_arch.sp.size / 8
         self.vm.set_mem(self.cpu.ESP, pck32(value))
diff --git a/miasm2/arch/x86/sem.py b/miasm2/arch/x86/sem.py
index e32b8001..70927435 100644
--- a/miasm2/arch/x86/sem.py
+++ b/miasm2/arch/x86/sem.py
@@ -693,17 +693,14 @@ def push_gen(ir, instr, src, size):
         raise ValueError('bad size stacker!')
     if src.size < size:
         src = src.zeroExtend(size)
-    elif src.size == size:
-        pass
-    else:
-        raise ValueError('strange arg size')
+    off_size = src.size
 
     sp = mRSP[instr.mode]
-    new_sp = sp - m2_expr.ExprInt(size / 8, sp.size)
+    new_sp = sp - m2_expr.ExprInt(off_size / 8, sp.size)
     e.append(m2_expr.ExprAff(sp, new_sp))
     if ir.do_stk_segm:
         new_sp = m2_expr.ExprOp('segm', SS, new_sp)
-    e.append(m2_expr.ExprAff(ir.ExprMem(new_sp, size),
+    e.append(m2_expr.ExprAff(ir.ExprMem(new_sp, off_size),
                              src))
     return e, []
 
@@ -722,7 +719,7 @@ def pop_gen(ir, instr, src, size):
         raise ValueError('bad size stacker!')
 
     sp = mRSP[instr.mode]
-    new_sp = sp + m2_expr.ExprInt(size / 8, sp.size)
+    new_sp = sp + m2_expr.ExprInt(src.size / 8, sp.size)
     # don't generate ESP incrementation on POP ESP
     if src != ir.sp:
         e.append(m2_expr.ExprAff(sp, new_sp))
@@ -937,25 +934,34 @@ def cmps(ir, instr, size):
     lbl_df_1 = m2_expr.ExprId(ir.gen_label(), ir.IRDst.size)
     lbl_next = m2_expr.ExprId(ir.get_next_label(instr), ir.IRDst.size)
 
-    s = instr.v_admode()
-    a = ir.ExprMem(mRDI[instr.mode][:s], size)
-    b = ir.ExprMem(mRSI[instr.mode][:s], size)
+    src1 = mRSI[instr.mode][:instr.v_admode()]
+    src2 = mRDI[instr.mode][:instr.v_admode()]
+
+    if ir.do_str_segm:
+        if instr.additional_info.g2.value:
+            raise NotImplementedError("add segm support")
+        src1_sgm = m2_expr.ExprOp('segm', DS, src1)
+        src2_sgm = m2_expr.ExprOp('segm', ES, src2)
+    else:
+        src1_sgm = src1
+        src2_sgm = src2
+
+    offset = m2_expr.ExprInt(size / 8, src1.size)
+
+    e, _ = l_cmp(ir, instr,
+                 ir.ExprMem(src1_sgm, size),
+                 ir.ExprMem(src2_sgm, size))
 
-    e, _ = l_cmp(ir, instr, b, a)
 
     e0 = []
-    e0.append(m2_expr.ExprAff(a.arg,
-                              a.arg + m2_expr.ExprInt(size / 8, a.arg.size)))
-    e0.append(m2_expr.ExprAff(b.arg,
-                              b.arg + m2_expr.ExprInt(size / 8, b.arg.size)))
+    e0.append(m2_expr.ExprAff(src1, src1 + offset))
+    e0.append(m2_expr.ExprAff(src2, src2 + offset))
     e0.append(m2_expr.ExprAff(ir.IRDst, lbl_next))
     e0 = IRBlock(lbl_df_0.name, [e0])
 
     e1 = []
-    e1.append(m2_expr.ExprAff(a.arg,
-                              a.arg - m2_expr.ExprInt(size / 8, a.arg.size)))
-    e1.append(m2_expr.ExprAff(b.arg,
-                              b.arg - m2_expr.ExprInt(size / 8, b.arg.size)))
+    e1.append(m2_expr.ExprAff(src1, src1 - offset))
+    e1.append(m2_expr.ExprAff(src2, src2 - offset))
     e1.append(m2_expr.ExprAff(ir.IRDst, lbl_next))
     e1 = IRBlock(lbl_df_1.name, [e1])
 
@@ -969,20 +975,28 @@ def scas(ir, instr, size):
     lbl_df_1 = m2_expr.ExprId(ir.gen_label(), ir.IRDst.size)
     lbl_next = m2_expr.ExprId(ir.get_next_label(instr), ir.IRDst.size)
 
-    s = instr.v_admode()
-    a = ir.ExprMem(mRDI[instr.mode][:s], size)
+    src = mRDI[instr.mode][:instr.v_admode()]
 
-    e, extra = l_cmp(ir, instr, mRAX[instr.mode][:size], a)
+    if ir.do_str_segm:
+        if instr.additional_info.g2.value:
+            raise NotImplementedError("add segm support")
+        src_sgm = m2_expr.ExprOp('segm', ES, src)
+    else:
+        src_sgm = src
+
+    offset = m2_expr.ExprInt(size / 8, src.size)
+    e, extra = l_cmp(ir, instr,
+                     mRAX[instr.mode][:size],
+                     ir.ExprMem(src_sgm, size))
 
     e0 = []
-    e0.append(m2_expr.ExprAff(a.arg,
-                              a.arg + m2_expr.ExprInt(size / 8, a.arg.size)))
+    e0.append(m2_expr.ExprAff(src, src + offset))
+
     e0.append(m2_expr.ExprAff(ir.IRDst, lbl_next))
     e0 = IRBlock(lbl_df_0.name, [e0])
 
     e1 = []
-    e1.append(m2_expr.ExprAff(a.arg,
-                              a.arg - m2_expr.ExprInt(size / 8, a.arg.size)))
+    e1.append(m2_expr.ExprAff(src, src - offset))
     e1.append(m2_expr.ExprAff(ir.IRDst, lbl_next))
     e1 = IRBlock(lbl_df_1.name, [e1])
 
@@ -1081,12 +1095,11 @@ pa_regs = [
 
 def pusha_gen(ir, instr, size):
     e = []
+    cur_sp = mRSP[instr.mode]
     for i, reg in enumerate(pa_regs):
-        stk_ptr = mRSP[instr.mode] + \
-            m2_expr.ExprInt(-(reg[size].size / 8) * (i + 1), instr.mode)
-        e.append(m2_expr.ExprAff(ir.ExprMem(
-            stk_ptr, reg[size].size), reg[size]))
-    e.append(m2_expr.ExprAff(mRSP[instr.mode], stk_ptr))
+        stk_ptr = cur_sp + m2_expr.ExprInt(-(size / 8) * (i + 1), instr.mode)
+        e.append(m2_expr.ExprAff(ir.ExprMem(stk_ptr, size), reg[size]))
+    e.append(m2_expr.ExprAff(cur_sp, stk_ptr))
     return e, []
 
 
@@ -1100,16 +1113,15 @@ def pushad(ir, instr):
 
 def popa_gen(ir, instr, size):
     e = []
+    cur_sp = mRSP[instr.mode]
     for i, reg in enumerate(reversed(pa_regs)):
         if reg == mRSP:
             continue
-        stk_ptr = mRSP[instr.mode] + \
-            m2_expr.ExprInt((reg[size].size / 8) * i, instr.mode)
-        e.append(m2_expr.ExprAff(reg[size], ir.ExprMem(stk_ptr, instr.mode)))
+        stk_ptr = cur_sp + m2_expr.ExprInt((size / 8) * i, instr.mode)
+        e.append(m2_expr.ExprAff(reg[size], ir.ExprMem(stk_ptr, size)))
 
-    stk_ptr = mRSP[instr.mode] + \
-        m2_expr.ExprInt((instr.mode / 8) * (i + 1), instr.mode)
-    e.append(m2_expr.ExprAff(mRSP[instr.mode], stk_ptr))
+    stk_ptr = cur_sp + m2_expr.ExprInt((size / 8) * (i + 1), instr.mode)
+    e.append(m2_expr.ExprAff(cur_sp, stk_ptr))
 
     return e, []
 
@@ -1726,29 +1738,34 @@ def movs(ir, instr, size):
     lbl_df_1 = m2_expr.ExprId(ir.gen_label(), ir.IRDst.size)
     lbl_next = m2_expr.ExprId(ir.get_next_label(instr), ir.IRDst.size)
 
-    a = mRDI[instr.mode][:instr.v_admode()]
-    b = mRSI[instr.mode][:instr.v_admode()]
+    dst = mRDI[instr.mode][:instr.v_admode()]
+    src = mRSI[instr.mode][:instr.v_admode()]
 
     e = []
-    src = b
-    dst = a
     if ir.do_str_segm:
         if instr.additional_info.g2.value:
             raise NotImplementedError("add segm support")
-        src = m2_expr.ExprOp('segm', DS, src)
-        dst = m2_expr.ExprOp('segm', ES, dst)
-    e.append(m2_expr.ExprAff(ir.ExprMem(dst, size),
-                             ir.ExprMem(src, size)))
+        src_sgm = m2_expr.ExprOp('segm', DS, src)
+        dst_sgm = m2_expr.ExprOp('segm', ES, dst)
+
+    else:
+        src_sgm = src
+        dst_sgm = dst
+
+    offset = m2_expr.ExprInt(size / 8, src.size)
+
+    e.append(m2_expr.ExprAff(ir.ExprMem(dst_sgm, size),
+                             ir.ExprMem(src_sgm, size)))
 
     e0 = []
-    e0.append(m2_expr.ExprAff(a, a + m2_expr.ExprInt(size / 8, a.size)))
-    e0.append(m2_expr.ExprAff(b, b + m2_expr.ExprInt(size / 8, b.size)))
+    e0.append(m2_expr.ExprAff(src, src + offset))
+    e0.append(m2_expr.ExprAff(dst, dst + offset))
     e0.append(m2_expr.ExprAff(ir.IRDst, lbl_next))
     e0 = IRBlock(lbl_df_0.name, [e0])
 
     e1 = []
-    e1.append(m2_expr.ExprAff(a, a - m2_expr.ExprInt(size / 8, a.size)))
-    e1.append(m2_expr.ExprAff(b, b - m2_expr.ExprInt(size / 8, b.size)))
+    e1.append(m2_expr.ExprAff(src, src - offset))
+    e1.append(m2_expr.ExprAff(dst, dst - offset))
     e1.append(m2_expr.ExprAff(ir.IRDst, lbl_next))
     e1 = IRBlock(lbl_df_1.name, [e1])
 
diff --git a/test/arch/x86/unit/mn_pushpop.py b/test/arch/x86/unit/mn_pushpop.py
index ffcc3fa5..7ac400c0 100755
--- a/test/arch/x86/unit/mn_pushpop.py
+++ b/test/arch/x86/unit/mn_pushpop.py
@@ -39,6 +39,7 @@ class Test_PUSHAD_32(Asm_Test_32):
     '''
 
     def check(self):
+        assert self.myjit.cpu.ESP == self.stk_origin - 0x4 * 8
         buf = self.myjit.vm.get_mem(self.myjit.cpu.ESP, 0x4 * 8)
         assert(buf == self.buf)
 
@@ -65,6 +66,7 @@ class Test_PUSHA_32(Asm_Test_32):
     '''
 
     def check(self):
+        assert self.myjit.cpu.ESP == self.stk_origin - 0x2 * 8
         buf = self.myjit.vm.get_mem(self.myjit.cpu.ESP, 0x2 * 8)
         assert(buf == self.buf)
 
@@ -91,6 +93,7 @@ class Test_PUSHA_16(Asm_Test_16):
     '''
 
     def check(self):
+        assert self.myjit.cpu.ESP == self.stk_origin - 0x2 * 8
         buf = self.myjit.vm.get_mem(self.myjit.cpu.SP, 0x2 * 8)
         assert(buf == self.buf)
 
@@ -117,12 +120,202 @@ class Test_PUSHAD_16(Asm_Test_16):
     '''
 
     def check(self):
+        assert self.myjit.cpu.ESP == self.stk_origin - 0x4 * 8
         buf = self.myjit.vm.get_mem(self.myjit.cpu.SP, 0x4 * 8)
         assert(buf == self.buf)
 
 
+class Test_PUSH_mode32_32(Asm_Test_32):
+    MYSTRING = "test push mode32 32"
+
+    def prepare(self):
+        self.myjit.ir_arch.symbol_pool.add_label("lbl_ret", self.ret_addr)
+
+    def test_init(self):
+        init_regs(self)
+        self.buf = ""
+        self.buf += pck32(0x11223344)
+
+    TXT = '''
+    main:
+       PUSH 0x11223344
+       JMP lbl_ret
+    '''
+
+    def check(self):
+        assert self.myjit.cpu.ESP == self.stk_origin - 0x4
+        buf = self.myjit.vm.get_mem(self.myjit.cpu.ESP, 0x4)
+        assert(buf == self.buf)
+
+
+class Test_PUSH_mode32_16(Asm_Test_32):
+    MYSTRING = "test push mode32 16"
+
+    def prepare(self):
+        self.myjit.ir_arch.symbol_pool.add_label("lbl_ret", self.ret_addr)
+
+    def test_init(self):
+        init_regs(self)
+        self.buf = ""
+        self.buf += pck16(0x1122)
+
+    TXT = '''
+    main:
+       PUSHW 0x1122
+       JMP lbl_ret
+    '''
+
+    def check(self):
+        assert self.myjit.cpu.ESP == self.stk_origin - 0x2
+        buf = self.myjit.vm.get_mem(self.myjit.cpu.ESP, 0x2)
+        assert(buf == self.buf)
+
+
+class Test_PUSH_mode16_16(Asm_Test_16):
+    MYSTRING = "test push mode16 16"
+
+    def prepare(self):
+        self.myjit.ir_arch.symbol_pool.add_label("lbl_ret", self.ret_addr)
+
+    def test_init(self):
+        init_regs(self)
+        self.buf = ""
+        self.buf += pck16(0x1122)
+
+    TXT = '''
+    main:
+       PUSHW 0x1122
+       JMP lbl_ret
+    '''
+
+    def check(self):
+        assert self.myjit.cpu.ESP == self.stk_origin - 0x2
+        buf = self.myjit.vm.get_mem(self.myjit.cpu.ESP, 0x2)
+        assert(buf == self.buf)
+
+
+class Test_PUSH_mode16_32(Asm_Test_16):
+    MYSTRING = "test push mode16 32"
+
+    def prepare(self):
+        self.myjit.ir_arch.symbol_pool.add_label("lbl_ret", self.ret_addr)
+
+    def test_init(self):
+        init_regs(self)
+        self.buf = ""
+        self.buf += pck32(0x11223344)
+
+    TXT = '''
+    main:
+       .byte 0x66, 0x68, 0x44, 0x33, 0x22, 0x11
+       JMP lbl_ret
+    '''
+
+    def check(self):
+        assert self.myjit.cpu.ESP == self.stk_origin - 0x4
+        buf = self.myjit.vm.get_mem(self.myjit.cpu.ESP, 0x4)
+        assert(buf == self.buf)
+
+
+class Test_POP_mode32_32(Asm_Test_32):
+    MYSTRING = "test pop mode32 32"
+
+    def prepare(self):
+        self.myjit.ir_arch.symbol_pool.add_label("lbl_ret", self.ret_addr)
+
+    def test_init(self):
+        self.value = 0x11223344
+        self.myjit.push_uint32_t(self.value)
+        init_regs(self)
+
+    TXT = '''
+    main:
+       POP EAX
+       JMP lbl_ret
+    '''
+
+    def check(self):
+        assert self.myjit.cpu.ESP == self.stk_origin + 0x4
+        assert self.myjit.cpu.EAX == self.value
+
+
+class Test_POP_mode32_16(Asm_Test_32):
+    MYSTRING = "test pop mode32 16"
+
+    def prepare(self):
+        self.myjit.ir_arch.symbol_pool.add_label("lbl_ret", self.ret_addr)
+
+    def test_init(self):
+        self.value = 0x1122
+        self.myjit.push_uint16_t(self.value)
+        init_regs(self)
+
+    TXT = '''
+    main:
+       POPW AX
+       JMP lbl_ret
+    '''
+
+    def check(self):
+        assert self.myjit.cpu.ESP == self.stk_origin + 0x2
+        assert self.myjit.cpu.AX == self.value
+
+
+class Test_POP_mode16_16(Asm_Test_16):
+    MYSTRING = "test pop mode16 16"
+
+    def prepare(self):
+        self.myjit.ir_arch.symbol_pool.add_label("lbl_ret", self.ret_addr)
+
+    def test_init(self):
+        self.value = 0x1122
+        self.myjit.push_uint16_t(self.value)
+        init_regs(self)
+
+    TXT = '''
+    main:
+       POPW AX
+       JMP lbl_ret
+    '''
+
+    def check(self):
+        assert self.myjit.cpu.ESP == self.stk_origin + 0x2
+        assert self.myjit.cpu.AX == self.value
+
+
+class Test_POP_mode16_32(Asm_Test_16):
+    MYSTRING = "test pop mode16 32"
+
+    def prepare(self):
+        self.myjit.ir_arch.symbol_pool.add_label("lbl_ret", self.ret_addr)
+
+    def test_init(self):
+        self.value = 0x11223344
+        self.myjit.cpu.SP -= 0x4
+        self.myjit.vm.set_mem(self.myjit.cpu.SP, pck32(self.value))
+        init_regs(self)
+
+    TXT = '''
+    main:
+       POP EAX
+       JMP lbl_ret
+    '''
+
+    def check(self):
+        assert self.myjit.cpu.ESP == self.stk_origin + 0x4
+        assert self.myjit.cpu.EAX == self.value
+
+
 if __name__ == "__main__":
     [test(*sys.argv[1:])() for test in [Test_PUSHA_16, Test_PUSHA_32,
-                                        Test_PUSHAD_16, Test_PUSHAD_32
+                                        Test_PUSHAD_16, Test_PUSHAD_32,
+                                        Test_PUSH_mode32_32,
+                                        Test_PUSH_mode32_16,
+                                        Test_PUSH_mode16_16,
+                                        Test_PUSH_mode16_32,
+                                        Test_POP_mode32_32,
+                                        Test_POP_mode32_16,
+                                        Test_POP_mode16_16,
+                                        Test_POP_mode16_32,
                                         ]
     ]