From 44e412cffb14c98806b1eb5a91bbccf5180b9dbb Mon Sep 17 00:00:00 2001
From: _xeroxz <xerox@back.engineering>
Date: Fri, 11 Jun 2021 14:29:33 -0700
Subject: [PATCH] added a param to the lifter_callback_t...

---
 include/vmlifters.hpp     |  2 +-
 src/vmlifters/add.cpp     |  6 +++---
 src/vmlifters/jmp.cpp     | 12 ++++++++++--
 src/vmlifters/lconst.cpp  | 27 ++++++++++++++++++---------
 src/vmlifters/lflags.cpp  |  4 ++--
 src/vmlifters/lreg.cpp    |  6 ++++--
 src/vmlifters/lvsp.cpp    |  4 ++--
 src/vmlifters/nand.cpp    |  9 ++++++---
 src/vmlifters/pushvsp.cpp |  4 ++--
 src/vmlifters/read.cpp    |  9 ++++++---
 src/vmlifters/shr.cpp     |  6 ++++--
 src/vmlifters/sreg.cpp    |  9 ++++++---
 src/vmlifters/vmexit.cpp  |  3 ++-
 13 files changed, 66 insertions(+), 35 deletions(-)

diff --git a/include/vmlifters.hpp b/include/vmlifters.hpp
index e0ec8b1..955f840 100644
--- a/include/vmlifters.hpp
+++ b/include/vmlifters.hpp
@@ -4,7 +4,7 @@
 
 namespace vm::lifters
 {
-    using lifter_callback_t = std::function< void( vtil::basic_block *, vm::instrs::virt_instr_t * ) >;
+    using lifter_callback_t = std::function< void( vtil::basic_block *, vm::instrs::virt_instr_t *, vmp2::v3::code_block_t* ) >;
     using lifter_t = std::pair< vm::handler::mnemonic_t, lifter_callback_t >;
 
     // taken from
diff --git a/src/vmlifters/add.cpp b/src/vmlifters/add.cpp
index f52bdc9..65fe6d7 100644
--- a/src/vmlifters/add.cpp
+++ b/src/vmlifters/add.cpp
@@ -4,7 +4,7 @@ namespace vm::lifters
 {
     vm::lifters::lifter_t addq = {
         // vsp[0] = vsp[1] + vsp[0];
-        vm::handler::ADDQ, []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr ) {
+        vm::handler::ADDQ, []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr, vmp2::v3::code_block_t* code_blk ) {
             auto [ t0, t1 ] = blk->tmp( 64, 64 );
             blk->pop( t0 );
             blk->pop( t1 );
@@ -15,7 +15,7 @@ namespace vm::lifters
 
     vm::lifters::lifter_t adddw = {
         // vsp[0] = vsp[1] + vsp[0];
-        vm::handler::ADDDW, []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr ) {
+        vm::handler::ADDDW, []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr, vmp2::v3::code_block_t* code_blk ) {
             auto [ t0, t1 ] = blk->tmp( 32, 32 );
             blk->pop( t0 );
             blk->pop( t1 );
@@ -26,7 +26,7 @@ namespace vm::lifters
 
     vm::lifters::lifter_t addw = {
         // vsp[0] = vsp[1] + vsp[0];
-        vm::handler::ADDW, []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr ) {
+        vm::handler::ADDW, []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr, vmp2::v3::code_block_t* code_blk ) {
             auto [ t0, t1 ] = blk->tmp( 16, 16 );
             blk->pop( t0 );
             blk->pop( t1 );
diff --git a/src/vmlifters/jmp.cpp b/src/vmlifters/jmp.cpp
index 538346c..4f07da6 100644
--- a/src/vmlifters/jmp.cpp
+++ b/src/vmlifters/jmp.cpp
@@ -4,7 +4,15 @@ namespace vm::lifters
 {
     vm::lifters::lifter_t jmp = {
         // jmp
-        vm::handler::JMP, []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr ) {
-            blk->jmp( vinstr->trace_data.vsp.qword[ 0 ] );
+        vm::handler::JMP,
+        []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr, vmp2::v3::code_block_t *code_blk ) {
+            if ( code_blk->jcc.has_jcc )
+            {
+                // TODO: figure out what bit the JCC is determined on...
+                blk->js( vtil::REG_FLAGS.select( 1, 11 ), code_blk->jcc.block_addr[ 0 ],
+                         code_blk->jcc.block_addr[ 1 ] );
+            }
+            else
+                blk->jmp( vinstr->trace_data.vsp.qword[ 0 ] );
         } };
 }
\ No newline at end of file
diff --git a/src/vmlifters/lconst.cpp b/src/vmlifters/lconst.cpp
index 2b246e3..964e3a2 100644
--- a/src/vmlifters/lconst.cpp
+++ b/src/vmlifters/lconst.cpp
@@ -4,55 +4,64 @@ namespace vm::lifters
 {
     vm::lifters::lifter_t lconstq = {
         // push imm<N>
-        vm::handler::LCONSTQ, []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr ) {
+        vm::handler::LCONSTQ,
+        []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr, vmp2::v3::code_block_t *code_blk ) {
             blk->push( vtil::operand( vinstr->operand.imm.u, 64 ) );
         } };
 
     vm::lifters::lifter_t lconstdw = {
         // push imm<N>
-        vm::handler::LCONSTDW, []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr ) {
+        vm::handler::LCONSTDW,
+        []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr, vmp2::v3::code_block_t *code_blk ) {
             blk->push( vtil::operand( vinstr->operand.imm.u, 32 ) );
         } };
 
     vm::lifters::lifter_t lconstw = {
         // push imm<N>
-        vm::handler::LCONSTW, []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr ) {
+        vm::handler::LCONSTW,
+        []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr, vmp2::v3::code_block_t *code_blk ) {
             blk->push( vtil::operand( vinstr->operand.imm.u, 16 ) );
         } };
 
     vm::lifters::lifter_t lconstbzxw = {
         // push imm<N>
-        vm::handler::LCONSTBZXW, []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr ) {
+        vm::handler::LCONSTBZXW,
+        []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr, vmp2::v3::code_block_t *code_blk ) {
             blk->push( vtil::operand( vinstr->operand.imm.u, 16 ) );
         } };
 
     vm::lifters::lifter_t lconstbsxdw = {
         // push imm<N>
-        vm::handler::LCONSTBSXDW, []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr ) {
+        vm::handler::LCONSTBSXDW,
+        []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr, vmp2::v3::code_block_t *code_blk ) {
             blk->push( vtil::operand( vinstr->operand.imm.u, 32 ) );
         } };
 
     vm::lifters::lifter_t lconstbsxq = {
         // push imm<N>
-        vm::handler::LCONSTBSXQ, []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr ) {
+        vm::handler::LCONSTBSXQ,
+        []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr, vmp2::v3::code_block_t *code_blk ) {
             blk->push( vtil::operand( vinstr->operand.imm.u, 64 ) );
         } };
 
     vm::lifters::lifter_t lconstdwsxq = {
         // push imm<N>
-        vm::handler::LCONSTDWSXQ, []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr ) {
+        vm::handler::LCONSTDWSXQ,
+        []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr, vmp2::v3::code_block_t *code_blk ) {
             blk->push( vtil::operand( vinstr->operand.imm.u, 64 ) );
         } };
 
     vm::lifters::lifter_t lconstwsxq = {
         // push imm<N>
-        vm::handler::LCONSTWSXQ, []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr ) {
+        vm::handler::LCONSTWSXQ,
+        []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr, vmp2::v3::code_block_t *code_blk ) {
             blk->push( vtil::operand( vinstr->operand.imm.u, 64 ) );
         } };
 
     vm::lifters::lifter_t lconstwsxdw = {
         // push imm<N>
-        vm::handler::LCONSTWSXDW, []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr ) {
+        vm::handler::LCONSTWSXDW,
+        []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr, vmp2::v3::code_block_t *code_blk ) {
             blk->push( vtil::operand( vinstr->operand.imm.u, 32 ) );
         } };
 } // namespace vm::lifters
