added call offset to method prologue/epilogue

QDucasse · Feb 13, 2024 · ae83c7a · ae83c7a
1 parent e3e4bc7
commit ae83c7a
Show file tree

Hide file tree

Showing 5 changed files with 20 additions and 12 deletions.
diff --git a/gigue/fixer/fixer_generator.py b/gigue/fixer/fixer_generator.py
@@ -71,7 +71,6 @@ def __init__(
 
         # Prologue/Epilogue offsets
         self.call_size += 3
-        print(self.call_size)
         self.method_epilogue_offset += 3
 
     def build_interpreter_prologue(

diff --git a/gigue/generator.py b/gigue/generator.py
@@ -127,6 +127,8 @@ def __init__(
         # Prologue/Epilogue offsets
         self.method_prologue_offset = 1  # Stack sizing
         self.method_epilogue_offset = 2  # Stack sizing + ret
+        # TODO: Rename and dissociate prologue epilogue
+        self.method_call_offset = 1  # ra load/store
 
         # Generation
         self.weights: List[int] = weights
@@ -229,6 +231,7 @@ def generate_method(self, address: int, *args, **kwargs) -> Method:
                 builder=self.builder,
                 prologue_offset=self.method_prologue_offset,
                 epilogue_offset=self.method_epilogue_offset,
+                call_offset=self.method_call_offset,
             )
             logger.debug(
                 f"{self.log_jit_prefix()} {method.log_prefix()} Method added with size"
@@ -265,6 +268,7 @@ def generate_leaf_method(self, address: int) -> Method:
                 builder=self.builder,
                 prologue_offset=self.method_prologue_offset,
                 epilogue_offset=self.method_epilogue_offset,
+                call_offset=self.method_call_offset,
             )
             logger.debug(
                 f"{self.log_jit_prefix()} {method.log_prefix()} Leaf method added with"

diff --git a/gigue/method.py b/gigue/method.py
@@ -27,6 +27,7 @@ def __init__(
         local_vars_nb: int = 2,
         prologue_offset: int = 1,
         epilogue_offset: int = 2,
+        call_offset: int = 1,
     ):
         self.address: int = address
         self.body_size: int = body_size
@@ -49,11 +50,11 @@ def __init__(
         self.is_leaf: bool = self.call_number == 0
         self.prologue_size: int = (
             # stack space + register saving + ra saving
-            prologue_offset + self.used_s_regs + (1 if not self.is_leaf else 0)
+            prologue_offset + self.used_s_regs + (call_offset if not self.is_leaf else 0)
         )
         self.epilogue_size: int = (
             # register restoring + ra restoring + stack space + ret
-            self.used_s_regs + (1 if not self.is_leaf else 0) + epilogue_offset
+            self.used_s_regs + (call_offset if not self.is_leaf else 0) + epilogue_offset
         )
 
         self.builder: InstructionBuilder = builder

diff --git a/gigue/rimi/rimi_generator.py b/gigue/rimi/rimi_generator.py
@@ -79,6 +79,9 @@ def __init__(
 
         self.shadow_stack_size = shadow_stack_size
 
+        # Prologue/Epilogue offsets
+        self.method_call_offset += 1
+
     def build_interpreter_prologue(
         self, used_s_regs: int, local_var_nb: int, contains_call: bool
     ):

diff --git a/tests/rimi/test_rimi_generator.py b/tests/rimi/test_rimi_generator.py
@@ -9,12 +9,13 @@
 from tests.conftest import (
     INTERPRETER_START_ADDRESS,
     JIT_START_ADDRESS,
+    RET_ADDRESS,
     TEST_DATA_REG,
     TEST_DATA_SIZE,
     cap_disasm_bytes,
     check_size,
 )
-from tests.rimi.conftest import TEST_RIMI_SSP_REG
+from tests.rimi.conftest import TEST_RIMI_SSP_REG, start_resumable_emulation
 
 logger = logging.getLogger("gigue")
 
@@ -42,6 +43,7 @@ def test_execute_shadow_stack_trampoline_generated_binaries(
     cap_disasm_custom_setup,
     rimi_handler_setup,
     rimi_uc_emul_full_setup,
+    log_trace,
 ):
     generator = RIMIShadowStackTrampolineGenerator(
         jit_start_address=JIT_START_ADDRESS,
@@ -90,12 +92,12 @@ def test_execute_shadow_stack_trampoline_generated_binaries(
     # Handler
     rimi_handler = rimi_handler_setup
     rimi_handler.hook_instr_tracer(uc_emul)
+    rimi_handler.hook_reg_tracer(uc_emul)
     rimi_handler.hook_handler(uc_emul)
 
-    # TODO: Something fishy!
-    # start_address = INTERPRETER_START_ADDRESS
-    # end_address = RET_ADDRESS
-    # start_resumable_emulation(uc_emul, start_address, end_address)
+    start_address = INTERPRETER_START_ADDRESS
+    end_address = RET_ADDRESS
+    start_resumable_emulation(uc_emul, start_address, end_address)
 
 
 @pytest.mark.parametrize(
@@ -173,7 +175,6 @@ def test_execute_full_trampoline_generated_binaries(
     rimi_handler.hook_exception_tracer(uc_emul)
     rimi_handler.hook_handler(uc_emul)
 
-    # TODO: Something fishy!
-    # start_address = INTERPRETER_START_ADDRESS
-    # end_address = RET_ADDRESS
-    # start_resumable_emulation(uc_emul, start_address, end_address)
+    start_address = INTERPRETER_START_ADDRESS
+    end_address = RET_ADDRESS
+    start_resumable_emulation(uc_emul, start_address, end_address)