\ No newline at end of file
diff --git a/src/vmlifters/lflags.cpp b/src/vmlifters/lflags.cpp
index 30b6c51..8ba753a 100644
--- a/src/vmlifters/lflags.cpp
+++ b/src/vmlifters/lflags.cpp
@@ -4,6 +4,6 @@ namespace vm::lifters
 {
     vm::lifters::lifter_t lrflags = {
         // push flags
-        vm::handler::LRFLAGS,
-        []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr ) { blk->push( vtil::REG_FLAGS ); } };
+        vm::handler::LRFLAGS, []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr,
+                                  vmp2::v3::code_block_t *code_blk ) { blk->push( vtil::REG_FLAGS ); } };
 }
\ No newline at end of file
diff --git a/src/vmlifters/lreg.cpp b/src/vmlifters/lreg.cpp
index 324b876..2f0fb23 100644
--- a/src/vmlifters/lreg.cpp
+++ b/src/vmlifters/lreg.cpp
@@ -4,13 +4,15 @@ namespace vm::lifters
 {
     vm::lifters::lifter_t lregq = {
         // push vregX
-        vm::handler::LREGQ, []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr ) {
+        vm::handler::LREGQ,
+        []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr, vmp2::v3::code_block_t *code_blk ) {
             blk->push( make_virtual_register( vinstr->operand.imm.u, 8 ) );
         } };
 
     vm::lifters::lifter_t lregdw = {
         // push vregX
-        vm::handler::LREGDW, []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr ) {
+        vm::handler::LREGDW,
+        []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr, vmp2::v3::code_block_t *code_blk ) {
             blk->push( make_virtual_register( vinstr->operand.imm.u, 4 ) );
         } };
 } // namespace vm::lifters
\ No newline at end of file
diff --git a/src/vmlifters/lvsp.cpp b/src/vmlifters/lvsp.cpp
index 3bc49b7..7943aa9 100644
--- a/src/vmlifters/lvsp.cpp
+++ b/src/vmlifters/lvsp.cpp
@@ -4,6 +4,6 @@ namespace vm::lifters
 {
     vm::lifters::lifter_t lvsp = {
         // vsp = vsp[0]
-        vm::handler::LVSP,
-        []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr ) { blk->pop( vtil::REG_SP ); } };
+        vm::handler::LVSP, []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr,
+                               vmp2::v3::code_block_t *code_blk ) { blk->pop( vtil::REG_SP ); } };
 }
\ No newline at end of file
diff --git a/src/vmlifters/nand.cpp b/src/vmlifters/nand.cpp
index 141c8f4..bdf4bb8 100644
--- a/src/vmlifters/nand.cpp
+++ b/src/vmlifters/nand.cpp
@@ -10,7 +10,8 @@ namespace vm::lifters
         // and vregX, vregY
         // push vregX
         // pushf
-        vm::handler::NANDQ, []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr ) {
+        vm::handler::NANDQ,
+        []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr, vmp2::v3::code_block_t *code_blk ) {
             auto [ t1, t2 ] = blk->tmp( 64, 64 );
             blk->pop( t1 );
             blk->pop( t2 );
@@ -31,7 +32,8 @@ namespace vm::lifters
         // and vregX, vregY
         // push vregX
         // pushf
-        vm::handler::NANDDW, []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr ) {
+        vm::handler::NANDDW,
+        []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr, vmp2::v3::code_block_t *code_blk ) {
             auto [ t1, t2 ] = blk->tmp( 32, 32 );
             blk->pop( t1 );
             blk->pop( t2 );
@@ -52,7 +54,8 @@ namespace vm::lifters
         // and vregX, vregY
         // push vregX
         // pushf
-        vm::handler::NANDW, []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr ) {
+        vm::handler::NANDW,
+        []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr, vmp2::v3::code_block_t *code_blk ) {
             auto [ t1, t2 ] = blk->tmp( 16, 16 );
             blk->pop( t1 );
             blk->pop( t2 );
diff --git a/src/vmlifters/pushvsp.cpp b/src/vmlifters/pushvsp.cpp
index 566a430..1de8465 100644
--- a/src/vmlifters/pushvsp.cpp
+++ b/src/vmlifters/pushvsp.cpp
@@ -4,6 +4,6 @@ namespace vm::lifters
 {
     vm::lifters::lifter_t pushvsp = {
         // push vsp
-        vm::handler::PUSHVSP,
-        []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr ) { blk->push( vtil::REG_SP ); } };
+        vm::handler::PUSHVSP, []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr,
+                                  vmp2::v3::code_block_t *code_blk ) { blk->push( vtil::REG_SP ); } };
 }
\ No newline at end of file
diff --git a/src/vmlifters/read.cpp b/src/vmlifters/read.cpp
index 171d1d8..6299524 100644
--- a/src/vmlifters/read.cpp
+++ b/src/vmlifters/read.cpp
@@ -6,7 +6,8 @@ namespace vm::lifters
         // pop vregX
         // ldd vregX, vregX, 0
         // push vregX
-        vm::handler::READQ, []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr ) {
+        vm::handler::READQ,
+        []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr, vmp2::v3::code_block_t *code_blk ) {
             auto [ t0, t1 ] = blk->tmp( 64, 64 );
             blk->pop( t0 );
             blk->ldd( t1, t0, vtil::make_imm( 0ull ) );
@@ -17,7 +18,8 @@ namespace vm::lifters
         // pop vregX
         // ldd vregX, vregX, 0
         // push vregX
-        vm::handler::READDW, []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr ) {
+        vm::handler::READDW,
+        []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr, vmp2::v3::code_block_t *code_blk ) {
             auto [ t0, t1 ] = blk->tmp( 64, 32 );
             blk->pop( t0 );
             blk->ldd( t1, t0, vtil::make_imm( 0ull ) );
@@ -28,7 +30,8 @@ namespace vm::lifters
         // pop vregX
         // ldd vregX, vregX, 0
         // push vregX
-        vm::handler::READW, []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr ) {
+        vm::handler::READW,
+        []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr, vmp2::v3::code_block_t *code_blk ) {
             auto [ t0, t1 ] = blk->tmp( 64, 16 );
             blk->pop( t0 );
             blk->ldd( t1, t0, vtil::make_imm( 0ull ) );
diff --git a/src/vmlifters/shr.cpp b/src/vmlifters/shr.cpp
index 8e88e17..f8610ac 100644
--- a/src/vmlifters/shr.cpp
+++ b/src/vmlifters/shr.cpp
@@ -8,7 +8,8 @@ namespace vm::lifters
         // shr vregX, vregY
         // push vregX
         // pushf
-        vm::handler::SHRQ, []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr ) {
+        vm::handler::SHRQ,
+        []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr, vmp2::v3::code_block_t *code_blk ) {
             auto [ t0, t1 ] = blk->tmp( 64, 8 );
             blk->pop( t0 );
             blk->pop( t1 );
@@ -23,7 +24,8 @@ namespace vm::lifters
         // shr vregX, vregY
         // push vregX
         // pushf
-        vm::handler::SHRW, []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr ) {
+        vm::handler::SHRW,
+        []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr, vmp2::v3::code_block_t *code_blk ) {
             auto [ t0, t1 ] = blk->tmp( 16, 8 );
             blk->pop( t0 );
             blk->pop( t1 );
diff --git a/src/vmlifters/sreg.cpp b/src/vmlifters/sreg.cpp
index c6d6698..253322f 100644
--- a/src/vmlifters/sreg.cpp
+++ b/src/vmlifters/sreg.cpp
@@ -4,19 +4,22 @@ namespace vm::lifters
 {
     vm::lifters::lifter_t sregq = {
         // pop vregX
-        vm::handler::SREGQ, []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr ) {
+        vm::handler::SREGQ,
+        []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr, vmp2::v3::code_block_t *code_blk ) {
             blk->pop( make_virtual_register( vinstr->operand.imm.u, 8 ) );
         } };
 
     vm::lifters::lifter_t sregdw = {
         // pop vregX
-        vm::handler::SREGDW, []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr ) {
+        vm::handler::SREGDW,
+        []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr, vmp2::v3::code_block_t *code_blk ) {
             blk->pop( make_virtual_register( vinstr->operand.imm.u, 4 ) );
         } };
 
     vm::lifters::lifter_t sregw = {
         // pop vregX
-        vm::handler::SREGW, []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr ) {
+        vm::handler::SREGW,
+        []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr, vmp2::v3::code_block_t *code_blk ) {
             blk->pop( make_virtual_register( vinstr->operand.imm.u, 2 ) );
         } };
 } // namespace vm::lifters
\ No newline at end of file
diff --git a/src/vmlifters/vmexit.cpp b/src/vmlifters/vmexit.cpp
index 55f7a9b..1fa8a9b 100644
--- a/src/vmlifters/vmexit.cpp
+++ b/src/vmlifters/vmexit.cpp
@@ -4,7 +4,8 @@ namespace vm::lifters
 {
     vm::lifters::lifter_t vmexit = {
         // ret
-        vm::handler::VMEXIT, []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr ) {
+        vm::handler::VMEXIT,
+        []( vtil::basic_block *blk, vm::instrs::virt_instr_t *vinstr, vmp2::v3::code_block_t *code_blk ) {
             blk->vexit( vtil::make_imm( vinstr->trace_data.vsp.qword[ 0x13 ] - vinstr->trace_data.regs.r13 ) );
         } };
 }
\ No newline at end of file